pytorch-fsdp2

• Prefer a recent stable PyTorch where the docs show FSDP2 and DCP updated recently.
• Use torchrun --nproc_per_node <gpus_per_node> ... and ensure RANK, WORLD_SIZE, LOCAL_RANK are visible.

Reference: references/pytorch_fsdp2_tutorial.md (launch commands and setup), references/pytorch_fully_shard_api.md (user contract).

---

Minimal, correct pattern:

• dist.init_process_group(backend="nccl")
• torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
• Optionally create a DeviceMesh to describe the data-parallel group(s)

Reference: references/pytorch_device_mesh_tutorial.md (why DeviceMesh exists & how it manages process groups).

---

For big models, initialize on meta, apply sharding, then materialize weights on GPU:

• with torch.device("meta"): model = ...
• apply fully_shard(...) on submodules, then fully_shard(model)
• model.to_empty(device="cuda")
• model.reset_parameters() (or your init routine)

Reference: references/pytorch_fsdp2_tutorial.md (migration guide shows this flow explicitly).

---

Do not only call fully_shard on the topmost module.

Recommended sharding pattern for transformer-like models:

• iterate modules, if isinstance(m, TransformerBlock): fully_shard(m, ...)
• then fully_shard(model, ...)

Why:

• fully_shard forms “parameter groups” for collective efficiency and excludes params already grouped by earlier calls. Bottom-up gives better overlap and lower peak memory.

Reference: references/pytorch_fully_shard_api.md (bottom-up requirement and why).

---

Default behavior:

• None means True for non-root modules and False for root modules (good default).

Heuristics:

• If you’re memory-bound: keep defaults or force True on many blocks.
• If you’re throughput-bound and can afford memory: consider keeping unsharded params longer (root often False).
• Advanced: use an int to reshard to a smaller mesh after forward (e.g., intra-node) if it’s a meaningful divisor.

Reference: references/pytorch_fully_shard_api.md (full semantics).

---

FSDP2 uses:

• mp_policy=MixedPrecisionPolicy(param_dtype=..., reduce_dtype=..., output_dtype=..., cast_forward_inputs=...)
• offload_policy=CPUOffloadPolicy() if you want CPU offload

Rules of thumb:

• Start with BF16 parameters/reductions on H100/A100-class GPUs (if numerically stable for your model).
• Keep reduce_dtype aligned with your gradient reduction expectations.
• If you use CPU offload, budget for PCIe/NVLink traffic and runtime overhead.

Reference: references/pytorch_fully_shard_api.md (MixedPrecisionPolicy / OffloadPolicy classes).

---

• Create the optimizer after sharding so it holds DTensor params.
• If you need gradient accumulation / no_sync:
  • use the FSDP2 mechanism (set_requires_gradient_sync) instead of FSDP1’s no_sync().

Gradient clipping:

• Use the approach shown in the FSDP2 tutorial (“Gradient Clipping and Optimizer with DTensor”), because parameters/gradients are DTensors.

Reference: references/pytorch_fsdp2_tutorial.md.

---

Two recommended approaches:

A) Distributed Checkpoint (DCP) — best default

• DCP saves/loads from multiple ranks in parallel and supports load-time resharding.
• DCP produces multiple files (often at least one per rank) and operates “in place”.

B) Distributed state dict helpers

• get_model_state_dict / set_model_state_dict with StateDictOptions(full_state_dict=True, cpu_offload=True, broadcast_from_rank0=True, ...)
• For optimizer: get_optimizer_state_dict / set_optimizer_state_dict

Avoid:

• Saving DTensor state dicts with plain torch.save unless you intentionally convert with DTensor.full_tensor() and manage memory carefully.

References:

• references/pytorch_dcp_overview.md (DCP behavior and caveats)
• references/pytorch_dcp_recipe.md and references/pytorch_dcp_async_recipe.md (end-to-end usage)
• references/pytorch_fsdp2_tutorial.md (DTensor vs DCP state-dict flows)
• references/pytorch_examples_fsdp2.md (working checkpoint scripts)

---

• Launch with torchrun and initialize the process group.
• Set the CUDA device from LOCAL_RANK; create a DeviceMesh if you need multi-dim parallelism.
• Build the model (use meta if needed), apply fully_shard bottom-up, then fully_shard(model).
• Create the optimizer after sharding so it captures DTensor parameters.
• Use model(inputs) so hooks run; use set_requires_gradient_sync for accumulation.
• Add DCP save/load via torch.distributed.checkpoint helpers.

Reference: references/pytorch_fsdp2_tutorial.md, references/pytorch_fully_shard_api.md, references/pytorch_device_mesh_tutorial.md, references/pytorch_dcp_recipe.md.

• Wrap state in Stateful or assemble state via get_state_dict.
• Call dcp.save(...) from all ranks to a shared path.
• Call dcp.load(...) and restore with set_state_dict.
• Validate any resharding assumptions when loading into a different mesh.

Reference: references/pytorch_dcp_recipe.md.