[Paper Review] Cerebras-GPT: Open Compute-Optimal Language Models Trained on the Cerebras Wafer-Scale Cluster

Nolan Dey, Gurpreet Gosal+7 more|arXiv (Cornell University)|Apr 6, 2023

Topic Modeling24 citations

TL;DR

Cerebras-GPT trains compute-optimal GPT-like models from 111M to 13B parameters on the Pile dataset using Chinchilla scaling and Maximal Update Parameterization (µP), releasing open models and scaling laws that show state-of-the-art training efficiency for open models.

ABSTRACT

We study recent research advances that improve large language models through efficient pre-training and scaling, and open datasets and tools. We combine these advances to introduce Cerebras-GPT, a family of open compute-optimal language models scaled from 111M to 13B parameters. We train Cerebras-GPT models on the Eleuther Pile dataset following DeepMind Chinchilla scaling rules for efficient pre-training (highest accuracy for a given compute budget). We characterize the predictable power-law scaling and compare Cerebras-GPT with other publicly-available models to show all Cerebras-GPT models have state-of-the-art training efficiency on both pre-training and downstream objectives. We describe our learnings including how Maximal Update Parameterization ($μ$P) can further improve large model scaling, improving accuracy and hyperparameter predictability at scale. We release our pre-trained models and code, making this paper the first open and reproducible work comparing compute-optimal model scaling to models trained on fixed dataset sizes. Cerebras-GPT models are available on HuggingFace: https://huggingface.co/cerebras.

Motivation & Objective

• Combine recent LLM efficiency scaling techniques to create compute-optimal open models and scaling laws.
• Pre-train models on the Pile following Chinchilla data-efficiency rules to maximize performance per compute.
• Evaluate upstream (Pile) and downstream tasks to establish Pareto-optimal frontiers for open models.
• Demonstrate how Maximal Update Parameterization (µP) improves stability, accuracy, and hyperparameter transfer across scales.
• Release pre-trained models and code to enable reproducibility and community use.

Proposed method

• GPT-3-like autoregressive transformer decoder architecture with dense attention in all decoder blocks.
• Pre-train models from 111M to 13B parameters on the Pile dataset using 20 tokens per parameter for compute efficiency (Chinchilla rule).
• Use AdamW optimization with linear learning-rate decay, gradient clipping, and mixed precision (FP16/bfloat16 with bf16 preferred).
• Evaluate with cross-entropy on the Pile test set and compare against public models; re-normalize vocab-induced losses to GPT-2 vocabulary when needed.
• Investigate Standard Parameterization (SP) and Maximal Update Parameterization (µP) to assess stability and transferability; apply µTransfer to scale hyperparameters.
• Document training on the Andromeda Cerebras Wafer-Scale Cluster and release models and code on HuggingFace.

Experimental results

Research questions

• RQ1 What are the compute-efficient scaling laws for pre-training language models on the Pile when following Chinchilla-like data efficiency?
• RQ2 Do Cerebras-GPT models achieve state-of-the-art training efficiency on upstream and downstream tasks among open models of comparable size?
• RQ3 Does Maximal Update Parameterization (µP) improve stability, accuracy, and hyperparameter transfer across model scales?
• RQ4 How do compute-optimal models compare to fixed-dataset-size models in terms of downstream performance and efficiency?
• RQ5 What are the implications of balancing pre-training compute with inference compute for total deployment cost?

Key findings

• Cerebras-GPT models (111M–13B) on the Pile establish the compute-efficient Pareto frontier for pre-training and downstream tasks under 20 tokens per parameter.
• µP models show ~0.43% better Pile test loss and ~1.7% higher downstream accuracy on average compared to SP models, with more predictable scaling.
• The 13B model often achieves the best average downstream performance among comparably-sized open models on a diverse set of tasks.
• Pre-training on 20 tokens per parameter yields compute-efficient scaling consistent with DeepMind’s Chinchilla findings, and compute-frontier extrapolations align with larger-scale expectations.
• When considering training plus inference costs, Cerebras-GPT variants can be more cost-effective than fixed-token baselines up to substantial inference token budgets (e.g., ~200B tokens).
• The paper releases pre-trained weights and code (HuggingFace, Cerebras Modelzoo) and provides detailed µP guidance for stable scaling.