Skip to content

[feat] Add SP and PP kernels for qwen_moe true on policy - #4

Draft
maocheng23 wants to merge 1 commit into
feat/true_on_policy_qwen_moefrom
feat/true_on_policy_qwen_moe_sppp
Draft

[feat] Add SP and PP kernels for qwen_moe true on policy#4
maocheng23 wants to merge 1 commit into
feat/true_on_policy_qwen_moefrom
feat/true_on_policy_qwen_moe_sppp

Conversation

@maocheng23

@maocheng23 maocheng23 commented May 7, 2026

Copy link
Copy Markdown
Owner

Stacked on top of #3.

Summary

Updates the SGLang contract metadata for the Qwen3-30B-A3B MoE SP+PP true-on-policy stack.

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 SGLang rollout parity.

Companion PRs:

Main Changes

  • Flip disable_megatron_sequence_parallel from True to False for the dense and MoE true-on-policy schemas.
  • Keep SGLang runtime behavior unchanged; the deterministic TP/EP/MoE rollout policy is already provided by the MoE base PR, and this PR only updates the shared contract so Megatron SP is allowed.

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

@maocheng23
maocheng23 force-pushed the feat/true_on_policy_qwen_moe branch 3 times, most recently from 6f7b617 to e22c682 Compare May 23, 2026 02:47
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant