Tensor parallelism
Compare kernel, parallelism and disaggregated-serving choices using memory, communication and latency constraints.
THE IDEA
Tensor parallelism
When a model does not fit on one device, tensor parallelism splits individual weight matrices across devices. Split a matrix by columns and each device produces its own slice of the output, which is concatenated. Split by rows and each device produces a partial sum of the full output, which must be added with an all-reduce. Either way devices must exchange data: splitting the weights saves memory per device but adds communication to every layer.
A transformer MLP splits the first matrix by columns and the second by rows, so a single all-reduce is needed at the end. Splitting the weights saves memory per device but never removes the communication.
Predict first. Is splitting weights enough to avoid all cross-device exchange?
Chapter context
Scaling serving changes how work is placed and moved. An operation can be mathematically unchanged while its memory traffic changes dramatically. Device count alone does not describe the performance of the resulting system.
TRY IT YOURSELF
Part 1. Do it hands-on
Work out memory and communication by hand, then write column- and row-parallel products and a tensor-parallel MLP that matches the unsplit one.
A transformer MLP splits the first matrix by columns and the second by rows, so a single all-reduce is needed at the end. Splitting the weights saves memory per device but never removes the communication.
- To do: A 14 GB model is split evenly 4 ways across devices (tensor parallelism). Weights per device in GB?
- To do: If the weight matrix is split by ROWS (each device holds some input rows), every device produces a partial result of the full output size that must be combined by adding. Type the operation: sum or concatenate
- To do: A ring all-reduce of N bytes over p devices sends about 2(p−1)/p · N bytes per device. For N = 16,384 bytes (8192 values of 2 bytes) and p = 4, how many bytes per device?
Part 2. Do it in code
- column_parallel(X, W, p): split the COLUMNS of W into p equal parts, compute X @ W_part on each 'device' and concatenate the results along the last axis. ValueError if the columns do not divide evenly.
- row_parallel(X, W, p): split the ROWS of W (and the matching columns of X) into p parts, compute each partial product and ADD them (the all-reduce). ValueError if the rows do not divide evenly.
- shard_bytes(total_bytes, p): bytes of weights per device.
- ring_allreduce_bytes(n_bytes, p): 2 · (p − 1) / p · n_bytes sent per device.
- mlp_tensor_parallel(X, W1, W2, p): the two-matrix MLP relu(X @ W1) @ W2 with W1 split by columns and W2 split by rows across p devices (one all-reduce at the end). It must equal the unsplit result.
Stuck, or something not working? You can come back to it any time.
CHECK YOUR UNDERSTANDING
Is splitting weights enough to avoid all cross-device exchange?
Chapter reference: math, worked example & code
This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.
Math and notation
A weight-traffic lower-bound estimate is
This assumes the stated bytes really traverse the bottleneck bandwidth. KV transfer adds roughly bytes_to_move/network_bandwidth before protocol and synchronization overhead. Batching and caching can change the bytes per produced token.
Worked example
A fictional 14-GB dense weight scan across a sustained 700-GB/s memory path takes at least 14/700=0.02 seconds, giving a 50-step/s bandwidth ceiling under the one-scan-per-step assumption. Real execution includes other memory traffic and work. A 512-MiB KV transfer over 32 GiB/s has a payload-only lower bound of 1/64 second.
Reference code
The following example uses the Python standard library and does not call a remote service.
from math import isclose
assert isclose(14/700,.02)
assert isclose(1/(14/700),50)
transfer=(512*1024**2)/(32*1024**3)
assert isclose(transfer,1/64)
Python examples are provided for study and local execution; this page does not execute Python.
Common mistakes
Peak device bandwidth is not guaranteed sustained application bandwidth. Exact attention can be computed with different memory layouts. Additional GPUs can add communication bottlenecks.