Megatron-LM 数据集管线全解析从 IndexedDataset 二进制格式到 GPTDataset 三索引机制与快速 DataLoader 初始化【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM本篇技术文章围绕 Megatron Core 的datasets包展开系统讲解其数据管线Data Pipeline的三层架构底层IndexedDataset/IndexedDatasetBuilder的二进制存取格式、中上层由BlendedMegatronDatasetConfig与BlendedMegatronDatasetBuilder驱动的分布式感知的 DataLoader 构建流程、GPTDataset的文档/样本/洗牌三索引查表机制以及离线缓存预生成tools/prepare_cache.py、Packing 调度器和三个加速 DataLoader 初始化的配置开关。读完后你可以独立完成数据集预处理、理解训练启动时索引缓存的生成与命中逻辑并针对大规模数据混合场景配置快速加载路径。一、整体架构三层数据接口与构建类Megatron Core 的数据管线是分层设计的核心类关系如下源码位于 megatron/core/datasets/层级类作用底层IndexedDataset/IndexedDatasetBuilder最低层数据接口读写.bin.idx二进制文件配置BlendedMegatronDatasetConfig参数化 Builder 与各级数据集可按训练/推理体制扩展如GPTDatasetConfig构建BlendedMegatronDatasetBuilder构建最高层数据接口是分布式感知的构建入口中层MegatronDataset抽象基类建立在IndexedDataset之上的高阶抽象不同任务有不同扩展如GPTDataset顶层BlendedDataset建立在多个MegatronDataset之上的混合数据集仅在多个数据分布共同贡献同一个 split 时才需要需要注意一个重要的分布式约定所有 rank 都必须尝试通过BlendedMegatronDatasetBuilder构建数据集否则程序会挂起哪些 rank 真正执行构建逻辑则由BlendedMegatronDatasetConfig控制。这一约定在源码中体现为 blended_megatron_dataset_builder.py 中build_generic_dataset的 rank 0 先构建 →torch.distributed.barrier()→ 其余 rank 再构建缓存命中 模式。二、数据预处理IndexedDatasetBuilder 与 IndexedDataset数据预处理围绕两个类展开IndexedDatasetBuilderIndexedDataset官方文档指出目前端到端的数据预处理实现留给用户完成详见类文档字符串实际入口可参考 tools/preprocess_data.py。2.1 IndexedDatasetBuilder构建与合并数据集IndexedDatasetBuilder位于 indexed_dataset.py约 930 行起能够构建并合并IndexedDataset实例。其核心方法包括add_item(tensor, mode0)向数据集追加单条序列并记录其长度add_document(tensor, lengths, modesNone)追加整篇文档lengths给出文档内各序列的长度同时更新document_indicesadd_documents(documents, modesNone, eod_tokenNone, chunk_size1_000_000)以批量方式写入 jagged 数组形式的文档列表依赖awkward库可自动在文档边界插入 EOD token分块写入以控制内存add_index(path_prefix)把另一个已存在的IndexedDataset整体合并进来拼接索引与数据finalize(idx_path)关闭数据文件并写入.idx索引文件。构造参数为bin_path数据文件路径、dtype默认numpy.int32与multimodal是否多模态决定是否记录每条序列的 mode。2.2 IndexedDataset.bin 与 .idx 的二进制布局IndexedDataset是 Megatron Core 中最低层的数据接口。一个实例引用两个二进制文件数据文件.bin保存文档/序列数据索引文件.idx保存文档/序列元数据。索引文件的内容分两部分。先存数据集级元数据索引文件头为向后兼容保留索引版本号为向后兼容保留一个数字编码对应写入数据文件所用的数据类型数据集中的序列数量数据集中的文档数量。随后存文档级与序列级元数据按顺序存每条序列的元素个数int32按顺序存每条序列的字节偏移指针int64按顺序存每个文档对应的连续序列索引区间[...)int64按顺序存每条序列的 mode仅多模态情形int8。这与源码中_IndexWriter.write()的落盘顺序完全一致header version dtype code sequence_count document_count sequence_lengths sequence_pointers document_indices sequence_modes其中_INDEX_HEADER bMMIDIDX\x00\x00版本号为小端8字节整型 1。数据类型编码由DType枚举定义int32对应编码 4、uint16对应编码 8 等。在读取侧_IndexReader会对.idx做numpy.memmap然后按偏移量依次还原sequence_lengthsint32、sequence_pointersint64、document_indicesint64多模态时再还原sequence_modesint8。.bin数据文件有四种读取器实现由IndexedDataset.__init__根据参数选择_MMapBinReadermmapTrue时内存映射读取默认本地文件场景_FileBinReadermmapFalse时用文件指针seek readinto内置指数退避重试默认 3 次重试起始睡眠 10s 翻倍_S3BinReaderS3 对象存储场景以bin_chunk_nbytes为块大小维护内存缓存按块范围Range拉取字节_MultiStorageClientBinReader基于 Multi-Storage Client 的范围读。此外还暴露get(idx, offset, length)方法支持只取序列一部分的读取这正是上层GPTDataset拼接跨文档样本时依赖的关键能力。IndexedDataset还支持fast_cache_load跳过文件存在性断言与sequences_per_dataset直接用计数信息初始化索引免开.idx文件配合--per-dataset-sequences-path使用。三、数据加载之构建BlendedMegatronDatasetConfig 与 BuilderDataLoader 的构建是一个分布式感知的过程围绕五个要素展开BlendedMegatronDatasetConfigBlendedMegatronDatasetBuilderIndexedDatasetMegatronDatasetBlendedDataset3.1 BlendedMegatronDatasetConfig可扩展该类blended_megatron_dataset_config.py参数化BlendedMegatronDatasetBuilder进而参数化MegatronDataset与BlendedDataset。不同训练/推理体制需要不同的扩展例如 GPT 训练使用的GPTDatasetConfig。配置类中的关键字段包括random_seed、sequence_length随机种子与序列长度必填blend[[prefix1, prefix2], [0.3, 0.7]]形式的混合定义权重为 None 时按底层数据集长度推断不能与blend_per_split同用blend_per_splittrain/valid/test 三个 split 各自独立的混合定义split从单一分布抽样时 train/valid/test 的权重字符串如99,1,0不能与blend_per_split同用path_to_cache所有可复用数据集索引的缓存目录mmap_bin_files默认 True.bin用 mmap 还是文件指针读取mid_level_dataset_surplus默认 0.005中层数据集构建时的样本冗余比例顶层数据集若超采中层数据集需要调大fast_cache_load/defer_npy_index_mmap两个快速加载开关见第七节都要求path_to_cache非空num_dataset_builder_threads构建数据集的线程数。__post_init__中完成了若干合法性约束fast_cache_load时禁止与blend同用应改用--per-split-data-args-path或 per-split data pathblend与blend_per_split互斥当blend非空时必须提供split两者都为空时自动进入mockTrue的模拟数据模式并取1,1,1的均匀 split。split 字符串经parse_and_normalize_split归一化后再由convert_split_vector_to_split_matrix转换为 split 矩阵各 split 在非重叠区间上的书端点区间例如[0.99, 0.01, 0.0] - [(0, 0.99), (0.99, 1.0), None]。GPTDatasetConfiggpt_dataset.py在基类之上扩展了reset_position_ids、reset_attention_mask、eod_mask_loss、create_attention_mask、add_extra_token_to_sequence、drop_last_partial_validation_sequence、hybrid_context_parallel、sequences_per_dataset等字段并在__post_init__中依据 tokenizer 词表大小自动推导token_dtype_code词表超过 uint16 上限取 int32 编码 4否则 uint16 编码 8。3.2 BlendedMegatronDatasetBuilder最高层构建器BlendedMegatronDatasetBuilderblended_megatron_dataset_builder.py注意正确路径为 megatron/core/datasets/blended_megatron_dataset_builder.py构建 Megatron Core 最高层的数据接口。其__init__接收四类参数cls要实例化的MegatronDataset子类sizes每个 split 要求的最少样本总数可为 Noneis_built_on_rank判断当前 rank 是否构建数据集的可调用对象必须感知 Megatron Core 并行策略全局 rank、组内 rank、virtual rank 都可能影响返回值且必须在全局 rank 0 上恰好返回 Trueconfig数据集配置对象。build()方法按配置分三种情形处理每个 splitsplit 为 None什么都不做单个贡献数据集size非 None 时按比例抽样子数据集size为 None 时不做超额抽样多个贡献数据集权重与 size 均给定 → 按权重与 size 构建中层与顶层数据集仅给权重不给 size → 报错仅给 size → 顶层长度取size以各中层长度之和为上限两者皆无 → 构建穷举索引。值得注意的实现细节均可在源码中验证构建过程使用ThreadPoolExecutor并行构建各 prefix 的MegatronDataset_build_megatron_datasets_parallel每个任务通过contextvars.copy_context()把 OTel 追踪上下文传播进工作线程分布式初始化后默认 rank 0 先行构建线程数还会按 GPU 数适度放大barrier之后其他 rank 再构建——此时必然命中缓存开启fast_cache_load时会跳过 rank 0 先行 barrier 的同步点各 rank 直接并行构建/加载这正是其提速原理见第七节指定了混合权重时每个 prefix 构建的样本数由_get_size_per_split_per_dataset计算会乘以(1 mid_level_dataset_surplus)冗余系数保证顶层数据集有足够样本可抽。3.3 MegatronDataset可扩展与 BlendedDataset顶层MegatronDatasetmegatron_dataset.py是建立在IndexedDataset之上的高阶抽象基类子类如GPTDataset需实现numel_low_level_dataset与build_low_level_dataset两个静态方法。构造时它会汇总类名、数据集路径、样本数、split、以及_key_config_attributes()返回的关键配置属性random_seed、sequence_length、split、split_matrix、tokenizer序列化为 JSON 描述并计算 MD5 得到unique_description_hash——缓存索引文件名的唯一性就锚定在这个 hash 上因此任何关键参数变化都会导致缓存失效重建。BlendedDatasetblended_dataset.py建立在多个MegatronDataset之上仅当需要混合多个数据分布来贡献某个 split 时才需要它混合比例通过配置控制。约束包括各子数据集必须同类、同 split、权重为正、数据集聚合数小于 32767 等若size非 None 则权重会被归一化。四、数据加载之实现GPTDataset 的三索引机制GPTDataset由以下变量参数化底层IndexedDataset实例indexed_dataset、split 索引indexed_indices用于训练/验证/测试的连续文档或序列索引子集、样本数N、序列长度S、随机种子R。它创建三个索引映射以支撑查表1文档索引 Do_idx一维数组把i映射到文档索引长度E * |indexed_indices|其中E是满足E * |indexed_indices| N的最少 epoch 数。文档索引按R洗牌。Given: N 15 indexed_indices [5, 6, 7, 8, 9] E 3 Then, for example: Do_idx [8, 8, 9, 6, 7, 5, 8, 5, 6, 6, 5, 9, 7, 7, 9]2样本索引 Sa_idx二维数组把j映射到(i, Do_idx[i] 的偏移)对形状[N 1, 2]。行j与j 1分别作为第j个样本的左、右边界。Given: S 1024 Then, for example: Sa_idx[0] (0, 0) Sa_idx[1] (0, 1024) Do_idx[0] has length greater than S Sa_idx[2] (1, 512) Do_idx[0] has length 1536 Sa_idx[3] (2, 0) Do_idx[1] has length 1536 Sa_idx[4] (5, 300) Do_idx[2:5] are shorter documents relative to Do_idx[0:2] Sa_idx[5] (6, 24) Do_idx[5] has length 13003洗牌索引 Sh_idx一维数组把k映射到j长度N按R洗牌。Given N 10 Then, for example: Sh_idx [4, 0, 2, 6, 1, 9, 5, 8, 7, 3]查询第k个样本的过程# 1. 用洗牌索引得到样本索引内的索引 j j Sh_idx[k] # 2. 用样本索引得到样本左右边界在文档索引中的位置及各自起始 token 偏移 i, offset Sa_idx[j] i_next, offset_next Sa_idx[j 1] # 3. 用文档索引从连续的文档中取出 S 个 token sample [] sample indexed_dataset[Do_idx[i]][offset:] if i ! i_next: sample indexed_dataset[Do_idx[i 1:i_next]] sample indexed_dataset[Do_idx[i_next]][:offset_next]从源码看gpt_dataset.py 的_query_document_sample_shuffle_indices约 388-473 行实际实现与伪代码一致并做了两点增强样本跨越单个文档时直接用dataset.get(idx, offset, length)一次取出跨越多个文档时逐文档get拼接。此外取出的 token 总数为S add_extra_token_to_sequence默认多取 1 个 token__getitem__中据此切分出tokens text[:-1]与labels text[1:]保证输入与标签都是完整长度不足时以 pad token 补齐并把 pad 位置的 loss mask 置零。索引构建的关键工程细节_build_document_sample_shuffle_indices约 475-701 行epoch 计算_get_num_epochs会不断累加 epoch 直至 token 总量满足N * S add_extra_token的需求末 epoch 分离若最后一个 epoch 的样本数不足完整 epoch 样本数的 80%threshold 0.80则把末 epoch 与前面 epoch 分开洗牌separate_final_epoch避免最后一个不完整的 epoch 被全局打散后与前面样本过度混合C 加速样本索引由 helpers.cpp 中的build_sample_idx构建当len(document_index) * 2 len(sequence_lengths)访问密度高时会先把 mmap 的sequence_lengths复制进内存源码注释解释了这样做的两个好处——顺序预读整个文件以及进入 C 时持有 GIL 提高并行度缓存三个索引与description.txt一并保存到缓存目录文件名为{unique_description_hash}-{ClassName}-{split}-document_index.npy等fast_cache_load时跳过文件存在性检查直接视为命中defer_npy_index_mmap时索引不在初始化时加载、延迟到首次访问时以 mmap 方式加载此时__len__会改用纯算术公式估算样本数复用 helpers.cpp 的样本计数逻辑。4.1 BlendedDataset 的混合索引BlendedDataset由数据集聚合D、权重W每个数据集一个与规模S参数化。它会按权重比例从各贡献数据集抽样直到达到目标规模每一步抽样时从抽样误差sampling error最大的那个数据集抽取一个样本。它创建两个混合索引数据集索引 Da_idx一维数组把i映射到数据集索引长度SGiven D [d0, d1, d2] W [1/2, 1/4, 1/4] S 4 Then, for example: Da_idx [0, 1, 2, 0]数据集样本索引 Sa_idx一维映射把i映射到数据集Da_idx[i]内的样本索引长度SGiven Da_idx [0, 1, 2, 0] Then, for example: Sa_idx [0, 0, 0, 1]查询第k个样本sample D[Da_idx[k]][Sa_idx[k]]同样为节省初始化时间各索引在单个 rank 上顺序构建/缓存再由其他 rank 并行加载缓存索引锚定在BlendedDataset.__init__生成的 hash 上。源码实现blended_dataset.py 的_build_indices中索引由 helpers.cpp 的build_blending_indicessize非 None或build_exhaustive_blending_indicessize为 None穷举模式构建构建后还会校验各子数据集是否被超采若超采会抛出明确提示增大mid_level_dataset_surplus的IndexError。五、离线缓存预生成tools/prepare_cache.py对于 GPT 风格训练上述数据集缓存可以用 tools/prepare_cache.py 提前准备好而不必等训练启动时 rank 0 构建。该脚本复用了pretrain_gpt.py与pretrain_mamba.py的正常数据集构建路径包括GPTDataset、BlendedDataset与BlendedMegatronDatasetBuilder。它接受常规的数据集参数支持 blend 与 per-split 数据集定义并要求--data-cache-path以便生成的缓存能被后续训练复用。对于大型 blend 或大量文件 prefix 的场景尤其有用构建 document、sample、shuffle 索引可能耗时数分钟期间所有 GPU 都处于空闲状态而 rank 0 只做纯 CPU 工作。如果后续训练任务没有指定--global-batch-size该参数用于确定数据集规模与 split应通过--prepare-cache-world-size显式指定缓存准备时使用的 world size脚本将其直接赋给args.world_size见 prepare_cache.py 的_normalize_prepare_cache_args。明确的限制tools/prepare_cache.py不支持--mock-data、--sft、--fim-data或--step-batch-size-schedule源码_validate_prepare_cache_args会对这些选项直接抛ValueError同时--data-cache-path为必填。脚本还会在准备阶段强制关闭--dataloader-fast-cache-load与--dataloader-defer-npy-index-mmap这两个开关的意义只在于消费已存在的缓存并在运行前打印生效的 world size、DP size、global batch size、缓存路径与各 split 目标样本数。六、Packing Scheduler跨 DP×CP rank 的变长序列重调度Packing 调度器把变长序列重新调度到 DP×CP 各 rank 上以提升 GPU 利用率。它围绕以下模块构建data_scheduledata_schedule.py 包含高层调度逻辑与入口点HybridCPDataLoaderWrapper混合上下文并行CP调度的包装类。每次__next__调用它会(1) 从各 DP rank 拉取一批 packed 样本(2) 在 DP 组内 all-gather 序列长度(3) 用BalancedCPScheduler来自 megatron/core/pipeline_parallel/ 的hybrid_cp_schedule调度子样本(4) 通过 all-to-all 通信把子样本重路由到正确的 DPxCP rank。BasePackingScheduler打包调度器的抽象基类定义了get_groups_and_subsamples()调度算法与run()完整调度流水线fetch、schedule、reroute、pack、broadcast 及 VPP 处理的接口。DpBalancedScheduler具体调度器按原始顺序打包序列直到达到每个 DPxCP rank 的最大序列长度限制支持把 microbatch 数对齐到 DP size 与 VPP stage 的整数倍。wrap_data_iterator()顶层入口包装已有的data_iterator。它创建合适的调度器、运行调度流水线、广播元数据与新的num_microbatches返回新的数据迭代器、更新后的 microbatch 数以及 FLOPs 统计。get_batch_on_this_rank_for_sequence_packing()为当前 rank 拉取并广播单个 packed microbatch。处理 TP/PP 广播构造PackedSeqParams含cu_seqlens、max_seqlen、qkv_formatthd并可选地用 Transformer Engine 的thd_get_partitioned_indices在 CP rank 间划分序列。data_schedule_utilsdata_schedule_utils.py 包含调度器使用的工具函数如broadcast_scalars、broadcast_tensor、build_packed_microbatches、reroute_samples_to_dcp_ranks等data_schedule.py顶部 import 即列明了全部依赖项。七、快速 DataLoader 初始化三个加速开关大规模训练中DataLoader 初始化可能耗时数分钟——因为要打开并内存映射大量文件还会显著施压文件系统。Megatron Core 提供了三个由配置开关控制的优化7.1 --dataloader-fast-cache-load假定数据集缓存已存在于指定的--data-cache-path中。启用后通过移除同步点与文件检查断言来加速创建过程。从源码看其具体生效点有三处配置层blended_megatron_dataset_config.py 断言必须提供--data-cache-path且不能与--data-pathblend 形式同用应改用--per-split-data-args-path或--train-data-path/--valid-data-path/--test-data-path构建层blended_megatron_dataset_builder.py 跳过 rank 0 先行构建 barrier各 rank 直接并行构建同时跳过indexed_indices的重复计算数据层indexed_dataset.py 跳过sequence_lengths.shape[0]系列的一致性断言。7.2 --dataloader-defer-npy-index-mmap同样假定缓存已存在。启用后把数据集索引.npy文件的内存映射延迟到首次访问时进行。官方推荐与--num-workers 0搭配使用让 DataLoader 预取下一批数据从而用后台预取掩盖索引 mmap 的开销。实现上GPTDataset/BlendedDataset在_build_indices阶段只记录缓存路径并返回 None 索引__getitem__首次调用时才safe_numpy_load(..., mmap_moder)见 gpt_dataset.py 与 blended_dataset.py__len__则改用 token 数算术公式直接推算。7.3 --per-dataset-sequences-path通过该配置指定 tools/build_sequences_per_dataset.py 生成的 JSON 文件。该脚本对 blend 中的每个文件 prefix 打开.idx读取序列数与文档数_IndexReader汇总为单一文件。此配置在处理数百乃至上千个文件 prefix 时尤其有用它只需要一次open操作而不是每个 prefix 一次。该 JSON 经GPTDatasetConfig.sequences_per_dataset传入IndexedDataset后_IndexReader可跳过解析 34 字节头部之后的完整索引读取直接用给定的(sequence_count, document_count)初始化见 indexed_dataset.py。脚本用法示例来自其模块 docstringpython3 tools/build_sequences_per_dataset.py --per-split-data-args-path my-training-dataset-blend.json --per-dataset-sequences-path my-training-dataset-blend-sequences-per-dataset.json八、小结一次训练启动中数据管线的工作顺序把以上内容串起来一次 GPT 训练启动时数据管线的工作顺序为训练脚本通过megatron.training的参数解析生成GPTDatasetConfig含 blend、split、path_to_cache等并计算各 split 目标样本数所有 rank 调用BlendedMegatronDatasetBuilder.build()rank 0 先构建或用--dataloader-fast-cache-load并行构建为每个 prefix 建IndexedDatasetmmap.bin解析.idx每个GPTDataset按unique_description_hash检查缓存命中则 mmap 加载三个.npy索引或按defer_npy_index_mmap延迟加载未命中则由 rank 0 构建C 加速并写缓存多个数据集的 split 由BlendedDataset以 最大抽样误差 策略生成混合索引运行时__getitem__依次经 shuffle → sample → document 三级查表拼接出定长样本再经 DataLoader及其可选的HybridCPDataLoaderWrapperpacking 调度送入模型。相关源码与工具入口速查主题路径底层二进制接口megatron/core/datasets/indexed_dataset.pyC 索引构建加速megatron/core/datasets/helpers.cpp / helpers.pyGPT 数据集与三索引megatron/core/datasets/gpt_dataset.py混合数据集megatron/core/datasets/blended_dataset.py分布式构建器megatron/core/datasets/blended_megatron_dataset_builder.py配置数据类megatron/core/datasets/blended_megatron_dataset_config.py抽象基类megatron/core/datasets/megatron_dataset.py离线缓存预生成tools/prepare_cache.py每数据集元数据生成tools/build_sequences_per_dataset.pyPacking 调度器megatron/core/datasets/data_schedule.py / data_schedule_utils.py适用前提提示以上行为均基于当前仓库版本快速加载类开关fast cache load、defer mmap、per-dataset sequences都要求缓存已预先构建且配置了--data-cache-pathtools/prepare_cache.py不支持 mock/SFT/FIM/step-batch-size-schedule 路径离线预生成缓存时应避免这些模式。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Megatron-LM 数据集管线全解析：从 IndexedDataset 二进制格式到 GPTDataset 三索引机制与快速 DataLoader 初始化
Megatron-LM 数据集管线全解析从 IndexedDataset 二进制格式到 GPTDataset 三索引机制与快速 DataLoader 初始化【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM本篇技术文章围绕 Megatron Core 的datasets包展开系统讲解其数据管线Data Pipeline的三层架构底层IndexedDataset/IndexedDatasetBuilder的二进制存取格式、中上层由BlendedMegatronDatasetConfig与BlendedMegatronDatasetBuilder驱动的分布式感知的 DataLoader 构建流程、GPTDataset的文档/样本/洗牌三索引查表机制以及离线缓存预生成tools/prepare_cache.py、Packing 调度器和三个加速 DataLoader 初始化的配置开关。读完后你可以独立完成数据集预处理、理解训练启动时索引缓存的生成与命中逻辑并针对大规模数据混合场景配置快速加载路径。一、整体架构三层数据接口与构建类Megatron Core 的数据管线是分层设计的核心类关系如下源码位于 megatron/core/datasets/层级类作用底层IndexedDataset/IndexedDatasetBuilder最低层数据接口读写.bin.idx二进制文件配置BlendedMegatronDatasetConfig参数化 Builder 与各级数据集可按训练/推理体制扩展如GPTDatasetConfig构建BlendedMegatronDatasetBuilder构建最高层数据接口是分布式感知的构建入口中层MegatronDataset抽象基类建立在IndexedDataset之上的高阶抽象不同任务有不同扩展如GPTDataset顶层BlendedDataset建立在多个MegatronDataset之上的混合数据集仅在多个数据分布共同贡献同一个 split 时才需要需要注意一个重要的分布式约定所有 rank 都必须尝试通过BlendedMegatronDatasetBuilder构建数据集否则程序会挂起哪些 rank 真正执行构建逻辑则由BlendedMegatronDatasetConfig控制。这一约定在源码中体现为 blended_megatron_dataset_builder.py 中build_generic_dataset的 rank 0 先构建 →torch.distributed.barrier()→ 其余 rank 再构建缓存命中 模式。二、数据预处理IndexedDatasetBuilder 与 IndexedDataset数据预处理围绕两个类展开IndexedDatasetBuilderIndexedDataset官方文档指出目前端到端的数据预处理实现留给用户完成详见类文档字符串实际入口可参考 tools/preprocess_data.py。2.1 IndexedDatasetBuilder构建与合并数据集IndexedDatasetBuilder位于 indexed_dataset.py约 930 行起能够构建并合并IndexedDataset实例。其核心方法包括add_item(tensor, mode0)向数据集追加单条序列并记录其长度add_document(tensor, lengths, modesNone)追加整篇文档lengths给出文档内各序列的长度同时更新document_indicesadd_documents(documents, modesNone, eod_tokenNone, chunk_size1_000_000)以批量方式写入 jagged 数组形式的文档列表依赖awkward库可自动在文档边界插入 EOD token分块写入以控制内存add_index(path_prefix)把另一个已存在的IndexedDataset整体合并进来拼接索引与数据finalize(idx_path)关闭数据文件并写入.idx索引文件。构造参数为bin_path数据文件路径、dtype默认numpy.int32与multimodal是否多模态决定是否记录每条序列的 mode。2.2 IndexedDataset.bin 与 .idx 的二进制布局IndexedDataset是 Megatron Core 中最低层的数据接口。一个实例引用两个二进制文件数据文件.bin保存文档/序列数据索引文件.idx保存文档/序列元数据。索引文件的内容分两部分。先存数据集级元数据索引文件头为向后兼容保留索引版本号为向后兼容保留一个数字编码对应写入数据文件所用的数据类型数据集中的序列数量数据集中的文档数量。随后存文档级与序列级元数据按顺序存每条序列的元素个数int32按顺序存每条序列的字节偏移指针int64按顺序存每个文档对应的连续序列索引区间[...)int64按顺序存每条序列的 mode仅多模态情形int8。这与源码中_IndexWriter.write()的落盘顺序完全一致header version dtype code sequence_count document_count sequence_lengths sequence_pointers document_indices sequence_modes其中_INDEX_HEADER bMMIDIDX\x00\x00版本号为小端8字节整型 1。数据类型编码由DType枚举定义int32对应编码 4、uint16对应编码 8 等。在读取侧_IndexReader会对.idx做numpy.memmap然后按偏移量依次还原sequence_lengthsint32、sequence_pointersint64、document_indicesint64多模态时再还原sequence_modesint8。.bin数据文件有四种读取器实现由IndexedDataset.__init__根据参数选择_MMapBinReadermmapTrue时内存映射读取默认本地文件场景_FileBinReadermmapFalse时用文件指针seek readinto内置指数退避重试默认 3 次重试起始睡眠 10s 翻倍_S3BinReaderS3 对象存储场景以bin_chunk_nbytes为块大小维护内存缓存按块范围Range拉取字节_MultiStorageClientBinReader基于 Multi-Storage Client 的范围读。此外还暴露get(idx, offset, length)方法支持只取序列一部分的读取这正是上层GPTDataset拼接跨文档样本时依赖的关键能力。IndexedDataset还支持fast_cache_load跳过文件存在性断言与sequences_per_dataset直接用计数信息初始化索引免开.idx文件配合--per-dataset-sequences-path使用。三、数据加载之构建BlendedMegatronDatasetConfig 与 BuilderDataLoader 的构建是一个分布式感知的过程围绕五个要素展开BlendedMegatronDatasetConfigBlendedMegatronDatasetBuilderIndexedDatasetMegatronDatasetBlendedDataset3.1 BlendedMegatronDatasetConfig可扩展该类blended_megatron_dataset_config.py参数化BlendedMegatronDatasetBuilder进而参数化MegatronDataset与BlendedDataset。不同训练/推理体制需要不同的扩展例如 GPT 训练使用的GPTDatasetConfig。配置类中的关键字段包括random_seed、sequence_length随机种子与序列长度必填blend[[prefix1, prefix2], [0.3, 0.7]]形式的混合定义权重为 None 时按底层数据集长度推断不能与blend_per_split同用blend_per_splittrain/valid/test 三个 split 各自独立的混合定义split从单一分布抽样时 train/valid/test 的权重字符串如99,1,0不能与blend_per_split同用path_to_cache所有可复用数据集索引的缓存目录mmap_bin_files默认 True.bin用 mmap 还是文件指针读取mid_level_dataset_surplus默认 0.005中层数据集构建时的样本冗余比例顶层数据集若超采中层数据集需要调大fast_cache_load/defer_npy_index_mmap两个快速加载开关见第七节都要求path_to_cache非空num_dataset_builder_threads构建数据集的线程数。__post_init__中完成了若干合法性约束fast_cache_load时禁止与blend同用应改用--per-split-data-args-path或 per-split data pathblend与blend_per_split互斥当blend非空时必须提供split两者都为空时自动进入mockTrue的模拟数据模式并取1,1,1的均匀 split。split 字符串经parse_and_normalize_split归一化后再由convert_split_vector_to_split_matrix转换为 split 矩阵各 split 在非重叠区间上的书端点区间例如[0.99, 0.01, 0.0] - [(0, 0.99), (0.99, 1.0), None]。GPTDatasetConfiggpt_dataset.py在基类之上扩展了reset_position_ids、reset_attention_mask、eod_mask_loss、create_attention_mask、add_extra_token_to_sequence、drop_last_partial_validation_sequence、hybrid_context_parallel、sequences_per_dataset等字段并在__post_init__中依据 tokenizer 词表大小自动推导token_dtype_code词表超过 uint16 上限取 int32 编码 4否则 uint16 编码 8。3.2 BlendedMegatronDatasetBuilder最高层构建器BlendedMegatronDatasetBuilderblended_megatron_dataset_builder.py注意正确路径为 megatron/core/datasets/blended_megatron_dataset_builder.py构建 Megatron Core 最高层的数据接口。其__init__接收四类参数cls要实例化的MegatronDataset子类sizes每个 split 要求的最少样本总数可为 Noneis_built_on_rank判断当前 rank 是否构建数据集的可调用对象必须感知 Megatron Core 并行策略全局 rank、组内 rank、virtual rank 都可能影响返回值且必须在全局 rank 0 上恰好返回 Trueconfig数据集配置对象。build()方法按配置分三种情形处理每个 splitsplit 为 None什么都不做单个贡献数据集size非 None 时按比例抽样子数据集size为 None 时不做超额抽样多个贡献数据集权重与 size 均给定 → 按权重与 size 构建中层与顶层数据集仅给权重不给 size → 报错仅给 size → 顶层长度取size以各中层长度之和为上限两者皆无 → 构建穷举索引。值得注意的实现细节均可在源码中验证构建过程使用ThreadPoolExecutor并行构建各 prefix 的MegatronDataset_build_megatron_datasets_parallel每个任务通过contextvars.copy_context()把 OTel 追踪上下文传播进工作线程分布式初始化后默认 rank 0 先行构建线程数还会按 GPU 数适度放大barrier之后其他 rank 再构建——此时必然命中缓存开启fast_cache_load时会跳过 rank 0 先行 barrier 的同步点各 rank 直接并行构建/加载这正是其提速原理见第七节指定了混合权重时每个 prefix 构建的样本数由_get_size_per_split_per_dataset计算会乘以(1 mid_level_dataset_surplus)冗余系数保证顶层数据集有足够样本可抽。3.3 MegatronDataset可扩展与 BlendedDataset顶层MegatronDatasetmegatron_dataset.py是建立在IndexedDataset之上的高阶抽象基类子类如GPTDataset需实现numel_low_level_dataset与build_low_level_dataset两个静态方法。构造时它会汇总类名、数据集路径、样本数、split、以及_key_config_attributes()返回的关键配置属性random_seed、sequence_length、split、split_matrix、tokenizer序列化为 JSON 描述并计算 MD5 得到unique_description_hash——缓存索引文件名的唯一性就锚定在这个 hash 上因此任何关键参数变化都会导致缓存失效重建。BlendedDatasetblended_dataset.py建立在多个MegatronDataset之上仅当需要混合多个数据分布来贡献某个 split 时才需要它混合比例通过配置控制。约束包括各子数据集必须同类、同 split、权重为正、数据集聚合数小于 32767 等若size非 None 则权重会被归一化。四、数据加载之实现GPTDataset 的三索引机制GPTDataset由以下变量参数化底层IndexedDataset实例indexed_dataset、split 索引indexed_indices用于训练/验证/测试的连续文档或序列索引子集、样本数N、序列长度S、随机种子R。它创建三个索引映射以支撑查表1文档索引 Do_idx一维数组把i映射到文档索引长度E * |indexed_indices|其中E是满足E * |indexed_indices| N的最少 epoch 数。文档索引按R洗牌。Given: N 15 indexed_indices [5, 6, 7, 8, 9] E 3 Then, for example: Do_idx [8, 8, 9, 6, 7, 5, 8, 5, 6, 6, 5, 9, 7, 7, 9]2样本索引 Sa_idx二维数组把j映射到(i, Do_idx[i] 的偏移)对形状[N 1, 2]。行j与j 1分别作为第j个样本的左、右边界。Given: S 1024 Then, for example: Sa_idx[0] (0, 0) Sa_idx[1] (0, 1024) Do_idx[0] has length greater than S Sa_idx[2] (1, 512) Do_idx[0] has length 1536 Sa_idx[3] (2, 0) Do_idx[1] has length 1536 Sa_idx[4] (5, 300) Do_idx[2:5] are shorter documents relative to Do_idx[0:2] Sa_idx[5] (6, 24) Do_idx[5] has length 13003洗牌索引 Sh_idx一维数组把k映射到j长度N按R洗牌。Given N 10 Then, for example: Sh_idx [4, 0, 2, 6, 1, 9, 5, 8, 7, 3]查询第k个样本的过程# 1. 用洗牌索引得到样本索引内的索引 j j Sh_idx[k] # 2. 用样本索引得到样本左右边界在文档索引中的位置及各自起始 token 偏移 i, offset Sa_idx[j] i_next, offset_next Sa_idx[j 1] # 3. 用文档索引从连续的文档中取出 S 个 token sample [] sample indexed_dataset[Do_idx[i]][offset:] if i ! i_next: sample indexed_dataset[Do_idx[i 1:i_next]] sample indexed_dataset[Do_idx[i_next]][:offset_next]从源码看gpt_dataset.py 的_query_document_sample_shuffle_indices约 388-473 行实际实现与伪代码一致并做了两点增强样本跨越单个文档时直接用dataset.get(idx, offset, length)一次取出跨越多个文档时逐文档get拼接。此外取出的 token 总数为S add_extra_token_to_sequence默认多取 1 个 token__getitem__中据此切分出tokens text[:-1]与labels text[1:]保证输入与标签都是完整长度不足时以 pad token 补齐并把 pad 位置的 loss mask 置零。索引构建的关键工程细节_build_document_sample_shuffle_indices约 475-701 行epoch 计算_get_num_epochs会不断累加 epoch 直至 token 总量满足N * S add_extra_token的需求末 epoch 分离若最后一个 epoch 的样本数不足完整 epoch 样本数的 80%threshold 0.80则把末 epoch 与前面 epoch 分开洗牌separate_final_epoch避免最后一个不完整的 epoch 被全局打散后与前面样本过度混合C 加速样本索引由 helpers.cpp 中的build_sample_idx构建当len(document_index) * 2 len(sequence_lengths)访问密度高时会先把 mmap 的sequence_lengths复制进内存源码注释解释了这样做的两个好处——顺序预读整个文件以及进入 C 时持有 GIL 提高并行度缓存三个索引与description.txt一并保存到缓存目录文件名为{unique_description_hash}-{ClassName}-{split}-document_index.npy等fast_cache_load时跳过文件存在性检查直接视为命中defer_npy_index_mmap时索引不在初始化时加载、延迟到首次访问时以 mmap 方式加载此时__len__会改用纯算术公式估算样本数复用 helpers.cpp 的样本计数逻辑。4.1 BlendedDataset 的混合索引BlendedDataset由数据集聚合D、权重W每个数据集一个与规模S参数化。它会按权重比例从各贡献数据集抽样直到达到目标规模每一步抽样时从抽样误差sampling error最大的那个数据集抽取一个样本。它创建两个混合索引数据集索引 Da_idx一维数组把i映射到数据集索引长度SGiven D [d0, d1, d2] W [1/2, 1/4, 1/4] S 4 Then, for example: Da_idx [0, 1, 2, 0]数据集样本索引 Sa_idx一维映射把i映射到数据集Da_idx[i]内的样本索引长度SGiven Da_idx [0, 1, 2, 0] Then, for example: Sa_idx [0, 0, 0, 1]查询第k个样本sample D[Da_idx[k]][Sa_idx[k]]同样为节省初始化时间各索引在单个 rank 上顺序构建/缓存再由其他 rank 并行加载缓存索引锚定在BlendedDataset.__init__生成的 hash 上。源码实现blended_dataset.py 的_build_indices中索引由 helpers.cpp 的build_blending_indicessize非 None或build_exhaustive_blending_indicessize为 None穷举模式构建构建后还会校验各子数据集是否被超采若超采会抛出明确提示增大mid_level_dataset_surplus的IndexError。五、离线缓存预生成tools/prepare_cache.py对于 GPT 风格训练上述数据集缓存可以用 tools/prepare_cache.py 提前准备好而不必等训练启动时 rank 0 构建。该脚本复用了pretrain_gpt.py与pretrain_mamba.py的正常数据集构建路径包括GPTDataset、BlendedDataset与BlendedMegatronDatasetBuilder。它接受常规的数据集参数支持 blend 与 per-split 数据集定义并要求--data-cache-path以便生成的缓存能被后续训练复用。对于大型 blend 或大量文件 prefix 的场景尤其有用构建 document、sample、shuffle 索引可能耗时数分钟期间所有 GPU 都处于空闲状态而 rank 0 只做纯 CPU 工作。如果后续训练任务没有指定--global-batch-size该参数用于确定数据集规模与 split应通过--prepare-cache-world-size显式指定缓存准备时使用的 world size脚本将其直接赋给args.world_size见 prepare_cache.py 的_normalize_prepare_cache_args。明确的限制tools/prepare_cache.py不支持--mock-data、--sft、--fim-data或--step-batch-size-schedule源码_validate_prepare_cache_args会对这些选项直接抛ValueError同时--data-cache-path为必填。脚本还会在准备阶段强制关闭--dataloader-fast-cache-load与--dataloader-defer-npy-index-mmap这两个开关的意义只在于消费已存在的缓存并在运行前打印生效的 world size、DP size、global batch size、缓存路径与各 split 目标样本数。六、Packing Scheduler跨 DP×CP rank 的变长序列重调度Packing 调度器把变长序列重新调度到 DP×CP 各 rank 上以提升 GPU 利用率。它围绕以下模块构建data_scheduledata_schedule.py 包含高层调度逻辑与入口点HybridCPDataLoaderWrapper混合上下文并行CP调度的包装类。每次__next__调用它会(1) 从各 DP rank 拉取一批 packed 样本(2) 在 DP 组内 all-gather 序列长度(3) 用BalancedCPScheduler来自 megatron/core/pipeline_parallel/ 的hybrid_cp_schedule调度子样本(4) 通过 all-to-all 通信把子样本重路由到正确的 DPxCP rank。BasePackingScheduler打包调度器的抽象基类定义了get_groups_and_subsamples()调度算法与run()完整调度流水线fetch、schedule、reroute、pack、broadcast 及 VPP 处理的接口。DpBalancedScheduler具体调度器按原始顺序打包序列直到达到每个 DPxCP rank 的最大序列长度限制支持把 microbatch 数对齐到 DP size 与 VPP stage 的整数倍。wrap_data_iterator()顶层入口包装已有的data_iterator。它创建合适的调度器、运行调度流水线、广播元数据与新的num_microbatches返回新的数据迭代器、更新后的 microbatch 数以及 FLOPs 统计。get_batch_on_this_rank_for_sequence_packing()为当前 rank 拉取并广播单个 packed microbatch。处理 TP/PP 广播构造PackedSeqParams含cu_seqlens、max_seqlen、qkv_formatthd并可选地用 Transformer Engine 的thd_get_partitioned_indices在 CP rank 间划分序列。data_schedule_utilsdata_schedule_utils.py 包含调度器使用的工具函数如broadcast_scalars、broadcast_tensor、build_packed_microbatches、reroute_samples_to_dcp_ranks等data_schedule.py顶部 import 即列明了全部依赖项。七、快速 DataLoader 初始化三个加速开关大规模训练中DataLoader 初始化可能耗时数分钟——因为要打开并内存映射大量文件还会显著施压文件系统。Megatron Core 提供了三个由配置开关控制的优化7.1 --dataloader-fast-cache-load假定数据集缓存已存在于指定的--data-cache-path中。启用后通过移除同步点与文件检查断言来加速创建过程。从源码看其具体生效点有三处配置层blended_megatron_dataset_config.py 断言必须提供--data-cache-path且不能与--data-pathblend 形式同用应改用--per-split-data-args-path或--train-data-path/--valid-data-path/--test-data-path构建层blended_megatron_dataset_builder.py 跳过 rank 0 先行构建 barrier各 rank 直接并行构建同时跳过indexed_indices的重复计算数据层indexed_dataset.py 跳过sequence_lengths.shape[0]系列的一致性断言。7.2 --dataloader-defer-npy-index-mmap同样假定缓存已存在。启用后把数据集索引.npy文件的内存映射延迟到首次访问时进行。官方推荐与--num-workers 0搭配使用让 DataLoader 预取下一批数据从而用后台预取掩盖索引 mmap 的开销。实现上GPTDataset/BlendedDataset在_build_indices阶段只记录缓存路径并返回 None 索引__getitem__首次调用时才safe_numpy_load(..., mmap_moder)见 gpt_dataset.py 与 blended_dataset.py__len__则改用 token 数算术公式直接推算。7.3 --per-dataset-sequences-path通过该配置指定 tools/build_sequences_per_dataset.py 生成的 JSON 文件。该脚本对 blend 中的每个文件 prefix 打开.idx读取序列数与文档数_IndexReader汇总为单一文件。此配置在处理数百乃至上千个文件 prefix 时尤其有用它只需要一次open操作而不是每个 prefix 一次。该 JSON 经GPTDatasetConfig.sequences_per_dataset传入IndexedDataset后_IndexReader可跳过解析 34 字节头部之后的完整索引读取直接用给定的(sequence_count, document_count)初始化见 indexed_dataset.py。脚本用法示例来自其模块 docstringpython3 tools/build_sequences_per_dataset.py --per-split-data-args-path my-training-dataset-blend.json --per-dataset-sequences-path my-training-dataset-blend-sequences-per-dataset.json八、小结一次训练启动中数据管线的工作顺序把以上内容串起来一次 GPT 训练启动时数据管线的工作顺序为训练脚本通过megatron.training的参数解析生成GPTDatasetConfig含 blend、split、path_to_cache等并计算各 split 目标样本数所有 rank 调用BlendedMegatronDatasetBuilder.build()rank 0 先构建或用--dataloader-fast-cache-load并行构建为每个 prefix 建IndexedDatasetmmap.bin解析.idx每个GPTDataset按unique_description_hash检查缓存命中则 mmap 加载三个.npy索引或按defer_npy_index_mmap延迟加载未命中则由 rank 0 构建C 加速并写缓存多个数据集的 split 由BlendedDataset以 最大抽样误差 策略生成混合索引运行时__getitem__依次经 shuffle → sample → document 三级查表拼接出定长样本再经 DataLoader及其可选的HybridCPDataLoaderWrapperpacking 调度送入模型。相关源码与工具入口速查主题路径底层二进制接口megatron/core/datasets/indexed_dataset.pyC 索引构建加速megatron/core/datasets/helpers.cpp / helpers.pyGPT 数据集与三索引megatron/core/datasets/gpt_dataset.py混合数据集megatron/core/datasets/blended_dataset.py分布式构建器megatron/core/datasets/blended_megatron_dataset_builder.py配置数据类megatron/core/datasets/blended_megatron_dataset_config.py抽象基类megatron/core/datasets/megatron_dataset.py离线缓存预生成tools/prepare_cache.py每数据集元数据生成tools/build_sequences_per_dataset.pyPacking 调度器megatron/core/datasets/data_schedule.py / data_schedule_utils.py适用前提提示以上行为均基于当前仓库版本快速加载类开关fast cache load、defer mmap、per-dataset sequences都要求缓存已预先构建且配置了--data-cache-pathtools/prepare_cache.py不支持 mock/SFT/FIM/step-batch-size-schedule 路径离线预生成缓存时应避免这些模式。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考