# FSDP2 [Fully Sharded Data Parallel (FSDP2)](https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html) は、モデル、勾配、オプティマイザの状態を GPU 間でシャード（分割）します。計算の前に、各 GPU はすべてのシャードから完全なパラメータセットを集約し、計算後にそれらを解放します。シャード化により、単一 GPU のメモリを超えるモデルの学習が可能になりますが、[DDP](./ddp) と比べて通信量が増加します。モデルやオプティマイザの状態が単一 GPU に収まらない場合は、FSDP を使用してください。 ```text ┌─────────────────┐ │ training data │ └────────┬────────┘ ┌──────────────────┼──────────────────┐ │ shard 0 │ shard 1 │ shard 2 ▼ ▼ ▼ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │ param │ │ param │ │ param │ │ shard 0 │ │ shard 1 │ │ shard 2 │ │ GPU 0 │ │ GPU 1 │ │ GPU 2 │ └──────┬──────┘ └──────┬──────┘ └──────┬──────┘ │ │ │ └──────── all-gather (params) ────────┘ │ full params on each GPU │ ┌──────────────────┼──────────────────┐ ▼ ▼ ▼ forward forward forward │ │ │ └───── reduce-scatter (grads) ────────┘ │ ┌──────────────────┼──────────────────┐ ▼ ▼ ▼ grad shard 0 grad shard 1 grad shard 2 optim shard 0 optim shard 1 optim shard 2 step step step ``` ## シャーディング戦略 FSDP2 は `~TrainingArguments.fsdp_config` でシャーディングを制御します。FSDP を有効にするには `fsdp=True` を設定し、FSDP 設定内の `reshard_after_forward` を設定して、メモリとスループットのトレードオフを選択します。 | `reshard_after_forward` | 動作 | |---|---| | `true` | フォワードパス後にパラメータを再シャーディングして、より多くのメモリを節約します | | `false` | forwardとbackwardの間でパラメータを保持したままにすることで、再all-gatherを回避できます。ただし、ピークメモリ使用量が増加するというトレードオフがあります。 | `auto_wrap_policy` は、モジュールを FSDP ユニットにラップする方法を制御します。デフォルト値は `"TRANSFORMER_BASED_WRAP"` で、モデルの transformer レイヤーをラップします。ラップを行わない場合（`"NO_WRAP"`）、モデル全体が 1 つの FSDP ユニットとなり、シャーディングによるメモリ上の利点を失います。 ## FSDPの構成 これらのフィールドは、FSDP2がモデルをラップ、シャード、ロードする方法を制御します。`reshard_after_forward`と`auto_wrap_policy`については、[シャーディング戦略](#sharding-strategies)で説明しています。 - `cpu_offload` は、パラメータと勾配が使用されていないときに CPU へオフロードして GPU メモリを節約します。 - `transformer_layer_cls_to_wrap` は、`auto_wrap_policy` が `"TRANSFORMER_BASED_WRAP"` の場合に、FSDP ユニットとしてラップする Transformer レイヤーを定義します。各ユニットは、自身の gather 操作と scatter 操作を管理します。フォワードパスでは、現在のユニットのパラメータのみが gather されます。メモリを節約するために、以前のユニットのパラメータは解放されます。 トップレベルのモデルのみをラップしても、GPU メモリの削減にはなりません。個々の `Linear` レイヤーをすべてラップすると、ユニット間の通信コストが非常に高くなります。このフィールドを空欄にすると、FSDP はモデル定義から値を読み取ります。 - `min_num_params` は、サイズベースのラッピングにおけるモジュールごとの最小パラメータ数を設定します。`auto_wrap_policy` が `"SIZE_BASED_WRAP"` の場合にのみ使用されます。 - `state_dict_type` はチェックポイントの形式を制御します。デフォルトは `"FULL_STATE_DICT"` で、単一の Transformers 互換チェックポイントになります。大規模モデルでは高速なため、ランクごとに 1 つのチェックポイントファイルを作成する `"SHARDED_STATE_DICT"` を使用できます。シャード化されたチェックポイントは FSDP にのみ読み戻せるため、共有または FSDP 外部で読み込む最終チェックポイントには `"FULL_STATE_DICT"` を保存してください。 - `cpu_ram_efficient_loading` は、チェックポイントをディスクから rank 0 のみで読み込みます。その他の GPU は空のモデルを初期化し、ブロードキャストによって重みを受け取ります。これにより、複数のプロセスが大きなモデルを CPU RAM に読み込むことを回避します。 - `activation_checkpointing` は、アクティベーションを保存する代わりに、バックワードパスの間にアクティベーションを再計算します。[TrainingArguments](/docs/transformers/main/ja/main_classes/trainer#transformers.TrainingArguments) には [gradient checkpointing](./grad_checkpointing) の代わりにこれを使用してください。両方を設定するとエラーが発生します。 FSDP 学習は、[Accelerate 設定ファイル](./accelerate#accelerate-config-file)または `fsdp_config` に渡す FSDP 設定ファイルのいずれかで設定します。 [accelerate config](https://huggingface.co/docs/accelerate/en/package_reference/cli#accelerate-config) コマンドを実行し、ハードウェアや学習環境に関する質問に答えてください。これにより、キャッシュに `default_config.yaml` ファイルが作成されます。 [accelerate launch](https://huggingface.co/docs/accelerate/en/package_reference/cli#accelerate-launch) を使用して、[Trainer](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer) ベースのスクリプトを実行します。`fsdp_config` は不要です。Accelerate の設定ファイルが同じ設定をカバーしているためです。 ```cli accelerate launch train.py ``` ```json { "version": 2, "reshard_after_forward": true, "cpu_offload": false, "auto_wrap_policy": "TRANSFORMER_BASED_WRAP", "transformer_layer_cls_to_wrap": ["LlamaDecoderLayer"], "state_dict_type": "FULL_STATE_DICT", "cpu_ram_efficient_loading": true, "activation_checkpointing": true } ``` `fsdp=True`を設定し、FSDP設定ファイルを`fsdp_config`に渡します。 ```py from transformers import TrainingArguments TrainingArguments( ..., fsdp=True, fsdp_config="path/to/fsdp.json", ) ``` > [!TIP] > Mixture-of-Expertsモデルでは、`fsdp_size`を`tp_size`および`ep_size`と組み合わせることで、エキスパートを1つのメッシュ次元に、それ以外の要素をもう1つの次元にシャーディングできます。[エキスパート並列化](./expert_parallelism#combining-with-fsdp2)を参照してください。 ## コンテキスト並列性を用いて評価 FSDP2 によるコンテキスト並列 (CP) は、Accelerate の `ParallelismConfig` の `cp_size` で設定します。起動時には、`--use_parallelism_config` と `--parallelism_config_cp_size` の両方を渡してください。Accelerate の YAML 内で `parallelism_config` が空でない場合、`use_parallelism_config` も有効になります。`--parallelism_config_cp_size` だけを指定しても無視され、`cp_size=1` のままになります。 FSDP2では、CPはシーケンスをランク間でシャードし、各GPUがコンテキストのスライスを保持します（`allgather`または`alltoall`）。これは、DeepSpeed ALSTを介して実行され、同じ設定で`sp_size`を使用する[Ulyssesシーケンス並列](./deepspeed_alst)とは異なります。CPとUlysses SPのいずれか一方を選択してください。`cp_size`と`sp_size`の両方を1より大きく設定しないでください。 [evaluate()](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer.evaluate) は [prediction_step()](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer.prediction_step) を通じて [train()](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer.train) と同じ context-parallel コンテキストを使用します。そのコンテキストが有効な間、各ランクが自身のシーケンスシャードを実行します。`eval_loss` は正しい損失値です（`cp_size` 倍にはなっていません）。早期停止と `load_best_model_at_end` はこの指標を使用します。 Logits はシーケンス方向にシャードされたままとなり、完全な長さの予測に再構成されることはありません。CP 下では、Logits に対する予測および `compute_metrics` はサポートされていません。 `eval_loss` が必要な場合、既存の FSDP2 + `cp_size` 起動を再利用できます。 [prediction_step()](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer.prediction_step) をサブクラス化してトークンを自分でカウントする場合、`num_items_in_batch` は context-parallel のコンテキストに入る前にカウントしてください。そのコンテキストが有効な間、シーケンスバッファはランク間でインプレースにシャードされます。その後でカウントすると、バッチ数が過小評価されます。デフォルトの [Trainer](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer) は最初からカウントしているため、`prediction_step` をオーバーライドする場合にのみこの処理が必要です。 ## 次のステップ - モデルが1つのGPUに収まる場合のデータ並列学習については、[DDP](./ddp)を参照してください。 - ZeRO最適化およびNVMeオフロードについては、[DeepSpeed](./deepspeed)を参照してください。 - PyTorch/XLA による TPU 上の FSDP では、`~TrainingArguments.fsdp_config` に `xla`、`xla_fsdp_settings`、`xla_fsdp_grad_ckpt` を設定します。 - Read the [FSDP chapter](https://nanotron-ultrascale-playbook.static.hf.space/index.html#zero-3:_adding_parameter_partitioning_(fsdp)) from The Ultra-Scale Playbook for more information about how FSDP works.