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 也能勝過資料並行訓練方法。
（來源：arXiv）
研究人員：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 研究人員使用最新的全域性模型。（注：留出集，是指從原始資料集中特意保留的一部分資料，用於評估模型的泛化效能。）
在谷歌張量處理單元 TPUv5e 和 TPUv6e 上進行實驗
在具體的實驗方法上，研究人員使用 Chinchilla-style 的僅解碼器 Transformer 架構。
同時，使用 QKLayerNorm 來降低對於學習率的敏感性。需要說明的是，QKLayerNorm 是一種改進的層歸一化技術，主要用於 Transformer 架構中的自注意力機制。另外，研究人員還使用 z 損失正則化來提高訓練穩定性。
本次研究使用的詞彙量為 32,768：其中有 32,000 個詞彙表內的單詞，另外還有一些額外的標記用於表示句子開頭和詞彙表外的內容。
實驗中，研究人員將多個序列打包到每個 batch 中，整個 batch 的最大序列長度為 2048。
這時，研究人員開始針對一系列的模型進行訓練，這些模型的 transformer 層數、注意力頭數量、QKV 維度和前饋層隱藏維度各不相同。
其中，QKV 維度是指查詢（Query）、鍵（Key）和值（Value）向量的維度。這些向量是自注意力機制的核心組成部分，用於計算輸入序列中不同位置之間的關係。
圖 | 每個模型比例的架構細節（來源：arXiv）
在大多數實驗中，研究人員全程使用 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 位浮點數動態範圍的同時，能夠減少儲存和計算開銷。
研究人員在谷歌張量處理單元 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 億引數量的模型訓練得擁有更低的評估損失。
（來源：arXiv）
結果顯示，DiLoCo 能夠實現更自然的水平擴充套件性。無論在任何情況，token 預算 D 僅僅是 N 的函式。
這意味著當使用 4 倍大的 batch 時，訓練步驟能夠減少到原來的 1/4。
對於 DiLoCo 來說，這會產生相當好的效能，並且可以同時使用更多資源，從而減少模型總訓練時間。
相比之下，資料並行訓練方法似乎需要更多的序列訓練。因為，訓練時間的減少與通訊的減少是相輔相成的。
為了展示這些效果，研究人員在下圖中繪製了在不同頻寬的網路下訓練時的理想化 wall-clock 時間。
（來源：arXiv）
藉此發現，DiLoCo 對更大 batch 的容忍度使其能夠更快地實現與資料並行方法大致相當的損耗。
在下圖中，研究人員繪製了不同過度訓練量下資料並行方法和 DiLoCo 的理想訓練時間（M = 2）。
（來源：arXiv）
其發現，DiLoCo 能夠透過降低通訊成本和利用更大的 batch 來加速過度訓練，因此需要的序列訓練步驟更少。
這表明 DiLoCo 對於過訓練來說是一個巨大的福祉，因為這意味著可以透過橫向可擴充套件性來攤銷計算時間。（注：過訓練，是指模型在訓練過程中過度擬合訓練資料，導致其在未見資料上的效能下降。）
與此同時，研究人員的結果表明，DiLoCo 和資料並行訓練這兩種方法通常都是有效的，儘管沒有明確的贏家，但是 M 之間的殘差存在顯著差異。
儘管如此，研究人員發現平均而言，雖然獨立擬合在預測損失和全域性批次大小方面略優於聯合擬合，但獨立擬合在預測內部學習率方面要明顯好得多。
總的來說，本次結果表明與資料並行方法一樣，DiLoCo 可以隨著模型大小的變化實現可預測的擴充套件，這使得在極大的尺度上調整超引數和訓練模型變得更加簡單。
此外，DiLoCo 相比資料並行方法可以帶來顯著優勢，包括：使用單個模型副本時具備更優的評估損失，在任何模型副本數量下最優 batch 大小都能增大。對於模型規模、過訓練和同步頻率來說，這將產生更好的魯棒性。
不過，眼下顯然需要可用於大規模部署 DiLoCo 等方法的系統和軟體，並能在真實的超大模型環境中實現其通訊效率優勢。
參考資料：
https://arxiv.org/pdf/2503.09799
排版：初嘉實