PyTorch完全にシャーディングされたデータパラレルを使用して、大規模モデルのトレーニングを加速する

使用することで、PyTorchは大規模モデルのトレーニングを加速します

この投稿では、Accelerate ライブラリを活用して大規模なモデルのトレーニングを行う方法について説明します。これにより、ユーザーは PyTorch FullyShardedDataParallel (FSDP) の最新機能を活用することができます。

機械学習 (ML) モデルのスケール、サイズ、およびパラメータがますます増加するにつれ、ML プラクティショナーは自身のハードウェア上でそのような大規模なモデルをトレーニングしたり、ロードしたりすることが困難になっています。 一方で、大規模なモデルは小さなモデルと比較して学習が速く（データと計算効率が高く）、パフォーマンスも著しく向上することがわかっています [1]。しかし、そのようなモデルをほとんどの利用可能なハードウェア上でトレーニングすることは困難です。

大規模なMLモデルをトレーニングするためには、分散トレーニングが重要です。 分散トレーニング の分野では、最近重要な進展がありました。最も注目すべき進展のいくつかは以下のとおりです：

1. ZeROを用いたデータ並列化 – Zero Redundancy Optimizer [2]
   1. ステージ1：データ並列ワーカー/ GPU間でオプティマイザーの状態を分割
   2. ステージ2：データ並列ワーカー/ GPU間でオプティマイザーの状態と勾配を分割
   3. ステージ3：データ並列ワーカー/ GPU間でオプティマイザーの状態、勾配、およびモデルパラメータを分割
   4. CPUオフロード：ZEROステージ2をベースに、勾配とオプティマイザーの状態をCPUにオフロードする [3]
2. Tensor Parallelism [4]：巨大なパラメータを持つ各レイヤーのパラメータをアクセラレータ/ GPU間で賢明に分割することで、並列計算を実現し、高コストな通信同期オーバーヘッドを回避するモデル並列化の形式です。
3. Pipeline Parallelism [5]：モデルの異なるレイヤーを異なるアクセラレータ/ GPUに配置し、パイプライン処理を使用してすべてのアクセラレータを同時に実行します。たとえば、2番目のアクセラレータ/ GPUは、最初のマイクロバッチで計算を行い、最初のアクセラレータ/ GPUは2番目のマイクロバッチで計算を行います。
4. 3D Parallelism [3]：ZEROを使用したデータ並列化 + Tensor Parallelism + Pipeline Parallelismを用いて、数百億のパラメータを持つ巨大なモデルをトレーニングします。たとえば、BigScienceの176Bパラメータの言語モデルはこれを使用しています [6]。

この投稿では、ZeROを使用したデータ並列化と具体的には最新のPyTorch機能 FullyShardedDataParallel (FSDP) について説明します。 DeepSpeed と FairScale はZERO論文の核心的なアイデアを実装しています。これらはすでに transformers Trainerに統合され、素晴らしいブログ「DeepSpeedとFairScaleを使用したZeROによるより多くのフィットと高速トレーニング」[10] によって補完されています。PyTorchは最近、Fairscale FSDPをPyTorch Distributedに追加の最適化と共に統合しました。

• 実際のデータなしで効率的なテーブルの事前学習：TAPEXへの導入
• 敵対的なデータを使用してモデルを動的にトレーニングする方法
• 🤗 Datasetsでの新しいオーディオとビジョンのドキュメンテーションを紹介します

GPT-2 Large (762M) および XL (1.5B) モデルバリアントを使用した因果言語モデリングのタスクを見ていきます。

以下は、GPT-2モデルの事前トレーニングのコードです。公式の因果言語モデリングの例と類似していますが、n_train (2000) および n_val (500) の2つの引数が追加されており、全データでの前処理/トレーニングを行わずに、クイックな概念実証ベンチマークを実行するためのものです。

run_clm_no_trainer.py

コマンド accelerate config を実行した後のサンプルのFSDP設定：

compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: FSDP
fsdp_config:
  min_num_params: 2000
  offload_params: false
  sharding_strategy: 1
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: 'no'
num_machines: 1
num_processes: 2
use_cpu: false

Multi-GPU FSDP

ここでは、シングルノードのマルチGPU環境で実験を行います。さまざまな設定で分散データ並列 (DDP) とFSDPのパフォーマンスを比較します。まず、GPT-2 Large (762M) モデルを使用し、DDPは特定のバッチサイズでメモリ不足のエラーが発生せずに動作します。次に、GPT-2 XL (1.5B) モデルを使用し、DDPはバッチサイズが1でもメモリ不足のエラーが発生します。FSDPを使用すると、GPT-2 Largeモデルのバッチサイズを大きくすることができ、DDPではできなかったGPT-2 XLモデルのトレーニングも可能になります。

ハードウェアのセットアップ：2X24GBのNVIDIA Titan RTX GPU。

GPT-2 Largeモデル（762Mパラメータ）のトレーニングコマンド：

export BS=#`OOMエラーが発生しないまで、異なるバッチサイズで試してください。
#つまり、大きなバッチサイズから始めて、GPUに収まるまで徐々に減らしてください`

time accelerate launch run_clm_no_trainer.py \
--model_name_or_path gpt2-large \
--dataset_name wikitext \
--dataset_config_name wikitext-2-raw-v1 \
--per_device_train_batch_size $BS 
--per_device_eval_batch_size $BS 
--num_train_epochs 1 
--block_size 12

サンプルのFSDP実行：

Table 1: GPT-2 Large（762M）モデルのFSDPのベンチマーク

表1からわかるように、DDPに比べて、FSDPはより大きなバッチサイズを可能にします。CPUオフロード設定なしとCPUオフロード設定ありの場合、それぞれ2倍から3倍までのバッチサイズを実現できます。トレーニング時間では、混合精度を使用したDDPが最も高速で、次にZEROステージ2とステージ3を使用したFSDPです。因果言語モデリングのタスクは常に固定のコンテキストシーケンス長（–block_size）を持つため、FSDPを使用したトレーニング時間の短縮はそれほど大きくありませんでした。動的バッチ処理を必要とするアプリケーションでは、より大きなバッチサイズを可能にするFSDPは、トレーニング時間の高速化が期待できます。FSDPの混合精度サポートは現在、トランスフォーマーにいくつかの問題があります。これがサポートされると、トレーニング時間の高速化がさらに大幅に改善されます。

GPUメモリに収まらない巨大なモデルのトレーニングを可能にするためのCPUオフロード

GPT-2 XLモデル（1.5Bパラメータ）のトレーニングコマンド：

export BS=#`OOMエラーが発生しないまで、異なるバッチサイズで試してください。
#つまり、大きなバッチサイズから始めて、GPUに収まるまで徐々に減らしてください`

time accelerate launch run_clm_no_trainer.py \
--model_name_or_path gpt2-xl \
--dataset_name wikitext \
--dataset_config_name wikitext-2-raw-v1 \
--per_device_train_batch_size $BS 
--per_device_eval_batch_size $BS 
--num_train_epochs 1 
--block_size 12

Table 2: GPT-2 XL（1.5B）モデルのFSDPのベンチマーク

表2からわかるように、DDP（fp16ありとなし）はバッチサイズ1でさえ実行できず、CUDA OOMエラーが発生します。FSDPのZero-Stage 3を使用すると、2つのGPUでバッチサイズ5（有効なバッチサイズ=10（5 X 2））で実行できます。2つのGPUを使用する場合、CPUオフロードを使用すると、最大バッチサイズを1つのGPUあたり14に増やすことができます。これにより、CPUオフロードを使用したFSDPは、1つのGPUでバッチサイズ10のGPT-2 1.5Bモデルのトレーニングを可能にします。これにより、最小限の計算リソースを持つML開発者でも、このような大きなモデルのトレーニングが可能になり、大規模モデルのトレーニングが民主化されます。

FSDP統合の機能と制限

AccelerateがFSDP統合に提供する現在のサポートと既知の制限について説明します。

FSDPサポートに必要なPyTorchバージョン : PyTorch Nightly（またはリリース後の1.12.0）です。FSDPを有効にしたモデルの保存は、最近の修正版でのみ利用可能です。

CLIを介した設定：

1. シャーディング戦略 : [1] FULL_SHARD、[2] SHARD_GRAD_OP
2. 最小パラメータ数 : FSDPのデフォルトの自動ラッピングのためのパラメータの最小数。
3. オフロードパラメータ : パラメータと勾配をCPUにオフロードするかどうかを決定します。

より詳細な制御のために、ユーザーはFullyShardedDataParallelPluginを利用することができます。このプラグインでは、auto_wrap_policy、backward_prefetch、およびignored_modulesを指定することができます。

このクラスのインスタンスを作成した後、ユーザーはそれを Accelerator オブジェクトの作成時に渡すことができます。

これらのオプションの詳細については、PyTorch FullyShardedDataParallel コードを参照してください。

次に、min_num_params の設定の重要性を見ていきます。以下は FSDP Auto Wrap Policy の重要性について説明した [8] からの抜粋です。

(出典: リンク)

default_auto_wrap_policy を使用する場合、ある層のパラメータ数が min_num_params よりも多い場合、その層は FSDP モジュールでラップされます。BERT-Large (330M) モデルを GLUE MRPC タスクでファインチューニングするためのコードは、ピークメモリ使用量の追跡のためのユーティリティを追加しながら FSDP 機能を正しく使用する方法を示す公式の完全な NLP の例です。

fsdp_with_peak_mem_tracking.py

私たちは Accelerate の追跡機能のサポートを活用して、トレーニングと評価のピークメモリ使用量と評価メトリクスをログに記録しています。以下は wandb run からのプロットのスナップショットです。

DDP は FSDP よりもメモリを2倍消費することが観察されます。オートラップのない FSDP はオートラップのある FSDP よりもメモリを多く消費しますが、DDP よりは少ないです。min_num_params=2k のオートラップを使用した FSDP は min_num_params=1M の設定と比較してわずかにメモリを少なく使用します。これは、FSDP Auto Wrap Policy の重要性を強調し、ユーザーは min_num_params を自由に調整してメモリを節約する設定を見つけるべきであり、多くの通信オーバーヘッドを引き起こしていないことを示しています。PyTorch チームは、この設定のための自動調整ツールに取り組んでいます（詳細は [8] を参照）。

注意事項

• PyTorch FSDP は、サブモジュールを自動的にラップし、パラメータをフラット化し、パラメータをその場でシャードします。そのため、モデルのラッピング前に作成されたオプティマイザは破損し、より多くのメモリを占有します。したがって、モデルを作成する前にモデルの準備をすることを強くお勧めします。単一のモデルの場合、Accelerate はモデルを自動的にラップし、警告メッセージとともにオプティマイザを作成します。

  FSDP の警告: FSDP を使用する場合、オプティマイザを作成する前にモデルの準備を呼び出すことが効率的でおすすめです

ただし、FSDP を使用する場合のモデルとオプティマイザの準備の推奨方法は以下の通りです:

model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+ model = accelerator.prepare(model)

optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)

- model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(model,
-        optimizer, train_dataloader, eval_dataloader, lr_scheduler
-    )

+ optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+         optimizer, train_dataloader, eval_dataloader, lr_scheduler
+        )

• 単一のモデルの場合、複数のパラメータグループでオプティマイザを作成し、それらをまとめて準備を呼び出した場合、パラメータグループは失われ、以下の警告が表示されます:

  FSDP の警告: FSDP を使用する場合、ネストされたモジュールのラッピングとパラメータのフラット化により、いくつかのパラメータグループが1つに統合されます。

これは、ラッピング前に作成されたパラメータグループが、ネストされた FSDP モジュールを1次元配列にフラット化するため、ラッピング後には意味を持たなくなるためです（これは多くのレイヤーを消費することができます）。例えば、以下は GPU 0 上の FSDP モデルの名前付きパラメーターです（2つの GPU を使用する場合）。ここでは、非ラップされた BERT-Base モデルの [bias、LayerNorm.weight] の名前付きパラメーターに対して重み減衰を適用していない場合、以下の FSDP ラップされたモデルにはこれらの文字列のいずれかを持つ名前付きパラメーターはなく、それらのレイヤーのパラメーターは他のさまざまなレイヤーのパラメーターと連結されます。詳細はこの issue に記載されています（ The original model parameters' .grads are not set, meaning that they cannot be optimized separately (which is why we cannot support multiple parameter groups) ）。

```
{
'_fsdp_wrapped_module.flat_param': torch.Size([494209]),

'_fsdp_wrapped_module._fpw_module.bert.embeddings.word_embeddings._fsdp_wrapped_module.flat_param': torch.Size([11720448]),

'_fsdp_wrapped_module._fpw_module.bert.encoder._fsdp_wrapped_module.flat_param': torch.Size([42527232])
}
```

• 複数のモデルの場合、エラーが発生しないように、最適化する前にモデルを準備する必要があります。

• FSDPでは、現在、混合精度はサポートされていません。PyTorchがサポートを修正するのを待っています。

(出典: リンク )

上記のワークフローは、FSDPがアクティブになったときに裏側で何が起こるかを概説しています。まず、DDPがどのように機能し、FSDPがそれをどのように改善するかを理解しましょう。DDPでは、各ワーカー/アクセラレータ/GPUには、モデルのパラメータ、勾配、オプティマイザの状態の完全なレプリカがあります。各ワーカーは異なるデータのバッチを取得し、順方向のパスを経て、損失が計算され、バックワードパスが行われて勾配が生成されます。次に、全ワーカーから勾配を受け取り、平均化するオールリデュース操作が行われます。これにより、各ワーカーは同じグローバル勾配を持つようになり、オプティマイザはモデルのパラメータを更新するために使用されます。完全なレプリカを持つことにより、各GPU上で多くの重複したメモリが消費され、バッチサイズやモデルのサイズが制限されることがわかります。

FSDPは、オプティマイザの状態、勾配、およびモデルのパラメータをデータ並列のワーカーに分散することによって、これを正確に対処します。さらに、これにより、すべてのテンソルのCPUオフロードが容易になり、利用可能なGPUメモリに収まらない大規模なモデルの読み込みが可能になります。DDPと同様に、各ワーカーは異なるデータのバッチを取得します。順方向のパス中、CPUオフロードが有効になっている場合、ローカルシャードのパラメータがまずGPU/アクセラレータにコピーされます。次に、各ワーカーは指定されたFSDPラップドモジュール/レイヤーのためにオールギャザ操作を実行し、必要なパラメータをすべて取得し、計算が行われ、他のワーカーのパラメータシャードが解放/空にされます。これはすべてのFSDPモジュールについて続きます。順方向のパス後、損失が計算され、バックワードパス中に再びオールギャザ操作が実行され、指定されたFSDPモジュールのすべての必要なパラメータが取得され、計算が行われ、他のワーカーのシャードが解放されます。次に、ローカル勾配が平均化され、リダクションスキャッタ操作を使用して各関連するワーカーにシャードされます。これにより、各ワーカーはローカルシャードのパラメータを更新することができます。CPUオフロードがアクティベートされている場合、勾配はCPUに渡され、CPU上で直接パラメータが更新されます。

PyTorch FSDPの動作やこの機能を使用した包括的な実験に関する詳細は、[7, 8, 9]を参照してください。

PyTorch FSDPの統合部分で問題が発生した場合は、accelerateでIssueを開いてください。

ただし、PyTorch FSDPの設定や展開に問題がある場合は、各ドメインの専門家にお問い合わせいただく必要がありますので、代わりにPyTorchのIssueを開いてください。

[1] Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers

[2] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

[3] DeepSpeed: Extreme-scale model training for everyone – Microsoft Research

[4] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

[5] Introducing GPipe, an Open Source Library for Efficiently Training Large-scale Neural Network Models

[6] Which hardware do you need to train a 176B parameters model?

[7] Introducing PyTorch Fully Sharded Data Parallel (FSDP) API | PyTorch

[8] Getting Started with Fully Sharded Data Parallel(FSDP) — PyTorch Tutorials 1.11.0+cu102 documentation

[9] Training a 1 Trillion Parameter Model With PyTorch Fully Sharded Data Parallel on AWS | by PyTorch | PyTorch | Mar, 2022 | VoAGI

[10] Fit More and Train Faster With ZeRO via DeepSpeed and FairScale

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