Hmm… for now, it seems there are several things that make the numbers difficult to interpret:
My direct answer would be:
- No, a dataset’s weight in the final validation objective does not have to equal its training sampling probability.
- I would not start by giving each dataset a separate optimizer or learning rate.
- I would first keep one global LR and separate the effects of:
- the target/evaluation weights,
- the source sampling probabilities,
- the per-source loss weights,
- the number of tokens that actually carry loss,
- the loss reduction across microbatches/devices,
- repetition and optimizer/scheduler state.
A real Pareto trade-off may remain after that: an update that improves one domain may worsen another. However, the curves alone do not yet show that the datasets intrinsically require different learning rates. The same observation can also be produced by unequal effective exposure, masking/truncation, repetition, different sequence-length distributions, optimizer history, or an objective that is intentionally willing to sacrifice a low-weight domain.
A reasonable default route
I would use the following order before changing the optimizer architecture.
-
Decide whether the weighted average is really the complete objective.
If a 2%-weighted domain is allowed to regress when the other 98% improves more, then that regression may be consistent with the objective.
If every domain must remain near its starting loss, then the actual goal is closer to:
- a weighted objective with per-domain regression constraints,
- a minimum replay/sampling floor,
- or a robust multi-objective objective,
rather than a plain weighted average.
-
Measure the realized objective, not only the configured mixture.
For each source, log at least:
- sampled microbatches/examples,
- non-padding input tokens,
- tokens where
labels != -100,
- unique examples/documents,
- estimated repeat count,
- truncation rate,
- and the exact loss denominator.
-
Inspect one complete optimizer step.
This is especially important with gradient accumulation or distributed training. Record:
- which sources appeared in each microbatch,
- valid target tokens per source,
- per-source loss sums and reported mean losses,
- the denominator used across the accumulated batch,
- whether valid-token counts are synchronized across ranks,
- and the gradient norm before/after clipping.
-
From the same checkpoint, compare only a few controlled choices.
Keep the optimizer-state policy, LR schedule, total token/compute budget, evaluation sets, and loss normalization fixed:
| Condition |
What it tests |
| Current mixture |
Baseline |
| Increase the problem source’s sampling |
Insufficient source exposure |
| Keep sampling fixed but increase its loss weight |
Insufficient gradient contribution |
| Optional simple staged mixture |
Whether the preferred mixture changes during training |
-
Only reconsider source-specific optimization if the effect remains.
If the same source still requires a materially different step scale after the accounting, preprocessing, repetition, and schedule effects are controlled, then source-dependent loss scaling, alternating phases, adapters/branches, or more advanced mixture methods become reasonable candidates.
I would still keep one global LR for the first controlled comparison. A separate LR per source in a fully shared dense model is difficult to interpret because every source updates the same parameters and Adam moments. Multiplying a source’s loss can change its update scale, but that is not generally identical to maintaining a separate optimizer state for that source.
A small sanity check
As a small diagnostic—not a production recipe—I tried a 135M proxy with a 95/5 FineWeb-Edu/OpenThoughts-114k source mixture.
The nominal 5% source probability resulted in approximately:
- 6.0% of sampled microbatches,
- 7.1% of input tokens,
- 2.4% of tokens with valid labels.
The difference came mainly from assistant-only masking and the much longer conversational examples. Doubling the source sampling brought the valid-label-token share to roughly 4.2%.
However, the custom loop averaged each microbatch’s mean loss with equal microbatch weight. Therefore, 2.4% was the share of valid target tokens, not necessarily 2.4% of the optimizer-level gradient weight. This distinction matters.
In one short seed:
- both 2x oversampling and 2x loss weighting improved the OpenThoughts validation loss;
- neither hurt the FineWeb-Edu validation loss during the short mixed run;
- source-only updates produced a small regression on the other source;
- mixed updates improved both;
- and the initial last-block gradient cosine was positive rather than strongly conflicting.
I would interpret that only as evidence that:
- serialization/masking can substantially change the apparent mixture;
- oversampling and loss weighting are both worth distinguishing;
- and “the smaller source is continually overwritten” should be tested rather than assumed.
It does not establish the production-optimal LR, sampling ratio, or optimizer. The proxy used a different model, short 512-token windows, a new optimizer state, one seed, two sources, and a short horizon.
What can “2–5%” mean?
There are several different percentages that can all be described informally as “the dataset ratio”:
| Quantity |
Meaning |
Why it can differ |
| Corpus share |
Fraction of the available corpus |
Original dataset sizes |
| Sampling share |
Fraction of examples/microbatches selected |
Sampler probabilities |
| Compute share |
Fraction of input tokens processed |
Sequence length, padding, packing |
| Supervised-token share |
Fraction of tokens with labels != -100 |
Prompt masking, truncation |
| Scalar loss contribution |
Contribution to the reduced training loss |
Reduction and loss weighting |
| Update contribution |
Influence on the optimizer update |
Gradient norm, alignment, clipping, Adam state |
For raw causal-LM text, input-token share and supervised-token share may be similar. For conversational data with assistant-only loss, they can be very different.
At minimum, I would report both:
- non-padding input tokens, and
- valid target tokens after all formatting, masking, packing, and truncation.
The first approximates compute exposure. The second tells you how many target positions directly contribute to the NLL.
Neither alone is the complete update weight, because the final contribution also depends on loss reduction and gradient statistics.
Loss reduction and gradient accumulation can change the effective mixture
Suppose one FineWeb microbatch contains 400 valid labels while one conversational microbatch contains 150.
These two implementations are not equivalent.
Averaging each microbatch mean
Conceptually:
mean_loss_for_microbatch / gradient_accumulation_steps
Each microbatch receives similar scalar weight, even if it contains a different number of valid labels. A token in the smaller-label microbatch may therefore receive more weight per token.
Normalizing over all valid labels in the accumulated batch
Conceptually:
sum_of_all_token_losses / total_valid_labels
This corresponds more closely to a token-level objective across the full effective batch.
Hugging Face documented this distinction when correcting causal-LM loss under gradient accumulation: averaging already-normalized batch losses is not generally equivalent to summing token losses and dividing by the total number of valid tokens. See the gradient accumulation explanation and current Trainer documentation.
If a custom compute_loss or compute_loss_func is used, I would check whether it handles num_items_in_batch. The current Trainer docs note that ignoring this can make gradient-accumulation loss slightly inaccurate.
For multi-GPU training, I would also check whether the denominator is global across devices. Current TrainingArguments includes average_tokens_across_devices; this matters if ranks receive different source mixtures or different numbers of valid labels.
So I would distinguish:
- configured source probability,
- valid labels per source,
- scalar loss contribution after reduction,
- and resulting gradient/update contribution.
The first two can be logged cheaply. The last two may only need occasional diagnostic logging.
Raw-text CPT and conversational reasoning data may not be equivalent sources
This may be particularly relevant if some of the eight datasets are raw documents while others are conversational or SFT-style reasoning data.
FineWeb-Edu is released as a large educational web-text pretraining corpus.
OpenThoughts-114k is released as a synthetic reasoning conversation dataset covering math, science, code, and puzzles, with ready-to-train data used for fine-tuning OpenThinker models.
That does not mean they cannot be mixed. It means the following choices become part of the mixture definition:
- full-sequence LM loss vs assistant-only loss,
- chat template and role markers,
- whether system/user tokens carry loss,
- truncation direction,
- sequence length,
- packing strategy,
- EOS/document boundaries,
- and whether raw-text CPT and conversational post-training should be separate phases.
For assistant-only training, I would inspect the final tokenized labels, not only the configuration flag.
Current TRL SFTTrainer documentation supports assistant_only_loss, but the chat template must be capable of returning the assistant-token mask. There is also a concrete TRL issue about assistant tokens being removed by truncation, leaving examples with all labels masked.
A useful invariant is:
(labels != -100).sum() > 0
checked after templating, packing, and truncation.
If examples from different sources are packed together, I would additionally check:
- whether source/document boundaries are recorded,
- whether EOS is inserted,
- whether positions reset,
- and whether one source can attend to the preceding source as context.
Packing is not only a throughput choice if it changes what context each target token sees.
I would also slice validation by sequence length. A dataset-level LR preference may partly be a length-distribution effect:
- short responses,
- long reasoning traces,
- highly truncated examples,
- or low-supervision-density examples
may behave differently even inside the same named dataset.
Depending on the intended behavior, a reasonable alternative is:
- raw-document CPT with full-token loss;
- then a shorter conversational/reasoning post-training stage with the appropriate template and mask;
- optionally followed by a small replay/mixed stage if retention is needed.
That is only one branch, not a requirement.
The weighted validation objective may need a more explicit definition
A weighted average is a valid objective, but its interpretation depends on how each per-domain loss is constructed.
I would make explicit whether each domain loss is:
- total NLL divided by valid target tokens,
- mean sequence loss,
- mean batch loss,
- full-sequence loss,
- assistant-only loss,
- or a normalized change from the starting checkpoint.
For example, combining full-token FineWeb loss and assistant-only conversational loss is mathematically valid if that is the chosen scalar objective. However, a weight of 5% no longer has a simple interpretation such as “5% of the model’s importance” or “5% of training tokens”.
Absolute losses across domains may also differ because of:
- entropy,
- formatting,
- tokenizer behavior,
- target-token definitions,
- and validation-set composition.
I would therefore retain both:
- the absolute domain loss, and
- the change from the same starting checkpoint.
A compact table could be:
| Domain |
Starting loss |
Current loss |
Change |
Target weight |
Weighted change |
If small changes are driving the decision, it is also worth estimating validation uncertainty, for example by bootstrapping documents or evaluating multiple fixed validation slices. A very small apparent LR preference may be within finite-validation-set variation.
Three possible objective interpretations
A. Plain weighted average
A low-weight domain may regress if the aggregate improves more elsewhere.
B. Weighted average with non-regression constraints
For important domains, define a maximum acceptable increase from the starting loss.
C. Robust or worst-domain objective
Give additional importance to the worst normalized regression.
These lead to different mixtures. Methods such as DoReMi are closer to robust/group-DRO-style mixture optimization than to an arbitrary fixed weighted average. DoGE is conceptually closer when a target mixture is explicitly specified, although it is still a research method involving proxy optimization rather than a drop-in CPT recipe.
How I would interpret the apparent learning-rate preference
I would not assume that a dataset has one intrinsic preferred LR.
What looks like a source-specific LR preference can be a preference for a particular combination of:
- current checkpoint,
- peak/current LR,
- remaining token horizon,
- warmup or rewarming,
- replay frequency,
- optimizer moments,
- batch/source composition,
- sequence length,
- and repetition count.
The paper Learning Dynamics in Continual Pre-Training for Large Language Models explicitly studies the interaction between distribution shift, LR annealing, peak LR, training steps, replay ratio, and general/domain losses.
Related CPT work also studies combinations of LR rewarming, re-decay, and replay rather than treating LR independently; see Simple and Scalable Strategies to Continually Pre-train Large Language Models.
Therefore, LR comparisons should ideally use:
- the same starting checkpoint,
- the same optimizer-state policy,
- the same scheduler history,
- the same total token or compute budget,
- the same mixture,
- and the same evaluation points.
In particular, resuming the full Trainer checkpoint and loading only model weights into a new Trainer are not equivalent. The current Trainer documentation describes checkpoint resumption of model, optimizer, and scheduler state.
A larger LR may also produce a temporary loss increase after a distribution shift or rewarming before recovering over a longer horizon. Conversely, in a very late-stage, short-horizon run, a smaller LR may genuinely be safer for a nearly converged domain.
So I would describe the present observation as:
evidence of different local LR sensitivities under the current mixture and schedule,
not yet evidence that separate optimizers are required.
How I would interpret different outcomes
| Observation |
More likely interpretation |
Reasonable next branch |
| More sampling improves target train and validation loss without hurting others |
Underexposure |
Keep a sampling floor or staged oversampling |
| More sampling improves train loss but worsens validation |
Repetition/overfitting or train-validation mismatch |
Reduce repeats, improve diversity, or use a cooldown |
| Loss upweighting helps similarly to oversampling |
Insufficient scalar contribution |
Prefer whichever has better stability/throughput |
| Oversampling helps more than upweighting |
Sampling variance/order may matter |
Retain oversampling, monitor repeat count |
| Only a smaller LR helps after exposure is controlled |
Genuine local LR sensitivity is more plausible |
Try phase-specific schedule or source-dependent scaling |
| One-source updates hurt other domains but mixed updates do not |
Replay/mixing is protective |
Avoid long source-homogeneous phases |
| All conditions fail on one source |
Preprocessing, masking, validation mismatch, or poor/noisy data |
Audit examples and validation slices |
| Aggregate improves while a low-weight domain regresses |
Possibly the stated scalar objective working as defined |
Decide whether a regression constraint is needed |
| Loss improves but the desired capability does not |
NLL is an incomplete proxy |
Add a small capability-specific evaluation |
| Results vary substantially by seed/order |
High stochastic variance |
Avoid choosing a mixture from one short run |
The comparison between oversampling and loss weighting has relevant precedent in Upsample or Upweight?. In its multilingual setting, the two are equivalent under idealized full-gradient conditions but diverge under SGD because of gradient variance. The paper reports faster convergence but a greater overfitting risk with strong upsampling and proposes reducing it later.
I would treat that as a useful mechanism and experimental option, not proof that the same schedule is optimal for this CPT setting.
Repetition and data order are separate from sampling probability
For a small source, increasing its probability also increases its repeat count unless more unique data are available.
I would log:
- cumulative source tokens,
- cumulative unique examples/documents,
- approximate epochs/repeats,
- repeat count at each validation point,
- and whether repeated iterable datasets are reshuffled.
This is important because two runs with the same source percentage but different total horizons can have very different repetition regimes.
Scaling Data-Constrained Language Models studies repeated data under constrained data availability and shows that repeated tokens can remain useful for some number of epochs, but their marginal value eventually declines. The exact repeat limit is not universal.
UniMax is another relevant example: it treats a maximum number of corpus repetitions as an explicit part of multilingual mixture design.
The order can matter as well. A fixed mixture is not the only simple option. Before adopting a fully adaptive sampler, I would consider a staged schedule:
- initial stable mixture,
- temporary increase of underperforming sources,
- then a mixed/replay or cooldown phase to limit regressions.
That is simpler to audit than a continuously changing optimizer policy.
Hugging Face implementation checks
Dataset interleaving
The current interleave_datasets documentation distinguishes:
first_exhausted: a subsampling-style stopping rule;all_exhausted: oversamples by restarting exhausted sources;all_exhausted_without_replacement: avoids repeated samples while exhausting all sources.
Therefore, the configured probabilities do not fully specify the resulting repetitions. The stopping strategy and finite source sizes matter.
There is also a relevant Datasets issue about reshuffling a small repeated iterable source. If a small dataset is restarted repeatedly, I would verify whether each pass is reshuffled or repeats the same order.
For distributed/streaming data, I would aggregate realized source counts from every rank and worker rather than infer the mixture from the configuration.
Per-domain evaluation
Trainer can receive a dictionary of evaluation datasets and report separately prefixed metrics. That is useful for retaining all per-domain losses while using a separate script or callback to compute the final weighted objective.
Multiple training sources and custom loss
Multiple independently controlled train sources and domain-specific loss composition are not fully first-class in the standard Trainer interface. There is an open Transformers feature request for multiple datasets and domain-specific losses.
Common implementation routes are:
- a custom sampler or iterable dataset that emits a
source_id,
- overriding
get_train_dataloader,
- overriding
compute_loss,
- or using a custom training loop.
If compute_loss is overridden, I would explicitly test gradient accumulation and num_items_in_batch.
Versions and one-batch inspection
Because masking, packing, and loss-normalization behavior has changed across library versions, I would record:
transformers,datasets,trl,accelerate,- DeepSpeed/FSDP configuration,
- and any Liger/FlashAttention integration.
For each source, saving one fully processed batch with:
input_ids,- decoded text,
labels,- attention/position IDs,
- and
source_id
can often explain more than another expensive training run.
Related research directions if the simple route is not enough
These are useful as a map, but I would not make them the first implementation step.
Proxy-based static mixture search
- Data Mixing Laws uses small mixture runs to predict performance at untested mixtures.
- RegMix similarly uses many smaller runs and regression to select a larger-run mixture.
Useful when an eight-dimensional grid search is impossible, but proxy-to-production transfer and repetition mismatch need care.
Robust/domain-weight optimization
- DoReMi learns domain weights with a proxy model and group DRO.
- DoGE optimizes training-domain weights for generalization to a specified target mixture.
These optimize different objectives, so the objective should be chosen before selecting the method.
Online/adaptive mixture allocation
Online methods can adjust sampling as source losses change, but they add:
- noisy reward/trajectory estimates,
- distributed sampler complexity,
- additional state,
- and more difficult reproducibility.
I would only move here after a fixed-mixture or staged-mixture comparison shows a clear need.
Modular alternatives
If strong, persistent interference is confirmed, alternatives include:
- source-specific adapters,
- short source-specific branches followed by merging,
- alternating phases,
- or keeping some conversational data in a separate post-training stage.
These change the modeling/training design, so they are later branches rather than default fixes.
Bottom line
I would not treat this as “find one LR that every dataset likes” yet.
I would treat it as:
- define the actual model-selection objective;
- verify the objective produced by the data pipeline and loss reduction;
- measure sampling, compute, valid-label, repetition, and per-domain loss separately;
- compare current sampling, oversampling, and loss weighting from the same state;
- try a simple staged mixture if the preferred allocation changes over time;
- only then consider source-specific optimizer state or more advanced mixture optimization.
The most useful first number is probably not the configured 2–5%. It is the complete path from:
source selection → serialized tokens → valid labels → reduced loss → optimizer update.
