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

作者：機(jī)器之心 2023-09-12 13:43:00

大語言模型的卓越表現(xiàn)一次又一次地超出了人們的想象。過去幾個(gè)月，LLaMA 及 LLaMA2 向開源社區(qū)的全面放開，使想訓(xùn)練一個(gè)自己的大語言模型的人們多了一種很好的選擇。

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

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

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

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

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

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

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

Megatron-LM 技術(shù)帶來的高性能 LLaMA 訓(xùn)練機(jī)會(huì)

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

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

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

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

圖示：DeepSpeed ZeRO Stage-2

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

Megatron-LM DistributedOptimizer

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

Megatron-LLaMA OverlappedDistributedOptimizer

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

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

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

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

Megatron-LLaMA 的未來計(jì)劃

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