Adam-mini: Use Fewer Learning Rates To Gain More

ML Adam等 POA
我们提出了Adam-mini优化器,其内存占用比AdamW少45%到50%,并且性能相当甚至更好。Adam-mini通过减少Adam中的学习率资源(即$1/\sqrt{v}$)来降低内存使用率。我们发现,如果我们(1)按照我们提出的Hessian结构原则仔细地将参数分成块;(2)为每个参数块分配一个单一但好的学习率,那么$v$中$\geq$ 90%的这些学习率可以被无害地去除。我们进一步发现,对于这些参数块中的每一个,存在一个高质量的单一学习率,如果有足够的资源来搜索它,它可以胜过Adam。然后,我们提供了一种成本效益的方法来找到好的学习率,并提出了Adam-mini。在经验上,我们验证了Adam-mini在从125M到7B的各种语言模型上进行预训练、监督微调和RLHF时,性能与AdamW相当甚至更好。Adam-mini的减少的内存使用率还减轻了GPU和CPU之间的通信开销,从而增加了吞吐量。例如,在使用$2\times$ A800-80GB GPU对Llama2-7B进行预训练时,Adam-mini的吞吐量比AdamW高49.6%,这节省了33%的预训练时间。
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/\sqrt{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 $2\times$ A800-80GB GPUs, which saves 33% wall-clock time for pre-training.
许愿