乐趣区

关于人工智能:如何使用-MegatronLM-训练语言模型

jiezi

在 PyTorch 中训练大语言模型不仅仅是写一个训练循环这么简略。咱们通常须要将模型散布在多个设施上，并应用许多优化技术以实现稳固高效的训练。Hugging Face 🤗 Accelerate 的创立是为了反对跨 GPU 和 TPU 的分布式训练，并使其可能非常容易的集成到训练代码中。🤗 Transformers 还反对应用 Trainer API 来训练，其在 PyTorch 中提供性能残缺的训练接口，甚至不须要本人编写训练的代码。

Megatron-LM 是钻研人员用于预训练大型 Transformer 模型的另一个风行工具，它是 NVIDIA 利用深度学习钻研团队开发的一个弱小框架。与 accelerate 和 Trainer 不同，Megatron-LM 应用起来并不简略，对于初学者来说可能难以上手。但它针对 GPU 上的训练进行了高度优化。在这篇文章中，你将学习如何应用 Megatron-LM 框架在 NVIDIA GPU 上训练语言模型，并将其与 transformers 联合。

咱们将合成在此框架中训练 GPT2 模型的不同步骤，包含:

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

为什么抉择 Megatron-LM?

在进入训练细节的解说之前，让咱们首先理解是什么让这个框架比其余框架更高效。本节的灵感来自这篇对于应用 Megatron-DeepSpeed 进行 BLOOM 训练的精彩 博客，请参阅该博客以获取更多详细信息，因为该博客旨在对 Megatron-LM 进行具体的介绍。

数据加载

Megatron-LM 带有一个高效的 DataLoader，其中数据在训练前被 tokenize 和 shuffle。它还将数据拆分为带有索引的编号序列，并将索引存储，因而 tokenize 只须要计算一次。为了构建索引，首先依据训练参数计算每个 epoch 的数量，并创立一个排序，而后对数据进行 shuffle 操作。这与大多数状况不同，咱们通常迭代整个数据集直到其用尽，而后反复第二个 epoch。这平滑了学习曲线并节俭了训练工夫。

交融 CUDA 内核

当一个计算在 GPU 上运行时，必要的数据会从内存中取出并加载到 GPU 上，而后计算结果被保留回内存。简略来说，交融内核的思维是: 将通常由 PyTorch 独自执行的相似操作组合成一个独自的硬件操作。因而能够将多个离散计算合并为一个，从而缩小在多个离散计算中的内存挪动次数。下图阐明了内核交融的思维。它的灵感来自这篇 论文，该论文具体探讨了这个概念。

当 f、g 和 h 交融在一个内核中时，f 和 g 的两头后果 x’ 和 y’ 存储在 GPU 寄存器中并立刻被 h 应用。然而如果不交融，x’ 和 y’ 就须要复制到内存中，而后由 h 加载。因而，交融 CUDA 内核显着放慢了计算速度。此外，Megatron-LM 还应用 Apex 的 AdamW 交融实现，它比 PyTorch 实现更快。

尽管咱们能够在 transformers 中自定义 Megatron-LM 中的 DataLoader 和 Apex 的交融优化器，但自定义交融 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

你还须要在容器的 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)

而后应用以下命令将数据 tokenize、shuffle 并解决成二进制格局以进行训练:

#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 选项指的是预处理中应用的线程数量和调配给每个线程的数据块大小。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 上并行运算，你须要将 tensor-model-parallel-size 参数更改为所需的 GPU 数量。第二种抉择是流水线并行，它将模型按层分成大小相等的几块。参数 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

请留神，如果你打算应用原始的分词器，你将须要在转换后将生成的词汇文件和合并表替换为咱们之前介绍的原始文件。

不要遗记将你的模型推送到 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 来训练语言模型。问题是咱们须要思考什么时候应用它。因为额定的预处理和转换步骤，这个框架显然减少了一些工夫开销。因而，重要的是你要思考哪个框架更适宜你的需要和模型大小。咱们倡议将其用于预训练模型或微调，但可能不适用于中型模型的微调。APITrainer 和 accelerate 库对于模型训练同样也十分不便，并且它们与设施无关，为用户提供了极大的灵活性。

祝贺 🎉 当初你学会了如何在 Megatron-LM 框架中训练 GPT2 模型并使其反对 transformers！

原文链接: https://hf.co/blog/megatron-training

作者: Loubna Ben Allal

译者: gxy-gxy

审校 / 排版: zhongdongy (阿东)

退出移动版