fix(transformers): gate _tie_weights_nemo on tie_word_embeddings flag

[yuhezhang-ai]

What does this PR do ?

Fixes #2941: _tie_weights_nemo re-tied lm_head.weight to the input embeddings for any model exposing
_nemo_tied_weights_keys, without consulting the controlling config.tie_word_embeddings flag.
_nemo_tied_weights_keys records candidate tied keys from pre-v5/list-form _tied_weights_keys; it does not prove
that the checkpoint is actually tied.

Impact is real for shipped untied remote-code force-HF recipes. Confirmed affected:

• examples/llm_finetune/nemotron/nemotron_nano_4b_squad.yaml
• examples/llm_finetune/nemotron/nemotron_nano_9b_squad.yaml
• examples/llm_finetune/nemotron/nemotron_nano_9b_squad_peft.yaml

Those recipes load Nemotron Nano 4B/9B through trust_remote_code: true + force_hf: true. The hub configs set
tie_word_embeddings=False, the checkpoints contain separate trained lm_head.weight and
backbone.embeddings.weight, and the recipe compatibility patch converts remote list-form
_tied_weights_keys = ["lm_head.weight"] into _nemo_tied_weights_keys. Before this PR, _tie_weights_nemo then
forced an alias and discarded the trained head.

The fix gates the re-tie on the canonical get_controlling_tie_word_embeddings resolver. Tied models, such as the
Nemotron-Flash case #1817 was written for, keep the existing re-tie behavior unchanged. Models without a config
attribute conservatively fall back to re-tying.

Before / after

Real recipe-path source load: nvidia/NVIDIA-Nemotron-Nano-9B-v2.

Using the same recipe condition that exposes the bug (apply_cache_compatibility_patches() + trust_remote_code=True

• force_hf=True), compare a raw HF source load against a NeMoAuto source load on a seeded 128-token prompt:

branch	lm_head aliased?	max abs diff vs HF	mean abs diff vs HF	cosine
main	true	22.5625	2.9893	-0.0224
this PR	false	0.25	0.0141	about 1.000

This is the silent manifestation: the model loads successfully, but logits diverge from the raw HF reference immediately
on main because the trained untied lm_head is aliased away.

Unit/mechanism coverage.

The regression tests fail on main and pass on this branch:

• test_untied_config_keeps_separate_lm_head
• test_untied_state_dict_roundtrip_is_lossless
• test_tied_config_reties
• test_missing_config_still_reties

Targeted result on this branch: 4 passed.

Resume smoke note.

A real 9B FSDP2 DCP recipe resume smoke was also run. This PR produced a checkpoint at step 3, then main and this PR
both resumed one step from that same checkpoint. Both resumed step 4 at loss=0.466808.

That result means the DCP checkpoint resume path was correct in this smoke. It does not cover the bug fixed here: on
main, the model is already wrong immediately after the initial NeMoAuto source load, before any training checkpoint is
saved or resumed. This is also why the current checkpoint robustness test can miss the issue: it captures reference
logits from the already-loaded NeMo model, then validates checkpoint reloads against that reference. We should add a
separate source-load parity check that compares raw HF source load vs NeMoAuto source load before training starts.

Changelog

• _tie_weights_nemo now returns early when the controlling tie_word_embeddings flag is False, preserving
  separately loaded lm_head.weight for untied models.
• Added TestTieWeightsNemoConfigGate regression tests: untied preservation, tied re-tie, missing-config fallback, and
  lossless construct-then-load round-trip.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed the Contributor
  Guide.
• Did you write any new necessary tests?
  • tests/unit_tests/_transformers/test_model_init.py::TestTieWeightsNemoConfigGate
• Did you add or update any necessary documentation?
  • No public API change.
• Does this PR introduce a breaking change?
  • No. This restores the checkpoint-config contract for untied models and preserves tied behavior.

Additional Information

• Related to refactor(models): single tie_word_embeddings guard with per-class declared TieSupport #2935 but intentionally kept as a minimal correctness fix instead of refactoring the full tie-weight
  pipeline.
• Introduced by the fix: tie weights outside _init_model #1817-era workaround for pre-v5/list-form _tied_weights_keys. That workaround is still needed for
  tied remote-code models; this PR only adds the missing config gate for untied checkpoints.
• Existing checkpoint robustness coverage does not exercise this exact source-load failure. It captures reference logits
  from the already-constructed NeMo model, then compares checkpoint reloads against that reference. If the source load is
  already corrupted, the corrupted model can become the baseline. We should add an opt-in raw-HF-source vs
  NeMoAuto-source parity check for remote-code/force-HF recipes.

[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.

[yuhezhang-ai]

/ok to test 9199d1e