DL S048 · Data Mixing and Curriculum
Design a data mix that works for a mixed-domain model. Part of the 'Deep Learning & LLMs From Scratch' 80-session self-study series.
🎯 Design a domain-weighted training mix and defend it with cheap ablations.
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.
- Explain the 'sampling weight' construction and how it differs from 'natural weight'.
- Design a 5-domain mix (web / code / books / arxiv / wiki) with defensible ratios.
- Implement weighted sampling across shards inside a PyTorch DataLoader.
- Run a mini-ablation (train a 20M model on 3 different mixes for 500M tokens each) and pick the winner in <$20 GPU cost.
- Explain when curriculum (change the mix over time) helps and when it hurts.
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:
| Source | Tokens available |
|---|---|
| Web | 5,000 M |
| Code | 300 M |
| Books | 200 M |
| Arxiv | 50 M |
| Wiki | 30 M |
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:
| Source | Tokens avail | Target sample % | Tokens seen | Effective epochs |
|---|---|---|---|---|
| Web | 5000M | 67% | 670M | 0.13× (subsample) |
| Code | 300M | 5% | 50M | 0.17× |
| Books | 200M | 4.5% | 45M | 0.22× |
| Arxiv | 50M | 2.5% | 25M | 0.5× |
| Wiki | 30M | 2% | 20M | 0.67× |
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, yWrap 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.
Instantiate MixedPretrainDataset with the RedPajama-style weights (web 0.67, code 0.05, books 0.045, arxiv 0.025, wiki 0.02, stackexchange 0.02, c4 0.15) using dummy shard lists (one 10 MB .bin per domain works). Pull 100,000 samples and count how many came from each domain. Print the empirical fraction next to the target weight. Then bump the code weight to 0.30, keep everything else proportional, re-run, and confirm the code count rises accordingly. This is the fastest way to catch a bug in your sampling code before a $200 real run notices it.
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 totalTwo 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:
- 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.
- 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:
| Source | Sampling % | Source of tokens |
|---|---|---|
| Web (S046 output) | 60% | Our custom CC pipeline |
| Code | 15% | The Stack v2, dedup subset |
| Books | 8% | Project Gutenberg + PG-19 |
| Wiki | 5% | HF wikipedia/20220301.en |
| Arxiv | 5% | HF RedPajamaArXiv |
| StackExchange | 4% | HF RedPajamaStackExchange |
| Instruction data (last 5% only) | 3% | Alpaca-clean |
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
I once ran a mix with books at 25% because I thought long-range coherence would help. It did — the model wrote beautifully. But it wrote books. Every completion was novelistic, third-person narrator, past tense. Even short prompts like "list three benefits of X" turned into two-paragraph stories. Ratios matter more than sizes.
Lesson: if a source is >10% of the mix, expect the model to inherit its stylistic idioms, not just its knowledge. Diversity of style is a design goal.
Trying to make a "science-strong" 300M model, I upsampled arxiv from 5% to 20%. My MMLU-STEM went up 3 points. My HellaSwag (common-sense) went down 5 points. Net negative on any composite eval. Turns out arxiv writing style has almost no overlap with general reasoning — you can't buy STEM without paying for general.
Lesson: the mix is a global tradeoff, not per-domain improvement. Optimize on a composite eval, not on the source you happen to care about.
"Mixture weights should reflect how much of each domain you want the model to be good at. If code matters as much as prose, use fifty-fifty."
Mixture weights control how many gradient steps a domain contributes, which is not the same as how much capability it produces. Domains differ enormously in how much signal a token carries and in how much of their structure transfers: some are internally repetitive, so their tokens are highly redundant and extra weight buys little; others are dense and diverse, so each token teaches something new. Domains also interact — code appears to help general reasoning, and multilingual data shares structure across languages — so the capability produced by a domain is not confined to that domain. The consequence is that the weight producing the best code ability is generally not the weight matching your desired code capability share, and the only way to find it is to measure per-domain, not to reason from intent.
Because proportional allocation is the correct answer to a huge number of resource problems, and it is the only allocation that needs no measurement — you can derive it from your preferences alone, which makes it enormously attractive when the alternative is a sweep you cannot afford. It is also not badly wrong: proportional weights give a reasonable model, so nothing forces a rethink. It breaks when a domain fails to improve despite a large weight increase, or when raising one domain's weight unexpectedly improves a different domain's evaluation. Both are routine, and both are unexplainable if weights map onto capability one-for-one.
Measure per-domain loss separately, and check whether weight actually moves it:
# Hold out a validation split PER DOMAIN, not one pooled split.
# A pooled loss is dominated by whichever domain has the most
# weight, so it cannot tell you the mixture is wrong.
for name, split in val_by_domain.items():
print(name, evaluate(model, split))
# Then run short paired runs at different weights and compare
# the per-domain curves:
# run A: code weight w
# run B: code weight 2w (same total tokens)
# If code loss barely moves between A and B, the domain is
# saturated and the extra weight is being spent for nothing --
# and the other domains paid for it.Why does upsampling a small high-quality domain stop helping — and start hurting — beyond a certain number of passes, when the data is genuinely good?
- 1Upsampling a domain does not create new examples. It causes the same tokens to be visited more often, so the model sees identical text repeatedly.forced by · the weight controls sampling frequency from a fixed pool, and the pool does not grow
- 2On the first pass, reducing loss on a document requires learning structure that generalises to it, because the model has no record of the specific string.forced by · an unseen document can only be predicted from patterns learned elsewhere
- 3On later passes a cheaper route opens: storing the specific continuations of that exact text lowers loss on it more directly than any further generalisation would.forced by · gradient descent takes whichever route lowers the loss fastest, and memorisation is a very direct route for text already seen
- 4Capacity spent on memorised specifics is capacity not spent on transferable structure, and the model has finite capacity — so the cost is not merely neutral, it displaces.forced by · parameters used to store one document's idiosyncrasies are unavailable for representing patterns that generalise
- 5Therefore the return on additional passes over a fixed pool declines and eventually turns negative, and the turning point comes sooner for smaller pools, since fewer distinct tokens are exhausted faster relative to model capacity.forced by · the number of passes before memorisation dominates scales with pool size relative to capacity, not with data quality
Therefore quality does not exempt a domain from repetition limits. A small excellent corpus upsampled hard becomes a memorisation target, and its excellence is exactly what makes that tempting.
This yields a concrete prediction and a concrete practice. The prediction: the pass count at which a domain turns unproductive should depend on the ratio of its distinct token count to model capacity, so the same mixture weight that works for a small model will be too aggressive for a larger one — a mixture is not transferable across scales without rechecking. The practice: track per-domain validation loss for the upsampled domain specifically, and watch for the point where it continues to fall on training data while flattening on held-out data of the same domain. That divergence, not the global loss, is your repetition limit, and it is measurable within a short run rather than requiring a full one.
Stop picturing the mixture as slices of a pie. Picture each domain as a pool of distinct tokens, and the mixture weight as deciding how many times you walk through that pool during training. A large pool at low weight might get one partial pass; a small pool at high weight might get many.
The question for each domain is then not "does it deserve this share" but "how many passes is this pool worth before it stops teaching and starts being memorised" — a question with a real, measurable answer, unlike the share question, which has only opinions.
- Convert weights into passes: (weight x total tokens) / (distinct tokens in domain). That number is the one with meaning; the percentage is not.
- Hold out validation per domain. A pooled validation loss is dominated by the heaviest domain and is blind to exactly the failure you are looking for.
- Domains transfer. Improvements from one show up in others, so per-domain weights cannot be tuned independently and a change to one requires re-measuring all of them.
- A mixture is tied to a scale. Re-derive it when model size changes materially rather than carrying it over.
Fire this model the moment you see: a mixture stated only as percentages · a domain whose evaluation refuses to improve despite more weight · a small curated corpus being upsampled heavily · a mixture copied from a paper at a different model scale · a single pooled validation loss used to judge a multi-domain run.
How do you set mixture weights when you cannot afford a full sweep at target scale?
Start from a published mixture, then adjust it using per-domain validation from a short proxy run — that combination gets most of the available value for a small fraction of a sweep's cost. Adjust in the direction the proxy indicates, but expect to upsample small domains somewhat more at larger scale than the proxy suggested, since the repetition limit moves with capacity.
The one thing that is not optional is per-domain held-out evaluation. Without it, every option on this list degrades into guessing, because a pooled loss cannot tell you which domain a change helped or hurt. Build that first; the mixture decision is comparatively easy once you can see its effects, and impossible when you cannot.
10 · Retention scaffold
Recall
1. Difference between natural weight and sampling weight?
Natural weight = the fraction of the total available tokens each source contributes. Sampling weight = the fraction of actually sampled training tokens per step. Upsampling ratio = sampling / natural.
2. What is the upsampling ceiling and why does it exist?
Roughly 4 epochs on any single source. Past that, the model starts memorizing and downstream performance degrades. Sets a hard cap on how much you can stretch a small high-quality source.
3. Why does adding code to a general mix improve non-code benchmarks?
Empirical: adding 5–15% code improves reasoning-style benchmarks (BBH, MMLU) even though the tasks aren't code. The hypothesis is that code's syntactic discipline forces the model to learn compositional structure.
4. What is "annealing" or "cooldown" and why does it help?
The last ~5–10% of training tokens uses a much cleaner mix (wiki/books/arxiv/textbooks heavy). Gives 1–3 point downstream gains for near-zero compute cost. Basically the closest thing to a free lunch in pretraining.
5. How much does a 3-mix ablation cost at 20M params × 500M tokens?
About $2–3 on a rented A100 (~25 min per run × 3). Cheap enough that you should always do it before committing to a real training run.
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:
- 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.
- 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:
- Web is the plurality, always — 45–75% of tokens.
- Code is 10–20%, up from 5% in the RedPajama era, because it lifts reasoning.
- Books + academic + wiki together are 5–15% — they carry the long-form and factual load.
- Annealing cooldown on the last ~5% with the cleanest, highest-quality mix you own is now standard practice, essentially free MMLU + reasoning points.
Ma et al. (2024, At Which Training Stage Does Code Data Help LLMs Reasoning?) showed a clean surprise: adding code to the pretraining mix helps reasoning benchmarks most in early training and hurts slightly if you add too much in the last stage. The mechanism: code teaches the model syntactic composition and long-range coreference (matching braces, variable scoping), which transfers to natural-language reasoning. But at the annealing stage, you want prose that resembles your target output distribution, and code isn't that. Practical rule now: front-load code (17–20% early), fade it out in the cooldown (<5%). This is a schedule, not a static mix.
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
MixedPretrainDatasetper-step domain draw.