LLM訓練通信量減少10000倍！全新分佈式優化器，整合世間算力訓練強大AI

新智元報道

編輯：alan

【新智元導讀】近日，Nous Research宣佈了一項重大突破，通過使用與架構和網絡無關的分佈式優化器，研究人員成功將訓練LLM時GPU間的通信量降低了1000到10000倍！

如果可以使用世界上所有的算力來訓練AI模型，會怎麼樣？

近日，憑藉發佈了開源的Hermes 3（基於Llama 3.1）而引起廣泛關注的Nous Research，再次宣佈了一項重大突破——DisTrO（分佈式互聯網訓練）。

通過使用與架構和網絡無關的分佈式優化器，研究人員成功將訓練LLM時GPU間的通信量降低了1000到10000倍！

初步技術報告：https://github.com/NousResearch/DisTrO/

在如此誇張的改進之下，大模型訓練的重要成本和瓶頸——帶寬，也就不再是問題。

使用DisTrO的方法，你可以將訓練負載分佈到互聯網上，而整個網絡世界也就成爲了一個巨大的異構的AI服務器集羣。

——任何有相關算力的設備都可以參與到訓練過程之中。

實驗證明，本文的方法基本不會導致模型性能下降，同時DisTrO-AdamW在收斂速度方面，也與標準的AdamW+All-Reduce相當。

分佈式互聯網訓練

一般來說，訓練大規模神經網絡涉及到大量的通信開銷。

比如做數據並行的時候，不同的訓練數據在不同的硬件（顯卡等）上進行前向和反向計算，之後，同一批數據計算出的梯度需要在顯卡之間先完成同步，才能進入下一個epoch。

如果是模型並行，那麼中間數據就需要通過All-Reduce進行拼接或者累加。

這些數據通信開銷如果不能overlap掉，就會成爲模型訓練的瓶頸。

而恰好，老黃的顯存和帶寬又很貴，甚至組多卡時候需要的硬件也很貴。

爲了解決這個問題，研究人員開發了DisTrO，在不依賴攤銷分析的情況下，將GPU間通信要求降低了四到五個數量級，從而能夠在慢速網絡上對大型神經網絡進行低延遲訓練。

DisTrO是通用、可擴展，並且時鐘同步的（與SGD、Adam等類似，每個訓練步驟使用相同的算術運算並花費相同的時間）。

另外，與之前的ad-hoc低通信優化器相比，DisTrO對電信網絡的拓撲和神經網絡架構不敏感，能夠以最小的開銷原生支持分佈式數據並行訓練（DDP）。

LLM預訓練

研究人員使用Nanotron作爲預訓練框架，且僅在DDP策略下運行（每個GPU都將整個模型加載到VRAM中）。

LLM選擇1.2B大小的Llama 2，模型和訓練所用的超參數如下：

訓練數據使用Dolma v1.7數據集，隨機選出的10%代表性樣本（前 105B個token）。

優化器採用AdamW，β1=0.9、β2=0.95，峰值學習率爲4×10e-4，使用餘弦衰減方案，權重衰減設置爲0.1。

作爲對比的另一組實驗，將AdamW替換爲DisTrO-AdamW，但不更改超參數，並禁用Nanotron中的All-Reduce操作。

與以前的分佈式訓練方法不同，DisTrO不同步優化器狀態（甚至可以無狀態）。

下圖是兩組實驗的訓練損失曲線，使用105B數據訓練25000步。可以看出，DisTrO的收斂能力與All-Reduce持平。

重要的是，在不影響訓練效果的情況下，DisTrO將通信量從74.4GB直接減到了86.8MB！相當於帶寬壓力減少了857倍。

作者還表示，這857倍只是初期測試，後面調調超參數，減少個1000倍到3000倍也不是問題。

如果是後訓練和微調，甚至可以實現高達10000倍的通信優化，且基本不影響訓練效果。

最後，爲了驗證訓練效果，作者在訓練後的模型上執行了GPT4All零樣本基準測試，並與在相同數量的token上訓練的TinyLlama（checkpoint）進行了比較。

結果如上表所示，TinyLlama的架構和訓練過程與本文的實驗非常相似，可以作爲對結果進行健全性檢查的衡量標準。

未來應用

數據流

在本實驗的場景中，32個節點使用最簡單的All-Reduce（全連接），每個節點平均傳輸86.8MB（2.8MB×31），並接收相同數量的數據。

如果使用專用服務器進行數據聚合，則每個節點只需上傳2.8MB數據（接收數據不變），通信量進一步減少。

另外，不對稱性是有優點的，因爲大多數消費互聯網的帶寬嚴重偏向於更高的下載速度。

假設穩定的網速爲100Mbps下載和10Mbps上傳，則最壞情況下的延遲僅爲下載6.94秒，上傳2.24秒，overlap一下則每步延遲爲6.94秒。

ps：以上的數據傳輸都是原始的向量，如果用上壓縮技術還能更快。

帶寬

作者表示，目前的實驗和研究還比較有限，無法斷定隨着模型變大，帶寬減少的比率是會增加、減少還是保持不變。

不過目前的1.2B似乎是DisTrO能夠良好工作的最小尺寸（再小就不收斂了），所以可以假設隨着模型大小的增長，需要的通信會相對越來越少。

不過也可能通信量與模型大小沒有關係，這時可以在不增加通信帶寬的情況下增加模型大小，觀察更大的模型是否會改善訓練和學習的效果。

如果後一種情況屬實，那麼未來GPU設計和製造的範式將會被改變（更大VRAM和更窄帶寬）。

恰好我們也更喜歡計算密集型負載（而不是I/O密集型），畢竟現在的帶寬要比計算貴得多。

聯邦學習

除了訓練LLM，DisTrO還能用來做什麼？

在互聯網上做分佈式訓練，讓人一下就想到了聯邦學習。

在允許模型協作訓練的同時，保持每個參與者的數據的私密性和去中心化，這在LLM被大公司掌握的當下，顯得越來越重要。

到目前爲止，聯邦學習一直缺乏在有限的互聯網帶寬上訓練大型模型的有效方法。

而DisTrO對如何處理數據，或將數據分配給各個GPU節點沒有任何要求，並且可以無狀態（類似於聯邦平均），因此適用於聯邦學習的未來。

虛擬異構GPU集羣

此外，DisTrO可以創建一個完全去中心化且無需許可的網絡來協作和共享資源。

實驗表明，DisTrO對於訓練期間少量降級或丟棄的節點具有顯著的彈性，並且可以輕鬆地適應新節點的加入。

在這種能力加持之下，一方面可以保障整個系統的安全性，降低不可信節點使用對抗性攻擊破壞運行的風險。

另一方面，也可以鼓勵機構和個人靈活貢獻自己的計算資源，釋放潛在的算力。

甚至一些內存或者算力不太夠的老卡，也能加入進來賺點外快，採用FSDP、SWARM Parallelism等策略與DisTrO協同工作。

能源

DisTrO的進一步大規模應用，可能會緩解建設大型數據中心所帶來的能源消耗、基礎設施成本和土地使用等相關問題。

Llama 3.1項目需要構建兩個大型整體超級集羣，每個集羣包含 24,000個H100 GPU，僅訓練過程就產生了相當於11,000噸的二氧化碳排放。

當今的LLM，除了模型參數大小的增長，訓練數據量也在不斷增大，導致AI相關的數據中心已經摸到了現代電網的極限。

DisTrO可用於自適應平衡多個使用過剩容量的小型模塊化數據中心，通過動態平衡訓練技術利用現有基礎設施，減輕訓練對環境的負面影響。

目前，DisTrO背後的理論還需要進一步探究，更嚴謹、更詳細的學術論文以及完整的代碼將在未來發布。

參考資料：

https://venturebeat.com/ai/this-could-change-everything-nous-research-unveils-new-tool-to-train-powerful-ai-models-with-10000x-efficiency/