Implementation of Distributed Training with FSDP and DeepSpeed in AI: A From-Scratch Tutorial

The AI Engineering From Scratch repository provides a self-contained Python implementation that reconstructs the core mechanics of PyTorch FSDP and DeepSpeed ZeRO-3, including process-group initialization, gradient all-reduction, and parameter sharding with all-gather reconstruction.

The rohitg00/ai-engineering-from-scratch repository contains a pedagogical implementation of distributed training that mirrors production-grade libraries without hiding complexity behind abstractions. This lesson, "Distributed Data Parallel and FSDP from Scratch," demonstrates how to implement FSDP (Fully Sharded Data Parallel) and DeepSpeed ZeRO concepts using raw collective communication primitives. By working through the source code in phases/19-capstone-projects/48-distributed-fsdp-ddp/code/main.py, you gain visibility into the exact mechanisms that enable large-scale AI training across multiple GPUs.

Core Building Blocks of Distributed Training

The implementation breaks distributed training into three reproducible stages: process initialization, gradient synchronization, and parameter sharding.

Process Group Initialization

Distributed training begins with establishing communication between ranks. In phases/19-capstone-projects/48-distributed-fsdp-ddp/code/main.py lines 55-62, the code configures a gloo backend for CPU-based testing (though nccl works identically on GPUs) and sets the required environment variables MASTER_ADDR and MASTER_PORT. The function torch.distributed.init_process_group establishes the distributed context, allowing subsequent collective operations to synchronize across the world size.

The DDP Contract

The Distributed Data Parallel pattern follows a strict lifecycle: broadcast, forward, all-reduce, optimizer step. At construction, broadcast_module (lines 77-80) synchronizes model parameters from rank 0 to all other ranks, ensuring every process starts from identical weights. After the backward pass, all_reduce_grads_ (lines 82-90) averages gradients across the cluster using dist.all_reduce with ReduceOp.SUM, followed by division by world_size. This manual gradient synchronization reproduces exactly what PyTorch’s DistributedDataParallel wrapper automates.

FSDP Parameter Sharding and Reconstruction

Fully Sharded Data Parallel reduces per-rank memory to 1/N of the model size by splitting parameters into equal shards. Before each layer’s forward pass, ranks reconstruct the full tensor via all_gather operations, then discard non-local shards immediately after use. The fsdp_round_trip_sketch function (lines 157-175) validates this cycle by verifying that the gathered tensor is bit-identical to the original on every rank, confirming the correctness of the sharding logic implemented in lines 94-101.

Hands-On Implementation from the Source Code

The following functions from phases/19-capstone-projects/48-distributed-fsdp-ddp/code/main.py demonstrate the exact implementation of these distributed training primitives.

Initialize the process group with CPU fallback support:

import os, sys, torch.distributed as dist
def init_process_group(rank, world_size, backend="gloo", port=29500):
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = str(port)
    # macOS uses lo0, Linux uses lo

    os.environ.setdefault("GLOO_SOCKET_IFNAME", "lo")
    dist.init_process_group(backend=backend, rank=rank, world_size=world_size)

Broadcast model parameters from the primary rank:

def broadcast_module(module, src=0):
    for t in list(module.parameters()) + list(module.buffers()):
        dist.broadcast(t.data, src=src)

All-reduce gradients and compute the global L2 norm:

def all_reduce_grads_(module, world_size):
    total_sq = 0.0
    for p in module.parameters():
        if p.grad is None:
            p.grad = torch.zeros_like(p.data)
        dist.all_reduce(p.grad.data, op=dist.ReduceOp.SUM)
        p.grad.data.div_(world_size)
        total_sq += float(p.grad.data.pow(2).sum())
    return total_sq ** 0.5

Validate FSDP sharding with a round-trip reconstruction test:

def fsdp_round_trip_sketch(module, world_size, rank):
    ok = True
    for p in module.parameters():
        full = p.data.clone()
        flat = full.flatten()
        per = (flat.numel() + world_size - 1) // world_size
        padded = torch.cat([flat, torch.zeros(per * world_size - flat.numel())])
        my_slice = padded[rank * per:(rank + 1) * per]
        gathered = [torch.empty(per, dtype=flat.dtype) for _ in range(world_size)]
        dist.all_gather(gathered, my_slice)
        rebuilt = torch.cat(gathered)[:flat.numel()].view_as(full)
        if not torch.allclose(rebuilt, full):
            ok = False
            break
    return ok

Memory Optimization and Production Comparisons

The educational implementation achieves the same memory scaling as production frameworks. According to the documentation in phases/19-capstone-projects/48-distributed-fsdp-ddp/docs/en.md lines 70-71, sharding parameters reduces per-rank memory to 1/N of the total model size, where N is the world size. The trade-off is communication overhead: each layer requires an all_gather operation before computation.

Production systems add optimizations atop this foundation. As noted in phases/19-capstone-projects/48-distributed-fsdp-ddp/docs/en.md lines 26-30, PyTorch’s native FSDP introduces a flat parameter view, overlaps communication with computation, and offers optional CPU offloading. Similarly, DeepSpeed’s ZeRO-3 is described in phases/19-capstone-projects/48-distributed-fsdp-ddp/docs/en.md lines 167-176 as "conceptually identical to FSDP," implementing the same sharding strategy with additional memory optimizations for optimizer states and gradients.

Running the Distributed Training Example

Execute the complete implementation by running the demo script:

python3 phases/19-capstone-projects/48-distributed-fsdp-ddp/code/main.py

The script spawns multiple processes, executes the DDP gradient averaging, runs the FSDP round-trip validation, and writes a JSON summary to outputs/ddp-demo.json. Successful execution confirms that gradients are correctly averaged across ranks and that parameter sharding reconstructs bit-identical tensors, validating the correctness of your distributed training implementation.

Summary

• Process initialization requires setting MASTER_ADDR and MASTER_PORT before calling init_process_group, typically with gloo for CPU testing or nccl for GPU training.
• DDP synchronization follows a strict contract: broadcast parameters from rank 0, compute locally, then all-reduce gradients before the optimizer step.
• FSDP sharding splits parameters into world_size shards, reconstructs full tensors via all_gather during forward passes, and reduces per-rank memory consumption to 1/N.
• Production parity exists in the core logic—PyTorch FSDP and DeepSpeed ZeRO-3 implement the same collective patterns with additional performance optimizations like communication overlap.

Frequently Asked Questions

What is the difference between DDP and FSDP in this implementation?

DDP keeps a full copy of the model on every rank and only synchronizes gradients via all_reduce, while FSDP shards parameters across ranks, using all_gather to reconstruct layers on demand. The DDP approach in phases/19-capstone-projects/48-distributed-fsdp-ddp/code/main.py lines 82-90 requires more memory (full model per GPU) but less communication, whereas the FSDP sketch in lines 141-157 reduces memory to 1/N but adds per-layer gather overhead.

Why does the example use the gloo backend instead of nccl?

The implementation defaults to gloo in init_process_group to enable testing on CPU-only machines and macOS systems where nccl is unavailable. As documented in phases/19-capstone-projects/48-distributed-fsdp-ddp/docs/en.md, switching to nccl requires only changing the backend parameter and ensuring CUDA-visible devices, as the collective communication logic remains identical.

How does this from-scratch implementation compare to PyTorch's native FSDP module?

This educational code implements the same mathematical operations as torch.distributed.fsdp.FullyShardedDataParallel—specifically parameter sharding and all-gather reconstruction—but lacks production optimizations. According to phases/19-capstone-projects/48-distributed-fsdp-ddp/docs/en.md lines 26-30, PyTorch’s native module adds a flat parameter view, overlaps communication with computation, and supports CPU offloading, whereas this implementation focuses on clarity and pedagogical correctness.

What is the "FSDP round-trip" and why does it matter?

The FSDP round-trip refers to the cycle of sharding a parameter, scattering it across ranks, then gathering the shards to reconstruct the original tensor. The fsdp_round_trip_sketch function in phases/19-capstone-projects/48-distributed-fsdp-ddp/code/main.py lines 157-175 validates that rebuilt tensors are numerically identical to the original full tensors, ensuring that sharding logic does not introduce precision errors or synchronization bugs that would corrupt model training.

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