Leistungsplaner: Ein Lernratenplaner unabhängig von Stapelgröße und Tokenanzahl

Power Scheduler: A Batch Size and Token Number Agnostic Learning Rate Scheduler

August 23, 2024

Autoren: Yikang Shen, Matthew Stallone, Mayank Mishra, Gaoyuan Zhang, Shawn Tan, Aditya Prasad, Adriana Meza Soria, David D. Cox, Rameswar Panda

cs.AI

Zusammenfassung

Die Ermittlung der optimalen Lernrate für das Vortrainieren von Sprachmodellen ist eine anspruchsvolle Aufgabe. Dies liegt nicht nur daran, dass eine komplexe Korrelation zwischen Lernrate, Batch-Größe, Anzahl der Trainings-Token, Modellgröße und anderen Hyperparametern besteht, sondern auch daran, dass es prohibitiv teuer ist, eine Hyperparameter-Suche für große Sprachmodelle mit Milliarden oder Billionen von Parametern durchzuführen. Aktuelle Studien schlagen vor, kleine Proxy-Modelle und kleine Korpora zu verwenden, um Hyperparametersuchen durchzuführen und die optimalen Parameter auf große Modelle und große Korpora zu übertragen. Während die Nullschuss-Übertragbarkeit theoretisch und empirisch für mit der Modellgröße zusammenhängende Hyperparameter wie Tiefe und Breite nachgewiesen ist, ist die Nullschuss-Übertragung von kleinen Korpora auf große Korpora noch wenig erforscht. In diesem Paper untersuchen wir die Korrelation zwischen optimaler Lernrate, Batch-Größe und Anzahl der Trainings-Token für den kürzlich vorgeschlagenen WSD-Scheduler. Nach Tausenden von kleinen Experimenten fanden wir eine Potenzgesetzbeziehung zwischen den Variablen und zeigten deren Übertragbarkeit über Modellgrößen hinweg. Basierend auf der Beobachtung schlagen wir einen neuen Lernraten-Scheduler vor, den Power-Scheduler, der unabhängig von der Anzahl der Trainings-Token und der Batch-Größe ist. Das Experiment zeigt, dass die Kombination des Power-Schedulers mit der Maximum Update Parameterization (muP) mit einem Satz von Hyperparametern unabhängig von der Anzahl der Trainings-Token, der Batch-Größe, der Modellgröße und sogar der Modellarchitektur beeindruckende Leistungen erzielen kann. Unsere 3B dichten und MoE-Modelle, die mit dem Power-Scheduler trainiert wurden, erreichen vergleichbare Leistungen wie modernste kleine Sprachmodelle. Wir stellen diese vortrainierten Modelle unter https://ibm.biz/BdKhLa als Open Source zur Verfügung.

English

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.