MegatronGPT2

概述

MegatronGPT2 模型由 Mohammad Shoeybi、Mostofa Patwary、Raul Puri、Patrick LeGresley、Jared Casper 和 Bryan Catanzaro 在论文 《Megatron-LM：使用模型并行主义训练数十亿参数的语言模型》 中提出。

论文摘要如下：

最近的语言建模工作表明，训练大型 Transformer 模型可以推动自然语言处理应用的最新技术水平。然而，由于内存限制，训练非常大的模型可能相当困难。在这项工作中，我们介绍了我们训练非常大的 Transformer 模型的技术，并实现了一种简单、高效的层内模型并行方法，该方法使得训练具有数十亿参数的 Transformer 模型成为可能。我们的方法不需要新的编译器或库的更改，与流水线模型并行是正交和互补的，并且可以通过在原生 PyTorch 中插入一些通信操作来完全实现。我们通过在 512 个 GPU 上收敛高达 83 亿参数的基于 Transformer 的模型来阐释这种方法。与一个强大的单 GPU 基线（维持 39 TeraFLOPs，即峰值 FLOPs 的 30%）相比，我们在整个应用中维持 15.1 PetaFLOPs，扩展效率达到 76%。为了证明大型语言模型可以进一步推动最先进技术（SOTA），我们训练了一个类似于 GPT-2 的 83 亿参数 Transformer 语言模型和一个类似于 BERT 的 39 亿参数模型。我们表明，随着模型规模的增长，仔细关注类 BERT 模型中层归一化的位置对于实现性能提升至关重要。使用 GPT-2 模型，我们在 WikiText103（困惑度为 10.8，而 SOTA 困惑度为 15.8）和 LAMBADA（准确率为 66.5%，而 SOTA 准确率为 63.2%）数据集上取得了 SOTA 结果。我们的 BERT 模型在 RACE 数据集上取得了 SOTA 结果（准确率为 90.9%，而 SOTA 准确率为 89.4%）。

该模型由 jdemouth 贡献。原始代码可以在这里找到。该代码库包含了 Megatron 语言模型的多 GPU 和多节点实现。特别是，它包含了一种使用“张量并行”和“流水线并行”技术的混合模型并行方法。

使用技巧

我们提供了预训练的 GPT2-345M 检查点，可用于评估或微调下游任务。

要访问这些检查点，首先需要注册并设置 NVIDIA GPU Cloud (NGC) 注册中心命令行工具。更多关于下载模型的文档可以在 NGC 文档中找到。

另外，你也可以使用以下命令直接下载检查点：

wget --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/megatron_lm_345m/versions/v0.0/zip -O
megatron_gpt2_345m_v0_0.zip

从 NVIDIA GPU Cloud (NGC) 获取检查点后，你需要将其转换为 Hugging Face Transformers GPT2 实现可以轻松加载的格式。

以下命令可以让你进行转换。我们假设 `models/megatron_gpt2` 文件夹包含 `megatron_gpt2_345m_v0_0.zip` 文件，并且该命令是在该文件夹内运行的。

python3 $PATH_TO_TRANSFORMERS/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py megatron_gpt2_345m_v0_0.zip