Date: 2025-06-27
在如今的人工智能浪潮中,大规模语言模型(上百亿乃至千亿参数)正迅速改变着我们的工作和生活。然而,训练这些庞大的模型往往面临“算力不足、显存不够用、通信太慢”等诸多挑战。为了让大模型的训练过程更顺畅、更高效,沐曦MXMACA软件平台(简称 MXMACA)具有无缝兼容CUDA的能力,科学兼容Megatron-LM[1]的绝大多数特性。此外,MXMACA进行多方面的优化,帮助科研人员和工程师能够快速在沐曦硬件环境中完成各类前沿模型的训练。下面,我们将从几个关键角度介绍MXMACA在大模型训练方面的改进思路和优化效果,让更多的读者轻松了解“大模型训练背后的那些事”。
通常,大模型采用「张量并行(Tensor Parallel, TP)+ 流水线并行(Pipeline Parallel, PP)+ 数据并行(Data Parallel, DP)+ 序列并行(Sequence Parallel, SP + 专家并行(Expert Parallel,EP)+ 上下文并行(Context Parallel, CP))」的多维并行策略,让成百上千张 GPU 同时参与训练。然而,随着模型参数量飙升(DeepSeek V3为6710亿参数),单靠原版 Megatron-LM 往往会遇到以下几个难题:
MoE 模型训练会出现「热门专家」被过度调用而导致计算和显存极不均匀,拖慢训练速度且易导致显存溢出。此外,跨节点的 AlltoAll 通信占据了较多的训练时间。
在分布式训练中,过细的模型切分虽然可以提升计算并行度,但往往会大幅增加跨节点通信开销。特别是在计算与通信需要共享硬件资源的原生并行架构中,计算操作和通信操作会相互竞争有限的带宽和计算单元,这种资源争用问题常常导致实际并行效率低于理论预期。
大模型需要存储很多“中间计算结果”(比如激活值、梯度、优化器状态)和大量参数,当模型规模上升时,很容易出现“显存不够用”的状况,导致训练中断,进而影响效率。
当你用成百上千块 GPU 训练一个模型时,如何把每一种并行方式合理组合,才能既不爆显存又能让计算满载?靠人工一遍遍尝试,不但耗时,还容易错过更优的组合。集群训练如何减少因故障导致的中断和资源浪费,如何快速定位慢节点,都是集群训练常遇的挑战。
大模型训练常受限于某些关键算子的低效实现,这些显存访问密集型算子是拉低模型MFU的一个重要因素。
为了解决这些痛点,我们结合沐曦曦云C系列GPU的硬件特点,做了多方面的“落地优化”。既保留了框架的灵活性,也在常见疑难场景中提供了“一键开关”式的配置。下面我们将从 MoE 优化、计算通信并行、显存优化、自动调优与集群训练、算子融合等几个重点模块,逐一展开。
Mixture-of-Experts(MoE)是日益流行的混合专家模型,通过路由让tokens选择相应的专家计算,能够显著提升模型容量与表达能力。然而,同时也带来了专家之间负载不均、显存爆炸等挑战。MXMACA 针对 MoE 提供了多种优化策略,帮助你在显存和吞吐之间找到更好的平衡。
在MoE 模型训练初期,某些专家会被大量 token 路由(“热门”),而其他专家几乎闲置。这导致频繁且不均匀的 AlltoAll 通信:热门专家所在的显卡要不断从多台显卡拉数据,通信开销巨大。
想象一个在线商店,某款商品突然在一两个城市爆火,下单量激增。如果所有订单都要从远在总部的仓库发货,就会出现“配送中心爆仓、快递车来回奔波”导致迟迟配送不到顾客手中。优化方式就好比在每个城市中心先存放几箱这款热销商品,顾客下单之后直接从本地仓库发货,大大缩短配送路径。等到热度消退、全国范围内需求趋于均衡时,再把多余的本地库存退回到总部或取消本地备货。一句话:把“最畅销的那几件”临时放到客户附近供他们随时取,就能避免每次都从很远的仓库拉货。
在 MoE 前向(Forward)/反向(Backward)时,Batch 内的某些 token 会被路由到热门专家(Expert),导致该专家对应的 GPU 需要处理大量激活,占用显存陡增,容易 OOM(Out-Of-Memory)。尤其是在训练早期,token 分配波动较大,很难预先调好“重计算参数”。
动态侦测:在每个训练 Step 之前,先统计当前各个“专家并行 Rank”所分配到的 token 总数量;
阈值触发:若某个 Rank 分配到的 token 数量超过预设阈值,则自动开启重计算逻辑;否则保持常规计算;
智能开关:对不同的moe dispatcher采用不同的重计算方式。
当你和同事们分摊搬东西时,如果A同事拿了特大箱子,其他人手都空着。这时,你会让A暂时把一些东西放地上(重新计算),等到他搬完一部分再回来挑起;等到大家分工均匀了,就恢复正常搬运。这样既能让大家都忙起来,也避免了某个人因超负荷工作而累倒。
DeepSeek提出的DualPipe[2] 方案需要在流水线并行(Pipeline Parallel)中为模型参数保留两份拷贝,这对显存要求极高,且在 Bubble 较大的场景下并行效率有限。
DualPipeV将模型在 PP 维度拆分为前半段(PP0 - PPN/2-1)与后半段(PPN/2 - PPN-1):
前半段按照PP顺序(PP0 - PPN/2-1),看做一个完整模型布置到所有节点上;
后半段按照PP逆序(PPN-1 - PPN/2),看做一个完整模型布置到所有节点上。
两组之间交替发送激活与梯度,充分减少空闲等待时间。这样,只需在每张显卡上保留一份参数拷贝,同时保持较高的流水线并行度。
就像工厂生产线:如果把生产过程切成两半,整个流水线是一个V字型,每组工人在处理前半个流水线的一道工序的同时,负责后半个流水线的一道工序。当其中一道工序需要等待时,可以处理另一道工序,甚至可以双管齐下,两侧工序同时进行。
除了前面提到的“某些专家突然超载”情况,整个专家(MoE)网络里还有很多“子环节”、各种小运算(例如:激活函数、向量重排、共享专家算子等),在显存吃紧的时候,也需要“分级”来处理。
a. 轻量级重计算:只对激活函数、向量重排、router路由这些小环节做重计算;
b. 中度重计算:在上面基础上,选择专家内部的某些全连接层和共享专家(Shared Expert)做重计算。
c. 全量重计算:MoE模型部分全量做重计算。
想象训练MoE模型像在沙漠探险。 显存是珍贵的骆驼运力(负重能力)。轻量: 只背少量必需品(省运力,稍慢)。中度: 选择性背大件(平衡运力与速度)。全量: 所有装备现用现造(运力最省,行动最慢)。分层选择,用最小速度代价换最大运力空间。
在显存紧张的思路下,多级内存优化相较于不优化时能节省 12% 左右的显存峰值,而整体训练速度仅损失3% 左右,为中小集群训练带来显著价值。
在 MoE模型 中,不同专家收到的输入token数量往往不同,这导致每个专家要做的矩阵乘法(GEMM)大小不一。GPU 在处理大小不一的矩阵运算时,可以采用groupgemm提升算力利用率,但相对于均匀计算,效率还是有所降低。
想象你有好几批货,大小差异很大,不利于装入标准箱进行一次性搬运。这时让较大的货物,拿去一部分,较小的货物,添加一部分,就可以一次性把好几个标准箱同时装车,搬运效率更高。
在大规模并行训练中,“算”与“传”往往会发生冲突:当 GPU 在做大矩阵计算时,却要停下来做 AllReduce/AlltoAll等通信,结果就是一边算,一边等。或者在已有的“算”与“传”并行场景中,两者发生硬件资源竞争,导致性能相互影响。MXMACA 主要通过 SDMA、通算融合算子等手段,尽量让“算”与“传”不再相互干扰。
在计算通信并行场景中,由于GPU核心既承担计算任务,又承担通信任务(如AllGather、ReduceScatter),容易导致资源竞争,使得通信与计算互相拖慢。
节点间使用CPU和网卡来实现通信传输。
通信最大程度减少对GPU的使用,可以有效减轻互相抢资源的情况。
SDMA通信引擎的实现,就好像生产车间里出现了“自动小推车”,一台机器算完半成品后,直接把它放到小推车上,小推车负责自动把零件送到下一个工作台;原来那台机器不用为送货而分摊精力。
在 TP(张量并行)切分场景下,计算与通信的依赖关系难以打通:
1. GEMM+ReduceScatter/AllGather 融合:
将 GEMM 计算与通信算子写入同一个 CUDA Kernel 中,直接将 GEMM 结果远写到其他 GPU,省去了显存读写与 kernel 启动开销。同时实现了通信和计算细粒度切分,使细粒度间的计算和通信任务不存在依赖关系,从而并行执行。
2. 无依赖算子 SDMA 传输:
对于无依赖关系的通信算子(如BWD中部分AllGather或者ReduceScatter),使用SDMA完成,从而避免与Compute算子争夺内存带宽和算力资源。
好比生产线上的半成品不再“先放到货架,再由叉车搬去下一个工序”;而是在同一个环节里边加工边传送,让“传送”像流水一样跟着“加工”一起走,省掉了中途的反复搬运。显卡就像流水线上的工人,既动手加工又顺手交接,效率显著提升。
在原生Megatron-LM的MoE 中,单层 Transformer 里前向会有两个 AlltoAll,反向也有两个AlltoAll。这些通信操作往往与专家(Expert)计算串行执行,导致并行度严重不足。
通过将 MoE 层划分为多个子单元,实现 AlltoAll 通信与专家计算的高度并行:
将其中两个 AlltoAll 与 Shared Expert 的前向和反向计算并行;
另外的AlltoAll与D/W分离后的专家计算并行。
理论上可达 75% 的全 Overlap 率,相比原生Overlap水平大幅提升。
MoE Comm Overlap,相当于原始MoE计算和通信都在一条路上,现在增加了一条路,通过计算通信分解,让AlltoAll通信单独走一条路,大大减少来回等待。
在 DeepSeek V3中,MoE Comm Overlap 使得AlltoAll通信与计算并行度提升约 3 倍:
训练时的显存就像钱包里的空间,装不下就会“爆卡”。 MXMACA提供了一系列显存优化策略,从 Granular Activation Offload到Granular Recompute,多管齐下帮你“花最少的钱,装最多的东西”,让有限的显存能撑起更大规模的训练任务。
在流水线并行中,不同阶段(Pipeline Stage)需要存储的“中间激活数据”数量并不一样。有些阶段需要保留很多激活,有些阶段只需要少量。若直接把所有激活都卸载到主机内存,势必增加大量数据传输,很难与计算相互掩盖,拖慢训练。
就像搬家时,你把最重的家具先搬到小车上存放,但把沙发、床这些需要马上用的常驻在家里。等到后面空间还紧张,再逐个决定把哪几个“没那么急”的物件先运出去。这样既不占用车的所有空间,也避免了一次性搬空再慢慢拉回来的低效。
重计算在显存紧张时非常有效,但如果把所有计算都重算一遍(全量重计算),会让整体训练速度大幅下降。很多时候,仅把“轻量”的那部分(例如归一化层或激活函数)重算,就能腾出不少显存,又影响不大。
Norm 重计算:只把归一化层(LayerNorm)相关的中间结果释放显存,反向时再重算。
激活函数重计算:只把激活函数(如 GELU、Swiglu 等)的中间结果释放显存,反向时再重算。
不均匀细粒度:对不同的PP stage,因为显存压力不同,使用的重计算方法和重计算力度也可以不同。
我们可以根据实际显存压力和性能需求,把“Norm 重算”和“激活重算”与传统的“全量重算”灵活组合。例如:在某些阶段只做 Norm 重计算,其他阶段保持全量;或者只做激活重计算……总之,以最小代价解决显存不足问题。
该比喻和细粒度激活卸载类似。
当训练规模从数十张GPU扩展到成百上千张GPU 时,手动在多维并行维度上逐个尝试,几乎是不可能在有限时间里搞定的工作。MXMACA 通过“Auto Search”引擎和“DLRover”[4]两大工具,实现了自动化调优与容错加速,让你更专注于算法设计与数据准备,而非配置参数。
在 Megatron-LM 里,你可能同时考虑张量并行(Tensor Parallel)、流水线并行(Pipeline Parallel)、数据并行(Data Parallel)、专家并行(MoE Parallel)等维度。不同组合下,显存占用和性能差别巨大,要人工一一尝试,既浪费时间,也容易错过更优解。
MXMACA 引入一套基于算子、显存与通信三大模块的自动调优(Auto Search)引擎:
1. 构建性能模型:
2. 全局搜索与预测:
就像你要组织一次大规模搬家,有几百个箱子、几十辆卡车,各卡车载重不同、路况也不同。传统做法是“卡车 A 多装点、卡车 B 少装点、卡车C 跑好点……”人工来回摸索。Auto Search 就是提前用模型测算好哪几种装载方案最经济,给你一个“前五优选”,你只需要挑个最方便的兑现即可。
大型分布式训练任务常因节点故障或网络抖动导致中断,因为故障导致的集群空闲和回滚训练,都会导致集群资源的浪费。另一方面,在训练过程中,往往需要向存储介质一次性写入数百上千 GB 数据,耗时数分钟甚至十几分钟,影响迭代效率。
借助DLRover Flash Checkpoint 机制,将训练状态(包括模型权重、优化器状态、学习率调度状态等)先写入CPU,再异步持久化到分布式文件系统。主要优势有:
1. 异步写入:
2. 故障恢复:
就像你在写文档时,Word 会自动把内容先存在缓存里,然后后台再慢慢写到硬盘;如果电脑突然关机,缓存里最后的内容会被紧急落盘,下次打开就能直接恢复到缓存时的状态。
在千卡级别集群上,DLRover Flash Checkpoint 将大小1T左右的 Checkpoint 写盘时间从 10 分钟缩减到 10秒以内;节省85%因为集群故障导致的训练回滚和空闲时间。
在大集群训练中,某台机器网络带宽突然变差、某块显卡温度过高降频、或者其他硬件异常等等,都会让那台节点训练速度变慢。可一旦出现“1 台慢”,整个训练“队伍”就会被拖慢,因为大家需要等待。
就像车队比赛时,摄影机会记录每辆车的圈速、过弯情况、进站时间等。一旦发现某辆车某圈速度比别人慢,就立即发出提示,帮助车队找出“哪里出现了瓶颈”(比如轮胎不行、油压不稳、驾驶员操作问题等),及时修正,保持整体队速。
除了上面提到的几个核心方向,MXMACA 还在算子融合、并行调度等细节方面做了许多打磨,让整体训练更顺滑。下面简要介绍两项常见的补充优化。
模型里有很多很常见但“零碎”(Memory bound)的操作,比如 “加偏置再激活再 Dropout 再相加” 这一连串动作,如果每一步都拆成单独算子去执行,就会大量占用显存带宽和启动内核的开销。
想象你去餐厅吃套餐,原本你要点“炸鸡”“薯条”“饮料”“沙拉”四样,如果每次都是厨师分散地一个一个做,出餐就会慢很多;而“套餐”把它们组合成一个连贯流程,一次性煎炸加热、打包好,效率便会高不少。
在DeepSeek V3 模型训练中,启用 Flash Fusion 后可带来 5%以上的性能提升,且降低了显存占用。
在 Pipeline Parallel(PP)中,1F1B的模型调度方式,容易产生“泡沫”(Bubble),即GPU闲置等待的时间,影响资源利用率。这种现象随着PP的增大,或者Global Batch Size(GBS)的减小,愈加严重。
借鉴 Zero-Bubble Pipeline Parallelism[3]的思想,MXMACA集成了ZBH1(Zero Bubble H1)方案:
通俗比喻
就像一条生产线,原本上下游有时会错拍,有时会轮到机器没料可做。Zero Bubble 就是优化排产计划,让前后几道工序更均匀衔接,减少待料时间。不需要给第一道工序额外加机器(显存),却能让整体产量更高。
优化效果
在启用VPP时,若模型总层数与 PP Stage 数量不再为整除关系(如质数层),常会出现无法均匀拆分为每 VPP Stage 保持相同层数的瓶颈。例如,61 层模型切为PP = 8 时,每个 Stage 无法平均分配层数;又比如 15 层模型切为PP = 3 时,每个 Stage 均为5层,质数层无法进一步采用VPP切分。
MXMACA 提供“空层插入”(Empty Layer) 功能:
实测在 PP=8、VPP=2 场景下,经过空层补齐的 61→64 层模型,与直接使用不均匀 PP 相比,训练速度提升 6% 以上。
通过一系列“计算通信并行”、“专家模型优化”、“显存优化”、“自动调优与集群训练工具”等手段,MXMACA 成功让 Megatron-LM 在沐曦硬件环境中实现了如下优势:
希望在有限硬件条件下训练上百亿大模型
想快速配置集群并行,不想一遍遍试命令行参数
想让训练过程有更强的容错、断点续训能力
想站在“技术巨人”肩膀上,用最少的工程成本跑出最大价值
那么不妨试试 MXMACA 提供的这些优化能力(https://developer.metax-tech.com/softnova/docker)。未来,我们也会持续迭代、不断打磨各种新功能,助力更多团队、更多应用场景,让“大模型训练”真正实现“大众化”、变得“人人可跑、人人可用”。
[1] https://github.com/NVIDIA/Megatron-LM
[2] DeepSeek-AI. (2025). DeepSeek-V3 Technical Report. arXiv:2412.19437.
Retrieved from https://arxiv.org/abs/2412.19437
[3] https://github.com/sail-sg/zero-bubble-pipeline-parallelism
[4] https://github.com/intelligent-machine-learning/dlrover