Tensor Parallelism

Tensor Parallelism operates within a single neural network layer, splitting large tensor operations like matrix multiplications. For example, a weight matrix W of size [M, N] can be split column-wise across P devices, with each device holding a shard of size [M, N/P]. The input is broadcast, each device performs a local GEMM (General Matrix Multiply), and the outputs are combined via a communication operation like All-Gather. This contrasts with Model Parallelism, which splits entire layers across devices.

The efficiency of Tensor Parallelism is governed by its communication patterns, which are required to combine partial results. Key operations include:

• All-Reduce: Sums partial results from all devices (common in linear layers).
• All-Gather: Collects sharded results from all devices to form a complete tensor.
• Reduce-Scatter: Reduces (e.g., sums) a tensor and scatters the result. The communication-to-computation ratio is critical. For large matrix multiplications, the high computational cost can often hide this overhead, but for smaller operations or slower interconnects, it becomes a bottleneck.

Tensor Parallelism is rarely used in isolation. It is typically combined with other strategies to train massive models:

• Data Parallelism: Different devices process different data batches. Combined with Tensor Parallelism, this forms a 2D parallelism scheme.
• Pipeline Parallelism: Different devices handle different layers (stages). Combined with Tensor and Data parallelism, this forms the powerful 3D parallelism paradigm used in frameworks like Megatron-LM. This hybrid approach allows scaling to models with hundreds of billions or trillions of parameters by distributing memory and compute load across thousands of accelerators.

Optimal performance requires mapping the logical parallel group to the physical hardware topology. Communication within a Tensor Parallel group is frequent and bandwidth-intensive. Therefore, devices in a group should be connected via the highest-bandwidth links available (e.g., NVLink within a node, InfiniBand across nodes). Placing the group across slow interconnects can cripple performance. This necessitates topology-aware scheduling in cluster managers to ensure low-latency, high-bandwidth connectivity for communicating devices.

A primary benefit of Tensor Parallelism is the reduction of activation memory per device. During the forward pass, large intermediate tensors (activations) are sharded across devices. For example, the output of a large linear layer is split, so each device only stores a fraction. This is crucial for training very large models, as activation memory often becomes the limiting factor before parameter memory. However, it does not reduce the memory footprint of the optimizer states, which is addressed by complementary techniques like the Zero Redundancy Optimizer (ZeRO).

This is the most prominent use case. The multi-billion parameter feed-forward layers (MLPs) and attention projections within Transformer decoders are split across devices.

• Enables serving models like GPT-4, Llama 3, and Mixtral on multi-GPU/TPU pods.
• Reduces per-device memory footprint, allowing for larger batch sizes or longer context lengths during inference.
• Communication overhead is incurred during the all-reduce operations required to combine partial results after each parallelized matrix multiplication.

Used in conjunction with Data Parallelism and Pipeline Parallelism to train models with trillions of parameters.

• Tensor Model Parallelism (as in Megatron-LM) splits individual layer operations.
• Applied to the most memory-intensive components: the GEMM operations in self-attention and MLP blocks.
• Allows the optimizer states and parameters for a single layer to be partitioned, overcoming the memory wall for foundation model pre-training.

Targets layers where the parameter size exceeds the high-bandwidth memory (HBM) of a single accelerator.

• Large Embedding Tables in recommendation models can be sharded across devices using tensor parallelism.
• Wide Linear Layers (e.g., in MoE experts or final output layers) with massive input/output dimensions.
• The technique transforms a memory capacity problem into a communication bandwidth problem, which is often easier to scale with specialized interconnects like NVLink.

Mapping parallelism strategy to physical hardware layout for minimal latency.

• On NVIDIA DGX systems, tensor-parallel groups are confined to nodes with NVLink for fast all-reduce.
• On TPU v4/v5 pods, leveraging the high-bandwidth ICI (Inter-Chip Interconnect) for cross-core communication within a slice.
• Contrast with Pipeline Parallelism, which is better suited for cross-node scaling with slower interconnects, as it requires less frequent communication.

Enables scaling the expert capacity within sparse MoE layers.

• Each expert's dense feed-forward network can be split via tensor parallelism across a subset of devices.
• This allows for larger, more powerful experts without being limited by the memory of a single device.
• Combined with Expert Parallelism (where different experts are placed on different devices), it facilitates the training of massive models like Google's Switch Transformer.

Addresses hardware constraints beyond just memory, such as register file size or thread block limits.

• Extremely large kernel operations (e.g., a massive matrix multiplication) may hit resource limits on a single GPU's streaming multiprocessor (SM).
• Splitting the operation via tensor parallelism creates smaller, more hardware-friendly kernels for each device.
• This can improve occupancy and instruction-level parallelism by reducing resource contention per device.

Model Parallelism is a distributed training strategy where the parameters and computational graph of a single neural network are partitioned across multiple accelerators. Unlike data parallelism, where each device holds a full model copy, model parallelism splits the model itself.

• Key Distinction: Tensor Parallelism is a form of model parallelism focused on splitting individual tensor operations.
• Use Case: Essential for models whose parameters exceed the memory of a single device (e.g., large language models with hundreds of billions of parameters).
• Communication Pattern: Requires frequent synchronization between devices during the forward and backward passes, as the output of one partitioned layer becomes the input for the next.

Pipeline Parallelism is a distributed training strategy where the sequential layers of a model are assigned to different devices, forming a processing pipeline. Micro-batches of data flow through these stages sequentially.

• Core Mechanism: Overlaps computation across devices by having different devices process different micro-batches simultaneously (e.g., GPU 1 processes the forward pass of micro-batch 2 while GPU 2 processes the backward pass of micro-batch 1).
• Bubble Problem: Inefficiency arises from pipeline 'bubbles'—idle time at the start and end of a batch as the pipeline fills and drains.
• Combination with Tensor Parallelism: Often used in conjunction with Tensor Parallelism in systems like Megatron-LM, where layers are split across devices via pipeline parallelism, and the attention/MLP blocks within a layer are split via tensor parallelism.

Data Parallelism is the most common distributed training strategy, where the same model is replicated across all available devices, and each device processes a different subset (shard) of the training data in parallel.

• Synchronization: Gradients are averaged across all devices after each backward pass, typically using an All-Reduce collective communication operation, to update the identical model replicas.
• Primary Limitation: The entire model must fit into the memory of a single device.
• Contrast with Tensor Parallelism: Data Parallelism replicates the model; Tensor Parallelism splits the model. They are orthogonal and often combined for extreme scaling.

The Zero Redundancy Optimizer (ZeRO) is a memory optimization paradigm for data-parallel training that partitions optimizer states, gradients, and model parameters across devices to eliminate memory redundancy.

• ZeRO Stages:
  • Stage 1: Partitions optimizer states.
  • Stage 2: Partitions gradients in addition to optimizer states.
  • Stage 3: Partitions model parameters in addition to the above.
• Relationship to Parallelism: ZeRO enables data parallelism for models that would otherwise require model/tensor parallelism by removing the memory overhead of replicating optimizer states and parameters. It is often used alongside tensor and pipeline parallelism in frameworks like DeepSpeed.

Operator Fusion is a compiler-level optimization critical for efficient execution, especially in parallel contexts. It combines multiple sequential neural network operations into a single, fused kernel.

• Objective: Minimize costly reads and writes of intermediate tensors to slow global memory (e.g., HBM).
• Example: Fusing a convolution, bias addition, and ReLU activation into one kernel. Instead of writing the convolution output to memory and reading it back for the next operation, the result is kept in fast registers or shared memory.
• Impact on Tensor Parallelism: Reduces the volume and frequency of communication required between devices when operations within a fused group are split, as only the final fused result may need to be communicated.

Collective Communication refers to coordinated communication patterns involving a group of processes (e.g., across multiple GPUs). It is the backbone of all distributed training strategies.

• Key Operations for Tensor Parallelism:
  • All-Reduce: Sums a tensor across all devices and distributes the result back to each. Used in data parallelism for gradient averaging.
  • All-Gather: Each device holds a shard of a tensor; the operation collects all shards, making the full tensor available on every device.
  • Reduce-Scatter: A tensor is summed across devices, but the resulting sum is split into shards, with each device receiving a unique shard.
• Performance Bottleneck: The latency and bandwidth of these collectives often determine the efficiency of tensor-parallel execution.