BAdam: A Memory Efficient Full Parameter Optimization Method for Large Language Models

URL

• paper: https://arxiv.org/pdf/2404.02827
• code: https://github.com/Ledzy/BAdam

TL;DR

• BAdam 全称是 Blockwise Adam，是一种内存高效的大型语言模型参数优化方法。
• BAdam 的核心思想是将模型参数分块，每次只更新一个块的参数，而且对每个块用 Adam 连续更新 K 次，然后再更新下一个块。

Algorithm

• BAdam 的算法流程如上图所示，其中 K 是一个超参数，表示每个块的参数用 Adam 更新次数。
• 当一个块的参数更新 K 次后，就丢掉 这个块的所有优化器信息，包括梯度、一阶动量、二阶动量。
• 假设模型有 M billion 参数，最低内存占用：
  • Adam：
    • 参数：2M GB（fp16）
    • 优化器：
      • 梯度：4M GB（fp32）
      • 参数：4M GB（fp32）
      • 一阶动量：4M GB（fp32）
      • 二阶动量：4M GB（fp32）
    • 总共：18M GB
  • BAdam：
    • 参数：2M GB（fp16）
    • 优化器：
      • 梯度：$\frac{4M}{D}$ GB（fp32）
      • 参数：$\frac{4M}{D}$ GB（fp32）
      • 一阶动量：$\frac{4M}{D}$ GB（fp32）
      • 二阶动量：$\frac{4M}{D}$ GB（fp32）
    • 总共：$2M+\frac{16M}{D}$ GB，其中 D 是块的数量。

Thoughts

• BAdam 的核心思想是将模型参数分块更新，但实际上和 Adam 的思想是 完全不一样
• 因为 Adam 想要追求的是全训练过程的步长自适应，即每一个 step 的步长都来自于只有所有历史信息。
• 而 BAdam 只能保证当前步长最多由之前 K 个 step 的信息决定，所以 BAdam 的收敛性和 Adam 是不一样的。