# How Distributed Training is Taught Using FSDP and DeepSpeed: A Complete Guide

> Learn distributed training with FSDP and DeepSpeed. This guide shows how ai-engineering-from-scratch teaches these advanced techniques, from DDP to ZeRO-3 equivalence.

- Repository: [Rohit Ghumare/ai-engineering-from-scratch](https://github.com/rohitg00/ai-engineering-from-scratch)
- Tags: deep-dive
- Published: 2026-07-31

---

**The ai-engineering-from-scratch curriculum teaches distributed training by implementing a toy Distributed Data Parallel (DDP) system on CPU using the Gloo backend, then demonstrating Fully Sharded Data Parallel (FSDP) mechanics and their direct equivalence to DeepSpeed's ZeRO-3 stage.**

The *ai-engineering-from-scratch* repository provides a ground-up approach to understanding distributed training concepts without requiring immediate access to multi-GPU clusters. By starting with CPU-based simulations and progressing to advanced sharding strategies, the curriculum demystifies how **FSDP** and **DeepSpeed** reduce memory consumption from `O(N)` to `O(1)` per GPU through strategic parameter sharding and collective communication operations.

## Building the Foundation with Toy DDP on CPU

The curriculum begins with a minimal yet complete **Distributed Data Parallel (DDP)** implementation that runs entirely on CPU using PyTorch's Gloo backend. This approach allows learners to experiment with distributed mechanics on a single machine before scaling to GPU clusters.

In [`phases/19-capstone-projects/48-distributed-fsdp-ddp/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/19-capstone-projects/48-distributed-fsdp-ddp/code/main.py), the implementation establishes several critical distributed training primitives:

- **Process group creation** – The `init_process_group` function configures the Gloo master address and port, spawning separate processes for each rank to simulate a multi-device environment.

- **Parameter synchronization** – The `broadcast_module` function ensures rank 0 sends initial weights to all other ranks, guaranteeing every process starts from identical model parameters.

- **Gradient aggregation** – After `loss.backward()`, the `all_reduce_grads_` function sums gradients across all ranks and averages them, keeping model weights synchronized during training.

Each rank executes its own forward pass through `MinimalDDP.forward`, processing local data slices while maintaining consistent model states across the distributed environment.

## Understanding FSDP Through the Round-Trip Sketch

The repository introduces **Fully Sharded Data Parallel (FSDP)** through a concrete `fsdp_round_trip_sketch` function that demonstrates the core sharding pattern without requiring production-grade infrastructure.

According to the source code in [`main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/main.py), the FSDP implementation follows this pattern:

1. **Sharding** – Model parameters are divided across ranks, with each GPU storing only a fraction of the total weights.

2. **All-gather** – Before each forward pass, the system performs an *all-gather* operation to reconstruct the full model tensor from sharded slices temporarily.

3. **Computation** – The forward pass executes using the gathered parameters.

4. **Cleanup** – After the forward pass completes, the extra parameter copies are dropped to free memory.

The same pattern applies to the backward pass, with gradients undergoing *reduce-scatter* operations to update only the relevant shards. The `fsdp_round_trip_sketch` includes verification through `assert result["fsdp_round_trip_all_ranks_ok"]` to ensure gathered tensors match the original on every rank.

```python

# FSDP round-trip sketch – low-level demonstration

import torch
import 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"

```

## Mapping FSDP to DeepSpeed ZeRO-3

The curriculum explicitly connects the FSDP sketch to **DeepSpeed ZeRO-3** (Zero Redundancy Optimizer stage 3), demonstrating that both approaches implement identical sharding strategies using different abstractions.

As documented in [`phases/10-llms-from-scratch/05-scaling-distributed/docs/en.md`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/10-llms-from-scratch/05-scaling-distributed/docs/en.md), DeepSpeed ZeRO-3 shards parameters, optimizer states, and gradients across data parallel ranks, then performs all-gather operations before the forward pass and reduce-scatter after the backward pass. This mirrors the FSDP sketch's behavior exactly.

The key insight taught in the repository is that both methods:
- Reduce per-GPU memory from `O(N)` (full model) to `O(1)` (sharded fraction)
- Incur communication overhead (all-gather and reduce-scatter operations)
- Reconstruct parameters only when needed during computation

The lesson page *"Scaling: Distributed Training, FSDP, DeepSpeed"* provides memory-budget tables that help calculate exactly how many GPUs are required for models of specific sizes, making the trade-off between memory and communication bandwidth concrete.

## Practical Implementation Details

The toy implementation in [`main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/main.py) demonstrates five critical distributed training layers:

**Process Group Creation** – Sets up the Gloo backend with `init_process_group` at line 55, establishing the communication fabric between ranks.

**Parameter Broadcast** – The `broadcast_module` function at line 77 ensures rank 0 broadcasts initial weights so every rank starts from the same model state.

**Forward Pass Execution** – Each rank runs its data slice through `MinimalDDP.forward` at line 24 independently.

**Gradient All-Reduce** – The `all_reduce_grads_` function at line 82 sums gradients across ranks after backward propagation.

**Verification** – The main execution block at line 106 asserts `result["fsdp_round_trip_all_ranks_ok"]` to validate that sharding and reconstruction work correctly.

## Code Examples for Distributed Training

The repository provides runnable examples that demonstrate distributed training concepts on CPU hardware:

### Simulating Multi-Rank Training

```python

# 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)

```

### Memory Budget Calculation

```python

# 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)

```

## Summary

- The **ai-engineering-from-scratch** curriculum teaches distributed training by implementing a toy DDP system on CPU using the Gloo backend, making distributed concepts accessible without GPU hardware.

- The **FSDP sketch** in [`main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/main.py) demonstrates parameter sharding through all-gather and reduce-scatter operations, verified by round-trip consistency checks.

- **DeepSpeed ZeRO-3** implements an identical sharding strategy to FSDP, reducing per-GPU memory from `O(N)` to `O(1)` while adding communication overhead.

- Key files include the toy implementation at [`phases/19-capstone-projects/48-distributed-fsdp-ddp/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/19-capstone-projects/48-distributed-fsdp-ddp/code/main.py) and the high-level concept documentation at [`phases/10-llms-from-scratch/05-scaling-distributed/docs/en.md`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/10-llms-from-scratch/05-scaling-distributed/docs/en.md).

- The curriculum provides memory calculators and verification utilities to experiment with the memory-communication trade-off before deploying to production clusters.

## Frequently Asked Questions

### What is the difference between FSDP and DeepSpeed ZeRO-3?

**FSDP (Fully Sharded Data Parallel)** and **DeepSpeed ZeRO-3** implement the same fundamental strategy: they shard model parameters, gradients, and optimizer states across data-parallel ranks, then use all-gather operations to reconstruct layers only during computation. According to the ai-engineering-from-scratch source code, FSDP is PyTorch's native implementation while DeepSpeed ZeRO-3 provides additional production features like optimizer state sharding and offloading, but both reduce per-GPU memory from `O(N)` to `O(1)`.

### Can you learn distributed training without access to multiple GPUs?

Yes. The repository specifically designs its **toy DDP implementation** to run on CPU using the Gloo backend, allowing learners to experiment with process groups, broadcast operations, and gradient all-reduce on a single machine. The `run_distributed_demo` function simulates multi-rank environments without requiring GPU hardware, making distributed training concepts accessible for local development and testing.

### How does the curriculum verify that FSDP sharding works correctly?

The `fsdp_round_trip_sketch` function in [`main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/main.py) implements a verification mechanism that shards a tensor across ranks, all-gathers it back together, and asserts that the reconstructed tensor matches the original on every rank. The main execution block checks `result["fsdp_round_trip_all_ranks_ok"]` to ensure the sharding logic preserves data integrity, providing immediate feedback when the distributed communication patterns function correctly.

### When should you choose FSDP over standard DDP?

Choose **FSDP** when model parameters exceed the memory capacity of individual GPUs. The curriculum explains that standard DDP replicates the full model on every rank (`O(N)` memory), while FSDP shards parameters across ranks (`O(1)` memory). Use FSDP when training large language models where the memory budget requires sharding, and accept the trade-off of increased communication overhead from all-gather and reduce-scatter operations.