Skip to content

Handle step key correctly in dp_reshardable checkpoint save with --optimizer-cpu-offload - #69

Open
artkorenev wants to merge 27 commits into
radixark:miles-mainfrom
artkorenev:graft/distopt-2874-step
Open

Handle step key correctly in dp_reshardable checkpoint save with --optimizer-cpu-offload#69
artkorenev wants to merge 27 commits into
radixark:miles-mainfrom
artkorenev:graft/distopt-2874-step

Conversation

@artkorenev

@artkorenev artkorenev commented Jul 16, 2026

Copy link
Copy Markdown

Grafts NVIDIA#2874 (f4502eb1c9).

Wraps the optimizer step in LocalNonpersistentObject in
sharded_param_state_dp_reshardable so it is not persisted into the rank-0
common state.

🤖 Generated with Claude Code

yueming-yuan and others added 27 commits February 25, 2026 19:22
@yueming-yuan @claude @guapisolo
[1/8] fix: misc compatibility fixes for PyTorch and TE (radixark#2)
b4d401c 
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
@yueming-yuan @claude @guapisolo
[2/8] feat: support partial checkpoint loading (radixark#3)
6c3fb8c 
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
@yueming-yuan @guapisolo
[3/8] feat: add post-attention and post-MLP layernorm support (radixa…
3f0fcac 
…rk#4)

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Co-authored-by: Yueming Yuan <yym022502@gmail.com>
@yueming-yuan @claude @guapisolo
[4/8] fix: MLA RoPE triton kernel head indexing and v_dim=0 support (r…
e82d9ee 
…adixark#5)

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
@yueming-yuan @claude @guapisolo
[5/8] feat: support MTP training in RL (radixark#6)
3011052 
- Detach output layer params to prevent MTP gradient flowing to output layer
- Add mtp_kwargs interface for flexible MTP label/loss_mask passing
- Roll mtp_labels and loss_mask for RL training compatibility

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: Yueming Yuan <yym022502@gmail.com>
@guapisolo @yueming-yuan
[6/8] feat: support rollout routing replay (R3) and bypass for MTP la…
228a364 
…yers (radixark#10)

- Add is_mtp flag to MoE layers and multi_token_prediction module
- Bypass routing replay for MTP layers (MTP uses fresh routing)
- Replace rdxa/dev's built-in RouterReplay with miles.utils.routing_replay:
  - moe_utils.py: use get_routing_replay_compute_topk() wrapper
  - router.py: use register_routing_replay() for initialization

Co-authored-by: Yueming Yuan <yym022502@gmail.com>
@yueming-yuan @claude @guapisolo
[7/8] feat: add INT4 fake QAT for MoE grouped linear (radixark#9)
a95ad14 
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: Yueming Yuan <yym022502@gmail.com>
@guapisolo @claude
[8/8] fix: CUDA IPC incompatibility from Megatron bump (radixark#11)
57784b8 
After bumping Megatron (rdxa/dev), colocated IPC weight update fails with
torch.AcceleratorError: CUDA error: invalid argument during
torch.multiprocessing serialization of CUDA tensors.

Root cause: Megatron's new TMS hook (PR NVIDIA#3048) alters allocator behavior
in training flow, causing allocations via cuMemCreate/cuMemMap which are
incompatible with CUDA IPC (_share_cuda_() fails).

Fix: resolve mapping.py and dynamic_context.py conflicts to isolate
hook side effects so TMS/allocator state remains IPC-compatible during
the weight update phase.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@guapisolo @claude
fix: dp_reshardable checkpoint backward compat in Megatron core
07220d2 
- merge(): truncate dp_reshardable padding on optimizer/param_state path
- load_parameter_state_from_dp_reshardable: tolerate missing 'padding' key
- ShardedTensor: relax flattened_range to deprecation warning

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@guapisolo @yueming-yuan
Upgrade Megatron from Dec 17 to Feb 13
038e8e5 
This PR rebases from radixark Megatron fork [miles-20260218](https://github.com/radixark/Megatron-LM/tree/miles-20260218) and resolve conflicts.

Upgrade Megatron from Dec 17 (3714d81) to Feb 13 (1dcf0da)

PR link: radixark#13

Co-authored-by: Yueming Yuan <yym022502@gmail.com>
Made-with: Cursor
@yueming-yuan
enable tms no cpu backup region for grad/optimizer's param buffers (r…
74f0b38 
…adixark#17)
@yueming-yuan @claude
add disable_param_buffers_cpu_backup method to keep cpu backup when u…
5243579 
…se `--disable-weight-backuper` in miles (radixark#18)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
@yueming-yuan @fzyzcjy
[fix] fix mixed precision parameters load from checkpoint (radixark#19)
cf014f9 
Co-authored-by: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com>
@yueming-yuan @fzyzcjy
[feat] distributed optimizer low GPU memory resume from checkpoint (r…
90f69fa 
…adixark#20)

Co-authored-by: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com>
@guapisolo
[fix] Add detach() to fp32 param shard for leaf-tensor consistency
32dbe9f
@Zhichenzzz
[fix] attention_output_gate TP slice when num_kv_heads < TP (radixark#22
923b85b 
)
@guapisolo
[fix] Enforce param dtype before wrap ddp (radixark#24)
d9a5080
@maocheng23 @zju-stu-lizheng
@zyxiyy02 @yizhang2077 @claude
[feat] Init true on policy with qwen_dense
23924a0 
Squash merge of the dense true-on-policy Megatron branch.

Co-authored-by: zju-stu-lizheng <lizheng.cs@zju.edu.cn>
Co-authored-by: zyxiyy02 <282300612+zyxiyy02@users.noreply.github.com>
Co-authored-by: Yi Zhang <1109276519@qq.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@maocheng23
Package miles_megatron_plugins with Megatron Core (radixark#49)
3a60706
@Shi-Dong
Broaden megatron-core packages.find to megatron* (namespaces), keep p…
a381c7d 
…lugins (radixark#50)
@yueming-yuan @fzyzcjy
DeepSeek V4 RL support (radixark#28)
ae20f41 
Co-authored-by: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com>
@yueming-yuan
Skip missing expert bias updates (radixark#55)
36ceb4f
@XinyuJiangCMU @JessicaJiang-123
[AMD] Fix QuantizedTensor import for ROCm TransformerEngine 2.8 (radi…
3a8f0c7 
…xark#58)

Co-authored-by: Zhiyao Jiang <jessicajiang324@gmail.com>
@Zhichenzzz
fix(dist-ckpt): handle bare BytesIO _extra_state in key recovery (TE …
87d1155 
…FP8) (radixark#53)
@guapisolo @zyzshishui @claude
Restore enforce_marked_param_dtypes after model materialization (radi…
79fc089 
…xark#60)

Co-authored-by: zyzshishui <82826991+zyzshishui@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@fzyzcjy
Wire witness support inside Megatron (radixark#21)
4716f75
@artkorenev @claude
Handle step key correctly in dp_reshardable checkpoint save with `-…
ef1b586 
…-optimizer-cpu-offload`

Graft of NVIDIA#2874 (f4502eb): wrap the optimizer `step` in
LocalNonpersistentObject in sharded_param_state_dp_reshardable.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@artkorenev
artkorenev marked this pull request as ready for review July 16, 2026 12:14
@yushengsu-thu

yushengsu-thu commented Jul 17, 2026 •
edited
Loading

Copy link
Copy Markdown

The graft matches upstream exactly, but I traced the load path with --optimizer-cpu-offload and hit two problems:

  1. Old checkpoints won't load anymore. Pre-PR saves leaked the raw step into common.pt (since [2/8] feat: support partial checkpoint loading #3), and the load request now puts a LocalNonpersistentObject at the same path, so merge() in serialization.py:126 throws on the duplicate leaf. No strictness flag gets around that.

  2. step gets reset to 1 on resume. The load skeleton is built after HDO's dummy_step(), LNO keeps that local value, and _set_main_param_and_optimizer_states then copies it over the step that load_state_dict just restored from param_groups (:863-873 vs :913). So Adam bias correction restarts on every resume. The old leak was actually what made resume work before. Upstream has the same bug, the code is identical there.

For 2 the easy fix is skipping step in _set_main_param_and_optimizer_states, same as load_parameter_state_from_fs_model_space already does (:1852). For 1, drop the key from the request when is_loading instead of wrapping it. If nobody has live cpu-offload + dp_reshardable checkpoints then 1 is just a release note, but I'd fix 2 before this lands since it breaks the exact resume flow the PR is for.

(Only cpu-offload / torch AdamW runs are affected — FusedAdam keeps step in param_groups, which is also why no existing test catches this.)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

9 participants

@artkorenev @yushengsu-thu @yueming-yuan @guapisolo @Zhichenzzz @maocheng23 @Shi-Dong @XinyuJiangCMU @fzyzcjy