← Back
AI Mechanics

Building GPT from Scratch — Part 8: FSDP — When the Model Won't Fit

Drew Zhu·1mo ago·11 min read··👀 813

TL;DR

DDP replicates the full model on every GPU — which fails when the model exceeds one GPU's memory. FSDP (Fully Sharded Data Parallel) shards parameters, gradients, and optimizer states across GPUs. Each GPU holds only 1/N of the model's memory footprint. During forward/backward, FSDP gathers parameters just-in-time and discards them after use. The tradeoff: more communication, but you can train models N× larger. Same convergence as DDP — the math is identical, only the memory layout changes.

#buildgptfromscratch#fsdp

## The Problem

Our 3.3M-parameter GPT fits trivially on any GPU. But as models scale:

A single A100 has 80 GB. LLaMA 7B barely fits for inference — and cannot be trained on one GPU because optimizer states (AdamW's momentum + variance) triple the memory.

DDP doesn't help here — it replicates the model, so every GPU needs the full footprint. With 4 GPUs and DDP, you still need 112 GB per GPU for LLaMA 7B training.

FSDP solves this: shard everything across GPUs so each one holds only 1/N.

The ZeRO Insight

FSDP implements Microsoft's ZeRO (Zero Redundancy Optimizer) idea. The observation:

In standard DDP, every GPU stores:

  • Full model parameters (redundant — all copies are identical)

  • Full gradients (redundant — all are identical after AllReduce)

  • Full optimizer states (redundant — all are identical)

This is 3× redundant across N GPUs. ZeRO eliminates the redundancy in stages:

FSDP is ZeRO Stage 3 — it shards everything.

How FSDP Works

The Lifecycle of a Parameter

In DDP, every parameter lives on every GPU permanently. In FSDP, parameters are sharded at rest and gathered temporarily during compute:

The sequence for each layer:

  1. AllGather — reconstruct full parameters from shards across all GPUs

  2. Forward — run computation using full parameters

  3. Discard — free the gathered parameters (keep only local shard)

  4. AllGather (again during backward) — need full parameters for gradient computation

  5. Backward — compute gradients

  6. ReduceScatter — each GPU gets the gradient shard for its parameter shard

  7. Optimizer step — update only the local shard using the local gradient shard

AllGather vs AllReduce

DDP uses AllReduce: every GPU ends up with the complete averaged gradient.

FSDP uses:

  • AllGather: assemble the full tensor from shards (before compute)

  • ReduceScatter: reduce gradients AND scatter — each GPU gets only its shard of the result

AllGather (before forward):
  GPU 0 has: [A, -, -, -]
  GPU 1 has: [-, B, -, -]
  GPU 2 has: [-, -, C, -]
  GPU 3 has: [-, -, -, D]
  
  After: ALL GPUs have [A, B, C, D]

ReduceScatter (after backward):
  GPU 0 has grad: [g0a, g0b, g0c, g0d]
  GPU 1 has grad: [g1a, g1b, g1c, g1d]
  GPU 2 has grad: [g2a, g2b, g2c, g2d]
  GPU 3 has grad: [g3a, g3b, g3c, g3d]
  
  After:
  GPU 0 gets: sum(g0a, g1a, g2a, g3a)  <- only its shard
  GPU 1 gets: sum(g0b, g1b, g2b, g3b)
  GPU 2 gets: sum(g0c, g1c, g2c, g3c)
  GPU 3 gets: sum(g0d, g1d, g2d, g3d)

ReduceScatter is more memory-efficient than AllReduce — each GPU only receives 1/N of the gradient, which is all it needs to update its 1/N parameter shard.

PyTorch FSDP Implementation

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy, MixedPrecision

# Define which layers to shard (typically per TransformerBlock)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from functools import partial

wrap_policy = partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={TransformerBlock},
)

# Mixed precision: compute in bf16, communicate in fp32
mixed_precision = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32,
    buffer_dtype=torch.float32,
)

# Wrap the model
model = GPT(config).to(device)
model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,  # ZeRO Stage 3
    mixed_precision=mixed_precision,
    auto_wrap_policy=wrap_policy,
    device_id=rank,
)

Sharding Strategies

PyTorch FSDP supports partial sharding for the throughput/memory tradeoff:

ShardingStrategy.FULL_SHARD     # ZeRO-3: shard params + grads + optimizer
ShardingStrategy.SHARD_GRAD_OP  # ZeRO-2: shard grads + optimizer only
ShardingStrategy.NO_SHARD       # = DDP (no sharding)
ShardingStrategy.HYBRID_SHARD   # Full shard within node, replicate across nodes

HYBRID_SHARD is particularly useful for multi-node training: full sharding within a node (where communication is fast over NVLink) and replication across nodes (where network is slower). This minimizes cross-node traffic.

Unified Training Script

Here's how the training script supports single-device, DDP, and FSDP:

def setup_training(config, strategy="auto"):
    """
    strategy: "auto" | "single" | "ddp" | "fsdp"
    auto: single if 1 GPU, DDP if model fits, FSDP if it doesn't
    """
    device, distributed, rank, world_size = setup_distributed()

    model = GPT(config)
    param_memory = sum(p.numel() * 4 for p in model.parameters())  # bytes in fp32
    training_memory = param_memory * 4  # params + grads + optimizer (2x)

    if not distributed:
        # MacBook or single GPU
        model = model.to(device)
    elif strategy == "fsdp" or (strategy == "auto" and training_memory > 70e9):
        # Model too big for one GPU — use FSDP
        model = FSDP(model.to(device), ...)
    else:
        # Model fits — DDP is simpler and slightly faster
        model = DDP(model.to(device), device_ids=[rank])

    return model, device, rank, world_size

The Communication Cost

FSDP communicates more than DDP because it gathers parameters before every forward/backward:

DDP per step:
  - 1 AllReduce of gradients = 2M bytes (M = model size)

FSDP per step:
  - N AllGathers of parameters (one per FSDP unit) = M bytes total
  - N ReduceScatters of gradients = M bytes total
  - 1 AllGather during backward (for gradient computation) = M bytes
  Total: ~3M bytes

FSDP moves ~50% more data than DDP. But it uses 1/N the memory — enabling models that literally cannot run with DDP.

Prefetching: Hiding the Cost

FSDP doesn't wait until a layer needs its parameters to start gathering. It prefetches — issuing the AllGather for layer K+1 while layer K is computing:

Time  →
GPU compute: [  Layer 0 forward  ][  Layer 1 forward  ][  Layer 2  ...
Communication: [AllGather layer 1][AllGather layer 2  ][AllGather ...
                 ↑ overlapped

With good prefetching, the communication is almost entirely hidden behind compute — especially for large models where each layer's compute time exceeds the AllGather time.

Checkpointing with FSDP

Saving checkpoints is trickier with sharded models — each GPU only has 1/N of the parameters:

from torch.distributed.fsdp import StateDictType, FullStateDictConfig

# Option 1: Gather full state dict on rank 0 (simple, needs memory)
with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT):
    if rank == 0:
        state_dict = model.state_dict()
        torch.save(state_dict, "checkpoint.pt")

# Option 2: Each GPU saves its shard (fast, resumable on same GPU count)
with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
    state_dict = model.state_dict()
    torch.save(state_dict, f"checkpoint_rank{rank}.pt")

Option 1 is portable (load on any GPU count). Option 2 is fast (no gathering) but requires the same number of GPUs to resume.

When to Use DDP vs FSDP

The rule of thumb:

  • Model < 1B params + A100 80GB → DDP

  • Model 1B–10B + A100 80GB → FSDP

  • Model > 10B → FSDP + tensor parallelism (next post)

Mixed Precision: Cutting Memory in Half

FSDP pairs naturally with mixed precision (bf16/fp16):

mixed_precision = MixedPrecision(
    param_dtype=torch.bfloat16,    # store params in bf16 (2 bytes vs 4)
    reduce_dtype=torch.float32,     # communicate in fp32 (numerical stability)
    buffer_dtype=torch.float32,
)

This halves the parameter memory — a 7B model goes from 28 GB to 14 GB per full copy, and sharded across 4 GPUs: 3.5 GB per GPU. Suddenly a 7B model trains on 4× consumer GPUs.

The optimizer master weights stay in fp32 for numerical stability, but those are sharded too.

Activation Memory: The Other Bottleneck

Parameters and optimizer states aren't the only memory consumers. During training, the forward pass stores activations (intermediate results needed for backward):

Layer 0 forward → store activations_0
Layer 1 forward → store activations_1
...
Layer N forward → store activations_N
  ← now backward starts, uses activations in reverse order

For a transformer with 32 layers, context_length=2048, embed_dim=4096, batch_size=4:

  • Activation memory ≈ 32 × 4 × 2048 × 4096 × 2 bytes ≈ 4 GB

FSDP doesn't help with activation memory (it's unique per GPU — different data = different activations). The solution: activation checkpointing (also called gradient checkpointing):

from torch.utils.checkpoint import checkpoint

class TransformerBlock(nn.Module):
    def forward(self, x):
        # Don't store activations — recompute during backward
        return checkpoint(self._forward, x, use_reentrant=False)

    def _forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ff(self.ln2(x))
        return x

This trades compute for memory: forward runs twice (once normally, once during backward to recompute activations), but memory drops from O(N_layers) to O(1). Combined with FSDP, it enables training models much larger than GPU memory would normally allow.

Going Multi-Node: FSDP Across Machines

FSDP across multiple machines works the same as single-node — but the network becomes critical. Here's why: DDP crosses the network once per training step (to average gradients). FSDP crosses the network every time a layer needs its parameters — which is every layer, twice (forward and backward).

Think of it this way: DDP ships a summary at the end of the day (gradients). FSDP ships the actual materials before every task and returns results after each one (parameters gathered, gradients scattered). Same final result, but far more trips across the wire.

The Communication Challenge

DDP across nodes: one AllReduce per step (gradient size = model size). Easily overlapped.

FSDP across nodes: AllGather + ReduceScatter per FSDP unit (per transformer block). For a 70B model with 80 layers:

  • 80 AllGathers during forward (each gathering one layer's parameters)

  • 80 AllGathers during backward (same layers again)

  • 80 ReduceScatters during backward (scattering gradients)

  • Total: 240 collective operations per training step

Each operation moves one layer's parameters across the network. For LLaMA 70B:

  • One layer's parameters: ~875M params × 2 bytes (bf16) = 1.75 GB

  • 240 operations × 1.75 GB = 420 GB per training step

Over InfiniBand (400 GB/s): ~1 second of communication. Over ethernet (25 GB/s): ~17 seconds. This is why InfiniBand is non-negotiable for multi-node FSDP.

HYBRID_SHARD: The Multi-Node Strategy

Full sharding across nodes means every AllGather/ReduceScatter crosses the slow inter-node network. That's 240 round trips over InfiniBand per step. The solution: only shard within a node, and replicate across nodes.

Each node shards the model across its 8 local GPUs (fast NVLink). Across nodes, each node holds a full copy of the model — just like DDP. This combines FSDP's memory savings (1/8 per GPU) with DDP's minimal cross-node communication (one AllReduce per step).

```python from torch.distributed.fsdp import ShardingStrategy

model = FSDP( model, sharding_strategy=ShardingStrategy.HYBRID_SHARD, ... )


Memory per GPU: model sharded 8 ways (within node) → 1/8 memory footprint. Cross-node traffic: one AllReduce per step (gradient size) — same as DDP. Best of both worlds: FSDP's memory savings + DDP's minimal cross-node communication.

### How to Launch Multi-Node FSDP

```bash
# Slurm (most HPC clusters)
srun --nodes=4 --ntasks-per-node=8 --gpus-per-node=8 \
  torchrun \
    --nnodes=4 \
    --nproc_per_node=8 \
    --rdzv_backend=c10d \
    --rdzv_endpoint=$MASTER_ADDR:29500 \
    train.py --strategy fsdp

# The code auto-detects multi-node via WORLD_SIZE > nproc_per_node

Shared Storage for Checkpoints

Multi-node FSDP needs all nodes to access checkpoints:

# Option 1: Rank 0 saves full state dict to shared filesystem
if rank == 0:
    with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT):
        torch.save(model.state_dict(), "/shared_nfs/checkpoint.pt")

# Option 2: Each rank saves its shard locally (faster, but needs all nodes to resume)
with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
    torch.save(model.state_dict(), f"/local/checkpoint_rank{rank}.pt")

# Option 3: Distributed checkpoint (PyTorch 2.0+) — best for large models
from torch.distributed.checkpoint import save, load
save(model.state_dict(), checkpoint_id="/shared_nfs/step_1000/")

Distributed checkpointing (Option 3) is the production choice — each GPU saves its shard in parallel to shared storage, avoiding the memory spike of gathering the full model on rank 0.

Multi-Node Performance Profile

4-node × 8-GPU (32 GPUs total), LLaMA 7B, HYBRID_SHARD:

  Intra-node AllGather (NVLink):   0.8ms per layer
  Intra-node ReduceScatter:        0.8ms per layer
  Cross-node AllReduce (IB):       12ms per step (once, overlapped)
  Compute per step:                180ms

  Efficiency vs single-node FSDP:  ~92%
  Efficiency vs single GPU:        ~88% (32 GPUs → ~28× speedup)

The 8% efficiency loss comes from the cross-node AllReduce and occasional straggler delays. At this scale, training that would take 30 days on one GPU finishes in ~26 hours.

Putting It All Together

A production training script handles the full spectrum:

# Development on MacBook — test logic, iterate fast
python train.py --batch_size 8

# Single GPU training — larger batch, real data
python train.py --batch_size 32

# Multi-GPU DDP — 4× throughput, model fits on each GPU
torchrun --nproc_per_node=4 train.py --batch_size 32

# FSDP — model too large for one GPU (single node)
torchrun --nproc_per_node=8 train.py --batch_size 32 --strategy fsdp

# Multi-node FSDP — 4 machines × 8 GPUs = 32 GPUs
torchrun --nnodes=4 --nproc_per_node=8 \
  --rdzv_endpoint=$MASTER:29500 \
  train.py --strategy fsdp

Same code path, different scale. The model architecture never changes — only the wrapper and launch command.

Key Takeaways

  1. FSDP shards everything. Parameters, gradients, optimizer states — each GPU holds only 1/N. This is how you train models larger than GPU memory.

  2. The tradeoff is communication vs. memory. FSDP communicates more than DDP (AllGather + ReduceScatter per layer vs. one AllReduce per step), but uses 1/N the memory. Prefetching hides most of the cost within a node.

  3. HYBRID_SHARD is the multi-node answer. Shard within a node (fast NVLink), replicate across nodes (infrequent InfiniBand). You get FSDP's memory savings without drowning the cross-node network in per-layer traffic.

  4. Same convergence as single-GPU. The math is identical — you're computing the exact same gradients, just assembling them differently. No approximation, no accuracy loss.

  5. Activation checkpointing complements FSDP. FSDP handles parameter/optimizer memory. Checkpointing handles activation memory. Together they push the boundary of what fits on a given GPU budget.

Next: tensor parallelism — splitting individual layers across GPUs for the largest models.

Full code: github.com/the-last-programmers/gpt-from-scratch

Comments

Loading comments...

Building GPT from Scratch — Part 8: FSDP — When the Model Won't Fit | The Last Programmers | The Last Programmers