- 轉自知乎
作者 張志強 螞蟻Ling模型研發負責人
螞蟻開源大模型的低成本訓練細節,疑似曝光!
這段時間,螞蟻一篇技術論文引發關注。論文中顯示,他們推出的兩款MoE大模型,能夠在國產GPU上完成與英偉達同效的訓練。一時間,該消息在技術圈發酵,登上了熱搜,甚至還傳出「計算成本低于DeepSeek」一些傳聞。
現在,螞蟻Ling模型研發負責人張志強在知乎上作出了回應
他發布長文《關于我們摳 FLOPS 的一些點滴》,分享了他們一些大模型訓練的經驗和教訓。
包括訓練正確性對齊、Router TP(Tensor Parallelism)bug 修復、訓練穩定性等問題的解決。
最后還回應了外界對于他們成本計算的誤解,并表示不管是在 GPU 還是在國產加速卡上,LLM 的訓練成本優化都是無止境的。
Ling 的訓練過程一定程度地說明,在我們做的這些技術努力上,國產加速卡的訓練成本與 GPU 相當甚至更低,同時可以保證 Loss 收斂一模一樣。
在不改變原意的基礎上,量子位做了如下整理在此分享給大家,希望能給大家帶來一定的啟發。
(量子位已獲原作者授權)
關于我們摳 FLOPS 的一些點滴
本周開始看到有媒體關注我們團隊的模型訓練成果,其實月初我們就在 GitHub 和 Hugging Face 上發布了 Ling 模型權重和技術報告(https://arxiv.org/abs/2503.05139),名字就叫「EVERY FLOP COUNTS」,關于使用非 NVIDIA 加速卡集群訓練 Ling 300B MoE 大模型的一些技術細節。我們的技術報告被外媒記者發現了,“出口轉內銷”地被關注到。其實我們本來就準備在月底的小型技術沙龍上分享經驗教訓的,既然被關注到了,就來提前說明一下吧。
從開源來,回社區去
即使如最近大熱的 DeepSeek,也受限于算力問題進行了很多精彩的優化,對于我們一線研發人員來說,克服環境的限制就是工作。眾所周知,和國外的大模型團隊相比,中國團隊面對了更多的異構加速卡的挑戰,我們并不是第一家面對異構問題的公司,比如智源研究院就發起了 FlagScale 項目,研發面向異構加速卡的訓練框架。有了開源社區,我們可以利用同行們的前期探索作為工作的基礎。
同樣,我們的實踐成果也回饋給社區,希望可以幫助社區減少不必要的重復勞動。螞蟻在去年開源 DLRover 項目(https://github.com/intelligent-machine-learning/dlrover),報告提到的輕量級選擇性跟蹤框架 XPUTimer 就集成在 DLRover 上,可以為不同算力平臺上的大規模訓練任務提供監控診斷功能。希望這些對社區的回饋,可以給大家帶來一些啟發。
一些收獲和經驗教訓
在寫這份技術報告時,我們希望分享 Ling 研發過程的一些關鍵 insight。Insight 可以是 novelty story,也可以是 bitter lesson。這里和大家聊聊我們得到的一些教訓。作為較早吃螃蟹的人,分享這些教訓并不是想吐槽,只是希望可以幫助其他同行避開一些問題,當然也希望可以促進國產加速卡的更快成熟。下面展開聊一聊幾個我印象深刻的 bitter lesson。
訓練正確性對齊
為了讓大規模 MoE LLM 可以在多個算力平臺上進行無縫切換訓練,訓練正確性對齊是必不可少又極其繁瑣的一個過程。對齊有不同的標準,比如在不同平臺訓練都可以正常收斂是一個標準,而算子精度、訓練框架、loss 完全對齊又是另外一個標準?!昂苌岛芴煺妗钡奈覀儽局夹g問題應該知其然又知其所以然的信念,定下了一個非常嚴格標準,基礎算子(除符合預期的精度誤差)完全對齊 + 分布式訓練框架前后向計算完全對齊 + 大規模訓練長跑 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(https://github.com/inclusionAI/AReaL)也開源了,強化學習也是通往 AGI 之路的重要一環。我們后續的更多工作也會陸續開源在 inclusionAI org(https://huggingface.co/inclusionAI)里。
每個 AI 研發工程師都相信 AGI 必將到來。我們相信 AGI 一定是普惠大眾的,感謝大家的關心,期待未來的工作也能受到持續關注。
知乎鏈接:
https://zhuanlan.zhihu.com/p/1888526583813350974
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.