FastChat部署多轮对话时上下文内存增长过快的调优方法

文章导读
多轮对话场景里,FastChat 会把这一轮的请求消息按原样拼成 prompt 发给后端。每增加一轮对话,token 数就增加一次;使用 vLLM 这类带 KV 缓存的后端时,显存占用还会在 token 长度达到某一阈值后出现明显抬升。这个增长不算异常,但不加限制任其增长,最后一定会在某一轮因为 OOM 被系统杀掉。调优的路径相对固定:在应用层截断长的历史消息,在 worker 层限制并发和显存上
📋 目录
  1. Ⅰ 先区分预期增长和异常泄漏
  2. Ⅱ 先在前端截断历史,而不是全部交给模型
  3. Ⅲ worker 层的显存和并发参数
  4. Ⅳ 用日志和显存曲线做收敛验证
A A

多轮对话场景里,FastChat 会把这一轮的请求消息按原样拼成 prompt 发给后端。每增加一轮对话,token 数就增加一次;使用 vLLM 这类带 KV 缓存的后端时,显存占用还会在 token 长度达到某一阈值后出现明显抬升。这个增长不算异常,但不加限制任其增长,最后一定会在某一轮因为 OOM 被系统杀掉。调优的路径相对固定:在应用层截断长的历史消息,在 worker 层限制并发和显存上限,用日志验证显存增长是否回到同步曲线。

多轮对话内存增长过快,通常不是单个原因。先看 KV 缓存是否持续占用,再看每轮是不是把全部历史消息原样送进模型。可以先做三件事:给 worker 设置显存上限、限制最大并发序列数、在应用层只保留最后若干轮对话。这三步不改变模型输出质量,但能压住显存峰值。若曲线仍持续上升,再检查调度参数和上下文剪枝。

先区分预期增长和异常泄漏

FastChat 部署一个模型后,模型权重占用的显存基本固定,真正随对话变化的是 KV 缓存。KV 缓存按“序列 × token 数”占用,多轮对话让 token 数不断上升,缓存占用也随之上升。如果你用 nvidia-smi dmon -c 10 -s mu 连续观察,会发现显存曲线跟着输入长度走,这属于预期增长。

真正要查的,是输入 token 不变时显存还在涨。这时可以打开 vLLM worker 日志看 gpu_cache_usage,它表示 KV 缓存当前使用率。如果长期接近 0.9 以上,说明缓存快被占满;如果请求长度没有变化却持续攀升,才需要怀疑异常泄漏。另一个验证办法是请求 /v1/chat/completions 时记录每条消息的 prompt 长度,把长度和显存曲线放在一起看。

先在前端截断历史,而不是全部交给模型

多轮对话上限失控的最常见起点,是前端只做追加不清理。FastChat 的接口会原样接收 messages 数组,你想要控制 token 增长,就需要自己在组装请求前截断历史。下面这个函数保留 system 消息,只保留最后 N 条对话,适合先做快速验证:

def trim_messages(messages, keep_last=10):
    if not messages:
        return []
    head = []
    if messages[0]['role'] == 'system':
        head = [messages[0]]
        messages = messages[1:]
    return head + messages[-keep_last:]

使用的时候,把原来的 messages 先经过这个函数再发给 FastChat。keep_last 是按条数裁剪,不是按 token 数;如果你希望按 token 边界控制,可以先用模型对应的 tokenizer 统计总长度,超限时从旧到新逐条移除。需要留意的是,截断的边界就是丢上下文:如果用户在第 3 轮被告知过关键约束,而你又只保留最后 10 轮,模型会“忘记”那条约束。所以建议把关键约束写进 system 消息,而不是放在对话历史里。

FastChat部署多轮对话时上下文内存增长过快的调优方法

worker 层的显存和并发参数

应用层截断解决的是“每轮送多少 token”的问题,worker 层还限制“并发多少序列”。使用 vLLM 后端时,worker 进程通常是 python -m fastchat.serve.vllm_worker,其命令行支持 `--max-gpu-memory`、`--max-num-seqs` 和 `--max-paddings`。`--max-gpu-memory` 控制 worker 能用到多少显存比例,`--max-num-seqs` 控制最大并发序列数,`--max-paddings` 控制 padding 的调度余量。可以参考以下启动方式:

python -m fastchat.serve.vllm_worker \
    `--model-path` /path/to/model \
    `--host` 0.0.0.0 `--port` 21002 \
    `--worker-address` http://127.0.0.1:21002 \
    `--max-gpu-memory` 0.7 \
    `--max-num-seqs` 32 \
    `--max-paddings` 64

启动后可以在 worker 日志里看到它打印的并发上限和缓存估算值。需要说明的是,`--max-gpu-memory` 0.7 不是“GPU 只用到 70%”,还需要包含权重和激活值;设得过高反而会让 vLLM 在分配 KV 缓存时先拿到过大的预配额,弱化其他任务的调度空间。`--max-num-seqs` 调小能减少缓存峰值,但会降低并发吞吐;如果业务上线有时间要求,必须结合自己的压力测试观察延迟,没有统一推荐值。

用日志和显存曲线做收敛验证

调整完参数后,不要只看一条命令的输出。建议做一个固定持续时间的压测:准备若干条 2000 token 左右的对话请求,循环调用 10 分钟,记录每轮 worker 日志中的 gpu_cache_usage 和 nvidia-smi 的显存占用。如果两个数值都落在稳定区间,说明参数生效。若 gpu_cache_usage 持续上升且请求长度不变,需要回到更早的设置排查。

还有一点,无论怎么调优,单 worker 能承载的上下文总是有上限。无限轮次的多轮对话最终还会触到边界。建议在上层服务里设置会话最大轮数或最长生命周期,时间到了就强制开启新会话。这样不用经常重启 worker,也可以让显存曲线形成“涨到一定高度后回落到初始值”的循环,而不是只涨不停。