螞蟻國產GPU訓練大模型細節曝光！Ling模型研發負責人發文詳解背后故事

來源：硬AI
作者：張雅琦

近日，媒體報道稱螞蟻集團Ling團隊成功在國產加速卡上訓練出3000億參數的MoE大語言模型，性能比肩$英偉達 (NVDA.US)$芯片。意味着國產大模型能夠尋找成本更低、效率更高的國產芯片或其他替代方案。

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 一定是普惠大眾的，感謝大家的關心，期待未來的工作也能受到持續關注。

編輯/Rocky