S114 · Scaling Laws — Chinchilla, Compute-Optimal Training

"""scaling_laws.py — compute-optimal LLM sizing calculator.
 
Given a compute budget in FLOPs, predict the compute-optimal model size
and training tokens using Chinchilla exponents. Also evaluate known LLMs
against the Chinchilla frontier and print how 'over-trained' each is.
"""
from __future__ import annotations
import math
from dataclasses import dataclass
 
# Chinchilla-fit constants (Hoffmann et al. 2022, Table A9).
# L(N,D) = A/N^ALPHA + B/D^BETA + E
A, ALPHA = 406.4, 0.34
B, BETA  = 410.7, 0.28
E        = 1.69
 
# Compute cost of one forward+backward pass through the whole dataset.
def compute_flops(n_params: float, n_tokens: float) -> float:
    """Kaplan/Chinchilla approximation: 6 × N × D."""
    return 6 * n_params * n_tokens
 
 
def chinchilla_loss(n_params: float, n_tokens: float) -> float:
    """Predicted loss under the Chinchilla parametric fit."""
    return A / n_params ** ALPHA + B / n_tokens ** BETA + E
 
 
def optimal_split(compute_flops: float) -> tuple[float, float]:
    """Given a training compute budget C, return compute-optimal (N, D).
 
    From Chinchilla Section 3: N_opt scales like C^0.5, D_opt = ~20 N_opt.
    """
    # Chinchilla's practical result: N_opt ≈ 0.6 * sqrt(C), D_opt ≈ 20 * N_opt.
    n_opt = 0.6 * math.sqrt(compute_flops)
    d_opt = 20 * n_opt
    return n_opt, d_opt
 
 
@dataclass
class LLM:
    name: str
    params: float
    tokens: float
 
    @property
    def flops(self) -> float:
        return compute_flops(self.params, self.tokens)
 
    @property
    def tokens_per_param(self) -> float:
        return self.tokens / self.params
 
    @property
    def chinchilla_optimal_params(self) -> float:
        return optimal_split(self.flops)[0]
 
    @property
    def over_train_ratio(self) -> float:
        """How many times MORE tokens/param than Chinchilla-optimal (20)."""
        return self.tokens_per_param / 20
 
 
MODELS = [
    LLM("GPT-3",       175e9,   300e9),
    LLM("Chinchilla",   70e9,   1.4e12),
    LLM("LLaMA-1-65B",  65e9,   1.4e12),
    LLM("LLaMA-2-70B",  70e9,   2.0e12),
    LLM("LLaMA-3-8B",    8e9,    15e12),
    LLM("LLaMA-3-70B",  70e9,    15e12),
    LLM("LLaMA-3-405B",405e9,    15e12),
    LLM("Qwen-2.5-7B",   7e9,    18e12),
]
 
print(f"{'model':<15}{'params':>10}{'tokens':>10}{'FLOPs':>12}"
      f"{'tok/param':>12}{'Chinch N*':>12}{'over×':>8}")
print("-" * 82)
for m in MODELS:
    print(f"{m.name:<15}"
          f"{m.params/1e9:>8.0f}B"
          f"{m.tokens/1e12:>8.1f}T"
          f"{m.flops:>12.1e}"
          f"{m.tokens_per_param:>12.1f}"
          f"{m.chinchilla_optimal_params/1e9:>10.1f}B"
          f"{m.over_train_ratio:>8.1f}×")
 
print("\n--- If YOU had 100 H100s × 30 days × 300 TFLOP/s BF16 ---")
c = 100 * 30 * 86400 * 300e12
n, d = optimal_split(c)
print(f"Budget       : {c:.2e} FLOPs")
print(f"Optimal N    : {n/1e9:.1f}B params")
print(f"Optimal D    : {d/1e12:.2f}T tokens")
print(f"Predicted loss: {chinchilla_loss(n, d):.3f}")