torch.export()-only export Llama arg

[jackzhxng]

Summary

Option to only save the torch.export()ed model and skip the to_edge and to_executorch passes.

Test plan

python -m examples.models.llama.export_llama --checkpoint /tmp/Llama-3.2-1B-Instruct/original/consolidated.00.pth  --params /tmp/Llama-3.2-1B-Instruct/original/params.json  --metadata '{"append_eos_to_prompt": 0, "get_bos_id":128000, "get_eos_ids":[128009, 128001], "get_n_bos": 0, "get_n_eos": 0}' --output_name="llama3_2_1B.pt2" --export_only

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit a6a7162 with merge base 8f9fb7e ():
💚 Looks good so far! There are no failures yet. 💚

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

[cccclai]

Should we check if we want the module or the exported_module? Asking because I think QAT is working on the module not exported_module, we can double check with @navsud

[navsud]

Should we check if we want the module or the exported_module? Asking because I think QAT is working on the module not exported_module, we can double check with @navsud

QAT works on torch.export.export(...).module().

[jackzhxng]

Not sure if you are referring to whether the torch.export()ed model vs. the eager model is needed or the exported_program vs. exported_program.module() needs to be exported - if the latter then you can only save an ExportedProgram, after that when you load you can call .module() on it

[kimishpatel]

are we not importing this stuff to phabricator and make sure it is not breaking anythng?

[jackzhxng]

@kimishpatel the way difftrain works on these is that a diff is created in phabricator post-merge, if there are internal tests that break then they are forward-fixed in a separate diff and exported, or they are reverted

[kimishpatel]

But that means diff train cannot land, right? if thats the case, thats ok. I just think that we should not land breaking changes internally

[jackzhxng]

Yup, so the if internal tests break then the diff train cannot land until it is reverted / a forward fix diff is stacked on top of it