Compute-Optimal Training › Compute-Optimal Recipes
Language AI HandbookAll chapters Compute-Optimal Training › Compute-Optimal RecipesOne section · Sections in this chapter6 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Based on scaling laws and practical experience, we can establish concrete recipes for different compute scales. These recipes synthesize insights from Chinchilla and subsequent work to provide starting points for training runs. They represent practical wisdom distilled from experience, not rigid prescriptions. Each recipe encodes choices for the architecture and training setup that have been validated across many training runs.
Small-Scale Training (10^20 - 10^21 FLOPs)
This regime corresponds to models in the 100M to 1B parameter range, trainable on a single node of 8 GPUs in days to weeks. These scales are accessible to academic researchers and small teams, making them important proving grounds for new techniques.
The key recommendations are:
- Architecture: 12-24 layers, 768-2048 hidden dimension
- Tokens: 2-20B tokens (maintaining 20:1 token-to-parameter ratio)
- Learning rate: 3e-4 to 6e-4 peak
- Batch size: 256K to 512K tokens
- Warmup: 1% of total steps
def small_scale_recipe(compute_budget: float) -> dict:
"""Generate training recipe for small-scale runs (10^20 - 10^21 FLOPs)."""
allocation = compute_optimal_allocation(compute_budget)
n_params = allocation["params"]
n_tokens = allocation["tokens"]
# Derive hyperparameters based on scale
if n_params < 300e6:
peak_lr = 6e-4
batch_tokens = 256 * 1024
elif n_params < 700e6:
peak_lr = 5e-4
batch_tokens = 384 * 1024
else:
peak_lr = 3e-4
batch_tokens = 512 * 1024
total_steps = int(n_tokens / batch_tokens)
warmup_steps = int(total_steps * 0.01)
return {
"params_b": n_params / 1e9,
"tokens_b": n_tokens / 1e9,
"peak_lr": peak_lr,
"batch_tokens": batch_tokens,
"total_steps": total_steps,
"warmup_steps": warmup_steps,
"min_lr": peak_lr * 0.1,
}
# Example recipes
compute_scales = [1e20, 3e20, 1e21]Small-Scale Training Recipes Compute Params Tokens Peak LR Steps ------------------------------------------------------------ 1e+20 0.91B 18.3B 3e-04 34,823 3e+20 1.58B 31.6B 3e-04 60,315 1e+21 2.89B 57.7B 3e-04 110,120
Medium-Scale Training (10^22 - 10^23 FLOPs)
This regime covers models from 3B to 30B parameters, requiring multi-node training over weeks. These scales represent the domain of well-resourced research labs and enterprise training efforts. The engineering challenges shift from single-machine optimization to distributed systems coordination.
The key recommendations are:
- Architecture: 32-48 layers, 3072-5120 hidden dimension
- Tokens: 60B-600B tokens
- Learning rate: 1.5e-4 to 3e-4 peak
- Batch size: 1M to 4M tokens
- Warmup: 0.5% of total steps
- Parallelism: Tensor parallelism across 4-8 GPUs, data parallelism across nodes
Large-Scale Training (10^24+ FLOPs)
For frontier models (70B+ parameters), compute-optimal training requires careful attention to factors that become critical only at extreme scale. Training runs at this scale can take months, making each decision consequential and recovery from errors expensive.
The key considerations are:
- Extended context: Longer sequences (4K-32K tokens) for quality
- Data quality: Heavy filtering and deduplication become essential
- Checkpoint strategy: Frequent saves due to long training times
- Gradient checkpointing: Trade compute for memory to fit larger models
import numpy as np
def training_recipe(compute_budget: float) -> dict:
"""Generate complete training recipe for any scale."""
allocation = compute_optimal_allocation(compute_budget)
n_params = allocation["params"]
n_tokens = allocation["tokens"]
# Scale-aware hyperparameters
if n_params < 1e9:
peak_lr = 5e-4
batch_tokens = 512 * 1024
seq_len = 2048
elif n_params < 10e9:
peak_lr = 3e-4
batch_tokens = 2 * 1024 * 1024
seq_len = 4096
elif n_params < 70e9:
peak_lr = 1.5e-4
batch_tokens = 4 * 1024 * 1024
seq_len = 4096
else:
peak_lr = 1e-4
batch_tokens = 8 * 1024 * 1024
seq_len = 8192
total_steps = int(n_tokens / batch_tokens)
warmup_steps = max(1000, int(total_steps * 0.005))
# Estimate architecture
# Very rough: larger models use more layers relative to width
if n_params < 1e9:
n_layers = 24
elif n_params < 10e9:
n_layers = 32
elif n_params < 70e9:
n_layers = 48
else:
n_layers = 80
# Estimate d_model from params and layers
# Approximate: params ≈ 12 * n_layers * d_model^2 (for transformer)
d_model_approx = int(np.sqrt(n_params / (12 * n_layers)))
# Round to multiple of 128 for efficiency
d_model = ((d_model_approx + 64) // 128) * 128
return {
"params_b": n_params / 1e9,
"tokens_b": n_tokens / 1e9,
"peak_lr": peak_lr,
"min_lr": peak_lr * 0.1,
"batch_tokens": batch_tokens,
"seq_len": seq_len,
"total_steps": total_steps,
"warmup_steps": warmup_steps,
"n_layers": n_layers,
"d_model": d_model,
"tokens_per_param": n_tokens / n_params,
}Training Recipes Across Scales Compute: 1e+21 FLOPs Model: 2.9B params (32L, 2688H) Data: 58B tokens (20 per param) LR: 3e-04 → 3e-05 Batch: 2.1M tokens, Seq: 4096 Steps: 27,530 (warmup: 1,000) Compute: 1e+22 FLOPs Model: 9.1B params (32L, 4864H) Data: 183B tokens (20 per param) LR: 3e-04 → 3e-05 Batch: 2.1M tokens, Seq: 4096 Steps: 87,058 (warmup: 1,000) Compute: 1e+23 FLOPs Model: 28.9B params (48L, 7040H) Data: 577B tokens (20 per param) LR: 1e-04 → 1e-05 Batch: 4.2M tokens, Seq: 4096 Steps: 137,651 (warmup: 1,000) Compute: 1e+24 FLOPs Model: 91.3B params (80L, 9728H) Data: 1826B tokens (20 per param) LR: 1e-04 → 1e-05 Batch: 8.4M tokens, Seq: 8192 Steps: 217,645 (warmup: 1,088)
