[WIP] Add index-based data packing and worker-side materialization

[YanhuiDua]

概述

本 PR 将 RL 训练数据准备重构为 controller 只做规划、worker 只做物化 的两层结构，并把数据准备与 agent loop 侧 meta-ready 的 RolloutState 契约打通：controller 全程只读写 write_train_meta 元数据（长度、flag），不接触任何张量；TrainingWorker.fit 按三段流水线完成张量构造与训练。总数据流：

rollout_groups: list[list[RolloutState]]
      │  TrainingController.fit（规划 + 分发）
      ▼
rollout_item_refs / advantages / pack_plan / batch_attr   （跨 Ray 边界，每 DP rank 一次）
      │  TrainingWorker.fit（物化 + 训练）
      ▼
RolloutState →convert→ (seq_ctx, loss_ctx) →pack→ [optimizer_step][pack] →_fit→ WorkerLogItem

TrainingController.fit pipeline

rollout_groups: list[list[RolloutState]]        # replay buffer 选出的 COMPLETED groups
  │
  ▼ ① _prepare_rollout_items
  │    按 session 聚类 reward、估计 advantage、单遍扫描推导 batch 级 flag，
  │    对 write_train_meta 字段（supervised_tokens / train_prompt_length / num_tokens）fail-fast
  ▼ ② _build_pack_plan
  │    RLDataPacker.pack(num_tokens) 生成 [dp][optimizer_step][pack][sample_index] 索引计划，
  │    按 DP 切分并把全局索引 remap 成本地索引（不 ray.put、不取张量），顺带产出 packing/* 指标
  ▼ ③ dispatch
  │    每个 DP rank ray.put 一次（data_replicate_size 个 replica 共享同一份序列化），
  │    worker.fit.remote(...)；finally 中释放 pixel_values object refs
  ▼ ④ _build_trainer_log_info
      合并数据统计 + packing/* + worker 耗时 max → data_info: dict[str, float]

各阶段输入输出：

阶段	输入	输出
_prepare_rollout_items	rollout_groups: list[list[RolloutState]]（只读 write_train_meta）	_PreparedBatch：{rollout_items, advantages, batch_attr, cluster_rewards, distillation_reward_observations}
_build_pack_plan	rollout_items（只读 num_tokens）、advantages、worker_cfg、data_replicate_size	_PackPlan：{dp_dispatches: dict[dp_rank, _DPRankDispatch], plan_log: dict[str, float]}
dispatch	dp_dispatches[dp_rank]	worker.fit.remote(rollout_item_refs, advantages, pack_plan, batch_attr)
_build_trainer_log_info	_PreparedBatch + _PackPlan + log_infos: list[WorkerLogItem]	data_info: dict[str, float]（packing/*、advantages/*、rewards/*、worker 耗时 max）

跨 Ray 边界的 worker 入参契约：

• rollout_item_refs: list[ray.ObjectRef]：第 0 个元素是本 DP rank 的 list[RolloutState]，由 data_replicate_size（= tp_size × sp_size）个 replica 共享，每 batch 只序列化一次
• advantages: list[float]：与解引用后的 state 列表逐一对齐
• pack_plan: DPRankPackIndices：本地 [optimizer_step][pack][sample_index] 索引计划，索引寻址解引用后的 state 列表
• batch_attr: TrainBatchAttr：{rollout_idx, use_3d_position_ids, pack_loss_keys, has_routed_experts}；batch 内所有 worker 看到同一份，保证跨 rank collective 形状对称

TrainingWorker.fit pipeline

fit(rollout_item_refs, advantages, pack_plan, batch_attr) -> WorkerLogItem
  │
  ▼ ① convert（纯 CPU）
  │    ray.get 解引用 → _convert_one_rollout_state：
  │    shift（input_ids[:-1] / labels[1:] / logprobs[1:]）、逐位置 advantage 广播（-100 位为 0）、
  │    3D MRoPE 补齐、get_train_seq_ctx、routed_experts 挂载、distillation teacher targets；
  │    loss_cfg.build(data, device="cpu") 构建 per-item loss_ctx
  ▼ ② pack（纯 CPU）
  │    _pack_train_items 按 plan 取 item；_pack_one_batch 对每个 pack：
  │    seq_ctx 沿序列维 cat + padding 到 pack_max_length（短 chunk padding / 3D position /
  │    dummy routed experts）；cat loss_kwargs 张量，模板键 all-present/all-absent 校验 +
  │    per-key padding 值，重建 pack 级 loss_ctx（device="cpu"）
  ▼ ③ _fit（上卡训练）
       seq_ctx.to(DEVICE) + sp_split；loss_kwargs.to(DEVICE) → MTP ctx（用 pre-split 的
       device labels）→ loss_kwargs.sp_split；old_logprobs / entropy / mismatch & IS metrics /
       ref KL → LossContext.build_batches → engine.train_step（含 SFT 混训）

各阶段输入输出：

阶段	输入	输出
convert	list[RolloutState] + advantages + batch_attr	list[(seq_ctx: SequenceContext, loss_ctx: BaseRLLossContext)]，全部 CPU
pack	上一步输出 + pack_plan + batch_attr	list[list[(seq_ctx, loss_ctx)]]，外层下标 = optimizer step，内层 = pack（空 pack 物化为全 padding pack）
_fit	上一步输出 + rollout_idx	WorkerLogItem（entropy、mismatch/IS metrics、train_metrics、packing_conversion_time_s / packing_pack_time_s）

生命周期约定：loss_ctx 与 seq_ctx 同步流转 —— convert / pack 阶段都在 CPU 上构建，只有 _fit 统一做 .to(DEVICE) 与 sp_split（sp_mesh.size() > 1 时）；MTP ctx 必须用 sp 切分前的 device labels + 切分后的 seq_ctx 构建。

其他改动

• BaseRLLossConfig.build / DistillationLossConfig.build 新增 device 参数（默认行为不变），支撑 CPU 侧构建 loss_ctx。
• pack 策略的 shuffle 决策（XTUNER_DETERMINISTIC gate）下沉到 _legacy_pack / _greedy_pack 各自实现内，RLDataPacker.pack() 入口不再感知策略细节。
• 删除 controller 侧 legacy packing 路径，以及 RLTrainItem / RLLossInputs / data.py 中间 proto；数据 proto 只剩 worker 公开的 TrainBatchAttr 和 controller 私有的 _DPRankDispatch / _PackPlan / _PreparedBatch。

测试

• tests/rl/test_pack.py：四种 pack 策略与 XTUNER_DETERMINISTIC gate 行为
• tests/rl/test_rollout_to_train_item.py：convert 阶段布局/对齐/advantage 广播/teacher targets contract
• tests/rl/test_training_worker_rank.py：pack 阶段拼接、padding 值表、模板键、空 pack 与 CPU loss_ctx contract
• tests/rl/test_prepare_train_data.py：controller 侧 meta-ready 契约（prepare / pack plan / trainer log info）

[jayhenry]

@claude review

[github-actions]

Claude:

Summary

本 PR 将 RL 训练的 pack 规划（RLDataPacker，仅产出索引）与实际 Tensor 拼接（TrainingWorker 本地 materialize）解耦，并新增 pack 耗时、训练耗时与 padding token 统计。整体分层方向合理，但 routed-experts 的全 padding pack 路径存在崩溃缺陷，且有两处存量测试因接口改名/校验收紧而必然失败。

ProduceBatchResult impact: not affected —— fit 仅消费 data_batches，新增计数写入 trainer 侧 data_info / step_timer_dict，不触及 ProduceBatchResult 任何字段。

RoutedExperts impact: affected —— padding 侧 routed-experts 构造从 controller 移入 TrainingWorker._create_padding_item，逐样本 ref 的 ownership 与释放路径不变，但全 padding 分支的 dummy tensor 在新的 SequenceContext.cat 路径下会触发 shape 断言失败（见 Critical 第 1 条）。

Ray concurrency impact: not affected —— TrainingWorker 未声明任何 concurrency group，新增的两个 @ray_method getter 只在 TrainingController.__init__ 调用一次。

Main Flowchart after this PR

flowchart TD
    A[BaseRLTrainer._prepare_train_data<br/>产出 WorkerInputItem<br/>advantages 改为 Tensor] --> B[TrainingController.fit]
    B --> C[提取 data_lengths]
    C --> D[RLDataPacker.pack<br/>仅索引规划 + padding 统计]
    D --> E[ray.put 整个 data_batches<br/>嵌套 ObjectRef 下发]
    E --> F[TrainingWorker.fit<br/>ray.get 全量 batch]
    F --> G[_materialize_packs<br/>按索引选样本]
    G --> H[_single_pack<br/>SequenceContext.cat + 拼接]
    H --> I[_create_padding_item<br/>全 padding 分支 routed-experts 出错]
    I --> J[按 packed_batch_num_per_step 执行 optimizer step]
    B --> K[返回 TrainingLogInfo<br/>pack_time / train_time / padding_tokens]

    style D fill:#cde4ff
    style G fill:#cde4ff
    style E fill:#ffe4b5
    style I fill:#ffb3b3

Loading

核心原理实现与单测

• RLDataPacker.pack 输出 [dp][optimizer_step][pack][sample_index] 索引规划与 padding token 数，并在 _count_padding_tokens 中强校验“每个样本索引恰好出现一次”，规划正确性有真实不变量兜底。
• tests/rl/test_pack.py 通过 public pack() 覆盖 greedy / balance / native 三种策略，断言索引全覆盖、单 pack 不超长、padding 账目自洽、balance 的 rank 间偏斜上界，并含 1024 样本随机 fuzz，属于真实代码路径覆盖。
• Materialize 侧（SequenceContext.cat、label/advantage/logprob 拼接、padding tensor、VLM position ids）落在 TrainingWorker._materialize_packs / _single_pack / _create_padding_item，新增两条断言（总长等于 pack_max_length；num_padding 等于 advantages 中 -100 计数），文本路径与空 pack 场景已被 TestTrainingWorkerPackMaterialization 覆盖。
• get_dp_rank 在 tests/rl/test_training_worker_rank.py 中按 tp/sp 组合参数化覆盖。
• 行为变更值得注意：advantages 现在与 input_ids 等长（response 路径改用 actual_advantages[:-1]），修正了此前每样本多一个元素的展平偏差，并由 _single_pack 的断言固化。

抽象与信息隐藏评估

• Warning xtuner/v1/rl/trainer/pack.py:79-96 RLDataPacker 把三个策略方法与 strategy_map 暴露为公开 Interface，而调用者只需要 pack()，Interface 几乎等于 Implementation，建议改为私有以加深 Module。
• Warning xtuner/v1/rl/trainer/worker.py:601-619 _set_pack_data_properties 在 __init__ 之外写入 _pack_* 实例属性，与 _single_pack / _create_padding_item 形成隐式调用顺序耦合，正是 routed-experts 缺陷得以藏身之处，建议改为显式传递不可变的 pack spec。

单测建议

• Warning tests/rl/test_pack.py:88-115 materialize 的单测只覆盖 model_cfg=None 的纯文本路径，rollout_routed_experts 与 3D position_ids（qwen3-vl）分支完全无覆盖，导致下方 Critical 缺陷无法被发现。
• Warning xtuner/v1/rl/trainer/worker.py:722-726 controller 下发的 packed_data_indices[dp_rank] 步数与 worker 侧 optimizer_steps 断言之间的契约没有直接测试，而 balance 恒定产出 optimizer_steps 步、greedy/native 可能更少，值得补一条契约测试。

其他 Issues

Critical

• xtuner/v1/rl/trainer/worker.py:687-690 全 padding pack 沿用 size=(1,1,1) 的 dummy routed-experts，但新路径必经 SequenceContext.cat 后走 list 分支，会触发 rollout_routed_experts.size(0) != input_ids.size(1) 断言，使开启 routed-experts 的 MoE RL 训练崩溃。
• tests/rl/test_rl_colocate_trainer_integration.py:280 该测试的 train_worker_cfg 使用 pack_max_length=2048，但仍以 pack_max_length=1024 调用 fit，会被 controller 新增的一致性校验直接抛 ValueError。
• tests/rl/test_prepare_train_data.py:105 _prepare_train_data 已将 advantage（list）改名为 advantages（Tensor），但该测试文件未同步更新，断言会 KeyError，且文件开头声明的“advantage 比 shifted_labels 多 1 个元素”契约已被本 PR 有意推翻。

Warning

• xtuner/v1/rl/trainer/controller.py:92-101 改为单次 ray.put 全量 batch 后每个 worker 都 ray.get 并反序列化整批数据，相比原先的 dp 分片下发，单 worker CPU 峰值内存约放大 dp_size 倍。

Verdict

REQUEST_CHANGES

[claude]

Claude: [正确性] [复杂] 全 padding pack（pack_indices 为空，由 _align_pack_count 与 greedy 的 total_pack_indices.extend([[] ...]) 产生）仍沿用旧的 size=(1,1,1) dummy。

旧代码里这个 dummy 就是该 item 自身的 seq_ctx，_add_rollout_routed_experts 走 tensor 分支并扩展到 pack_max_length；现在 padding item 一定会经过 SequenceContext.cat（worker.py:640），而 cat 把 rollout_routed_experts 收集成 list（sequence_context.py:356-357），于是 fit 走 list 分支只生成 size=(1, num_hidden_layers, num_experts_per_tok)，随后 worker.py:559-561 的 rollout_routed_experts.size(0) == input_ids.size(1) 断言必然失败，开启 routed-experts 的 MoE RL 训练会崩溃。

建议去掉该特判，统一按 pad_len 构造：

Suggested change

	if pad_len == self.config.pack_max_length:
	pad_rand_index = torch.randint(low=0, high=1, size=(1, 1, 1))
	else:
	pad_rand_index = torch.randint(low=0, high=self._pack_n_routed_experts, size=(pad_len, 1, 1))
	pad_rand_index = torch.randint(low=0, high=self._pack_n_routed_experts, size=(pad_len, 1, 1))

RoutedExperts impact: 全 padding pack 的 routed-experts 张量长度与 input_ids 不一致，训练直接中断。

[claude]

Claude: [测试] 该测试的 train_worker_cfg 用的是 pack_max_length=2048（本文件 L135），但 L280/286/316 仍传 pack_max_length=1024，会被 controller 新增的一致性校验（controller.py:81-84）直接抛 ValueError，测试必然失败。