引言

MoE大模型的参数和所需要的AI集群规模持续摸高，模型训练的内存占用大和通信时延高等问题成为大参数模型训练的核心瓶颈，MindSpeed基于昇腾硬件特点，实现了Swap内存交换和DualPipeV流水排布特性，让内存在CPU和NPU之间高效交换，All2All通信与无依赖关系的计算相互掩盖，支撑更大规模的模型训练对内存和通信掩盖的需求。

MoE模型面临内存占用和通信耗时的挑战

在大模型蓬勃发展的当下，MoE凭借其模型效果成为了主要模型结构，稀疏性MoE模型总参数量的内存占用的稠密模型的4~8倍，激活值的内存需求量增加30%以上，模型训练容易出现OOM内存不足问题；同时MoE模型每层前向和反向计算有4次All2All通信，模型专家数量较多时需做EP（专家并行）切分，EP切分涉及节点间数据传输通信，大规模集群中All2All的通信耗时占单次训练迭代时间高达50%，且通信耗时受限于节点间带宽限制。当模型的训练参数量飙升至千亿级别乃至更高时，即使配备了顶尖的AI 集群，内存和通信的瓶颈效应仍然存在。简单地扩大集群规模，不仅会导致成本的显著增加，还会引入更多的通信开销，从而导致端到端的性能收益大打折扣。

MindSpeed训练加速特性Swap、DualPipeV，充分挖掘内存、通信性能

MindSpeed支持Swap内存交换和DualPipeV流水排布，在传统节点上能实现内存节省20%+，通信掩盖比例60%+的性能效果。昇腾超节点打破以CPU为中心的冯诺依曼架构，创新提出了对等计算架构，用高速总线互联替代传统以太，通信带宽提升了15倍。Swap交换涉及NPU与CPU内存拷贝通信，基于超节点高带宽的硬件特点，Swap内存交换在保证内存优化的基础上，训练迭代更快、吞吐性能更高；同时DualPipeV流水排布在超节点的低时延通信下，流水掩盖可以高达85%+，支撑训练更大参数的模型。

支持Swap内存交换，实现内存优化20%+，提升训练性能。

在模型训练阶段，MoE训练过程的张量数据（如模型权重、优化器状态、激活值等）内存占用相比于稠密模型增加30%以上。MindSpeed开发了Swap特性，对于NPU上内存非常大的张量，通过将NPU暂时不使用的内存高效卸载到 CPU内存，在重新使用时再从CPU高效加载到NPU，在有限的硬件资源下训练更大规模的模型。结合基于昇腾驱动的虚拟内存实现的创新优化器内存Swap方法，可以进一步降低性能开销。DeepSeek V3模型上实测，整网内存可节省20%+，在超节点高带宽的加持下，NPU和CPU间的内存交换速度更快，保证内存交换效果的同时训练迭代时间更短、吞吐性能更高。

支持DualPipeV流水排布，实现通信掩盖比例高达85%+，支撑训练更大参数模型

MoE类模型一直以来存在All2All通信耗时长这一关键问题，为了更好地进行All2All通信掩盖，MindSpeed基于昇腾硬件特点实现DualPipeV[1]流水排布方案，并尽可能多的在warmup和cooldown阶段实现P2P通信掩盖。在Deepseek V3模型上的DualPipeV特性效果测试中，传统节点能实现60%+的通信掩盖比例。借助超节点的高通信带宽大大提升NPU卡间通信效率，DualPipeV的流水掩盖比例可高达85%以上。

传统节点的DualPipeV掩盖效果

超节点的DualPipeV掩盖效果

MindSpeed Swap内存交换和DualPipeV流水排布的特性使用指导

安装依赖

详见链接的“安装依赖”章节：https://gitee.com/ascend/MindSpeed/tree/2.0.0_core_r0.8.0

安装MindSpeed

下载源码安装：

git clone -b 2.0.0_core_r0.8.0 https://gitee.com/ascend/MindSpeed.git
pip install -e MindSpeed

如需使用Ascend Transformer Boost（ATB）加速库算子，请先安装 CANN-NNAL 并初始化添加环境，例如：

# CANN-NNAL默认安装路径为：/usr/local/Ascend/nnal
# 运行CANN-NNAL默认安装路径下atb文件夹中的环境配置脚本set_env.sh
source /usr/local/Ascend/nnal/atb/set_env.sh

获取 Megatron-LM 并指定分支

见链接的“安装”章节的第3步：https://gitee.com/ascend/MindSpeed/tree/2.0.0_core_r0.8.0

快速上手使用MindSpeed加速特性Swap和DualPipeV

仅仅一行代码就可以轻松使能 MindSpeed 的各项功能。以 GPT 模型为例：在 Megatron-LM 目录下修改pretrain_gpt.py文件，在import torch下新增一行：import mindspeed.megatron_adaptor，即如下修改：

import os
import torch
import mindspeed.megatron_adaptor
from functools import partial
from typing import Union

在 Megatron-LM 目录下，准备好训练数据和训练脚本，训练脚本格式见MindSpeed/tests_extend/system_tests/pretrain_base.sh，在训练脚本中同时添加参数--schedules-method dualpipev、--moe-fb-overlap、--swap-optimizer并执行

…
GPT_ARGS="
    --moe-fb-overlap \                        <-- 1f1b掩盖
    --schedules-method dualpipev \            <-- dualpipe流水
    --swap-optimizer \                        <-- 开启 swap optimizer 特性
    --tensor-model-parallel-size ${TP} \
    --pipeline-model-parallel-size ${PP} \
    --expert-model-parallel-size ${EP} \
    --sequence-parallel \
    --context-parallel-size ${CP} \
    --context-parallel-algo  ${CP_TYPE} \
--num-layers ${NUM_LAYERS} \0
…

5      结语

本期为大家介绍了MindSpeed训练加速特性中的Swap内存交换和DualPipeV流水排布，在训练性能加速上，MindSpeed将持续探索演进，欢迎访问MindSpeed开源社区：https://gitee.com/ascend/MindSpeed/

参考链接：

[1]https://github.com/deepseek-ai/DualPipe