Megatron-LM 训练数据管线深度解析：FIM 数据增强与变长打包（THD/Varlen）数据集实战指南

Megatron-LM 训练数据管线深度解析FIM 数据增强与变长打包THD/Varlen数据集实战指南【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM本文聚焦 Megatron-LM 数据管线的两大核心组件面向代码/预训练语料的Fill-in-the-MiddleFIM数据增强数据集GPTFIMDataset与面向指令微调/长上下文预训练的变长打包THD数据集VarlenDataset。读者将掌握FIM 的 PSM/SPM 两种变换格式与全部配置参数、VarlenDataset的四种输入 schema 自动检测与规范化规则、命令行启用方式以及二者在源码层面的底层实现与打包调度调用链可直接迁移到自己的训练脚本中。总览Megatron-LM 训练侧的两条数据路径Megatron-LMMegatron-Core 训练侧封装的训练数据管线位于 megatron/training/datasets/ 目录包含__init__.py、data_samplers.py、fim_dataset.py、sft_dataset.py、utils.py与varlen_dataset.py六个文件。其中 README.md 文档化了其中两个最具代表性的数据集数据集基类用途触发入口GPTFIMDatasetMegatron-Core 的GPTDataset将预训练样本按概率变换为 FIM 格式代码模型、续写预训练--fim-dataVarlenDatasetSFTDataset家族将长度差异极大的 SFT/长上下文语料打包为 THD 变长格式--use-varlen-dataset二者共同支撑了「预训练数据增强」与「指令/长文本高效训练」两条主路径是理解 Megatron-LM 训练侧数据流的关键切入点。FIM 数据集为代码与大模型预训练注入中间填充能力什么是 FIM 与为什么需要它Fill-in-the-MiddleFIM中间填充是代码大模型预训练的重要数据增强手段将一段文本随机拆成 prefix前缀、middle中间、suffix后缀三部分并重排让模型学习根据上下文补全中间内容的能力。它最初由 FIM 相关研究论文提出如 Codex 的 PSM 变体可显著提升代码补全与长程依赖建模能力。Megatron-LM 通过 fim_dataset.py 实现这一增强GPTFIMDatasetConfig继承自 Megatron-Core 的GPTDatasetConfig的配置对象fim_dataset.pyGPTFIMDataset继承GPTDataset的数据集类从IndexedDataset加载 token 序列在返回每个样本前应用 FIM 变换fim_dataset.py。配置属性与语义GPTFIMDatasetConfigGPTFIMDatasetConfig的每个属性都有明确的概率与 token 语义属性字段名CLI含义rate--fim-rate默认0.5将样本转换为 FIM 示例的概率。1.0表示总是应用 FIM0.0表示从不应用spm_rate--fim-spm-rate默认0.5使用 SPM 模式相对 PSM的概率剩余概率1 - spm_rate选择 PSMprefix-suffix-middle。例如spm_rate 0.3表示 30% SPM、70% PSMextra_tokens--fim-prefix-token/--fim-middle-token/--fim-suffix-token/--fim-pad-token/--fim-eod-token包含 FIM 特殊 token 的字典{prefix, middle, suffix, pad, eod}split_sample--fim-split-sample可选的切分 token应用 FIM 前按该 token 切分样本。若提供输入序列在每一处该 token 处被切分FIM 独立应用于每个片段。例如A B C SPLIT_SAMPLE D E F SPLIT_SAMPLE G H→FIM(Fragment 1) SPLIT_SAMPLE FIM(Fragment 2) SPLIT_SAMPLE FIM(Fragment 3)fragment_rate--fim-fragment-rate使用split_sample时对每个片段应用 FIM 的概率no_prefix--fim-no-prefix若解码后的序列以此前缀开头则跳过 FIM对应的 CLI 参数注册在 arguments.py默认 FIM 特殊 token 为fim_prefix、fim_middle、fim_suffix、fim_pad、|endoftext|。PSM 与 SPM 两种变换格式文档明确给出两种格式的 token 排列PSM 格式prefix-suffix-middle[prefix_tok] prefix [suffix_tok] suffix [middle_tok] middleSPM 格式suffix-prefix-middleFIM 论文中的 variant 2[prefix_tok, suffix_tok] suffix [middle_tok] prefix middle源码中两种格式的拼接实现在_permute方法中fim_dataset.pySPM 将prefix_tok_id与suffix_tok_id两个哨兵 token 相邻放置随后是 suffix、middle_tok、prefix 与 middlePSM 则按 prefix、suffix_tok、suffix、middle_tok、middle 的顺序排列。是否选择 SPM 由一次np_rng.binomial(1, fim_spm_rate)抽样决定。特殊边界情况文档列出的两个特殊情形在源码中均有对应实现序列以no_prefix开头 → 跳过 FIM_permute中先将样本解码为文本ids_to_text(..., remove_special_tokensTrue)若以no_prefix开头则原样返回fim_dataset.py未应用 FIM → 样本原样返回当伯努利抽样binomial(1, fim_rate)结果为 0 时_permute直接返回原始样本fim_dataset.py。源码级实现细节切分、片段级 FIM 与长度保持GPTFIMDataset的核心变换逻辑值得深入理解EOD 段级切分_query_document_sample_shuffle_indices在取回样本后先按eod_tok_id定位文档边界对每个非空段独立执行 FIM 变换并重新拼接fim_dataset.py保证跨文档的样本不会把两个不同文档混进一次 FIM片段级 FIMsplit_sample 模式当配置了fim_split_sample样本先在该 token 处被切分为多个文件级片段每个片段以fragment_fim_rate概率独立 FIM切分 token 保留在片段之间fim_dataset.py。这非常适合仓库级语料一个样本是仓库、每个片段是一个文件的场景片段内没有切分 token 时回退为整段 FIM。fim_split_sample字符串在__init__中被一次性转换为 token idfim_dataset.py长度保持变换后若超出原长度则右侧截断sample[:sample_len]不足则用pad_tok_id补齐fim_dataset.py。_permute内部同样对三个片段做长度平衡新增 3 个哨兵 token 后若过长优先截断 suffix过短则用 pad token 填充fim_dataset.py可复现性FIM 的随机状态使用config.random_seed初始化独立的np.random.RandomStatefim_dataset.py保证多轮训练/断点续训的 FIM 随机序列可复现。命令行启用方式# 基础启用50% 样本做 FIM其中 50% 用 SPM --fim-data --fim-rate 0.5 --fim-spm-rate 0.5 # 仓库/文件级 FIM按 file_sep 切分每个文件片段 100% FIM --fim-data --fim-split-sample file_sep --fim-fragment-rate 1.0 # 跳过以 // 开头的注释片段 --fim-no-prefix // 需要说明的是FIM 的启用还依赖训练侧参数校验assert args.fim_rate与assert args.fim_spm_rate见 arguments.py且GPTFIMDatasetConfig中的extra_tokens字典在__init__中被按prefix/middle/suffix/pad/eod顺序批量转换为 token idfim_dataset.py因此这些 FIM token 必须已加入 tokenizer 词表。Varlen 数据集变长指令数据的 THD 打包设计目标与定位VarlenDataset将长度差异极大的 SFT 式指令数据打包为 THDTHD 即 variable-length packed sequence格式。它扩展SFTDataset家族因此复用相同的打包 /cu_seqlens/ 上下文并行CP填充逻辑并且独立于--sft标志选择——文档明确强调二者无隐式耦合见 varlen_dataset.py 的模块 docstring。其入口是--use-varlen-dataset标志arguments.py且与--sft互斥arguments.py。与SFTDataset相比VarlenDataset新增三大能力多源加载接受 HuggingFace Hub 仓库 idowner/repo、本地.parquet文件、本地.jsonl/.json文件Schema 自动检测按列名推断四种输入布局Mock 变体MockVarlenDataset端到端镜像MockSFTDataset合成对数正态序列长度分布 / 定长文件 / 基于IndexedDataset的验证模式通过--varlen-mock-dataset-config-json配置。Schema 自动检测输入格式推断文档给出了完整的检测优先级表布局由数据集列名推断最显式者优先首个匹配生效Schema检测依据归一化为openai-messages存在messages列直接透传sharegpt存在conversations列messages 列表alpaca/dolly同时存在 instruction 列与 output 列3 轮 messages 列表pretrain-text存在text列原始字符串无聊天模板alpaca/dolly 布局接受的列名源码中的同义词元组定义于 varlen_dataset.pyinstructioninstruction、prompt、query、questionoutputoutput、response、completion、answer可选额外用户轮上下文inputStanford Alpaca、contextDolly选择逻辑实现在_select_convertervarlen_dataset.py依次检查messages→conversations→ instructionoutput →text。若四种布局均不匹配数据集构造直接抛出ValueError错误信息中列出实际看到的列名与支持的 schema。归一化规则Normalization Rules三种指令微调布局都会转换为父类SFTDataset期望的 messages 列表保证 leadingsystem轮样本未以 system 轮开始时自动前置一个空 system 轮varlen_dataset.py、varlen_dataset.py这样SFTDataset._split_conversations会把每个样本当作一个对话其实现按 system 消息切分对话见 sft_dataset.pyShareGPT 说话人 → 角色映射通过from字段映射_SHAREGPT_ROLE_MAPvarlen_dataset.pyhuman/user→usergpt/assistant/model/chatgpt/bing/bard→assistantsystem→systemtool/function/observation→tool。无法识别的说话人回退为user而不是报错Alpaca/Dolly 上下文并入 user 轮存在input/context时用空行与 instruction 拼接f{instruction}\n\n{extra_input}见_alpaca_to_messagesvarlen_dataset.py丢弃非role/content键messages样本中的name、tool_calls等键被剔除varlen_dataset.py因为它们不属于聊天模板输入。pretrain-text是例外它原样返回text列的原始字符串数据集据此跳过聊天模板与 prompt masking。VarlenDataset.__getitem__通过返回值类型分派str走普通tokenizer.tokenize路径并把每个 token 都作为目标无 prompt 掩码messages 列表则走tokenizer.tokenize_conversation(..., return_targetTrue)的聊天模板路径varlen_dataset.py。这支持 Dolma、OLMo midtraining 这类长上下文预训练语料走与 SFT 相同的 THD 打包路径。局限性明确的失败语义文档列出的限制在源码中都以显式ValueError表达而非静默错误处理轮次内容必须是普通字符串以图片/文本 part 列表形式承载内容的多模态样本会抛ValueError_ensure_str_contentvarlen_dataset.pyinstruction/output 字段的非字符串值抛ValueError而非强制转换_first_presentvarlen_dataset.py此外树状结构OpenAssistant oasst1与偏好数据集chosen/rejected不在支持范围内HF Hub 仓库仅加载splittrain见 varlen_dataset.py。VarlenDataset 底层实现多源加载、样本结构与 THD 打包多源加载与列名检测VarlenLowLevelDataset继承SFTLowLevelDatasetvarlen_dataset.py按如下规则解释dataset_pathHF Hub 仓库 id含/且本地文件系统不存在如Yukang/LongAlpaca-12kload_dataset(path, splittrain)本地.parquetload_dataset(parquet, data_filespath, splitall)parquet 的 footer schema 使分块加载安全其他情况按 jsonl/json 处理用 stdlibjson逐行解析后经Dataset.from_list一次构建。刻意避开datasets.load_dataset(json, ...)——其 pyarrow JSON 读取器按并行分块推断 schema当行间字段并集变化时如 LongAlpaca-12k会抛CastErrorfrom_list则在单次 pass 内统一整个文件的 schema。路径判定启发式_looks_like_hf_idvarlen_dataset.py要求含/、不是绝对/相对文件路径、且本地不存在。列名到转换器的映射在构造时一次性完成访问时才应用。样本结构单条未打包对话 original_seq_len/padded_seq_len与父类SFTDataset.__getitem__在内部完成多对话预打包并输出固定sequence_length长度见 sft_dataset.py不同VarlenDataset.__getitem__每次只返回一条已 tokenize 的未打包对话包含tokens/labels/loss_mask/position_ids其长度等于样本实际 token 数仅填充到pad_granularity不填充到sequence_length外加original_seq_len/padded_seq_len张量——由上游打包调度器BasePackingScheduler.get_required_sample_keys定义的 schema直接消费varlen_dataset.py并经get_batch_and_global_seqlensdata_schedule_utils.py收集各 DP rank 的序列长度后跨 DP×CP 网格调度打包避免逐样本填充浪费。处理流程要点截断超过config.sequence_length的样本右侧截断到sequence_length 1若截断去掉了 EOD 则补一个varlen_dataset.py短样本末尾无条件保证 EODvarlen_dataset.py空 tokenization 守卫空白pretrain-text行 tokenize 结果为空时用[eod, eod]表示保证 next-token shift 后仍产生合法的 1-token 样本varlen_dataset.pypad token 回退tokenizer 无显式 pad token如 Qwen3 原始预训练 tokenizer时回退用eod填充因loss_mask会把 pad 位置置零而不影响 lossvarlen_dataset.pypad_granularity由_calculate_padding_divisor计算sft_dataset.py即 CP 对齐标准 CP 下cp_size * 2hybrid CP 下dp_size * cp_size * 2乘以序列并行sp_size这是 CP 切片所需的最小对齐故意不对齐到sequence_lengthloss_mask右侧填充尾与IGNORE_INDEX-100prompt 位置均置零varlen_dataset.py。SBHD 验证模式--varlen-sbhd-validation为验证 THD 路径数值正确性VarlenDataset实现了参考用的 SBHD 模式右填充到sequence_length 1、丢弃打包元数据、输出[sequence_length]形状作为与 THD 路径对比的数值参考varlen_dataset.py。该模式与--sequence-packing-scheduler不兼容且不支持--mock-dataMockVarlenDataset未实现该模式由参数校验守卫见 arguments.py。Mock 变体与参数配置# 启用变长数据集自动选择 dp_balanced 打包调度器 --use-varlen-dataset --data-path hf_repo_or_local_file # 数值验证SBHD 参考路径不启用打包调度器 --use-varlen-dataset --varlen-sbhd-validation # Mock 数据 默认对数正态分布 --use-varlen-dataset --mock-data # Mock 数据 自定义分布/文件/IndexedDataset 验证模式 --use-varlen-dataset --mock-data \ --varlen-mock-dataset-config-json \ {mode:distribution,type:lognormal,min_seq_len:1024,max_seq_len:2048,mean_seq_len:1536,lognormal_sigma:1.1}--use-varlen-dataset未显式给定打包调度器时会自动选择dp_balanced--varlen-mock-dataset-config-json接受内联 JSON 或文件路径schema 与--sft-mock-dataset-config-json相同mode: file/distribution/verification未指定时默认对数正态分布min_seq_lenseq_length//2、max_seq_lenseq_length、mean_seq_lenseq_length*3//4、lognormal_sigma1.1arguments.py。MockVarlenDataset的输出形状与VarlenDataset.__getitem__完全一致THD 模式一条未打包样本 original_seq_len/padded_seq_len确保 mock 与真实数据路径锻炼完全相同的下游管线varlen_dataset.py。与打包调度器的衔接VarlenDataset的价值最终体现在与序列打包调度器的配合上。get_batch_and_global_seqlensdata_schedule_utils.py从data_iterator拉取num_microbatches个批次并展平为未打包子样本收集每个子样本的padded_seq_len通过_get_global_seqlens_and_ids在所有 DP rank 间 gather产出global_id_seqlens、global_ids_this_rank、offsets与seqlens_gathered四类调度输入调度器据此把变长样本跨 DP×CP 网格组合进微批次cu_seqlens/max_seqlen在随后的_pack_sequences中生成。该调用链在 data_schedule.py 处被消费注释也明确指出sft_dataset.py中序列已预打包因此需要先解包再重新调度data_schedule_utils.py——这也是VarlenDataset刻意输出未打包单样本的原因把打包决策完全交给上游调度器消除逐样本填充浪费。实践建议与使用要点FIM 与代码预训练仓库级语料建议配合--fim-split-sample做文件级切分与--fim-fragment-rate片段级 FIM记得为fim_prefix/fim_middle/fim_suffix/fim_pad及 EOD 保留词表位置且所有 FIM token 必须存在于 tokenizer 词表中__init__会直接查 token id变长 SFT 与长上下文预训练--use-varlen-dataset独立于--sftLLaMA 后训练类messages数据、OpenOrca/Vicuna 类conversations数据、Alpaca/Dolly 类 instruction 数据以及 Dolma/OLMo 类text语料均可自动检测字段缺失时构造会抛带列名清单的ValueError便于快速定位 schema 问题可复现与验证FIM 随机性由random_seed控制THD 路径正确性可用--varlen-sbhd-validation与参考实现数值比对Mock 数据可脱离真实数据快速跑通全管线并做吞吐基准测试约束提醒--use-varlen-dataset与--sft互斥多模态 content 列表与偏好/树状数据集不在 Varlen 支持范围HF Hub 仅加载train分片。以上两条数据路径均可在 megatron/training/datasets/ 下找到完整实现与单元测试对应物读者可结合 data_schedule.py、data_schedule_utils.py 与 sft_dataset.py 继续深入阅读构建符合自身语料形态的高效训练数据管线。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考