Skip to content

[feat] Add Qwen3 MoE true-on-policy SP/PP support - #26317

Open
maocheng23 wants to merge 47 commits into
sgl-project:sglang-milesfrom
maocheng23:feat/true_on_policy_qwen_moe_sppp
Open

[feat] Add Qwen3 MoE true-on-policy SP/PP support#26317
maocheng23 wants to merge 47 commits into
sgl-project:sglang-milesfrom
maocheng23:feat/true_on_policy_qwen_moe_sppp

Conversation

@maocheng23

@maocheng23 maocheng23 commented May 25, 2026

Copy link
Copy Markdown
Collaborator

Based on maocheng23#4 and retargeted to sgl-project/sglang:sglang-miles.

Summary

Adds the SGLang rollout-side contract and deterministic kernel plumbing for Qwen3-30B-A3B MoE true-on-policy training, including the SP+PP contract update used by the coupled Miles and Megatron-LM PRs.

This PR is one of three coupled PRs that should be reviewed and landed together because they extend the qwen3_moe_true_on_policy_v1 contract to allow Megatron sequence parallel and pipeline parallel training while preserving exact SGLang rollout parity.

Companion PRs:

Main Changes

  • Add Qwen3 MoE true-on-policy runtime policy/schema wiring under python/sglang/srt/true_on_policy/.
  • Add batch-invariant and TP-invariant ops used by deterministic rollout/recompute parity paths.
  • Route Qwen3 MoE, layernorm, communicator, model runner, CUDA graph, and TP worker paths through the true-on-policy runtime policy where needed.
  • Disable NVLS on the deterministic TP decode path so TP reductions match the Megatron parity contract.
  • Remove earlier debug/weight-audit and legacy true-on-policy target paths from the final branch.
  • Flip the shared schema flag so Megatron sequence parallel is allowed for the dense and MoE true-on-policy schemas.

Validation

Cleaned 8-GPU ion7 real-workload run:

  • Run id: moe_sppp_pr_tp2_pp2_ep2_onpolicy_real3_cleaned_noci_nosave_260524_ion7
  • Ray job: raysubmit_a9NCtiu6K5dRVCsc
  • Topology: Megatron TP=2, PP=2, EP=2, ETP=1, sequence parallel enabled; SGLang 4 rollout engines, each TP=2, EP=2.
  • Result: Ray job succeeded, all 8 GPUs returned to 0 MiB.
  • Step 1: train/train_rollout_logprob_abs_diff=0.0, train/train_rollout_kl=0.0, grad_norm=0.0377418305, weight version 1.0, mixed version 0.0.
  • Step 2: train/train_rollout_logprob_abs_diff=0.0, train/train_rollout_kl=0.0, grad_norm=0.0389085777, weight version 2.0, mixed version 0.0.
  • Timing: cleaned step times 485.2571s and 462.8948s, about 2.6-2.9% slower than the previous on-policy no-debug timing run and in the same overhead band versus off-policy.

Focused checks after cleanup:

  • git diff --check
  • Miles fast checks: 53 passed
  • Megatron targeted extension/MoE checks: 3 passed
  • Megatron 8-rank tensor-parallel mapping check passed on all ranks
  • Full SP+PP 8-GPU E2E exact-logprob rerun on ion7

Local record:
recovery/qwen3_moe_sppp_clean/journal/2026-05-23-sppp-e2e.md


CI States

Latest PR Test (Base): ❌ Missing run-ci label -- add it to run CI tests.
Latest PR Test (Extra): ❌ Blocked -- run-ci is required first.

hzh0425 and others added 30 commits May 15, 2026 14:30
Co-authored-by: ybyang <10629930+whybeyoung@users.noreply.github.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>
Co-authored-by: b8zhong <b8zhong@users.noreply.github.com>
…g kv_args (sgl-project#25380)

Co-authored-by: JoeLee314 <liqichao@baidu.com>
…gl-project#24704)

Co-authored-by: Shangming Cai <csmthu@gmail.com>
Co-authored-by: xuyongfei <xuyongfei.xyf@antgroup.com>
layer_mapping was changed from a compact list to a sparse array
containing None for out-of-stage layers by PP support (sgl-project#24704),
but the HiCache consumer was not updated to handle None entries.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

deterministic Issues on deterministic inference/kernels documentation Improvements or additions to documentation npu

Projects

None yet

Development

Successfully merging this pull request may close these issues.