How to Set Up Distributed Training for Deep Learning

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:

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:

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:

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:

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

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

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