Skip to the page
Tecyfy

What is tensor parallelism?

Splitting each layer's weight matrices across several GPUs, so every GPU does part of every layer and they combine their results after each one.

advanced · Also called TP, Intra-Layer Model Parallelism

A 70B model at 16 bits needs about 140 GB for its weights, more than one 80 GB GPU holds. Something has to split it. Tensor parallelism splits it the finest way: every layer is cut into slices, and each GPU holds one slice of every layer.

How a layer is split

Most of a transformer layer is matrix multiplication: the input times a large grid of weights. A grid can be cut into columns. Give GPU 1 the left half of the columns and GPU 2 the right half, and each can multiply the same input by its half at the same time. Each ends up with half of the output.

The Megatron-LM paper from NVIDIA, in 2019, showed how to pair the cuts so that a whole attention block or feed-forward block needs only one exchange between GPUs at the end. Attention splits naturally too: different attention headsMulti‑Head AttentionParallel attention layers allowing a model to focus on different subspaces go to different GPUs. That paper trained an 8.3-billion-parameter model across 512 GPUs at 76% of perfect scaling.

The price: talking after every layer

Each split layer finishes with the GPUs combining their partial results, an "all-reduce", before the next layer can start. That happens twice per layer, on every token. So tensor parallelism only works well over very fast links, such as NVLink between GPUs in one server. Across servers on an ordinary network, the waiting eats the gain.

That is why the usual pattern is tensor parallelism inside a server (2, 4 or 8 GPUs) and pipeline parallelismPipeline ParallelismPutting different layers of a model on different GPUs and passing work from one to the next like an assembly line, with small batches keeping every stage busy. or data parallelismData ParallelismTraining on several GPUs by giving each a full copy of the model and a different slice of each batch, then averaging their gradients so every copy stays the same. between servers.

Why it helps when serving

When you run a model, tensor parallelism does more than make it fit. Every GPU reads its own slice of the weights from its own memory at the same time, so the work per token is split too, and responses come faster. Each GPU also holds its share of the KV cacheKV CacheWork a model saves from the text it has already read, so it is not redone for every new token. Prompt caching keeps that work between requests and charges less for it., which leaves room for more conversations at once.

In vLLMvLLMAn open-source engine for serving language models fast, built on PagedAttention, that answers requests through the same kinds of API the big model makers offer. it is one setting: --tensor-parallel-size 4 spreads a model across 4 GPUs. The number has to divide the model's attention heads evenly, which is why it is almost always 2, 4 or 8.

When to use it

Use it when a model does not fit on one GPU, or when you need each answer faster, and your GPUs are in one server with fast links between them. If the model fits on one GPU and you need more total traffic rather than faster answers, run several copies of it instead.

Sources

Where each fact comes from

  1. 1Megatron-LM, training multi-billion parameter language models using model parallelism
  2. 2Parallelism and scaling