VLM 数据集转换#

1. 数据集格式与处理#

考虑到多模态数据集的多样性，本项目采用 Energon 加载器来提升数据处理性能，该加载器要求数据集以标准 WebDataset 格式存储。WebDataset 以原生文件格式（jpg、mp4 等）存储数据，允许各种原生多模态数据集简单地压缩并转换为 WebDataset 格式，然后由 Energon 读取。

参考文档：

Energon：https://nvidia.github.io/Megatron-Energon/
WebDataset：https://huggingface.co/docs/hub/datasets-webdataset

本目录提供 tools/data_preprocess/vlm/convert_to_webdataset.py 用于将 .json/.jsonl 标注文件 + 原始媒体文件（图像/视频）转换为 Energon 可直接读取的 WebDataset 目录（同时生成 Energon 所需的索引和 dataset.yaml）。

2. 支持的数据类型（`--sample_type`）#

根据 --sample_type 写入不同的 dataset.yaml 文件，决定了 tar 文件内样本的字段组织方式：

常见预训练数据集可使用 caption 格式，SFT 数据集可使用 VQA 格式。对于多图像或混合 SFT 需求，推荐使用 multi_mix_qa 格式。

sample_type	适用场景	说明
`vqa`	单图 VQA	生成 `VQASample` 映射，图像字段为 `jpg`，文本从 `json[...]` 中提取
`caption`	单图描述	生成 `CaptioningSample` 映射，图像字段为 `jpg`，文本从 `json[...]` 中提取
`multi_mix_qa`	多图/视频混合 QA	使用 `CrudeWebdataset`，通过 `subflavors.sample_type` 传递给下游 cooker 进行解析
`multi_vid_vqa`	多视频 VQA	同上
`packed_captioning` / `packed_vqa` / `packed_multi_mix_qa`	离线打包后的数据	通常由 `offline_packing` 工作流生成（见第 2 节）
其他字符串	自定义场景	仍写入 `CrudeWebdataset`，但需确保下游实现了对应的 `sample_type` 解析逻辑

注意事项：

--media 仅用于写入数据集元数据（用于区分 image/video/mix）。实际样本是否包含图像/视频取决于每条数据是否包含 image(s) / video(s) 字段。
如果一条数据既没有 image(s) 也没有 video(s)，将被写入为"纯文本样本"（仅包含 json）。

3. 转换脚本使用方法#

支持的输入文件：

--json_file：.json（list[dict]）或 .jsonl（每行一个 dict）
--image_dir / --video_dir：原始媒体文件根目录（条目中存储的是相对路径）

python tools/data_preprocess/vlm/convert_to_webdataset.py \
  --output_dir /workspace/wds_data/ \
  --json_file tests/datasets/vlm/mllm_demo.json \
  --image_dir tests/datasets/vlm/ \
  --video_dir tests/datasets/vlm/ \
  --media mix \
  --columns_messages messages \
  --maxcount 10000 \
  --maxsize 3000000000 \
  --sample_type multi_mix_qa

参数说明：

参数	默认值	说明
`--output_dir`	-	输出目录（生成 `pretrain-*.tar` + Energon 元数据目录）
`--json_file`	-	输入 `.json/.jsonl` 文件
`--image_dir`	-	图像根目录（样本包含 `image(s)` 或 `sample_type=vqa/caption` 时必需）
`--video_dir`	-	视频根目录（样本包含 `video(s)` 时必需）
`--media`	`image`	`image/video/mix`
`--columns_messages`	`messages`	条目中对话/文本字段的键名
`--maxcount`	`10000`	每个分片（tar）的最大样本数
`--maxsize`	`3000000000`	每个分片（tar）的最大字节数
`--sample_type`	-	数据类型（见上表）

输出说明：

输出目录包含 pretrain-0.tar、pretrain-1.tar…（每个 tar 根据 WebDataset 规范按照 __key__ 存储若干文件，如 xxx.jpg/xxx.json/xxx.0_a.mp4 等）
同时生成 Energon 所需的元数据目录（通常名为 .wds/），包含 dataset.yaml 和索引文件；训练时 --data-path 通常指向 --output_dir

4. 输入 JSON 约定（常用字段）#

每条数据支持以下字段组合（均为相对路径，将与 --image_dir/--video_dir 拼接以读取二进制文件）：

图像：image: "a/b.jpg" 或 images: ["a/b.jpg", "c/d.jpg"]
视频：video: "a/b.mp4" 或 videos: ["a/b.mp4", "c/d.mp4"]
文本/对话：默认读取 messages（可通过 --columns_messages 修改）

不同 sample_type 的文本字段要求（与脚本生成的 dataset.yaml 对齐）：

vqa：messages 应支持读取 json[0][content] 和 json[1][content]（通常为长度 >= 2 的列表，元素包含 content）
caption：messages 应支持读取 json[captions][0][content]（例如 dict 包含 captions: [{content: ...}]）
multi_mix_qa / multi_vid_vqa 等：脚本写入结构化的 json（包含 texts/media/name），下游根据对应的 sample_type cooker 进行解析

5. 离线 Packing 数据处理#

在多模态场景中，提供了序列离线 packing 处理方法

详情请参阅离线数据打包指南

VLM 数据集转换

目录