fix(ckpt): reduce distributed optimizer resume memory

[cuichenx]

What does this PR do ?

Prevents a transient CUDA OOM while resuming standard FP32 distributed-Adam training checkpoints by reusing Megatron-Core's already-correct optimizer-state scaffold instead of asking Transformer Engine FusedAdam to allocate a second copy.

The defect is pre-existing shared checkpoint behavior, not a regression introduced by Nemotron Omni PR #4776. That PR's near-full Omni configuration made the duplicate-state peak reproducible.

Changelog

• Add a checkpoint-load-scoped fast path for standard Transformer Engine FusedAdam instances with FP32 parameters and FP32 exp_avg/exp_avg_sq state.
• Retain Transformer Engine's loader for precision-aware, master-weight, remainder, CPU-offload/HDO, MCore FSDP, stub, non-TE, mixed-precision, and otherwise incompatible optimizer states.
• Restore each optimizer's original loader on success, load failure, or partial chained-optimizer setup failure and release superseded allocator cache only after the checkpoint-load frame exits.
• Cover state-tensor reuse, fallback modes, chained dense/expert optimizers, cleanup, optional TE absence, LayerWise exclusion, and exception handling with focused tests.

Root cause and measured result

During distributed checkpoint resume, MCore creates FP32 Adam moment tensors for the sharded-state scaffold. TE FusedAdam's loader allocates optimizer-owned tensors and copies those same moments, so both sets coexist before checkpoint payload I/O. In the reproducer, the unfixed process OOMed there at a 80,997 MiB whole-device peak.

The fixed run loaded the same complete 415.69 GB model/optimizer/RNG checkpoint and completed the first resumed optimizer step:

Measurement	Unfixed resume	Fixed resume	Uninterrupted control
Checkpoint-load peak (nvidia-smi, 200 ms)	80,997 MiB / OOM	69,707 MiB	n/a
Training peak (nvidia-smi, 200 ms)	not reached	79,651 MiB	79,673 MiB
Rank-0 allocator max	79.43 GB at OOM	69.48 GB through resumed step	69.71 GB through control

This shows the change removes the transient restore duplicate without retaining extra state into training.

Validation

• DFW Slurm 14115430: 8xH100, TP8/EP8, 48/52 trainable language layers, seq4096, full native torch_dist optimizer/RNG resume; checkpoint load and resumed step 2 completed (0:0).
• DFW Slurm 14115632: 22 focused tests passed; genuine TE CUDA smoke verified incoming moment data_ptr reuse, no measurable replacement allocation, value correctness, and loader restoration; targeted strict mypy passed (0:0).
• DFW Slurm 14116185: post-review validation; 25 focused tests passed, genuine TE CUDA smoke passed, and targeted strict mypy passed (0:0).
• uv run --active --no-sync pre-commit run --all-files
• git diff --check origin/main...HEAD

Runtime image: final NeMo 26.06, manifest sha256:64fcec59b0eeee2853761d16767c603e03e0aa4ba03becc9a7793bb0c46545e7, squashfs SHA-256 0dbe69a640877451d19a576628295082cc4892f405db290f4d425588f629f41a.

GitHub Actions CI

This is a shared checkpoint-load path, so L1 checkpoint/training coverage is recommended. No broad CI tier was manually triggered while the PR is draft.

Before your PR is "Ready for review"

Pre checks:

• Read and followed the contributor guidelines.
• Added focused regression coverage.
• Documentation is not required: there is no new API, config, checkpoint format, or user workflow.
• Transformer Engine availability is guarded; unsupported optimizer modes retain their existing loader.
  • Reviewer: Does the PR have correct import guards for all optional libraries?

Additional Information

• Reproduced from the Nemotron Omni training configuration validated in [data] fix: isolate and unify Nemotron Omni collation #4776, but the failing optimizer restore chain predates that PR.
• Full training-state resume is covered. Model-weight-only load was already successful and is deliberately excluded from the new fast path.
• Current [data] fix: isolate and unify Nemotron Omni collation #4776 Energon code has no dataloader-state restore integration, so no resumed Energon trajectory-parity claim is made.

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

[claude]

Light review — LGTM

The memory_efficient_fp32_optimizer_state_loading context manager is well-guarded and fallback-safe. Verified points:

• Fallback safety is correct. If the incoming state dict is not an exact FP32 Adam scaffold ({exp_avg, exp_avg_sq} only, all float32), _load_state_dict_without_fp32_reallocation delegates to the TE original loader. Worst case is no speedup, never incorrect optimizer state. The eligibility gates (master_weights, store_param_remainders, use_precision_aware_optimizer, optimizer_cpu_offload, use_megatron_fsdp, name_to_dtype_map, param dtype) all bail conservatively.
• No late-binding closure bug. original_load_state_dict is rebound each loop iteration and captured per-instance via the _fallback default arg, so each patched optimizer restores/falls back to its own method.
• Restore path is complete. The finally block distinguishes a pre-existing instance override (setattr restore) from no override (delattr), and torch.cuda.empty_cache() only runs when something was patched. Rollback on mid-setup exception is covered by tests.
• Call-site gating is sound. checkpoint_optimizer = optimizer if cfg.checkpoint.load_optim and not cfg.checkpoint.finetune else None correctly disables the fast path when optimizer state is not being loaded (PEFT/finetune), yielding a no-op.

Minor observation (non-blocking): _optimizer_state_is_fp32_adam requires the per-param state keys to be exactly {exp_avg, exp_avg_sq}. If a checkpoint ever carries a per-param step tensor, the fast path silently disengages and falls back to the TE loader — correct, but a short comment noting the exact-match is intentional would help future readers.

Test coverage is thorough: dtype/master-weight/state-dtype variants, incompatible distributed modes, chained optimizers, LayerWise children, TE-unavailable no-op, and rollback on exceptions during both setup and loading.

Suggested test cases

• No perf tests impacted (change is confined to training/optim.py, training/setup.py, and their unit tests; no scripts/performance/configs/ entries touched).
• tests/unit_tests/training/test_optim.py::TestMemoryEfficientFp32OptimizerStateLoading (all cases in this class).