Power Scheduler: A Batch Size and Token Number Agnostic Learning Rate
  Scheduler
Authors: name: Yikang Shen; name: Matthew Stallone; name: Mayank Mishra; name: Gaoyuan Zhang; name: Shawn Tan; name: Aditya Prasad; name: Adriana Meza Soria; name: David D. Cox; name: Rameswar Panda
Published: 2024-08-23 20:22:20

Abstract: Finding the optimal learning rate for language model pretraining is a
challenging task. This is not only because there is a complicated correlation
between learning rate, batch size, number of training tokens, model size, and
other hyperparameters but also because it is prohibitively expensive to perform
a hyperparameter search for large language models with Billions or Trillions of
parameters. Recent studies propose using small proxy models and small corpus to
perform hyperparameter searches and transposing the optimal parameters to large
models and large corpus. While the zero-shot transferability is theoretically
and empirically proven for model size related hyperparameters, like depth and
width, the zero-shot transfer from small corpus to large corpus is
underexplored. In this paper, we study the correlation between optimal learning
rate, batch size, and number of training tokens for the recently proposed WSD
scheduler. After thousands of small experiments, we found a power-law
relationship between variables and demonstrated its transferability across
model sizes. Based on the observation, we propose a new learning rate
scheduler, Power scheduler, that is agnostic about the number of training
tokens and batch size. The experiment shows that combining the Power scheduler
with Maximum Update Parameterization (muP) can consistently achieve impressive
performance with one set of hyperparameters regardless of the number of
training tokens, batch size, model size, and even model architecture. Our 3B
dense and MoE models trained with the Power scheduler achieve comparable
performance as state-of-the-art small language models. We open-source these
pretrained models at https://ibm.biz/BdKhLa.