# How to Set Up Distributed Training for Deep Learning

> Master distributed training for deep learning. Learn to set up single-node multi-GPU with DataParallel or multi-node scaling with DistributedDataParallel and torch distributed.

- Repository: [Harvard Edge Computing/cs249r_book](https://github.com/harvard-edge/cs249r_book)
- Tags: how-to-guide
- Published: 2026-02-19

---

**You can set up distributed training using `torch.nn.DataParallel` for single-node multi-GPU setups or `torch.nn.parallel.DistributedDataParallel` with `torch.distributed` for multi-node high-performance scaling.**

The **cs249r_book** repository provides a comprehensive guide to distributed deep learning, covering both the mathematical foundations and practical PyTorch implementations. This article distills the essential patterns from the repository's training documentation, giving you copy-paste ready code to scale your models from a single workstation to a multi-node cluster.

## Single-Node Multi-GPU with DataParallel

For small-to-medium scale training on workstations with 2–4 GPUs, `torch.nn.DataParallel` offers the fastest path to parallelization. This approach replicates your model on every available GPU of a single node, automatically splitting mini-batches and gathering gradients in a parameter-server-style architecture.

### Implementation Details

According to the source in `book/quarto/contents/core/training/training.qmd` (lines 3575–3582), wrapping your model requires minimal code changes:

```python
import torch
import torch.nn as nn

model = MyModel()                     # any nn.Module

model = torch.nn.DataParallel(model)  # <-- wraps the model

# training loop stays unchanged

for xb, yb in dataloader:
    out = model(xb)                   # automatically splits xb across GPUs

    loss = loss_fn(out, yb)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

```

**How it works**: `DataParallel` creates a *master* copy on `cuda:0` and *worker* copies on the remaining GPUs. During the forward pass, input batches are scattered across devices, and gradients are collected on the master GPU before the optimizer updates the shared parameters.

**Limitations**: All gradients flow through the master GPU, creating a communication bottleneck that limits scaling beyond 4–8 GPUs. This makes `DataParallel` unsuitable for large-scale distributed training across multiple nodes.

## Multi-Node Scalable Training with DistributedDataParallel

For large-scale training across many GPUs or multiple nodes, `torch.nn.parallel.DistributedDataParallel` (DDP) provides near-linear scaling. This approach runs one dedicated training process per GPU, using the **NCCL** backend for high-performance GPU-to-GPU communication and **All-Reduce** collectives to average gradients.

### Process Group Initialization

As implemented in `book/quarto/contents/core/training/training.qmd` (lines 3585–3594), you must explicitly initialize a process group before wrapping your model:

```python
import torch
import torch.distributed as dist
import torch.nn.parallel as DDP

def init_ddp():
    # Use the NCCL backend for GPU-GPU communication

    dist.init_process_group(backend="nccl")      # <-- explicit process-group init

    torch.cuda.set_device(local_rank)            # each process gets its own GPU

model = MyModel().to(device)
model = DDP.DistributedDataParallel(model)      # <-- wraps the model

# Inside each process the loop looks the same as before

for xb, yb in dataloader:
    xb = xb.to(device); yb = yb.to(device)
    out = model(xb)
    loss = loss_fn(out, yb)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

```

**Why this scales better**: Each replica runs in its own process with a dedicated GPU, eliminating the master-GPU bottleneck. `DistributedDataParallel` internally performs **ring-AllReduce** (or other efficient algorithms) to average gradients, overlapping the reduction with the backward pass for maximum throughput.

### Launch Configuration with torchrun

To execute distributed training, launch your script using `torchrun` (or `python -m torch.distributed.launch`) with the `--nproc_per_node` flag set to the number of GPUs:

```bash
torchrun --nproc_per_node=2 example_ddp.py

```

For multi-node clusters, set the `MASTER_ADDR` and `MASTER_PORT` environment variables to the head node's address and launch the script on each host with appropriate `--nnodes` and `--node_rank` flags. Ensure **NCCL** (or Gloo) is installed and the machines have a fast interconnect such as InfiniBand.

## Complete Implementation Examples

The cs249r_book repository provides stand-alone scripts demonstrating both patterns. Below are adapted versions you can run immediately.

### DataParallel Example

This script converts a single-GPU training loop to use `DataParallel`:

```python

# example_dp.py

import torch
import torch.nn as nn
import torch.optim as optim
from tinytorch.core.tensor import Tensor   # TinyTorch tensor wrapper

from tinytorch.core.layers import Linear
from tinytorch.core.losses import MSELoss

# ------- Model -------------------------------------------------

class SimpleMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = Linear(2, 1)

    def forward(self, x):
        return self.fc.forward(x)

model = SimpleMLP()
model = torch.nn.DataParallel(model)   # <‑‑ wrap for data-parallel

# ------- Data --------------------------------------------------

X = Tensor([[1.0, 0.5], [0.5, 1.0]])   # tiny toy data

y = Tensor([[2.0], [1.5]])

optimizer = optim.SGD(model.parameters(), lr=0.01)
loss_fn   = MSELoss()

# ------- Training ----------------------------------------------

for epoch in range(5):
    optimizer.zero_grad()
    preds = model(X)
    loss  = loss_fn(preds, y)
    loss.backward()
    optimizer.step()
    print(f"epoch {epoch}: loss={loss.item():.4f}")

```

Run with: `python example_dp.py` (requires a machine with ≥2 GPUs to observe parallel speed-up).

### DistributedDataParallel Example

This script demonstrates proper process initialization and distributed sampling:

```python

# example_ddp.py

import os, torch, torch.distributed as dist, torch.nn as nn, torch.optim as optim
from tinytorch.core.tensor import Tensor
from tinytorch.core.layers import Linear
from tinytorch.core.losses import MSELoss
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, TensorDataset, DistributedSampler

def setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "127.0.0.1"      # replace with master node IP

    os.environ["MASTER_PORT"] = "29500"
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def cleanup():
    dist.destroy_process_group()

class SimpleMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = Linear(2, 1)

    def forward(self, x):
        return self.fc.forward(x)

def main(rank, world_size):
    setup(rank, world_size)

    # ----- Model -------------------------------------------------

    model = SimpleMLP().to(rank)
    model = DDP(model, device_ids=[rank])

    # ----- Data --------------------------------------------------

    X = Tensor([[1.0, 0.5], [0.5, 1.0], [0.2, 0.8], [0.9, 0.1]])
    y = Tensor([[2.0], [1.5], [1.0], [2.2]])
    dataset = TensorDataset(X, y)

    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    loader  = DataLoader(dataset, batch_size=2, sampler=sampler)

    optimizer = optim.SGD(model.parameters(), lr=0.01)
    loss_fn   = MSELoss()

    # ----- Training ---------------------------------------------

    for epoch in range(5):
        sampler.set_epoch(epoch)            # shuffle each epoch

        for xb, yb in loader:
            xb = xb.to(rank); yb = yb.to(rank)
            optimizer.zero_grad()
            preds = model(xb)
            loss  = loss_fn(preds, yb)
            loss.backward()
            optimizer.step()
        if rank == 0:
            print(f"[rank {rank}] epoch {epoch}: loss={loss.item():.4f}")

    cleanup()

if __name__ == "__main__":
    # torchrun automatically passes rank/world_size; manual launch example:

    # torchrun --nproc_per_node=2 example_ddp.py

    pass

```

Run on a single node with 2 GPUs: `torchrun --nproc_per_node=2 example_ddp.py`.

## Theoretical Foundations in the cs249r Book

The repository connects these implementations to their mathematical underpinnings. In `book/quarto/contents/core/training/training.qmd` (lines ~3620–3630), the **Parallelism Opportunities** section explains the communication cost of **All-Reduce** operations, demonstrating why ring-based reductions achieve **O(N)** scaling instead of the **O(N²)** complexity of naive parameter-server architectures.

The **Framework Integration** and **Data Parallel Framework APIs** subsections (lines ~3650–3680) tie this theory to the concrete PyTorch code shown above, helping you understand the trade-offs between the two APIs. The book also references Goyal et al.'s 2017 paper on linear scaling of ImageNet training (arXiv:1706.02677), which established the learning rate scaling rules essential for effective distributed training.

## Summary

- **`torch.nn.DataParallel`** provides the simplest distributed training setup for single-node, multi-GPU workstations, but suffers from GPU 0 bottlenecks at scale.
- **`torch.nn.parallel.DistributedDataParallel`** requires explicit process group initialization with `dist.init_process_group(backend="nccl")`, but eliminates bottlenecks through All-Reduce collectives.
- Launch DDP scripts using `torchrun --nproc_per_node=NUM_GPUS` to automatically handle process spawning and rank assignment.
- Use `DistributedSampler` to ensure each process trains on disjoint data shards without duplication.
- The cs249r_book repository contains the complete theoretical background and working code in `book/quarto/contents/core/training/training.qmd`.

## Frequently Asked Questions

### When should I use DataParallel versus DistributedDataParallel?

Use **DataParallel** only for quick experiments on single nodes with 2–4 GPUs where you want minimal code changes. Use **DistributedDataParallel** for any production training, multi-node setups, or when scaling beyond 4 GPUs, as it provides significantly better performance through overlapping communication and computation.

### What is the NCCL backend and why is it required?

**NCCL** (NVIDIA Collective Communications Library) is a high-performance GPU communication library optimized for NVIDIA GPUs. When you call `dist.init_process_group(backend="nccl")`, PyTorch uses NCCL to execute ring-AllReduce operations across GPUs, achieving bandwidth-optimal gradient synchronization that Gloo or MPI backends cannot match for GPU workloads.

### How do I configure the learning rate for distributed training?

When increasing the batch size linearly with the number of GPUs, you should scale the learning rate proportionally (the "linear scaling rule" cited in the cs249r_book). If you double the effective batch size by adding GPUs, double your base learning rate to maintain similar convergence dynamics, though you may need gradual warmup to prevent early training instability.

### Can I run DistributedDataParallel on a single machine with multiple GPUs?

Yes. While designed for multi-node clusters, DDP works efficiently on single machines with multiple GPUs. Set `MASTER_ADDR` to `localhost` and launch with `torchrun --nproc_per_node=NUM_GPUS`. This still provides better performance than DataParallel because it eliminates the parameter-server bottleneck, even within a single node.