Building DDP, FSDP, TP, and PP using only Python multiprocessing and shared memory

WIP - Final cleanups and figures are missing.

Training large neural networks requires distributing computation across multiple devices. But understanding how this distribution works is often obscured by complex frameworks and GPU-specific APIs. In this post, we strip away the abstractions and implement the four major parallelism paradigms from scratch using only Python’s multiprocessing module and shared memory-no GPUs, no torch.distributed, no NCCL. The goal isn’t production code.

It’s to see the mechanics clearly: what data moves where, when synchronization happens, and why each approach exists. What we’ll build:

  • Data Parallel (DDP): Replicate the model, shard the data, synchronize gradients
  • Fully Sharded Data Parallel (FSDP/ZeRO-3): Shard everything-parameters, gradients, optimizer states
  • Tensor Parallel (TP): Split individual layers across devices
  • Pipeline Parallel (PP): Split the model into sequential stages

All implementations train a simple MLP on MNIST, running as parallel CPU processes communicating through shared memory.

Setup: The Base Model

Out test model is a simple MLP for MNIST.

class MNISTMLP(nn.Module):

    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(28 * 28, 256),
            nn.LayerNorm(256),
            nn.ReLU(),
            nn.Linear(256, 64),
            nn.LayerNorm(64),
            nn.ReLU(),
            nn.Linear(64, 10),
        )

    def forward(self, x):
        bs = x.shape[0]
        x = x.reshape(bs, -1)
        x = self.model(x)
        return x

This has just enough structure (multiple linear layers, normalization) to make the parallelism patterns meaningful without drowning in model complexity.

Data Distributed Parallel

Data Distributed Parallel (DDP) is the standard method for scaling up the effective batch size of a training run when you can’t scale the batch size on a single device. Each process (or “rank”) holds a full replica of the model but receives a unique shard of the training data. After the forward and backward passes, the gradients are All-Reduced (averaged) and broadcast back to all nodes ensuring all model instances apply an identical optimizer update.

As a consequence of the increased effective batch size, the total wall-time of training a model decreases roughly linear with number of GPUs (with some communication cost).

In the following sub-sections, we implement the basic building blocks needed for DP before training for different world size.

Data Sharding

This ensure each process trains on a unique subset of the data for an epoch using a custom sampler. This is a fairly standard implementation for such data sampler.

@dataclass()
class ShardSampler(torch.utils.data.Sampler):
    total_size: int
    rank: int
    world: int
    shuffle: bool = True
    seed: int = 0
    epoch: int = field(default=0, init=False, repr=False)

    def set_epoch(self, epoch: int) -> None:
        self.epoch = epoch

    def __iter__(self):
        g = torch.Generator()
        g.manual_seed(self.seed + self.epoch)
        idxs = (
            torch.randperm(self.total_size, generator=g)
            if self.shuffle
            else torch.arange(self.total_size)
        )
        idxs = idxs[self.rank::self.world]
        return iter(idxs.tolist())

    def __len__(self) -> int:
        return (self.total_size + self.world - 1) // self.world

The key detail: every rank uses the same seed + epoch to generate the same permutation, then slices it with a different offset (self.rank). This guarantees non-overlapping data shards without any inter-process communication.

Manual Gradient All-Reduce with Shared Memory

The core of DDP lies in the communication of gradients. Our custom DDPWrapper wrapper handles the gradient synchronization manually using a shared memory buffer. A few things to notice:

  1. SHMReducer is essentially an interface for reducing the gradients across different ranks. We implement this in the next step.
  2. flat_grad keeps the flatten gradient of all parameters. While _views maps each parameter to the chunk in that flattened gradient to identify where gradient will be reduced for each parameter across ranks.
  3. We register a hook on each parameter to write the local gradient into the corresponding view after the backward pass. This hook is called by PyTorch’s autograd engine automatically.
  4. In allreduce_grads we first reduce the gradient, and then copy back reduced gradient (which is there in view) back to param.grad of each process for optimizer step.
class DDPWrapper(nn.Module):

    def __init__(self, module, rank, world_size, shm_name, barrier):
        super().__init__()
        self.module = module
        self.rank = rank
        self.world_size = world_size

        # flatten param list and build a flat grad view
        self.params = [p for p in self.module.parameters() if p.requires_grad]
        self.total = flat_numel(self.params)

        # open SHM buffer created by parent: shape [world, total]
        self.reducer = SHMReducer(shm_name, world_size, self.total, barrier)

        # create a single flat grad tensor view over params
        # (no extra alloc on step)
        self._views = []
        offset = 0
        self.flat_grad = torch.zeros(
            self.total,
            device=self.params[0].device,
        )

        for p in self.params:
            n = p.numel()
            view = self.flat_grad[offset : offset + n].view_as(p)
            self._views.append((p, view))
            offset += n

        # register per-param hooks to write into the flat view
        for p, view in self._views:

            def _make_hook(v):
                def hook(grad):
                    v.copy_(grad)
                    return grad

                return hook

            p.register_hook(_make_hook(view))

    def forward(self, *args, **kwargs):
        return self.module(*args, **kwargs)

    def allreduce_grads(self) -> None:
        self.reducer.reduce(self.rank, self.flat_grad)

        for p, view in self._views:
            if p.grad is None:
                p.grad = torch.empty_like(p)
            p.grad.copy_(view)

The SHMReducer facilitates a centralized reduction: each process writes its flattened gradient to its allocated row in the SHM buffer, waits for all others, and then calculates the average across the rows before updating its local flat gradient view.

class SHMReducer:
    def __init__(self, name, world_size, total_numel, barrier):
        self.world_size = world_size
        self.total_numel = total_numel
        self.shm = shared_memory.SharedMemory(create=False, name=name)
        self.barrier = barrier

        count = world_size * total_numel
        storage_1d = torch.frombuffer(
            self.shm.buf,
            dtype=torch.float32,
            count=count,
        )
        self.buf = storage_1d.view(world_size, total_numel)

    def reduce(self, rank, flat_grad_view):
        g = flat_grad_view.view(-1)
        n = g.numel()

        # Gather: write the gradients to shared memory
        self.buf[rank, :n].copy_(g)

        # Sync: wait for all processes to write their gradient
        self.barrier.wait()

        # Reduce: mean of gradients
        avg = self.buf[:self.world_size, :n].mean(dim=0)
        self.barrier.wait()

        # Broadcast: back the reduced gradient to each process
        g.copy_(avg)

Note on initialization: To ensure all model replicas are identical at the start, an initial broadcast of the parameters must occur. In our shared-memory setup, this is implicit-the parent process initializes the shared memory and ensures all child processes receive the same initial model state by copying rank 0’s state to all others before the training loop starts. In a real torch.distributed DDP environment, an explicit broadcast operation handles this.

Train Loop

The training loop explicitly calls model.allreduce_grads() to perform the synchronization after the loss.backward() call.

for x, y in train_loader:
    pred = model(x)
    y = torch.nn.functional.one_hot(y, num_classes=10).float()
    loss = loss_fn(pred, y)

    opt.zero_grad(set_to_none=True)
    loss.backward()

    # manual sync
    model.allreduce_grads()
    opt.step()

We validate our implementation by training with different world sizes and observing how total wall time decreases while accuracy remains consistent:

World Size Total Wall Time (sec) Final Acc
2 41.62 0.89
4 26.55 0.89
6 24.81 0.89
8 22.08 0.89

The speedup is sub-linear (as expected-our barrier-based all-reduce has non-trivial synchronization cost on CPU), but the trend is clear: more workers → less wall time, identical convergence.

What’s Missing from Real DDP

There are two important concepts that we did not cover which make DDP work in real life.

  1. Ring-All Reduction: Real DDP uses a Ring All-Reduce algorithm. Instead of centralizing communication through a single shared buffer (which creates a bottleneck), Ring All-Reduce chains the ranks into a ring structure, allowing efficient, simultaneous communication between neighbors. Each rank sends and receives data only to/from its two neighbors, and after 2 × (world_size - 1) steps, every rank has the fully reduced result. This distributes bandwidth load evenly and scales well.
  2. Bucketing + Overlap: To hide communication latency, real DDP employs gradient bucketing. Gradients are grouped into fixed-size buckets, and the All-Reduce for a bucket is kicked off asynchronously as soon as that bucket’s gradients are ready-while the backward pass is still computing gradients for earlier layers. This overlaps communication with computation, significantly improving efficiency.

Fully Sharded Data Parallel (ZeRO-3)

FSDP (or ZeRO-3) is an extension of DDP that addresses a fundamental limitation: in DDP, every rank holds a full copy of the model parameters, gradients, and optimizer states. For large models, this redundancy becomes the memory bottleneck-not the data. FSDP eliminates this redundancy by sharding all three across ranks. Each rank owns only a slice (1/Nth) of the flat parameter vector. Before each forward pass, the full parameters are reconstructed via an all-gather.

After the backward pass, gradients are reduced and scattered back so each rank only stores and updates its own shard. The wrapper gets updated accordingly:

  1. Each rank maintains a local parameter shard and gathers the full parameters only when needed for compute.
  2. After backward, gradients are reduce-scattered: averaged across ranks, then each rank retains only the gradient slice corresponding to its parameter shard for the optimizer step.
class MyFSDPNoDist(nn.Module):
    """
    FSDP-ish:
      - Shard the flat parameter vector across ranks.
      - Per step: all-gather full params into the module for forward;
        reduce-scatter averaged grads back to the local shard;
        step only local shard.
    """

    def __init__(
        self,
        module,
        rank,
        world,
        shm_params_full,
        shm_grads_all,
        barrier,
    ):
        super().__init__()
        self.module = module
        self.rank = rank
        self.world = world
        self.barrier = barrier

        # flat param views (same order every time)
        self.params = [p for p in self.module.parameters() if p.requires_grad]
        self.total = flat_numel(self.params)

        # offsets for each param in the flat vector
        self._views = []
        off = 0
        for p in self.params:
            n = p.numel()
            self._views.append((p, off, n))
            off += n
        assert off == self.total

        # my shard (contiguous slice of the flat vector)
        self.my_off, self.my_len, self.all_offs, self.all_sizes = split_even(
            self.total,
            world,
            rank,
        )

        # local shard parameter tensor + optimizer will own only this
        with torch.no_grad():
            flat_init = self._flatten_params()
            local_init = flat_init[
                self.my_off:self.my_off + self.my_len
            ].clone()

        self.local_shard = nn.Parameter(local_init)

        # map SHM buffers
        self._shm_params = shared_memory.SharedMemory(
            create=False,
            name=shm_params_full,
        )
        self.params_full_buf = torch.frombuffer(
            self._shm_params.buf,
            dtype=torch.float32,
            count=self.total,
        )

        self._shm_grads = shared_memory.SharedMemory(
            create=False,
            name=shm_grads_all,
        )
        self.grads_all_buf = torch.frombuffer(
            self._shm_grads.buf,
            dtype=torch.float32,
            count=self.world * self.total,
        ).view(self.world, self.total)

    # --- helpers ---

    def _flatten_params(self):
        flat = torch.empty(self.total, dtype=torch.float32)

        for p, off, n in self._views:
            flat[off:off+n].copy_(p.data.view(-1))

        return flat

    def _unflatten_to_module(self, flat_vec):
        # write flat vector into module param .data
        i = 0

        for p, off, n in self._views:
            p.data.copy_(flat_vec[off:off+n].view_as(p))
            i += 1

    def _flatten_grads(self):
        gflat = torch.zeros(self.total, dtype=torch.float32)

        for p, off, n in self._views:
            if p.grad is not None:
                gflat[off:off+n].copy_(p.grad.view(-1))

        return gflat

    # --- FSDP-ish steps ---

    def all_gather_params_into_module(self):
        # each rank writes its local shard into the
        # shared "full params" vector
        self.params_full_buf[
            self.my_off:self.my_off+self.my_len
        ].copy_(self.local_shard.data)

        self.barrier.wait()

        # everyone reads full vector and loads into module for forward
        self._unflatten_to_module(self.params_full_buf)
        self.barrier.wait()

    def reduce_scatter_grads_to_local(self):
        # each rank writes its full grad vector into grads_all[rank]
        gflat = self._flatten_grads()
        self.grads_all_buf[self.rank, :].copy_(gflat)
        self.barrier.wait()

        # average over world -> only take my shard slice
        shard_slice = slice(self.my_off, self.my_off + self.my_len)
        avg_shard = self.grads_all_buf[:, shard_slice].mean(dim=0)

        # set grad on local shard param for optimizer step
        if self.local_shard.grad is None:
            self.local_shard.grad = torch.empty_like(
                self.local_shard.data
            )

        self.local_shard.grad.copy_(avg_shard)
        self.barrier.wait()

    def forward(self, *a, **kw):
        return self.module(*a, **kw)

Notice that the optimizer only receives the local shard-not the full model.

opt = torch.optim.SGD([model.local_shard], lr=1e-2)

And the training loop reflects the gather-compute-scatter rhythm:

for x, y in tqdm.tqdm(
    train_loader,
    desc=f"[R{rank}] epoch {epoch}",
    leave=False,
):
    # 1) all-gather: build full params into the module
    model.all_gather_params_into_module()

    # 2) forward/backward on full params
    pred = model(x)
    yoh = torch.nn.functional.one_hot(
        y,
        num_classes=10,
    ).float()
    loss = loss_fn(pred, yoh)

    opt.zero_grad(set_to_none=True)

    for p in model.module.parameters():
        if p.grad is not None:
            p.grad = None

    loss.backward()

    # 3) reduce-scatter grads to my shard & step
    model.reduce_scatter_grads_to_local()
    opt.step()

What’s Missing from Real FSDP

  1. Activation Sharding: Our implementation shards parameters, gradients, and optimizer states-but the activations produced during the forward pass are still fully replicated on every rank. In a true ZeRO-3 / FSDP setup, activations are also partitioned across ranks: each rank stores only its slice of the activation tensor, and an all-gather reconstructs the full activation when needed during the backward pass.

    This is significant because for large models with long sequences, activation memory often dominates over parameter memory. Without activation sharding, FSDP only solves half the memory problem.

  2. Per Layer Sharding: Our implementation flattens all parameters into a single vector and shards it contiguously. Real FSDP operates per-layer (or per “FSDP unit”), wrapping individual layers or groups of layers independently. This enables more fine-grained control: parameters are gathered just before a layer’s forward pass and discarded immediately after, so the full model never resides in memory simultaneously.

Tensor Parallel

Once the activation memory, gradients and optimizer parameters have been distributed - the next bottleneck in terms of memory can be the model itself. When a single layer is too large to fit in memory-common with the enormous embedding tables and attention projections in modern LLMs-we need to split the layer itself.

Tensor Parallelism (TP) does exactly this: it partitions individual weight matrices across devices so that each rank computes only a slice of the layer’s output. The partial results are then combined (via all-reduce or concatenation) to produce the correct full output.

We demonstrate this with two complementary primitives: ColumnParallelLinear and RowParallelLinear, which split a linear layer along the output feature and input feature dimensions respectively.

ColumnParallelLinear

This layer splits the output dimension across TP ranks. Each rank holds a slice of the weight matrix with shape [in_features, out_features // tp_size] and produces a local output [B, out_local]. No communication is needed, the outputs are simply kept local.

class ColumnParallelLinear(nn.Module):
    """
    OUT features are split across TP ranks.
    No gather; returns local y [B, out_local].
    """

    def __init__(
        self,
        in_features,
        out_features,
        tp_rank,
        tp_size,
        bias=True,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.tp_rank = tp_rank
        self.tp_size = tp_size

        self.out_local = split_out_features(
            out_features,
            tp_size,
            tp_rank,
        )

        self.weight = nn.Parameter(
            torch.empty(in_features, self.out_local)
        )
        self.bias = (
            nn.Parameter(torch.empty(self.out_local))
            if bias
            else None
        )

        self.reset_parameters()

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=5 ** 0.5)

        if self.bias is not None:
            bound = 1 / self.weight.size(0) ** 0.5
            nn.init.uniform_(self.bias, -bound, bound)

    def forward(self, x):
        y = x.matmul(self.weight)  # [B, out_local]

        if self.bias is not None:
            y = y + self.bias

        return y

RowParallelLinear

This layer splits the input dimension. Each rank holds a weight slice [in_features // tp_size, out_features] and computes a partial output. The partial outputs are then summed across ranks via a shared-memory reduce to produce the correct full result.

class RowParallelLinear(nn.Module):
    """
    IN features are split across TP ranks;
    computes local partial = x_local @ W_local (bias=None),
    then reduce-sum across TP ranks via a shared-memory
    scratchpad + Barrier.

    Bias is replicated and added after the sum.
    """

    def __init__(
        self,
        in_features,
        out_features,
        tp_rank,
        tp_size,
        shm_name,
        max_bs,
        tp_barrier,
        bias=True,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.tp_rank = tp_rank
        self.tp_size = tp_size
        self.tp_barrier = tp_barrier

        self.in_local = split_out_features(
            in_features,
            tp_size,
            tp_rank,
        )

        self.weight = nn.Parameter(
            torch.empty(self.in_local, out_features)
        )
        self.bias = (
            nn.Parameter(torch.empty(out_features))
            if bias
            else None
        )

        self.reset_parameters()

        # shared scratchpad: [tp, max_bs, out_features]
        self._shm = shared_memory.SharedMemory(
            create=False,
            name=shm_name,
        )

        self.max_bs = max_bs
        total = tp_size * max_bs * out_features

        self.buf = torch.frombuffer(
            self._shm.buf,
            dtype=torch.float32,
            count=total,
        ).view(tp_size, max_bs, out_features)

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=5 ** 0.5)

        if self.bias is not None:
            bound = 1 / self.weight.size(0) ** 0.5
            nn.init.uniform_(self.bias, -bound, bound)

    def forward(self, x_local):
        B = x_local.size(0)

        partial = x_local.matmul(
            self.weight
        )  # [B, out_features]

        # write my partial for peers (detached in SHM)
        self.buf[self.tp_rank, :B, :].copy_(
            partial.detach()
        )

        # sync: ensure all partials written
        self.tp_barrier.wait()

        # sum across TP ranks:
        # keep autograd path only for my partial
        parts = []

        for r in range(self.tp_size):
            if r == self.tp_rank:
                parts.append(partial)
            else:
                parts.append(
                    self.buf[r, :B, :].clone().detach()
                )

        y_full = torch.stack(
            parts,
            dim=0,
        ).sum(dim=0)  # [B, out_features]

        # sync: ensure everyone finished reading
        # before next write overwrites buffer
        self.tp_barrier.wait()

        if self.bias is not None:
            y_full = y_full + self.bias

        return y_full

The TP Model

With these two primitives, our model becomes:

class MNIST_TP(nn.Module):
    def __init__(
        self,
        tp_rank,
        tp_size,
        shm_rowbuf_name,
        tp_barrier,
        max_bs,
    ):
        super().__init__()

        self.fc1_cp = ColumnParallelLinear(
            28 * 28,
            512,
            tp_rank,
            tp_size,
            bias=True,
        )

        self.act1 = nn.ReLU()

        self.fc2_rp = RowParallelLinear(
            512,
            128,
            tp_rank,
            tp_size,
            shm_name=shm_rowbuf_name,
            max_bs=max_bs,
            tp_barrier=tp_barrier,
            bias=True,
        )

        self.act2 = nn.ReLU()
        self.head = nn.Linear(128, 10)  # replicated

    def forward(self, x):
        B = x.size(0)
        x = x.view(B, -1)

        x_local = self.fc1_cp(x)       # [B, 512_local]
        x_local = self.act1(x_local)

        x_full = self.fc2_rp(x_local)  # [B, 128]
        x_full = self.act2(x_full)

        logits = self.head(x_full)
        return logits

Why Column-then-Row and Not the Reverse?

The Column → Row ordering is standard in TP because it minimizes communication:

ColumnParallelLinear (FC1) takes the full replicated input and splits the output feature dimension. The result is a local activation [B, out_local] on each rank-no communication needed. RowParallelLinear (FC2) takes those local activations as input (each rank already has the right slice), computes a partial result, and then all-reduces the partials to get the full output.

Communication happens only once: at the RowParallel layer. If you reversed the order-Row first, then Column-you’d need an all-gather after the Row layer to reconstruct the full activation, and then the Column layer’s output would still be local, requiring yet another communication step downstream. The Column → Row pattern keeps the intermediate activation local and defers the single required communication to the natural aggregation point.

What’s Missing from Real TP

  1. Backward-Pass Gradient Synchronization: Our implementation handles forward communication (summing partials in RowParallelLinear), but the backward pass introduces hidden complexity. The gradients for ColumnParallelLinear’s weights must be all-reduced across the TP group-because the same input was used by all ranks, the gradient for each rank’s weight shard is only a partial gradient.

Real TP libraries (like Megatron-LM) handle this transparently through custom autograd functions that insert the necessary communication in the backward graph.

Pipeline Parallel

Another method to fit big models into memory is by splitting the model into chunks or stages. Each process holds a chunk of the model and hence, both the forward pass activations and backward pass gradients have to be shared from one GPU to another and vice-versa.

To illustrate this, we split our basic MLP model into two stages namely Stage0 and Stage1 in that order as shown below. As you will notice, combined Stage0 and Stage1 in order will give you our base MNISTMLP model.

class Stage0(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(28 * 28, 512),
            nn.LayerNorm(512),
            nn.ReLU(),
        )

    def forward(self, x):
        B = x.size(0)
        return self.net(x.view(B, -1))  # [B, 512]


class Stage1(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(512, 128),
            nn.LayerNorm(128),
            nn.ReLU(),
            nn.Linear(128, 10),
        )

    def forward(self, a):
        return self.net(a)  # [B, 10]

Baseline Train Loop

The forward pass has to carefully synchronise to pass the output from the Stage0 to Stage1. Inversely for backward pass, the Stage1 has to calculate the gradients first before they are sent to Stage0 module for its gradient calculation based on gradients of activation passed.

for ep in range(epochs):
    if pp_rank == 0:
        for x, y in tqdm.tqdm(
            train_loader,
            desc=f"[PP0] ep{ep}",
            leave=False,
        ):
            A_graphs = []

            for m in range(micro_batches):
                xs = x[
                    m * micro_bs:(m + 1) * micro_bs
                ]
                ys = y[
                    m * micro_bs:(m + 1) * micro_bs
                ]

                a = stage(xs)
                A_graphs.append(a)

                act_buf[m, :, :].copy_(a.detach())
                lbl_buf[m, :].copy_(ys.to(torch.int64))

            pp_barrier.wait()  # activations ready
            pp_barrier.wait()  # wait for grad wrt activations

            opt.zero_grad(set_to_none=True)

            for m in reversed(range(micro_batches)):
                dA = gad_buf[m, :, :].clone()
                A_graphs[m].backward(dA)

            opt.step()
            pp_barrier.wait()

    else:
        for _ in range(steps_per_epoch):
            pp_barrier.wait()  # wait activations
            opt.zero_grad(set_to_none=True)

            for m in range(micro_batches):
                a = (
                    act_buf[m, :, :]
                    .clone()
                    .detach()
                    .requires_grad_(True)
                )

                logits = stage(a)
                y = lbl_buf[m, :].clone()

                loss = nn.functional.cross_entropy(
                    logits,
                    y,
                )
                loss.backward()

                gad_buf[m, :, :].copy_(
                    a.grad.detach()
                )

            pp_barrier.wait()  # grad ready
            opt.step()
            pp_barrier.wait()  # end of batch

The problem is clear: there’s a pipeline bubble. Stage 0 does all its forward passes, then sits idle while Stage 1 runs forward and backward. Then Stage 1 sits idle while Stage 0 does backward. At any given moment, roughly half the compute is wasted.

Parallel Pipeline with 1F1B

In our previous implementation, a bottleneck exist because both process sit idle for different stages of training: Stage0 waits after sending activations to Stage1 till gradients are calculated and sent back to Stage0. Similarly, Stage0 does the backward pass and do the forward pass for the next mini-batch while Stage1 model waits.

The schedule runs for M + P - 1 cycles (where M = micro-batches and P = pipeline stages). Each rank calculates which micro-batch to run forward and backward for in each cycle:

for batch_idx, (x, y) in enumerate(
    tqdm.tqdm(
        train_loader,
        desc=pbar_desc,
        leave=False,
    )
):
    # --- Initialization for the batch ---
    opt.zero_grad(set_to_none=True)

    # Used by Stage 0 to hold FWD graphs
    A_graphs = []

    # Used by Stage 1 to hold (a, loss) for BWD
    stage1_graphs = []

    # The core 1F1B loop runs for M + P - 1 cycles
    # (M=micro_batches, P=pp_size)
    for cycle in range(num_pipeline_cycles):

        # --- Calculate micro-batch indices for this cycle ---

        # Forward pass micro-batch index:
        # m_fwd = cycle - p
        #
        # This determines which micro-batch (m_fwd)
        # this rank (p) should process in the current cycle.
        m_fwd = cycle - pp_rank
        is_fwd_step = (
            (m_fwd >= 0)
            and (m_fwd < micro_batches)
        )

        # Backward pass micro-batch index:
        # m_bwd = cycle - p - (P - 1)
        #
        # Since P=2, this simplifies to:
        # m_bwd = cycle - p - 1
        m_bwd = cycle - pp_rank - 1
        is_bwd_step = (
            (m_bwd >= 0)
            and (m_bwd < micro_batches)
        )

        # --- STAGE 0 (Rank 0) Logic: F0 and B0 ---
        if pp_rank == 0:

            if is_fwd_step:
                # Forward Pass (F0)
                xs = x[
                    m_fwd * micro_bs:
                    (m_fwd + 1) * micro_bs
                ]
                ys = y[
                    m_fwd * micro_bs:
                    (m_fwd + 1) * micro_bs
                ]

                a = stage(xs)

                # Store graph for BWD
                A_graphs.append(a)

                # Copy data to SHM for Stage 1
                # (Activations A and Labels Y)
                act_buf[m_fwd, :, :].copy_(
                    a.detach()
                )
                lbl_buf[m_fwd, :].copy_(
                    ys.to(torch.int64)
                )

            if is_bwd_step:
                # Backward Pass (B0)

                # Get dL/dA from SHM
                # (written by Stage 1 in a previous cycle)
                dA = gad_buf[m_bwd, :, :].clone()

                # Backward pass through Stage 0.
                # A_graphs[m_bwd] retrieves the graph
                # from the corresponding FWD pass.
                A_graphs[m_bwd].backward(dA)

        # --- STAGE 1 (Rank 1) Logic: F1 and B1 ---
        else:  # pp_rank == 1

            if is_fwd_step:
                # Forward Pass (F1)

                # Clone A from SHM and enable gradient tracking
                a = (
                    act_buf[m_fwd, :, :]
                    .clone()
                    .detach()
                    .requires_grad_(True)
                )

                logits = stage(a)
                y = lbl_buf[m_fwd, :].clone()

                loss = nn.functional.cross_entropy(
                    logits,
                    y,
                )

                # Store necessary items for BWD pass
                # in a future cycle
                stage1_graphs.append({
                    "a": a,
                    "loss": loss,
                })

            if is_bwd_step:
                # Backward Pass (B1)

                # Retrieve the stored loss graph and 'a' tensor
                graph_data = stage1_graphs[m_bwd]
                loss = graph_data["loss"]

                # 'a' is needed so a.grad can be populated
                a = graph_data["a"]

                # Calculate gradients through Stage 1
                loss.backward()

                # Copy dL/dA to SHM for Stage 0
                gad_buf[m_bwd, :, :].copy_(
                    a.grad.detach()
                )

        # SYNC 1:
        # Barrier at the end of each pipeline cycle.
        #
        # This ensures ranks stay in lock-step.
        # e.g., Rank 1 can't start cycle c+1
        # (and read SHM) until Rank 0 finishes
        # cycle c (and writes SHM).
        pp_barrier.wait()

    # --- Optimization Step ---
    # After all FWD/BWD cycles for the batch
    opt.step()

    # SYNC 2:
    # Ensure all ranks have finished opt.step()
    # before starting the next batch.
    pp_barrier.wait()
Method Total Wall Time (sec) Final Acc
PP (with bubble) 43.75 0.97
PP (1F1B) 35.75 0.94

With the interleaved schedule, we save roughly 18% of training time by reducing idle cycles in the pipeline.

What’s Missing from Real PP

  1. Pipeline Bubbles at Scale: Our implementation splits the model into just two stages, which keeps things simple but doesn’t show the full picture. With more stages (4, 8, 16+), the pipeline bubble in the G-Pipe schedule grows proportionally-the first and last stages sit idle for P - 1 cycles at the start and end of each batch.

    Production systems use more sophisticated schedules like Interleaved 1F1B (assigning multiple non-contiguous stages to each rank), Chimera (bidirectional pipelines), or Zero-Bubble PP (which carefully reorders computation to eliminate the bubble entirely at the cost of increased memory for in-flight micro-batches).

  2. Activation Memory and Re-materialization: The activations passed between stages (act_buf in our SHM setup) can be a massive memory and bandwidth bottleneck for large models with big hidden dimensions and long sequences. Real systems often use activation checkpointing (also called re-materialization): instead of storing the full activation tensor, they discard it after the forward pass and recompute it at the start of the backward pass.

    This trades compute for memory-a worthwhile exchange when activation memory is the binding constraint.

Putting It All Together

This exercise demonstrates the fundamental communication patterns behind the four major parallelism paradigms:

Paradigm What’s Replicated What’s Sharded Communication Pattern
DDP Model, Optimizer Data All-Reduce (gradients)
FSDP Nothing Params, Grads, Optimizer, Data All-Gather (params) + Reduce-Scatter (grads)
TP Some layers Weight matrices (per-layer) All-Reduce (partial outputs)
PP Nothing Model stages Point-to-point (activations, gradients)

While our implementations use shared memory and barriers-a far cry from NCCL over NVLink-the core patterns are identical. The shared memory buffer is the communication channel, and the barrier is the synchronization primitive. Replace shm.buf with a GPU-to-GPU transfer and barrier.wait() with a CUDA stream sync, and the structure is the same.

What We Didn’t Cover

  1. Distributed Inference: A crucial, unaddressed topic-particularly for large language models. Inference often uses the same TP and PP concepts but without backward-pass overhead, and introduces new concerns like KV-cache partitioning and speculative decoding across devices.
  2. Communication Overlap: The single biggest lever for production performance. Every paradigm benefits from overlapping communication with computation-Ring All-Reduce with bucketing (DDP), prefetching the next layer’s params while computing the current layer (FSDP), or overlapping the reduce with downstream compute (TP). Our barrier-based approach forces strict synchronization; real systems carefully pipeline these operations.
  3. Combining Paradigms: Production training runs at scale (e.g., training a 70B+ parameter LLM) rarely use just one paradigm. A typical setup might use TP within a node (8 GPUs connected by NVLink), PP across nodes in a rack, and FSDP across racks-a 3D parallelism configuration where each dimension uses the communication pattern best suited to the available bandwidth.

All code is available in the standalone scripts for each parallelism method.