Llama2 model cleanup (#5859)
Summary:
- Removes redundant steps in the Llama2 export
- Factors out checkpointing to be shared with future Llama models (namely 3.2 multimodal)
- Comments and orders code more clearly
PR chain:
- [Add kwarg example inputs to eager model base](https://github.com/pytorch/executorch/pull/5765)
- **YOU ARE HERE ~>** [Llama2 model cleanup](https://github.com/pytorch/executorch/pull/5859)
- [Accept model type parameter in export_llama](https://github.com/pytorch/executorch/pull/5910)
- [Export TorchTune llama3_2_vision in ET](https://github.com/pytorch/executorch/pull/5911)
- [Add et version of TorchTune MHA for swapping with custom op](https://github.com/pytorch/executorch/pull/5912)
Pull Request resolved: https://github.com/pytorch/executorch/pull/5859
Test Plan:
Ensure export + eval is similar before and after for Stories 110M:
```
python -m examples.models.llama2.eval_llama -c <checkpoint.pth> -p <params.json> -t <tokenizer.model/bin> -d fp32 --max_seq_len 2048 --limit 1000
```
Before:
```
wikitext: {'word_perplexity,none': 14464.645927166595, 'word_perplexity_stderr,none': 'N/A', 'byte_perplexity,none': 5.99788806086652, 'byte_perplexity_stderr,none': 'N/A', 'bits_per_byte,none': 2.5844545973083983, 'bits_per_byte_stderr,none': 'N/A', 'alias': 'wikitext'}
```
After:
```
wikitext: {'word_perplexity,none': 14464.299192404438, 'word_perplexity_stderr,none': 'N/A', 'byte_perplexity,none': 5.997861173678705, 'byte_perplexity_stderr,none': 'N/A', 'bits_per_byte,none': 2.584448130015399, 'bits_per_byte_stderr,none': 'N/A', 'alias': 'wikitext'}
```
Reviewed By: malfet, dbort
Differential Revision: D64145852
Pulled By: dvorjackz
fbshipit-source-id: daeee834955e154e7c8262ce776bd3039991027d J
Jack Zhang committed
4745070d02822b68cc480b9d00fe6cf10dc51aaa
Parent: 4e8f609
Committed by Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
on 10/15/2024, 4:18:22 AM