本文作者:张雅琦
来源:硬AI
近日,媒体报道称蚂蚁集团Ling团队成功在国产加速卡上训练出3000亿参数的MoE大语言模型,性能比肩英伟达芯片。意味着国产大模型能够寻找成本更低、效率更高的国产芯片或其他替代方案。
27日,Ling团队在知乎发文,称原计划在月底的小型技术沙龙上分享这些经验,但由于媒体提前关注,他们决定提前公开“抠 FLOPS 的一些点滴”。
在使用国产加速卡训练大模型时,Ling团队面临着训练正确性对齐和国产加速卡稳定性两大挑战。
团队表示,他们通过基础算子对齐、loss和grad突刺处理机制,实现了大规模训练长跑loss差异低于0.1%的目标,并验证了小尺寸模型训练结果预测大尺寸模型训练效果的可靠性。这为国产加速卡在大模型训练中的应用奠定了基础。
国产加速卡的适配挑战
Ling团队表示,在使用国产加速卡训练大模型过程中,研发团队遇到了训练正确性对齐的严峻挑战。
团队设定了极其严格的标准:基础算子完全对齐、分布式训练框架前后向计算完全对齐,以及大规模训练长跑loss差异低于0.1%。
"这换来了无数个通宵debug的难忘体验,"他们如此描述这一过程。团队将不同平台的基础算子进行了完全对齐实现,包括matmul、linear等关键算子。
与GPU相比,国产加速卡在稳定性方面确实存在较大差距。团队坦言:
"时常会遇到由于机器不稳定带来的loss以及grad异常,从而引发突刺,影响模型的收敛过程。"
为解决这一问题,团队设计了两种突刺处理机制:
对于loss突刺,他们会将历史最近的一部分loss作为参考,当当前loss明显高于历史均值时,会跳过该步训练或降低学习率。
对于更深层次的问题,团队开发了grad突刺处理机制,通过监控模型的梯度变化,在异常出现时及时干预。
"通过grad+loss突刺处理机制,可以自动处理大部分的loss异常。"
值得注意的是,团队在对齐过程中同时进行了scaling law的研究,发现通过设计合理的外推拟合方法,可以用一系列小尺寸模型的训练结果预测大尺寸模型的训练效果,误差低于0.5%。这一发现也验证了他们对跨平台训练loss差异控制在0.1%以内的要求是合理的。
Router TP bug修复与NormHead迁移
在框架层面,团队遇到的技术难题更加复杂。FSDP向MindSpeed(Megatron)对齐引入tensor parallelism特性导致了一系列模型收敛问题,特别是在MoE相关的router部分表现尤为严重。
团队详细解释了问题所在:
"在router的前向计算上,由于sp(sequence parallel)在Megatron中对router的输入进行了切分,导致其输入并不完整。"
为解决这一问题,团队为当时使用的MindSpeed 0.6.0版本开发了额外补丁,修复了router反向计算中的梯度问题。同时,他们在每次scatter操作后添加了对应的gradient_scale实现,保证梯度计算的正确性。
蚂蚁技术团队强调,Ling模型的发布只是他们工作的一个里程碑。他们从DeepSeek团队使用FP8训练大模型的经验中获得启发,同时也关注到兄弟团队基于强化学习的AReaL项目。
"每个AI研发工程师都相信AGI必将到来。我们相信AGI一定是普惠大众的。"
以下为Ling团队发布在知乎的全文:
关于我们抠 FLOPS 的一些点滴
本周开始看到有媒体关注我们团队的模型训练成果,其实月初我们就在 GitHub 和 Hugging Face 上发布了 Ling 模型权重和技术报告,名字就叫「EVERY FLOP COUNTS」,关于使用非 NVIDIA 加速卡集群训练 Ling 300B MoE 大模型的一些技术细节。我们的技术报告被外媒记者发现了,“出口转内销”地被关注到。其实我们本来就准备在月底的小型技术沙龙上分享经验教训的,既然被关注到了,就来提前说明一下吧。
从开源来,回社区去
即使如最近大热的 DeepSeek,也受限于算力问题进行了很多精彩的优化,对于我们一线研发人员来说,克服环境的限制就是工作。众所周知,和国外的大模型团队相比,中国团队面对了更多的异构加速卡的挑战,我们并不是第一家面对异构问题的公司,比如智源研究院就发起了 FlagScale 项目,研发面向异构加速卡的训练框架。有了开源社区,我们可以利用同行们的前期探索作为工作的基础。
同样,我们的实践成果也回馈给社区,希望可以帮助社区减少不必要的重复劳动。蚂蚁在去年开源 DLRover 项目,报告提到的轻量级选择性跟踪框架 XPUTimer 就集成在 DLRover 上,可以为不同算力平台上的大规模训练任务提供监控诊断功能。希望这些对社区的回馈,可以给大家带来一些启发。
一些收获和经验教训
在写这份技术报告时,我们希望分享 Ling 研发过程的一些关键 insight。Insight 可以是 novelty story,也可以是 bitter lesson。这里和大家聊聊我们得到的一些教训。作为较早吃螃蟹的人,分享这些教训并不是想吐槽,只是希望可以帮助其他同行避开一些问题,当然也希望可以促进国产加速卡的更快成熟。下面展开聊一聊几个我印象深刻的 bitter lesson。
训练正确性对齐
为了让大规模 MoE LLM 可以在多个算力平台上进行无缝切换训练,训练正确性对齐是必不可少又极其繁琐的一个过程。对齐有不同的标准,比如在不同平台训练都可以正常收敛是一个标准,而算子精度、训练框架、loss 完全对齐又是另外一个标准。“很傻很天真”的我们本着技术问题应该知其然又知其所以然的信念,定下了一个非常严格标准,基础算子(除符合预期的精度误差)完全对齐 + 分布式训练框架前后向计算完全对齐 + 大规模训练长跑 loss 差异低于 0.1%,当然这也换来了无数个通宵 debug 的难忘体验。
有趣的是,在做正确性对齐的过程中,我们同步也在做关于 scaling law 的研究。我们发现,通过设计一个合理的外推拟合方法,在不进行真实训练的情况下,一个尺寸较大(比如 20B、80B)的模型在正式训练较长时间(比如 2T token)后的 loss,可以被一系列 1B 以下的小尺寸模型的训练外推预测,其预测误差低于 0.5%。这样看来,跨平台训练的 loss 差异低于 0.1% 其实是一个合理的要求。
在算子对齐上,我们将不同平台的基础算子进行了完全对齐实现,比如 matmul、linear 等。
Router TP(Tensor Parallelism)bug 修复
在框架上,FSDP 向 MindSpeed(Megatron)对齐引入 tensor parallelism 特性会导致一系列模型收敛问题,尤其是在 MoE 相关的 router 部分非常严重。这里展开讲一下我们的工作。
在 router 的前向计算上,由于 sp(sequence parallel)在 Megatron 中对 router 的输入进行了切分,导致其输入并不完整,因此在 router 相关 loss 计算(包括 load_balance_loss 和 z_loss)时会额外使用 gather 操作将不同 sp rank 上的数据同步到一起,以进行完整 batch 计算。这个过程并没有专门针对反向进行对应的 reduce 实现,会导致回传梯度重复,需要手动对 router 相关的 loss 系数进行放缩。值得注意的是该 bug 已经在 Megatron 0.7.0 版本修复;当时 MindSpeed 支持到 0.6.0 版本,因此需要进行额外 patch 修复。
在 router 的反向计算上,Megatron 对 router 通过 gather 操作获取了完整的 logits,而 MindSpeed 在后续的 permute/unpermute 操作中需要强制使用 local logits,因此额外进行一次 scatter 操作来进行切分,出现了 loss 收敛性问题。经过排查,我们发现是 scatter_to_sequence_parallel_region在反向实现中进行了一次 _gather_along_first_dim操作导致梯度比正常梯度更大。最终我们在每一次 scatter 操作之后添加了对应的 gradient_scale 实现以保证梯度的正确性,从而满足 loss 收敛的需求。
NormHead 迁移
参考百川的训练经验,我们也采用了 NormHead 来保证训练的稳定(虽然初衷是为了保证训练稳定,但是后来通过 scaling law 分析,我们发现 NormHead 在 loss 上也会带来一些优势)。NormHead 从 FSDP 迁移到多 D 并行的 MindSpeed/Megatron 上也遇到了问题。
FSDP 上的参数在逻辑上是没有被切分的,因此 NormHead 的实现非常简单高效,通过 Torch 原生自带的 torch.nn.functional.normalize 即可完成对 lm_head.weight 标准化操作。在 MindSpeed/Megatron 中,由于涉及到了多 D 并行,因此需要修改 NormHead 的实现方法进行适配。最直接简单的方案就是结合 torch.nn.functional.normalize 的实际计算过程,将本地设备上的 lm_head.weight 先进行标准化计算,最后使用 reduce 对标准化后的 lm_head.weight 值进行同步。遗憾的是我们发现这样实现无法保证 loss 收敛,分析其原因主要是由于在不同机器上进行数据同步采用 Megatron.core.tensor_parallel.mappings._ReduceFromModelParallelRegion,而该方案没有在反向传播过程中实现对应的梯度同步,最终导致 loss 上升;于是我们重写了一版_ReduceFromModelParallelRegionForNormHead并实现了对应的反向以保证loss收敛。
另一方面,国产加速卡的某些算子可能不支持 BF16 计算,而 FP32 的算子计算效率远低于 BF16 算子,为了防止在多 D 并行中阻塞住模型的整体计算,需要对 NormHead 性能进行优化。我们设计了基于 all2all 通信的 NormHead 实现以及 HeadNormCache 等方案,以在国产加速卡上达到更优的计算效率。
训练稳定性
与 GPU 相比,国产加速卡在稳定性上确实存在不少问题,时常会遇到由于机器不稳定带来的 loss 以及 grad 异常,从而引发突刺,影响模型的收敛过程。为了缓解这些问题,我们设计了两种不同的突刺处理机制。
对于 loss 突刺,我们会把历史最近的一部分 loss 作为参考,如果当前 loss 与参考的历史 loss 均值相比有明显的上升,我们就会跳过这一步的训练直接开始下一步,或直接降低这一步的学习率来减少影响。这种方法在大多数情况下是有效的,可以很好地缓解训练不稳定问题。
但我们在实验观察中发现,loss 突刺处理机制并不能解决所有的训练不稳定问题,因为 loss 是模型训练过程的一个很宏观的表现,模型的状态在 loss 产生突刺之前可能已经出现了不稳定。Grad 会直接作用于模型参数,对其监控相比于 loss 更加迅速,因此我们也开发了 grad 突刺处理机制。参考 loss 突刺的实现,我们在自研的 ATorch 框架中对所有的 _ParamAndGradBuffer 进行处理,从而实现对模型 grad 的监控。如果 grad 出现异常就跳过这一步训练。通过 grad+loss 突刺处理机制,可以自动处理大部分的 loss 异常。
成本的计算
这次大家的一些误解也源于对成本计算的方式,其实我们在成本计算上使用了学术界比较通行的计算方法,这里也简单介绍一下。
根据在不同平台上对 Ling-Plus 的真实训练记录,我们可以观察到某个平台在 K 张加速卡上持续一段时间(比如一周)的 token 数,再根据技术报告表 1 上提到的不同加速卡的单位时间成本,就可以很简单地计算出对应平台上训练单位 token 量(报告里以 1 万亿 token 为单位)的成本。
表1:AI加速器特性与单位成本(估算)
事实上,不管是在 GPU 还是在国产加速卡上,LLM 的训练成本优化都是无止境的。Ling 的训练过程一定程度地说明,在我们做的这些技术努力上,国产加速卡上的训练成本与 GPU 相当甚至更低,同时可以保证 loss 收敛一模一样。
未来的工作
Ling 模型的发布只是我们工作的一个里程碑,后续我们还会进一步改进自己的工作。DeepSeek 为我们对训练经济性的提升带来了启发,DeepSeek 在训练中使用了 FP8 证明了这样的低精度浮点数是可以训练出来优秀的大模型的;同样我们兄弟团队基于强化学习的 AReaL也开源了,强化学习也是通往 AGI 之路的重要一环。我们后续的更多工作也会陆续开源在 inclusionAI org里。
每个 AI 研发工程师都相信 AGI 必将到来。我们相信 AGI 一定是普惠大众的,感谢大家的关心,期待未来的工作也能受到持续关注。
本文来自微信公众号“硬AI”,关注更多AI前沿资讯请移步这里