Article

LLM Pretraining Data Mixture in Production: Controlling Capability Bias with Mixture Manifest, Token Budget, and Domain Loss Feedback

A systematic guide to engineering data mixtures for LLM pretraining: using immutable Mixture Manifests to freeze domain weights and corpus versions, Token Budgets for precise accounting, and proxy model search, domain loss feedback, starvation protection, and replay gates to continuously calibrate capability structure. Turn mixture ratios from probabilistic parameters into auditable, reproducible training contracts.

Background: Same Total Data, Radically Different Capability Profiles

Large language model pretraining typically mixes web pages, code, papers, books, encyclopedias, math, and multilingual corpora. The most underestimated variable in engineering isn’t “whether you have this data,” but how many effective training tokens each domain ultimately contributes.

With the same one trillion tokens: adjusting the code ratio from 5% to 20% can reallocate the model’s programming ability, natural language fluency, and multilingual performance. In continual pretraining, an overly high proportion of specialized domains can even cause general capability forgetting.

Therefore, data mixture cannot be just a list of weights in the launch parameters. It must be governed as an auditable, reproducible, grayscale-deployable, and rollback-capable training artifact.

Core Principle: Elevate “Mixture” from Probabilistic Parameters to a Training Contract

1. Freeze the Input Space with an Immutable Mixture Manifest

Every training run must be bound to a single, unique Mixture Manifest. The Manifest should at minimum record:

  • Domain identifier and semantic definition;
  • Data snapshot, filtering version, and tokenizer version;
  • Target weight, minimum weight, and maximum weight;
  • Planned token count, allowed repeat epochs, and exhaustion strategy;
  • Sampling algorithm, random seed, and training stage;
  • Proxy experiment, approval record, and artifact hash.
mixture_id: pretrain-mix-2026-07-v3
total_token_budget: 1000000000000
sampler: weighted_token_sampler_v2
seed: 20260728
domains:
  - id: web_en
    snapshot: web-en-2026-06-clean-v4
    tokenizer: tokenizer-sha256-7a91
    target_weight: 0.42
    min_weight: 0.30
    max_weight: 0.50
  - id: code
    snapshot: code-2026-06-license-v2
    tokenizer: tokenizer-sha256-7a91
    target_weight: 0.18
    min_weight: 0.10
    max_weight: 0.25
  - id: math
    snapshot: math-2026-05-dedup-v3
    tokenizer: tokenizer-sha256-7a91
    target_weight: 0.10
    min_weight: 0.05
    max_weight: 0.16

Once published, a Manifest must never be modified in place. Any change to weights, snapshots, or sampling rules must produce a new version.

2. Account with Effective Tokens, Not Document Counts

Let the total training budget be T, and the weight for domain i be p_i. Its planned budget is:

T_i = round(T × p_i)

The production system should also track:

MetricMeaning
consumed_tokens_iTokens consumed
accepted_tokens_iTokens that passed filtering
repeated_tokens_iTokens read multiple times
remaining_unique_tokens_iRemaining unique tokens

Document count is unsuitable as the primary ledger. Average document length, filtering ratio, language token density, and sequence packing efficiency vary significantly across domains. Sampling by document can easily cause the actual token ratio to drift from the configuration.

3. Distinguish Static vs. Dynamic Mixtures

Static mixtures fix weights for the entire training run. They are easy to reproduce and audit, making them suitable for first-time large-scale training.

Dynamic mixtures adjust weights by training stage—e.g., prioritizing general web coverage early on, then increasing code, math, or high-quality long texts in later stages. This is more flexible but requires:

  • Stage boundaries;
  • Weight change rate limits;
  • Domain weight upper and lower bounds;
  • A fixed replay set;
  • A general capability forgetting gate.

Recent work like RegMix-D uses the full loss trajectory from proxy training (not just the final point) to predict different mixtures for different stages. This direction is worth watching, but paper results cannot be blindly adopted as default production configurations for arbitrary models.

DoReMi: Focus on Per-Domain Excess Loss

DoReMi uses the difference in domain loss between a reference model and a proxy model for reweighting. Simplified:

excess_loss_i = loss_proxy_i - loss_reference_i

The system focuses on domains that are not yet well-learned, rather than simply upweighting high-entropy, noisy domains. In the paper, the proxy model is used to search for weights, which are then transferred to a larger target model.

For production deployment, you must bind:

  • Reference Model fingerprint;
  • Proxy Model configuration;
  • Domain validation set version;
  • Loss aggregation method;
  • Smoothing window and weight clipping rules.

Without this information, saving only the final weights makes it impossible to explain why the next run didn’t produce the same result.

RegMix and Data Mixing Laws: Predict Large-Scale Results with Small-Scale Experiments

RegMix trains multiple small models with different data mixtures, fits a “mixture-to-performance” mapping, and then searches for candidate ratios. Data Mixing Laws study the predictable relationship between mixture, model scale, training steps, and performance.

These methods are suitable for costly large-model pretraining, but proxy experiments must cover real uncertainty:

  • Different random seeds;
  • Different token horizons;
  • Different model widths or depths;
  • Interaction terms between domains;
  • Trade-offs between target tasks and general tasks.

Proxy models should provide candidate regions, not a single point answer that requires no review.

Engineering Implementation: Building a Replayable Data Mixture Control Plane

1. Define Domains Before Optimizing Weights

Data domains must have stable semantics. Do not equate a domain directly with a file directory or vendor name. It is recommended to maintain multiple dimensions:

DimensionExample
Source domainWeb, books, code repos, papers
Capability domainMath, code, law, multilingual
Quality tierHigh, medium, low confidence
Compliance tierLicense type, regional restrictions, sensitivity level

A single document can carry multiple tags, but the primary domain used for sampling must be unambiguous; otherwise, weight statistics will double-count.

2. The Sampler Must Be Deterministically Replayable

Given a Manifest, data snapshot, seed, global step, and data parallel size, the sampler must produce a consistent sequence of domain choices and samples.

from dataclasses import dataclass
import random

@dataclass(frozen=True)
class Domain:
    name: str
    weight: float

def choose_domain(domains: list[Domain], seed: int, global_sample_id: int) -> str:
    rng = random.Random(f"{seed}:{global_sample_id}")
    names = [d.name for d in domains]
    weights = [d.weight for d in domains]
    return rng.choices(names, weights=weights, k=1)[0]

A real implementation must also handle multi-rank sharding, checkpoint recovery, and data exhaustion, but the core requirement remains: sample selection must not depend on untraceable, process-local random state.

3. Maintain a Token Ledger for Each Domain

Training monitoring should at least include:

  • Planned vs. actual weight;
  • Deviation per hour, per billion tokens;
  • Domain data exhaustion rate and repeat rate;
  • Domain-level training loss and fixed validation loss;
  • Data read failure, filter failure, and decode failure rates;
  • Effective token cost per domain.

Correct weights do not guarantee correct execution. If a domain frequently fails to read, the system may silently compensate with other domains, causing the actual training ratio to drift.

4. Implement Domain Starvation and Oversampling Protection

Low-weight domains may not appear in a batch for a long time, causing their metrics to become stale. Small, high-weight domains may be repeated many times quickly.

The following protection mechanisms should be in place:

ParameterPurpose
min_tokens_per_windowMinimum token supply within a time window
max_repeat_epochsMaximum repeat epochs
max_weight_deltaMaximum weight change between adjacent stages
staleness_limitMaximum staleness for domain metrics
fallback_policyAction when a domain is unavailable: stop, reduce weight, or switch to a backup snapshot

Do not default to “proportionally redistribute tokens from unavailable domains to others.” This compensation breaks the training contract.

5. Dynamic Weight Adjustment Must Follow a Release Process

Dynamic data scheduling should not allow monitoring scripts to directly modify training parameters. The recommended workflow:

Domain-level signal collection → Candidate weight calculation → Upper/lower bound and rate-of-change validation
→ Simulation on a fixed replay set → Generate new Mixture Manifest → Approval
→ Switch at a stage boundary → Monitor and allow rollback

A running training job should only consume published Manifests and must not accept unversioned temporary weights.

Applicable Scenarios

This governance framework is particularly suitable for:

  • Multi-domain general model pretraining;
  • Continual pretraining for specialized domains like math, code, and finance;
  • Language ratio control for multilingual models;
  • Training platforms with continuously added data and frequently changing quality tiers;
  • Enterprise model engineering that requires auditing data contributions and capability change sources.

Small-scale, single-domain fine-tuning usually does not require a complex dynamic mixing system, but it is still recommended to use a Manifest to freeze the data snapshot and token budget.

Common Misconceptions

Misconception 1: More high-quality data is always better

High quality in a single domain does not guarantee optimal overall mixture. Domains exhibit complementarity, competition, and stage-dependent effects. This must be confirmed through proxy experiments and multi-capability evaluation.

Misconception 2: Only monitor the overall training loss

The overall loss may decrease while a low-weight domain has already degraded. You must retain domain-level loss, fixed validation sets, and capability group metrics.

Misconception 3: Configuring by sample ratio equals token ratio

Token density varies significantly between long texts, code, and different languages. The actual contribution must be verified against the Token Ledger.

Misconception 4: Optimal proxy model weights are unconditionally transferable

Changes in architecture, scale, token horizon, tokenizer, and optimizer can all shift the optimal region. Proxy results can only narrow the search space.

Misconception 5: More frequent dynamic adjustment is smarter

The loss signal itself is noisy. Overly frequent adjustments make both the data distribution and the optimization process non-stationary, making diagnosis and reproduction difficult.

Launch Checklist

  • Each domain has a stable ID, definition, and owner;
  • Data snapshots, filters, and tokenizers are all hashed and frozen;
  • Mixture Manifest is immutable and rollback-capable;
  • Sum of weights, upper/lower bounds, and stage plans pass validation;
  • Planned vs. actual tokens can be reconciled per domain;
  • Sampler is replayable across multiple ranks and after checkpoint recovery;
  • Starvation, repeat rate, and data exhaustion gates are set;
  • Proxy experiments include multiple seeds and token horizons;
  • General, specialized, multilingual, and safety capabilities are regressed separately;
  • Dynamic weight adjustments only occur at defined stage boundaries;
  • The effective Manifest can be restored when training anomalies occur;
  • Every weight change has an approval, rationale, and evidence link.

References

  1. DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining
  2. Stanford CRFM: DoReMi
  3. RegMix: Data Mixture as Regression for Language Model Pre-training
  4. RegMix-D: Dynamic Data Mixing via Proxy Training Trajectories
  5. Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance
  6. Megatron-LM Data Arguments and Blended Dataset Support
  7. DataComp-LM

FAQ

Should data mixture be calculated by document count, sample count, or token count?
In production, use effective training tokens as the primary ledger. Document or sample counts can severely misrepresent actual training contribution due to differences in document length, filtering rates, and packing efficiency.
Can domain weights be modified dynamically during training?
Yes, but only through a versioned stage plan with min/max weights, smooth transitions, and replay validation. Do not adjust weights online based on single-step loss fluctuations.
Can optimal mixture ratios found on a small model be directly applied to a large model?
Only as candidates. Proxy models significantly reduce search cost, but you still need short-run validation and capability regression on the target architecture and target token horizon.
How do you determine if a domain is being oversampled?
Monitor actual token proportion, unique token coverage, repeat epochs, near-duplicate rate, and validation loss for that domain simultaneously. Weight configuration alone cannot reveal the problem of small corpora being cyclically consumed.