feat(rollout): add SGLang native group sampling - #269
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds an opt-in rollout path that uses SGLang’s native parallel sampling (sampling_params.n = group_size) to generate an entire prompt group via a single /generate request, improving multimodal prefix reuse and reducing request overhead while keeping the existing per-sample fallback behavior.
Changes:
- Add
--sglang-native-group-samplingCLI flag and launcher wiring to couple native sampling with--group-rmand a balanced default router policy. - Implement native group sampling eligibility checks and the single-request generation/mapping path in
sglang_rollout.py, reusing existing token/logprob/status handling. - Add unit tests covering eligibility, payload shaping (including multimodal batch-of-one), output contract validation, routing key behavior, and launcher activation rules.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
relax/engine/rollout/sglang_rollout.py |
Adds native group-sampling request path, eligibility gating, and refactors shared token/logprob application logic. |
relax/utils/arguments.py |
Introduces the --sglang-native-group-sampling rollout CLI flag (default off). |
scripts/training/multimodal/run-qwen3-vl-4B-8xgpu.sh |
Enables env-driven activation that also sets --group-rm and defaults router policy to round_robin for native mode. |
tests/engine/rollout/test_sglang_native_group_sampling.py |
Adds contract and behavior tests for native group sampling and launcher rules. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| parser.add_argument( | ||
| "--sglang-native-group-sampling", | ||
| action="store_true", | ||
| default=False, | ||
| help=( | ||
| "Use one SGLang /generate request with sampling_params.n equal to the prompt group size " | ||
| "for standard first-turn group-RM rollouts." | ||
| ), | ||
| ) |
| if _native_group_sampling_eligible(args, group, evaluation): | ||
| group = await _generate_native_group(args, state, group, sampling_params) | ||
| else: | ||
| tasks = [] |
|
补充下正确性测试吧，跑 baseline 和 experiment，看看那些指标（loss、grad norm、reward 等等）是否正确。 |
|
别的我没啥问题了。 |
|
我将默认的multimodal reasoning 数据清洗后放大分辨率，以证明该方案在 prefix 阶段的优势。此方案在多图大分辨率 OCR-VLM 训练中具有一定的优势。 如代码需要修改，或 CLI 参数需要对齐沟通，请告诉我。 |
# 🐛 Bug Fix ## Prevent pre-encoded media leak into sample dicts - Wrap the group generate branch in try/finally so the temporary pre-encoded media attributes are always dropped before samples reach Sample.to_dict(), covering the per-sample fallback, abort and error paths alike - Previously the base64 media blobs could leak into the data buffer via __dict__ on early returns ## Keep aborted groups eligible for the native path - Revert tokens, rollout_tokens and multimodal_train_inputs to their pre-call state when the request is aborted before any output, so retried groups no longer silently degrade to per-sample requests ## Exclude speculative decoding from native eligibility - The draft/verify path is not validated with parallel sampling, so keep the per-sample fanout when a speculative algorithm is set ## Add visibility for the native/fallback decision - Log once when native group sampling activates and warn once when the flag is enabled but a training group is ineligible - Warn at startup when flag combinations make every group ineligible --- # ♻️ Refactor ## Remove dead code and document semantics - Drop the unreachable routed_experts branch: the native payload never requests it and routing replay is excluded by eligibility - Document group-shared timing duplication and session-permit concurrency semantics --- # ✅ Tests ## Cover abort, cleanup and eligibility edge cases - Mid-flight abort maps n outputs to aborted samples with partial tokens retained - Abort before dispatch restores a fresh, still-eligible group - Pre-encoded media is stripped on success and on fallback failure - Speculative decoding joins the incompatible-modes matrix - Ineligible-group warning fires only once
由于rollout 引擎卡死导致 baseline 实验重启，重启后baseline 模型行为在前 50step 迅速坍缩，仅供参考，如时间充裕会重新再测。 |
|
第一张图片看起来是正常的；第二张图片的 baseline 曲线看起来有很大问题，后续有再测吗？ |
|
以及你的图片只展示了 reward，clearml 或者日志里面应该也有 输出数量、response token length、rollout_log_probs length、finish_reason、reward、advantage、loss、grad norm 这些指标的对比，也可以一并贴近来。 |
第二轮测试还在进行，目前到 50step 了，这次没有快速坍缩，100step 的时候我更新到这里。 |
好的 |
此处依旧是旧实验，在进行三轮新实验的过程中，训练会被偶发 SGLang 502 Bad Gateway 干扰重启，于是没有新的完整曲线（新的训练在 30step 和 60step 的时候分别遇到 BUG 重启）。 #282 此处为 BUG 描述和初步解决方案，确认后可提 PR 修复。 |
|
那要不你补充下第一次测试的这些正确性曲线？ |
这些均为我跑过的实验，供参考，曲线的图例已经标明。其中 完整的baseline 只成功跑了一次，其余baseline 的训练均由我上述提到的issue中的问题被迫终止。若此 issue 确认为BUG 且初步解决方案被认同，我可以在后续的开发中fix。 |
|
可以，我觉得先验收通过了；你可以先把 bug 提新的 社区 PR 修了，然后再跑吧…… |
# 🐛 Bug Fix ## Guard shared multimodal encoding - Skip group-level pre-encoding when media payloads are empty - Preserve text-only rollout behavior for empty multimodal mappings --- # ✅ Tests ## Cover empty media groups - Verify empty multimodal payloads bypass the encoder - Confirm reward assignment continues through the fallback path
ffbaa3f to
4984fb3
Compare
这是目前跑的对比。其中 Native 和 baseline 都有坍缩的现象，但坍缩的 Native 由于 OOM 没有训完。坍缩的 Native 表现和 坍缩的baseline 一致（完成的 step 内）。同时未坍缩的 Baseline 和未坍缩的 Native表现也几乎一致（完成的 step 内）。后面有需要可以多次实验去收集完整的坍缩的 Native 版本的实验。
|
|
感谢提交和详细的测试。合入的主要顾虑：1）native 路径依赖 SGLang n>1 group sampling，官方对多模态/投机解码等组合支持有限，导致 fallback 分支多、维护成本高；2）收益场景（多模长输入重复计算）后续计划通过 EPD 解决。建议先不合入。不影响此前验收结论哈～ |
感谢团队的 review。后面有机会继续向团队学习，贡献高质量代码。 |
标题
feat(rollout): 使用 SGLang 原生采样生成 prompt groupclose #268
变更概述
本 PR 为标准 group-RM rollout 增加一条默认关闭的 native group sampling 路径。
对于同一个 prompt/image 的 8 条 sample，Relax 不再发送 8 个独立
/generate请求，而是发送一个sampling_params.n=8的请求。SGLang 只需 prefill 一次公共多模态前缀，然后将其展开为 8 条 decode 分支。原有逐 sample 路径继续作为 fallback，native 功能默认关闭。
具体改动
--sglang-native-group-samplingCLI 参数。n>1时被误判为多个 batch entry。--group-rm。round_robin；可以通过SGLANG_ROUTER_POLICY显式覆盖。--group-rm，保持训练脚本原有默认行为。兼容与回退条件
以下情况不会使用 native group sampling：
这些情况继续使用现有逐 sample 路径。
启动脚本激活规则
训练脚本现在使用以下规则：
之所以默认使用 round-robin，是因为 native 会把每个 rollout 的 HTTP 请求数从约 512 降至 64。请求数变少、单请求 decode 工作量变大后，cache-aware 路由更容易产生长尾负载失衡。
Cache-aware 问题的发现与修复过程
本 PR 的路由默认值来自实际 A/B 排查，不是预先假设。
第一阶段：验证 native 机制是否命中
Baseline 四轮共发送 2048 个
/generate请求，即每轮约 512 个独立 sample 请求。开启 nativen=8后，四轮请求数降到 256，即每轮约 64 个 prompt group 请求。请求数严格接近
512 -> 64，同时每轮仍然产出 64 groups / 512 samples、零 abort，说明 native 合并和输出回填都已生效。第二阶段：发现 cache-aware 下性能反而下降
最初只开启 native，保持默认/cache-aware 路由。稳态结果为：
76.7 -> 82.7秒，增加 7.8%；9,678.52 -> 8,848.82token/s，下降 8.6%；1449.26 -> 1428.72，只下降约 1.4%，不足以解释性能回退。这说明请求合并已经成功，但系统出现了新的瓶颈。
进一步统计 SGLang worker 访问日志：
Native/cache-aware 中有一个 engine 收到 67 个完整 group，而最少的 engine 只有 24 个。一个 native group 又包含 8 条 decode，因此这种请求数倾斜会被放大为明显的 decode 长尾。
第三阶段：为什么转向 round-robin
Native
n=8已经在单个请求内部完成公共 prefix 复用，各 group 之间又是不同的 prompt/image，所以跨请求 cache affinity 不再是主要收益来源。此时更重要的是把 64 个粗粒度 decode group 均匀铺到 8 个 engine。从实验现象推断，cache-aware 的请求级 cache/load 信号没有充分计入一个
n=8请求背后的 8 条 decode 工作量；而 round-robin 不依赖这个估计，可以直接保证 group 数量均匀。因此先采用最小改动：保持 native
n=8不变，只把路由切到 round-robin。如果仍然存在严重的 group 内长度长尾，再考虑拆成两个n=4请求，在 prefix 复用和调度粒度之间折中。本轮 round-robin 已经解决主要问题，无需拆分n=4。切换后的四轮请求分布为
31–33，CV 降至约 1.6%。稳态 response-token throughput 相对 baseline 提升 14.3%，生成墙钟时间降低 12.6%。自动化测试
Native 专项测试：
相邻 rollout 和 argument 回归：
两个 deselected 是当前沙箱中会阻塞的既有 Deepeyes permit 集成测试，不涉及本改动。
以下静态检查全部通过：
py_compilebash -ngit diff --check完整命令、验收矩阵和实验日志见
/root/paddlejob/workspace/env_run/research/TASK24_PR_TEST.md。8-GPU A/B 实验
实验配置：
--group-rm，避免 reward 调用时机成为额外变量Native + round-robin 相对 baseline：
请求负载方面，默认路由的每卡请求数范围为
24–67、CV 约 41.9%；round-robin 后收敛到31–33、CV 约 1.6%。该结果是启动脚本把 round-robin 作为 native 默认策略的直接依据。风险与限制
检查清单