Prhub

#2261 fix(rollout): restore partial continuation token budget

原始 PR 作者 Daucloud 合并时间 2026-08-12 13:26 文件变更 2 提交数 1 评论 0 代码增减 +4 / -0

执行摘要

修复 partial rollout 续跑时重复使用完整响应预算

Issue #2260 报告:启用 --partial-rollout 时,中断样本恢复后,SGLang 请求使用完整 max_new_tokens,导致请求 token 数超过上下文限制(如 19202 输入 + 24576 生成 > 40960),SGLang 返回 400,且通用 HTTP 重试机制对确定性错误无效重试。

该 PR 值得精读,虽然改动小,但涉及正确性修复,且回归源于 1eed2493 重构。建议关注其根因与测试缺失,未来应补充单元测试覆盖 partial rollout 续跑预算计算。

讨论亮点

无 Review 评论或讨论。

实现拆解

  1. 在 slime/rollout/sglang_rollout.py 的 generate() 函数中,在调用 _prepare_prompt_ids() 之后、构造 payload 之前,增加 sampling_params["max_new_tokens"] -= sample.response_length
  2. 在 slime/rollout/sglang_streaming_rollout.py 的 generate_streaming() 函数中做同样的修改,保证流式路径行为一致。
  3. 因为响应预算在每次请求前独立扣除,不影响初始请求(此时 response_length 为 0);当预算耗尽时,现有的零预算分支将样本标记为 TRUNCATED,不发送请求。
  4. 未修改重试策略或上下文长度裁剪逻辑,超出历史行为范围,以最小改动恢复原有语义。
文件 模块 状态 重要度
slime/rollout/sglang_rollout.py 训练引擎 modified 4.54
slime/rollout/sglang_streaming_rollout.py 训练引擎 modified 4.54

关键符号

generate generate_streaming

关键源码片段

slime/rollout/sglang_rollout.py core-logic

主变更文件,在 generate() 中恢复响应预算减法,修复续跑超上下文问题。

async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, Any]) -> Sample:
    """Generate using traditional SGLang router with token-based workflow"""
    # ... 状态初始化、URL 构造等 ...
​
    # 准备 prompt,包含已生成 token 的累积(partial rollout 场景)
    prompt_ids = _prepare_prompt_ids(sample, state.tokenizer, state.processor)
​
    # 关键修复:扣除已生成的 response token,避免续跑超预算
    # 初始请求时 response_length 为 0,行为不变;续跑时只请求剩余预算
    sampling_params["max_new_tokens"] -= sample.response_length
​
    # 预算为 0 时标记 TRUNCATED 并返回,不发送请求
    assert sampling_params["max_new_tokens"] >= 0, \
        f"max_new_tokens: {sampling_params['max_new_tokens']} should not be less than 0"
    if sampling_params["max_new_tokens"] == 0:
        sample.status = Sample.Status.TRUNCATED
        return sample
​
    # 构造请求 payload(略)...
slime/rollout/sglang_streaming_rollout.py core-logic

流式路径的对应修复,保证行为与普通路径一致。

async def generate_streaming(args: Namespace, sample: Sample, sampling_params: dict[str, Any]) -> Sample:
    """Streaming counterpart to generate,基于 SSE 流式输出"""
    # ... 状态初始化、URL 构造等 ...
​
    # 准备 prompt,与普通路径相同
    prompt_ids = _prepare_prompt_ids(sample, state.tokenizer, state.processor)
​
    # 同样扣除已生成 token,保证流式续跑遵守总响应预算
    sampling_params["max_new_tokens"] -= sample.response_length
​
    # 预算为 0 时标记 TRUNCATED 并返回
    assert sampling_params["max_new_tokens"] >= 0, \
        f"max_new_tokens: {sampling_params['max_new_tokens']} should not be less than 0"
    if sampling_params["max_new_tokens"] == 0:
        sample.status = Sample.Status.TRUNCATED
        return sample
​
    # 构造流式 payload(略)...

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

该改动简单且局部,风险较低。但缺少单元测试,可能未覆盖以下场景:sample.response_length 大于 max_new_tokens 时,断言会失败(负值预算未处理,仍是原有断言行为);流式路径中 SSE 中断后样本的 response_length 计数的准确性依赖上游状态维护,若在中断前未正确更新样本,可能导致预算扣除不足。

影响范围限于 partial rollout 功能,特别是 oversampling 产生的中断样本续跑。正面影响:避免超出上下文限制的请求及无效重试,提升训练稳定性。负面影响:若中断样本的 response_length 统计不准确,可能导致续跑预算错误(过多或过少),但现有断言会暴露负值情况。团队影响小,无需配置或部署变更。

缺少测试覆盖 核心路径变更 回归风险(源于重构)

关联 Issue

#2260 [Bug] Partial rollout continuation reuses the full response budget

完整报告

参与讨论