rl_dataset.py — RLHF 数据集

文件路径: verl/utils/dataset/rl_dataset.py

文件概述

RLHF 训练的核心数据集类。从 Parquet 文件加载 prompt 数据，支持 tokenization、多模态输入（图像/视频）、数据过滤、动态采样等。

核心类

RLDataset

class RLDataset(Dataset):
    def __init__(self, parquet_files, tokenizer, config, processor=None):
        """
        Args:
            parquet_files: Parquet 数据文件路径
            tokenizer: HuggingFace tokenizer
            config: 数据配置（max_prompt_length、filter 等）
            processor: 多模态处理器（处理图像/视频）
        """

数据处理流程

Parquet 文件 → 读取 DataFrame → 提取消息 → apply_chat_template → tokenize
    → 裁剪/padding 到固定长度 → 返回 {input_ids, attention_mask, position_ids, ...}

关键特性

1. 数据过滤: 支持按 prompt 长度过滤，去除过长/过短的样本
2. 多模态: 支持图像和视频输入（Qwen2-VL 格式）
3. 左侧 padding: RLHF 中通常使用左 padding，便于 response 对齐
4. HDFS 支持: 自动从 HDFS 下载数据文件

与其他模块的关系

• 依赖 tokenizer.py 加载 tokenizer
• 依赖 chat_template.py 处理聊天模板
• 依赖 vision_utils.py 处理多模态输入
• 依赖 fs.py 处理 HDFS 文件
• 被 PPO trainer 的 DataLoader 使用

小结

rl_dataset.py 是 RLHF 数据管道的起点，负责将原始文本数据转化为模型可接受的 tensor 格式。