NVIDIA/Megatron-Energon

Megatron Energon – 用于大规模多模态模型训练的数据加载器

简介 – 这是一个Python包（以及一个小型CLI），用于为NVIDIA的Megatron‑LM训练框架准备、索引和流式传输大规模多模态数据集（图像、视频、文本等）。它也可以独立使用。

重要性 – 训练当今大规模视觉语言或视频语言模型需要跨多个存储系统分布的数PB数据。Megatron Energon处理了这些繁重的工作：

• 构建可复现的WebDataset或JSON‑Lines索引，
• 将众多异构源合并为单个逻辑数据集，
• 将工作分片到多个节点/进程，
• 提供可恢复的流水线，使失败的任务可以从上次中断的地方继续。

主要特性（README中描述）

特性	功能
CLI (energon)	准备WebDataset索引、收集媒体元数据、将数据集挂载为只读虚拟文件系统等。
Python API	get_train_dataset 和 get_loader 等函数返回PyTorch兼容的数据集/加载器，可在训练循环中迭代。
分布式支持	WorkerConfig允许指定排名、世界大小和工作进程数，因此同一代码可以在单GPU或完整的多节点集群上运行。
多模态友好	支持图像、视频、文本等，使用在TaskEncoder中注册的“cookers”将原始样本转换为模型期望的张量。
可复现性和可恢复性	索引存储在.nv-meta文件夹中；加载器可以无需重新处理整个数据集即可重新启动。
安装	pip install megatron-energon（稳定版）或直接从仓库安装。

典型工作流

1. 准备数据 – 运行 energon prepare /path/to/dataset（或针对JSONL文件）以生成.nv-meta索引和可选的数据集划分。
2. （可选）收集媒体统计信息 – energon prepare-media … 收集图像大小、视频时长等，可用于过滤或分桶。
3. 在代码中加载 – 使用Python API获取train_ds和train_loader，然后在训练脚本中迭代批次。
4. 与Megatron‑LM结合 – 将批次提供给Megatron的模型并行训练循环，进行大规模LLM或多模态模型训练。

适用人群 – 构建大规模视觉语言、视频语言或其他多模态模型的研究人员和工程师，他们需要可扩展且可复现的数据流水线。这是研究级代码，因此API可能会演变。

了解更多 – 完整文档请访问 https://nvidia.github.io/Megatron-Energon/ ，GitHub仓库提供问题、测试和示例。