DL S048 · Data Mixing and Curriculum

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 48 / 80 · Module M08 · ~2 hours

The story we're starting with

You have 1B tokens of clean web from S046. You could pretrain on that. Your model would learn "the internet" — fluent-ish English, some markdown, a lot of listicles, medium-quality reasoning. It would be bad at code, bad at math, bad at long-form reasoning, and roughly average at everything else.

Now look at what the actual open frontier models are trained on. RedPajama-V1: 67% CommonCrawl, 15% C4, 4.5% GitHub, 4.5% books, 2.5% arxiv, 2% wiki, 2% StackExchange. LLaMA-1: same buckets, near-identical ratios. Mistral, MPT, Falcon — all riffs on the same mix. Why? Because a decade of empirical evidence says that the right mixture of sources at pretraining time makes the model capable of things no single source could produce. Web gives you fluency. Code gives you reasoning. Books give you long-range coherence. Wiki gives you facts. Arxiv gives you dense technical writing. Mix them and you get a general model.

The tricky bit isn't the sources — those are commodities now. It's the weights. Should code be 4.5% or 15%? Should you upsample arxiv 4× because there's not enough of it? Does putting books at the end of training (curriculum) beat uniform mixing? This session teaches you the framework, gives you the ratios that work, and hands you the cheap-ablation methodology to defend your own.

Prerequisites

• S046 (Pretraining dataset) — you have a web corpus. Now we mix it with other sources.
• S047 (Chinchilla) — you need the tokens-per-parameter framework to budget ablations.

1 · Domains — what they are and what they contribute

Pretraining sources fall into ~7 domains. Each has a distinctive fingerprint on what the model becomes good at.

Note code punches way above its weight for reasoning. Multiple papers (Codex, PaLM, subsequent internal Anthropic/OpenAI ablations) show that adding 5–15% code to a general model improves non-code benchmarks like BBH, GSM8K, and MMLU. The going hypothesis: code's syntactic discipline forces the model to learn compositional structure it otherwise wouldn't.

2 · Natural weight vs sampling weight — the distinction that trips everyone

Say you have:

Natural weight is what you'd get if you just concatenated everything: web ≈ 90%, code ≈ 5.4%, everything else <5%. That's the supply distribution.

Sampling weight is what you actually want the model to see per step. RedPajama's target sampling weight is web ≈ 67%, code ≈ 5%, books ≈ 4.5%, and so on. To hit these numbers you have to upsample the small sources — sample from arxiv more often per step than its natural share.

The upsampling factor f_i = sampling_weight_i / natural_weight_i tells you how many "epochs through this source" the model will do per epoch through the corpus.

For our example, aiming for the RedPajama mix over 1B training tokens:

For a 1B-token training run, nothing hits >1 epoch. For a 15B-token run (Llama-3-style over-training), arxiv would hit 7.5× and wiki 10× — both dangerously close to overfitting. That's a real constraint on scaling.

3 · Implementing weighted sampling in a DataLoader

The cleanest pattern: one shard-list per domain, a per-step domain draw, then a within-domain random shard read.

import numpy as np, torch
from torch.utils.data import IterableDataset
 
class MixedPretrainDataset(IterableDataset):
    def __init__(self, domain_shards: dict[str, list[str]], weights: dict[str, float],
                 seq_len: int = 1024, seed: int = 0):
        self.domains = list(domain_shards.keys())
        self.shards  = domain_shards               # {"web": ["/data/web-0000.bin", ...], ...}
        self.weights = np.array([weights[d] for d in self.domains])
        self.weights /= self.weights.sum()
        self.seq_len = seq_len
        self.rng = np.random.default_rng(seed)
 
    def __iter__(self):
        while True:
            # Step 1: pick a domain
            d_idx = self.rng.choice(len(self.domains), p=self.weights)
            domain = self.domains[d_idx]
            # Step 2: pick a shard from that domain
            shard_path = self.rng.choice(self.shards[domain])
            arr = np.memmap(shard_path, dtype=np.uint16, mode="r")
            # Step 3: pick a random offset for one seq_len window
            start = self.rng.integers(0, len(arr) - self.seq_len - 1)
            chunk = arr[start:start + self.seq_len + 1].astype(np.int64)
            x = torch.from_numpy(chunk[:-1])
            y = torch.from_numpy(chunk[1:])
            yield x, y

Wrap with a DataLoader(dataset, batch_size=B, num_workers=8) and you have a working weighted-mix stream. Each worker inherits the RNG seed offset (torch.utils.data.get_worker_info().id) so workers don't produce identical batches.

4 · How do you pick the weights?

Three approaches, in order of increasing effort and returns.

4.1 Copy the reference mix

For a first run, use RedPajama-V1's weights exactly. It's a defensible baseline everyone knows.

WEIGHTS = {"web": 0.67, "code": 0.045, "books": 0.045, "arxiv": 0.025,
           "wiki": 0.02, "stackexchange": 0.02, "c4": 0.15}

4.2 Ablate — the cheap experimenter's path

Train a small proxy model (20M–50M params) on 500M–1B tokens with each candidate mix. Compare final loss + a few downstream evals (HellaSwag, ARC-easy). Pick the mix that wins.

Rule of thumb: a mix that wins at 20M usually wins at 100M. Cross-scale transfer of mix rankings is high (~85% agreement). This is the DoReMi paper's core empirical claim.

Costs to run a 3-way mix ablation at 20M params × 500M tokens:

C = 6 · 20e6 · 500e6 = 6e16 FLOPs per run
   × 3 runs           = 1.8e17 FLOPs total
   / (312e12 · 0.4)   = 1.4e3 seconds per run ≈ 25 min
   × 3                = ~75 min on one A100
   × $2/hr            = ~$2.50 total

Two and a half dollars to pick a defensible mix. Do this. It pays for itself in the first day of your real run.

4.3 DoReMi — let a proxy model learn the weights

DoReMi (Xie et al., 2023) trains a small model with initial mix W₀, then a reference model trained on some other mix, and updates W adaptively based on per-domain excess loss. The intuition: domains where your proxy model has high loss relative to the reference need more weight.

The algorithm is essentially group-DRO (distributionally robust optimization) applied to domains. In their experiments DoReMi weights beat hand-picked weights by 2–3× training speedup at fixed final loss.

Implementation is nontrivial (~200 lines of PyTorch + a reference model to compare against). Worth it only if you're scaling past 1B params.

5 · Curriculum — do you change the mix over time?

Two schools of thought:

Uniform mix throughout (LLaMA, GPT-3, most published models). Simple, reproducible, works fine.

Curriculum — start with "easy" data (web, code) and add "hard" data (papers, books) later. Or the reverse. The published evidence is mixed — several papers show ~1-2% gains from thoughtful curricula, several show gains disappear with better hyperparams.

Where curriculum reliably helps:

1. End-of-training high-quality dump. Take your last 10% of training and feed only the highest-quality sources (Wikipedia, textbooks, papers). Every serious model does this. Called "annealing". Gains: 1–3 points on downstream evals.
2. Instruction data at the end. Blend in a small fraction of instruction-tuned examples in the last 5% — makes the base model much easier to SFT afterward. Called "midtraining" or "cooldown" — coined by the Phi-2 authors.

Where curriculum reliably doesn't help:

• Sorting by document length ("short first, then long").
• Sorting by perplexity ("easy first").
• Any hand-crafted difficulty ordering. The literature has largely killed these; they don't beat uniform mixing after learning rate tuning.

6 · A defensible 5-domain mix for our 100M model

Here's what I'd actually run for the M09 100M pretraining:

Web dominates for fluency. Code oversized (15% > RedPajama's 5%) because our downstream evals include HumanEval. Books smaller (8% vs 4.5%) because we care about long-context. Wiki upsized 2.5× beyond natural weight for fact density. And the final annealing phase (last 100M tokens of a 2B-token run) shifts to {wiki: 25%, arxiv: 15%, books: 25%, web: 30%, code: 5%}.

7 · The mini-ablation protocol

Concrete recipe I use before committing to a real run:

# ablate_mix.py — evaluate 3 candidate mixes at 20M params, 500M tokens each.
 
MIXES = {
    "baseline":  {"web": 0.90, "code": 0.10},
    "balanced":  {"web": 0.60, "code": 0.15, "books": 0.10, "wiki": 0.05,
                  "arxiv": 0.05, "stackexchange": 0.05},
    "code-heavy":{"web": 0.50, "code": 0.30, "books": 0.10, "wiki": 0.05,
                  "arxiv": 0.05},
}
 
for name, weights in MIXES.items():
    model  = TinyGPT(n_layer=6, n_head=6, n_embd=384)  # ~20M params
    ds     = MixedPretrainDataset(SHARDS, weights, seq_len=512)
    losses = train(model, ds, tokens=500_000_000, lr=6e-4)
    evals  = eval_downstream(model)  # HellaSwag, ARC-e, MMLU-tiny
    print(name, "final_loss", losses[-1], "avg_eval", evals["avg"])

Rank by (a) final validation loss on a held-out mixed set and (b) average downstream eval. If they agree, use the winner. If they disagree, prefer downstream eval — training loss is often gamed by whatever mix is easiest, not what generalizes.

8 · Diagram — the whole per-step pipeline

The whole system: choose domain → choose shard → choose offset. Three RNG draws per training example. O(1) memory. Easy to change weights mid-run just by editing a config.

9 · War stories

10 · Retention scaffold

Recall

What frontier mixes actually look like in 2025

The RedPajama-V1 mix (67% web, 15% C4, 4.5% code, 4.5% books, ...) was state of the art in early 2023. It's now archaic. Here's how the frontier evolved through 2024–2025:

Llama 3 (Meta, 2024) — the tech report reveals a two-stage approach:

1. Stage 1 (majority of 15T tokens): a mix heavier in code (~17%) and reasoning-adjacent data than any prior public mix. General web is still the plurality but the code share doubled from Llama 2.
2. Stage 2 ("annealing", last ~5% of tokens): switches to high-quality math + code + reasoning data with a decaying learning rate. Meta reports this alone lifts MMLU by <a couple points and code benchmarks noticeably.

DeepSeek-V3 (Dec 2024) — 14.8T tokens, mix explicitly tilted toward math + code (the whole model is built for the R1 reasoning line). They also do a stage-2 annealing dump with instruction-adjacent data.

DCLM-Baseline mix (2024) — 3.8T tokens, ~all filtered CC. No books, no arxiv, no explicit code bucket. Their argument: a well-filtered web scrape already contains enough code, enough math, and enough long-form content that explicit domain buckets are second-order. Provocative and, on their eval suite, correct at 1.4B scale.

FineWeb-Edu (HuggingFace, 2024) — 1.3T tokens, single source, filtered for "educational" content by a Llama-3-70B judge. On matched compute, models trained purely on FineWeb-Edu beat models trained on 10× more raw CC. A one-source "mix".

The 2025 consensus is not a specific ratio; it's four principles:

1. Web is the plurality, always — 45–75% of tokens.
2. Code is 10–20%, up from 5% in the RedPajama era, because it lifts reasoning.
3. Books + academic + wiki together are 5–15% — they carry the long-form and factual load.
4. Annealing cooldown on the last ~5% with the cleanest, highest-quality mix you own is now standard practice, essentially free MMLU + reasoning points.

Stretch

Take a public mix (e.g., RedPajama-V1) and design 3 variants: baseline, code-heavy, book-heavy. Estimate the annealing dump (last 10%) for each. Compute the effective epochs per source at (a) 1B total tokens, (b) 15B tokens. Which variants blow through the 4-epoch ceiling? Which sources would you need more of?

In your own words

"The sampling weight of a domain matters more than its absolute size because ______."

Spaced-review pointer

• S046 — the web bucket comes from there.
• S047 — total token count comes from Chinchilla; per-source token count = weight × Chinchilla D.

Next-session teaser

Session 049 gets into the deep magic: dedup + quality filtering at production scale. MinHash + LSH, near-dup detection, quality classifiers. These are the machinery that turn "raw web" into "usable web" — the pipeline half we handwaved in S046.

Bring back tomorrow

• Concept: natural weight vs sampling weight.
• Number: 4-epoch upsampling ceiling.
• Snippet: the MixedPretrainDataset per-step domain draw.

Previous: ← DL S047 · Next: DL S049 →