1.利用TPU-MLIR实现LLM INT8量化部署
利用TPU-MLIR实现LLM INT8量化部署
在年7月,源码我们已成功将静态设计应用于ChatGLM2-6B在BMX单芯片部署,编译采用F量化模式,源码模型大小为GB,编译cryptsetup 源码平均速度为3 token/s。源码为提升效率与降低存储需求,编译我们进一步对模型执行了INT8量化部署。源码
传统TPU-MLIR的编译INT8量化方案并不适合LLM。这主要是源码由于LLM中PTQ校准或QAT训练成本过高,一轮校准可能需1-2天,编译且量化误差导致模型精度大量损失。源码jqury2048源码基于此,编译我们沿用了ChatGLM2的源码W8A策略,对GLMBlock中Linear Layer权重进行per-channel INT8量化存储,编译运算时反量化至F,源码以确保精度损失几乎为零。vol指标公式源码
在编译器的Top至Tpu层lowering阶段,TPU-MLIR自动替换MatMul算子,将权重矩阵切分为W8AMatMul,以区分具有不同矩阵输入的算子。以ChatGLM2中某个MatMul算子为例,豆瓣app源码vue量化后权重从MB减至MB,额外的Scale使用了0.MB存储,实现近一半的存储空间节省。相关源码可在TPU-MLIR仓库查询。
性能提升主要源于W8AMatMul后端算子优化。公式源码线加粗TPU架构下,W8A的计算过程分为5步,通过GDMA与BDC指令并行执行数据搬运与运算,将Local Memory分为两部分,确保效率。当左矩阵数据量较小时,性能瓶颈在于右矩阵数据加载,W8A量化减少数据搬运总量,额外运算时间被覆盖,性能影响可忽略。
从LLM角度看,推理流程包括prefill与decode。prefill阶段输入词向量补位至最大文本长度,decode阶段固定取前一轮生成的token作为输入。因此,prefill阶段GLMBlock接收数据量大时,W8A性能提升有限,而decode阶段$L_{ row}$恒为1,能实现显著性能提升。
应用W8A量化后,ChatGLM2-6B整体性能得到优化。具体结果展示如下: