feat(models): complete tie_word_embeddings guards for #2512 families

[Achyuthan-S]

What

• Reject tie=True on the untied-default families — mistral3_vlm, step3p5, step3p7, nemotron_omni — via the existing reject_unsupported_tied_word_embeddings guard.
• Reject tie=False on the tied-default gemma4_moe, via a new symmetric reject_unsupported_untied_word_embeddings helper. Gemma4 is now tied-only; the fix(gemma4_moe): re-tie lm_head to active embed_tokens on MoE path #2601 untied tests are flipped to assert the rejection (untying Gemma4 was never a tested load/conversion path — fresh construction leaving lm_head separate doesn't prove a tied checkpoint can materialize a correct separate lm_head.weight).

Verified per family (config inspection only, no weights)

model	checkpoint	observed	guard
mistral3_vlm	mistralai/Mistral-Medium-3.5-128B	tie_word_embeddings=False (top & nested) → untied	reject tie=True
step3p5	stepfun-ai/Step-3.5-Flash	untied (config load warns on validate_layer_type; flag False per inspection)	reject tie=True
step3p7	stepfun-ai/Step-3.7-Flash	no top-level flag, separate lm_head → untied	reject tie=True
nemotron_omni	nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16	nested llm_config tie=False → untied	reject tie=True
gemma4_moe	Gemma default	tied	reject tie=False

Tests

• New unit tests for the symmetric reject_unsupported_untied_word_embeddings helper.
• gemma4_moe untied tests updated to assert the rejection.
• Guards run at the start of __init__ on the top-level config, so they fail fast before any build.

Closes #2512.

cc @yuhezhang-ai

[copy-pr-bot]

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 completes the tie_word_embeddings audit by adding a symmetric guard for tied-default architectures and wiring fast-fail validation into additional model families so unsupported tying/untying configurations are rejected at construction time.

Changes:

• Added reject_unsupported_untied_word_embeddings() to explicitly reject tie_word_embeddings=False for tied-default models (Gemma4 MoE).
• Applied reject_unsupported_tied_word_embeddings() to additional untied-default model families (mistral3_vlm, step3p5, step3p7, nemotron_omni) to reject tie_word_embeddings=True.
• Updated and expanded unit tests to cover the new guard and to flip Gemma4 MoE “untied” tests to assert rejection.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/unit_tests/utils/test_checkpoint_utils.py	Adds unit tests for the new untied-guard helper and minor formatting cleanup.
tests/unit_tests/models/gemma4_moe/test_gemma4_moe_tied_weights.py	Updates Gemma4 MoE tests to assert tie_word_embeddings=False is rejected and adds pytest dependency.
nemo_automodel/components/models/step3p7/model.py	Adds a constructor-time guard rejecting tie_word_embeddings=True for an untied-default family.
nemo_automodel/components/models/step3p5/model.py	Adds a constructor-time guard rejecting tie_word_embeddings=True for an untied-default family.
nemo_automodel/components/models/nemotron_omni/model.py	Adds a constructor-time guard rejecting tie_word_embeddings=True for an untied-default family.
nemo_automodel/components/models/mistral3_vlm/model.py	Adds a constructor-time guard rejecting tie_word_embeddings=True for the supported untied Mistral3 VLM checkpoint.
nemo_automodel/components/models/gemma4_moe/model.py	Adds a constructor-time guard rejecting tie_word_embeddings=False for tied-default Gemma4 MoE.
nemo_automodel/components/checkpoint/utils.py	Introduces the new symmetric reject_unsupported_untied_word_embeddings() helper.

---

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

[Achyuthan-S]

Fixed — updated the docstring to say it raises when the controlling flag evaluates to False.

[yuhezhang-ai]

/ok to test b37ec0f

[yuhezhang-ai]

Thanks @Achyuthan-S. Will approve and merge after ci passed. I verified the guards against the actual checkpoint configs (including the nemotron_omni nested llm_config case) and the behavior is correct.

One follow-up I've decided on, tracked in #2935: replace the two-helper API with a single constructor-facing guard — each model class declares a three-state tie_word_embeddings_support (TIED_ONLY / UNTIED_ONLY / BOTH), and one helper compares the declared value against the config's requested value. That removes the per-call-site choice of which reject helper to use (easy to get wrong for composite configs), and a registry-level test will enforce that every registered model declares it. A second, smaller piece adds a flip check in the from_pretrained bridge for BOTH models (rejecting untied-checkpoint → tie=True, which would discard a trained lm_head).

You're welcome to pick it up if you're interested. Otherwise I'll take it. Either way, no action needed on this PR.

[Achyuthan-S]

Thanks @yuhezhang-ai — glad the checkpoint verification lined up, and thanks for the merge. I'll pick up #2935 as the follow-up (plan posted over there , on the issue page).