verl数据格式转换:parquet适配实战
1. 背景与问题引入
在使用 verl 框架进行大型语言模型的强化学习训练时,数据准备是关键的第一步。很多用户在实际操作中会遇到一个常见问题:原始数据为 Arrow 格式,但 verl 的默认数据加载机制仅支持 Parquet 格式。
这并不是框架的限制,而是出于性能和兼容性的设计选择。Arrow 是一种高效的内存数据格式,适合快速读写;而 Parquet 是列式存储格式,更适合大规模离线数据处理和分布式训练场景。verl 基于 HuggingFace 的datasets库构建数据流,默认通过load_dataset("parquet")加载训练集。
本文将带你从零开始,解决verl 数据格式不兼容的实际问题,重点聚焦于如何高效完成Arrow 到 Parquet 的格式转换,并确保字段映射正确、训练流程无缝衔接。
我们还会探讨两种解决方案:推荐的“格式转换”方案和灵活的“自定义数据集类”方案,帮助你在不同工程约束下做出最优选择。
2. 环境准备与 verl 验证
在开始数据处理前,首先要确认 verl 已正确安装并可调用。
2.1 进入 Python 环境
python2.2 导入 verl 模块
import verl如果导入失败,请检查是否已通过 pip 或源码方式正确安装 verl。建议使用虚拟环境以避免依赖冲突。
2.3 查看 verl 版本号
print(verl.__version__)输出类似如下内容即表示安装成功:
0.1.0提示:建议使用最新稳定版本,确保兼容性。若版本过旧,可通过
pip install --upgrade verl更新。
3. 数据格式适配核心策略
verl 的RLHFDataset类默认使用以下方式加载数据:
dataframe = datasets.load_dataset("parquet", data_files=parquet_file)["train"]这意味着它期望传入的是.parquet文件路径。如果你的数据是.arrow(即.feather)格式,则需要进行适配。
主要有两种解决方案:
- 方案一:文件格式转换(推荐)
- 方案二:自定义数据集类(高级用法)
我们将依次展开说明。
4. 方案一:文件格式转换(推荐做法)
这是最简单、最直接、也最稳定的解决方案——将 Arrow 格式批量转换为 Parquet 格式。
4.1 使用 HuggingFace Datasets 转换
HuggingFace 的datasets库天然支持多种格式之间的互转,代码简洁且高效。
from datasets import load_dataset import os # 加载原始 Arrow 格式数据集 ds = load_dataset("PRIME-RL/Eurus-2-RL-Data") # 定义输出目录 output_dir = "/data/oss_bucket_0/seadawn/openlm_hub/eurus-2-rl-data-parquet" os.makedirs(output_dir, exist_ok=True) # 分别保存训练集和验证集为 Parquet 格式 ds["train"].to_parquet(os.path.join(output_dir, "train.parquet")) ds["validation"].to_parquet(os.path.join(output_dir, "validation.parquet"))这段代码完成了以下几步:
- 从 HuggingFace Hub 加载原始数据集(自动下载)
- 创建本地输出目录
- 将
train和validation子集分别导出为 Parquet 文件
优势:
- 无需修改任何训练配置
- 转换后可复用,节省后续重复处理成本
- Parquet 更适合分布式训练 I/O 性能优化
4.2 验证转换结果
转换完成后,建议手动检查文件是否存在,并查看基本统计信息:
from datasets import Dataset # 重新加载 Parquet 文件验证 loaded_ds = Dataset.from_parquet(os.path.join(output_dir, "train.parquet")) # 打印样本数量和字段 print(f"Loaded dataset size: {len(loaded_ds)}") print(f"Features: {list(loaded_ds.features.keys())}")预期输出应包含prompt,data_source,reward_model等字段。
4.3 在训练命令中引用新格式
一旦完成转换,只需在启动训练时指向新的 Parquet 文件路径即可:
python3 -m verl.trainer.main_fastrl \ data.train_files=/data/oss_bucket_0/seadawn/openlm_hub/eurus-2-rl-data-parquet/train.parquet \ data.val_files=/data/oss_bucket_0/seadawn/openlm_hub/eurus-2-rl-data-parquet/validation.parquet这样,verl 就能正常加载数据并进入训练流程。
5. 方案二:自定义数据集类(适用于无法转换场景)
如果你由于权限、存储或流程限制无法转换数据格式,可以选择扩展 verl 的数据加载逻辑,实现对 Arrow 格式的原生支持。
5.1 继承 RLHFDataset 并重写加载方法
创建一个自定义数据集类,继承自verl.utils.dataset.RLHFDataset,并覆盖_read_files_and_tokenize方法。
from verl.utils.dataset import RLHFDataset from datasets import load_dataset import os class EurusArrowDataset(RLHFDataset): def _read_files_and_tokenize(self): dataframes = [] # 遍历所有输入文件 for file_path in self.data_files: # 显式指定 arrow 格式加载 ext = os.path.splitext(file_path)[1].lower() if ext == ".arrow" or ext == ".feather": split = "train" if "train" in file_path else "validation" df = load_dataset("arrow", data_files=file_path)[split] dataframes.append(df) else: raise ValueError(f"Unsupported file extension: {ext}") # 合并所有分片 self.dataframe = datasets.concatenate_datasets(dataframes) # 可选:过滤过长 prompt self.dataframe = self.maybe_filter_out_long_prompts(self.dataframe) print(f"Final dataset length: {len(self.dataframe)}")5.2 注册自定义数据集类
在训练配置中指定使用你的自定义类。通常通过 YAML 配置文件实现:
data: custom_cls: path: /path/to/your/custom_dataset.py name: EurusArrowDataset train_files: - /data/oss_bucket_0/seadawn/openlm_hub/eurus-2-rl-data/eurus-2-rl-data-train-00000-of-00004.arrow - /data/oss_bucket_0/seadawn/openlm_hub/eurus-2-rl-data/eurus-2-rl-data-train-00001-of-00004.arrow - /data/oss_bucket_0/seadawn/openlm_hub/eurus-2-rl-data/eurus-2-rl-data-train-00002-of-00004.arrow - /data/oss_bucket_0/seadawn/openlm_hub/eurus-2-rl-data/eurus-2-rl-data-train-00003-of-00004.arrow val_files: /data/oss_bucket_0/seadawn/openlm_hub/eurus-2-rl-data/eurus-2-rl-data-validation.arrowverl 会在运行时动态加载该类,并验证其是否继承自torch.utils.data.Dataset。
5.3 注意事项
- 自定义类必须放在可导入路径下
- 文件路径需准确无误
- 若使用多机训练,确保每个节点都能访问该模块
- 建议添加日志打印,便于调试加载过程
6. 字段映射与数据结构验证
无论采用哪种方案,都必须确保数据字段与 verl 的默认配置匹配。
6.1 默认字段配置解析
根据 verl 源码中的legacy_data.yaml配置:
prompt_key: prompt reward_fn_key: data_source这意味着:
prompt字段用于输入提示文本data_source字段用于决定使用哪个奖励函数(支持多奖励模型)
你的数据集中已有这两个字段,因此无需额外映射。
6.2 其他可用字段说明
| 字段名 | 是否必需 | 用途说明 |
|---|---|---|
prompt | 是 | 用户输入提示 |
data_source | 是 | 奖励函数选择键 |
reward_model | 否 | 指定具体奖励模型名称 |
ability | 否 | 任务能力标签(如写作、推理) |
extra_info | 否 | 附加元信息,可用于日志记录 |
这些字段都会被保留并在训练过程中传递,可用于条件奖励计算或多任务调度。
7. 多文件加载机制详解
即使你不做格式转换,也可以直接传入多个 Arrow 文件路径,verl 会自动合并它们。
7.1 内部实现原理
在RLHFDataset._read_files_and_tokenize中有如下逻辑:
if not isinstance(data_files, list): data_files = [data_files] dataframes = [] for parquet_file in self.data_files: dataframe = datasets.load_dataset("parquet", data_files=parquet_file)["train"] dataframes.append(dataframe) self.dataframe = datasets.concatenate_datasets(dataframes)这表明:
- 支持单个文件或文件列表
- 自动合并多个分片
- 使用
concatenate_datasets保证高效拼接
7.2 直接使用多 Arrow 文件(配合自定义类)
结合前面的自定义类,你可以轻松支持多 Arrow 文件输入:
data_files = [ "part-00000-of-00004.arrow", "part-00001-of-00004.arrow", "part-00002-of-00004.arrow", "part-00003-of-00004.arrow" ]只要在自定义类中循环加载并合并即可。
8. 缓存与性能优化建议
8.1 数据缓存路径
verl 默认会将处理后的数据缓存到本地:
self.cache_dir = os.path.expanduser(config.get("cache_dir", "~/.cache/verl/rlhf"))首次运行较慢,后续可直接读取缓存,提升启动速度。
8.2 过滤过长 prompt
默认开启:
self.filter_overlong_prompts = config.get("filter_overlong_prompts", True)可在配置中关闭:
data: filter_overlong_prompts: false建议保持开启,防止 OOM 错误。
8.3 推荐最佳实践
| 实践项 | 推荐做法 |
|---|---|
| 数据格式 | 统一转为 Parquet |
| 存储位置 | 使用高速 SSD 或分布式文件系统 |
| 文件数量 | 控制在 1~8 个之间,避免过多小文件 |
| 字段命名 | 严格遵循prompt,data_source等标准 |
| 自定义类 | 仅在必要时使用,优先格式转换 |
9. 总结
本文系统梳理了在使用 verl 框架时,面对 Arrow 格式数据的适配方案。核心结论如下:
- 首选方案是格式转换:将 Arrow 转为 Parquet,简单、稳定、兼容性强。
- 支持多文件自动合并:无论是 Parquet 还是 Arrow,verl 都能处理文件列表。
- 字段映射清晰明确:
prompt和data_source是关键字段,必须存在。 - 自定义数据集类提供灵活性:当无法转换格式时,可通过继承
RLHFDataset实现原生支持。 - 缓存机制提升效率:首次处理后数据会被缓存,加快后续训练启动速度。
通过合理选择方案,你可以轻松打通数据链路,让 verl 快速进入训练状态,专注于算法迭代而非数据搬运。
获取更多AI镜像
想探索更多AI镜像和应用场景?访问 CSDN星图镜像广场,提供丰富的预置镜像,覆盖大模型推理、图像生成、视频生成、模型微调等多个领域,支持一键部署。