Megatron-LMを使用して言語モデルをトレーニングする方法

'Megatron-LMの言語モデルトレーニング方法'

PyTorchで大規模な言語モデルをトレーニングするには、単純なトレーニングループだけでは不十分です。通常、複数のデバイスに分散しており、安定した効率的なトレーニングのための多くの最適化技術があります。Hugging Face 🤗 Accelerateライブラリは、トレーニングループに非常に簡単に統合できるように、GPUとTPUを跨いで分散トレーニングをサポートするために作成されました。🤗 TransformersもTrainer APIを介して分散トレーニングをサポートしており、トレーニングループの実装を必要とせずにPyTorchでの完全なトレーニングを提供します。

大規模なトランスフォーマーモデルを事前トレーニングするための研究者の間でのもう一つの人気ツールはMegatron-LMです。これはNVIDIAのApplied Deep Learning Researchチームによって開発された強力なフレームワークです。🤗 AccelerateとTrainerとは異なり、Megatron-LMの使用は直感的ではなく、初心者には少し抵抗があるかもしれません。しかし、これはGPU上でのトレーニングに最適化されており、いくつかの高速化を提供することができます。このブログ記事では、Megatron-LMを使用してNVIDIAのGPU上で言語モデルをトレーニングし、それをtransformersと一緒に使用する方法を学びます。

このフレームワークでGPT2モデルをトレーニングするためのさまざまなステップを紹介します。これには以下が含まれます。

• 環境のセットアップ
• データの前処理
• トレーニング
• モデルの🤗 Transformersへの変換

なぜMegatron-LMを選ぶのか？

トレーニングの詳細に入る前に、他のフレームワークよりもこのフレームワークが効率的である理由を理解しましょう。このセクションは、Megatron-DeepSpeedでのBLOOMトレーニングについての素晴らしいブログから着想を得ています。詳細については参照してください。このブログ記事はMegatron-LMへの優しい入門を提供することを目的としています。

• SetFit プロンプトなしで効率的なフューショット学習
• 非常に大規模な言語モデルとその評価方法
• 🤗評価による言語モデルのバイアスの評価

データローダー

Megatron-LMには、データがトークン化され、トレーニング前にシャッフルされる効率的なデータローダーが付属しています。また、データは番号付きのシーケンスに分割され、それらは計算が必要な場合にのみ計算されるようにインデックスで保存されます。インデックスを作成するために、エポック数はトレーニングパラメータに基づいて計算され、順序が作成され、その後シャッフルされます。これは通常の場合とは異なり、データセット全体を繰り返し処理してから2番目のエポックのために繰り返すというものです。これにより、学習曲線が滑らかになり、トレーニング中の時間が節約されます。

組み込みCUDAカーネル

GPU上で計算を実行する場合、必要なデータはメモリから取得され、計算が実行され、結果がメモリに保存されます。簡単に言えば、組み込みカーネルのアイデアは、通常はPyTorchによって別々に実行される類似の操作を、単一のハードウェア操作に統合することです。そのため、複数の個別の計算で行われるメモリ移動の回数を減らします。以下の図は、カーネルフュージョンのアイデアを示しています。これは、詳細について説明しているこの論文からインスピレーションを受けています。

f、g、hが1つのカーネルで結合された場合、fとgの中間結果x’とy’はGPUレジスタに保存され、hによって即座に使用されます。しかし、フュージョンがない場合、x’とy’はメモリにコピーされ、hによって読み込まれる必要があります。したがって、カーネルフュージョンは計算に著しいスピードアップをもたらします。Megatron-LMはまた、PyTorchの実装よりも高速なApexのFused AdamWの実装を使用しています。

Megatron-LMのようにDataLoaderをカスタマイズし、transformersでApexのFused optimizerを使用することもできますが、カスタムのFused 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形式に変換する必要があります。1つのjsonに1つのテキストサンプルが含まれるようにします。🤗 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-eod

workersとchunk_sizeオプションは、前処理に使用されるワーカーの数と各ワーカーに割り当てられるデータのチャンクサイズを指します。dataset-implは、インデックス化されたデータセットの実装モードを指します。[‘lazy’、’cached’、’mmap’]から選択できます。これにより、2つのファイル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時間かかります。

このセットアップでは、データ並列処理を使用していますが、1つのGPUに収まらない非常に大きなモデルの場合は、モデル並列処理を使用することも可能です。最初のオプションは、テンソル並列処理で、単一のトランスフォーマーモジュールの実行を複数のGPUに分割します。 tensor-model-parallel-sizeパラメータを希望するGPUの数に変更する必要があります。 2番目のオプションは、パイプライン並列処理で、トランスフォーマーモジュールを均等なサイズのステージに分割します。 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

注意してください。トークナイザをそこからロードする予定の場合は、変換後に生成された語彙ファイルとマージテーブルを、前述のオリジナルのもので置き換える必要があります。

モデルをハブにプッシュしてコミュニティと共有するのを忘れないでください。たった3行のコードで完了します🤗：

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は大規模なモデルの推論も効率的に処理します。モデル並列処理を使用して非常に大きなモデル（例：Model Parallelismを使用）をトレーニングした場合、次のコマンドで推論に簡単に使用できます：

from transformers import AutoModelForCausalLM

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

これにより、加速ライブラリが内部で使用され、利用可能なデバイス（GPU、CPU RAM）にモデルの重みが自動的にディスパッチされます。

免責事項：Megatron-LMを使用して言語モデルをトレーニングできることを示しました。問題は、いつ使用するかです。このフレームワークは、追加の前処理と変換ステップにより、明らかに時間のオーバーヘッドが発生します。したがって、ケースとモデルのサイズに応じて、より適切なフレームワークを選択することが重要です。事前トレーニングモデルや拡張されたファインチューニングに使用することをお勧めしますが、VoAGIサイズのモデルの短いファインチューニングにはおそらく向いていません。 Trainer APIと accelerate ライブラリもモデルトレーニングに非常に便利で、デバイスに依存せず、ユーザーに大きな柔軟性を提供します。

おめでとうございます🎉 これで、Megatron-LMでGPT2モデルをトレーニングし、transformersでサポートされるようにする方法を知りました！

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!