fix(checkpoint): use global FSDP v1 pipeline layer keys

[cuichenx]

What does this PR do?

Prevent FSDP DTensor checkpoint collisions between pipeline stages by applying Megatron-Core's global-layer naming handler on both save and load. Without it, stages can write different model weights and Adam moments under identical stage-local keys; a successful load can then silently resume from incorrect state. This affects dense GPT as well as hybrid Nemotron models.

Depends on NVIDIA/Megatron-LM#8159, which contains the shared checkpoint handler, FSDP v1 PP mesh fix, and shared-embedding gradient reduction fixes. That Core PR depends on NVIDIA/Megatron-LM#8010 for meta-device embedding initialization. Keep this draft until that dependency is available.

Changelog

• Apply the global-layer handler after transformations that resolve live modules by local name, preserving tensor aliases used by in-place loading.
• Preserve v2 preprocessing; its optimizer keys are already global.
• Retain compatibility with older Core for PP=1; raise a clear error for PP>1 when the handler is unavailable.
• Test handler invocation and both old-Core compatibility cases.

PP=1 checkpoint keys remain unchanged. Checkpoints already affected by stage-local collisions have omitted stage-specific tensors and need to be regenerated; this change cannot reconstruct them.

Validation

• Six focused Bridge checkpoint tests passed with the updated Core source.
• All-files pre-commit passed. Strict mypy was run but is not clean on the baseline (missing local dependencies and existing typing errors); no clean type-check claim.
• Revised Core #8159 gradient regression: 96 cases passed on each of eight H100 ranks, in addition to 26 unchanged mesh and 16 checkpoint cases.
• Lightning PP=2/EP=2/MTP=2 with #8010 (5920bd95) plus Core #8159: checkpoint weights and Adam state restore exactly, next loss/gradient norm match exactly, and all 287 post-step model/Adam tensors match numerically (maximum absolute difference 4.331e-8).
• Full 52-layer Lightning with #8010 plus Core #8159: three forward/backward/optimizer steps completed on eight H100s, PP=2/EP=2/TP=1/CP=1, original widths/128 experts/vocabulary, MTP=2 and synthetic sequence length 256. Loss 12.176648 → 10.490705 → 9.382988, no skipped/NaN iterations; peak allocated memory 69.898 GiB. BF16 Adam moments fit this full-size smoke on one node; checkpoint restart was validated with the smaller FP32-moment run above.

Training probes use function-level patch injection into NeMo 26.08.01 and include #8010 plus the existing upstream embedding-gradient reduction fix from NVIDIA/Megatron-LM#6486.

The training probes set model.deallocate_pipeline_outputs=False, required for the output views returned by FSDP v1 hooks. Short synthetic-data runs establish correctness for the tested configurations, not convergence or performance.

Independent subagent review completed before publication; no blocking findings.

Before this PR is ready for review

• Followed contributor guidelines; signed and signed-off commit.
• Added focused regression coverage.
• Documented dependency and checkpoint compatibility above.
• Guarded the optional Core handler import.
• Core dependency merged and available in Bridge's Core version.

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

[github-actions]

Automatic Claude reviews have been retired. To request a pull-request review, post a comment containing:

/review

Add model=claude to use a Claude reviewer (the default is model=codex). mode=light|strict selects the review depth; for example, /review model=claude mode=strict. Comment /review help for all options.

[github-actions]

🌿 Preview your docs: https://nvidia-preview-pr-6424-d608c62f907b.docs.buildwithfern.com/nemo/megatron-bridge