Adam有了mini版：内存占用少一半，吞吐量提升50%

2024-07-06 06:07

机器之心报道

编辑：panda

在训练大型语言模型（LLM）时，Adam(W) 基本上已经成为了人们默认使用的优化器。

Adam 尽管性能优异，但使用成本很高。具体来说，Adam 需要内存来保存其优化器状态：一阶动量 m 和二阶动量 v^2。这总共需要模型大小至少 2 倍的内存。这样的内存消耗已经成为了 LLM 训练的一大主要负担。

举个例子，要训练一个 7B 模型，只是 Adam 就需要每张卡有大约 56 GB 来保存 m 和 v；而如果再加上梯度，则总共需要 86 GB。即使使用最先进的 A100-80GB，成本也过高了。

为了支持这样的高内存算法，实践中必须要使用 CPU 卸载与分片，但这又会增加延迟，减慢训练速度。在训练 PaLM （有 5400 亿参数）这样的更大型模型时，情况还会更糟。在这种情况下，Adam 自身就要占用超过 50 GB，并且这也是预训练阶段的一大主要开销。

因此，人们希望设计出内存需求更少又有效的优化器。首先，减少内存可以减轻 CPU 卸载的负担并能减轻对模型参数执行分片的需求。这些都能减少 GPU 和 CPU 之间的通信量，并进一步提升训练过程的吞吐量和速度。其次，这允许实践者使用更少的 GPU 来训练所需大小的模型，从而极大地节省成本和能源。第三，这能降低训练 LLM 的门槛并鼓励更多 GPU 资源有限的研究者参与进来。

但是，修改 Adam 但不影响其性能的难度非常大。其中一个主要原因是我们仍然不太理解 Adam 的 m 和 v 的作用。我们还不清楚 Adam 的哪些组件对其卓越性能而言必不可少，也就更不知道可以重新设计和改进哪些组件了。

Adafactor 是一个颇受欢迎的尝试，其能降低在 v 上低秩分解所用的内存。但是，很多研究发现 Adafactor 训练 LLM 的性能不佳。可能的原因有两个。第一，当前 Adam 中的 v 可能对有效性来说非常关键，已经没有缩减空间。第二，v 是有可能缩减的，只是 Adafactor 使用的方法并不是最合适的：矩阵分解是一种可以广泛应用的通用方法，但它没有利用太多特定于问题的结构，因此它在特定的神经网络任务上效果不佳。

近日，香港中文大学（深圳）、深圳市大数据研究院、杜克大学和斯坦福大学的一个联合研究团队发现：可以通过一个简单技巧来降低 v 的使用量。

目前而言，Adam 的 v 会为每个参数单独分配一个学习率，即第 i 个参数获得的学习率为

，，其中 v_i 是 v 的第 i 个组件。对于十亿参数量的模型，Adam 就要设计十亿个学习率。

该团队认为这些学习率资源可以大幅降低，同时还不会影响模型性能，甚至能让模型获得更优性能。该团队注意到，Transformer 的 Hessian 有一种接近块对角线的结构，其由不同大小的密集子块构成。他们发现，对于每一个这样的密集子块，都存在性能优于 Adam 的单个高质量学习率——只要有足够的资源将其搜索出来。由于密集子块的数量远少于参数数量，这个发现表明使用远远更少的学习率是有可能取得优良性能的。剩下的问题如何高效地找到它们。

为了找到性能足以比肩甚至优于 Adam 的优良学习率，该团队提出了一种低成本的简单方法 Adam-mini。另外，该团队也发布了 Adam-mini 的一种实现。

论文标题：Adam-mini: Use Fewer Learning Rates To Gain More
论文地址：https://arxiv.org/pdf/2406.16793
实现代码：https://github.com/zyushun/Adam-mini

在实验中，Adam-mini 的表现非常优秀。如图 1 所示，在预训练 Llama2-7B 时，Adam-mini 的内存占用可以大幅降低，同时吞吐量也有明显提升。在保证性能与 Adam 相当的同时，其能带来 33% 的速度提升。

图 2 展示了 Adam-mini 的大致图示和学习率使用情况。可以看到，Adam-mini 是根据 Hessian 结构来分配学习率，因此其使用的学习率比 SGD 多，但远少于 Adam。

本研究的通讯作者，香港中文大学（深圳）数据科学学院副教授孙若愚（Ruoyu Sun）表示：「如果你喜欢 Adam(W)，那 Adam-mini 就是你的不二之选！」

方法

动机和观察

为了应对特征值的异构性，Transformer 的每一参数块都需要不同的学习率。Adam 的 v 可以提供这一点。

论文《Why Transformers Need Adam: A Hessian Perspective》发现：必须为每个块使用不同的学习率。但是，Adam 所做的却远不止此：其并不只是为每个块分配不同的学习率，而是为每个参数都分配不同的学习率。而参数量（可能超过数十亿）要远远多于块的数量（通常数百计）。那问题就来了：

有必要为每个参数都使用单独的学习率吗？如果没必要，又能节省多少？

该团队进行了一番探索，得到了如下发现：

如图 4 (a) 和 (b) 所示，Adam 的表现优于最优的单学习率方法。这符合预期。
如图 4 (c) 和 (d) 所示，这里研究了其 Hessian 是 (a) 的一个密集子块的新问题。该团队探究了针对这个问题的最优单学习率方法，结果发现其表现优于 Adam，即使 Adam 分配的学习率要多得多。(a) 的所有三个子块都有类似的现象。
如果收集 (2) 中的那些最优学习率，并将其用于梯度下降的一个「逐块」版本，那么其在原始问题上的速度会比 Adam 快，见图 4(b) 中的绿线。

也就是说，对于带有块对角 Hessian 的一般问题而言，许多学习率并不一定会带来额外收益。

该团队也在 Transformer 中观察到了类似的现象。

总结起来，该团队发现：对于 Transformer，有可能使用比 Adam 远远更少的学习率实现与之相当或更好的性能。剩下的问题就是在不使用网格搜索的情况下如何找到这些学习率。

新方法：Adam-mini

基于上述讨论，该团队提出了 Adam-mini，见算法 1。Adam-mini 的目标是降低 Adam 中的学习率资源，而无需费力地对学习率执行网格搜索。Adam-mini 包含两个步骤。步骤 1 只在初始化时进行。

步骤 1-1：将模型参数切分为块。对于 Transformer，该团队选择的策略是「Partition for Transformers」，即根据头来切分所有的 Query 和 Key，并使用 PyTorch 来切分其余部分。对于其它网络，则使用默认的 PyTorch 切分即可，他们将其称之为「Partition for non-Transformers」。

步骤 1-2：使用 embd_blocks。对于 Transformer，其包含嵌入层和输出层。对于其它网络，不选择任何参数。

步骤 2：对于每个位于 embd_blocks 之外的参数块，都使用单个学习率。为了高效地为每个块选择合适的学习率，Adam-mini 的做法是直接将原始 Adam 中的 g◦g 替换成其均值。Adam-mini 像 Adam 一样在这些均值上使用了移动平均值。

一个简单示例。为了说明 Adam-mini 的关键设计，该团队在论文中给出了一个简单的示例。假设有一个问题，其有 5 个参数 w ∈ ℝ^5，Adam 和 Adam-mini 都会执行 w = w − u ◦ m，其中 m 是一阶动量，而 u 有不同的形式，如下所示：

对于 Adam：
‍‍
‍‍

对于 Adam-mini：假设分片方式为 (1, 2, 3) 和 (4, 5)，则

请注意，有效元素 u_mini 的数量就等于块的数量，这远远小于 u_Adam，其等于参数数量。结果表明，这能让 LLM 的 v 中的元素数量减少 90% 以上。

分片策略的原则

现在讨论如何为 Adam-mini 选择参数分片。该团队基于前述分析得出了一个广义原则：

原则 1：应当将参数切分成不同的参数块，使得每个参数块都关联了 Hessian 中最小的密集子块。

由于参数的块对角结构，PyTorch 的默认分片方法是一个合理的候选项。实际上，这种分片确实能很好地用于非 Transformer 任务，比如 ResNet、扩散模型和图模型。算法 3 Partition for non-Transformers 展示了该策略。

但不幸的是，该团队发现，默认的 PyTorch 切片并不总是能很好地应对 Transformer。比如他们发现 Adam-mini 在 1B 模型上会出现训练不稳定问题（见图 6(d)）。

他们猜测这是因为 PyTorch 切片无法完整地理解 Hessian结构。经过一番探索后，他们发现 Hessian 子模块分为两类：

第一类：类似于 Hessian 整体，这种 Hessian 子块本身也有进一步的块对角结构，由更小的密集矩阵构成。这一类包含 Query 和 Key。他们通过实验发现，小型密集子块的数量就等于多头注意力中头的数量。

第二类：这种 Hessian 子块有密集的结构，无法进一步分成更小的块。这一类包含 Value、注意力投射和 MLP 层。请注意 Value 的 Hessian 结构不同于 Query 和 Key 的 Hessian 结构，但它们全都由 4 个头组成。这是因为 Value 位于自注意力设计的 softmax 算子之外，而 Query 和 Key 却不是。

基于上述发现，该团队发现默认的 PyTorch 分片确实不是最适合 Transformer 的。可以根据头将 Query 和 Key 进一步切分成不同的块。根据原则 1，不同的头应当属于不同的块。根据直觉就能知道，不同的头在理解 token 时的作用也不同，因此它们需要不同的学习率就很合理了。

由此，该团队得到了算法 2：Partition for Transformers，其可根据头来切分 Query 和 Key。该过程遵循基于最小 Hessian 子块的原则如图 6 (d) 所示。该策略确实能稳定训练过程并提升性能。

Adam-mini 的一些特点

减少内存用量

根据设计，Adam-mini 可为 Transformer 减少学习率的数量——从参数总数减少到嵌入层的大小、输出层的大小和其它层中块的数量的总和。因此，减少的内存比例取决于模型中非嵌入参数的比例。

在 Llama2-7B 上，这个比例是 96.2%。对于主流的语言模型，这个比例通常 ≥ 90%。请注意，如果层数更多，这个比例也会更加接近 100%；比如对于 Llama3-70B，这个比例为 99.25%。因此，Adam-mini 可将 v 减少至少 90%，由此可为 Adam 带来 45% 到 50% 的内存节省。

吞吐量更高

Adam-mini 可取得比 AdamW 更高的吞吐量，尤其是当硬件资源有限时。

基于此，Adam-mini 可以减少预训练的总时间。如表 2 所示。当在 2 台 A800-80GB 上预训练 Llama2-7B 时，Adam-mini 的吞吐量比 AdamW 高 49.6%。召回吞吐量 (↑) 指的是每秒处理的 token 数量，因此在预训练时，Adam-mini 处理相同数量的 token 可节省 33.1% 的时间。

实验

预训练

设置

该团队预训练了一些开源 LLM，包括 GPT2 系列和 LLM 系列。他们在主流英语语料库上从头开始训练了这些模型。

具体来说，他们在 Openwebtext 上训练了 GPT2 系列，在 CommonCrawl 上训练了 TinyLlama-1B 和 Llama2-7B。他们使用 2B、3B 和 25B token 训练了模型。模型配置（比如上下文长度）是根据标准协议选择的。

该团队将 Adam-mini 与 AdamW 以及 Adafactor、CAME 和 SM3 等常用内存高效型方法进行了比较。对于 Adafactor 和 SM3，为了确保与其它方法进行公平的比较，这里整合了 β_1 = 0.9 的动量。他们在同样的预算内对所有方法的学习率进行了调节，并报告了最佳性能。

Llama 系列

图 7 (a) 展示了预训练 TinyLlama-1B 的验证损失曲线。Llama2-7B 的训练曲线已经在图 1(c) 中给出。对于 TinyLlama-1B 和 Llama2-7B，可以发现 Adam-mini 使用更少的内存便能达到与 AdamW 相近的效果。

GPT2 系列

图 7 (b) 展现了 GPT2-125M 的验证损失曲线，图 8 展示了大小从 330M 到 1.5B 的 GPT2 的验证损失曲线。可以看到，Adam-mini 使用更少的内存就能取得与 AdamW 相当的优良性能，而其它方法的表现要更差一些。

敏感度分析

该团队在 GPT2-125M 预训练任务上测试了 Adam-mini 对超参数的敏感度。这里报告了使用 2.5B token 训练之后的验证损失。如图 7 所示，Adam-mini 看起来对超参数并不敏感。

监督式微调和 RLHF

该团队评估了 Adam-mini 对下游微调任务的有效性。

具体来说，他们考虑了两种有代表性的任务：监督式微调（SFT）和根据人类反馈的强化学习（RLHF）。这里使用的预训练模型是 Llama-2-7B。

结果见表 3，可以看到 Adam-mini 的表现优于 AdamW，即便 Adam-mini 使用了单个学习率并且内存效率更高。

非 LLM 任务

该团队也在非 LLM 任务评估了 Adam-mini。表 4 给出了在 ImageNet 上训练 ResNet18、在 CelebA 上训练扩散模型、在 OGB-arxiv 上训练图卷积网络（GCN）和图注意力网络（GAT）的结果。结果发现 Adam-mini 使用更少的内存便能取得与 AdamW 相当或更好的性能。

留住用户，AIGC如何通过个性化提升转化率？

7月10日，《AIGC体验派》第四期，邀请到火山引擎直播洞察产品负责人刘晨晨和火山引擎销售助手产品负责人梁正，一起聊聊用户留存转化率：

提升直播间转化率，AIGC复盘如何留住用户？
聊到客户心坎里，金牌话术生成让小白秒变“老司机”？
销售Copilot，基于大模型的销售助手究竟是噱头还是真香？

识别海报二维码或点击阅读原文，立即报名直播。

转载请联系本公众号获得授权

投稿或寻求报道：[email protected]

微信扫码关注该文公众号作者

戳这里提交新闻线索和高质量文章给我们。

来源: qq

点击查看作者最近其他文章