执行摘要
- 一句话:新增 fanout 端到端测试,验证单 prompt 变长多样本训练链路
- 推荐动作:值得阅读以理解 slime 中 'compact rollout' 模式的设计约束和集成方式。特别是
rollout_id 共享对 step splitter、reducer、loss aggregation 的影响。
功能与动机
此前 CPU 单元测试已覆盖 rollout-aware step splitter、per-rollout mean reducer 等独立环节,但始终缺少一个端到端的训练验证。这是首次在真实训练流程中检验 custom_generate 返回 list[Sample] 并共享 rollout_id 的完整链路,确保框架不会因为样本数不均匀而 silent 地退化到错误的分组归一化。
实现拆解
- 新增
slime/rollout/_fanout_test_helpers.py,包含 compact_generate(异步生成函数,对每个 prompt 调用一次 sglang 生成 base_sample,再 deepcopy 出 N = 1 + (index % 3) 个 siblings,全部设置相同 rollout_id)和 grpo_normalize_by_group_index(替代默认的 _post_process_rewards,按 sample.group_index 分组进行 GRPO 标准化)。
- 新增
tests/test_qwen2.5_0.5B_fanout_short.py,定义 prepare(下载模型和数据集)和 execute(运行训练命令,使用 --custom-generate-function-path 和 --custom-reward-post-process-path 注入辅助函数)。测试通过环境变量计数器验证自定义 generate 被调用了 num_rollout * rollout_batch_size 次,没有静默回退。
- 修改
.github/workflows/pr-test.yml 和 .github/workflows/pr-test.yml.j2,在 e2e-test-sglang 矩阵中加入该测试。
- 修改
pyproject.toml,在 [tool.ruff] 下添加空的 [tool.ruff.lint] section,以适配较新版本的 ruff 配置格式。
关键文件:
slime/rollout/_fanout_test_helpers.py(模块 rollout 辅助;类别 source;类型 dependency-wiring;符号 compact_generate, grpo_normalize_by_group_index): 核心辅助模块,提供 compact_generate 和 grpo_normalize_by_group_index 两个自定义钩子函数,是端到端测试的基础设施。
tests/test_qwen2.5_0.5B_fanout_short.py(模块 fanout 测试;类别 test;类型 test-coverage;符号 prepare, execute): 端到端测试文件,验证完整 fanout 训练链路。测试使用真实模型和数据集,通过自定义钩子检查 rollout_id 共享和分组归一化正确性。
.github/workflows/pr-test.yml(模块 CI 配置;类别 infra;类型 infrastructure): CI 工作流配置,在 e2e-test-sglang 矩阵中新增本测试条目,确保每次提交都自动运行。
pyproject.toml(模块 项目配置;类别 config;类型 configuration): 修复 ruff 配置格式,添加 [tool.ruff.lint] section 以适配新版 ruff。
.github/workflows/pr-test.yml.j2(模块 CI 模板;类别 infra;类型 infrastructure): CI 工作流模板,与 pr-test.yml 同步修改,确保生成正确。
关键符号:compact_generate, grpo_normalize_by_group_index, prepare, execute
关键源码片段
slime/rollout/_fanout_test_helpers.py
核心辅助模块,提供 compact_generate 和 grpo_normalize_by_group_index 两个自定义钩子函数,是端到端测试的基础设施。
import copy
import os
from collections import defaultdict
MAX_FANOUT = 3
COUNTER_FILE_ENV = "SLIME_FANOUT_TEST_COUNTER_FILE"
async def compact_generate(args, sample, sampling_params):
"""One prompt to N siblings, deterministic N = 1 + (index % MAX_FANOUT).
Strategy: call sglang once, deepcopy N-1 times. Bounded GPU cost —
we're pinning the framework's per-rollout handling, not generation
diversity.
"""
from slime.rollout.sglang_rollout import generate
# Counter file is used by the test to verify that every prompt
# actually went through this function (no silent fallback).
counter_path = os.environ.get(COUNTER_FILE_ENV)
if counter_path:
try:
with open(counter_path, "a") as f:
f.write(f"{sample.index}\n")
except OSError:
# best-effort
pass
base_sample = await generate(args, sample, sampling_params)
# Number of siblings: 1, 2, or 3 — varies per prompt to stress uneven fanout.
n = 1 + (sample.index % MAX_FANOUT)
siblings = []
for _ in range(n):
s = copy.deepcopy(base_sample)
# Critical: all siblings share rollout_id so the per-rollout reducer
# aggregates them as ONE rollout (not N) and the rollout-aware step
# splitter keeps them in the same step. group_index is inherited
# via deepcopy and is used by grpo_normalize_by_group_index below.
s.rollout_id = sample.index
siblings.append(s)
return siblings
tests/test_qwen2.5_0.5B_fanout_short.py
端到端测试文件,验证完整 fanout 训练链路。测试使用真实模型和数据集,通过自定义钩子检查 rollout_id 共享和分组归一化正确性。
import os
import tempfile
import slime.utils.external_utils.command_utils as U
TIGHT_DEVICE_MEMORY = U.get_bool_env_var("SLIME_TEST_TIGHT_DEVICE_MEMORY", "1")
MODEL_NAME = "Qwen2.5-0.5B-Instruct"
MODEL_TYPE = "qwen2.5-0.5B"
NUM_GPUS = 4
FANOUT_COUNTER_FILE = os.environ.get(
"SLIME_FANOUT_TEST_COUNTER_FILE",
os.path.join(tempfile.gettempdir(), "slime_fanout_test_counter.log"),
)
def prepare():
"""Download model and dataset, clear the test counter."""
U.exec_command("mkdir -p /root/models /root/datasets")
U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}")
U.hf_download_dataset("zhuzilin/dapo-math-17k")
# Ensure previous run invocations don't bleed into this test.
try:
os.remove(FANOUT_COUNTER_FILE)
except FileNotFoundError:
pass
def execute():
"""Run the training command with custom generate and reward post-process hooks."""
ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ "
# Shape: rollout_batch_size=8, n_samples_per_prompt=1 (all fanout is
# owned by compact_generate), global_batch_size=4, num_rollout=3.
rollout_args = (
"--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl "
"--input-key prompt --label-key label --apply-chat-template "
"--rollout-shuffle --rm-type deepscaler --num-rollout 3 "
"--rollout-batch-size 8 --n-samples-per-prompt 1 "
"--rollout-max-response-len 8192 --rollout-temperature 0.8 "
"--global-batch-size 4 --balance-data "
"--custom-generate-function-path slime.rollout._fanout_test_helpers.compact_generate "
"--custom-reward-post-process-path slime.rollout._fanout_test_helpers.grpo_normalize_by_group_index "
)
# ... (train command and post-train counter assertion follow)
评论区精华
无 review 评论。但设计上有几点值得注意:辅助函数放在 slime/rollout/ 而非 tests/ 是因为 --custom-generate-function-path 通过 importlib.import_module 解析点路径,测试文件文件名包含点(test_qwen2.5_0.5B_fanout_short.py)无法被正确 import;因此采用 underscored 前缀的模块且放置在源码目录下。
风险与影响
- 风险:风险较低。测试运行在 4 卡 GPU 上,增加了 CI 耗时;pyproject.toml 的 ruff 配置改动可能影响 linting 行为,但目前仅添加空 section,无实际影响。若测试在 CI 中不稳定(如 GPU 资源竞争或超时),可能需要调整超时设置。
- 影响:仅影响 CI 测试矩阵和开发调试流程。对用户功能无直接影响。但此测试覆盖了核心路径,能预防未来回归。
- 风险标记:CI 新增测试耗时, ruff 配置改动可能意外
关联脉络
- PR #1926 Move micro-batch scheduling from training side to rollout side: 引入了 rollout-aware 的 step splitter,本测试的 fanout 模式依赖此调度正确性。
- PR #1930 [1/N] Support training with variable global batch size: 支持动态 global batch size,fanout 中样本数变化依赖该基础设施。
- PR #1933 [2/N] Support training with variable global batch size: 继续 variable global batch size 工作,fanout 测试是对其的集成验证。
- PR #1921 Add example for streaming output: 展示了自定义 generate 函数的设计模式,本测试采用了类似方式。
- PR #1920 Move fully_async example to main codebase: 将自定义 rollout 模式纳入主代码库,fanout 测试是其延伸。
参与讨论