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