Tensor Parallelism

Tensor parallelism operates by splitting individual weight matrices within a neural network layer (e.g., the linear projections in a transformer's attention or MLP block) across multiple devices. For a matrix multiplication Y = XW, the input X is broadcast, and the weight matrix W is partitioned along its column or row dimension. Each GPU computes a partial result, which must then be aggregated via a collective communication operation (like an all-reduce) to produce the full output. This fine-grained splitting distinguishes it from pipeline parallelism, which partitions whole layers.

This strategy introduces significant inter-GPU communication overhead during both forward and backward passes. Key operations include:

• All-reduce: Used to sum partial results from partitioned matrix multiplications.
• All-gather: Used to collect partitioned tensors from all devices to reconstruct a full tensor.
• Reduce-scatter: Used to distribute and reduce a tensor across devices. The frequency of this communication, often happening within each transformer block, makes network bandwidth and latency critical bottlenecks. Performance scales well only when computation time outweighs communication time, which is true for very large tensors.

Tensor parallelism is most effective for models where individual layers exceed the memory of a single GPU. It provides a more balanced memory footprint across devices compared to pipeline parallelism, which can leave some GPUs idle (pipeline bubbles). It is particularly suited for the feed-forward networks (FFNs) and attention projections in transformers, where the weight matrices can be many gigabytes in size. The technique allows the total model size to scale almost linearly with the number of GPUs used for the tensor-parallel group.

In practice, tensor parallelism is rarely used alone for training models with hundreds of billions of parameters. It is combined in a 3D parallelism strategy:

• Tensor Parallelism (intra-layer) across a small group of GPUs (e.g., 4-8) within a node with high-bandwidth NVLink.
• Pipeline Parallelism (inter-layer) across multiple nodes to handle depth.
• Data Parallelism across multiple replicas of the combined model for batch processing. This hybrid approach maximizes GPU utilization and memory efficiency, with tensor parallelism handling the widest layers that don't fit on one chip.

While crucial for training, tensor parallelism is a foundational technique for serving massive models that cannot fit on a single inference GPU. Serving frameworks like vLLM and TGI support tensor-parallel distributed inference. The communication pattern is simplified during inference (forward pass only), but latency is sensitive to the all-reduce steps between generating each token. Therefore, the tensor-parallel group is typically confined to a single node with ultra-fast interconnects (NVLink) to minimize this latency penalty for user-facing requests.

Tensor parallelism operates by sharding the large weight matrices of a neural network layer (e.g., the linear projections in a transformer's attention or MLP blocks) across multiple devices. For a matrix multiplication Y = XW, if the weight matrix W is split column-wise across two GPUs, each GPU computes a partial result. The full output Y is then obtained via an all-reduce communication operation across the GPUs. This splitting can be applied to the feed-forward network (FFN) and attention layers, with the specific sharding strategy (column-wise, row-wise, or a combination) determining the communication pattern.

In practice, tensor parallelism is almost always combined with pipeline parallelism to train models with hundreds of billions or trillions of parameters. The hybrid strategy works as follows:

• Tensor Parallelism splits individual layers horizontally across a small group of GPUs (often within a single server node due to high bandwidth requirements).
• Pipeline Parallelism splits the model's layers vertically into sequential stages across many such groups. This combination allows systems to scale to thousands of GPUs. The communication-intensive tensor parallel groups are kept local (via NVLink), while pipeline stages communicate over the network, optimizing for both intra-node speed and inter-node scalability.

Sequence parallelism is a complementary technique that splits the input sequence dimension (batch size × sequence length) across devices. It is used to handle very long sequences that exceed GPU memory when computing operations like attention and layer normalization. When combined with tensor parallelism, it allows for scaling along both the model dimension (weights) and the sequence dimension (activations). This is essential for training on long-context data, as it reduces the activation memory per device, which is a major bottleneck alongside parameter memory.

It is crucial to distinguish tensor parallelism from the more common data parallelism:

• Data Parallelism: Replicates the entire model on each GPU and splits the training batch. Gradients are averaged via all-reduce after the backward pass. Efficient for models that fit on one GPU.
• Tensor Parallelism: Splits the model itself across GPUs, with each device holding a portion of the parameters. Processes a single batch collaboratively, with communication during the forward and backward passes. Used when the model does not fit on one GPU. They are orthogonal and often used together: data parallelism across nodes, tensor + pipeline parallelism within a node.

A model parallelism strategy that partitions a model's sequential layers into different stages, each assigned to a separate GPU. Activations are passed forward like an assembly line, while gradients flow backward. This is complementary to tensor parallelism and is often combined with it in 3D parallelism to train models with trillions of parameters.

• Key Difference: Splits the model by layer (vertical), whereas tensor parallelism splits within a layer (horizontal).
• Communication Pattern: Only required at the boundaries between stages, leading to lower bandwidth requirements but potential pipeline bubbles (idle time).

The most common parallelism strategy, where the entire model is replicated across multiple GPUs, and each GPU processes a different subset (mini-batch) of the training data. Gradients are averaged across all devices after each backward pass.

• Contrast with Tensor Parallelism: Data parallelism replicates the model; tensor parallelism splits a single model instance.
• Typical Use: Used for training when the model fits on a single GPU. It scales efficiently with batch size but requires significant inter-GPU bandwidth for gradient synchronization (e.g., via NCCL).

A hybrid approach that combines data parallelism, tensor parallelism, and pipeline parallelism to achieve optimal scaling for the largest models (e.g., GPT-3, Megatron-Turing NLG). Each dimension addresses a different scaling constraint.

• Data Parallelism: Scales with the size of the dataset.
• Tensor Parallelism: Scales with the size of individual model layers.
• Pipeline Parallelism: Scales with the depth (number of layers) of the model.
• Framework: Implemented in libraries like Megatron-LM and DeepSpeed.

A compression technique that reduces the numerical precision of a model's weights and activations (e.g., from 32-bit floating-point to 8-bit integers). While distinct from parallelism, it is a critical companion technology.

• Synergy with Tensor Parallelism: Quantization reduces the memory footprint and communication volume of tensors sharded across GPUs, improving efficiency.
• Types: Includes Post-Training Quantization (PTQ) for quick deployment and Quantization-Aware Training (QAT) for higher accuracy.

A memory optimization technique that trades compute for memory by selectively saving only certain intermediate activations during the forward pass. Non-checkpointed activations are recomputed during the backward pass.

• Relationship to Parallelism: Enables training larger models or using larger batch sizes within a fixed GPU memory budget, which is a prerequisite for effective data or tensor parallelism.
• Impact: Can reduce activation memory by up to 80% at the cost of ~30% more compute time.

The low-level communication operations used to synchronize data across GPUs in parallel training. Tensor parallelism heavily relies on these.

• All-Reduce: Sums a tensor across all GPUs and returns the result to each (critical for data parallelism).
• All-Gather: Collects shards of a tensor from all GPUs and assembles the full tensor on each device.
• Reduce-Scatter: Sums a tensor across GPUs, then splits and distributes the result.
• Implementation: Optimized libraries like NVIDIA NCCL are essential for performance.