fix(megatron): tolerate model _extra_state keys absent from the checkpoint (QAD)

[ChenhanYu]

What does this PR do?

Type of change: Bug fix (+ new example)

Fixes a CheckpointException when resuming a sharded ModelOpt checkpoint for Megatron QAD/QAT (load a quantized NVFP4 checkpoint, then train with --modelopt-enabled):

torch.distributed.checkpoint ... default_planner.py: create_default_local_load_plan
RuntimeError: Missing key in checkpoint state_dict: decoder.final_layernorm._extra_state/shard_0_1.

Root cause: _load_extra_state_from_sharded_checkpoint builds the extra-state request from the model's sharded_state_dict. When the model declares an _extra_state entry the checkpoint lacks — e.g. a TransformerEngine version bump adds decoder.final_layernorm._extra_state that an older/quantized checkpoint never saved — torch DCP's planner hard-raises before Megatron's StrictHandling runs (LOG_UNEXPECTED does not drop ShardedObject extra_state keys).

Fix: restrict the request to _extra_state keys actually present in the checkpoint (via load_sharded_metadata). A genuinely-missing norm extra_state is safe to leave freshly initialized — it holds fp8/RNG metadata, not learned weights.

Also adds a launcher QAD example (tools/launcher/examples/meta-llama/Llama-3.2-1B-Instruct/megatron_lm_qad.yaml): PTQ → MMLU gate → QAD/SFT train, each task calling the Megatron-LM ModelOpt example scripts (quantize.sh / mmlu.sh / train.sh) directly.

Usage

uv run launch.py --yaml examples/meta-llama/Llama-3.2-1B-Instruct/megatron_lm_qad.yaml --yes

Testing

• slurm.py --yaml ... --dryrun resolves all 3 tasks (EXIT 0).
• End-to-end GPU validation of the _extra_state fix on the QAD path is pending (needs the quantize→train run in the TRT-LLM container) — flagged as a draft until then.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ (only drops requests for keys absent from the checkpoint)
• If you copied code from any other sources or added a new PIP dependency: N/A
• Did you write any new necessary tests?: ❌ — the plugin has no unit tests yet (# TODO: Add unit tests for this plugin); the fix is exercised by the new QAD example / GPU e2e.
• Did you update Changelog?: ❌ (can add if desired)
• Did you get Claude approval on this PR?: N/A (draft)

Additional Information

Draft — pending GPU e2e validation. Surfaced while triaging the nmm-sandbox Megatron QAD CI test (Llama-3.2-1B-Instruct).

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 6901c2ea-7bcb-43bb-8f89-5ce84914b355

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch chenhany/megatron-qad-extra-state

---

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-1931/
Built to branch gh-pages at 2026-07-06 18:40 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.62%. Comparing base (795c589) to head (d411508).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
...lopt/torch/opt/plugins/mcore_dist_checkpointing.py	0.00%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1931      +/-   ##
==========================================
- Coverage   77.62%   77.62%   -0.01%     
==========================================
  Files         515      515              
  Lines       57301    57304       +3     
==========================================
  Hits        44482    44482              
- Misses      12819    12822       +3     

Flag	Coverage Δ
unit	54.94% <0.00%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.