本文涉及到的详细测试代码和测试步骤放置于：
https://github.com/xinyuwei-david/david-share.git下的：Deel-Learning/Optimizers
本文中不再赘述代码实现。欢迎给repo点亮Star，您的点赞是作者持续创作的动力。
一、常见的优化器
常见的优化器，包括但不限于：
SGD（随机梯度下降）：最基本的梯度优化方法。
Momentum SGD：在SGD基础上加入动量，加速收敛。
Adagrad：具有自适应学习率，适应不同参数的学习速度。
Adam（自适应动量估计）：结合动量和自适应学习率，是一种广泛使用的优化器。
AdamW：Adam的改进版，解决了权重衰减问题。
二、AdamW 优化器的原理与内存消耗
AdamW 在训练期间会跟踪两个关键参数：
一阶矩（动量，记作 ( m_t )）：是过去梯度的指数移动平均，有助于引导优化朝着正确的方向发展。通过累积之前的梯度信息，动量使参数更新更平滑、更快速。
二阶矩（方差，记作 ( v_t )）：是梯度平方的指数移动平均，用于调整每个参数的学习率，避免在任何方向上出现过大的更新。
为每个模型参数存储这两个额外的状态（动量和二阶矩）会导致内存消耗的大幅增加。默认情况下，这些优化器状态使用float32数据类型，意味着优化器的内存占用是模型参数的两倍。
示例：以拥有 8.03B 个参数的 Llama 3.1 8B 模型为例，AdamW 优化器需要创建并存储：一阶矩：8.03B 个参数的动量。
二阶矩：8.03B 个参数的二阶矩。
总计需要存储 16.06B 个额外参数。由于每个参数占用 4 字节（float32），优化器状态共消耗 16.06B × 4 = 64.24 GB 的内存。
三、Paged AdamW 优化器
Paged AdamW 是 AdamW 的一个变体，主要区别在于其内存管理策略。它允许优化器状态在 GPU RAM 和 CPU RAM 之间进行切换，以节省 GPU 内存。
内存管理策略
分页机制：当 GPU 内存不足时，优化器状态会被“分页”到 CPU 内存。这类似于操作系统的内存分页，只有在需要时才进行数据传输。
优势：这使得即使在 GPU 内存有限的情况下，也能训练大型模型。由于 NVIDIA GPU 的统一内存技术，分页过程对性能影响较小，确保了训练速度。
四、代码中可指定的优化器列表
在训练过程中，可以指定的优化器包括：
['adamw_hf', 'adamw_torch', 'adamw_torch_fused', 'adamw_torch_xla', 'adamw_torch_npu_fused', 'adamw_apex_fused', 'adafactor', 'adamw_anyprecision', 'sgd', 'adagrad', 'adamw_bnb_8bit', 'adamw_8bit', 'lion_8bit', 'lion_32bit', 'paged_adamw_32bit', 'paged_adamw_8bit', 'paged_lion_32bit', 'paged_lion_8bit', 'rmsprop', 'rmsprop_bnb', 'rmsprop_bnb_8bit', 'rmsprop_bnb_32bit', 'galore_adamw', 'galore_adamw_8bit', 'galore_adafactor', 'galore_adamw_layerwise', 'galore_adamw_8bit_layerwise', 'galore_adafactor_layerwise', 'lomo', 'adalomo']
五、解决内存消耗的技术：量化与分页
为了有效管理 AdamW 优化器的高内存消耗，引入了以下两种技术：
1. 8-bit 量化 AdamW
原理：将优化器状态从
float32（32 位精度）降为int8（8 位精度），大幅减少内存占用。效果：对于 Llama 3.1 模型，优化器状态的内存需求从约 64 GB 降至 16 GB。尽管量化可能引入一些数值误差，但对训练稳定性和损失的影响微乎其微。
2. 分页 AdamW
原理：利用 CUDA 的统一内存特性，当 GPU 内存不足时，将优化器状态自动分页到 CPU 内存。
优势：
高效性：相比手动的 CPU 卸载，分页机制更加高效，无需显式管理数据传输。
灵活性：如果 GPU 内存充足，优化器状态将完全保留在 GPU 内存中，几乎没有性能开销。
六、实验结果与分析
在 A100 GPU 上对不同配置的优化器进行了实验，结果如下：
学习曲线对比
8-bit 量化 vs. 32-bit 精度：使用 8-bit 量化的优化器状态与 32-bit 版本的学习曲线几乎没有差异。这表明了 8-bit 量化在显著节省内存的同时，并未对训练过程产生负面影响。
内存消耗
标准 AdamW：由于内存限制（超出 40 GB），无法在单一 GPU 上完整训练 Llama 3.1 模型。
分页优化器：只有使用分页优化器的配置能够避免内存不足的问题。通过分页技术，可以在一块 40 GB 的 GPU 上成功完成训练。
训练时间
分页优化器的性能：分页优化器对训练时间的影响不显著。由于 8-bit 量化状态占用更少内存，Paged AdamW 8-bit 甚至比 Paged AdamW 32-bit 运行得更快。
总结
AdamW 优化器：在大型模型训练中表现优秀，但高内存消耗是主要瓶颈。
8-bit 量化：通过降低优化器状态的精度，可以大幅减少内存占用，且对训练效果影响很小。
分页技术：利用 GPU 的统一内存，实现了优化器状态在 GPU 和 CPU 间的自动分页，有效解决了内存限制问题。
实践意义：结合 8-bit 量化和分页技术，可以在有限的硬件资源下训练大型模型，为模型训练的普及和推广提供了可行方案。
参考：https://newsletter.kaitchup.com/p/fine-tuning-llms-with-32-bit-8-bit
