首頁 > AI資訊 > 最新資訊 > 螞蟻國(guó)產(chǎn)GPU訓(xùn)練大模型細(xì)節(jié)曝光!Ling模型研發(fā)負(fù)責(zé)人回應(yīng):關(guān)于我們摳FLOPS的一些點(diǎn)滴

螞蟻國(guó)產(chǎn)GPU訓(xùn)練大模型細(xì)節(jié)曝光!Ling模型研發(fā)負(fù)責(zé)人回應(yīng):關(guān)于我們摳FLOPS的一些點(diǎn)滴

新火種    2025-05-22

作者 張志強(qiáng) 螞蟻Ling模型研發(fā)負(fù)責(zé)人

螞蟻開源大模型的低成本訓(xùn)練細(xì)節(jié),疑似曝光!

這段時(shí)間,螞蟻一篇技術(shù)論文引發(fā)關(guān)注。論文中顯示,他們推出的兩款MoE大模型,能夠在國(guó)產(chǎn)GPU上完成與英偉達(dá)同效的訓(xùn)練。一時(shí)間,該消息在技術(shù)圈發(fā)酵,登上了熱搜,甚至還傳出「計(jì)算成本低于DeepSeek」一些傳聞。

現(xiàn)在,螞蟻Ling模型研發(fā)負(fù)責(zé)人張志強(qiáng)在知乎上作出了回應(yīng)。

他發(fā)布長(zhǎng)文《關(guān)于我們摳 FLOPS 的一些點(diǎn)滴》,分享了他們一些大模型訓(xùn)練的經(jīng)驗(yàn)和教訓(xùn)。

包括訓(xùn)練正確性對(duì)齊、Router TP(Tensor Parallelism)bug 修復(fù)、訓(xùn)練穩(wěn)定性等問題的解決。

最后還回應(yīng)了外界對(duì)于他們成本計(jì)算的誤解,并表示不管是在 GPU 還是在國(guó)產(chǎn)加速卡上,LLM 的訓(xùn)練成本優(yōu)化都是無止境的。

Ling 的訓(xùn)練過程一定程度地說明,在我們做的這些技術(shù)努力上,國(guó)產(chǎn)加速卡的訓(xùn)練成本與 GPU 相當(dāng)甚至更低,同時(shí)可以保證 Loss 收斂一模一樣。

在不改變?cè)獾幕A(chǔ)上,量子位做了如下整理在此分享給大家,希望能給大家?guī)硪欢ǖ膯l(fā)。

(量子位已獲原作者授權(quán))

關(guān)于我們摳 FLOPS 的一些點(diǎn)滴

本周開始看到有媒體關(guān)注我們團(tuán)隊(duì)的模型訓(xùn)練成果,其實(shí)月初我們就在 GitHub 和 Hugging Face 上發(fā)布了 Ling 模型權(quán)重和技術(shù)報(bào)告(https://arxiv.org/abs/2503.05139),名字就叫「EVERY FLOP COUNTS」,關(guān)于使用非 NVIDIA 加速卡集群訓(xùn)練 Ling 300B MoE 大模型的一些技術(shù)細(xì)節(jié)。我們的技術(shù)報(bào)告被外媒記者發(fā)現(xiàn)了,“出口轉(zhuǎn)內(nèi)銷”地被關(guān)注到。其實(shí)我們本來就準(zhǔn)備在月底的小型技術(shù)沙龍上分享經(jīng)驗(yàn)教訓(xùn)的,既然被關(guān)注到了,就來提前說明一下吧。

從開源來,回社區(qū)去

即使如最近大熱的 DeepSeek,也受限于算力問題進(jìn)行了很多精彩的優(yōu)化,對(duì)于我們一線研發(fā)人員來說,克服環(huán)境的限制就是工作。眾所周知,和國(guó)外的大模型團(tuán)隊(duì)相比,中國(guó)團(tuán)隊(duì)面對(duì)了更多的異構(gòu)加速卡的挑戰(zhàn),我們并不是第一家面對(duì)異構(gòu)問題的公司,比如智源研究院就發(fā)起了 FlagScale 項(xiàng)目,研發(fā)面向異構(gòu)加速卡的訓(xùn)練框架。有了開源社區(qū),我們可以利用同行們的前期探索作為工作的基礎(chǔ)。

同樣,我們的實(shí)踐成果也回饋給社區(qū),希望可以幫助社區(qū)減少不必要的重復(fù)勞動(dòng)。螞蟻在去年開源 DLRover 項(xiàng)目(https://github.com/intelligent-machine-learning/dlrover ),報(bào)告提到的輕量級(jí)選擇性跟蹤框架 XPUTimer 就集成在 DLRover 上,可以為不同算力平臺(tái)上的大規(guī)模訓(xùn)練任務(wù)提供監(jiān)控診斷功能。希望這些對(duì)社區(qū)的回饋,可以給大家?guī)硪恍﹩l(fā)。

一些收獲和經(jīng)驗(yàn)教訓(xùn)

在寫這份技術(shù)報(bào)告時(shí),我們希望分享 Ling 研發(fā)過程的一些關(guān)鍵 insight。Insight 可以是 novelty story,也可以是 bitter lesson。這里和大家聊聊我們得到的一些教訓(xùn)。作為較早吃螃蟹的人,分享這些教訓(xùn)并不是想吐槽,只是希望可以幫助其他同行避開一些問題,當(dāng)然也希望可以促進(jìn)國(guó)產(chǎn)加速卡的更快成熟。下面展開聊一聊幾個(gè)我印象深刻的 bitter lesson。

訓(xùn)練正確性對(duì)齊

為了讓大規(guī)模 MoE LLM 可以在多個(gè)算力平臺(tái)上進(jìn)行無縫切換訓(xùn)練,訓(xùn)練正確性對(duì)齊是必不可少又極其繁瑣的一個(gè)過程。對(duì)齊有不同的標(biāo)準(zhǔn),比如在不同平臺(tái)訓(xùn)練都可以正常收斂是一個(gè)標(biāo)準(zhǔn),而算子精度、訓(xùn)練框架、loss 完全對(duì)齊又是另外一個(gè)標(biāo)準(zhǔn)。“很傻很天真”的我們本著技術(shù)問題應(yīng)該知其然又知其所以然的信念,定下了一個(gè)非常嚴(yán)格標(biāo)準(zhǔn),基礎(chǔ)算子(除符合預(yù)期的精度誤差)完全對(duì)齊 + 分布式訓(xùn)練框架前后向計(jì)算完全對(duì)齊 + 大規(guī)模訓(xùn)練長(zhǎng)跑 loss 差異低于 0.1%,當(dāng)然這也換來了無數(shù)個(gè)通宵 debug 的難忘體驗(yàn)。

有趣的是,在做正確性對(duì)齊的過程中,我們同步也在做關(guān)于 scaling law 的研究。我們發(fā)現(xiàn),通過設(shè)計(jì)一個(gè)合理的外推擬合方法,在不進(jìn)行真實(shí)訓(xùn)練的情況下,一個(gè)尺寸較大(比如 20B、80B)的模型在正式訓(xùn)練較長(zhǎng)時(shí)間(比如 2T token)后的 loss,可以被一系列 1B 以下的小尺寸模型的訓(xùn)練外推預(yù)測(cè),其預(yù)測(cè)誤差低于 0.5%。這樣看來,跨平臺(tái)訓(xùn)練的 loss 差異低于 0.1% 其實(shí)是一個(gè)合理的要求。

在算子對(duì)齊上,我們將不同平臺(tái)的基礎(chǔ)算子進(jìn)行了完全對(duì)齊實(shí)現(xiàn),比如 matmul、linear 等。

Router TP(Tensor Parallelism)bug 修復(fù)

在框架上,F(xiàn)SDP 向 MindSpeed(Megatron)對(duì)齊引入 tensor parallelism 特性會(huì)導(dǎo)致一系列模型收斂問題,尤其是在 MoE 相關(guān)的 router 部分非常嚴(yán)重。這里展開講一下我們的工作。

在 router 的前向計(jì)算上,由于 sp(sequence parallel)在 Megatron 中對(duì) router 的輸入進(jìn)行了切分,導(dǎo)致其輸入并不完整,因此在 router 相關(guān) loss 計(jì)算(包括 load_balance_loss 和 z_loss)時(shí)會(huì)額外使用 gather 操作將不同 sp rank 上的數(shù)據(jù)同步到一起,以進(jìn)行完整 batch 計(jì)算。這個(gè)過程并沒有專門針對(duì)反向進(jìn)行對(duì)應(yīng)的 reduce 實(shí)現(xiàn),會(huì)導(dǎo)致回傳梯度重復(fù),需要手動(dòng)對(duì) router 相關(guān)的 loss 系數(shù)進(jìn)行放縮。值得注意的是該 bug 已經(jīng)在 Megatron 0.7.0 版本修復(fù);當(dāng)時(shí) MindSpeed 支持到 0.6.0 版本,因此需要進(jìn)行額外 patch 修復(fù)。

在 router 的反向計(jì)算上,Megatron 對(duì) router 通過 gather 操作獲取了完整的 logits,而 MindSpeed 在后續(xù)的 permute/unpermute 操作中需要強(qiáng)制使用 local logits,因此額外進(jìn)行一次 scatter 操作來進(jìn)行切分,出現(xiàn)了 loss 不斂性問題。經(jīng)過排查,我們發(fā)現(xiàn)是 scatter_to_sequence_parallel_region在反向?qū)崿F(xiàn)中進(jìn)行了一次 _gather_along_first_dim操作導(dǎo)致梯度比正常梯度更大。最終我們?cè)诿恳淮?scatter 操作之后添加了對(duì)應(yīng)的 gradient_scale 實(shí)現(xiàn)以保證梯度的正確性,從而滿足 loss 收斂的需求。

NormHead 遷移

參考百川的訓(xùn)練經(jīng)驗(yàn),我們也采用了 NormHead 來保證訓(xùn)練的穩(wěn)定(雖然初衷是為了保證訓(xùn)練穩(wěn)定,但是后來通過 scaling law 分析,我們發(fā)現(xiàn) NormHead 在 loss 上也會(huì)帶來一些優(yōu)勢(shì))。NormHead 從 FSDP 遷移到多 D 并行的 MindSpeed/Megatron 上也遇到了問題。

FSDP 上的參數(shù)在邏輯上是沒有被切分的,因此 NormHead 的實(shí)現(xiàn)非常簡(jiǎn)單高效,通過 Torch 原生自帶的 torch.nn.functional.normalize 即可完成對(duì) lm_head.weight 標(biāo)準(zhǔn)化操作。在 MindSpeed/Megatron 中,由于涉及到了多 D 并行,因此需要修改 NormHead 的實(shí)現(xiàn)方法進(jìn)行適配。最直接簡(jiǎn)單的方案就是結(jié)合 torch.nn.functional.normalize 的實(shí)際計(jì)算過程,將本地設(shè)備上的 lm_head.weight 先進(jìn)行標(biāo)準(zhǔn)化計(jì)算,最后使用 reduce 對(duì)標(biāo)準(zhǔn)化后的 lm_head.weight 值進(jìn)行同步。遺憾的是我們發(fā)現(xiàn)這樣實(shí)現(xiàn)無法保證 loss 收斂,分析其原因主要是由于在不同機(jī)器上進(jìn)行數(shù)據(jù)同步采用 Megatron.core.tensor_parallel.mappings._ReduceFromModelParallelRegion,而該方案沒有在反向傳播過程中實(shí)現(xiàn)對(duì)應(yīng)的梯度同步,最終導(dǎo)致 loss 上升;于是我們重寫了一版_ReduceFromModelParallelRegionForNormHead并實(shí)現(xiàn)了對(duì)應(yīng)的反向以保證loss收斂。

另一方面,國(guó)產(chǎn)加速卡的某些算子可能不支持 BF16 計(jì)算,而 FP32 的算子計(jì)算效率遠(yuǎn)低于 BF16 算子,為了防止在多 D 并行中阻塞住模型的整體計(jì)算,需要對(duì) NormHead 性能進(jìn)行優(yōu)化。我們?cè)O(shè)計(jì)了基于 all2all 通信的 NormHead 實(shí)現(xiàn)以及 HeadNormCache 等方案,以在國(guó)產(chǎn)加速卡上達(dá)到更優(yōu)的計(jì)算效率。

訓(xùn)練穩(wěn)定性

與 GPU 相比,國(guó)產(chǎn)加速卡在穩(wěn)定性上確實(shí)存在不少問題,時(shí)常會(huì)遇到由于機(jī)器不穩(wěn)定帶來的 loss 以及 grad 異常,從而引發(fā)尖刺,影響模型的收斂過程。為了緩解這些問題,我們?cè)O(shè)計(jì)了兩種不同的尖刺處理機(jī)制。

對(duì)于 loss 尖刺,我們會(huì)把歷史最近的一部分 loss 作為參考,如果當(dāng)前 loss 與參考的歷史 loss 均值相比有明顯的上升,我們就會(huì)跳過這一步的訓(xùn)練直接開始下一步,或直接降低這一步的學(xué)習(xí)率來減少影響。這種方法在大多數(shù)情況下是有效的,可以很好地緩解訓(xùn)練不穩(wěn)定問題。

但我們?cè)趯?shí)驗(yàn)觀察中發(fā)現(xiàn),loss 尖刺處理機(jī)制并不能解決所有的訓(xùn)練不穩(wěn)定問題,因?yàn)?loss 是模型訓(xùn)練過程的一個(gè)很宏觀的表現(xiàn),模型的狀態(tài)在 loss 產(chǎn)生尖刺之前可能已經(jīng)出現(xiàn)了不穩(wěn)定。Grad 會(huì)直接作用于模型參數(shù),對(duì)其監(jiān)控相比于 loss 更加迅速,因此我們也開發(fā)了 grad 尖刺處理機(jī)制。參考 loss 尖刺的實(shí)現(xiàn),我們?cè)谧匝械?ATorch 框架中對(duì)所有的 _ParamAndGradBuffer 進(jìn)行處理,從而實(shí)現(xiàn)對(duì)模型 grad 的監(jiān)控。如果 grad 出現(xiàn)異常就跳過這一步訓(xùn)練。通過 grad+loss 尖刺處理機(jī)制,可以自動(dòng)處理大部分的 loss 異常。

成本的計(jì)算

這次大家的一些誤解也源于對(duì)成本計(jì)算的方式,其實(shí)我們?cè)诔杀居?jì)算上使用了學(xué)術(shù)界比較通行的計(jì)算方法,這里也簡(jiǎn)單介紹一下。

根據(jù)在不同平臺(tái)上對(duì) Ling-Plus 的真實(shí)訓(xùn)練記錄,我們可以觀察到某個(gè)平臺(tái)在 K 張加速卡上持續(xù)一段時(shí)間(比如一周)的 token 數(shù),再根據(jù)技術(shù)報(bào)告表 1 上提到的不同加速卡的單位時(shí)間成本,就可以很簡(jiǎn)單地計(jì)算出對(duì)應(yīng)平臺(tái)上訓(xùn)練單位 token 量(報(bào)告里以 1 萬億 token 為單位)的成本。

表1:AI加速器特性與單位成本(估算)

事實(shí)上,不管是在 GPU 還是在國(guó)產(chǎn)加速卡上,LLM 的訓(xùn)練成本優(yōu)化都是無止境的。Ling 的訓(xùn)練過程一定程度地說明,在我們做的這些技術(shù)努力上,國(guó)產(chǎn)加速卡上的訓(xùn)練成本與 GPU 相當(dāng)甚至更低,同時(shí)可以保證 loss 收斂一模一樣。

未來的工作

Ling 模型的發(fā)布只是我們工作的一個(gè)里程碑,后續(xù)我們還會(huì)進(jìn)一步改進(jìn)自己的工作。DeepSeek 為我們對(duì)訓(xùn)練經(jīng)濟(jì)性的提升帶來了啟發(fā),DeepSeek 在訓(xùn)練中使用了 FP8 證明了這樣的低精度浮點(diǎn)數(shù)是可以訓(xùn)練出來優(yōu)秀的大模型的;同樣我們兄弟團(tuán)隊(duì)基于強(qiáng)化學(xué)習(xí)的 AReaL(https://github.com/inclusionAI/AReaL)也開源了,強(qiáng)化學(xué)習(xí)也是通往 AGI 之路的重要一環(huán)。我們后續(xù)的更多工作也會(huì)陸續(xù)開源在 inclusionAI org(https://huggingface.co/inclusionAI)里。

每個(gè) AI 研發(fā)工程師都相信 AGI 必將到來。我們相信 AGI 一定是普惠大眾的,感謝大家的關(guān)心,期待未來的工作也能受到持續(xù)關(guān)注。

相關(guān)推薦
免責(zé)聲明
本文所包含的觀點(diǎn)僅代表作者個(gè)人看法,不代表新火種的觀點(diǎn)。在新火種上獲取的所有信息均不應(yīng)被視為投資建議。新火種對(duì)本文可能提及或鏈接的任何項(xiàng)目不表示認(rèn)可。 交易和投資涉及高風(fēng)險(xiǎn),讀者在采取與本文內(nèi)容相關(guān)的任何行動(dòng)之前,請(qǐng)務(wù)必進(jìn)行充分的盡職調(diào)查。最終的決策應(yīng)該基于您自己的獨(dú)立判斷。新火種不對(duì)因依賴本文觀點(diǎn)而產(chǎn)生的任何金錢損失負(fù)任何責(zé)任。

熱門文章