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

[jefcoder]

Hi everyone

I am in the late stages of pretraining on a mixture of eight datasets, aiming to minimize a weighted validation loss. I am facing two primary challenges:

1. Learning Rate Disparity: Large datasets continue to improve with a high LR, whereas smaller datasets require an extremely low LR to avoid regression. How can I balance these conflicting LR requirements within a single training run?

2. Sampling Trade-offs: Low-weight datasets show increasing validation loss when sampled at their target weight, suggesting they are being "overwritten" by larger updates. However, oversampling them slows the optimization of the overall objective. What are the best practices for managing this sampling tension?

Seeking advice from experts in large-scale multi-domain pretraining.

[akoumpa]

Hi @jefcoder ,

Apologies for the delay, AutoModel already supports fixed weighted sampling for datasets. However, it currently lacks is a first-class domain-mixture objective that separates:

• sampling probability (q_i)
• target objective weight (w_i)
• per-domain validation and model selection
• optional time-varying sampling

I'm working on a PR to address these, I'll follow up here, when it's ready.

Thank you

[akoumpa]

I have a draft PR here that i am validating, please take a look @jefcoder and let me know if you find any issues. thank you.