Skip to content

Swap mha #6719

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Nov 11, 2024
Merged

Swap mha #6719

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b32f26c
Swap mha
jackzhxng Nov 5, 2024
c6ac3c6
Fix tests
jackzhxng Nov 8, 2024
b373248
Delete old file
jackzhxng Nov 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 13 additions & 10 deletions extension/llm/modules/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,14 +1,17 @@
## Export Friendly Modules
## Export-friendly Modules

Modules in this directory are:
* Extending `torch.nn.Module`.
* Guranteed to work out of the box with `torch.export.export()` and `torch.aot_compile()`.
* Guranteed to be able to work with ExecuTorch.
Modules in this directory:
* Extend `torch.nn.Module`.
* Are guaranteed to work out of the box with `torch.export.export()`.
* Should work out of the box with `torch.aot_compile()`.
* Should be able to workt with ExecuTorch.

All modules should be covered by unit tests to make sure they are:
1. giving the same output as the reference implementation in PyTorch or torchtune
2. export friendly
3. AOTI friendly
4. ExecuTorch friendly
1. Give the output as the reference eager model in PyTorch or TorrchTune
2. Export-friendly

Notice that these modules are subject to change (may upstream to torchtune) so proceed with caution.
Additionally, we aim to make these modules:
3. AOTI-friendly
4. ExecuTorch-friendly

These modules are subject to change (may upstream to TorchTune) so proceed with caution.
Loading
Loading