NVIDIA Megatron-LM 于 2019 年首次推出，在 AI 社區引發了一波創新浪潮，使研究人員和開發者能夠利用該開源庫的基礎進一步推動大型語言模型（LLM）的進步。如今，許多最流行的 LLM 開發者框架都受到 Megatron-LM 庫的啟發并使用其構建，從而引發了一系列基礎模型和 AI 初創公司的浪潮。基于 Megatron-LM 構建的一些最流行的 LLM 框架包括?Colossal-AI,?Hugging Face Accelerate, 和?NVIDIA NeMo。
為便于輕松遷移，并使研究人員和模型開發者能夠訪問分布式訓練中的最新研究成果，NVIDIA 最近對 Megatron-LM 進行了改進。這樣便推出了 NVIDIA Megatron-Core，一個基于 PyTorch 的開源庫，包含一系列 GPU 優化技術、先進的系統級創新和模塊化 API，以便在大規模訓練模型。
Megatron-Core 繼續推進大規模分布式訓練。本文將重點介紹一些近期的進展，其中包括用于多模態訓練的新的大型語言和視覺助手（LLaVA）流程。
NVIDIA Megatron-Core
Megatron-Core 包含具有先進系統級創新技術的 GPU 優化技術。它將這些技術抽象為可組合和模塊化的 API，為框架開發者和研究人員提供充分的靈活性，以便在 NVIDIA 加速計算基礎設施上大規模訓練自定義的 Transformer。
Megatron-Core 庫提供 Transformer 模型的核心構建塊，例如注意力機制、Transformer 塊和層、歸一化層和嵌入技術。此外，該庫還原生內置了其他功能，包括激活重新計算和分布式檢查點。
可以使用 Megatron-Core 在大規模計算中高效構建熱門 LLM 架構，如 GPT、BERT、T5 和 RETRO。此外，Megatron-Core 與所有 NVIDIA Tensor Core GPU 兼容，并且可以利用 NVIDIA Hopper 架構 支持的 FP8 數據格式，進一步提高計算吞吐量并減少顯存占用。借助 Megatron-Core，客戶如 Reka AI 和 Codeium 能夠大規模訓練模型。
“Megatron-Core 的模塊化、可組合設計無縫集成到我們的多模態 LLM 架構中，”Reka AI 的技術人員 Deyu Fu 說，“它借助優化的 GPU 內核和并行技術，幫助我們輕松處理超大型模型和廣泛的上下文。同時，它支持密集和稀疏訓練，以便在集群級別高效擴展。”
Codeium 的機器學習工程師 Devin Chotzen-Hartzell 說：“通過使用 Megatron-Core，我們只需在庫中打開一個標志，就能夠始終處于訓練大型語言模型的技術前沿。這使我們能夠專注于數據和對齊方面的差異化。”
Megatron-Core 現已支持多模態訓練
隨著視覺指令調整的引入，大型多模態模型引起了研究人員和行業的廣泛興趣。這些模型利用各種類型的數據來生成全面的、上下文感知的響應，使用多個感官輸入來了解環境并與環境進行交互。這一進步使生成式 AI 模型更接近人類處理世界的方式。
我們很高興地宣布，Megatron-Core v0.7 現已支持多模態。如需了解使用 LLaVA 的完整多模態參考工作流，請訪問 GitHub 上的 NVIDIA/Megatron-LM。模型開發者可以使用 Megatron 下的開源 多模態數據加載器，輕松混合具有確定性和可再現性的多模態數據集。同時，這也適用于檢查點保存和加載。
LLaVA 工作流示例 將向您展示如何：
- 為基于 Megatron 網絡數據集格式，準備預訓練和監督式微調（SFT）數據集。
- 利用 Megatron Core 的并行性和內存節省技術來訓練一個從 Mistral 和 CLIP 初始化的 LLaVA 架構模型。
- 使用不同的任務（例如 COCO 字幕和 VQAv2）來進行評估。
Megatron-Core v0.7 版本專注于 LLaVA 管線的功能方面，其 Massive Multidiscipline Multimodal Understanding（MMMU） 評分為 38，這在基于 7B 參數的 LLM-based LLaVA 架構的預期范圍內。
此外，借助 Megatron-Core (Mcore) 規格系統，研究人員可以輕松自定義 PyTorch 模型定義中的子模塊。在接下來的幾個版本中，Megatron-Core 將支持針對不同模型使用異構并行策略，這對視覺模型特別有用，因為它們通常較小，需要的分片技術不那么復雜，相比于多模態訓練中的大型語言模型。
包括多模態數據加載程序在內的所有 Megatron-Core 多模態訓練功能將很快集成到 NVIDIA NeMo 中，以增強 NeMo 中當前用于 NeVa 等模型的多模態功能。
針對專家混合模型的訓練吞吐量優化
在快速發展的生成式 AI 領域，混合專家（MoE）模型已成為一種有吸引力的選擇，因為它們可以預訓練以在不增加浮點運算次數的情況下實現更高的準確性。在 MoE 模型中，密集 FFN 層被替換為 MoE 層，其中每個令牌都會路由到路由器選擇的幾個專家。
Megatron-Core v0.7 擴展了 MoE 功能，并添加了各種訓練速度和內存優化，使其成為大規模訓練 MoE 的最全面的解決方案。具體來說，現在 Megatron-Core 支持在 GShard 中使用令牌丟棄的 MoE 訓練，并具有多種訓練速度優化，包括具有多 CUDA 流計算和梯度累積融合的增強型 GroupedGEMM。
表 1 顯示，在 BF16 精度下訓練時，Megatron-Core 可實現超過 400 TFLOP/s 的每個 GPU 吞吐量，且采用 all-to-all 分配器，序列長度為 4096。每個令牌都會路由到 (-moe-router-topk) 兩位專家。我們將繼續優化 MoE 的 FP8 方法，并將在即將發布的 Megatron-Core 版本中提供這些方法。
Megatron-Core 還支持 MoE 的專家并行性，這些并行性可以與 Megatron-Core 已支持的其他并行性技術（例如張量并行性、數據并行性、序列并行性和管道并行性）相結合。有關更多詳細信息，請參閱 用戶指南。
| 模型 | 精度 | GPU 數量 | MBS | GBS | TP | EP | PP | 梯度累加 | 每個 GPU 的吞吐量 (TFLOP/s/GPU) |
| Mistral 7B (密集模型基準) | BF16 | 128 | 4 | 256 | 2 | 不適用 | 1 | 1 | 492 |
| Mixtral 8x7B | BF16 | 128 | 1 | 256 | 1 | 8 | 4 | 8 | 402 |
快速分布式檢查點以提高訓練彈性
分布式檢查點對于在大規模訓練中保持彈性至關重要。PyTorch 原生解決方案 torch.save 通常缺乏效率和可擴展性，這導致開發更高效的解決方案。例如，Azure Nebula 和 AWS Gemini 提供異步檢查點，而 PyTorch 分布式檢查點（DCP） 使用線程保存每個秩的檢查點。這些方法通過利用并行和異步，速度比普通的 torch.save 更快。但是，實現高效異步檢查點仍然存在挑戰。
具體來說，這些解決方案可以在多 GPU 服務器中執行異步檢查點，或者通過使用 Python 線程（由于 Python 全局解釋器鎖定，這種線程效率低下）來執行異步檢查點。這會增加檢查點時間并降低訓練速度。這些解決方案還強制用戶加載與存儲檢查點相同的并行配置（例如 PP 和 TP 大小），從而防止在長時間訓練期間輕松進行動態并行重構。
Megatron-Core v0.7 通過引入完全并行和異步保存功能來解決這些問題。借助完全并行保存 (FPS)，數據并行副本執行并行寫入，從而更好地利用可用的文件系統帶寬。異步并行保存進一步加速分布式檢查點，首先將模型參數復制到 CPU (或未來的本地存儲)，然后將檢查點堅持到后臺的穩定存儲中，盡可能減少對主訓練過程的中斷。
最重要的是，Megatron-Core 允許用戶從使用不同張量并行度和管線并行度保存的檢查點恢復訓練，這樣在訓練期間可以根據需要靈活更改訓練配置。
Megatron-Core 中的保存和加載 API 與 PyTorch 原生 API 高度相似，這使得采用 Megatron-Core 分布式檢查點變得輕松。 在不使用分布式優化器的情況下，此改進可將 Nemotron-4 340B 的檢查點開銷減少 26 倍，并將 Nemotron-4 15B 的檢查點開銷減少 50 倍。 借助分布式優化器，用戶可以將 Nemotron-4 340B 的檢查點開銷減少 42 倍（圖 2）。
提高可擴展性
自 v0.5 版本以來，Megatron-Core 支持通過向后傳遞對數據并行性梯度 all-reduce 進行細粒度重疊。為實現此功能，將參數分組到數據桶。在數據桶參數的所有梯度準備就緒時，為數據桶啟動異步通信集合。這提高了 Megatron-Core 的吞吐量，因為它減少了公開的數據并行通信量。尤其是在每個 GPU 運行小批量配置 (并減少梯度累積) 時，這項功能尤其有用。
圖 3 顯示了弱擴展實驗中 Nemotron-4 15B 模型的每 GPU 吞吐量，其中批量大小隨著數據并行大小的增加而增加（全局批量大小為 3*data_parallel_size），張量并行大小為 8。我們觀察到，在數據并行大小為 32、批量大小為 96 時，此優化將吞吐量提高了 34%。使用數據并行時，--overlap-grad-reduce 標志可以啟用這種重疊技術。有關更多詳細信息，請參閱 Megatron-Core 文檔。
--overlap-grad-reduce 優化的效果。在 Megatron-Core v0.6 版本中，我們引入了分布式優化器支持，其中優化器狀態在數據并行副本上分割，從而減少峰值內存占用。分布式優化器還將之前需要的梯度 all-reduce 分解為梯度 reduce-scatter（RS）和參數 all-gather（AG）。Megatron-Core 通過將 reduce-scatter 與反向傳遞計算重疊，并將 all-gather 與前向傳遞計算重疊，實現了這些優化。這些優化有助于 Nemotron-4 15B 的近線性擴展。圖 4 顯示了在使用與圖 3 類似的實驗設置時，啟用這些優化后，15B 模型的每個 GPU 吞吐量。
此外，v0.7 版本通過提供更好的內置啟發式算法來確定存儲桶的大小，從而進一步提高了 Megatron-Core 在處理大數據并行大小時的速度。這使得通信能夠保持帶寬限制，而不是延遲限制，即使在數據并行大小超過 300（GPU 數量超過 3,000）時也是如此。圖 5 顯示了 15B 模型的每個 GPU 吞吐量，在數據并行大小高達 384 的情況下，并具有相同的弱擴展實驗設置，所有優化均已啟用（包括具有 reduce-scatter 和 all-gather 重疊的分布式優化器）。
Megatron-Core 優化還可為需要其他并行性維度的大型模型開箱即用，包括管道并行性。 Nemotron-4 340B 模型使用這些優化在 Megatron-Core 中，以在使用 BF16 的大型 GPU 規模下實現高訓練吞吐量。表 2 顯示了 Nemotron-4 340B 基礎模型上 Megatron-Core 的每個 GPU 吞吐量，具體來說，TP 大小為 8，PP 大小為 12，虛擬管線級的數量為 8，序列長度為 4096。有關更多詳細信息，請參閱 Nemotron-4 340B 技術報告。其他并行性配置可能會導致吞吐量略高。
| 精度 | GPU 數量 (H100) | 數據并行大小 | 批量大小 | 每個 GPU 的吞吐量 (TFLOP/s/GPU) |
| BF16 | 1536 | 16 | 768 | 419.3 |
| BF16 | 3072 | 32 | 1536 | 418.3 |
| BF16 | 6144 | 64 | 2304 | 405.0 |
開始使用
Megatron-Core 可在 NVIDIA/Megatron-LM 資源庫中以開源形式提供，并可與 Megatron-LM 或 NVIDIA NeMo 一起使用。Megatron-LM 是一種輕量級訓練框架，提供了可定制的原生 PyTorch 訓練循環，非常適合偏好較少抽象層的用戶。它為探索 Megatron-Core 提供了直接的切入點。有關更多詳細信息，請參閱 Megatron-Core 文檔。
Megatron-Core 已深度集成到 NVIDIA NeMo，NVIDIA NeMo 是一個企業級 AI 軟件平臺，具有安全性、穩定性、可管理性和支持。它包含適用于 LLM 功能的 Megatron-Core，并為多模態和語音 AI 提供了一套更豐富的工具。如需了解更多信息，請參閱 NVIDIA NeMo 框架文檔。
?
