执行摘要

• 一句话：支持可变 global batch size 训练，调度迁至 rollout 侧
• 推荐动作：值得精读。这是 miles 训练数据管线的架构级变更：调度、计账、指标三条链路同时重构且保持严格 no-op。重点读 build_dp_schedule 的「先打包后分发」两阶段设计、rollout_mask_sums 分母广播的数学，以及 can_schedule_on_rollout_side 的资格判定列表（它决定了未来新 backend 的接入方式）。建议配合 THUDM/slime 系列 PR 一起看设计演进，并关注后续 #2135/#2197 的清理与修复。

功能与动机

PR body 明确说明动机：RL rollout 生产天然是变长的——dynamic-sampling 过滤会丢样本、agentic/多轮 rollout 每次执行产生的训练样本数不定、采集到的批次几乎不会恰好是 global_batch_size 的整数倍。legacy 管线假设固定 global_batch_size（按样本计数）、每个 DP rank 持有相同样本数、per-sample loss 加权，后果是：拖尾样本被裁剪（浪费生成）、一次产出 N 个样本的 rollout 在 loss 里被计 N 次（长执行被过加权）、兄弟样本可能被拆到不同 step 或裁剪边界。设计跟随 THUDM/slime 的 variable-global-batch-size 系列（slime #1930/#1933），FLOPs 均衡打包移植 slime #2017/#2029，FFD 打包与 partial final step 是 miles 自己的扩展。

实现拆解

实现分 5 步拆解：

1. 调度侧迁移（commit 1）：新增 miles/utils/dp_schedule.py 纯 Python 调度模块，build_dp_schedule 一次返回 partitions / micro_batch_indices / num_microbatches / num_rollouts；has_full_schedule_config 检查训练后端是否广告完整并行配置（dp/cp/vpp/mb_group）。train_data_conversion.py 的 split_train_data_by_dp 签名从 dp_size 改为 train_parallel_config，通过 can_schedule_on_rollout_side 分流：可调度则走 split_train_data_by_dp_scheduled_raw 并把布局写入 shard；否则保留 legacy split_train_data_by_dp_raw（multi-LoRA、fsdp/torchtitan、indep_dp/FT、多模态）。训练侧 get_data_iterator 改为消费预计算布局，删掉原来跨 rank 的 all_reduce(MAX) 同步点。

2. 以 rollout 为单位的 step 规划与打包：样本按 Sample.rollout_id（缺省回落 sample.index）聚合，global_batch_size 按 rollout 计数，同一 rollout 的兄弟样本永不跨 step。动态路径下用新增 first_fit_decreasing_pack 全局 FFD 打包到 max_tokens_per_gpu * cp 预算，再通过 expand_bins_by_splitting 把 micro-batch 数对齐到 dp_size * mb_group（VPP 场景需要）；随后按 strided、Karmarkar-Karp（--balance-data）或 FLOPs 均衡（--balance-by-flops）分发到 DP rank。静态路径按固定 micro_batch_size 切块，若无法对齐则启动即失败（commit 5 / #2197）。

3. 每 rollout 的 loss 计账：convert_samples_to_train_data 新增 rollout_ids 与 rollout_mask_sums（_compute_rollout_mask_sums 把整次 rollout 的 mask 总数广播给每个兄弟样本）；cp_utils.get_sum_of_sample_mean 新增 denominators 参数，loss.py 用 rollout 分母替换 per-sample 分母，使散落在不同 micro-batch、DP、CP rank 的 sum/分母 在既有 all-reduce 下精确重建为每 rollout 一次 token 加权均值；metric 侧 log_utils 改为 (sum, count) 元组 + reduce_gathered_log_dict 的 count-weighted 归约，aggregate_train_losses 新增 num_rollouts 分支。

4. 参数与资格校验：miles/utils/arguments.py 新增 --allow-partial-train-step、--balance-by-flops 等参数；can_schedule_on_rollout_side 明确排除 multi-LoRA、多模态、配置不全的 backend；rollout_data_conversion.py 新增 validate_compact_rollout_ids 保证 compact 形状下兄弟样本必须共享 rollout_id。

5. 测试与 CI 配套：CPU 侧覆盖调度不变式（覆盖完备、无重叠、PP 同步、超大样本独享 bin、VPP 对齐、尾部裁剪、partial step、rollout 分组）、FFD 性质、denominator 等价性、归约语义；GPU e2e 三个：test_qwen2.5_0.5B_compact_rollout.py（compact 拆分）、test_qwen3_4B_variable_gbs.py（gbs=24 训成 [24,8]）、test_qwen3_30B_A3B/test_dp_attention.py（dp-attention 顺序不敏感）。

关键文件：

• miles/utils/dp_schedule.py（模块 调度器；类别 source；类型 core-logic；符号 build_dp_schedule, has_full_schedule_config, _calculate_workloads）: 新增的纯 CPU 调度器核心：按 rollout 聚合样本、规划每步 rollout 数、FFD 打包、对齐到 dp*vpp、分发到各 DP rank，是整个 PR 的地基。
• miles/ray/rollout/train_data_conversion.py（模块 数据转换；类别 source；类型 core-logic；符号 split_train_data_by_dp, can_schedule_on_rollout_side, split_train_data_by_dp_scheduled_raw, _compute_rollout_mask_sums）: 调度资格判定与数据契约扩展的枢纽：新增 rollout_ids / rollout_mask_sums / num_microbatches / micro_batch_indices / num_rollouts 五个键，并实现 rollout 侧调度与 legacy 路径的分流。
• miles/utils/seqlen_balancing.py（模块 打包层；类别 source；类型 core-logic；符号 first_fit_pack, first_fit_decreasing_pack, _first_fit, expand_bins_by_splitting）: 新增 first-fit / FFD 装箱与基于 LPT 的 bin 拆分扩容，是动态打包与 align 对齐的底层工具，被 build_dp_schedule 直接依赖。
• miles/backends/training_utils/log_utils.py（模块 指标归约；类别 source；类型 core-logic；符号 reduce_gathered_log_dict, gather_log_data, log_rollout_data, aggregate_train_losses）: 指标归约语义从 mean-of-rank-means 改为 count-weighted (sum, count)，aggregate_train_losses 支持按 num_rollouts 归一，是 per-rollout 计账在日志侧的落地。
• miles/backends/training_utils/cp_utils.py（模块 并行工具；类别 source；类型 data-contract；符号 get_sum_of_sample_mean）: get_sum_of_sample_mean 新增 denominators 参数，loss 分母从 per-sample 改为 per-rollout，是计账修正的数学核心。
• miles/backends/training_utils/loss.py（模块 损失计算；类别 source；类型 data-contract）: loss reducer 接入 rollout 分母，使跨 micro-batch 的部分贡献重建为每 rollout 一次 token 加权均值。
• miles/backends/training_utils/data.py（模块 数据接口；类别 source；类型 dependency-wiring；符号 get_num_rollouts）: 新增 get_num_rollouts，把真实 per-step rollout 数从 shard 穿到归一化与 LR 调度。
• miles/utils/arguments.py（模块 参数解析；类别 source；类型 configuration）: 新增 --allow-partial-train-step、--balance-by-flops 等参数开关，并承载相关参数校验。
• miles/ray/rollout/rollout_data_conversion.py（模块 数据校验；类别 source；类型 core-logic；符号 validate_compact_rollout_ids）: 新增 validate_compact_rollout_ids，在 flatten 前强制 compact 兄弟样本共享 rollout_id，保证计账分组正确。
• tests/fast/utils/test_dp_schedule.py（模块 调度测试；类别 test；类型 test-coverage；符号 assert_invariants, test_static_stride_single_step, test_dynamic_with_vpp_rounds_to_mb_group, test_rollout_grouping_keeps_samples_together）: 440 行 CPU 单测覆盖调度不变式与静态/动态/平衡/VPP/裁剪/分组场景，是调度器正确性的主要保证。
• tests/fast/backends/training_utils/test_rollout_mask_sums.py（模块 计账测试；类别 test；类型 test-coverage；符号 TestRolloutMaskSumsPrecompute, TestDenominators, TestAggregateTrainLossesNumRollouts）: 验证 rollout_mask_sums 预计算、denominator 与 legacy per-sample 的严格等价、per-rollout loss 归约，是计账数学的回归防线。
• tests/e2e/short/test_qwen2.5_0.5B_compact_rollout.py（模块 紧凑回滚；类别 test；类型 test-coverage；符号 prepare, execute）: compact 拆分场景的端到端验证：每个 execution 拆成两个共享 rollout_id 的训练样本，验证按 rollout 计 step 与 loss 加权。
• tests/e2e/megatron/test_qwen3_4B_variable_gbs.py（模块 可变批大小；类别 test；类型 test-coverage；符号 prepare, execute）: 可变 gbs 主路径 e2e：gbs=24 采集 32 个 rollouts，配合 --allow-partial-train-step 与 cp2 训练成 [24, 8] 两个不均匀 step。

关键符号：build_dp_schedule, has_full_schedule_config, _calculate_workloads, can_schedule_on_rollout_side, split_train_data_by_dp, split_train_data_by_dp_scheduled_raw, split_train_data_by_dp_raw, _compute_rollout_mask_sums, _package_shards, first_fit_pack, first_fit_decreasing_pack, _first_fit, expand_bins_by_splitting, _split_bin_by_tokens, reduce_gathered_log_dict, gather_log_data, log_rollout_data, aggregate_train_losses, get_sum_of_sample_mean, get_num_rollouts, validate_compact_rollout_ids

关键源码片段

miles/ray/rollout/train_data_conversion.py

调度资格判定与数据契约扩展的枢纽：新增 rollout_ids / rollout_mask_sums / num_microbatches / micro_batch_indices / num_rollouts 五个键，并实现 rollout 侧调度与 legacy 路径的分流。

def can_schedule_on_rollout_side(args, data, train_parallel_config):
    """判定是否能在 rollout 侧预计算完整 DP/microbatch 调度。"""
    if not has_full_schedule_config(train_parallel_config):
        return False # fsdp/torchtitan 只暴露 dp_size，走 legacy
    if is_multi_lora_enabled(args):
        return False # slot 连续性约束，训练侧调度保留
    if "multimodal_train_inputs" in data:
        return False # 多模态序列长度在训练侧才会展开
    if "rollout_ids" not in data:
        return False
    global_batch_size = data.get("dynamic_global_batch_size", args.global_batch_size)
    return len(set(data["rollout_ids"])) >= global_batch_size # 至少凑满一个 step
​
​
def split_train_data_by_dp(args, data, train_parallel_config):
    """按 DP 切分并写入对象存储；能预计算调度时走 rollout 侧，否则回退 legacy。"""
    if can_schedule_on_rollout_side(args, data, train_parallel_config):
        shards = split_train_data_by_dp_scheduled_raw(args, data, train_parallel_config=train_parallel_config)
    else:
        shards = split_train_data_by_dp_raw(args, data, dp_size=train_parallel_config["dp_size"])
    store = object_store.get_instance()
    return [store.put(value=shard, value_spec=ROLLOUT_DATA_VALUE_SPEC) for shard in shards]
​
​
def split_train_data_by_dp_scheduled_raw(args, data, *, train_parallel_config):
    """在 rollout 侧跑调度器，把 micro-batch 布局打进每个 rank 的 shard。"""
    total_lengths = [len(t) for t in data["tokens"]]
    data["total_lengths"] = total_lengths
    global_batch_size = data.get("dynamic_global_batch_size", args.global_batch_size)
    partitions, micro_batch_indices, num_microbatches, num_rollouts = build_dp_schedule(
        args, train_parallel_config, total_lengths,
        global_batch_size=global_batch_size,
        rollout_indices=data["rollout_ids"],
    )
    logger.info(f"Rollout-side DP schedule: num_samples={len(total_lengths)}, "
                f"num_rollouts={num_rollouts}, num_microbatches={num_microbatches}")
    shards = _package_shards(args, data, partitions)
    for rank, shard in enumerate(shards):
        shard["num_microbatches"] = num_microbatches
        shard["micro_batch_indices"] = micro_batch_indices[rank]
        shard["num_rollouts"] = num_rollouts
    return shards
​
​
def _compute_rollout_mask_sums(rollout_ids, loss_masks):
    """每个样本携带其整次 rollout 的 loss mask 总和，作为 loss 分母。
​
    同一 rollout 的兄弟样本即使散落在不同 micro-batch、不同 DP/CP rank，
    各自贡献的 sum/分母 也能在既有 all-reduce 下重新拼成一次 rollout 均值。
    """
    totals = {}
    for rid, mask in zip(rollout_ids, loss_masks, strict=True):
        totals[rid] = totals.get(rid, 0) + sum(mask)
    return [totals[rid] for rid in rollout_ids]

评论区精华

review 中最有价值的交锋集中在两点：

• 命名歧义（已解决）：guapisolo 指出 mbs is confusing "micro batch size" or "micro batches"，并亲自提交 commit f66b78a 将 mbs/step_mbs/num_mbs/rank_mbs_idx 等全部改为 micro_batch 全称，同步清理 assert 文案与单测。

• 静态路径硬 assert 的脆弱性（已解决）：guapisolo 评论 The hard assert is required, but it's very fragile. There is no way to ensure len(step_micro_batches) % align_to == 0 or fit in target，建议直接禁止非 multi-LoRA 路径混用 dynamic gbs；后续由 commit ee9e54d（#2197）改为「启动即失败」并给出可操作报错，而非运行期静默出错。

• 后续清理：guapisolo 在 DISMISSED review 中说明 recover some key cmts to improve readability. LGTM now, but better see #2135，提示还存在一个后续清理 PR。

• 下游消费方：yushengsu-thu 在关联 Issue 下留言 I need this pr for enabling variable global batch size in multi-lora #2137, waiting for this pr merging，说明本 PR 是 multi-LoRA 可变 gbs 的前置依赖。

  • dp_schedule 中 mbs 标识符歧义 (style): guapisolo 亲自提交 commit f66b78a，将 mbs/step_mbs/num_mbs/rank_mbs_idx 等全部改为 micro_batch 全称，并同步清理 assert 文案与单测。
  • 静态路径对齐硬 assert 的脆弱性 (design): 后续 commit ee9e54d（#2197）实现：静态路径无法对齐时启动即失败，并给出可操作报错（step_size % (dp_size * micro_batch_size * mb_group) == 0），避免运行期静默出错。
  • multi-LoRA 可变 gbs 的依赖关系 (question): 本 PR 将 multi-LoRA 明确排除在 rollout-side 调度之外（保留 legacy 路径），#2137 作为下游消费方等待合并。
  • 后续清理 PR #2135 (other): 存在一个后续清理 PR #2135 需跟进；本 PR 主体已获 APPROVED。

风险与影响

• 风险：具体风险点如下：

• 静态路径启动失败：build_dp_schedule 静态分支在 step_size % (dp_size * micro_batch_size * mb_group) != 0 时直接 AssertionError，老配置可能突然失败。这是刻意设计（启动即失败换取正确性），但团队需按报错提示调整配置。

• 计账重建依赖 all-reduce：rollout_mask_sums 分母广播到所有兄弟样本，靠现有 DP*CP all-reduce 把 sum/分母 重新拼出每 rollout 均值；任何新的切分维度（如未来 CP 扩展）都必须透传 denominator 一致性，否则数值偏差且难定位。
• TIS decoupled 语义变化：修改后的 response masks 只修正 numerator，denominator 仍用预计算的 rollout 总数；若 mask 修改面积很大，loss 会被系统性低估。
• metric 口径变化：reduce_gathered_log_dict 从 mean-of-rank-means 变成 count-weighted，(sum,count) 元组与旧 scalar 并存，跨版本对比指标时需注意口径不同。
• 性能：调度是纯 Python 且每 step 执行一次 Karmarkar-Karp / FFD，当样本量极大（数万 samples）时 CPU 耗时与内存占用可能成为 rollout 管理器的瓶颈。
• 数据契约扩展：ROLLOUT_DATA_VALUE_SPEC 新增 5 个键，老 trainer 与新版 rollout 混跑会出现 spec 不匹配；依赖 gbs 数值做显存估算的脚本需按 rollout 口径重新校准。

• 影响：影响范围与程度：

• 用户/实验者：新增 --allow-partial-train-step、--balance-by-flops 参数；compact/agentic 负载的 loss 计账被修正（不再按样本数过加权）；拖尾 rollouts 可选训练而非丢弃。对经典「1 rollout = 1 sample」用户无感知（已有 30B 数值验证为严格 no-op）。

• 系统：调度职责从训练侧迁移到 rollout 侧，训练侧 get_data_iterator 简化（去掉 all_reduce(MAX) 同步点）；对象存储 spec 扩展 5 个键；指标日志从 per-rank mean 变为 (sum, count)。
• 团队：本 PR 是后续多特性的地基（multi-LoRA 可变 gbs #2137、fully-async 精确计账、dp-attention 调度验证），且改变了训练管线对 gbs 的语义理解，需要同步更新文档与实验脚本。
• 风险标记：核心训练管线重构, gbs 口径从样本改为 rollout, 静态路径启动即失败, 跨 DP/CP 计账依赖 all-reduce, TIS 解耦分母语义变化, 指标归约口径变化

关联脉络

• PR #1673 [feat] support sample rollout submission granularity to keep fully async concurrency: 同属「rollout 计账粒度」演进线：fully-async 下样本数不定正是可变 gbs 的核心动机场景，本 PR 的 rollout_id 分组机制建立在其提交粒度重构之上。
• PR #1717 Rewrite fully-async rollout as FullyAsyncRolloutFn on the class-based rollout API: fully-async/agentic 多轮 rollout 每次执行产出样本数不定，是 variable gbs 的动机来源；本 PR 的 per-rollout 计账补上了该路径的 loss 归一。
• PR #2041 Address engines by base URL when the router is dp-aware: 本 PR 新增 tests/e2e/megatron/test_qwen3_30B_A3B/test_dp_attention.py，验证 rollout-side 调度在 dp-aware 请求重排下保持顺序不敏感（按 rollout_id 分组），与 dp-attention 路由修复直接交叉。
• PR #2012 router: enable dp-aware routing under dp-attention: dp-aware 路由开启后样本完成顺序被打乱，本 PR 的调度必须不依赖完成顺序，二者构成同一 dp-attention x rollout 调度的验证链。