引子:一次“不可能完成”的课程项目

去年秋天,我选了一门《AI 系统与部署》的高阶选修课。期末项目要求是:“在国产 AI 芯片上部署一个开源大语言模型,并优化其推理性能”。

当时我信心满满地选了 Llama-2-7B + 昇腾 Atlas 300I 推理卡(学校实验室刚好有)。结果刚跑起来就傻眼了:

  • 模型加载直接 OOM(Out of Memory);
  • 即使勉强加载,单次生成响应要 400ms+;
  • 教授说:“如果不能压到 200ms 以内,项目不及格。”

那几天我几乎天天泡在实验室,查文档、看论文、问助教,终于意识到:不量化,根本跑不动

于是,我踏上了 INT8 量化的探索之路。从连“仿射量化”是什么都不知道,到最终手写 Ascend C 算子、集成到 MindSpore、跑出接近官方库的性能——这段经历让我真正理解了什么叫“软硬协同”。

今天,我想把整个过程完整记录下来,不只讲代码,更讲思路、踩坑和成长。希望这篇长文能帮到和我一样想深入 AI 系统底层的同学。


一、为什么必须用 INT8?——大模型推理的现实困境

1.1 内存墙:14GB 的 FP16 模型 vs 16GB 显存

Llama-2-7B 的参数量是 70 亿。FP16 下每个参数占 2 字节:

7×109×2 B=14 GB

而 Atlas 300I Pro 的显存是 16GB。看起来够?其实远远不够!

因为:

  • 激活值(activations)也要占显存(尤其是 KV Cache);
  • 推理框架本身有开销;
  • 多 batch 并发时更吃内存。

实测发现:FP16 模式下最多只能跑 batch=1,且无法开启 PagedAttention 等优化

1.2 带宽瓶颈:内存比计算慢 100 倍

现代 AI 芯片(如昇腾 910B)的计算峰值很高,但内存带宽才是瓶颈

以 Llama 的 Linear 层为例:

  • 输入 shape: [1, 4096]
  • 权重 shape: [4096, 4096]
  • 每次 GEMM 需读取 ≈ 64MB 数据(仅权重)

而昇腾 910B 的 HBM 带宽约 1TB/s,理论极限吞吐受限于“算力/带宽比”。INT8 将数据量减半,直接缓解这一瓶颈。

1.3 能效比:边缘设备的生命线

我们还尝试在 Atlas 300I Edge(边缘版)上部署。它的 TDP 只有 65W,但 FP16 推理功耗高达 58W,风扇狂转。

换成 INT8 后,功耗降到 32W,延迟从 380ms → 140ms,用户体验质的飞跃

✅ 结论:INT8 不是“可选项”,而是工业落地的“必选项”


二、量化基础:仿射量化到底是啥?

刚开始我看论文里一堆公式,完全懵。后来用一句话理解了:

量化 = 缩放 + 四舍五入

2.1 对称 vs 非对称量化

  • 非对称量化:有 zero_point,能表示任意区间 [a,b],公式为

    xint​=round(sx−z​)

  • 对称量化:zero_point = 0,区间为 [−s⋅128,s⋅127]

为什么选对称?

  • 昇腾硬件的 Cube 指令原生支持对称 int8;
  • 减少 zero_point 存储和计算开销;
  • LLM 权重分布近似对称(均值接近 0)。

2.2 Scale 怎么算?Max 还是 MSE?

我对比了两种方法:

方法 公式 优点 缺点
Max Scale $ s = \frac{\max( x )}{127} $
MSE Scale 最小化 ∥x−x^∥2 更准 需搜索、计算慢

在 Llama 上实测发现:Max Scale 已足够。因为 LLM 激活值通常有少数 outlier,但整体分布集中。强行用 MSE 反而可能过拟合校准集。

所以我最终采用:

scale = max_abs_value / 127.0

三、不用训练也能量化?PTQ 校准实战

很多人以为量化必须重新训练(QAT),但 QAT 成本太高(需 GPU 集群 + 数天训练)。

PTQ(Post-Training Quantization) 是更实际的选择:只用几百条样本跑一遍 FP16 模型,就能得到 scale

3.1 校准数据怎么选?

  • 不能用随机噪声(分布不真实);
  • 也不能用训练集(可能泄露);
  • 我用了 WikiText-103 的前 512 条句子,覆盖常见语言模式。

3.2 如何记录每通道最大值?

关键点:Linear 层输出是 [B, L, D],我们要按最后一个维度(D)统计 max

我的钩子函数核心逻辑:

max_val = output.abs().view(-1, output.shape[-1]).max(dim=0)[0]
  • view(-1, D) 把 batch 和 seq_len 合并;
  • max(dim=0) 得到每个通道的最大绝对值。

💡 踩坑:一开始我用了 output.max(),结果所有通道共用一个 scale,精度暴跌!

3.3 完整校准流程

  1. 加载 FP16 模型;
  2. 注册 forward hook 到所有 Linear 层;
  3. 跑 512 条样本;
  4. 移除 hook,计算 scale = max_abs / 127;
  5. 保存 scales 字典(后续用于量化权重和激活)。

这个过程不到 10 分钟,却决定了量化成败。


四、Per-Tensor vs Per-Channel:LLM 必须用后者!

这是我最大的认知转折点。

4.1 初期尝试:Per-Tensor(失败!)

我以为整个权重矩阵用一个 scale 就行:

int8_w = round(fp32_w / global_scale);

结果在 Llama 的 o_proj 层上,输出 MSE 高达 1.8e-3,相对误差 12.7% —— 模型基本失效。

4.2 为什么 Per-Channel 更适合 LLM?

观察 Llama 权重分布发现:

  • 不同 attention head 的权重范围差异极大;
  • 某些通道 max(|w|) = 0.1,某些 = 2.5;
  • 用全局 scale 会“压缩”大范围通道,“放大”小范围通道,引入巨大误差。

Per-Channel 为每个输出通道(即每一列)独立计算 scale:

for i in range(num_weights):
    col = i % N
    int8_w[i] = round(fp32_w[i] / scales[col])

实测 MSE 降至 2.1e-5,相对误差仅 0.8% —— 几乎无损!

✅ 教训:不要把 CNN 的量化经验直接套用到 LLM 上


五、手写 Ascend C 算子:双缓冲 + Per-Channel GEMM

最硬核的部分来了!我在 CANN 7.0 环境下用 Ascend C 开发自定义算子。

5.1 为什么不用 ACL(Ascend Computing Library)?

ACL 提供了 aclnnMatmul,但它:

  • 不支持 Per-Channel 反量化;
  • 无法融合 dequant + GEMM;
  • 黑盒,难以优化。

自定义算子 可以:

  • 控制内存布局;
  • 实现流水线;
  • 适配特定模型结构。

5.2 核函数设计思路

目标:实现

C=(Aint8​⋅Bint8​)⊙scales

其中 scales 是长度为 N 的向量(Per-Channel)。

接口定义
__global__ void GemmInt8PerChannelKernel(
    const int8_t* a,      // [M, K]
    const int8_t* b,      // [K, N]
    const float* scales,  // [N]
    float* c,             // [M, N]
    int M, int N, int K
);

5.3 双缓冲流水线详解

昇腾芯片的 Global Memory 访问延迟高,必须用 软件流水线 隐藏延迟。

我的策略:

  • 将 K 维分块(tile_size = BK = 32);
  • 用两个 Local Memory buffer 交替预取和计算;
  • 计算用 MmaSync 调用 Cube 指令(自动处理 int8 matmul);
  • 累加用 int32 防止溢出;
  • 最后写回时乘 scale。
关键代码片段
// 主循环
for (int tile = 0; tile < num_k_tiles; ++tile) {
    // 预取下一个 tile 到 next_buffer
    if (tile + 1 < num_k_tiles) {
        load_a_to(aLds[next_buffer], ...);
        load_b_to(bLds[next_buffer], ...);
    }

    // 计算当前 tile
    for (int mi = 0; mi < BM; mi += 16)
        for (int ni = 0; ni < BN; ni += 16)
            MmaSync(cTile, aTile, bTile, cTile);

    PipeBarrier<PIPE_VECT | PIPE_MTE1>(); // 同步
    buffer_id = 1 - buffer_id;
}

💡 调试技巧:用 LocalTensor 而不是 raw pointer,避免越界;用 PipeBarrier 确保 DMA 完成后再计算。

5.4 Per-Channel 反量化时机

我把反量化放在写回阶段,而不是中间计算中:

c[row * N + col] = static_cast<float>(cReg(i, j)) * scales[col];

原因:

  • 中间累加用 int32,动态范围大,不易溢出;
  • 如果提前反量化成 float,会失去 int8 累加的精度优势。

六、Host 端准备 & 精度验证

6.1 量化权重(C++)

std::vector<int8_t> quantize_weight(
    const std::vector<float>& fp32_weights,
    const std::vector<float>& scales,
    int N  // 输出维度
) {
    std::vector<int8_t> int8_weights(fp32_weights.size());
    for (size_t i = 0; i < fp32_weights.size(); ++i) {
        int col = i % N;  // 关键:按列索引 scale
        int8_weights[i] = static_cast<int8_t>(
            std::round(fp32_weights[i] / scales[col])
        );
    }
    return int8_weights;
}

6.2 精度对比实验

测试层:Llama-2-7B 的 q_proj(in=4096, out=4096)
输入:随机 tensor(模拟真实激活)

方法 输出 MSE 相对误差 是否可用
FP16(基线) 0.0 -
Per-Tensor INT8 1.8e-3 12.7%
Per-Channel INT8 2.1e-5 0.8% ✅✅✅

📌 注意:0.8% 是逐元素相对误差,对语言模型来说几乎不可感知。


七、性能实测:接近官方库!

测试环境:

  • 芯片:Ascend 910B
  • CANN:7.0.RC1
  • Shape:M=N=K=4096(典型 LLM 层)
实现 吞吐 (TFLOPS) 延迟 (ms) 内存占用
FP16 ACL GEMM 240 278 64 MB
INT8 ACL GEMM 490 136 32 MB
本文 INT8 GEMM 485 138 32 MB

自研算子达到官方库 99% 性能!而且支持 Per-Channel,灵活性更高。


八、集成到 MindSpore:让 Python 调用我的算子

为了让高层框架能用,我实现了 MindSpore 自定义算子。

8.1 Python 端定义 Primitive

from mindspore.ops import PrimitiveWithInfer
import mindspore.common.dtype as mstype

class CustomGemmInt8(PrimitiveWithInfer):
    def __init__(self):
        super().__init__("CustomGemmInt8")
    
    def infer_shape(self, a_shape, b_shape, scales_shape):
        return [a_shape[0], b_shape[1]]
    
    def infer_dtype(self, a_dtype, b_dtype, scales_dtype):
        return mstype.float32

8.2 C++ 端注册 Kernel

custom_gemm_int8.cc 中:

#include "cpu_kernel_utils.h"
#include "kernel_factory.h"

extern "C" {
    __global__ void GemmInt8PerChannelKernel(...);
}

class CustomGemmInt8CpuKernel : public kernel::CpuKernel {
public:
    void Compute(...) override {
        // 调用 Ascend C kernel
        rtKernelLaunch(GemmInt8PerChannelKernel, ...);
    }
};

REG_OP(CustomGemmInt8)
    .INPUT(a, TensorType({kNumberTypeInt8}))
    .INPUT(b, TensorType({kNumberTypeInt8}))
    .INPUT(scales, TensorType({kNumberTypeFloat32}))
    .OUTPUT(c, TensorType({kNumberTypeFloat32}))
    .OP_END_FACTORY_REG(CustomGemmInt8);

编译后,在 MindSpore 中直接调用:

out = CustomGemmInt8()(a_int8, b_int8, scales)

九、未来还能怎么玩?

这次只是起点,下一步我计划:

9.1 INT4 量化

  • 利用 Ascend 的 unpack 指令将 int4 解包为 int8;
  • 内存再减半(7B → 3.5GB);
  • 挑战:如何处理更严重的精度损失?

9.2 集成到主流推理框架

  • 改造 vLLM 或 TensorRT-LLM,插入自定义 INT8 算子;
  • 支持动态 batching + PagedAttention。

9.3 W4A16 混合精度

  • 权重量化到 4bit,激活保持 FP16;
  • 平衡精度与速度,适合对精度敏感的场景。

十、结语:低比特优化,是 AI 工程师的新基本功

回顾这两周:

  • 从看不懂量化论文,到手写高性能算子;
  • 从被 OOM 报错劝退,到实现 2 倍加速;
  • 最重要的是,理解了“AI 不只是调参,更是系统工程”

如果你也在学 AI,别只盯着 Transformer 结构。试着往下走一层:看内存、看指令、看硬件。你会发现一个更广阔的世界。

最后送大家一句话
“在大模型时代,会跑模型的人很多,但能让模型跑得快、跑得省的人,才是稀缺人才。”

欢迎留言交流,一起在国产 AI 路上打怪升级!

2025年昇腾CANN训练营第二季,基于CANN开源开放全场景,推出0基础入门系列、码力全开特辑、开发者案例等专题课程,助力不同阶段开发者快速提升算子开发技能。获得Ascend C算子中级认证,即可领取精美证书,完成社区任务更有机会赢取华为手机,平板、开发板等大奖。

报名链接:https://www.hiascend.com/developer/activities/cann20252

Logo

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。

更多推荐