Adam有了mini版：内存占用少一半，吞吐量提升50%
内容提要
研究团队提出了一种新的优化器Adam-mini，通过减少学习率的数量来降低内存使用，并在预训练任务中取得了优秀的性能。Adam-mini在内存占用和吞吐量方面优于AdamW，对超参数不敏感。在监督式微调和强化学习任务中，Adam-mini表现更好。
延伸解读
内存节省的实际意义
Adam-mini通过为每个参数块分配单个学习率，将优化器状态v的元素数量减少90%以上，从而为Adam带来45%到50%的内存节省。这意味着在训练7B模型时，优化器状态所需内存从约56GB降至约28GB，使得在有限硬件上训练更大模型成为可能，降低了LLM训练的门槛。
吞吐量提升与训练加速
在2台A800-80GB上预训练Llama2-7B时，Adam-mini的吞吐量比AdamW高49.6%，处理相同数量token可节省33.1%的时间。这主要得益于内存占用降低减少了CPU卸载和分片通信开销，对于硬件资源有限的研究者，能显著缩短预训练总时间。
基于Hessian结构的分片策略
Adam-mini根据Transformer的Hessian块对角结构，将参数切分为与最小密集子块对应的块。对于Query和Key，按头进一步切分，因为每个头对应一个密集子块；而Value、注意力投射和MLP层则保持默认分片。这种策略稳定了训练并提升了性能，避免了PyTorch默认分片在1B模型上的不稳定问题。
超参数不敏感与下游任务表现
在GPT2-125M预训练中，Adam-mini对超参数不敏感，减少了调参负担。在监督式微调（SFT）和基于人类反馈的强化学习（RLHF）任务中，使用Llama-2-7B，Adam-mini的表现优于AdamW，即使它使用单个学习率且内存效率更高，显示了其在下游任务中的有效性。
Q&A
Adam-mini优化器的主要优势是什么？
Adam-mini在内存占用和吞吐量方面优于AdamW，能够显著降低内存使用，同时提升训练速度。
Adam-mini是如何减少内存使用的？
Adam-mini通过为每个参数块分配单个学习率，显著减少学习率的数量，从而降低内存使用。
在什么情况下使用Adam-mini会更有利？
在硬件资源有限的情况下，使用Adam-mini可以获得更高的吞吐量，减少预训练总时间。
Adam-mini与传统的Adam优化器相比有什么不同？
Adam-mini减少了每个参数的学习率数量，优化了学习率分配，而Adam为每个参数分配单独的学习率。
Adam-mini在监督式微调任务中的表现如何？
在监督式微调任务中，Adam-mini的表现优于AdamW，尽管使用了单个学习率。
研究团队是如何发现Adam-mini的优化方法的？
研究团队通过观察Transformer的Hessian结构，发现可以通过简单技巧降低v的使用量，从而提出了Adam-mini。
