🎓 A1Learner
🔍

PyTorch Why Distributed

Data parallel training copies one model across several workers, with each copy training on its own slice of the batch. The copies average their gradients after every step, so every copy stays the same model.

Split the batch

A batch of four rows splits into two shards of two rows. Each worker takes one shard.

x = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
print("Worker 1 rows:", x[:2].shape[0])
print("Worker 2 rows:", x[2:].shape[0])

Average the gradients

Two shard gradients average into one shared gradient. Every copy then takes the same step.

grads = torch.tensor([-10.5, -44.5])
print("Averaged weight gradient:", grads.mean().item())

Example

import torch
from torch import nn

torch.manual_seed(0)
x = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
y = torch.tensor([[3.0], [5.0], [7.0], [9.0]])
loss_fn = nn.MSELoss()


def gradient(xs, ys):
    model = nn.Linear(1, 1)
    with torch.no_grad():
        model.weight.fill_(0.5)
        model.bias.fill_(0.0)
    loss = loss_fn(model(xs), ys)
    loss.backward()
    return model.weight.grad.item(), model.bias.grad.item()


full_w, full_b = gradient(x, y)
w1, b1 = gradient(x[:2], y[:2])
w2, b2 = gradient(x[2:], y[2:])
avg_w = (w1 + w2) / 2
avg_b = (b1 + b2) / 2

print("Idea only: two workers, one process, no cluster")
print(f"Worker 1 gradient: weight {w1:.1f}, bias {b1:.1f}")
print(f"Worker 2 gradient: weight {w2:.1f}, bias {b2:.1f}")
print(f"Averaged gradient: weight {avg_w:.1f}, bias {avg_b:.1f}")
print(f"Full-batch gradient: weight {full_w:.1f}, bias {full_b:.1f}")
print("Averaged matches full batch:", avg_w == full_w and avg_b == full_b)

Key points

  • Worker: one process that holds one copy of the model
  • World size: the number of workers training together
  • Shard: the slice of the batch that one worker trains on
  • Each worker runs forward and backward on its own shard alone
  • Averaging the shard gradients gives the full-batch gradient when the shards are equal in size
  • The averaged gradient is the step every copy takes together

In this example, one process plays both workers, and the first printed line says so plainly. The data is four rows from the line y = 2x + 1, split into two shards of two rows. The gradient function builds a fresh Linear model with weight 0.5 and bias 0.0 each time. It runs one backward pass and returns the weight and bias gradients. Worker 1 prints weight -10.5 and bias -6.5 from its shard. Worker 2 prints weight -44.5 and bias -12.5 from its shard. The average of the two shard gradients prints as weight -27.5 and bias -9.5. The full-batch gradient prints the same two numbers. The match holds because the two shards are equal in size. The next lesson runs the averaging step for real, across two processes.

Note

The closing program runs both workers inside one process, so it shows the idea and not a cluster run

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

In data parallel training, what do the workers average after each backward pass?