淘天集團與愛橙科技合作發表開源大型模型訓練架構Megatron-LLaMA

9 月12 日，淘天集團聯合愛橙科技正式對外開源大模型訓練框架－Megatron-LLaMA，旨在讓技術開發者們能夠更方便的提升大語言模型訓練性能，降低訓練成本，並且保持和LLaMA 社區的兼容性。測試顯示，在32 卡訓練上，相較於HuggingFace 上直接獲得的程式碼版本，Megatron-LLaMA 能夠取得176% 的加速；在大規模的訓練上，Megatron-LLaMA 相比較32 卡擁有幾乎線性的擴展性，而且對網路不穩定表現出高容忍度。目前 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 社區開源上
  ##L#L#L#L#L#L#L#L#11.中一項重要工作。 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 Stage-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 被放置在一個本機佇列（Local grad bucket queue）中，從而確保通訊順序。在訓練計算的同時，資料並行組間以 Bucket 為單位，透過集合通訊交換各自所需的梯度。 Megatron-LLaMA 中 Bucket 的實作盡可能採用了位址索引，只在有需要值變更時才新分配空間，避免了顯存浪費。
  
  
  上述的設計，再結合大量的工程優化，使得在大規模訓練時，Megatron-LLaMA 可以很充分地使用硬件，實現了比原生Megatron-LM更好的加速。從32張A100卡擴展到512張A100卡的訓練，Megatron-LLaMA在常用混部的網路環境中仍能取得0.85的擴展比。
  
                                             Megatron-中最佳的未來計畫上#3) 未來計畫中提供愛#3/Hega-Tegas#Luno#SMega 的未來計畫#Vegacis 的未來計畫#3 月維護支援的訓練框架，在內部已有廣泛的應用。隨著越來越多的開發者湧入 LLaMA 的開源社群並貢獻可以相互借鑒的經驗，相信未來在訓練框架層面會有更多的挑戰和機會。 Megatron-LLaMA 將會緊密關注社群的發展，並與廣大開發者共同推進以下方向：
  
  
  自適應最優配置選擇
更多模型結構或局部設計改變的支援
• 在更多不同類別硬體環境下的極致效能訓練解決方案
• 專案位址：https://github.com/alibaba/ Megatron-LLaMA
• #

以上是淘天集團與愛橙科技合作發表開源大型模型訓練架構Megatron-LLaMA的詳細內容。更多資訊請關注PHP中文網其他相關文章！