Adam-mini: Gebruik Minder Leersnelheden Om Meer Te Bereiken

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

Samenvatting

Wij stellen Adam-mini voor, een optimalisator die vergelijkbare of betere prestaties levert dan AdamW met 45% tot 50% minder geheugengebruik. Adam-mini vermindert het geheugengebruik door de leersnelheidsbronnen in Adam (d.w.z. 1/v) te beperken. Wij ontdekken dat ≥ 90% van deze leersnelheden in v zonder problemen kunnen worden verwijderd als we (1) de parameters zorgvuldig verdelen in blokken volgens ons voorgestelde principe over de Hessiaanse structuur; (2) een enkele maar goede leersnelheid toewijzen aan elk parameterblok. Wij ontdekken verder dat er voor elk van deze parameterblokken een enkele hoogwaardige leersnelheid bestaat die Adam kan overtreffen, mits er voldoende middelen beschikbaar zijn om deze te vinden. Wij bieden vervolgens een kosteneffectieve manier aan om goede leersnelheden te vinden en stellen Adam-mini voor. Empirisch verifiëren wij dat Adam-mini vergelijkbare of betere prestaties levert dan AdamW op verschillende taalmodelgroottes van 125M tot 7B voor voorafgaande training, gesuperviseerde fine-tuning en RLHF. Het verminderde geheugengebruik van Adam-mini verlicht ook de communicatieoverhead tussen GPU's en CPU's, waardoor de doorvoer wordt verhoogd. Zo bereikt Adam-mini bijvoorbeeld 49,6% hogere doorvoer dan AdamW bij het vooraf trainen van Llama2-7B op 2× A800-80GB GPU's, wat 33% wandkloktijd bespaart voor de voorafgaande training.

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.