Hybrid Parallelism for Training Large Models

Introduction

In our last lesson, we established that data, pipeline, and tensor parallelism each have distinct strengths and weaknesses. Data parallelism is simple but doesn't solve the model memory problem. Pipeline parallelism divides the model's memory footprint but introduces an efficiency-killing "bubble." Tensor parallelism also reduces memory needs and avoids the bubble, but demands extremely high-speed interconnects, making it impractical across machine boundaries.

No single strategy is sufficient for training the enormous models that define the state of the art. This lesson addresses the natural next question: how do we combine these strategies to leverage their strengths while mitigating their weaknesses?

We will explore hybrid parallelism, the standard practice for large-scale training. Specifically, we will focus on 3D parallelism, a powerful technique that organizes GPUs into a three-dimensional grid, with each dimension corresponding to data, pipeline, and tensor parallelism. You will learn the architectural reasoning for this approach and the design principles that guide its application in real-world systems like those used to train Llama 3 and other frontier models.

The Concept of 3D Parallelism

At its core, 3D parallelism is a "divide and conquer" strategy applied to the hardware itself. Instead of treating a cluster of GPUs as a flat, one-dimensional pool, we structure it as a 3D grid where a GPU's position in the grid determines its role.

Each axis of this grid corresponds to a different mode of parallelism:

1. Tensor Parallelism (TP) Axis: GPUs along this axis work together to execute a single model layer. They partition the layer's weight matrices and perform intra-layer parallel computations.
2. Pipeline Parallelism (PP) Axis: GPUs along this axis form a pipeline. Each GPU (or group of TP'd GPUs) is a stage, holding a different subset of the model's layers.
3. Data Parallelism (DP) Axis: This axis represents replication. Each "slice" of the PPxTP grid is a full replica of the model, processing a different chunk of the global data batch.

This structure allows us to map the parallelism strategy to the physical network topology of a GPU cluster. For a deeper understanding of this mapping and the formal definition of the 3D mesh, please study the following material.

The key insight is that this hybrid approach isn't random; it's a deliberate design choice that aligns with the communication costs we discussed previously.

• Tensor Parallelism, with its frequent, high-volume communication, is used within a node where GPUs are connected by high-bandwidth NVLink.
• Pipeline Parallelism, with its less frequent, smaller point-to-point transfers, is used across nodes connected by slower interconnects like InfiniBand.
• Data Parallelism, which requires an AllReduce of all gradients, is layered on top of the entire model replica (a TP x PP group), often across nodes that may have even higher latency between them.

A Design Heuristic for 3D Parallelism

Given a large model and a GPU cluster, how do you decide the size of each parallelism dimension (e.g., TP=8, PP=4, DP=32)? There is a well-established heuristic that directly follows from the trade-offs we've discussed.

The following video from a Stanford lecture on large-scale training provides a clear, practical rule of thumb for applying these strategies sequentially.

Let's formalize the rule of thumb you just heard:

1. Step 1: Apply Tensor Parallelism (TP) to Fit Layers.

   • Goal: Reduce the memory footprint of the largest layers so they fit on a single GPU.
   • Action: Increase the TP degree up to the limit of the high-speed interconnect domain (e.g., typically 8 GPUs within a node connected by NVLink). Stop here, as going beyond this incurs prohibitive communication costs.

2. Step 2: Apply Pipeline Parallelism (PP) to Fit the Model.

   • Goal: If the model, even with TP, is still too large for the aggregate memory of a TP-group of GPUs, we must partition it further.
   • Action: Increase the PP degree, distributing groups of layers across different nodes. You continue increasing the PP degree until one full replica of the sharded model fits within the memory of the TP_size * PP_size GPUs that constitute it.

3. Step 3: Apply Data Parallelism (DP) to Scale Throughput.

   • Goal: Now that the model fits, use all remaining GPUs to process more data in parallel and accelerate training.
   • Action: Use the remaining GPUs as data-parallel replicas. The DP degree will be Total GPUs / (TP_size * PP_size).

This hierarchy of decisions provides a robust starting point for configuring distributed training. It prioritizes solving the fundamental memory constraint first, then scaling for performance.

3D Parallelism in Practice: Case Studies

This heuristic isn't just theoretical; it's precisely how today's largest models are trained. Let's look at some real-world examples.

The following clip, from the same Stanford lecture, analyzes the parallelism strategies used to train several well-known large language models.

Let's summarize the key observations from these examples:

• Megatron-LM (1 Trillion Parameters): The configuration table showed that TP_size was capped at 8 (the intra-node limit). As the model grew, the PP_size was increased to accommodate the memory requirements, while the DP_size was adjusted based on the remaining available GPUs. This perfectly matches the heuristic.
• Llama 3: The training report for Llama 3 explicitly states the hierarchy: Tensor Parallelism (TP) and Context Parallelism (CP, a specialized form of sequence parallelism) are used for the highest bandwidth links. Pipeline Parallelism (PP) is used next. Finally, Data Parallelism (DP) is used over the remaining resources, as it's most tolerant of network latency.
• BLOOM (176B): The Hugging Face blog post (LINK) you read earlier describes training on 384 A100 GPUs. The final configuration used TP=4, PP=8, and DP=12 (4 * 8 * 12 = 384). The choice of TP=4 was likely due to the specific node hardware having 4-way NVLink connectivity. This again shows a configuration tailored to the hardware topology.

These examples demonstrate that 3D parallelism is not just one strategy but a flexible framework for composing our three fundamental parallelism techniques to match the demands of the model and the characteristics of the hardware cluster.

A Glimpse into Implementation

You don't need to implement these complex strategies from scratch. Frameworks like Megatron-DeepSpeed, PyTorch FSDP, and JAX handle the intricate details. However, given your background, it's useful to understand the core programming concepts.

A key idea is the device mesh, a logical representation of the 3D hardware grid. The programmer defines named axes for the mesh (e.g., 'data', 'pipe', 'tensor'). Then, data and model parameters are sharded by specifying which axes they should be split across.

The following tutorial provides a concise explanation of how this works in JAX, a popular framework for high-performance ML research.

The key implementation takeaways are:

• The system needs to manage multiple communication groups simultaneously (a TP group, a PP group, a DP group).
• Gradients must be synchronized across all three axes (data, pipe, tensor) before an optimizer step, as shown in the train_step function in the resource.
• Frameworks abstract this by allowing you to define sharding at the module or parameter level, then automatically generate the necessary communication and computation steps.

Conclusion

This lesson has demonstrated how to move beyond individual parallelism strategies to the powerful hybrid approaches used in modern large-scale AI. By composing data, tensor, and pipeline parallelism, we can overcome their individual limitations and achieve efficient training at massive scales.

Key Takeaways:

• 3D Parallelism is the standard hybrid strategy, combining Data (DP), Pipeline (PP), and Tensor Parallelism (TP) by organizing GPUs into a logical 3D grid.
• The mapping is topology-aware: TP is used for high-speed intra-node links, PP for lower-speed inter-node links, and DP is layered on top to utilize all available hardware.
• A common design heuristic is to apply parallelism in order of necessity: first TP to fit layers, then PP to fit the whole model, and finally DP to scale throughput.
• Real-world systems like those for training Llama 3 and BLOOM use precisely these hybrid strategies, with parallelism degrees tuned to the model size and hardware configuration.

Preview of the Next Lesson:

Throughout our discussion of data parallelism, we have assumed a synchronous update rule: all GPUs compute gradients, average them together with an AllReduce, and then apply the same update. This ensures consistency but can lead to straggler problems, where the entire system must wait for the slowest worker. In the next lesson, we will explore the alternative: asynchronous stochastic gradient descent (SGD), analyzing its trade-offs in terms of hardware utilization, system complexity, and model convergence.