fix(checkpoint): preserve tied lm_head on resume

[yuhezhang-ai]

What does this PR do ?

Fix tied-LM-head checkpoint/resume and sharded-training behavior by making lm_head.weight actually alias the input embedding weight when the model config requests tied word embeddings, and by only omitting lm_head.weight from checkpoints when that alias is real.

Bug Details

This is a checkpoint and model-state correctness bug for tied word embeddings.

For models with config.tie_word_embeddings=true, AutoModel treated lm_head.weight as redundant and omitted it from the saved checkpoint state. That is only safe when lm_head.weight and the input embedding weight actually share storage. In the failing Llama pretraining reproduction, the model was built from config on meta tensors and materialized/initialized later, after __init__ had already run. The config claimed tied embeddings, but the resident tensors were not physically tied at checkpoint-save time.

On resume, the checkpoint loader reconstructed the missing lm_head.weight from embed_tokens.weight. Because the trained lm_head.weight had diverged as an untied parameter before save, resume silently changed the model state and produced a large first-step loss spike.

The original confirmed reproduction is Llama 3.2 1B from-config pretraining with tie_word_embeddings=true. The issue is not specific to Slurm SIGTERM/enforced checkpoints: normal checkpoint resume showed the same spike.

The affected area is broader than pretraining. Any LLM/VLM path can be affected when:

• the model/text config says tie_word_embeddings=true
• the local model or pipeline stage owns both lm_head.weight and the input embedding weight
• initialization, checkpoint load, or FSDP2/DTensor sharding leaves those tensors as separate parameters before optimizer construction or checkpoint save

Confirmed or directly exercised cases include:

• Llama from-config pretraining with tied embeddings
• Gemma3 VLM / Gemma3ForConditionalGeneration, whose HF config declares tied text embeddings and whose sharded local model can need the alias re-applied after FSDP2/DCP state handling

This PR should be a no-op for models with tie_word_embeddings=false, intentionally shape-mismatched heads such as speculative draft heads, pipeline stages that do not own both tensors, and models that already keep the alias intact.

Changelog

• Add ensure_tied_lm_head() to re-establish a local lm_head / input-embedding alias when the config requests tied word embeddings.
• Prefer model-native tie_weights() and add a direct local assignment fallback for wrapped/sharded models where generic tying does not reach both tensors.
• Re-apply tied LM-head aliases after meta/from-config initialization and after base-model checkpoint load.
• Re-apply tied LM-head aliases immediately after non-PP FSDP2/DTensor sharding, before optimizer construction.
• Refresh ModelState tied-head metadata immediately before save-time key filtering, so DCP/FSDP alias normalization cannot leave a stale has_local_tied_lm_head=false and accidentally persist duplicate lm_head.weight.
• Re-apply tied LM-head aliases after DCP/full-state model loads, so compatibility materialization cannot leave the restored model split.
• Make tied-LM-head checkpoint detection require actual local storage aliasing before dropping lm_head.weight.
• Keep shape-mismatch guards so asymmetric heads are saved as independent parameters.
• Add explicit tie_weights() support for the custom Llama causal LM implementation.
• Update model-onboarding guidance and LLM implementation patterns for explicit tied-head support.
• Add/extend unit tests for tied, untied, storage-untied, stale metadata, shape-mismatched, and repair behavior.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

Documentation note: this updates contributor/model-onboarding guidance rather than user-facing docs because there is no new user API or config field.

Validation:

• After latest CI follow-up: container run of python -m pytest tests/unit_tests/utils/test_checkpoint_utils.py tests/unit_tests/checkpoint/test_checkpointing.py tests/unit_tests/models/llama/test_llama_tied_weights_cpu.py -q passed with 111 passed, 2 skipped.
• After latest CI follow-up: container run of ruff format --check and ruff check on the touched checkpoint files/tests passed.
• Tiny two-GPU Gemma3 VLM DCP roundtrip probe passed: ModelState omitted lm_head.weight, optimizer state stayed under model.language_model.embed_tokens.weight, restored loss matched source loss, and the restored model remained tied.
• Earlier after merging latest main: container run of python -m pytest tests/unit_tests/utils/test_checkpoint_utils.py tests/unit_tests/checkpoint/test_checkpointing.py tests/unit_tests/models/llama/test_llama_tied_weights_cpu.py -q passed with 109 passed, 2 skipped.
• source work/runs/_shared/env.sh && uv run --no-sync pytest tests/unit_tests/checkpoint/test_checkpointing.py tests/unit_tests/models/llama/test_llama_tied_weights_cpu.py -q
• pytest tests/unit_tests/utils/test_checkpoint_utils.py -q in the Automodel container: 11 passed
• ruff format / ruff check on the touched checkpoint, infrastructure, model, skill, and test files
• Reproduced normal-checkpoint resume loss spike before the fix and verified the first post-resume loss spike is gone after the fix
• Verified with targeted FSDP2 probes that an already-tied model remains tied after sharding, and that a Gemma3 VLM-style sharded model with a broken local alias is repaired by ensure_tied_lm_head()

Additional Information

• Related to the tied-embedding resume loss-spike report discussed in the PR thread.

[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 9920aa0

[yuhezhang-ai]

/ok to test fca98ed

[yuhezhang-ai]

/ok to test dbeb5c8

[yuhezhang-ai]

/ok to test 50f781e

[akoumpa]

/claude review

[akoumpa]

/ok to test 406d593

[akoumpa]

full CI has passed https://github.com/NVIDIA-NeMo/Automodel/actions/runs/27562491560/job/81515688479?pr=2511

[akoumpa]

FMing; codecov is hang in status report.