Model Parallelism

Tensor parallelism splits individual tensor operations, like matrix multiplications, across devices. For a linear layer Y = XW, the weight matrix W is partitioned along its column or row dimension. Each device holds a shard of W and computes a partial result, requiring an all-reduce communication operation to combine outputs. This is highly effective for models with extremely large feed-forward or attention layers, as used in models like GPT-3 and Megatron-LM. Key considerations include:

• Minimizing communication overhead by partitioning the dimension with the largest size.
• Balancing computational load across devices.
• Optimizing for high-bandwidth inter-device links like NVLink.

Pipeline parallelism partitions the model vertically by assigning different layers (or groups of layers) to different devices. The training batch is split into smaller microbatches that are fed into the pipeline sequentially. While one device processes a microbatch for its assigned layers, the next device can process the previous microbatch, creating an assembly-line effect. This strategy introduces pipeline bubbles—periods of device idle time. Techniques like GPipe and PipeDream aim to minimize these bubbles through careful scheduling. It is ideal for models with many sequential layers, such as vision transformers or deep ResNets.

Sequence parallelism is a specialized form of tensor parallelism designed for long-sequence models. It partitions the sequence length dimension of activations and associated computations (like attention) across devices. For example, in the attention mechanism, the Query, Key, and Value projections for different chunks of the sequence are computed on different devices. This requires communication for operations like all-gather and reduce-scatter. It is crucial for training models like LLaMA on very long contexts, as it reduces the peak memory per device needed to store activations for the entire sequence.

Expert parallelism is the native strategy for Mixture of Experts (MoE) models. Each expert—a dedicated feed-forward sub-network—is placed on a different device. A gating network routes each input token to the top-k most relevant experts. This requires an all-to-all communication operation to send tokens to the devices hosting their assigned experts and another to return the processed outputs. This strategy enables massive model capacity (e.g., models with trillions of parameters like Switch Transformers) while activating only a sparse subset of parameters per token, keeping computational cost manageable.

3D parallelism combines multiple strategies to train models at an extreme scale. It is a hybrid approach integrating:

• Data Parallelism: Replicates the entire model across groups of devices, splitting the batch.
• Tensor (Model) Parallelism: Splits layers within a model replica.
• Pipeline Parallelism: Splits layers across different model replicas. Frameworks like DeepSpeed and Megatron-DeepSpeed implement this to efficiently distribute trillion-parameter models across thousands of GPUs. The configuration (e.g., 4-way tensor, 8-way pipeline, 64-way data) is tuned based on the global number of devices, model size, and cluster interconnect bandwidth.

The efficiency of any model parallelism strategy hinges on optimizing inter-device communication. Key communication primitives include:

• All-Reduce: Aggregates values from all devices and distributes the result back (used in tensor parallelism).
• All-Gather: Collects shards from all devices to form a complete tensor on each device.
• Reduce-Scatter: Reduces values across devices and scatters the result in shards.
• All-to-All: Every device sends a distinct chunk of data to every other device (used in expert parallelism). Performance is maximized by overlapping communication with computation. While a GPU is computing, its DMA engines can asynchronously handle data transfers for the next required communication, hiding latency. This requires careful kernel scheduling and use of CUDA streams.

The primary driver for model parallelism is training foundation models like GPT-4, Llama, and Claude, whose parameter counts (e.g., hundreds of billions) far exceed the memory of even the largest GPUs (e.g., 80GB H100).

• Pipeline Parallelism: Layers are partitioned across devices in a sequential pipeline, with micro-batches flowing through to keep all devices utilized.
• Tensor Parallelism: Individual layers, like the attention heads or feed-forward networks in a transformer block, are split across devices, with communication required for each layer's operations.
• Combined Strategies: Modern frameworks like Megatron-LM and DeepSpeed combine pipeline and tensor parallelism (3D parallelism) to scale across thousands of GPUs.

During inference, model parallelism enables serving giant multimodal models (e.g., GPT-4V, Gemini Ultra) that process images, audio, and text. A single user query activates a model too large for one inference chip.

• Horizontal Partitioning: Different model components (vision encoder, language model, fusion network) are placed on separate devices.
• Memory-Constraint Solution: This is essential for providing low-latency API access to massive models, as loading the full model into a single device's memory is impossible.
• Dynamic Load Balancing: Some systems can shift partitions based on the input modality to optimize throughput.

In fields like climate modeling, drug discovery, and cosmology, models integrate massive physics-based simulations with deep neural networks, creating uniquely large architectures.

• Hybrid Model Partitioning: The physics simulation component may run on CPUs or specialized hardware, while the deep learning surrogate model is split across GPU clusters.
• Example - AlphaFold 2: While primarily using data parallelism, its complex Evoformer and structure module could be partitioned for memory-intensive protein folding predictions.
• Communication Patterns: These models often require custom parallelism strategies tailored to their specific computational graph and data flow.

Model parallelism is used in embodied AI and robotics, where a single model is partitioned across a robot's heterogeneous compute units: a GPU, an NPU, and a microcontroller.

• Sensor-Specific Branches: A vision branch of a neural network runs on a dedicated NPU, while the policy network runs on a GPU.
• Reducing Latency: By processing sensor inputs in parallel on dedicated hardware, overall pipeline latency is reduced compared to a sequential single-device setup.
• System-on-Chip (SoC) Design: This mirrors the hardware design of modern robotics SoCs, which contain multiple specialized accelerators.

Implementing model parallelism manually is complex. Several frameworks abstract the distribution logic.

• PyTorch Fully Sharded Data Parallel (FSDP): A hybrid of data and model parallelism that shards model parameters, gradients, and optimizer states across data-parallel workers.
• NVIDIA Megatron-LM: Provides efficient tensor and pipeline parallelism implementations specifically optimized for transformer models.
• Microsoft DeepSpeed: Offers the ZeRO family of optimizations, which can be viewed as a form of automated model parallelism across data-parallel processes.
• Google's JAX & Pathways: Uses a single-program, multiple-data (SPMD) model to express parallelism declaratively, which can compile to efficient model-parallel executions.

A distributed training strategy where identical copies of the entire model are placed on multiple devices (e.g., GPUs). Each device processes a different subset of the training data (a mini-batch) in parallel. Gradients are synchronized and averaged across all devices after each backward pass to update the model weights.

• Key Difference from Model Parallelism: Data parallelism replicates the model; model parallelism splits the model.
• Use Case: Ideal for training models that fit on a single device but require larger effective batch sizes for faster training.
• Framework Support: Native in PyTorch (DistributedDataParallel) and TensorFlow (tf.distribute.MirroredStrategy).

A specialized form of model parallelism where the model's layers are partitioned sequentially across devices, forming a processing pipeline. A single mini-batch is split into smaller micro-batches that are fed into the pipeline sequentially to keep all devices busy, overcoming the idle time inherent in naive model parallelism.

• Bubble Overhead: The challenge of pipeline 'bubbles' where some devices are idle during pipeline flush.
• GPipe & Interleaved Schedules: Techniques like those introduced in the GPipe paper to improve hardware utilization.
• Primary Use: Enables training of extremely deep models (e.g., transformers with hundreds of layers) that cannot fit layer-by-layer on one device.

A fine-grained model parallelism technique where individual tensor operations (like matrix multiplications) are split across devices. For example, the weight matrix of a large linear layer can be partitioned along its rows or columns, with the computation and communication orchestrated to produce the correct result.

• Megatron-LM: A seminal implementation for transformer models, splitting the attention heads and MLP layers across GPUs.
• Communication Intensive: Requires frequent all-reduce operations between devices during the forward and backward passes.
• Use Case: Critical for models where even a single layer's parameters exceed the memory of one accelerator, such as large dense feed-forward layers in transformers.

A sparsely-activated architecture that enables massive model capacity. The model consists of many expert sub-networks (e.g., feed-forward layers). A gating network dynamically routes each input token to only a few (e.g., 2) relevant experts. This is a form of conditional computation and can be combined with model parallelism by placing different experts on different devices.

• Sparsity: Only a fraction of the total parameters are used per input, keeping computational cost manageable.
• Load Balancing Challenge: Requires careful routing to ensure experts receive balanced workloads.
• Examples: Switch Transformers and GShard popularized MoE in language models, allowing for models with trillions of parameters.

A combined strategy that integrates Data, Pipeline, and Tensor Parallelism to train models at the extreme scale of hundreds to thousands of accelerators. It represents the state-of-the-art approach for training foundation models (e.g., 175B+ parameters).

• Hierarchical Partitioning: Data parallelism is applied across groups of nodes, pipeline parallelism across stages within a node, and tensor parallelism within a stage.
• Maximizes Scale: Allows the model size, batch size, and training speed to scale nearly linearly with the total number of devices.
• Orchestration Complexity: Requires sophisticated frameworks like DeepSpeed, Megatron-DeepSpeed, or Alibaba's Whale to manage the complex communication patterns.

A memory optimization technique for data parallelism that partitions the model states (optimizer states, gradients, parameters) across devices instead of replicating them. It eliminates memory redundancy, allowing for the training of much larger models with a fixed number of GPUs.

• ZeRO Stages: Stage 1 partitions optimizer states, Stage 2 adds gradient partitioning, Stage 3 adds parameter partitioning, requiring communication to gather parameters before use.
• Not Pure Model Parallelism: Works within a data-parallel group but uses model-parallel-like strategies for memory management.
• Impact: A cornerstone of the DeepSpeed library, enabling 10x larger model training on the same hardware.