PyTorch DDP、Accelerate、Transformers Trainer を用いた分散トレーニング

Hugging Face は、PyTorch における分散学習の実装のために、ネイティブの Distributed Data Parallelism (DDP) から Accelerate ライブラリ、最終的に Transformers Trainer API へと至る 3 つの抽象化レベルの進化を詳述しています。この進化により、開発者はトレーニングループに対する細かな制御と、分散設定の高レベルな自動化の間で選択できるようになります。

ネイティブ PyTorch Distributed Data Parallelism (DDP)

PyTorch DDP は、単一マシン上の複数 GPU（マルチ GPU）でも、複数マシンにまたがる複数 GPU（マルチノード）でも、トレーニングを可能にします。モデルを各 GPU にコピーし、loss.backward() 呼び出し時にすべてのコピー間で勾配を平均化することで、デバイス間の重みの一貫性を保ちます。

ネイティブ DDP を実装するには、いくつかの手動設定ステップが必要です：

• Process Group Setup: 開発者は setup と cleanup 関数を定義し、プロセスグループを初期化および破棄し、通信のために MASTER_ADDR と MASTER_PORT を指定する必要があります。
• Model Wrapping: モデルは DistributedDataParallel (DDP) モジュールでラップする必要があり、オプティマイザはこのラップされたモデルに基づいて宣言し、勾配が正しく計算されるようにします。
• Execution: スクリプトは通常、torchrun コマンドラインモジュールを使用して起動し、ノード数とノードあたりのプロセス数を指定します。

🤗 Accelerate で分散を簡素化

Accelerate は pytorch.distributed の軽量ラッパーで、同じコードを単一 GPU、複数 GPU、または TPU 上で最小限の変更で実行できるようにします。手動でのプロセスグループ設定が不要になり、デバイス配置が簡素化されます。

技術的な改善

Accelerate は Accelerator クラスを導入し、accelerator.prepare() の一呼び出しでモデル、オプティマイザ、データローダーの分散を処理します。これにより、手動の .to(rank) 呼び出しや DDP のラップが不要になります。

簡素化に加えて、Accelerate はカスタムサンプラーを用いてメモリ効率を向上させます。デバイス間でデータローダーの完全コピーを複数作成する代わりに、Accelerate は元のデータセットのフルコピーをメモリ上に一つだけ保持し、データのサブセットを利用可能なノードに分割します。これにより、非常に大規模なデータセットでのトレーニング時にメモリが爆発するのを防ぎます。

ノートブック統合

Jupyter Notebook で作業するユーザー向けに、Accelerate は notebook_launcher ユーティリティを提供します。これにより、トレーニング関数と使用するプロセス数を渡すだけで、ノートブックから直接マルチ GPU トレーニングを起動できます。

🤗 Trainer による高レベル抽象化

Transformers の Trainer API は最も高い抽象化レベルを提供し、分散トレーニングに関連するほぼすべてのボイラープレートコードを排除します。ユーザーが明示的なトレーニングループを書く必要なく、基盤となる分散ロジックを自動的に処理します。

Trainer を使用するには、開発者はハイパーパラメータを管理するために TrainingArguments を定義し、カスタム compute_loss 関数を実装するために Trainer をサブクラス化できます。Trainer はその後、分散システム全体でのトレーニングと評価プロセスを自動的に管理します。Accelerate の例と同様に、Trainer も notebook_launcher と統合でき、インタラクティブ環境での迅速な実験が可能です。