使用Megatron-Deepspeed訓(xùn)練GPT-2并生成文本
本文介紹如何使用GPU云服務(wù)器，使用Megatron-DeepSpeed框架訓(xùn)練GPT-2模型并生成文本。
背景信息
GPT-2模型是OpenAI于2018年在GPT模型的基礎(chǔ)上發(fā)布的新的無監(jiān)督NLP模型，當(dāng)時被稱為“史上最強通用NLP模型”。該模型可以生成連貫的文本段落，并且能在未經(jīng)預(yù)訓(xùn)練的情況下，完成閱讀理解、問答、機器翻譯等多項不同的語言建模任務(wù)。GPT-2模型尤其在文本生成上有著驚艷的表現(xiàn)，其生成的文本在上下文連貫性和情感表達上都超過了人們的預(yù)期。
GPT-2模型由多層單向Transformer的解碼器部分構(gòu)成，根據(jù)模型層次規(guī)模，GPT-2模型有以下幾個規(guī)格：
Megatron-Deepspeed框架結(jié)合了兩種主要技術(shù)：
Megatron-LM是由NVIDIA開源的Transformer 模型框架。
DeepSpeed是由Microsoft開源的一個深度學(xué)習(xí)優(yōu)化庫。
DeepSpeed團隊通過將DeepSpeed庫中的ZeRO分片（ZeRO sharding）數(shù)據(jù)并行（Data Parallelism）和管道并行（Pipeline Parallelism）與Megatron-LM中的張量并行（Tensor Parallelism）相結(jié)合，開發(fā)了一種基于3D并行的實現(xiàn)，這就是Megatron-DeepSpeed，它使得千億級參數(shù)量以上的大規(guī)模語言模型（LLM）的分布式訓(xùn)練變得更簡單、高效和有效。
本文基于阿里云GPU服務(wù)器，使用Megatron-DeepSpeed框架訓(xùn)練GPT-2模型并生成文本。
阿里云不對第三方模型“GPT-2”的合法性、安全性、準確性進行任何保證，阿里云不對由此引發(fā)的任何損害承擔(dān)責(zé)任。
您應(yīng)自覺遵守第三方模型的用戶協(xié)議、使用規(guī)范和相關(guān)法律法規(guī)，并就使用第三方模型的合法性、合規(guī)性自行承擔(dān)相關(guān)責(zé)任。
操作步驟
創(chuàng)建ECS實例
在ECS實例創(chuàng)建頁面，創(chuàng)建ECS實例。
關(guān)鍵參數(shù)說明如下，其他參數(shù)的配置，請參見自定義購買實例。
實例規(guī)格：選擇ecs.gn7i-c8g1.2xlarge（單卡NVIDIA A10）。
鏡像：使用云市場鏡像，名稱為aiacc-train-solution，您可以直接通過名稱搜索該鏡像，選擇最新版本即可。
公網(wǎng)IP：選中分配公網(wǎng)IPv4地址，帶寬計費方式選擇按使用流量，帶寬峰值選擇100 Mbps，以加快模型下載速度。
安裝Megatron-Deepspeed框架
使用root用戶遠程登錄ECS實例。具體操作，請參見通過密碼或密鑰認證登錄Linux實例。
執(zhí)行以下命令，啟動容器。
docker run -d -t --network=host --gpus all --privileged --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --name megatron-deepspeed -v /etc/localtime:/etc/localtime -v /root/.ssh:/root/.ssh nvcr.io/nvidia/pytorch:21.10-py3執(zhí)行以下命令，進入容器終端。
docker exec -it megatron-deepspeed bash執(zhí)行以下命令，下載Megatron-DeepSpeed框架。
git clone https://github.com/bigscience-workshop/Megatron-DeepSpeed執(zhí)行以下命令，安裝Megatron-DeepSpeed框架。
cd Megatron-DeepSpeed pip install -r requirements.txt
處理數(shù)據(jù)
本指南使用1GB 79K-record的JSON格式的OSCAR數(shù)據(jù)集。
執(zhí)行以下命令，下載數(shù)據(jù)集。
wget https://huggingface.co/bigscience/misc-test-data/resolve/main/stas/oscar-1GB.jsonl.xz wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt執(zhí)行以下命令，解壓數(shù)據(jù)集。
xz -d oscar-1GB.jsonl.xz執(zhí)行以下命令，預(yù)處理數(shù)據(jù)。
python3 tools/preprocess_data.py \ --input oscar-1GB.jsonl \ --output-prefix meg-gpt2 \ --vocab gpt2-vocab.json \ --dataset-impl mmap \ --tokenizer-type GPT2BPETokenizer \ --merge-file gpt2-merges.txt \ --append-eod \ --workers 8如果回顯信息類似如下所示，表示預(yù)處理數(shù)據(jù)完成。
執(zhí)行以下命令，新建data目錄。
mkdir data執(zhí)行以下命令，將處理好的數(shù)據(jù)移動到data目錄下。
mv meg-gpt2* ./data mv gpt2* ./data
預(yù)訓(xùn)練
本示例使用單機單卡的GPU實例完成GPT-2 MEDIUM模型的預(yù)訓(xùn)練。
創(chuàng)建預(yù)訓(xùn)練腳本文件。
執(zhí)行以下命令，創(chuàng)建預(yù)訓(xùn)練腳本文件。
vim pretrain_gpt2.sh按
i鍵，進入編輯模式，在文件中添加以下信息。#! /bin/bash # Runs the "345M" parameter model GPUS_PER_NODE=1 # Change for multinode config MASTER_ADDR=localhost MASTER_PORT=6000 NNODES=1 NODE_RANK=0 WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES)) DATA_PATH=data/meg-gpt2_text_document CHECKPOINT_PATH=checkpoints/gpt2 DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE --nnodes $NNODES --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT" python -m torch.distributed.launch $DISTRIBUTED_ARGS \ pretrain_gpt.py \ --tensor-model-parallel-size 1 \ --pipeline-model-parallel-size 1 \ --num-layers 24 \ --hidden-size 1024 \ --num-attention-heads 16 \ --micro-batch-size 4 \ --global-batch-size 8 \ --seq-length 1024 \ --max-position-embeddings 1024 \ --train-iters 5000 \ --lr-decay-iters 320000 \ --save $CHECKPOINT_PATH \ --load $CHECKPOINT_PATH \ --data-path $DATA_PATH \ --vocab-file data/gpt2-vocab.json \ --merge-file data/gpt2-merges.txt \ --data-impl mmap \ --split 949,50,1 \ --distributed-backend nccl \ --lr 0.00015 \ --lr-decay-style cosine \ --min-lr 1.0e-5 \ --weight-decay 1e-2 \ --clip-grad 1.0 \ --lr-warmup-fraction .01 \ --checkpoint-activations \ --log-interval 10 \ --save-interval 500 \ --eval-interval 100 \ --eval-iters 10 \ --fp16按
Esc鍵，輸入:wq后，按Enter鍵保存文件。
修改測試代碼。
Megatron源碼有一個斷言需要注釋掉，以保證代碼正常運行。
執(zhí)行以下命令，打開測試代碼文件。
vim /workspace/Megatron-DeepSpeed/megatron/model/fused_softmax.py +191按
i鍵，進入編輯模式，在assert mask is None, "Mask is silently ignored due to the use of a custom kernel"前加#。
按
Esc鍵，輸入:wq后，按Enter鍵保存文件。
預(yù)訓(xùn)練。
執(zhí)行以下命令，開始預(yù)訓(xùn)練。
nohup sh ./pretrain_gpt2.sh &
執(zhí)行如下命令，可以持續(xù)地查看nohup.out的輸出，達到監(jiān)控程序的效果。
tail -f nohup.out如果回顯信息類似如下所示，表示預(yù)訓(xùn)練完成。
說明預(yù)訓(xùn)練完成大概需要1小時30分鐘，如果超時斷開了ECS連接，重新遠程登錄ECS實例后，執(zhí)行以下命令，繼續(xù)查看預(yù)訓(xùn)練進度。預(yù)訓(xùn)練完成后，可以執(zhí)行
Ctrl+Z命令退出。docker exec -it megatron-deepspeed bash cd Megatron-DeepSpeed tail -f nohup.out
（可選）執(zhí)行以下命令，查看生成的模型checkpoint路徑。
本示例生成的模型checkpoint路徑設(shè)置在
/workspace/Megatron-DeepSpeed/checkpoints/gpt2。ll ./checkpoints/gpt2
使用GPT-2模型生成文本
執(zhí)行以下命令，安裝相關(guān)依賴。
說明由于網(wǎng)絡(luò)原因，執(zhí)行命令后可能會失敗，建議您多次嘗試。
pip install mpi4py如果回顯信息類似如下所示，表示依賴安裝完成。
創(chuàng)建文本生成腳本。
執(zhí)行以下命令，創(chuàng)建文本生成腳本。
vim generate_text.sh按
i鍵，進入編輯模式，在文件中增加以下內(nèi)容。#!/bin/bash CHECKPOINT_PATH=checkpoints/gpt2 VOCAB_FILE=data/gpt2-vocab.json MERGE_FILE=data/gpt2-merges.txt python tools/generate_samples_gpt.py \ --tensor-model-parallel-size 1 \ --num-layers 24 \ --hidden-size 1024 \ --load $CHECKPOINT_PATH \ --num-attention-heads 16 \ --max-position-embeddings 1024 \ --tokenizer-type GPT2BPETokenizer \ --fp16 \ --micro-batch-size 2 \ --seq-length 1024 \ --out-seq-length 1024 \ --temperature 1.0 \ --vocab-file $VOCAB_FILE \ --merge-file $MERGE_FILE \ --genfile unconditional_samples.json \ --num-samples 2 \ --top_p 0.9 \ --recompute按
Esc鍵，輸入:wq后，按Enter鍵保存文件。
執(zhí)行以下命令，生成文本。
sh ./generate_text.sh如果回顯信息類似如下所示，表示生成文本完成。
執(zhí)行以下命令，查看生成的JSON格式的文本文件。
vim unconditional_samples.json回顯信息類似如下所示。
了解更多AIGC實踐和GPU優(yōu)惠
活動入口：立即開啟AIGC之旅
