DataProto：训练流水线里的 batch 信封

verl/protocol.py 里的 DataProto 是 verl trainer/data 工程的核心抽象。你可以把它理解成一个“能被切分、拼接、搬运、合并的 batch 信封”。它既装 tensor，也装 Python 对象，还装控制信息，因此能在 trainer、Ray worker、rollout server、reward manager 之间传递。

先验知识

读 DataProto 前，需要先知道：

LLM RL 的 batch 不只有 tensor。reward 可能要看原始答案、数据来源、工具调用、图片输入、评测函数名，这些不能自然塞进普通 torch tensor。
分布式训练需要把一个全局 batch 切成多个 worker shard，worker 返回后再拼回去。
Ray 远端调用返回的是 object ref。driver 不一定应该立刻 ray.get()，有时需要把未来结果继续传给下一个 worker。
TensorDict 是 PyTorch 生态里的 dict-like tensor 容器。它能保证多个 tensor 共享同一个 batch 维度，并支持整体切片、拼接、搬设备。

本页原先不适合小白的地方

原说明已经讲了 batch、non_tensor_batch、meta_info 三块，但还缺几层学习支架：

没有明确解释为什么 tensor 和 non-tensor 必须分开。
没有把 TensorDict 的 batch 维度讲成“第 0 维是样本维”。
DataProtoFuture 和 Ray object refs 容易被看成高级细节，但它其实解释了为什么 worker 调用可以异步衔接。
split、chunk、concat、union、reorder 的用途相似，初学者容易混。需要把它们绑定到 trainer 里的具体场景。

三个区域

区域	类型	放什么	例子
`batch`	`TensorDict`	所有第 0 维是样本维的 tensor	`prompts`、`responses`、`attention_mask`、`old_log_probs`、`advantages`
`non_tensor_batch`	`dict[str, np.ndarray]`	每条样本对应的非 tensor 对象	`uid`、`data_source`、`reward_model`、`extra_info`、`multi_modal_inputs`
`meta_info`	`dict`	整个 batch 共享的控制信息	`temperature`、`global_steps`、`timing`、`metrics`

DataProto.check_consistency() 会要求 batch.batch_size[0] 和每个 non_tensor_batch[key].shape[0] 对齐。也就是说，如果 tensor batch 有 512 条 response，uid 也必须有 512 个。

TensorDict 是什么

DataProto.batch 不是裸 dict，而是 TensorDict：

python

TensorDict(
    {
        "prompts": Tensor[batch, prompt_len],
        "responses": Tensor[batch, response_len],
        "attention_mask": Tensor[batch, prompt_len + response_len],
    },
    batch_size=(batch,),
)

学习时先抓一个规则：第 0 维是样本维，其他维才是 token 维、特征维或模型内部维度。

所以：

batch[0] 是第 0 条样本。
batch[:8] 是前 8 条样本。
DataProto.chunk(8) 是沿样本维切给 8 个 worker。
DataProto.concat([...]) 是把多个 worker 的样本维拼回来。

`non_tensor_batch` 为什么存在

reward 和多模态数据经常不是纯 tensor：

reward_model 里可能有 ground truth、style tag、test cases。
data_source 决定用哪个 reward function。
extra_info 可能保存题目 id、原始 prompt、工具配置。
multi_modal_inputs 可能是 processor 产生的对象字典。
uid 是 group-based advantage 的分组 id。

这些对象需要按样本同步切分、repeat、reorder、concat。DataProto 让 tensor 和 non-tensor 一起移动，避免“tensor 顺序变了，metadata 没变”的灾难。

`meta_info` 放什么

meta_info 不是每条样本一份，而是整个 batch 或这次调用共享的信息。

常见例子：

rollout 前写入 global_steps、temperature。
rollout 输出带回 meta_info["timing"]，trainer 取出后写进 metrics。
worker 返回的训练指标放在 meta_info["metrics"]。
DataProto.concat() 对 metrics 有特殊聚合逻辑，避免多个 worker 的指标互相覆盖。

构造和转换

最常见入口：

python

batch = DataProto.from_single_dict(batch_dict)

from_single_dict() 会把 torch tensor 放入 batch，把 numpy array 放入 non_tensor_batch。如果传入普通 Python list，需要先变成 np.ndarray，或者走 from_dict() 的 non_tensors 逻辑。

worker 调用前常见转换：

python

batch_td = batch.to_tensordict()

to_tensordict() 会把 non_tensor_batch 转成 NonTensorStack，把 meta_info 转成 non-tensor data，方便 unified worker 只收一个 TensorDict。

`union()`：同一批样本的横向合并

union() 用于“同一批样本，新增字段”：

python

batch = batch.union(gen_batch_output)
batch = batch.union(old_log_prob)
batch = batch.union(ref_log_prob)
batch = batch.union(values)

它的要求是两个 DataProto 的 batch size 一致。它不会纵向增加样本数，只会增加字段。

学习用途：

rollout 输出和原 prompt batch 合并。
actor logprob、ref logprob、critic value 逐步贴回同一个训练样本。
reward model 输出贴回 batch。

如果两个输入里有同名 key，源码会检查值是否一致。这能尽早发现两个模块写出了冲突字段。

`repeat()`：为每个 prompt 扩展多条 response

GRPO、RLOO、best-of-n 类训练常需要一个 prompt 采样多条 response：

python

gen_batch = gen_batch.repeat(repeat_times=rollout.n, interleave=True)

如果 uid=[A, B]，repeat(3, interleave=True) 后是：

text

A, A, A, B, B, B

这让后续 compute_grpo_outcome_advantage(index=uid) 能知道哪些 response 属于同一个 prompt group。

`chunk()` / `split()`：纵向切 batch

chunk(chunks=N) 用于 dispatch 场景：

python

parts = data.chunk(chunks=worker_group.world_size)

默认要求能等分。如果打开 auto padding，会先补齐再切，collect 后再去掉 padding。

split(split_size) 是按固定大小切：

python

micro_batches = data.split(split_size=64)

学习区别：

chunk() 更像“切给 N 个 worker”。
split() 更像“按每块大小切成若干段”。

`concat()`：纵向拼 worker 结果

worker group collect 时常用：

python

output = DataProto.concat(worker_outputs)

它把多个 DataProto 沿样本维拼接。batch 用 torch.cat，non_tensor_batch 用 np.concatenate，meta_info["metrics"] 会被聚合成 list-of-values 结构。

学习用途：

多个 Ray worker 各自返回 shard，driver 拼回全局 batch。
validation 需要把多个 padded/unpadded 结果拼起来。
REMAX 路径把 sampled rollout 和 greedy baseline 合成一次请求。

`select()` / `pop()`：选择字段

select() 返回某些字段的视图或拷贝：

python

sub = batch.select(batch_keys=["responses", "attention_mask"])

pop() 会从原对象移除字段，并返回被移除字段组成的新 DataProto：

python

gen_batch = batch.pop(non_tensor_batch_keys=[...])

trainer 的 _get_gen_batch() 用 pop() 控制哪些 non-tensor 字段送去 rollout，哪些字段保留在训练 batch 上。

`reorder()`：保持 tensor 和 metadata 同步重排

_balance_batch() 会按 token 工作量重排样本：

python

batch.reorder(global_idx)

reorder() 同时重排 batch 和 non_tensor_batch。如果只重排 tensor，不重排 uid、reward_model、multi_modal_inputs，reward 和 advantage 会对错样本。

`DataProtoFuture`：Ray object refs 的延迟信封

DataProtoFuture 解决的是“远端结果先不拉回 driver”的问题。它保存：

futures: list[ray.ObjectRef]
collect_fn
可选 dispatch_fn

当 .get() 被调用时，它才 ray.get(self.futures)，然后 DataProto.concat(output)。

学习用途：

理解 @register(blocking=False) 为什么能返回未来结果。
理解异步 worker 调用不是立刻 materialize 到 driver。
理解 BatchData.concat() 看到 ray.ObjectRef 时会返回 DataProtoFuture.concat(data)。

`BatchData`：统一处理 DataProto、TensorDict 和 transfer queue metadata

single_controller 的 dispatch/collect 不想到处写：

python

if isinstance(x, DataProto): ...
elif isinstance(x, TensorDict): ...
elif isinstance(x, KVBatchMeta): ...

所以 BatchData 统一提供：

BatchData(arg).chunk(chunks=N)
BatchData(output_list).concat()

这就是为什么 worker 方法可以有的收 DataProto，有的收 TensorDict，有的收 Ray object refs，但调度层仍然能复用同一套切分/拼接逻辑。

读 trainer 时的 DataProto 心智图

python

batch = DataProto.from_single_dict(batch_dict)
batch.non_tensor_batch["uid"] = make_uid(batch)

gen_batch = _get_gen_batch(batch)
gen_batch = gen_batch.repeat(rollout.n)
gen_output = rollout.generate(gen_batch)

batch = batch.repeat(rollout.n).union(gen_output)
batch.batch["response_mask"] = compute_response_mask(batch)

batch = batch.union(actor_old_logprob(batch))
batch = batch.union(ref_logprob(batch))
batch = batch.union(critic_values(batch))

batch.batch["token_level_rewards"] = reward_or_reward_minus_kl(batch)
batch = compute_advantage(batch)
actor.update_actor(batch)

本节参考与延伸阅读

源码：verl/protocol.py，重点读 DataProto、DataProtoFuture、BatchData、union_tensor_dict()、union_numpy_dict()。
源码：verl/single_controller/base/decorator.py，看 dispatch_dp_compute_data_proto()、collect_dp_compute_data_proto() 如何调用 BatchData。
源码：verl/trainer/ppo/ray_trainer.py，搜索 DataProto.from_single_dict、repeat、union、reorder。
官方 docs：docs/api/data.rst、docs/single_controller.rst。
外部资料：PyTorch TensorDict 文档；HybridFlow: A Flexible and Efficient RLHF Framework, arXiv:2409.19256。

DataProto：训练流水线里的 batch 信封 ​

先验知识 ​

本页原先不适合小白的地方 ​

三个区域 ​

TensorDict 是什么 ​

non_tensor_batch 为什么存在 ​

meta_info 放什么 ​

构造和转换 ​

union()：同一批样本的横向合并 ​

repeat()：为每个 prompt 扩展多条 response ​

chunk() / split()：纵向切 batch ​

concat()：纵向拼 worker 结果 ​

select() / pop()：选择字段 ​

reorder()：保持 tensor 和 metadata 同步重排 ​

DataProtoFuture：Ray object refs 的延迟信封 ​

BatchData：统一处理 DataProto、TensorDict 和 transfer queue metadata ​

读 trainer 时的 DataProto 心智图 ​

本节参考与延伸阅读 ​

DataProto：训练流水线里的 batch 信封

先验知识

本页原先不适合小白的地方

三个区域

TensorDict 是什么

`non_tensor_batch` 为什么存在

`meta_info` 放什么

构造和转换

`union()`：同一批样本的横向合并

`repeat()`：为每个 prompt 扩展多条 response

`chunk()` / `split()`：纵向切 batch

`concat()`：纵向拼 worker 结果

`select()` / `pop()`：选择字段

`reorder()`：保持 tensor 和 metadata 同步重排

`DataProtoFuture`：Ray object refs 的延迟信封

`BatchData`：统一处理 DataProto、TensorDict 和 transfer queue metadata

读 trainer 时的 DataProto 心智图

本节参考与延伸阅读