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

[YanhuiDua]

概述

本 PR 为 RL 训练引入了可配置的数据 Pack 能力，并将 Pack 规划与实际 Tensor 拼接进行了解耦。RLDataPacker 现在只生成如下结构的索引规划：[dp_rank][optimizer_step][pack][sample_index]。实际的数据获取、Tensor 拼接和 Padding 操作由各个 TrainingWorker 在本地完成。

主要改动

• 新增 RLDataPacker，支持三种 Pack 策略：
  • greedy
  • balance
  • native
• RLDataPacker 只负责：
  • 根据样本长度生成 Pack 索引
  • 将样本分配到不同 DP Rank
  • 按 Optimizer Step 对 Pack 进行划分
  • 统计 Padding Token 数量
• 将实际 Pack 逻辑移动到 TrainingWorker：
  • 根据索引获取原始样本
  • 执行 SequenceContext.cat
  • 拼接 Labels、Advantages 和 Rollout Logprobs
  • 创建 Padding Tensor
  • 补齐 VLM Position IDs
  • 补齐 Routed Experts 数据
• 原始训练 Batch 只写入 Ray Object Store 一次，并通过嵌套的 ObjectRef 共享给各个 Training Worker。
• TrainingController.fit 新增以下统计信息：
  • Pack 耗时
  • 训练耗时
  • Padding Token 数量

数据流程

BaseRLTrainer
│
├─ 生成 WorkerInputItem 列表
│
▼
TrainingController
├─ 提取样本长度
├─ 生成仅包含索引的 Pack Plan
└─ 将原始 Batch 写入 Ray Object Store
│
▼
TrainingWorker
├─ 获取原始 Batch
├─ 根据分配到的索引选择样本
├─ 执行实际 Pack 和 Padding
└─ 按规划执行 Optimizer Step

[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，测试必然失败。