FSDP (Fully Sharded Data Parallel) is the PyTorch-native strategy for training large models across multiple GPUs by sharding model parameters, gradients, and optimizer states across all workers — reducing per-GPU memory by up to Nx (where N is GPU count) compared to standard data parallelism, enabling training of models that would not fit in a single GPU's memory.

Why Not Standard Data Parallel?

• DDP (DistributedDataParallel): Full model replica on every GPU.
• 7B parameter model in fp32: 28GB parameters + 28GB gradients + 56GB optimizer (Adam) = 112GB per GPU.
• Even 80GB A100 cannot hold this.
• FSDP: Shards all three across GPUs.
• With 8 GPUs: ~14GB per GPU — fits easily.

FSDP Memory Savings

How FSDP Works

1. Initialization: Model parameters are sharded — each GPU holds only 1/N of parameters. 2. Forward Pass: Before computing a layer, FSDP all-gathers that layer's parameters from all GPUs. 3. Compute: Forward computation using full parameters. 4. Free: After forward, full parameters freed — only shard retained. 5. Backward Pass: Same all-gather for each layer, compute gradients, then reduce-scatter gradients. 6. Optimizer Step: Each GPU updates only its shard of parameters.

PyTorch FSDP API

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    mixed_precision=MixedPrecision(param_dtype=torch.bfloat16),
    auto_wrap_policy=size_based_auto_wrap_policy,
)

Key Configuration

• Sharding Strategy: FULL_SHARD (ZeRO-3), SHARD_GRAD_OP (ZeRO-2), NO_SHARD (DDP).
• Auto Wrap Policy: Controls which modules are FSDP-wrapped — affects communication granularity.
• Mixed Precision: bfloat16 params + float32 reduce → further memory savings.
• Activation Checkpointing: Combined with FSDP for maximum memory efficiency.

FSDP vs. DeepSpeed ZeRO

• PyTorch FSDP is the native implementation inspired by DeepSpeed ZeRO.
• DeepSpeed: Third-party library with ZeRO-1/2/3, offloading to CPU/NVMe.
• FSDP: First-class PyTorch citizen — tighter integration with PyTorch ecosystem.
• Both achieve similar memory savings; choice depends on ecosystem preference.

FSDP is the standard approach for training large language models on GPU clusters — it democratizes large model training by making billion-parameter models trainable on commodity multi-GPU setups that would otherwise require expensive model parallelism engineering.