1. 背景介绍
在2023年7月，我们通过静态设计方案成功将ChatGLM2-6B模型部署到单颗BM1684X芯片上，采用F16量化模式，模型大小为12GB，平均推理速度约为3 token/s。详情可参考《算丰技术揭秘｜探索ChatGLM2-6B模型与TPU部署》。为了进一步提升模型的推理效率并降低存储空间，我们后续对模型进行了INT8量化部署，整体性能提升70%以上，模型大小降至6.4GB，推理速度达到6.67 token/s。
2. 量化方案
首先，TPU-MLIR原有的INT8量化方案并不适合直接应用于LLM。主要原因有两点：
- 无论是PTQ校准还是QAT训练，对于LLM来说成本都过高。例如，对LLM进行一轮PTQ校准可能需要1-2天时间。
- 量化带来的误差在LLM上无法收敛，最终会导致模型精度大量损失。
因此，我们沿用了ChatGLM2使用的W8A16策略，即只对GLMBlock中Linear Layer的权重进行per-channel量化存储，在实际运算时将其反量化回F16进行计算。由于LLM中Linear Layer的权重数值之间差异非常小，对INT8量化较为友好，量化后的结果与F16计算结果在余弦相似度上仍能保持99%以上，精度几乎可以做到0损失。
W8A16 MatMul
小提示：W8A16的核心思想是“用空间换时间”——权重存储用INT8节省一半空间，计算时再转回F16，既保持了精度又减少了内存占用。
3. TPU-MLIR实现
在TPU-MLIR编译器的Top到Tpu层的lowering阶段，编译器会自动搜寻模型中右矩阵输入为权重、且该矩阵维数为2的MatMul算子，并将其替换为W8A16MatMul算子。这样做的目的是为了与左右矩阵均为Activation的MatMul算子（如mm、bmm）区分开——在编译器中，mm、bmm以及Linear Layer都会被统一转换为MatMul算子。
以ChatGLM2中的一个MatMul算子为例：
- 原始：L = (max_length x 4096 x f16)，R = (4096 x 27392 x f16)
- 量化后：权重由原来的214MB降为107MB
- 额外产生的Scale (4096 x f16)仅占0.008MB的存储空间
可以看出，权重量化后存储空间基本可以达到减半的效果。算子替换源码与权重量化源码可在TPU-MLIR仓库中查看。
Op Replacement in TPU-MLIR
常见问题：为什么编译器要区分“右矩阵为权重”和“右矩阵不为权重”？
答案：因为只有右矩阵是权重（即固定参数）时，才能提前进行INT8量化。如果右矩阵也是动态激活值，则需要动态量化，实现复杂且精度损失更大。
4. 后端性能提升原理
前一节介绍的量化只实现了存储空间减半的效果，而性能提升主要依赖于W8A16MatMul后端算子的实现。如果您对TPU架构不熟悉，建议通过两期视频了解：TPU原理介绍(1)和TPU原理介绍(2)（可关注b站“算能开发者”进行观看）。
按照算能当前的TPU架构，W8A16的计算过程主要分为以下5个步骤：
- 从Global Memory中加载数据到Local Memory
- 将INT8权重Cast为F16
- 与Scale数据相乘完成反量化操作
- 与Input Activation进行矩阵乘运算
- 将计算结果存储回Global Memory
W8A16Matmul Computation on TPU
由于Local Memory空间有限，对于大型数据通常需要进行切分，分批对数据进行加载、运算与存储。为了提升效率，我们通常利用GDMA与BDC指令并行执行，同时进行数据搬运与运算操作。因此，Local Memory大致需要被划分为两部分区域：同一个循环内，一个区域用于数据运算，另一个区域存储上一循环计算好的结果以及加载下一循环需要用到的数据，如下图所示。
Local Memory Partition
矩阵乘等式如下：
当矩阵乘运算中左矩阵数据量较小时，性能瓶颈主要在于右矩阵的数据加载上——即数据加载时间远大于数据运算时间。W8A16通过量化将右矩阵的数据搬运总量缩小为原来的一半，额外多出的Cast与Scale运算时间可以被数据搬运时间覆盖，因此并不会影响整体runtime，如下图所示。
GDMA and BDC parallel
总而言之，从后端角度来说，当左矩阵的维度M越小、右矩阵的维度K越大时，W8A16带来的性能提升收益越大。
从LLM的角度来看，以ChatGLM2为例，一次推理的完整流程分为一轮prefill与多轮decode：
- 在prefill阶段，基于我们当前的静态设计方案，输入词向量会被补位为当前模型所支持的最大文本长度max_length（例如512、1024、2048等）。
- 在decode阶段，则固定只取前一轮生成的一个token作为输入。
ChatGLM2 Inference
因此，max_length越长，GLMBlock接收的输入数据量越大，Linear Layer的M也就越大，这会导致W8A16的性能提升越有限。而decode阶段，M始终保持为1，此时W8A16就能带来明显的性能提升。
MatMuls in ChatGLM2 prefill and decode phase
常见问题：为什么prefill阶段性能提升不如decode阶段明显？
答案：prefill阶段输入序列长度较大（M较大），此时数据加载时间相对于运算时间占比降低，W8A16节省的一半搬运时间带来的收益相对变小；而decode阶段M=1，数据搬运成为主要瓶颈，W8A16的减半效果非常显著。
5. 效果展示
将W8A16量化应用于ChatGLM2-6B后，整体性能如下：
- 性能：整体性能得到70%以上的提升
- 精度：与F16下的回答略有不同，但答案正确性仍然可以保证
- 模型大小：由12GB降为6.4GB
Result Comparison
小提示：在实际部署时，建议根据业务场景对精度和速度的容忍度进行权衡。如果对精度要求极高，可保留F16；如果对速度或存储有更高要求，W8A16是极佳选择。