編輯：alan【新智元導讀】近日，Nous Research宣布了一項重大突破，通過使用與架構(gòu)和網(wǎng)絡(luò)無關(guān)的分布式優(yōu)化器，研究人員成功將訓練LLM時GPU間的通信量降低了1000到10000倍！如果可以使用世界上所有的算力來訓練AI模型，會怎么樣？
近日，憑借發(fā)布了開源的Hermes 3（基于Llama 3.1）而引起廣泛關(guān)注的Nous Research，再次宣布了一項重大突破DisTrO（分布式互聯(lián)網(wǎng)訓練）。
通過使用與架構(gòu)和網(wǎng)絡(luò)無關(guān)的分布式優(yōu)化器，研究人員成功將訓練LLM時GPU間的通信量降低了1000到10000倍！
在如此夸張的改進之下，大模型訓練的重要成本和瓶頸帶寬，也就不再是問題。
使用DisTrO的方法，你可以將訓練負載分布到互聯(lián)網(wǎng)上，而整個網(wǎng)絡(luò)世界也就成為了一個巨大的異構(gòu)的AI服務(wù)器集群。
任何有相關(guān)算力的設(shè)備都可以參與到訓練過程之中。
實驗證明，本文的方法基本不會導致模型性能下降，同時DisTrO-AdamW在收斂速度方面，也與標準的AdamW+All-Reduce相當。
分布式互聯(lián)網(wǎng)訓練
一般來說，訓練大規(guī)模神經(jīng)網(wǎng)絡(luò)涉及到大量的通信開銷。比如做數(shù)據(jù)并行的時候，不同的訓練數(shù)據(jù)在不同的硬件（顯卡等）上進行前向和反向計算，之后，同一批數(shù)據(jù)計算出的梯度需要在顯卡之間先完成同步，才能進入下一個epoch。
如果是模型并行，那么中間數(shù)據(jù)就需要通過All-Reduce進行拼接或者累加。
這些數(shù)據(jù)通信開銷如果不能overlap掉，就會成為模型訓練的瓶頸。
而恰好，老黃的顯存和帶寬又很貴，甚至組多卡時候需要的硬件也很貴。
為了解決這個問題，研究人員開發(fā)了DisTrO，在不依賴攤銷分析的情況下，將GPU間通信要求降低了四到五個數(shù)量級，從而能夠在慢速網(wǎng)絡(luò)上對大型神經(jīng)網(wǎng)絡(luò)進行低延遲訓練。
DisTrO是通用、可擴展，并且時鐘同步的（與SGD、Adam等類似，每個訓練步驟使用相同的算術(shù)運算并花費相同的時間）。
另外，與之前的ad-hoc低通信優(yōu)化器相比，DisTrO對電信網(wǎng)絡(luò)的拓撲和神經(jīng)網(wǎng)絡(luò)架構(gòu)不敏感，能夠以最小的開銷原生支持分布式數(shù)據(jù)并行訓練（DDP）。
LLM預訓練研究人員使用Nanotron作為預訓練框架，且僅在DDP策略下運行（每個GPU都將整個模型加載到VRAM中）。
LLM選擇1.2B大小的Llama 2，模型和訓練所用的超參數(shù)如下：
訓練數(shù)據(jù)使用Dolma v1.7數(shù)據(jù)集，隨機選出的10%代表性樣本（前 105B個token）。
優(yōu)化器采用AdamW，β1=0.9、β2=0.95，峰值學習率為4×10e-4，使用余弦衰減方案，權(quán)重衰減設(shè)置為0.1。
作為對比的另一組實驗，將AdamW替換為DisTrO-AdamW，但不更改超參數(shù)，并禁用Nanotron中的All-Reduce操作。
與以前的分布式訓練方法不同，DisTrO不同步優(yōu)化器狀態(tài)（甚至可以無狀態(tài)）。
下圖是兩組實驗的訓練損失曲線，使用105B數(shù)據(jù)訓練25000步�？梢钥闯�，DisTrO的收斂能力與All-Reduce持平。
重要的是，在不影響訓練效果的情況下，DisTrO將通信量從74.4GB直接減到了86.8MB！相當于帶寬壓力減少了857倍。
作者還表示，這857倍只是初期測試，后面調(diào)調(diào)超參數(shù)，減少個1000倍到3000倍也不是問題。
如果是后訓練和微調(diào)，甚至可以實現(xiàn)高達10000倍的通信優(yōu)化，且基本不影響訓練效果。
最后，為了驗證訓練效果，作者在訓練后的模型上執(zhí)行了GPT4All零樣本基準測試，并與在相同數(shù)量的token上訓練的TinyLlama（checkpoint）進行了比較。
結(jié)果如上表所示，TinyLlama的架構(gòu)和訓練過程與本文的實驗非常相似，可以作為對結(jié)果進行健全性檢查的衡量標準。
未來應(yīng)用
數(shù)據(jù)流
在本實驗的場景中，32個節(jié)點使用最簡單的All-Reduce（全連接），每個節(jié)點平均傳輸86.8MB（2.8MB×31），并接收相同數(shù)量的數(shù)據(jù)。
如果使用專用服務(wù)器進行數(shù)據(jù)聚合，則每個節(jié)點只需上傳2.8MB數(shù)據(jù)（接收數(shù)據(jù)不變），通信量進一步減少。
另外，不對稱性是有優(yōu)點的，因為大多數(shù)消費互聯(lián)網(wǎng)的帶寬嚴重偏向于更高的下載速度。
假設(shè)穩(wěn)定的網(wǎng)速為100Mbps下載和10Mbps上傳，則最壞情況下的延遲僅為下載6.94秒，上傳2.24秒，overlap一下則每步延遲為6.94秒。
ps：以上的數(shù)據(jù)傳輸都是原始的向量，如果用上壓縮技術(shù)還能更快。
帶寬
作者表示，目前的實驗和研究還比較有限，無法斷定隨著模型變大，帶寬減少的比率是會增加、減少還是保持不變。
不過目前的1.2B似乎是DisTrO能夠良好工作的最小尺寸（再小就不收斂了），所以可以假設(shè)隨著模型大小的增長，需要的通信會相對越來越少。
不過也可能通信量與模型大小沒有關(guān)系，這時可以在不增加通信帶寬的情況下增加模型大小，觀察更大的模型是否會改善訓練和學習的效果。
如果后一種情況屬實，那么未來GPU設(shè)計和制造的范式將會被改變（更大VRAM和更窄帶寬）。
恰好我們也更喜歡計算密集型負載（而不是I/O密集型），畢竟現(xiàn)在的帶寬要比計算貴得多。
聯(lián)邦學習
除了訓練LLM，DisTrO還能用來做什么？
在互聯(lián)網(wǎng)上做分布式訓練，讓人一下就想到了聯(lián)邦學習。
在允許模型協(xié)作訓練的同時，保持每個參與者的數(shù)據(jù)的私密性和去中心化，這在LLM被大公司掌握的當下，顯得越來越重要。
到目前為止，聯(lián)邦學習一直缺乏在有限的互聯(lián)網(wǎng)帶寬上訓練大型模型的有效方法。
而DisTrO對如何處理數(shù)據(jù)，或?qū)?shù)據(jù)分配給各個GPU節(jié)點沒有任何要求，并且可以無狀態(tài)（類似于聯(lián)邦平均），因此適用于聯(lián)邦學習的未來。
虛擬異構(gòu)GPU集群
此外，DisTrO可以創(chuàng)建一個完全去中心化且無需許可的網(wǎng)絡(luò)來協(xié)作和共享資源。
實驗表明，DisTrO對于訓練期間少量降級或丟棄的節(jié)點具有顯著的彈性，并且可以輕松地適應(yīng)新節(jié)點的加入。
在這種能力加持之下，一方面可以保障整個系統(tǒng)的安全性，降低不可信節(jié)點使用對抗性攻擊破壞運行的風險。
另一方面，也可以鼓勵機構(gòu)和個人靈活貢獻自己的計算資源，釋放潛在的算力。
甚至一些內(nèi)存或者算力不太夠的老卡，也能加入進來賺點外快，采用FSDP、SWARM Parallelism等策略與DisTrO協(xié)同工作。
能源
DisTrO的進一步大規(guī)模應(yīng)用，可能會緩解建設(shè)大型數(shù)據(jù)中心所帶來的能源消耗、基礎(chǔ)設(shè)施成本和土地使用等相關(guān)問題。
Llama 3.1項目需要構(gòu)建兩個大型整體超級集群，每個集群包含 24,000個H100 GPU，僅訓練過程就產(chǎn)生了相當于11,000噸的二氧化碳排放。
當今的LLM，除了模型參數(shù)大小的增長，訓練數(shù)據(jù)量也在不斷增大，導致AI相關(guān)的數(shù)據(jù)中心已經(jīng)摸到了現(xiàn)代電網(wǎng)的極限。
DisTrO可用于自適應(yīng)平衡多個使用過剩容量的小型模塊化數(shù)據(jù)中心，通過動態(tài)平衡訓練技術(shù)利用現(xiàn)有基礎(chǔ)設(shè)施，減輕訓練對環(huán)境的負面影響。
目前，DisTrO背后的理論還需要進一步探究，更嚴謹、更詳細的學術(shù)論文以及完整的代碼將在未來發(fā)布。
