蚂蚁开源大模型的低成本训练细节,疑似曝光!
这段时间,蚂蚁一篇技术论文引发关注。论文中显示,他们推出的两款 MoE 大模型,能够在国产 GPU 上完成与英伟达同效的训练。一时间,该消息在技术圈发酵,登上了热搜,甚至还传出「计算成本低于 DeepSeek」一些传闻。
现在,蚂蚁 Ling 模型研发负责人张志强在知乎上作出了回应。
他发布长文**《关于我们抠 FLOPS 的一些点滴》**,分享了他们一些大模型训练的经验和教训。
包括训练正确性对齐、Router TP(Tensor Parallelism)bug 修复、训练稳定性等问题的解决。
最后还回应了外界对于他们成本计算的误解,并表示不管是在 GPU 还是在国产加速卡上,LLM 的训练成本优化都是无止境的。
Ling 的训练过程一定程度地说明,在我们做的这些技术努力上,国产加速卡的训练成本与 GPU 相当甚至更低,同时可以保证 Loss 收敛一模一样。
在不改变原意的基础上,量子位做了如下整理在此分享给大家,希望能给大家带来一定的启发。
(量子位已获原作者授权)
关于我们抠 FLOPS 的一些点滴
本周开始看到有媒体关注我们团队的模型训练成果,其实月初我们就在 GitHub 和 Hugging Face 上发布了 Ling 模型权重和技术报告(arxiv.org/abs/2503.05... FLOP COUNTS」**,关于使用非 NVIDIA 加速卡集群训练 Ling 300B MoE 大模型的一些技术细节。我们的技术报告被外媒记者发现了,"出口转内销" 地被关注到。其实我们本来就准备在月底的小型技术沙龙上分享经验教训的,既然被关注到了,就来提前说明一下吧。
从开源来,回社区去
即使如最近大热的 DeepSeek,也受限于算力问题进行了很多精彩的优化,对于我们一线研发人员来说,克服环境的限制就是工作。众所周知,和国外的大模型团队相比,中国团队面对了更多的异构加速卡的挑战,我们并不是第一家面对异构问题的公司,比如智源研究院就发起了 FlagScale 项目,研发面向异构加速卡的训练框架。有了开源社区,我们可以利用同行们的前期探索作为工作的基础。
同样,我们的实践成果也回馈给社区,希望可以帮助社区减少不必要的重复劳动。蚂蚁在去年开源 DLRover 项目(github.com/intelligent... 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(github.com/inclusionAI... AGI 之路的重要一环。我们后续的更多工作也会陆续开源在 inclusionAI org(huggingface.co/inclusionAI...
每个 AI 研发工程师都相信 AGI 必将到来。我们相信 AGI 一定是普惠大众的,感谢大家的关心,期待未来的工作也能受到持续关注。
知乎链接:
zhuanlan.zhihu.com/p/188852658...
欢迎在评论区留下你的想法!
--- 完 ---