AxoNN：透過四維混合並行計算推進大型語言模型訓練

深度神經網路（DNN）的訓練因大型語言模型（LLM）和生成式人工智慧的興起而經歷了前所未有的增長。這些模型的有效性與其規模成正比，這一發展得益於GPU技術的進步以及像PyTorch和TensorFlow這樣的框架。然而，訓練擁有數十億參數的神經網路面臨著重大技術挑戰，因為這些模型超出了單個GPU的處理能力。這需要將模型分散到多個GPU上，並平行化矩陣乘法操作。訓練效率受到多種因素的影響，包括持續的計算性能、子通信者之間的集體通信操作，以及計算與非阻塞集體操作的重疊。

最近，訓練LLM的努力推動了基於GPU的集群利用的邊界，儘管效率仍然是一個挑戰。Meta使用2000個NVIDIA A100 GPU訓練了Llama 2，而Megatron-LM的管道並行性在對3072個GPU進行1000B參數模型基準測試時達到了52%的峰值性能。Megatron-LM和DeepSpeed的組合在對4480個A100 GPU訓練530B參數模型時達到了36%的峰值性能。MegaScale在12288個A100 GPU上對175B參數模型達到了55.2%的峰值性能。在AMD系統上，FORGE訓練在2048個MI250X GPU上達到了28%的峰值性能，而其他研究在1024個MI250X GPU上對1T參數模型的基準測試中達到了31.96%的峰值。

來自美國馬里蘭大學（University of Maryland, College Park）、德國馬克斯·普朗克智能系統研究所（Max Planck Institute for Intelligent Systems, Tübingen）和美國加州大學伯克利分校（University of California, Berkeley）的研究人員提出了一種名為AxoNN的新型四維混合並行算法，這一算法在一個高度可擴展、可攜帶的開源框架中實現。研究人員在AxoNN中引入了幾項性能優化，以提高矩陣乘法內核的性能，有效地將非阻塞集體操作與計算重疊，並使用性能建模來識別最佳配置。除了性能之外，他們還解決了由於LLM訓練數據記憶而引發的隱私和版權問題，調查了“災難性記憶化”。一個擁有4050億參數的LLM在Frontier上使用AxoNN進行了微調。

AxoNN在三個領先的超級計算平台上進行評估：NERSC/LBL的Perlmutter搭載NVIDIA A100 GPU（每個40GB DRAM）、OLCF/ORNL的Frontier搭載AMD Instinct MI250X GPU（每個128GB DRAM，分為兩個獨立管理的64GB圖形計算晶片），以及CSCS的Alps搭載GH200超級晶片（每個H100 GPU 96GB DRAM）。所有系統每個節點都使用四個HPE Slingshot 11網路介面卡，每個提供25 GB/s的雙向鏈路速度。性能測量遵循嚴格的方法論，運行十次迭代並平均最後八次，以考慮預熱變異性。基準測試是根據理論峰值性能值進行的，報告達到的峰值百分比和總持續的bf16 flop/s。

AxoNN在所有三個超級計算機上對GPT風格的變壓器顯示出卓越的弱擴展性能。在所有平台上，對於4096個GPU/GCD，幾乎達到了理想的擴展，涵蓋了大規模LLM訓練的典型硬體範圍。雖然在Alps的6144個H100 GPU上運行60B模型時，效率略微下降至76.5%，相比於1024個GPU的性能，Frontier的廣泛GPU可用性使得前所未有的擴展測試成為可能。AxoNN在Frontier上對8192個GCD保持了近乎完美的弱擴展，效率相對於512 GCD性能為88.3%。在Perlmutter上，AxoNN始終達到或超過每個GPU廣告的312 Tflop/s峰值的50%。線性性能擴展的證據是持續浮點運算的幾乎8倍增長，從512個GPU的80.8 Pflop/s增長到4096個GPU的620.1 Pflop/s。

總之，研究人員引入了AxoNN，其對機器學習的貢獻超越了性能指標，提供了可擴展、易於使用和可攜帶的模型並行性訪問。它使在普通計算約束下訓練和微調更大模型成為可能，允許順序LLM訓練代碼庫有效利用分散資源。此外，通過使得在特定領域數據上微調大型模型的能力民主化，AxoNN擴展了各個領域從業者的能力。因此，理解和解決記憶風險的緊迫性日益增加，因為越來越多的研究人員現在可以處理前所未有的規模和複雜性的模型，這些模型可能無意中捕獲了訓練數據中的敏感信息。

查看論文。所有對這項研究的讚譽都歸功於這個項目的研究人員。此外，隨時在Twitter上關注我們，並別忘了加入我們的80k+ ML SubReddit。

🚨 推薦閱讀 – LG AI研究發布NEXUS：一個先進的系統，整合代理AI系統和數據合規標準，以解決AI數據集中的法律問題。