LLM訓練通信量減少10000倍！全新分布式優化器，整合世間算力訓練強大AI
如果可以使用世界上所有的算力來訓練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背后的理論還需要進一步探究，更嚴謹、更詳細的學術論文以及完整的代碼將在未來發布。
參考資料：
