最近在折腾 MindSpore Transformers 的 LLM 预训练链路卡了一周的数据预处理终于跑通了。整个过程让我对 Megatron 风格数据集设计有了更深的体会它不只是把文本切成一段段 token而是把语料、tokenizer、样本边界、文档信息打包成一套可以高效随机读取的二进制格式。这篇文章把整套思路和可落地的代码写出来给正在做 LLM 预训练数据管线的同学一个参考尤其是打算在 MindSpore 生态里复现 GPT、Llama 这类模型训练的人。1. 项目背景为什么 LLM 训练前要单独做数据预处理1.1 预训练数据流的痛点接触过大模型训练的人都知道预训练数据集动辄几十 GB、几百 GB原始语料往往是纯文本、JSON、Parquet 混合在一起。如果每次训练都在数据迭代器里现 tokenize再把文本拼成固定长度序列不仅慢而且每个 worker 都重复调用 tokenizerCPU 开销直接爆炸。更麻烦的是训练框架需要随机访问任意位置的样本如果数据是压缩包或动态拼接shuffle 和分布式并行根本做不好。所以工程上通用的做法是在训练启动前离线把原始语料统一加工成一个内存友好、按偏移量随机读取的二进制格式。Megatron 风格数据集就是这类方案的代表它把 token 序列和索引信息分开存储训练时用内存映射直接读速度快、省内存而且天然支持不同 worker 取不同样本。1.2 Megatron 风格到底指什么Megatron-LM 是业界做大规模模型并行训练时绕不开的参考实现。它的数据预处理脚本会把语料变成两个文件.bin和.idx。.bin里面是一长串 token id本质上是 NumPy 数组的二进制内容.idx里记录文档和样本的偏移索引。训练时用np.memmap打开.bin借助.idx里的偏移量就能在任意位置拿到一个完整样本。这个设计妙在没有 padding也不要求每个文档长度对齐。所有文档先拼成一条 token 流再切分成固定长度的样本。如果某个文档结尾不够一个样本长度就从下一个文档继续接保证样本长度始终是seq_length 1。这种思路后来被很多框架沿用MindSpore Transformers 在支持大模型预训练时也很适合用同样的数据格式减少迁移成本。1.3 这套加工流程适合什么场景如果你的项目满足下面几个条件用 Megatron 风格预处理会非常舒服训练目标是自回归语言模型比如 GPT、Llama、Bloom 这类 Causal LM。语料规模大训练过程中需要频繁随机访问不同位置的样本而不仅仅是顺序读取。需要做多机多卡训练希望每个 rank 能独立、高效地取数据不依赖全局共享存储上的重复解析。对数据 pipeline 的性能有要求不能让数据读取成为训练瓶颈。反过来如果你的语料非常小只做微调或者做的是纯分类、检索这类任务那这个方案就有点重直接用 JSON 加载也够了。2. 核心概念拆解token、document、sample 与索引文件2.1 Token 级数据的基本单位在预训练里文本最终都要变成 token id 数组。这里有个容易忽视的点每个文档 tokenize 完之后并不只是一个简单的数组它末尾还要带上一个特殊的结束符比如eos_token_id。作用有两个一是让模型学会“文本到此结束”二是作为文档边界标记方便后续按文档再组织索引。对常见 tokenizer 来说你要特别留意它是否默认添加特殊 token。比如 Llama 的 tokenizer如果直接对一段文本调用tokenizer(text)它可能会自动加上 bos/eos。而做预训练数据预处理时我们往往希望手动控制这些特殊 token 的位置所以通常把add_special_tokensFalse然后在文档末尾显式加上eos_token_id。这样可以避免每个文档开头都带 bos浪费序列位置。2.2 Document 与 sample 的映射关系Megatron 风格的数据结构里有两类索引文档索引和样本索引。文档索引告诉你某个原始文档在整条 token 流中的起止位置样本索引告诉你每个训练样本对应 token 流的哪个区间。两个索引都记录了“边界”但服务的目标不同。文档索引主要用于数据混合和采样。比如你想按照不同来源的语料做比例混合或者训练时希望按文档随机抽样而不是纯粹按 token 顺序抽就需要知道每个文档的偏移。样本索引则直接服务于训练数据加载器。每个样本在 token 流中占用seq_length 1个 token前seq_length个作为输入 ids后seq_length个作为 labels这样正好完成自回归的移位。2.3 bin idx 双文件设计的价值为什么不用一个大 JSON 或者 HDF5核心原因是随机访问效率。训练时每个 step 要从几十 GB 数据里随机取几千个样本如果用 JSON你需要一次性把所有 token 读入内存或者频繁随机 seek 文本解析代价都很高。用.bin .idx你可以用内存映射方式把整块 token 流映射到进程地址空间取值时只是做一次 NumPy 切片基本是内存拷贝级别没有磁盘 IO 瓶颈。另外.idx文件本身是一个小型的索引结构可以在启动时完整加载到内存。训练时先在内存里随机生成样本编号再通过索引定位到.bin中具体偏移整个过程 O(1) 级别非常干净。3. 预处理方案设计从原始文本到固定长度序列3.1 技术选型tokenizer 的选择在 MindSpore Transformers 场景下我推荐直接使用 HuggingFacetransformers加载 tokenizer。原因是生态成熟、模型词表一致、特殊 token 配置灵活。你训练用的模型如果是mindspore_transformers里加载的权重通常也是从 HF 转换来的所以词表完全兼容。如果你不希望引入 HF也可以直接用 MindSpore 自己实现的 tokenizer 接口但要确保它产出的 token id 和训练阶段的模型词表完全一致。这里最容易踩坑的是 tokenizer 版本不一致导致预训练时同一个词的 token id 对不上。所以预处理前一定要固定 tokenizer 版本和词表文件。3.2 预处理流程总览整个预处理流程其实可以分成四步读取原始语料按约定规则切分为 document。对每个 document 调用 tokenizer得到 token id 序列并补上 eos。将所有 document 的 token id 序列首尾拼接形成一条超长的 token 流。按seq_length 1切分样本记录样本索引同时记录每个 document 的偏移形成文档索引。流程不复杂难的是在大规模语料上实现时不把内存撑爆。最简单粗暴的方式是先把所有 token 加到 Python list 里最后转成 NumPy 写文件。这种做法在几 GB 语料内问题不大但到了几百 GB 就必须换思路用np.memmap边 tokenize 边写入。3.3 数据格式与元信息约定我习惯在输出目录里额外放一个config.json作为元信息记录seq_length、tokenizer_path、vocab_size、dtype、doc_cnt、sample_cnt等关键参数。这样训练脚本启动时可以先读配置避免手动传错参数。.bin文件的 dtype 选择要看词表大小。如果vocab_size 65535用np.uint16存储能省一半空间否则用np.uint32。Llama 词表一般是 32000GPT 系列也差不多所以工程上默认uint16即可但保险起见还是动态判断一下。4. 实操用 MindSpore Transformers 实现 Megatron 风格预处理4.1 搭建最小可运行环境我实际实验的环境是 Python 3.10 MindSpore 2.2 CPU 版本。数据处理阶段其实不需要 GPU纯 CPU 就能跑唯一要求是内存别太小。依赖库如下pip install mindspore transformers numpy tqdm如果你已经有一份训练用的环境不需要额外装太多东西。数据预处理脚本可以独立运行不依赖 MindSpore 的模型定义。4.2 预处理主脚本实现下面给出一个可用版本。为了突出重点我做了简化但核心逻辑是完整的。假设原始语料是一个大文本文件多个 document 之间用空行分隔。import json import os from pathlib import Path import numpy as np from tqdm import tqdm from transformers import AutoTokenizer def read_documents(input_path): with open(input_path, r, encodingutf-8) as f: content f.read() docs [doc.strip() for doc in content.split(\n\n) if doc.strip()] return docs def tokenize_document(doc, tokenizer, eos_token_id): tokens tokenizer(doc, add_special_tokensFalse)[input_ids] tokens.append(eos_token_id) return tokens def preprocess(input_path, output_prefix, tokenizer_path, seq_length2048): tokenizer AutoTokenizer.from_pretrained(tokenizer_path) eos_token_id tokenizer.eos_token_id if eos_token_id is None: raise ValueError(tokenizer must have eos_token_id) docs read_documents(input_path) vocab_size tokenizer.vocab_size dtype np.uint16 if vocab_size 65535 else np.uint32 bin_path f{output_prefix}.bin idx_path f{output_prefix}.idx # 先用 memmap 创建 bin 文件初始大小可设置一个估计值后面再截断 # 这里为了简单先把所有 token 放内存再统一写文件 all_tokens [] doc_offset [0] sample_offset [0] for doc in tqdm(docs, desctokenizing): toks tokenize_document(doc, tokenizer, eos_token_id) all_tokens.extend(toks) doc_offset.append(len(all_tokens)) # 按 seq_length 1 切分样本 num_samples (len(all_tokens) - 1) // seq_length max_sample_start num_samples * seq_length for sample_idx in range(num_samples): start sample_idx * seq_length sample_offset.append(start) # 截断到最后一个完整样本结束位置 final_token_len num_samples * seq_length 1 all_tokens all_tokens[:final_token_len] token_arr np.asarray(all_tokens, dtypedtype) token_arr.tofile(bin_path) # 保存索引 idx_dict { doc_offset: doc_offset, sample_offset: sample_offset, vocab_size: vocab_size, dtype: np.dtype(dtype).name, seq_length: seq_length, } with open(idx_path, w, encodingutf-8) as f: json.dump(idx_dict, f) # 保存元信息 config_path f{output_prefix}_config.json config { bin_path: bin_path, idx_path: idx_path, seq_length: seq_length, vocab_size: vocab_size, dtype: np.dtype(dtype).name, num_docs: len(doc_offset) - 1, num_samples: num_samples, } with open(config_path, w, encodingutf-8) as f: json.dump(config, f, indent2) print(fpreprocess done. samples{num_samples}, docs{len(doc_offset)-1})这段代码用 JSON 保存索引而不是 Megatron 原版二进制 idx。好处是可读性强调试方便训练端加载也简单。如果你的数据量极大比如超过内存上面直接把所有 token 放在 list 里的做法就不太合适。更好的做法是先用 memmap 初始化一个很大的 bin 文件然后边 tokenize 边把 token 写入 memmap。我后面在避坑部分会展开讲。4.3 生成数据校验结果运行完之后建议立刻做一次自检随机取出几个样本还原成文本看是否合理同时检查 labels 是不是 input_ids 的右移一位。import numpy as np import json def verify_sample(bin_path, idx_path, sample_id, seq_length2048): idx json.load(open(idx_path)) token_arr np.fromfile(bin_path, dtypeidx[dtype]) sample_start idx[sample_offset][sample_id] sample_tokens token_arr[sample_start : sample_start seq_length 1] input_ids sample_tokens[:-1] labels sample_tokens[1:] assert len(input_ids) seq_length assert len(labels) seq_length return input_ids, labels这一步非常值得做别偷懒。我碰到过好几种问题比如eos_token_id没加上、文档切分导致的 token 数量不对、sample_offset算错都是靠这种自检捞出来的。5. 数据集加载与训练接入5.1 自定义 Dataset 加载 bin 文件预处理完的数据必须在训练脚本里用高效方式加载。MindSpore 的GeneratorDataset可以包一个自定义迭代器每次随机返回一个样本。为了支持几十 GB 的 bin 文件我用np.memmap而不是np.fromfile这样不会把整个文件读进内存。下面是一个兼容 Memmap 的 Dataset 类import numpy as np import json from mindspore.dataset import GeneratorDataset class MegatronStyleDataset: def __init__(self, bin_path, idx_path, seq_length2048, seed42): idx json.load(open(idx_path)) self.seq_length seq_length self.dtype np.dtype(idx[dtype]) self.sample_offset np.asarray(idx[sample_offset], dtypenp.int64) self.num_samples len(self.sample_offset) - 1 self.bin_arr np.memmap(bin_path, dtypeself.dtype, moder) self.rng np.random.default_rng(seed) def __len__(self): return self.num_samples def __getitem__(self, _): # 用随机索引采样方便做全局 shuffle sample_id self.rng.integers(0, self.num_samples) start self.sample_offset[sample_id] sample_tokens self.bin_arr[start : start self.seq_length 1] input_ids sample_tokens[:-1].astype(np.int32) labels sample_tokens[1:].astype(np.int32) return input_ids, labels def create_dataset(bin_path, idx_path, seq_length, batch_size): dataset MegatronStyleDataset(bin_path, idx_path, seq_length) ds GeneratorDataset( sourcedataset, column_names[input_ids, labels], num_parallel_workers8, ) ds ds.batch(batch_size, drop_remainderTrue) return ds注意__getitem__里的sample_id是随机生成的这样每个 epoch 看到的样本顺序都不同省去了额外做 shuffle 的麻烦。如果你希望严格保证每个样本在一个 epoch 内只出现一次那最好在GeneratorDataset内置的shuffle参数上做文章而不是在__getitem__里随机。两种模式各有取舍后面细说。5.2 与 MindSpore 训练流程对接拿到GeneratorDataset之后训练循环就没什么特殊了。假设你已经定义了模型和优化器只需要把原来的数据迭代器替换成这个数据集。import mindspore as ms from mindspore import Model, nn # 构建数据 train_ds create_dataset( bin_pathyour_prefix.bin, idx_pathyour_prefix.idx, seq_length2048, batch_size8, ) # 假设 net 是定义好的 Causal LM net build_model() loss_fn nn.CrossEntropyLoss() optimizer nn.AdamWeightDecay(paramsnet.trainable_params()) model Model(net, loss_fnloss_fn, optimizeroptimizer) model.train(epochs1, train_datasettrain_ds)如果你的模型输出 logits 维度是[batch, seq, vocab]那么 labels 也是[batch, seq]。MindSpore 的交叉熵损失会自动做 flatten 和 mask。要是你想忽略 padding 部分可以再给 loss 传一个ignore_index。5.3 epoch、shuffle 与多 worker 的坑在GeneratorDataset里如果num_parallel_workers 1多个 worker 会并发调用__getitem__。如果你在__getitem__里用np.random.default_rng一定要小心每个 worker 的随机序列会不会重复。我踩过的坑是直接给Dataset.__init__传一个固定的 seed然后每个 worker 共享同一个 rng结果不同 epoch 采样的顺序高度相似。解决方法是让每个 worker 拿到不同的 seed。可以在Dataset.__init__里根据当前进程 ID 或者线程 ID 生成一个新的 rng或者直接把随机采样逻辑放到外面由GeneratorDataset的shuffle机制处理。更稳妥的做法是__getitem__只根据传入的index返回固定样本然后靠ds.shuffle(buffer_size)做打乱。这样每个样本在一个 epoch 内只会被访问一次并且不同 worker 之间不会重复随机。6. 常见问题与避坑实录6.1 文档边界丢失与样本跨度过大如果你在 tokenize 时没有给每个文档补 eos那拼接后的 token 流就无法区分文档边界。一些模型很在意这一点尤其是需要按文档做 attention mask 的场景。如果不在预处理阶段记录文档索引后期想按文档做加权采样就得重新扫描全量 token代价极高。我的建议是一开始就把 doc_offset 保存下来哪怕当下训练用不到以后做数据配比评估也用得上。6.2 超大文件内存管理改造成追加写入前面给的示例为了简洁把所有 token 放在内存里。实际我处理 20GB 语料时token 数量是几十亿Python list 直接撑不住。改造成 memmap 追加写入的思路很简单先创建一个足够大的空 binary 文件比如预估 token 总量乘以 dtype 字节数再在预处理循环里维护一个写入游标。核心代码大概是estimated_tokens 10_000_000_000 fp np.memmap(bin_path, dtypedtype, modew, shape(estimated_tokens,)) cursor 0 doc_offset [0] for doc in docs: toks tokenize_document(doc, tokenizer, eos_token_id) fp[cursor : cursor len(toks)] toks cursor len(toks) doc_offset.append(cursor) # token 流写完后截断文件到实际长度 fp.flush() fp._mmap.close() with open(bin_path, rb) as f: f.truncate(cursor * np.dtype(dtype).itemsize)这样做的好处是内存占用几乎为 0坏处是如果预估 token 量太大会瞬间占用大块磁盘。经验法则是先写一个快速脚本统计语料字符总量按每个字符约 0.3~0.5 个 token 估算。估算有偏差没关系后面会 truncate。6.3 tokenizer 版本不一致导致错位这个坑很隐蔽。如果预处理时用的 tokenizer 是LlamaTokenizer训练时模型配套的词表文件却是另一个 commit 版本token id 可能完全对不上。你训练出来 loss 很低但生成结果全是乱码就是因为 id 映射错位。所以一定要把 tokenizer 的tokenizer_config.json、vocab.json或merges.txt一起快照保存预处理和训练必须锁定同一份。6.4 自检工具样本一致性验证最后分享一个我写的小工具脚本每次预处理完跑一遍能发现绝大多数问题def full_validation(bin_path, idx_path, sample_nums100): idx json.load(open(idx_path)) token_arr np.fromfile(bin_path, dtypeidx[dtype]) seq_len idx[seq_length] for _ in range(sample_nums): sid np.random.randint(0, len(idx[sample_offset]) - 1) start idx[sample_offset][sid] sample token_arr[start : start seq_len 1] assert len(sample) seq_len 1, fsample length mismatch: {len(sample)} input_ids sample[:-1] labels sample[1:] for i in range(seq_len): if labels[i] ! input_ids[i]: pass # 因为跨文档可能存在不同 token跨文档样本中 labels 和 input_ids 除了右移一位还有可能在文档相接处出现 token 跳跃所以不能硬等。自检重点放在长度和索引越界上。最后分享一点个人体会这套预处理方式一开始会觉得绕但真正把几十 GB 语料跑成 .bin 文件之后训练时的数据读取性能会给你非常踏实的反馈。我在 MindSpore Transformers 上复现 1.3B 模型预训练时换成 Megatron 风格数据后step 时间直接降了将近 30%瓶颈从数据读取变成了计算。自己做一遍预处理你对 token、样本、文档的理解也会深很多后面排查 loss 异常会更有方向感。希望这篇记录能帮你少走几个坑特别是在内存和索引这两块值得多花时间验证。