NVIDIA Megatron Core

大规模训练生成式 AI 模型。

NVIDIA Megatron Core 是一款用于大规模生成式 AI 的可组合训练库。它为训练和后训练工作流提供经过 GPU 优化的基础构件，使团队能够构建具备现代 LLM、MoE 和多模态开发所需性能、灵活性与可扩展性的自定义系统。Megatron Core 与 Megatron Bridge、Transformer Engine 以及生态系统中的各类库集成，可为生产级训练与科研提供完善的整体解决方案。

立即下载开始使用

探索 NVIDIA Megatron-Core 的特性和优势

并行技术

Megatron-Core 库为大规模训练提供先进的模型并行技术，包括张量、序列、工作流、上下文和 MoE 专家并行。

Megatron-Core 的用户可以灵活地结合不同的模型并行技术，以更适合其训练工作负载。此外，Megatron-Core 还提供节省内存的功能，包括激活重新计算、分布式优化器和分布式检查点。

如需了解详情，请参阅 API 文档。

可定制的基础模组

Megatron-Core 通过模块化和可组合的 API 提供可定制的基础模组。对于 Transformer 模型，它提供注意力机制、归一化层、嵌入技术等。

如需详细了解 MCore 规格系统，请参阅文档。

可扩展性和训练弹性

通过自动重启、故障/卡顿检测以及快速分布式 checkpoint 等训练弹性特性，在大规模环境中高效训练大型模型。

了解更多，查看 Megatron-Core 如何在多达 6K+ H100 GPU 的规模下支持 Nemotron-4 340B 模型训练，并在每个 GPU 上实现较高的吞吐量。

在此可扩展性基准测试中查看性能详情。有关弹性功能的更多信息，请参阅 NVIDIA Resiliency Extension。

尖端研究

只需升级到最新的 Megatron-Core，即可利用 NVIDIA 的前沿研究，始终处于分布式训练的前沿。

自 2019 年以来，Megatron-Core 率先推出了大型模型训练，并继续引领大规模训练的创新。

阅读此博客，了解一些近期进展。

使用多专家模型进行训练

使用多专家模型 (MoE) 预训练模型，这是一种热门技术，可在不增加计算资源的情况下实现更高的准确性。

Megatron-Core 为无令牌丢弃和令牌丢弃用例提供高性能功能，并优化了训练速度。

在我们的代码库中详细了解 MoE 功能。

超越 Transformer：混合模型

Megatron-Core 将其支持范围从基于 Transformer 的模型扩展到结合了状态空间模型、状态空间对偶和递归神经网络的混合模型。

混合模型克服了一些注意力限制，已成为用于序列建模任务的引人注目的模型架构。

在我们的论文和代码示例中了解有关训练基于 Mamba 的混合模型的更多信息。

多模态训练

使用 Megatron-Core 的并行技术及其多模态数据加载器库进行多模态训练，在混合多模态数据集时具有确定性和可重复性。

开始使用 LLaVA (大型语言和视觉助手) 训练流程。

开始使用

将 Megatron Core 与 NVIDIA NeMo 结合使用

Megatron Bridge 是 NeMo Framework 中的一个库，用于在 Hugging Face 与 Megatron Core 之间建立桥接，并为最新模型提供经过验证的脚本。它支持这两种格式之间的双向 checkpoint 转换，使其他项目既可以利用 Megatron Core 的并行能力，也可以将模型导出到不同的推理引擎中使用。

了解详情

将 Megatron-Core 与 NVIDIA Megatron-LM 结合使用

Megatron-LM 是一个开源轻量级训练框架，具有原生 PyTorch 训练循环，用于探索 Megatron-Core。它易于定制，适合偏好基于 Megatron-Core 训练技术的最小抽象层的研究人员。

了解详情

出色的训练速度和可扩展性

Megatron-Core 能够利用其并行技术高效训练大语言模型。在以下弱扩展实验中，GPT 模型的参数范围在 20 亿到 4620 亿之间，Megatron-Core 演示了高达 6144 个 H100 GPU 的超线性扩展。

模型大小
Tensor MP 大小
工作流 MP 大小
数据并行大小
GPU 数量
批量大小
每个 GPU teraFLOP/s
MFU
21 亿
1
1
16
64
16
64
256
441
412
45%
42%
42 亿
2
1
16
64
32
128
256
431
415
44%
42%
83 亿
4
1
16
64
64
256
256
457
426
46%
43%
197 亿
8
1
16
64
128
512
512
439
429
44%
43%
410 亿
8
1
32
128
256
1024
768
469
439
47%
44%
780 亿
8
2
32
96
512
1536
960
446
418
45%
42%
148B
8
4
24
72
768
2304
1152
456
432
46%
44%
314B
8
8
16
48
1024
3072
1152
490
464
50%
47%
509B
8
20
8
24
1280
3840
1440
473
426
48%
43%

聚合吞吐量 (弱扩展)

A graph showing NVIDIA Megatron-Core aggregate throughput for weak scaling

聚合吞吐量 (强扩展)

在 1770 亿参数 GPT-3 模型的强扩展设置中，使用相同的批量大小 (共 1152 个序列) ，Megatron-Core 展示了从 96 个 GPU 到 4608 个 H100 GPU 的近乎线性扩展。

A graph showing NVIDIA Megatron-Core aggregate throughput for strong scaling
查看基准测试详情

资源

以空前速度使用 Megatron 训练大规模模型

开始使用