今天小編(嬴覓晴)要和大家分享的是螞蟻國產GPU訓練大模型細節曝光!Ling模型研發負責人回應:關于我們摳FLOPS的一些點滴,歡迎閱讀~
螞蟻開源大模型的低成本訓練細節,疑似曝光!
這段時間,螞蟻一篇技術論文引發關注。論文中顯示,他們推出的兩款 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 完全對齊又是另外一個标準。" 很傻很天真 " 的我們本着技術問題應該知其然又知其所以然的信念,定下了一個非常嚴格标準,基礎算子(除符合預期的精度誤差)完全對齊 + 分布式訓練框架前後向計算完全對齊 + 大規模訓練長跑 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
一鍵三連「點贊」「轉發」「小心心」
歡迎在評論區留下你的想法!
— 完 —
速搶席位!中國 AIGC 產業峰會觀眾報名通道已開啟 ♀️
首批嘉賓曝光啦 百度、無問芯穹、數勢科技、生數科技、像素綻放等十數位 AI 領網域創變者将齊聚峰會,讓更多人用上 AI、用好 AI,與 AI 一同加速成長~
4 月 16 日,就在北京,一起來深度求索 AI 怎麼用
一鍵星标
科技前沿進展每日見
關于螞蟻國產GPU訓練大模型細節曝光!Ling模型研發負責人回應:關于我們摳FLOPS的一些點滴就分享完了,您有什麼想法可以聯系小編(嬴覓晴)。