如何使用 Megatron-LM 训练语言模型

在 PyTorch 中训练大型语言模型需要的不仅仅是一个简单的训练循环。它通常分布在多个设备上，并采用许多优化技术以实现稳定高效的训练。Hugging Face 🤗 Accelerate 库旨在支持跨 GPU 和 TPU 的分布式训练，并可非常轻松地集成到训练循环中。🤗 Transformers 也通过 Trainer API 支持分布式训练，该 API 在 PyTorch 中提供功能完整的训练，甚至无需实现训练循环。

研究人员预训练大型 Transformer 模型的另一个常用工具是 Megatron-LM，这是一个由 NVIDIA 深度学习应用研究团队开发的强大框架。与 `accelerate` 和 `Trainer` 不同，使用 Megatron-LM 不那么直接，对初学者来说可能有点不知所措。但它针对 GPU 训练进行了高度优化，可以提供一些加速。在这篇博客文章中，您将学习如何使用 Megatron-LM 在 NVIDIA GPU 上训练语言模型，并将其与 `transformers` 一起使用。

我们将尝试分解此框架中训练 GPT2 模型的不同步骤，包括：

• 环境设置
• 数据预处理
• 训练
• 模型转换为 🤗 Transformers

为什么选择 Megatron-LM？

在深入训练细节之前，我们首先了解是什么让这个框架比其他框架更高效。本节受这篇关于使用 Megatron-DeepSpeed 训练 BLOOM 的精彩博客启发，请参阅该博客了解更多详细信息，因为这篇博客旨在对 Megatron-LM 进行温和的介绍。

数据加载器 (DataLoader)

Megatron-LM 配备了一个高效的 DataLoader，其中数据在训练前被标记化和打乱。它还将数据分割成带索引的编号序列，这些索引被存储起来，这样只需计算一次。为了构建索引，根据训练参数计算 epoch 数量，然后创建并打乱排序。这与大多数情况下我们遍历整个数据集直到耗尽，然后再重复第二个 epoch 的方式不同。这可以平滑学习曲线并在训练期间节省时间。

融合 CUDA 内核 (Fused CUDA Kernels)

当在 GPU 上运行计算时，所需数据从内存中获取，然后运行计算，结果再保存回内存。简单来说，融合内核的思想是，通常由 PyTorch 单独执行的相似操作被组合成一个单一的硬件操作。因此，通过将多个离散计算合并为一个，它们减少了内存移动的次数。下图说明了内核融合的思想。它受到这篇论文的启发，该论文详细讨论了这一概念。

当 f、g 和 h 融合在一个内核中时，f 和 g 的中间结果 x' 和 y' 存储在 GPU 寄存器中，并立即被 h 使用。但如果没有融合，x' 和 y' 将需要复制到内存中，然后由 h 加载。因此，内核融合可以显著加速计算。Megatron-LM 还使用了来自 Apex 的 AdamW 融合实现，它比 PyTorch 实现更快。

虽然可以像 Megatron-LM 那样定制 DataLoader，并使用 Apex 的融合优化器与 `transformers`，但构建自定义的融合 CUDA 内核对于初学者来说并不友好。

现在您已经熟悉了这个框架以及它为什么具有优势，让我们深入了解训练细节！

如何使用 Megatron-LM 进行训练？

设置

设置环境最简单的方法是从 NGC 拉取一个预装所有必要组件的 NVIDIA PyTorch 容器。更多详情请参阅文档。如果您不想使用此容器，则需要安装最新的 PyTorch、CUDA、NCCL 和 NVIDIA APEX 版本以及 `nltk` 库。

安装 Docker 后，您可以使用以下命令运行容器（`xx.xx` 表示您的 Docker 版本），然后在其内部克隆 Megatron-LM 仓库

docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:xx.xx-py3
git clone https://github.com/NVIDIA/Megatron-LM

您还需要将 tokenizer 的词汇文件 `vocab.json` 和合并表 `merges.txt` 添加到容器的 Megatron-LM 文件夹中。这些文件可以在模型的存储库中与权重一起找到，请参阅 GPT2 的此存储库。您也可以使用 `transformers` 训练自己的 tokenizer。您可以查看 CodeParrot 项目以获取实际示例。现在，如果您想将这些数据从容器外部复制到容器内部，可以使用以下命令：

sudo docker cp vocab.json CONTAINER_ID:/workspace/Megatron-LM
sudo docker cp merges.txt CONTAINER_ID:/workspace/Megatron-LM

数据预处理

在本教程的其余部分，我们将使用 CodeParrot 模型和数据作为示例。

训练数据需要一些预处理。首先，您需要将其转换为松散的 JSON 格式，每行一个 JSON 包含一个文本样本。如果您使用的是 🤗 Datasets，下面是如何操作的示例（始终在 Megatron-LM 文件夹内部）

from datasets import load_dataset

train_data = load_dataset('codeparrot/codeparrot-clean-train', split='train')
train_data.to_json("codeparrot_data.json", lines=True)  

然后使用以下命令将数据进行标记化、打乱并处理成二进制格式以进行训练：

#if nltk isn't installed
pip install nltk
python tools/preprocess_data.py \
       --input codeparrot_data.json \
       --output-prefix codeparrot \
       --vocab vocab.json \
       --dataset-impl mmap \
       --tokenizer-type GPT2BPETokenizer \
       --merge-file merges.txt \
       --json-keys content \
       --workers 32 \
       --chunk-size 25 \
       --append-eod

`workers` 和 `chunk_size` 选项分别指预处理中使用的 worker 数量以及分配给每个 worker 的数据块大小。`dataset-impl` 指的是索引数据集的实现模式，来自 ['lazy', 'cached', 'mmap']。这将输出两个文件 `codeparrot_content_document.idx` 和 `codeparrot_content_document.bin`，它们用于训练。

训练

您可以按照如下所示配置模型架构和训练参数，或将其放入您将运行的 bash 脚本中。此命令在 8 个 GPU 上运行 1.1 亿参数 CodeParrot 模型的预训练。请注意，数据默认按 969:30:1 的比例划分为训练/验证/测试集。

GPUS_PER_NODE=8
MASTER_ADDR=localhost
MASTER_PORT=6001
NNODES=1
NODE_RANK=0
WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))
DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE --nnodes $NNODES --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT"
CHECKPOINT_PATH=/workspace/Megatron-LM/experiments/codeparrot-small
VOCAB_FILE=vocab.json
MERGE_FILE=merges.txt
DATA_PATH=codeparrot_content_document
GPT_ARGS="--num-layers 12
--hidden-size 768
--num-attention-heads 12
--seq-length 1024
--max-position-embeddings 1024
--micro-batch-size 12
--global-batch-size 192
--lr 0.0005
--train-iters 150000
--lr-decay-iters 150000
--lr-decay-style cosine
--lr-warmup-iters 2000
--weight-decay .1
--adam-beta2 .999
--fp16
--log-interval 10
--save-interval 2000
--eval-interval 200
--eval-iters 10
"
TENSORBOARD_ARGS="--tensorboard-dir experiments/tensorboard"
python3 -m torch.distributed.launch $DISTRIBUTED_ARGS \
        pretrain_gpt.py \
        --tensor-model-parallel-size 1 \
        --pipeline-model-parallel-size 1 \
        $GPT_ARGS \
        --vocab-file $VOCAB_FILE \
        --merge-file $MERGE_FILE \
        --save $CHECKPOINT_PATH \
        --load $CHECKPOINT_PATH \
        --data-path $DATA_PATH \
        $TENSORBOARD_ARGS

在此设置下，训练大约需要 12 小时。

此设置使用数据并行，但对于不适合单个 GPU 的超大型模型，也可以使用模型并行。第一种选择是张量并行，它将单个 Transformer 模块的执行分散到多个 GPU 上，您需要将 `tensor-model-parallel-size` 参数更改为所需的 GPU 数量。第二种选择是流水线并行，其中 Transformer 模块被分成大小相等的阶段。参数 `pipeline-model-parallel-size` 决定了将模型分成多少个阶段。有关更多详细信息，请参阅这篇博客

将模型转换为 🤗 Transformers

训练完成后，我们希望在 `transformers` 中使用该模型，例如用于评估或部署到生产环境。您可以按照此教程将其转换为 `transformers` 模型。例如，训练结束后，您可以复制最后一次迭代的权重（150k），并将 `model_optim_rng.pt` 文件转换为 `transformers` 支持的 `pytorch_model.bin` 文件，使用以下命令：

# to execute outside the container:
mkdir -p nvidia/megatron-codeparrot-small
# copy the weights from the container
sudo docker cp CONTAINER_ID:/workspace/Megatron-LM/experiments/codeparrot-small/iter_0150000/mp_rank_00/model_optim_rng.pt nvidia/megatron-codeparrot-small
git clone https://github.com/huggingface/transformers.git
git clone https://github.com/NVIDIA/Megatron-LM.git
export PYTHONPATH=Megatron-LM
python transformers/src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py nvidia/megatron-codeparrot-small/model_optim_rng.pt

请注意，如果您打算从那里加载 tokenizer，则需要在转换后用我们之前介绍的原始文件替换生成的词汇文件和合并表。

别忘了将您的模型推送到 Hub 并与社区分享，只需三行代码 🤗

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("nvidia/megatron-codeparrot-small")
# this creates a repository under your username with the model name codeparrot-small
model.push_to_hub("codeparrot-small")

您还可以轻松地用它来生成文本

from transformers import pipeline

pipe = pipeline("text-generation", model="your_username/codeparrot-small")
outputs = pipe("def hello_world():")
print(outputs[0]["generated_text"])

 def hello_world():   
    print("Hello World!")

Transformers 也能高效处理大型模型推理。如果您训练了一个非常大的模型（例如使用模型并行），您可以使用以下命令轻松地将其用于推理：

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("your_username/codeparrot-large", device_map="auto")

这将在幕后使用 accelerate 库，自动将模型权重分配到您可用的设备（GPU、CPU RAM）上。

免责声明：我们已经展示了任何人都可以使用 Megatron-LM 训练语言模型。问题是什么时候使用它。由于额外的预处理和转换步骤，这个框架显然会增加一些时间开销。因此，您决定哪个框架更适合您的情况和模型大小很重要。我们建议将其用于模型预训练或扩展微调，但可能不适用于中等大小模型的较短微调。`Trainer` API 和 `accelerate` 库也非常方便模型训练，它们是设备无关的，并为用户提供了显著的灵活性。

恭喜 🎉 现在您知道如何使用 Megatron-LM 训练 GPT2 模型并使其受 `transformers` 支持了！