你有沒有想過，訓練一個幾百億參數的AI模型時，最先崩潰的往往不是算力，而是顯存？
這事說起來挺反直覺的。大家平時聊起大模型訓練,總覺得瓶頸是算力不夠、GPU太少。但實際情況往往是,一台訓練任務跑得好好的,突然就因為某個環節內存爆了,整個訓練直接中斷。更麻煩的是,這個"某個環節"每次可能都不一樣。有時候是路由矩陣撐爆了顯存,有時候是詞表投影層吃掉了所有空間,有時候又是優化器狀態占滿了內存。
Salesforce AI Research 的研究者們在這篇論文裡指出了一個很樸素但常被忽視的道理：訓練能不能成功,不取決於平均顯存占用,而取決於每一個組件的峰值占用有沒有超過設備容量。哪怕你把三個瓶頸都壓下去了,只要第四個還在瘋長,訓練照樣會在某個臨界點崩掉。
這就好比一個裝修隊給一棟樓做防水,四面牆分別請了四個不同的師傅施工,前三面牆都做得天衣無縫,第四面牆卻隨便糊了糊。下雨的時候,水照樣會從第四面牆滲進來,前面三面牆做得多好都沒用。樓漏水這件事只需要一個薄弱點,不需要四個。如果這四個師傅不能同時把關,樓遲早要出問題。
MoE模型
訓練里的四個"薄弱牆面"
混合專家模型
MoE：Mixture-of-Experts，混合專家模型，一種讓不同"專家"網路分別處理不同輸入的架構，可以在不顯著增加計算量的前提下大幅擴展模型參數量
之所以特別容易在長上下文或大批量訓練時出問題,是因為它天生帶著四個會隨著配置變化而瘋長的"計算包袱"。
第一個包袱是專家調度。在MoE架構里,每個GPU只負責一部分專家,但輸入的token（可以理解成文本被切分後的最小處理單元）要先被路由到它該去的專家那裡。如果路由不均衡,某幾個專家突然被大量token選中,負責這些專家的GPU瞬間就要處理海量數據,內存占用直接飆升。
第二個包袱是詞表投影層。模型輸出的時候,需要把每個token的表示映射到整個詞表上,算出每個詞的概率。這一步產生的張量大小是"token數乘以詞表大小",詞表通常有幾萬甚至幾十萬個詞,長上下文一來,這個張量能輕鬆撐爆顯存。
第三個包袱是梯度檢查點。
梯度檢查點：訓練時為了省顯存,不保留每一層的中間計算結果,反向傳播時需要用到就重新算一遍,用計算時間換內存空間
這個技術本身是為了省內存的,但每個檢查點邊界處仍然要保留一個"活著"的輸入張量,直到反向傳播真正用上它。層數越深、序列越長,這些懸而未決的張量加起來也是一筆不小的開銷。
第四個包袱是優化器狀態。像AdamW這類優化器,需要給每個參數額外存一份動量、一份二階矩估計,參數量一大,這部分狀態占的內存能輕鬆超過模型本身。
論文裡說得很直白：這四個包袱誰先"爆表"完全取決於具體配置。詞表大、上下文長,詞表投影先出問題;路由本身不均衡,專家調度先崩;網路很深,檢查點邊界先撐不住;參數量巨大但設備數量有限,優化器狀態先超標。
這也是為什麼之前很多方案總是"治標不治本"。你解決了一個問題,換個場景另一個問題立刻冒出來,像打地鼠遊戲一樣,永遠打不完。
四個針對性的解決方案
論文提出了四個操作,分別對準上面四個包袱,每一個都只改變計算的順序和粒度，不碰模型本身、不碰精度、不碰損失函數的計算方式。換句話說,訓練出來的結果和標準全參數BF16訓練完全一致,不是那種"用精度換內存"的取巧做法。
先說專家調度這塊。之前已經有一個叫LLEP的方法（Least-Loaded Expert Parallelism,最閒專家並行）,通過把熱門專家的溢出任務挪給閒著的GPU,來解決負載不均衡的問題。但LLEP有個問題：它雖然把任務挪走了,卻還是一次性把整批數據都塞進內存,該來的那批數據有多大,內存占用就有多大。
論文提出的PipelinedLLEP做了一個關鍵改動：把整批數據切成一個個小塊（chunk）,規定每個GPU往每個小塊里最多能塞多少token,然後用流水線的方式,讓傳輸和計算重疊進行。這樣一來,無論路由多不均衡,單個小塊的大小是被嚴格限制住的,不會因為某個專家突然爆紅就讓內存跟著炸開。
這就像是一個食堂打飯,如果讓所有想吃紅燒肉的人一次性衝上來打飯,窗口肯定被擠爆。但如果規定每一輪最多放20個人進來打飯,不管這一輪里有多少人想吃紅燒肉,這一輪的壓力都是可控的。要是不設這個限流規則,某一天突然全食堂的人都想吃紅燒肉,窗口直接被衝垮,沒人能吃上飯。
不過這裡有個細節值得說說：如果單純把批次切塊,並行執行計算和通信,內存峰值其實降不下來,因為一個循環調用多次的時候,每次調用產生的計算圖（autograd圖）都會一直掛在內存里,直到反向傳播算完才釋放。論文用了一個嵌套的梯度檢查點技巧解決了這個問題,把每個小塊的專家計算單獨包一層可重入檢查點,讓前一個塊的中間結果能及時釋放,後一個塊才分配自己的空間。這個細節聽起來技術性很強,但本質就是"用一次就扔",而不是"攢著最後一起處理"。
實測數據很直接：在65536個token、128個專家、top-8路由的配置下,標準專家並行一旦碰上嚴重的路由不均衡直接內存溢出崩潰,而PipelinedLLEP相比LLEP能省下56.9%到59.3%的峰值內存,速度還基本不受影響,甚至有的配置下反而更快。
詞表投影：讓巨大的logit張量"隱形"
第二個方案叫Ring-DTP,針對的是詞表投影這個環節。
要理解這個方案,得先明白為什麼詞表投影這麼費內存。計算交叉熵損失（cross-entropy loss,衡量模型預測和真實答案差距的常用損失函數）的時候,理論上每個token只需要三個數：一個最大值、一個指數和、以及目標詞對應的那個logit值。但傳統做法是先把整個"token數乘詞表大小"的logit矩陣算出來,再從裡面提取需要的那幾個數。這就好比為了知道一群人里誰最高,先把每個人的詳細體檢報告全部列印出來堆在桌上,而實際上你只需要一個數字。
之前已經有一些單機上的融合交叉熵核心（比如Cut Cross-Entropy和Liger Kernel）,通過在線維護log-sum-exp的方式避免生成完整的logit矩陣。但這些方法有個前提：整批數據要在同一塊GPU上。而Megatron的做法雖然把權重分片了,卻要求每個GPU持有相同的批次數據,這直接砍掉了有效批量大小。
Ring-DTP的做法是讓每個GPU保留自己獨有的一批數據,同時也持有詞表的一部分權重分片,然後讓數據或權重像接力賽一樣在GPU之間循環傳遞,直到每一批數據都和每一片權重"見過面"。每次相遇只計算一小條logit,提取出需要的三個統計量後立刻釋放這條logit,絕不讓完整的大矩陣成型。
這有點像一個大型聯誼活動,如果要求所有參與者同時擠進一個大廳認識每一個人,場地壓力巨大。但如果換成輪轉制,每一輪只有一小批人和一小批人碰頭,聊完記錄一下關鍵資訊就換下一輪,場地壓力就小多了,而且最終每個人還是能認識到所有該認識的人。如果不這麼設計,非要一次性湊齊所有人,場地根本裝不下這麼多人。
實驗數據顯示,在8路分片、16384個token的配置下,標準方法的峰值內存是42.5GB,而Ring-DTP只用了7.3GB,省了82.8%,時間只多花了5.1%。當token數翻倍到32768時,標準方法的內存幾乎翻倍漲到79.5GB,Ring-DTP卻只漲到10.6GB,說明這個方法在長上下文場景下的優勢會越來越明顯。這正是讓詞表投影層能在百萬token上下文裡跑起來的關鍵。
檢查點邊界：把"睡著的"張量挪到CPU去
第三個方案叫SCO（Selective Checkpoint Offload,選擇性檢查點卸載）,解決的是梯度檢查點技術留下的一個尾巴問題。
梯度檢查點省內存的方式是"不存中間結果,需要時重新算",但每一層的輸入張量必須一直留在顯存里,從前向傳播開始一直等到反向傳播真正重新計算這一層為止。這個等待期可能很長,尤其是層數很多的時候,相當於很多個"半成品"張量同時占著顯存位置,誰也不動,就是干占著。
SCO的思路很簡單：把這些暫時用不上但又必須留著的張量,先挪到CPU內存里存著,快用到的時候再提前一層取回來。這樣GPU顯存里同時最多只需要留兩個正在恢復的張量,而不是所有層的輸入張量全部囤在那裡。
這就像是搬家公司打包東西,如果所有箱子都堆在客廳里等著裝車,客廳會被堵得進出不了。但如果把還沒輪到裝車的箱子先挪到走廊或者陽台上,等快輪到了再搬回客廳門口,客廳始終只有一兩個箱子占地,整個流程一樣能順利完成,只是多了一次搬運的功夫。要是不做這個中轉,客廳堆滿了箱子,連走路的空間都沒有,搬家反而更慢。
論文用gpt-oss-20b模型做了測試,給CPU內存不同的預算,GPU的顯存峰值確實隨著預算的增加單調下降,吞吐量的變化不到2%,但省下的顯存換來的是能跑更大的批次,最大批次提升了17.7%。這說明SCO這種"分批卸載"的策略,在幾乎不影響速度的前提下,實實在在地擴大了訓練的可承受空間。
優化器更新：別讓GPU乾等著
第四個方案叫OffloadStreamAdamW,針對的是優化器狀態卸載這件事本身帶來的新問題。
前面提到,把優化器狀態（比如AdamW需要的動量和二階矩）放到CPU內存里能省顯存,這是ZeRO-Offload等已有方法的做法。但問題是,CPU算這個更新的速度遠不如GPU,更新期間GPU完全閒著,而且此時GPU的顯存本來就因為激活值已經釋放而處於空閒狀態,這段時間等於是資源雙重浪費。
OffloadStreamAdamW的解法是反過來利用這個"空閒期"：把CPU上存的參數狀態分成一個個小批次（bucket）,輪流傳到GPU上,讓GPU來做真正的計算,算完再傳回CPU保存。整個過程用三條並行的流水線（傳輸、計算、寫回）交替進行,讓GPU不再是乾等待著,而是變成真正幹活的角色。
這就好比一個倉庫管理員,原本的做法是把所有貨物一次性搬到隔壁小屋去逐個稱重登記,稱重的人手速很慢,搬運工在旁邊干站著沒事幹。改進後的做法是讓搬運工按批次不停地搬運貨物,稱重的人也不停地稱重,兩邊同時忙碌,誰都不閒著。如果繼續用老辦法一次性全搬過去再乾等,搬運工的體力（GPU算力）就白白浪費了。
實測效果是,相比CPU上的AVX向量化AdamW實現,OffloadStreamAdamW把優化器更新這一步從3.95秒降到了1.93秒,快了2.05倍。而且論文還發現,增加緩衝槽位（staging slots）數量並不能進一步提速,說明這個流程本身已經被傳輸頻寬卡住了,這是理論上能達到的速度上限附近。
四個方案合體：一百萬token上下文成真
把這四個方案組合起來放進一個叫MoP（Mixture-of-Parallelisms,混合併行）的整體架構里,論文在120B、241B、667B三個不同規模的MoE模型上做了端到端測試。
結果相當亮眼：對比精心調優過的FSDP2（一種常見的分布式訓練框架配置）基線,組合方案能訓練一百萬token的上下文長度,是基線能達到的上下文長度的8到32倍。而在兩者都能跑的最長上下文長度上做比較,新方案的吞吐量還能達到基線的7.6到10.4倍。
這個數字差距挺誇張的。FSDP2基線在241B模型這個規模上,超過32K token就直接內存溢出,而新方案的最短測試配置就是128K,是基線極限的四倍起步。批量大小方面同樣有優勢,最大能跑的全局批次是基線的3到12倍。
論文還專門做了個訓練質量的驗證實驗,用gpt-oss-20b模型在數學題數據集上做微調,對比新方案和FSDP2基線訓練出來的模型,在AIME 2025測試集上的準確率分別是59.8%和59.6%,幾乎沒有差異。這說明省內存這件事沒有偷工減料,模型學到的東西是一樣的。
這幾個方案分別針對四個不同的內存瓶頸,而且論文特意強調每一個都可以單獨啟用,不依賴其他三個。這個設計思路挺聰明的：不同的訓練任務遇到的瓶頸不一樣,有的模型詞表特別大,有的路由特別不均衡,讓用戶按需開啟相應的方案,而不是被迫承擔一整套複雜系統的全部開銷。
寫在後面
讀這篇論文的過程中,最觸動我的其實不是某一個具體的技術細節,而是論文標題里那句"flattening every memory peak"背後的思維方式。
大部分工程優化文章講的是"我們把X降低了多少",但這篇論文一開篇就先給你講清楚一個殘酷的事實：降低一個峰值,如果留著另一個峰值不管,訓練照樣跑不起來。這種"整體約束"的視角,其實比單點優化更難做,因為你得同時理解四個完全不同的系統組件,還要保證它們之間不互相衝突。
論文裡有個細節我覺得值得單獨說一說:在講PipelinedLLEP的時候,他們發現光靠"流水線重疊"這個手段並不足以真正降低內存峰值,還需要配合嵌套的梯度檢查點技巧才能把每個塊的中間結果及時釋放掉。這說明很多看似樸素的"分塊處理"思路,實際落地時往往藏著一層容易被忽略的坑,通信重疊解決的是速度問題,內存峰值的問題需要另外的機制去處理,這兩件事看起來相關,其實是兩個獨立的維度。
另外一個有意思的地方是,論文附錄里專門討論了一個"延遲出現雙峰分布"的現象:在特定的分塊數量下,同樣的配置有時候跑得快,有時候莫名其妙慢了一大截,而且這種慢的模式一旦觸發就會穩定重現。研究者最後發現根源出在主機端的啟動模式上,插入一次同步操作就能讓這個慢模式轉移到別的配置去。這種"玄學問題"最後被系統性地定位出來,而不是簡單歸因於"隨機波動",這種排查思路本身也挺值得學習的。
這篇論文沒有解決的問題其實也挺明顯：四個參數(序列並行度、專家並行度、投影並行度、chunk大小、bucket大小)目前還是靠人工調參和實測曲線來選,論文自己也在最後承認這一點。如果未來能有一個自動化的方式根據模型結構和硬體配置直接推薦這套參數,那才是真正把這套方法變成"開箱即用"的工具,而不是需要專家經驗才能用好的精密儀器。
Q&A
Q1：論文提出的四個內存優化方法分別解決什麼問題？
A：分別對應MoE訓練中四個會不受控增長的內存瓶頸：PipelinedLLEP限制專家調度時的token緩衝區大小，Ring-DTP讓詞表投影不用生成完整的巨大logit矩陣，SCO把梯度檢查點的部分張量卸載到CPU內存，OffloadStreamAdamW用GPU流水線加速CPU優化器狀態的更新過程。
Q2：這套方法相比傳統訓練方案能帶來多大的提升？
A：在120B到667B參數規模的MoE模型上，組合使用這四個方法後能訓練一百萬token長度的上下文，是調優後FSDP2基線能達到長度的8到32倍，在相同上下文長度下吞吐量最高能達到基線的10.4倍，最大批量能達到基線的12倍。
Q3：使用這些內存優化方法會不會影響模型訓練效果？
A：不會。論文特意做了對比實驗，用同樣的數據和超參數分別訓練模型，新方案和傳統FSDP2基線在AIME 2025數學測試集上的準確率分別是59.8%和59.6%，幾乎沒有差異，說明省內存並沒有犧牲訓練質量。
