“国产加速卡低成本训练取得突破”刷屏！蚂蚁负责人发文详解：关于我们抠 FLOPS 的一些点滴

本文作者：张雅琦

来源：硬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 一定是普惠大众的，感谢大家的关心，期待未来的工作也能受到持续关注。