8 位优化器

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

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

为了克服由此带来的计算、量化和稳定性挑战，8 位优化器包含三个部分

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

有了这些组件，使用 8 位状态执行优化器更新就变得很简单。在执行更新之前，8 位优化器状态会反量化为 32 位，然后状态会再次量化回 8 位进行存储。

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

稳定的嵌入层

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

稳定的嵌入层有三个特点

• 初始化：利用 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 上更新