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 number of learning rates in Adam: Instead of assigning an
individual learning rate for each parameter using $1/\sqrt{v}$, Adam-mini uses
the average of $v$ within a pre-defined parameter block as the learning rate
for that block. Such a design is inspired by two empirical findings. First, the
Hessian of Transformers exhibits a near-block diagonal structure with different
sizes of dense sub-blocks. Second, for each of these dense sub-blocks, there
exists a single high-quality learning rate that can outperform Adam, provided
that sufficient resources are available to search it out. Adam-mini provides
one cost-effective way to find these good learning rates and manage to cut down
$\geq 90% v$ in Adam. 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 2x A800-80GB GPUs, which saves 33%
wall-clock time for pre-training.

Adam-mini 是一个优化器，通过减少内存占用量（相较于 AdamW 减少了 45％至 50％）来实现与之相当或更好的性能。Adam-mini 通过减少 Adam 中的学习率数量来降低内存使用，使用预定义参数块中 $v$ 的平均值作为该块的学习率。通过在各种语言模型上进行实证验证，我们发现 Adam-mini 在预训练、监督微调和增强学习等方面显示出与 AdamW 相当或更好的性能，同时还减少了通信开销，并提高了吞吐量。