Tensor Parallelism › Limitations and Practical Considerations

Tensor parallelism is powerful, but it comes with real constraints that practitioners must account for when designing training infrastructure.

The most immediate limitation is that every dimension you want to shard must be divisible by the tensor-parallel degree. This sounds like a minor implementation detail, but it shapes model architecture choices from the ground up. A model with a hidden dimension of 4096 and 32 attention heads can be tensor-parallelized with degrees 1, 2, 4, 8, 16, and 32 (all of which divide both 4096 and 32 evenly), but a model with a hidden dimension of 3840 and 30 heads is more constrained. If you want degree 8, you need all key dimensions to be divisible by 8.

Large research teams deploying Megatron-style training consistently choose hidden dimensions and head counts that are highly divisible: powers of 2 or multiples of 128 are standard. GPT-3's hidden dimension of 12288 is $96 \times 128$, divisible by 1, 2, 3, 4, 6, 8, 12, 16, 24, 32, and more. This kind of divisibility is not accidental; it is designed to accommodate tensor parallelism at multiple degrees without architectural changes.

FFN dimensions are similarly constrained. The standard FFN expansion factor of 4 gives a hidden dimension of $4d$, and the product $4d$ must be divisible by the tensor-parallel degree. In practice, since $4$ itself is divisible by 2 and 4, any hidden dimension divisible by 8 will work with degrees up to 8. At degree 16, you need $4d$ to be divisible by 16, so $d$ must be divisible by 4, which is true for all powers of 2.

The second major limitation is communication bandwidth dependency. The all-reduce operations that make tensor parallelism correct are not free; they consume network bandwidth proportional to the activation tensor size at every layer. On modern DGX nodes, NVLink provides 600 GB/s bidirectional bandwidth, which is typically sufficient for tensor-parallel degrees up to 8 within a node. Crossing node boundaries over InfiniBand introduces a 10-20x bandwidth reduction, making tensor parallelism inefficient at those scales.

This is why Megatron-style systems typically combine tensor parallelism (within a node) with pipeline parallelism (across nodes). Pipeline parallelism communicates only at pipeline stage boundaries, transmitting only the activations for one microbatch at the boundary layer rather than all-reducing full activations at every layer. This reduces cross-node communication volume by orders of magnitude compared to applying tensor parallelism across nodes. The combination of intra-node tensor parallelism and inter-node pipeline parallelism became the standard architecture for large-scale training, codified in the Megatron v2 and v3 papers.

Tensor parallelism increases the degree of synchronization required within a group. With $N$ GPUs in a tensor-parallel group, all $N$ devices must synchronize at every all-reduce. If one GPU stalls due to a hardware issue, memory pressure, or kernel scheduling delay, all other GPUs in the group wait. This tight coupling means that hardware failures or stragglers have immediate impact on training throughput.

In contrast, data parallelism is much more loosely coupled: each GPU processes its own microbatch independently, and gradient averaging happens once per step rather than twice per transformer layer. Pipeline parallelism falls in between: stages must synchronize at pipeline boundaries, but each stage can compute independently otherwise. The tight synchronization of tensor parallelism is one reason it is restricted to within a single node, where communication is fast enough that synchronization delays are measured in microseconds rather than milliseconds.

Implementing tensor parallelism correctly requires careful attention to weight initialization. When a weight matrix is split across $N$ GPUs and each shard is initialized independently from a normal distribution, the effective initialization variance seen by the full layer computation differs from what a single-GPU initialization would produce.

Consider column parallelism: each GPU initializes its shard $W_i \in \mathbb{R}^{d \times (h/N)}$ from $\mathcal{N}(0, \sigma^2)$. The full output $Y = XA$ has variance proportional to $d \sigma^2$ (by standard variance analysis of matrix-vector products). This is the same as if the full matrix had been initialized with variance $\sigma^2$ on a single GPU. So far, so good.

For row parallelism, the situation is different. Each GPU computes a partial sum $Z_i = X_i A_i$, and the all-reduce computes $Y = \sum_i Z_i$. The variance of $Y$ is $N$ times the variance of a single $Z_i$ (since the $Z_i$ are independent). To maintain the correct variance, Megatron-LM scales the initialization of row-parallel weight shards by $1/\sqrt{N}$. This ensures that the aggregated output has the same distribution as the single-GPU case. Without this adjustment, deeper layers would see activations with inflated variance, potentially destabilizing training.

Choosing the right tensor-parallel degree in practice involves balancing several competing factors. The key considerations are:

• Memory capacity: Use a degree large enough that the model fits on each GPU with room for optimizer states and activations.
• Divisibility: All key dimensions (hidden size, number of heads, FFN intermediate size) must be divisible by the degree.
• Communication budget: On NVLink hardware, degrees 2-8 are typically efficient. Higher degrees diminish returns because the per-layer communication overhead grows while the per-GPU compute shrinks.
• Node boundary alignment: Keep tensor parallelism within a single node to exploit NVLink. Use pipeline or data parallelism for cross-node scaling.

The interaction between these factors means that choosing a tensor-parallel degree is fundamentally a hardware-aware decision. A model that trains efficiently with degree 8 on an A100 DGX node might train better with degree 4 on older hardware with slower NVLink, because the same communication volume represents a larger fraction of compute time.

Despite these limitations, tensor parallelism has proven essential for training the largest language models. Models like GPT-3, Megatron-Turing NLG (530 billion parameters), and subsequent systems from Google, Meta, and Anthropic all rely on tensor parallelism as a core component of their training infrastructure. The technique scales naturally with the number of high-bandwidth links available and integrates cleanly with data and pipeline parallelism to enable distributed training at scales that were previously considered impractical.

We will explore how tensor parallelism integrates with pipeline parallelism and data parallelism in the next chapter, which covers 3D parallelism strategies for training at extreme scale.