大模型如何訓練百萬 Token 上下文：上下文并行與 Ring Attention

只用了幾年時間，上下文窗口就從 4k 膨脹到 1000 萬。Meta 發布的 Llama 4 Scout 的時候說這個模型支持 1000 萬 Token，是 Llama 3 那 128k 的 78 倍。而Google Gemini 3 Pro 是 100 萬，Claude 4 也桐鄉市100萬。

一次推理跑完整個代碼庫、幾百篇論文、連續好幾天的對話記錄在技術上可行了，但問題是硬件跟不上。

405B 參數的模型，32 位精度下光權重就要 6.5TB 內存。再算上梯度、狀態、激活值，后者還隨上下文長度二次方增長。單臺 NVIDIA HGX B300 配了 2.3TB HBM3e都不夠。

這就逼著必須做多節點分布式訓練和推理，幾十上百塊 NVIDIA Blackwell GPU 、NVLink 再加上 InfiniBand，就成了數據中心的標配。所以難點就變味了 GPU 之間的通信瓶頸。

并行化基礎

模型或數據集超出單卡容量，就得上并行策略，但是每種策略本質上都是拿通信開銷換內存空間。

數據并行是最直接的方案：整個模型復制到每張卡上，訓練數據切開，每張卡跑不同的 batch跑完一步同步梯度。適合小模型，計算是瓶頸、內存不是問題的場景。

模型并行針對大模型：單卡裝不下，就把模型拆開，不同的層放不同的卡上，按順序跑。405B 這種規模只能這樣，并且下游的卡得等上游算完中間是有空轉的。

張量并行更極端：連單個矩陣乘法都塞不進一張卡。就需要把矩陣按行或按列切開，分到各卡上算，再通過 all-reduce 合起來。

但這些都有共同的局限。模型大、上下文又長到幾百萬 Token，張量并行也頂不住。因為注意力的二次方內存增長太兇，激活值直接占滿顯存。128k 上下文的激活值內存是 8k 的 16 倍，這個目前沒辦法，因為就是這么夸張。

上下文并行與序列并行

序列并行和上下文并行都是在設備間切序列來省內存，但切法不一樣。

序列并行配合張量并行使用，只切那些非矩陣乘法的操作，比如層歸一化、dropout。張量并行管不到的地方，序列并行接手，每張卡處理一部分激活值。兩者配合能把序列撐長一些，但到 128k 以上還是會有問題，因為注意力的二次方增長是繞不過去。

上下文并行更徹底：整個序列在所有模塊里都切開，包括注意力。每個操作拿到的都是分區后的序列。百萬級上下文的訓練就靠這個，把激活值的內存占用分攤到各卡上。

注意力一直是最麻煩的問題，因為模型的其他操作基本都是逐 Token 獨立處理并行起來很自然。但注意力不行，每個 Token 都要"看"序列里所有其他 Token。序列切到多張卡上之后，GPU 1 的 Token 怎么看 GPU 2 的 Token？直接等數據傳完再算，整個流水線就卡住了。

Ring Attention 就是來解決這個問題的，讓多節點多卡的大模型訓練和推理能在大規模數據中心里跑起來。

Zig Zag Ring Attention：通信和計算重疊

Ring Attention 把 GPU 組織成環形拓撲。每張卡的工作流程是這樣的：持有序列中 Q、K、V 張量的一個分塊；用本地的 K 和 V 給自己的 Q 分塊算注意力；把 K 和 V 傳給環里的下一張卡；從上一張卡接收 K 和 V；循環往復，直到所有 Q Token 都跟所有 K/V Token 算完注意力。

關鍵在于計算和通信是重疊的。GPU 1 拿著當前的 K/V 分塊算注意力的時候，同時在從 GPU 0 接收下一批分塊。通信延遲減少了，因為不用干等數據全到了再開算。

GPT 這類自回歸模型有個額外的麻煩：Token 只能看前面的 Token不能看后面的。所以會導致負載不均衡有些卡會空轉，Zig-Zag Ring Attention 解決這個問題的辦法是交錯分配，不是按順序切塊而是 GPU 0 拿 Token [0, 4, 8...]，GPU 1 拿 [1, 5, 9...]，以此類推。每張卡都拿到早期和晚期 Token 的混合，因果注意力計算時負載就均衡了環里不會有卡閑著。

但是代價是索引邏輯稍微復雜一點，不過大規模場景下性能收益很可觀，因果掩碼下也能做到接近滿 GPU 利用率。

上下文并行與 Ring Attention 常見問題

上下文并行把輸入序列切到多張 GPU 上，突破訓練時的內存限制。跟張量并行、數據并行不同，它在所有模型模塊里都切序列維度。單卡裝不下的百萬級 Token 上下文，只有靠這個才能訓。

Ring Attention 把 GPU 排成環，每張卡一邊算當前數據的注意力，一邊把鍵值對往下傳。通信和計算重疊，全對全的注意力計算不用等完整序列數據到齊，GPU 不會干等。

而序列并行只切非矩陣乘法操作（層歸一化之類的），配合張量并行用。上下文并行在所有模塊里都切序列，包括注意力。超過 128k Token 的上下文必須用后者，因為激活值內存二次方增長太猛了。

為什么 Zig-Zag Ring Attention 比標準 Ring Attention 更好？

Zig-Zag 用交錯分配代替順序分配，因果掩碼計算時各卡負載更均衡。標準 Ring Attention 會讓后面的卡等前面的分塊，造成計算空閑。Zig-Zag 把早期和晚期 Token 均勻撒到各卡上，避免這個問題。

那么訓練百萬級 Token 上下文的模型需要什么硬件？

多節點 GPU 集群，配 HBM 內存，加高速互連——NVIDIA NVLink 1.8TB/s 或者 InfiniBand。405B 參數模型 32 位精度從頭訓練加推理，4 臺 NVIDIA HGX B300 的機架部署是個不錯的起點。

總結

上下文并行本質上是拿通信開銷換內存空間，而網絡帶寬是最要命的瓶頸。Ring Attention 要在 GPU 之間不停交換鍵值對，傳輸時間一旦超過計算時間，各卡就會從"邊算邊傳"退化成"等數據"。NVIDIA NVLink 1.8TB/s 加 InfiniBand 的高速互連，在多機架部署里不是可選項是必需品。互連帶寬必須匹配 GPU 計算吞吐量，否則上下文并行的效果會大打折扣。

https://avoid.overfit.cn/post/fd6022b9196942ffb737ba306925b6db

by Khang Pham