PaperlayerView exact revision on arXiv

Source-linked AI summary

Cerebras-GPT: Open Compute-Optimal Language Models Trained on the Cerebras Wafer-Scale Cluster

Nolan Dey, Gurpreet Gosal, Zhiming, Chen, Hemant Khachane, William Marshall, Ribhu Pathria, Marvin Tom, Joel Hestness

arXiv:2304.03208v1cs.LGcs.CL

TL;DR

Open models had not focused on compute-efficient scaling, leaving a need for reproducible comparisons to guide foundation-model training. This paper trains Cerebras-GPT models on the Pile using Chinchilla scaling and reports strong efficiency across pre-training and downstream objectives, while releasing models and code.

  • Problem

    Prior open models had not targeted compute efficiency, and the research community needed reproducible scaling efforts to guide compute-efficient training decisions for large foundation models.

    Source: §1 · p. 1

  • Method

    The authors train an open family of 111M-to-13B-parameter Cerebras-GPT models on the Pile following Chinchilla scaling rules, and also test µP for scaling stability and hyperparameter transfer.

    Source: §1 · pp. 1–2, §3 · p. 4

  • Results

    Cerebras-GPT defines state-of-the-art compute-optimal Pareto frontiers for pre-training and downstream objectives, and µP improves compute-optimal frontier loss by 0.4%.

    Source: §3 · p. 4

  • Takeaways & Limitations

    The authors release pre-trained models, code, and training details to support community use and reproduction of their results.

    Source: §8 · p. 13

    The models have not been extensively tested in downstream or real application settings, including for factual accuracy, profanity, toxicity, or other socially undesirable text generation.

    Source: §7 · p. 13

Abstract

from arXiv · show

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 ($\mu$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.

1 Introduction

Cerebras-GPT combines compute-efficient scaling techniques in an open model family and reports strong efficiency across pre-training and downstream objectives. The paper also presents scaling laws, documents its training experience, and releases models and code for reproducibility.

1 Introduction

  • Earlier open models had not targeted compute efficiency, motivating reproducible scaling efforts to guide compute-efficient training decisions for large foundation models.Source: §1 · p. 1
  • The open Cerebras-GPT family spans 111M to 13B parameters and is trained on the Pile following DeepMind’s Chinchilla scaling rules.The models are presented as compute-optimal and intended to produce scaling laws.Source: §1 · pp. 1–2
  • Cerebras-GPT models form compute-optimal Pareto frontiers for pre-training and downstream objectives, with scaling laws to estimate benefits from further model and dataset scaling.The authors recommend considering pre-training and expected inference compute together when choosing model and dataset sizes.Source: §1 · p. 2
  • Cerebras-GPT shows state-of-the-art training efficiency across pre-training and downstream objectives, while its largest model achieves state-of-the-art performance on most downstream tasks against similarly sized public models.The contribution summary describes this as the first such open effort.Source: §1 · pp. 1–2
  • The work documents training Cerebras-GPT on the Andromeda AI Cluster, which comprises 16 Cerebras CS-2 systems.Source: §1 · p. 2
  • The authors release pre-trained models and code and share training details so the community can use and reproduce the results.Models are available on HuggingFace and code in the Cerebras Modelzoo.Source: §1 · p. 2

2 Methodology

The methodology combines GPT-3-like model designs and Pile pre-training with Chinchilla-style compute-efficient scaling. It also contrasts standard parameterization with µP, which is used to improve stability and transfer hyperparameters across model scales.

2.1 Model Architecture

  • Cerebras-GPT uses an autoregressive GPT-3-like decoder with dense attention in every block, model dimensions based on aspect ratio or GPT-3 shapes, and a maximum sequence length of 2048 tokens.Source: §2.1 · p. 3

2.2 Pre-training Corpus

  • Pre-training uses the Pile’s provided train, test, and validation splits and GPT-2 byte-pair encoding; evaluations use Pile test cross-entropy and are run on checkpoints for fair comparison.The Pile combines 22 data sources, and the authors do not deduplicate it.Source: §2.2 · p. 3

2.3 Model Training

  • The training setup follows Chinchilla scaling, finding roughly 20 tokens per parameter most compute-efficient, and uses linear learning-rate decay in most runs after a 375M-token warmup.For the 13B model, batch size rose from 720 to 1080 sequences after 84B tokens when the validation–training loss gap grew.Source: §2.3 · pp. 3–4
  • Released Cerebras-GPT models use bfloat16, which the authors found more stable than FP16 because of its extra exponent range.Source: §2.3 · p. 4

2.4 Standard (SP) and Maximal Update Parameterization (µP)

  • The main models use standard parameterization, but SP can become unstable at large scales as weights and activations approach numerical representation limits, potentially requiring costly restarts or tuning.Source: §2.4 · p. 4
  • µP controls initialization, layer-wise learning rates, and activation magnitudes to stabilize training across widths and improve transfer of hyperparameters from smaller to larger models.The authors tune a 40M-parameter µP model and transfer settings along the µP scaling law up to 2.7B parameters, with changes to activation scaling, initializers, and layer-wise learning rates.Source: §2.4 · p. 4

3 Results

Cerebras-GPT models establish compute-optimal frontiers for pre-training and downstream objectives, with scaling laws that characterize these results. The section also reports that µP improves scaling stability, predictability, and performance.

3 Results

  • Cerebras-GPT defines state-of-the-art compute-optimal Pareto frontiers on pre-training and downstream objectives, and the study releases an openly reproducible Pile scaling law.The models span 111M–13B parameters.Source: §3 · p. 4

3.1 Pre-training Results

  • Across model scales, Cerebras-GPT sets the Pile-loss efficiency frontier; its 13B model falls within 0.5% of projected loss.The 13B loss projection was estimated using a scaling law fit to models up to 6.7B parameters.Source: Figure 2 · p. 5, §3.1 · p. 5
  • Cerebras-GPT’s roughly 20 tokens per parameter is compute-optimal on the Pile; Pythia’s 12B model is 0.3% above the frontier at 25.3 tokens per parameter.Loss degradation varies predictably with tokens per parameter.Source: Figure 3 · p. 5, §3.1 · pp. 5–6

3.2 Downstream Results

  • Cerebras-GPT forms the downstream accuracy frontier against FLOPs, while smaller parameter-matched GPT-J, OPT, and Pythia models score higher than Cerebras-GPT.The Cerebras-GPT parameter-size trend continues, suggesting competitiveness with GPT-NeoX 20B if scaled to that size.Source: Figure 4 · p. 6, §3.2 · pp. 6–7
  • Deduplicated Pythia models typically improve downstream accuracy by 1.8% on average, though their Pile test-loss results may not be directly comparable.The passage attributes the comparison caveat to differences in dataset preparation.Source: §3.2 · p. 7, Table 2 · p. 7

3.3 Maximal Update Parameterization (µP) and µTransfer

  • µP models average 0.43% better Pile test loss and 1.7% higher downstream accuracy than SP models, with lower loss-scaling variability.Relative to the SP scaling-law fit, µP models have 0.04% standard deviation versus 0.66% for SP models; downstream results are robust across scales except at 2.7B.Source: §3.3 · pp. 7–8, Figure 5 · p. 8
  • µP stabilizes training across scales and lets learning-rate hyperparameters tuned on a 40M model transfer to all model sizes.With µP, weight and gradient norms behave similarly at different scales.Source: §3.3 · p. 7

4 Trading Off Training and Inference FLOPs

The paper extends compute-optimal model selection to account for both pre-training and expected inference costs. Under this combined metric, most Cerebras-GPT models compare favorably with Pythia up to roughly 200B inference tokens.

4 Trading Off Training and Inference FLOPs

  • Inference cost scales with model size and number of inferences, giving smaller models an overall inference-cost advantage proportional to their size.Earlier analysis considered pre-training compute, which is proportional to squared model size when models train on a constant number of tokens per parameter.Source: §4 · pp. 8–9
  • The proposed total-cost metric adds pre-training FLOPs to per-token inference FLOPs multiplied by expected inference tokens.The authors use it to estimate when inference costs bring total compute in line with models trained on different token counts.Source: §4 · p. 9
  • 200B inference tokens: most Cerebras-GPT models provide better Pile test-loss-per-compute-FLOP than Pythia until all models reach roughly this level.Models between the Cerebras-GPT and Pythia frontiers may achieve better loss for the same total compute budget.Source: Figure 6 · p. 9, §4 · p. 9
  • The combined-cost formulation can help organizations and governments assess training budgets and can also be applied to monetary, energy, or carbon-footprint costs.A model trained with too many data samples may require many inferences before its pre-training cost is amortized.Source: §4 · p. 9

5 Cerebras Stack

The Cerebras-GPT scaling study ran on Andromeda, a 16-system Cerebras cluster, using a software stack designed to simplify model training and scaling. The cluster showed near-linear performance scaling and less than 8% utilization deviation across model scales.

5 Cerebras Stack

  • Andromeda’s near-linear performance scaling across its 16 CS-2 systems supported the study’s compute-efficient scaling-law experiments.The study describes the cluster as capable of high-performance scale-out and reports scaling its training runs to the full system count.Source: §5 · p. 9

5.1 Andromeda AI Supercomputer

  • Andromeda is a 16-CS-2 Cerebras cluster; each system has a WSE-2 with 40 GB of high-bandwidth SRAM and 7.5 PFLOP/s peak half-precision throughput.The cluster's stated peak throughput is 120 PFLOP/s.Source: §5.1 · p. 10

5.2 CSoft Platform and Weight Streaming Mode

  • CSoft automates training configuration, while Weight Streaming moves weights to the wafer and gradients back without model parallelism.The authors report testing models beyond GPT-3's 175B parameters with no changes outside model configurations.Source: §5.2 · p. 10

Our Language Model Scaling Experience

  • The authors report that Weight Streaming made model development and scaling easier, with single-system tests, near-linear cluster scaling, and system-count-only configuration changes.CSoft automatically selects data-parallel configurations for larger runs.Source: Our Language Model Scaling Experience · p. 11

5.3 Performance Scalability

  • Within 9%: weak-scaling performance was linear across all tested model sizes and CS-2 system counts, up to 16 systems.Table 4 reports weak-scaling performance relative to one CS-2; strong-scaling tests also showed consistent scalability for batch sizes used in training.Source: Table 4 · p. 11, §5.3 · p. 11, Table 5 · p. 11
  • Less than 8%: Andromeda’s FLOP/s utilization deviated by under 8% across model scales, indicating consistent performance across model and batch sizes.Table 6 reports utilization relative to a 111M-parameter model running on one CS-2; larger values mean higher utilization.Source: §5.3 · p. 11, Table 6 · p. 12

6 Related Work

Prior scaling studies often increased model size without correspondingly scaling dataset size, while many influential large-model results and training data were not publicly available. More recent work emphasized compute-optimal data scaling, but open and reproducible evidence remained limited.

6 Related Work

  • Many large-model studies scaled model size while training on roughly 300B tokens, and their datasets and resulting models were not publicly released.Such models improved downstream performance but could be trained only by select organizations with large compute clusters.Source: §6 · p. 12
  • Open releases such as GPT-J, GPT-NeoX, OPT, Pythia, and BLOOM made models and datasets available for community testing and comparison.These releases were typically smaller than the largest models described in prior work.Source: §6 · p. 12
  • Hoffmann et al. (2022) found that compute-optimal pre-training scales dataset size linearly with model parameters, but their dataset and models were not publicly available.LLaMa also trained large models on open datasets, though most were trained with more data than compute-optimal for their size.Source: §6 · p. 12
  • µP analytically controls width-related training instabilities and lets small-model optimal hyperparameters transfer to very large models.The authors say its comprehensive approach simplifies training efforts.Source: §6 · p. 12

7 Limitations

The paper identifies unexamined modeling choices and limited downstream and application testing as boundaries of its work. It also notes that additional safety-related evaluation and mitigation are needed before presenting the models to users.

7 Limitations

  • The study uses established architectures but does not explore newer architectural features, downstream-task tuning procedures, or dataset-cleaning approaches.The authors identify position embeddings, activation functions, denoising objectives, instruction fine-tuning, and further dataset cleaning as areas for future exploration.Source: §7 · p. 13
  • The models have not been extensively tested on downstream tasks or real applications, including factual accuracy, profanity, and toxicity.The authors call for further safety testing, mitigations, and output curation before user presentation.Source: §7 · p. 13

8 Conclusion

Cerebras-GPT combines compute-optimal pre-training with µP and reports efficient scaling from 111M to 13B parameters. The authors release models, code, and reproduction instructions, alongside a model card for the 13B model.

8 Conclusion

  • Cerebras-GPT models are reported to achieve state-of-the-art pre-training efficiency on pre-training and downstream objectives, with µP further improving scaling results.The family spans 111M to 13B parameters and was trained compute-optimally on the Pile.Source: §8 · p. 13

Cerebras-GPT Open-Source References

  • The paper provides reproduction instructions and releases pretrained Cerebras-GPT checkpoints and code for community use.The initial release includes seven standard-parameterization models from 111M to 13B parameters under Apache 2.0.Source: §8 · p. 13, Cerebras-GPT Open-Source References · p. 19

Model Card

Author Contributions and Acknowledgements

  • The manuscript credits contributors for experimental strategy, training, data preparation, evaluation, writing, and project leadership, and thanks additional Pile contributors and reviewers.Cerebras engineers are also acknowledged.Source: Author Contributions and… · p. 19

A.1 Pile Dataset Preprocessing

The preprocessing and stability appendix describes dataset-wide shuffling and numerical choices used to train larger models reliably. It reports improved validation loss from shuffling across documents and identifies mixed-precision underflow and Adam epsilon as stability concerns.

A.1 Pile Dataset Preprocessing

  • Dataset-wide shuffling of Pile training documents improves validation loss by 0.7–1.5% over aggressive shuffling within windows of contiguous documents.The tokenized training set contains roughly 371B tokens.Source: §A.1 · p. 19
  • The Pile is known to contain toxic, gender-biased, pejorative, and racially sensitive content.Source: §A.1 · p. 19

A.2 Ensuring Stable Training

  • For models larger than roughly 1.3B parameters, loss scaling alone did not prevent FP16 instability: large gradients can force smaller scales, leaving other gradients vulnerable to underflow.Underflow can yield missing or low-precision gradients and exacerbate instability.Source: §A.2 · pp. 19–20
  • The authors recommend carrying attention softmax gradients in FP32 through the query-key product and query/key projection gradients to address underflow.Key-projection bias weights can grow rapidly when partially underflowed gradients become eccentric relative to their near-zero expected values early in training.Source: §A.2 · p. 20
  • Bfloat16's wider exponent range reduces underflow risk and removes the need for dynamic loss scaling; the authors use it for all final models.They report that it improves stability without eliminating all low-precision training concerns.Source: §A.2 · p. 20
  • A large Adam epsilon can slow small-weight updates and degrade model quality, so the authors choose epsilon below √µv/1000 for their 6.7B and 13B models.Here, mt is the running average of gradients and vt the running average of squared gradients.Source: §A.2 · p. 20

B Downstream Task Details

The appendix describes six zero- and few-shot downstream tasks and reports evaluation comparisons across open models. It notes that Cerebras-GPT 13B improves over publicly available models of comparable size, while Pythia’s training trajectories vary with tokens per parameter.

B Downstream Task Details

  • The six downstream evaluations are HellaSwag and PIQA for common-sense reasoning, WinoGrande for pronoun resolution, Lambada for contextual word prediction, ARC for science questions, and OpenBookQA for common-sense QA.HellaSwag examples are selected to be difficult for language models while remaining easy for humans.Source: B Downstream Task Details · pp. 20–21
  • OpenBookQA is a multiple-choice common-sense question-answering task.Source: B Downstream Task Details · p. 21

C.1 Pre-training Losses Throughout Training

  • Across model sizes and compute budgets, Cerebras-GPT’s Pile-loss trajectories approach final results along the scaling law; Pythia models trained on more tokens per parameter trend away from it.Pythia models trained closer to 20 tokens per parameter align more closely with Cerebras-GPT trajectories.Source: §C.1 · p. 22, Figure 8 · p. 22

C.2 Complete Downstream Task Testing

  • The appendix provides Pile and downstream results for GPT-J, GPT-NeoX, OPT, Pythia, and Cerebras-GPT, including five-shot results and plots.For comparisons with prior work, the authors also report results using minimum length-normalized NLL, which tends to favor longer continuations with moderate probability.Source: §C.2 · pp. 22–23, Table 8 · p. 24, Table 9 · p. 25

C.3 Differences Between Cerebras-GPT and Other Models

  • Cerebras-GPT 13B improves over other publicly available models of comparable size.Source: §C.3 · p. 23

C.4 Bias

The paper evaluates bias with CrowS-Pairs and compares Cerebras-GPT with Pythia and larger model baselines. Cerebras-GPT has lower average bias than the larger baselines, but its 13B model exceeds several of them on most tested categories.

C.4 Bias

  • Cerebras-GPT models show less average bias than the larger model baselines, but Cerebras-GPT 13B has greater bias than GPT-3, OPT, or LLaMA on six of nine categories.The comparison uses CrowS-Pairs, where higher values indicate higher bias.Source: §C.4 · p. 28, Table 11 · p. 28
  • Bias tends to grow with Cerebras-GPT and Pythia model size, though the range of bias values varies across models and categories.The authors suggest compute-efficient training might mitigate some bias relative to training smaller models on more data.Source: §C.4 · p. 28
  • At similar compute budgets, Cerebras-GPT models show roughly 1–2% higher average bias than the corresponding Pythia models.The comparison pairs Cerebras-GPT 1.3B, 2.7B, 6.7B, and 13B with Pythia 160M, 410M, 2.8B, and 12B, respectively.Source: §C.4 · p. 28
  • The authors recommend further bias evaluation and mitigation for Cerebras-GPT and larger models intended for production.Source: §C.4 · p. 28

D Additional Tokens-per-parameter Experiments

Additional experiments support 20 tokens per parameter as nearly compute-optimal for the tested Pile-trained models. The appendices detail the loss-degradation estimate, model-computation conventions, and µP tuning and batch-size guidance.

Estimated Chinchilla Losses

  • The Chinchilla loss-degradation estimate was fit from points at three FLOP levels and validated with low error on a fourth held-out level.The authors note the fitted curve may be erroneous because the true functional form is unknown.Source: Estimated Chinchilla Losses · p. 29
  • The degradation formulation agreed with tests, and Chinchilla and Pythia results suggest common scaling characteristics above 20 tokens per parameter across MassiveText and the Pile.Source: Estimated Chinchilla Losses · p. 29

Our Tokens-per-parameter Experiments

F Number of Training FLOPs

G.1 µP Implementation

  • µP implementation adjusts embedding and logit activations, fully connected layer initialization and learning rates, and key-query dot-product scaling.The key-query scaling choice assumes significant query-key alignment later in training; Table 14 summarizes implementation changes.Source: §G.1 · p. 31, Table 14 · p. 32

G.2 µP Hyperparameter Search Details

  • A 200-sample random search on a 40M proxy model yielded ηbase = 6e-3, σbase = 0.08, and memb = 10 for transfer to 111M–2.7B models.These tuned hyperparameters closely match those used by Yang et al. (2021).Source: §G.2 · p. 31, Figure 13 · p. 31

Use Large Enough Batch Sizes For Each Model Size

  • Proxy-model tuning need not use the largest target model’s batch size; a batch above the proxy’s critical batch size supports good hyperparameter transfer.Batch size should then scale appropriately with model size.Source: Use Large Enough Batch Sizes For Each… · pp. 31, 33

Simple Batch Size + Learning Rate Scaling with µP

Related papers

Loading 2304.03208v1…