Model Parallelism

This is the most common form of model parallelism, where sequential layers of a neural network are distributed across different devices. For example, in a 100-layer transformer, layers 1-50 might be placed on Device A and layers 51-100 on Device B. The activations (the output of one layer) must be communicated to the device holding the next layer in the sequence. This strategy is straightforward but can lead to significant idle time (bubbles) as devices wait for data from preceding stages, especially in synchronous execution.

This strategy splits individual tensor operations, such as large matrix multiplications, across devices. For a linear layer Y = XW + b, the weight matrix W can be partitioned:

• Column-wise (Split along output features): Each device computes a portion of the output channels.
• Row-wise (Split along input features): Requires an all-reduce operation to combine partial results. This method is essential for models with massive layers (e.g., large feed-forward networks in transformers) that exceed a single device's memory. It is often combined with data parallelism for maximum scalability.

Pipeline parallelism is a hybrid strategy that combines layer-wise partitioning with a scheduling technique to improve hardware utilization. The model is split into stages (groups of layers), each assigned to a device. Instead of processing one sample at a time, the system processes a stream of microbatches. While Device 2 processes the first microbatch through its stage, Device 1 can begin processing the second microbatch. This overlaps computation across devices, reducing idle time. The pipeline bubble—the time spent filling and draining the pipeline—remains a key performance challenge.

A specialized strategy for sparsely-activated models like Mixture of Experts (MoE). In an MoE layer, the model has many sub-networks ("experts"), but for a given input token, only a small subset (e.g., 2 out of 128) are activated. Experts are distributed across devices. The implementation requires:

• A gating network to select experts per token.
• An all-to-all communication operation to route tokens to the devices hosting their selected experts.
• Another all-to-all to gather the processed tokens. This allows for models with trillions of parameters while keeping the computational cost per token manageable.

The efficiency of model parallelism is dictated by inter-device communication. Key patterns include:

• Point-to-Point: Sending activations/gradients between specific devices (common in layer-wise).
• Collective Operations: All-reduce (summing gradients across devices) and all-gather (collecting partitioned tensors) are critical for tensor and data-parallel hybrid setups.
• Synchronization Points: Devices must often synchronize via barriers to ensure correctness, creating performance bottlenecks. Optimizations like overlapping communication with computation (using non-blocking operations) are essential to hide latency.

Implementing model parallelism manually is complex. Major frameworks provide abstractions:

• PyTorch: torch.distributed with FullyShardedDataParallel (FSDP) for hybrid data/model parallelism and PipelineParallel for pipeline strategies.
• TensorFlow/Mesh TensorFlow: Declarative APIs for specifying tensor partitions across a device mesh.
• Megatron-LM (NVIDIA): A specialized library for efficient tensor and pipeline parallelism of large language models, providing optimized kernels for communication.
• DeepSpeed (Microsoft): Offers ZeRO-Offload and 3D parallelism (combining data, tensor, and pipeline parallelism) for extreme model scale. These tools automate gradient synchronization, loss calculation, and optimizer steps across partitions.

Unlike data parallelism, which replicates the entire model and splits the dataset, model parallelism splits the model itself. The primary goal is to overcome memory limitations. Common partitioning strategies include:

• Layer-wise (Pipeline) Parallelism: Assigning different layers or groups of layers to different devices.
• Tensor (Intra-layer) Parallelism: Splitting individual tensor operations (e.g., a large matrix multiplication) across devices.
• Expert Parallelism: Used in Mixture-of-Experts (MoE) models, where different "expert" sub-networks are placed on different devices. The choice depends on the model architecture and the communication cost between devices.

Model parallelism is driven by the exponential growth of model parameters, which has far outstripped the memory capacity of individual accelerators.

• Memory Walls: A single NVIDIA H100 GPU has 80GB of HBM. Modern LLMs like GPT-4 or Claude 3 Opus have parameter counts in the hundreds of billions, requiring terabytes of memory for training.
• Specialized Hardware: NPUs and other accelerators often have constrained on-chip memory (SRAM) compared to GPU HBM, making intra-chip model partitioning critical for large layers.
• Interconnect Bottlenecks: The efficiency of model-parallel training is gated by the bandwidth of inter-device links (e.g., NVLink, InfiniBand).

Implementing model parallelism requires deep integration with the deep learning framework's execution engine.

• PyTorch: Offers torch.nn.parallel.DistributedDataParallel for data parallelism and more manual APIs (e.g., torch.distributed.rpc) for model parallelism. Frameworks like FairScale and DeepSpeed (with its ZeRO-3 optimizer) provide advanced automated model-parallel strategies.
• TensorFlow/Mesh TensorFlow: Google's Mesh TensorFlow allows users to specify a layout for tensors across a mesh of devices, abstracting the parallelism.
• JAX: With its pjit (parallel jit) and shard_map primitives, JAX allows explicit specification of how arrays are sharded across hardware, enabling sophisticated model-parallel layouts.
• Megatron-LM (NVIDIA): A seminal framework for efficient tensor-model-parallel training of large language models.

Splitting the model introduces new communication points that dominate performance if not managed.

• Forward Pass: Activations must be sent from the device holding layer N to the device holding layer N+1.
• Backward Pass: Gradients must be passed backwards through the same chain. This creates a pipeline bubble in naive implementations.
• Optimizer Step: With parameters distributed, optimizer states may also be sharded (as in ZeRO-3).
• All-Reduce vs. Point-to-Point: Tensor parallelism often uses all-reduce collectives to combine partial results, while pipeline parallelism uses point-to-point sends/receives. Overlapping communication with computation is critical.

In practice, model parallelism is almost always combined with other forms to train massive models efficiently. The state-of-the-art approach is 3D Parallelism, popularized by DeepSpeed and Megatron:

1. Data Parallelism (DP): Replicates the model across groups of devices, each processing a different data batch.
2. Pipeline Parallelism (PP): Splits model layers vertically across devices within a DP group.
3. Tensor Parallelism (TP): Splits individual layers horizontally across a subset of devices. This combination allows scaling to thousands of GPUs/NPUs by balancing memory savings (PP, TP) with statistical efficiency (DP).

Primary Use Cases:

• Training Large Language Models (LLMs) and Vision Transformers: Essential for models with >10B parameters.
• Inference for Massive Models: Deploying a model too large for a single device chip.
• Leveraging Heterogeneous Hardware: Placing different parts of a model on hardware optimized for specific operations (e.g., attention on NPU, embeddings on CPU).

Key Trade-offs:

• Complexity: Dramatically increases system and code complexity.
• Communication Overhead: Can become the performance bottleneck.
• Load Imbalance: Inefficient if partitions are not computationally balanced.
• Reduced Device Utilization: Idle time due to pipeline bubbles or synchronization points.

A parallel computing paradigm where the same model is replicated across multiple processors or devices, with each replica processing a different subset of the training data. Gradients are then synchronized (e.g., via all-reduce) to update a global model. This is the most common strategy for scaling batch training.

• Key Mechanism: Replicated model, partitioned data.
• Primary Goal: Scale training by increasing effective batch size.
• Typical Use: Training models that fit on a single device but require larger batches.

A strategy that partitions a model's sequential layers across multiple devices. Different devices process different layers for a continuous stream of data microbatches, forming an execution pipeline. This technique is essential for models with long sequential dependencies, like large transformers.

• Key Mechanism: Vertical splitting of the model graph (layer-by-layer).
• Primary Goal: Handle models too tall (deep) for a single device's memory.
• Challenge: Pipeline bubbles caused by idle stages during startup and wind-down.

A fine-grained form of model parallelism that splits individual tensor operations (e.g., large matrix multiplications within a layer) across multiple devices. For example, the weight matrix of a linear layer can be partitioned column-wise or row-wise, with communication required to combine partial results.

• Key Mechanism: Horizontal splitting of individual layers/operators.
• Primary Goal: Distribute the compute and memory load of massive individual layers.
• Typical Use: Enabling extremely wide layers, such as large feed-forward networks in transformers.

A parallel computing model where different, independent tasks or functions are executed concurrently on multiple processing units. Unlike data or model parallelism, the tasks may operate on different data and execute different code. This is common in heterogeneous workflows.

• Key Mechanism: Concurrent execution of distinct functions.
• Primary Goal: Increase throughput by exploiting functional independence.
• Example: A server concurrently handling inference requests, model pre-processing, and logging.

The combined application of multiple parallelism strategies (e.g., data + pipeline + tensor) to train models at extreme scale. This is necessary for modern foundation models that exceed the memory and compute capacity of any single device or simple strategy.

• Key Mechanism: 3D parallel composition (data, pipeline, tensor dimensions).
• Primary Goal: Achieve optimal utilization across thousands of accelerators.
• Framework Example: Megatron-LM uses tensor parallelism within a node and pipeline parallelism across nodes.

Low-level mechanisms that coordinate execution and memory state across parallel threads or processes. They are the foundational building blocks for implementing higher-level parallelism strategies.

• Barrier: Forces all threads to reach a point before any proceed.
• Atomic Operations: Indivisible read-modify-write ops (e.g., Compare-and-Swap).
• Mutex/Semaphore: Enforce mutual exclusion or limit concurrent access.
• Memory Fence: Enforces ordering of memory operations for consistency.