Tensor Parallelism

Unlike pipeline parallelism, which splits a model across layers (inter-layer), tensor parallelism splits operations within a single layer (intra-layer). This typically involves partitioning the weight matrices of linear layers (e.g., Feed-Forward Network or Attention projections) across GPUs. Each GPU holds a shard of the full weight matrix and computes on a shard of the input, requiring an all-reduce communication operation to combine results before the next layer.

Tensor parallelism imposes a synchronous communication pattern. For a standard matrix multiplication Y = XW split across devices, the operation requires:

• An all-gather of input shards or
• An all-reduce of output shards. This makes it highly sensitive to inter-GPU bandwidth. It is most effective on systems with NVLink or InfiniBand connections, as network latency can become the dominant bottleneck, especially for smaller batch sizes.

A primary benefit is the proportional reduction in GPU memory required per device. Splitting a layer's parameters across N GPUs reduces the memory footprint for that layer's weights and its associated KV cache by approximately a factor of N. This is crucial for deploying models whose single-layer parameters exceed the memory of one GPU, enabling the inference of models with hundreds of billions of parameters.

In large-scale training and inference, tensor parallelism is rarely used alone. It is commonly combined with:

• Pipeline Parallelism: To split different model stages across GPU groups.
• Data Parallelism: To process different data batches concurrently. Frameworks like Megatron-LM popularized the 3D parallelism approach, which integrates all three strategies to maximize hardware utilization for trillion-parameter models.

Common splitting strategies include:

• Column Parallelism: Splits the weight matrix along its column dimension. The input is broadcast, and outputs are concatenated.
• Row Parallelism: Splits the weight matrix along its row dimension. The input is split, and outputs are summed via all-reduce.
• Megatron's 1D Scheme: Applies column parallelism to one linear layer and row parallelism to the next, optimizing communication. More advanced 2D, 2.5D, and 3D tensor parallelism schemes exist for extreme-scale systems.

While essential for training massive models, tensor parallelism presents specific trade-offs for inference:

• Latency Overhead: The required all-reduce synchronization adds fixed latency per layer, which can degrade Time-To-First-Token (TTFT).
• Optimal Scale: The performance benefit often only materializes when model size makes alternative strategies (like single-GPU inference) impossible.
• Serving Complexity: Requires coordinated multi-GU execution, complicating serving system design compared to data-parallel replication.

Tensor parallelism operates by partitioning the weight matrices and activations of a neural network layer across multiple devices. For a linear layer Y = XW, the input tensor X and weight matrix W are split along specific dimensions.

• Column-wise Parallelism: The weight matrix W is split column-wise. Each GPU holds a slice of W and computes a partial result. An all-gather operation is required to combine outputs.
• Row-wise Parallelism: The weight matrix W is split row-wise. The input X is broadcast to all GPUs, each computes on its slice, and an all-reduce (summation) combines the results.

This fine-grained splitting occurs within a single transformer layer, making communication a critical bottleneck that demands ultra-high-bandwidth interconnects like NVLink.

Modern high-performance inference servers integrate tensor parallelism to serve models larger than a single GPU's memory.

• vLLM: Supports tensor parallelism natively, allowing it to partition the KV cache and model weights across GPUs. Its PagedAttention algorithm works in concert with tensor parallelism to manage memory efficiently.
• TensorRT-LLM: Uses the same model-parallel strategies as Megatron-LM. During its ahead-of-time (AOT) compilation process, it generates highly optimized kernels for the specific tensor-parallel configuration, fusing communication and computation.
• Triton Inference Server: Can deploy tensor-parallel models via its Model Ensemble feature, where different model partitions are loaded on different GPUs and Triton manages the inter-process communication.

These engines combine tensor parallelism with continuous batching and KV caching for maximum throughput.

The primary challenge of tensor parallelism is balancing the increased computational throughput with the overhead of inter-GPU communication.

• Communication Overhead: Each layer requires synchronization points (all-reduce, all-gather). The volume of data communicated is proportional to the activation size and batch size.
• Interconnect Dependency: Performance is gated by inter-GPU bandwidth. NVLink (600 GB/s+) is essential; PCIe (32 GB/s) often becomes a severe bottleneck.
• Optimal Scale: Efficiency diminishes as more GPUs are added to a single tensor-parallel group due to Amdahl's Law. It is typically combined with pipeline parallelism for models spanning hundreds of GPUs.

Rule of Thumb: Tensor parallelism is most effective for layers with very large weight matrices (e.g., FFN layers in dense models) and when GPUs are connected via NVLink/NVSwitch.

Tensor parallelism is one of two primary model parallelism techniques, often used alongside pipeline parallelism.

• Granularity: Tensor parallelism splits within a layer (intra-layer). Pipeline parallelism splits across layers (inter-layer).
• Communication Pattern: Tensor parallelism requires fine-grained, synchronous communication at each layer. Pipeline parallelism uses coarse-grained, asynchronous communication between stages (micro-batches).
• Idle Time (Bubbles): Pipeline parallelism suffers from pipeline bubbles where devices are idle. Pure tensor parallelism has no bubbles but has constant communication overhead.
• Typical Use: In practice, large-scale training and inference use a 3D parallelism approach: Data Parallelism across nodes, Pipeline Parallelism across layers, and Tensor Parallelism within a node of tightly coupled GPUs.

Deploying tensor-parallel models in production requires specific infrastructure and monitoring.

• Hardware Homogeneity: All GPUs in a tensor-parallel group must be identical in model and memory capacity.
• Failure Resilience: The failure of a single GPU in the group causes the entire model instance to fail, requiring robust health checks and restart mechanisms.
• Load Balancing: Incoming requests must be routed to a full set of GPUs hosting one model instance. Load balancers must be aware of the model-parallel group topology.
• Latency Profile: Adds a deterministic communication delay to each layer's processing time. This increases P99 latency (tail latency) and must be accounted for in Service Level Agreements.

It is a powerful tool for serving massive models, but introduces significant operational complexity compared to single-GPU deployment.

A model parallelism strategy that partitions a model's sequential layers across multiple devices. Data flows through these stages in a pipeline using micro-batches, keeping all devices active but introducing pipeline bubbles (idle time) at the start and end of each batch. It is highly effective for models with many sequential layers and is often combined with tensor parallelism for extremely large models.

A neural network architecture designed for sparse activation. It consists of many specialized sub-networks (experts), and a routing network dynamically selects a small subset (e.g., 2 out of 128) for each input token. This enables massive model capacity (trillions of parameters) with a manageable, fixed computational cost per token. Unlike tensor parallelism, which splits individual operations, MoE activates different parameters conditionally.

An inference optimization technique that dynamically groups incoming requests of varying sequence lengths into a single batch to maximize GPU utilization. Instead of waiting for a fixed batch size (static batching), it schedules new requests as previous ones finish, significantly improving throughput. This is a complementary technique to tensor parallelism, optimizing how already-parallelized model instances process real-time traffic.

A critical inference optimization that stores the computed Key and Value tensors for previously processed tokens during autoregressive generation. This eliminates the need to recompute these tensors for the entire prompt and prior context on each new token generation step, dramatically reducing latency and compute. Efficient management of the KV cache memory is a major concern in systems using tensor-parallel models.

A model 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). This decreases the memory footprint and increases computational speed, which directly reduces the communication overhead and cost of tensor parallelism. Techniques include Post-Training Quantization (PTQ) and Quantization-Aware Training (QAT).

vLLM is a high-throughput inference serving engine. Its key innovation, PagedAttention, treats the KV cache as non-contiguous memory pages. This allows for efficient memory sharing (e.g., for identical prompts) and dynamic allocation, virtually eliminating memory fragmentation. This is especially valuable for serving tensor-parallel models, as it maximizes usable context length and improves overall system efficiency.