ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

昇思大模型数据预处理全指南:从数据管线到MindRecord优化

昇思大模型数据预处理全指南:从数据管线到MindRecord优化 在昇思上跑大模型大家第一眼看的往往是模型结构、loss 设计、分布式并行策略但真正让训练卡住的地方却经常藏在数据层里。我印象很深的第一次大模型实验模型是标准的 decoder-only 结构并行方案也没问题可每轮迭代 GPU 利用率就是上不去盯了半天网络和优化器最后发现是 mindspore.dataset 里一个朴素的数据变换写法把整个流水线堵死了。这篇博客我打算把大模型场景下基于 mindspore.dataset 的数据变换与预处理方案完整串一遍从最基础的管线骨架到文本 tokenizer 和掩码构造再到 MindRecord 和多卡吞吐优化最后讲讲那些线上排查才能真正体会到的坑。这套内容适合正在做昇思大模型训练或微调、想系统化搭建数据预处理流程的人看。不管你是刚把某个开源模型跑通还是已经在调多卡训练脚本里面提到的很多细节应该都能直接帮上忙。1. 数据管线是大模型训练的隐形瓶颈1.1 一个 step 里 GPU 到底在等谁很多人习惯把训练性能问题一刀切归结为显存不够或并行策略不对但等你用 profiler 看一遍迭代时间分布会发现数尤其是微调场景GPU 每一步要等的时长里数据准备经常能占到 30% 以上。训练日志里显示的 step 时间包含了数据从磁盘、内存、经过一系列 map 算子转换成 batch 张量再送入设备的全过程。如果 pipeline 里写了一个特别耗时的 Python 预处理函数这个等待会被进一步放大。昇思的数据集模块说白了就是一条生产流水线源头是文件或者 Python 生成器中间串联若干数据变换map 算子末端是 batch 和 repeat。跟手写 for 循环里一点点预处理相比这套机制天然支持并行 worker、预取队列和算子融合。如果只是把 Gym 里搬来的数据处理脚本原封不动塞进 .map()那你实际上还在用单线程做预处理Gad 多卡并行优势完全发挥不出来。1.2 什么时候应该认真看待数据变换判断瓶颈是不是数据侧其实有个很粗糙但好用的方法先把 .map() 里的算子去掉只保留数据读取和 batch看训练能跑到多少迭代每秒。恢复算子后再跑一次如果速度明显下降说明预处理环节偏重如果几乎没变化那瓶颈在别处。大模型预训练语料动辄几十 GB 甚至上 TB每个 step 要从原始文本里切出几百个样本还要在线做 tokenize、截断、mask 拼接数据处理开销不会小值得好好优化。那为什么很多人忽略了这条呢因为训练脚本里写数据部分通常只占 50 行模型部分五六百行目光自然集中在模型和训练循环上。但恰恰是这 50 行的设计决定了 64 卡训练时每个 step 的喂数速度。2. 搭一条可扩展的 mindspore.dataset 管线骨架2.1 数据源选择内置 Reader 还是 GeneratorDatasetmindspore.dataset 提供了多种数据读取方式大模型场景下最常见的两个选择是内置文件读取器和 GeneratorDataset。先看内置的 TextFileDataset、MindDataset 这套。它们的好处是走的是 C 底层实现读取效率高并且天然对接 shuffle、shard 等分布式能力。缺点也很明显它对文件格式有要求要么是文本行要么是 MindRecord原始的 JSONL 语料不能直接读。再看 GeneratorDataset。它接收一个 Python 可迭代对象灵活度非常高无论数据在 JSONL、CSV、数据库还是内网接口里只要你能写个生成器产出样本就能接进 dataset 管线。代价是 Python 到 C 之间有一层包装如果生成器里写了太多复杂逻辑会成为新的瓶颈。我自己的经验是小规模微调、数据格式快速变化阶段优先用 GeneratorDataset大规模预训练或数据格式稳定之后一定迁到 MindRecord。前者保证开发效率后者保证长时间训练不拖后腿。2.2 列设计从样本到模型输入的字段规划跟手写训练脚本不同mindspore.dataset 是列式处理的每个样本是一行包含若干 column。大模型微调的最小数据列通常这样设计列名含义示例prompt用户输入原文请总结这段文本: ...response期望输出这段文本主要讨论...input_idstokenize 后的输入 id 序列[101, 2043, 2001, ...]attention_mask需要模型关注的 token 位置[1, 1, 1, 0, ...]labels计算 loss 的目标 id通常对 answer 部分保留真实 id其余为 -100[-100, -100, 2043, 2001]列设计要在管线搭建前想好因为后边的 map 算子、batch 逻辑、模型前向的数据对齐全部依赖它。很多人图省事把 prompt 和 response 拼成一个大字符串存成一列然后 tokenize 之后再想办法切分绕一大圈还容易错位。正确做法是原始字段各自保留在变换阶段按需拼接和生成新列。2.3 shuffle、repeat、batch 的正确排列这是数据管线里最容易踩坑的地方。我见过不少训练脚本写成这样先 .batch() 再 .shuffle()或者先 .repeat() 再 .shuffle()最终效果千奇百怪。官方文档里给出的推荐顺序是基于实际数据流语义推导出来的shuffle 打乱样本顺序batch 按打乱后的顺序切分成组repeat 对整个 epoch 重复如果先 batch 再 shuffle那么 batch 内的样本永远是相邻的模型看到的 batch 之间相关性很高如果先 repeat 再 shuffle那么 shuffle 的随机性会被 repeat 稀释每个 epoch 看到的顺序差异变小。大模型训练最忌讳这种结构化的数据顺序因为模型很容易学到序列位置带来的假规律导致验证集表现虚高。除了顺序shuffle 的 buffer_size 也值得注意。它不是把所有数据读进来再打乱而是维护一个滑动窗口窗口越大随机性越好但内存开销也越大。对于几十 GB 的大语料窗口设成几千到几万是常见的折中。import mindspore.dataset as ds dataset ds.GeneratorDataset( sourcemy_generator(), column_names[prompt, response] ) dataset dataset.shuffle(buffer_size10000) dataset dataset.batch(batch_size8) dataset dataset.repeat(num_epochs)3. 文本变换tokenizer、截断与掩码的构造细节3.1 在 transform 算子层做文本处理的选择文本预处理无非两件事把字符串变成 token id 序列再把序列整理成模型需要的定长结构。落地到 mindspore.dataset 时有个核心选择用框架自带的 text 变换算子还是自己写 Python 函数塞进 map。mindspore.dataset.text 下面提供了 BertTokenizer、SentencePieceTokenizer、Vocab 等工具。好处是底层可能是 C 实现并且能直接跟 dataset 的并行机制融合跑起来更快。缺点是需要多花时间理解它的调用方式对于 huggingface 的 tokenizer 生态兼容性一般。如果你的模型用的是自定义词表或者已经在 huggingface tokenizer 上做了大量工作那就没必要把分词逻辑重写。更现实的做法是在 GeneratorDataset 的源生成器里做 tokenize直接产出 id 序列或者写一个 Python 函数作为变换算子在 map 阶段调用第二种方式实操更多因为很多大模型微调框架的 tokenize 逻辑依赖 huggingface tokenizer 的分词细节重写到 text 算子层反而容易引入不一致。比如 BERT 式的 WordPiece 分词、GPT 式的 BPE 分词这些细节在两个生态之间迁移时经常出现微妙差异一旦 tokenization 不一致同一个句子的 id 序列就对不齐下游训练结果根本没法复现。3.2 截断、填充与 bucket_batch_by_length大模型输入通常是定长的。短的要填充到固定长度超长的要截断。填充的位置也有讲究decoder-only 架构一般左侧填充因为要保证最后一个 token 是真实文本末尾方便自回归生成encoder-decoder 架构则通常是右侧填充。直接用 pad 到 max_length 是最简单的方案但存在严重浪费。比如语料平均长度 300你却 pad 到 1024那约七成算力都被 pad token 浪费了。工程上常用 bucket_batch_by_length先按样本长度分桶每个桶内部 pad 到该桶的最大长度桶之间 batch 后会形成动态 shape。昇思的 mindspore.dataset 里提供了现成接口凡是没 max_length 硬约束的大模型训练我都建议用它通常能省下 20% 到 40% 的无效计算。3.3 标签掩码微调场景的 -100 规则大模型微调跟预训练有个关键区别预训练是整段文本都用来预测下一个 token所以 labels 就是 input_ids 顺移一位微调往往只有 answer 部分参与 loss 计算prompt 部分不计算否则模型只需要机械背题指令跟随能力反而下降。常规做法是把不参与 loss 的位置设成 -100 或某个忽略标记。在 mindspore.dataset 里这个逻辑可以写成一个 map 函数def build_labels(input_ids, prompt_len): labels [-100] * len(input_ids) labels[prompt_len:] input_ids[prompt_len:] return labels dataset dataset.map( operationsbuild_labels, input_columns[input_ids, prompt_len], output_columns[labels] )这里有个容易出错的地方input_columns 和 output_columns 的对应关系必须写清楚。map 默认会把所有输入列转换成输出列如果你的函数只返回一个 labels那原来的 input_ids 列可能被覆盖。踩过这个坑之后我习惯在 map 里显式声明 input_columns 和 output_columns即使看起来冗余也比留隐患强。4. 参数化调优并行度、算子合并与缓存策略4.1 num_parallel_workers 设多少合适mindspore.dataset 的 map、batch、shuffle 都支持 num_parallel_workers 参数控制该算子使用多少并行线程或进程。很多人图省事直接设个 8 或 16但实际效果可能出乎意料。并行度不是越高越好。大模型训练机器 CPU 核数通常几十个数据管线算子又是多级串联的每一级都开太多 worker反而导致 CPU 频繁上下文切换整体吞吐下降。我一般用以下方式粗调看机器 CPU 核数和内存从 num_parallel_workers4 起步逐步上调每档跑 100 个 step 看耗时找到耗时最低点后再加一档确认防止随机波动同时观察内存占用防止 worker 数过大导致 OOM需要注意的是如果 map 里跑的是 Python 函数尤其是 huggingface tokenizer受 GIL 限制开再多的线程也可能还在单核上挤牙膏。这种情况建议使用 py_transforms 或让 GeneratorDataset 本身启用多进程。4.2 把多个变换塞进一个 map数据变换经常是一串算子tokenize、截断、build labels、转 numpy 数组。新手喜欢写成dataset dataset.map(tokenize) dataset dataset.map(truncate) dataset dataset.map(build_labels) dataset dataset.map(to_array)这种写法逻辑清晰但性能并不好。每个 map 都是一条独立的调度边界数据要经过多次内部队列和线程切换。更优的做法是把强耦合的操作合并进一个函数用一次 map 完成多步变换。从模型角度tokenize 和 mask 构造本来就是一套连贯逻辑中间出现中间态 token 序列没有保留价值合并后内存也更省。我实践下来一个 map 里做完 tokenize 截断 label 构造 转 ndarray比拆成四个 map 能减少约 15% 的端到端耗时。不过也有例外tf 之类的重变换因为调用外部库且结果可复用拆开反而方便配合 cache 使用。这种情况下合并与否要看具体算子特性不能一概而论。4.3 cache什么时候能缓存什么时候不能mindspore.dataset 提供 .cache() 能力把经过某个算子之后的数据缓存到内存或本地磁盘后面若干个 epoch 直接复用。这个机制非常适合数据增强开销大的扩展领域但在大模型场景下要谨慎。如果数据变换是无状态的比如 tokenize、build labels那么缓存非常有效。训练进行到第 2 个 epoch 时不需要重新做分词直接从缓存读取CPU 压力大幅下降。但如果你的训练流程里依赖了随机增强、随机采样或者样本顺序每个 epoch 都要变化那么缓存范围就要仔细界定否则会破坏随机性等于变相降低了数据多样性。另外还要注意缓存和分布式训练的相互作用。多卡场景下如果每张卡各自缓存一份完整数据显存之外的内存会被迅速吃光。通常的做法是让每张卡只缓存自己分片那部分或者干脆在数据量够大的情况下放弃缓存直接用 MindRecord 配合高并行读取。5. 分布式训练与 MindRecord从文件到多卡数据流5.1 为什么大模型最终绕不开 MindRecord大模型预训练数据量大、原始格式五花八门相比之下 MindRecord 有几个很硬的理由成为最终的存储格式。首先是读取效率。MindRecord 是一种二进制格式内部按列存储读取期间不需要做 JSON 解析也不需要每次启动都重新跑一遍 Python 生成器。其次是随机访问友好它天然支持按样本索引读取分布式训练每个 shard 可以各自定位到对应数据块不需要每张卡都遍历全量文件。再次是内存占用可控训练脚本在数据侧不保留全量样本MR 文件被按需读取整体内存曲线非常平稳。这意味着如果你的数据管线每次启动训练都要把几十 GB 的 JSON 重新解析一遍那说明还没到真正的大模型工程阶段。到了这个阶段我一定会花半小时做一次离线转换把语料变成 MindRecord训练时直接读省下的是每天的重复 CPU 消耗。下面是我的转换脚本骨架核心目的是把 GeneratorDataset 产生的样本写入 MindRecord 文件import mindspore.dataset as ds from mindspore.mindrecord import FileWriter def generator(): for rec in json_loader(train.jsonl): yield rec[input_ids], rec[attention_mask], rec[labels] dataset ds.GeneratorDataset( sourcegenerator, column_names[input_ids, attention_mask, labels] ) FileWriter.open(train.mindrecord).write_dataset(dataset)实际操作时要注意把 source 定义成可重复调用的生成器或具备iter的类。FileWriter 在底层会反复调用它来读取样本如果生成器只能跑一次转换会在中途报错。5.2 多卡训练下的 shuffle 与分片一致性分布式场景数据侧最大的坑往往在 shuffle 和分片的关系上。8 卡训练时如果每张卡各自读全量数据再自行 shuffle那同一份样本可能被多张卡重复训练梯度更新变相放大学习率如果每张卡只读自己的 shard 但没有统一的 shuffle 策略那么每张卡长期看到的都是固定顺序模型可能出现局部过拟合。昇思的处理方式是在 dataset 上设置 num_shards 和 shard_id。每张卡只读自己的数据分片同时在每个 epoch 内 shuffle。这样可以做到全局样本不重复且每个 epoch 的样本顺序重新打乱。关键在于 sharding 必须发生在 shuffle 之前否则每张卡拿到的分片内容是随机切分的样本之间会出现交集或漏读。实际环境中因为多卡启动器对 dataset 的实例化次数不一致容易出现一个 epoch 结束后各卡分片漂移的情况。我在排查这类问题时会在每个 epoch 结束时打印每张卡的最后一个样本 id对比是否按预期循环。这个调试手段虽然笨但能快速定位数据顺序错乱。5.3 动态 shape 与 MindRecord 的配合大模型训练对数据侧还有个特殊要求shape 最好静态至少在一个 step 内稳定。MindRecord 在写入时如果样本长度不一致读取后再做 bucket_batch_by_length 也挺自然。但有一个细节我吃过亏MindRecord 的 shape 信息是写入时就定好的如果后续你的模型改 max_length旧文件就不适用了得重新生成。所以我的习惯是第一次生成 MindRecord 就固定好 max_length 或固定 bucket 策略不要中期随意改变长度配置。真要调宁可重新离线生成也要避免在训练循环里动态改采样逻辑那会显著拖慢数据读取。6. 线上排查用 10 分钟定位数据侧问题6.1 先看单个样本链路再谈优化训练卡住或 loss 不收敛时第一件事不是改模型而是用 dataset 的 create_dict_iterator 跑一个 step把拿到的 batch 打印出来。这一步能暴露大部分问题列对不对、shape 是否合理、数据是否为空、mask 是否全零。for batch in dataset.create_dict_iterator(output_numpyTrue): print(batch[input_ids].shape) print(batch[attention_mask][:2]) break如果这一步能看到完整 batch 且数值合理那问题基本在模型侧的数据消费方式上如果连 batch 都拼不出来那问题一定在数据管线里。我十次排障有八次靠这一步定位方向根本不需要立刻上 profiler。6.2 用 profiler 看 wait 时间昇思的 profiler 可以输出数据集算子耗时能够看到每个 map 算子、batch 算子、shuffle 算子分别花了多少时间以及训练任务在等数据时的 wait 时间。数据分析的几个重点某个 map 算子的耗时远高于其他算子优先优化这个 Python 函数或对它做缓存batch 算子的耗时突然升高可能是动态 shape 导致的频繁重分配wait 时间占比高说明数据供给跟不上 GPU 消费需要加大 prefetch 或提高并行 worker这些数据比任何猜测都靠谱。后续微调参数时每一次改动都用 profiler 对比而不是拍脑袋。6.3 动态 shape 导致前向报错的排查链路大模型场景用 bucket_batch_by_length 之后batch 内长度不一部分网络结构可能直接报 shape 不匹配。遇到这种情况我很少去改模型的前向逻辑而是先检查 dataset 输出打印一个 epoch 里所有 batch 的 shape 集合确认它们分布在哪些长度区间必要时调整 bucket 的边界值把过长的样本单独分桶实在无法处理就在 bucket 算子之前加一个 filter把超长样本过滤掉这里有个隐藏问题filter 掉样本会改变数据量影响 epoch 总步数。如果多卡环境里各卡过滤条件一致问题不大但如果过滤逻辑依赖随机状态各卡过滤量不一致就会产生梯度步数不齐。所以 filter 要写纯确定性逻辑不能依赖随机数。6.4 损失不降先怀疑数据最后一句话送给所有在昇思上做大模型的同行也是我踩了无数坑之后最值钱的体会模型不收敛永远先怀疑数据再怀疑超参最后才怀疑模型结构。数据顺序是否被意外固化、mask 是否错位、label 是否出现越界、tokenizer 是否在某个 step 被错误截断这些问题的排查成本往往远低于改一次网络再去重新训练。我见过有人在一版数据里连续训了两天最后发现 attention_mask 全部置 1prompt 和 response 一起参与 loss 计算模型学会了复读机行为。这个问题如果在训练第一天用上面第一招打印 batch十分钟就能抓到根本不用等两天后的结果。数据预处理看起来是杂活但它才是整个大模型训练工程里杠杆率最高的环节。
返回列表