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

来源: 投稿
2024-04-24 11:51:00

经过社区开发者们几个月的开发与贡献,现正式发布昇思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将集成自动拆分副本的逻辑,达成更易用的细粒度多副本并行。

参考链接:https://www.mindspore.cn/tutorials/experts/zh-CN/master/parallel/multiple_copy.html

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倍+。

参考链接:https://www.mindspore.cn/lite/docs/zh-CN/r2.3.0rc1/use/cloud_infer/runtime_distributed_python.html

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执行模式下可以有更好的执行性能。

参考链接:https://www.mindspore.cn/docs/zh-CN/r2.3.0rc1/api_python/mindspore/mindspore.JitConfig.html?highlight=jitconfig

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

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

参考链接:https://www.mindspore.cn/docs/zh-CN/r2.3/design/dynamic_graph_and_static_graph.html#%E5%8A%A8%E6%80%81%E5%9B%BE%E8%BD%AC%E9%9D%99%E6%80%81%E5%9B%BE%E6%8A%80%E6%9C%AF

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反演示例图(上图:传统反演;下图:我们的工作)

参考链接:https://gitee.com/mindspore/mindscience/tree/master/MindElec

展开阅读全文
点击引领话题📣 发布并加入讨论🔥
0 评论
2 收藏
分享
返回顶部
顶部