Base pretraining¶

Goal¶

Train the model to predict the next token over a huge corpus.

The objective is simple:

Given tokens x₁, x₂, …, x_{t-1}, predict xₜ.

This teaches grammar, facts, code, style, reasoning patterns, world knowledge, and latent task-solving ability.

Scaling laws and token budget¶

The Chinchilla paper (Hoffmann et al., 2022) showed that many earlier large LMs were undertrained for their size and that, under a fixed compute budget, model size and number of training tokens should scale together. Chinchilla, a 70B model trained on more data, outperformed much larger models trained on fewer tokens.

The headline rule of thumb from Chinchilla:

For compute-optimal training, ~20 tokens per parameter.

In practice, modern open models often train far past the compute-optimal point. Llama 3 8B was trained on 15T+ tokens (~1900 tokens per parameter) — wildly "overtrained" by Chinchilla but optimal for inference cost: a smaller model trained longer is cheaper to serve.

Practical implications:

Do not simply maximize parameter count.
Choose model size and token budget jointly.
A smaller model trained longer can beat a larger undertrained model — and is cheaper to serve.
Data quality becomes increasingly important as token count grows.

Data mixture¶

A pretraining mixture might include:

Category	Purpose
High-quality web	General knowledge and breadth
Books	Long-form coherence and style
Code	Reasoning, formatting, tool use
Math	Symbolic reasoning
Academic text	Technical knowledge
Q&A	Answer-like patterns
Multilingual text	Cross-lingual capability
Synthetic text	Targeted coverage
Domain data	Specialization

The mixture is not static. Many teams change the mixture during training.

Example phases:

Broad early pretraining — maximum diversity.
Quality-upweighted middle phase — less spam, more high-value text.
Annealing / cooldown phase — high-quality data, instruction-like text, code, math, long-form text.

OLMo 2 describes a post-pretraining "mid-training" approach where domain-specific data mixtures are generated and training is restarted with a changed schedule, highlighting how data mixture and learning-rate schedule matter even after initial pretraining.

The annealing trick

Many open recipes (Llama 3, OLMo 2, MiniCPM) end pretraining with a short, sharp learning-rate decay over a high-quality data mix — math, code, instruction-like text. This "annealing" or "cooldown" phase consistently improves downstream evals at low marginal cost. See MiniCPM, arxiv.org/abs/2404.06395, for an explicit description of the technique.

Training mechanics¶

Loss¶

Usually cross-entropy loss over next-token prediction.

Optimizer¶

Common choices:

AdamW — current default. Adds decoupled weight decay to Adam.
Adam variants — Adafactor, LAMB, etc.
Lion — momentum-only, lower memory; used in some experiments and for some models.

AdamW is the safe default. Lion saves optimizer-state memory and trains slightly differently; both work.

Learning-rate schedule¶

Common pattern:

Warmup — linear ramp from 0 to peak LR over first 0.5–2% of steps.
Long cosine or linear decay — slowly decay to ~10% of peak.
Cooldown / annealing — sharp decay to near-zero on a high-quality data mix at the end.

Learning rate is one of the most important stability knobs. Modern recipes use a peak LR roughly proportional to 1/sqrt(d_model) — see µP (Yang et al., 2022) for a principled treatment.

Batch size¶

Large global batch sizes are common. Batch size is often measured in tokens, not examples.

global_batch_tokens = sequence_length
                   × micro_batch_size
                   × gradient_accumulation_steps
                   × number_of_devices

Typical global batch sizes for modern open models: 4M–16M tokens per step.

Sequence packing¶

Short documents are packed together to avoid wasting context space.

Important detail: you usually need attention masks or document-boundary handling so unrelated documents do not leak into each other as if they were one continuous text.

Packing without document attention masks

Naïve packing (just concatenating documents with EOS in between, no attention mask) is a common silent bug. The model learns weird cross-document statistics. Use document-aware FlashAttention variants or block-diagonal masks.

Precision¶

Common formats:

FP32 for some optimizer states (m, v in Adam)
BF16 for activations/weights
FP16 in older systems (less stable than BF16; usually requires loss scaling)
FP8 in newer large-scale systems (H100+)

BF16 is popular because it has the same exponent range as FP32 — far more numerically stable than FP16 — at half the memory.

Parallelism¶

Large LLM training combines several forms of parallelism:

Method	What it does
Data parallelism (DP)	Replicates model across workers; each sees different data
Tensor parallelism (TP)	Splits matrix operations across devices
Pipeline parallelism (PP)	Splits layers across devices
Sequence parallelism (SP)	Splits the sequence dimension
Expert parallelism (EP)	For MoE expert routing
ZeRO / FSDP	Shards optimizer states, gradients, parameters across DP workers

In practice you compose these — e.g., 3D parallelism = DP × TP × PP, sometimes plus EP for MoE.

Frameworks

Megatron-LM / Megatron-Core — NVIDIA. The reference for TP/PP/SP. github.com/NVIDIA/Megatron-LM
DeepSpeed — Microsoft. ZeRO stages ½/3. github.com/microsoft/DeepSpeed
FSDP — PyTorch native fully-sharded data parallel.
Torchtitan — PyTorch's reference distributed training stack. github.com/pytorch/torchtitan
nanotron — HuggingFace's clean training library. github.com/huggingface/nanotron

Checkpointing¶

You save:

Model weights
Optimizer state
LR scheduler state
Random seeds
Data loader position
Training metrics
Tokenizer and config
Code version
Data mixture version

Intermediate checkpoints are valuable for ablations, recovery, and post-training.

Save more checkpoints than you think

Frontier-grade pretraining runs save a checkpoint every few thousand steps. Disk is cheap; rerunning a training month is not. Intermediate checkpoints also let you do ablations like "what if we'd switched the data mix here?"

Stability issues¶

Common failure modes:

Loss spikes
Divergence
NaNs
Tokenizer bugs
Bad data shard
Optimizer instability
Incorrect masking
Repeated data
Bad learning-rate warmup
Distributed training desynchronization
Silent data corruption
Evaluation contamination

Mitigations:

Gradient clipping (typical: global norm 1.0)
Stable normalization (RMSNorm pre-norm)
Careful initialization (e.g., scaled init for deep models)
BF16/FP32 mixed precision
Loss spike detectors
Automatic bad-batch skipping (with caution — silent skipping hides bugs)
Frequent validation
Small-scale recipe validation before full run

Don't auto-skip silently

A loss-spike detector that auto-skips bad batches without logging exactly which batch and why will hide real bugs (data corruption, optimizer NaN, distributed desync). Always log skipped batches; if you skip more than a handful per run, something is broken.

Practical tips¶

Train small models first. Validate your recipe at 100M, 500M, 1B before launching the full run. The cost is tiny relative to the main run, and you will find bugs.
Reproduce a known result first. Re-train a published 1B model from a known recipe. If your numbers match within 0.5%, your stack works. If they don't, find out why before launching anything bigger.
Watch validation loss by domain. Aggregate train loss is too coarse. Per-domain validation loss (web, code, math, books, multilingual) reveals data-mixture problems early.
Run benchmark snapshots during training. Loss correlates with downstream tasks but not perfectly. Run cheap benchmarks every few thousand steps so you can catch regressions before the run finishes.
Plan for hardware failures. Multi-thousand-GPU runs will lose nodes. Resumable checkpointing and elastic schedulers are not optional.