Tensor parallelism — splitting one matrix multiplication across GPUs

Yesterday we saw fully sharded data parallelism (FSDP).

FSDP solves a storage problem:

A GPU does not need to permanently hold the entire model.

But during computation, FSDP typically gathers the parameters needed for a layer so that each worker can execute that layer on its local data.

That leaves another problem.

What if one layer itself is too large, or we want more compute than one GPU can provide for a single token stream?

Then we can split the layer's actual mathematics across GPUs.

That is tensor parallelism.

Its central idea is:

Instead of GPU 0 computing an entire matrix multiplication, several GPUs compute pieces of it simultaneously and communicate only when the pieces must be combined.

Start with the operation hiding inside almost every Transformer layer¶

Much of a Transformer is matrix multiplication.

For example:

Suppose:

and:

Then:

In an LLM, \(B\) here can represent many token positions flattened together.

If:

and:

then \(W\) contains:

parameters.

At BF16, that one matrix is roughly:

A single Transformer block contains several large matrices like this.

Tensor parallelism asks:

Why should one GPU own and multiply the whole matrix?

Split the matrix by columns¶

Suppose we have two GPUs.

Write:

Each half contains different output columns.

Then:

So GPU 0 can store \(W_1\) and compute:

GPU 1 stores \(W_2\) and computes:

Then:

Visually:

                    W

          ┌─────────┬─────────┐
          │   W1    │   W2    │
          └─────────┴─────────┘
               ↓         ↓
             GPU 0     GPU 1

X ───────────→ XW1      XW2

               ↓         ↓

              Y1         Y2

               └────┬────┘
                    ↓

              [ Y1 | Y2 ]

Each GPU stores only half of \(W\) and performs roughly half of the matrix multiplication.

This is column-parallel linear algebra.

A tiny numerical example¶

Let:

and:

The full multiplication is:

Now split \(W\) across two GPUs:

GPU 0 computes:

GPU 1 computes:

Concatenate them:

Exactly the same answer.

The mathematics has not changed.

Only the physical execution has.

But the next layer creates a problem¶

Suppose the next operation is:

Our \(Y\) is split:

GPU 0: Y1
GPU 1: Y2

One possibility is to gather the full \(Y\) onto both GPUs.

But if we gather after every matrix multiplication, communication can overwhelm the benefit.

A better trick is to choose the partitioning of the next matrix so that it naturally consumes the split output.

This is where tensor parallelism becomes more interesting than simply "cut every matrix in half."

Split the next matrix by rows¶

Since:

partition \(V\) by rows:

Then:

Multiplying gives:

Now GPU 0 computes:

while GPU 1 computes:

Then the GPUs only need to sum:

That sum can be implemented with a collective such as an all-reduce.

The pattern is:

column-parallel layer

GPU 0: XW1 → Y1
GPU 1: XW2 → Y2

          ↓

row-parallel layer

GPU 0: Y1V1 → Z1
GPU 1: Y2V2 → Z2

          ↓

      all-reduce

          ↓

       Z1 + Z2

Notice what we avoided:

gather full Y

between the two large matrix multiplications.

This maps beautifully onto the Transformer MLP¶

Recall the simplified feed-forward network:

Ignoring notation orientation, the first projection expands the hidden dimension:

The second projection contracts it:

Suppose:

and:

With four tensor-parallel GPUs, the expanded hidden dimension can be split:

GPU 0 → FFN channels 0–8191
GPU 1 → FFN channels 8192–16383
GPU 2 → FFN channels 16384–24575
GPU 3 → FFN channels 24576–32767

Each GPU performs the activation function on its own channels.

Then the second projection consumes those local channels and produces a partial contribution to the final hidden state.

Those partial contributions are summed.

Conceptually:

             x
             │
     ┌───────┼───────┬────────┐
     ↓       ↓       ↓        ↓
   GPU0    GPU1    GPU2     GPU3
   W1_0    W1_1    W1_2     W1_3
     ↓       ↓       ↓        ↓
    h0      h1      h2       h3
     ↓       ↓       ↓        ↓
   GELU    GELU    GELU     GELU
     ↓       ↓       ↓        ↓
   W2_0    W2_1    W2_2     W2_3
     ↓       ↓       ↓        ↓
    z0      z1      z2       z3
     └───────┼───────┴────────┘
             ↓
          all-reduce
             ↓
       z0+z1+z2+z3

The large MLP computation is physically spread across four accelerators.

Attention can be split too¶

Multi-head attention gives us another convenient dimension.

Suppose a Transformer has:

attention heads.

With eight tensor-parallel GPUs, we can 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

Each GPU can compute Q, K, V and attention for its local heads.

The outputs eventually need to participate in the output projection, where another carefully chosen partition and collective reconstruct the correct result.

This is one reason the multi-head structure of Transformers is friendly to parallel hardware.

The restaurant-kitchen mental model¶

Imagine a restaurant receives one enormous order.

Data parallelism is like opening four identical kitchens and sending a different customer's order to each kitchen.

Kitchen 0 → customer A
Kitchen 1 → customer B
Kitchen 2 → customer C
Kitchen 3 → customer D

Tensor parallelism is different.

One customer's order is too large for one kitchen.

So:

Kitchen 0 → appetizers
Kitchen 1 → pasta
Kitchen 2 → vegetables
Kitchen 3 → desserts

They are all working on the same order.

At certain points, their outputs must be combined before service can continue.

That synchronization is the price of splitting one computation across kitchens.

Tensor parallelism reduces parameter memory per GPU¶

If a weight matrix contains:

parameters and we split it evenly across \(N\) tensor-parallel GPUs, each stores roughly:

For our 537 MB matrix across four GPUs:

More importantly, the matrix multiplication itself is also divided.

Each GPU performs roughly one quarter of the relevant FLOPs.

So tensor parallelism can provide both:

and:

But again, there is a price.

Communication happens inside every Transformer block¶

With ordinary data parallelism, GPUs can independently execute almost an entire forward/backward pass before synchronizing gradients.

Tensor-parallel GPUs cannot.

They cooperate on the same layer.

That means communication may occur repeatedly:

attention projection
↓
collective
↓
attention output
↓
collective
↓
MLP projection
↓
collective
↓
next layer

The exact pattern depends on the architecture and implementation, but the principle is:

That makes network latency and bandwidth extremely important.

Why tensor-parallel groups are often kept local¶

Suppose eight GPUs inside one server have extremely fast links.

Communication between servers is slower.

Because tensor parallelism communicates so frequently, it is often attractive to keep a tensor-parallel group within the fastest local fabric:

server 0:

GPU 0
GPU 1
GPU 2
GPU 3
GPU 4
GPU 5
GPU 6
GPU 7

← high-bandwidth tensor-parallel communication →

Then slower inter-server links can be used for parallelism patterns that communicate less frequently.

This is not a universal rule, but it illustrates a major design principle:

The physical topology of the cluster influences the logical topology of the model.

Tensor parallelism is not free scaling¶

Suppose one GPU takes:

to execute a large layer.

With two GPUs, the arithmetic might ideally fall toward:

But suppose combining partial results costs:

Then total time becomes roughly:

The speedup is:

not \(2\times\).

With more GPUs, each GPU gets less arithmetic while communication remains.

Eventually:

can dominate:

So there is usually a practical tensor-parallel degree rather than "use as many GPUs as possible."

Small matrix pieces can also hurt GPU efficiency¶

GPUs love large matrix multiplications.

They can keep many compute units busy and amortize overhead.

If we split a matrix across too many GPUs, each local multiplication becomes smaller.

At some point:

huge GEMM
→ excellent GPU utilization

turns into:

many smaller GEMMs
→ worse utilization

So increasing tensor parallelism creates two possible penalties:

and:

This is another reason parallelism degree must be tuned.

Tensor parallelism during inference¶

Tensor parallelism is not just a training technique.

It is extremely useful for inference.

Suppose a model's weights require:

With 80 GB GPUs, the model cannot fit on one accelerator.

With four-way tensor parallelism:

before accounting for KV cache and overhead.

The model's layers can execute across the four GPUs as one logical inference engine.

During autoregressive decoding, however, each generated token repeatedly traverses those tensor-parallel collectives.

So inference latency becomes sensitive to interconnect performance too.

This is why putting tensor-parallel inference across distant machines is usually much less attractive than keeping the group tightly connected.

Compare FSDP and tensor parallelism carefully¶

Yesterday's FSDP and today's tensor parallelism both shard parameters.

But they answer different questions.

FSDP asks:

Where should the parameters live when they are not being used?

Tensor parallelism asks:

How should the layer's actual computation be divided while it is running?

A useful simplified comparison is:

FSDP

persistent:
GPU 0 → W0
GPU 1 → W1

compute:
gather W
GPU 0 → full layer on data A
GPU 1 → full layer on data B

versus:

TENSOR PARALLELISM

persistent:
GPU 0 → W0
GPU 1 → W1

compute:
GPU 0 → part of layer
GPU 1 → part of layer
        ↓
combine partial results

So:

Combine tensor and data parallelism¶

Suppose we have:

GPUs.

We could organize them into eight groups of four:

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

Within each group:

The four GPUs cooperate to execute one model replica.

Across the eight groups:

Each group processes different training data.

Conceptually:

                    different data
                         ↓

DP replica 0 → [GPU GPU GPU GPU]
DP replica 1 → [GPU GPU GPU GPU]
DP replica 2 → [GPU GPU GPU GPU]
...
DP replica 7 → [GPU GPU GPU GPU]

                 ↔ tensor parallel ↔

So total GPU count is:

Here:

This is the beginning of multi-dimensional parallelism.

Add yesterday's FSDP idea¶

Now imagine the data-parallel replicas also shard optimizer state and perhaps parameters using FSDP.

We now have several overlapping decompositions:

tensor parallelism
→ split one layer's computation

data parallelism
→ split training examples

FSDP
→ split persistent training state

gradient accumulation
→ split global batch across time

activation checkpointing
→ split activation lifetime across time

Each solves a different bottleneck.

Modern large-model training stacks combine them because no single dimension scales indefinitely.

The deeper mental model: choose where the equals sign happens¶

Return to:

GPU 0 can independently compute:

GPU 1 can independently compute:

The GPUs only need to communicate at the point where mathematics says:

must happen.

This gives a powerful way to think about distributed tensor computation:

Keep operations local for as long as algebra permits, and communicate only when the mathematical pieces must meet.

Good tensor-parallel layouts are therefore partly exercises in algebra.

We partition matrices so that consecutive operations can consume local pieces without constantly reconstructing full tensors.

The best partition is not merely the one that divides memory evenly.

It is the one that minimizes expensive synchronization while preserving large, efficient local matrix multiplications.

The progression so far¶

We can now see a hierarchy of increasingly aggressive decomposition.

Ordinary data parallelism:

whole model
different data

FSDP:

sharded model storage
temporarily reconstructed layers
different data

Tensor parallelism:

sharded model storage
sharded layer computation
same token stream crosses multiple GPUs

The conceptual shift is important.

With data parallelism, a GPU can mostly pretend the others do not exist until synchronization.

With tensor parallelism, the GPUs are parts of a single virtual accelerator.

A token moving through the Transformer is effectively moving through a machine assembled from several physical GPUs.

That gives us today's main principle:

But that illusion only works if the GPUs can exchange partial results fast enough.

So once again, the limiting resource moves.

We started with a model too large for one GPU.

We divided the computation.

Now the challenge becomes moving the boundaries between those pieces efficiently.

That is the recurring story of large-scale LLM systems:

Every time we split something to solve a memory or compute problem, we create a communication problem at the cut.