Multi-GPU Training Explained: Model Sharding and Performance Trade-offs (Part 2)

Understanding Different Model Sharding Strategies for Accelerating Model Training in a Multi-GPU Environment

Photo by Javier Ortiz on Unsplash

This is the part two of our two-part blog series on multi-GPU model training. In the previous blog LINK, we learned how we can divide the input data and process them in parallel in multiple GPUs to accelerate training. It was predominantly used for training classical ML models and many smaller computer vision and natural language processing tasks. We also learned how we can modify it to context parallelism and use it for LLM training. In this blog, we will learn about model sharding. This method became vital as the models could no longer fit in the GPUs. Let’s get started.

Model Sharding

It is a set of methods that divide the model and spread it across GPUs. The core idea has been prevalent for a long time; GPipe came in 2018 — the concept became more mainstream with the advent of LLMs and VLMs in 2023 during the generative AI boom, when models became enormous with demand for more GPU RAM and compute, forcing researchers to heavily rely on model sharding for training. We will look at a few such methods, which are designed primarily for LLMs and VLMs.

Tensor Parallelism (TP)

It is one of the most granular forms of model parallelisation. It is used for model training when the model size is larger than the GPU memory. The core idea is to vertically divide the model’s layers and spread them across multiple GPUs. This method is primarily used for large language models (LLMs) as they have very large single layers like attention heads and feed-forward layers. It was popularised by NVIDIA’s Megatron-LM framework and became a core part of “3D parallelism” (combining Tensor Parallelism with data and pipeline parallelism).

Tensor Parallelism

Algorithm Breakdown

The primary function of any deep learning task is matrix multiplication, Y=XW, where X is the input and W is the weight matrix. We divide the weight matrix into multiple chunks, and each chunk is placed on a GPU. There are 2 primary ways to do it.

• Column parallelism: if W=[W1, W2, …, Wn] (vertical split into n columns) then Y=[XW1, XW2,…,XWn]. So, each GPU holds a slice of weights with the whole input and computes a partial output tensor, which is concatenated later to get the final output.
• Row parallelism: It is often used immediately after column parallelism to maintain the sequential flow. For example, if X=[X1, X2, …, Xn] (vertical split into n columns) and each GPU holds a horizontal slice of the input. Later, GPUs perform an all-reduce operation to sum their results together.

Tensor parallelism looks mathematically elegant and intuitive — but physically implementing it is altogether a different ball game. Unlike other multi-GPU training algorithms, it requires syncing multiple times within a single layer, and a huge amount of data needs to be moved for the All-reduce step. Due to this constraint, NVLink (designed primarily for GPU-to-GPU communication, with around 900 GB/s in the Hopper series) is essential for communication, as PCIe bandwidth (PCIe 6.0 is only 128 GB/s) is too low.

Pros

1. Can work for models with individual layers too large to fit in a single GPU
2. If implemented properly, GPU idle time is extremely low

Cons

1. Requires high-bandwidth interconnects like NVLink, which in turn requires specialised hardware like NVSwitch.
2. Hard to scale beyond 16 GPUs as communication overhead outweighs the computational speedup.
3. Only useful for very wide models.

Pipeline Parallelization

It is used when the model is too big to fit into a single GPU’s memory — but individual layers could easily fit into one GPU. The idea is to split the layers across multiple GPUs. It creates data flow like a pipeline where different computations are done at different stages and passed on to the next stage. It is primarily used by LLMs, VLMs, and other deeper models. It is often used in tandem with Tensor Parallelism by leveraging frameworks like PyTorch, DeepSpeed, Megatron-LM, etc.

Algorithm Breakdown

Model Sharding: The model is split into multiple stages/partitions, and each of these stages is loaded into a GPU. Each stage consists of a contiguous sequence of layers. For example, layers 1–10 in GPU 0, layers 11–20 in GPU 1, etc. The number of stages is generally equal to the number of GPUs, but it can be adjusted. Other model training parameters, like optimisation states, gradients, etc., are also partitioned per stage.

1. Forward Pass: Each stage of the pipeline processes the forward pass and passes the output to the next stage. While the next stage is processing the output, a new batch is being processed by the current stage to reduce GPU wait time. The process continues until all micro-batches (explained below) are done.
2. Backward Pass: Similarly, the backward pass processes the data in the reverse order of the forward pass. The backward pass overlaps with many micro-batches doing their forward pass, this implementation design helps to keep GPU utilisation high.
3. Parameter Update: Gradients from all micro-batches are accumulated until the backward pass of a full batch is completed. Then, each GPU updates its local parameters using the optimiser. No all-reduce is used, as each GPU is independent.

Primary Problem: Pipeline Bubble

In a naive implementation, if one sends one batch of data through the system, only one GPU will work at a time while others wait, causing a compute bubble (idle time). To solve this problem, micro-batching was introduced. We divide a large batch into smaller sub-batches, called a micro-batch, which are passed to the pipeline in sequential order. Due to the micro-batches being small, they are computed quickly and passed to the next stage, reducing waiting time for GPUs in the later stages of the pipeline.

Another primary bottleneck is intra-GPU communications. We need a very high bandwidth like NVLink, although PCIe can work as a good starting point. There is a similar model sharding mechanism called sequence parallelism; it works in tandem with Tensor Parallelism. It tries to parallelise the activation functions, batch normalisation, etc. along the depth of the model.

The strategy mentioned above (micro-batching) is one of the strategies; there are other, more advanced strategies like GPipe, PipeDream, etc.

Pros

1. Can train models that do not fit into one GPU
2. Can scale seamlessly for larger models
3. Lesser computational overhead as compared to Tensor Parallelism
4. One can make it work efficiently even if all GPUs are not identical

Cons

1. Pipeline Bubble (idle time): Micro-batching and other strategies mitigate them — but cannot reduce beyond a certain margin
2. Complexity of implementation: Though there are libraries for its off-the-shelf implementation, one needs to manually set different layers for different GPUs for optimal performance.
3. There can be a bubble anywhere in the pipeline, making it difficult to debug and load balance.
4. Intra-GPU bottleneck dependent

Expert Parallelization

It is a distributed training method designed primarily for Mixture of Experts (MoE) architectures. MoE architecture contains a base where attention is applied to each token, which is later routed to an expert, which predicts the output token. Experts are FFNs designed for a specific task; experts don’t talk with each other directly. Models like Google’s Switch Transformer, Mistral 8x7B, DeepSeek, etc. are popular open-source models that use MoE.

Algorithm Breakdown

1. Each MoE part is sharded from the model and put in a GPU.
2. Forward Pass: Input tokens arrive in each GPU after attention and routing. A forward pass is performed on each GPU and the output is passed to the master GPU. Then, the next tokens are processed.
3. Backward Pass: Similar to the forward pass — but in reverse order. Gradients are calculated locally and stored in the GPU.
4. It is often used in conjunction with Tensor Parallelism or data parallelism; in that case, additional steps will be included.

Pros

1. Enables large parameter model training with significantly fewer FLOPs
2. Memory efficient
3. High throughput

Cons

1. Intra-GPU bottleneck dependent
2. Load imbalance risks may result in bubbles

There are many other variants of model sharding methods like Zero Redundancy Optimizer (ZeRO), Parameter Server, etc., which are designed for specific use cases; we will be skipping them for the sake of brevity.

Key Takeaways

• Tensor Parallelism: It vertically partitions a model’s individual layers (such as Attention heads or Feed-Forward networks) across multiple GPUs. By splitting the weights of a single operation, it allows GPUs to compute partial results in parallel, enabling the training of massive layers that would otherwise exceed the VRAM of a single device.
• Pipeline Parallelism: It partitions a neural network horizontally by layers, assigning different stages of the model to separate GPUs in a sequential chain. To maximize efficiency, it splits training batches into smaller “micro-batches” that flow through the GPU pipeline concurrently, significantly reducing idle time through an overlapped execution schedule.
• Expert Parallelism: It is used in Mixture-of-Experts (MoE) models where different “expert” layers are distributed across separate GPUs while the shared network components are replicated. During training, a gating mechanism routes specific tokens to the relevant GPU-hosted experts, allowing for massive parameter scaling without a proportional increase in the computational load per device.

Thank you for reading!

From the Author: Apurva Bhatt

If you found this article insightful and beneficial, please consider following me and leaving a clap for more in-depth content! Your support helps me continue producing content that aids our collective understanding.