FSDP2

Fully Sharded Data Parallel (FSDP2) は、モデル、勾配、オプティマイザの状態を GPU 間でシャード（分割）します。計算の前に、各 GPU はすべてのシャードから完全なパラメータセットを集約し、計算後にそれらを解放します。シャード化により、単一の GPU のメモリを超えるモデルの学習が可能になりますが、DDP と比べて通信量が増加します。モデルやオプティマイザの状態が単一の GPU に収まらない場合は、FSDP を使用してください。

                      ┌─────────────────┐
                      │  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 を設定して、メモリとスループットのトレードオフを選択します。

auto_wrap_policy は、モジュールを FSDP ユニットにラップする方法を制御します。デフォルト値は "TRANSFORMER_BASED_WRAP" で、モデルの transformer レイヤーをラップします。ラップしない場合（"NO_WRAP"）、モデル全体が 1 つの FSDP ユニットとなり、シャーディングによるメモリ上の利点を失います。

FSDPの構成

これらのフィールドは、FSDP2がモデルをラップ、シャード、ロードする方法を制御します。reshard_after_forwardとauto_wrap_policyについては、シャーディング戦略で説明しています。

• 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 には gradient checkpointing の代わりにこれを使用してください。両方を設定するとエラーが発生します。

FSDP 学習は、Accelerate 設定ファイルまたは fsdp_config に渡す FSDP 設定ファイルのいずれかで設定します。

accelerate launch train.py

{
  "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
}

from transformers import TrainingArguments

TrainingArguments(
    ...,
    fsdp=True,
    fsdp_config="path/to/fsdp.json",
)

Mixture-of-Expertsモデルでは、fsdp_sizeをtp_sizeおよびep_sizeと組み合わせることで、エキスパートを1つのメッシュ次元に、それ以外の要素をもう1つの次元にシャーディングできます。エキスパート並列化を参照してください。

コンテキスト並列性を用いて評価

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シーケンス並列とは異なります。CPとUlysses SPのいずれか一方を選択してください。cp_sizeとsp_sizeの両方を1より大きく設定しないでください。

evaluate() は prediction_step() を通じて train() と同じ context-parallel コンテキストを使用します。そのコンテキストが有効な間、各ランクが自身のシーケンスシャードを実行します。eval_loss は正しい損失値です（cp_size 倍にはなっていません）。早期停止と load_best_model_at_end はこの指標を使用します。

Logits はシーケンス方向にシャードされたままとなり、完全な長さの予測に再構成されることはありません。CP 下では、Logits に対する予測および compute_metrics はサポートされていません。

eval_loss が必要な場合、既存の FSDP2 + cp_size 起動を再利用できます。

prediction_step() をサブクラス化してトークンを自分でカウントする場合、num_items_in_batch は context-parallel のコンテキストに入る前にカウントしてください。そのコンテキストが有効な間、シーケンスバッファはランク間でインプレースにシャードされます。その後でカウントすると、バッチ数が過小評価されます。デフォルトの Trainer は最初からカウントしているため、prediction_step をオーバーライドする場合にのみこの処理が必要です。

次のステップ

• モデルが1つのGPUに収まる場合のデータ並列学習については、DDPを参照してください。
• ZeRO最適化およびNVMeオフロードについては、DeepSpeedを参照してください。
• PyTorch/XLA による TPU 上の FSDP では、~TrainingArguments.fsdp_config に xla、xla_fsdp_settings、xla_fsdp_grad_ckpt を設定します。
• Read the FSDP chapter from The Ultra-Scale Playbook for more information about how FSDP works.