做 LLM 预训练的朋友应该都有这个体会真正把大模型跑起来之前数据侧的功夫往往比模型侧还要多。最近我在 MindSpore 环境里复现一个 Transformers 架构的 GPT 类模型参考了 Megatron 风格的数据预处理方案把原始文本语料转成可以直接喂给训练循环的二进制数据集。整个过程踩了不少坑也把 Megatron 那套 bin/idx 格式的设计逻辑摸了个底朝天。这篇东西就围绕“MindSpore Transformers LLM Megatron 风格数据集预处理”这条主线把从语料清洗、tokenizer 编码、二进制落盘到 MindSpore 侧数据加载的完整链路讲清楚适合正在做 LLM 预训练、微调或数据处理模块自研的工程师参考。很多人第一次接触 Megatron 风格预处理时会觉得这不就是把文本转成 token id 再存下来么有什么好讲的。实际动手后才发现里面有一堆细节决定训练能不能跑得稳文本要不要拼接、document 之间用什么符号隔开、每个样本到底切多长、标签到底是什么、数据文件要不要乱序、和 MindSpore 的 Dataset 接口怎么对接。这些点任何一个出错轻则 loss 不收敛重则训练到一半崩溃。下面我把完整方案拆开讲。1. 项目背景与整体设计思路1.1 为什么要在 MindSpore 里复刻 Megatron 的预处理逻辑先回答一个最基础的问题MindSpore 生态里明明有 MindRecord 这种自研数据格式为什么还要绕一圈去复刻 Megatron 风格的 bin/idx我的答案很简单因为大部分公开的 LLM 预训练语料、开源模型权重和训练框架之间的数据对齐默认就是按 Megatron 那套约定做的。举个例子很多人从公开渠道下载 RedPajama、The Pile 这类语料时会发现它们常常已经被人预处理成 bin/idx 文件分发。如果训练侧只认 MindRecord就得先把 bin/idx 转回去再转过来中间既浪费时间又容易丢元信息。更麻烦的是团队里如果有多套卡、多套框架并行做实验HuggingFace Transformers 侧读的是 tokenized 文本Megatron 侧读的是 bin/idxMindSpore 侧如果又是另一套那数据集的一致性根本没法保证。所以我的选择是直接用 Megatron 风格数据作为统一中间格式MindSpore 侧通过自定义 Dataset 加载器去读 bin/idx。这样语料生产一次三套框架都能吃模型侧评测对比时也能保证喂进去的数据完全一致。1.2 从原始文本到训练样本要经历的完整链路整条预处理的链路可以分成四段。第一段是语料准备与清洗解决“文本里有什么不该有”的问题第二段是 tokenizer 编码解决“文本怎么变成数字”的问题第三段是样本切分与二进制落盘解决“数字怎么组织成训练样本”的问题第四段是数据加载与 shuffle解决“训练循环怎么高效读取”的问题。前三段通常在离线脚本里完成输出就是 bin 和 idx 两个文件第四段在训练进程里完成由数据加载器负责。这个划分我觉得非常舒服因为离线阶段可以随便跑大内存任务在线阶段只需要做 IO 和采样不会因为数据预处理占用训练机的 CPU 资源。设计时需要特别留意一个原则预处理时不要做任何跟模型结构强绑定的决策。比如要不要做 attention mask、样本内部的 key-value 怎么排、micro batch 怎么切这些应该留到训练侧去处理。数据侧只负责产出“干净的 token 序列”和“document 边界信息”最多再提供一个 sequence 级别的切分维度这是 Megatron 数据格式兼容多种并行策略的关键。2. Megatron 风格数据格式的核心原理2.1 bin/idx 文件到底在存什么Megatron 风格数据集由两个文件组成.bin和.idx。bin 文件是纯二进制内容里面按顺序存了一串 uint16 或 int32 的 token id。idx 文件是索引文件记录 bin 文件里每个 document 和每个样本的边界信息让加载器可以随机访问任意一个样本而不必把整个 bin 读进内存。bin 文件的组织方式非常直白多个 document 的 token id 首尾相连拼接成一个长数组后整体写入。document 之间不额外存分隔符因为在编码阶段就已经把 eod tokenend of document比如|endoftext|对应的 id当成普通 token 拼进序列里了。所以 bin 文件里实际存的是一个连续的、带有特殊 eod token 标记的 token 流。idx 文件的结构稍微复杂一点可以理解成一个小型索引和大型索引的组合。小型索引保存版本号、整个数据集包含的 doc 数量、token 总数量以及每个 doc 的 token 数量和它在 bin 文件内的起始字节偏移。大型索引则更进一步把每个 doc 内部按 sequence 切好的样本也登记下来记录每个样本的 token 数量和起始偏移。这样训练循环拿到一个全局样本编号后通过 idx 就能立刻定位到对应的 token 范围不用扫描全文件。2.2 document 拼接与 eod token 的角色LLM 预训练语料里不同的文章或文档天然是独立语义单元。如果直接把每篇文章单独切成固定长度样本会有两个问题短文章会造成大量 padding 浪费长文章又会被硬切成毫无语义关联的碎片。Megatron 风格采用的做法是把多个 document 首尾相接拼成长流遇到文章边界就插入一个 eod token然后在长流上按固定窗口滑动切样本。eod token 在这里有多重身份。第一它是 document 边界指示器训练时模型看到它就知道前面的文本单元结束了第二它是自然语言语义的软分隔符跨 document 拼接的样本即便前半段讲科技、后半段讲美食模型也会因为 eod 的存在学到“这是两段独立内容”第三它是 loss 计算的天然分界点很多训练脚本会在计算 loss 时屏蔽 eod 之后的跨 document 部分避免模型学一些无意义的“跨界预测”。我在实际操作中默认使用 tokenizer 自带的 eod token id。如果用 HuggingFace 的 GPT2Tokenizer就是|endoftext|对应的 id如果用 LlamaTokenizer则要格外小心因为 Llama 语义里更强调s和/s端到端一致性的要求更高。总之用什么 tokenizer 编码就一定用同一份 tokenizer 来解析和映射这是后续所有环节都不出错的前提。2.3 样本切分与标签生成链条走到这一节长 token 流还只是一个一维数组。预训练时模型输入需要的是类似[batch_size, seq_length]的张量而标签是每个输入 token 右移一位的 next token。Megatron 风格数据会在预处理阶段就把样本边界切好省去训练时反复做切分的开销。具体切法是token 流每seq_length 1个 token 组成一个原始块前seq_length个作为输入后seq_length个作为标签从第二个 token 到最后一个 token。这也意味着预训练时配置的seq_length必须和预处理时完全一致。如果你预处理用 2048训练配置改成了 4096那数据加载器切出来的样本语义就全乱了基本等于重新训练。切分还要考虑样本之间的对齐方式。有些实现是从整个流开头按固定步长切有些实现会为每个 document 单独设置偏移。更精细的做法是让每个样本尽量从一个 document 的起始 token 开始而不是让一个样本的前半个是上一篇文章的结尾。可这样做会带来大量的尾部 padding 浪费。实际工程里主流方案还是“整流滑动切分”用 document 内的随机偏移来缓解跨文档问题既能保证样本完整度又不牺牲数据利用率。这个细节在后续实操代码里能看到。3. 数据预处理详细实操3.1 环境准备与依赖安装先把实验环境说清楚。我用的是 MindSpore 2.2 及以上版本模型侧是基于 MindFormers 里的 Transformers 架构改造的 GPT 类模型。数据预处理脚本纯粹是 Python 实现不依赖 MindSpore这样可以在任何机器上先跑出数据来。需要安装的基础依赖如下pip install mindspore pip install mindformers pip install transformers pip install datasets pip install tiktoken如果你用的是 LLaMA 类模型还需要安装对应 tokenizer 所需的依赖。我这次样例用的是 GPT2 tokenizer它对应的 merge 文件由 HuggingFace transformers 自动下载。注意这里tiktoken不是必须的因为transformers内部的 tokenizer 已经够用但tiktoken在编码超长文本时速度更快如果你语料是 TB 级建议做一下编码性能对比。3.2 语料清洗与 JSONL 标准格式原始语料格式五花八门我建议无论源头是 Common Crawl、维基百科还是内部爬虫统一转成 JSONL 落地每行一个 JSON 对象至少包含text字段。这个格式对多进程处理、断点续跑、后续 tokenize 都很友好。{text: The quick brown fox jumps over the lazy dog., meta: {source: example}} {text: Another document example with different content., meta: {source: example}}清洗阶段我强烈建议做下面几件事。第一统一全文编码为 UTF-8去除控制字符第二过滤超短的噪音文本比如短于 50 个字符第三去掉重复行和近似重复文档这一步可以显著减少模型的记忆负担对困惑度指标影响很大第四不必要的 HTML 标签、URL、乱码符号要根据场景决定是否去除。我踩过的一个坑是语料里存在大量“几乎重复”的文本比如新闻网站同一事件的多篇转载只改了几个字。用简单的集合去重去不掉必须配合 MinHash 之类的近似去重算法。如果条件有限可以先做前缀哈希去重能挡掉一批模板重复内容。3.3 编码进程的并行加速当语料达到几十 GB 甚至更大时单进程 tokenize 会慢到怀疑人生。这里我非常推荐用multiprocessing配合批次读取实现并行编码。基本思路是把 JSONL 文件按行切分成多个分片每个分片交给一个子进程子进程内部加载 tokenizer、逐行编码、把结果按 document 粒度写入独立的临时 bin 文件并记录索引最后在主进程里把所有分片合并成一个总 bin 和总 idx。这样能利用上多核 CPU实测单机 64 核可以把一个 100GB 语料的编码时间从几小时压到几十分钟。下面是一个简化版的并行编码骨架重点看 tokenizer 调用和数据组织方式进程中尽量不要共享 tokenizer 对象因为多进程复制会带来额外开销。import json import os import numpy as np from multiprocessing import Pool from transformers import GPT2Tokenizer def encode_document(text, tokenizer): tokens tokenizer.encode(text, add_special_tokensFalse) # 结尾追加 eod token标记文档结束 tokens tokens [tokenizer.eos_token_id] return np.asarray(tokens, dtypenp.uint16) def process_shard(shard_path, output_prefix, tokenizer_name): tokenizer GPT2Tokenizer.from_pretrained(tokenizer_name) tokenizer.add_special_tokens({pad_token: |endoftext|}) # 为简化逻辑这里演示单文档编码实际可按行循环 bin_path output_prefix .bin idx_path output_prefix .idx doc_ends [] token_count 0 with open(shard_path, r, encodingutf-8) as fin, \ open(bin_path, wb) as fbin: for line in fin: line line.strip() if not line: continue obj json.loads(line) arr encode_document(obj[text], tokenizer) arr.tofile(fbin) doc_ends.append(token_count len(arr)) token_count len(arr) with open(idx_path, w, encodingutf-8) as fout: json.dump({doc_ends: doc_ends, total_tokens: token_count}, fout)上面只是演示结构真正的生产级脚本还需要处理进程内聚合、二进制索引格式、错误文档跳过和日志输出。特别是tokenizer.eos_token_id和tokenizer.pad_token_id初始化的时候就要确认它们的值是不是同一个很多 tokenizer 默认 eos 是|endoftext|pad 是 None如果不处理后面补 pad 的时候会直接报错。3.4 完整生成 bin 与 idx 文件单一分片输出的是临时格式还需要一个合并阶段把多个分片的 token 流串联起来同时生成标准 idx 文件。标准 Megatron idx 文件内部是一个头部加两个索引块我用一个简单的 NumPy 版本示例来说明过程。这里的关键数据结构是三个一维数组sizes每个 document 编码后的 token 数含 eod tokenoffsets每个 document 的第一个 token 在总 token 流里的起始位置total_size所有 token 数之和然后按顺序把每个 document 的 token 数组写入 bin 文件再用np.cumsum计算 offsets这样索引信息就规整了。import numpy as np import os def write_megatron_style(path_prefix, doc_token_lists): bin_path path_prefix .bin idx_path path_prefix .idx sizes [] offsets [] total 0 with open(bin_path, wb) as f: for doc_tokens in doc_token_lists: arr np.asarray(doc_tokens, dtypenp.uint16) arr.tofile(f) sizes.append(len(arr)) offsets.append(total) total len(arr) sizes np.asarray(sizes, dtypenp.int32) offsets np.asarray(offsets, dtypenp.int64) np.savez_compressed(idx_path, sizessizes, offsetsoffsets, totaltotal)看起来很简单但实际生产环境里你需要注意几个点。第一bin 文件动辄几十 GB不能用np.save的格式因为加载时需要一次性读入内存第二idx 文件如果也是几十 MB 的 NumPy npz加载还行如果更大就要用内存映射mmap方式。第三写入 bin 的 dtype 必须是固定长度整数。uint16可以覆盖大多数词表大小65536 以内如果 tokenizer 词表超过这个范围就要改用int32。选 dtype 直接影响 bin 文件体积uint16 能比 int32 省一半空间所以不是越大越好。3.5 离线预检验证 bin 数据的正确性生成完 bin/idx 之后千万别急着扔给训练脚本。我先跑一个离线预检函数随机抽几个位置把 token id 解码回文本确认没有乱码或异常的越界 id。这一步成本很低但能解决后续调式时至少一半的“为什么 loss 是 NaN”的疑惑。def sample_check(bin_path, idx_path, tokenizer, num_samples5): data np.fromfile(bin_path, dtypenp.uint16) sizes np.load(idx_path .npz)[sizes] offsets np.load(idx_path .npz)[offsets] rng np.random.default_rng(42) total_docs len(sizes) for i in range(num_samples): doc_id rng.integers(0, total_docs) start offsets[doc_id] end start sizes[doc_id] ids data[start:end] text tokenizer.decode(ids.tolist()) print(fdoc {doc_id}: {text[:200]})如果 decode 结果里出现了大段 pad token 或者|endoftext|位置异常基本可以定位到编码阶段 eod token 处理有误。如果 decode 出来的是乱码大概率是 dtype 声明和落盘数据不一致或者二进制偏移算错了。这个预检函数建议固定随机种子方便以后回归对比。4. MindSpore Transformers 侧数据接入4.1 MindSpore 数据加载器设计训练侧要做的第一件事就是写一个继承自mindspore.dataset.GeneratorDataset或mindspore.dataset.MindDataset的加载器。因为 bin/idx 不是 MindSpore 原生格式所以最简单的是基于 NumPy 数组实现一个可迭代的数据源。核心设计思路是启动时通过内存映射np.memmap把 bin 文件映射到虚拟内存不真正读入物理内存再用 idx 加载 offsets 和 sizes作为样本索引表。每次迭代按全局样本编号找到它所属的 document然后从 memmap 里切出所需 token 序列。注意每次只切一段不要一次性把所有样本加载进来否则多卡训练时内存会吃紧。下面是一个简化版的 MindSpore 加载伪代码import numpy as np import mindspore as ms from mindspore.dataset import GeneratorDataset class MegatronDataset: def __init__(self, bin_path, idx_path, seq_length, eod_token_id): self.bin np.memmap(bin_path, dtypenp.uint16, moder) with np.load(idx_path) as data: self.sizes data[sizes] self.offsets data[offsets] self.seq_length seq_length self.eod_token_id eod_token_id self._build_sample_boundaries() def _build_sample_boundaries(self): self.sample_boundaries [] for doc_id, size in enumerate(self.sizes): start self.offsets[doc_id] length size - 1 # 去掉末尾的 eod避免样本纯边界 num_samples (length // self.seq_length) or 1 for i in range(num_samples): sample_start start i * self.seq_length self.sample_boundaries.append((doc_id, sample_start)) def __len__(self): return len(self.sample_boundaries) def __getitem__(self, idx): _, start self.sample_boundaries[idx] tokens self.bin[start: start self.seq_length 1].astype(np.int32) input_ids tokens[: self.seq_length] label_ids tokens[1: self.seq_length 1] return input_ids, label_ids这里有一个细节_build_sample_boundaries会给每个 document 至少保留一个样本即便 document 本身长度不够seq_length。为什么因为如果完全按整块切分过短的 document 会被完全丢掉导致语料浪费。保留一个短样本后训练侧可以自定义 mask 策略来决定是否让模型学习这部分内容。4.2 与 mindformers 训练流程的对接方式如果你用的是 MindFormers 的高层训练接口可以把MegatronDataset传进GeneratorDataset再设置batch_size和shuffle。MindSpore 的GeneratorDataset支持基于索引的随机采样所以 shuffle 层面不需要自己去实现。dataset MegatronDataset( bin_pathyour_data.bin, idx_pathyour_data.idx.npz, seq_length2048, eod_token_idtokenizer.eos_token_id, ) ms_dataset GeneratorDataset( sourcedataset, column_names[input_ids, label_ids], shuffleTrue, num_parallel_workers8, ) ms_dataset ms_dataset.batch(batch_size8, drop_remainderTrue)这里要特别注意num_parallel_workers的配置。它不是越大越好过大的并发会反复切分 memmap造成 IO 抖动。我在 64 核机器上测试8 到 16 个 worker 通常能打满单卡数据带宽再往上提升有限。如果你用多卡并行每张卡一个进程各自读同一份 bin 文件依靠操作系统页缓存也能扛住但如果是几百 TB 级别的语料建议把数据预分成多份 shard每张卡只映射其中一部分减少单机 IO 压力。如果你需要更细粒度的控制比如不同 epoch 用不同随机偏移可以在__getitem__里做样本级随机偏移。这样每个 epoch 喂给模型的样本不完全相同能提升训练效果。但注意样本级随机偏移需要非常小心 eod token 的处理不能随机到把 eod 硬切成两个样本中间去否则语义边界就碎了。4.3 seq_length 对齐与 batch 组装数据预处理时设定的seq_length必须与训练配置中的seq_length一致这一点我前面强调过。MindFormers 的配置通常写在 YAML 文件里比如model: seq_length: 2048 train: batch_size: 8 dataset: data_path: ./your_data.binYAML 里的data_path指向 bin 路径MindFormers 会尝试用自定义数据集读取器去加载 idx 文件。如果你的版本没有内置 Megatron 读取器就要像我上面写的那样自己实现MegatronDataset并传给Trainer。这块最容易犯的错是预处理脚本用的seq_length是 2049因为多取了一个 token 做 label训练配置却写了 2048。对不齐之后每个样本的 label 会跨越到下一个样本loss 曲线看着像模像样实际上学的内容是乱的。batch 组装也是一个大坑。因为每个样本的长度是固定的seq_length所以 batch 组装不需要 padding但如果你的语料里有超短 document边界保留的样本长度不足则必须统一用一个特殊的 pad token 去补。我的建议是在预处理阶段就把过短的 document 直接过滤掉或者把它和下一个 document 一起拼接成完整样本不要让训练侧去处理参差不齐的序列。5. 常见问题与排查技巧实录5.1 tokenizer 不一致导致乱码与崩溃这是最隐蔽的问题。离线预处理时用 GPT2Tokenizer 编码训练侧加载模型时用了 LLaMA 的 tokenizer 做 embedding 对齐两边词表顺序不一样模型拿到的 token id 对应的是完全不同含义的文本训练 loss 自然下不去。排查方法很简单训练启动前打印几个 token 的 id 映射确认模型 config 里的 vocab_size、pad_token_id、bos_token_id、eos_token_id 与预处理脚本完全一致。尤其是 eod token很多框架管它叫eod_token_id它就是 tokenizer 的eos_token_id但有些英文语料里eos被特殊符号占用必须显式指定|endoftext|。我自己的经验是预处理脚本和训练脚本共用同一个 tokenizer 配置文件不要靠直觉复制粘贴 vocab.json 和 merges.txt最好在代码里用 argparse 同时传给两边从源头杜绝不一致。5.2 内存不足的一个隐藏凶手np.fromfile不少人在验证阶段习惯用np.fromfile(bin_path, dtypenp.uint16)把整个 bin 文件读进内存。这个操作在数据量小的时候没问题但一旦 bin 文件超过 30GB内存立刻爆掉。正确做法有两个。第一验证阶段用np.memmap代替np.fromfile只映射不加载第二加载阶段用np.load的mmap_moder读 idx 文件保证索引数组也是按需加载。整个训练进程启动后常驻内存应该只包含 idx 索引的一小部分和 bins 的页缓存这样才不会把训练机的显存和内存一起拖垮。5.3 idx 文件损坏后的定位与修复idx 文件如果中途写坏常见的表现是训练到某一步突然崩掉报 index out of range 或者奇怪的 token 越界。我碰到过一次合并阶段用多进程写临时 idx但没有等所有子进程结束就启动合并导致最后 offset 对不上。定位问题的方法是做一个全量扫描从 idx 里依次读取每个样本的 offsets 和 sizes验证它们在 bin 文件长度范围内并且 token id 最大值小于词表大小。如果某处越界就能立刻确定是哪个 shard 出了问题。修复方式很简单找到损坏的 document把它从最终数据里剔除或者重新编码一次然后重新生成 idx。关键是要在生成脚本里加入“写完后校验文件大小和 document 数是否匹配”的步骤从流程上杜绝这种问题。5.4 数据吞吐量不足与训练瓶颈如果你的训练代码一切正常但 GPU 利用率上不去先看数据侧。Megatron 风格 bin/idx 在加载时是顺序大块读取随机访问性很好通常不会成为瓶颈但你依然要注意三点。第一确认 bin 文件所在磁盘是 SSD 或本地 NVMe网络文件系统 HDD 在随机读取性能上很差。第二适当调大GeneratorDataset的 prefetch 缓冲区和num_parallel_workers让数据生产速度略快于训练消费速度。第三多进程读同一份 memmap 时Linux 的 page cache 会自然优化但如果多卡之间共享一个 bin 文件并且卡数很多我建议把数据预切成分片每张卡单独映射一个分片减少锁竞争和页表开销。5.5 一个容易被忽略的细节随机种子与全局乱序预训练对数据的均匀随机性非常敏感。如果每轮 epoch 都按固定的 doc 顺序切样本模型会产生顺序记忆导致评估时 loss 波动。Megatron 风格数据在 idx 里记录的是 document 索引和 offset真正的随机化发生在训练加载阶段而不是预处理阶段。所以在 MindSpore 侧GeneratorDataset的shuffleTrue一定要开并且设置合理的global_seed多卡训练时每张卡拿到不同的 shuffle 序列保证整体数据被均匀消费。如果嫌GeneratorDataset的 shuffle 不够彻底可以在__getitem__里加一个随机偏移让每个 epoch 的样本边界都略有变化。注意加了随机偏移后同一个 token 位置在不同 epoch 可能被当成 input 也可能被当成 label这对模型没有坏处反而能提升数据多样性。6. 离线到在线链路的一点个人经验这一路做下来我最深的体会是数据预处理模块的工程质量往往比模型结构改动更能决定预训练实验的成败。很多团队在调模型上下大功夫却忽略了一个问题——喂进去的数据是不是真的干净、一致、可复现。我给你一个最实用的建议在一开始就为整套数据链路写一个“数据指纹”模块。每次预处理完成把 tokenizer 版本、词表大小、总 token 数、doc 数量、seq_length、eod token id 全部打成一个哈希连同 bin/idx 文件一起保存。训练脚本在启动时校验这个指纹如果不匹配直接报错。这样既能在多卡、多机环境下保证所有 worker 用完全一致的数据也能在后续复现实验结果时快速定位是不是数据变了。还有一个小技巧预处理脚本里给每个 document 保留meta信息比如原始来源、清洗规则版本。当模型出现异常行为时可以从样本一路回溯到原始语料快速判断是数据问题还是模型问题。这些元信息不一定要写进 bin/idx可以单独存一份 JSONL按 doc_id 对齐即可。如果你后续想在 MindSpore 里跑 MoE、长上下文模型这套数据方案也依然能复用只需要在加载器里增加更丰富的 sample 边界信息或者在离线阶段按更长窗口重新切分。数据基础设施是预训练最值得持续投入的部分我强烈建议第一次做就把格式、校验、回归流程搭好后面换模型、换语料都是几个小时的事。