Agendador de Potência: Um Agendador de Taxa de Aprendizado Agnóstico ao Tamanho do Lote e ao Número de Tokens

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

August 23, 2024

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

cs.AI

Resumo

Encontrar a taxa de aprendizado ótima para o pré-treinamento de modelos de linguagem é uma tarefa desafiadora. Isso se deve não apenas à complexa correlação entre a taxa de aprendizado, tamanho do lote, número de tokens de treinamento, tamanho do modelo e outros hiperparâmetros, mas também porque é proibitivamente caro realizar uma busca de hiperparâmetros para grandes modelos de linguagem com bilhões ou trilhões de parâmetros. Estudos recentes propõem o uso de pequenos modelos proxy e pequenos corpora para realizar buscas de hiperparâmetros e transpor os parâmetros ótimos para grandes modelos e grandes corpora. Embora a transferibilidade de zero-shot seja teoricamente e empiricamente comprovada para hiperparâmetros relacionados ao tamanho do modelo, como profundidade e largura, a transferência de zero-shot de pequenos corpora para grandes corpora é pouco explorada. Neste artigo, estudamos a correlação entre a taxa de aprendizado ótima, tamanho do lote e número de tokens de treinamento para o programador WSD recentemente proposto. Após milhares de pequenos experimentos, encontramos uma relação de lei de potência entre as variáveis e demonstramos sua transferibilidade entre tamanhos de modelo. Com base na observação, propomos um novo programador de taxa de aprendizado, o programador Power, que é agnóstico em relação ao número de tokens de treinamento e tamanho do lote. O experimento mostra que combinar o programador Power com a Parametrização Máxima de Atualização (muP) pode alcançar consistentemente um desempenho impressionante com um conjunto de hiperparâmetros, independentemente do número de tokens de treinamento, tamanho do lote, tamanho do modelo e até mesmo arquitetura do modelo. Nossos modelos densos de 3B e MoE treinados com o programador Power alcançam desempenho comparável aos modelos de linguagem pequenos de última geração. Disponibilizamos esses modelos pré-treinados em código aberto em 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.