Tensor parallelism is a distributed computing technique used in deep learning to split individual tensors (multi-dimensional arrays) across multiple accelerators (e.g., GPUs or TPUs), enabling the training and inference of models that exceed the memory capacity of a single device. Unlike data parallelism, where each device holds a full copy of the model, tensor parallelism partitions the model's weight matrices and operations along specified dimensions (e.g., rows or columns), allowing large layers to be computed in parallel while minimizing communication overhead. It is a key component of "3D parallelism" (combining data, pipeline, and tensor parallelism) often employed in large language models and other scale-intensive neural architectures.

1 Background

1.1 The memory wall problem in large models

The increasing size of deep neural networks, particularly transformer‑based language models, has led to models with hundreds of billions or even trillions of parameters. A single high‑end GPU typically provides 40–80 GB of memory, which is insufficient to store the parameters, gradients, and optimizer states of such models. This “memory wall” forces practitioners to distribute model data across multiple devices. Simple techniques such as swapping to CPU memory introduce unacceptable latency, motivating the need for parallelization strategies that keep data on‑device.

1.2 Evolution from data parallelism to model parallelism

Early distributed training relied on data parallelism, where each device holds a complete copy of the model and processes a different mini‑batch of data. While effective for small models, data parallelism fails when the model itself does not fit on a single device. Model parallelism addresses this by splitting the model across devices. Initially, layer‑wise (or pipeline) parallelism was used, assigning entire consecutive layers to different devices. However, for very large layers (e.g., attention and feed‑forward networks with huge hidden dimensions), even a single layer may exceed device memory. Tensor parallelism emerged as a finer‑grained variant that partitions the tensors within a single layer.

1.3 Role of tensor parallelism in 3D parallelism

Modern large‑scale training typically combines three orthogonal parallelism strategies: data parallelism (replicating the model across devices and splitting data), pipeline parallelism (splitting layers across devices), and tensor parallelism (splitting individual layers across devices). This combination, known as 3D parallelism, allows models of arbitrary size to be trained by exploiting both model replication and intra‑layer decomposition. Tensor parallelism reduces the memory footprint per device for the largest layers, while pipeline and data parallelism handle the overall throughput and scaling.

2 Mechanism

2.1 Partitioning strategies

Tensor parallelism divides the weight matrices of a neural network layer along one or more dimensions. The choice of splitting dimension affects communication patterns and computational load balance. The two most common strategies are row‑wise and column‑wise splitting.

2.1.1 Row‑wise splitting

In row‑wise splitting, a weight matrix \(W\) of shape \(M \times N\) is divided along its rows into \(p\) chunks, each of size \(\lceil M/p \rceil \times N\). Each device stores a subset of rows. During forward propagation, the input is broadcast to all devices; each device computes a partial output, and an all‑reduce operation sums the contributions. This approach is typically used for the second linear layer in a transformer’s feed‑forward block.

2.1.2 Column‑wise splitting

Column‑wise splitting partitions \(W\) along its columns into \(p\) chunks of size \(M \times \lceil N/p \rceil\). Each device holds a subset of columns and computes a portion of the output. The outputs from all devices are concatenated via an all‑gather operation to produce the full output. Column‑wise splitting is often applied to the first linear layer of a feed‑forward network or to the query, key, and value projections in attention.

2.2 Collective communication primitives

Tensor parallelism relies on efficient collective communication operations to exchange partial results between devices. The choice of operation depends on the partitioning strategy and the computation step.

2.2.1 All‑reduce and reduce‑scatter

All‑reduce sums tensors across all devices and broadcasts the result back to each device. It is used after row‑wise splitting to combine partial outputs. Reduce‑scatter is a variant that reduces (sums) data and then scatters the result, which can be more efficient when followed by an all‑gather. Some implementations use reduce‑scatter instead of all‑reduce to reduce memory pressure.

2.2.2 All‑gather

All‑gather collects data from all devices and concatenates them, distributing the full concatenated tensor to every device. It is employed after column‑wise splitting to reconstruct the complete output. All‑gather is often combined with reduce‑scatter in a two‑phase communication pattern (e.g., in the Megatron‑LM style).

2.3 Linear layer decomposition

Tensor parallelism decomposes the computation of a linear layer into parallel subtasks. Two representative styles are widely adopted.

2.3.1 Megatron‑LM style

Introduced by the NVIDIA Megatron‑LM framework, this approach splits the weight matrices of both the first and second linear layers in a transformer’s feed‑forward block. The first layer uses column‑wise splitting, the second uses row‑wise splitting. For a feed‑forward network with two linear layers, the forward pass proceeds as: all‑gather the input (or use a fused kernel) after the first split, then all‑reduce the output of the second split. This arrangement reduces the total communication volume compared to naive splitting.

2.3.2 Block‑wise tensor parallelism (e.g., in GSPMD)

The GSPMD (Generalized SPMD) system, used in JAX/XLA, generalizes tensor parallelism by allowing arbitrary block‑wise partitioning of tensors. Instead of splitting only along rows or columns, GSPMD can partition tensors along multiple axes (e.g., 2D block partitioning). This flexibility supports more complex models, such as those with spatial dimensions or mixture‑of‑experts, and allows the compiler to automatically choose a partitioning strategy that minimizes communication.

3 Implementation

3.1 Software frameworks

Several frameworks provide built‑in support for tensor parallelism, each with different design philosophies.

3.1.1 NVIDIA Megatron‑LM

Megatron‑LM is a framework specifically designed for training large transformer models using tensor and pipeline parallelism. It implements the row/column splitting described in §2.3.1 with efficient fused kernels (e.g., fused bias‑gelu). Megatron‑LM also supports model parallelism across multiple nodes via NCCL communicators, and it is widely used for models such as GPT‑3 and Megatron‑Turing NLG.

3.1.2 DeepSpeed

DeepSpeed, developed by Microsoft, offers tensor parallelism as part of its “ZeRO‑Infinity” and “DeepSpeed‑Speed” libraries. It integrates with PyTorch and provides automatic tensor partitioning for transformer layers. DeepSpeed’s implementation supports both 1D and 2D tensor parallelism, and it can be combined with ZeRO‑stage optimizations to further reduce memory usage.

3.1.3 JAX / XLA (GSPMD)

JAX, together with the XLA compiler, uses the GSPMD approach to automatically partition computations across devices. Users annotate tensors with sharding specifications (e.g., using jax.sharding), and XLA compiles a distributed program with optimal collective communications. GSPMD supports arbitrary tensor parallelism and is used in Google’s PaLM and Gemini models.

3.2 Hardware considerations

The efficiency of tensor parallelism is heavily influenced by the underlying hardware topology.

Tensor parallelism requires frequent all‑reduce and all‑gather operations. High‑bandwidth interconnects such as NVIDIA NVLink (up to 900 GB/s per GPU pair in H100) and InfiniBand (up to 400 Gb/s per link) reduce communication latency. On systems with lower bandwidth (e.g., Ethernet), tensor parallelism may become bottlenecked by communication, limiting scalability.

3.2.2 Memory hierarchy and compute-to-communication ratio

The ratio of computation to communication per tensor operation determines whether tensor parallelism is beneficial. For layers with large hidden dimensions (e.g., 12288 in GPT‑3), the compute load is high relative to the data volume communicated, making tensor parallelism efficient. Conversely, small layers may see overhead dominate. Modern accelerators with high compute throughput (e.g., H100 Tensor Cores) shift the balance toward communication optimization.

3.3 Kernel optimizations

To mitigate communication overhead, frameworks employ specialized kernels.

3.3.1 Fused kernels for communication and compute

Fused kernels combine communication operations with adjacent computations. For example, a kernel may perform a matrix multiplication and an all‑reduce in one step, reducing kernel launch overhead and memory traffic. Megatron‑LM’s “fused bias gelu” kernel and DeepSpeed’s “fused communication” kernels are examples.

3.3.2 Overlapping communication with computation

By scheduling communication operations to run concurrently with computation (e.g., using CUDA streams), the effective latency of all‑reduce or all‑gather can be hidden. This technique is particularly effective when the computation graph allows pipelining of multiple layers. Overlapping can improve throughput by 10–30% in large models.

4 Comparison with Other Parallelism Strategies

4.1 Data parallelism

Data parallelism replicates the entire model on each device and splits the training data. It is simple to implement and scales well for small models. However, it does not reduce per‑device memory for the model itself, making it unsuitable when a model exceeds device memory. Tensor parallelism complements data parallelism by reducing memory per device, and both are often used together.

4.2 Pipeline parallelism

Pipeline parallelism partitions the model by layers, assigning consecutive layers to different devices. It reduces the memory footprint for activations but introduces idle time (bubbles) due to pipeline flushes. Tensor parallelism, by contrast, splits a single layer across devices, which increases communication frequency but can be combined with pipeline parallelism to further reduce memory and improve utilization.

4.3 Sequence parallelism

Sequence parallelism splits the input sequence dimension across devices, often used in conjunction with tensor parallelism for long‑context transformers. While tensor parallelism splits the hidden dimension, sequence parallelism splits the token dimension, enabling training of extremely long sequences without exceeding device memory. Both can be applied simultaneously.

4.4 Hybrid parallelism (3D)

Hybrid parallelism combines data, pipeline, and tensor parallelism to exploit all dimensions of the hardware topology. Typically, tensor parallelism is applied within a node (where inter‑device bandwidth is high), pipeline parallelism is applied across nodes, and data parallelism is applied across multiple nodes or replicas. This three‑way decomposition maximizes throughput and memory efficiency for the largest models, as demonstrated in training runs of GPT‑3, PaLM, and LLaMA‑2.

5 Applications

5.1 Training large language models (GPT‑3, LLaMA, PaLM)

Tensor parallelism is a cornerstone of training large language models. OpenAI’s GPT‑3 (175B parameters) used tensor parallelism across 8 GPUs per node with NVLink. Meta’s LLaMA models employed tensor parallelism in combination with pipeline parallelism. Google’s PaLM (540B parameters) used GSPMD‑based tensor parallelism across TPU v4 pods. In all cases, tensor parallelism enabled the storage and computation of massive weight matrices that would otherwise be impossible on a single accelerator.

5.2 Inference serving and on‑device deployment

During inference, tensor parallelism allows large models to be served across multiple devices with low latency. For example, a 175B‑parameter model can be split across 8 GPUs, with each GPU holding a fraction of the parameters. Inference frameworks such as NVIDIA Triton Inference Server and vLLM support tensor parallelism to reduce per‑device memory and improve throughput. On devices with limited memory (e.g., smartphones), tensor parallelism can enable quantized models with split layers to run across multiple NPUs.

5.3 Scientific computing (e.g., physics simulations with neural operators)

Beyond language modeling, tensor parallelism is used in scientific machine learning. Neural operators (e.g., Fourier Neural Operators) that process large 3D grids can be parallelized across GPUs by splitting spatial or channel dimensions. Tensor parallelism reduces memory for each device and allows scaling to higher‑resolution simulations. Frameworks like NVIDIA Modulus and JAX‑based PDE solvers leverage this technique.

6 Challenges and Future Directions

6.1 Communication bottlenecks and topology‑aware partitioning

As model sizes grow and accelerator counts increase, communication becomes a primary bottleneck. Current tensor parallelism uses simple row/column splits, which may not align with the underlying network topology. Future work explores topology‑aware partitioning that groups devices with high‑bandwidth connections (e.g., NVLink domains) for tensor parallelism, while using lower‑bandwidth interconnects for pipeline or data parallelism. Adaptive algorithms that dynamically adjust sharding based on measured communication costs are also under development.

6.2 Dynamic tensor parallelism for adaptive workloads

Many modern models have variable‑sized layers (e.g., heterogeneous transformer blocks) or dynamic computation graphs. Static tensor parallelism, where the partitioning is fixed at compile time, may lead to load imbalance. Dynamic tensor parallelism allows re‑partitioning during training or inference, redistributing tensors based on current memory and compute loads. This is an active research area, particularly for mixture‑of‑experts models where different experts may have varying sizes.

6.3 Integration with sparsity and mixture‑of‑experts models

Sparse models, such as Mixture‑of‑Experts (MoE), introduce additional dimensions of parallelism. Tensor parallelism can be applied to both the dense layers and the expert layers. However, sparsity complicates partitioning because different inputs activate different experts. Efficient tensor parallelism for MoE requires balancing the communication of router outputs and expert parameter slices. Techniques like expert‑sharding (a form of tensor parallelism across experts) and combined data‑expert parallelism are being integrated with traditional tensor parallelism to scale models beyond a trillion parameters.