Adam(W)目前为训练LLM的主流优化器，但其内存开销较大，这是因为Adam优化器需要存储一阶动量m和二阶动量v，总内存占用至少是模型大小的两倍，这对现有的高端显卡也是一种负担。论文提出一种新的优化器Adam-mini，在不牺牲性能的情况下减少Adam优化器的内存占用。
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是一个有效且高效的优化器，适用于大规模模型的训练和微调。
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不仅在内存占用和计算效率方面具有优势，还能在多种任务中保持或提升模型性能，是一个有效且高效的优化器。
https://avoid.overfit.cn/post/fdf7a6cfd34f4158aaac31a6ed3cc9b6
