Power Scheduler: Un Pianificatore del Tasso di Apprendimento Indipendente dalla Dimensione del Batch e dal Numero di Token

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

August 23, 2024

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

cs.AI

Abstract

Trovare il tasso di apprendimento ottimale per il pretraining dei modelli linguistici è un compito impegnativo. Ciò non solo perché esiste una complessa correlazione tra tasso di apprendimento, dimensione del batch, numero di token di addestramento, dimensione del modello e altri iperparametri, ma anche perché è estremamente costoso eseguire una ricerca degli iperparametri per modelli linguistici di grandi dimensioni con miliardi o trilioni di parametri. Studi recenti propongono di utilizzare modelli proxy di piccole dimensioni e un corpus ridotto per eseguire ricerche degli iperparametri e trasporre i parametri ottimali a modelli e corpus di grandi dimensioni. Sebbene la trasferibilità zero-shot sia stata dimostrata teoricamente ed empiricamente per iperparametri legati alla dimensione del modello, come profondità e ampiezza, la trasferibilità zero-shot da un corpus piccolo a uno grande è ancora poco esplorata. In questo articolo, studiamo la correlazione tra tasso di apprendimento ottimale, dimensione del batch e numero di token di addestramento per il recentemente proposto scheduler WSD. Dopo migliaia di piccoli esperimenti, abbiamo trovato una relazione di legge di potenza tra le variabili e ne abbiamo dimostrato la trasferibilità attraverso diverse dimensioni del modello. Sulla base di questa osservazione, proponiamo un nuovo scheduler per il tasso di apprendimento, chiamato Power scheduler, che è agnostico rispetto al numero di token di addestramento e alla dimensione del batch. L'esperimento mostra che combinare il Power scheduler con la Maximum Update Parameterization (muP) può ottenere costantemente prestazioni impressionanti con un unico set di iperparametri, indipendentemente dal numero di token di addestramento, dalla dimensione del batch, dalla dimensione del modello e persino dall'architettura del modello. I nostri modelli densi e MoE da 3B addestrati con il Power scheduler raggiungono prestazioni comparabili con i migliori modelli linguistici di piccole dimensioni. Abbiamo reso open-source questi modelli pretrainati all'indirizzo 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.