Audit `tie_word_embeddings` handling in custom model implementations

[yuhezhang-ai]

Some custom CausalLM/VLM model classes expose or pass through tie_word_embeddings,
but not all of them explicitly implement weight tying between the input embeddings
and lm_head.

This is usually fine for current released configs because many of these models
intentionally default to tie_word_embeddings=False, and some architectures may
require separate embedding/head weights, especially with pipeline parallelism or
asymmetric model structure.

However, the behavior is confusing when a user sets tie_word_embeddings=True.
It may not be obvious whether the model should actually tie lm_head.weight to
embed_tokens.weight, or whether the config value should be rejected or
documented as unsupported.

Suggested cleanup

Audit custom nn.Module model implementations that define lm_head and input
embeddings:

• If the architecture supports tied embeddings, add explicit _tied_weights_keys
  and tie_weights() handling.
• If the architecture intentionally requires separate weights, document that
  tie_word_embeddings=True is unsupported or validate against it.
• Add small tied/untied config tests where tying is supported.

This is not currently considered an urgent checkpoint-resume correctness issue
after the tied-lm-head checkpoint guard in #2511, but it would make model behavior clearer
and safer for future configs.

[Achyuthan-S]

Hi @yuhezhang-ai, I’d like to take #2512.
I can do a scoped pass:
1 .audit model classes with lm_head + input embeddings,
2. add explicit tie/unsupported handling,
3. add small tied/untied tests where supported.
If any of this is already covered by #2571 or internal work, I can adjust scope.

[yuhezhang-ai]

Hi @Achyuthan-S Thanks, that would be great. #2571 is already closed, so feel free to take #2512 as the broader follow-up audit. Appreciate the help!

[Achyuthan-S]

I went through the custom model classes that define an lm_head + input embeddings. Around 31 of them expose tie_word_embeddings but never alias lm_head.weight to the embedding when it's True — only llama and qwen2 handle tying end-to-end, including re-tying on checkpoint load in their state-dict adapters.

It's mostly masked today: almost all of these default to tie=False, and the FSDP + checkpoint paths re-tie anyway (moe/parallelizer.py, the #2511 guard). The gap shows up on single-GPU / non-FSDP runs and when loading a genuinely tied HF checkpoint (which omits lm_head.weight).

Proposed split (audit only — no code yet):

Actually tie-able — fix properly:

• gemma4_moe (Gemma4ForConditionalGeneration) — Gemma defaults tie=True, so this is the real one. It inherits HF's lm_head, with no explicit tie and no adapter re-tie.
• gpt2 — already hard-ties unconditionally; just needs to declare it.

Separate heads by design — reject tie=True with a clear error (like the existing pipelining guard):

• The rest (~29) — the DeepSeek / GLM / Qwen3 MoE variants, the VLMs, etc.
• For VLM wrappers like qwen3_vl_moe / qwen3_omni_moe, the guard has to check the top-level flag, not nested text_config (checkpoint code already treats their top-level head as untied).

I'd start with gemma4_moe (explicit tie_weights() + adapter re-tie + tied/untied test), then the reject guard as a follow-up — re-checking per-family tie defaults via AutoConfig first, since a small dense variant could legitimately ship tie=True.

@yuhezhang-ai — sound right? I have the full per-model list if useful.

[yuhezhang-ai]

@Achyuthan-S Thanks, this sounds directionally right.

Could you take Gemma4 MoE first? That looks like the real correctness issue: HF/Gemma defaults to tie_word_embeddings=True, but our MoE path replaces the language_model after HF init, so we should explicitly re-tie lm_head.weight to the active embed_tokens.weight and add tied/untied tests.

For GPT-2, I agree it is different: our implementation already hard-ties unconditionally and does not expose a supported tie_word_embeddings config path. So I do not think we need to make it configurable or reject tie_word_embeddings=False. If the audit found a concrete metadata gap, then a small declaration/test is fine; otherwise GPT-2 can probably be left as-is.

For the remaining models, I agree with the follow-up direction: if the HF/default architecture is untied, reject tie_word_embeddings=True with a clear error instead of silently pretending to support it. I would prefer we confirm HF defaults/checkpoint behavior per family before landing a broad reject list.

[Achyuthan-S]

@yuhezhang-ai Took Gemma4 MoE first. Confirmed the bug: HF's __init__ ties lm_head to the text embed_tokens, but the MoE path replaces language_model right after with a fresh embed_tokens, so the head ends up aliased to the orphaned embedding — and since Gemma defaults tie_word_embeddings=True, it silently runs untied.

Fix is a re-tie right after the swap: when tie_word_embeddings is set, point lm_head.weight at the now-active embed_tokens.weight. The shared param survives the bf16 cast in initialize_weights, and it makes the #2511 checkpoint guard actually recognize the tie (has_local_tied_lm_head → True), so save/resume drops and reconstructs the head correctly instead of letting two copies drift apart.

Added CPU tied/untied tests: tie shares storage after construction and survives init; untied keeps separate storage. PR incoming.

On GPT-2: agreed, leaving as-is. It hard-ties unconditionally and doesn't expose a tie_word_embeddings config path, and the audit didn't surface a concrete metadata gap that breaks anything — the checkpoint guard keys off storage sharing, not _tied_weights_keys — so no declaration needed.

For the rest: I'll confirm HF defaults / checkpoint behavior per family before any reject list and post findings as a separate follow-up rather than one broad change.

[Achyuthan-S]

@yuhezhang-ai — went through the rest of the families before starting the reject-guard PR. I checked each one's default tie_word_embeddings against transformers 5.8.1 and the in-repo configs. Short version: a blanket reject across the separate-head set would've been wrong in a couple of spots, so here's where things landed.

Most are straightforward. They all default to untied, so rejecting tie=True is safe:

deepseek_v3, deepseek_v32, deepseek_v4, glm4_moe, glm4_moe_lite, glm_moe_dsa, gpt_oss, qwen3_moe, qwen3_next, qwen3_5_moe, hy_mt2, hy_v3, ling_v2, mimo_v2_flash, minimax_m2, minimax_m3_vl, nemotron_v3, mistral4, kimivl, kimi_k25_vl, llava_onevision, nemotron_parse.

A couple of caveats in there: minimax_m2 and nemotron_v3 load via trust_remote_code, so that False is the upstreamed transformers default — I'd still spot-check the actual checkpoints. And the VLMs (kimivl, kimi_k25_vl, llava_onevision) are keyed on text_config (DeepSeek/Qwen backbone), which is also False.

A few need the flag read carefully rather than off the top-level config:

• mistral3_vlm has top-level tie=True but text_config=False, and the head follows the text config, so the guard has to look at text_config, not the top.
• qwen3_vl_moe and qwen3_omni_moe are the opposite case: top-level is the real flag (False), and the nested text_config says True but it's ignored. That's already how checkpoint/utils.py:114-117 hardcodes them, so we should match that and key on top-level.

The one that surprised me is qwen2_5_omni. Its text_config actually defaults tie=True, and since the wrapper is just super().__init__(thinker_config) with no language_model swap, HF's tie holds. So it's genuinely tied, not a separate-head model. I don't think it should get a reject guard at all — maybe a small tied-storage test later, like we did for gemma4.

Three I still need to confirm against their published checkpoints, since the configs aren't in transformers 5.8.1 yet: step3p5, step3p7, nemotron_omni. I'll check those before they go in.

For the guard, I'm not going to gate on is_tied_word_embeddings() directly — it returns True for anything tied (qwen2_5_omni included), so that would over-reject. Instead I'll mirror its flag-resolution (text_config first, top-level fallback, plus the hardcoded non-tied set in checkpoint/utils.py) and only apply the raise inside the separate-head model classes. That way qwen2_5_omni never reaches the guard, and the per-model nuances above fall out of the shared resolution.

I'll hold the PR until #2601 merges. Does this split look right, especially dropping qwen2_5_omni from the reject list?

[yuhezhang-ai]

@Achyuthan-S

Thanks for the detailed pass. I think the rule should be simpler for now: NeMo should follow the original HF model's actual tied/untied behavior. If a user tries to change the controlling tie_word_embeddings flag away from that HF-supported behavior, we should reject it unless we add an explicit tested conversion path.

I checked two special cases with tiny HF/NeMo instantiation under transformers 5.8.1:

• qwen2_5_omni: default construction is not tied. The Thinker/top-level config has tie_word_embeddings=False, nested text_config=True, and lm_head.weight is model.embed_tokens.weight is false. If the Thinker/top-level flag is set to true, HF does tie it. So this one should follow the Thinker/top-level flag, not nested text config.
• mistral3_vlm: default construction is tied. Top-level Mistral3Config.tie_word_embeddings=True, nested text config false, and lm_head.weight is model.language_model.embed_tokens.weight is true. So I would not classify this as separate-head based on text_config=False.

For this follow-up, I would avoid supporting tied/untied conversions that HF does not use by default. Changing an HF-tied model to untied requires materializing a separate lm_head.weight from the tied embedding; changing an HF-untied model to tied may discard a distinct trained lm_head.weight. We should not silently do either without explicit model-specific support and tests.

Implementation-wise, I think we need a small model-aware resolver instead of each call site guessing between top-level and text_config. The important part is: do not prefer text_config when the top-level config already has tie_word_embeddings.

Something like:

def get_controlling_tie_word_embeddings(config, model_class_name: str) -> bool:
    # Known composite models where the top-level / thinker-level config owns
    # the top-level lm_head tying decision.
    if "Qwen2_5OmniThinkerForConditionalGeneration" in model_class_name:
        return bool(getattr(config, "tie_word_embeddings", False))

    if "Mistral3FP8VLMForConditionalGeneration" in model_class_name:
        return bool(getattr(config, "tie_word_embeddings", False))

    if "Qwen3VLMoeForConditionalGeneration" in model_class_name:
        return bool(getattr(config, "tie_word_embeddings", False))

    if "Qwen3OmniMoeThinkerForConditionalGeneration" in model_class_name:
        return bool(getattr(config, "tie_word_embeddings", False))

    # General rule: top-level config wins when present.
    if hasattr(config, "tie_word_embeddings"):
        return bool(config.tie_word_embeddings)

    # Fallback only for models that do not expose a top-level tie flag.
    text_config = getattr(config, "get_text_config", lambda: None)()
    return bool(getattr(text_config, "tie_word_embeddings", False))

Then use this resolver in two places:

1. constructor/config guards: reject when the controlling flag requests a tied/untied mode NeMo does not support for that family
2. checkpoint tied-head utilities: replace the current text_config-first is_tied_word_embeddings() logic with this resolver

Keep has_local_tied_lm_head() as the storage guard for save filtering; it should still only drop lm_head.weight when the tensors actually share storage.

[Achyuthan-S]

Thanks @yuhezhang-ai — that clarifies it. I'll treat HF default tied/untied behavior (verified by construction + storage) as the source of truth, and reject when the controlling flag is changed away from that default unless we add an explicit tested conversion path.

Noted on the two cases I had wrong:

1. qwen2_5_omni: default untied at Thinker/top-level; follow top-level, not nested text_config.
2. mistral3_vlm: default tied at top-level; not separate-head on text_config=False.

I'll implement get_controlling_tie_word_embeddings(config, model_class_name) per your sketch, use it in constructor guards and in is_tied_word_embeddings(), and keep has_local_tied_lm_head() storage-based for save filtering. Holding the PR until #2601 merges.