【实践】VLLM显存暴增 | 多卡推理 | 批量推理

显存占用问题

用两张A800跑Llama3-8B，按理说显存占用应该在16G以上，浮动2~3G，但是发现显存直接拉满了，一共花了140G

发现通过vllm.LLM创建模型时，有个参数控制GPU的使用量gpu_memory_utilization，其默认值为0.9，正好是140/160
将gpu_memory_utilization改成0.15，也就是160*0.15 =21G时，模型运行正常，且显存占用压到了21G

但问题就来了，为什么16G的模型在设置gpu_memory_utilization=0.9的时候，会占用140G的显存呢，除了模型本身的参数之外，显存里面还装了什么东西？

在 vLLM 中，我们发现 LLM 服务的性能受到内存的瓶颈。在自回归解码过程中，LLM 的所有输入标记都会生成其注意键和值张量，并且这些张量保存在 GPU 内存中以生成下一个标记。这些缓存的键和值张量通常称为 KV 缓存。这些缓存特别大，LLaMA-13B 中的单个序列最多占用 1.7GB。而且其大小取决于序列长度，序列长度变化很大且不可预测，vllm引入PagedAttention来有效管理KV缓存，这是一种受操作系统中虚拟内存和分页的经典思想启发的注意力算法。与传统的注意力算法不同，PagedAttention 允许在不连续的内存空间中存储连续的键和值。具体来说，PagedAttention 将每个序列的 KV 缓存划分为块，每个块包含固定数量令牌的键和值。

现在确定这些多出来的数据是KV缓存，但是缓存大到占用140G的显存吗？

多卡和多机推理

单卡
gpu_memory_utilization=0.3

gpu_memory_utilization=0.25

单卡要比双卡占内存更少，推理更快

批量推理

    "model":"/data/models/Meta-Llama-3-8B-Instruct",
    "tensor_parallel_size": 2,
    "gpu_memory_utilization": 0.9,
    "temperature": 0.2,
    "top_p":0.95,
    "max_tokens":100,

    "model":"/data/models/Meta-Llama-3-8B-Instruct",
    "tensor_parallel_size": 2,
    "gpu_memory_utilization": 0.15,
    "temperature": 0.2,
    "top_p":0.95,
    "max_tokens":100,