fix(checkpoint): resolve tie_word_embeddings top-level-first to match HF tying by Achyuthan-S · Pull Request #2732 · NVIDIA-NeMo/Automodel

Achyuthan-S · 2026-06-23T13:09:59Z

First of two PRs for the tied-embedding audit (#2512). This one is the shared resolver plus the is_tied_word_embeddings() swap; the constructor reject-guards across the separate-head model classes will follow in a second PR.

What

Adds get_controlling_tie_word_embeddings(config, model_class_name) and routes is_tied_word_embeddings() through it. The old logic read text_config.tie_word_embeddings first, which disagrees with how HF actually ties lm_head — HF ties on the top-level config flag. Verified by construction under transformers 5.8.1:

mistral3_vlm: top-level True / text False → HF builds it tied (old code reported untied)
qwen2_5_omni: top-level False / text True → HF builds it untied (old code reported tied)

So the resolver is top-level-first, falling back to text_config only when there's no top-level flag. has_local_tied_lm_head() is unchanged (still the storage-based guard for save filtering).

One thing to confirm

Your sketch returned the top-level flag for all four composite models. I kept Qwen3OmniMoe / Qwen3VLMoe force-untied instead, because following top-level for them flips the existing test_model_state_disables_tied_embeddings_for_non_tied_models guarantee (it keeps lm_head.weight in the saved state dict for those). Their real configs default top-level untied, so it's identical in practice — I just went conservative to preserve that checkpoint safety. Qwen2_5Omni / Mistral3 follow top-level via the general rule, as intended. Happy to switch those two to follow top-level too if you'd prefer — it'd just mean updating that test.

Tests

Updated the resolver / is_tied unit tests to the top-level-first contract and added resolver coverage. Full checkpoint suite stays green.

pytest tests/unit_tests/utils/test_checkpoint_utils.py tests/unit_tests/checkpoint/ -q
# 13 + 208 passed

Resolver semantics (resolved in review)

Per @yuhezhang-ai's review, the resolver returns the actual top-level flag for all composite models, including Qwen3 Omni/VL MoE (no force-untied). Checkpoint save safety stays storage-based via has_local_tied_lm_head() — a config-tied model with a separate lm_head still keeps lm_head.weight on save. Added ModelState tests for both directions (kept when storage is separate; dropped when shared).
(Also bump the test line from # 13 + 208 passed to # 13 + 209 passed.)

Part of #2512 — reject-guard PR to follow.

cc @yuhezhang-ai

… HF tying Add get_controlling_tie_word_embeddings() and route is_tied_word_embeddings() through it. HF ties lm_head on the top-level config flag, not nested text_config (verified by construction for mistral3_vlm and qwen2_5_omni under transformers 5.8.1). Keep Qwen3OmniMoe/Qwen3VLMoe force-untied to preserve the ModelState save guarantee; has_local_tied_lm_head() is unchanged. Refs NVIDIA-NeMo#2512 Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

copy-pr-bot · 2026-06-23T13:10:02Z

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 updates NeMo AutoModel’s checkpoint utilities to resolve tie_word_embeddings using Hugging Face’s effective semantics (top-level config flag controls tying), and adds unit tests to lock in the new contract. This supports the tied-embedding audit (#2512) by making tied/untied detection consistent with HF construction behavior, which affects checkpoint save/load filtering decisions around lm_head.weight.

Changes:

Add get_controlling_tie_word_embeddings(config, model_class_name) with top-level-first resolution and a small force-untied exclusion list.
Route is_tied_word_embeddings() through the new resolver (replacing the prior text_config-first behavior).
Update and extend unit tests to cover the resolver contract and exclusion behavior.

Reviewed changes

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

File	Description
`nemo_automodel/components/checkpoint/utils.py`	Introduces the shared tie-flag resolver and updates `is_tied_word_embeddings()` to use top-level-first semantics.
`tests/unit_tests/utils/test_checkpoint_utils.py`	Updates existing tests to the new top-level-first contract and adds resolver-focused coverage.

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

Achyuthan-S · 2026-06-23T13:18:34Z

    """
-    non_tied_lm_head_models = {
-        "Qwen3OmniMoeThinkerForConditionalGeneration",  # complicated config structure
-        "Qwen3VLMoeForConditionalGeneration",  # top-level lm_head is untied despite nested text config
-    }
-    model_class_name = type(model).__name__
-    for m in non_tied_lm_head_models:
-        if m in model_class_name:
-            return False
    config = getattr(model, "config", None)
-    text_config = getattr(config, "get_text_config", lambda: None)()
-    return bool(getattr(text_config, "tie_word_embeddings", getattr(config, "tie_word_embeddings", False)))
+    if config is None:
+        return False
+    return get_controlling_tie_word_embeddings(config, type(model).__name__)


Good catch, though worth noting this is pre-existing (the old code matched on type(model).__name__ too), and the top-level-first change actually improves the wrapped case: old text-first would mis-report these as tied when the exclusion missed, whereas top-level-first returns untied (their real default) regardless of wrapping. The exclusion only adds value for a synthetic top=True. If we keep the force-untied exclusion (the open question above), I'll add wrapper-unwrapping so the class-name match is reliable.

Achyuthan-S · 2026-06-23T13:18:51Z

 def test_is_tied_word_embeddings_respects_qwen3_vl_moe_exclusion():
+    """Qwen3VLMoe stays force-untied (separate top-level head), ignoring nested text_config=True."""
+
    class DummyTextConfig:
        tie_word_embeddings = True


Agreed — if the force-untied exclusion stays (the open question in the description), I'll set the top-level flag to True here so the test actually guards the exclusion rather than passing via the general rule. Holding until we settle whether these two follow top-level or stay force-untied.

yuhezhang-ai · 2026-06-24T01:21:24Z

@Achyuthan-S Thanks for the PR!

I’m okay with the resolver-first approach, but I don’t think we should merge this as-is.

Two points:

I’d like the resolver to represent the actual controlling flag, even for Qwen3 Omni/VL MoE. So for those models it should still return the top-level tie_word_embeddings value, not force-return False. The checkpoint save safety should continue to come from has_local_tied_lm_head() being storage-based. Otherwise the follow-up constructor guards won’t be able to detect and reject the unsupported top_level=True case.
Since this PR is mostly shared plumbing, I’d like either a small amount of standalone coverage showing the checkpoint behavior it fixes now, or we can just fold this into the reject-guard PR. I don’t want to merge a resolver-only change whose main consumer is still future work unless the contract and immediate behavior are clearly pinned down.

So my preference is: update this PR to follow the top-level flag for the Qwen3 cases and add/adjust tests around the actual checkpoint behavior; then we can keep the broader audit/reject guards in a second PR. Otherwise, please combine this with the guard work so we review the end-to-end behavior together.

…eview Per review: resolver returns the actual top-level flag for Qwen3 Omni/VL MoE instead of forcing untied, so the constructor guard can detect/reject unsupported top-level=True. Save safety stays storage-based via has_local_tied_lm_head(). Adds ModelState coverage for the storage-gated save behavior. Refs NVIDIA-NeMo#2512 Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

Achyuthan-S · 2026-06-24T05:08:26Z

Hi @yuhezhang-ai , updated per your review: Qwen3 Omni/VL MoE now return the top-level flag (force-untied removed). Added two ModelState tests — config-tied but separate storage keeps lm_head.weight; shared storage drops it. Guards still planned as PR 2.

yuhezhang-ai · 2026-06-24T16:50:16Z

/ok to test f45cbee

Achyuthan-S · 2026-06-25T16:11:05Z

Hey @yuhezhang-ai — thanks for the review. CI is green on f45cbee. Let me know if anything else is needed from my side before merge; constructor guards will follow as PR 2 for #2512.

…families (#2805) * feat(models): reject tie_word_embeddings=True on separate-head model families Add reject_unsupported_tied_word_embeddings() (built on the #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 #2512 Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com> * fix(models): run tie guard on top-level config before unwrap/super (review) 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 #2512 Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com> * fix(checkpoint): resolve Omni wrapper tie flag from thinker_config (review) 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 #2512 Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com> --------- Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

Achyuthan-S requested a review from a team as a code owner June 23, 2026 13:09

Copilot AI review requested due to automatic review settings June 23, 2026 13:09

github-actions Bot added the community-request label Jun 23, 2026

Copilot started reviewing on behalf of Achyuthan-S June 23, 2026 13:10 View session

Copilot AI reviewed Jun 23, 2026

View reviewed changes

Achyuthan-S mentioned this pull request Jun 23, 2026

Audit tie_word_embeddings handling in custom model implementations #2512

Closed

copy-pr-bot Bot temporarily deployed to nemo-ci June 24, 2026 16:50 Inactive

copy-pr-bot Bot temporarily deployed to test June 24, 2026 16:50 Inactive

copy-pr-bot Bot temporarily deployed to public June 24, 2026 16:50 Inactive

copy-pr-bot Bot temporarily deployed to public June 24, 2026 16:54 Inactive

copy-pr-bot Bot temporarily deployed to public June 24, 2026 16:55 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci June 24, 2026 17:41 Inactive

yuhezhang-ai approved these changes Jun 25, 2026

View reviewed changes

yuhezhang-ai merged commit 27f9c34 into NVIDIA-NeMo:main Jun 25, 2026
70 checks passed

Achyuthan-S mentioned this pull request Jun 26, 2026

feat(models): reject tie_word_embeddings=True on separate-head model families #2805

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

fix(checkpoint): resolve tie_word_embeddings top-level-first to match HF tying#2732

fix(checkpoint): resolve tie_word_embeddings top-level-first to match HF tying#2732
yuhezhang-ai merged 2 commits into
NVIDIA-NeMo:mainfrom
Achyuthan-S:Achyuthan-S/feat/tie-embeddings-reject-guard

Achyuthan-S commented Jun 23, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jun 23, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Achyuthan-S Jun 23, 2026

Uh oh!

Achyuthan-S Jun 23, 2026

Uh oh!

yuhezhang-ai commented Jun 24, 2026

Uh oh!

Achyuthan-S commented Jun 24, 2026

Uh oh!

yuhezhang-ai commented Jun 24, 2026

Uh oh!

Achyuthan-S commented Jun 25, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

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

Uh oh!

What

One thing to confirm

Tests

Resolver semantics (resolved in review)

Uh oh!

copy-pr-bot Bot commented Jun 23, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Achyuthan-S Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

Achyuthan-S Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

yuhezhang-ai commented Jun 24, 2026

Uh oh!

Achyuthan-S commented Jun 24, 2026

Uh oh!

yuhezhang-ai commented Jun 24, 2026

Uh oh!

Achyuthan-S commented Jun 25, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Achyuthan-S commented Jun 23, 2026 •

edited

Loading