8 位优化器

有状态的优化器会随着时间的推移维护梯度统计信息，例如，过去梯度值的指数平滑和（带动量的 SGD）或平方和 (Adam)。与普通的随机梯度下降相比，此状态可用于加速优化，但会使用可能原本分配给模型参数的内存。因此，这限制了实践中可以训练的模型最大尺寸。现在来看看可以使用 8 位优化器训练的最大模型。

bitsandbytes 优化器使用 8 位统计信息，同时保持使用 32 位优化器状态的性能水平。

为了克服由此产生的计算、量化和稳定性挑战，8 位优化器具有三个组件

1. 分块量化：将输入张量分成更小的块，这些块被独立量化，隔离异常值并将误差更均匀地分布在所有位上。每个块在内核之间并行处理，从而实现更快的优化和高精度量化。
2. 动态量化：以高精度量化小值和大值。
3. 稳定嵌入层：提高了具有词嵌入模型的优化期间的稳定性。

有了这些组件，使用 8 位状态执行优化器更新非常简单。在执行更新之前，8 位优化器状态被反量化为 32 位，然后状态被量化回 8 位以进行存储。

8 位到 32 位的转换在寄存器中逐元素发生，这意味着执行量化和反量化不需要缓慢地复制到 GPU 内存或额外的临时内存。对于 GPU，这使得 8 位优化器比常规 32 位优化器快得多。

稳定嵌入层

稳定嵌入层提高了 NLP 任务标准词嵌入层的训练稳定性。它解决了非均匀输入分布的挑战并减轻了极端的梯度变化。这意味着稳定嵌入层可以支持更激进的量化策略，而不会损害训练稳定性，并且它可以帮助实现稳定的训练结果，这对于处理多样化和复杂语言数据的模型尤其重要。

稳定嵌入层有三个特点

• 初始化：利用 Xavier 均匀初始化来保持一致的方差，从而降低大梯度的可能性。
• 归一化：在添加位置嵌入之前合并层归一化，有助于输出稳定性。
• 优化器状态：专门为此层采用 32 位优化器状态以增强稳定性，而模型的其余部分可以使用标准的 16 位精度。

分页优化器

分页优化器构建在 CUDA 的统一内存功能之上。统一内存提供了一个 GPU 和 CPU 可以轻松访问的单一内存空间。虽然 PyTorch 不支持此功能，但它已添加到 bitsandbytes 中。

分页优化器的工作方式类似于常规的 CPU 分页，这意味着它仅在 GPU 内存耗尽时才激活。当发生这种情况时，内存会从 GPU 逐页传输到 CPU。内存是映射的，这意味着页面在 CPU 上预先分配，但它们不会自动更新。只有在访问内存或启动交换操作时才会更新页面。

统一内存功能不如常规的异步内存传输有效，并且您通常无法获得完整的 PCIe 内存带宽利用率。如果您执行手动预取，传输速度可能会很高，但仍然只有 PCIe 内存带宽的一半或更差（在 16x 通道的 PCIe 3.0 上测试）。

这意味着性能高度依赖于特定的用例。例如，如果您在每个前向-后向-优化器循环中逐出 1 GB 的内存，那么在最佳情况下，您可以预期大约 50% 的 PCIe 带宽作为时间开销。因此，对于具有 16x 通道的 PCIe 3.0 的 1 GB 内存将以 16 GB/s 的速度运行，这意味着每个优化器步骤的开销为 1/(16*0.5) = 1/8 = 125ms。考虑到 PCIe 接口、通道以及每次迭代中逐出的内存，可以估计其他开销。

与 CPU 卸载相比，如果所有内存都适合设备，则分页优化器具有零开销；如果需要逐出某些内存，则只有一些开销。对于卸载，您通常卸载模型的固定部分，并且需要在每次迭代模型时卸载和加载所有这些内存（有时正向和反向传播都需要两次）。

< > 在 GitHub 上更新