Google團隊再次證明Scaling Law的準確性，即使總通信量減少99%，仍能超越數據並行訓練

Scaling Law 由 OpenAI 團隊於 2020 年正式提出，並在其論文《神經語言模型的擴展定律》（Scaling Laws for Neural Language Models）中進行了詳細闡述。Scaling Law 揭示了模型性能與模型規模、數據量和計算資源之間的數學關係。但在這兩年，由於實際應用局限性、理論解釋不完善等原因，Scaling Law 遭到了一些質疑。而 Google 旗下兩支團隊的最新研究，似乎又給 Scaling Law 投上了一張支持票。

在大模型領域，數據並行訓練（Data-Parallel）方法佔據著核心地位，它是訓練大模型的基礎技術之一。然而，隨著模型規模的不斷增大，數據並行訓練方法的通信開銷和內存限制等缺點日益凸顯。

為此，Google DeepMind 此前研發出一種名為 DiLoCo（分佈式低通信優化，Distributed Low-Communication Optimization）的方法，其具備減少通信開銷和提高擴展性等優勢。

近日，來自 Google Research 和 Google DeepMind 的研究人員發現 DiLoCo 竟然具備更多優勢。

研究中，他們研究了在固定計算預算之下訓練大模型時 DiLoCo 的 Scaling Law 行為，借此揭示了 DiLoCo 一系列更通用的優點：包括其能夠增加最佳 batch 大小、能夠提高下遊規模的泛化能力、能夠在固定 token 預算之下改善評估損失。（註：評估損失，是指模型在驗證集或測試集上計算出的損失值。）

研究人員發現隨著模型尺寸的增大，DiLoCo 會呈現出可預測的穩健擴展。如果調整得當，DiLoCo 的模型規模擴展性優於數據並行訓練方法，即使在小模型規模下 DiLoCo 也能勝過數據並行訓練方法。

研究人員：Scaling Law 是準確的

而本次研究中，基於相關數據研究人員為數據並行訓練方法和 DiLoCo 分別建立了評估損失和最優超參數的 Scaling Law。

借此證明，當 Scaling Law 擴展到更大的模型規模時，其具備更好的評估損失和最優超參數。

令人十分驚訝的是：許多情況下對於相同的 token 預算，在通信效率更高的 DiLoCo 上，這些 Scaling Law 預測會比數據並行訓練方法產生更低的評估損失。

為了利用研究人員的 Scaling Law 來預測 DiLoCo 的超參數，他們分別在訓練具有 40 億參數和 100 億參數的模型時加以測試。

事實證明，Scaling Law 是準確的。即使總通信量減少 99% 以上，DiLoCo 的表現也能優於數據並行訓練。

對於研究中的每個實驗，在不同帶寬和不同延遲的網絡下，研究人員都能實現理想的端到端 wall-clock 訓練時間。（註：wall-clock time，指的是從任務開始到結束所經歷的實際時間，包括所有等待時間和計算時間。）

研究人員的實驗包括 M = 1 的 DiLoCo（即僅使用單個副本）。而在這些數據中一個令人驚訝的現像是：即使沒有通信瓶頸，DiLoCo 也能改進訓練效果。

由於 DiLoCo 實際上是另一種優化算法 Lookahead optimizer 的增強版本，因此它不會導致任何通信的減少。

然而，通過使用低頻動量操作，無論在怎樣的模型規模下，DiLoCo 在評估損失和對更大 batch 的容忍度上，都要優於數據並行訓練方法。

值得注意的是，M=1 的 DiLoCo 在評估損失和訓練時間方面都優於數據並行訓練方法。

研究人員發現，M=1 的 DiLoCo 在所有模型尺度上都能實現較低的評估損失，並且能對更大的 batch 產生更好的魯棒性，從而能夠大大減少 wall-clock 訓練時間。

研究中，研究人員分別通過數據並行訓練方法和 DiLoCo 來開展模型訓練。

在數據並行訓練方法中，在每一步研究人員都會獲取一個 batch 大小為 B 的數據。

然後，研究人員計算批次梯度（batch gradient），並使用學習率 γ 進行優化。

在比較數據並行方法和 DiLoCo 時，研究人員始終確保模型大小 N 和總 token 預算 D 保持相同。

為了計算某些留出集的評估損失 L，針對數據並行方法研究人員使用當前模型，針對 DiLoCo 研究人員使用最新的全局模型。（註：留出集，是指從原始數據集中特意保留的一部分數據，用於評估模型的泛化性能。）

在Google張量處理單元 TPUv5e 和 TPUv6e 上進行實驗

在具體的實驗方法上，研究人員使用 Chinchilla-style 的僅解碼器 Transformer 架構。

同時，使用 QKLayerNorm 來降低對於學習率的敏感性。需要說明的是，QKLayerNorm 是一種改進的層歸一化技術，主要用於 Transformer 架構中的自注意力機制。另外，研究人員還使用 z 損失正則化來提高訓練穩定性。

本次研究使用的詞彙量為 32,768：其中有 32,000 個詞彙表內的單詞，另外還有一些額外的標記用於表示句子開頭和詞彙表外的內容。

實驗中，研究人員將多個序列打包到每個 batch 中，整個 batch 的最大序列長度為 2048。

這時，研究人員開始針對一系列的模型進行訓練，這些模型的 transformer 層數、注意力頭數量、QKV 維度和前饋層隱藏維度各不相同。

其中，QKV 維度是指查詢（Query）、鍵（Key）和值（Value）向量的維度。這些向量是自注意力機制的核心組成部分，用於計算輸入序列中不同位置之間的關係。

在大多數實驗中，研究人員全程使用 C4 數據集的訓練集來訓練模型。

此外，還計算了 3 項任務的下遊零樣本評估指標：HellaSwag、Piqa 和 Arc-Easy。在進行過度訓練消融時，使用了 Dolma 數據集。

其還使用了一個改進版的 NanoDO，它利用 DrJAX 在副本之間並行化內部訓練步驟，並將模型副本軸暴露出來以便進行顯式編程。這對於 JAX 中更好的擴展性能至關重要，因為 DrJAX 提供了 jax.vmap 的豐富版本，能夠提供有關 DiLoCo 副本的更明確的分片信息。

與此同時，研究人員使用 all-reduce 進行外部優化。all-reduce 是一種分佈式計算中的通信操作，用於在多個 GPU 或多個 TPU 之間高效地聚合數據。

研究中，其始終使用 bfloat16 來表示模型權重和梯度。bfloat16 是一種 16 位浮點數格式，它由 Google 團隊開發，它在保留 32 位浮點數動態範圍的同時，能夠減少存儲和計算開銷。

研究人員在Google張量處理單元 TPUv5e 和 TPUv6e 上進行了大部分實驗，並在 TPUv-5 上進行了最大規模的實驗。

對於每個實驗，其還計算了一個理想化的 wall-clock 訓練時間，並在這一時間中考慮了理想化的計算時間和通信時間。同時，其還專門測量了端到端的 wall-clock 時間。

其發現，更大的水平並行化例如通過將 batch 大小加倍，將能夠減少 wall-clock 時間。

實驗中，研究人員假設其模型正在跨多個數據中心開展訓練。當在數據中心之內的時候，能夠擁有一個高帶寬網絡。當跨數據中心的時候，則分別擁有高帶寬、中帶寬或低帶寬的網絡。

當通信時間處於理想狀態的時候，研究人員始終將高帶寬網絡用於數據中心內的網絡，以及將三種網絡中的任意一種用於跨數據中心網絡。

DiLoCo：能實現更自然的水平擴展

利用相關數據，研究人員推導出了 Scaling Law，從而能夠預測較大模型的評估損失和最優超參數。

根據 Chinchilla Scaling Law，研究人員假設最優 token 預算由 D=20N 給出。

這意味著在模型大小固定的情況下，如果將 batch 大小 B 加倍，則訓練步驟的數量將會減半。

實驗中，研究人員通過使用 Scaling Law 設置的超參數來訓練 4B 和 10B 模型來驗證上述觀點。

儘管下圖展示的是「插值」區域的情況，同時這也是大量掃瞄的結果，但是這些發現從定性角度來看，同樣適用於「外推」區域。這使得研究人員在 M = 1、2 時，能夠使用 DiLoCo 將 40 億參數和 100 億參數量的模型訓練得擁有更低的評估損失。

結果顯示，DiLoCo 能夠實現更自然的水平擴展性。無論在任何情況，token 預算 D 僅僅是 N 的函數。

這意味著當使用 4 倍大的 batch 時，訓練步驟能夠減少到原來的 1/4。

對於 DiLoCo 來說，這會產生相當好的性能，並且可以同時使用更多資源，從而減少模型總訓練時間。

相比之下，數據並行訓練方法似乎需要更多的串行訓練。因為，訓練時間的減少與通信的減少是相輔相成的。

為了展示這些效果，研究人員在下圖中繪製了在不同帶寬的網絡下訓練時的理想化 wall-clock 時間。

借此發現，DiLoCo 對更大 batch 的容忍度使其能夠更快地實現與數據並行方法大致相當的損耗。

在下圖中，研究人員繪製了不同過度訓練量下數據並行方法和 DiLoCo 的理想訓練時間（M = 2）。

其發現，DiLoCo 能夠通過降低通信成本和利用更大的 batch 來加速過度訓練，因此需要的串行訓練步驟更少。

這表明 DiLoCo 對於過訓練來說是一個巨大的福祉，因為這意味著可以通過橫向可擴展性來攤銷計算時間。（註：過訓練，是指模型在訓練過程中過度擬合訓練數據，導致其在未見數據上的性能下降。）

與此同時，研究人員的結果表明，DiLoCo 和數據並行訓練這兩種方法通常都是有效的，儘管沒有明確的贏家，但是 M 之間的殘差存在顯著差異。

儘管如此，研究人員發現平均而言，雖然獨立擬合在預測損失和全局批量大小方面略優於聯合擬合，但獨立擬合在預測內部學習率方面要明顯好得多。

總的來說，本次結果表明與數據並行方法一樣，DiLoCo 可以隨著模型大小的變化實現可預測的擴展，這使得在極大的尺度上調整超參數和訓練模型變得更加簡單。

此外，DiLoCo 相比數據並行方法可以帶來顯著優勢，包括：使用單個模型副本時具備更優的評估損失，在任何模型副本數量下最優 batch 大小都能增大。對於模型規模、過訓練和同步頻率來說，這將產生更好的魯棒性。

不過，眼下顯然需要可用於大規模部署 DiLoCo 等方法的系統和軟件，並能在真實的超大模型環境中實現其通信效率優勢。

參考資料：

https://arxiv.org/pdf/2503.09799

排版：初嘉實