Challenges in balancing learning rate and data sampling for multi-dataset continued pretraining

A user seeks advice on managing continued pretraining across a mixture of eight open-source datasets, specifically addressing the trade-offs between learning rate selection and dataset sampling ratios.

• Large datasets like FineWeb-Edu-scale (over 1T tokens) require relatively large learning rates for validation loss improvement, while smaller datasets (around 50B tokens) need extremely small learning rates.
• Undersampling small datasets that contribute only 2–5% of the weighted objective causes their validation loss to increase due to insufficient exposure and overwriting by larger datasets.
• Oversampling these smaller datasets necessitates reducing the sampling rate of larger ones, which may slow down the optimization of the overall weighted objective.

The author is looking for training strategies or optimizer modifications that can handle these conflicting requirements to minimize the weighted validation loss efficiently.