Adam-mini: Utilize Menos Taxas de Aprendizado Para Obter Mais

Adam-mini: Use Fewer Learning Rates To Gain More

June 24, 2024

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

cs.AI

Resumo

Propomos o Adam-mini, um otimizador que alcança desempenho equivalente ou superior ao AdamW com 45% a 50% menos uso de memória. O Adam-mini reduz a memória ao diminuir os recursos da taxa de aprendizado no Adam (ou seja, 1/v). Descobrimos que mais de 90% dessas taxas de aprendizado em v poderiam ser removidas sem prejuízo se (1) particionarmos cuidadosamente os parâmetros em blocos seguindo nosso princípio proposto na estrutura Hessiana; (2) atribuirmos uma única, porém boa, taxa de aprendizado a cada bloco de parâmetros. Além disso, descobrimos que, para cada um desses blocos de parâmetros, existe uma única taxa de aprendizado de alta qualidade que pode superar o Adam, desde que recursos suficientes estejam disponíveis para encontrá-la. Em seguida, apresentamos uma maneira econômica de encontrar boas taxas de aprendizado e propomos o Adam-mini. Empiricamente, verificamos que o Adam-mini tem desempenho equivalente ou superior ao AdamW em vários modelos de linguagem com tamanhos de 125M a 7B para pré-treinamento, ajuste fino supervisionado e RLHF. O menor uso de memória do Adam-mini também reduz as sobrecargas de comunicação entre GPUs e CPUs, aumentando assim a taxa de transferência. Por exemplo, o Adam-mini alcança uma taxa de transferência 49,6% maior do que o AdamW ao pré-treinar o Llama2-7B em 2 vezes A800-80GB GPUs, o que economiza 33% do tempo de relógio para o pré-treinamento.

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.