Power Scheduler: Een Leerritmeplanner Onafhankelijk van Batchgrootte en Aantal Tokens

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

Samenvatting

Het vinden van de optimale leerrate voor het vooraf trainen van taalmodelen is een uitdagende taak. Dit komt niet alleen omdat er een complexe correlatie bestaat tussen leerrate, batchgrootte, aantal trainings tokens, modelgrootte en andere hyperparameters, maar ook omdat het buitengewoon kostbaar is om een hyperparameterzoektocht uit te voeren voor grote taalmodelen met miljarden of biljoenen parameters. Recente studies stellen voor om kleine proxy-modellen en een kleine corpus te gebruiken om hyperparameterzoektochten uit te voeren en de optimale parameters over te dragen naar grote modellen en een grote corpus. Hoewel de zero-shot overdraagbaarheid theoretisch en empirisch is bewezen voor hyperparameters gerelateerd aan modelgrootte, zoals diepte en breedte, is de zero-shot overdracht van een kleine corpus naar een grote corpus nog onderbelicht. In dit artikel bestuderen we de correlatie tussen optimale leerrate, batchgrootte en aantal trainings tokens voor de recent voorgestelde WSD-planner. Na duizenden kleine experimenten hebben we een machtswetrelatie tussen variabelen gevonden en de overdraagbaarheid ervan over modelgroottes aangetoond. Op basis van deze observatie stellen we een nieuwe leerrateplanner voor, de Power-planner, die onafhankelijk is van het aantal trainings tokens en de batchgrootte. Het experiment toont aan dat het combineren van de Power-planner met Maximum Update Parameterization (muP) consistent indrukwekkende prestaties kan bereiken met één set hyperparameters, ongeacht het aantal trainings tokens, de batchgrootte, de modelgrootte en zelfs de modelarchitectuur. Onze 3B dense en MoE-modellen, getraind met de Power-planner, bereiken vergelijkbare prestaties als state-of-the-art kleine taalmodelen. We hebben deze vooraf getrainde modellen open-source gemaakt op 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.