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

本文涉及的产品
实时数仓Hologres,5000CU*H 100GB 3个月
实时计算 Flink 版,5000CU*H 3个月
检索分析服务 Elasticsearch 版,2核4GB开发者规格 1个月
简介: 论文提出一种新的优化器Adam-mini,在不牺牲性能的情况下减少Adam优化器的内存占用。

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

Adam-mini

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

  1. 参数分块:Adam-mini首先将模型参数按照Hessian矩阵的结构划分为多个块。Hessian矩阵通常具有近似块对角结构,每个块代表一组参数。论文提出的分块策略基于Hessian结构,将每个块内的参数视为一个整体进行处理。
  2. 块内平均学习率:对于每个参数块,Adam-mini不再为每个参数单独分配学习率,而是为整个块分配一个平均的学习率。具体方法是计算块内所有参数的梯度平方的平均值,然后基于这个平均值来计算该块的学习率。这一过程显著减少了所需的学习率数量。
  3. 内存节省:由于Adam-mini使用的学习率数量大大减少,所需的二阶动量(即Adam中的v参数)的存储也相应减少。论文中的实验表明,这种方法可以减少45%到50%的内存占用。
  4. 具体算法:- 在初始化时,将模型参数分块。- 对于每个参数块,计算块内梯度平方的平均值,并更新块的学习率。- 使用更新后的学习率进行参数更新。

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

算法示例

 # Adam-mini 的伪代码
 defadam_mini(params, grads, lr, beta1, beta2, epsilon, weight_decay):
     # 初始化动量和二阶动量
     m= {}
     v= {}
     forparaminparams:
         m[param] =np.zeros_like(param)
         v[param] =np.zeros_like(param)

     # 参数分块
     param_blocks=partition_parameters(params)

     forblockinparam_blocks:
         # 获取当前块的梯度
         grad_block= [grads[param] forparaminblock]

         # 更新动量
         m_block= (1-beta1) *grad_block+beta1*m_block
         m_block_hat=m_block/ (1-beta1**t)

         # 更新二阶动量(平均值)
         v_block= (1-beta2) *np.mean([g**2forgingrad_block]) +beta2*v_block
         v_block_hat=v_block/ (1-beta2**t)

         # 更新参数
         forparaminblock:
             param_update=lr*m_block_hat/ (np.sqrt(v_block_hat) +epsilon)
             params[param] -=param_update+weight_decay*params[param]

     returnparams

通过这种方法,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不仅在内存占用和计算效率方面具有优势,还能在多种任务中保持或提升模型性能,是一个有效且高效的优化器。

https://avoid.overfit.cn/post/fdf7a6cfd34f4158aaac31a6ed3cc9b6

相关实践学习
部署Stable Diffusion玩转AI绘画(GPU云服务器)
本实验通过在ECS上从零开始部署Stable Diffusion来进行AI绘画创作,开启AIGC盲盒。
目录
相关文章
|
3月前
|
存储 缓存 监控
|
2月前
|
监控 JavaScript 算法
如何使用内存监控工具来定位和解决Node.js应用中的性能问题?
总之,利用内存监控工具结合代码分析和业务理解,能够逐步定位和解决 Node.js 应用中的性能问题,提高应用的运行效率和稳定性。需要耐心和细致地进行排查和优化,不断提升应用的性能表现。
204 77
|
2月前
|
存储 缓存 JavaScript
如何优化Node.js应用的内存使用以提高性能?
通过以上多种方法的综合运用,可以有效地优化 Node.js 应用的内存使用,提高性能,提升用户体验。同时,不断关注内存管理的最新技术和最佳实践,持续改进应用的性能表现。
147 62
|
2月前
|
存储 缓存 监控
如何使用内存监控工具来优化 Node.js 应用的性能
需要注意的是,不同的内存监控工具可能具有不同的功能和特点,在使用时需要根据具体工具的要求和操作指南进行正确使用和分析。
82 31
|
1月前
|
存储 缓存 监控
Docker容器性能调优的关键技巧,涵盖CPU、内存、网络及磁盘I/O的优化策略,结合实战案例,旨在帮助读者有效提升Docker容器的性能与稳定性。
本文介绍了Docker容器性能调优的关键技巧,涵盖CPU、内存、网络及磁盘I/O的优化策略,结合实战案例,旨在帮助读者有效提升Docker容器的性能与稳定性。
153 7
|
1月前
|
机器学习/深度学习 算法框架/工具 网络架构
深度学习中的正则化技术及其对模型性能的影响
本文深入探讨了深度学习领域中正则化技术的重要性,通过分析L1、L2以及Dropout等常见正则化方法,揭示了它们如何帮助防止过拟合,提升模型的泛化能力。文章还讨论了正则化在不同类型的神经网络中的应用,并指出了选择合适正则化策略的关键因素。通过实例和代码片段,本文旨在为读者提供关于如何在实际问题中有效应用正则化技术的深刻见解。
|
2月前
|
并行计算 算法 测试技术
C语言因高效灵活被广泛应用于软件开发。本文探讨了优化C语言程序性能的策略,涵盖算法优化、代码结构优化、内存管理优化、编译器优化、数据结构优化、并行计算优化及性能测试与分析七个方面
C语言因高效灵活被广泛应用于软件开发。本文探讨了优化C语言程序性能的策略,涵盖算法优化、代码结构优化、内存管理优化、编译器优化、数据结构优化、并行计算优化及性能测试与分析七个方面,旨在通过综合策略提升程序性能,满足实际需求。
80 1
|
2月前
|
缓存 Ubuntu Linux
Linux环境下测试服务器的DDR5内存性能
通过使用 `memtester`和 `sysbench`等工具,可以有效地测试Linux环境下服务器的DDR5内存性能。这些工具不仅可以评估内存的读写速度,还可以检测内存中的潜在问题,帮助确保系统的稳定性和性能。通过合理配置和使用这些工具,系统管理员可以深入了解服务器内存的性能状况,为系统优化提供数据支持。
67 4
|
2月前
|
机器学习/深度学习 测试技术 PyTorch
深度学习之测量GPU性能的方式
在深度学习中,测量GPU性能是一个多方面的任务,涉及运行时间、吞吐量、GPU利用率、内存使用情况、计算能力、端到端性能测试、显存带宽、框架自带性能工具和基准测试工具等多种方法。通过综合使用这些方法,可以全面评估和优化GPU的性能,提升深度学习任务的效率和效果。
150 5
|
2月前
|
机器学习/深度学习 算法 开发者
探索深度学习中的优化器选择对模型性能的影响
在深度学习领域,优化器的选择对于模型训练的效果具有决定性作用。本文通过对比分析不同优化器的工作原理及其在实际应用中的表现,探讨了如何根据具体任务选择合适的优化器以提高模型性能。文章首先概述了几种常见的优化算法,包括梯度下降法、随机梯度下降法(SGD)、动量法、AdaGrad、RMSProp和Adam等;然后,通过实验验证了这些优化器在不同数据集上训练神经网络时的效率与准确性差异;最后,提出了一些基于经验的规则帮助开发者更好地做出选择。