Scaling Model Training with Distributed Computing

When training takes too long on one machine, the obvious move is to add machines. Distributed training does exactly that, but every worker has to stay in step with the others, and that coordination is not free. Whether more workers make training faster depends on how much work each does between synchronisations.

This guide explains the main ways to split training, shows how data parallelism keeps every copy of the model identical, and measures when adding workers helped and when it hurt.

Ways to split training

Most teams start with data parallelism. When the model outgrows a single device, PyTorch's FSDP2 offers fully sharded data parallelism with per-parameter sharding (PyTorch: fully_shard).

How data parallelism stays correct

Each worker runs the same model on a different slice of each batch. After the backward pass, gradients are averaged across workers, so every copy applies the same update. Each worker also needs the same starting weights and the same data order; the script fixes both with shared random seeds, which is why its replicas match exactly. PyTorch's DistributedDataParallel groups gradients into buckets and reduces one bucket at a time, overlapping that communication with the rest of the backward pass (PyTorch: DDP notes).

The communication runs over a backend. PyTorch's rule of thumb is NCCL for NVIDIA CUDA GPUs, XCCL for Intel GPUs and Gloo for CPUs; the distributed package is stable on Linux and macOS and a prototype on Windows (PyTorch: distributed).

The experiment

The model is a three-layer network with 1.6 million parameters, trained on smartphone activity data from the UCI HAR dataset, under CC BY 4.0 (UCI: HAR). One, two and four worker processes run on one laptop over Gloo, each limited to two CPU threads so that it stands in for a separate machine. The global batch stays fixed, so more workers means smaller slices each.

import os
import time

import numpy as np
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn.functional as F
from torch.nn.parallel import DistributedDataParallel as DDP

THREADS_PER_WORKER = 2                                           # each worker stands in for one machine


def load(split):                                                 # UCI HAR: 561 features, six activities
    x = torch.from_numpy(np.loadtxt(f"har/{split}/X_{split}.txt", dtype=np.float32))
    y = torch.from_numpy(np.loadtxt(f"har/{split}/y_{split}.txt", dtype=np.int64) - 1)
    return x, y


def build():
    return torch.nn.Sequential(torch.nn.Linear(561, 1024), torch.nn.ReLU(),
                               torch.nn.Linear(1024, 1024), torch.nn.ReLU(), torch.nn.Linear(1024, 6))


def worker(rank, world, global_batch, epochs, results):
    os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT="29511", USE_LIBUV="0")
    dist.init_process_group("gloo", rank=rank, world_size=world)    # Gloo: PyTorch's backend for CPUs
    torch.set_num_threads(THREADS_PER_WORKER)
    torch.manual_seed(0)                                         # identical starting weights on every worker
    (x, y), (test_x, test_y) = load("train"), load("test")
    model = DDP(build())
    optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
    share = global_batch // world                                # each worker takes its slice of every global batch
    dist.barrier()
    start = time.perf_counter()
    for epoch in range(epochs):
        order = torch.randperm(len(x), generator=torch.Generator().manual_seed(epoch))   # same order everywhere
        for begin in range(0, len(x) - global_batch + 1, global_batch):
            rows = order[begin + rank * share: begin + (rank + 1) * share]
            optimizer.zero_grad()
            F.cross_entropy(model(x[rows]), y[rows]).backward()  # DDP averages gradients across workers here
            optimizer.step()
    seconds = time.perf_counter() - start
    checksum = sum(p.detach().double().sum().item() for p in model.parameters())
    with torch.no_grad():
        accuracy = (model(test_x).argmax(1) == test_y).float().mean().item()
    results[rank] = (seconds, checksum, accuracy)
    dist.destroy_process_group()


def run(world, global_batch, epochs):
    with mp.Manager() as manager:
        results = manager.dict()
        mp.spawn(worker, args=(world, global_batch, epochs, results), nprocs=world)
        return [results[r] for r in range(world)]


if __name__ == "__main__":
    print(f"{sum(p.numel() for p in build().parameters()):,} parameters, {THREADS_PER_WORKER} CPU threads per worker")
    print("\n1. Same training, split across workers (global batch 256, 5 epochs)")
    for world in (1, 2, 4):
        runs = run(world, 256, 5)
        same = len({round(r[1], 6) for r in runs}) == 1
        print(f"  {world} worker(s): every replica identical {same}; test accuracy {runs[0][2]:.2%}")
    print("\n2. Speed-up over one worker (varies)")
    for global_batch, epochs in ((256, 10), (2048, 40)):
        times = {world: max(r[0] for r in run(world, global_batch, epochs)) for world in (1, 2, 4)}
        print(f"  global batch {global_batch:>4}: " + ", ".join(
            f"{w} workers {times[1] / times[w]:.2f}x" for w in (2, 4)) + f" (one worker {times[1]:.1f} s) (varies)")

1,631,238 parameters, 2 CPU threads per worker

1. Same training, split across workers (global batch 256, 5 epochs)
  1 worker(s): every replica identical True; test accuracy 71.60%
  2 worker(s): every replica identical True; test accuracy 71.60%
  4 worker(s): every replica identical True; test accuracy 71.60%

2. Speed-up over one worker (varies)
  global batch  256: 2 workers 0.39x, 4 workers 0.54x (one worker 4.4 s) (varies)
  global batch 2048: 2 workers 1.01x, 4 workers 1.49x (one worker 15.0 s) (varies)

What the output shows

• The copies stayed identical. With one, two or four workers, every replica ended with the same weights as the others, and test accuracy was 71.60% each time. Data parallelism with a fixed global batch is the same training, divided up.
• Small steps got slower with more workers. With a global batch of 256, two workers took about two and a half times as long as one, and four took nearly twice as long. Each step averaged roughly 6.5 MB of gradients, 1.6 million four-byte numbers, while the computation per step was tiny.
• Large steps began to pay. With a global batch of 2,048, two workers only broke even, and four trained 1.49 times faster, an efficiency of about 37%.

This is the trade-off Goyal and colleagues describe: for distributed training to be efficient, each worker's share of work must be large, which means larger minibatches. Large minibatches cause optimisation difficulties, which they solved with a linear rule for scaling the learning rate and a warm-up period. That let them train ResNet-50 with minibatches of 8,192 on 256 GPUs in one hour at about 90% scaling efficiency from 8 to 256 GPUs (Goyal et al.: large minibatch SGD).

The laptop here is a harsh test: its workers talk over the loopback network. Real clusters use dedicated interconnects and NCCL, and scale far better, but the principle is the same everywhere: compute per step has to outweigh communication per step.

Before going distributed

1. Profile the single-worker job. Slow data loading or preprocessing will not improve with more workers.
2. Use the biggest single device you can, which is often simpler and cheaper; see GPU vs CPU for ML.
3. Raise the batch size, adjusting the learning rate and warm-up, before adding workers.
4. Add workers when each will have plenty to do, and measure scaling efficiency, not just speed.
5. Shard when the model will not fit, rather than shrinking it to fit.

Distributed jobs usually run on a cluster scheduler; see Kubernetes for ML. For what the extra machines cost, see cost optimisation for ML workloads.

Common mistakes

• Adding workers to a job bottlenecked on data loading.
• Keeping the per-worker batch tiny, so communication dominates.
• Raising the global batch without adjusting the learning rate.
• Reporting speed-up without efficiency per worker.
• Different random seeds or data orders per worker, which breaks replica consistency.

Distributed training basics are part of the Advanced MLOps Concepts module of the MLOps course.

Sources

• Distributed communication package (PyTorch documentation) — PyTorch
• Distributed Data Parallel (PyTorch documentation) — PyTorch
• torch.distributed.fsdp.fully_shard (PyTorch documentation) — PyTorch
• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour (arXiv) — Goyal et al.
• Human Activity Recognition Using Smartphones (UCI Machine Learning Repository) — UCI Machine Learning Repository