Планировщик мощности: Планировщик скорости обучения, независимый от размера партии и количества токенов

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

August 23, 2024

Авторы: Yikang Shen, Matthew Stallone, Mayank Mishra, Gaoyuan Zhang, Shawn Tan, Aditya Prasad, Adriana Meza Soria, David D. Cox, Rameswar Panda

cs.AI

Аннотация

Поиск оптимальной скорости обучения для предварительного обучения языковой модели представляет собой сложную задачу. Это связано не только с тем, что существует сложная корреляция между скоростью обучения, размером пакета, количеством обучающих токенов, размером модели и другими гиперпараметрами, но также с тем, что проведение поиска гиперпараметров для крупных языковых моделей с миллиардами или триллионами параметров является чрезмерно дорогостоящим. Недавние исследования предлагают использовать небольшие прокси-модели и небольшие корпуса для поиска гиперпараметров и транспонирования оптимальных параметров на крупные модели и большие корпуса. Хотя нулевая передача теоретически и эмпирически доказана для гиперпараметров, связанных с размером модели, таких как глубина и ширина, нулевая передача от небольшого корпуса к большому корпусу недостаточно исследована. В данной статье мы изучаем корреляцию между оптимальной скоростью обучения, размером пакета и количеством обучающих токенов для недавно предложенного планировщика WSD. После тысяч небольших экспериментов мы обнаружили степенную зависимость между переменными и продемонстрировали ее передаточную способность на модели различных размеров. На основе наблюдений мы предлагаем новый планировщик скорости обучения, планировщик Power, который не зависит от количества обучающих токенов и размера пакета. Эксперименты показывают, что комбинация планировщика Power с максимальной параметризацией обновления (muP) последовательно достигает впечатляющих результатов с одним набором гиперпараметров независимо от количества обучающих токенов, размера пакета, размера модели и даже архитектуры модели. Наши модели 3B dense и MoE, обученные с использованием планировщика Power, достигают сравнимой производительности с передовыми небольшими языковыми моделями. Мы предоставляем исходный код этих предварительно обученных моделей по ссылке 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.