Model Parallelism

Tensor Parallelism splits individual tensor operations, like large matrix multiplications within a single layer, across multiple devices. For example, in a linear layer Y = XW, the weight matrix W is partitioned column-wise. Each device computes a partial result, which must then be combined via an all-reduce communication operation.

• Key Use Case: Extremely wide layers (e.g., feed-forward networks in Transformers).
• Communication Pattern: High frequency, as communication is required within layers.
• Hardware Requirement: Very high-bandwidth interconnects (e.g., NVLink) to mitigate communication overhead.

Pipeline Parallelism partitions the model's layers (stages) across different devices. The network is split into sequential chunks, forming a processing pipeline. Micro-batches of data flow through these stages sequentially. To keep all devices busy, the pipeline must be bubbles.

• Key Technique: GPipe and its variants use gradient accumulation and synchronous weight updates.
• Communication Pattern: Point-to-point between adjacent stages, lower frequency than tensor parallelism.
• Optimization: 1F1B (One Forward pass followed by One Backward pass) scheduling to reduce the pipeline bubble and memory footprint.

Expert Parallelism is a strategy designed for Mixture-of-Experts (MoE) models. In an MoE layer, different specialized sub-networks (experts) are placed on different devices. A gating network routes each input token to a small subset of experts (e.g., top-2).

• Core Mechanism: Only the experts responsible for a given token are activated, creating sparse computation.
• Communication Challenge: Requires an all-to-all collective to route tokens to their designated expert devices and gather results.
• Scale: Enables models with trillions of parameters, as the total parameter count grows with the number of experts while active compute remains roughly constant.

3D Parallelism combines Data Parallelism, Pipeline Parallelism, and Tensor Parallelism to train models at the extreme scale of hundreds or thousands of accelerators. Each dimension addresses a different scaling constraint.

• Data Parallelism: Replicates the entire model across groups of devices for different data batches.
• Pipeline Parallelism: Splits the model layers across devices within a data-parallel group.
• Tensor Parallelism: Splits individual layers across a final subset of devices.
• Implementation: Frameworks like Megatron-LM and DeepSpeed use this hybrid approach to train models like GPT-3 and beyond, optimizing the trade-off between compute, memory, and communication efficiency.

Sequence Parallelism partitions the sequence dimension (e.g., tokens in a sentence) of the input tensor across devices. This is crucial for training Transformer models on very long sequences that do not fit in a single device's memory.

• Applies To: Operations like layer normalization, dropout, and especially the attention mechanism.

• Attention Implementation: Requires careful communication to assemble the full QK^T attention scores matrix across devices before applying the softmax, often using an all-gather operation.

• Benefit: Enables training with context lengths in the hundreds of thousands of tokens by distributing the massive activation memory for the sequence dimension.

While not pure model parallelism, the Zero Redundancy Optimizer (ZeRO) is a foundational memory optimization that enables it. ZeRO partitions the optimizer states, gradients, and model parameters across devices in a data-parallel group, eliminating memory redundancy.

• ZeRO Stage 3: Partitions the entire model state, meaning each device only stores a shard of the parameters it is responsible for updating. This effectively enables model parallelism from a memory perspective.
• Communication Cost: Requires all-gather operations to collect parameters before forward/backward passes and reduce-scatter to aggregate gradients, trading communication for massive memory savings.
• Result: Allows training models far larger than the aggregate memory of a single device.

This is the quintessential use case. Models like GPT-4, Llama, and Claude have hundreds of billions of parameters, far exceeding the memory capacity of any single GPU or TPU. Model parallelism is combined with data parallelism and pipeline parallelism in frameworks like Megatron-LM to enable training. Layers are split across devices, with activations and gradients communicated between them during the forward and backward passes.

Deploying a multi-hundred-billion parameter model for real-time inference requires distributing it across many inference chips. Model parallelism allows the model to be loaded into the aggregate memory of a cluster. Specialized serving systems like TensorRT-LLM or vLLM manage the partitioning and orchestrate the execution flow across devices, handling the communication of intermediate key-value (KV) caches for autoregressive generation.

Large Vision Transformers (ViTs) and multimodal architectures (e.g., CLIP, Flamingo) with massive embedding dimensions and attention heads also exceed single-device memory. Tensor parallelism is used to split the large matrix multiplications within the multi-head attention and feed-forward layers across devices. This allows training on high-resolution image and video data where activations are already memory-intensive.

Model parallelism can be applied in clusters with mixed accelerator types (e.g., combining GPUs with NPUs). Layers can be assigned to the hardware best suited for their operation type. For instance, convolution-heavy early vision layers could be placed on optimized NPUs, while attention-heavy transformer layers are placed on GPUs. This requires sophisticated runtime schedulers to manage the cross-device data flow.