SIGN IN SIGN UP

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