Data mixing: sampling, exposure, proxy ablations and curriculum

A mixture is a sampling policy

Available data are the pantry; the training sampler is the diet. A source's fraction of storage does not determine how often its tokens contribute gradients. Nor does a desired capability percentage determine the right sampling weight: domains differ in redundancy, noise, task relevance and transfer. Treat “code teaches reasoning” or “books teach coherence” as hypotheses to test, not guaranteed capabilities attached to source labels.

Bring the versioned source groups/token shards of DL046 and the budget arithmetic of DL047. This lesson derives natural weight, sampling weight and exposure; implements a length-aware seeded shard sampler; runs a complete bounded CPU mixture comparison; explains DoReMi's actual reference/proxy/main-model order; and treats curriculum as an experiment with a token clock, not a free gain.

1 · Domains, rights and evaluation targets

Web, code, books, scientific papers, Wikipedia, Q&A and educational text are useful corpus categories. They are not mutually exclusive semantic labels: a web crawl can contain code and textbook-like prose, and a paper can contain both. A directory calledbooks does not establish long-context quality or rights. Record source identity, permissions, language, extraction quality and duplicate groups within and across buckets.

Do not assume Books3, Gutenberg, PG-19, The Stack, arXiv, Wikipedia or an instruction collection share a blanket license. A model family using a named source is not a grant of permission for a new use. Synthetic instructions also require generator/prompt/version/usage lineage and contamination checks.

Define held-out evaluation before choosing weights: per-domain loss under the same frozen tokenizer; a fixed target mixture for aggregate loss; and task metrics that reflect the intended deployment. A pooled loss measured under each candidate's own training mixture compares different objectives and can reward an easier domain without improving target performance.

2 · Natural weights, sampling weights and effective passes

Let U-i be unique eligible tokens in domain i, U their total, p-i=U-i/U the natural share, w-i the target token sampling share, and D the total processed training tokens. Then expected processed domain tokens are w-i D and effective passes are:

The upsampling factor f equals effective passes only when D=U. Sampling with replacement means an effective pass is an expected exposure count, not a guarantee that every token was seen once. Duplicate-rich pools overstate unique U and hide repeated exposure.

For available millions of tokensweb5000, code300, books200, arxiv50, wiki30, total U=5580M. ArXiv natural share is50/5580≈.8961%. A2.5% sample share has an upsampling factor2.79, but at D=1B its expected passes are.5; at15B they are7.5. Wiki at2% has.6667 and10 passes. This is an exposure warning, not a universal four-epoch prohibition.

The original seven-source weights summed to.975, and its five-row table silently omitted other sources. A sampler that normalizes them changes every displayed percentage. Use an explicit validated complete vector; do not label a seven-bucket plan “five domains.”

import math, numpy as np
 
def normalize_weights(weights, domains):
    if set(weights)!=set(domains):raise ValueError('weights must name every domain exactly')
    values=np.array([weights[d] for d in domains],dtype=float)
    if not values.size or not np.isfinite(values).all() or (values<0).any() or values.max()<=0:
        raise ValueError('finite nonnegative weights with positive total required')
    # Scale first: two valid 1e308 weights must not overflow their sum.
    scaled=values/values.max()
    return scaled/scaled.sum()
 
supply={'web':5000e6,'code':300e6,'books':200e6,'arxiv':50e6,'wiki':30e6}
# Explicit five-domain teaching proposal, not an attributed published recipe.
base={'web':.70,'code':.15,'books':.08,'arxiv':.04,'wiki':.03}
assert math.isclose(sum(base.values()),1.)
assert math.isclose((.025/(50/5580)),2.79)
for total in (1e9,15e9):
    print('processed tokens',total,'effective passes',
          {d:base[d]*total/supply[d] for d in supply})
old={'web':.67,'code':.045,'books':.045,'arxiv':.025,'wiki':.02,'qa':.02,'c4':.15}
assert not math.isclose(sum(old.values()),1.)
print('old weight sum',sum(old.values()))

Changing vocabulary/tokenization changes the token unit, so exposure comparisons require a common tokenizer. With variable-length examples, drawing domain i with probability w-i gives a document share, not necessarily a token share. Use fixed valid-token blocks, length-aware quotas or explicit accounting to achieve a token budget.

3 · Correct weighted sampling across local shards

Choose domain by weights, then choose a shard proportional to its number of eligible windows, then a uniformly random valid start. Equal-probability shard choice overweights small shards. A file of length L has L−T valid T+1-token windows, with starts0 through L−T−1; the high endpoint for NumPy's exclusive integer draw is L−T. The original code subtracted one too many and rejected the last start, including the only window in an exactly T+1-long file.

This implementation uses raw little-endianuint16 document-local shards. A real packed shard needs the DL046 document-offset table to enumerate only intended windows. It maps files once per process and validates lengths. Worker RNGs derive from seed, epoch, distributed rank and worker ID. That makes streams distinct, not disjoint samples—sampling with replacement can legitimately overlap.

from pathlib import Path
import copy, torch
from torch.utils.data import IterableDataset,DataLoader,get_worker_info
 
torch.set_num_threads(1)
 
class DomainStream:
    def __init__(self,domain_shards,weights,context=4,seed=0,epoch=0,rank=0,worker=0):
        if context<=0:raise ValueError('positive context required')
        self.domains=sorted(domain_shards);self.context=context
        self.probs=normalize_weights(weights,self.domains)
        self.rng=np.random.default_rng(np.random.SeedSequence([seed,epoch,rank,worker]))
        self.arrays={};self.counts={};self.draws=0
        for d in self.domains:
            if not domain_shards[d]:raise ValueError('empty domain')
            arrs=[]
            for filename in domain_shards[d]:
                if Path(filename).stat().st_size % 2:raise ValueError('truncated uint16 shard')
                arr=np.memmap(filename,dtype='<u2',mode='r')
                if len(arr)<context+1:raise ValueError('shard too short')
                arrs.append(arr)
            self.arrays[d]=arrs
            self.counts[d]=np.array([len(a)-context for a in arrs],dtype=float)
    def __next__(self):
        index=int(self.rng.choice(len(self.domains),p=self.probs));d=self.domains[index]
        counts=self.counts[d]
        shard=int(self.rng.choice(len(counts),p=counts/counts.sum()))
        start=int(self.rng.integers(0,int(counts[shard])))
        chunk=np.array(self.arrays[d][shard][start:start+self.context+1],dtype=np.int64)
        self.draws+=1
        return torch.from_numpy(chunk[:-1]),torch.from_numpy(chunk[1:]),d,shard,start
    def state_dict(self):
        return dict(rng=copy.deepcopy(self.rng.bit_generator.state),draws=self.draws)
    def load_state_dict(self,state):
        self.rng.bit_generator.state=copy.deepcopy(state['rng']);self.draws=state['draws']
 
class MixedPretrainDataset(IterableDataset):
    def __init__(self,shards,weights,context=4,seed=0,epoch=0,rank=0):
        super().__init__();self.args=(shards,weights,context,seed,epoch,rank)
    def __iter__(self):
        worker=get_worker_info(); wid=0 if worker is None else worker.id
        stream=DomainStream(*self.args,worker=wid)
        while True:yield next(stream)
 
folder=Path('batch7-dl-pretraining-fixtures/mixing048');folder.mkdir(parents=True,exist_ok=True)
shards={}
for i,d in enumerate(sorted(base)):
    shards[d]=[]
    for j,length in enumerate((5,13)):
        path=folder/f'{d}-{j}.bin';(np.arange(length,dtype=np.uint16)+i*20).astype('<u2').tofile(path)
        shards[d].append(str(path))
stream=DomainStream(shards,base,seed=48)
counts={d:0 for d in base}; within={d:[0,0] for d in base};starts=[]
for _ in range(10000):
    x,y,d,s,start=next(stream);counts[d]+=1;within[d][s]+=1;starts.append(start)
    assert torch.equal(x[1:],y[:-1])
for d in base:
    assert abs(counts[d]/10000-base[d])<.025
    assert abs(within[d][0]/counts[d]-.1)<.08  # 1 versus9 eligible windows
saved=stream.state_dict(); expected=[next(stream) for _ in range(5)]
stream.load_state_dict(saved); actual=[next(stream) for _ in range(5)]
assert all(torch.equal(a[0],b[0]) and a[2:]==b[2:] for a,b in zip(expected,actual))
a=DomainStream(shards,base,seed=48,worker=0)
b=DomainStream(shards,base,seed=48,worker=1)
assert [next(a)[2:] for _ in range(12)] != [next(b)[2:] for _ in range(12)]
batch=next(iter(DataLoader(MixedPretrainDataset(shards,base),batch_size=3,num_workers=0)))
assert batch[0].shape==(3,4)
print('domain counts',counts,'shard-window weighting',within,'single-stream resume pass')

A controlled weight-change check keeps the same files and seed, doubles only code's target share from .15 to .30, and rescales the other shares proportionally. Rebuilding the stream applies the new policy; this is not a live mutation of worker copies.

raised={d:(.30 if d=='code' else w*.70/.85) for d,w in base.items()}
assert math.isclose(sum(raised.values()),1.)
code_fractions=[]
for policy in (base,raised):
    probe_stream=DomainStream(shards,policy,seed=481)
    code_fractions.append(sum(next(probe_stream)[2]=='code' for _ in range(4000))/4000)
assert abs(code_fractions[0]-.15)<.025 and abs(code_fractions[1]-.30)<.025
assert code_fractions[1]>code_fractions[0]
print('controlled code-share change',code_fractions)

The state dictionary resumes one stream with the same manifest/config. Multiworker prefetch introduces buffered samples and multiple cursors; saving only a parent RNG cannot reproduce its exact future batches. For strict replay use a counter-indexed global sample plan or checkpoint every worker cursor and consumed/buffered order. Recreate workers when changing epoch/weights unless using an explicitly synchronized shared schedule; mutating a parent object does not update existing worker copies.

Sampling need not physically duplicate files. But memory is not O(1) in every variable: file maps, shard metadata, worker buffers and indexes have costs. Concatenate-and-index alternatives can also avoid copying through virtual indexes. Choose by throughput and reproducibility, not an invented impossibility claim.

4 · Sampling probability is not loss weighting

Suppose target objective is a weighted sum of per-domain expected losses with weights π-i, while the sampler draws domains with q-i. Multiplying each sampled example's loss byπ-i/q-i is an unbiased importance-weighted estimator when q is positive whereverπ is positive. If q=π, no extra domain multiplier is required; multiplying byπ again would square the effect. Variance can explode when q is tiny. Token masks and variable lengths require a valid-token denominator, not a naive mean over examples.

This is why a50/50 sampling policy is not automatically50/50 capability and not the same as drawing natural data then multiplying domain losses. Log realized tokens by domain, gradient weighting and validation target weights separately.

5 · Complete bounded CPU ablation

A real20M-model/500M-token sweep was not run. The following complete synthetic bigram experiment demonstrates the selection protocol: identical initialization, equal updates/batch/context, fixed training-token budget, independent held-out generation per domain, and a fixed aggregate validation objective. It is not evidence that one mixture will win for a transformer, a language or a benchmark.

import torch.nn.functional as F
 
def synthetic_tokens(domain,batch,length,generator):
    # Two authored Markov-like domains with incompatible next-token tendencies.
    x=torch.randint(0,7,(batch,length),generator=generator)
    y=(x+(1 if domain=='prose' else 3))%7
    return x,y
 
def compare_mixes(steps=100):
    candidates={'prose-heavy':{'prose':.9,'code':.1},
                'balanced':{'prose':.5,'code':.5},
                'code-heavy':{'prose':.2,'code':.8}}
    torch.manual_seed(48);initial=torch.randn(7,7)*.01
    target={'prose':.5,'code':.5};results={}
    for name,weights in candidates.items():
        logits=torch.nn.Parameter(initial.clone());optim=torch.optim.SGD([logits],lr=.8)
        domain_rng=np.random.default_rng(48);train_rng=torch.Generator().manual_seed(480)
        counts={'prose':0,'code':0}
        for step in range(steps):
            d='code' if domain_rng.random()<weights['code'] else 'prose'
            x,y=synthetic_tokens(d,16,8,train_rng)
            loss=F.cross_entropy(logits[x].reshape(-1,7),y.reshape(-1))
            optim.zero_grad(set_to_none=True);loss.backward();optim.step();counts[d]+=x.numel()
        scores={}
        with torch.no_grad():
            for j,d in enumerate(target):
                g=torch.Generator().manual_seed(4800+j)
                x,y=synthetic_tokens(d,64,8,g)
                scores[d]=F.cross_entropy(logits[x].reshape(-1,7),y.reshape(-1)).item()
        results[name]=dict(valid=scores,target_loss=sum(target[d]*scores[d] for d in target),
                           processed_tokens=sum(counts.values()),tokens_by_domain=counts)
    return results
 
ablations=compare_mixes()
assert len({v['processed_tokens'] for v in ablations.values()})==1
assert ablations['code-heavy']['valid']['code']<ablations['prose-heavy']['valid']['code']
print('synthetic mixture ablation only',ablations)

For the original proposed large sweep, compute per-run work6*20e6*500e6=6e16 and total three-run work1.8e17 FLOPs. At an assumed124.8TFLOP/s sustained rate the total is1442 seconds, about24 minutes—not24 minutes per run followed by another multiplication by3. Actual throughput for a small model and task evaluation may be much lower. There is no verified under20-dollar or2.50-dollar completion promise.

A complete future large-model protocol is: freeze corpus/tokenizer/manifests and all candidate mixtures; allocate a sweep budget separate from the main run; build identical architecture/initialization and schedules; train each candidate for the same valid-token count; checkpoint by token clock; evaluate fixed per-domain held-out sets and fixed downstream task versions; repeat seeds; select on validation only; test the chosen policy once on untouched tests; confirm ranking at an intermediate scale before the expensive run. No unexplained train()/eval_downstream() fragment stands in for this protocol.

6 · DoReMi: reference, robust proxy, then main model

The primary paper's order matters:

1. Train a small reference model with an initial mixture.
2. Train a small proxy under a domain distributionally robust objective using excess token loss relative to the reference.
3. Average learned domain weights over training and use them to train a fresh larger main model with the standard training procedure.

Algorithm1 forms per-domain means of clipped token-level differencesmax(proxy_loss-reference_loss,0), applies exponentiated-gradient domain weight updates, normalizes and smooths toward uniform, then updates proxy parameters under the weighted objective. It does not simply upweight whichever domain has the largest absolute perplexity: inherently hard/noisy domains can have high reference loss too. Taking the difference of domain means and clipping afterward is not the same operation.

def doremi_weight_step(weights,proxy_losses,reference_losses,eta=.1,smoothing=.001):
    domains=list(weights)
    if not math.isfinite(eta) or eta<0 or not math.isfinite(smoothing) or not 0<=smoothing<=1:
        raise ValueError('invalid update parameters')
    if set(proxy_losses)!=set(domains) or set(reference_losses)!=set(domains):
        raise ValueError('losses must name every domain exactly')
    p=normalize_weights(weights,domains)
    if (p<=0).any():raise ValueError('this update requires strictly positive domain weights')
    excess=[]
    for d in domains:
        a=np.asarray(proxy_losses[d],float);b=np.asarray(reference_losses[d],float)
        if a.ndim!=1 or a.shape!=b.shape or not a.size:
            raise ValueError('matching nonempty one-dimensional token losses required')
        if not np.isfinite(a).all() or not np.isfinite(b).all() or (a<0).any() or (b<0).any():
            raise ValueError('finite nonnegative token losses required')
        clipped=np.maximum(a-b,0)
        excess.append(float(np.sum(clipped/clipped.size)))
    with np.errstate(over='ignore',invalid='ignore'):
        logw=np.log(p)+eta*np.array(excess)
    if not np.isfinite(logw).all():raise ValueError('weight update overflow')
    logw-=logw.max()
    updated=np.exp(logw);updated/=updated.sum()
    updated=(1-smoothing)*updated+smoothing/len(domains)
    return dict(zip(domains,updated)),dict(zip(domains,excess))
 
w,e=doremi_weight_step({'a':.5,'b':.5},{'a':[4,1],'b':[2,2]},
                       {'a':[2,3],'b':[2,2]})
assert e['a']==1 and e['b']==0 and w['a']>.5
assert max(np.mean([4,1])-np.mean([2,3]),0)==0  # wrong aggregation differs
print('DoReMi weight-update mechanism only',w,e)

The function accepts finite nonnegative per-token negative log-likelihoods and strictly positive domain priors. It rejects NaNs, absent domains and overflowing updates rather than writing a corrupt mixture. Token masks must be applied before supplying these one-dimensional loss arrays. This validates the weight update, not a full DoReMi reproduction. The paper's small-to-large results motivate proxy experiments but do not establish85% ranking agreement across arbitrary scales or adoption by private labs. Proxy/reference training and averaging cost compute; decide whether the expected benefit justifies it for the actual run.

7 · Curriculum on a token clock

A curriculum changes weights over training. “Annealing” can refer to learning-rate decay and often accompanies data-mix changes; those are two variables, not synonyms. To isolate effects compare unchanged mix/changed LR, changed mix/unchanged LR and the combined recipe under the same total budget. A tail branch costs real tokens unless it replaces part of a planned run.

A two-stage proposal uses a broad early mix and a smaller high-quality/domain-focused tail. Count total exposure as the sum over phases ofphase_tokens * phase_weight / unique_tokens. If instructions have3% weight only during the last5%, their whole-run share is.15%, not3%. Normalize every phase independently; specify a transition token index and save it with checkpoint state.

def phase_exposure(phases,supply):
    if not supply or any(not math.isfinite(u) or u<=0 for u in supply.values()):
        raise ValueError('finite positive unique-token supplies required')
    seen={d:0. for d in supply}
    for tokens,weights in phases:
        if not math.isfinite(tokens) or tokens<0:raise ValueError('finite nonnegative token budget required')
        p=normalize_weights(weights,list(supply))
        for d,w in zip(supply,p):seen[d]+=tokens*w
    return {d:seen[d]/supply[d] for d in supply}
 
tail={'web':.30,'code':.05,'books':.25,'arxiv':.15,'wiki':.25}
print('two-stage exposure at2B',phase_exposure([(1.8e9,base),(.2e9,tail)],supply))
assert math.isclose(.03*.05,.0015)
variants={'baseline':base,
          'code-heavy':{'web':.55,'code':.30,'books':.08,'arxiv':.04,'wiki':.03},
          'book-heavy':{'web':.53,'code':.15,'books':.25,'arxiv':.04,'wiki':.03}}
for total in (1e9,15e9):
    for name,weights in variants.items():
        exposure=phase_exposure([(.9*total,weights),(.1*total,tail)],supply)
        print('stretch phase-aware passes',name,total,exposure,
              'over four passes, diagnostic not a hard limit',[d for d,e in exposure.items() if e>4])
assert math.isclose(phase_exposure([(.9e9,base),(.1e9,tail)],supply)['wiki'],.052e9/30e6)

The Llama3 report describes approximately50% general-knowledge,25% mathematical/reasoning,17% code and8% multilingual categories and explicitly studies high-quality code and math annealing. It does not support a universal instruction to fade code below5% at the end. Ma et al.'s stage study compares pretraining/instruction-tuning and dynamic mixtures; it is not a universal20%-early/5%-late law.

A fixed mixture remains a legitimate baseline. Length, perplexity and difficulty ordering can interact with optimization and coverage; neither “all curricula are dead” nor “every serious model gets free1–3 points” is defensible. Check forgetting and every domain, not only the tail's favored benchmark.

8 · Reading modern mixture reports critically

LLaMA/RedPajama's named sources provide a historical reference, not a universally optimal pantry. RedPajama-V1 collection proportions, LLaMA sampling proportions and actual token exposure need not be identical. Mistral/MPT/Falcon recipes cannot be inferred from the same seven labels, and a web-only corpus can itself contain diverse material.

DCLM demonstrates strong results from a filtered Common Crawl source under its benchmark setup and examines additional-source mixing; it does not prove explicit buckets never matter. FineWeb-Edu changes which web records are retained, so it is a selection policy even when called a one-source mix. DeepSeek-V3 reports14.8T tokens and changes in mathematical/programming content; do not invent an undisclosed full ratio table or universal annealing percentage from that description.

Repeated high-quality text can still have diminishing returns. But the exact number of useful passes depends on corpus, model, optimizer, task and schedule. The small-pool memorization story is a risk model, not a proof that the same weight must always become too high—or too low—when model size grows. Re-evaluate cross-scale transfer instead of teaching contradictory rules in different paragraphs.

If a book-heavy model becomes overly narrative, examine style-conditioned evaluations and sampling; if arXiv weighting improves STEM but harms commonsense, it is a tradeoff to measure, not proof that scientific prose cannot transfer general reasoning. The original workplace anecdotes and exact score changes are replaced by these diagnostic cases.

9 · Recall and complete answers

1. Natural versus sampling weight? p=U-i/U describes available unique tokens; w describes desired processed-token share. Effective passes are wD/U-i. For arXiv50M/5580M, natural share≈.896%, a2.5% target gives factor2.79 but only.5 passes in a1B run.
2. What is the upsampling ceiling? There is no universal four-epoch ceiling. Monitor repeated exposure, clean per-domain validation and train/validation divergence under the actual scale/schedule.
3. Why may code help non-code reasoning? Studies observe transfer in specified settings; syntax/composition is a proposed mechanism. The effect must be tested with controlled mixtures and non-code tasks, not inferred from a fixed5–15% recipe.
4. What went wrong if STEM rose3 points and commonsense fell5? A target tradeoff changed. Inspect confidence intervals, contamination, domain loss and a predeclared target evaluation; neither metric alone determines deployment utility. Do not invent causality from the two numbers.
5. Is cooldown free? It consumes or reallocates tokens and interacts with LR. Count phase exposures and evaluate forgetting; no guaranteed gain or universal5–10% tail exists.
6. Why unequal shard probabilities? A shard with one valid window should not be as likely as one with nine if the goal is uniform valid-window sampling. Weight shards by eligible-window count, then include the last valid offset.
7. What does a worker seed fix? It separates RNG streams; it does not ensure disjoint data or exact multiworker resume. Parent mutations do not propagate automatically to worker copies.
8. How much does the proposed sweep cost? Its simplified training arithmetic is1.8e17 FLOPs total. Convert using measured sustained throughput and actual overhead/price; the local bigram test is not a paid GPU timing result.
9. Stretch solution: define baseline/code-heavy/book-heavy weights that each sum to1; compute phase exposures at1B and15B; identify heavily reused small sources without declaring a hard cutoff; run a fixed-objective proxy comparison and test scale transfer. Store the selected policy and token-clock schedule with the model.

Tomorrow derive passes from a pie chart; after a week explain why training loss from different mixtures cannot be compared as if the validation distribution were unchanged.

Sources and optional navigation

DoReMi Algorithm1, LLaMA, Llama3, DeepSeek-V3, data-constrained scaling, code-stage study, DCLM, FineWeb. The old FineWeb URL2402.16827 was not the FineWeb paper; use2406.17557. Optional inherited resources: State of GPT, RedPajama-V1 announcement.

Course map · Previous DL047 · Next DL049.