ARTICLE DETAIL

资讯详情

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

昇思MindSpore大模型训练评估体系与性能优化实践

昇思MindSpore大模型训练评估体系与性能优化实践 大模型训练跑起来了loss也在降但我现在反而更关心另一组问题当前这步训练到底算快还是慢显存余量离OOM还有多远模型并行切得合不合理这些要是答不上来那训练过程就还不在掌控之中。我这两年在昇思 MindSpore 上折腾过从亿级到千亿级参数的模型训练从单卡调试到集群并行都踩过不少坑最深的体会是性能优化的前提是先有一套可量化、可对比的评估体系。没有这套东西所有优化都是拍脑袋今天调个学习率明天改个并行度最后连自己都不知道哪个改动真正起了作用。这篇文章不打算讲教科书式的原理而是把我在昇思上建立训练评估体系和做性能优化的完整思路写出来。你可能是刚开始用 MindSpore 训练大模型的人也可能已经在调参但总感觉卡在瓶颈上又或者纯粹想看看别人是怎么把 profiling、并行策略、显存优化串起来的。无论哪种我相信这套“先评估、再定位、后优化”的流程值得你抄作业。1. 为什么大模型训练要先建立评估体系1.1 只看 Loss 远远不够很多人觉得训练跑起来、loss 在下降就是成功了。这个判断在大模型场景下非常危险。Loss 下降只代表模型在朝着正确的方向优化参数但它完全不告诉你计算资源的利用效率、通信开销占了多少、显存还有没有冗余、数据加载有没有拖后腿。我举个例子。两次训练都降到同样的 loss一次用了 10 小时一次用了 6 小时区别可能不在于模型结构而在于并行切分、通信拓扑和数据管线。如果你没有记录吞吐、MFU、通信占比这些指标你根本说不出 6 小时是为什么快。下次换个集群或改个模型规模你依旧无从下手。1.2 评估体系的分层设计我习惯把大模型训练评估体系分成四层健康层、效率层、资源层、稳定性层。每一层回答不同的问题。第一层是健康层核心看 loss 曲线和验证集精度。它回答“模型学没学过”。但要注意观察 loss 的方式不能只看绝对值要看趋势、噪声和与理论基线的差距。如果 loss 下降速度比预期慢那可能是学习率策略不对或者数据喂错了。第二层是效率层核心看吞吐量、MFUMachine Flop Utilization实际有效计算量与硬件理论峰值算力的比值。它回答“算力有没有被用起来”。吞吐可以用 tokens/s 或者 samples/s 衡量MFU 是更严格的指标。比如一张 A100 的理论 BF16 算力大约是 312 TFLOPS如果你训练一个 7B 模型达到的 MFU 只有 12%那说明有大量算力被浪费在等待、通信或低效算子上了。第三层是资源层看显存占用、内存带宽、通信带宽、磁盘 IO 等。它回答“硬件资源配置是否平衡”。显存接近上限意味着你可能随时 OOM显存占用过低则暗示并行切分或 batch size 有调整空间。通信带宽占用高但不一定有问题要结合通信占比和是否重叠来判断。第四层是稳定性层主要记录训练是否频繁中断、是否出现溢出、是否发生梯度异常、是否能从 checkpoint 顺利续训。大模型训练动辄几周甚至几个月稳定性在某种程度上比单步性能更重要。一次失败回滚可能浪费几天时间。这四层缺一不可。我见过很多团队只盯着 loss结果显存 OOM 了才意识到要做重计算也有团队只盯吞吐结果优化了一通loss 不收敛回头发现数据 pipeline 打乱顺序影响了训练正确性。1.3 建立基线和对比表评估体系不是拍脑袋定几个指标就完了必须建立基线。我在每个新模型上会先不做任何花哨优化用最简单的数据并行或单一并行策略跑 200 步记录下上述四层指标作为基线。之后每做一次改动都重跑同样的步数对比基线和优化后的差异。这里有个很重要的小技巧对比时尽量保持相同的 step 数不要比墙钟时间。因为两个训练任务如果吞吐不同跑到相同 wall time 时看到的 step 数不同loss 的状态也不一致对比没有意义。反过来固定 step 数能让你同时对比优化前后的每步耗时和 loss数据才干净。2. 昇思 MindSpore 评估体系实操基础2.1 用 MindInsight 做可视化监控昇思 MindSpore 自带 MindInsight这是我在训练评估里用得最多的工具。训练过程中可以通过 summary 算子把 loss、学习率、梯度范数、权重统计等标量写到日志目录然后启动 MindInsight 服务在浏览器里看曲线和计算图。我一般会在训练脚本里加这样的记录逻辑from mindspore.train.callback import SummaryCollector summary_collector SummaryCollector(summary_dir./summary_dir) ... model.train(epochs, train_dataset, callbacks[summary_collector])然后启动mindinsight start --port 8080 --summary-base-dir ./summary_dir浏览器打开后最常用的是标量面板。我会同时展示 loss、learning rate、grad_norm 三条曲线这样能快速看出学习率衰减和梯度变化是否匹配。如果 grad_norm 突然爆掉loss 大概率也会跟着飞。另外一个容易被忽略的能力是计算图可视化。MindInsight 可以展示算子的连接关系排查模型结构和数据流问题。我之前遇到过模型里某个 tensor 的 shape 没对齐在跑动态图时被广播机制悄悄修正了导致结果一直不对后来就是靠计算图查出来的。静态图模式下的异常在这种视图里会很显眼。2.2 自定义 Callback 记录性能指标MindInsight 已经覆盖了训练可视化和部分性能数据但我在大规模训练时还是会写自定义 Callback把吞吐、MFU、通信占比这些信息集中落到一个 CSV 或日志里方便后期脚本化分析。MindSpore 的 Callback 机制非常灵活可以在每个 step、epoch 或者训练阶段插入逻辑。我这里给一个简化思路不纠结具体 API 版本import time from mindspore.train.callback import Callback class PerfEvalCallback(Callback): def __init__(self, tokens_per_step, log_interval20): self.tokens_per_step tokens_per_step self.log_interval log_interval self.step_start_time None self.step_count 0 def step_begin(self, run_context): if run_context.original_args().cur_step_num % self.log_interval 0: self.step_start_time time.time() def step_end(self, run_context): if self.step_start_time and run_context.original_args().cur_step_num % self.log_interval 0: duration time.time() - self.step_start_time throughput self.tokens_per_step / duration print(step:, run_context.original_args().cur_step_num, throughput(tokens/s):, throughput)注意这里tokens_per_step要算准。如果你的 batch size 是 32每句序列长度是 4096 tokens那么每个 step 的 token 数就是 131072不管使用多卡还是单卡这个值应该按全局视角计算。多卡并行时每个 rank 看到的微 batch 可能只占一部分但你统计的吞吐通常是全局吞吐所以要把并行维度考虑进去。2.3 性能分析器的正确打开姿势MindSpore 提供了 Profiler 工具来采集算子耗时、通信耗时和 GPU/昇腾利用率。我一般不会让 Profiler 在整个训练过程中一直开那会增加额外开销影响真实性能。正确姿势是在几次小规模测试训练或正式训练刚开始的 100 步里打开 Profiler采集到时序数据后立即关掉保存 profile 文件再交给 MindInsight 查看 timeline。from mindspore.profiler import Profiler profiler Profiler(output_path./profile_data) # 启动后做一定步数的训练 profiler.end()Timeline 视图特别适合治“玄学性能问题”。你会看到每个 step 里哪些算子占了大部分时间哪些时间片浪费在通信等待上数据的 copy 和 transform 是不是卡了流程。注意第一次开启 profiler 后训练速度会明显变慢这是正常的。你只需要用 profile 出来的相对时间占比来定位瓶颈而不是作为性能基准。3. 大模型训练的性能优化路径3.1 并行策略选型从简单到混合评估体系建立以后就可以针对指标做优化了。大模型训练里最先要考虑的是并行策略。很多初学者一上来就想上最复杂的混合并行其实不一定是好事。我一般遵循“先数据并行再模型并行最后混合并行”的渐进路径。数据并行的实现成本最低只要确认模型能放进单卡显存把 batch 切到不同卡上每个 step 做梯度 AllReduce 即可。昇思 MindSpore 里可以通过设置并行模式来启用import mindspore as ms ms.set_auto_parallel_context(parallel_modems.ParallelMode.DATA_PARALLEL)但到了百亿参数级别单卡肯定放不下就需要张量并行和流水线并行。张量并行的思想是把一个算子内部的权重矩阵按行或列切开多卡协作完成同一个算子优点是通信只发生在单个算子内缺点是通信频率高。流水线并行则是按层切分把不同层的计算放在不同设备上数据像流水线一样一趟趟穿过各层优点是通信频率相对低缺点是会出现流水线气泡也就是某些卡在等待上游数据时闲下来。MindSpore 的自动并行和混合并行能力很强可以在set_auto_parallel_context(parallel_modems.ParallelMode.AUTO_PARALLEL)下让框架自动搜索切分策略。但在大模型场景下我建议先手动画出主卡通信拓扑再让框架自动填充单算子策略。全自动并行在模型结构复杂时搜索空间巨大可能算半天也得不到最优解。更实用的做法是理解模型结构里哪里是计算密集型算子、哪里是访存密集型算子再手动指定关键算子的切分方式。3.2 显存优化重计算、Offload与冗余状态显存是训练大模型最容易触顶的资源。这里要区分几种显存占用模型参数、优化器状态、中间激活值、梯度、通信临时缓冲。大模型训练里激活值会随 batch size、序列长度和层数线性增长经常成为压垮显存的最后一根稻草。最常见的优化手段是重计算Recompute/Gradient Checkpointing。思路很多读者应该不陌生在前向传播时只保留部分层的激活值反向传播需要用到某层激活时重新计算一次。这样可以大幅降低显存但代价是额外的计算开销。我在实践中的体验是重计算让单 step 时间增加 10%-25%但能把可用 batch size 扩大一倍以上整体吞吐还是划算的。MindSpore 里对 Cell 开启重计算很简单在定义的模块上调用对应设置即可。比较保险的用法是选择性地对部分层开启重计算而不是全模型无脑打开。一般优先重计算注意力层和 MLP 层中激活较大的算子。另一种手段是 Offload把一部分优化器状态或激活值搬到 CPU 内存。这适合 CPU 内存比 GPU/昇腾显存宽裕的服务器。但是要小心 CPU 与加速卡之间的 PCIe 传输会成为瓶颈。我做过一次实验直接把全部优化器状态放到 CPU吞吐掉了四成后来只把一阶动量放到 CPU吞吐损失控制在十个点内。这类取舍你不能通过阅读文档凭空判断必须在你自己的物理集群上测一轮。还有一个容易被忽略的隐性问题Adam 优化器每个参数要保存两份动量加上参数本身和梯度一个 float32 参数至少要占 16 字节。如果你的模型是 7B 参数单卡仅参数状态就要 112GB 以上显然放不下。所以大模型训练一定要用类似 ZeRO 的优化器状态分片方案把优化器状态切到不同 rank 上。MindSpore 生态里也支持这类方案虽然不是只有一条路但思路都是一样的让每张卡只保存它负责的那部分优化器状态需要时通过通信获取。3.3 算子融合与编译模式选择MindSpore 支持两种运行模式PyNative 模式和 Graph 模式。PyNative 模式适合调试Graph 模式适合跑性能。我在刚接触 MindSpore 时也纠结过为什么要区分两种模式本质原因是动态图灵活但逐算子执行有大开销静态图可以把整个模型编译成优化后的计算图再对算子做融合与内存复用。用静态图跑大模型训练时像Softmax、LayerNorm这种小算子可以通过融合减少 kernel 启动次数。昇思的图编译会在底层自动做一部分算子融合但你也可以在模型定义里尽量组织算子减少不必要的中间 tensor。我建议大模型训练从一开始就把set_context(modems.GRAPH_MODE)打开调试时再切到 PyNative。如果你用的是 Jupyter 或 VSCode 里的 MindSpore 内核静态图下报错位置有时不那么直观但只要你会看 traceback 的步数信息还是能定位到具体算子的。3.4 数据管线与通信优化数据管线是大模型训练里最容易被低估的性能瓶颈。很多团队把优化重点放在模型并行上结果一开 profiler 发现数据加载占了 30% 的 step 时间。MindSpore 提供了强大的 MindData 组件你需要关注几个参数num_parallel_workers、prefetch_size、shuffle是否合理、repeat和batch顺序等。我的习惯是让数据加载进程和计算进程充分并行。例如在训练启动时让多个 worker 同时做数据读盘、解码和预处理然后通过异步队列把数据预取到训练设备。监听指标很简单如果 GPU/昇腾计算单元利用率低同时数据队列频繁为空就是数据管线没喂够。如果计算单元已经跑满再加 worker 数收益也不大反而增加 CPU 和内存争抢。通信优化方面AllReduce 是所有数据并行训练里绕不开的通信模式。每个 step 都要把各卡的梯度聚合成一个全局梯度。流水线并行的通信也不能小看切分点之间会频繁传递张量。MindSpore 支持通信算子融合会把多个小的通信请求合并成一个大的通信请求减少通信次数。还可以调整通信算子的执行时机让它们尽可能和计算重叠。昇腾硬件上还有专门的通信加速能力设置好拓扑亲和性之后通信耗时可以明显下降。4. 实战案例从评估到优化的闭环4.1 用评估发现瓶颈这里我虚拟一个贴近真实情况的案例模型规模 7B采用 32 卡混合并行训练序列长度 4096每步训练 256k tokens。硬件环境假设是常见的高端加速卡集群。第一步是跑基线1000 步记录各项指标。基准数据显示平均每步耗时 1.2 秒吞吐 213k tokens/sMFU 只有 14%。看一眼资源层的数据显存峰值接近 85%不算危险但也不宽裕。再看 profiler 的 timeline发现每步中有约 40% 的空闲时间通信占比 28%算子计算只占了 32%。这个结论非常明显问题不在单一算子的计算效率而在于通信等待和并行切分不合理。如果你只看 loss 曲线会觉得一切正常但看到 MFU 14%就知道这台集群其实被浪费了八成以上的算力。这就是评估体系的价值。4.2 优化动作与效果对比针对上面发现的问题我做了三个调整。第一调整并行策略。原来 32 卡是 8 路数据并行乘 4 路张量并行但张量并行的卡间通信过于频繁且因为序列变长激活值很大导致部分卡在通信等待。我改成 4 路数据并行乘 8 路模型并行后通信步长更集中单算子内部切分更均匀通信占比从 28% 降到 17%。第二开启局部重计算。显存峰值从 85% 降到 62%余量被用来增大微 batch让每步能处理的 token 数略微提升。虽然重计算带来了一些计算开销但整体的 MFU 反而提升了。第三融合梯度通信。MindSpore 的通信算子融合把多个小梯度打包发送减少通信启动次数同时把梯度通信和下一层的前向计算尽量重叠起来。这一步之后每步平均耗时降到 0.85 秒吞吐提升到 301k tokens/sMFU 提高到 22%。虽然没有一步登天但训练总时间缩短了接近三分之一。你以为这就完了没有。我用同一套评估流程继续看数据发现 MFU 22% 依然不算高但瓶颈已经转移到数据加载和部分融合算子的 kernel 实现细节上。这种“优化一层、再评估一层、暴露下一层”的过程才是性能优化的真正节奏。4.3 评估报表怎么设计为了让优化动作可追溯我习惯在每次实验后生成一张简表记录改动项和关键指标。表格不一定复杂但信息必须完整。下面是一个简化的模板改动项每步耗时(ms)吞吐(k tokens/s)MFU显存峰值通信占比loss500步基线8DP×4TP120021314%85%28%2.214DP×8TP102025017%82%20%2.194DP×8TP重计算110023316%62%19%2.204DP×8TP重计算通信融合85030122%63%17%2.18这张表能让我们快速看出每个改动的独立贡献。注意我特意记录了 loss500步目的是验证性能优化没有恶化收敛性。如果改成重计算后 loss 发生了异常那就说明重计算逻辑有问题需要及时回退。5. 常见问题与排查技巧实录5.1 典型问题速查现象可能原因排查工具解决思路Loss 完全不动数据标签错乱、学习率过小、模型初始化问题Loss曲线、数据可视化先用小规模过拟合测试确认模型可学调大一倍学习率看趋势检查 shuffle 是否破坏标签对齐Loss 发散突然 NaN/Inf学习率过大、梯度飞涨、除零、混合精度溢出grad_norm、Loss曲线、MindInsight 参数直方图降低学习率开梯度裁剪检查 loss scaling 配置检查数据是否有异常值显存 OOMbatch过大、激活值积压、优化器状态过多profiler、显存监控调小 batch开启重计算减少冗余状态Offload 到 CPU 内存计算单元利用率低数据加载阻塞、通信等待、小算子过多timeline、数据队列监控提高num_parallel_workers和prefetch_size融合通信算子融合检查同步点多卡吞吐不随卡数线性提升通信开销过大、并行切分不平衡通信占比、单卡吞吐对比尝试通信融合调整张量并行与流水线并行比例检查数据并行 global batch 是否过大step 时间波动大集群共享干扰、IO 抖动、其他任务抢占多次采样、机器监控调整数据预取缓存使用独立存储避免训练节点上跑其他重负载任务5.2 从 Profiler 里读出的真实坑我在一次训练中发现了非常诡异的现象所有卡都在忙碌但卡与卡之间有不少空闲间隙。起初以为是通信拓扑不好换了网络拓扑也没改善。后来仔细看 timeline发现是某个模型并行切分把一个大算子切到了单卡上其他卡必须等它算完才能继续。这个“长尾”效应在并行训练里很常见解决方式是进一步拆分该算子或者把计算负载均衡地分配到多卡。另一个坑是数据随机性掩盖了优化效果。有些优化动作本质上不影响计算量但改变了随机种子或执行顺序导致 loss 提前下降。你要区分是真正的性能提升还是随机波动。我的做法是对每次优化在同一批随机种子下跑对照组。比如把优化开关设成一个 flag跑两次基线两次优化看均值而不是单次实验得出结论。5.3 小规模试跑的价值很多人喜欢全量跑一次长训练然后用长时间验证效果。我建议反过来先缩到 100 步到 200 步的短测试。配合 Profiler 和自定义回调你可以快速完成一轮评估优化循环。我通常会在真正大规模训练启动前花半天时间跑短测试虽然看起来耽误了时间但它能帮我避开一整天无效调参。曾经有次我调整了并行策略后连续两次实验都炸 OOM最后靠短测试发现是某个 PyNative 模式下的临时算子没有被静态图正确释放。这种问题如果等到全量训练跑起来才发现损失的时间和算力都不可估量。所以短测试不仅是性能评估的手段也是稳定性的保险。6. 在昇思上做性能优化的一点工具心得6.1 VSCode 与 MindSpore 内核的配合我日常开发环境会用 VSCode 远程连接训练服务器编译期的提示、跳转、断点调试非常趁手。如果你用的是 Jupyter也可以选择 MindSpore 作为内核在 notebook 里快速验证自定义算子。不过要提醒一点在静态图模式下VSCode 的断点不是每个张量都能随时查看因为图的执行是延迟或融合的。遇到这种情况最好的办法是切换到 PyNative 模式做逻辑调试调通后再切回 Graph 模式跑性能。还有一个细节MindSpore 版本更新较快接口偶有调整。不同版本的 Profiler 输出格式和 summary 目录结构可能不同。我建议把项目依赖的 MindSpore 版本固定下来并在代码里写明版本号。多人协作时也要保证大家用同一个版本的容器镜像避免“在我这能跑到你那就报错”的问题。6.2 内存管理思想与 Julia 的类比我知道昇思 MindSpore 在国内外的 AI 框架里显存管理策略和 Julia 那种“显式管理内存”的思想有几分相似。Julia 里你可以控制 GC 时机和数组内存布局MindSpore 里你也可以通过静态图和内存池机制减少频繁的显存申请释放。在 PyNative 模式下每个张量对象产生和销毁都带来额外开销在静态图模式下模型编译器会预先规划好中间 tensor 的复用显存波动小训练更稳。做性能优化时如果你能形成“显存是有限的资源池需要像工程内存一样精细化规划”的意识会比盲目抄优化配置有效得多。移动端或手游性能优化圈子里的很多经验也可以迁移过来比如“避免频繁创建临时对象”“把高频小操作合并成大操作”“延迟加载不必要资源”。大模型训练里的显存临时 buffer 管理、通信算子融合、数据预取本质都是一样的逻辑。优化的核心是减少等待和浪费而不是盲目堆硬件。7. 最后再说两个实用的“土办法”7.1 用自带日志画出实时看板有时我不太想把所有指标都推到完整可视化系统里训练集群的网络策略也未必允许开额外服务。这时候就用最笨的办法在训练脚本里定时把关键指标追加到一个 CSV 文件训练结束之后用 Python 一次性画图。我在实际项目中就是这么干的简单可靠无需额外部署。如果你需要实时查看也可以让脚本同时打印到 stdout再用tail -f跟踪。这个方法虽然土但在排查问题时效率很高。7.2 保存多次 checkpoint 和评估数据性能优化之外我在大模型训练里还养成了定期保存 checkpoint 的习惯。这个习惯看似跟性能无关但它能让你快速回退到某个 step 重新评估。比如你发现 10000 步之后 loss 开始异常发散如果没有 checkpoint只能从头再来如果有每 1000 步的 checkpoint就可以从 9000 步回去换参数重新训练省下巨额算力。这个经验往往在项目初期不明显等训练真正跑上几天后就会特别珍贵。我一向认为性能优化不是一锤子买卖而是一套“评估-优化-再评估”的循环。昇思 MindSpore 给了很多底层控制力比如静态图、自动并行、Profiler、MindInsight但工具只是前提真正决定你能不能把速度调上去的是你有没有一套清晰可复用的评估指标体系。我建议你下一次训练开始时别急着改并行配置先把评估代码写好跑一个相对稳定的基线。有了基线后面的每一步优化都变得有据可依。这就是我个人在大模型训练里吃过不少亏之后形成的习惯也是今天我愿意写这么多字分享出来的原因。
返回列表