Skip to content

[feat] Add Qwen3 MoE true-on-policy parity - #30

Open
maocheng23 wants to merge 40 commits into
miles-mainfrom
feat/true_on_policy_qwen_moe
Open

[feat] Add Qwen3 MoE true-on-policy parity#30
maocheng23 wants to merge 40 commits into
miles-mainfrom
feat/true_on_policy_qwen_moe

Conversation

@maocheng23

@maocheng23 maocheng23 commented May 1, 2026 •
edited
Loading

Copy link
Copy Markdown

Summary

Adds the Megatron-LM side of the Qwen3-30B-A3B MoE true-on-policy stack.

This PR is one of three coupled PRs that must land together because they share the qwen3_moe_true_on_policy_v1 contract and were validated as one end-to-end stack.

Companion PRs:

Target

Bit-identical rollout/train logprob parity for Qwen3-30B-A3B MoE under deterministic decode, with differentiable Megatron backward.

SGLang remains the numerical source of truth. Megatron reproduces the SGLang MoE numerics in training, including deterministic routing/top-k and EP reduction behavior required by the contract.

Main Changes

  • Adds the Qwen3-MoE true-on-policy extension path at the MoE layer boundary.
  • Uses SGLang fused MoE behavior for the direct parity path and local autograd support for training.
  • Updates the SGLang fused MoE imports to the rebased moe_runner.triton_utils layout.
  • Keeps the MoE-specific parity path narrow; generic Megatron router/dispatcher fallback edits and debug-only helpers were cleaned out.
  • Preserves dense-stack behavior outside the explicit qwen3_moe_true_on_policy_v1 contract.

Validation

Latest rebased EP4 E2E on ion7, Qwen3-30B-A3B, 8x H200, run qwen3-moe-top-ep4-rebased-e2e-260522-ion7-r4:

  • Ray job: raysubmit_yc4WK8aCKzP4tx1Q
  • Train: TP=1, EP=4, ETP=1, CP=2
  • Rollout: rollout_num_gpus=8, rollout_num_gpus_per_engine=4, SGLang TP=4, SGLang EP=4
  • Runtime confirmed: recompute_logprobs_via_prefill=False, use_rollout_logprobs=False
  • rollout logprob: -2.8848648071289062e-05
  • ref logprob: -2.8848648071289062e-05
  • train logprob: -2.8848648071289062e-05
  • train/train_rollout_logprob_abs_diff = 0.0
  • train/train_rollout_kl = 0.0

Local record:
recovery/qwen3_moe_clean/journal/2026-05-22-qwen3-moe-clean-e2e.md

Test Plan

  • Rebased onto current radixark/Megatron-LM:miles-main stack.
  • Local syntax/diff sanity checks for touched Python paths.
  • Remote import/syntax checks in the ion7 container.
  • EP4 E2E exact-zero logprob validation on ion7.
  • CI replay for the three-PR stack.
  • Longer on/off-policy throughput comparison after reviewer signoff on the cleaned deterministic path.

yueming-yuan and others added 17 commits February 25, 2026 19:22
@yueming-yuan @claude @guapisolo
[1/8] fix: misc compatibility fixes for PyTorch and TE (#2)
b4d401c 
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
@yueming-yuan @claude @guapisolo
[2/8] feat: support partial checkpoint loading (#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 (#4)
3f0fcac 
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 (#5)
e82d9ee 
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
@yueming-yuan @claude @guapisolo
[5/8] feat: support MTP training in RL (#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 (#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 (#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 (#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: #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 (#17)
74f0b38
@yueming-yuan @claude
add disable_param_buffers_cpu_backup method to keep cpu backup when u…
5243579 
…se `--disable-weight-backuper` in miles (#18)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
@yueming-yuan @fzyzcjy
[fix] fix mixed precision parameters load from checkpoint (#19)
cf014f9 
Co-authored-by: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com>
@yueming-yuan @fzyzcjy
[feat] distributed optimizer low GPU memory resume from checkpoint (#20)
90f69fa 
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 (#22)
923b85b
@guapisolo
[fix] Enforce param dtype before wrap ddp (#24)
d9a5080
This was referenced May 1, 2026
Open
Closed
@maocheng23
maocheng23 force-pushed the feat/true_on_policy_qwen_moe branch from c63c77f to 9c390a1 Compare May 5, 2026 01:51
@maocheng23 maocheng23 mentioned this pull request May 7, 2026
Open
5 tasks
@maocheng23
maocheng23 force-pushed the feat/true_on_policy_qwen_dense branch 2 times, most recently from 9546575 to 57258c8 Compare May 18, 2026 06:15
@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
maocheng23 force-pushed the feat/true_on_policy_qwen_moe branch 3 times, most recently from ab103f8 to 9c6c2a9 Compare May 19, 2026 18:44
@maocheng23 maocheng23 changed the title [feat] Init true on policy with qwen_moe [feat] Add Qwen3 MoE true-on-policy parity May 21, 2026
maocheng23 and others added 2 commits May 22, 2026 19:02
@maocheng23 @zju-stu-lizheng
@zyxiyy02 @yizhang2077 @claude
Add Qwen3 MoE true-on-policy Megatron parity
f7a520a 
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 @zju-stu-lizheng
@zyxiyy02 @yizhang2077 @claude
Fix MoE true-on-policy weight refresh path
9fe48e6 
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 and others added 8 commits May 22, 2026 19:02
@maocheng23
Align MoE contract gate with direct SGLang path
a143c75
@maocheng23
Remove true-on-policy debug dump plumbing
dcf88ac
@maocheng23
Keep GRPO loss unchanged for MoE parity
5954723
@maocheng23
Remove batch-invariant escape hatch
1ef3a0e
@maocheng23
Share SGLang MoE forward path
ae93354
@maocheng23 @cursoragent
Simplify MoE: one forward, remove defensive bloat
248d04b 
- forward.py: single canonical sglang_moe_forward (102 lines)
- autograd.py: pure backward wrapper calling shared forward (182 lines)
- moe_experts.py: weight-layout adapter only, no forward logic (36 lines)
- moe_layer_ext.py: one linear orchestration path (301 lines)
- Remove verbose RuntimeError guards, replace with asserts
- Remove weight caching (premature optimization)
- Consolidate two parallel forward paths into one

Net: -574 lines, structurally impossible for forward paths to diverge.
Co-authored-by: Cursor <cursoragent@cursor.com>
@maocheng23
Fix SGLang MoE weight adapter init
1e562d7
@maocheng23
Update SGLang fused MoE imports
b0679e2
@maocheng23
maocheng23 force-pushed the feat/true_on_policy_qwen_moe branch from 7ac2a80 to b0679e2 Compare May 23, 2026 02:47
@maocheng23 maocheng23 mentioned this pull request May 23, 2026
Open
6 tasks
@maocheng23
maocheng23 marked this pull request as ready for review May 23, 2026 21:22
@maocheng23
maocheng23 changed the base branch from feat/true_on_policy_qwen_dense to miles-main May 23, 2026 22:20
Zhichenzzz
Zhichenzzz approved these changes May 25, 2026
View reviewed changes

@Zhichenzzz Zhichenzzz left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great work, LGTM! Just a small possible change for me, Thanks @maocheng23 !

Comment thread miles_megatron_plugins/true_on_policy/moe_layer_ext.py
torch.distributed.all_gather(gathered, padded, group=ep_group)
return torch.cat(gathered, dim=0)

@staticmethod

@Zhichenzzz Zhichenzzz May 25, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The list-of-tensors form forces torch to assemble a single contiguous input buffer before handing off to NCCL. The input here is already a single contiguous tensor grad_output.contiguous(), could we pass it directly to reduce_scatter_tensor to reduce comm overhead?

  torch.distributed.reduce_scatter_tensor(
      grad_local,
      grad_output,
      group=ctx.ep_group,
  )

This was referenced May 26, 2026
Open
Open
Open
Open
Open
Open
Open
Open
Open
Open
Open
Open
Open
Open
Open
@zianglih zianglih mentioned this pull request Aug 6, 2026
Draft
7 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@Zhichenzzz Zhichenzzz Zhichenzzz approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@maocheng23 @Zhichenzzz @yueming-yuan @guapisolo