Data parallelism — why copying the whole model can still make training faster

Over the last few lessons, we have built up the memory side of LLM training:

• AdamW adds optimizer state.
• Gradient accumulation lets a huge statistical batch exist across time.
• Activation checkpointing trades extra compute for less activation memory.

Today we will add the first major form of distributed training: data parallelism.

Its central idea is almost embarrassingly simple:

Each GPU independently computes gradients on its own slice of the batch. Then the GPUs combine those gradients so that every copy of the model makes the same update.

This lets us process more training tokens at once without changing the mathematics of minibatch gradient descent.

Start with one GPU¶

Suppose our model has parameters:

A batch contains 8 sequences:

A B C D E F G H

One GPU performs:

forward on all 8
        ↓
compute loss
        ↓
backward
        ↓
gradient g
        ↓
AdamW update

If those 8 sequences contain \(T\) total tokens, the gradient is the average contribution from those tokens.

Now suppose we have four GPUs.

The naive thought is:

Put a quarter of the model on each GPU.

That is one possible strategy, but it is not data parallelism.

Data parallelism does the opposite.

Each GPU gets the entire model.

Four copies of the same model¶

At the beginning of a training step:

GPU 0: θ
GPU 1: θ
GPU 2: θ
GPU 3: θ

All four copies are identical.

Now split the batch:

GPU 0 → A B
GPU 1 → C D
GPU 2 → E F
GPU 3 → G H

Each GPU runs forward and backward independently.

So we obtain four local gradients:

If each GPU processed the same number of training tokens, the global batch gradient is:

Every GPU then uses that same \(g\) for the optimizer step.

Afterward:

GPU 0: θ'
GPU 1: θ'
GPU 2: θ'
GPU 3: θ'

The copies remain synchronized.

Why this is mathematically equivalent to one larger batch¶

Suppose each GPU processes two examples.

GPU 0 sees examples 1 and 2:

GPU 1 sees 3 and 4:

And similarly for the other GPUs.

Averaging the four GPU gradients gives:

That simplifies to:

Which is exactly the gradient we would have obtained by processing all eight examples as one batch on one sufficiently large GPU.

So data parallelism changes where the computation happens, not the intended gradient.

The classroom mental model¶

Imagine four teaching assistants grading different stacks of exams.

Each TA has the same grading rubric.

TA 0 → exams 1–25
TA 1 → exams 26–50
TA 2 → exams 51–75
TA 3 → exams 76–100

Each TA computes statistics about the mistakes in their stack.

Then they meet and combine their observations:

"Students struggled with question 4."
"Question 7 was mostly fine."
"Question 9 had a systematic misconception."

Only after combining all four stacks does the professor revise the lesson plan.

The model parameters are the lesson plan.

Each GPU has the same plan.

The local gradients are each TA's observations.

The synchronization step combines those observations into one global update.

The synchronization problem¶

The important new systems operation is:

One simple implementation would send every gradient to GPU 0:

GPU 1 ─┐
GPU 2 ─┼──→ GPU 0 → average gradients
GPU 3 ─┘

Then GPU 0 sends the result back:

GPU 0 → GPU 1
      → GPU 2
      → GPU 3

But this creates a bottleneck at GPU 0.

Large models may have billions of gradient values.

We need a distributed way to combine them.

That operation is commonly called all-reduce.

What all-reduce means¶

"Reduce" means combine values across workers, often by summing:

"All" means every worker receives the result.

So:

before:

GPU 0 → g0
GPU 1 → g1
GPU 2 → g2
GPU 3 → g3

after all-reduce:

GPU 0 → g0+g1+g2+g3
GPU 1 → g0+g1+g2+g3
GPU 2 → g0+g1+g2+g3
GPU 3 → g0+g1+g2+g3

Divide by \(N\), if needed by the chosen convention, and every GPU has the same global gradient.

Ring all-reduce: nobody becomes the central server¶

A common mental model is to arrange GPUs in a ring:

GPU 0 → GPU 1 → GPU 2 → GPU 3
  ↑                         ↓
  └─────────────────────────┘

Instead of sending the entire gradient tensor everywhere at once, divide it into chunks.

For four GPUs:

gradient = [chunk A | chunk B | chunk C | chunk D]

During a reduce-scatter phase, chunks circulate around the ring and are progressively summed.

Eventually:

GPU 0 owns summed chunk A
GPU 1 owns summed chunk B
GPU 2 owns summed chunk C
GPU 3 owns summed chunk D

Then an all-gather phase circulates those completed chunks so every GPU reconstructs the full reduced gradient.

Conceptually:

reduce-scatter:
everyone helps compute different pieces of the sum

all-gather:
everyone shares the completed pieces

No single GPU must ingest every other GPU's full gradient as a central coordinator.

Communication is now part of the training step¶

Suppose a model has:

If gradients are represented with 2 bytes per parameter, the gradient tensor alone is roughly:

A full data-parallel synchronization therefore involves enormous amounts of data movement.

The exact network traffic depends on the collective algorithm and topology, but the important point is immediate:

A GPU can finish backward and then sit idle waiting for gradients to arrive from other GPUs.

If that happens often enough, adding more GPUs may barely help.

The scaling ideal¶

Suppose one GPU processes:

In a perfect world:

1 GPU  → 100k tokens/s
2 GPUs → 200k tokens/s
4 GPUs → 400k tokens/s
8 GPUs → 800k tokens/s

That would be linear scaling.

Define scaling efficiency as:

If eight GPUs achieve:

then:

So scaling efficiency is:

The missing 10% may come from communication, synchronization, load imbalance, kernel inefficiencies, or other overhead.

Why faster GPUs make networking more important¶

Suppose backward computation takes:

and gradient synchronization takes:

Then one iteration needs roughly:

Now imagine GPU arithmetic becomes twice as fast.

Backward falls to:

but the network still needs:

Total time becomes:

Compute improved by \(2\times\), but total iteration speed improved only:

The communication fraction became more important.

This is why high-end AI clusters need extremely fast GPU interconnects.

Accelerator performance cannot be considered independently from network performance.

Fortunately, communication can overlap computation¶

Backward does not produce every gradient simultaneously.

Recall backpropagation moves from the final layer toward the first:

layer 80 gradient
        ↓
layer 79 gradient
        ↓
layer 78 gradient
        ↓
...
layer 1 gradient

Once the gradient for layer 80 is ready, we do not necessarily need to wait for layer 1 before starting communication.

We can begin synchronizing earlier gradients while the GPU continues computing later backward work.

Conceptually:

time →

GPU compute:
[L80 backward][L79 backward][L78 backward][L77 backward]...

network:
              [sync L80]
                           [sync L79]
                                        [sync L78]

The goal is to hide communication behind computation.

If done well, much of the all-reduce cost disappears from the critical path.

Gradient buckets make overlap practical¶

Synchronizing every tiny parameter tensor separately would create too many network operations.

So frameworks often group gradients into buckets.

For example:

bucket 1:
layer 80 gradients
layer 79 gradients
layer 78 gradients

bucket 2:
layer 77 gradients
layer 76 gradients
...

When a bucket becomes ready, its all-reduce can begin.

This creates a pipeline:

compute gradients
      ↓
fill bucket
      ↓
launch communication
      ↓
continue backward

Bucket size itself becomes a tuning parameter.

Large buckets improve communication efficiency but may delay the moment communication can begin.

Small buckets start earlier but create more overhead.

Again we see a systems tradeoff.

The slowest GPU controls the step¶

Data parallel training is synchronous.

Before the optimizer update, every worker must contribute its gradient.

Suppose four GPUs take:

GPU 0: 92 ms
GPU 1: 93 ms
GPU 2: 94 ms
GPU 3: 130 ms

The step cannot finish in 94 ms.

Everyone waits for GPU 3.

This is the straggler problem.

Conceptually:

GPU 0: ██████████ done ───────── wait
GPU 1: ███████████ done ──────── wait
GPU 2: ███████████ done ──────── wait
GPU 3: ███████████████████ done

The effective step time is determined by the slowest participant.

At thousands of GPUs, even small variations can matter.

Data loading can create stragglers too¶

The GPU may not be the culprit.

Imagine one worker receives its training batch late because storage or preprocessing stalls.

Then:

GPU 0 → ready
GPU 1 → ready
GPU 2 → ready
GPU 3 → waiting for data

The other GPUs may eventually wait at synchronization.

So high-performance LLM training also requires:

fast storage
data sharding
prefetching
tokenization pipelines
host-to-device transfer

to keep every worker fed.

The training system is only as fast as its slowest recurring bottleneck.

Connect data parallelism to gradient accumulation¶

Yesterday we had:

Now the \(N_{\text{devices}}\) term has a concrete meaning.

Suppose:

and:

Then:

sequences per optimizer update.

The two scaling dimensions do different things:

more data-parallel GPUs
→ process more microbatches simultaneously

more gradient accumulation
→ process more microbatches sequentially

Both enlarge the effective batch.

One uses more space.

The other uses more time.

But ordinary data parallelism duplicates memory¶

Here is the major limitation.

Every GPU holds the entire model:

And in ordinary data-parallel training, each worker may also hold a full copy of:

gradients
Adam first moments
Adam second moments

Suppose the training state for a model requires:

If each GPU has:

of HBM, copying the full state to every GPU is impossible.

Adding more data-parallel GPUs does not solve that.

You would have:

GPU 0: cannot fit model state
GPU 1: cannot fit model state
GPU 2: cannot fit model state
...

The model must fit on each data-parallel worker.

This is the crucial distinction between data parallelism and model sharding.

Data parallelism versus model parallelism¶

Data parallelism divides:

while replicating:

Model parallel techniques divide:

across GPUs.

So:

DATA PARALLELISM

GPU 0: full model + data A
GPU 1: full model + data B
GPU 2: full model + data C
GPU 3: full model + data D

versus:

MODEL PARALLELISM

GPU 0: part of model
GPU 1: part of model
GPU 2: part of model
GPU 3: part of model

same training examples flow through the pieces

Large LLM training usually combines multiple forms of parallelism because neither alone is enough.

ZeRO blurs the boundary¶

In our AdamW lesson, we mentioned ZeRO-style sharding.

ZeRO begins with the data-parallel setup but asks:

If every data-parallel worker ultimately needs the same update, do all of them really need permanent copies of every piece of training state?

Different ZeRO stages shard progressively more state:

optimizer states
      ↓
gradients
      ↓
parameters

across data-parallel workers.

This reduces memory duplication.

The price is additional communication because a worker may need to fetch parameter or state pieces it does not permanently own.

This turns the simple picture:

every GPU has everything

into a more sophisticated one:

the logical model is replicated,
but physical storage of its training state is distributed.

That idea will lead naturally into fully sharded data parallelism.

The deeper mental model: parallel training is a synchronized vote¶

Think of each GPU as independently asking:

"Given the examples I just saw, how should the model change?"

GPU 0 votes:

GPU 1 votes:

And so on.

The optimizer is not allowed to act until the votes are combined:

Then every model copy receives the same decision.

This gives us the essential rhythm of synchronous data-parallel training:

copy same model everywhere
        ↓
show each copy different data
        ↓
compute local gradients independently
        ↓
communicate and average gradients
        ↓
apply identical optimizer update
        ↓
repeat

The beauty is that most expensive Transformer arithmetic happens independently on every GPU.

The difficulty is the synchronization boundary.

Put the last four lessons together¶

We can now see four different levers for making LLM training feasible and fast.

AdamW tells us why training has large parameter-sized state.

Gradient accumulation spreads a large statistical batch across time.

Activation checkpointing spreads activation storage across time, recomputing what we choose not to remember.

Data parallelism spreads batch computation across space, using many GPUs simultaneously.

So the emerging picture is:

one giant training step on paper

                ↓ decompose

many GPUs
×
many microbatches
×
recomputed activation segments
×
synchronized gradients

                ↓

one optimizer update

The important principle is:

A single-GPU equation might say:

A giant training cluster still wants essentially that gradient.

It just obtains it by splitting the work across thousands of processors and then carefully putting the pieces back together.

That is the recurring trick behind large-scale LLM training:

The mathematics describes one model learning from one batch. The systems engineering makes that illusion survive across an entire data center.