昇思MindSpore大模型训练：从评估体系到性能优化的完整实践

昇思 MindSpore 大模型训练跑通只是第一步真正头疼的是后面7B 或 13B 的模型在昇腾集群上训练Loss 在降但吞吐量始终上不去显存明明还有余量数据加载却时不时卡一下同样的卡数别人的 MFU 能跑到 40% 以上自己怎么调都卡在 20%。这些问题的根源往往不是某个优化技巧没学会而是从一开始就没建立起一套完整的评估体系。大模型训练是个系统工程没有评估就没有方向优化也只能靠拍脑袋。这篇就把我在昇思 MindSpore 上做训练评估和性能优化的实践完整拆开讲适合正在用 MindSpore 跑大模型、或者准备从单卡转到集群训练的同学参考。1. 为什么大模型训练必须先有评估体系1.1 训练是系统工程不能只看 Loss很多人训练大模型的时候眼睛只盯着 Loss。Loss 在降就觉得万事大吉Loss 抖动就觉得模型出了问题。但大模型训练和单卡小模型训练有本质区别小模型几分钟跑完错了重来成本低大模型一次训练任务要跑几天甚至几周数据和算力成本都很高如果等到最后才发现性能有瓶颈、显存有浪费那已经亏大了。大模型训练的完整链路包括数据加载与预处理、前向计算、反向传播、梯度聚合、参数更新、Checkpoint 保存等多个环节。任何一个环节出问题都会拖慢整体训练。但麻烦的是这些环节是叠加在一起的从外部很难直接看出到底是哪一环卡住了。我之前就遇到过这种情况训练跑起来之后发现 AI Core 利用率很低但计算量明明很大后来查了半天才发现是数据队列的空转率太高GPU 都在等数据算力全部浪费了。没有评估体系这种问题根本无从查起。评估体系的核心价值就是把训练过程中看不见摸不着的状态变成可量化的数字。只有量化了才能定位瓶颈只有定位了瓶颈才能针对性地优化。这就像工厂生产一样你得先知道哪条产线产能不足才能去改造哪条产线。1.2 评估体系的三层结构模型质量、训练状态、资源效率我在 MindSpore 上的实践下来评估体系至少要拆成三层来看缺一层都不行。第一层是模型质量层。这层回答的是模型学得怎么样指标包括训练 Loss、验证集 Loss、下游任务的效果等。这一层是最终目标但它不能告诉我们训练过程本身有没有问题。Loss 降得慢可能是学习率设置问题也可能是数据质量差原因藏得很深。第二层是训练状态层。这层回答的是训练过程稳不稳指标包括梯度范数、Loss 尖峰频率、参数更新幅度等。这层是承上启下的关键模型质量出了问题往往能在训练状态上找到蛛丝马迹。比如 Loss 突然飙升如果同时看到梯度范数爆炸那基本可以断定是梯度问题如果梯度范数正常而 Loss 飙升那就要怀疑数据质量了。第三层是资源效率层。这层回答的是算力用得值不值指标包括吞吐量tokens/s、FLOPs 利用率MFU、显存占用率、通信开销占比等。这一层是大模型训练特有的评估维度因为训练集群的算力资源昂贵利用率低等于烧钱。这三层的关系是层层支撑的资源效率影响训练速度训练状态影响模型收敛模型质量决定最终效果。很多人在大模型训练中踩坑就是只看了第一层忽略了第二层和第三层。我自己最开始也是这样后来把三层指标都拉起来看之后很多问题的定位速度明显加快了。2. 核心评估指标与 MindSpore 实战2.1 吞吐量训练效率的硬指标吞吐量是大模型训练最直观的效率指标一般用每秒处理的 token 数来衡量。计算公式很简单一个 Step 内处理的 token 数除以这个 Step 的耗时。一个 Step 内处理的 token 数取决于全局 Batch Size 和序列长度。全局 Batch Size 不是简单地等于单卡 Batch Size而是要乘以卡数和梯度累积步数全局Batch Size 单卡微批大小 × 梯度累积步数 × 并行设备卡数举个例子如果单卡 Batch Size 是 4序列长度是 4096用了 8 张卡梯度累积 4 步那么一个 Step 实际处理的 token 数就是 4 × 4096 × 8 × 4 524288 tokens。如果这个 Step 耗时 52 秒吞吐量就是 524288 / 52 ≈ 10082 tokens/s。在实际操作中我习惯在训练脚本里加一个简单的计时回调把每个 Step 的耗时和吞吐量实时打印出来。在 MindSpore 里可以这样写import time from mindspore.train.callback import Callback class ThroughputCallback(Callback): def __init__(self, global_batch_tokens, log_interval10): self.global_batch_tokens global_batch_tokens self.log_interval log_interval self.t0 time.time() def step_end(self, run_context): cb_params run_context.original_args() cur_step cb_params.cur_step_num if cur_step % self.log_interval 0: dt time.time() - self.t0 throughput self.global_batch_tokens * self.log_interval / dt print(fstep: {cur_step}, 耗时: {dt:.2f}s, 吞吐量: {throughput:.1f} tokens/s) self.t0 time.time()这里要注意的一点是计时回调在第一次调用时会包含模型初始化和编译的时间所以通常跳过前几个 Step 的数据从稳定后的第 5 个 Step 开始统计。MindSpore 的图模式在第一次迭代时需要编译计算图这个时间可能比较长不算入稳态吞吐量。吞吐量的参考值因模型规模、硬件和并行配置差异很大。以我在 8 卡 Atlas 800T A2 训练 7B 模型的经验来说纯数据并行、序列长度 4096 的情况下稳态吞吐量在 8000 到 12000 tokens/s 之间属于正常范围。低于这个区间就要警惕了后面第 4 节的排查思路可以帮你定位。2.2 FLOPs 利用率MFU算力用得值不值单纯看吞吐量还不够因为吞吐量受全局 Batch Size 影响很大。同样一个模型把全局 Batch Size 翻倍吞吐量可能能提升 80%但这是靠堆显存换来的不代表算力效率变好了。这时候就要看 MFUModel FLOPs Utilization也就是模型计算效率。MFU 的公式是MFU 实际完成的 FLOPs /硬件理论峰值 FLOPs × 运行时间这里的实际 FLOPs 对 LLM 训练来说有个比较通用的估算方式。每处理一个 token前向传播的计算量大约是 2 倍的参数量反向传播大约是 4 倍加在一起大概是 6 倍参数量。所以实际 FLOPs 6 × 模型参数量 × 处理的 token 总数举个例子一个 7B 模型处理 1000 万个 token理论计算量大约是 6 × 7e9 × 1e7 4.2e17 FLOPs。如果单卡理论峰值算力是 313 TFLOPS昇腾 910B 的 FP16 算力8 卡跑这个任务用了 600 秒那么MFU 4.2e17 /313e12 × 8 × 600≈ 0.28也就是 28%这个数字在真实训练场景里属于中等水平。大模型训练的 MFU 能到 40% 以上就算非常优秀了30% 左右是常见水平低于 20% 就说明有重大问题需要排查。这里要说一个特别容易踩的坑不同框架和工具统计 FLOPs 的口径不一样有的只算前向有的把激活重计算也算进去了有的不包含通信空闲时间。所以跟别人对比 MFU 之前先确认大家对口径的理解是否一致不然数字差了 10 个百分点都不知道为什么。2.3 显存与通信开销资源健康度监测显存是训练资源的另一个硬约束。要评估显存是否健康首先得知道显存花在了哪里。大模型训练时显存主要被五部分占用模型参数FP16 下每个参数占 2 字节7B 模型就是 14GB梯度同样 FP16也是 14GB优化器状态AdamW 需要额外的一阶动量FP32 4字节和二阶动量FP32 4字节加上参数副本FP32 4字节每个参数最多需要 12 字节7B 模型就是 84GB激活值取决于 Batch Size、序列长度和层数这个是大头通信缓冲区和临时张量一个 7B 模型不开重计算的情况下激活值很容易撑爆显存。我实测单卡 Batch Size 为 1、序列长度 2048、32 层 Transformer 的激活显存大约在 12 到 20GB 之间如果序列长度翻倍到 4096激活显存会线性增长到 24 到 40GB。这在单卡 64GB 显存上已经非常吃紧了。MindSpore 的 Profiler 工具可以采集显存使用情况输出每个算子的显存分配和释放记录。我在排查 OOM 时第一步就是看 Profiler 报告把前十个显存占用最高的算子列出来就能快速判断激活值是不是主因。昇腾场景下使用npu-smi info也能实时查看每卡的显存用量和算力负载适合做粗粒度监控。通信开销也是评估体系里不可忽略的一环。数据并行训练在每轮梯度计算后都要执行一次 AllReduce把各卡的梯度聚合起来。梯度同步的时间受两个因素影响带宽和消息大小。通信占比可以通过 MindSpore Profiler 的 Step Trace 界面直接看它会显示每个 Step 中通信算子的耗时占比。如果通信占比超过 20%就说明梯度聚合开始拖累整体效率了这时候就该考虑梯度压缩或者减少通信频率了。判断一个训练任务到底是计算瓶颈还是通信瓶颈有个很实用的技巧把 Batch Size 减半看吞吐量是否成比例下降。如果吞吐量也差不多减半说明算力是瓶颈如果吞吐量下降不明显说明大量时间花在了固定开销上比如通信等待或者数据加载上。这个简单的实验能帮你快速锁定优化的方向。3. 从评估结果到性能优化的落地链路3.1 数据管道从 CPU 到 NPU 的第一道坎评估做完优化就该上场了。首先要查的是数据管道。大模型训练的数据量巨大数据读取、解码、图像变换、tokenize 这些操作如果在 CPU 上执行效率不高NPU 就会大量时间在等数据。我在 MindSpore 上遇到过一次这种情况AI Core 利用率只有 15%怎么调并行策略都没用后来发现数据加载的耗时已经超过了计算耗时的三倍。MindSpore 的数据管道主要靠mindspore.dataset来构建。几个关键参数的配置直接影响管道效率num_parallel_workers控制每个操作并行线程数。多核机器上建议设置为 CPU 核数的四分之一到二分之一过高反而会因为线程切换导致性能下降。我一般在 64 核的机器上设 16 到 24 个 worker。prefetch_size控制预取到内存的 batch 数量。默认值可能偏小如果发现数据队列空转率高可以加大这个值。我习惯设为 2 到 4让数据管道提前准备好两个到四个 batch避免计算单元空闲。shuffle机制也要注意。大模型训练需要保证数据的随机性但全局 shuffle 的开销很大。MindSpore 的mindspore.dataset.config.set_seed配合shuffle_buffer_size可以实现一个较优的随机策略每个 epoch 开始时用较大的 shuffle buffer 打乱一次数据训练过程中再用较小的 buffer 做局部扰动。这个策略在我的实践中能保持准确率不下降的同时减少接近 40% 的数据重排耗时。还有一点很容易忽略的是数据落地格式。明文本地存储的 TFRecord 和 MindRecord 格式在读取效率上有数量级差异。MindRecord 是昇思专门为高性能数据读取设计的格式底层做了块预读和多线程解压优化。把训练数据从 JSON 转成 MindRecord 之后我的数据加载耗时直接减掉了大约 60%。评估数据管道是否健康可以看训练日志里打印的数据队列利用率。MindSpore 在 Profiler 报告的具体数据中会显示每个 Step 中数据加载和计算的重叠情况。如果数据加载时间完全暴露在计算时间之外没有重叠那就说明数据管道的并行度不够。3.2 并行策略让卡数翻倍收益不衰减大模型训练很少只用单卡所以并行策略的选择和配置是性能优化的重心。常见的并行维度有数据并行DP、张量并行TP和流水线并行PP三者的适用场景完全不同。数据并行最简单每卡持有完整模型副本只同步梯度。但问题是单卡装不下大模型而且全局 Batch Size 太大时收敛效果会变差。实测下来数据并行虽然显存效率最高但也是通信开销最大的方式因为每个 Step 都要全量同步梯度。张量并行把单个 Transformer 层按矩阵维度切开分布到多卡上分别计算。好处是显存压力减小坏处是两个相邻算子之间要频繁做 AllReduce 和 AllGather 通信。张量并行适合单机多卡场景因为 NVLink 或 HCCS 的带宽够用跨机时通信开销会显著上升。昇腾的 HCCS 总线在单机 8 卡范围内带宽很高所以 8 卡以内部署张量并行非常合适。流水线并行把模型按层切成多个阶段每个阶段放在不同的卡上。这样节省显存效果很明显但会产生气泡Bubble也就是某些卡在等待上游阶段输出的时候处于空闲状态。气泡率和切分的数量有关理论上流水线段数越多气泡越多。不过把微批次切小一点让流水线尽量填满可以把气泡控制在可接受范围内。MindSpore 的并行配置方式已经有了比较成熟的封装。数据并行直接设置context.set_auto_parallel_context(parallel_modedata_parallel)就能用。张量并行和流水线并行需要通过mindspore.rewrite或model.parallel_optimizer配合算子级别的标注来实现。我的建议是能用框架自动并行方案解决的别自己去手工切分模型手工切分虽然理解起来更直观但容易在边界条件上出错调试成本很高。从我踩过的坑来看并行策略的选择没有绝对最优只有针对当前场景的相对最优。8 卡以内的 7B 模型我用的是数据并行加梯度累积13B 模型单机放不下我就切成 2 路张量并行加 4 路数据并行更大规模的模型就要引入流水线并行组合方式需要根据显存带宽和通信延迟的实测值来定。这个类比的逻辑很简单就像搬家一样小房子自己搬就行别墅就得请吊车但吊车也分大小得看现场条件。3.3 算子与图优化榨干最后一点算力数据管道的冗余清掉了并行策略也定下来了接下来就要往算子层面抠性能了。昇思 MindSpore 在昇腾上的性能和有没有正确使用图模式、有没有开启算子融合关系很大。MindSpore 有两种执行模式PyNative 模式动态图和 Graph 模式静态图。PyNative 模式方便调试但算子间有大量的框架调度开销训练吞吐量一般会打六到七折。我在实际训练中总会在确认代码逻辑稳定之后切到 Graph 模式跑正式训练任务。Graph 模式会把整个网络编译成一张静态计算图框架可以基于整图做算子融合、内存复用和调度优化。MindSpore 的图优化能力里最实用的是算子融合。典型的融合场景包括把连续的多个小算子融合成一个大算子减少内核启动次数把归一化、激活函数和 Dropout 融合到主算子内部把矩阵乘法和后面的加偏置、激活融合成一体。这些融合操作不需要手工改模型代码MindSpore 编译期会自动处理你只需要确保用的是 Graph 模式。另一个重要的优化手段是重计算Recompute。Transformer 层在前向传播时会保存大量中间激活值用于反向传播这是显存消耗的大头。重计算的思路很简单前向传播时不保存中间激活反向传播时重新算一遍。这会让计算量增加约 30%但能显著减少显存占用——以 70B 规模的模型来说不开重计算根本训不动开了之后就能装进显存。MindSpore 里开启重计算也很方便在网络的 Cell 上调用cell.recompute()就可以了。不过要注意重计算不能盲目全开因为算力开销是实打实的。我建议只对 Transformer 层中显存占用最大的那几个模块开启重计算比如 Attention 和 FeedForward 的中间结果这样能平衡显存节省和计算开销。混合精度也是一个性价比很高的优化手段。昇腾 910B 对 FP16 计算的加速比通常能达到 FP32 的好几倍所以大模型训练几乎都是混合精度。MindSpore 的amp.auto_mix_precision可以自动识别哪些算子适合降精度计算。我实测下来开启混合精度后训练速度大概能提升 60% 到 80%而精度损失在合理容差范围内。唯一需要注意的是一些数值稳定性敏感的算子比如 Layernorm 和 Softmax框架一般会自动保留高精度计算不需要你手动干预。4. 踩坑实录评估与优化中的典型问题4.1 Loss 不降反升是模型问题还是数据问题训练中间发现 Loss 不降反升这是大模型训练里最让人紧张的场景之一。我在一次训练对话模型的实验里就遇到过这种情况跑到第 3000 步左右Loss 曲线突然从 1.8 飙到 2.6然后又慢慢回到 1.7 附近过几百步又飙一次。遇到这种问题我先查了梯度范数。MindSpore 的train.GradMonitor回调可以实时记录每个 Step 的梯度范数发现 Loss 飙升的时刻梯度范数确实出现了 10 倍以上的暴涨。这基本可以判断是梯度不稳定的问题。接下来要排查的是学习率是否设置过大。我的配置是固定学习率 3e-4在 7B 模型上这个值本身就偏大配合大 Batch Size 更容易触发梯度爆炸。把学习率改为前 2000 步线性 warmup 到 1.5e-4之后用余弦衰减降到 1e-5Loss 尖峰出现的频率明显降低了。这个操作背后的逻辑是大模型训练初期参数随机性大梯度幅度也大此时过大的学习率很容易让参数更新步长超出稳定范围。但有一次 Loss 不降却查不出梯度问题梯度范数完全正常参数更新幅度也很平稳。后来细查数据集才发现训练数据里混入了大量重复样本和少量标注错误的样本导致模型在部分 batch 上学到的信息互相矛盾。清洗数据之后Loss 曲线立刻变得平滑。这是数据层的问题评估体系里如果只看模型质量层指标很难定位到这一步。所以我在实践里养成了一个习惯每次训练之前先花时间做数据质量抽检跑一个小的样本集验证数据分布是否符合预期。4.2 吞吐量上不去从 Profiler 的数据里找答案吞吐量上不去的排查我有一套固定的操作流程。第一步是打开 MindSpore Profiler 的 Step Trace 报告看一个 Step 的耗时构成。这里能看到数据处理耗时、算子计算耗时、通信耗时和框架调度耗时的占比。如果数据处理耗时占比很高优先优化数据管道采用前面提到的方法调整num_parallel_workers和prefetch_size或者转换数据格式。如果算子计算耗时占比高就要看具体是哪些算子拖延了时间。Profiler 的算子耗时列表会按执行时间降序排列重点关注 Top 10 的算子。如果在 Top 10 里看到大量的小算子说明图优化和算子融合没有生效可以检查是否误用了 PyNative 模式如果看到某个大算子特别慢就要检查它的输入是否出现了频繁的格式转换。如果通信耗时占比高先看通信算子的类型。如果 AllReduce 占比很高可以尝试梯度压缩或梯度累积来降低通信频率如果 AllGather 占比很高通常是张量并行维度过大导致的需要调整并行策略。我在一次多机训练中遇到的麻烦是单机内耗时正常但拆到两台机器后吞吐量掉了一半。Profiler 显示通信耗时占比达到 35% 以上而且通信算子集中在跨机链路上。梯度的全量 AllReduce 产生的消息量太大跨机带宽成了瓶颈。后面改用梯度压缩方案把 FP32 梯度量化为 FP16 传输通信耗时降到了 18%吞吐量恢复到了原来的八成。这里要强调的是梯度的量化会增加部分误差需要观察 Loss 收敛曲线来判断是否可接受。我在这个案例里观察了 5000 个 Step收敛曲线和 FP32 传输没有明显差异。4.3 显存 OOM 的快速定位方法OOMOut of Memory应该是每个跑大模型训练的人都遇到过的问题。MindSpore 在显存不足时会直接报错但错误信息往往不直接显示是哪个算子导致的。我的排查方法是先把模型的最大显存占用降下来再逐步加回 Batch Size。具体来说先把 Batch Size 设为 1序列长度保持不变跑一个 Step 看显存占用。如果单卡显存占用已经超过 70%说明模型结构和激活值本身就是显存大头需要使用重计算、减少模型层数或改用更小的模型配置来释放空间。如果单卡 Batch Size 1 时显存占用在 50% 以下OOM 大概率是安全和缓存临时张量过多导致的适当增大prefetch_size反而可能通过减少重复申请张量来降低峰值的显存压力这个操作看似矛盾但在连续训练的场景下确实能缓解碎片化问题。MindSpore 还提供了显存池设置可以通过环境变量调整缓存池大小。昇腾场景下这个设置经常被忽视但实际影响很大。我遇到过一次卡上明明有 20GB 空闲显存但 OOM 报错的案例后来排查发现是显存池的缓存策略导致碎片化严重调整相关参数之后问题就消失了。这类问题的难点是复现困难所以建议在训练开始前就预估好显存分配多留一点余量不要把手头的每字节显存都压满。5. 实操心得与建议做这套评估与性能优化下来我最大的感受是评估体系的建立要趁早不要等训练任务跑起来之后再去补。第一次跑大模型任务前先把吞吐量回调、Profiler 采集、显存监控这些工具都配置好让每个指标在训练第一天就有数据。性能问题越早发现定位和修复的成本就越低。第二个体会是优化一定要一次只改一个变量。把学习率、Batch Size、并行策略、数据管道参数一起调整出问题了根本不知道是哪个改动导致的。我养成的一个习惯是每次只改一个参数跑 200 个 Step 看指标变化确认没有异常再改下一个。这看起来进度慢但实际上比盲目调参快得多。第三个建议是性能优化的目标要量化。不要用让训练更快一些这类模糊目标要定成把 MFU 从 22% 提升到 30%或把数据加载耗时占比从 30% 降到 10%。只有量化目标才能验证优化是否有效也才能判断下一步往哪个方向走。最后再分享一个小技巧把评估指标做成自动化的日志和看板不要靠人肉盯训练过程。我后来把吞吐量、Loss、梯度范数、显存占用都输出到日志并在终端实时刷屏训练任务跑起来的第一个小时就在盯这些指标发现问题马上介入。这套流程跑熟之后性能问题基本都能在训练早期被发现不用等到跑了一半才发现模型根本训不动。