FSDP — Fully Sharded Data Parallel

Overview#

FSDP is PyTorch's answer to ZeRO-3. Each worker owns 1/N of the parameters; before each forward layer (the 'wrapping unit'), the worker AllGathers the full layer parameters, runs the layer, then discards the gathered weights. Backward proceeds symmetrically with a ReduceScatter on gradients. The optimiser step is local on each worker's parameter shard. Memory per rank drops from O(P) for vanilla DDP to O(P/N) — the same factor as ZeRO-3 — at the cost of one extra AllGather per wrapping unit per pass.

FSDP first appeared inside Meta's FairScale library in 2021, was rewritten and upstreamed into `torch.distributed.fsdp` for PyTorch 1.11 (March 2022), and was rewritten again as FSDP2 in PyTorch 2.4 (2024). The FSDP2 rewrite replaced the FlatParameter abstraction with per-parameter DTensor sharding, which made FSDP composable with tensor parallelism (HSDP — Hybrid Sharded Data Parallel), pipeline parallelism, and expert parallelism without the special-case glue the FlatParameter design required.

By 2026 FSDP is the dominant sharded-data-parallel strategy in the PyTorch ecosystem. HuggingFace Trainer wires it up via `accelerate launch --use-fsdp`; HuggingFace Accelerate exposes it through `FullyShardedDataParallelPlugin`; torchtune (Meta's lightweight LLM fine-tuning library) is FSDP2-native; Lightning Fabric exposes it as a one-flag strategy; and TRL's SFT / DPO / PPO trainers all support FSDP. Yobitel NeoCloud customers fine-tuning Llama 3.1, Qwen 2.5, Mistral, and DeepSeek-V3 derivatives on 8-128 GPU H100/H200 pods commonly use FSDP2 as the default engine, and the Yobibyte managed fine-tune service wraps FSDP (under the hood) for QLoRA and full-parameter recipes on 7B-70B targets.

This entry documents the production surface: the FSDP2 API and its predecessor FSDP1, the four sharding strategies, mixed-precision and CPU-offload configuration, the auto_wrap_policy for transformer blocks, sizing tables for common Llama-class workloads, observability hooks, and the migration paths to and from DeepSpeed ZeRO. This entry helps you stand up FSDP for sharded training and fine-tuning on Yobitel NeoCloud or your own multi-GPU PyTorch cluster, with the right wrapping, sharding, and precision choices.

Quick start#

The example below fine-tunes Llama 3.1 8B on 8x H100 SXM5 using FSDP2 with FULL_SHARD, BF16 mixed precision, transformer-block-granularity auto-wrap, and `backward_prefetch=BACKWARD_PRE`. The first block installs PyTorch 2.4+ and writes a minimal training script that uses the FSDP2 API. The second block launches with `torchrun`. The third block shows the equivalent HuggingFace Accelerate path that picks up the same shape via YAML.

# 1. Install PyTorch 2.4+ (FSDP2) and the model libraries
pip install "torch>=2.4" transformers accelerate datasets

cat > train_fsdp2.py <<'PY'
import os
import functools
import torch
import torch.distributed as dist
from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy
from torch.distributed.device_mesh import init_device_mesh
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.models.llama.modeling_llama import LlamaDecoderLayer

dist.init_process_group(backend="nccl")
rank = dist.get_rank(); world = dist.get_world_size()
torch.cuda.set_device(rank % torch.cuda.device_count())

mesh = init_device_mesh("cuda", (world,), mesh_dim_names=("dp",))
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3.1-8B",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)

mp_policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32,   # keep gradient reduction in FP32 for stability
)

# Shard every transformer block, then shard the root model
for layer in model.model.layers:
    fully_shard(layer, mesh=mesh, mp_policy=mp_policy)
fully_shard(model, mesh=mesh, mp_policy=mp_policy)

optim = torch.optim.AdamW(model.parameters(), lr=2e-5, weight_decay=0.0,
                          fused=True)
# ... training loop: loss.backward() and optim.step() use the sharded params.
PY

# 2. Launch on 8x H100 SXM5
torchrun --standalone --nproc_per_node=8 train_fsdp2.py

# 3. Equivalent HuggingFace Accelerate config (fsdp.yaml)
cat > fsdp.yaml <<'YAML'
compute_environment: LOCAL_MACHINE
distributed_type: FSDP
mixed_precision: bf16
num_machines: 1
num_processes: 8
fsdp_config:
  fsdp_version: 2
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
  fsdp_backward_prefetch: BACKWARD_PRE
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_use_orig_params: true
  fsdp_offload_params: false
YAML
accelerate launch --config_file fsdp.yaml sft.py

How it works#

FSDP shards three pieces of training state across the data-parallel group: parameters, gradients, and optimiser state. The compute graph is unchanged from DDP; what changes is which rank owns which bytes at which moment and the collective operations that move bytes in and out of GPU memory just in time. FSDP organises parameters into 'wrapping units' (usually transformer blocks) — the AllGather / ReduceScatter granularity is per-unit.

Forward pass: before unit U executes, FSDP issues `all_gather` on U's sharded parameters to reconstruct the full layer weights on every rank. The unit runs the forward, producing activations as usual. Immediately after the forward of U, the gathered copies are freed; only the local shard plus the activations remain in memory. If `backward_prefetch=BACKWARD_PRE` is set (the recommended default), the AllGather for unit U+1 is launched asynchronously while U is still computing, so most of the communication overlaps with compute.

Backward pass: each unit's full parameters are AllGathered again (FSDP2 can cache them when activations are still live, FSDP1 always re-gathers). The backward produces gradients of the full unit; FSDP then ReduceScatters the gradients across the FSDP group so each rank ends up with 1/N of the gradient. Gathered parameters are freed again. Once the entire backward completes, each rank has the gradient slice corresponding to its parameter slice.

Optimiser step: AdamW (or the chosen optimiser) runs locally on each rank's slice. There is no AllReduce; the per-rank update is mathematically equivalent to the global step because the work has been partitioned. After the step, FSDP holds only the updated parameter shards — no AllGather is required until the next forward.

Four sharding strategies are exposed: FULL_SHARD (true ZeRO-3, the default at FSDP2), SHARD_GRAD_OP (shard gradients and optimiser state but replicate parameters — ZeRO-2 equivalent, lower memory savings but no per-forward AllGather), NO_SHARD (degenerates to DDP, useful as a sanity baseline), and HYBRID_SHARD (FULL_SHARD inside a node-shaped sub-group, DDP-style replication across that sub-group — gives intra-node ZeRO-3 plus inter-node DDP, the canonical multi-node-pod choice).

Mixed precision is handled by `MixedPrecisionPolicy` (FSDP2) or `MixedPrecision` (FSDP1). Three dtypes are configurable: `param_dtype` (the working precision of gathered weights during forward/backward, usually BF16), `reduce_dtype` (the precision of the gradient ReduceScatter, often FP32 for numerical stability), and `output_dtype` (the dtype the model returns to user code). The Adam master weights remain in FP32 on the sharded slice regardless.

• FULL_SHARD: params + grads + optimiser state sharded; ~N-fold memory cut; AllGather per forward + per backward + ReduceScatter per backward.
• SHARD_GRAD_OP: grads + optimiser state sharded, params replicated; ~50 percent memory cut vs DDP; no per-forward AllGather.
• HYBRID_SHARD: FULL_SHARD within node-shaped subgroup, replicated across subgroups; the recommended default for multi-node pods.
• NO_SHARD: equivalent to DDP — a sanity-check baseline, not a production choice.
• MixedPrecisionPolicy: BF16 working dtype + FP32 reduce dtype is the safe default; FP16 with loss scaling is only for Volta/Turing hardware.
• Auto-wrap policy: wrap at transformer-block granularity (`LlamaDecoderLayer`, `MistralDecoderLayer`); module-by-module wrapping makes AllGather overhead dominate.
• CPU offload: `CPUOffload(offload_params=True)` moves parameter shards to CPU when not in use — frees GPU memory at the cost of PCIe traffic.

Reference and specifications#

FSDP is configured by the call to `fully_shard()` (FSDP2) or the `FullyShardedDataParallel` constructor (FSDP1) and through the surrounding mixed-precision / wrapping / offload policy objects. The table below documents the surface that matters for production runs as of PyTorch 2.4 (June 2026). Fields marked FSDP2-only do not exist on FSDP1; FSDP1 equivalents (`MixedPrecision`, `sharding_strategy=ShardingStrategy.FULL_SHARD`) are noted where the migration is non-obvious.

Workload patterns#

Three workload shapes cover the bulk of FSDP production usage: full-parameter SFT of 7B-13B models on a single 8-GPU node, full-parameter SFT or continued pretraining of 70B on a 4-node HSDP pod, and QLoRA fine-tuning of a 70B model on a single 8x H100 node with CPU offload. Each maps to a different config emphasis, and each maps to a different Yobitel NeoCloud training-pod shape — Pattern A on a single 8x H100 SXM5 node, Pattern B on a 32-GPU (4-node) NeoCloud HSDP pod, Pattern C on a single 8x H100 node with the Yobibyte managed fine-tune service.

Pattern A — FULL_SHARD SFT of Llama 3.1 8B on a single 8x H100 node, the standard small-model fine-tuning shape. Pattern B — HYBRID_SHARD continued-pretrain of Llama 3.1 70B on a 32-GPU NeoCloud pod, with FULL_SHARD inside each node and replicated AllReduce across nodes. Pattern C — QLoRA on Llama 3.1 70B with CPU offload, the cheapest 70B fine-tune shape on a single 8x H100 box; this is the shape Yobibyte's managed fine-tune service offers as a turnkey recipe for customers who do not want to operate FSDP themselves.

# A — FULL_SHARD SFT of Llama 3.1 8B on a single 8x H100 node
from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy
from torch.distributed.device_mesh import init_device_mesh

mesh = init_device_mesh("cuda", (8,), mesh_dim_names=("dp",))
mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16,
                          reduce_dtype=torch.float32)
for layer in model.model.layers:
    fully_shard(layer, mesh=mesh, mp_policy=mp)
fully_shard(model, mesh=mesh, mp_policy=mp)
# Optimiser sees sharded params; AdamW(fused=True) is the standard choice.


# B — HYBRID_SHARD continued-pretrain of Llama 3.1 70B on 32 GPUs (4 nodes)
#     FULL_SHARD inside each 8-GPU node, replicated across the 4 nodes.
mesh = init_device_mesh("cuda", (4, 8), mesh_dim_names=("replicate", "shard"))
for layer in model.model.layers:
    fully_shard(layer, mesh=mesh["replicate", "shard"], mp_policy=mp,
                reshard_after_forward=True)
fully_shard(model, mesh=mesh["replicate", "shard"], mp_policy=mp)
# Effective layout: each node holds the full model in sharded form;
# inter-node AllReduce only runs on gradients (DDP-style).


# C — QLoRA on Llama 3.1 70B with CPU offload (Yobibyte managed shape)
from torch.distributed.fsdp import CPUOffload, FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from transformers.models.llama.modeling_llama import LlamaDecoderLayer
import functools

# Apply 4-bit quantisation + LoRA adapters first, then wrap with FSDP1.
auto_wrap = functools.partial(transformer_auto_wrap_policy,
                              transformer_layer_cls={LlamaDecoderLayer})
model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    auto_wrap_policy=auto_wrap,
    mixed_precision=MixedPrecision(param_dtype=torch.bfloat16,
                                   reduce_dtype=torch.float32),
    cpu_offload=CPUOffload(offload_params=True),
    use_orig_params=True,           # required for PEFT/LoRA
    limit_all_gathers=True,
)

Sizing and capacity planning#

The questions that drive FSDP sizing are: which sharding strategy fits the model, what per-GPU throughput will I get, and where does CPU offload help versus hurt? The table below gives reference configs and observed throughput for BF16 AdamW-trained Llama-family architectures at 4K sequence length, with `mp_policy=BF16/FP32`, `backward_prefetch=BACKWARD_PRE`, FlashAttention-2 on, and activation checkpointing on. Throughput is per-GPU sustained tokens-per-second from Yobitel-operated NeoCloud fleets and published Meta / HF benchmarks.

• FULL_SHARD is the right default through 70B on 8 GPUs; above 70B, compose with TP via FSDP2 + DTensor.
• HYBRID_SHARD beats FULL_SHARD on multi-node pods because the cross-node AllGather is the main cost — keep AllGather inside the NVLink island.
• CPU offload costs 2-3x throughput; only worth it when no other capacity is available or when running QLoRA at the limit.
• Activation checkpointing roughly halves activation memory at ~30 percent FLOP cost — required for 70B at long context on 80 GB cards.
• FlashAttention-2/3 is mandatory at scale; without it, attention activations dominate before FSDP can help.
• Effective batch = micro_batch x DP x grad_accum; aim for 1-4M tokens per step for stable AdamW.

Limits and quotas#

FSDP itself imposes few hard limits; the constraints come from GPU memory, NCCL configuration, the auto-wrap policy, and the checkpoint pipeline. The table below summarises the operational ceilings worth knowing before designing an FSDP run.

Observability#

FSDP emits diagnostic information through standard PyTorch hooks plus a few FSDP-specific ones: `FSDP.summon_full_params()` for inspection, `set_state_dict_type` for checkpoint shape, and the optional `TORCH_DISTRIBUTED_DEBUG=DETAIL` env var that traces per-collective timing. In production you pair this with TensorBoard / W&B for loss curves, DCGM exporter for GPU-level metrics, and NCCL profiling for collective-level diagnostics. The signals that detect 90 percent of FSDP production issues are throughput drift, AllGather p99 latency, and grad-norm sanity.

• samples_per_second / tokens_per_second: holds within +-5 percent in steady state; drops correlate with AllGather under-overlap.
• AllGather p99 / ReduceScatter p99: visible under TORCH_DISTRIBUTED_DEBUG=DETAIL; >100ms inside an NVLink island means topology regression.
• grad_norm: spikes above 10x baseline indicate divergence — almost always LR or data, not FSDP.
• GPU memory peak per step: log with torch.cuda.max_memory_allocated() to validate wrap policy.
• DCGM_FI_PROF_NVLINK_TX_BYTES vs DCGM_FI_PROF_PCIE_TX_BYTES: separates FSDP collectives (NVLink) from CPU offload traffic (PCIe).
• torch.profiler with NCCL traces: the right tool for diagnosing 'why is AllGather not overlapping'.
• Checkpoint save duration: a sudden jump usually means FULL_STATE_DICT accidentally got re-enabled.

# Prometheus rules for an FSDP training job
groups:
  - name: fsdp-training
    interval: 60s
    rules:
      - alert: FSDPStepTimeRegression
        expr: |
          avg_over_time(fsdp:step_time_seconds[5m]) >
          1.10 * avg_over_time(fsdp:step_time_seconds[1h] offset 30m)
        for: 10m
        labels: { severity: warning, team: training }
        annotations:
          summary: "FSDP step time +10% vs 1h baseline on {{ $labels.job_name }}"

      - alert: FSDPAllGatherP99
        expr: histogram_quantile(0.99,
                rate(fsdp:allgather_seconds_bucket[5m])) > 0.2
        for: 10m
        annotations:
          summary: "FSDP AllGather p99 > 200ms — fabric or wrap-policy regression"

      - alert: FSDPGradNormSpike
        expr: fsdp:grad_norm > 10 * avg_over_time(fsdp:grad_norm[1h] offset 30m)
        for: 5m
        labels: { severity: warning }
        annotations:
          summary: "FSDP grad-norm spike — investigate LR or data corruption"

Cost and FinOps#

FSDP cost is a function of throughput tax for the chosen sharding strategy and the hardware mix. FULL_SHARD on a single 8-GPU node is within ~5 percent of DDP throughput where DDP fits; CPU offload typically costs 2-3x in wall-clock; HSDP across nodes preserves most of FULL_SHARD's per-GPU throughput. The table below uses Yobitel NeoCloud list pricing (June 2026) for the canonical 70B fine-tuning shape (Llama 3.1 70B SFT on 8x H100 SXM5) and the alternative shapes.

• Single-node FSDP with FULL_SHARD beats DeepSpeed ZeRO-3 on the same hardware by ~5-10 percent on routine 7B-13B fine-tunes — same math, tighter PyTorch integration.
• QLoRA is the cheapest 70B fine-tune path on a single 8x H100 node; the Yobibyte managed fine-tune service exposes this as a turnkey recipe at the same NeoCloud node rate.
• HSDP across 4 nodes finishes the same job ~5x faster than single-node FULL_SHARD at the same USD-per-token — buy more nodes for time-to-result, not for cost.
• Activation checkpointing trades ~30 percent FLOPs for ~5x activation memory — usually still cheaper than scaling out.
• Reserved capacity (Yobitel NeoCloud 1-3 year terms) drops the rate 35-40 percent for predictable multi-month runs.

Security and compliance#

FSDP is a training library; security and compliance apply at the cluster boundary the same way they do for DeepSpeed or Megatron. The library itself is BSD-licensed and ships inside PyTorch with no telemetry call-home. For UK and EU sovereign deployments, FSDP runs inside the same NCSC Cloud Security Principles- and G-Cloud 14-aligned Yobitel NeoCloud tenancies as the other training engines — the framework choice does not change the sovereignty posture.

Two operational notes for compliance audits. First, FSDP sharded checkpoints (`__N_N.distcp` files in the torch_dist format) are not portable across world sizes without an explicit re-shard step; use `torch.distributed.checkpoint.load` or `torch.distributed.checkpoint.format_utils.dcp_to_torch_save` to materialise a single state dict for audit retention. Second, CPU-offload spills sensitive weights to host RAM and (when pin_memory is on) keeps them in a pinned region that survives container restart; ensure the host environment is dedicated and that ephemeral storage is encrypted on Yobitel NeoCloud sovereign tenancies.

Migration and alternatives#

FSDP and DeepSpeed ZeRO have converged functionally at the math layer — FSDP FULL_SHARD and ZeRO-3 give the same memory and throughput on the same workload. The meaningful differences are ergonomic and ecosystem. FSDP has tighter PyTorch core integration, native DTensor composition with TP, and a cleaner extension story via FSDP2 (`torch.distributed._composable.fsdp.fully_shard`). DeepSpeed has the richer offload story (CPU and NVMe via ZeRO-Infinity), the Megatron-DeepSpeed hybrid, and an established HuggingFace Trainer config surface. For frontier 100B+ training, Megatron-LM with 3D parallelism is the production default; FSDP composes with TP via DTensor but the recipe ecosystem is thinner.

Troubleshooting#

The error patterns below cover the common FSDP failure modes observed on Yobitel-operated training fleets and the public PyTorch issue tracker. Each row maps an observable symptom to the underlying mechanism and the minimum-viable fix.

Where this fits in the Yobitel stack#

FSDP is one of the supported training engines on Yobitel NeoCloud sovereign tenancies. For PyTorch-native customers running 7B-70B fine-tunes on 8-128 GPU pods, the canonical path is FSDP2 with FULL_SHARD inside an 8-GPU node and HYBRID_SHARD across multi-node pods, BF16 mixed precision, FlashAttention-2/3, and SHARDED_STATE_DICT checkpoints landing on the pod's parallel filesystem. Pre-built NGC-derived containers ship PyTorch 2.4+, the NCCL and Transformer Engine versions tested against the Yobitel NeoCloud H100 / H200 / B200 fleets, and DCGM observability dashboards pre-applied.

For customers who want the QLoRA / SFT recipe without operating FSDP directly, Yobibyte's managed fine-tune service wraps FSDP under the hood and exposes a higher-level Llama / Qwen / Mistral recipe surface backed by the same NeoCloud capacity — same engine, no cluster operation. Trained checkpoints flow into Yobibyte's inference path (industry-standard runtimes) after the standard HuggingFace conversion, and customer-tuned weights stay in the customer's sovereign region throughout the lifecycle.