如何看待FAIR提出的8-bit optimizer：效果和32-bit optimizer相当？ 第1页

谢邀。上午学习了一下这篇文章，顺便看了看代码，在这里抛砖引玉一下~

本文主要是提出的一种对 optimizer 进行量化的方法，在不修改超参，不影响模型精度的情况下，把 adam / momentum 的状态量量化至 int8，从而缓解训练时的显存压力。

这个问题的背景大概是随着模型越来越大，尤其是预训练模型规模指数级增长，对显存的需求也就越来越高，而原始的 adam 优化器（因为感觉在 nlp 中 adam 比 sgd/momentum 用的更多一些，所以后文主要讨论 adam）对于每个参数都需要 m 和 v 两个 fp32 的参数，相当于每 1B 的参数都需要 8G 的存储空间，占了整体的很大一部分。所以如果能够把 optimizer state 量化下来，就能适当缓解显存的压力。

先要对优化器量化的流程做一个简单的介绍。一个常规的流程是这样的：

       低精度优化器 --> 高精度优化器状态 --> 结合梯度更新参数 --> 重新量化回低精度参数     

毕竟直接少了 3/4 的信息，所以为了避免精度损失，作者主要提出了 3 个 trick。前两个是针对量化这个过程的，最后一个对 Embedding 结构的一个针对性调整。

Block-wise Quantization

作者把参数划分为了小 Block（在实践中使用的是每 2048 个参数一个 block），在进行量化的时候，按照 block 内绝对值最大的数对这个 block 进行归一化，使得所有参数都落在 [-1, 1] 这个范围。相较于之前的整个参数层一起归一，有 3 点好处：

1. 经过观察，在正态分布下，绝对值很大的参数的比例会很少，所以一起归一会使得大多数参数变得很小，从而使得量化过程中的一些数字范围对应的 int8 没有被充分利用，导致更多的信息丢失。而以 block 为单位会使这个影响限制在 block 中。

2. 一般来说，1 中提到的不到 1% 的这些”大数“ 往往是（arguably）更重要的，而量化过程可以保证最大的数的精度没有损失，所以划分为多个 block 之后可以保证更多“大数”的精度，从而保留重要部分的信息。

3. 分成小 block 有利于显卡的并行计算，提升计算效率。

Dynamic Quantization

第二条则是调整量化映射的方式。从 fp32 转至 int8，一般不会直接截断 cast，因为往往较小的数需要保留更多的小数位上的信息。所以之前作者提出过 Dynamic Tree Quantization，就是把 int8 也表示为类似于 fp32/fp16 的形式，分为指数部分和小数部分，如下图。这个结构表示的是 [-1, 1] 之间的数，分为 4 小部分：

1. 符号位；
2. 又连续多少 0 表明起始位是 1e-n；
3. 第一个标注为 1 的数为标注位，表示后面就是小数位了；
4. 后面的地方就是表示的一个线性的数值，例如下图中后 4 位是 9，而最大值是 15，所以为 9/15。

在本文中，因为观察到 adam 的 v 和 m 基本都在固定的 3~5 个数量级上，所以改成了固定的小数位数。并且因为 adam 的 v 项是恒正的，所以对于它去掉了指示符号的一位。

Stable Embedding Layer

最后是一个对 embedding layer 的一个改进。在实验中，他们发现 emebddign layer 经常出现梯度溢出等问题，所以在 embedding 中多加了个 layer norm，并且调小了初始值。文章宣称这种方法对原先 fp32 的训练也有效果。

具体实现

文章配了一个开源的 github，实现了高效版的 8bit Adam：

我去简单看了一下里面的实现，主要有这样几点。

• 基本就是每个 optimizer 实现了几个 kernel，分别是 fp32 版，int8 w/o blockwise, int8 w blockwise，都是 inplace 运算。里面比较广泛地使用了 Nvidia 的 cub 库：https://github.com/NVIDIA/cub，用来做 load, store 和 reduce（用来求最大值）。
• 在做量化方面，正向的查表（int8 -> fp32）就是在 python 中预先做好表再传入 kernel 的，反向的是通过类似 2 分法的方式完成的，具体可以看一下 dQuantize 和 quantize_2D 2 个函数。里面有配置一个随机的量化选项，我不太清楚这是干啥的...

有的朋友可能要问了，DeepSpeed 不都已经说了可以把 optimizer 移到 CPU 上去做了吗？那这个工作的意义在哪里呢？实际上，随着模型规模的不断提升，我们慢慢会把 CPU 内存也都用上，所以这个方法也可以起到降低 CPU 内存压力的效果。尤其是对于我们团队最新开源的派大星（PatrickStar），我们可以做到只把马上要进行计算的参数放在 GPU 上，其余部分全部动态 offload 至 CPU。那么这个工作可能也可以让派大星能支持的规模进一步提升（目前的规模基本在单张 V100，240G 内存，训练 12B 参数的模型）。对我们这个工作有兴趣的朋友，可以看看这里：

以上。