Planificateur de puissance : un planificateur de taux d'apprentissage agnostique en termes de taille de lot et de nombre de jetons

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

August 23, 2024

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

cs.AI

Résumé

Trouver le taux d'apprentissage optimal pour la pré-entraînement des modèles de langage est une tâche difficile. Ceci est dû non seulement à la corrélation complexe entre le taux d'apprentissage, la taille du lot, le nombre de jetons d'entraînement, la taille du modèle et d'autres hyperparamètres, mais aussi au coût prohibitif de réaliser une recherche d'hyperparamètres pour de grands modèles de langage avec des milliards ou des billions de paramètres. Des études récentes proposent d'utiliser de petits modèles proxy et un petit corpus pour effectuer des recherches d'hyperparamètres et de transposer les paramètres optimaux à de grands modèles et de grands corpus. Alors que la transférabilité zéro-shot est théoriquement et empiriquement prouvée pour les hyperparamètres liés à la taille du modèle, comme la profondeur et la largeur, le transfert zéro-shot du petit corpus au grand corpus est peu exploré. Dans cet article, nous étudions la corrélation entre le taux d'apprentissage optimal, la taille du lot et le nombre de jetons d'entraînement pour l'ordonnanceur WSD récemment proposé. Après des milliers de petites expériences, nous avons trouvé une relation de loi de puissance entre les variables et démontré sa transférabilité à travers les tailles de modèle. Sur la base de l'observation, nous proposons un nouvel ordonnanceur de taux d'apprentissage, l'ordonnanceur Power, qui est agnostique quant au nombre de jetons d'entraînement et à la taille du lot. L'expérience montre que la combinaison de l'ordonnanceur Power avec la Paramétrisation de Mise à Jour Maximale (muP) peut atteindre de manière cohérente des performances impressionnantes avec un ensemble d'hyperparamètres quel que soit le nombre de jetons d'entraînement, la taille du lot, la taille du modèle, voire l'architecture du modèle. Nos modèles denses 3B et MoE entraînés avec l'ordonnanceur Power atteignent des performances comparables à celles des petits modèles de langage de pointe. Nous mettons ces modèles pré-entraînés en open source sur https://ibm.biz/BdKhLa.

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.