32卡176%訓練加速，開源大模型訓練框架Megatron-LLaMA來了

淘天集團和愛橙科技於9月12日正式發布了開源的大模型訓練框架－Megatron-LLaMA。該框架的目標是使技術開發者能夠更方便地提升大語言模型的訓練效能，降低訓練成本，並與LLaMA社群保持相容性。測試結果顯示，在32卡訓練上，與HuggingFace上直接取得的程式碼版本相比，Megatron-LLaMA能夠實現176%的加速；在大規模訓練上，Megatron-LLaMA幾乎線性擴展，且對網路不穩定表現出較高的容忍度。目前，Megatron-LLaMA已在開源社群上線

開源位址：https://github.com/alibaba/Megatron-LLaMA

#大語言模型的卓越表現一次又一次地超越了人們的想像。在過去幾個月裡，LLaMA和LLaMA2向開源社群全面開放，為那些想要訓練自己的大語言模型的人們提供了一個很好的選擇。在開源社群中，已經有許多基於LLaMA開發的模型，包括續訓/SFT（如Alpaca、Vicuna、WizardLM、Platypus、StableBegula、Orca、OpenBuddy、Linly、Ziya等）和從零開始訓練（如Baichuan、QWen 、InternLM、OpenLLaMA）的工作。這些工作不僅在各種大模型能力客觀評測榜單上表現出色，同時也展示了在長文本理解、長文本生成、代碼編寫、數學求解等實際應用場景中的出色性能。此外，業界也出現了許多有趣的產品，例如LLaMA結合Whisper的語音聊天機器人、LLaMA結合Stable Diffusion的繪畫軟體，以及醫學/法律領域的輔助諮詢機器人等

#儘管從HuggingFace 上可以拿到LLaMA 的模型代碼，但用自己的數據訓練一個LLaMA 模型對個人用戶或中小型組織並不是一件低成本且簡單的工作。大模型的體積和資料的規模，使得在普通的計算資源上無法完成有效的訓練，算力和成本成為嚴重的瓶頸。 Megatron-LM 社群的使用者在這方面的訴求非常急迫。

淘天集團和愛橙科技在大模型應用上有著非常廣闊應用場景，在大模型的高效訓練上進行了非常多的投入。 LLaMA 的問世，在數據處理、模型設計、微調及強化學習反饋調整等方面都給予了包括淘天集團和愛橙科技在內的許多公司非常多的啟示，也助力業務應用場景取得了新的突破。因此，為了回饋整個LLaMA 開源社群、促進中文預訓練大模型開源社群的發展，讓開發者們能夠更方便地提升大語言模型的訓練性能，降低訓練成本，淘天集團聯合愛橙科技將部分內部優化技術開源，發表Megatron-LLaMA，期望與每位合作夥伴共建Megatron 及LLaMA 生態。

Megatron-LLaMA 提供了一套標準的Megatron-LM 實現的LLaMA，並提供了##與HuggingFace 格式自由切換的工具，方便與社區生態工具相容。 Megatron-LLaMA 重新設計了Megatron-LM 的反向流程，使得無論在節點數較少需要開較大梯度聚合（Gradient Accumulation, GA）、或是節點數較多必須使用小GA 的場景，都能夠取得卓越的訓練表現。

• 在32 卡訓練上，相較於HuggingFace 上直接取得的程式碼版本，Megatron-LLaMA 能夠取得176% 的加速；即便是採用DeepSpeed 及FlashAttention 優化過的版本，Megatron-LLaMA 仍然能減少至少19% 的訓練時間。
• 在大規模的訓練上，Megatron-LLaMA 相比較 32 卡擁有著幾乎線性的擴展性。例如使用 512 張 A100 復現 LLaMA-13B 的訓練，Megatron-LLaMA 的反向機制相對於原生 Megatron-LM 的 DistributedOptimizer 能夠節約至少兩天的時間，且沒有任何精確度損失。
• Megatron-LLaMA 對網路不穩定表現出高容忍度。即使是現在性價比較高的4x200Gbps 通訊頻寬的8xA100-80GB 訓練集群（這種環境通常是混部環境，網路只能使用一半的頻寬，網路頻寬是嚴重的瓶頸，但租用價格相對低廉）上， Megatron-LLaMA 仍能取得0.85 的線性擴展能力，然而在這個指標上Megatron-LM 僅能達到不足0.7。

Megatron-LM 技術帶來的高效能LLaMA 訓練機會

LLaMA是目前大語言模型開源社群中一項重要的工作。 LLaMA在LLM的結構中引入了BPE字元編碼、RoPE位置編碼、SwiGLU活化函數、RMSNorm正則化以及Untied Embedding等最佳化技術，在許多客觀和主觀評測中取得了卓越的效果。 LLaMA提供了7B、13B、30B、65B/70B的版本，適用於各類大模型需求的場景，也受到廣大開發者的青睞。和許多其他開源大模型一樣，由於官方只提供了推理版的程式碼，如何以最低成本開展高效訓練，並沒有一個標準的範式

Megatron-LM 是一種優雅的高性能訓練解決方案。 Megatron-LM 中提供了張量並行（Tensor Parallel，TP，把大乘法分配到多張卡並行計算）、流水線並行（Pipeline Parallel，PP，把模型不同層分配到不同卡處理）、序列並行（Sequence Parallel， SP，序列的不同部分由不同卡處理，節約顯存）、DistributedOptimizer 優化（類似DeepSpeed Zero Stage-2，切分梯度和優化器參數至所有計算節點）等技術，能夠顯著減少顯存佔用並提升GPU利用率。 Megatron-LM 經營著一個活躍的開源社區，持續有新的最佳化技術、功能設計合併進框架中。

然而，開發基於Megatron-LM並不簡單，特別是在昂貴的多卡機器上進行調試和功能驗證更是十分昂貴的。 Megatron-LLaMA首先提供了一套基於Megatron-LM框架實現的LLaMA訓練程式碼，支援各種規模的模型版本，並且可以簡單地適配支援LLaMA的各種變種，包括直接支援HuggingFace格式的Tokenizer。因此，Megatron-LLaMA可以輕鬆應用於現有的離線訓練連結中，無需進行過多的適應。在中小規模訓練/微調LLaMA-7b和LLaMA-13b的場景中，Megatron-LLaMA能夠輕鬆達到業界領先的54%以上的硬體利用率（MFU）

#Megatron- LLaMA 的反向流程優化

需要重新書寫的內容是：圖示：DeepSpeed ZeRO 階段-2

#DeepSpeed ZeRO 是微軟推出的一套分散式訓練框架，其中提出的技術對許多後來的框架都有非常深遠的影響。 DeepSpeed ZeRO Stage-2（後文簡稱 ZeRO-2）是此框架中一項節約顯存佔用且不增加額外運算量與通訊量的技術。如上圖所示，由於計算需要，每個 Rank 都需要擁有全部的參數。但對於優化器狀態而言，每個 Rank 只負責其中的一部分即可，不必所有 Rank 同時執行完全重複的操作。於是ZeRO-2 提出將優化器狀態均勻地切分在每個Rank 上（注意，這裡並不需要保證每個變數被均分或完整保留在某個Rank 上），每個Rank 在訓練過程中只負責對應部分的最佳化器狀態和模型參數的更新。在這種設定下，梯度也可以以此方式進行切分。預設情況下，ZeRO-2 在反向時在所有Rank 間使用Reduce 方式聚合梯度，而後每個Rank 只需要保留自身所負責的參數的部分，既消除了冗餘的重複計算，又降低了顯存佔用。

Megatron-LM DistributedOptimizer

#

原生 Megatron-LM 透過 DistributedOptimizer 實現了類似 ZeRO-2 的梯度和優化器狀態切分，以減少訓練中的顯存佔用。如上圖所示，DistributedOptimizer 在每次獲得預設的梯度聚合過的所有梯度後，透過 ReduceScatter 算子，將先前累積的全部梯度分發到不同的 Rank。每個 Rank 只獲得自己需要處理的部分梯度，而後進行優化器狀態的更新和對應參數的更新。最後各個 Rank 透過 AllGather 的方式從其他節點上取得更新的參數，最終取得全部的參數。實際訓練的結果顯示，Megatron-LM 的梯度和參數通訊與其他運算串列進行，對於大規模預訓練任務，為了確保總批資料大小不變，通常無法開啟較大的 GA。於是通訊佔比會伴隨機器增加上升，這時候串行通訊的特性導致擴展性很弱。在社群內，這方面的需求也很迫切

Megatron-LLaMA OverlappedDistributedOptimizer

#為了解決這個問題，Megatron-LLaMA 改進了原生Megatron-LM 的DistributedOptimizer，使其梯度通訊的算符能夠可以和計算相並行。特別的，相較於 ZeRO 的實現，Megatron-LLaMA 在並行的前提下，透過巧妙的優化優化器分區策略，使用了更具有擴展性的集合通訊方式來提升擴展性。 OverlappedDistributedOptimizer 的主要設計保證瞭如下幾點：a) 單一集合通訊算子資料量足夠大，充分利用通訊頻寬；b) 新切分方式所需通訊資料量應等於資料並行所需的最小通訊資料量； c) 完整參數或梯度與切分後的參數或梯度的轉換過程中，無法引入過多顯存拷貝。

具體而言，Megatron-LLaMA對DistributedOptimizer進行了改進，提出了OverlappedDistributedOptimizer，用於在訓練的反向流程中結合新的切分方式進行最佳化。如圖所示，在初始化OverlappedDistributedOptimizer時，會預先為所有參數指派它們所屬的Bucket。每個Bucket中的參數是完整的，一個參數只屬於一個Bucket，一個Bucket中可能有多個參數。在邏輯上，每個Bucket將被連續地等分為P份（P為資料並行組的數量），資料並行組中的每個Rank將負責其中的一份

Bucket被放置在本地梯度桶佇列中，以確保通訊順序。在進行訓練計算的同時，資料並行組之間以桶為單位，透過集合通訊交換各自所需的梯度。在Megatron-LLaMA中，Bucket的實作盡可能地採用了位址索引，只有在需要變更值時才會分配新的空間，以避免顯存浪費

透過結合大量的工程優化，上述設計使得在大規模訓練時，Megatron-LLaMA能夠充分利用硬件，實現比原生的Megatron-LM更好的加速效果。在常用的網路環境中，將訓練規模從32張A100卡擴展到512張A100卡，Megatron-LLaMA仍能夠達到0.85的擴展比

Megatron-LLaMA 的未來計劃

Megatron-LLaMA是一個由淘天集團和愛橙科技共同開源並提供後續維護支援的訓練框架，在內部已經廣泛應用。隨著越來越多的開發者加入LLaMA的開源社群並貢獻可供互鑑的經驗，相信未來在訓練框架層面會面臨更多挑戰和機會。 Megatron-LLaMA將密切關注社區的發展，並與廣大開發者共同推動以下方向的發展：

• #自適應最優配置選擇
• ##更多模型結構或局部設計變更的支援
• 在更多不同類別硬體環境下的極致效能訓練解決方案

#專案地址：https://github.com/alibaba/Megatron-LLaMA

以上是32卡176%訓練加速，開源大模型訓練框架Megatron-LLaMA來了的詳細內容。更多資訊請關注PHP中文網其他相關文章！