Adam-mini: Use Fewer Learning Rates To Gain More
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2024
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866912243094913024 |
|---|---|
| author | Zhang, Yushun Chen, Congliang Li, Ziniu Ding, Tian Wu, Chenwei Kingma, Diederik P. Ye, Yinyu Luo, Zhi-Quan Sun, Ruoyu |
| author_facet | Zhang, Yushun Chen, Congliang Li, Ziniu Ding, Tian Wu, Chenwei Kingma, Diederik P. Ye, Yinyu Luo, Zhi-Quan Sun, Ruoyu |
| contents | We propose Adam-mini, an optimizer that achieves on par or better performance than AdamW with 50% less memory footprint. Adam-mini reduces memory by cutting down the learning rate resources in Adam (i.e., $1/\sqrt{v}$). By investigating the Hessian structure of neural nets, we find Adam's $v$ might not function at its full potential as effectively as we expected. We find that $\geq$ 99.9% of these learning rates in $v$ could be harmlessly removed if we (1) carefully partition the parameters into blocks following our new principle on Hessian structure; (2) assign a single but good learning rate to each parameter block. We then provide one simple 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 39M to 13B for pre-training, supervised fine-tuning, and RLHF. The reduced memory footprint of Adam-mini also alleviates communication overheads among GPUs, thereby increasing throughput. For instance, Adam-mini achieves 49.6% higher throughput than AdamW when pre-training Llama 2-7B on $2\times$ A800-80GB GPUs, which saves 33% wall-clock time for pre-training. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2406_16793 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | Adam-mini: Use Fewer Learning Rates To Gain More Zhang, Yushun Chen, Congliang Li, Ziniu Ding, Tian Wu, Chenwei Kingma, Diederik P. Ye, Yinyu Luo, Zhi-Quan Sun, Ruoyu Machine Learning Artificial Intelligence We propose Adam-mini, an optimizer that achieves on par or better performance than AdamW with 50% less memory footprint. Adam-mini reduces memory by cutting down the learning rate resources in Adam (i.e., $1/\sqrt{v}$). By investigating the Hessian structure of neural nets, we find Adam's $v$ might not function at its full potential as effectively as we expected. We find that $\geq$ 99.9% of these learning rates in $v$ could be harmlessly removed if we (1) carefully partition the parameters into blocks following our new principle on Hessian structure; (2) assign a single but good learning rate to each parameter block. We then provide one simple 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 39M to 13B for pre-training, supervised fine-tuning, and RLHF. The reduced memory footprint of Adam-mini also alleviates communication overheads among GPUs, thereby increasing throughput. For instance, Adam-mini achieves 49.6% higher throughput than AdamW when pre-training Llama 2-7B on $2\times$ A800-80GB GPUs, which saves 33% wall-clock time for pre-training. |
| title | Adam-mini: Use Fewer Learning Rates To Gain More |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2406.16793 |