多GPU部署ChatGLM-3时KV Cache显存分配不足的优化配置

多GPU部署ChatGLM-3时遇到KV Cache显存分配不足，优先查的不是KV Cache参数，而是多卡显存的实际分布。常见情况是accelerate的device_map把权重均衡到各卡，但KV Cache在生成时仍集中在设备0，或者max_model_len和batch_size设置的预分配区间超出了剩余显存。以下步骤按前置条件、参数调整、验证、失败回退来组织，供实际部署时逐项排查。

前置条件

开始配置前，需要先确认四件事：

GPU数量和每卡总显存，至少要知道哪张卡剩余显存最低。通过nvidia-smi可看到实时占用。
推理框架是transformers原生生成，还是vLLM等服务化框架。两者的调参入口不同。
业务最长输入和输出长度。KV Cache的显存占用与序列长度是线性关系，若长度可压缩，问题会简单很多。
模型路径中是否包含ChatGLM-3的原始代码，加载时通常需要trust_remote_code=True，这会影响部分配置参数的写法。

这些信息不需要一次找齐，但能帮你判断后面调整参数时先动哪一项。

调整多卡参数：控制显存上限和分卡策略

使用transformers直接推理时，先用max_memory给每张卡设定上限，再把device_map设为"balanced"，避免自动分卡把KV Cache集中到第一张卡。下面的代码骨架放在模型加载前，max_memory的值会按当前可用显存动态计算：

import torch
from transformers import AutoModel, AutoTokenizer

model_path = "/path/to/chatglm3"  # 替换为你的模型路径
max_memory = {}
for i in range(torch.cuda.device_count()):
    total, free = torch.cuda.mem_get_info(i)
    max_memory[i] = f"{int(free // 1024**3 - 1)}GiB"  # 预留约1GiB给激活值和输出

tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModel.from_pretrained(
    model_path,
    trust_remote_code=True,
    device_map="balanced",
    max_memory=max_memory,
    torch_dtype=torch.float16,
    use_cache=True,
)

代码中预留约1GiB只是示例，具体预留量要按生成时batch大小和输出长度调整。如果这一步后仍然报KV Cache不足，再压缩max_new_tokens或max_length。这两个参数直接决定KV Cache的最长预分配区间。

使用服务化框架时：显存利用率和模型并行度

如果部署时使用vLLM这类服务化框架，KV Cache的分配方式与transformers不同。你通常只需要关注两个参数：tensor_parallel_size和gpu_memory_utilization。前者把单份权重和KV Cache切到多卡，后者控制每卡预留给模型和KV Cache的总显存比例。可先按以下命令启动：

python -m vllm.entrypoints.openai.api_server \
    `--model` /path/to/chatglm3 \
    `--tensor-parallel-size` 2 \
    `--max-model-len` 8192 \
    `--gpu-memory-utilization` 0.85

其中max_model_len要小于模型训练的最大长度，也要小于每卡剩余显存能支撑的长度；gpu_memory_utilization不建议直接拉满，先从0.8或0.85试，再根据启动日志调整。不同版本的vLLM对ChatGLM-3的支持程度和默认配置有差别，实际以框架报错提示为准。

验证KV Cache是否仍然不足

验证分两步。第一步确认模型如何分布在多卡上：

print(model.hf_device_map)

如果发现某个模块，特别是编码层或decoder层集中在一张卡，说明device_map的分卡粒度不够细，可考虑手动切分层。第二步观察生成过程中的显存变化，用nvidia-smi -l 1持续刷新，留意每卡显存在生成前和生成后各增加多少。如果启动时没有报OOM，但生成到一半报KV Cache不足，那么需要降低max_len或batch size，而不是单纯调整max_memory或gpu_memory_utilization。

失败时的回退顺序

优化配置不会一次成功，出现OOM时建议按以下顺序回退：

先缩减max_new_tokens或max_model_len，确认输出长度缩到多少时不再OOM，再反向查找失效项。
若缩减长度不可接受，把use_cache改为False。这会显著降低显存占用，但每次生成都要重新计算KV Cache，速度会明显下降，只能作为止血手段，不算优化。
vLLM场景下，把gpu_memory_utilization调低0.05～0.1，有时启动参数设置过高会导致每卡预留空间不足，系统报错反而指向KV Cache。
如果所有参数调低后仍不够，说明当前单卡总显存与模型并行策略不匹配，需要减少请求并发数或改用显存更大的GPU，这是硬件层面的边界。

每调整一个参数，都重新跑一次可复现的最小请求，记录能否通过以及显存峰值。确认不了的时候，宁可保留保守值，也不要为了更高吞吐硬调高显存上限。