9 月 12 日，淘天集團聯(lián)合愛橙科技正式對外開源大模型訓練框架 ——Megatron-LLaMA，旨在讓技術(shù)開發(fā)者們能夠更方便的提升大語言模型訓練性能，降低訓練成本，并且保持和 LLaMA 社區(qū)的兼容性。測試顯示，在 32 卡訓練上，相比 HuggingFace 上直接獲得的代碼版本，Megatron-LLaMA 能夠取得 176% 的加速；在大規(guī)模的訓練上，Megatron-LLaMA 相比較 32 卡擁有幾乎線性的擴展性，而且對網(wǎng)絡不穩(wěn)定表現(xiàn)出高容忍度。目前 Megatron-LLaMA 已在開源社區(qū)上線。

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

大語言模型的卓越表現(xiàn)一次又一次地超出了人們的想象。過去幾個月，LLaMA 及 LLaMA2 向開源社區(qū)的全面放開，使想訓練一個自己的大語言模型的人們多了一種很好的選擇。在開源社區(qū)中，已有非常多的基于 LLaMA 產(chǎn)出的模型，包括進行續(xù)訓 / SFT（如 Alpaca、Vicuna、WizardLM、Platypus、StableBegula、Orca、OpenBuddy、Linly、Ziya 等）和從零開始訓練（Baichuan、QWen、InternLM、OpenLLaMA）的工作。這些工作不僅在目前各種大模型能力客觀評測榜單上表現(xiàn)卓越，同時展現(xiàn)了在長文本理解、長文本生成、代碼編寫、數(shù)學求解等實用場景的優(yōu)秀性能。更進一步，業(yè)界出現(xiàn)了很多有意思的產(chǎn)品，例如 LLaMA 結(jié)合 Whisper 的語音聊天機器人、LLaMA 結(jié)合 Stable Diffusion 的繪畫軟件、醫(yī)學 / 法律領域的輔助咨詢機器人等。

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

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

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

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

Megatron-LM 技術(shù)帶來的高性能 LLaMA 訓練機會

LLaMA 是目前大語言模型開源社區(qū)中一項重要工作。LLaMA 在 LLM 的結(jié)構(gòu)中引入了 BPE 字符編碼、RoPE 位置編碼、SwiGLU 激活函數(shù)、RMSNorm 正則化以及 Untied Embedding 等優(yōu)化技術(shù)，在許多客觀和主觀評測中取得了卓越的效果。LLaMA 提供了 7B、13B、30B、65B/70B 的版本，適用于各類大模型需求的場景，也受到廣大開發(fā)者的青睞。同諸多開源大模型一樣，由于官方只提供了推理版的代碼，如何以最低成本開展高效訓練，并沒有一個標準的范式。

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

然而，基于 Megatron-LM 進行開發(fā)并不簡單，在昂貴的多卡機上調(diào)試及功能性驗證更是十分昂貴的。Megatron-LLaMA 首先提供了一套基于 Megatron-LM 框架實現(xiàn)的 LLaMA 訓練代碼，支持各種規(guī)模的模型版本，并且可以很簡單地適配支持 LLaMA 的各類變種，包括對 HuggingFace 格式的 Tokenizer 的直接支持。于是，Megatron-LLaMA 可以很便捷地應用在已有的離線訓練鏈路中，無需進行過多的適配。在中小規(guī)模訓練 / 微調(diào) LLaMA-7b 和 LLaMA-13b 的場景，Megatron-LLaMA 能夠輕松達到業(yè)界領先的 54% 及以上的硬件利用率（MFU）。

Megatron-LLaMA 的反向流程優(yōu)化

圖示：DeepSpeed ZeRO Stage-2

DeepSpeed ZeRO 是微軟推出的一套分布式訓練框架，其中提出的技術(shù)對很多后來的框架都有非常深遠的影響。DeepSpeed ZeRO Stage-2（后文簡稱 ZeRO-2）是該框架中一項節(jié)約顯存占用且不增加額外計算量和通信量的技術(shù)。如上圖所示，由于計算需要，每個 Rank 都需要擁有全部的參數(shù)。但對于優(yōu)化器狀態(tài)而言，每個 Rank 只負責其中的一部分即可，不必所有 Rank 同時執(zhí)行完全重復的操作。于是 ZeRO-2 提出將優(yōu)化器狀態(tài)均勻地切分在每個 Rank 上（注意，這里并不需要保證每個變量被均分或完整保留在某個 Rank 上），每個 Rank 在訓練進程中只負責對應部分的優(yōu)化器狀態(tài)和模型參數(shù)的更新。在這種設定下，梯度也可以按此方式進行切分。默認情況下，ZeRO-2 在反向時在所有 Rank 間使用 Reduce 方式聚合梯度，而后每個 Rank 只需要保留自身所負責的參數(shù)的部分，既消除了冗余的重復計算，又降低了顯存占用。

Megatron-LM DistributedOptimizer

原生 Megatron-LM 通過 DistributedOptimizer 實現(xiàn)了類似 ZeRO-2 的梯度和優(yōu)化器狀態(tài)切分，以減少訓練中的顯存占用。如上圖所示，DistributedOptimizer 在每次獲得預設的梯度聚合過的所有梯度后，通過 ReduceScatter 算子，將之前累積的全部梯度分發(fā)到不同的 Rank。每個 Rank 只獲得自己需要處理的部分梯度，而后進行優(yōu)化器狀態(tài)的更新和對應參數(shù)的更新。最后各個 Rank 通過 AllGather 的方式從其他節(jié)點上獲取更新過的參數(shù)，最終取得全部的參數(shù)。實際訓練的結(jié)果顯示，Megatron-LM 的梯度和參數(shù)通信與其他計算串行進行，對于大規(guī)模預訓練任務，為了保證總批數(shù)據(jù)大小不變，通常無法開啟較大的 GA。于是通信占比會伴隨機器增加上升，這時候串行通信的特點導致擴展性很弱。在社區(qū)內(nèi)，這方面的需求也很迫切。

Megatron-LLaMA OverlappedDistributedOptimizer

為了解決這一問題，Megatron-LLaMA 改進了原生 Megatron-LM 的 DistributedOptimizer，使其梯度通信的算子能夠可以和計算相并行。特別的，相比于 ZeRO 的實現(xiàn)，Megatron-LLaMA 在并行的前提下，通過巧妙的優(yōu)化優(yōu)化器分區(qū)策略，使用了更具有具有擴展性的集合通信方式來提升擴展性。OverlappedDistributedOptimizer 的主要設計保證了如下幾點：a) 單一集合通信算子數(shù)據(jù)量足夠大，充分利用通信帶寬；b) 新切分方式所需通信數(shù)據(jù)量應等于數(shù)據(jù)并行所需的最小通信數(shù)據(jù)量；c) 完整參數(shù)或梯度與切分后的參數(shù)或梯度的轉(zhuǎn)化過程中，不能引入過多顯存拷貝。

具體而言，Megatron-LLaMA 改進了 DistributedOptimizer 的機制，提出了 OverlappedDistributedOptimizer，用于結(jié)合新的切分方式優(yōu)化訓練中的反向流程。如上圖所示，在 OverlappedDistributedOptimizer 初始化時，會預先給所有參數(shù)分配其所屬的 Bucket。Bucket 中的參數(shù)是完整的，一個參數(shù)僅屬于一個 Bucket，一個 Bucket 中可能有多個參數(shù)。邏輯上，每個 Bucket 將會被連續(xù)等分成 P（P 為數(shù)據(jù)并行組的數(shù)量）等份，數(shù)據(jù)并行組中的每個 Rank 負責其中的一份。

Bucket 被放置在一個本地隊列（Local grad bucket queue）中，從而保證通信順序。在訓練計算的同時，數(shù)據(jù)并行組間以 Bucket 為單位，通過集合通訊交換各自需要的梯度。Megatron-LLaMA 中 Bucket 的實現(xiàn)盡可能采用了地址索引，只在有需要值更改時才新分配空間，避免了顯存浪費。

上述的設計，再結(jié)合大量的工程優(yōu)化，使得在大規(guī)模訓練時，Megatron-LLaMA 可以很充分地使用硬件，實現(xiàn)了比原生 Megatron-LM 更好的加速。從32張A100卡擴展到512張A100卡的訓練，Megatron-LLaMA在常用混部的網(wǎng)絡環(huán)境中仍然能夠取得0.85的擴展比。

Megatron-LLaMA 的未來計劃

Megatron-LLaMA 是由淘天集團和愛橙科技共同開源并提供后續(xù)維護支持的訓練框架，在內(nèi)部已有廣泛的應用。隨著越來越多的開發(fā)者涌入 LLaMA 的開源社區(qū)并貢獻可以相互借鑒的經(jīng)驗，相信未來在訓練框架層面會有更多的挑戰(zhàn)和機會。Megatron-LLaMA 將會緊密關(guān)注社區(qū)的發(fā)展，并與廣大開發(fā)者共同推進以下方向：

• 自適應最優(yōu)配置選擇
• 更多模型結(jié)構(gòu)或局部設計改動的支持
• 在更多不同類硬件環(huán)境下的極致性能訓練解決方案

項目地址：https://github.com/alibaba/Megatron-LLaMA