Efficient Large-Scale GPT Training Using a Cerebras Wafer-Scale Cluster

State-of-the-art language models are extremely challenging to train: they require huge compute budgets, complex distributed compute techniques and large teams with deep expertise in machine learning and parallel programming. At Cerebras, we have built a simple, data-parallel weight streaming architecture that allows users and researchers to get high quality models efficiently and with push-button scaling. This blog highlights the ease of using Cerebras to train large-scale GPT models and reproduce a state-of-the-art external baseline.

Introduction

Large Language Models (LLMs) like OpenAI’s GPT-3 [3] show impressive language understanding and text generation capabilities. However, the ability and expertise to train these models from scratch and at scale are limited to a handful of large organizations. Training language models from scratch has a range of benefits, including data privacy, domain customization, and finer-grained control over model updates. Today, there are many challenges to training LLMs at scale, including and not limited to intricate modeling issues (e.g., activation checkpointing, 3D parallelism for training) and hardware and system issues (e.g., hardware failures, multi-node orchestration).

At Cerebras, our mission is to reduce the cost of curiosity for our customers. Anyone should be able to harness the power of large language models, whether general-purpose models or specialized domain models in fields such as genomics [8], and not worry about spending months setting up infrastructure or hiring large teams with deep technical expertise to fulfill their business requirements. With this in mind, we designed our weight streaming execution mode from the ground up to build on the flexibility of the underlying architecture. For the ML practitioner, the Cerebras AI Model Studio cloud service provides simple push-button access for anyone to the performance of the Cerebras Wafer-Scale Cluster.

This blog highlights how easy it is to use Cerebras to train large-scale GPT models and reproduce a state-of-the-art external baseline. As a candidate for our replication, we picked a configuration from the exhaustive study done by the BigScience Architecture and Scaling group [1], which consists of authors from multiple organizations, including Hugging Face, EleutherAI, New York University and the Allen Institute for AI. In that study, the authors evaluated the impact of various modeling choices on the zero-shot performance of GPT models. In our work, we replicated the pre-training of a GPT-3 XL (1.3 billion parameters) model on the Pile [2] dataset and reported the zero-shot evaluation accuracy across 24 downstream tasks.

Scaling Pre-Training on a Cerebras Wafer-Scale Cluster

Following the setup and hyperparameters from BigScience [1], we pre-trained a GPT-3 XL (1.3B) model on the Pile [2] dataset for 112B tokens. The model uses learned position embeddings with GELU [9] activation and no layer norm [10] after the embedding layer. An implementation of the model can be found in our public Cerebras Model Zoo repository.

While a single CS-2 system can seamlessly pre-train GPT models up to 175 billion parameters, to speed up training for our experiment, we leverage a Cerebras Wafer-Scale Cluster equipped with multiple CS-2 systems. We scaled pre-training efficiently on an 8 x CS-2 cluster without being forced to worry about the significant complexities of different parallelism strategies on GPUs. The remarkable ease of scaling is shown in Figure 1. A more detailed discussion of our scaling properties can be found in this blog post.

Zero-Shot Generalization

We evaluated the quality of our pre-trained model using following the tasks and setup described in BigScience [1]. Table 1 reports the average accuracy using the EleutherAI (EAI) harness [6] results across 24 tasks. Note we did not evaluate the model on SST-2 [7] since the evaluation metric on this task exhibits high variance across different pre-training setups (see Appendix E in [1]).

Figure 2 shows a detailed comparison of each task’s accuracy between the BigScience [1] reference and our trained model. Our zero-shot results are comparable to the results (within standard error) from the reference for most . We note that our results are different from the published metrics for some tasks. This is primarily due to minor differences in our pre-training setup, such as training on a different subset of 112B tokens.