【论文阅读】S3: Increasing GPU Utilization during Generative Inference for Higher Throughput

原文链接：
https://arxiv.org/pdf/2306.06000.pdf

Abstract

名字：scheduling sequences with speculation

• 除了已经很大的模型参数之外，保存序列中先前标记信息的键/值 (KV) 缓存可能会变得比模型本身还要大。
• 它为KV缓存保留了内存的最大序列长度，以保证在不知道输出序列长度的情况下生成完整的序列。这限制了我们使用较小的批量大小，从而导致 GPU 利用率较低，最重要的是吞吐量较低。
• 设计一个系统预测输出序列的长度根据长度生成query

1 Intruduction

• 内存容量和带宽。凸显了内存限制以及高效内存利用以提高 GPU 计算资源利用率的需求。、
• 提高 GPU 利用率和吞吐量的常见方法是增加批处理大小：批次内输入共享模型权重
• 因此，GPU 只需要将模型权重从其高带宽内存 (HBM) 加载到片上 SRAM 一次，然后将其重新用于批次内的所有输入——通常用于服务卷积和全连接神经网络
• KV cache：
  • the self-attention layer in Transformer-based text generation LLMs presents a challenge to this simple optimization due to its autoregressive nature. Specifically, when generating a new token in a sequence, the model needs to attend to all previous tokens in the sequence, requiring the model to retain all information from previous tokens and store them in HBM. We call this region in the HBM holding the information key/value cache (KV cache).
• KV 缓存的大小随着批处理大小和序列的增加而增长，这限制了最大批处理大小，从而降低了 GPU 利用率，最终降低了吞吐量。

我们提出了 S3，通过推测来调度序列，这是一个通过预测输出序列长度和减少内存浪费来最大化吞吐量的框架。

• S3 不会影响模型的困惑度，因为它不会改变模型的架构

本文贡献：

• 我们增加了较长延迟 SLO 下可实现的批量大小，并允许服务提供商通过使用更大的批量大小在在线和离线场景中提供更高的吞吐量。
• 我们微调 Distillbert 模型来预测给定输入提示的输出序列长度
• 我们提供了一种从错误预测中恢复的机制。 S3 抢占超出其分配内存的序列并重新训练预测器以从错误中学习

2 Background

• 基于 Transformer 的生成模型是自回归的

• 由于模型一次生成一个token，因此它必须迭代自身 n 次才能生成 n 个 token长的序列。

• 一次迭代涉及一个输入token遍历模型，该模型是一堆transformer layer:

  • 其中包含一个attention layer、两层layernorm和两个feedforward layer。
  • the self-attention layer uses information on the past tokens to generate the next token.

• 第($i^{th}$) 次迭代的模型使用当前标记 ($t_i$) 及其在自注意力层中已生成的每个标记 ($t_0\dots \dots t_{i-1}$)。我们可以将自注意力层表示为:
  $$h_{out}=softmax(\frac{q_i{K}^T}{\sqrt{d_h}})\ast V$$

• *where $d_h$ is **hidden dimension of the model,$h_{out}$***当前隐藏输出向量，$q^i$是当前查询向量

• K和V矩阵中的第j行代表$t_j$的键和值向量，

• 该模型将 K 和 V 矩阵存储为键/值 (KV) 缓存，以避免在每次迭代时生成键和值向量

• 否则，它必须存储每个先前标记的隐藏状态，并在每个转换器层将其与权重矩阵 $W_{QKV}$∈ $R^{d_h\ast 2d_h}$相乘。

• 这将需要每层 2(i − 1)dh FLOPs，几乎与 ti 处自注意力层的 2idh FLOPs 相同。

• 使用半精度数字时，KV 缓存的大小为每个token $4l{d}_h$字节

• 缓存为键和值缓存的每个数字使用 2 个字节，其中 l 是模型中transformer层的数量

  • a 20 billion parameter model, has 44 layers and 6144 hidden dimensions and thus uses 1MB per KV cache per token

  • 2.2 KV Cache Management on GPUs

• KV 缓存相对较小（例如几个 MB），并且当序列较短时可以轻松存储在 GPU HBM（高带宽内存）中，因为缓存存储有关序列中先前token的信息

• huggingface 这种方法允许库为每个缓存分配精确的内存量，但代价是频繁的内存访问

• NVIDIA的FasterTransformer库为每个序列保留了内存的最大序列长度

• 它通过简单地以仅**附加方式填充保留内存来删除冗余内存访问。**然而，这种方法有其自身的缺点，因为它为序列保留了超出严格必要的内存

2.3 Observation

• 语言模型受内存限制
• batch size越大并且sequence length越小，GPU的利用率越低

3 S3 Design

• 主要构成是三部分：
  • scheduler
  • predictor
  • supervisor
• 文本生成查询到达主机 DRAM 中的请求池
• 然后预测器预测调度程序用于批处理请求的输出序列长度
• 调度程序将批处理分派到 GPU，文本生成器模型生成文本。
• 主管监督 GPU 利用率并处理错误预测。
• Output sequence length predictor：
  • 将序列长度进行分桶
    $$\frac{maxsequencelength}{numberofbuckets}$$
  • Length-aware sequence scheduler:
  • 调度器根据预测结果对序列进行批处理和调度，以在不超过 GPU HBM 容量的情况下最大化 GPU 利用率
  • 自注意力层的输入不共享权重不会从batching中获益
  • 如果模型无法容纳在一个 GPU 中，S3 会使用管道并行性并以 Transformer 层粒度对模型进行分片。
• Supervisor
  • 主管负责监督 GPU 利用率并处理错误预测。
  • the supervisor doubles the assigned memory for the evicted sequences to fix the short misprediction.
• 文本生成请求到达请求池。预测器预测池中序列的输出序列长度。Supervisor在后台运行并检查当前的HBM使用情况，接下来，调度程序使用预测和可用的 HBM 来批量请求以获得最大 GPU 利用率。它通过将该批次调度到生成调度序列的 GPU 来完成其工作。

4 Evaluation

• 前馈中的输入共享相同的权重，而自注意力中的输入则关注自己的序列。
• 并行部分：
  • GPT-3一共有96层：
  • We partition GPT-3 into 6 and 8 GPUs in a pipelineparallel manner, allocating 16 and 12 transformer layers per GPU for each configuration, respectively.
  • We also evaluate on 10 GPU setting where we allocate 10 layers to 9 GPUs and 6 to the remaining GPU
• More GPUs shard the model into finer pieces and leave more space for storing the KV cache, allowing us to increase the batch size.
• 面向吞吐量的系统使用内存层次结构将部分模型存储在较慢的内存中并增加批量大小,然而，它们都为每个序列分配相同的内存，并且不考虑基于预测抢占序列。
• 减少KV 缓存开销:
  • The issue of attention layers in Transformers requiring quadratic computation（二次计算）and memory with respect to the sequence length（记忆序列长度）：
    • 低阶近似
    • 利用稀疏性
    • multi-query attention：每个键值 (KV) 缓存使用一个注意力头而不是采用多个注意力头
    • 模型压缩和量化