Core ML

Data Parallelism

If you have 8 GPUs, put an identical copy of the model on every GPU. Give each GPU a different slice of the dataset. They each calculate how the model should change, average their answers together, and update simultaneously. You just trained 8x faster.

Why Does This Exist?

When training a neural network, processing a massive dataset (like 300 billion tokens of text) takes an excruciating amount of time. If a single GPU can process a batch of 32 sequences per second, training the model might take 3 years.

To speed this up, you buy 8 GPUs. But how do you actually make them work together on the same model?

Data Parallelism (DP) is the most intuitive and widespread form of distributed training. It scales the training process horizontally by splitting the Data, not the Model. As long as the entire model can physically fit into the VRAM of a single GPU, Data Parallelism allows you to scale your effective batch size infinitely, slashing your training time from years to weeks.

Think of It Like This

Imagine you have to grade 8,000 math exams. You have an Answer Key (the Model). It takes you 80 hours to grade them all yourself.

You hire 7 assistants, so there are 8 of you in total. You print 8 identical copies of the Answer Key (Model Replicas) and hand one to each person. You divide the 8,000 exams into 8 stacks of 1,000 (Data Shards) and give one stack to each person. You all grade at the same time. At the end of the day, you meet, sum up the total number of mistakes everyone found, update your rubrics together, and start the next day. You just finished 80 hours of work in 10 hours.

How It Actually Works

The math behind Data Parallelism is elegant. Because the gradient of a sum is equal to the sum of the gradients, calculating the gradient on 8 small batches and adding them together is mathematically identical to calculating the gradient on 1 massive batch.

Here is the exact step-by-step loop of PyTorch's DistributedDataParallel (DDP):

1. Replicate and Shard

At step 0, the orchestrator copies the exact same Model Weights to all 8 GPUs. It also configures the DistributedSampler to ensure that if GPU 1 reads data rows 1-32, GPU 2 reads rows 33-64. No two GPUs will see the exact same data in a single step.

2. Forward and Backward Pass (Independent)

Each GPU runs a standard forward pass on its unique data slice. It calculates the loss, and runs a backward pass to calculate the Gradients. At this exact millisecond, GPU 1 and GPU 2 have different gradients, because they looked at different data.

3. The `All-Reduce` Operation (Synchronized)

Before the Optimizer is allowed to update the weights, the GPUs pause. Over the high-speed NVLink or InfiniBand network, they perform an operation called All-Reduce. All-Reduce takes the gradients from GPU 1, GPU 2, ... GPU 8, adds them all together, and broadcasts the sum back to every single GPU.

4. The Optimizer Step (Identical)

Now, every single GPU has the exact same summed gradient. They each run optimizer.step(). Because they started with identical weights, and they just applied identical gradients, they end up with identical updated weights. The loop repeats.

Show Me the Code

You do not write the network synchronization code yourself. PyTorch DDP handles the All-Reduce completely under the hood by hooking into the .backward() call.

1import torch2import torch.nn as nn3from torch.nn.parallel import DistributedDataParallel as DDP4
5# Assume this script was launched with torchrun on 8 GPUs.6# local_rank is 0 on GPU 1, 1 on GPU 2, etc.7local_rank = int(os.environ["LOCAL_RANK"])8torch.cuda.set_device(local_rank)9
10# 1. Create the model and copy it to this specific GPU11model = ResNet50().to(local_rank)12
13# 2. Wrap it in DDP. 14# DDP automatically broadcasts the initial weights from GPU 0 to ensure 15# everyone starts perfectly synced.16ddp_model = DDP(model, device_ids=[local_rank])17
18optimizer = torch.optim.Adam(ddp_model.parameters(), lr=1e-3)19
20for batch in dataloader:21    # 3. Independent Forward Pass22    inputs, labels = batch[0].to(local_rank), batch[1].to(local_rank)23    outputs = ddp_model(inputs)24    loss = loss_fn(outputs, labels)25    26    # 4. Independent Backward Pass + Automatic All-Reduce!27    # As the gradients flow backwards, DDP intercepts them and sums them 28    # across all 8 GPUs over the network in the background.29    loss.backward()30    31    # 5. Identical Optimizer Step32    # Every GPU applies the exact same summed gradient.33    optimizer.step()34    optimizer.zero_grad()

Watch Out For

The Batch Size Scaling Trap. If 1 GPU processes a batch of 32, then 8 GPUs running Data Parallelism process an Effective Batch Size of $32 \times 8 = 256$ . If you scale up to 1,024 GPUs, your effective batch size becomes 32,768! As batch size increases, the gradient becomes less noisy, which is usually good. But if the batch size becomes too massive, the model loses the stochastic noise it needs to escape local minima, and final accuracy will drop. You must adjust your Learning Rate (usually scaling it up linearly: new LR = old LR $\times$ number of GPUs) to compensate.

The Quick Version

Data Parallelism scales training horizontally across multiple GPUs.
Every GPU holds an identical copy of the model.
The dataset is sharded so each GPU processes different rows of data.
After calculating gradients independently, the GPUs pause and sum their gradients together over the network using an All-Reduce operation.
Because every GPU applies the exact same averaged gradient, the model copies stay perfectly synchronized.
It only works if the entire model can fit inside the VRAM of a single GPU.

Related concepts

ZeRO & FSDP