Adam-mini: Verwenden Sie weniger Lernraten, um mehr zu gewinnen.

Adam-mini: Use Fewer Learning Rates To Gain More

June 24, 2024

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

cs.AI

Zusammenfassung

Wir schlagen Adam-mini vor, einen Optimierer, der eine vergleichbare oder bessere Leistung als AdamW mit 45% bis 50% weniger Speicherplatz erreicht. Adam-mini reduziert den Speicherbedarf, indem er die Ressourcen für die Lernrate in Adam (d.h. 1/v) reduziert. Wir stellen fest, dass über 90% dieser Lernraten in v unbedenklich entfernt werden könnten, wenn wir (1) die Parameter sorgfältig in Blöcke unterteilen, die unserem vorgeschlagenen Prinzip der Hesse'schen Struktur folgen; (2) jeder Parametergruppe eine einzige, aber gute Lernrate zuweisen. Wir stellen weiter fest, dass für jede dieser Parametergruppen eine einzige hochwertige Lernrate existiert, die Adam übertreffen kann, vorausgesetzt ausreichende Ressourcen stehen zur Verfügung, um sie zu ermitteln. Wir bieten dann einen kostengünstigen Weg, um gute Lernraten zu finden, und schlagen Adam-mini vor. Empirisch bestätigen wir, dass Adam-mini bei verschiedenen Sprachmodellen mit Größen von 125M bis 7B für Vor-Training, überwachtes Feintuning und RLHF eine vergleichbare oder bessere Leistung als AdamW erbringt. Der reduzierte Speicherbedarf von Adam-mini verringert auch die Kommunikationsüberlastung zwischen GPUs und CPUs und erhöht somit die Durchsatzleistung. Adam-mini erreicht beispielsweise eine 49,6% höhere Durchsatzleistung als AdamW beim Vor-Training von Llama2-7B auf 2-mal A800-80GB GPUs, was eine 33%ige Zeitersparnis beim Vor-Training bedeutet.

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.