Distributed Training

This closing note of the Production DL & MLOps category covers distributed training — spreading a training job across multiple GPUs or machines, necessary when a single device's compute or memory can't handle a model or dataset at the desired scale.

The Two Main Distributed Training Strategies

Code — Data Parallel Training with DistributedDataParallel

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def train(rank, world_size):
    setup(rank, world_size)
    model = MyModel().to(rank)
    model = DDP(model, device_ids=[rank])

    sampler = torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    loader = torch.utils.data.DataLoader(dataset, batch_size=32, sampler=sampler)

    for x_batch, y_batch in loader:
        optimizer.zero_grad()
        loss = loss_fn(model(x_batch.to(rank)), y_batch.to(rank))
        loss.backward()   # gradients are automatically synchronized across all GPUs here
        optimizer.step()

DistributedDataParallel (DDP) is PyTorch's standard, efficient approach — each process runs on one GPU with its own model replica, and gradient synchronization happens automatically and efficiently during backward(), keeping every replica's weights consistent after each step.

Why Distributed Training Isn't Simply "More Is Always Better"

Synchronizing gradients across many GPUs/machines introduces communication overhead — beyond a certain point, adding more devices yields diminishing returns as communication cost starts to rival or exceed the computation time saved. Effective distributed training requires attention to network bandwidth between devices, efficient synchronization strategies, and appropriately scaling other hyperparameters (like learning rate) as effective batch size grows with more parallel workers.

When Distributed Training Is Actually Necessary

• The model doesn't fit in a single GPU's memory at all — model parallelism is required, not optional.
• Training time on a single GPU is impractically long for the project's timeline, and the added engineering complexity of distributed training is justified by the time saved.
• Very large-scale datasets where single-GPU throughput becomes the limiting factor for how much data can be processed in a reasonable time.

Common Mistakes

• Reaching for distributed training as a default for models that would train perfectly well, and much more simply, on a single GPU — the added engineering and debugging complexity should be justified by an actual need.
• Not adjusting the learning rate when scaling up to more parallel workers — since the effective batch size grows with data parallelism, learning rate often needs to be scaled up correspondingly (see Learning Rate Tuning) to maintain similar training dynamics.

Interview Relevance

Q: "What's the difference between data parallelism and model parallelism in distributed training, and when would you need model parallelism specifically?" Data parallelism replicates the full model across multiple GPUs, with each processing a different slice of the batch and synchronizing gradients afterward — appropriate when the model fits comfortably in a single GPU's memory but faster training is needed. Model parallelism instead splits the model itself across multiple GPUs, which is necessary specifically when a model is too large to fit in a single GPU's memory at all, regardless of how much training speed is needed — the two strategies address fundamentally different constraints (throughput vs memory capacity) and are sometimes combined for very large-scale training.

Key Takeaways — Production DL & MLOps

• ML pipelines formalize the project lifecycle into reproducible, automatable stages, with data validation catching problems before they propagate into expensive training runs.
• Experiment tracking (e.g. MLflow) and a model registry make experimentation and deployment history systematically queryable, rather than relying on memory or informal conventions.
• Data and model versioning close reproducibility gaps that code versioning alone leaves open.
• Data drift, concept drift, and model drift together explain why a deployed model's performance can silently degrade over time without any code change — ongoing monitoring and periodic retraining address this.
• A/B testing validates a new model's real-world impact on live traffic before full rollout, complementing offline evaluation.
• Inference latency, throughput, GPU utilization, and memory are interrelated production efficiency concerns, each with distinct diagnostic approaches and optimization levers.

Next: Modern AI covers the newest wave of deep learning concepts — embeddings, vector databases, RAG, AI agents, and more.

Practice Question

Why does adding more GPUs to a distributed training job eventually yield diminishing returns, rather than continuing to speed up training proportionally?

Related DL Notes

• GPU Utilization
• Memory Optimization
• Mixed Precision Training