pytorch-fsdp2

スキル：トレーニングスクリプトで PyTorch FSDP2（fully_shard）を正しく使用する

このスキルでは、初期化、シャーディング、混合精度／オフロードの設定、チェックポイント処理を正しく行いながら、コーディングエージェントがトレーニングループに PyTorch FSDP2 を追加する方法を説明します。

PyTorch の FSDP2 は、主に torch.distributed.fsdp.fully_shard と、それによってモジュールにインプレースで追加される FSDPModule メソッドを通じて提供されます。参照先：references/pytorch_fully_shard_api.md、references/pytorch_fsdp2_tutorial.md。

---

このスキルを使用する場面

次の場合は FSDP2 を使用します。

• モデルが 1 台の GPU に収まらない場合（パラメーター＋勾配＋オプティマイザーの状態）。
• FSDP1 よりも、eager モードでのシャーディング方式として パラメーターごとの DTensor ベースのシャーディング（より調査しやすく、シャーディングされた state dict もシンプル）を使用したい場合。
• 後から DeviceMesh を使って、DP と Tensor Parallel を組み合わせる可能性がある場合。

次の場合は使用を避けるか、慎重に扱ってください。

• PyTorch のバージョンをまたいで厳密な後方互換性のあるチェックポイントが必要な場合（DCP はこれを推奨していません）。
• FSDP2 スタックに対応していない古い PyTorch バージョンを使用せざるを得ない場合。

代替手段（FSDP2 が最適ではない場合）

• DistributedDataParallel（DDP）：標準的なデータ並列ラッパーを使った、従来型の分散データ並列トレーニングが必要な場合に使用します。
• FullyShardedDataParallel（FSDP1）：データ並列ワーカー間でパラメーターをシャーディングする、従来の FSDP ラッパーを使用します。

参照先：references/pytorch_ddp_notes.md、references/pytorch_fsdp1_api.md。

---

エージェントが守るべき契約

1. torchrun で起動し、プロセスごとに CUDA デバイスを設定します（通常は LOCAL_RANK を使用します）。
2. fully_shard() はボトムアップで適用します。つまり、ルートモジュールより先にサブモジュール（Transformer ブロックなど）をシャーディングします。
3. FSDP2 のフックを実行するため、model.forward(input) ではなく model(input) を呼び出します（明示的に unshard() を呼び出すか、forward メソッドを登録する場合を除きます）。
4. シャーディング後にオプティマイザーを作成し、DTensor パラメーター（fully_shard の適用後）を対象として構築されていることを確認します。
5. 分散チェックポイント（DCP）または分散 state dict ヘルパーを使ってチェックポイントを保存します。全テンソルに意図的に集約する場合を除き、単純な torch.save(model.state_dict()) は使用しないでください。

（これらのルールはすべて、公式 API ドキュメント／チュートリアルで直接説明されています。参照先を確認してください。）

---

手順

0）バージョンと環境の確認

• FSDP2 と DCP のドキュメントが最近更新されている、比較的新しい安定版 PyTorch を推奨します。
• torchrun --nproc_per_node <gpus_per_node> ... を使い、RANK、WORLD_SIZE、LOCAL_RANK が参照可能であることを確認します。

参照先：references/pytorch_fsdp2_tutorial.md（起動コマンドとセットアップ）、references/pytorch_fully_shard_api.md（ユーザー契約）。

---

1）分散環境を初期化してデバイスを設定する

最小限で正しいパターン：

• dist.init_process_group(backend="nccl")
• torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
• 必要に応じて、データ並列グループを記述する DeviceMesh を作成します。

参照先：references/pytorch_device_mesh_tutorial.md（DeviceMesh が必要な理由と、プロセスグループの管理方法）。

---

2）meta デバイス上でモデルを構築する（非常に大きなモデルに推奨）

大規模モデルの場合は、meta 上で初期化し、シャーディングを適用してから GPU 上に重みを実体化します。

• with torch.device("meta"): model = ...
• サブモジュールに fully_shard(...) を適用し、その後 fully_shard(model) を適用します。
• model.to_empty(device="cuda")
• model.reset_parameters()（または独自の初期化ルーチン）

参照先：references/pytorch_fsdp2_tutorial.md（移行ガイドにこの手順が明示されています）。

---

3) fully_shard() をボトムアップで適用する（ラップ方針＝「必要な箇所に適用」）

最上位のモジュールだけに fully_shard を呼び出すのは避けてください。

Transformer のようなモデルに推奨されるシャーディングパターン：

• モジュールを反復処理し、if isinstance(m, TransformerBlock): fully_shard(m, ...) を実行する
• その後、fully_shard(model, ...) を実行する

理由：

• fully_shard は集合通信の効率化のために「パラメーターグループ」を作成し、先行する呼び出しですでにグループ化されたパラメーターは除外します。ボトムアップで適用すると、オーバーラップが向上し、ピークメモリ使用量が低下します。

参考：references/pytorch_fully_shard_api.md（ボトムアップで適用する要件とその理由）。

---

4) メモリとパフォーマンスのトレードオフに応じて reshard_after_forward を設定する

デフォルトの動作：

• None は、ルート以外のモジュールでは True、ルートモジュールでは False を意味します（適切なデフォルト値です）。

目安：

• メモリがボトルネックの場合：デフォルトを維持するか、多くのブロックで True を明示的に設定します。
• スループットがボトルネックで、メモリに余裕がある場合：シャーディングされていないパラメーターをより長く保持することを検討します（ルートでは通常 False です）。
• 高度な設定：意味のある除数であれば、int を使って forward 後により小さなメッシュ（例：ノード内）に再シャーディングできます。

参考：references/pytorch_fully_shard_api.md（動作の詳細）。

---

5) 混合精度とオフロード（任意ですが、よく使用されます）

FSDP2 では以下を使用します：

• mp_policy=MixedPrecisionPolicy(param_dtype=..., reduce_dtype=..., output_dtype=..., cast_forward_inputs=...)
• CPU オフロードを行う場合は offload_policy=CPUOffloadPolicy()

経験則：

• H100/A100 クラスの GPU では、モデルの数値安定性を確認したうえで、パラメーターとリダクションに BF16 を使用するところから始めます。
• reduce_dtype は、勾配のリダクションに関する想定に合わせます。
• CPU オフロードを使用する場合は、PCIe/NVLink の通信量と実行時のオーバーヘッドを考慮してください。

参考：references/pytorch_fully_shard_api.md（MixedPrecisionPolicy / OffloadPolicy クラス）。

---

6) オプティマイザー、勾配クリッピング、勾配蓄積

• シャーディング後にオプティマイザーを作成し、DTensor パラメーターを保持するようにします。
• 勾配蓄積や no_sync が必要な場合：
  • FSDP1 の no_sync() ではなく、FSDP2 の仕組み（set_requires_gradient_sync）を使用します。

勾配クリッピング：

• パラメーターと勾配は DTensor であるため、FSDP2 チュートリアルの「Gradient Clipping and Optimizer with DTensor」で示されている方法を使用します。

参考：references/pytorch_fsdp2_tutorial.md。

---

7) チェックポイント：DCP または分散状態辞書ヘルパーを推奨

推奨される方法は次の 2 つです。

A) Distributed Checkpoint（DCP）— 最適なデフォルト

• DCP は複数のランクから並列に保存・読み込みを行い、読み込み時の再シャーディングをサポートします。
• DCP は複数のファイル（多くの場合、少なくともランクごとに 1 つ）を生成し、「インプレース」で処理します。

B) 分散状態辞書ヘルパー

• StateDictOptions(full_state_dict=True, cpu_offload=True, broadcast_from_rank0=True, ...) とともに get_model_state_dict / set_model_state_dict を使用します。
• オプティマイザーには get_optimizer_state_dict / set_optimizer_state_dict を使用します。

避けること：

• DTensor.full_tensor() を使って意図的に変換し、メモリを慎重に管理する場合を除き、DTensor の状態辞書を通常の torch.save で保存しないでください。

参考：

• references/pytorch_dcp_overview.md（DCP の動作と注意事項）
• references/pytorch_dcp_recipe.md および references/pytorch_dcp_async_recipe.md（エンドツーエンドの使用方法）
• references/pytorch_fsdp2_tutorial.md（DTensor と DCP の状態辞書の処理フロー）
• references/pytorch_examples_fsdp2.md（動作するチェックポイントスクリプト）

---

ワークフローのチェックリスト（コピーしてすぐに使えます）

ワークフロー A：既存のトレーニングスクリプトに FSDP2 を導入する

• torchrun で起動し、プロセスグループを初期化します。
• LOCAL_RANK から CUDA デバイスを設定します。多次元並列処理が必要な場合は、DeviceMesh を作成します。
• モデルを構築し（必要に応じて meta を使用）、下位から順に fully_shard を適用してから、fully_shard(model) を呼び出します。
• シャーディング後にオプティマイザーを作成し、DTensor パラメーターを取り込ませます。
• フックが実行されるよう model(inputs) を使用します。勾配累積には set_requires_gradient_sync を使用します。
• torch.distributed.checkpoint ヘルパーを使って DCP の保存／読み込みを追加します。

参考：references/pytorch_fsdp2_tutorial.md、references/pytorch_fully_shard_api.md、references/pytorch_device_mesh_tutorial.md、references/pytorch_dcp_recipe.md。

ワークフロー B：DCP の保存／読み込みを追加する（最小パターン）

• 状態を Stateful でラップするか、get_state_dict を使って状態を組み立てます。
• すべてのランクから、共有パスに対して dcp.save(...) を呼び出します。
• dcp.load(...) を呼び出し、set_state_dict で復元します。
• 別のメッシュに読み込む場合は、再シャーディングに関する前提を検証します。

参考：references/pytorch_dcp_recipe.md。

デバッグチェックリスト（エージェントが最初に確認する項目）

1. すべてのランクが別々の GPU を使用していますか？
   そうでない場合は、torch.cuda.set_device(LOCAL_RANK) と torchrun のフラグを確認します。
2. 誤って forward() を直接呼び出していませんか？
   model(input) を使用するか、明示的に unshard() を呼び出すか、forward を登録します。
3. fully_shard() は下位から順に適用されていますか？
   ルートのみをシャーディングした場合、メモリ使用量やパフォーマンスが悪化し、混乱が生じる可能性があります。
4. オプティマイザーは適切なタイミングで作成されていますか？
   シャーディング後に DTensor パラメーターを使って作成する必要があります。
5. チェックポイントの保存先は一貫していますか？
   • DCP を使用する場合、変換方法を理解していない限り、その場しのぎの torch.save と混在させないでください。
   • DCP に関する PyTorch バージョン互換性の警告に注意してください。

---

よくある問題と解決方法

• forward フックが実行されない → model.forward(...) ではなく、model(inputs) を呼び出します（または unshard() を明示的に呼び出します）。
• オプティマイザーが非 DTensor パラメーターを認識する → すべての fully_shard 呼び出しの後にオプティマイザーを作成します。
• ルートモジュールのみがシャーディングされている → ルートより先に、サブモジュールに下位から順に fully_shard を適用します。
• forward の後にメモリ使用量が急増する → より多くのモジュールで reshard_after_forward=True を設定します。
• 勾配累積の同期がずれる → FSDP1 の no_sync() ではなく、set_requires_gradient_sync を使用します。

参考：references/pytorch_fully_shard_api.md、references/pytorch_fsdp2_tutorial.md。

---

最小限のリファレンス実装の概要（エージェント向け）

コーディングエージェントは、以下のラベル付きブロックを含むスクリプトを実装します。

• init_distributed(): プロセスグループを初期化し、デバイスを設定
• build_model_meta(): meta 上にモデルを構築し、fully_shard を適用して重みを実体化
• build_optimizer(): シャーディング後にオプティマイザーを作成
• train_step(): model(inputs) と DTensor 対応のパターンを使った forward／backward／step
• checkpoint_save/load(): DCP または分散状態辞書ヘルパー

具体例は references/pytorch_examples_fsdp2.md と公式チュートリアルのリファレンスにあります。

---

参考資料

• references/pytorch_fsdp2_tutorial.md
• references/pytorch_fully_shard_api.md
• references/pytorch_ddp_notes.md
• references/pytorch_fsdp1_api.md
• references/pytorch_device_mesh_tutorial.md
• references/pytorch_tp_tutorial.md
• references/pytorch_dcp_overview.md
• references/pytorch_dcp_recipe.md
• references/pytorch_dcp_async_recipe.md
• references/pytorch_examples_fsdp2.md
• references/torchtitan_fsdp_notes.md（任意、本番環境向けの注意事項）
• references/ray_train_fsdp2_example.md（任意、統合例）