LLM訓練通信量減少10000倍,全新分佈式優化器,整合世間算力訓練強大AI
近日,Nous Research宣佈了一項重大突破,通過使用與架構和網絡無關的分佈式優化器,研究人員成功將訓練LLM時GPU間的通信量降低了1000到10000倍!
如果可以使用世界上所有的算力來訓練AI模型,會怎麼樣?
近日,憑藉發佈了開源的Hermes 3(基於Llama 3.1)而引起廣泛關注的Nous Research,再次宣佈了一項重大突破——DisTrO(分佈式互聯網訓練)。
通過使用與架構和網絡無關的分佈式優化器,研究人員成功將訓練LLM時GPU間的通信量降低了1000到10000倍!
在如此誇張的改進之下,大模型訓練的重要成本和瓶頸——帶寬,也就不再是問題。
使用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/
本文來自微信公眾號「新智元」,作者:新智元,36氪經授權發佈。