[Executorch][llm] Enable local global attention in export_llama script

[kimishpatel]

Stack from ghstack (oldest at bottom):

• -> [Executorch][llm] Enable local global attention in export_llama script #10612
• [Executorch][llm] Enable leveraging ring kv cache via module swap #10611
• [Executorch][llm] Make custom update cache op operate on indices #10610
• [Executorch][llm] Add ring buffer based kv cache and mask calculation to MHA #10609
• [Executorch][llm] Add support for ring kv cache and ring attention #10608

Added a new option of --local_global_attention that takes in pattern of sizes to determine which layers are using local sliding window attention.
For example, [0, 256, 256, 0, 256, 256] can be used for 6 layers transformer. Or you can also use [0, 256, 256] as pattern you want
to repeat.

Differential Revision: D73891423

cc @larryliu0820 @mergennachin @cccclai @helunwencser @jackzhxng

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/10612

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 3 New Failures

As of commit 22472c1 with merge base 1ae8c2c ():

NEW FAILURES - The following jobs have failed:

• .github/workflows/build-presets.yml (gh)
• pull / unittest / linux / linux-job (gh)
  examples/models/llama/tests/test_export_llama_lib.py::ExportLlamaLibTest::test_has_expected_ops_and_op_counts
• pull / unittest-editable / linux / linux-job (gh)
  examples/models/llama/tests/test_export_llama_lib.py::ExportLlamaLibTest::test_has_expected_ops_and_op_counts

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D73891423

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D73891423

[jackzhxng]

From my understanding local global attention is tied closely to the model, and isn't something you can adjust given a particular checkpoint unlike these other export options such as qmode, enable_kv_cache, etc. I think it would be better to have this in model_args instead so that we can represent it in the config json file, so it is part of model configuration, not export configuration

[kimishpatel]

From my understanding local global attention is tied closely to the model, and isn't something you can adjust given a particular checkpoint unlike these other export options such as qmode, enable_kv_cache, etc. I think it would be better to have this in model_args instead so that we can represent it in the config json file, so it is part of model configuration, not export configuration

You can also make all the layers just do sliding window attention for example, so you can view this also as pure optimization.

I havent thought a lot of about making this a model arg vs. export arg. I think your point definitely makes sense but I think it can be configurable as well if you want to support infinite generation as many runners do.

[jackzhxng]

Hmm can a model trained without local(/global) attention run fine with local attention to take advantage of performance optimizations? If so and we view this as an export optimization then it might make sense and we can apply it to any arbitrary transformer. Main concern for me is that when we have this as an export option, then models that require local/global attention become coupled to this export CLI option, instead of just being able to represent it in the config json file, like is done in Huggingface

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D73891423

[kimishpatel]

Hmm can a model trained without local(/global) attention run fine with local attention to take advantage of performance optimizations? If so and we view this as an export optimization then it might make sense and we can apply it to any arbitrary transformer. Main concern for me is that when we have this as an export option, then models that require local/global attention become coupled to this export CLI option, instead of just being able to represent it in the config json file, like is done in Huggingface

I think you can even when the model is not trained. Using sliding window (e.g. attention sink https://arxiv.org/abs/2309.17453) is used to support much larger context. e.g. you may have a model trained with 128k context but you cant fully support that with such large kv cache

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D73891423

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D73891423