Adam-mini : Utilisez moins de taux d'apprentissage pour en tirer davantage

Adam-mini: Use Fewer Learning Rates To Gain More

June 24, 2024

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

cs.AI

Résumé

Nous proposons Adam-mini, un optimiseur qui atteint des performances équivalentes ou supérieures à AdamW tout en réduisant l'empreinte mémoire de 45% à 50%. Adam-mini réduit la mémoire en diminuant les ressources dédiées au taux d'apprentissage dans Adam (c'est-à-dire 1/v). Nous constatons que geq 90% de ces taux d'apprentissage dans v peuvent être supprimés sans conséquence si nous (1) partitionnons soigneusement les paramètres en blocs selon notre principe proposé sur la structure du Hessien ; (2) attribuons un seul mais bon taux d'apprentissage à chaque bloc de paramètres. Nous découvrons en outre que, pour chacun de ces blocs de paramètres, il existe un seul taux d'apprentissage de haute qualité qui peut surpasser Adam, à condition que des ressources suffisantes soient disponibles pour le trouver. Nous proposons ensuite une méthode rentable pour identifier de bons taux d'apprentissage et introduisons Adam-mini. Empiriquement, nous vérifions qu'Adam-mini performe au moins aussi bien qu'AdamW sur divers modèles de langage allant de 125M à 7B pour le pré-entraînement, le fine-tuning supervisé et le RLHF. La réduction de l'empreinte mémoire d'Adam-mini atténue également les surcharges de communication entre les GPU et les CPU, augmentant ainsi le débit. Par exemple, Adam-mini atteint un débit 49,6% plus élevé qu'AdamW lors du pré-entraînement de Llama2-7B sur 2 fois des GPU A800-80GB, ce qui permet d'économiser 33% du temps réel pour le pré-entraînement.

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.