Tensor Parallelism

Tensor parallelism operates within a single neural network layer, splitting the weight matrices and activations of that layer across multiple devices. This is distinct from pipeline parallelism, which splits different layers across devices.

• For a linear layer Y = XW, the weight matrix W can be split column-wise or row-wise.
• A column-wise split requires an all-gather operation after the distributed matrix multiplication to combine results.
• A row-wise split requires the input X to be broadcast or split accordingly.
• This fine-grained splitting allows the training of layers whose parameters exceed the memory of a single device.

Because operations are split within a layer, tensor parallelism requires frequent synchronization between devices during the forward and backward passes. The communication pattern is characterized by collective operations.

• All-reduce is commonly used during the backward pass to sum gradients from all partitions.
• All-gather is used to collect sharded outputs in the forward pass.
• Reduce-scatter can be used to distribute and sum gradients efficiently.
• The communication overhead scales with the activation size and the degree of parallelism, making it most efficient for layers with very large hidden dimensions where compute time dominates communication.

This strategy is particularly effective for layers with massive weight matrices, such as the feed-forward networks (FFNs) and attention projections in modern transformers. The efficiency gain comes from distributing the computationally intensive matrix multiplications.

• In a transformer's FFN layer (e.g., with a hidden dimension of 4096 expanding to 16384), the large intermediate matrix is an ideal candidate for splitting.
• The Megatron-LM approach famously applies tensor parallelism to both the self-attention and FFN modules.
• The benefit diminishes for layers with small hidden sizes, where communication overhead can negate computational gains.

Performance is highly dependent on the interconnect bandwidth and latency between devices. Optimal deployment requires careful mapping of model partitions to the physical hardware topology.

• NVLink or NVSwitch connections between GPUs provide the high-bandwidth, low-latency communication essential for efficient tensor parallelism.
• Placing partitions across a slower PCIe bus or network interconnect can create a severe communication bottleneck.
• For multi-node setups, tensor parallelism is often combined with other strategies (like pipeline parallelism) to confine its high-bandwidth requirements to within a single node.

In practice, tensor parallelism is rarely used alone. It is combined with data parallelism and pipeline parallelism in a 3D parallelism strategy to train trillion-parameter models.

• Data Parallelism: Replicates the entire model across device groups, splitting the batch. Handles sample-level parallelism.
• Tensor Parallelism (intra-layer): Splits individual layers. Handles model-component-level parallelism.
• Pipeline Parallelism (inter-layer): Splits different layers of the model. Handles model-depth parallelism.
• This hybrid approach, exemplified by DeepSpeed and Megatron-DeepSpeed, allows each form of parallelism to address different scaling constraints (memory, compute, communication).

Implementing efficient tensor parallelism requires deep integration with the model execution runtime and compiler stack. Major frameworks provide specialized APIs and automated strategies.

• PyTorch: Supports it via torch.distributed.tensor (DTensor) and the parallelize_module API, allowing sharding annotations on nn.Modules.
• DeepSpeed: Offers tensor parallelism through its inference and training engines, often in conjunction with its ZeRO memory optimizations.
• JAX: Enables tensor parallelism via the pjit (parallel jit) transformation and sharding specifications on arrays.
• The compiler's role is to lower the annotated sharded operations to efficient kernel launches and the necessary collective communication primitives.

On specialized Neural Processing Units (NPUs), tensor parallelism is implemented through hand-optimized kernels that leverage hardware-specific matrix multiplication units (MXUs) and high-bandwidth on-chip memory.

• Hardware Mapping: The sharded matrix multiplications are mapped directly to the NPU's systolic arrays or tensor cores, with communication between cores handled via dedicated on-chip networks.
• Memory Efficiency: By splitting tensors, each NPU core operates on a smaller block, reducing its local memory footprint and allowing larger effective models to run.
• Vendor SDKs: Implementation relies on low-level APIs in vendor SDKs (e.g., NVIDIA's CUDA, Google's TPU API, AMD's ROCm) to manage the distributed computation and synchronization.

The most fundamental use case is to overcome the hard memory wall of a single accelerator. When a model's layer is too large to load, tensor parallelism provides a direct solution.

• Problem: A linear layer with shape [Hidden_In, Hidden_Out] where Hidden_In * Hidden_Out * dtype_size > Device Memory.
• Solution: Split the weight matrix along its rows or columns. For a column-wise split, the input is broadcast, and each device computes a partial output. An all-gather operation then reconstructs the full output.
• Trade-off: Introduces communication overhead proportional to the size of the activations, making it most efficient for layers with very large hidden dimensions where computation dominates.

Data parallelism is a parallel computing paradigm where the same operation (e.g., a forward pass) is applied concurrently to different subsets (batches) of a dataset across multiple processing units (e.g., GPUs). Each device holds a complete copy of the model. Gradients are synchronized across devices after processing each batch, typically using an All-Reduce operation.

• Primary Use: Training models where the model fits on a single device but the dataset is large.
• Key Mechanism: Synchronous or asynchronous gradient aggregation.
• Example: Training a ResNet-50 on 8 GPUs, where each GPU processes 32 images from a total batch size of 256.

Model parallelism is a technique for partitioning the computational graph or parameters of a neural network across multiple processors or devices to handle models that are too large to fit on a single unit's memory. Unlike tensor parallelism, which splits individual operations, model parallelism typically splits the network by layers or sub-graphs.

• Primary Use: Running or training models whose parameters exceed the memory of a single accelerator.
• Key Mechanism: Different devices execute different parts of the model's sequential layers.
• Example: Placing the first 24 transformer decoder layers of a large language model on GPU 0 and the remaining 24 layers on GPU 1.

Pipeline parallelism is a strategy that partitions a model's layers across multiple devices and processes different microbatches of data simultaneously in a staged assembly line. It introduces bubbles (idle time) into the pipeline but allows for high throughput by keeping all devices active.

• Primary Use: Training very large models where both data and model parallelism are insufficient.
• Key Mechanism: Overlapping computation across devices by scheduling microbatches.
• Scheduling Schemes: GPipe (synchronous, large bubbles) and PipeDream (asynchronous, 1F1B).

SIMD (Single Instruction, Multiple Data) and SIMT (Single Instruction, Multiple Threads) are parallel processing architectures at the hardware instruction level that tensor parallelism leverages.

• SIMD: A single instruction controls multiple processing elements to perform the same operation on multiple data points simultaneously. Common in CPU vector units (AVX, NEON).
• SIMT: The execution model of GPUs. A single instruction is issued to a warp (typically 32 threads), where each thread executes it on its own data. It handles control flow divergence by masking threads.
• Relation to Tensor Parallelism: Splitting a large matrix multiplication across devices effectively creates a larger, distributed SIMD/SIMT operation.

A memory consistency model defines the formal rules for the observable order of memory operations (loads and stores) performed by different threads or processes in a parallel system. It is critical for correctness when implementing tensor parallelism across devices with shared or partitioned memory.

• Sequential Consistency: The simplest model; the result of any execution is as if all operations were executed in some sequential order consistent with program order.
• Weaker Models: Modern hardware (GPUs, NPUs) often employ weaker models (e.g., release-acquire semantics) for performance, requiring explicit memory barriers or fences to enforce ordering for correctness.

Amdahl's Law and scaling laws provide the theoretical framework for analyzing the benefits of parallelism like tensor parallelism.

• Amdahl's Law: States the maximum speedup of a program is limited by its serial fraction. If 10% of a program is serial, maximum speedup is 10x, regardless of processors.
• Strong Scaling: Measures time reduction for a fixed problem size with added processors. Tensor parallelism aims for strong scaling on large layer computations.
• Weak Scaling: Measures throughput increase when problem size grows proportionally with processors. Ideal for scenarios where tensor size grows with model capacity.