大模型Token打骨折：DeepSeek百万Token只需1元——谈谈大模型中常用的量化操作对成本影响

AI教程

2024-05-09

大模型Token打骨折：DeepSeek百万Token只需1元——谈谈大模型中常用的量化操作对成本影响

1 大模型中的量化操作

在 GPT-3 之后, NLP 领域进入了由 pretrain + prompt 的范式代替 pretrain + fine-tuning 的 Bert 范式的时代.

从目前 ChatGPT 按 token 收费的情况来看, 每 token 推理成本的摊销将会决定产品定价和背后的营利空间. 如果按目前收费单位为: xx 美分/1k token, 推理成本可以简单转化为模型每 1k token 的平均推理时间 x 功耗 x 电费计价(当然, 这里面没有统计 GPU 的购买成本/ 人力成本等). 尤其对于像国内有些公司是希望借助 GPT 的能力来优化自身产品以实现营利(我个人将其理解为"隐式收费"), 那其 GPT 的运营成本将更为敏感。

最近DeepSeek出了最新的API，将价格打骨折了，在某些任务上效果和GPT4-turbo不相上下。下面是一个整理的价格对比表

我们关注的是模型的平均推理时间和功耗, 平均推理时间可以用 latency 或 throughput 来衡量, 而功耗则可以用参考生成 token 过程中所用到 GPU 的功耗(因为如果是用 TP/PP 等方法就会引入多个 GPU)来近似。

量化(Quantization), 可以很好地通过将 float 模型表征为低位宽模型实现减小模型存储空间, 加速模型推理的目标。

量化, 在深度学习领域中, 可以将其定义为

a technique that mapping of a k-bit integer to a float element, which saves space and speedup computation by compressing the digital representation

量化可以按不同角度对其进行归类: 按量化执行的阶段,：

• 可以分为训练中量化(QAT, Quantization-Aware-Training)
• 训练后量化(PTQ, Post-Training-Quantization); 按量化间隔是否等距, 可以其将分为：
• 均匀量化
• 非均匀量化

在 LLMs 背景下, QAT 的研究目前仍未有机构做出相关的靠谱研究, 主要受限于 QAT 需要引入模拟量化的操作, 会引起显存&计算量进一步上涨以及梯度 mismatch 的问题, 从而增加训练成本以及影响 Scaling Laws. 非均匀量化除非有特殊硬件支持, 否则在 GPU 上目前多数只能通过 Look-Up-Table 或移位等方式来实现, 速度和精度没法得到同时保证.

2 背景

这两个公式表达了量化和反量化的过程。量化通常用于数据压缩，在数字信号处理中非常常见，尤其是在模拟信号转换为数字信号的过程中。

Q(X) = \text{clamp}\left(\left\lfloor \frac{X}{S} \right\rfloor + Z, 0, 2^b - 1\right)

\hat{X} = (Q(X) - Z) \times S

第一个公式定义了量化操作 Q(X)：这里：

• clamp 函数将值限制在指定的范围内，这里是 0 到 2^b - 1。
• X 是要量化的输入值，代表tensor。
• S 是量化的比例因子，即为 scale。
• Z 是零点偏移，用于确保量化的对称性。
• b 是量化位宽。

第二个公式定义了反量化操作 X^：这里：

• X^ 是反量化后的值，应接近原始值 X。
• Q(X) 是量化后的值。
• Z 和 S 分别是与量化过程中相同的零点偏移和比例因子。

这两个操作通常在数字信号处理、图像和声音压缩等领域中使用。量化可以减少所需的存储空间，但也会引入量化误差。反量化则是在处理或存储之后，试图重建原始信号的过程。

多数的量化算法可以理解为就是为了找到更好的 S 和 Z 使得量化模型的结果尽可能逼近原模型的结果.

目前关于 LLMs 模型的量化算法可以理解为仅关心Linear算子的量化, 对于其他部件(如: norm, embedding, softmax, add...)都没有进行研究. 所以, 下面的所有讨论均以 Linear 为核心研究对象, 所有示例结果将以 LLaMA 模型进行展示。

LLMs 模型的推理可以大致分为两个 stage: context and generation. 两个 stage 有其各自的鲜明的特点,

• 在 context 阶段走的是 causal attention, 其行为可以类比训练的前向过程;
• 而 generations 阶段 sequence length 恒等于 1 这就要求推理框架需要支持两套计算逻辑(在 FasterTransformer 中可以看出)以适配其不同的特点. 在多数情况下, context 阶段是 compute bound(这不一定, 需要 seqlen 大于计算强度), 而 generation 是 IO bound. 很多情况下, generation 较 context 在应用中出现频率更高, 而量化模型由于其低位宽的权重表征, 可以大大缓解 IO bound 现象. (当然如果在服务时使得 batch 化技术来加大一次推理的 batch 的话, 量化的效果可能会退化为节约模型存储(功耗)下降).

IO 操作：

1. 从文件中读取数据
2. 向文件中写入数据
3. 从网络上下载数据
4. 向网络上发送数据
5. 与数据库交互读取或写入数据
6. 与外部设备进行通信，如打印机、扫描仪等

3 量化方法

3.1 LLM.int8()

由于观察到 input 的 outliers 只会固定在几个特定的 hidden-dim 的特点(这点我在 LLaMA 模型中也观察到该现象, 且随着模型加深越发严重. 具体原因个人分析为 RMSNorm 引起), 且 outliers 占据的 dims 很少(不到 1%). 故提出将 Linear 拆成两部分, 一部分为 int8, 一部分为 fp16, 分别计算后相加. 该方法得到广泛的应用, 有两个方面, 一个是因为被 huggingface 集成, 另一个是因为其几乎不掉点. 但该方法的缺点也是比较明显: 模型量化仅到 8bit, 仍是 4bit 的 2 倍大; Linear 的 latency 大幅上升, 原因在于它拆成两个 matmul kernel, 而且后续为了 fp16 相加引入外积操作等, 即计算流程更为复杂多步

3.2 Zero Quant 系列（v1,v2)

首次提次对采用 input token-wise quantization 并结合 weight group-wise quantization; 另外设计 LKD(Layerwise Knowledge Distillation, 使用随机生成的数据); 同时, 还做了一些 kernel fused 的工作, 实现了一个适配于 int8 的 backend. 这系列的工作都比较像 technical report, 且适用的模型尺寸比较小, 均在 20B 以下. 方法的 scaling 效果较差, 建议 follow 其量化粒度的设计.

3.3 SmoothQuant

同样是为了解决 input outlier 的问题, 韩松团队提供将 input 的动态范围除上 scale(该 scale > 1 即可以实现动态范围减小, 从而改善量化结果), 并将该 scale 吸到下一层的 weight 内, 利用 weight 的细粒度量化来承担该量化困难(因为 input 往往使用 token-wise quantization, 而 weight 通常使用 channel-wise quantization 或 group-wise quantization). 相较于 LLM.int8(), 由于 input 和 weight 全都是 int8, 并不会出现复杂的计算逻辑, 可以调用 CUTLASS 默认实现的 int8 gemm 来加速. 其缺点为: 精度没有 LLM.int8()有保证, 且容易受到 calibration-set 的影响), 同时一旦 weight 精度调至 4bit, 则模型精度下滑严重)

3.4 GPTQ

经典之作, 目前几乎是 4bit/3bit 方案的默认首选, 但也仅限于开源世界的娱乐可用, 离落地认定的靠谱精度还是有比较大的距离. 该方法来源于同一团队在 nips22 的工作(Optimal Brain Compression)延伸, 其同样将方法泛化到剪枝领域(也是大模型剪枝领域的经典, SparseGPT). 该方法的思路大致为: 利用 hessian 信息作为准则判定每个权重量化后对输出 loss(通常定义为 MSE)造成的影响, 量化影响最大的权重(即最敏感)挑选出来先进行量化, 然后对其他权重进行更新来补偿该权重量化导致的影响, 如此往复, 直至全部量化结果. 当然, 在 GPTQ 中作了一些简化, 比如是基于列元素进行量化循环, 来减少算法的运行时间. 该方法的优点: 首次将 4bit/3bit 权重量化在 176B 的模型上做 work, 同时也提出对应的 kernel(但比较糙, 优化空间大, 有不少团队做了优化). 缺点: 4bit/3bit 的方案原始 kernel 由于有 unpack 操作, 导致 gemv 操作的计算时间低于 fp16), 且精度距离落地有明显距离. 注: 从它开始, 很多人只开始研究 4w16f 的方案(即 weight-only quantization), 因为在 batch=1 的 gemv 计算中, 只需要控制权重的读入时间即可, 且 input 的动态范围过大, 量化掉点过大.

3.5 AWQ, Activation-aware Weight Quantization

2023 - Lin 等 - AWQ activation-aware weight quantization for LLM compression and acceleration

SmoothQuant 的续作, 从源代码来看, 它对 SmoothQuant 中计算 scale 时需要的超参 alpha, 增加了一步通过 grid search 得到每个 scale 的最优参数, 但论文的故事包装得很好, 同时取得的效果也是十分显著的, 符合大道至简的准则. 该方案是也是 4-bit weight-only quantization, 其 kernel 实现凭借对 PTX 的深刻理解和应用, 取得了目前这些 weight-only quantization 的方案的第一. 在此基础上稍加优化即可以得到一个不错的 baseline.

3.6 SqueezeLLM

| https://arxiv.org/abs/2306.07629

| 2023 - Kim 等 - SqueezeLLM Dense-and-Sparse Quantization

通过观察到部分权重决定了最终模型的量化性能, 提出以非均匀量化的方式缩小这些敏感权重的量化误差. 即通过 loss 的二阶 hessian 信息来确定量化敏感的权重, 将量化点安置在这些敏感权重附近, 其它点以 MSE 最小来安置. 该方法以少量的存储空间换来了目前最优的 4-bit weight 精度, 但其缺点也是极其明显: 由于采用 LUT 来实现非均匀量化, 导致其 kernel 在 batch > 1(文中的 batch 我均定义为 batch * seqlen)的情况下, Linear 的执行速度急剧下滑。参考资料：

3.7 QLoRA

这里顺带简单介绍一下 QLoRA. 该方法提出 4-bit NormalFloat, 一种新的数制(属于非均匀量化), 从理论角度上证明是 4bit 最优数制。利用该方法量化模型的 backbone 得到 4-bit 的 backbone, 然后基于 lora 进行 SFT, 在只需要 4-bit 模型权重的情况下完成 SFT, 从而使得许多人可以实现在单张消费级卡(i.e. 3080)上玩 LLaMA。但当时我跑它的时候, 其缺点就是明显的 kernel 速度慢, 原因同样是因为它需要通过 LUT 来实现, 不知道现在情况怎么样了.

4 总结

4-bit weight-only quantization 是一个相对比较均衡的方案。在这个 setting 下, 量化的研究工作应更多集中在模型的精度提升的层面上, 尽可能地减少对模型智能的影响. 但对于如果想进一步得到更轻更快更强的模型, 可以从其他小型化策略入手. 在这些策略中, 个人认为蒸馏是一个最值得往前走的方案. 在 LLaMA-2 的 tecnical report 中就有多处地方使用了蒸馏, 比如: 在 RLHF 阶段仅用 70B 的 reject sampling dataset 来 fine-tuning 其他几个小尺寸的模型, 以及很多人都会尝试去用 GPT4 的 SFT 数据来 fine-tuning 自己的模型. 剪枝不太推荐, 因为至少从 SparseGPT 的复现结果来看, 除了非结构化剪枝精度还算有保证外, 其余方案精度下滑明显, 包括 NV 的 2:4 和 4:8 方案, 距离落地还有些距离, 且和量化结合后并不能进一步拿到 50%的压缩收益。最后, 再提几点我认为有可能可以做的方向吧:

• 更加系统全面地推理优化,包括: 更深度更大粒度的 kernel-fusion, 其他部件优化(i.e. long context 下 kv-cache 的存储和 IO 时间, attention 计算优化), system2 的推理路径的优化
• 在模型训练中引入量化友好的策略, 来使得模型的权重和激活可以变得对量化不敏感, 从而实现 4w4f
• 尝试引入 QAT 方案, 达到所见即所得, 拥抱极限 -- 但这个有点太激进, 还是需要在模型有足够理解后去尝试.
• 端云推理的协同优化, 即手机端和 GPU 之间如何交互, 利用手机端训个人 SFT, 分配算力等