Training at Scale

The largest models are trained on tens of thousands of G P Us, all working on one model at the same time. How can so many machines cooperate, without getting in each other's way? The core idea is much simpler than you might think.

• Tens of thousands of GPUs
• One single model
• The core idea is simple

Sharing gradients

Each G P U gets a different slice of the batch. Each computes its own gradient. Then they all share and average those gradients, and every G P U applies the very same update. So their copies of the model stay identical. Here, we simulate it with four G P Us. Run it. It's simply an average. In PyTorch, distributed data parallel does all of this for you, with a few extra lines.

import torch

def average_gradients(grads_per_gpu):
    return torch.stack(grads_per_gpu).mean(dim=0)

g0 = torch.tensor([0.2, -0.4])
g1 = torch.tensor([0.4, -0.2])
g2 = torch.tensor([0.3, -0.6])
g3 = torch.tensor([0.5, -0.2])

avg = average_gradients([g0, g1, g2, g3])
print("Each GPU saw different data")
print("Shared gradient:", avg.tolist())
print("All GPUs apply the same update, so weights stay identical")

Each GPU saw different data
Shared gradient: [0.3500000238418579, -0.3500000238418579]
All GPUs apply the same update, so weights stay identical

What makes it hard

Scale brings its own trouble. Gradients have to be shared quickly, which is why fast networks matter so much. Hardware fails, so regular checkpoints are essential. If you use many G P Us, you usually have to raise the learning rate. And sometimes models are so large that they have to be split across devices. This is why the biggest training runs need teams of specialists.

Summary