Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling

In plain words

Researchers who want to know how a language model learns, and how that changes with size, have had few models to study. Most public models lacked the controls they need: few saved checkpoints, and no consistent record of the data order. That makes it hard to tell what caused a behaviour.

Pythia is a set of 16 models built to be compared. There are eight sizes, from tiny to about 12 billion parameters. Each size is trained twice: once on one copy of the training data and once on a copy with near-duplicates removed. Every run sees the data in the same order, and the authors saved many intermediate versions of each model. Because the order is fixed, a researcher can ask what a model had seen at any saved point.

The paper shows three uses. It changes the gender words in the last part of training and checks whether bias measures move. It checks how memorized text relates to where it sits in training. It checks when a model's accuracy starts to track how often the relevant facts appear in its training data.

TL;DR

Pythia is a controlled suite of 8 sizes (70M to 12B), each trained on the Pile and on a deduplicated Pile, with the same data order and 154 checkpoints per model §1 §2.4. The authors release the models, checkpoints, training code and data-loader tools §2.7. Use it when you need a public LM whose training history you can reproduce. Its three case studies are author-reported single runs.

Problem

Public LLMs did not, at the time, meet three requirements the authors set for a scientific suite: public access, training provenance (checkpoints, consistent data order and documented hyperparameters), and consistency across scale §2.1. The authors compared suites on these points, and the fewest-checkpoint model in each suite set the count §2.1 tab. 2. Studying how behaviour emerges during training, or what training data causes it, is hard without that control §1 §3.1.

Method

Pythia is a training recipe applied to eight model sizes, plus the data and checkpoints.

Architecture §2.3:

1. Decoder-only autoregressive transformer with fully dense attention in every layer, not alternating sparse and dense layers §2.3.
2. Flash Attention during training §2.3.
3. Rotary position embeddings §2.3, with rotary-pct 0.25 app. E.
4. Parallel attention and feed-forward sublayers, plus the initialisation methods of GPT-J §2.3.
5. Untied input embedding and output unembedding matrices §2.3.

Data §2.2:

1. Train on the Pile, an English corpus of about 300B tokens with the GPT-NeoX BPE tokenizer trained on the Pile §2.2.
2. Second suite: the Pile after near-duplicate removal (MinHash LSH, threshold 0.87), about 207B tokens §2.2.
3. Monolingual by choice; the authors cite dataset quality and evaluation coverage §2.2.

Training §2.4 app. E:

1. GPT-NeoX library with DeepSpeed ZeRO stage 1, data and tensor parallelism §2.4.
2. Global batch of 1024 sequences of 2048 tokens (2,097,152 tokens per step) for every size §2.4 app. B.
3. 143,000 steps, which is about 300B tokens §2.4 app. E.
4. Adam with betas (0.9, 0.95) and epsilon 1e-8 app. E.
5. Cosine learning-rate decay to 0.1 times the peak rate app. E app. B.
6. Checkpoint at initialisation, every 1000 steps (2,097,152,000 tokens), and at steps 1, 2, 4, ..., 512 §2.4 app. B.

# Training schedule for one Pythia run  [§2.4] [app. E]
# Input:  corpus in a fixed, recorded order; size-specific config
# Output: final model and checkpoints at fixed steps

def train_pythia(corpus, config):
    # Same batches, in the same order, for every model size  [§2.4]
    loader = fixed_order_loader(corpus, batch_size=1024, seq_len=2048)
    # 143,000 steps in total  [app. E]
    for step in range(143000):
        batch = next(loader)
        # Learning rate: peak per size, cosine decay to 0.1 x peak  [app. B]
        learning_rate = cosine_schedule(step, config.lr)
        loss = next_token_loss(model, batch)
        # Gradient clipping at norm 1.0  [app. E]
        apply_adam_update(model, loss, learning_rate, clip_norm=1.0)
        # Save at steps 0, 1, 2, 4, ..., 512 and every 1000 steps  [§2.4]
        if is_checkpoint_step(step):
            save(model, optimizer, step)
    return model

Deduplicated suite §2.4: the deduplicated Pile is about 207B tokens, so the 300B-token run goes through it about 1.5 times. The paper says the second epoch shows no sign of hurting benchmark scores §2.4.

Naming app. B: models are named by parameter count, but the authors recommend non-embedding counts for analysis tab. 1.

Hyperparameters

Per-size values tab. 1 tab. 3 tab. 5. Non-embedding parameter counts are from tab. 1. GPU counts are from the appendix hardware table tab. 5; the body table tab. 3 agrees for the sizes it lists.

Sources: sizes, layers, dims, heads, peak LR tab. 1; GPUs and GPU-hours app. D. The 1.0B row is a discrepancy; see Paper vs code.

Training hyperparameters shared by all sizes app. E unless noted:

Case-study settings:

Results

Author-reported. Setup: evaluation with the LM Evaluation Harness on public benchmarks, comparing Pythia to OPT and BLOOM of similar size §2.5 app. G. The paper states Pythia performs very similarly to OPT and BLOOM on these benchmarks §2.5.

Zero-shot scores for the fully trained suite on four tasks, from app. G. The paper gives ±standard errors that are omitted here.

Case studies (author-reported):

Gender-bias measures are described in words in the paper, and there is no table of numbers for them in the main text.

Paper vs code

Official code: EleutherAI/pythia (Apache-2.0), commit a19eecb807ec2c79a39ebf18108816e6ffffc1d5. The paper links this repository §2.7 app. E.

• Weight decay: the appendix table says 0.01 app. E pythia-70m.yml has weight-decay: 0.1. The same 0.1 appears in pythia-12b.yml and pythia-1b.yml. Use 0.1 unless you check a newer config.
• 1.0B peak learning rate: paper says 3.0e-4 tab. 1; code says 0.00025 pythia-1b.yml and the same in the deduplicated config.
• Eval interval: appendix lists 143000 app. E; the 70M config has 100000 pythia-70m.yml. The paper does not say which value its runs used.
• Memorization decoding: the paper defines memorization as generating the next tokens correctly §3.2 without naming a decoding method. The code uses greedy decoding (temperature 0, top-k 0, top-p 0) eval_memorization.py.
• Memorization threshold: the paper defines a string as (k,l)-memorized when the next l tokens are generated correctly §3.2. The code counts a sequence as memorized when its accuracy is exactly 1, meaning all 32 tokens match mem_graph.py.
• Gender intervention: the paper describes the procedure §3.1. The repository has no code for it; case-studies/README.md marks the intervention as a TODO and points to a NeoX branch README.md.
• WinoBias prompting: the paper describes the prompt templates and scoring app. C.1. The repository's own README marks this as a TODO README.md.

Pitfalls

• Two suite versions. The first release (v0) used a different batch size for 160M, 410M and 1.4B (4M tokens), different checkpoint steps, and a different minimum learning rate for 6.9B and 12B. The authors overwrote v0 on 31 March 2023. The v1 suite used here is uniform at 2M tokens, and its minimum rate is 0.1 times peak for every size app. B. Results from v0 are not covered by this card.
• Data order versus raw Pile. The standard Pile is 334B tokens with the GPT-NeoX tokenizer, so training at 300B tokens does not see all of it. Use the authors' data-loader tools to recover exactly what each batch contained §2.4 §2.7.
• Deduplicated run repeats data. The deduplicated Pile is about 207B tokens and the run is about 1.5 epochs §2.4. Checkpoints after the epoch boundary see repeated data.
• One seed per model. Each size was trained once per dataset. The authors say multi-seed runs would cost about 10 million USD for 25 seeds app. F.2. Differences of a few points on a benchmark, or a single gender-bias curve, may be seed noise. The paper gives no seed-variance measure.
• Intervention is not a clean counterfactual. The intervention replaces masculine pronouns with feminine ones in the training stream §3.1. The paper acknowledges that the CrowS-Pairs measure is disputed and that bias measures' reliability is an open question §3.1. The paper's own wording on the LAMBADA side effect is ambiguous ("marginal decrease in model perplexity" §3.1); read the appendix figure app. C.1 instead.
• Bias benchmarks were adapted. WinoBias and CrowS-Pairs were built for other model types. The paper rewrites WinoBias as prompt-and-log-probability scoring and CrowS-Pairs as a perplexity comparison app. C.1. Do not compare these numbers to published scores without the same setup.
• Term-frequency counts are approximate. The counting method for "salient entities" in the pretraining data is not stated in enough detail to reimplement §3.3. The count covers only data the checkpoint has seen §3.3.
• Memorization eval uses short prompts. It uses only the first 64 tokens of each training sequence §3.2, so it measures memorization of the start of documents, not the whole corpus.
• Compute. The GPU-hours in the table are for one run of each size; the paper's total was 544,280 A100-hours including two suites and the v1 re-run app. D.
• Figure and table numbers follow the order of floats in main.tex. Appendix letters follow the \appendix order. Check them against the compiled PDF before quoting.

Lineage

• Compared against: GPT-Neo, OPT, BLOOM, GPT-2, GPT-3 and T5 suites, judged on public access, data and order, checkpoints, and scale consistency §2.1 tab. 2 app. F.1.
• Relies on: GPT-NeoX library and DeepSpeed for training §2.4; Flash Attention §2.3; rotary embeddings §2.3; the Pile §2.2; LM Evaluation Harness for evaluation §2.5.
• Superseded by: not stated in this card's sources.