随着對 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