今天小編分享的互聯網經驗：超越英偉達！字節跳動MegaScale如何實現大規模GPU的高效利用？，歡迎閲讀。
在 2024 年的 2 月末。字帖跳動發布了一篇論文，叫做《MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs》。标題中的 MegaScale 指的是是一個大語言模型的生產框架，類似于英偉達的開源框架 Megatron-LM。
英偉達這套框架主要是通過數據并行（data parallelism）、張量并行（tensor parallelism）、流水線并行（pipeline parallelism）這三點來提高大規模 GPU 的效率。我們曾在硅星人的文章《Gemini 背後，谷歌真正可怕之處并不在模型本身……》中科普到，谷歌提出了一個概念叫做 MFU，全稱為模型 FLOPs 利用率（Model FLOPs Utilization），數字越高代表訓練速度越快。
MegaScale 在處理 1750 億參數規模的語言模型時，在 12288 個 GPU 上實現了 55.2% 的 MFU，相比于 Megatron-LM 提高了大約 1.34 倍。為了保證訓練過程中的高效率和穩定性，團隊特别注重深入觀測性，并開發了一系列診斷工具，用于監控系統組件深層事件，識别問題根源，從而實現容錯性和減緩落後節點的影響。
當選擇了一個事件時，其依賴關系就會顯現出來
那麼接下來我們就來仔細拆解一下，説説 MegaScale 主要的實現方法以及這個 " 診斷工具 "。MegaScale 使用了一種叫做 "3D 并行 " 的設計。我們剛才説了，Megatron-LM 使用了三種并行手段，而 MegaScale 則是将這三種并行合在了一起。一句話概括一下這麼做的好處：減少 GPU 之間的通信消耗，進而增加并行計算能力。
在數據并行方面，MegaScale 将模型參數在數據并行度（data parallel rank）内的各個 worker 之間進行分發，每個 worker 持有模型的一份參數副本。在前向傳播階段，通過 all-gather 操作收集來自其他 worker 的最新模型參數。在反向傳播階段，使用 reduce-scatter 操作聚合各 worker 計算出的梯度。MegaScale 借鑑了 PyTorch FSDP 的做法，通過提前預取初始的 all-gather 操作，并讓它與數據加載操作重疊。這麼做最大的好處是減少了每次迭代中的通信時間。
在張量并行性和序列并行性中實現與并行 transformer 塊結合的重疊通信
所謂 all-gather 和 reduce-scatter，其實是分布式計算和深度學習中常用的通信原語，主要用于在并行環境中實現數據的聚合和分發。在深度學習訓練的背景下，all-gather 常用于數據并行場景中，當模型在多個 GPU 或節點上并行計算時，每個 GPU 計算完各自的梯度或者更新後的模型參數後，需要通過 all-gather 操作把所有節點上的梯度或參數聚合到每個節點上，以便後續進行全局梯度平均或參數同步。
那 reduce-scatter 則是用于梯度的聚合。當各個節點完成了本地的梯度計算後，通過 reduce-scatter 操作将所有節點的梯度匯總，然後将匯總得到的全局梯度均等地分散到所有節點，這樣每個節點都可以用同一個全局梯度來更新自己的模型參數。
可能你還會有點暈，不過别急，等我把這三種并行計算都梳理完，我會用很一個故事來解釋。在張量并行中，MegaScale 基于模型塊的基礎來實現一種通信重疊。單個設備可能托管多個模型片段，對每個片段獨立進行前向傳播和反向傳播，并在此基礎上安排通信操作。all-gather 操作在模型片段的前向傳播之前觸發，而 reduce-scatter 操作在其反向傳播之後開始。針對無法隐藏首次 all-gather 和最後一次 reduce-scatter 操作的情況，MegaScale 采用優化手段盡可能減少這部分開銷。
在流水線并行上，MegaScale 采用交錯 1F1B 調度方法，在熱身階段（warm-up phase），前向傳播僅依賴于之前的 receive 操作，因此可以将 send 和 receive 解耦合，打破原本可能存在的阻塞關系，使得 send 操作可以在等待 compute 的同時進行，實現計算與通信的重疊。到了冷卻階段（cool-down phase）則是熱身階段的逆過程，同樣可以應用類似的技術。
1F1B
接下來我換種簡單的方式解釋，想象一下，老師罰你抄課文，但是一個人抄太慢了，于是你把這本書抽成許多頁，讓你的同學們每個人抄一部分，這就是數據并行的基本概念。如果你發現一本書裏的某一頁太大，一個人都抄不完，這時你可以選擇橫向拆分這張大紙，每人抄一部分，再拼接回去，這就是張量并行。然後老師看你抄課文這麼快，于是罰你不僅要抄課文，還要畫插圖，而且不同環節的工作必須有序銜接。為此，你組織了一個流水線作業，每個人負責一個工序，比如一人打草稿，另一人描邊，第三人上色，這樣就構成了流水線并行。
而 3D 并行呢，就是你對此次抄課文事情的統籌規劃工作，讓系統能夠有效平衡計算資源的利用和通信成本，解決單個 GPU 無法承載超大規模模型的問題，從而在保證訓練效率的同時，大幅度提升了模型訓練的可擴展性。
再説説 MegaScale 的診斷工具，它用來記錄了每個機器排名在運行過程中關鍵代碼段的執行時間。與之前的工具（如 torch 分析器或 Megatron-LM 計時器）不同，這個診斷工具的作用是基于 GPU 裏面的 CUDA 事件方法計時事件。這種方法最大程度地減少了對 CUDA 同步的需求，從而防止了性能下降，且允許在生產訓練作業中始終運行它。該工具提供了兩種可視化模式，并可以從不同的角度分析收集到的數據。
顏色的深度表示在不同計算節點（rank）上代碼段的運行時間
工具記錄了在統一時間線上采集的事件流，當選擇特定事件時，顯示相關的依賴關系，尤其是在三 D 并行訓練場景中，可以直觀展現各 GPU 工作者間的邏輯拓撲結構、數據流向及其涉及的通信操作。若發生錯誤，可以直接查看選定 GPU 工作者的錯誤信息，從而輔助快速定位訓練異常。而當某個 GPU 工作者因為執行 NCCL 通信操作故障導致整個集群陷入阻塞狀态時，該工具通過讓每個 GPU 在通信超時時記錄自身的當前事件，據此構建基于 3D 并行設定邏輯拓撲的數據依賴可視化表示，迅速鎖定問題節點。
通過 CUDA 事件監控器記錄每個計算節點上關鍵代碼段的執行時間，對比分析不同機器之間的性能差異。該工具利用 CUDA 事件方法精準測量，規避了過多 CUDA 同步造成的性能損失，可直接應用于生產環境中的訓練任務。可視化模式如熱力圖清晰展示各設備間的計算階段耗時差異，幫助找出拖慢訓練速度的 " 落後者 "。
是的，即便是相同的 GPU，他們在不同的算法中表現是不同的。就好像同樣接受 9 年制義務教育，有人是考上了清華北大，而我只記得我怎麼被老師罰抄課文一樣。尤其是在超過 1 萬顆 GPU 并行的場合下，通過診斷工具能快速發現問題，進而提高整體并行計算的效率。
而且團隊考慮的非常周全，甚至連交換機都要 " 壓榨 "。研究團隊專門設計了一個網絡拓撲結構，并安排網絡流量以減少 ECMP 散列衝突。首先，在頂部交換機（ToR）級别，一個 400G 的下行端口被抽成兩個 200G 的下行端口，并使用特定的 AOC 電纜。由于每個上行鏈路的帶寬是下行鏈路的兩倍，衝突概率降低了。其次，伺服器上的八個 200G NIC 以多線方式連接到八個不同的交換機。通過相同的 ToR 交換機集合連接的 GPU 伺服器數量可以達到 64 個。我們策略性地安排了來自我們訓練任務的數據密集節點在同一個頂部交換機（ToR）下運行。這種方法顯著減少了通信所需的交換機跳數，并進一步減少了 ECMP 散列衝突的概率。
MegaScale 帶給我們怎樣的啓示？
從設計邏輯上來説，MegaScale 的确是很有新意。同時他們也驗證了 12288 個 GPU 的情況下，框架的運行效率遠遠超過 2020 年的 Megatron-LM。但問題就在于，我們熟知的大語言模型，比如 ChatGPT4，它至少得有 20000 個英偉達 A100，而 GPT5 估計大約要使用 50000 個 H100 的算力。MegaScale 如果想要獲得市場認可，團隊需要帶着它走出實驗室，感受實際環境的壓力。
Megatron-LM 在 2024 年 3 月 4 日仍在更新
此外 Megatron-LM 也好，DeepMind 的 Mesh-TensorFlow 也好，他們都是開源工具，而且 Megatron-LM 甚至在我寫這篇稿子的時候（2024 年 3 月 4 日），還在更新，内容是給 ChatGPT 模型添加 --rotary-base 參數（這個參數通常用于指定基準旋轉角度，可以指定旋轉操作相對于特定的基準角度進行，而不是相對于默認的基準角度，通過這個參數，可以指定旋轉操作相對于特定的基準角度進行，而不是相對于默認的基準角度）。MegaScale 目前還只是停留在實驗室閉源階段，沒有開源。當然，由于是 MegaScale 本身是由大學和企業共同開發的，這種涉及商業合作的項目往往需要知識產權保護，因此不會開源。
