refactor(models): single tie_word_embeddings guard via TieSupport - #2998
Conversation
…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>
There was a problem hiding this comment.
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
TieSupportandreject_unsupported_tie_word_embeddings(model_cls, config)innemo_automodel/components/checkpoint/utils.py, removing the prior split helper interface. - Updates affected model classes to declare
tie_word_embeddings_supportand 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>
|
Thanks for already tracking the Qwen3.5 post-swap re-tie item. One concrete data point: Qwen2 and The supported dense Qwen3.5 VLM configs are mixed by size ( |
|
For the final 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 |
|
Thanks — agreed on both, and I ran the swap/re-tie audit across the custom VLMs. Confirmed your Qwen3.5 finding: VLM audit table (the risky pattern = HF parent ties on init, then language_model/
Every UNTIED_ONLY VLM is safe because the constructor guard rejects tie=True On the contract: I'll add explicit |
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>
|
For Since Mistral3FP8VLM does not swap |
|
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. |
|
Sounds good — keeping Mistral3FP8VLM on HF's inherited tie_weights + the alias And agreed on making the declaration mandatory: the registry-iterating test will |
…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 Thanks for the latest update. Since this PR includes correctness fixes we’d like to land soon (Mistral3FP8VLM BOTH and Qwen3.5 post-swap re-tie), do you think you’ll have time to finish the remaining checklist items shortly? The main remaining pieces seem to be Layer 2 |
…gemma/gemma4_drafter TIED_ONLY; bagel/mistral4-VLM UNTIED_ONLY) Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
…pport Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
…d; reword PP validator Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
…flip check Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
|
Thanks @yuhezhang-ai ! No need to take over — I've pushed the rest of the checklist:
479 green across the tie suite. Two flagged for review: diffusion_gemma/gemma4_drafter |
Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
| hf_config: The resolved config with user overrides applied (the requested value). | ||
| kwargs: The from_pretrained kwargs (hub-locating keys are reused for the raw load). | ||
| """ | ||
| if not isinstance(pretrained_model_name_or_path, str): |
There was a problem hiding this comment.
This skips the new flip guard for local checkpoints passed as pathlib.Path, even though from_pretrained accepts str | os.PathLike. In that case a user can still flip tie_word_embeddings silently. Can we normalize path-like inputs with os.fspath() and run the same raw-config check for both str and os.PathLike sources?
There was a problem hiding this comment.
Oh yes you are right , good catch , thank you — fixed . Normalized path-like sources with os.fspath() so the raw-config flip check runs for both str and os.PathLike, and added pathlib.Path coverage (rejected flip, matching value, and an assertion that the Path is normalized to str before the raw read).
from_pretrained accepts str | os.PathLike; the flip guard only handled str, so a pathlib.Path local checkpoint could silently flip tie_word_embeddings. Normalize path-like sources with os.fspath() and run the same raw-config check for both. Add pathlib.Path coverage to the flip-guard tests. Addresses review feedback. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
|
Hi @akoumpa , could you help run Also, if you want to take a look, this PR makes tied-embedding handling more robust by requiring models with an lm_head to explicitly declare their tie support and define tie_weights() when they own the embedding/head relationship. It also reduces reliance on ambiguous nested versus top-level config fields, particularly for VLMs. Scoped CI: |
…-contract Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> # Conflicts: # nemo_automodel/components/models/diffusion_gemma/model.py
|
/ok to test b965726 |
|
Thanks for pushing this over the line, @yuhezhang-ai — the runtime tie-support wiring and the extra tie_weights() hooks are nice hardening on top of the guard. CI results look great (AM-657 FT+PEFT, Qwen3.5, Mistral4 VLM all green). Happy with the refinements — nothing further from me; ready whenever the code-owner review lands. |
|
/nvskills-ci |
What does this PR do ?
Replaces the two-helper
tie_word_embeddingsguard API with a single per-classdeclaration + one constructor guard, per the design agreed in #2935.
Each registered model that builds an
lm_headdeclares aTieSupportpolicy(
BOTH/TIED_ONLY/UNTIED_ONLY) as a class attribute; a singlereject_unsupported_tie_word_embeddings(model_cls, config)reads that policy andthe controlling flag (
get_controlling_tie_word_embeddings) and rejects anunsupported 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):
TieSupportenum +reject_unsupported_tie_word_embeddings(model_cls, config)in
checkpoint/utils.py; removereject_unsupported_tied_word_embeddingsandreject_unsupported_untied_word_embeddings.UNTIED_ONLY, gemma4_moeTIED_ONLY— via a per-classtie_word_embeddings_supportdeclaration.BOTH/ undeclared-default coverage.Remaining (subsequent commits in this PR):
UNTIED_ONLY; add ernie4_5 / diffusion_gemma / gemma4_drafter →TIED_ONLY; promote llama / qwen2 / qwen3_5 →BOTH.text_config/llm_config) value; wirethe guard into
build_bagel_from_hf_backbones(registry bypass).VLM containers, …), one comment per exemption.
tie=True rejection + drop the unconditional
_tied_weights_keys; add qwen2untied test; add qwen3_5 tied+untied tests.
tie_weights()override; qwen3_5 post-swap re-tie verify.getattr(config, "tie_word_embeddings", True)with theFalsedefault in__init__/tie_weights.NeMoAuto*.from_pretrainedsymmetric flip guard (raise on anymismatch between the checkpoint's raw
tie_word_embeddingsand therequested value, both directions); reword the PP validator at
pipelining/hf_utils.py:742.instruct per-model
from_pretrainedimplementers to call the shared flipcheck.
Before your PR is "Ready for review"
Pre checks:
If you haven't finished some of the above items you can still open "Draft" PR.
Additional Information
_tie_weights_nemogate the ernie4_5 demotion lines up with).cc @yuhezhang-ai