Li L, Qian S, Lu J, et al. Transformer-Lite: High-efficiency Deployment of Large Language Models on Mobile Phone GPUs. arxiv preprint arxiv:2403.20041, 2024.

无代码，技术报告

⭐️⭐️⭐️

http://arxiv.org/abs/2403.20041

1. 背景 & 动机

这篇论文是OPPO AI Center发表的，其提出了Transformer-Lite框架来缓解移动设备GPU上部署大型语言模型（LLM）时存在的性能问题。大型语言模型（如ChatGLM2 6B和Gemma 2B）被广泛应用于智能助手、文本摘要、翻译和多模态任务等，但它们在移动设备上的部署面临着几个关键问题：

  1. 计算能力和内存带宽的需求：LLMs需要大量的计算能力和内存带宽，这些资源在移动设备上是有限的。
  2. 用户体验：当前的在端侧设备上部署LLM的方法通常有着较慢的推理速度，这会严重影响用户体验。
  3. 成本：在云上大规模部署LLM有着可观的成本，而且随着移动设备性能的不断提升，本地部署LLM不仅可以减少与云部署相关的高成本，还可以扩大LLM在移动设备上的应用前景以及保护用户隐私。

为了解决上述问题。这篇文章文章提出了以下优化技术：

  • 符号表达式的方法：支持动态形状模型推理，包括动态形状推导、内存重用和执行调度等。
  • 算子优化和执行优先级设置：加快推理速度并减少手机延迟。
  • FP4量化方法：称为M0E4，减少dequantize开销，使得矩阵乘法更高效。
  • 基于子张量的技巧：避免了在LLM推理后复制KV缓存的内存压力。

作者团队还实现了一个名为Transformer-Lite的移动推理引擎，该引擎与高通和联发科处理器兼容，并通过一系列实验评估了其性能。通过这些技术，Transformer-Lite引擎在prefill和decoding方面相比基于CPU的FastLLM和基于GPU的MLC-LLM取得了显著的速度提升。

2. 方法

2.1 对于动态形状推理的符号表达方法

不同于多数静态尺寸的CV模型，LLM 是一种动态形状输入的场景，每次迭代的输入形状都会发生变化。这导致模型中一些激活张量的形状发生变化，给内存重用和算子性能优化带来了巨大挑战。 例如，将形状[“sumN-N”,1,2,128]与axis 0上的[“N”,1,2,128]连接，输出形状为[“sumN”,1,2,128]。为此，作者团队利用类似于 Nimble 和 DISC 的方法，将深度学习模型中的算子分为两类：张量计算算子和形状计算算子。后一类算子包含形状算子和那些其输入取决于形状算子结果的算子。形状计算算子在 CPU 上执行，计算形状信息。相反，负责计算激活的张量和权重的张量的计算算子则在 GPU 上执行。

from http://arxiv.org/abs/2403.20041
Fig 1. 形状推理

from http://arxiv.org/abs/2403.20041

作者团队使用了SymPy包且结合了ONNX来实现了这一能力。

对于内存复用，实现内存复用的先决条件是获得张量之间的内存大小关系，这在静态形状推理中很简单。符号表达式便于轻松确定这种关系。首先，计算每个张量的内存大小，即元素数乘以数据类型的字节数，从而为每个张量的内存大小创建一个符号表达式。**然后结合减法和除法来辨别张量之间的内存大小关系。**例如，用 “N ”*4096 除以 “N ”32128，得到的整数为 1，表示大小相等，可以重复使用内存。另一方面，当 “N” * 4096 与 “sumN” * 2 * 128 相比较时，减法和除法都会产生另一个符号表达式，而不是一个确定的整数。因此，除非探索所有潜在的输入形状，否则它们之间的大小关系是不确定的，这意味着无法进行内存重用。 作者团队使用了OpenCL工具来实现了内存复用。

通常，LLM的每次推理迭代都需要更新输入形状。在形状更新过程中，需要重新计算动态张量的形状，并更新具有动态形状输入或输出张量的算子参数。因此，频繁的输入形状更新通常会导致不可忽略的性能开销。为了避免这个问题，作者团队采用了Mask机制，在解码阶段将模型输入序列长度填充为64或128的倍数。这种方法可以在实际计数超过填充长度时更新形状，从而使每次推理迭代的形状更新时间开销可以忽略不计。

2.2 算子优化和执行优先级设置

作者实现了半精度和4bit量化的混合精度矩阵乘法。此外，prefill和decoding阶段对矩阵乘法的形状是不同的。在prefill阶段，矩阵乘法进行矩阵-矩阵计算，而在decoding阶段，这些运算符被简化为向量-矩阵计算。作者对这两种方式的乘法都做了优化，文中没有涉及到算子优化的细节。

GPU在执行LLM时还需要负责用户界面渲染，LLM的执行可能会造成界面渲染延迟。为了解决这个问题，作者利用高通和ARM提供的OpenCL扩展，将深度学习模型算子的执行优先级设置为最低级别。这种方法有效地缓解了移动设备的渲染延迟。

2.3 FP4量化方法：M0E4

在传统的量化方法中，如GPTQ和AWQ，模型权重从浮点数被量化为4位整数，以减少模型大小和内存带宽需求。然而，为了保持高精度，计算的时候通常使用半精度进行。在模型推理时，需要将量化的权重dequantize为浮点数以执行矩阵乘法，这个过程涉及到将整数转换为浮点数，会引起显著的性能开销。在将整数转换为浮点数的过程中需要使用复杂的算法或指令来实现。因此，dequantize会引发不可忽略的性能开销。

M0E4中的"M0"表示指数部分使用0位。这意味着指数部分被设置为一个常量值，通常是用于表示半精度浮点数的最小正指数。“E4"表示尾数部分使用4位。这些4位用于存储量化的尾数信息，这与浮点数表示中的尾数（或称为小数部分）相对应。

作者采用group-wise量化，并通过缩放和偏置系数将每个group的张量(表示为w0)转换为一个新的浮点张量(表示为w1)。并且确保所有w1元素都在$[2^n, 2^{n+1} - eps]$的范围内。$n$的取值为1，2。在这个范围内，就可以保证这些元素的指数部分是一致的。

为了提高量化的准确性，作者为舍入分离了一个额外的位(表示为y2，如图所示):y1 = min (y1+ y2,15)。这里，15是一个4位无符号整数所能表示的最大值。

from http://arxiv.org/abs/2403.20041
Fig 2. M0E4 量化

from http://arxiv.org/abs/2403.20041

在dequantize时候，用如下公式即可

$$ w1_{half} = as_{half} (\text{ConstExpBinPart} | (w1_{fp4} \ll \text{ConstBitShiftNum})) $$

2.4 基于subtensor来缓解KV Cache压力

一张图就可以说明，结合SWA应该可以工作的不错。

from http://arxiv.org/abs/2403.20041
Fig 3. KV Cache 优化

from http://arxiv.org/abs/2403.20041

发现这个做法不如vLLM种的Page Attention来的统一、简洁。

3. 实验

3.1 各个LLM上的性能

from http://arxiv.org/abs/2403.20041
Fig 4. 各个LLM上的性能

from http://arxiv.org/abs/2403.20041

3.2 和FastLLM和MLC-LLM比较

from http://arxiv.org/abs/2403.20041
Fig 5. 和MLC-LLM比较

from http://arxiv.org/abs/2403.20041

from http://arxiv.org/abs/2403.20041
Fig 6. 和FastLLM比较

from http://arxiv.org/abs/2403.20041