Distributed Training with FSDP and DeepSpeed: From Toy DDP to Production-Grade Sharding

The ai-engineering-from-scratch repository teaches distributed training with FSDP and DeepSpeed through a hands-on curriculum that begins with a CPU-based toy DDP implementation and progresses to demonstrating how FSDP achieves the same memory efficiency as DeepSpeed ZeRO-3 via parameter sharding and collective operations.

Distributed training with FSDP and DeepSpeed enables scaling language models that exceed single-GPU memory capacity. The rohitg00/ai-engineering-from-scratch repository provides a ground-up educational framework that implements these concepts from scratch, allowing practitioners to experiment with distributed strategies on CPU before deploying to GPU clusters. This approach bridges the gap between theoretical sharding algorithms and production frameworks like PyTorch Fully Sharded Data Parallel and DeepSpeed.

From DDP to FSDP: The Learning Progression

The curriculum follows a pedagogical sequence that builds complexity incrementally. It starts with a toy Distributed Data Parallel (DDP) implementation using the Gloo backend on CPU, then introduces Fully Sharded Data Parallel (FSDP) concepts that mirror DeepSpeed’s ZeRO-3 stage.

This progression demonstrates four key paradigms:

• Data Parallelism – Full model replication across ranks with gradient synchronization
• Tensor Parallelism – Splitting individual layers across devices
• Pipeline Parallelism – Distributing sequential model stages across workers
• Fully Sharded Data Parallel – Sharding parameters, gradients, and optimizer states across all ranks

Core Implementation in main.py

The hands-on code resides in phases/19-capstone-projects/48-distributed-fsdp-ddp/code/main.py, which implements a minimal DDP wrapper and an FSDP sketch.

Process Group Initialization

The implementation begins with init_process_group at line 55, which configures the Gloo master address and port, then spawns separate processes for each rank. This establishes the communication backbone required for collective operations.

# Conceptual flow from main.py line 55

dist.init_process_group(
    backend="gloo",
    rank=rank,
    world_size=world_size,
    init_method=f"tcp://{master_addr}:{master_port}"
)

Parameter Synchronization

Before training begins, broadcast_module at line 77 ensures rank 0 broadcasts initial weights to all ranks, guaranteeing consistent starting parameters. During the backward pass, all_reduce_grads_ at line 82 sums gradients across ranks and averages them, maintaining mathematical equivalence to single-device training.

The FSDP Round-Trip Mechanism

The fsdp_round_trip_sketch function at line 41 demonstrates the core FSDP pattern. Before each forward pass, parameter shards are all-gathered from across ranks to reconstruct the full tensor. After the forward computation, these extra copies are dropped to free memory. The same pattern applies during backpropagation.

# From main.py line 41 - conceptual FSDP sketch

def fsdp_round_trip_sketch(model, world_size, rank):
    # Shard parameters across ranks

    # All-gather before forward

    # Compute forward/backward

    # Drop gathered copies (reduce-scatter for grads)

    return all_ranks_ok

Verification occurs at line 106 via assert result["fsdp_round_trip_all_ranks_ok"], which confirms that gathered tensors match the original data on every rank.

FSDP vs. DeepSpeed ZeRO-3: Architectural Parallels

According to the curriculum documentation in phases/10-llms-from-scratch/05-scaling-distributed/docs/en.md at line 66, the FSDP sketch directly mirrors DeepSpeed ZeRO-3 (Zero Redundancy Optimizer, stage 3). Both approaches shard parameters, optimizer states, and gradients across all GPUs, then use collective communication to reconstruct parameters only when needed.

The memory reduction follows a clear progression:

• Standard DDP: Per-GPU memory is O(N) – each rank stores the full model, gradients, and optimizer states
• FSDP/ZeRO-3: Per-GPU memory is O(1) – each rank stores only 1/world_size of parameters and optimizer states

Both implementations add communication overhead through all-gather operations before the forward pass and reduce-scatter operations after the backward pass. The curriculum provides memory-budget tables that help calculate exactly how many GPUs are required for models of specific sizes.

Practical Code Examples

Simulating Distributed Training on CPU

Run a complete 2-rank distributed simulation without GPUs to verify the implementation:

# 1️⃣  Simulate a 2‑rank distributed run (CPU‑only)

from phases_19_capstone_projects_48_distributed_fsdp_ddp.code import main as ddp

result = ddp.run_distributed_demo(
    world_size=2,
    in_dim=16,
    hidden=12,
    out_dim=3,
    batch_size=4,
    num_steps=3,
    seed=11,
)
print(result["fsdp_round_trip_all_ranks_ok"])   # → True

print(result["param_sum_spread"])               # ≈ 0 (parameters stay in sync)

Testing the FSDP Round-Trip

Verify sharding correctness with the low-level sketch:

# 2️⃣  FSDP round‑trip sketch – low‑level demonstration

import torch, torch.distributed as dist
from torch import nn

model = nn.Linear(8, 8)

# ... spawn processes, init_process_group …

# Inside each rank:

ok = ddp.fsdp_round_trip_sketch(model, world_size=dist.get_world_size(), rank=dist.get_rank())
assert ok, "sharding round‑trip failed"

Calculating Memory Requirements

Determine GPU requirements before training:

# 3️⃣  Memory calculator (from the lesson) – decide how many GPUs you need

from phases_10_llms_from_scratch_05_scaling_distributed.code import memory_calculator

calc = memory_calculator(params_billions=70, num_gpus=8, sharding="fsdp")
print(f"Per‑GPU total: {calc['per_gpu_total_gb']:.1f} GB")  # → ≈105 GB (needs >8 GPUs)

Key Files and Repository Structure

File	Purpose
phases/19-capstone-projects/48-distributed-fsdp-ddp/code/main.py	Core implementation containing init_process_group, broadcast_module, and fsdp_round_trip_sketch
phases/19-capstone-projects/48-distributed-fsdp-ddp/code/test_main.py	Unit tests verifying sharding integrity, gradient norms, and the FSDP round-trip
phases/19-capstone-projects/48-distributed-fsdp-ddp/docs/en.md	Lesson documentation with build steps and exercises
phases/10-llms-from-scratch/05-scaling-distributed/docs/en.md	High-level comparison of FSDP and DeepSpeed ZeRO-3 with memory tables

Summary

• Distributed training with FSDP and DeepSpeed reduces per-GPU memory from O(N) to O(1) by sharding parameters, gradients, and optimizer states across ranks.
• The ai-engineering-from-scratch repository implements a toy DDP wrapper using the Gloo backend, enabling CPU-based experimentation with distributed concepts.
• fsdp_round_trip_sketch in main.py demonstrates the core mechanism: all-gather parameters before forward, compute, then drop copies to save memory.
• DeepSpeed ZeRO-3 implements the identical sharding strategy as FSDP, using the same collective communication patterns (all-gather and reduce-scatter).
• The curriculum includes memory calculators and verification utilities that let you validate sharding correctness before deploying to expensive GPU clusters.

Frequently Asked Questions

What is the difference between DDP and FSDP?

DDP (Distributed Data Parallel) replicates the entire model on every GPU and only synchronizes gradients during the backward pass, resulting in O(N) memory usage per device. FSDP (Fully Sharded Data Parallel) shards the model parameters across all GPUs, gathering them only when needed for computation, which reduces per-GPU memory to O(1) while adding communication overhead.

How does FSDP relate to DeepSpeed ZeRO-3?

DeepSpeed ZeRO-3 and FSDP implement the same fundamental algorithm: they partition optimizer states, gradients, and parameters across data parallel ranks, then use all-gather operations to reconstruct parameters before the forward pass. According to the source documentation in phases/10-llms-from-scratch/05-scaling-distributed/docs/en.md, both approaches achieve identical memory savings, with DeepSpeed providing additional production features like optimizer offloading.

Can I run these distributed training examples without GPUs?

Yes. The ai-engineering-from-scratch implementation uses the Gloo backend (via init_process_group) which runs on CPU, allowing you to test distributed logic, verify the fsdp_round_trip_sketch, and understand collective operations without specialized hardware. This makes the curriculum accessible for learning before moving to GPU clusters.

What memory savings does FSDP provide compared to standard training?

FSDP reduces per-GPU memory consumption proportionally to the number of ranks. For a model with N parameters, standard training requires O(N) memory per GPU, while FSDP requires only O(N/world_size) for parameters plus O(N) communication buffers temporarily during all-gather operations. The repository's memory calculator demonstrates that a 70-billion parameter model requiring terabytes of memory in standard training can fit on 8-16 GPUs using FSDP sharding.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}