Model Parallelism

Tensor parallelism splits individual weight matrices or tensors across devices, performing the computation for a single layer in parallel. This is often applied to the large linear layers within transformer blocks.

• Key Mechanism: For a matrix multiplication Y = XW, the weight matrix W is partitioned column-wise (or row-wise) across devices. Each device computes a partial result, and communication (like an all-reduce operation) is required to combine outputs.
• Primary Use: Essential for models where even a single layer's parameters exceed a device's memory, such as in massive Mixture of Experts (MoE) layers or dense feed-forward networks in large language models.
• Example: In the Megatron-LM framework, the self-attention and MLP layers of transformers are split across GPUs using tensor parallelism.

Pipeline parallelism partitions the model's layers (vertical slicing) across devices. Each device holds a consecutive subset of the model's layers, forming a processing pipeline.

• Key Mechanism: A mini-batch of data is split into smaller micro-batches. These micro-batches are fed into the pipeline sequentially. While one device processes a micro-batch for its set of layers, other devices can be processing different micro-batches for their respective stages, aiming to keep all devices busy.
• Challenge: Pipeline bubbles—idle time created when the pipeline is filling or draining—reduce hardware utilization. Techniques like 1F1B (One Forward pass followed by One Backward pass) scheduling are used to mitigate this.
• Primary Use: Training extremely deep models (e.g., models with hundreds of layers) where the sequential depth is the limiting factor.

Expert parallelism is a specialized strategy for Mixture of Experts (MoE) models. It distributes the different 'expert' sub-networks across the available devices.

• Key Mechanism: The gating/router network selects which experts are relevant for a given input token. The tokens are then dispatched (via an all-to-all communication operation) to the devices hosting those specific experts for processing. The results are gathered and combined.
• Primary Use: Enables the training of models with a massive number of parameters (e.g., trillions) by activating only a small, conditional subset of parameters per token. This provides large model capacity without a proportional increase in FLOPs per token.
• Example: Used extensively in models like Switch Transformers and GLaM, where thousands of experts are sharded across many devices.

3D parallelism combines data parallelism, tensor parallelism, and pipeline parallelism to train models at the extreme scale of hundreds or thousands of accelerators.

• Key Mechanism:
  • Data Parallelism replicates the entire model across groups of devices, splitting the batch of data.
  • Tensor Parallelism splits individual layers within a device group.
  • Pipeline Parallelism splits layers across different device groups.
• Communication Patterns: Each parallelism type introduces a different communication cost (all-reduce for data, all-reduce for tensor, point-to-point for pipeline). The combined strategy must carefully orchestrate these to minimize bottlenecks.
• Primary Use: Training state-of-the-art foundation models (e.g., GPT-3, GPT-4, BLOOM) where the model size far exceeds the memory of any single device or even a single node.

Sequence parallelism partitions the sequence length dimension (the batch of tokens) across devices. This is crucial for training models with very long context windows.

• Key Mechanism: For operations like layer normalization or dropout, which require statistics computed across the sequence dimension, an all-gather communication is needed. For attention, it enables the distribution of the massive Key-Value (KV) cache.
• Primary Use: Directly addresses the memory bottleneck of the activations (not just parameters) during training for long sequences. It reduces the peak activation memory per device by a factor proportional to the degree of sequence parallelism.
• Relation to Other Techniques: Often used in conjunction with tensor or pipeline parallelism. It is conceptually similar to splitting the batch dimension in data parallelism, but applied to the sequence length within a single sample or micro-batch.

The Zero Redundancy Optimizer (ZeRO) is a memory optimization paradigm that enables efficient data-parallel training by partitioning the optimizer states, gradients, and model parameters across devices, eliminating memory redundancy.

• Key Stages:
  • ZeRO-Stage 1: Partitions optimizer states (e.g., momentum, variance) across devices.
  • ZeRO-Stage 2: Also partitions gradients, with an all-gather only for the parameters needed for the local layer update.
  • ZeRO-Stage 3: Also partitions the model parameters themselves, gathering them just-in-time for forward/backward computation.
• Primary Use: While not pure model parallelism, ZeRO is a foundational technique that makes large-scale data parallelism feasible. It is often combined with model parallelism (e.g., pipeline or tensor) in frameworks like DeepSpeed to achieve optimal memory efficiency and enable the training of models with over a trillion parameters.

A distributed training strategy where identical copies of the entire model are placed on multiple devices (GPUs/TPUs). Each device processes a different subset of the training data (a mini-batch) in parallel. The key distinction from model parallelism is that gradients are synchronized across devices after each forward/backward pass via an all-reduce operation to update all model copies consistently. This is highly effective for scaling batch size when the model fits on a single device.

A memory optimization paradigm for data-parallel training that partitions the optimizer states, gradients, and model parameters across devices instead of replicating them. This eliminates memory redundancy, allowing for the training of models much larger than what can fit on any single device's memory. ZeRO operates in stages:

• Stage 1: Partitions optimizer states.
• Stage 2: Partitions gradients.
• Stage 3: Partitions model parameters. It's implemented in DeepSpeed and integrated into PyTorch as Fully Sharded Data Parallel (FSDP).