Adam-mini: Utilizzare Meno Tassi di Apprendimento per Ottenere di Più

Adam-mini: Use Fewer Learning Rates To Gain More

June 24, 2024

Autori: Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Yinyu Ye, Zhi-Quan Luo, Ruoyu Sun

cs.AI

Abstract

Proponiamo Adam-mini, un ottimizzatore che raggiunge prestazioni pari o superiori a AdamW con un'impronta di memoria ridotta del 45% al 50%. Adam-mini riduce la memoria diminuendo le risorse dedicate al tasso di apprendimento in Adam (ovvero, 1/v). Abbiamo scoperto che ≥90% di questi tassi di apprendimento in v può essere rimosso senza conseguenze negative se (1) partizioniamo attentamente i parametri in blocchi seguendo il nostro principio proposto sulla struttura dell'Hessiano; (2) assegniamo un singolo ma efficace tasso di apprendimento a ciascun blocco di parametri. Abbiamo inoltre osservato che, per ciascuno di questi blocchi di parametri, esiste un singolo tasso di apprendimento di alta qualità che può superare Adam, a condizione che siano disponibili risorse sufficienti per individuarlo. Forniamo quindi un metodo economicamente vantaggioso per trovare buoni tassi di apprendimento e proponiamo Adam-mini. Empiricamente, abbiamo verificato che Adam-mini performa almeno quanto o meglio di AdamW su vari modelli linguistici con dimensioni che vanno da 125M a 7B per pre-addestramento, fine-tuning supervisionato e RLHF. La ridotta impronta di memoria di Adam-mini allevia anche i sovraccarichi di comunicazione tra GPU e CPU, aumentando così il throughput. Ad esempio, Adam-mini raggiunge un throughput del 49,6% superiore rispetto a AdamW durante il pre-addestramento di Llama2-7B su 2× GPU A800-80GB, risparmiando il 33% del tempo di clock per il pre-addestramento.

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.