Gemma3 is Torch Exportable

[guangy10]

What does this PR do?

Initial effort to add torch.expport support for the Gemma3 model!
Gemma3 provides a 1b variant that is suitable for ExecuTorch to bring it for on-device use-case. This PR is focusing on creating the export recipe and validate the exported model can produce same output as eager.

Expand support to other models that utilize HybridCache as well including gemma2 and cohere2.

End2end Test Validation with Exported Graph

RUN_SLOW=1 pytest tests/models/gemma3/test_modeling_gemma3.py -s -v -k test_export_text_only_with_hybrid_cache

Export generated texts: 'What is the capital of France?

The capital of France is Paris.

Final Answer: The final answer is $\boxed{Paris'


Eager generated texts: 'What is the capital of France?

The capital of France is Paris.

Final Answer: The final answer is $\boxed{Paris'

PASSED
======================================================= 1 passed, 319 deselected, 108 warnings in 26.95s =======================================================

Ene2end Validation in Optimum-ExecuTorch

Before submitting

• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case. Gemma3 is ExecuTorch compatible #37727
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@ArthurZucker @gante @qubvel

[github-actions]

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the Ready for review button (at the bottom of the PR page). This will assign reviewers and trigger CI.

[guangy10]

CC: @tugsbayasgalan for review as well

[guangy10]

removed unused buffer

[guangy10]

Fixed in the ExportableModule to make gemma3 lowerable to ExecuTorch

[guangy10]

I've ran make fixup but it doesn't fixed the linter.

[Cyrilvallez]

Hey @guangy10! Thanks for the PR! The changes LGTM, however the gemma3 change should be reflected in modular_gemma3.py, which is the source of modeling_gemma3.py! 🤗
Any way you could downstream the small change to other models using Hybridcache as well? They should be the models that were changed in #37447 👌

[Cyrilvallez]

It's the source of the issue in check_repo_consistency 😉

[Cyrilvallez]

cc @gante as well for viz!

[guangy10]

Hey @guangy10! Thanks for the PR! The changes LGTM, however the gemma3 change should be reflected in modular_gemma3.py, which is the source of modeling_gemma3.py! 🤗 Any way you could downstream the small change to other models using Hybridcache as well? They should be the models that were changed in #37447 👌

Done. Could you take another look?

[Cyrilvallez]

LGTM, thanks a lot! Great work, super clean 🤗
Merging