Tensor Parallelism

Unlike pipeline parallelism which splits the model by layers (inter-layer), tensor parallelism splits individual layers (intra-layer). The weights and computations of a single layer, such as a transformer's Multi-Head Attention or Feed-Forward Network, are distributed across multiple GPUs. Each GPU holds a shard of the weight matrix and processes a portion of the input, requiring an all-reduce communication operation to combine results before proceeding to the next layer.

Tensor parallelism introduces significant inter-GPU communication overhead within each forward and backward pass. Key operations include:

• All-Reduce: Used to sum partial results from each GPU shard.
• All-Gather: Used to collect sharded data from all devices.
• Reduce-Scatter: Used to distribute and reduce data across devices. The frequency of these operations makes high-bandwidth interconnects like NVLink essential for performance, as communication can become the bottleneck, especially for smaller batch sizes.

A primary benefit is the reduction of memory pressure per GPU. By sharding the model's parameters, gradients, and optimizer states, each GPU only needs to store a fraction of the full model's weights. This is critical for deploying models whose size exceeds the VRAM of a single accelerator. For example, a 70B parameter model with FP16 weights (~140GB) can be split across 8 GPUs, requiring only ~18GB of model memory per device.

Specific patterns define how tensors are split:

• Megatron-LM Style (Column/Row Parallelism): Splits linear layers by column (for the first GEMM in a feed-forward network) or by row (for the second GEMM). This minimizes communication by strategically placing all-reduce operations.
• Tensor Model Parallelism in PyTorch: Implemented via torch.distributed and torch.nn.parallel.DistributedDataParallel with custom communication primitives.
• Framework Integration: Deeply integrated into serving engines like vLLM and Text Generation Inference (TGI), which combine it with other optimizations like PagedAttention and continuous batching.

Tensor parallelism is rarely used in isolation. It is typically combined with:

• Data Parallelism: For scaling across many nodes, where each data-parallel group contains a full model replica split via tensor parallelism.
• Pipeline Parallelism: For extremely large models, where different groups of layers are placed on different devices (pipeline parallelism), and individual layers within a group are further split via tensor parallelism. This hybrid approach is used in frameworks like DeepSpeed and Megatron-DeepSpeed.

Ideal for:

• Serving or training single, massive models that don't fit on one GPU.
• Scenarios where low latency per request is critical, as it avoids the pipeline bubbles of pipeline parallelism.

Key Limitations:

• Communication Overhead: Performance degrades without ultra-fast interconnects.
• Limited Scalability: Efficiency drops when scaling across too many GPUs for a single layer due to increased communication. It is most effective within a single node or high-bandwidth pod.
• Implementation Complexity: Requires careful manual partitioning or reliance on specialized frameworks.

PyTorch's Fully Sharded Data Parallel (FSDP) is a hybrid parallelism technique that combines data parallelism with tensor and pipeline parallelism concepts. It shards model parameters, gradients, and optimizer states across data-parallel workers, significantly reducing per-GPU memory footprint. While not pure tensor parallelism, FSDP's parameter sharding is a form of intra-layer model parallelism that enables training models larger than the memory of a single GPU.

• Key Mechanism: Shards each model parameter across the data-parallel dimension.
• Communication Pattern: Uses all-gather to collect parameters for computation and reduce-scatter to aggregate gradients.
• Use Case: Primarily for training extremely large models where memory, not just compute, is the primary constraint.

A model parallelism technique where consecutive layers (stages) of a neural network are placed on different GPUs. Activations are passed sequentially between stages. To improve hardware utilization, it often employs micro-batching, where small batches are processed in an interleaved manner across stages.

• Key Difference from Tensor Parallelism: Splits the model by layer depth, not by splitting individual layers.
• Communication Pattern: Point-to-point between adjacent stages.
• Use Case: Effective for models with many sequential layers where a single layer fits on one GPU.

The most common distributed training strategy, where identical model replicas are placed on multiple devices (GPUs). Each replica processes a different subset of the training data (a mini-batch) in parallel. Gradients are averaged across all replicas after each backward pass to synchronize model updates.

• Key Difference: Replicates the entire model; parallelism comes from data, not model structure.
• Communication Overhead: Requires all-reduce operations to synchronize gradients, which can become a bottleneck for very large models.
• Primary Use: Scaling training batch size, often combined with model parallelism techniques like tensor or pipeline parallelism.

A memory optimization paradigm for data parallel training, developed by Microsoft as part of DeepSpeed. ZeRO eliminates memory redundancy by partitioning the optimizer states, gradients, and model parameters across data parallel processes instead of replicating them.

• ZeRO Stages:
  • Stage 1: Partitions optimizer states.
  • Stage 2: Partitions gradients + optimizer states.
  • Stage 3: Partitions parameters, gradients, and optimizer states (full model partitioning).
• Relationship to Tensor Parallelism: ZeRO-3 can be viewed as a form of parameter partitioning across data parallel workers, complementary to tensor parallelism's intra-layer splitting.