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

> Implement distributed training FSDP and DeepSpeed in AI from scratch. Our tutorial reconstructs PyTorch FSDP and DeepSpeed ZeRO-3 mechanics for efficient AI model training.

- Repository: [Rohit Ghumare/ai-engineering-from-scratch](https://github.com/rohitg00/ai-engineering-from-scratch)
- Tags: tutorial
- Published: 2026-06-24

---

**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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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:

```python
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:

```python
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:

```python
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:

```python
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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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:

```bash
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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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.