Adam-mini: Usa menos tasas de aprendizaje para obtener más

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

Resumen

Proponemos Adam-mini, un optimizador que logra un rendimiento igual o mejor que AdamW con una huella de memoria entre un 45% y un 50% menor. Adam-mini reduce la memoria disminuyendo los recursos de tasa de aprendizaje en Adam (es decir, 1/v). Descubrimos que geq 90% de estas tasas de aprendizaje en v podrían eliminarse sin perjuicio si (1) particionamos cuidadosamente los parámetros en bloques siguiendo nuestro principio propuesto sobre la estructura del Hessiano; (2) asignamos una única pero buena tasa de aprendizaje a cada bloque de parámetros. Además, encontramos que, para cada uno de estos bloques de parámetros, existe una única tasa de aprendizaje de alta calidad que puede superar a Adam, siempre que se disponga de recursos suficientes para encontrarla. Luego, proporcionamos una forma rentable de encontrar buenas tasas de aprendizaje y proponemos Adam-mini. Empíricamente, verificamos que Adam-mini tiene un rendimiento igual o mejor que AdamW en varios modelos de lenguaje que van desde 125M hasta 7B para pre-entrenamiento, ajuste fino supervisado y RLHF. La reducida huella de memoria de Adam-mini también alivia la sobrecarga de comunicación entre GPUs y CPUs, aumentando así el rendimiento. Por ejemplo, Adam-mini logra un 49.6% más de rendimiento que AdamW al pre-entrenar Llama2-7B en 2 veces GPUs A800-80GB, lo que ahorra un 33% del tiempo de reloj en el pre-entrenamiento.

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.