[1/N][BUGFIX][Sync sglang-miles] True on-policy training support for FSDP2 - #23543
[1/N][BUGFIX][Sync sglang-miles] True on-policy training support for FSDP2 #23543zyzshishui wants to merge 14 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces several optimizations and fixes for RL on-policy inference across Qwen2 and Qwen3 models. Key changes include refactoring RMSNorm to handle residual additions and precision more consistently, implementing a deterministic inference path for MoE kernels, and significantly optimizing position embedding interpolation in Qwen3-VL by replacing NumPy-based CPU loops with vectorized Torch operations. Additionally, the PR refactors logit processing in the sampler and adds regression tests for on-policy logprob alignment. Review feedback highlighted a scaling bug in RMSNorm during all-reduce fusion, precision loss when maintaining FP32 residuals, and a performance bottleneck in the Qwen3-VL interpolation logic caused by data movement between GPU and CPU.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 25f0eefdd3
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
…roject#18639) (cherry picked from commit 74c9eab) Co-authored-by: Zilin Zhu <zhuzilinallen@gmail.com> Co-authored-by: Nan Jiang <willjiang2018@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b3bc14b55a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
zyzshishui
changed the title
2 participants
Motivation
Synced from #18639 and add tests
Modifications
Accuracy Tests
Manually ported logprobs from FSDP2 Qwen3-0.6B https://huggingface.co/datasets/zyzshishui0627/qwen3-0.6b-fsdp2 and confirmed diff is 0.
Speed Tests and Profiling
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ci