cp: fix(gemma4_moe): re-tie lm_head to active embed_tokens on MoE path (2601) into r0.5.0 - #2709
Merged
Merged
Conversation
…2601) * fix(gemma4_moe): re-tie lm_head to active embed_tokens on MoE path The MoE path replaces language_model after HF __init__, orphaning the lm_head<->embed_tokens tie that HF set up. Re-tie lm_head to the active embed_tokens when tie_word_embeddings is set (Gemma defaults True). Add CPU tied/untied tests. Refs #2512 Signed-off-by: Achyuthan-S <as21154@nyu.edu> * refactor(gemma4_moe): expose MoE re-tie via tie_weights() Wrap the post-swap lm_head re-point in a tie_weights() override so AutoModel and ensure_tied_lm_head() re-tie to the active MoE embedding. Add hook tests for re-tie and untied no-op. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com> * fix(gemma4_moe): make lm_head re-tie follow top-level config flag Per review: HF Gemma4 ties on the top-level Gemma4Config.tie_word_embeddings regardless of the nested text_config (verified by construction on transformers 5.8.1). Read top-level first with a text_config fallback, set the top-level flag in tests, and add a case proving top-level wins when the two disagree. Refs #2512 Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com> --------- Signed-off-by: Achyuthan-S <as21154@nyu.edu> Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com> Signed-off-by: NeMo Bot <nemo-bot@nvidia.com>
Contributor
Author
|
/ok to test 8e289ee |
akoumpa
approved these changes
Jun 22, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
beep boop [🤖]: Hi @Achyuthan-S 👋,