Accelerate ND-Parallel: A guide to Efficient Multi-GPU Training

Training large models across multiple GPUs can be challenging due to the complexities of different parallelism strategies. In Accelerate, together with Axolotl , we have integrated a quick and easy way to use any combination of parallelism strategies in your training script!

We've also included a more comprehensive end-to-end training script in the Accelerate repo which demonstrates how to setup your dataloader, optimizer, and training loop, and how to save your model after training.

To further streamline fine-tuning models at scale and compose parallelism strategies with a variety of fine-tuning techniques, we've also integrated this technique into Axolotl. To help you get started right away we've tested some example configs which you can modify to suit your needs - try one out with:

We've made it easy to configure the degrees of different parallelism strategies and how they are combined through the ParallelismConfig class in Accelerate, or through config fields in Axolotl, but how do we know which configuration will work best for our use case? As we scale to training models with tens or even hundreds of billions of parameters, the primary challenge comes from understanding the different parallelism strategies and how they interact to minimise communication overhead across devices. In this post, we'll walk through how the different parallelism strategies work, and when and how you might want to compose them.

Data parallelism (DP) is the most common technique for training models across multiple GPUs, and involves replicating the model, gradients and optimizer states across each device, whilst evenly distributing data batches between GPUs, and synchronising gradients across devices before updating parameters. This can significantly increase throughput compared to single-device training, but requires that your model is able to fit on a single device.

We can control the number of replicas of the model with the dp_replicate_size parameter in Accelerate's ParallelismConfig or config field in Axolotl. It's worth noting that DP is a top-most-level parallelism strategy, meaning that if we use dp_replicate_size=2 and we compose it with other parallelism strategies, there would be 2 replicas of the model, each also influenced by the other parallelism strategies. For example, if we use dp_replicate_size=2 and tp_size=2 , we would have 2 replicas of the model, each with 2 tensor parallel shards.

What if our model is too large to fit on a single device? Fully sharded data parallel (FSDP) addresses this issue by sharding (distributing evenly) the model’s weights, gradients, and optimizer states across GPUs (this is inspired by DeepSpeed’s ZeRO-3), whilst each device still receives its portion of the full batch of data. As you may notice from the diagram above, rather than requiring a full copy of the entire model on each device, we only gather the weights for a single layer at a time before the forward pass, after which the weights may be sharded again.

You can also check out the Axolotl ND-Parallelism docs for more details - adding ND parallel techniques to your existing configs is as simple as adding one or more of the following fields to your Axolotl config file:

We use the term shard to describe data on a single device which is a partition of a larger piece of data.

We use the term node to refer to a single machine which hosts multiple GPUs (up to a maximum of 8), with fast intra-node communication channels using e.g. NVLink between GPUs. When using multiple nodes for training, we rely on relatively slower inter-node communication channels between machines using e.g. Infiniband. We also refer to the total number of devices in the process pool as the world size - e.g. a single node with 8 GPUs represents a world size of 8, and 4 nodes would represent a world size of 32.