CANN-ATB量化推理-昇腾NPU上W8A8量化为什么比W4A16更实用
Llama2-70B 权重 140GB,8 卡 TP 刚好放得下但没什么余量给 KV Cache。W8A8 量化把权重从 fp16 压到 int8,权重体积减半,4 卡就能跑 70B。W4A16 理论上压得更狠(4 倍压缩),但精度损失在实际业务里往往不可接受。ATB 的量化推理模块支持两种方案,这篇帮你选。
量化方案对比
| 方案 | 权重精度 | 激活精度 | 压缩比 | 精度损失 |
|---|---|---|---|---|
| W16A16 | fp16 | fp16 | 1× | 无(baseline) |
| W8A8 | int8 | int8 | 2× | 0.1-0.5% |
| W4A16 | int4 | fp16 | 4× | 1-5% |
| W8A16 | int8 | fp16 | 2× | 0.05-0.2% |
W4A16 压缩比最高但精度损失大。W8A8 精度好但需要激活也做 int8 量化,实现复杂度更高。W8A16 是最保守的选择——权重 int8 但激活保持 fp16。
ATB 的 W8A8 实现
fromatbimportLLM model=LLM("meta-llama/Llama-2-70b-hf",device="npu:0,1,2,3",tensor_parallel_size=4,quantize="w8a8"# W8A8 量化)ATB 的 W8A8 量化流程:
1. 权重离线量化:fp16 权重 → int8 权重 + scale 2. 运行时反量化:int8 权重 × scale → fp16 → MatMul 3. 激活量化:fp16 激活 → int8 激活 + scale 4. int8 MatMul:int8 权重 × int8 激活 → int32 结果 5. 反量化:int32 结果 × weight_scale × act_scale → fp16步骤 2 的反量化走 DMA 引擎,零计算开销。步骤 4 的 int8 MatMul 在昇腾NPU的 Cube 单元上有原生支持——int8 GEMM 的算力是 fp16 GEMM 的 2 倍(因为 Cube 单元的 int8 吞吐量是 fp16 的 2 倍)。
ATB 的 W4A16 实现
model=LLM("meta-llama/Llama-2-70b-hf",device="npu:0,1",tensor_parallel_size=2,quantize="w4a16"# W4A16 量化)W4A16 的权重是 4bit,但 MatMul 不支持 4bit 输入。所以需要先反量化到 fp16 再算:
1. 权重离线量化:fp16 → int4 + scale + zero_point 2. 运行时反量化:int4 → fp16(每个权重需要查表+缩放) 3. fp16 MatMul:反量化后的 fp16 权重 × fp16 激活步骤 2 是瓶颈。4bit 反量化需要逐组(group_size=128)做查表和缩放,走 Vector 单元。反量化时间约占 MatMul 总时间的 20-30%。4bit 省了权重存储和 HBM 读取,但反量化的额外计算把一部分收益吃回去了。
实测精度
Llama2-70B 在 MMLU 和 GSM8K 上的精度对比:
| 方案 | MMLU (5-shot) | GSM8K (8-shot) |
|---|---|---|
| W16A16 (baseline) | 69.8% | 57.4% |
| W8A16 | 69.5% | 57.0% |
| W8A8 | 69.0% | 56.2% |
| W4A16 (GPTQ) | 67.2% | 51.8% |
W4A16 在 GSM8K 上掉了 5.6 个点——数学推理对精度特别敏感。W8A8 只掉了 1.2 个点,大部分业务可以接受。
性能数据
Atlas 800I A2 × 4,Llama2-70B:
| 方案 | decode 速度 (tokens/s) | 单卡显存 (GB) | 最大并发 |
|---|---|---|---|
| W16A16 (TP=8) | 2,400 | 35 | 16 |
| W8A8 (TP=4) | 2,800 | 42 | 12 |
| W8A16 (TP=4) | 2,600 | 42 | 12 |
| W4A16 (TP=2) | 1,800 | 38 | 8 |
W8A8 的 decode 速度反而比 W16A16 快——int8 GEMM 的算力翻倍,在 compute-bound 的 prefill 阶段收益最大。W4A16 看似省显存但速度最慢,因为反量化开销。
建议
- 显存够用:不量化,fp16 跑
- 显存紧但精度不能降:W8A16,最安全的量化
- 想省卡 + 接受微小精度损失:W8A8,2 倍压缩 + 2 倍 GEMM 算力
- 极端省显存场景:W4A16,但要对精度做充分评估
量化不是免费的午餐。W8A8 是当前昇腾NPU上性价比最高的方案——精度损失小、int8 GEMM 有硬件加速。W4A16 除非你真的放不下模型,否则不建议。仓库在这里:
https://atomgit.com/cann/ATB
