Model Parallelism

Tensor parallelism splits individual tensor operations (like matrix multiplications within a layer) across multiple devices. It operates at a finer granularity than pipeline parallelism, partitioning the weights and activations of a single layer.

• Key Mechanism: For a linear layer Y = XW, the weight matrix W is split along its column or row dimension. The input X is broadcast, partial matrix multiplications are performed in parallel, and the results are synchronized (e.g., via an All-Reduce operation).
• Primary Use: Critical for scaling the Feed-Forward Network (FFN) and Attention blocks within transformer models like GPT-3 and Llama, where a single layer's parameters exceed GPU memory.
• Example: In the Megatron-LM framework, tensor parallelism is used to split the multi-head attention and MLP layers across GPUs, enabling the training of models with hundreds of billions of parameters.

ZeRO is a memory optimization technique for data parallelism that eliminates memory redundancy by partitioning the model states (weights, gradients, optimizer states) across devices, fetching them as needed.

• ZeRO Stages:
  • Stage 1: Partitions optimizer states (e.g., momentum, variance) across devices, reducing memory proportional to the number of data parallel processes.
  • Stage 2: Partitions gradients in addition to optimizer states.
  • Stage 3 (ZeRO-Offload/Infinity): Partitions the model parameters across all devices. Parameters are gathered via communication before use and released afterward. ZeRO-Infinity extends this to offload partitioned states to CPU and NVMe memory.
• Impact on Parallelism: While primarily a data parallel technique, ZeRO Stage 3 enables a form of parameter parallelism. It is often combined with pipeline and tensor parallelism in frameworks like DeepSpeed to achieve optimal memory efficiency for large model training and inference.

PyTorch FSDP is a native training and inference technique that shards model parameters, gradients, and optimizer states across data-parallel processes. It is a hybrid approach that combines data and model parallelism.

• Core Mechanism: It shards each model parameter across the available devices, significantly reducing the per-device memory footprint compared to standard Data Parallel (DP).
• Communication: Employs an all-gather collective operation to reconstruct the full parameters for the forward or backward pass of a given layer, then uses reduce-scatter to aggregate gradients.
• Use Case: Primarily designed for training extremely large models that don't fit on a single GPU, but its sharding strategy is also applicable to memory-constrained inference scenarios.

TensorFlow supports model parallelism through manual device placement with tf.device and via higher-level libraries.

• Manual Placement: Developers can explicitly assign specific model layers or operations to different devices (e.g., /GPU:0, /GPU:1) using tf.device() context managers.
• Mesh TensorFlow: A library built on top of TensorFlow that generalizes distributed tensor computation. It allows developers to specify a logical mesh of processors and how tensors and computations are split across it, enabling sophisticated tensor parallelism and pipeline parallelism patterns.
• Use Case: Suitable for models with natural partitions (e.g., placing encoder/decoder on different devices) or for implementing custom, fine-grained parallelism strategies.

Megatron-LM is a persistent, research-focused framework developed by NVIDIA for training and serving large transformer language models with efficient model parallelism.

• Core Innovation: Implements tensor parallelism (intra-layer model parallelism), where matrix multiplications within a single transformer layer (e.g., the feedforward network or attention heads) are split across multiple GPUs.
• Communication Pattern: Heavily relies on high-bandwidth all-reduce operations between GPUs within a layer to combine partial results.
• Pipeline Parallelism: Also integrates pipeline parallelism inter-layer partitioning to scale to models with thousands of layers. It uses the GPipe schedule or the more memory-efficient interleaved schedule.
• Use Case: The de facto standard for parallelizing the largest class of transformer-based models (e.g., GPT, T5) across GPU clusters.

DeepSpeed is an optimization library that makes distributed training and inference easy, efficient, and effective. It provides multiple forms of model parallelism.

• ZeRO (Zero Redundancy Optimizer): A memory optimization technology that partitions optimizer states, gradients, and parameters across devices (similar to FSDP), eliminating memory redundancy.
• Pipeline Parallelism: Implements a robust pipeline parallelism scheduler that supports 1F1B (One Forward pass followed by One Backward pass) and other schedules for efficient training.
• 3D Parallelism: DeepSpeed seamlessly combines ZeRO-powered data parallelism, tensor-slicing model parallelism (like Megatron), and pipeline parallelism into 3D parallelism, enabling the scaling of models to trillions of parameters.
• Use Case: Essential for training and serving the world's largest models by combining all major parallelism strategies.

Hugging Face provides user-friendly abstractions for model parallelism, lowering the barrier to entry.

• Accelerate Library: Allows the same PyTorch code to run seamlessly on any distributed configuration (CPU, single GPU, multi-GPU, TPU). It supports device_map="auto" for large models, which automatically places model layers across available devices (CPU/GPU) to fit memory constraints—a form of naive model parallelism for inference.
• Transformers Integration: The transformers library has built-in support for loading very large models using frameworks like Accelerate and bitsandbytes, facilitating model-parallel inference out-of-the-box.
• Use Case: Ideal for practitioners who want to quickly load and run inference with large pre-trained models (like LLMs) without deep expertise in distributed systems.

Alpa is a system for automating the parallelization of large-scale neural networks, representing the next generation of model parallelism tools.

• Core Idea: Instead of manually choosing between data, operator, or pipeline parallelism, developers provide a single-device model definition. Alpa's compiler automatically generates an optimal hybrid parallelization plan tailored to the specific model and cluster configuration.
• Underlying Tech: Built on JAX and XLA, and uses Ray for cluster orchestration. It treats the parallelization problem as a joint optimization for intra-operator (tensor) and inter-operator (pipeline) parallelism.
• Use Case: For research teams and organizations that want to scale complex models without becoming experts in hand-tuning distributed execution plans.