LLM pretraining

Deep dive

Decision table

Data

1. Sources: filtered web crawl (Common Crawl via FineWeb-style pipelines), code, math, papers, books, multilingual text. Llama 3's final mix: about 50% general knowledge, 25% math and reasoning, 17% code, 8% multilingual.
2. Cleaning: URL/language filters, exact dedup plus MinHash near-dedup, heuristic and classifier quality filters (FineWeb-Edu keeps 1.3T of 15T tokens via an educational-quality classifier), PII scrubbing, n-gram decontamination against eval sets.
3. Mixing: sample sources by weight, not by size; tune weights on small proxy models.
4. Packing: tokenize (BPE, 32k-256k vocab), concatenate documents with EOS, cut into fixed sequences (4k-8k). Mask attention across document boundaries (Llama 3 does) so tokens never attend into an unrelated document.
5. Annealing: in the last few percent of training, upsample the highest-quality and math/code data while the LR decays to near zero.
6. Long context: train at 4-8k, then extend in stages with a larger RoPE base (Llama 3: base 500,000; 8k → 128k in six stages, about 800B tokens).

Optimizer, batch and LR

• AdamW β₁ 0.9, β₂ 0.95 (not 0.999: faster adaptation reduces spikes), ε 1e-8 to 1e-5, weight decay 0.1 on matrices only (not norms, biases, often not embeddings), global grad-norm clip 1.0.
• Peak LR falls with size: Llama 2 3e-4 (7B, 13B) and 1.5e-4 (34B, 70B); Llama 3 405B 8e-5. μP lets you tune LR on a small proxy and transfer it across width.
• Batch: 1-16M tokens. Ramp it up early (GPT-3: 32k → 3.2M tokens over the first 4-12B tokens; Llama 3 405B: 4M → 8M → 16M tokens).
• Schedule: linear warmup (order of 2000 steps), then cosine or warmup-stable-decay with the decay in the last 10-20% of steps.
• Init: normal(0, 0.02), output projections scaled by 1/√(2L).

Numerics and parallel layout

• bf16 GEMMs with fp32 master weights, optimizer state and gradient reduction (mixed precision); fp8 GEMMs at the largest scales.
• 16 B/param of state: a 7B model needs 112 GB, more than one 80 GB GPU, so shard it (ZeRO/FSDP). Activations: FlashAttention plus selective recomputation.
• Order the dimensions by bandwidth need: tensor parallel innermost on NVLink (≤ 8), then context parallel, pipeline parallel across nodes, data parallel outermost over the slowest links (all-reduce or reduce-scatter overlapped with backward).
• Measure MFU = (6N + 12·L·T·d) × tokens/s ÷ peak FLOPs. 35-45% is good for dense models (Llama 3 405B: 38-43%).

Checkpointing and failures

• Save model, optimizer, LR-scheduler step, RNG states and the data-loader position; restarting must replay exactly the same token stream.
• Every rank writes its own shard; write to local NVMe or host memory first, upload asynchronously, keep the last k plus periodic long-lived ones.
• Interval by Young's rule: τ ≈ √(2·δ·M), δ = time to write a checkpoint, M = mean time between job failures. δ = 1 min and M = 3 h gives τ ≈ 19 min.
• Detect hangs (NCCL watchdog timeouts), stragglers and silent data corruption (loss or grad-norm outliers on one rank); drain the bad node and restart from the last checkpoint automatically.

Loss spikes

Causes: attention-logit growth, output-logit drift, Adam's second-moment estimate lagging a sudden gradient change, bad data batches. Mitigations, cheapest first:

1. Skip the step when the grad norm or loss is far above its running average (code below); never apply NaN/inf.
2. Repeated spikes: roll back to a checkpoint before the spike and skip the offending data. PaLM restarted about 100 steps before the spike and skipped 200-500 batches.
3. Architecture: QK-norm, z-loss 1e-4 · log²Z on the softmax normalizer (PaLM), no biases, β₂ = 0.95, lower peak LR.

Code

Budget planner (params, FLOPs, GPU-hours, model-state memory), the two LR schedules, and a spike guard on the pre-clip gradient norm.

// Pretraining planner + LR schedule + spike guard: the arithmetic you redo for every run.
pub struct Model { pub layers: u64, pub d: u64, pub d_ff: u64, pub n_kv: u64, pub n_head: u64, pub vocab: u64, pub seq: u64 }

impl Model {
    pub fn params(&self) -> u64 { // no biases, SwiGLU MLP, untied embedding + head
        let dh = self.d / self.n_head;
        let attn = 2 * self.d * self.d + 2 * self.d * self.n_kv * dh;
        self.layers * (attn + 3 * self.d * self.d_ff + 2 * self.d) + 2 * self.vocab * self.d + self.d
    }
    /// Training FLOPs per token: 6N for the weights + 12 L T d for attention scores (PaLM's MFU count).
    pub fn flops_per_token(&self) -> f64 {
        6.0 * self.params() as f64 + 12.0 * (self.layers * self.seq * self.d) as f64
    }
}

pub struct Layout { pub tp: u64, pub pp: u64, pub dp: u64, pub zero_shard: bool }

/// Bytes of weights + grads + AdamW state per GPU (bf16 weights/grads, fp32 master, m, v = 16 B/param).
pub fn state_bytes_per_gpu(n: u64, l: &Layout) -> f64 {
    let per_model_shard = n as f64 / (l.tp * l.pp) as f64;
    let opt = 12.0 * per_model_shard / if l.zero_shard { l.dp as f64 } else { 1.0 }; // ZeRO-1: fp32 master, m, v sharded over DP
    4.0 * per_model_shard + opt
}

/// Warmup-stable-decay: linear warmup, flat peak, then linear decay to `floor` over the last `decay` steps.
pub fn wsd_lr(step: u64, peak: f64, floor: f64, warmup: u64, total: u64, decay: u64) -> f64 {
    if step < warmup { return peak * (step + 1) as f64 / warmup as f64; }
    let start = total - decay;
    if step < start { return peak; }
    let f = ((step - start) as f64 / decay as f64).min(1.0);
    peak + (floor - peak) * f
}

/// Cosine from peak to `floor` after linear warmup (Llama 2: floor = 0.1 * peak).
pub fn cosine_lr(step: u64, peak: f64, floor: f64, warmup: u64, total: u64) -> f64 {
    if step < warmup { return peak * (step + 1) as f64 / warmup as f64; }
    let p = ((step - warmup) as f64 / (total - warmup) as f64).min(1.0);
    floor + 0.5 * (peak - floor) * (1.0 + (std::f64::consts::PI * p).cos())
}

pub enum Verdict { Apply, SkipBatch, Rollback }

/// Spike guard on the pre-clip global grad norm and the loss: skip isolated outliers,
/// roll back to the last checkpoint if they keep coming.
pub struct SpikeGuard { ema_norm: f64, ema_loss: f64, strikes: u32, pub seen: u64 }
impl SpikeGuard {
    pub fn new() -> Self { SpikeGuard { ema_norm: 0.0, ema_loss: 0.0, strikes: 0, seen: 0 } }
    pub fn check(&mut self, grad_norm: f64, loss: f64) -> Verdict {
        self.seen += 1;
        if !(grad_norm.is_finite() && loss.is_finite()) { return Verdict::SkipBatch; } // never apply NaN/inf
        if self.seen <= 100 { self.update(grad_norm, loss); return Verdict::Apply; }   // warm up the averages
        let spike = grad_norm > 6.0 * self.ema_norm || loss > self.ema_loss + 0.5;
        if !spike { self.strikes = 0; self.update(grad_norm, loss); return Verdict::Apply; }
        self.strikes += 1;
        if self.strikes >= 5 { self.strikes = 0; Verdict::Rollback } else { Verdict::SkipBatch }
    }
    fn update(&mut self, g: f64, l: f64) {
        let a = if self.ema_norm == 0.0 { 1.0 } else { 0.02 }; // first sample seeds the EMA
        self.ema_norm += a * (g - self.ema_norm);
        self.ema_loss += a * (l - self.ema_loss);
    }
}

fn main() {
    // Llama-2-7B shape trained on 2T tokens.
    let m = Model { layers: 32, d: 4096, d_ff: 11008, n_kv: 32, n_head: 32, vocab: 32000, seq: 4096 };
    let n = m.params();
    assert_eq!(n, 6_738_415_616);
    let tokens = 2.0e12;
    let flops = m.flops_per_token() * tokens;
    let (peak, mfu) = (989e12, 0.40); // H100 SXM dense bf16
    let gpu_hours = flops / (peak * mfu) / 3600.0;
    println!("N = {:.2}B, tokens/param = {:.0}, FLOPs = {:.2e}, H100-hours at 40% MFU = {:.0}",
             n as f64 / 1e9, tokens / n as f64, flops, gpu_hours);
    let l = Layout { tp: 1, pp: 1, dp: 64, zero_shard: true };
    println!("model state per GPU: {:.1} GB", state_bytes_per_gpu(n, &l) / 1e9);
    assert!(cosine_lr(2000, 3e-4, 3e-5, 2000, 500_000) > 2.99e-4);
    assert!((wsd_lr(499_999, 3e-4, 3e-6, 2000, 500_000, 50_000) - 3e-6).abs() < 1e-8);
    let mut g = SpikeGuard::new();
    for i in 0..200 { assert!(matches!(g.check(1.0 + 0.01 * (i % 7) as f64, 2.5), Verdict::Apply)); }
    assert!(matches!(g.check(40.0, 2.5), Verdict::SkipBatch));
    assert!(matches!(g.check(f64::NAN, 2.5), Verdict::SkipBatch));
}

For the Llama 2 7B shape this prints N = 6.74B, 297 tokens/param, 9.4e22 FLOPs and about 66k H100-hours at 40% MFU. Meta reported 184,320 A100-hours for that run, which is about 40% of A100 peak by the 6N count.

Efficiency

• Cost: C ≈ 6·N·D FLOPs: a 7B model on 2T tokens is 9.4e22 FLOPs, about 66k H100-hours at 40% MFU (above). Bound: compute, then interconnect.
• Levers:
  • parallel layout by bandwidth (TP in-node, PP across nodes, DP outermost) — 52% of peak on 3072 GPUs (Megatron-LM) — 100B+
  • bf16, FlashAttention, recomputation — higher MFU, less memory — every run
  • spike guard, rollback — no wasted steps
• Naive → standard → state of the art:
  • plain data parallel, fp32 · low MFU
  • FSDP, bf16, FlashAttention · 35-45% MFU on dense models
  • 4D parallelism, automatic restart · 38-43% MFU for Llama 3 405B

Pitfalls

• Counting only 6N for MFU at long sequence lengths: the attention term is 12·L·T·d per token and is no longer small at 32k+ context.
• Resuming without the data-loader position: the model sees some data twice and some never, and loss curves are not comparable.
• Weight decay on norm gains and embeddings.
• Eval-set contamination from web crawl: inflated benchmarks that collapse on fresh test sets.
• Restarting on a different world size without fixing the global batch: the effective LR and the token schedule change silently.
• Cross-document attention in packed sequences: the model learns to use unrelated context.
• A single slow or faulty GPU throttles the whole synchronous job; monitor per-rank step time.

Implementations

Sources

• Kaplan et al., Scaling Laws for Neural Language Models, 2020. paper
• Hoffmann et al., Training Compute-Optimal Large Language Models, 2022. paper
• Brown et al., Language Models are Few-Shot Learners, NeurIPS 2020. paper
• Chowdhery et al., PaLM: Scaling Language Modeling with Pathways, 2022. paper
• Touvron et al., Llama 2: Open Foundation and Fine-Tuned Chat Models, 2023. paper
• Llama Team, AI @ Meta, The Llama 3 Herd of Models, 2024. paper
• Muennighoff et al., Scaling Data-Constrained Language Models, NeurIPS 2023. paper
• Penedo et al., The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale, 2024. paper
• Yang et al., Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer, 2022. paper
• Hu et al., MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies, 2024. paper
• Wortsman et al., Small-scale proxies for large-scale Transformer training instabilities, ICLR 2024. paper
• Narayanan et al., Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC 2021. paper · code
• Young, A first order approximation to the optimum checkpoint interval, CACM 1974. paper