在 PyTorch 中訓(xùn)練大語言模型不僅僅是寫一個(gè)訓(xùn)練循環(huán)這么簡單。我們通常需要將模型分布在多個(gè)設(shè)備上，并使用許多優(yōu)化技術(shù)以實(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 模型的另一個(gè)流行工具，它是 NVIDIA 應(yīng)用深度學(xué)習(xí)研究團(tuán)隊(duì)開發(fā)的一個(gè)強(qiáng)大框架。與?accelerate?和?Trainer?不同，Megatron-LM 使用起來并不簡單，對(duì)于初學(xué)者來說可能難以上手。但它針對(duì) 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é)的講解之前，讓我們首先了解是什么讓這個(gè)框架比其他框架更高效。本節(jié)的靈感來自這篇關(guān)于使用 Megatron-DeepSpeed 進(jìn)行 BLOOM 訓(xùn)練的精彩 博客，請(qǐng)參閱該博客以獲取更多詳細(xì)信息，因?yàn)樵摬┛椭荚趯?duì) Megatron-LM 進(jìn)行詳細(xì)的介紹。

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

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

融合 CUDA 內(nèi)核

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

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

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

現(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ì)信息，請(qǐng)參閱 文檔。如果你不想使用此容器，則需要安裝最新的 pytorch、cuda、nccl 和 NVIDIA APEX 版本和?nltk?庫。

在安裝完 Docker 之后，你可以使用以下命令運(yùn)行容器 (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)重的模型倉庫中找到，請(qǐng)參閱 GPT2 庫。你還可以使用?transformers?訓(xùn)練自己的分詞器。你可以查看 CodeParrot 項(xiàng)目 以獲取實(shí)際示例?，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ù)作為示例。

我們需要對(duì)預(yù)訓(xùn)練數(shù)據(jù)進(jìn)行預(yù)處理。首先，你需要將其轉(zhuǎn)換為 json 格式，一個(gè) json 的一行包含一個(gè)文本樣本。如果你正在使用 ?? Datasets，這里有一個(gè)關(guān)于如何做到這一點(diǎn)的例子 (請(qǐng)?jiā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?選項(xiàng)指的是預(yù)處理中使用的線程數(shù)量和分配給每個(gè)線程的數(shù)據(jù)塊大小。dataset-impl?指的是索引數(shù)據(jù)集的實(shí)現(xiàn)方式，包括 ['lazy', 'cached', 'mmap']。這將輸出?codeparrot_content_document.idx?和 ?codeparrot_content_document.bin?兩個(gè)文件用于訓(xùn)練。

訓(xùn)練

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

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

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

訓(xùn)練結(jié)束后，我們希望使用?transformers?庫對(duì)該模型進(jìn)行評(píng)估或?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


請(qǐng)注意，如果你打算使用原始的分詞器，你將需要在轉(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)練了一個(gè)非常大的模型 (例如訓(xùn)練時(shí)使用了模型并行)，你可以通過以下命令輕松地進(jìn)行推理:

from?transformers?import?AutoModelForCausalLM

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


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

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

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

文中提到的鏈接請(qǐng)?jiān)L問?閱讀原文?查看哦~

英文原文:?https://hf.co/blog/megatron-training

作者: Loubna Ben Allal

譯者: gxy-gxy

審校/排版: zhongdongy (阿東)