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.