verl数据准备全流程:RLHFDataset使用详解
在大型语言模型(LLM)的强化学习后训练中,高质量、结构清晰、格式统一的数据是训练稳定性和效果上限的关键前提。verl 作为专为 LLM 后训练设计的高效 RL 框架,其数据处理流程并非简单加载 JSONL 文件,而是一套兼顾工程鲁棒性、模板一致性与 token 级控制的端到端流水线。其中,RLHFDataset是整个数据准备环节的核心载体——它不只负责读取,更承担了提示构造、分词对齐、长度裁剪、填充规整等关键预处理任务。
本文将完全聚焦于RLHFDataset的实际使用,不讲抽象原理,不堆砌 API 列表,而是以一个真实可复现的端到端流程为线索,带你从原始对话数据出发,一步步构建出 verl 可直接消费的训练数据集。你会看到:如何组织数据文件、怎样配置 tokenizer 和模板、哪些参数真正影响训练质量、常见报错如何定位、以及为什么某些“看似合理”的写法反而会导致训练崩溃。所有内容均基于 verl 官方实现和生产级实践提炼,目标明确:让你第一次调用RLHFDataset就能成功跑通,第二次就能调优出好结果。
1. 数据准备的本质:不是“读文件”,而是“建协议”
在 verl 中,数据准备远不止torch.utils.data.Dataset的简单继承。它的核心目标是建立一个跨角色、跨设备、跨阶段的 token 级数据契约。这意味着:
- Actor rollout 生成的 response 必须与 Critic 评估的 value 序列严格对齐;
- Reference policy 计算的 log_prob 必须与 reward model 打分的 token 位置一一对应;
- 所有角色(Actor、Critic、Ref、RM)接收到的
input_ids、attention_mask、position_ids必须来自同一份预处理逻辑,否则后续的 KL 散度计算、优势函数估计都会失效。
RLHFDataset正是这个契约的落地执行者。它不输出原始字符串,而是输出一个结构化的DataProto对象(或其底层 dict 表示),其中每个字段都已按 RLHF 流程要求完成标准化。
1.1 输入数据格式:Parquet 是唯一推荐格式
verl 明确推荐使用 Parquet 文件作为数据源,而非常见的 JSONL 或 CSV。原因很实际:
- 列式存储:可按需读取
prompt、chosen、rejected等特定列,避免加载整行文本; - 类型安全:Schema 强制约束字段类型(如
prompt: string,chosen: string),防止空值或类型错乱; - 高效压缩:相比纯文本,体积通常减少 60% 以上,IO 压力显著降低;
- 分布式友好:Spark、Dask、Polars 等均可原生读取,便于大规模数据清洗。
一个标准的 RLHF 训练数据 Parquet 文件应至少包含以下三列(DPO/GRPO 场景)或两列(PPO 场景):
| 字段名 | 类型 | 说明 |
|---|---|---|
prompt | string | 用户输入的指令或问题,不含任何系统提示词 |
chosen | string | 期望模型生成的优质回复(正样本) |
rejected | string | 质量较差或存在事实错误的回复(负样本,DPO/GRPO 必需) |
注意:
prompt字段必须是纯文本指令,不能包含<|user|>、[INST]等模板标记;这些由RLHFDataset在内部动态注入。
1.2 创建最小可行 Parquet 文件(实操)
我们用pyarrow快速生成一个仅含 3 条样本的测试文件,用于后续验证:
import pyarrow as pa import pyarrow.parquet as pq # 构造原始数据(模拟人工标注) data = { "prompt": [ "请用一句话解释量子纠缠。", "写一首关于春天的七言绝句。", "对比 Python 和 Rust 在 Web 后端开发中的适用场景。" ], "chosen": [ "量子纠缠是指两个或多个粒子形成一种关联状态,即使相隔遥远,测量其中一个粒子的状态会瞬间决定另一个的状态。", "春山暖日和风,阑干楼阁帘栊。杨柳秋千院中,啼莺舞燕,小桥流水飞红。", "Python 适合快速迭代、生态丰富、适合中小型服务;Rust 适合高并发、低延迟、内存安全要求严苛的核心服务,但开发成本更高。" ], "rejected": [ "量子纠缠就是粒子之间有神秘联系,科学家也不太懂。", "春天来了,花开了,鸟叫了,真美。", "Python 和 Rust 都是编程语言,都能写代码。" ] } # 写入 Parquet table = pa.table(data) pq.write_table(table, "sample_rlhf_data.parquet") print(" Parquet 文件已生成:sample_rlhf_data.parquet")运行后,你将得到一个约 2KB 的sample_rlhf_data.parquet文件。这就是RLHFDataset的起点。
2. 初始化 RLHFDataset:4 个必传参数与 3 个关键配置项
RLHFDataset的初始化签名如下(精简版):
from verl.data import RLHFDataset dataset = RLHFDataset( data_files="sample_rlhf_data.parquet", # 必传:路径或路径列表 tokenizer=tokenizer, # 必传:HuggingFace Tokenizer 实例 config=data_config # 必传:OmegaConf 配置对象 )表面看只有三个参数,但config内部隐藏着决定数据质量的“开关”。下面逐一分解。
2.1 必传参数一:data_files
支持单文件、文件列表、glob 模式:
# 单文件 data_files = "data/train.parquet" # 多文件(推荐:便于分片) data_files = ["data/part_000.parquet", "data/part_001.parquet"] # Glob 模式(自动匹配) data_files = "data/train_*.parquet"重要提醒:路径必须是本地绝对路径或可被 PyArrow 直接读取的 URI(如
s3://bucket/path/)。相对路径在分布式环境下极易出错。
2.2 必传参数二:tokenizer
必须是已加载的 HuggingFacePreTrainedTokenizerBase实例,且需满足:
- 已调用
tokenizer.pad_token_id和tokenizer.eos_token_id被正确设置; - 若模型无
pad_token,需手动添加:tokenizer.add_special_tokens({"pad_token": "[PAD]"}); - 推荐使用与训练模型完全一致的 tokenizer,避免 vocab mismatch。
from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf") if tokenizer.pad_token is None: tokenizer.add_special_tokens({"pad_token": "[PAD]"}) # 注意:添加 special tokens 后,务必 resize model embedding2.3 必传参数三:config—— 数据行为的总控台
这是最易被忽略、却最关键的参数。config是一个OmegaConf对象,必须包含以下子字段:
| 配置项 | 类型 | 必填 | 说明 | 典型值 |
|---|---|---|---|---|
max_prompt_length | int | prompt 最大 token 数(截断阈值) | 512 | |
max_response_length | int | response 最大 token 数(含 EOS) | 512 | |
chat_template | str | ❌(但强烈建议) | HuggingFace 格式聊天模板 | "llama-2"或自定义 Jinja2 模板 |
padding_side | str | ❌(默认right) | 填充方向(影响 attention mask) | "left"(对 PPO rollout 更友好) |
一个最小可用的data_config示例:
from omegaconf import OmegaConf data_config = OmegaConf.create({ "max_prompt_length": 512, "max_response_length": 512, "chat_template": "llama-2", # 自动加载 transformers 内置模板 "padding_side": "right" })
chat_template的作用:RLHFDataset会将prompt+chosen/rejected按照指定模板拼接成完整对话字符串,再进行分词。例如llama-2模板会生成:[INST] 请用一句话解释量子纠缠。 [/INST] 量子纠缠是指...
3. 数据处理全流程解析:从字符串到 DataProto 的 5 个关键步骤
当你调用dataset[i]时,RLHFDataset.__getitem__内部会依次执行以下操作。理解每一步,是调试数据问题的基石。
3.1 步骤一:读取原始行并提取字段
从 Parquet 中读取第i行,提取prompt、chosen、rejected字段。若字段为空或非字符串,立即抛出ValueError。
3.2 步骤二:应用聊天模板(Chat Template Application)
这是RLHFDataset区别于普通 Dataset 的核心。它调用tokenizer.apply_chat_template(),将原始字段组装为标准对话格式:
# 伪代码示意 messages = [ {"role": "user", "content": row["prompt"]}, {"role": "assistant", "content": row["chosen"]} ] formatted_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) # 输出:"[INST] 请用一句话解释量子纠缠。 [/INST] 量子纠缠是指..."优势:确保所有角色看到的输入格式完全一致,避免因手写模板导致的 token 错位。
3.3 步骤三:分词与截断(Tokenization & Truncation)
对formatted_text进行分词,并严格按max_prompt_length和max_response_length截断:
input_ids:完整对话的 token ID 序列;attention_mask:对应位置的掩码(1=有效,0=padding);position_ids:从 0 开始递增的位置索引(对 RoPE 模型至关重要)。
关键逻辑:
- 若总长度 ≤
max_prompt_length + max_response_length,保留全部; - 若超长,则优先截断 prompt 部分,保证 response 至少保留
max_response_length个 token(含 EOS)。
3.4 步骤四:填充(Padding)与对齐
根据padding_side进行填充,使所有样本长度统一为max_prompt_length + max_response_length:
padding_side="right":在序列末尾补pad_token_id;padding_side="left":在序列开头补pad_token_id(对 actor rollout 生成更友好,因 KV cache 可复用)。
同时生成label字段:将input_ids中 prompt 部分设为-100(忽略 loss),response 部分保留原值,供后续 loss 计算。
3.5 步骤五:构建 DataProto 兼容字典
最终返回一个 dict,结构如下(已简化):
{ "input_ids": [1, 2, 3, ..., 0, 0], # shape: (seq_len,) "attention_mask": [1, 1, 1, ..., 0, 0], # shape: (seq_len,) "position_ids": [0, 1, 2, ..., n, n+1], # shape: (seq_len,) "labels": [-100, -100, ..., 123, 456], # shape: (seq_len,), prompt 位置为 -100 "prompt_lengths": [23], # list[int], 每个样本 prompt token 数 }该 dict 可直接传入DataLoader,并在训练循环中被封装为DataProto。
4. 常见问题排查指南:5 个高频报错及解决方案
4.1 报错:KeyError: 'prompt'
现象:初始化RLHFDataset时抛出KeyError: 'prompt'
原因:Parquet 文件中不存在prompt列,或列名大小写不匹配(如Prompt、PROMPT)
解决:
- 用
pq.read_table("file.parquet").schema查看实际列名; - 确保列名严格为小写
prompt、chosen、rejected; - 如需重命名,用 PyArrow 修改:
table = pq.read_table("old.parquet") table = table.rename_columns(["prompt", "chosen", "rejected"]) pq.write_table(table, "new.parquet")
4.2 报错:ValueError: Input is not valid for the chosen chat template
现象:apply_chat_template失败,提示模板不兼容
原因:chat_template名称错误,或 tokenizer 未注册该模板
解决:
- 检查
tokenizer.chat_template是否为None; - 使用
tokenizer.init_kwargs.get("chat_template")查看内置模板; - 显式指定完整模板字符串(Jinja2):
tokenizer.chat_template = "{% if messages[0]['role'] == 'system' %}{{ messages[0]['content'] }}{% endif %}{% for message in messages %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token }}{% endif %}{% endfor %}"
4.3 报错:RuntimeError: expected scalar type Long but found Int
现象:DataLoader 返回 batch 后,在模型 forward 时报 tensor 类型错误
原因:RLHFDataset默认返回np.int64,而 PyTorch 期望torch.long
解决:在DataLoader中启用collate_fn自动转换:
from torch.utils.data import DataLoader from verl.data.collator import RLHFCollator collator = RLHFCollator(tokenizer.pad_token_id, return_tensors="pt") dataloader = DataLoader(dataset, batch_size=8, collate_fn=collator)4.4 现象:训练 loss 为 NaN 或剧烈震荡
可能原因:max_prompt_length设置过大,导致大量 padding,attention mask 未正确屏蔽
验证方法:打印一个 batch 的attention_mask:
for batch in dataloader: print("Attention mask sum:", batch["attention_mask"].sum().item()) break若sum远小于batch_size * seq_len,说明 padding 过多。
解决:调小max_prompt_length/max_response_length,或改用padding_side="left"。
4.5 现象:Actor rollout 生成内容与 prompt 不匹配
原因:RLHFDataset生成的input_ids包含完整对话(user+assistant),但 rollout 时只应输入 user 部分
解决:在 PPO 训练循环中,务必使用batch.pop(...)分离:
# 正确:只将 prompt 部分送入 rollout gen_batch = batch.pop(batch_keys=['input_ids', 'attention_mask', 'position_ids']) gen_output = actor_rollout_wg.generate_sequences(gen_batch) # gen_batch 只含 prompt5. 进阶技巧:3 种提升数据质量的实战方法
5.1 方法一:动态长度控制(Per-sample Length)
RLHFDataset支持为每个样本单独指定长度,避免一刀切截断损失信息。只需在 Parquet 中增加prompt_length和response_length列:
# Parquet 新增列 data["prompt_length"] = [128, 256, 192] data["response_length"] = [384, 256, 320]然后在data_config中启用:
data_config.dynamic_length = True # 启用动态长度5.2 方法二:多轮对话支持(Multi-turn Chat)
RLHFDataset原生支持多轮对话。只需将prompt和chosen字段改为消息列表:
# Parquet 中存为 JSON 字符串 data["prompt"] = [ '[{"role": "user", "content": "你好"}, {"role": "assistant", "content": "你好!"}]', ... ] data["chosen"] = ['[{"role": "user", "content": "今天天气如何?"}]']RLHFDataset会自动解析 JSON 并应用模板。
5.3 方法三:混合数据源(Mixing Datasets)
训练时常需混合 SFT、DPO、KTO 等多种格式数据。RLHFDataset支持通过data_type字段自动路由:
# Parquet 新增列 data["data_type"] = ["sft", "dpo", "kto"]在data_config中配置各类型权重:
data_config.mixing_weights = {"sft": 0.3, "dpo": 0.5, "kto": 0.2}6. 总结:数据准备不是前置步骤,而是训练策略本身
回顾全文,RLHFDataset的价值远不止于“把文件变成 tensor”。它是一个可编程的数据协议引擎:
- 通过
chat_template,你定义了模型“看到什么”; - 通过
max_prompt_length和padding_side,你控制了 KV cache 的效率与显存占用; - 通过
dynamic_length和data_type,你实现了多任务、多阶段的联合优化。
因此,不要把它当作一个黑盒工具调用完就丢弃。相反,把它看作训练配方的一部分——每次调整数据配置,都应像调整学习率一样谨慎,并通过验证集指标(如 reward score、KL 散度)来量化其影响。
当你下次启动 verl 训练时,不妨花 10 分钟检查你的 Parquet Schema、打印一个 batch 的input_ids形状、验证attention_mask的求和值。这些微小动作,往往比调参更能决定一次训练的成败。
--- > **获取更多AI镜像** > > 想探索更多AI镜像和应用场景?访问 [CSDN星图镜像广场](https://ai.csdn.net/?utm_source=mirror_blog_end),提供丰富的预置镜像,覆盖大模型推理、图像生成、视频生成、模型微调等多个领域,支持一键部署。