NVIDIA Megatron-Core

从头开始大规模训练生成式 AI 模型。

NVIDIA Megatron-Core 是一个基于 PyTorch 的开源库，可在数千个 GPU 上以惊人的速度大规模训练大型模型。它采用 GPU 优化的训练技术和先进的系统级创新技术，所有这些创新均可通过可组合的 API 实现。Megatron-Core 与 NVIDIA NeMo™ 无缝集成，提供端到端云原生解决方案，用于构建、定制和部署大语言模型 (LLM) 。

立即下载开始使用

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

并行技术

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

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

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

可定制的基础模组

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

借助 Megatron-Core (Mcore) 规格系统，研究人员可以在 PyTorch 模型定义中轻松自定义所需的抽象级别的子模块。

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

可扩展性和训练弹性

借助训练弹性功能 (例如快速分布式检查点。
xi
详细了解 Megatron-Core 如何在 H100 GPU 规模高达 6K+ 的情况下训练 Nemotron-4 340B 模型，同时实现高每个 GPU 吞吐量。

详情请参阅可扩展性基准测试。

尖端研究

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

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

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

使用多专家模型进行训练

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

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

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

超越 Transformer：混合模型

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

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

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

多模态训练

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

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

开始使用

将 Megatron-Core 与 NVIDIA NeMo 结合使用

NVIDIA NeMo 是一个端到端平台，用于开发自定义生成式 AI，包括 LLM 以及多模态、视觉和语音 AI。NeMo 基于 NVIDIA NeMo 构建，适合开发者构建企业就绪型生成式 AI 应用。

了解详情

将 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 以惊人的速度训练大型模型

开始使用