Adam-mini: Используйте меньшее количество скоростей обучения для получения большего выигрыша.

Adam-mini: Use Fewer Learning Rates To Gain More

June 24, 2024

Авторы: Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Yinyu Ye, Zhi-Quan Luo, Ruoyu Sun

cs.AI

Аннотация

Мы предлагаем Adam-mini, оптимизатор, который достигает такой же или даже лучшей производительности, чем AdamW, с уменьшением объема памяти на 45% до 50%. Adam-mini уменьшает использование памяти, сокращая ресурсы скорости обучения в Adam (т.е. 1/v). Мы обнаружили, что более 90% этих скоростей обучения в v могут быть безопасно удалены, если (1) тщательно разделить параметры на блоки в соответствии с нашим предложенным принципом структуры гессиана; (2) назначить одну, но хорошую скорость обучения для каждого блока параметров. Мы также обнаружили, что для каждого из этих блоков параметров существует одна высококачественная скорость обучения, которая может превзойти Adam, при условии наличия достаточных ресурсов для ее поиска. Затем мы предлагаем один эффективный способ найти хорошие скорости обучения и предлагаем Adam-mini. Эмпирически мы проверяем, что Adam-mini работает так же или лучше, чем AdamW, на различных языковых моделях от 125M до 7B для предварительного обучения, надзорного дообучения и RLHF. Уменьшенный объем памяти Adam-mini также снижает накладные расходы на коммуникацию между GPU и ЦП, тем самым увеличивая пропускную способность. Например, Adam-mini достигает на 49.6% большей пропускной способности, чем AdamW, при предварительном обучении Llama2-7B на 2 раза A800-80GB GPU, что экономит 33% времени на предварительное обучение.

English

We propose Adam-mini, an optimizer that achieves on-par or better performance than AdamW with 45% to 50% less memory footprint. Adam-mini reduces memory by cutting down the learning rate resources in Adam (i.e., 1/v). We find that geq 90% of these learning rates in v could be harmlessly removed if we (1) carefully partition the parameters into blocks following our proposed principle on Hessian structure; (2) assign a single but good learning rate to each parameter block. We further find that, for each of these parameter blocks, there exists a single high-quality learning rate that can outperform Adam, provided that sufficient resources are available to search it out. We then provide one cost-effective way to find good learning rates and propose Adam-mini. Empirically, we verify that Adam-mini performs on par or better than AdamW on various language models sized from 125M to 7B for pre-training, supervised fine-tuning, and RLHF. The reduced memory footprint of Adam-mini also alleviates communication overheads among GPUs and CPUs, thereby increasing throughput. For instance, Adam-mini achieves 49.6% higher throughput than AdamW when pre-training Llama2-7B on 2times A800-80GB GPUs, which saves 33% wall-clock time for pre-training.