在 PyTorch 中訓(xùn)練大語言模型不僅僅是寫一個訓(xùn)練循環(huán)這么簡單。我們通常需要將模型分布在多個設(shè)備上，并使用許多優(yōu)化技術(shù)以實現(xiàn)穩(wěn)定高效的訓(xùn)練。Hugging Face ???Accelerate 的創(chuàng)建是為了支持跨 GPU 和 TPU 的分布式訓(xùn)練，并使其能夠非常容易的集成到訓(xùn)練代碼中。???Transformers 還支持使用 Trainer API 來訓(xùn)練，其在 PyTorch 中提供功能完整的訓(xùn)練接口，甚至不需要自己編寫訓(xùn)練的代碼。

Megatron-LM 是研究人員用于預(yù)訓(xùn)練大型 Transformer 模型的另一個流行工具，它是 NVIDIA 應(yīng)用深度學(xué)習(xí)研究團(tuán)隊開發(fā)的一個強大框架。與 accelerate 和 Trainer 不同，Megatron-LM 使用起來并不簡單，對于初學(xué)者來說可能難以上手。但它針對 GPU 上的訓(xùn)練進(jìn)行了高度優(yōu)化。在這篇文章中，你將學(xué)習(xí)如何使用 Megatron-LM 框架在 NVIDIA GPU 上訓(xùn)練語言模型，并將其與 transformers 結(jié)合。

我們將分解在此框架中訓(xùn)練 GPT2 模型的不同步驟，包括:

• 環(huán)境設(shè)置
• 數(shù)據(jù)預(yù)處理
• 訓(xùn)練
• 將模型轉(zhuǎn)化為 ?? Transformers

為什么選擇 Megatron-LM?

在進(jìn)入訓(xùn)練細(xì)節(jié)的講解之前，讓我們首先了解是什么讓這個框架比其他框架更高效。本節(jié)的靈感來自這篇關(guān)于使用 Megatron-DeepSpeed 進(jìn)行 BLOOM 訓(xùn)練的精彩 博客，請參閱該博客以獲取更多詳細(xì)信息，因為該博客旨在對 Megatron-LM 進(jìn)行詳細(xì)的介紹。

數(shù)據(jù)加載

Megatron-LM 帶有一個高效的 DataLoader，其中數(shù)據(jù)在訓(xùn)練前被 tokenize 和 shuffle。它還將數(shù)據(jù)拆分為帶有索引的編號序列，并將索引存儲，因此 tokenize 只需要計算一次。為了構(gòu)建索引，首先根據(jù)訓(xùn)練參數(shù)計算每個 epoch 的數(shù)量，并創(chuàng)建一個排序，然后對數(shù)據(jù)進(jìn)行 shuffle 操作。這與大多數(shù)情況不同，我們通常迭代整個數(shù)據(jù)集直到其用盡，然后重復(fù)第二個 epoch 。這平滑了學(xué)習(xí)曲線并節(jié)省了訓(xùn)練時間。

融合 CUDA 內(nèi)核

當(dāng)一個計算在 GPU 上運行時，必要的數(shù)據(jù)會從內(nèi)存中取出并加載到 GPU 上，然后計算結(jié)果被保存回內(nèi)存。簡單來說，融合內(nèi)核的思想是: 將通常由 PyTorch 單獨執(zhí)行的類似操作組合成一個單獨的硬件操作。因此可以將多個離散計算合并為一個，從而減少在多個離散計算中的內(nèi)存移動次數(shù)。下圖說明了內(nèi)核融合的思想。它的靈感來自這篇 論文，該論文詳細(xì)討論了這個概念。

當(dāng) f、g 和 h 融合在一個內(nèi)核中時，f 和 g 的中間結(jié)果 x' 和 y' 存儲在 GPU 寄存器中并立即被 h 使用。但是如果不融合，x' 和 y' 就需要復(fù)制到內(nèi)存中，然后由 h 加載。因此，融合 CUDA 內(nèi)核顯著加快了計算速度。此外，Megatron-LM 還使用 Apex 的 AdamW 融合實現(xiàn)，它比 PyTorch 實現(xiàn)更快。

雖然我們可以在 transformers 中自定義 Megatron-LM 中的 DataLoader 和 Apex 的融合優(yōu)化器，但自定義融合 CUDA 內(nèi)核對新手來說太不友好了。

現(xiàn)在你已經(jīng)熟悉了該框架及其優(yōu)勢，讓我們進(jìn)入訓(xùn)練細(xì)節(jié)吧！

如何使用 Megatron-LM 框架訓(xùn)練？

環(huán)境設(shè)置

設(shè)置環(huán)境的最簡單方法是從 NGC 拉取附帶所有所需環(huán)境的 NVIDIA PyTorch 容器。有關(guān)詳細(xì)信息，請參閱 文檔。如果你不想使用此容器，則需要安裝最新的 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。這些文件可以在帶有權(quán)重的模型倉庫中找到，請參閱 GPT2 庫。你還可以使用 transformers 訓(xùn)練自己的分詞器。你可以查看 CodeParrot 項目 以獲取實際示例?，F(xiàn)在，如果你想從容器外部復(fù)制這些數(shù)據(jù)，你可以使用以下命令:

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

數(shù)據(jù)預(yù)處理

在本教程的其余部分，我們將使用 CodeParrot 模型和數(shù)據(jù)作為示例。

我們需要對預(yù)訓(xùn)練數(shù)據(jù)進(jìn)行預(yù)處理。首先，你需要將其轉(zhuǎn)換為 json 格式，一個 json 的一行包含一個文本樣本。如果你正在使用 ?? Datasets，這里有一個關(guān)于如何做到這一點的例子 (請在 Megatron-LM 文件夾中進(jìn)行這些操作):

from datasets import load_dataset

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

然后使用以下命令將數(shù)據(jù) tokenize、shuffle 并處理成二進(jìn)制格式以進(jìn)行訓(xùn)練:

#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 選項指的是預(yù)處理中使用的線程數(shù)量和分配給每個線程的數(shù)據(jù)塊大小。dataset-impl 指的是索引數(shù)據(jù)集的實現(xiàn)方式，包括 ['lazy', 'cached', 'mmap']。這將輸出 codeparrot_content_document.idx 和 codeparrot_content_document.bin 兩個文件用于訓(xùn)練。

訓(xùn)練

你可以使用如下所示配置模型架構(gòu)和訓(xùn)練參數(shù)，或?qū)⑵浞湃肽銓⑦\行的 bash 腳本中。該命令在 8 個 GPU 上參數(shù)為 110M 的 CodeParrot 模型進(jìn)行預(yù)訓(xùn)練。請注意，數(shù)據(jù)默認(rèn)按 969:30:1 的比例劃分為訓(xùn)練/驗證/測試集。

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

使用以上設(shè)置，訓(xùn)練大約需要 12 個小時。

該設(shè)置使用數(shù)據(jù)并行，但也可以對無法放在單個 GPU 的超大模型使用模型并行。第一種選擇是設(shè)置張量并行，它將模型中的張量拆分到多個 GPU 上并行運算，你需要將 tensor-model-parallel-size 參數(shù)更改為所需的 GPU 數(shù)量。第二種選擇是流水線并行，它將模型按層分成大小相等的幾塊。參數(shù) pipeline-model-parallel-size 表示將模型分成的塊數(shù)。有關(guān)詳細(xì)信息，請參閱此 博客

將模型轉(zhuǎn)換為 ?? Transformers

訓(xùn)練結(jié)束后，我們希望使用 transformers 庫對該模型進(jìn)行評估或?qū)⑵洳渴鸬缴a(chǎn)環(huán)境中。你可以按照 教程 將其轉(zhuǎn)換為 transformers 模型。例如，在訓(xùn)練完成后，你可以復(fù)制第 150k 次迭代的權(quán)重，并使用以下命令將文件 model_optim_rng.pt 轉(zhuǎn)換為 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

請注意，如果你打算使用原始的分詞器，你將需要在轉(zhuǎn)換后將生成的詞匯文件和合并表替換為我們之前介紹的原始文件。

不要忘記將你的模型推送到 hub 并與社區(qū)分享，只需三行代碼 ??:

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 還可以有效地處理大模型推理。如果你訓(xùn)練了一個非常大的模型 (例如訓(xùn)練時使用了模型并行)，你可以通過以下命令輕松地進(jìn)行推理:

from transformers import AutoModelForCausalLM

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

這將在內(nèi)部調(diào)用 accelerate 庫 自動在你可用的設(shè)備 (GPU、CPU RAM) 之間分配模型權(quán)重。

免責(zé)聲明: 我們已經(jīng)證明任何人都可以使用 Megatron-LM 來訓(xùn)練語言模型。問題是我們需要考慮什么時候使用它。由于額外的預(yù)處理和轉(zhuǎn)換步驟，這個框架顯然增加了一些時間開銷。因此，重要的是你要考慮哪個框架更適合你的需求和模型大小。我們建議將其用于預(yù)訓(xùn)練模型或微調(diào)，但可能不適用于中型模型的微調(diào)。 APITrainer 和 accelerate 庫對于模型訓(xùn)練同樣也非常方便，并且它們與設(shè)備無關(guān)，為用戶提供了極大的靈活性。

恭喜 ?? 現(xiàn)在你學(xué)會了如何在 Megatron-LM 框架中訓(xùn)練 GPT2 模型并使其支持 transformers！

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

作者: Loubna Ben Allal

譯者: gxy-gxy

審校/排版: zhongdongy (阿東)文章來源地址http://www.zghlxwxcb.cn/news/detail-463603.html

到了這里，關(guān)于如何使用 Megatron-LM 訓(xùn)練語言模型的文章就介紹完了。如果您還想了解更多內(nèi)容，請在右上角搜索TOY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！