Skip to content

[sglang-miles] True on-policy training support for FSDP2 - #18639

Merged
yueming-yuan merged 5 commits into
sglang-milesfrom
sglang-miles-trueon
Feb 12, 2026
Merged

[sglang-miles] True on-policy training support for FSDP2#18639
yueming-yuan merged 5 commits into
sglang-milesfrom
sglang-miles-trueon

Conversation

@yueming-yuan

@yueming-yuan yueming-yuan commented Feb 11, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Summary

Cherry-pick true on-policy training changes from sglang-slime-trueon to sglang-miles, enabling deterministic inference that exactly matches the FSDP2 training forward pass.

Changes across 11 files in 3 commits:

Core layer changes:

  • layernorm.py: Change fp32_residual default to True, remove weight_dtype/override_orig_dtype params, add two-branch forward_native (bf16 vs fp32 residual paths), add post_residual_addition to x in allreduce fusion path
  • communicator.py: Thread explicit post_residual_addition parameter through prepare_attn and prepare_attn_and_communicate
  • logits_processor.py: Remove special bf16 matmul path for on-policy tied-weight lm_head
  • rotary_embedding.py: Remove torch.compile(dynamic=True) wrapper when rl_on_policy_target is set
  • sampler.py: Compute log_softmax after standard temperature division (instead of separate bf16 div path)
  • fused_moe.py: Gate torch.compile MoE sum reduce with enable_deterministic_inference check

Model changes:

  • qwen3.py: Use fp32_residual=False in on-policy norm_kwargs, explicit post_residual_addition in decoder forward
  • qwen2.py: Remove bf16 cast in MLP, remove fp32 embed params_dtype, use fp32_residual=False norm_kwargs
  • qwen2_moe.py: Add on-policy norm_kwargs to final norm
  • qwen3_moe.py: Manual softmax→topk→renormalize routing (bypass fused topk), disable fused_kv_buffer and fused_qk_norm_rope for on-policy
  • qwen3_vl.py: Rewrite fast_pos_embed_interpolate with pure-torch manual bilinear interpolation (torch.compile compatible), inline deepstack computation

Infrastructure:

  • parallel_state.py: Add try/except safety for get_tensor_model_parallel_rank() when TP group not initialized

Origin

yueming-yuan and others added 4 commits February 10, 2026 13:42
@yueming-yuan
int4
e957263
@yueming-yuan
true on policy
5e06f4e
@yueming-yuan
true on policy 2: layernorm, communicator
e4f2066
@yueming-yuan @claude
true on policy 3: parallel_state, qwen3 decoder, qwen3_vl VLM
602b62c 
- parallel_state.py: add try/except to get_tensor_model_parallel_rank()
  (returns 0 on exception, needed for on-policy weight loading)
- qwen3.py: Qwen3DecoderLayer.forward uses explicit post_residual_addition
  parameter instead of **kwargs for type safety
- qwen3_vl.py: rewrite fast_pos_embed_interpolate to use manual bilinear
  interpolation (torch.compile compatible), inline deepstack_embeds
  computation for VLM true on-policy support

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@gemini-code-assist

gemini-code-assist Bot commented Feb 11, 2026

Copy link
Copy Markdown
Contributor

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

@yueming-yuan yueming-yuan changed the title True on-policy training support for FSDP2 [sglang-miles] True on-policy training support for FSDP2 Feb 11, 2026
@yueming-yuan
yueming-yuan merged commit e236af5 into sglang-miles Feb 12, 2026
1 check passed
@yueming-yuan
yueming-yuan deleted the sglang-miles-trueon branch February 12, 2026 21:14
yueming-yuan added a commit that referenced this pull request Feb 13, 2026
@yueming-yuan @claude
[sglang-miles] True on-policy training support for FSDP2 (#18639)
e48416c 
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
zianglih pushed a commit to zianglih/sglang that referenced this pull request Feb 28, 2026
@yueming-yuan @claude Ziang Li
[sglang-miles] True on-policy training support for FSDP2 (sgl-project…
42a9887 
…#18639)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
yueming-yuan added a commit that referenced this pull request Mar 5, 2026
@yueming-yuan
[1/8] [sglang-miles] True on-policy training support for FSDP2 (#18639)
0257c73
yueming-yuan added a commit that referenced this pull request Mar 27, 2026
@yueming-yuan
[1/8] [sglang-miles] True on-policy training support for FSDP2 (#18639)
40cef40
yueming-yuan added a commit that referenced this pull request Apr 6, 2026
@yueming-yuan
[1/8] [sglang-miles] True on-policy training support for FSDP2 (#18639)
74c9eab
@zyzshishui zyzshishui mentioned this pull request Apr 23, 2026
Draft
5 tasks
zyzshishui pushed a commit to zyzshishui/sglang that referenced this pull request Apr 24, 2026
@yueming-yuan @zhuzilin
@nanjiangwill @zyzshishui
[1/8] [sglang-miles] True on-policy training support for FSDP2 (sgl-p…
7b38b25 
…roject#18639)

(cherry picked from commit 74c9eab)

Co-authored-by: Zilin Zhu <zhuzilinallen@gmail.com>
Co-authored-by: Nan Jiang <willjiang2018@gmail.com>
yueming-yuan added a commit that referenced this pull request May 21, 2026
@yueming-yuan
[1/14] [sglang-miles] True on-policy training support for FSDP2 (#18639)
f5b3fd2
JessicaJiang-123 pushed a commit to JessicaJiang-123/sglang that referenced this pull request Jun 21, 2026
@yueming-yuan
[1/14] [sglang-miles] True on-policy training support for FSDP2 (sgl-…
6569384 
…project#18639)
nanjiangwill pushed a commit to nanjiangwill/sglang that referenced this pull request Jul 7, 2026
@yueming-yuan
[1/14] [sglang-miles] True on-policy training support for FSDP2 (sgl-…
4fae135 
…project#18639)
XinyuJiangCMU pushed a commit to XinyuJiangCMU/sglang that referenced this pull request Jul 17, 2026
@yueming-yuan
[1/24] [sglang-miles] True on-policy training support for FSDP2 (sgl-…
6d24a34 
…project#18639)
yueming-yuan added a commit that referenced this pull request Jul 25, 2026
@yueming-yuan @maocheng23
[1/23] [sglang-miles] True on-policy training support (FSDP2 + qwen_d…
59c8394 
…ense) (#18639)

Squashes the qwen_dense on-policy follow-up into the original FSDP2 commit:
the follow-up rewrote the norm_kwargs / rl_on_policy_target call sites this
commit introduced, so they cannot be applied independently.

Rebased onto v0.5.16: rl_on_policy_target predicates now go through the
sglang.srt.true_on_policy helpers, and every legacy get_global_server_args()
call site uses runtime_context.get_server_args() so v0.5.16's
test_legacy_global_ratchet baseline still holds.

Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
yueming-yuan added a commit that referenced this pull request Jul 25, 2026
@yueming-yuan @maocheng23
[1/24] [sglang-miles] True on-policy training support (FSDP2 + qwen_d…
f9e7b88 
…ense) (#18639)

Squashes the qwen_dense on-policy follow-up into the original FSDP2 commit:
the follow-up rewrote the norm_kwargs / rl_on_policy_target call sites this
commit introduced, so they cannot be applied independently.

Rebased onto v0.5.16: rl_on_policy_target predicates now go through the
sglang.srt.true_on_policy helpers, and every legacy get_global_server_args()
call site uses runtime_context.get_server_args() so v0.5.16's
test_legacy_global_ratchet baseline still holds.

Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
yueming-yuan added a commit that referenced this pull request Jul 25, 2026
@yueming-yuan @maocheng23
[1/25] [sglang-miles] True on-policy training support (FSDP2 + qwen_d…
809db90 
…ense) (#18639)

Squashes the qwen_dense on-policy follow-up into the original FSDP2 commit:
the follow-up rewrote the norm_kwargs / rl_on_policy_target call sites this
commit introduced, so they cannot be applied independently.

Rebased onto v0.5.16: rl_on_policy_target predicates now go through the
sglang.srt.true_on_policy helpers, and every legacy get_global_server_args()
call site uses runtime_context.get_server_args() so v0.5.16's
test_legacy_global_ratchet baseline still holds.

Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
yueming-yuan added a commit that referenced this pull request Jul 25, 2026
@yueming-yuan @maocheng23
[1/27] [sglang-miles] True on-policy training support (FSDP2 + qwen_d…
5338b11 
…ense) (#18639)

Squashes the qwen_dense on-policy follow-up into the original FSDP2 commit:
the follow-up rewrote the norm_kwargs / rl_on_policy_target call sites this
commit introduced, so they cannot be applied independently.

Rebased onto v0.5.16: rl_on_policy_target predicates now go through the
sglang.srt.true_on_policy helpers, and every legacy get_global_server_args()
call site uses runtime_context.get_server_args() so v0.5.16's
test_legacy_global_ratchet baseline still holds.

Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@merrymercy merrymercy Awaiting requested review from merrymercy merrymercy is a code owner
@Ying1123 Ying1123 Awaiting requested review from Ying1123 Ying1123 is a code owner
@hnyls2002 hnyls2002 Awaiting requested review from hnyls2002 hnyls2002 is a code owner
@Fridge003 Fridge003 Awaiting requested review from Fridge003 Fridge003 is a code owner
@ispobock ispobock Awaiting requested review from ispobock ispobock is a code owner
@xiezhq-hermann xiezhq-hermann Awaiting requested review from xiezhq-hermann xiezhq-hermann is a code owner
@ch-wan ch-wan Awaiting requested review from ch-wan ch-wan is a code owner
@BBuf BBuf Awaiting requested review from BBuf BBuf is a code owner
@Edwardf0t1 Edwardf0t1 Awaiting requested review from Edwardf0t1 Edwardf0t1 is a code owner
@FlamingoPg FlamingoPg Awaiting requested review from FlamingoPg FlamingoPg is a code owner
@AniZpZ AniZpZ Awaiting requested review from AniZpZ AniZpZ is a code owner
@HaiShaw HaiShaw Awaiting requested review from HaiShaw HaiShaw is a code owner
@CatherineSue CatherineSue Awaiting requested review from CatherineSue CatherineSue is a code owner
@slin1237 slin1237 Awaiting requested review from slin1237 slin1237 is a code owner
@JustinTong0323 JustinTong0323 Awaiting requested review from JustinTong0323 JustinTong0323 is a code owner
@yizhang2077 yizhang2077 Awaiting requested review from yizhang2077 yizhang2077 is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@yueming-yuan