2024/07/08 15:44

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

在训练大型语言模型（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 相当或更好的性能。

工程Adam-miniAdam

相关数据

网格搜索技术

网格搜索是一项模型超参数优化技术，常用于优化三个或者更少数量的超参数，本质是一种穷举法。对于每个超参数，使用者选择一个较小的有限集去探索。然后，这些超参数笛卡尔乘积得到若干组超参数。网格搜索使用每组超参数训练模型，挑选验证集误差最小的超参数作为最好的超参数。

来源：Deep Learning Book

自注意力技术

自注意力（Self-attention），有时也称为内部注意力，它是一种涉及单序列不同位置的注意力机制，并能计算序列的表征。自注意力在多种任务中都有非常成功的应用，例如阅读理解、摘要概括、文字蕴含和语句表征等。自注意力这种在序列内部执行 Attention 的方法可以视为搜索序列内部的隐藏关系，这种内部关系对于翻译以及序列任务的性能非常重要。

参数技术

在数学和统计学裡，参数（英语：parameter）是使用通用变量来建立函数和变量之间关系（当这种关系很难用方程来阐述时）的一个数量。

来源：维基百科

数据科学技术

数据科学，又称资料科学，是一门利用数据学习知识的学科，其目标是通过从数据中提取出有价值的部分来生产数据产品。它结合了诸多领域中的理论和技术，包括应用数学、统计、模式识别、机器学习、数据可视化、数据仓库以及高性能计算。数据科学通过运用各种相关的数据来帮助非专业人士理解问题。

来源：维基百科

学习率技术

在使用不同优化器（例如随机梯度下降，Adam）神经网络相关训练中，学习速率作为一个超参数控制了权重更新的幅度，以及训练的速度和精度。学习速率太大容易导致目标（代价）函数波动较大从而难以找到最优，而弱学习速率设置太小，则会导致收敛过慢耗时太长

来源：Liu, T. Y. (2009). Learning to rank for information retrieval. Foundations and Trends® in Information Retrieval, 3(3), 225-331. Wikipedia

超参数技术

在机器学习中，超参数是在学习过程开始之前设置其值的参数。相反，其他参数的值是通过训练得出的。不同的模型训练算法需要不同的超参数，一些简单的算法（如普通最小二乘回归）不需要。给定这些超参数，训练算法从数据中学习参数。相同种类的机器学习模型可能需要不同的超参数来适应不同的数据模式，并且必须对其进行调整以便模型能够最优地解决机器学习问题。在实际应用中一般需要对超参数进行优化，以找到一个超参数元组（tuple），由这些超参数元组形成一个最优化模型，该模型可以将在给定的独立数据上预定义的损失函数最小化。

来源：Wikipedia

神经网络技术

（人工）神经网络是一种起源于 20 世纪 50 年代的监督式机器学习模型，那时候研究者构想了「感知器（perceptron）」的想法。这一领域的研究者通常被称为「联结主义者（Connectionist）」，因为这种模型模拟了人脑的功能。神经网络模型通常是通过反向传播算法应用梯度下降训练的。目前神经网络有两大主要类型，它们都是前馈神经网络：卷积神经网络（CNN）和循环神经网络（RNN），其中 RNN 又包含长短期记忆（LSTM）、门控循环单元（GRU）等等。深度学习是一种主要应用于神经网络帮助其取得更好结果的技术。尽管神经网络主要用于监督学习，但也有一些为无监督学习设计的变体，比如自动编码器和生成对抗网络（GAN）。

来源：机器之心

梯度下降技术

梯度下降是用于查找函数最小值的一阶迭代优化算法。要使用梯度下降找到函数的局部最小值，可以采用与当前点的函数梯度（或近似梯度）的负值成比例的步骤。如果采取的步骤与梯度的正值成比例，则接近该函数的局部最大值，被称为梯度上升。

来源：Vapnik V. N. (2000). The Nature of Statistical Learning Theory. Information Science and Statistics. Springer-Verlag.Wikipedia

语料库技术

语料库一词在语言学上意指大量的文本，通常经过整理，具有既定格式与标记；事实上，语料库英文 "text corpus" 的涵意即为"body of text"。

来源：维基百科

动量技术

优化器的一种，是模拟物理里动量的概念，其在相关方向可以加速SGD，抑制振荡，从而加快收敛

来源：An overview of gradient descent optimization algorithms

强化学习技术

强化学习是一种试错方法，其目标是让软件智能体在特定环境中能够采取回报最大化的行为。强化学习在马尔可夫决策过程环境中主要使用的技术是动态规划（Dynamic Programming）。流行的强化学习方法包括自适应动态规划（ADP）、时间差分（TD）学习、状态-动作-回报-状态-动作（SARSA）算法、Q 学习、深度强化学习（DQN）；其应用包括下棋类游戏、机器人控制和工作调度等。

来源：机器之心

优化器技术

优化器基类提供了计算梯度loss的方法，并可以将梯度应用于变量。优化器里包含了实现了经典的优化算法，如梯度下降和Adagrad。优化器是提供了一个可以使用各种优化算法的接口，可以让用户直接调用一些经典的优化算法，如梯度下降法等等。优化器（optimizers）类的基类。这个类定义了在训练模型的时候添加一个操作的API。用户基本上不会直接使用这个类，但是你会用到他的子类比如GradientDescentOptimizer, AdagradOptimizer, MomentumOptimizer（tensorflow下的优化器包）等等这些算法。

来源：维基百科

语言模型技术

语言模型经常使用在许多自然语言处理方面的应用，如语音识别，机器翻译，词性标注，句法分析和资讯检索。由于字词与句子都是任意组合的长度，因此在训练过的语言模型中会出现未曾出现的字串(资料稀疏的问题)，也使得在语料库中估算字串的机率变得很困难，这也是要使用近似的平滑n元语法(N-gram)模型之原因。

来源：维基百科

矩阵分解技术

矩阵分解是一种将矩阵简化为其组成部分的方法。这种方法可以简化更复杂的矩阵运算，这些运算可以在分解的矩阵上执行，而不是在原始矩阵本身上执行。它的衍生Non-negative matrix factorization也被用于降维等操作上。

来源：矩阵分解

图卷积网络技术

假设有一张图，要做分类，传统方法需要手动提取一些特征，比如纹理啊，颜色啊，或者一些更高级的特征。然后再把这些特征放到像随机森林等分类器，给到一个输出标签，告诉它是哪个类别。而深度学习是输入一张图，经过神经网络，直接输出一个标签。特征提取和分类一步到位，避免了手工提取特征或者人工规则，从原始数据中自动化地去提取特征，是一种端到端（end-to-end）的学习。相较于传统的方法，深度学习能够学习到更高效的特征与模式。

来源：图卷积网络