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 框架文档。
