Core ML

Tensor Parallelism

When a single mathematical matrix is too large to fit in one GPU's memory, you have to split the math itself. Tensor Parallelism cuts a matrix in half, gives half to GPU 1 and half to GPU 2, and merges their answers at the end of the calculation.

Why Does This Exist?

Data Parallelism works beautifully, but it has a hard limit: The entire model must fit on a single GPU. A 70-Billion parameter model (like Llama 3) requires about 140GB of VRAM just to hold the weights in 16-bit precision. The largest standard GPU (an Nvidia H100) only has 80GB of VRAM. You cannot use Data Parallelism on an H100 for a 70B model, because you cannot even load the model.

We have to split the model itself. The most aggressive way to split a model is Tensor Parallelism (TP). Instead of assigning Layer 1 to GPU 1 and Layer 2 to GPU 2 (that's Pipeline Parallelism), Tensor Parallelism literally slices a single layer (a single Tensor) in half.

Think of It Like This

Imagine you need to multiply two massive, 1,000-digit numbers together by hand. The whiteboard (your GPU memory) is too small to hold all the intermediate math.

You hire a friend. You split the problem vertically. You multiply the first number by the first 500 digits of the second number. Your friend multiplies the first number by the last 500 digits. When you are both done, you shout your answers to each other and add them together. You just performed Tensor Parallelism.

How It Actually Works (The Megatron-LM Approach)

Tensor Parallelism was popularized by Nvidia's Megatron-LM paper. It specifically targets the two heaviest parts of a Transformer architecture: the Self-Attention block and the Feed-Forward Network (FFN).

Let's look at how it splits the FFN (which consists of two massive matrix multiplications: $Y = \text{GeLU}(X A) B$ ).

1. Column-Parallel (The First Matrix)

The first matrix $A$ is sliced vertically (Column-wise).

GPU 1 gets $A_{\text{left}}$
GPU 2 gets $A_{\text{right}}$ When the input $X$ arrives, both GPUs receive a full copy of $X$ . GPU 1 calculates $Y_{\text{left}} = X A_{\text{left}}$ . GPU 2 calculates $Y_{\text{right}} = X A_{\text{right}}$ . At this point, neither GPU has the full answer, but they do not need to communicate yet.

2. Row-Parallel (The Second Matrix)

The second matrix $B$ is sliced horizontally (Row-wise).

GPU 1 gets $B_{\text{top}}$
GPU 2 gets $B_{\text{bottom}}$ GPU 1 immediately multiplies its partial answer from step 1 by its half of matrix B: $Z_1 = Y_{\text{left}} B_{\text{top}}$ . GPU 2 does the same: $Z_2 = Y_{\text{right}} B_{\text{bottom}}$ .

3. The All-Reduce (The Sync)

Now, both GPUs have a partial sum. The true final answer is simply $Z_1 + Z_2$ . The GPUs trigger an All-Reduce operation over the network to add their matrices together. Now, both GPUs have the full, exact final output $Z$ , and they are ready for the next layer.

The Hardware Constraint: Why NVLink is Mandatory

In Data Parallelism, GPUs only communicate once per training step (during the backward pass). In Tensor Parallelism, the GPUs have to communicate constantly—often multiple times per layer, during both the forward and backward pass.

If you try to run Tensor Parallelism across standard Ethernet cables, the network latency will destroy your training speed. The GPUs will spend 99% of their time waiting for the All-Reduce to finish. Therefore, Tensor Parallelism is almost exclusively restricted to GPUs within the same physical server box, connected by ultra-fast copper interconnects like Nvidia NVLink (which can transfer 900 GB/second). If you have an 8-GPU server, your maximum Tensor Parallelism degree is 8.

Show Me the Code

You rarely implement the matrix slicing manually. DeepSpeed or Megatron-LM handles the math. But conceptually, here is what is happening inside the DeepSpeed engine:

1import torch2
3# The full, massive weight matrix (too big for one GPU)4# W = torch.randn(10000, 10000)5
6# We shard it. GPU 0 gets the left half, GPU 1 gets the right half.7if local_rank == 0:8    W_shard = torch.randn(10000, 5000) # GPU 0 memory9elif local_rank == 1:10    W_shard = torch.randn(10000, 5000) # GPU 1 memory11
12# The input X arrives (copied to both GPUs)13X = torch.randn(1, 10000)14
15# 1. Independent Math (No communication overhead!)16# GPU 0 calculates the left 5000 outputs. GPU 1 calculates the right 5000 outputs.17Y_partial = torch.matmul(X, W_shard) 18
19# 2. Network Synchronization (All-Gather)20# We need to stitch the two halves back together.21Y_full = torch.zeros(1, 10000)22# This command forces the GPUs to send their partial tensors to each other23torch.distributed.all_gather(tensor_list=[Y_partial_from_0, Y_partial_from_1], 24                             tensor=Y_partial)25
26# Now both GPUs have Y_full and can proceed to the next layer!

Watch Out For

The "Too Much TP" Trap. Because Tensor Parallelism requires such intense communication, you should only use it when you absolutely have to. If a model fits into 2 GPUs, you should use a TP degree of 2, not 8. Using a TP degree of 8 when it isn't strictly necessary will slow down your training due to the excessive All-Reduce operations. The golden rule of 3D Parallelism: Keep TP as small as possible (usually $\le 8$ ), use Pipeline Parallelism for the rest of the model splitting, and use Data Parallelism for the rest of the cluster.

The Quick Version

Tensor Parallelism (TP) splits individual mathematical operations (like Matrix Multiplication) across multiple GPUs.
It is required when a single layer of a model is too large to fit in one GPU's memory.
The standard approach splits the first matrix vertically (Column-Parallel) and the second matrix horizontally (Row-Parallel) to minimize network synchronization.
Because it requires massive communication bandwidth, TP is usually restricted to GPUs physically wired together inside the same server node (using NVLink).