The Complete Guide to Scale-Out on Cerebras Wafer-Scale Clusters

Abstract

Training large ML models at scale is difficult because of the complexities involved in 3D parallelism techniques. Moreover, standard scaling methods suffer from bottlenecks at various points in the system, and therefore are unable to reach linear scaling.  In this article, I’ll show how Cerebras has solved these challenges through “appliance mode”, enabling push-button linear scaling on Cerebras Wafer-Scale Clusters of up to 192 CS-2 systems.

Picture this: you’re sitting at your desk, it’s the Friday before a long weekend, and you need to kick off a training run for a GPT-style large NLP model.  It’s a multi-billion parameter model, and you’ve just started working on getting it to train on a NVIDIA cluster. The command used to run training is arcane, and your model code is filled with details about chunk sizes and GPU mappings for pipeline parallel. You pray that there are no bugs.  The hours march on.

Meanwhile, your colleague, who used a Cerebras Wafer-Scale Cluster, has long gone home.

Getting large language models to train at scale is no easy feat on GPUs.

Often, for experimentation, models are written to train on a single server multi-GPU setup.  Once it’s time to scale out, in order to speed up training, multiple servers are involved, and things get tricky. Inter-process communications libraries have to be considered. We have to use batch-scheduling systems. And, because model sizes don’t fit on a single device, and even individual layers don’t fit on a single device, we have to resort to complex training strategies like pipelined or tensor model parallelism (Figure 1).  For a deeper dive on parallelism strategies, I highly recommend this post on our developer blog by my colleagues Natalia Vassilieva.

Many papers have been written about efforts to try to make distributed cluster training faster, such as this one by Narayanan et al which describes their efforts to train Megatron-LM across more than 3,000 GPUs (!) by using complex interleaving techniques, but latency at each stage can cause bottlenecks in the entire system. Moreover, debugging latency in such a setup is time-consuming and especially so as the number of systems increases.

Figure 2, which is drawn from the Microsoft DeepSpeed blog, elegantly illustrates the complexity involved in training large models on GPUs. It is necessary to leverage 3-D parallelism, data, model and pipeline parallel. Training is limited by the slowest speed between any of the elements, and this is the source of much difficulty in distributed training.

Scaling models is turning out to be a full-time engineering problem for many organizations. Standard approaches include using the Torch Pipe API, which allows out-of-the-box pipeline parallelism. Unfortunately, even something simple like adding skip connections becomes difficult. There are limitations to input-output shapes, and conditional flow is impossible, as described in this very clear article from Hugging Face.

Even once pipeline parallel strategies are in place, problems at the cluster-scale arise.  If your cluster is in the cloud, it’s difficult to ensure that the connectivity between machines is the same, and speed of training will be bottlenecked by the slowest connection. This is, incidentally, why very large NLP models are typically trained using dedicated hardware, outside of traditional cloud structures.

These issues make it difficult to even start training large language models, and often, tasks like hyperparameter tuning become even more challenging due to the brittleness of the system.  These systems make debugging training latency and restarting from checkpoints a difficult process.