Adam-mini：内存占用减半，性能更优的深度学习优化器

Adam(W)目前为训练LLM的主流优化器，但其内存开销较大，这是因为Adam优化器需要存储一阶动量m和二阶动量v，总内存占用至少是模型大小的两倍，这对现有的高端显卡也是一种负担。论文提出一种新的优化器Adam-mini，在不牺牲性能的情况下减少Adam优化器的内存占用。

Adam-mini

Adam-mini通过减少学习率资源来降低内存占用的具体方法如下：

参数分块：Adam-mini首先将模型参数按照Hessian矩阵的结构划分为多个块。Hessian矩阵通常具有近似块对角结构，每个块代表一组参数。论文提出的分块策略基于Hessian结构，将每个块内的参数视为一个整体进行处理。

块内平均学习率：对于每个参数块，Adam-mini不再为每个参数单独分配学习率，而是为整个块分配一个平均的学习率。具体方法是计算块内所有参数的梯度平方的平均值，然后基于这个平均值来计算该块的学习率。这一过程显著减少了所需的学习率数量。

内存节省：由于Adam-mini使用的学习率数量大大减少，所需的二阶动量（即Adam中的v参数）的存储也相应减少。论文中的实验表明，这种方法可以减少45%到50%的内存占用。

具体算法：

在初始化时，将模型参数分块。

对于每个参数块，计算块内梯度平方的平均值，并更新块的学习率。

使用更新后的学习率进行参数更新。

这种方法不仅减少了内存占用，还通过减少GPU和CPU之间的通信开销，提高了训练效率。例如，在Llama2-7B模型的预训练中，Adam-mini在两块A800-80GB GPU上实现了比AdamW高49.6%的吞吐量，并节省了33%的训练时间。

算法示例

# Adam-mini 的伪代码

def adam_mini(params, grads, lr, beta1, beta2, epsilon, weight_decay):

   # 初始化动量和二阶动量

   m = {}

   v = {}

   for param in params:

       m[param] = np.zeros_like(param)

       v[param] = np.zeros_like(param)

   # 参数分块

   param_blocks = partition_parameters(params)

   for block in param_blocks:

       # 获取当前块的梯度

       grad_block = [grads[param] for param in block]

       # 更新动量

       m_block = (1 - beta1) * grad_block + beta1 * m_block

       m_block_hat = m_block / (1 - beta1**t)

       # 更新二阶动量（平均值）

       v_block = (1 - beta2) * np.mean([g**2 for g in grad_block]) + beta2 * v_block

       v_block_hat = v_block / (1 - beta2**t)

       # 更新参数

       for param in block:

           param_update = lr * m_block_hat / (np.sqrt(v_block_hat) + epsilon)

           params[param] -= param_update + weight_decay * params[param]

   return params

通过这种方法，Adam-mini成功地减少了学习率资源的使用，从而大幅降低了内存占用，并在多种任务中表现出色。

性能表现

1、内存和吞吐量性能

Adam-mini在预训练Llama2-7B模型时的性能：

内存占用：Adam-mini显著降低了内存占用。例如，在Llama2-7B预训练时，Adam-mini减少了45%到50%的内存消耗。

吞吐量：由于内存减少，Adam-mini能够支持更大的每GPU批次大小，从而提高了吞吐量。在两块A800-80GB GPU上，Adam-mini实现了比AdamW高49.6%的吞吐量，节省了33%的训练时间。

2、预训练性能

TinyLlama-1B：图7(a)显示了TinyLlama-1B的验证损失曲线。Adam-mini的表现与AdamW相当，但内存占用更低。

GPT2系列：图8展示了GPT2不同规模模型的训练曲线，包括GPT2-125M、GPT2-330M、GPT2-770M和GPT2-1.5B。Adam-mini在所有这些模型上的表现均与AdamW相当，而Adafactor、CAME等方法在这些任务中的表现较差。

通过这些实验结果，论文证明了Adam-mini不仅在内存占用和计算效率上优于传统的AdamW，还能在不同任务中保持或提升模型性能。这些结果表明，Adam-mini是一个有效且高效的优化器，适用于大规模模型的训练和微调。

非LLM任务的表现

Adam-mini在多种非LLM任务中均表现出色，能够在减少内存占用的同时，保持或提升模型性能。这些结果证明了Adam-mini在图像识别、扩散模型训练和图卷积网络等任务中的广泛适用性和有效性。

图像分类：在ImageNet上训练ResNet18，Adam-mini的测试精度与AdamW相当。

扩散模型训练：在CelebA数据集上训练扩散模型，Adam-mini的训练损失低于AdamW。

图神经网络：在OGB-arxiv数据集上训练Graph Convolution Network (GCN)和Graph Attention Network (GAT)，Adam-mini的验证精度优于或相当于AdamW。

总结

Adam-mini基于Hessian矩阵的结构，将模型参数划分为多个块，每个块使用单一的平均学习率，从而大幅减少了需要存储的学习率数量。在非LLM任务中的实验结果进一步验证了Adam-mini的广泛适用性。Adam-mini不仅在内存占用和计算效率方面具有优势，还能在多种任务中保持或提升模型性能，是一个有效且高效的优化器。

MORE

kaggle比赛交流和组队

喜欢就关注一下吧：

点个在看你最好看！