Megatron-LM を用いた言語モデルのトレーニング

概要

Megatron-LM は、NVIDIA の Applied Deep Learning Research チームが開発した高性能フレームワークで、GPU 上で大規模トランスフォーマーモデルを効率的に事前学習することを目的としています。Hugging Face Trainer API や accelerate ライブラリに比べて設定はやや複雑ですが、Megatron-LM は専門的なデータ処理とハードウェアレベルの最適化により大幅な高速化を実現します。

Megatron-LM の技術的最適化

Megatron-LM は、最適化されたデータロードとカーネル融合という 2 つの主要な仕組みで、標準的な PyTorch ループに比べて優れたトレーニング効率を実現します。

効率的な DataLoader

Megatron-LM は、トレーニング開始前にデータをトークン化しシャッフルする DataLoader を使用します。インデックス付きのシーケンスを一度だけ計算し、トレーニングパラメータに基づいてエポック数を決定します。このアプローチにより、次のエポックに入る前にデータセット全体を反復する必要がなくなり、学習曲線が滑らかになりトレーニング時間が短縮されます。

結合された CUDA カーネル

メモリオーバーヘッドを最小化するため、Megatron-LM は結合された CUDA カーネルを使用します。標準的な PyTorch の操作では、データがメモリから取得され、計算され、各操作ごとに再び保存されます。カーネル融合により、類似した操作が単一のハードウェア操作に統合され、中間結果が GPU レジスタに保持され、メインメモリへコピーされることがなくなります。

さらに、Megatron-LM は NVIDIA Apex ライブラリの結合された AdamW 実装を使用しており、標準的な PyTorch 実装よりも高速なパフォーマンスを提供します。

実装ワークフロー

Megatron-LM でモデル（例: GPT-2）をトレーニングするには、環境設定、データ前処理、トレーニング、モデル変換の 4 つのステップからなるパイプラインが必要です。

環境設定

推奨される設定は NGC から提供される NVIDIA PyTorch コンテナを使用することで、必要な依存関係がすべて含まれています。手動でインストールする場合は、最新バージョンの PyTorch、CUDA、NCCL、NVIDIA Apex、そして nltk ライブラリが必要です。また、トークナイザーの vocab.json と merges.txt ファイルを Megatron-LM ディレクトリ内に配置する必要があります。

データ前処理

データは処理前に、1 行につき 1 サンプルの緩やかな JSON 形式に変換する必要があります。その後、Megatron-LM は tools/preprocess_data.py を使用してデータをトークン化、シャッフルし、バイナリ形式に変換します。このプロセスで .idx と .bin ファイルが生成され、トレーニングフェーズで必要となります。dataset-impl は 'lazy'、'cached'、'mmap' のいずれかに設定できます。

トレーニング実行

torch.distributed.launch を使用してトレーニングを開始します。110M パラメータのモデル（例: CodeParrot-small）を 8 GPU でトレーニングすると、以下の構成で約 12 時間かかります。

  • Architecture: 12 層、隠れサイズ 768、注意ヘッド 12。
  • Sequence Length: 1024。
  • Optimizer: コサイン学習率減衰を伴う AdamW。
  • Batch Size: マイクロバッチサイズ 12、グローバルバッチサイズ 192。

モデル並列化戦略

単一 GPU に収まらないほど大きなモデルに対して、Megatron-LM は 2 種類のモデル並列化をサポートしています：

  1. Tensor Parallelism（テンソル並列）: tensor-model-parallel-size パラメータを使用して、単一のトランスフォーマーモジュールの実行を複数の GPU に分割します。
  2. Pipeline Parallelism（パイプライン並列）: pipeline-model-parallel-size パラメータを使用して、トランスフォーマーモジュールを同等のサイズのステージに分割し、GPU 間で実行します。

Hugging Face Transformers との統合

Megatron-LM で学習したモデルを評価や本番環境で使用するには、transformers ライブラリがサポートする形式に変換する必要があります。これは、convert_megatron_gpt2_checkpoint.py スクリプトを使用して model_optim_rng.pt チェックポイントファイルを pytorch_model.bin ファイルに変換することで実現します。

変換後は、AutoModelForCausalLM を使用してモデルをロードできます。モデル並列化で学習した非常に大規模なモデルの場合、device_map="auto" 引数を使用すると、accelerate ライブラリを介して利用可能な GPU と CPU の RAM に重みを自動的に割り当てることができます。

フレームワーク選択の指針

Megatron-LM は高度に最適化されているため、大規模モデルの事前学習や長時間のファインチューニングに最適です。ただし、前処理や変換ステップでオーバーヘッドが発生します。中規模モデルの短時間のファインチューニングには、デバイスに依存せず柔軟性の高い Hugging Face Trainer API と accelerate ライブラリの使用が推奨されます。

Sources