本文围绕蚂蚁集团Ling团队在国产加速卡上训练3000亿参数MoE大语言模型展开，详细阐述了团队在训练过程中遇到的挑战、采取的解决办法，以及后续的工作计划等内容。
近日，有媒体报道，蚂蚁集团Ling团队成功利用国产加速卡训练出具有3000亿参数的MoE大语言模型，而且其性能可与英伟达芯片相媲美。这一成果意义重大，意味着国产大模型有了新的发展方向，能够去探寻成本更低、效率更高的国产芯片或者其他替代方案。在27日，Ling团队在知乎发布文章。他们原本计划在月底的小型技术沙龙上分享训练的经验，但由于媒体提前关注，便决定提前公开“抠FLOPS的一些点滴”。### 国产加速卡训练大模型的挑战与应对在使用国产加速卡训练大模型时，Ling团队遭遇了两大挑战，分别是训练正确性对齐和国产加速卡稳定性问题。#### 训练正确性对齐挑战团队为自己设定了极为严格的标准：基础算子完全对齐、分布式训练框架前后向计算完全对齐，并且大规模训练长跑loss差异要低于0.1%。团队成员形容这一过程“换来了无数个通宵debug的难忘体验”。他们将不同平台的基础算子，如matmul、linear等关键算子进行了完全对齐实现。同时，团队在对齐过程中还进行了scaling law的研究。研究发现，通过设计合理的外推拟合方法，能够用一系列小尺寸模型的训练结果预测大尺寸模型的训练效果，误差低于0.5%。这一发现也验证了他们对跨平台训练loss差异控制在0.1%以内的要求是合理的。#### 国产加速卡稳定性挑战和GPU相比，国产加速卡在稳定性方面存在较大差距。团队表示“时常会遇到由于机器不稳定带来的loss以及grad异常，从而引发突刺，影响模型的收敛过程”。为了解决这一问题，团队设计了两种突刺处理机制。对于loss突刺，他们会把历史最近的一部分loss作为参考，当当前loss明显高于历史均值时，会跳过该步训练或降低学习率。对于更深层次的问题，团队开发了grad突刺处理机制，通过监控模型的梯度变化，在异常出现时及时干预。团队称“通过grad loss突刺处理机制，可以自动处理大部分的loss异常”。### 框架层面的技术难题及解决#### Router TP bug修复在框架层面，FSDP向MindSpeed（Megatron）对齐引入tensor parallelism特性导致了一系列模型收敛问题，特别是在MoE相关的router部分表现得尤为严重。团队详细解释道，在router的前向计算上，由于sp（sequence parallel）在Megatron中对router的输入进行了切分，导致其输入并不完整。为解决这一问题，团队为当时使用的MindSpeed 0.6.0版本开发了额外补丁，修复了router反向计算中的梯度问题。同时，他们在每次scatter操作后添加了对应的gradient_scale实现，保证梯度计算的正确性。#### NormHead迁移参考百川的训练经验，团队采用了NormHead来保证训练的稳定。但NormHead从FSDP迁移到多D并行的MindSpeed/Megatron上遇到了问题。经过分析和调试，团队重写了一版_ReduceFromModelParallelRegionForNormHead并实现了对应的反向以保证loss收敛。另一方面，为了防止在多D并行中阻塞住模型的整体计算，团队设计了基于all2all通信的NormHead实现以及HeadNormCache等方案，以在国产加速卡上达到更优的计算效率。### 后续展望蚂蚁技术团队强调，Ling模型的发布只是他们工作的一个里程碑。他们从DeepSeek团队使用FP8训练大模型的经验中获得启发，同时也关注到兄弟团队基于强化学习的AReaL项目。团队成员表示“每个AI研发工程师都相信AGI必将到来。我们相信AGI一定是普惠大众的”。后续他们还会进一步改进自己的工作，更多工作也会陆续开源在inclusionAI org里。成本计算方面，团队使用了学术界比较通行的计算方法。根据在不同平台上对Ling – Plus的真实训练记录，可以计算出对应平台上训练单位token量的成本。事实证明，在他们所做的技术努力下，国产加速卡上的训练成本与GPU相当甚至更低，同时可以保证loss收敛一模一样。
表1：AI加速器特性与单位成本（估算）
本文介绍了蚂蚁集团Ling团队利用国产加速卡训练3000亿参数MoE大语言模型的成果，深入分析了训练过程中遇到的训练正确性对齐、加速卡稳定性、框架层面技术难题等挑战及相应的解决办法，还提及了成本计算方式和后续工作计划。这一成果不仅展示了国产加速卡在大模型训练中的潜力，也为国产大模型的发展提供了新的思路和方向。
原创文章，作者：小耀，如若转载，请注明出处：https://www.xiaoyaoxin.com/archives/7017.html
