Tensor Parallelism

Unlike pipeline parallelism which splits by layer, tensor parallelism partitions the weight matrices of a single layer across multiple GPUs. Each device holds a slice of the tensor and performs a partial computation. The results are then aggregated using collective communication primitives like all-reduce or all-gather to reconstruct the full output activation. This allows a single transformer attention head or large feed-forward network to span multiple accelerators, directly addressing the memory wall for massive models.

Tensor parallelism is extremely communication-intensive. During every forward and backward pass, devices must exchange intermediate activations and gradients. This requires high-bandwidth, low-latency interconnects such as NVLink or InfiniBand. The technique is typically confined to a single node with 8 GPUs because cross-node latency from standard networking (like Ethernet) can become a severe bottleneck, negating the benefits of the added compute.

In transformer models, tensor parallelism is often implemented as a pair of complementary strategies:

• Column Parallelism: The weight matrix is split along its columns. Each GPU computes a portion of the output, which is then concatenated.
• Row Parallelism: The weight matrix is split along its rows. Each GPU computes against a full input, and the partial outputs are summed via an all-reduce operation. This pairing ensures that the full layer computation is distributed without requiring a single device to hold the complete intermediate tensor.

Tensor parallelism is orthogonal to data parallelism and often combined with it in a 3D parallelism strategy. While data parallelism replicates the entire model across devices and shards the data batch, tensor parallelism shards the model itself. ZeRO (Zero Redundancy Optimizer) stages can be seen as an evolution that partitions optimizer states and gradients across data-parallel workers, reducing memory without the intense communication overhead of full tensor parallelism. Modern frameworks like Megatron-LM and DeepSpeed integrate all three techniques.

Effective tensor parallelism demands:

• High-bandwidth interconnects: NVLink (900 GB/s for H100) or InfiniBand NDR (400 GB/s) are essential.
• Intra-node confinement: Scaling beyond a single node (typically 8 GPUs) introduces a steep communication tax that often makes additional parallelism strategies like pipeline parallelism more efficient.
• Balanced partitioning: The model architecture must be carefully partitioned to ensure each GPU receives an equal computational load, avoiding straggler bottlenecks.

During inference, tensor parallelism is primarily used to serve models that exceed a single GPU's memory. The communication overhead is less critical than during training because there is no backward pass. However, it still adds latency to every forward pass. For training, the technique is essential for fitting massive models but introduces a significant synchronization barrier. The all-reduce operations required for gradient synchronization can become the dominant factor in iteration time if the interconnect is insufficient.

A complementary strategy that partitions a model by layer depth rather than by tensor width. While Tensor Parallelism splits individual weight matrices, Pipeline Parallelism assigns sequential groups of layers to different devices, passing activations between them. This introduces pipeline bubbles—idle time while devices wait for data—which are mitigated by splitting mini-batches into micro-batches for staggered execution. The GPipe and PipeDream schedulers are canonical implementations. Pipeline Parallelism is often combined with Tensor Parallelism in 3D parallelism regimes for models exceeding hundreds of billions of parameters.

The simplest distributed training paradigm where the entire model is replicated across multiple devices, and each device processes a different subset of the input batch. Gradients are synchronized across replicas using an all-reduce collective operation before each optimizer step. Unlike Tensor Parallelism, Data Parallelism does not reduce the memory required to house a single model instance—it only increases throughput. The ZeRO optimizer (Zero Redundancy Optimizer) extends Data Parallelism by partitioning optimizer states, gradients, and parameters across devices to eliminate memory redundancy.

A technique that distributes the sequence dimension of activations and their corresponding computation across multiple devices. This is critical for long-context training where the activation memory of a single sequence exceeds device capacity. Sequence Parallelism splits the input along the token dimension before LayerNorm and attention blocks, then gathers results. It is often paired with Tensor Parallelism in the Megatron-LM framework, where Tensor Parallelism handles the attention and MLP blocks while Sequence Parallelism handles the normalization and dropout regions, significantly reducing activation memory footprint.

A memory optimization technology from Microsoft's DeepSpeed library that eliminates data redundancy in Data Parallelism. ZeRO operates in three stages:

• Stage 1: Partitions optimizer states across devices
• Stage 2: Adds gradient partitioning
• Stage 3: Partitions model parameters themselves

Unlike Tensor Parallelism, which requires high-bandwidth interconnects for frequent communication, ZeRO-3 can operate effectively over standard interconnects by gathering parameters on-demand during forward and backward passes. The combination of ZeRO-3 with Tensor Parallelism enables training of trillion-parameter models.

The low-level operations that make Tensor Parallelism possible. Key primitives include:

• All-Reduce: Aggregates and distributes tensors across all devices (used for gradient sync)
• All-Gather: Collects sharded tensors so every device has a full copy
• Reduce-Scatter: Reduces tensors and scatters results (inverse of All-Gather)
• Point-to-Point Send/Recv: Direct device-to-device transfers for pipeline stages

Tensor Parallelism relies heavily on All-Reduce and All-Gather operations after partial matrix multiplications. The efficiency of these collectives, governed by NCCL (NVIDIA Collective Communications Library), directly determines scaling efficiency.

A specialized parallelism strategy for Mixture of Experts (MoE) architectures. In MoE models, each input token is routed to only a subset of expert sub-networks. Expert Parallelism places different experts on different devices, so tokens are dynamically dispatched and results gathered based on routing decisions. This creates an all-to-all communication pattern distinct from Tensor Parallelism's all-reduce. Expert Parallelism is often combined with Tensor Parallelism when individual experts are themselves too large for a single device, creating a hybrid strategy for models like Mixtral 8x7B and GPT-4.