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

在训练大型语言模型(LLM)时,Adam(W) 基本上已经成为了人们默认使用的优化器。Adam 尽管性能优异,但使用成本很高。具体来说,Adam 需要内存来保存其优化器状态:一阶动量 m 和二阶动量 v^2。这总共需要模型大小至少 2 倍的内存。这样的内存消耗已经成为了 LLM 训练的一大主要负担。举个例子,要训练一个 7B 模型,只是 Adam 就需要每张卡有大约 56 GB 来保存 m 和 v;而如果再加上梯度,则总共需要 86 GB。即使使用最先进的 A100-80GB,成本也过高了。为了支持这样的高内存算法

在训练大型语言模型(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 个参数获得的学习率为 Adam有了mini版:内存占用少一半,吞吐量提升50%,,其中 v_i 是 v 的第 i 个组件。对于十亿参数量的模型,Adam 就要设计十亿个学习率。

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

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

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

论文标题: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% 的速度提升。

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

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

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

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

方法

动机和观察

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

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

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

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

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

如图 4 (a) 和 (b) 所示,Adam 的表现优于最优的单学习率方法。这符合预期。

如图 4 (c) 和 (d) 所示,这里研究了其 Hessian 是 (a) 的一个密集子块的新问题。该团队探究了针对这个问题的最优单学习率方法,结果发现其表现优于 Adam,即使 Adam 分配的学习率要多得多。(a) 的所有三个子块都有类似的现象。

如果收集 (2) 中的那些最优学习率,并将其用于梯度下降的一个「逐块」版本,那么其在原始问题上的速度会比 Adam 快,见图 4(b) 中的绿线。

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

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

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

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

新方法:Adam-mini

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

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

步骤 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版:内存占用少一半,吞吐量提升50%

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

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

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

分片策略的原则

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

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

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

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

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

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

他们猜测这是因为 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 减少学习率的数量——从参数总数减少到嵌入层的大小、输出层的大小和其它层中块的数量的总和。因此,减少的内存比例取决于模型中非嵌入参数的比例。

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

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

吞吐量更高

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

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

基于此,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 系列

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

图 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 相当的优良性能,而其它方法的表现要更差一些。

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

敏感度分析

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

监督式微调和 RLHF

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

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

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

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

非 LLM 任务

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

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

给TA打赏
共{{data.count}}人
人已打赏
工程

RAGFlow开源Star量破万,是时候思考下RAG的未来是什么了

2024-7-8 15:36:00

工程

开源视频版GPT-4o?快速记忆,实时问答,拿下CVPR'24长视频问答竞赛冠军

2024-7-8 16:00:00

0 条回复 A文章作者 M管理员
    暂无讨论,说说你的看法吧
个人中心
今日签到
搜索