大模型 AI 框架昇思 MindSpore 2.3.RC1 发布 训练、推理性能大幅提升 JIT 编译强化

经过社区开发者们几个月的开发与贡献，现正式发布昇思 MindSpore2.3.RC1 版本，通过多维混合并行以及确定性 CKPT 来实现超大集群的高性能训练，支持大模型训推一体架构，大模型开发训练推理更简、更稳、更高效，并在训推一体框架的基础上通过多样的大模型推理优化技术，进一步降低大模型推理成本；通过使能 kernel by kernel 调度执行，进一步提升静态图调试调优能力；持续升级 MindSpore TransFormers 大模型套件和 MindSpore One 生成式套件，全流程开箱即用，一周即可完成大模型全流程的开发、验证；创新 AI + 科学计算（科学智能）范式，孵化科学领域基础大模型；下面就带大家详细了解下 2.3.RC1 版本的关键特性。

1、大模型训练：细粒度多副本并行，有效提升计算通信并发度，显著提升大模型训练性能

大模型训练下，为了降低显存开销，广泛的使用算子级并行技术，其中引入了大量的模型并行的通信，极大地影响了大模型的训练效率。模型并行的通信，从网络的结构上来看，其处于正反向计算过程中，阻塞正反向计算，无法与正反向的计算进行互相掩盖。为了解决模型并行通信的掩盖问题，MindSpore 提出了多副本并行技术。

在旧版本的 MindSpore 上，通过将网络从数据开始进行拆分，如下图所示，在单张卡内，通过 slice 算子将 Batch 维度进行拆分，进而产生多个分支，这多个分支的计算与通信互相之间没有依赖，存在并发的空间，通过执行序调度算法，控制多个分支的计算与通信进行并发。

随着网络规模的增大，受限于显存限制，当一张卡内的 BatchSize 仅支持为 1 时，上述对整网进行 Batch 拆分的方案不再可行。因此，考虑到模型并行通信的位置，昇思 MindSpore2.3.RC1 版本将 Transformer 模型中的 AttentionProjection 层以及 FFN 层进行拆分，产生多个分支，通过执行序调度算法控制细粒度的多分支的并行，其中拆分从 AttentionProjection 开始，到下一个 Layer 的 QKV 计算前结束。

上图描述了序列并行场景下的细粒度多副本拆分与掩盖基本思路，拆分为两个副本，在正向可以达成 50%+ 的通信掩盖；而在反向，结合计算梯度的分支的计算与 TP 通信的掩盖，可达成 90% 的通信的掩盖。当前细粒度多副本并行仅在 MindSpore Transformers 的 LLAMA 网络进行了实现，需要对模型结构进行手动改造为多个副本。后续版本昇思 MindSpore 将集成自动拆分副本的逻辑，达成更易用的细粒度多副本并行。

2、大模型推理全栈升级

大模型大规模商用之后，推理消耗的算力规模将十分庞大，相应地带来高昂的成本，商业闭环依赖推理规模突破。在降低大模型推理的成本的同时，要兼顾模型精度和计算时延，不能影响用户的体验。昇思 MindSpore 2.3.RC1 版本，从最上层推理服务到模型脚本优化到推理引擎 LLM Serving，为用户提供端到端的高效推理解决方案。

2.1 训推一体：大模型训 / 推统一脚本，大幅简化部署流程，提高效率

模型脚本默认使能了增量推理、FlashAttention/PagedAttention 等推理加速技术，避免了模型导出、切分、推理脚本开发等一系列工作，训练到推理加速平滑迁移，部署周期下降到天级。

2.2 极致性能：持续提升融合大算子、并行推理、模型小型化的关键能力

融合大算子：新增 10 + 业界最新的推理融合大算子接口，模型开发人员可以快速使能推理融合算子实现加速。

并行推理：训练推理并行策略接口一致，提供训练并行到推理并行 ckpt 重切分接口，支持动态 shape 模型切分。

模型压缩：昇思 MindSpore 金箍棒升级到 2.0 版本，提供了针对大模型的业界 SOTA 以及华为诺亚自研的量化、减枝等算法，实现千亿大模型 10 倍 + 压缩。

以上技术均可泛化的应用于 Transformer 结构的大模型中，经过验证，在盘古、Llama 2 的 8 卡模型推理中，首 token 时延做到百 ms 级，平均 token 时延小于 50ms，保持业界领先水平。

2.3 服务化高吞吐

通过连续批调度、Prefill/Decoding 混合部署等手段，尽可能的消除掉冗余计算，确保算力不闲置，实现大模型推理吞吐提升 2 倍 +。

3、静态图优化：支持 O (n) 多级编译，使能 kernel by kernel 调度执行，提升静态图调试调优能力

整图下沉执行性能最优，但大模型的规模和参数量发展得更为庞大，整图下沉执行方式在整图编译过程中耗时较长，一个千亿级别的大模型的编译时间为 30 分钟 - 60 分钟，调试调优效率低下。为解决上述问题，昇思 MindSpore2.3.RC1 版本中，提供了多级编译技术，O0 原生构图不优化、O1 增加自动算子融合优化、O2 整图下沉执行优化。在 O0 的编译选项下，通过原生图编译和 kernel by kernel（KBK）的执行技术，可以将编译时间提升到 15 分钟以内，同时我们在新版本中还开发了 DryRun 技术，用户可以直接在离线的情况进行内存瓶颈分析和并行策略调优，结合这两大技术可以使得大模型调试效率倍增。在 O0 这种编译条件下，我们使能了 SOMAS/LazyInline/ 控制流 Inline 来提升内存复用率，使能了多流并行 / 流水异步调度，可以提升执行性能；在 O1 这种编译条件下，通过使能算子融合技术，KBK 执行模式下可以有更好的执行性能。

4、JIT 兼具易用性和性能，动静统一，提供灵活高效开发

昇思 MindSpore 支持图模式（静态图）和 PyNative 模式（动态图）两种运行方法。动态图易于调试，开发灵活，易用性好；静态图语法支持有限，但执行性能好。JIT 兼顾性能和易用性，通过对 Python 字节码进行分析 & 调整、执行流进行图捕获 & 图优化，支持入图的 Python 代码做静态图方式执行，不支持的进行子图切分以动态图方式执行，自动地做到动静统一，实现方法如下图所示。

5、MindSpore Elec：新增大地电磁智能反演模型

MindSpore Elec 电磁仿真套件升级至 0.3 版本，联合清华大学李懋坤教授团队、华为先进计算与存储实验室共同打造了基于昇思 MindSpore 的大地电磁（Magnetotelluric，MT）智能反演模型。该模型通过变分自编码器（VAE）灵活嵌入了多物理先验知识，达到了业界 SOTA。该成果已被国际顶级勘探地球物理期刊《Geophysics》收录，同时也在昇思人工智能框架峰会 2024 上发布亮相。

（1）基础 MT 反演：反演区域水平长度为 10km，深度为 1km。下图 1 中目标电阻率分布（第一列）与传统大地电磁反演（第二列）、大地电磁智能反演（第三列），可以看出大地电磁智能反演相比传统反演精度显著提升（前者残差为 0.0056 和 0.0054；后者为 0.023 和 0.024 ）；下图 2 中，大地电磁智能反演性能也优于传统反演方法（前者收敛步数为 4 和 4；后者为 6 和 4）。

图 1 大地电磁反演精度对比

图 2 大地电磁反演收敛速度对比（Pixel-based：传统反演；Feature-based：我们的工作）

（2）南部非洲 MT 反演：大地电磁智能反演模型也在南部非洲开源数据集（SAMTEX）上做了验证。该反演区域位于南部非洲西海岸附近，长度约为 750km，深度选定为 80km。该测区显著特征为在水平方向 100km 至 400km 之间，深度 20km 以浅的区域存在的高导结构。由于低频电磁波在导体结构中的衰减，MT 方法对高导结构下部区域的敏感度很低，因此无先验知识约束的传统 MT 反演难以准确重建高导地层的下边界位置。大地电磁智能反演对高导地层的下边界重建较为清晰准确，较好地将地层厚度的先验知识融入了反演。

图 3 南部非洲 MT 反演示例图（上图：传统反演；下图：我们的工作）