昇思 MindSpore 大模型自动优化实战:从配置到踩坑全记录
2026/9/19 6:18:02 网站建设 项目流程

昇思 MindSpore 这套名字刚出现的时候,很多人以为它只是一个要“国产替代”的深度学习框架,直到我自己动手在 MindSpore 上跑大模型微调和训练,才意识到它真正值钱的地方在于:把大模型训练和推理里那些又脏又累的优化工作,做成了尽量自动化的工具链。大模型自动优化这几个字听起来很虚,实际落地之后,能省下的是以周为单位的调参时间。这篇文章我想用一次完整的实战经历,把昇思 MindSpore 大模型自动优化技术怎么配置、怎么用、有哪些坑,一次讲清楚。适合正在做大模型训练、微调或推理加速,又不想把时间全耗在底层优化细节上的朋友。

1. 项目背景与大模型优化痛点

1.1 大模型训练和微调中最耗时间的环节

先说说大多数人在做 7B、13B 这类大模型时会遇到的问题。模型能跑起来只是一个起点,跑得快、跑得稳、显存不爆,才是真正磨人的地方。

我自己最开始从单卡小模型切到大模型时,最大的感觉是“处处都要管”。模型并行策略要自己设计,每张卡放多少层、哪些层需要重计算,都需要手动编排;混合精度怎么开,哪些算子保持 FP32,哪些可以降到 FP16 或 BF16,调不好就直接 Loss 变成 NaN;图编译阶段一旦算子不支持,排查报错的成本也很高;显存占用不是单纯看模型参数量,还要看激活值、梯度、优化器状态和通信缓冲区。

这些问题如果全部手动去做,每个环节都会有大量的试错。比较典型的时间消耗是:

  • 并行切分策略反复试,训练一两个 step 就 OOM,来回调整可能要一两天。
  • 混合精度开了之后精度崩了,但不确定是哪一层导致的,需要逐层排查。
  • 算子执行效率不高,同样一个模型,换个框架或者换份优化配置,性能可能差一倍。
  • 通信瓶颈被忽略,多卡训练时计算利用率不高,但 GPU 或 NPU 看起来已经跑满了。

这也是我做这个项目时最核心的诉求:能不能在尽可能少改业务代码的前提下,让框架本身帮我把编译、精度、并行、内存、算子调度这些事情自动优化掉。

1.2 为什么选昇思 MindSpore 来做自动优化

选型阶段我也对比过其他方案,最后选择昇思 MindSpore,主要是因为它在“自动优化”这条路径上做得比较完整,不是靠某一个单点工具去解决一个问题,而是从下到上形成了一套闭环。

核心的自动优化能力包括这么几个层次:

  • 图编译层:把 Python 代码编译成 MindIR 图,在图上做常量折叠、公共子表达式消除、算子融合等优化。
  • 并行策略层:支持数据并行、模型并行、流水并行和混合并行,并且可以基于策略搜索做自动切分。
  • 精度层:提供自动混合精度能力,根据算子类型和输入数据分布,自动决定哪些计算用 FP16、哪些保持 FP32。
  • 内存层:做内存复用、内存池管理和显存整理,避免训练过程中频繁分配和释放带来的碎片问题。
  • 算子层:提供 AOE 这类自动调优工具,针对目标硬件做算子级调优,并把结果回写到模型中。

我当时看到这套体系,其实想到的是一个类比:手动优化就像手动挡开车,省油但是费人;MindSpore 的自动优化更像是自动挡加自适应巡航,你不用管每一步换挡逻辑,只管好方向盘和油门,剩下的交给系统去匹配。

当然,自动优化不是完全没有代价。图模式编译时间会更长,自动并行搜索也有额外的计算开销,AOE 调优甚至会先跑很多组算子配置。但和手动折腾几天相比,这些前置成本是完全值得的。后面我会详细展开每一步怎么做。

2. 自动优化能力全景:先搞清楚有哪些“自动”

2.1 图模式编译与 MindIR 中间表示

MindSpore 有 PyNative 模式和 Graph 模式,大模型训练和推理我强烈建议走 Graph 模式。PyNative 模式适合调试,但执行时按 Python 算子逐行下发,优化空间很有限。

Graph 模式的核心在于先把模型转换成 MindIR。这个中间表示有点像给模型做了一次“静态体检”,框架能看到整个计算图的全局信息,也就能在图上做很多运行时看不出来的优化。

我在实际使用中,比较明显的收益来自图优化带来的算子融合。比如一个典型的大模型子图,可能先是一个矩阵乘,紧接着是加偏置、激活函数、Dropout,如果按 PyNative 模式一个算子一个算子地执行,每个算子都会触发一次内核启动和显存读写。但在 Graph 模式下,框架会把这些算子融合成更少的内核,减少中间结果的落地和读取。配合昇腾或 GPU 上的融合算子,性能提升非常大。

打开图模式其实只需要一行配置:

import mindspore as ms from mindspore import context context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=0)

不过要注意,图模式对动态 shape 的支持不如 PyNative 那么灵活。如果你的数据里有大量的动态维度,编译时可能会报错,或者出现比较严重的“图重建”开销。后面踩坑部分我会单独讲这个问题。

2.2 自动混合精度:不是简单把所有层都降成 FP16

混合精度是做大模型绕不开的话题。很多人刚接触时以为混合精度就是把 FP32 换成 FP16,其实这里面的讲究比想象中多。

FP16 能减少一半显存占用,同时利用硬件上的 Tensor Core 或者昇腾 AI Core 加速矩阵运算,但它的问题是动态范围小,容易出现溢出。BF16 动态范围和 FP32 更接近,但精度位数少,某些场景下收敛行为会不一样。

MindSpore 的自动混合精度方案会做两件事:一是根据算子的类型,决定哪些算子适合低精度执行,比如矩阵乘、卷积这类计算密集算子可以让它用 FP16;二是对那些精度敏感的算子,比如某些归一化层、损失计算,保持 FP32 或者提供keep_batchnorm_fp32这类选项去控制。

在训练脚本里,我一般是这样配置的:

from mindspore import Model, amp model = Model( network=net, loss_fn=loss, optimizer=optimizer, metrics={"acc"}, amp_level="O2", keep_batchnorm_fp32=True, )

amp_level的取值不一样,优化的侧重点也不一样。O0 表示全 FP32,O1 是白名单方式,O2 是黑名单方式。大模型场景我一般用 O2,让绝大多数算子走 FP16 或 BF16,再对少数敏感算子做保护。

推理和训练不太一样的地方在于,推理时还可以叠加量化,比如把权重从 FP16 进一步压缩到 INT8。MindSpore 也有相应的量化工具,但我个人体会是:训练阶段先把混合精度玩明白,推理再考虑量化,不要一上来就两者一起上,否则出了问题很难定位。

2.3 并行策略与自动并行

大模型单卡放下是前提,放不下就一定要并行。MindSpore 支持的并行粒度比较丰富,包括数据并行、模型并行、流水并行,以及组合起来的混合并行。

数据并行最简单,每张卡都放一份完整模型,只把数据切分到不同卡上,训练过程中做梯度同步。当模型大到单卡装不下时,就需要模型并行,把不同层或者同一层的参数切到多张卡。如果是几十B甚至上百B的模型,流水并行会更合适,把模型按层切成多个 stage,每个 stage 放在一张或一组卡上。

MindSpore 的自动并行能力,会尝试根据模型结构和硬件拓扑,自动搜索一个相对合理的切分策略,而不用我们手工去指定每一行matmul做哪种切分。

一行关键配置是:

context.set_auto_parallel_context( parallel_mode="semi_auto_parallel", device_num=8, gradients_mean=True, parameter_broadcast=True, )

我用semi_auto_parallel比较多。它的意思是:框架自动推导大部分算子的切分,但允许我们对关键算子手动指定策略。相比纯auto_parallel,这种方式可解释性更强,也更容易人工介入调优;相比纯数据并行,又能处理单卡放不下的模型。

真正训练时,数据加载部分也要配合并行策略,否则会出现每张卡读到相同数据的问题。MindSpore 的数据并行通常会由框架自动把数据集按 rank 切分,但如果你自己写自定义 Dataset,就需要确认num_shardsshard_id是否正确。

2.4 AOE 算子级自动调优工具

并行策略解决的是“怎么把计算分到多张卡上”的问题,而 AOE(Ascend Optimization Engine)解决的是“单个算子在这块硬件上怎么跑最快”的问题。

AOE 做的事情,简单说就是针对模型里的算子,在目标硬件上尝试不同的切分、调度、tiling 策略,然后选择性能最好的配置,保存下来。整个过程是自动化的,我们只需要把模型文件交给它。

我在训练脚本里会先用export导出 MindIR 文件,然后跑 AOE:

aoe --model=llama2_7b.mindir --framework=1 --job_type=1 --output_path=./aoe_output

不同版本参数名可能略有差异,执行前可以先aoe --help看一下。跑完 AOE 后,它会在输出目录生成优化后的配置,再通过 MindSpore 的接口把这些配置应用到下次图编译中。

AOE 看起来像是一个“锦上添花”的工具,但实际效果非常可观。尤其是在昇腾硬件上,某些关键算子手工写死的一种实现,往往不是最优的;AOE 会自动尝试多种 tiling 组合,经常会带来百分之二三十的端到端性能提升。唯一要提前做好心理准备的是,AOE 运行时间可能很长,特别是大模型,需要预留出足够的算力资源。

2.5 运行时内存复用与显存整理

最后这个“自动优化”很多人容易忽略,那就是内存管理。大模型的显存占用通常不是模型参数一个因素决定的,激活值、梯度、优化器状态、通信缓冲区,每一项都不可小觑。

MindSpore 在图模式下做内存复用,核心思路是分析张量的生命周期,对不同时间点使用的张量,尽量复用同一块物理内存。这有点像搬家的纸箱反复用,而不是每装一轮东西就买一个箱子。配合显存池化,也能减少频繁申请释放带来的碎片问题。

开启方式一般是在 context 里设置内存优化级别:

context.set_context(memory_optimize_level="O1", max_device_memory="31GB")

max_device_memory是用来限制模型最多使用的显存大小,给通信和框架预留一部分空间。这个值设置得太大会导致 OOM,设置得太小则浪费硬件资源,一般建议留出总显存的百分之十到二十作为余量。

再加上重计算技术,比如把部分算子的激活值不保存,反向传播时重新算一遍,能进一步压低峰值显存。虽然重计算会增加一点点计算量,但很多时候能让原本放不下的模型跑起来,属于关键时刻救命的优化手段。

3. 实战:从零搭建一个大模型自动优化训练脚本

3.1 环境准备与版本选择

先说版本。MindSpore 的版本迭代比较快,不同小版本之间的 API 和算子支持范围有差异,所以我建议你安装时尽量选择接近官方长期支持版的稳定版本,不要盲目追最新。

我这次使用的环境大概是:

  • Python 3.9
  • MindSpore 2.2 以上的稳定版
  • mindformers 配套版本
  • 昇腾 910B 或类似硬件,多卡环境

安装没有什么特殊之处,按官方指引用 pip 安装即可。需要注意的一点是,MindSpore 在不同硬件平台上对应不同的安装包,安装前一定确认好device_target是 Ascend、GPU 还是 CPU,装错了后面跑起来第一步就会报错。

mindformers 是大模型训练、微调和推理的高层套件,内置了很多常见模型结构,比如 LLaMA、Bloom、GLM 等。直接用底层 MindSpore API 去搭建一个 7B 模型很费劲,而且并行策略、优化器配置都已经有人帮你踩过坑了,直接用 mindformers 会更稳。

安装好之后,先用一个小的模型结构做端到端验证,不要一上来就 7B。我习惯先跑通一个几十M参数的配置,确认环境、数据、编译链路没问题,再切回真正的大模型。这一步能帮你把“环境问题”和“代码问题”分开,省掉很多无意义的查错时间。

3.2 从单卡跑到图模式:先打开编译优化

第一个实战阶段,我建议先把单卡跑通,再上多卡。单卡阶段的核心是把自动优化里和编译、精度相关的开关都打开。

基础配置如下:

import mindspore as ms from mindspore import context, set_seed from mindspore import Model, nn set_seed(42) context.set_context( mode=context.GRAPH_MODE, device_target="Ascend", device_id=0, ) # 开启内存优化级别 context.set_context(memory_optimize_level="O1")

然后定义模型、损失函数、优化器和回调:

network = create_model() loss = nn.CrossEntropyLoss() optimizer = nn.AdamWeightDecay(network.trainable_params(), learning_rate=1e-5) model = Model( network=network, loss_fn=loss, optimizer=optimizer, amp_level="O2", )

这里的关键点有两个。第一是GRAPH_MODE,没有特别强的调试需求就不要用 PyNative。第二是amp_level="O2",让自动混合精度接管精度策略。第一次编译时,图模式会花不少时间,这是正常的,MindSpore 会把编译之后的优化图缓存下来,后续重复执行会快很多。

如果训练过程中出现 Loss 不下降或者直接变成 NaN,优先检查是不是混合精度的问题,把amp_level降到 O1 或者 O0 对比一下。我后文会专门讲这个排查过程。

3.3 多卡并行:数据并行与半自动并行

单卡跑通后,如果模型尺寸还没超出单卡显存,最简单的扩展方式是数据并行。MindSpore 里配合mpirun或者rank_table做多卡启动,脚本里需要声明并行上下文。

以 8 卡数据并行为例:

context.set_auto_parallel_context( parallel_mode="data_parallel", device_num=8, gradients_mean=True, )

数据并行实现简单,但每张卡都要保存一份完整模型。一旦模型大到单卡放不下,就要切到半自动并行:

context.set_auto_parallel_context( parallel_mode="semi_auto_parallel", device_num=8, gradients_mean=True, parameter_broadcast=True, )

切换到半自动并行后,建议再做一件事:给模型里最关键、最耗时的算子手动指定切分策略。虽然框架可以自动推导,但某些场景下,人工指定能显著减少通信量。

比如一个大模型的 Linear 层,在数据并行下权重是每卡一份,但在模型并行下权重会被按列或按行切分。手动指定策略时,可以通过ops.MatMul().shard()等方式去配置。这部分的内容比较深,新手可以先不碰,先让框架自动推导跑通,再逐个人工介入。

多卡训练还需要特别注意数据集切分。特别是在semi_auto_parallel下,如果每张卡都读到了同样的数据,相当于全局 batch size 被缩小到了单卡 batch size,而且多卡之间的梯度同步会互相抵消部分收益。MindSpore 的数据集接口本身支持num_shardsshard_id,一定要检查好。

3.4 用 AOE 做算子级自动调优

当训练能稳定跑起来之后,就可以考虑算子级调优了。我比较推荐先把一个训练 epoch 或验证集上的性能数据记录下来,再跑 AOE,最后对比优化前后的吞吐。

大致流程分三步。

第一步,导出模型。这里我用验证模式下导出 MindIR:

from mindspore import export, Tensor input_ids = Tensor(shape=[1, 2048], dtype=ms.int32) export(network, input_ids, file_name="llama2_7b", file_format="MINDIR")

如果模型太大,或者导出时对输入 shape 有要求,需要按模型实际情况配一个合理的静态 shape。导出成功后会生成llama2_7b.mindir文件。

第二步,执行 AOE。命令大概长这样:

aoe --model=llama2_7b.mindir --framework=1 --job_type=1 --output_path=./aoe_output

framework=1表示输入是 MindIR 格式,job_type=1表示做性能调优。AOE 会逐个分析算子,在硬件上尝试若干候选实现并计时。模型越大,算子越多,这个过程越长。

第三步,把调优结果应用回训练脚本。AOE 输出目录里会有优化后的配置文件,在set_context里指定它,让图编译时优先使用 AOE 搜索到的最优实现:

context.set_context(optimize_kernel_cfg_path="./aoe_output")

这里要提醒一句,AOE 调优的目标硬件和训练硬件要一致,比如都是同一型号的昇腾芯片。换了一个芯片型号,之前调出来的算子配置不一定仍然最优。

3.5 训练过程中的性能监控与验证

很多人在大模型训练里只盯着 Loss,但“训练跑起来了”和“训练跑得好”完全是两回事。我在这次实战中同时使用了几种手段做性能监控,避免闷头训练了一天,最后发现效果很差。

第一是打开 Profiler。MindSpore 提供 Profiler 工具,可以采集算子的执行耗时、通信耗时、内存占用等信息。

from mindspore.profiler import Profiler profiler = Profiler(output_path="./profiler") # 训练若干 step profiler.stop()

跑完后再用 MindSpore Insight 或直接把 profiling 数据导出,分析热点算子在哪里,通信和计算有没有重叠。

第二是盯梯度范数和 Loss 曲线。如果 Loss 正常下降,但梯度范数出现异常跳变,多半是精度策略或学习率出了问题。

第三是非常朴素的计时。记录数据加载时间、单 step 时间、端到端吞吐,用这些数字去判断整体优化是否有进展。我一般会用一个步骤数除以总耗时,得到每秒样本数,作为一个简单但有效的性能指标。

优化不是一次性的动作,而是一个“基线 → 调整 → 对比 → 再调整”的闭环。建议每次只改一个变量,比如这次只改并行切分,下次只改内存优化级别,这样出了问题才找得到原因。

4. 踩坑实录与排查技巧

4.1 图编译失败,或者编译时间过长

图模式虽然性能好,但它对模型的“可编译性”要求更高。我遇到最多的报错是算子不支持、动态 shape 导致图重建,以及 Python 控制流没法完全转成静态图。

排查思路可以从这几个方向展开:

  • 先看报错日志里第一个 ERROR,后面的信息大多是连带错误,不用全部读完。
  • 检查模型里是否有动态 shape 的操作,比如不同 batch 的输入,或者依赖运行时结果的if分支。Graph 模式下这类写法容易触发问题。
  • 如果使用了自定义算子,先确认该算子是否有对应的高阶实现或 CPU 回退实现。
  • 编译时间过长时,检查是否关闭了缓存。图编译缓存第一次生成会比较慢,后面应该能明显提速。

有个小技巧是:先用一个 mini 版本的模型验证编译链路,再替换成完整模型。mini 版模型把隐藏层维度、层数都缩小,比如从 7B 缩到几十M,编译速度和问题定位效率都会高很多。

4.2 混合精度训练出现 NaN 或精度下降

混合精度是大模型训练里最容易翻车的点。我见过好几种 NaN 场景:有的是开 O2 后直接爆,有的是跑到几千步之后才偶发爆掉,还有的是梯度回传时溢出但 Loss 表面看起来正常。

排查时,我一般按顺序做这几件事:

  1. 先把amp_level降到 O1 或 O0,确认问题是否由混合精度引起。
  2. 检查学习率。混合精度下经常会对学习率敏感,调低学习率有时能直接避免溢出。
  3. 检查权重初始化。某些初始化方式产生的数值范围太大,FP16 下更容易溢出。
  4. 开启或检查动态 loss scaling。MindSpore 的amp模块一般会处理动态 scale,但要确认它确实生效。
  5. 在关键位置打印输入输出数值,比如第一个 Linear 层的输入、最后几层的梯度,锁定溢出的第一现场。

如果确定是某个算子导致的溢出,还可以用更精细的手段,给特定算子指定更高精度类型,或者把它排除在混合精度范围外。这属于比较进阶的用法,但遇到模型规模大、算子多的时候,比全局降精度有用得多。

4.3 显存不足 OOM,以及内存碎片问题

显存不足是做大模型的家常便饭。如果 OOM 发生在刚开始训练时,多半是模型本身超出了单卡或当前并行策略的承载能力。如果 OOM 发生在训练中间,可能和动态 shape 导致的内存峰值波动有关。

我处理 OOM 的优先级是这样的:

  • 先调小 batch size,最简单直接,但别只靠它。
  • 开启内存复用和重计算,能压缩峰值显存。
  • 用梯度累积模拟更大的 batch size,在不增加显存峰值的前提下,保持全局 batch size 不变。
  • 检查并行策略,确认权重、优化器状态、梯度是否真的被合理切分。
  • 把不需要的临时张量显式释放,尽量避免在 Python 层保留过多中间结果。

关于内存碎片,我印象最深的一次是训练前几十 step 没问题,越往后越卡,最后直接 OOM。后来分析发现是频繁分配和释放不规整大小的张量,导致显存碎片化。开启内存池和memory_optimize_level之后有改善,但要从根本上缓解,还是得让模型尽量使用静态 shape,减少运行时张量形状的剧烈变化。

4.4 自动并行后收敛效果变差

有一次我把模型从数据并行切到半自动并行后,Loss 下降变得很慢,心里很慌。一开始以为是精度策略变了,排查了一圈才发现问题出在数据加载上。

并行策略切换后,数据集的切分逻辑变了,但我的自定义 Dataset 里没有正确处理shard_id,导致部分卡读到的数据重叠严重,等效训练数据量变少。修正数据集切分后,收敛恢复正常。

另一个常见问题是多卡训练时随机种子没有按 rank 区分。如果每张卡初始化模型时用的种子都一样,虽然初始权重一致,但某些涉及随机丢弃或采样路径的操作会让各卡产生不同的随机状态,破坏并行一致性。建议给每张卡设置不同的全局种子偏移,并同步好数据 shuffle 的种子。

自动并行后的收敛问题,很多时候不是并行策略本身的问题,而是“并行引入的副作用”导致的。排查时不要只盯着模型和优化器,也要盯数据流和随机状态。

4.5 性能瓶颈定位与调优闭环

性能问题不像 OOM 那么显性,但它更值得花时间。我常用的定位方法是,先在 Profiler 里看几个关键指标:

  • 计算密集算子占总耗时比例。如果占比很低,说明模型没有把算力吃满。
  • 通信耗时占总耗时比例。多卡训练里通信经常是隐藏瓶颈。
  • 是否存在严重的 kernel 启动开销。算子太碎、太小时,启动开销会吃掉大量收益。
  • CPU 数据加载是否成为瓶颈。GPU 或 NPU 在等待数据时,利用率会掉得很明显。

针对不同瓶颈,解决办法也不同。算子太碎就靠图编译和算子融合去合并;通信瓶颈就重新考虑并行切分,减少跨设备通信量;数据加载慢就把数据预处理放到线程中,或者用mindrecord等高效格式。

调优闭环里,我最看重的一点是一次只改一个变量。因为大模型训练每个 step 的时间很敏感,哪怕你同时改了批次大小、并行策略和精度等级,性能提升了,你也不知道到底是哪个改动起了决定性作用。

5. 个人体会与扩展建议

5.1 自动优化不是黑盒,理解原理才敢放手

用了一段时间昇思 MindSpore 的大模型自动优化工具链后,我有一个很明显的感觉:工具越自动,越要理解工具背后的逻辑。因为你总会在某一天遇到自动策略不生效、性能不达标或者精度异常的情况,那时候如果只知道开关名称,不知道怎么排查,会非常被动。

我建议新上手的朋友,可以先按“全自动”模式跑通一条最小链路,也就是图模式 + 自动混合精度 + 半自动并行 + AOE,让框架把所有能优化的点都自动做一遍。跑通之后,再一个个工具去拆开研究。比如关掉 AOE 对比性能,改一下内存优化级别看显存变化,手动指定一个并行切分看看通信开销。这种“自动跑通 + 手动对比”的方式,比直接啃源码高效得多。

昇思 MindSpore 的自动优化能力,本质上是把工程师多年积累的经验固化到了框架里。你不需要从零发明轮子,但至少要会判断轮子当前的状态。

5.2 后续还能扩展的方向

训练和微调跑顺之后,我接下来想做的事情是把它进一步扩展到推理优化和多模态模型。推理阶段还有 KV Cache 管理、量化、动态 batch 等问题,和训练优化的关注点很不一样。多模态模型里除了文本,还有图像、音频编码器,并行和精度策略会更复杂。

另外,芬 SOP “自动优化”的思路其实可以迁移到业务里:不只是模型本身的算子或并行要优化,数据准备、checkpoint 保存、实验管理这些流程也值得自动化。把重复劳动减到最少,人才能真正把精力放到模型效果和改进上。

最后再分享一个小技巧:无论用多少自动优化工具,都一定要保留一个固定不变的标准配置作为基线。每次升级框架版本、换硬件或者改模型结构,先拿基线配置跑一遍,用真实数据判断新版到底有没有变快变稳。我在实际项目里靠这个习惯,避开了好几次“升级后性能倒退”的坑。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询