字節萬卡集群技術細節公開:2天搞定GPT-3訓練,算力利用率超英偉達Megatron-LM
隨著對Sora技術分析的展開,AI基礎設施的重要性愈發凸顯。
來自字節和北大的一篇新論文在此時吸引關注:
文章披露,字節搭建起的萬卡集群,能在1.75天內完成GPT-3規模模型(175B)的訓練。
具體來說,字節提出了一個名為MegaScale的生產系統,旨在解決在萬卡集群上訓練大模型時面臨的效率和穩定性挑戰。
在12288塊GPU上訓練1750億參數大語言模型時,MegaScale實現了55.2%的算力利用率(MFU),是英偉達Megatron-LM的1.34倍。
論文還透露,截止2023年9月,字節已建立起超過1萬張卡的Ampere架構GPU(A100/A800)集群,目前正在建設大規模Hopper架構(H100/H800)集群。
適用于萬卡集群的生產系統
大模型時代,GPU的重要性已無需贅述。
但大模型的訓練,并不是把卡的數量拉滿就能直接開干的——當GPU集群的規模來到“萬”字級別,如何實現高效、穩定的訓練,本身就是一個頗具挑戰的工程問題。
第一重挑戰:效率。
訓練大語言模型并非簡單的并行任務,需要在多個GPU之間分布模型,并且這些GPU需要頻繁通信才能共同推進訓練進程。通信之外,操作符優化、數據預處理和GPU內存消耗等因素,都對算力利用率(MFU)這個衡量訓練效率的指標有影響。
MFU是實際吞吐量與理論最大吞吐量之比。
第二重挑戰:穩定性。
我們知道,訓練大語言模型往往需要花費非常長的時間,這也意味著,訓練過程中失敗和延遲的現象并不鮮見。
失敗的成本是高昂的,因此如何縮短故障恢復時間變得尤為重要。
為了應對這些挑戰,字節跳動的研究人員構建了MegaScale,并已將其部署到字節的數據中心中,用以支持各種大模型的訓練。
MegaScale是在英偉達Megatron-LM的基礎上改進的。
具體改進包括,算法和系統組件的共同設計、通信和計算重疊的優化、操作符優化、數據流水線優化以及網絡性能調優等:
- 算法優化:研究人員在模型架構中引入并行化的Transformer塊、滑動窗口注意力機制(SWA)和LAMB優化器,來提高訓練效率而不犧牲模型的收斂性。
- 通信重疊:基于對3D并行(數據并行、流水線并行、張量并行)中各個計算單元操作的具體分析,研究人員設計技術策略有效地減少了非關鍵執行路徑上操作所帶來的延遲,縮短了模型訓練中每一輪的迭代時間。
- 高效操作符:對GEMM操作符進行了優化,對LayerNorm和GeLU等操作進行了融合,以減少啟動多個內核的開銷,并優化內存訪問模式。
- 數據流水線優化:通過異步數據預處理和消除冗余的數據加載器,來優化數據預處理和加載,減少GPU空閑時間。
- 集體通信群初始化:優化了分布式訓練中英偉達多卡通信框架NCCL初始化的過程。在未經優化的情況下,2048張GPU的集群初始化時間是1047秒,優化后可降至5秒以下;萬卡GPU集群的初始化時間則可降至30秒以下。
- 網絡性能調優:分析了3D并行中的機器間流量,設計技術方案提高網絡性能,包括網絡拓撲設計、減少ECMP哈希沖突、擁塞控制和重傳超時設置。
- 故障容忍:在萬卡集群中,軟硬件故障難以避免。研究人員設計了一個訓練框架,來實現自動故障識別和快速恢復。具體包括,開發診斷工具來監控系統組件和事件、優化checkpoint高頻保存訓練進程等。
論文提到,MegaScale能夠自動檢測和修復超過90%的軟硬件故障。
實驗結果表明,MegaScale在12288個GPU上訓練175B大語言模型時,實現了55.2%的MFU,是Megatrion-LM算力利用率的1.34倍。
訓練530B大語言模型的MFU對比結果如下:
One More Thing
就在這篇技術論文引發討論之際,字節類Sora產品也傳出了新消息:
剪映旗下類似Sora的AI視頻工具已經啟動邀請內測。
看樣子地基已經打好,那么對于字節的大模型產品,你期待嗎?
論文地址:https://arxiv.org/abs/2402.15627