Introduce export_llama to extension/llm as export_llm

[jackzhxng]

🚀 The feature, motivation and pitch

To generalize export_llama better to transformers as export_llm, the following directories and files in examples/models/llama will be moved to extension/llm:

• examples/models/llama/config
• examples/models/llama/source_transformation
• examples/models/llama/tests
• attention.py -> Will go into a new attention directory
• export_llama_hydra.py -> this will be renamed to export_llm.py and will be the main entrypoint to the new API
• export_llama_lib.py -> this will be renamed to export_llm_lib.py
• fairseq2.py
• hf_download.py
• llama_transformer.py -> this will be renamed to transformer.py
• model.py -> this will have Llama2Model renamed to TransformerModel. Also not a big fan of the naming of this file, open to take this opportunity to rename to eager.py or something like that to its intent better, which is to instantiate the model.
  • We could remove Llama-specific components from the version in extension/llm such as Fairseq2 and have Llama2Model inherit from it
• model_args.py
• norm.py
• rope.py
• static_attention.py -> Will go into a new attention directory

Key notes:

• export_llama.py stays put, export_llama will still exist to preserve BC, especially internally
  • export_llama_args.py then also stays put, supporting the legacy CLI
  • export_llama_hydra.py will be the new export_llm powered by Hydra, and it will be in extension/llm. Note that this means that the export_llm API will not support the legacy CLI
• Lots of internal callsites for functions in export_llama_lib such as _prepare_for_llama_export and _export_llama. These will need to all be renamed and refactored.

Alternatives

No response

Additional context

No response

RFC (Optional)

No response

cc @larryliu0820 @mergennachin @cccclai @helunwencser

[larryliu0820]

cc @kimishpatel

[larryliu0820]

The name model.py is becoming kinda like standard I feel. I don't have a strong opinion but it seems to be descriptive enough as what's inside that file.

[mergennachin]

It's gonna be hard to do this big migration. The devil's in the details -- it's hard to review this RFC and give feedback, without seeing the actual code example.

I'd recommend this intermediate milestone:

Have only public in API, called, /extension/llm/export_llm.py

In this intermediate stage, the implementation detail of export_llm.py still calls components in examples directly.

Over-time, we want to be in a state where examples calls extension/llm. And it's one-way street aka there's no call from extensions/llm to example. But in this intermediate stage, we can careful and surgical about gradual migration.

I really want to be careful on what we put in extensions/llm. I don't want to be yet another dumping ground of all the abstraction layers. We already have LLMEdgeManager, which I really don't like that.

Cc @jackzhxng

[jackzhxng]

@mergennachin down to start with that intermediate milestone, then we can look into the rest of the files. I doubt there's going to too many code changes for the rest of the files though, but I do agree that we shouldn't dump everything into extension/llm. Was thinking we can organize it better by having all of the model code (e.g. llama_transformer.py, rope.py, attention.py) go into extension/llm/modeling. With regards to abstractions layers, the only one here really is model.py, which is why I was suggesting to rename it to something else to reflect its purpose better. Long term though it would be good to get rid of the https://github.com/pytorch/executorch/blob/main/examples/models/model_factory.py abstraction behind it, I know we use it for CI but not sure why we need it for any other reason.