feat(models): reject tie_word_embeddings=True on separate-head model families - #2805
Conversation
…families Add reject_unsupported_tied_word_embeddings() (built on the NVIDIA-NeMo#2732 resolver) and wire it into the __init__ of the 25 verified untied-default model classes, so setting tie_word_embeddings=True raises a clear error instead of a silently-untied head. Excludes mistral3_vlm (HF default tied) and step3p5/step3p7/nemotron_omni (pending hub verification). Refs NVIDIA-NeMo#2512 Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
There was a problem hiding this comment.
Pull request overview
This PR strengthens NeMo AutoModel’s handling of Hugging Face’s tie_word_embeddings semantics by adding an explicit constructor-time reject guard for “separate-head” model families (HF default: untied), preventing silent construction with a randomly initialized lm_head when users set tie_word_embeddings=True.
Changes:
- Added
reject_unsupported_tied_word_embeddings(config, model_class_name)to checkpoint utilities to raise a clearNotImplementedErrorwhen the controllingtie_word_embeddingsflag isTruefor untied-default architectures. - Wired the reject guard into the constructors of many separate-head model families (mechanical one-import + one-call changes).
- Added unit coverage for the new helper and a CPU-safe integration test ensuring
qwen3_moerejectstie_word_embeddings=Trueat construction time.
Reviewed changes
Copilot reviewed 28 out of 28 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| nemo_automodel/components/checkpoint/utils.py | Adds the shared reject helper built on the existing controlling-flag resolver. |
| nemo_automodel/components/models/deepseek_v3/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/deepseek_v32/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/deepseek_v4/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/glm4_moe/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/glm4_moe_lite/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/glm_moe_dsa/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/gpt_oss/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/hy_mt2/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/hy_v3/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/kimi_k25_vl/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/kimivl/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/ling_v2/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/llava_onevision/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/minimax_m2/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/minimax_m3_vl/model.py | Adds constructor guards for both classes in the module. |
| nemo_automodel/components/models/mimo_v2_flash/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/mistral4/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/nemotron_parse/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/nemotron_v3/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/qwen2_5_omni/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/qwen3_5_moe/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/qwen3_moe/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/qwen3_next/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/qwen3_omni_moe/model.py | Adds constructor guard against tie_word_embeddings=True. |
| nemo_automodel/components/models/qwen3_vl_moe/model.py | Adds constructor guard against tie_word_embeddings=True. |
| tests/unit_tests/models/qwen3_moe/test_qwen3_moe_tie_guard.py | New CPU-safe integration test asserting construction rejects tie=True. |
| tests/unit_tests/utils/test_checkpoint_utils.py | Adds unit tests for the new reject helper, including composite top-level resolution. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| thinker_config = _resolve_thinker_config(config) | ||
| super().__init__(thinker_config) | ||
| reject_unsupported_tied_word_embeddings(self.config, type(self).__name__) |
There was a problem hiding this comment.
Fixed — moved the guard to the top of init on the original config, before _resolve_thinker_config() and super(), so it sees the controlling top-level flag and fails fast.
| # Extract text_config if this is a multimodal wrapper config | ||
| config = getattr(config, "text_config", config) | ||
| self.config = config | ||
| reject_unsupported_tied_word_embeddings(config, type(self).__name__) |
There was a problem hiding this comment.
Fixed — the guard now runs on the original config before the text_config unwrap, so a top-level tie_word_embeddings=True is caught.
| super().__init__(config) | ||
| reject_unsupported_tied_word_embeddings(config, type(self).__name__) |
There was a problem hiding this comment.
Fixed — moved before super().init(config) to fail fast before building the HF module.
| # Initialize HF parent (creates self.model, self.lm_head, vision encoder, etc.) | ||
| super().__init__(config) | ||
| reject_unsupported_tied_word_embeddings(config, type(self).__name__) |
There was a problem hiding this comment.
Fixed — moved before super().init(config), same reason.
…eview) Per review: in composite models the guard was reading the unwrapped text_config/thinker_config, which could miss a top-level tie_word_embeddings=True. Run it on the original config before any unwrap or super().__init__(), which also fails fast. Refs NVIDIA-NeMo#2512 Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
yuhezhang-ai
left a comment
yuhezhang-ai
left a comment
There was a problem hiding this comment.
Left one inline comment on the Omni wrapper config guard path.
| ): | ||
| # Check the controlling top-level flag on the original config before | ||
| # resolving to thinker_config and building the HF parent. | ||
| reject_unsupported_tied_word_embeddings(config, type(self).__name__) |
There was a problem hiding this comment.
Could you also cover the full Omni wrapper config case? Moving the guard before _resolve_thinker_config() makes it early, but Qwen2_5OmniConfig itself does not expose tie_word_embeddings; the controlling flag lives on config.thinker_config.tie_word_embeddings. With the current resolver this still returns False and the guard does not raise:
cfg = Qwen2_5OmniConfig()
cfg.thinker_config.tie_word_embeddings = True
get_controlling_tie_word_embeddings(cfg, "Qwen2_5OmniThinkerForConditionalGeneration") # FalseSame shape for full Qwen3OmniMoeConfig. Since this PR is specifically adding constructor reject guards, I think the resolver should handle thinker_config.tie_word_embeddings here rather than leaving this hole to a follow-up. Please add unit coverage for the full-wrapper config path too.
There was a problem hiding this comment.
Ohokay got it thats a good catch — fixed. The resolver now unwraps config.thinker_config for the Omni thinker models, since the full Qwen2_5OmniConfig/Qwen3OmniMoeConfig doesn't expose tie_word_embeddings at the top level. Verified with a real Qwen2_5OmniConfig: setting thinker_config.tie_word_embeddings=True now resolves to True and the guard raises. Added unit coverage for the full-wrapper path (both Omni classes) and the direct-thinker-config case.
| @@ -0,0 +1,56 @@ | |||
| # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. | |||
There was a problem hiding this comment.
Nit: this is a new Python file, so the repo guidance says the NVIDIA copyright header should use the current year, 2026.
There was a problem hiding this comment.
Done — updated the copyright header to 2026.
…eview) Qwen2_5OmniConfig/Qwen3OmniMoeConfig don't expose tie_word_embeddings at the top level; the controlling flag is on config.thinker_config. Unwrap to it in the resolver so the constructor guard catches a tied request via the full wrapper config. Add wrapper-path unit coverage, clarify the resolver docstring, and fix the copyright year to 2026. Refs NVIDIA-NeMo#2512 Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
|
@yuhezhang-ai , is everything as per you request ? Is it good to be merged now? |
|
/ok to test ad7ed14 |
@Achyuthan-S Look good to me. Thanks again for the contribution. Just triggered the CI test. |
Hello @yuhezhang-ai , it was great contributing. Thank you again for the opportunity. If you have any other issues that i can contribute to and solve .., it would be great if you tag me to it whenever . Thanks again ! |
svcnvidia-nemo-ci
removed
the
waiting-on-maintainers
Second of two PRs for #2512 — the constructor reject-guards. The shared resolver landed in #2732; this wires it into the model classes.
What
Adds
reject_unsupported_tied_word_embeddings(config, model_class_name)incheckpoint/utils.py, and calls it in the__init__of the separate-head model classes (one import + one line each). It raises a clearNotImplementedErrorwhen the controllingtie_word_embeddingsflag isTruefor a model whose HF default is untied, instead of silently leaving a randomly-initializedlm_head.It reuses
get_controlling_tie_word_embeddings()from #2732, so composite VLM/omni configs are read from the top-level flag rather than a nestedtext_config.Applied to the 25 verified untied-default classes:
deepseek_v3, deepseek_v32, deepseek_v4, glm4_moe, glm4_moe_lite, glm_moe_dsa, gpt_oss, hy_mt2, hy_v3, ling_v2, mimo_v2_flash, minimax_m2, minimax_m3_vl (both classes), nemotron_v3, mistral4, kimivl, kimi_k25_vl, llava_onevision, nemotron_parse, qwen2_5_omni, qwen3_moe, qwen3_5_moe, qwen3_next, qwen3_vl_moe, qwen3_omni_moe.
Scope
Deliberately left out, to keep this focused on the reject-
tie=True(untied-default) direction:mistral3_vlm— HF default is tied, so it's the opposite case (would rejecttie=False). Needs its own handling.step3p5,step3p7,nemotron_omni— their configs aren't in transformers 5.8.1 yet, so I want to confirm HF defaults against the published checkpoints before guarding them.Tests
test_checkpoint_utils.py: raises when tied, no-op when untied, and composite models resolve from the top-level flag.test_qwen3_moe_tie_guard.py: CPU integration test — constructingqwen3_moewithtie=Trueraises. The guard runs at the very start of__init__, before any GPU build, so this is CPU-safe even though the full model suites are CUDA-only.Existing model suites stay green; the per-model change is the same two lines everywhere, so the diff is large but mechanical.
Refs #2512.
cc @yuhezhang-ai
Does not fully close #2512 until tied-default flip guards + hub-verified families land.