refactor(models): single tie_word_embeddings guard via TieSupport by Achyuthan-S · Pull Request #2998 · NVIDIA-NeMo/Automodel

Achyuthan-S · 2026-07-09T17:17:48Z

What does this PR do ?

Replaces the two-helper tie_word_embeddings guard API with a single per-class
declaration + one constructor guard, per the design agreed in #2935.

Each registered model that builds an lm_head declares a TieSupport policy
(BOTH / TIED_ONLY / UNTIED_ONLY) as a class attribute; a single
reject_unsupported_tie_word_embeddings(model_cls, config) reads that policy and
the controlling flag (get_controlling_tie_word_embeddings) and rejects an
unsupported request at construction, so the per-call-site choice between two
helpers — easy to get wrong for composite configs — no longer exists.

Opening as Draft: the layers land as separate commits (agreed in #2935).

Changelog

Landed (commit 1 — mechanical, no behavior change):

Add TieSupport enum + reject_unsupported_tie_word_embeddings(model_cls, config)
in checkpoint/utils.py; remove reject_unsupported_tied_word_embeddings and
reject_unsupported_untied_word_embeddings.
Migrate the 30 existing guard call sites 1:1 — 29 UNTIED_ONLY, gemma4_moe
TIED_ONLY — via a per-class tie_word_embeddings_support declaration.
Port the checkpoint-utils and qwen3_moe guard tests to the new guard; add
BOTH / undeclared-default coverage.

Remaining (subsequent commits in this PR):

Declarations per the audit: demote baichuan / mistral3 / bagel →
UNTIED_ONLY; add ernie4_5 / diffusion_gemma / gemma4_drafter →
TIED_ONLY; promote llama / qwen2 / qwen3_5 → BOTH.
bagel: read the controlling nested (text_config/llm_config) value; wire
the guard into build_bagel_from_hf_backbones (registry bypass).
Registry-iterating test with an explicit exempt list (retrieval encoders,
VLM containers, …), one comment per exemption.
Demotion test debt: flip ernie4_5 untied tests → rejection; baichuan
tie=True rejection + drop the unconditional _tied_weights_keys; add qwen2
untied test; add qwen3_5 tied+untied tests.
BOTH gaps: qwen2 tie_weights() override; qwen3_5 post-swap re-tie verify.
llama micro-cleanup: align the adapter's
getattr(config, "tie_word_embeddings", True) with the False default in
__init__/tie_weights.
Layer 2: NeMoAuto*.from_pretrained symmetric flip guard (raise on any
mismatch between the checkpoint's raw tie_word_embeddings and the
requested value, both directions); reword the PP validator at
pipelining/hf_utils.py:742.
Docs: onboarding skill §2.3 + checklist require the declaration, and
instruct per-model from_pretrained implementers to call the shared flip
check.

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

Closes refactor(models): single tie_word_embeddings guard with per-class declared TieSupport #2935
Builds on merged feat(models): complete tie_word_embeddings guards for #2512 families #2896 (the two-helper API this replaces) and fix(transformers): gate _tie_weights_nemo on tie_word_embeddings flag #2942 (the
_tie_weights_nemo gate the ernie4_5 demotion lines up with).

cc @yuhezhang-ai

…guard Replace reject_unsupported_tied/untied_word_embeddings with a single reject_unsupported_tie_word_embeddings(model_cls, config) driven by a per-class TieSupport declaration. Migrate the 30 existing guard call sites 1:1 (29 UNTIED_ONLY, gemma4_moe TIED_ONLY) with no behavior change; update the checkpoint-utils and qwen3_moe guard tests. Part of NVIDIA-NeMo#2935. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

copy-pr-bot · 2026-07-09T17:17:51Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Copilot

Pull request overview

This PR refactors the tie_word_embeddings constructor-guard API from two separate helpers into a single, per-model policy declaration (TieSupport) plus a unified guard (reject_unsupported_tie_word_embeddings), reducing the risk of incorrect guard selection for composite configs.

Changes:

Introduces TieSupport and reject_unsupported_tie_word_embeddings(model_cls, config) in nemo_automodel/components/checkpoint/utils.py, removing the prior split helper interface.
Updates affected model classes to declare tie_word_embeddings_support and to call the unified guard early in __init__.
Updates and extends unit tests to validate UNTIED_ONLY, TIED_ONLY, BOTH, and the undeclared-default behavior.

Reviewed changes

Copilot reviewed 33 out of 33 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
tests/unit_tests/utils/test_checkpoint_utils.py	Ports guard unit tests to the unified helper and adds coverage for `BOTH` + undeclared default.
tests/unit_tests/models/qwen3_moe/test_qwen3_moe_tie_guard.py	Updates docstring to reflect the new `TieSupport` + unified guard behavior.
nemo_automodel/components/checkpoint/utils.py	Adds `TieSupport` enum and replaces two guard helpers with `reject_unsupported_tie_word_embeddings(model_cls, config)`.
nemo_automodel/components/models/step3p7/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/step3p5/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/qwen3_vl_moe/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/qwen3_omni_moe/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/qwen3_next/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/qwen3_moe/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/qwen3_5_moe/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/qwen2_5_omni/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/nemotron_v3/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/nemotron_parse/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/nemotron_omni/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/mistral4/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/mistral3_vlm/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/minimax_m3_vl/model.py	Declares `tie_word_embeddings_support` (for both relevant classes) and switches constructor guard calls to unified helper.
nemo_automodel/components/models/minimax_m2/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/mimo_v2_flash/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/llava_onevision/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/ling_v2/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/kimivl/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/kimi_k25_vl/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/hy_v3/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/hy_mt2/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/gpt_oss/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/glm4_moe/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/glm4_moe_lite/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/glm_moe_dsa/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/gemma4_moe/model.py	Declares `tie_word_embeddings_support` as `TIED_ONLY` and switches constructor guard call to unified helper.
nemo_automodel/components/models/deepseek_v4/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/deepseek_v32/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.
nemo_automodel/components/models/deepseek_v3/model.py	Declares `tie_word_embeddings_support` and switches constructor guard call to unified helper.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

…BOTH) Mistral3FP8VLMForConditionalGeneration loads both the tied Ministral-3 (lm_head not serialized) and untied Mistral-Medium-3.5-128B checkpoints; it inherits HF's tie_weights() and does not swap the language_model backbone, so the tied path works. Declaring UNTIED_ONLY rejected the tied load (AM-657). Declare TieSupport.BOTH, drop the construction reject, and let the from_pretrained flip guard enforce per-checkpoint semantics. Add tied/untied construction tests. Part of NVIDIA-NeMo#2935. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

yuhezhang-ai · 2026-07-09T19:35:06Z

Thanks for already tracking the Qwen3.5 post-swap re-tie item. One concrete data point: Qwen2 and Qwen3_5ForCausalLM already alias/separate correctly in a tiny tied/untied construction probe. The remaining failure I see is specifically Qwen3_5ForConditionalGeneration with tie_word_embeddings=True: after replacing self.model.language_model with Qwen3_5DenseTextBackbone, lm_head.weight is no longer aliased to model.language_model.embed_tokens.weight. Calling tie_weights() after the swap fixes it.

The supported dense Qwen3.5 VLM configs are mixed by size (0.8B/2B/4B tied, 9B/27B/Qwen3.6-27B untied), so I think this class should land as TieSupport.BOTH with tiny tied/untied alias tests.

yuhezhang-ai · 2026-07-09T19:35:26Z

For the final BOTH declarations, can we keep the model-onboarding contract consistent: explicit _tied_weights_keys, explicit tie_weights() using the actual local FQNs, and tiny tied/untied alias tests? Qwen2 currently behaves correctly via the HF base implementation, so this is not fixing a known runtime bug, but making the declared BOTH support robust and self-documenting.

Could you also include a small VLM audit table before marking this ready? The risky pattern is HF parent init tying weights, followed by replacing language_model / lm_head. Gemma4 already re-ties after the swap; Qwen3.5 dense VLM still needs that. Most other supported custom VLM examples look untied-only, but Qwen3.5 dense VLM and Mistral3FP8VLM are mixed/BOTH.

Achyuthan-S · 2026-07-09T19:49:03Z

Thanks — agreed on both, and I ran the swap/re-tie audit across the custom VLMs.

Confirmed your Qwen3.5 finding: Qwen3_5ForConditionalGeneration swaps in
Qwen3_5DenseTextBackbone (model.py:845) with no re-tie, so tie=True orphans the
alias; the tie_weights() at :670 belongs to the ForCausalLM sibling. I'll land
this class as BOTH with a tie_weights() that re-ties after the swap, plus tiny
tied/untied alias tests.

VLM audit table (the risky pattern = HF parent ties on init, then language_model/
lm_head is replaced):

Class	HF parent ties	Replaces LM/head	Re-ties after	Verdict
Mistral3FP8VLMForConditionalGeneration	yes	no swap	n/a (HF tie holds)	BOTH
Qwen3_5ForConditionalGeneration	yes	yes (:845)	NO → adding	BOTH
Qwen3_5ForCausalLM	no	yes + own head	yes (:670)	BOTH
Gemma4ForConditionalGeneration	yes	yes (:1021)	yes (:1050)	TIED_ONLY
Qwen3VLMoeForConditionalGeneration	yes	yes + own head	no	UNTIED_ONLY
Qwen3_5MoeForConditionalGeneration	yes	yes + own head	no	UNTIED_ONLY
MiniMaxM3 / KimiVL / KimiK25VL / LLaVAOneVision / NemotronOmni / NemotronParse / Step3p7	no	own separate head	n/a	UNTIED_ONLY
Qwen2_5Omni / Qwen3OmniMoe thinkers	HF thinker	no re-tie	n/a	UNTIED_ONLY

Every UNTIED_ONLY VLM is safe because the constructor guard rejects tie=True
before the swap can orphan the head; gemma4 already re-ties; mistral3_vlm doesn't
swap so HF's tie holds.

On the contract: I'll add explicit _tied_weights_keys + tie_weights() (local
FQNs) + tied/untied tests to qwen2, Qwen3_5ForCausalLM, and
Qwen3_5ForConditionalGeneration. For Mistral3FP8VLM — since it doesn't swap and
uses HF's inherited tie — do you want an explicit local override there too, or is
the inherited contract + the alias tests I added sufficient?

Declare TieSupport.BOTH on llama, qwen2, and the dense qwen3_5 classes with the onboarding contract (explicit _tied_weights_keys + tie_weights + tied/untied tests). Qwen3_5ForConditionalGeneration replaced self.model.language_model with the NeMo backbone after HF's post_init tie, orphaning lm_head under tie=True; add a tie_weights() that re-ties to the active embedding and call it in __init__. Add an explicit tie_weights() to qwen2 (HF base ties unreliably under transformers v5) and align the llama adapter's tie default to False. Part of NVIDIA-NeMo#2935. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

yuhezhang-ai · 2026-07-09T20:39:23Z

For Mistral3FP8VLMForConditionalGeneration, I think the inherited HF tie_weights() plus the tied/untied alias tests are sufficient. The explicit local override is important for our custom classes that own or swap the embedding/head path, like Qwen2 and Qwen3.5, especially Qwen3_5ForConditionalGeneration after the language-model replacement.

Since Mistral3FP8VLM does not swap language_model/lm_head, a local override would just duplicate HF’s contract and may be more brittle than useful. The alias tests are the right proof here.
Thanks!

yuhezhang-ai · 2026-07-09T21:08:43Z

btw, after the tasks are done, on unit test side, Can we make the TieSupport declaration mandatory via the registry-iterating test? The helper currently defaults undeclared classes to BOTH, so review/ci test is the only thing catching a missing declaration. I think each registered class that owns/builds an lm_head should explicitly declare TieSupport, with an explicit commented exemption list for classes that do not own the language head / are not CausalLM-style.

Achyuthan-S · 2026-07-10T08:00:16Z

Sounds good — keeping Mistral3FP8VLM on HF's inherited tie_weights + the alias
tests, no local override.

And agreed on making the declaration mandatory: the registry-iterating test will
require every MODEL_ARCH_MAPPING class that owns/builds an lm_head to declare
TieSupport, with an explicit commented exemption list for the non-CausalLM /
no-head classes (retrieval encoders, VLM containers, …). I'll add that as its own
commit so the undeclared→BOTH default can never silently hide a missing verdict.

…ONLY Neither arch serves a tied checkpoint (both configs default tie=False). Declare UNTIED_ONLY + reject tie=True at construction; add rejection tests. Drop baichuan's unconditional _tied_weights_keys since its NormHead lm_head cannot be meaningfully tied to the input embedding. Part of NVIDIA-NeMo#2935. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

Achyuthan-S requested a review from a team as a code owner July 9, 2026 17:17

Copilot AI review requested due to automatic review settings July 9, 2026 17:17

Copilot started reviewing on behalf of Achyuthan-S July 9, 2026 17:18 View session

Copilot AI reviewed Jul 9, 2026

View reviewed changes

github-actions Bot added the community-request label Jul 9, 2026

Achyuthan-S mentioned this pull request Jul 9, 2026

refactor(models): single tie_word_embeddings guard with per-class declared TieSupport #2935

Open

9 tasks

svcnvidia-nemo-ci added the waiting-on-customer Waiting on the original author to respond label Jul 10, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

refactor(models): single tie_word_embeddings guard via TieSupport#2998

refactor(models): single tie_word_embeddings guard via TieSupport#2998
Achyuthan-S wants to merge 4 commits into
NVIDIA-NeMo:mainfrom
Achyuthan-S:Achyuthan-S/refactor/tie-word-embeddings-support

Achyuthan-S commented Jul 9, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jul 9, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

yuhezhang-ai commented Jul 9, 2026

Uh oh!

yuhezhang-ai commented Jul 9, 2026

Uh oh!

Achyuthan-S commented Jul 9, 2026

Uh oh!

yuhezhang-ai commented Jul 9, 2026

Uh oh!

yuhezhang-ai commented Jul 9, 2026

Uh oh!

Achyuthan-S commented Jul 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

Achyuthan-S commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Changelog

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Jul 9, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

yuhezhang-ai commented Jul 9, 2026

Uh oh!

yuhezhang-ai commented Jul 9, 2026

Uh oh!

Achyuthan-S commented Jul 9, 2026

Uh oh!

yuhezhang-ai commented Jul 9, 2026

Uh oh!

yuhezhang-ai commented Jul 9, 2026

Uh oh!

Achyuthan-S commented Jul 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Achyuthan-S commented Jul 9, 2026 •

edited

Loading