Megatron-LM 中的 Megatron Energon：面向文本、图像、视频与音频的大规模多模态数据加载器实战指南

Megatron-LM 中的 Megatron Energon面向文本、图像、视频与音频的大规模多模态数据加载器实战指南【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LMMegatron Energon 是 Megatron-LM 生态中面向多模态训练的高性能数据加载库支持文本、图像、视频与音频样本的分布式读取、打包packing、分组grouping与多数据集混合。本文以 docs/user-guide/features/megatron_energon.md 为核心骨架结合本仓库examples/multimodal下的真实集成代码讲解其安装、核心 API、Worker 配置、数据混合与断点续训机制帮助你直接在 Megatron-LM 中搭建可复现的多模态数据流水线。概述Megatron Energon 能解决什么问题多模态大模型训练如 VLM、音视频模型的数据流远比纯文本复杂一条样本可能同时包含高分辨率图像、视频帧、音频与多轮对话文本且数据往往以 TB 级规模存放在对象存储中。Megatron Energon 面向这一场景提供了以下能力多模态支持文本、图像、视频和音频样本的统一读取与解码分布式加载面向多节点训练设计配合数据并行组在多个 worker 间切分与预取数据数据混合Blending以可配置权重混合多个数据集支持按Metadataset配置WebDataset 格式直接从 S3、GCS、Azure Blob Storage 等云端对象存储流式读取状态管理保存与恢复训练位置dataloader checkpoint支撑断点续训。Megatron Energon 以独立库的形式提供megatron-energon同时 Megatron-LM 在examples/multimodal/中提供了完整的集成层包括数据加载器 provider、任务编码器TaskEncoder与自定义样本类型。安装pip install megatron-energon安装后即可在代码中通过megatron.energon命名空间导入数据集、加载器与配置类。仓库中的多模态示例还提供了专用 Dockerfileexamples/multimodal/Dockerfile可在容器内直接运行 llava 架构 VLM 的预训练与指令微调。核心 API 与基本用法创建训练数据集get_train_dataset是入口函数返回一个可迭代的训练数据集对象from megatron.energon import get_train_dataset, get_loader, WorkerConfig # 创建数据集 ds get_train_dataset( /path/to/dataset, batch_size32, shuffle_buffer_size1000, worker_configWorkerConfig.default_worker_config(), ) # 创建 loader 并迭代 for batch in get_loader(ds): # Training step pass核心参数说明参数含义典型值数据集路径指向 megatron-energon 格式WebDataset dataset.yaml的目录/path/to/datasetbatch_size每个 batch 的样本数32、64shuffle_buffer_size用于打乱顺序的缓冲区大小100、1000max_samples_per_sequence打包进一条序列的最大样本数100worker_config并行加载的 worker 配置WorkerConfig(...)多模态示例对于图像-文本image-text数据集加载后每个 batch 中可直接按字段名取回张量# 加载图像-文本数据集 ds get_train_dataset( /path/to/multimodal/dataset, batch_size32, worker_configWorkerConfig(num_workers8, prefetch_factor2), ) for batch in get_loader(ds): images batch[image] # 图像张量 texts batch[text] # 文本描述 # Process batch在多模态场景下batch的字段结构由所用样本类型与任务编码器决定。仓库的 VLM 示例中batch 会进一步被编码为包含tokens、labels、imgs、num_tiles、cu_lengths、max_lengths等字段的结构化对象详见下文任务编码器一节。Worker 配置与常用参数WorkerConfig分布式加载依赖WorkerConfig描述当前 rank 在数据并行组中的位置WorkerConfig( num_workers8, # 并行 worker 数量 prefetch_factor2, # 每个 worker 预取的 batch 数 persistent_workersTrue, # 保持 worker 跨 epoch 存活 )在 Megatron-LM 的实际集成中WorkerConfig还接收rank、world_size与data_parallel_group从而将数据并行组内的每个 rank 绑定到不同的数据分片见 dataloader_provider.pyworker_config WorkerConfig( rankrank, # 数据并行 rank world_sizeworld_size, # 数据并行世界大小 num_workersargs.num_workers, # 每 worker 的加载线程数 data_parallel_groupdata_parallel_group, # 数据并行通信组 worker_debug_pathworker_debug_path, # worker 调试输出目录可选 worker_log_levelworker_log_level, # worker 日志级别 )常用参数一览参数描述batch_size每个 batch 的样本数shuffle_buffer_size打乱顺序的缓冲区大小max_samples_per_sequence打包进一条序列的最大样本数worker_config并行加载的 worker 配置数据混合Blender 与 MetadatasetPython 侧Blender需要按权重混合多个数据集时可直接使用Blenderfrom megatron.energon import Blender blended_ds Blender([ (/path/to/dataset1, 0.6), # 60% (/path/to/dataset2, 0.3), # 30% (/path/to/dataset3, 0.1), # 10% ])配置侧Metadataset生产环境中更推荐使用dataset.yamlMetadataset 配置将混合权重与数据集路径声明化。仓库中 examples/multimodal/pretrain_dataset.yaml 展示了最小配置__module__: megatron.energon __class__: Metadataset splits: train: datasets: - weight: 1. path: path_to_pretraining_dataset_in_energon_format subflavors: augmentation: false val: datasets: - weight: 1. path: path_to_pretraining_dataset_in_energon_format subflavors: augmentation: false而 NVLM 预训练混合配置 examples/multimodal/nvlm/pretrain_blend.yaml 展示了多数据集按规模加权混合的真实样例权重之和为 1权重按各数据集规模占比设定__module__: megatron.energon __class__: Metadataset splits: train: datasets: - weight: 0.579 # 数据集按规模加权权重之和为 1 path: path to laion dataset subflavors: augmentation: False - weight: 0.02 path: path to coco subflavors: augmentation: False - weight: 0.01 path: path to vqav2 dataset subflavors: augmentation: False val: datasets: - weight: 1. path: path to validation dataset subflavors: augmentation: Falsesubflavors用于给样本打标签任务编码器会在编码时读取例如 NVLM 示例中通过augmentation控制是否对图像做增强。训练入口只需把data_path指向这个 yaml 文件即可。与 Megatron-LM 的完整集成训练循环中的最小集成from megatron.energon import get_train_dataset, get_loader from megatron.training import get_args args get_args() train_ds get_train_dataset( args.data_path, batch_sizeargs.micro_batch_size, ) for iteration, batch in enumerate(get_loader(train_ds)): loss train_step(batch)仓库中的真实实现dataloader_providerexamples/multimodal示例通过train_valid_test_dataloaders_provider把 Energon 接入 Megatron-LM 的训练框架其关键设计包括仅在特定 rank 上构建 dataloader只让 tensor 并行 rank 0 且位于 pipeline 首/尾 stage 的 rank 创建数据加载器is_dataloader_rank避免冗余加载与重复打乱训练集get_train_dataset(..., virtual_epoch_length1000, max_samples_per_sequence100, shuffle_buffer_size100, image_decodepil)见 dataloader_provider.py其中virtual_epoch_length定义虚拟 epoch 长度image_decodepil指定用 PIL 解码图像验证集get_val_datasets返回验证数据集列表再用RepeatDataset包裹并LimitDataset截断到eval_iters * num_microbatches长度保证评估步数可预测可保存的 loader训练集使用get_savable_loader构建以支持 dataloader 状态保存与恢复EnergonDataloader 封装将 dataloader 迭代器包装为可循环cyclic的迭代对象对齐 Megatron-LM 训练循环的消费方式EnergonDataloader。断点续训与数据加载状态get_savable_loader返回的 loader 支持save_state_rank()与restore_state_rank()。在 dataloader_provider.py 中恢复流程为加载阶段若指定了args.dataloader_save目录则按数据并行 rank 读取名为train_dataloader_dprank{dp_rank:03d}.pt的 checkpoint 文件并调用train_dataloader.restore_state_rank(...)恢复采样位置每次训练后保存对应状态。这样中断后可从精确的样本位置续跑而不会重复消费或遗漏数据。数据准备把原始数据转换为 Energon 格式examples/multimodal/README.md给出了完整的实操链路以 LLaVA-Pretrain 为例下载 LLaVA-Pretrain 数据集并解压图像目录运行 convert_llava_pretrain_to_wds.py 将数据转换为 WebDataset 格式在 wds 目录下执行energon prepare ./按交互提示生成dataset.yaml与数据集划分 Please enter a desired train/val/test split like 0.5, 0.2, 0.3 or 8,1,1: 9,1,0 Do you want to create a dataset.yaml interactively? [Y/n]: Y Please enter a number to choose a class: 9 (VQASample) Do you want to set a simple field_map[Y] (or write your own sample_loader [n])? [Y/n]: Y Please enter a webdataset field name for image (class torch.Tensor): jpg Please enter a webdataset field name for context (class str): json[0][value] Please enter a webdataset field name for answers (typing.Optional[typing.List[str]], default: None): json[1][value] Please enter a webdataset field name for answer_weights (typing.Optional[torch.Tensor], default: None):更新pretrain_dataset.yaml中的path指向转换后的 wds 目录运行examples/multimodal/pretrain_mistral_clip.sh开始预训练SFT 阶段则使用examples/multimodal/sft_mistral_clip.sh。任务编码器把原始样本变成训练张量get_train_dataset可接收task_encoder参数。仓库在 examples/multimodal/dataset_helpers.py 中实现了TaskEncoder(DefaultTaskEncoder[...])展示了对不同 Energon 样本类型的编码分发CaptioningSample图像描述任务encode_captioningVQASample单轮视觉问答按__subflavors__[is_llava_training]或 key 中是否含llava分流到 LLaVA 预训练编码或通用 VQA 编码OCRSamplePDF/OCR 任务支持定位引用referring grounding与 grounded OCR 等任务模板SimilarityInterleavedSample多轮交错图文对话SFTMultiChoiceVQASample多选题式 VQASampleListSample离线打包好的样本列表。编码过程中会进行图像切块tiling、缩略图、token 化与序列长度计算_get_total_seq_length并支持上下文并行CP所需的 padding。若开启打包packing_buffer_size 0select_samples_to_pack通过贪心背包算法greedy_knapsack把多条短样本拼进一条序列以充分利用序列长度容量见 dataset_helpers.py 与 dataset_helpers.py。此外仓库还定义了自定义的 Energon 样本类型examples/multimodal/energon_util.pySampleListSample离线打包的样本列表与OfflineTargetAspectRatioSample离线预计算目标宽高比SFT 时可直接使用而无需在线计算可扩展性良好。生产使用要点总结分布式加载通过WorkerConfig传入数据并行 rank/组跨 worker 与节点并行读取、预取配合persistent_workersTrue复用进程打包与分组max_samples_per_sequence与packing_buffer_size/packing_seq_length控制序列打包shuffle_buffer_size控制打乱程度对象存储WebDataset 格式天然支持 S3、GCS、Azure Blob Storage 流式读取无需预先全量下载断点续训用get_savable_loadersave_state_rank/restore_state_rank保存/恢复采样位置混合数据用Blender代码或Metadatasetyaml配置按权重混合多源数据。下一步查看 多模态模型文档 了解本仓库支持的 VLM/多模态模型架构参考 训练示例 与 examples/multimodal/README.md 获取端到端的预训练、SFT 与评估流程。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考