Fully Sharded Data Parallel

大きなバッチサイズで巨大モデルの学習を高速化するために、完全シャードされたデータ並列モデルを使用できます。 このデータ並列のパラダイムは、オプティマイザの状態、勾配、パラメータをシャードすることで、より多くのデータとより大きなモデルを収容できるようにします。 詳細とその利点については、Fully Sharded Data Parallel ブログをご覧ください。 最新の PyTorch の Fully Sharded Data Parallel (FSDP) 学習機能を統合しました。 設定でこれを有効にするだけです。

デフォルトの動作

お使いのマシンで、次のコマンドを実行してください:

accelerate config

and answer the questions asked. This will generate a config file that will be used automatically to properly set the default options when doing

accelerate launch my_script.py --args_to_my_script

例えば、リポジトリのルートからexamples/nlp_example.pyをFSDPを有効にして実行するには、次のようにします。

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch_policy: BACKWARD_PRE
  fsdp_forward_prefetch: false
  fsdp_cpu_ram_efficient_loading: true
  fsdp_offload_params: false
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_transformer_layer_cls_to_wrap: BertLayer
  fsdp_use_orig_params: true
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

accelerate launch examples/nlp_example.py

現在、AccelerateはCLIを通じて以下の設定をサポートしています:

fsdp_sharding_strategy: [1] FULL_SHARD（オプティマイザの状態、勾配、パラメータをシャード化）、[2] SHARD_GRAD_OP（オプティマイザの状態と勾配をシャード化）、[3] NO_SHARD（DDP）、[4] HYBRID_SHARD（ノード内でオプティマイザの状態、勾配、パラメータをシャード化し、各ノードが完全なコピーを保持）、[5] HYBRID_SHARD_ZERO2（ノード内でオプティマイザの状態と勾配をシャード化し、各ノードが完全なコピーを保持）。詳細については、公式のPyTorch ドキュメントを参照してください。

fsdp_offload_params : パラメータと勾配をCPUにオフロードするかどうかを決定します

fsdp_auto_wrap_policy: [1] TRANSFORMER_BASED_WRAP、[2] SIZE_BASED_WRAP、[3] NO_WRAP

fsdp_transformer_layer_cls_to_wrap: Transformers にのみ適用されます。fsdp_auto_wrap_policy=TRANSFORMER_BASED_WRAP を使用する場合、ラップ対象の Transformer レイヤーのクラス名を、カンマ区切りの文字列として指定できます(大文字と小文字を区別します)。例: BertLayer、GPTJBlock、T5Block、BertLayer,BertEmbeddings,BertSelfOutput。これは重要です。重みを共有するサブモジュール(埋め込みレイヤーなど)が、異なる FSDP ラップ単位に分割されてしまうことを防ぐためです。このポリシーを使用すると、Multi-Head Attention を含み、その後ろに数層の MLP レイヤーが続く各ブロックに対してラップが行われます。残りのレイヤー(共有埋め込みを含む)は、最も外側の FSDP 単位にまとめてラップされるため、Transformer ベースのモデルにはこの設定を使用してください。Transformer モデルでは、Do you want to use the model's_no_split_modulesto wrap. It will try to usemodel._no_split_modules` を使用するかを yes と答えることで、model._no_split_modules を利用できます。

fsdp_min_num_params: fsdp_auto_wrap_policy=SIZE_BASED_WRAP を使用する場合の最小パラメータ数。

fsdp_backward_prefetch_policy: [1] BACKWARD_PRE、[2] BACKWARD_POST、[3] NO_PREFETCH

fsdp_forward_prefetch: True の場合、FSDP はフォワードパスの実行中に、次の all-gather を明示的にプリフェッチします。プリフェッチは最初のイテレーションの実行順序に従うため、静的グラフのモデルでのみ使用してください。つまり、モデルの実行中にサブモジュールの順序が動的に変更される場合は、この機能を有効にしないでください。

fsdp_state_dict_type: [1] FULL_STATE_DICT、[2] LOCAL_STATE_DICT、[3] SHARDED_STATE_DICT

fsdp_use_orig_params: True の場合、init 時に非一様な requires_grad を許可し、凍結パラメータと学習可能パラメータが混在する構成をサポートします。この設定は、この投稿で議論されているパラメータ効率の良いファインチューニングのような場合に有用です。また、このオプションにより複数のオプティマイザのパラメータグループを持つことも可能になります。FSDP でモデルを準備/ラップする前にオプティマイザを作成する場合は、これを True に設定する必要があります。

fsdp_cpu_ram_efficient_loading: Transformers モデルでのみ適用されます。True の場合、最初のプロセスのみが事前学習済みモデルのチェックポイントをロードし、他のすべてのプロセスは重みが空になります。from_pretrained メソッドで事前学習済み Transformers モデルをロードする際にエラーが発生する場合は、これを False に設定する必要があります。この設定が True の場合、fsdp_sync_module_states も True にする必要があります。そうしないと、メインプロセス以外のすべてのプロセスがランダムな重みを持つことになり、トレーニング中に予期しない動作が発生します。これを機能させるには、Transformers の from_pretrained メソッドを呼び出す前に分散プロセスグループが初期化されていることを確認してください。Trainer API を使用する場合、TrainingArguments クラスのインスタンスを作成する際に分散プロセスグループが初期化されます。

fsdp_sync_module_states: True の場合、個別にラップされた各 FSDP ユニットが、ランク 0 からモジュールのパラメータをブロードキャストします。

より詳細かつ繊細な制御を行うには、FullyShardedDataParallelPlugin を介して他の FSDP パラメータを指定できます。 FullyShardedDataParallelPlugin オブジェクトを作成する際は、accelerate の設定に含まれていなかったパラメータ、または上書きしたいパラメータを渡します。 FSDP パラメータは、accelerate の設定ファイルまたは起動コマンドの引数に基づいて選択され、FullyShardedDataParallelPlugin オブジェクトに直接渡す他のパラメータは、それらをセット/上書きします。

以下に例を示します。

from accelerate import FullyShardedDataParallelPlugin
from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig

fsdp_plugin = FullyShardedDataParallelPlugin(
    state_dict_config=FullStateDictConfig(offload_to_cpu=False, rank0_only=False),
    optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=False, rank0_only=False),
)

accelerator = Accelerator(fsdp_plugin=fsdp_plugin)

保存と読み込み

FSDPモデルを使用する際の新しい推奨されるチェックポイント保存方法は、accelerateの設定を行う際にStateDictTypeとしてSHARDED_STATE_DICTを使用することです。 以下は、accelerateのsave_stateユーティリティを使用して保存するためのコードスニペットです。

accelerator.save_state("ckpt")

チェックポイントフォルダを確認して、プロセスごとのモデルとオプティマイザーのシャードを表示します:

ls ckpt
# optimizer_0  pytorch_model_0  random_states_0.pkl  random_states_1.pkl  scheduler.bin

cd ckpt

ls optimizer_0
# __0_0.distcp  __1_0.distcp

ls pytorch_model_0
# __0_0.distcp  __1_0.distcp

学習を再開するためにそれらを再度読み込むには、accelerate の load_state ユーティリティを使用します。

accelerator.load_state("ckpt")

transformersのsave_pretrainedを使用する際は、モデルのstate dictを保存するためにstate_dict=accelerator.get_state_dict(model)を渡します。 以下に例を示します:

  unwrapped_model.save_pretrained(
      args.output_dir,
      is_main_process=accelerator.is_main_process,
      save_function=accelerator.save,
+     state_dict=accelerator.get_state_dict(model),
)

State Dict

accelerator.get_state_dict は、FullStateDictConfig(offload_to_cpu=True, rank0_only=True) コンテキストマネージャーを使用して、model.state_dict の実装を呼び出し、rank 0 の state dict のみを取得します。この state dict は CPU にオフロードされます。

その後、state を save_pretrained メソッドに渡すことができます。StateDictType と FullStateDictConfig には、state_dict の動作を制御するために使用できるいくつかのモードがあります。詳細については、PyTorch のドキュメントを参照してください。

StateDictType.SHARDED_STATE_DICT を選択した場合、Accelerator.save_state 実行時のモデルの重みは、モデルの各サブスプリットごとに n 個のファイルに分割されます。学習後にモデルへ再度読み込むためにこれらを単一の辞書にマージするには、merge_weights ユーティリティを使用できます。

from accelerate.utils import merge_fsdp_weights

# Our weights are saved usually in a `pytorch_model_fsdp_{model_number}` folder
merge_fsdp_weights("pytorch_model_fsdp_0", "output_path", safe_serialization=True)

最終的な出力は、model.safetensors または pytorch_model.bin（safe_serialization=False が渡された場合）に保存されます。

これはCLIからも呼び出すことができます:

accelerate merge-weights pytorch_model_fsdp_0/ output_path

FSDPのシャーディング戦略とDeepSpeed ZeROステージの対応

• FULL_SHARDはDeepSpeedのZeRO Stage-3に対応します。オプティマイザーの状態、勾配、パラメータをシャードします。
• SHARD_GRAD_OP は DeepSpeed の ZeRO Stage-2 に相当します。オプティマイザの状態と勾配をシャードします。
• NO_SHARD は ZeRO Stage-0 にマッピングされます。シャーディングを行わず、各 GPU がモデル、オプティマイザー状態、勾配の完全なコピーを保持します。
• HYBRID_SHARD は ZeRO++ Stage-3 にマッピングされ、zero_hpz_partition_size=<num_gpus_per_node> となります。ここでは、各ノード内でオプティマイザの状態、勾配、パラメータがシャード化され、各ノードが完全なコピーを保持します。

留意すべきいくつかの注意点

• 複数のモデルを使用する場合、対応するモデルと同じ順序でオプティマイザーをprepare呼び出しに渡してください。そうしないと、accelerator.save_state()およびaccelerator.load_state()が誤った、または予期しない動作を引き起こします。
• この機能は、Transformers ライブラリの run_translation.py スクリプトにおける --predict_with_generate と互換性がありません。

より詳細な制御を行うには、FullyShardedDataParallelPluginを利用できます。このクラスのインスタンスを作成した後、Acceleratorクラスのインスタンス化に渡すことができます。 これらのオプションの詳細については、PyTorchのFullyShardedDataParallelコードを参照してください。

FSDPとDeepSpeedの類似点と相違点にご興味のある方は、こちらのコングラントガイドをご覧ください!