最近在做一个基于LLM的简单Agent,用PyTorch加载了Qwen2.5-7B,每次调用模型进行工具调用或回复时,显存都会增加几百MB。我查了一下,可能是每次推理时都重新创建了计算图,或者历史对话拼接后没有释放缓存。我试了torch.no_grad()和清空缓存,但感觉还是不对。有没有大佬遇到过类似问题?是用kv cache复用,还是直接把历史对话截断到固定长度?另外,如果Agent需要调用外部API(比如搜索),返回结果再喂回模型,这中间怎么管理显存才不会崩?求指点,谢谢!
楼主
2026-07-19
用PyTorch写Agent时,多轮对话的显存一直涨怎么优化?
请 登录 后发表回复
全部回复
共 163 条
2楼
2026-09-25
你这种情况大概率是每次推理都重新算了历史token的KV,没用past_key_values缓存。把历史对话的KV cache存下来,每轮只喂新增的token,显存能稳很多。工具返回的结果别整段拼进去,做一下摘要或者只保留关键字段再喂回模型。另外Qwen2.5-7B本身权重就占不少,建议用device_map分一下或者量化,不然多轮下来缓存迟早撑爆。
3楼
5天前
这个思路不错,收藏了。
4楼
2天前
我遇到过类似情况,大概率是每次把完整历史拼进去导致KV cache没复用,白白重复算。你可以看看是不是每次调用都重新encode了整段对话,那样显存肯定涨。建议把历史对话截断到固定长度,或者用vLLM这类支持cache复用的推理框架,比裸PyTorch省很多。外部API返回结果喂回去之前记得detach掉中间张量,别把没用的计算图挂着。