大模型数据预处理实战：从数据源到动态Padding的流水线方案

「昇思 MindSpore 大模型」这句话里真正决定一次微调能不能顺利跑通的环节往往不是模型结构而是那串看起来平平无奇的mindspore.dataset。我在好几个大模型项目里见过同一种症状训练脚本没问题模型结构没问题数据预处理用 NumPy 硬写也能把流程 push 通但只要语料一上规模、进程一开多要么内存暴涨要么 CPU 空转要么每张卡的样本莫名其妙重复。最后排查下来问题全落在“数据怎么进模型”这条管线上。大模型的数据变换和预处理真不是把文本encode一下再转成 Tensor 那么简单。它涉及原始语料加载、清洗、去重、分词、截断、padding、mask 构造、shuffle、batch、多卡分片还要跟训练框架的静态图/动态图无缝衔接。这篇分享我会从生产环境里实际踩过的坑出发把基于mindspore.dataset的数据变换与预处理全方案拆开讲包含可以直接抄走的代码、参数选择背后的原因以及新手最容易掉进去的几个隐蔽陷阱。适合谁看用昇思 MindSpore 做 GPT、LLaMA 这类大模型预训练或微调的同学以及被数据加载卡到怀疑人生的朋友。1. 数据预处理为什么是大模型的“隐形胜负手”1.1 一条未经工程化处理的典型失败链路很多人第一次写大模型微调时会习惯性这么做data_list [] for file in files: for line in file: data_list.append(tokenizer.encode(line)) train_data np.array(data_list, dtypeobject) train_tensor Tensor(train_data)这个写法在小数据集上没问题但一放大就崩。原因有几个一次性把所有文本全部读进内存10 万条样本可能还能扛一千万条样本直接内存爆炸np.array(..., dtypeobject)会退化成 Python 对象数组后面每次取 batch 都要走一遍 Python 循环慢到怀疑人生没有并行加载、没有缓冲、没有预取GPU 只能在 CPU 算完一个 batch 后干等在多卡场景下如果不做显式分片每张卡读到的都是全量数据的随机子集训练效果会变成“假随机”评估指标也跟着乱。这不是代码风格的问题是数据流水线这个环节被严重低估了。大模型的语料通常是海量短文本一条样本经过分词后长度可能从几十到几千不等如果预处理链路本身没有异步、没有并行、没有合理的 padding 策略模型训练根本吃不满算力。1.2 mindspore.dataset 在整个训练链路中的位置MindSpore 里的mindspore.dataset是一整套数据引擎它不止是“数据集类”而是把数据源、变换、批处理、随机化、分布式切分、预取缓冲都串成一张计算图。你可以把它理解成一条流水线数据源 - shuffle - map 变换 - batch - repeat - 预取 - 模型其中Generatordataset负责把任意 Python 生成器变成可迭代的数据源map负责对每个样本做变换batch负责把样本装成 batchshuffle负责打乱顺序repeat负责控制 epoch。每个算子还支持并行 worker 和预取缓冲这些组合在一起才是大模型训练需要的“数据吞吐”。在 MindFormers 这类大模型套件里底层的数据集构建最终也会落到mindspore.dataset上。所以不管你是直接写训练循环还是通过MindFormers的 Trainer 微调掌握这套 API 都是通用的。2. 核心 API 与数据变换的底层逻辑2.1 数据源GeneratorDataset 是承载一切后续变换的起点GeneratorDataset的核心作用是把普通 Python 生成器包装成框架能高效消费的数据集。它接受一个可迭代对象或生成器函数并要求每个样本以 numpy 数组或 Python 标量的形式返回同时用column_names声明每一列的名字。看一个最小示例import numpy as np import mindspore.dataset as ds def my_generator(): for i in range(1000): yield np.array([i, i 1], dtypenp.int32) dataset ds.GeneratorDataset( sourcemy_generator(), column_names[sample_pair] )写成生成器而不是一次性 List最大好处是内存可控。生成器是“随取随产”map、batch是在数据流上下一个接一个地处理不会把所有样本同时堆进内存。这里有几个实操教训生成器内部不要做重 I/O。比如每次 yield 都去读磁盘、请求网络、打印日志这些操作会阻塞整个数据 pipeline。返回类型必须用 numpy 数组不要返回 Pythonlist。框架在底层会做类型检查返回 list 经常导致后续map阶段出现dtype报错。如果你有多个字段yield时按元组返回并让column_names的数量和元组长度一致。def my_generator(): for text in texts: ids tokenizer.encode(text) labels np.concatenate([ids[1:], [eos_id]]) mask np.ones_like(ids) yield ids, labels, mask dataset ds.GeneratorDataset( sourcemy_generator(), column_names[input_ids, labels, attention_mask] )2.2 map 变换把分词、截断、掩码都写成纯函数map是数据变换的核心算子。它会在每个样本进入 batch 之前对样本做一次或多次变换。这里的“变换”就是一个普通 Python 函数接收 numpy 数组返回 numpy 数组或数组元组。调用方式def build_labels(ids): ids np.asarray(ids, dtypenp.int32) labels np.concatenate([ids[1:], np.array([eos_id], dtypenp.int32)]) mask np.ones_like(ids) return ids, labels, mask dataset dataset.map( operations[build_labels], input_columns[input_ids], output_columns[input_ids, labels, attention_mask], num_parallel_workers4 )这里要理解三件事input_columns指定函数接收哪些列output_columns指定函数返回后对应哪些列。返回数量必须和output_columns数量一致。map默认是逐样本执行的适合无状态变换。像“截断到固定长度”“构造掩码”“把 label 偏移一位”这类操作都放在这里。如果是分词这种重 CPU 操作要格外注意并行模式和函数序列化问题。map默认开启多进程函数需要能被 pickle分词器对象如果不可序列化就会在多个 worker 启动时报错。常见解法是设置python_multiprocessingFalse改用多线程或者把分词器定义为全局单例。还有一个容易踩的坑不要在map函数里调用print或写日志。多进程下这些操作会重复刷屏严重影响性能还会干扰正常的输出观察。2.3 batch 与动态 padding决定显存余量的关键细节batch的作用是把样本打包成固定批量。大模型场景下每条样本长度不一样直接batch()会因为 shape 不一致而报错。你需要做 padding。最省钱的方案是“动态 padding”即只把同一个 batch 内的样本 padding 到当前 batch 的最长长度而不是把所有样本都 padding 到最大序列长度。这个操作要放在per_batch_map里完成。per_batch_map接受一个自定义函数函数签名是def per_batch_pad(col_1, col_2, ..., batch_info): ...其中每个col_i是当前 batch 中某一列的列表列表里每一项是一个样本的 numpy 数组batch_info里包含当前 batch 大小等信息。函数返回多个 numpy 数组数组的第一维就是 batch 大小。动态 padding 能省多少显存举个例子如果最大长度是 4096但样本平均长度只有 800固定 padding 会让 80% 的 token 都是垃圾填充动态 padding 则让每个 batch 只按最长那条样本算显存利用率可能差好几倍。不过要注意如果你跑的是静态图模式且模型侧不支持动态 shape那跨 batch 的长度变化会导致编译/运行时 shape 不匹配。这种情况建议把所有样本统一 padding 到固定max_seq_len或者开启动态 shape 支持。2.4 shuffle、repeat、分布式切分的正确组合方式shuffle、batch、repeat的顺序不是随意的。典型推荐顺序是dataset dataset.shuffle(buffer_size10000) dataset dataset.batch(batch_size8, drop_remainderTrue) dataset dataset.repeat(epochs)为什么shuffle放在batch前因为shuffle是针对单样本做缓冲打乱越早打乱后续 batch 划分越均匀。如果先batch再shuffle打乱的是 batch 的顺序而每个 batch 内部样本仍然保持了原始序列相邻关系这对训练稳定性是不利的。repeat放在batch后面是因为repeat的作用是“重复整条流水线”。如果先repeat再shuffle会把多份数据混在一起打乱内存占用翻倍先shuffle后repeat每个 epoch 边界更清晰内存也更友好。另一个关键点是分布式场景的分片。多卡训练时你需要让每张卡只读一部分数据否则数据会重复。GeneratorDataset提供了现成的num_shards和shard_id参数from mindspore.communication import init, get_rank, get_group_size init() rank_id get_rank() world_size get_group_size() dataset ds.GeneratorDataset( sourcemy_generator(), column_names[...], num_shardsworld_size, shard_idrank_id )不要自己在生成器里按进程号去切文件列表那样容易出错直接用框架分片简单可靠。3. 实操一套可直接抄走的大模型预处理流水线3.1 准备一个最小可跑的数据集与分词器为了让你能直接跑通我用一个简单的玩具分词器做演示。生产环境请换成真实分词器比如transformers.AutoTokenizer.from_pretrained(xxx)本质上都是调用tokenizer.encode(text)拿到 token id 列表。import re import numpy as np import mindspore.dataset as ds class SimpleTokenizer: def __init__(self): self.pad_id 0 self.unk_id 1 self.eos_id 2 self.bos_id 3 words [hello, world, mindspore, model, tokenizer, data, transform, how, are, you, i, love, nlp, preprocessing, batch, shuffle, map, epoch, loss, train, large, language, demo] self.vocab [pad, unk, eos, bos] words self.word2id {w: i for i, w in enumerate(self.vocab)} def encode(self, text, max_len128): words re.findall(r[a-z0-9], text.lower()) ids [self.word2id.get(w, self.unk_id) for w in words][:max_len] return np.array(ids, dtypenp.int32) tok SimpleTokenizer() raw_texts [ hello mindspore model tokenizer data transform, I love nlp preprocessing batch shuffle map epoch, how are you doing with mindspore data pipeline, this is a simple demo for large language model data preprocessing, ]3.2 完整流水线代码含动态 padding先定义生成器然后依次接上map、shuffle、batch、repeatdef sample_generator(): for text in raw_texts: ids tok.encode(text, max_len128) yield ids dataset ds.GeneratorDataset( sourcesample_generator(), column_names[input_ids], shuffleFalse ) def build_labels(input_ids): input_ids np.asarray(input_ids, dtypenp.int32) labels np.concatenate( [input_ids[1:], np.array([tok.eos_id], dtypenp.int32)] ).astype(np.int32) attention_mask np.ones_like(input_ids) return input_ids, labels, attention_mask dataset dataset.map( operations[build_labels], input_columns[input_ids], output_columns[input_ids, labels, attention_mask], num_parallel_workers4 ) dataset dataset.shuffle(buffer_size1000) def per_batch_pad(input_ids, labels, attention_mask, batch_info): bs len(input_ids) max_len max(arr.shape[0] for arr in input_ids) pad_ids np.zeros((bs, max_len), dtypenp.int32) pad_labels np.full((bs, max_len), -100, dtypenp.int32) pad_mask np.zeros((bs, max_len), dtypenp.int32) for i in range(bs): n input_ids[i].shape[0] pad_ids[i, :n] input_ids[i] pad_labels[i, :n] labels[i] pad_mask[i, :n] attention_mask[i] return pad_ids, pad_labels, pad_mask dataset dataset.batch( batch_size4, per_batch_mapper_batch_pad, drop_remainderFalse ) # 如果你是用 model.train 控制 epoch这里就不要再 repeat dataset dataset.repeat(3) it dataset.create_dict_iterator(num_epochs1) for idx, batch in enumerate(it): print(batch[input_ids].shape, batch[labels].shape, batch[attention_mask].shape) if idx 2: break这段代码里几个细节值得展开build_labels做的是“把输入 ids 整体右移一位”的标签构造。很多 GPT 类模型要求labels是下一个 token 的 id所以labels[i] input_ids[i1]最后补一个 eos 结束符。per_batch_pad里的pad_labels用-100填充这是为了配合CrossEntropyLoss(ignore_index-100)使用。在 loss 计算时padding 位置会被自动忽略这样不会把无意义的填充 token 算进损失。drop_remainderFalse表示最后一个 batch 大小不足时也保留。如果你的训练循环对 batch size 有强要求可以改成True但最后一点数据会丢掉。repeat(3)只在你自己写循环迭代时使用如果你把dataset直接传给mindspore.Model.train(epochs3, ...)就不要写.repeat(3)否则每个 epoch 会重复 3 遍数据。3.3 接入 mindspore.Model 与 MindFormers 的衔接如果你使用mindspore.Model来训练接入方式非常简单from mindspore import Model, nn net MyLLM() loss_fn nn.CrossEntropyLoss(ignore_index-100) optimizer nn.AdamWeightDecay(net.trainable_params()) model Model(net, loss_fnloss_fn, optimizeroptimizer) model.train(epochs3, train_datasetdataset, dataset_sink_modeTrue)这里dataset_sink_modeTrue会把数据以 sink 方式直接送入计算设备减少 CPU 到 GPU/昇腾设备之间的搬运开销。但要注意sink 模式对数据 pipeline 的输出 shape 有较强约束动态长度批次很容易报错。所以在真实的大模型预训练中更常见的做法是把序列统一 pad 到固定长度比如 2048 或 4096然后配合attention_mask屏蔽 padding 部分。在 MindFormers 中你也可以把自定义GeneratorDataset包装成 MindFormers 的数据配置来用。核心思路还是这套 pipeline先产出 token id再做 map 变换再 batch再接入训练器。理解了mindspore.dataset之后你就能看懂 MindFormers 源码里那些build_dataset到底做了什么。4. 常见问题与排查技巧实录4.1 CPU 跑满但显存没动静pipeline 卡住了我遇到过最典型的现象是训练进程启动后CPU 核心全部拉满但每一轮 GPU 利用率都只有个位数。排查下来基本都是数据生成端出了问题比如生成器里有大量正则分词、每次都重新加载文件、或者num_parallel_workers设得太高反而导致大量线程切换。排查思路先用num_parallel_workers1跑一遍如果速度和num_parallel_workers4差不多说明瓶颈不在并行度而在于单个 worker 内部耗时。把生成器里的 I/O 尽量提前到数据预处理阶段完成不要在流水线里实时做。用ds.config.set_prefetch_size(16)或更大值加大预取缓冲让流水线更平滑。在训练前单独跑一遍for batch in dataset:用time.perf_counter统计每个 batch 的平均耗时定位是 map 慢还是 batch 慢。4.2 dtype 或 shape 不匹配的报错怎么查这类报错在大模型预处理里很常见。我整理了一个速查表报错现象常见原因解决办法dtype不匹配GeneratorDataset 返回了 Python list统一转成np.int32/np.float32返回列数量不一致map的output_columns和实际返回数量不一致检查map的输入输出列数量是否对齐shape不一致同一个 batch 里样本长度不同用per_batch_map做动态 paddingtokenizer 在 worker 启动时报错多进程下分词器对象不可 picklepython_multiprocessingFalse或全局单例numpy转换失败混合了字符串列和整数列在生成器里先把文本转成 token id每次报错时先单独从 dataset 里取一个 batch 打印出来看看每列的真实 dtype 和 shape这是最快的方式。4.3 各卡数据重复或样本丢失很多人在多卡训练时不使用num_shards/shard_id而是自己在生成器里通过某个随机因子切片结果两个卡的数据重叠或者最后一张卡数据不完整。正确做法是用get_rank()和get_group_size()直接传给GeneratorDataset让框架按样本粒度切割保证互斥且完整。还有一个容易忽视的点如果你在生成器里自己维护了全局随机种子且每个 rank 都用同一个种子那么每张卡看到的 shuffle 结果会完全一样这会让“随机采样”失去意义。shuffle应该让 dataset 的算子管理不要在生成器里手工打乱。4.4 静态图模式下动态 shape 的坑大模型训练时如果用了dataset_sink_modeTrue框架在编译计算图时就会确定输入的 shape。动态 padding 会导致不同 batch 的实际序列长度不一样编译时容易报错。我的建议分层处理如果做预训练直接固定max_seq_len所有样本在 map 里就 padding 到该长度mask 标注真实长度。如果做 PEFT 微调且想省显存可以尝试动态 padding但要在 PyNative 模式或开启动态 shape 支持的前提下运行。如果你既想动态 padding 又想顺畅跑 Graph 模式可以给所有样本统一设置一个接近真实分布的max_seq_len例如 512再让per_batch_pad在一个 batch 内部做长度对齐但保证输出 shape 始终是(batch_size, 512)这样图编译能通过内存和显存也不会浪费太多。写在最后的小经验我自己用过很长一段时间的 PyTorch DataLoader后来切到 MindSpore 做大模型微调最不习惯的就是per_batch_map和GeneratorDataset的配合。踩了几天坑后才想明白一件事mindspore.dataset不是一个“数据集容器”它是一套数据流图每条算子都有明确的执行位置和并行方式你把数据变换拆得越细整条链路的可调优空间就越大。如果只记一条经验我会说不要为了快省掉shuffle也不要图省事把所有样本 pad 到最大长度。前者影响训练稳定性后者影响显存上限。动态 padding 合理的num_parallel_workers 正确分布式分片这三件事做好大模型数据预处理已经能扛住绝大多数生产场景。