如何使用Megatron-LM训练语言模型
使用Megatron-LM训练语言模型的方法
在PyTorch中训练大型语言模型需要比简单的训练循环更多的工作。通常会在多个设备上进行分布式训练，并使用许多优化技术来实现稳定高效的训练。Hugging Face 🤗 Accelerate库旨在支持在GPU和TPU上进行分布式训练，并且可以很容易地集成到训练循环中。🤗 Transformers还通过Trainer API支持分布式训练，提供了在PyTorch中完整的训练功能，甚至无需实现训练循环。
研究人员在预训练大型Transformer模型时，另一个流行的工具是Megatron-LM，这是NVIDIA应用深度学习研究团队开发的一个强大框架。与accelerate和Trainer不同，使用Megatron-LM并不简单，对初学者来说可能有些困惑。但它在GPU上的训练经过了高度优化，可以提高一些速度。在本博文中，您将学习如何在NVIDIA GPU上使用Megatron-LM训练语言模型，并与transformers一起使用。
我们将尝试分解在该框架中训练GPT2模型的不同步骤，包括：
- 环境设置
- 数据预处理
- 训练
- 模型转换为🤗 Transformers
为什么选择Megatron-LM？
在深入了解训练细节之前，让我们先了解一下这个框架相比其他框架更高效的原因。这部分内容受到了这篇关于使用Megatron-DeepSpeed进行BLOOM训练的精彩博文的启发，请参考该博文以获取更多细节，因为本博文旨在对Megatron-LM进行简要介绍。
数据加载器
Megatron-LM提供了一个高效的数据加载器，在训练之前会对数据进行分词和洗牌。它还将数据分割为带有索引的编号序列，这些索引仅需计算一次即可存储。为了构建索引，基于训练参数计算了周期数，并创建了一个顺序，然后对其进行洗牌。这与大多数情况不同，大多数情况下我们会遍历整个数据集直到耗尽，然后再进行第二个周期的遍历。这样可以平滑学习曲线并节省训练时间。
融合的CUDA核心
当在GPU上运行计算时，需要从内存中提取必要的数据，然后运行计算并将结果保存回内存。简单来说，融合核心的思想是将通常由PyTorch分别执行的类似操作合并为单个硬件操作。因此，它们通过将多个离散计算中的内存移动合并为一个来减少内存移动的次数。下图展示了融合核心的思想，该思想受到了详细讨论该概念的这篇论文的启发。
当f、g和h在一个核心中融合时，f和g的中间结果x’和y’存储在GPU寄存器中，并且立即被h使用。但是，如果没有融合，x’和y’需要被复制到内存中，然后由h加载。因此，核心融合可以显著加速计算。Megatron-LM还使用了Apex中的AdamW的融合实现，该实现比PyTorch的实现更快。
虽然您可以像Megatron-LM一样自定义数据加载器，并使用transformers的Apex融合优化器，但构建自定义的融合CUDA核心对初学者来说并不容易。
现在您对该框架及其优势有了一定了解，让我们进入训练细节！
如何使用Megatron-LM进行训练？
环境设置
设置环境的最简单方法是拉取一个带有所有所需安装的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您还需要在容器的Megatron-LM文件夹中添加词汇文件vocab.json和分割表merges.txt。这些文件可以在模型的代码库中找到，其中包括权重，参见GPT2的这个代码库。您也可以使用transformers训练自己的分词器。您可以查看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) 然后，数据被分词、洗牌并且被处理成用于训练的二进制格式，使用以下命令：
#如果未安装nltk
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-eodworkers和chunk_size选项指的是预处理中使用的worker数量和分配给每个worker的数据块大小。dataset-impl指的是索引数据集的实现模式[‘lazy’, ‘cached’, ‘mmap’]。这将输出两个文件codeparrot_content_document.idx和codeparrot_content_document.bin，用于训练。
训练
您可以根据下面所示配置模型架构和训练参数，或将其放入一个您将运行的bash脚本中。此命令在8个GPU上运行110M参数的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 的非常大的模型。第一个选项是使用张量并行性，在多个 GPU 上分割执行单个 Transformer 模块，您需要将 tensor-model-parallel-size 参数更改为所需的 GPU 数量。第二个选项是使用流水线并行性，将 Transformer 模块分割成大小相等的阶段。参数 pipeline-model-parallel-size 确定将模型分割为的阶段数量。有关更多详细信息，请参阅此博客
将模型转换为 🤗 Transformers
训练完成后，我们希望在 transformers 中使用该模型，例如进行评估或将其部署到生产环境中。您可以按照此教程将其转换为 transformers 模型。例如，在训练完成后，您可以复制最后一次迭代 150k 的权重，并使用以下命令将 model_optim_rng.pt 文件转换为 pytorch_model.bin 文件，该文件由 transformers 支持：
# 在容器外执行：
mkdir -p nvidia/megatron-codeparrot-small
# 从容器中复制权重
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注意，如果您计划从那里加载分词器，您将需要在转换后将生成的词汇表文件和合并表替换为之前介绍的原始文件。
不要忘记将您的模型推送到 hub 并与社区共享，这只需要三行代码 🤗：
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("nvidia/megatron-codeparrot-small")
# 这将在您的用户名下创建一个名为 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!")Tranfsormers 还可以高效地处理大模型的推理。如果您训练了一个非常大的模型（例如使用模型并行性），您可以使用以下命令轻松地进行推理：
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("your_username/codeparrot-large", device_map="auto")这将在后台使用 accelerate 库，自动将模型权重分派到您可用的设备（GPU、CPU RAM）上。
免责声明：我们已经展示了任何人都可以使用 Megatron-LM 训练语言模型。问题是何时使用它。由于额外的预处理和转换步骤，此框架显然会增加一些时间开销。因此，重要的是您确定哪个框架对于您的情况和模型大小更为适合。我们建议尝试用于预训练模型或扩展微调，但可能不适用于 VoAGI 大小的模型的短期微调。Trainer API 和 accelerate 库也非常适用于模型训练，它们是设备无关的，并为用户提供了重要的灵活性。
恭喜 🎉，现在您知道如何在 Megatron-LM 中训练 GPT2 模型并使其受到 transformers 的支持！
