fix(transformers): _tie_weights_nemo force-ties untied models (ignores tie_word_embeddings)

[yuhezhang-ai]

Summary

_tie_weights_nemo (nemo_automodel/_transformers/model_init.py:994-1007) re-ties lm_head.weight to the input embedding for any model exposing _nemo_tied_weights_keys — without checking config.tie_word_embeddings. For untied models this silently discards or corrupts trained weights.

Introduced in #1817 ("fix: tie weights outside _init_model") together with the patched post_init in _transformers/utils.py that sets _nemo_tied_weights_keys. #1817's intent was correct — re-tie tied pre-v5 remote-code models (Nemotron-Flash-1B) after checkpoint loading breaks the construction-time alias — but the re-tie fires unconditionally. Pre-dates the #2512/#2896 guard work; surfaced during the tie-support audit for #2935.

Affected paths

Two ways a model gets _nemo_tied_weights_keys:

1. Static class attribute: ernie4_5 dense + MoE (components/models/ernie4_5/model.py:416,520). An ERNIE config with tie_word_embeddings=False branches correctly in __init__, then gets force-tied anyway at the end of _init_model (model_init.py:1040-1041).
2. Dynamic (broader): the patched post_init (_transformers/utils.py) sets it for every pre-v5 remote-code model with list-form _tied_weights_keys. In pre-v5 HF convention, declaring the keys does not mean the model is tied — HF's own tie_weights() checks the config flag. Many models declare the keys but ship untied checkpoints; all of them are force-tied when loaded through NeMoAuto*.

Observable failure modes

• Fresh from_pretrained of an untied model (silent): HF loads the trained lm_head.weight, then _tie_weights_nemo aliases it away to embed_tokens.weight. No error, no spike — logits simply diverge from the HF reference from the first forward; the model is effectively tied although the checkpoint (and the class it will be served with) is untied.
• Resume/load of a checkpoint with distinct lm_head/embedding weights (loud): after construction the two params share one storage; load_state_dict/DCP copies both checkpoint tensors into the same storage — last writer wins, corrupting embeddings and/or head → loss spike on resume.

Proposed fix

Gate the re-tie on the controlling config flag (reusing the canonical resolver from components/checkpoint/utils.py):

def _tie_weights_nemo(model):
    if not hasattr(model, "_nemo_tied_weights_keys"):
        return
    if not get_controlling_tie_word_embeddings(model.config, type(model).__name__):
        return  # untied config: the loaded separate lm_head is authoritative
    ...

Tied models (Nemotron-Flash) keep the #1817 behavior unchanged. Add regression tests for both manifestations: (a) untied list-form remote-code model keeps distinct storage + exact HF logits parity through the NeMoAuto* path; (b) save-then-reload of distinct weights round-trips losslessly (no spike); (c) tied model still re-ties (no #1817 regression).

Relationship to #2935

Independent and complementary: #2935's constructor guards cover registered classes only. The dynamic path here affects unregistered HF/remote-code classes, which no per-class declaration can reach — this gate is the only line of defense for them. (The ERNIE static case additionally becomes unreachable once #2935 demotes ERNIE to tied-only.)