大模型推理优化技术之显存优化Permalink

大模型推理可能会遇到的性能问题Permalink

推理大模型时充分榨干 GPU 的性能是每个程序员所追求的。不过，要做到这一点，我们首先需要知道大模型推理的具体步骤，并分析它的性能瓶颈是什么，是受到算力限制还是内存限制，从而方便我们下一步的优化。

计算给定 GPU 上每个字节可能的操作数，并将其与模型注意力层的算术强度进行比较，可以揭示瓶颈所在：计算或内存。我们可以利用这些信息来选择合适的 GPU 进行模型推理，如果我们的用例允许，还可以使用批处理等技术来更好地利用我们的 GPU 资源。

KV CachePermalink

这是大模型推理性能优化的最常用技术。该技术可以在不影响计算精度的前提下，以空间换时间，提高推理性能。目前业界主流 LLM 推理框架均默认支持并开启了该功能。在咱们之前介绍 huggingface llama 实现的博客中有提到过，当 use_cache = True 时，KV cache 功能就默认打开。那么什么是 KV cache 呢？它又是如何加速大模型推理性能的呢？

原理Permalink

首先，回顾一下 transformer 中计算注意力权重的过程：

decoder 的输入在进入译码层后，就会分别经过 query key value 三个线性层，变为 Q K V 矩阵，然后进入 Masked multi head attention 做下图的运算。图中的 mask 是为了防止当前生成的 token 根据未来 token 的信息产生。

注意到，Transformer 模型具有自回归推理的特点，即每次推理都会将上次推理输出的结果作为输入的一部分，而每次推理会输出下一个 token，反复执行多次后得到最终的输出。因此，推理时，Transformer 的前后两轮的输出只相差一个 token，存在重复计算！

比如，图上例子中，若不使用 KV cache 技术，在推理生成“学”和“生”字时，前三行的 QK 矩阵和注意力矩阵的值都是一样的。以第一行为例，注意力矩阵的向量值取决于 Q 矩阵的第一行（K V 咱不关心）。每次计算得到的注意力矩阵向量，只有最后一行是全新的，用于产生新的 output。

因此，KV Cache 的办法就是将之前已经计算过的张量保存下来，下次推理时直接计算一行 Q 矩阵就行了，前面的数据可以直接使用之前保存的张量，避免重复计算。

代码Permalink

KV cache 技术已经应用在 huggingface transformer 库中：代码如下：

query = self._split_heads(query, self.num_heads, self.head_dim)
key = self._split_heads(key, self.num_heads, self.head_dim)
value = self._split_heads(value, self.num_heads, self.head_dim)

if layer_past is not None:
  past_key, past_value = layer_past
  key = torch.cat((past_key, key), dim=-2)
  value = torch.cat((past_value, value), dim=-2)

if use_cache is True:
  present = (key, value)
else:
  present = None

if self.reorder_and_upcast_attn:
  attn_output, attn_weights = self._upcast_and_reordered_attn(query, key, value, attention_mask, head_mask)
else:
  attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)

准备好 query key 和 value 矩阵后，当 layer_past 非空（事实上第一个词生成后就是非空的），那么直接将 past_key 和 past_value 取出，然后concat连接起来组成这次推理的 key 和 value 值。

KV cache 可在内存空间充足的场景下使用，能有效地提升推理性能。但需要注意的是，如果生成的序列非常长（大于 2048 ）那么就很有可能出现 out of Memory。