Planificador de Potencia: Un Programador de Tasa de Aprendizaje Agnóstico al Tamaño del Lote y al 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

Resumen

Encontrar la tasa de aprendizaje óptima para el preentrenamiento de modelos de lenguaje es una tarea desafiante. Esto se debe no solo a la complicada correlación entre la tasa de aprendizaje, el tamaño del lote, el número de tokens de entrenamiento, el tamaño del modelo y otros hiperparámetros, sino también porque resulta prohibitivamente costoso realizar una búsqueda de hiperparámetros para grandes modelos de lenguaje con miles de millones o billones de parámetros. Estudios recientes proponen utilizar modelos proxy pequeños y un corpus reducido para realizar búsquedas de hiperparámetros y trasladar los parámetros óptimos a modelos grandes y corpus extensos. Si bien la transferibilidad de cero disparos está teórica y empíricamente demostrada para hiperparámetros relacionados con el tamaño del modelo, como la profundidad y el ancho, la transferencia de cero disparos de un corpus pequeño a un corpus grande está poco explorada. En este artículo, estudiamos la correlación entre la tasa de aprendizaje óptima, el tamaño del lote y el número de tokens de entrenamiento para el programador WSD propuesto recientemente. Después de miles de experimentos pequeños, encontramos una relación de ley de potencia entre las variables y demostramos su transferibilidad entre tamaños de modelos. Basándonos en la observación, proponemos un nuevo programador de tasa de aprendizaje, el programador Power, que es agnóstico sobre el número de tokens de entrenamiento y el tamaño del lote. El experimento muestra que combinar el programador Power con la Parametrización Máxima de Actualización (muP) puede lograr consistentemente un rendimiento impresionante con un conjunto de hiperparámetros independientemente del número de tokens de entrenamiento, tamaño del lote, tamaño del modelo e incluso la arquitectura del modelo. Nuestros modelos densos de 3B y MoE entrenados con el programador Power logran un rendimiento comparable a los modelos de lenguaje pequeños de última generación. Ponemos a disposición estos modelos preentrenados en 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.