【本文正在参与 AI.x社区AIGC创新先锋者征文大赛】
作者 | Muhammad Saad Uddin
编译 | 岳扬
将 LLMs 投入生产环境使用，会遇到诸多挑战，尤其是想要为 LLMs 高效提供所需的计算资源时。有过此类经验的人可能深有体会，GPU 内存是支持 LLMs 运行的一项关键资源。由于这些模型体积巨大，且推理过程具有动态性质，因此对 GPU 内存使用的规划和优化提出了更高的要求。
Image by Author via DallE
出于以下几个原因，准确估算 GPU 内存的需求至关重要：
然而，计算 LLMs 服务所需的 GPU 内存并非一件简单的事。模型的大小（model size）、序列长度（sequence lengths）、批处理数量（batch sizes）以及解码算法（decoding algorithms）等多种因素，都会以复杂的方式影响内存使用。而且，传统的内存分配方法常因内存碎片和键值（KV）缓存等动态内存组件的低效管理而造成大量浪费。
在本文中，我将尽可能详细地解释如何计算 LLMs 服务所需的 GPU 内存。我将分析影响内存使用的各部分，并根据模型参数和工作负载特征（workload characteristics），逐步介绍如何估算内存占用大小。同时，我还会探讨 Paged Attention 和 vLLM 等先进的优化技术，这些技术能显著降低内存消耗并提升处理能力。通过阅读本文，你将能够全面了解如何规划和优化 LLMs 的 GPU 内存使用，从而在实际应用中实现高效且低成本的 LLMs 部署。
01 了解 LLM 推理过程中，主要消耗 GPU 内存的几部分
要想掌握 GPU 内存的计算方法，最关键的是了解各部分如何占用 GPU 内存的。了解内存的使用去向有助于我们更好地规划与优化资源。在 LLMs 推理过程中，主要消耗 GPU 内存的几部分包括权重（模型参数）、键值缓存内存（Key-Value Cache Memory）、激活值（Activations）与临时缓冲区（Temporary Buffers），以及系统开销（Overheads）（如果你对并行处理或分布式计算有所研究，可能对这个概念已有一定的认识）。
1.1 模型参数（权重）
模型参数是神经网络在训练过程中学到的数值（权重（weights）和偏置（biases））。这些参数定义了模型如何处理输入数据生成输出。
模型大小对 GPU 内存的影响
让我们看看这些模型：
请记住，对于 GPT-3 及其之后的模型，使用模型并行化（model parallelism）将模型参数分布到多个 GPU 上是十分必要的。
1.2 键值（KV）缓存内存
KV缓存存储生成序列中每个 token 所需的中间表示。简单来说，当模型每次生成一个 token 时，它需要记住之前的 tokens 以保持上下文。KV缓存存储了目前为止生成的每个 token 的键（key）和值（value）向量，使模型能够高效地处理过去的 tokens ，而无需重新计算。
工作原理：
序列长度（Sequence Length）和并发请求（Concurrent Requests）的影响：
计算每个 token 的 KV 缓存大小
让我们来分析一下如何得出每个 token 的 KV 缓存大小：
键向量（每层一个键向量）和值向量（每层一个值向量）
再次以 llama-13b 模型为例，假设模型具有以下特征：
i. 键向量：
ii. 值向量：
iii. 每个 token 的总KV缓存：
现在考虑输出内容为2000个 tokens 的情况：
KV缓存随着序列长度和并发请求数量的增加而线性增长。我从一篇论文[1]中了解到，KV缓存可以消耗多达 30% 甚至更多的GPU内存。
1.3 激活值和临时缓冲区
激活值（Activations）是指推理过程中神经网络层的输出，而临时缓冲区（temporary buffers）用于中间计算。激活值和缓冲区通常消耗的内存比模型权重和 KV 缓存要少。
尽管它们的容量较小，但激活值对于模型计算每一层的输出是十分必要的。它们在前向传递过程（forward pass）中被创建和丢弃，但仍需要足够的内存分配。
1.4 内存开销
额外的内存使用开销来自于内存分配和使用的低效率。下面是对其的简要介绍：
内存碎片：
计算过程中产生的中间步骤：
低效内存管理的影响：
示例：如果内存碎片在 40 GB GPU 上浪费了 20 %的内存，那么就有 8 GB 的内存本可以用来处理更多请求，但是现在被浪费了。
02 计算 GPU 内存需求
既然我们已经对关键内容有了足够的了解，那么就不再拖延，直接计算完整的 GPU 内存需求！
逐步计算：
要计算任何模型的内存需求，几乎以下内容都需要：了解模型权重、KV缓存、激活值和临时缓冲区以及系统内存开销。以 llama-2 13B 模型为例，公式为：
所需内存总量：模型权重 + KV缓存 + 激活值和系统内存开销
对于 13 B 模型来说：
模型权重 = 参数数量 × 每个参数的字节数
总 KV 缓存内存 = 每个 token 的 KV 缓存内存 × 输出序列长度 × 输出序列数量
激活值和系统内存开销 = GPU总内存的 5–10 %
模型权重 = 130 亿 × 2 字节 = 26 GB
总 KV 缓存内存 = 800 KB × 8192* tokens × 10* 并发请求= 66 GB
激活值和系统内存开销 = 0.1 × (26 GB + 66GB) = 9.2 GB
所需内存总量：26 GB + 66 GB + 9.2 GB = 101.2 GB
所以，运行 llama-2 7B 模型至少需要 3 个 A100 40GB GPU。
如果我想要托管一个 GPT-3 模型（我知道这很疯狂；D），计算方法与此类似，但这次我会假设每次只处理一个请求，并使用 OPT-175B[2] 模型的大小（ 96 层和每层 12288 维度）作为参考。
所需总内存：350 GB + 36 GB + 38.6 GB = 424.6 GB 几乎需要 11 个 A100