Adam-mini: 더 적은 학습률로 더 많은 성과를 얻기

Adam-mini: Use Fewer Learning Rates To Gain More

June 24, 2024

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

cs.AI

초록

우리는 AdamW와 동등하거나 더 나은 성능을 달성하면서도 메모리 사용량을 45%에서 50%까지 줄인 최적화 도구인 Adam-mini를 제안합니다. Adam-mini는 Adam의 학습률 리소스(즉, 1/v)를 줄여 메모리를 절약합니다. 우리는 v에 있는 이러한 학습률 중 90% 이상이 무해하게 제거될 수 있음을 발견했는데, 이는 (1) 제안된 헤시안 구조 원칙에 따라 매개변수를 블록으로 신중하게 분할하고, (2) 각 매개변수 블록에 단일하지만 우수한 학습률을 할당하는 경우에 가능합니다. 또한, 이러한 각 매개변수 블록에 대해 Adam을 능가할 수 있는 단일 고품질 학습률이 존재하며, 이를 탐색하기에 충분한 리소스가 제공된다면 가능하다는 것을 발견했습니다. 그런 다음, 우리는 우수한 학습률을 찾는 비용 효율적인 방법을 제시하고 Adam-mini를 제안합니다. 실험적으로, Adam-mini가 125M에서 7B 크기의 다양한 언어 모델에서 사전 학습, 지도 미세 조정, RLHF에 대해 AdamW와 동등하거나 더 나은 성능을 보임을 검증했습니다. Adam-mini의 감소된 메모리 사용량은 GPU와 CPU 간의 통신 오버헤드를 완화하여 처리량을 증가시킵니다. 예를 들어, Adam-mini는 2개의 A800-80GB GPU에서 Llama2-7B를 사전 학습할 때 AdamW보다 49.6% 더 높은 처리량을 달성하며, 이는 사전 학습에 소요되는 벽시계 시간을 33% 절약합니다.

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.