How should I balance learning rate and data sampling during CPT on multiple datasets?

[jefcoder]

Hi everyone,

I'm doing continued pretraining on a mixture of eight open-source datasets.

The training is now in the late stage where the validation loss decreases very slowly. My objective is to minimize a weighted validation loss across all datasets.

I'm facing two related challenges.

1. Choosing the learning rate

Different datasets seem to prefer very different learning-rate regimes.

• For very large datasets (e.g., FineWeb-Edu-scale datasets with more than 1T tokens), the validation loss continues to improve with a relatively large learning rate.
• For smaller datasets (around 50B tokens), the validation loss only improves when the learning rate is extremely small.

Since my goal is to reduce the overall weighted loss as quickly as possible, I'm unsure how to choose a single learning rate that works well across all datasets.

Is this simply an unavoidable trade-off, or are there training strategies that work better in this setting? For example, do people modify the optimizer or training procedure to handle datasets that require different effective learning rates?

2. Dataset sampling

Another issue is the sampling ratio.

Some datasets contribute only 2–5% of the total weighted objective. When I sampled those datasets at only 2–5% of the training tokens, their validation loss gradually increases. My guess is that they receive too little exposure and are continually overwritten by updates from the larger datasets.

However, if I oversample those smaller datasets, I must reduce the sampling rate of the larger datasets, which could slow down optimization of the overall weighted objective.

How can I handle this trade-off?

I'd really appreciate hearing from anyone who has experience with large-scale continued pretraining or multi-domain LLM pretraining. I'm also open to paid consulting or collaboration. Feel free to DM me.

[tianyu-l]

I personally don't have experience in data mixing recipe.

@felipemello1 @slimfrkha @rakkit @haydn-jones in case you could help

[felipemello1]

Some datasets contribute only 2–5% of the total weighted objective. When I sampled those datasets at only 2–5% of the training tokens, their validation loss gradually increases. My guess is that they receive too little exposure and are continually overwritten by updates from the larger datasets.

can you confirm that they are actually being sampled at this rate? In our design, we went with having an entire packed row being of just one dataset. e.g.

[
[math, math, math]
[ code, code, code]
]

Looking back, i think that this was not ideal, and we should have allowed:

[
[math, code, math]
[art, code, math]
]

So i wonder if your cfg matches the true distribution

If the dataset is small and batchsize is relatively small, then one train step would be heavily weighted towards that dataset, but then multiple training steps wouldnt have that dataset at all, possibly leading to forgetting.

We will have a refactor soon and we can address that there. Meanwhile, perhaps ask some AI to refactor the packing, to allow multiple datasets in the same row. You can check if that solves your issue with learning/forgetting.

[jefcoder]

Hi @felipemello1, I really appreciate your valuable feedback.
I have set the effective batch size to 256 sequences. I first sample the required number of sequences from each database, then mix and shuffle them to ensure a uniform distribution in the final batch.

[slimfrkha]

I would do a WSD LR regime where with a mixture focused on dominant datasets (eg fine web edu) on the warm+constant regime then switch in decay stage to data mix focused on the small HQ datasets (keeping small portion of fine web edu to avoid catastrophic forgetting).
I would make sure the dataset are globally shuffled as a lot of open source datasets aren't (so traversal order would yield ordered domains, e.g math then art then physics etc).
I'd watch the training loss only during the stable regime and start monitoring the weighted validation loss during the decay regime.

[jefcoder]

Hi @slimfrkha,

Thanks for the valuable feedback. I’ll try your suggestion and share the results.

For continued pretraining, I currently use warmup followed by a constant learning rate and save checkpoints regularly, so I haven’t used a decay phase yet. I saw the WSD method in the Nanotron / Hugging Face playbook, but I’m not sure how to apply it in practice and implement effectively.

If I want to evaluate checkpoints regularly, should I periodically run a decay phase from selected checkpoints while keeping the main run in the stable phase, as shown in the figure?

I’d also appreciate your guidance on a few points:

• Are there existing libraries or frameworks that support this workflow?
• How should I choose the ratio between the stable and decay phases?
• If I train small datasets only in decay phase, is it sufficient to reduce the loss for small datasets?
• Is this approach what they are called "curriculum training" - train main db first then other db in decay phase?

I’m new to WSD, so any guidance or recommended materials would be greatly appreciated.
Thank you.

[slimfrkha]

1. you can run decay periodically if it is not expensive for you and you can afford it. trade off is to run some mini decay just to see the trend and compare ckpts between themselves (before the final decay)
2. about the ratio, original WSD paper says something like 10%, of data but recent models training dont care about this anymore and perform decay for 200 to 500B tokens. in your case specifically, you can perform the stable regime on fine edu 1T then decay on the whole HQ small datasets (50B) + some fine edu tokens.
3. no framework supporting this workflow afaik. launch everything manually.
4. I also think the evaluation procedure is improvable: the loss (even validation loss) is a training signal but it does not mean your model would be better. I'd scrap the validation loss and use standard benchmarks: like if you inject a math dataset, look for a bench that tests similar math subjects. you'd have a set of benchmarks that you fix and for each ckpt run all the benchs (or a selected subset mid run if too expensive) to assess it.

[jefcoder]

Hi @slimfrkha,

I really appreciate your guidance — it was extremely helpful. That answers everything for me.

Thank you!

[haydn-jones]

Hey @jefcoder. I don't think there's a magic solution here, but @felipemello1 and @slimfrkha have made the right points separately so I'll (try) to tie them together. There are roughly three "easy" approaches to a multi dataset recipe:

1. Blind mixing and praying. Sounds like this isn't working, which isn't surprising when the data you care about is 2-5% of the mix.

2. Rewriting / augmentation. Use another LLM to rephrase your small datasets into new examples (fun demo here from huggingface https://huggingface.co/spaces/HuggingFaceFW/finephrase). I think this works well but it's not what I'd try first.

3. Scheduling, which is basically what @slimfrkha described. I think this is the righ approach. Run bulk data (like FineWeb-Edu or the pile) during the high-LR stable phase, then as LR decays upsample the small high-quality datasets above their natural rate (e.g. 20% instead of 2-5%), keeping some bulk data in the mix to avoid forgetting. Llama 3 and OLMo 2 definitely did this (and I'm sure many others). Avoid purely sequential dataset stages (all of dataset 1, then all of 2, ...) as that can cause forgetting.

For your LR question specifically, I don't think there's a single LR that serves both regimes. "Which LR does this dataset want" should be "at which point in the LR decay does this dataset get sampled." Your observation that small datasets only improve at tiny LR is the right observation, so merge it into the LR schedule.

TLDR: Try (3).

[haydn-jones]

Ok actually two more things.

1. If your goal is minimizing a weighted sum of datasets, you're going to have a pareto frontier of solutions, and some of them are probably just going to involve the smaller datasets having their val loss going up. If your goal is "the small datasets cannot degrade", that's a different objective.
2. You have a current checkpoint (the one where loss went up for your small datasets). Build a new dataset that has a mixture where the data is like 20% of your small datasets (the ones where loss went up), and 80% a random sample of your bigger ones. Start a training run using the previous checkpoint AND the previous optimizer states, and decay the LR linearly to zero during training. Make sure the schedule starts at the checkpoint's current LR with no re-warmup, otherwise the warmup spike will undo exactly what you're trying to fix. This can be short, so it's much cheaper than a full run, and you may be able to cheaply experiment with multiple epochs over the smaller data (it doesn't always hurt https://arxiv.org/abs/2305.16264). You should be paying attention to how the big-dataset and small-dataset val loss changes. If you can recover a lot of performance on the small sets, while only hurting the big ones a little bit, you're in a good spot.

[jefcoder]

Hi @haydn-jones,

Thank you for your valuable response. I will definitely follow your suggestions, as they will be a great help to me.

By the way, could you provide a practical guide for this approach? Since most training frameworks do not natively support a decay phase, I would appreciate your advice on the best way to implement it.

Additionally, I would be very grateful if you could look into the question I raised in @slimfrkha 's response?

Thank you.

[rakkit]

i think its better to rephrase question 2)
i will assume you want to max the "performance" on for example 8 downstream task (or dataset), and i would suggest to decouple the data mix and val data mix.
what i mean is, for example your 8 dataset is math, code, reasoning, french, spanish etc, and you will need to confirm your model's performance will be like 50% math, 10% code, 5% french, 10% spanish etc, this defines your val metrics. (and it does not have to be limited you 8 datasets but you can just add more val dataset for these "capbility")

then i think what you really want to solve is, how to mix of your 8 datasets to minimal val loss of your "averaged target", for that there is no magic solution, if you do pre-train, you can go e.g. remix or olmix, train proxy model and predict combination for minimal loss. there are also some methods suggest you can train model on each dataset and merge checkpoint to see data mix. For your CPT case, i think its not avoidable to run lots of short runs and fiting the data mix predict like loss = (pct of data1, pct of data 2 ...) and you can usually use linear model to solve it to get best pct s.t lowest loss.

for you question 1, its easier to solve.

for you "production" run, saying you solved 2 already and know the mix. then you just need to tune LR and BS (you need proper combination of these hyper-partermers instead of tune LR or BS seprately).
the general rules and suggestion is, you can treate it as some mid-training as continue of base model pretrain.

1. if you know BS/LR of base model, you can keep same BS and start your init LR as end of base model decay, usually dont need warmup, but usually you dont want very higher LR, ideally you want continue LR from end of pre-train. this is most of now-day LLM pre train recipe, (you can check LR scheduler of mimo-v2 or stepflash 3.5, the LR is mono-decrease).

one of the concern of higher LR is forgotten, so your base model's pre train knowledge get affected.

2. again, you will need to do 2d sweep of <BS, LR>, this is only way to keep it optimal. but you can do this on small scale and extrapolate to large scale, so e.g. when you want to run actual "200B" token, you can set your "sweep" to 5-10B token, and run it with same LR sheduler (stable-decay, maybe with warmup), but 2d grid (BS, LR) and find the best BS, LR combination give you lowest loss.

then you can extrapolate this to your target training horizon (e.g. 200B here), by LRT^(-1/4) and BT(1/2) rules (you can check step-law or https://arxiv.org/abs/2410.21676 or norm transfer, you can just ask chatgpt of key words how BS/LR scaling with tokens budget)

[jefcoder]

Hi @rakkit
Thanks for your detailed response. I really appreciate it.
Thanks again.