PoR: align Megatron-LM and Megatron-Bridge checkpoints #4595

Summary

Track checkpoint alignment between Megatron-LM (MLM) and Megatron-Bridge. The goal is for each project to load and resume from checkpoints produced by the other without unifying their training loops.

Both projects already use the torch distributed checkpoint format for sharded model weights. The remaining work is to align configuration metadata, mutable training state, tokenizer assets, and compatibility behavior.

Active scope

• Save Bridge-compatible run_config.yaml metadata from MLM checkpoints.
• Translate checkpoint configuration between Bridge ConfigContainer fields and MLM argparse.Namespace fields, including explicit handling for differently named settings such as RoPE and YARN options.
• Align train_state.pt contents for iteration, consumed samples, and other mutable resume state.
• Define the required tokenizer assets and loading behavior.
• Preserve legacy MLM checkpoint loading.
• Add bidirectional interoperability tests:
  • MLM save -> Bridge load and resume.
  • Bridge save -> MLM load and resume.
  • Legacy MLM checkpoint fallback.

Acceptance criteria

• Model weights load successfully in both directions.
• Configuration and mutable training state are restored without silent defaults or dropped fields.
• A short resumed run advances from the expected iteration and consumed-sample state in both frameworks.
• Legacy MLM checkpoints remain supported.

Out of scope

• Unifying the MLM and Bridge training loops.
• Replacing both implementations with shared save_checkpoint() or load_checkpoint() APIs.
• Broad removal of MLM global args or training-loop global state.

Background

Earlier work upstreamed shared builders and removed corresponding Bridge duplication: NVIDIA/Megatron-LM#4550, NVIDIA/Megatron-LM#4741, and #3988. Remaining general training refactors are tracked separately and are not the active focus of this issue.

Tracking

• Upstream parent: Align Megatron-LM and Megatron-Bridge checkpoint metadata NVIDIA/Megatron-LM#2319
• Project board: https://github.com/orgs/NVIDIA/projects/146/views/1
• 26.08 PoR tracker: https://linear.app/nvidia/issue/MB-585/megatron-bridge-2608-release-por-tracker
• PIC: @maanug-nv