Tensor parallelism — how several GPUs cooperate on one Transformer layer

Yesterday we saw how ZeRO lets a training job shard persistent model state instead of storing a complete copy on every data-parallel worker.

But ZeRO leaves us with a different question:

What if a single Transformer layer is itself too large — or too slow — for one GPU?

Then we can split the computation inside the layer across multiple accelerators.

This is tensor parallelism.

The core idea is:

That sounds straightforward. The interesting part is deciding where to split the tensors so that the GPUs do useful local work before they have to communicate.

Start with one linear layer¶

A large fraction of Transformer computation is matrix multiplication.

Suppose a layer computes:

where:

and:

Here \(B\) can represent the collection of token positions being processed together.

If \(W\) is enormous, one GPU must both:

1. store the entire matrix,
2. perform the entire multiplication.

Tensor parallelism asks:

Can GPU 0 own part of \(W\), GPU 1 own another part, and both compute at the same time?

Yes.

But there are two especially useful ways to cut the matrix.

Column parallelism: split the outputs¶

Suppose we split \(W\) by columns:

Then:

So two GPUs can compute:

and:

Conceptually:

                  ┌→ GPU 0: W1 → Y1
X on both GPUs ───┤
                  └→ GPU 1: W2 → Y2

final Y = [Y1 Y2]

Each GPU owns only half of the output features.

This is often called column-parallel linear algebra because we partition the columns of the weight matrix.

The useful property is that each GPU can compute its output shard independently once it has \(X\).

A concrete example¶

Suppose:

and:

Split the columns:

GPU 0 computes:

GPU 1 computes:

Concatenate them:

The result is exactly what one GPU would have computed with the full \(W\).

We changed the physical execution, not the mathematics.

Row parallelism: split the inputs¶

Now split \(W\) by rows:

To match that partition, split the input features:

Then:

GPU 0 computes:

while GPU 1 computes:

But neither GPU has the final answer.

We must combine:

Conceptually:

GPU 0: X1 × W1 → Z1 ─┐
                      ├→ SUM → Y
GPU 1: X2 × W2 → Z2 ─┘

That sum requires communication across the tensor-parallel group.

This is the central tension in tensor parallelism:

The jigsaw-puzzle mental model¶

Imagine two people assembling a giant picture.

A bad strategy is:

After placing every puzzle piece, stop and compare the entire puzzle.

They spend more time coordinating than assembling.

A better strategy is:

Give each person a large region they can work on independently, then synchronize only when the regions must connect.

Tensor parallelism works the same way.

The goal is not merely to divide arithmetic evenly.

The goal is to create large chunks of local computation between communication points.

That is why the exact tensor partition matters.

Transformers contain a natural column-then-row pattern¶

Consider a simplified feed-forward network:

Typically \(W_1\) expands the hidden dimension and \(W_2\) projects it back down.

For example:

A clever tensor-parallel layout can split \(W_1\) by columns.

Each GPU produces only part of the expanded hidden representation:

                 ┌→ GPU 0: H0
X replicated ────┤
                 └→ GPU 1: H1

Now notice something important.

The activation function \(\phi\), such as GELU or SiLU, is elementwise.

Therefore:

We do not need to gather the full \(H\) before applying the nonlinearity.

Each GPU can continue locally.

Then \(W_2\) can be row-partitioned so that each GPU consumes its local hidden shard:

GPU 0: H0 → local W2 shard → Z0 ─┐
                                 ├→ SUM → Y
GPU 1: H1 → local W2 shard → Z1 ─┘

We delayed communication until after substantial useful work.

That is much better than synchronizing after every operation.

This is why tensor-parallel design follows the computation graph¶

Suppose we partition a tensor, then immediately run an operation requiring the full tensor.

We would have to gather it right away.

The sharding accomplished little.

A good partition instead tries to preserve a tensor's distributed layout across several compatible operations.

Conceptually:

BAD:

split
 ↓
matmul
 ↓
GATHER
 ↓
elementwise op
 ↓
split again


BETTER:

split
 ↓
matmul
 ↓
elementwise op
 ↓
next sharded matmul
 ↓
REDUCE

The second design gets more computation per communication event.

This ratio is analogous to arithmetic intensity inside a GPU.

At cluster scale, we care about something like:

Attention can also be tensor-parallel¶

Multi-head attention gives us another natural partition.

Suppose attention has:

query heads and we use 8 GPUs.

A simple strategy is to assign:

heads to each GPU.

Conceptually:

GPU 0 → heads  0–7
GPU 1 → heads  8–15
GPU 2 → heads 16–23
...
GPU 7 → heads 56–63

Many computations for those heads can happen locally.

This is appealing because attention heads are partly independent before their outputs are combined.

Again, model structure gives us a useful place to cut the computation.

GQA makes the partition more subtle¶

Recall grouped-query attention:

Suppose:

Across 8 tensor-parallel GPUs, we might naturally assign one K/V head and its associated query heads to each GPU.

But if:

and we still want 8-way tensor parallelism, there are fewer K/V heads than GPUs.

Now the serving or training system may need to replicate some K/V information or choose a different partition.

So model architecture can constrain efficient parallelism.

This is a recurring theme:

Why tensor parallelism needs fast links¶

Imagine splitting one Transformer layer across 8 GPUs.

Every forward pass may require collective communication between those GPUs.

Then backward requires additional communication.

And this happens layer after layer:

layer 1:
compute → communicate

layer 2:
compute → communicate

layer 3:
compute → communicate

...

Unlike ZeRO parameter gathering, tensor-parallel communication sits directly inside the computation of each layer.

Latency and bandwidth therefore matter enormously.

This is why tensor-parallel groups are often placed across GPUs with very fast interconnects.

A rough performance model for a collective is:

For large tensors, bandwidth dominates.

For many small collectives, latency can matter greatly.

More GPUs do not imply proportional speedup¶

Suppose one GPU takes:

to compute a layer.

With two GPUs, perhaps each performs roughly half the arithmetic:

But now add:

of communication.

The layer takes approximately:

not 50 ms.

Speedup is:

not:

With still more GPUs, local compute per GPU shrinks while communication may not shrink proportionally.

Eventually:

can dominate:

That creates a practical limit to tensor-parallel degree.

Strong scaling eventually runs out¶

This is an example of strong scaling.

We keep the model and workload fixed while adding more processors.

Ideally:

In reality:

As \(N\) increases, the first term falls.

The other terms do not necessarily fall.

At some point adding GPUs yields little benefit — or can even make execution slower.

This is why "use more GPUs" is not a complete performance strategy.

Tensor parallelism and ZeRO are fundamentally different¶

Yesterday's distinction is now worth making precise.

With ZeRO-3, parameters are sharded at rest, but a layer's parameters may be gathered so a worker can perform the layer computation.

Tensor parallelism says:

Do not reconstruct the whole layer on one worker. Keep the layer partitioned and have multiple GPUs compute it together.

Conceptually:

ZeRO-3:

distributed weights
      ↓
gather needed layer
      ↓
local layer computation


Tensor parallelism:

distributed weights
      ↓
distributed layer computation
      ↓
collective communication

They solve different bottlenecks.

ZeRO primarily attacks redundant training-state memory.

Tensor parallelism lets one layer's storage and arithmetic span multiple devices.

Tensor parallelism also matters for inference¶

Suppose a 400-billion-parameter model simply does not fit on one GPU.

We can shard its weight matrices across several GPUs and execute each inference layer cooperatively.

So tensor parallelism is useful in both:

and:

But the performance tradeoffs differ.

During inference decode, each step may involve relatively little arithmetic per request.

That can make communication especially painful.

A tensor-parallel decode path may repeatedly perform:

small local computation
      ↓
collective
      ↓
small local computation
      ↓
collective

This is another reason batching matters: more active sequences can increase local work between synchronization points.

Tensor parallelism changes the memory picture too¶

If a weight matrix contains:

bytes and we shard it evenly across \(N\) GPUs, each GPU stores approximately:

So tensor parallelism reduces per-GPU parameter memory for the sharded layers.

But unlike ZeRO, that memory reduction is inseparable from distributed execution.

Every forward pass uses all members of the tensor-parallel group.

The GPUs are not independent replicas.

They collectively behave like one larger accelerator.

That gives us a useful mental distinction:

data-parallel GPUs:
multiple workers doing different examples

tensor-parallel GPUs:
multiple pieces of one logical worker

What happens when we combine them?¶

Suppose we have 32 GPUs.

We could form tensor-parallel groups of 4 GPUs:

TP group 0: GPUs  0–3
TP group 1: GPUs  4–7
TP group 2: GPUs  8–11
...
TP group 7: GPUs 28–31

Each group behaves like one logical model worker.

Then the 8 groups can process different data in parallel.

So:

and:

We can then add ZeRO-style sharding across the data-parallel dimension.

Now the topology matters.

Fast intra-node links may be used for tensor parallelism, while slower inter-node networking handles less latency-sensitive data-parallel communication.

This is why physical cluster topology becomes part of the training algorithm.

Pipeline parallelism gives us another axis¶

If the model is extremely deep, we might also split layers into pipeline stages.

For example:

pipeline stage 0:
layers  1–20
uses 4-way tensor parallelism

pipeline stage 1:
layers 21–40
uses 4-way tensor parallelism

pipeline stage 2:
layers 41–60
uses 4-way tensor parallelism

Then replicate the entire pipeline for data parallelism.

A large training job can therefore have coordinates like:

A GPU's role is determined by its position in that multidimensional grid.

This is sometimes called 3D parallelism.

Additional techniques can introduce still more axes, such as sequence or context parallelism.

Why communication topology becomes algorithmic¶

Suppose 8 GPUs inside one server communicate through extremely fast links, while communication between servers is slower.

Then an efficient arrangement might choose:

tensor parallelism
→ mostly inside a server

data parallelism
→ across servers

because tensor parallelism communicates frequently inside every layer.

We want its collectives on the fastest links.

So two clusters with the same number and type of GPUs can have different training performance if their network topology differs.

The hardware graph matters:

which GPU connects to which GPU
through what link
at what bandwidth
with what latency

The logical parallelism should map intelligently onto that graph.

The deeper connection to yesterday¶

ZeRO and tensor parallelism both begin from the same observation:

But they exploit that observation differently.

ZeRO says:

Partition state to remove redundant storage, then communicate when a worker needs missing state.

Tensor parallelism says:

Partition the operation itself, so several GPUs jointly execute one layer.

One is largely about where state resides.

The other is about where arithmetic happens.

At very large scale, those questions merge into a single systems problem.

The key mental model¶

Think of a giant matrix multiplication as moving a piano.

Data parallelism says:

Give several teams identical pianos and let each team move a different one.

ZeRO says:

Stop making every team permanently store an identical warehouse of equipment.

Tensor parallelism says:

This particular piano is too heavy for one team. Put several teams on the same piano.

Adding more people helps only if they can coordinate efficiently.

If everyone must stop every few seconds to negotiate who holds which corner, coordination overwhelms the benefit.

That gives us today's main principle:

And that is why tensor parallelism is powerful but not free.

The art is not simply splitting a tensor across GPUs.

It is choosing splits that let each GPU perform as much independent, useful computation as possible before the next synchronization.

At large scale, the Transformer architecture, tensor shapes, collective algorithms, GPU topology, and network bandwidth all become parts of the same computation.