最近在做一个基于Qwen2.5的简单Agent,就是在工具调用循环里反复让模型推理、解析动作、执行工具再拼回对话历史。用的vLLM离线推理(OfflineBatch),但跑十几个回合之后显存就涨到爆,只能重启进程。我已经把history列表做了截断,只保留最近几轮,token长度也控制了。怀疑是vLLM的cache机制在长循环里没释放旧序列,或者是我对OfflineBatch的调用方式不对?有没有大佬遇到过类似情况,是应该换用HF原生pipeline还是给vLLM加什么参数?救救孩子,真不想每个Agent任务都手动清显存。
PyTorch写Agent循环时显存越用越多,是推理框架的坑还是我代码的问题?
全部回复
共 56 条这问题我踩过一模一样的坑,vLLM的cache确实不会自动清掉旧序列,尤其OfflineBatch在循环里每次调用都相当于累积状态。你可以试试在每轮推理后手动调一下llm.offline_batch的reset或者干脆重建一个LLM实例,代价小但能解决。另外检查下是不是把对话历史整个拼进prompt了,截断后也要确保实际传入的token数真降下来了,有时候是系统prompt没算进去。换HF的话慢很多但没这毛病,看你要速度还是要省心。
之前跑ReAct循环也踩过这个坑,vLLM的cache确实不会跟着你的history截断自动清,它按sequence ID缓存KV,你每次拼新对话历史其实等于在造新序列,旧的那些如果没被evict就一直占着。OfflineBatch的坑在于它默认复用显存池,但你循环里如果每次传入的token长度波动大,预分配的buffer会越撑越大。可以试试在每次推理前手动调一下LLM的clear_cache或者用vllm的abort请求接口,不过离线模式这个不太好搞。另一个思路是干脆固定最大序列长度,把history截断到更短,让vLLM每次都能复用同一块显存区域,而不是动态扩。我之前换过HF原生pipeline,虽然慢但显存是真稳,不过Agent场景来回切换也挺麻烦的。你现在的history截断是截到多少轮?如果模型本身支持长上下文,其实可以考虑把工具结果做摘要压缩,而不是只截断,这样既能减少token也能避免频繁触发vLLM的重新分配。
遇到过一模一样的现象,当时也怀疑是vLLM的问题,后来排查下来发现锅大概率不在vLLM身上。你history截断只保留了最近几轮,但注意Qwen这类模型会把system prompt和工具定义也一起算进context,如果Agent循环里每次拼对话历史时都重新把完整的工具schema塞进去,那即便token长度看着没涨,实际KV cache的shape可能因为序列总长波动而一直重新分配。另外OfflineBatch有个细节,它默认会缓存每个序列的KV cache以备复用,但你的场景是不断生成新序列,旧序列的cache不会自动释放,除非你显式调用llm.destroy()或者把offline_batch的结果对象丢掉后手动gc.collect()。我试过在每轮循环末尾把output变量置空再torch.cuda.empty_cache(),有效但治标不治本。更靠谱的做法是直接改用LLM(..., max_num_batched_tokens=...)限制单次推理的token上限,同时把enable_prefix_caching=False关掉,这样至少不会让前缀缓存无限膨胀。如果不想折腾vLLM,用HF原生pipeline配torch.compile加cache_implementation="static"反而更稳,但速度会慢一半,看你对时延的容忍度了。你当前用的vLLM版本是多少?0.6.x和0.7.x之间的内存管理差异挺大的。
vLLM的prefix cache在长循环里确实容易这样,试试加--disable-prefix-caching或者手动清一下cache,我之前就是这么解决的。
这题我熟,之前做多轮tool calling也踩过同样的坑。vLLM的prefix cache在离线推理时确实容易把旧序列的KV block留在显存里,尤其是对话历史反复拼接时,建议试试把整个序列的request_id固定住,或者显式调用llm.finish_requests()来清状态。另一个思路是每轮单独创建LLM实例,虽然慢点但内存能稳定释放,HF pipeline反而没这问题。你history截断时记得保留系统提示和最新工具结果,不然效果会崩。
我之前跑类似的多轮agent也踩过这坑,大概率不是你的代码逻辑问题,vLLM的prefix cache在OfflineBatch里确实不会自动释放旧序列,你截断history只是控制输入长度,但缓存还是按完整序列存的。可以试试给LLM类传个enable_prefix_caching=False,或者干脆每次推理后手动调llm.cache_config清一下,不过最省事的方案是直接换HF pipeline,慢一点但稳。另外你看下是不是工具返回的content里有特殊token没处理干净,有时候隐式padding也会让显存缓慢涨。
vLLM的prefix cache在长对话里确实容易累积,试试加--max-num-batched-tokens或手动清理seq_group,比换HF省事。
我遇到过类似,OfflineBatch要每次重建LLM实例或调reset,你检查下是不是旧序列没被detach。
我之前搞类似的agent循环也踩过这坑,vLLM的显存增长多半不是代码逻辑问题,而是它的prefix cache在长对话里不断累积旧序列的KV,虽然你截断了history,但vLLM内部可能还保留着之前轮次的中间状态。你可以试试把OfflineBatch的调用改成每次创建新的LLM实例,或者手动调一下max_num_batched_tokens和gpu_memory_utilization,给cache留个上限。另外如果回合数实在多,建议直接换HF原生pipeline,虽然慢点但内存释放干净,省得半夜爬起来清显存。
之前做Toolformer类似的循环也踩过这个坑,vLLM的OfflineBatch在连续generate时确实会缓存KV cache和过去的sequence metadata,就算你截断history,它内部可能还保留着旧序列的block,尤其是当每次调用都传完整对话历史给同一个LLM实例时,显存碎片化会越来越严重。我当时试过给每个回合新建一个LLM实例再销毁,虽然慢但显存能稳,后来发现是vLLM对动态序列长度的释放策略比较保守。你试试在循环里显式调用一下torch.cuda.empty_cache(),然后给vLLM的SamplingParams里加上ignore_eos=False或者调小max_num_batched_tokens,看能不能缓解。另外如果你用的是LLM.generate而不是OfflineBatch的底层接口,建议看看vLLM的use_v2_block_manager参数,新版有这个选项能更激进地回收空闲block。不过说实话,如果agent回合数超过20,我最后直接换HF的pipeline加torch.inference_mode了,虽然慢一倍但起码不用老盯着显存,或者你也可以考虑给每个工具调用单独开一个子进程跑推理,这样内存回收最干净。
vLLM的OfflineBatch确实有这个问题,它内部会为每个序列维护KV cache,即便你截断了history列表,只要旧的sequence对象还在显存里没被回收,cache就不会释放。我之前写Agent也踩过这个坑,后来发现vLLM默认会缓存所有已处理过的prompt,哪怕你不再用它,它也不会主动清掉。你可以试试在每次迭代后显式调用llm.destroy()或者llm.offline_batch的返回结果用完之后把相关变量del掉,再配合torch.cuda.empty_cache(),但说实话治标不治本。更靠谱的做法是改用vLLM的在线API模式,每次请求是独立的,服务端会自己管理cache生命周期,虽然多一点点网络开销,但至少不会越积越多。还有个小技巧,把max_model_len调小一点,这样每个序列能占用的KV cache上限就被卡死了,超出部分会强制截断,虽然可能影响长对话的连贯性,但能救命。另外你如果只是做工具调用这种短上下文任务,HF原生pipeline其实更省心,它每次前向传播用完就释放,不会有这种累积问题,就是速度慢一些。我最后是换成了异步请求vLLM的HTTP接口,然后自己写了个简单的显存监控,超过阈值就重启worker,虽然粗暴但至少不用手动干预了。
vLLM的OfflineBatch确实会缓存KV cache,但sequence级别的不释放通常是因为你每次都把完整对话历史传进去了,虽然截断了history列表,但vLLM内部可能还保留着旧的block。你可以试试显式调用llm.free()或者每次生成后清一下tokenizer_cache,或者直接给SamplingParams加个max_tokens限制。另外我建议你换HF的pipeline跑一下同一段逻辑对比显存曲线,如果HF没问题那就是vLLM的调度策略问题,可以考虑用--enable-prefix-caching或者调低--max-model-len来缓解。我之前也踩过这坑,最后是改成每轮重建一个LLM实例才解决的,虽然慢点但至少不爆显存。
vLLM的prefix cache在长循环里确实容易积压,试试加--enable-prefix-caching=false或者手动清下cache,我之前这么解决过。
我遇到过,多半是OfflineBatch的sequence没释放,建议每轮迭代调一下llm.destroy()或者换AsyncLLMEngine试试。
我之前也踩过这个坑,vLLM的cache确实会在长循环里攒着旧序列不释放,尤其OfflineBatch反复调用时特别明显。建议试试给每个回合单独创建LLM实例,或者用完显式调一下llm.cache_config相关清理,虽然麻烦但有效。另外你确认一下是不是把整个对话历史都传进去了,vLLM对多轮拼接的显存占用比HF更敏感。如果不想折腾,HF的pipeline在小任务上反而稳,就是慢点,但至少不会爆显存。
vLLM的prefix cache在agent场景确实容易累积,试试加--enable-prefix-cache=false或者手动调gpu_memory_utilization。
之前跑类似multi-turn agent也遇到过,vLLM的prefix cache在离线batch里确实容易把旧序列的KV block留着,特别是工具调用这种长上下文反复拼接的场景。建议试试设置enable_prefix_caching=False,或者每个回合用新的LLM实例而不是复用同一个,虽然慢点但能稳。另外检查下是不是没调max_num_seqs,把它调小点强制淘汰旧block可能也行。我后来干脆换回HF的pipeline,配合torch.no_grad和手动清cache,虽然慢但至少不会爆显存。
我之前跑React Agent也踩过一模一样的坑,vLLM在离线批处理长轮次时确实不会自动释放之前的KV cache,尤其是你每次把整段历史都塞进同一个请求里,它内部会一直维护那些序列的状态。你光截断history列表没用,因为vLLM的cache是按sequence id管理的,新请求如果复用同一个序列id,旧token的显存块可能还在池子里占着,得显式调一下reset或者重新初始化一个LLM实例才行。我当时试过把每个循环都新建一个vLLM引擎,虽然慢但至少显存是平的,后来换成了HF的pipeline加torch.inference_mode,反而稳定很多,因为HF每次都会重新计算,不存在跨请求缓存残留。不过你要是追求速度,可以试试vLLM的prefix caching参数,或者用AsyncLLMEngine手动控制序列的释放,但那个API挺绕的。还有个取巧的办法,就是把工具调用结果直接截断到固定长度,比如只保留前200个字符,这样历史增长速度慢很多,能撑更久。你现在的token长度控制是只看输入还是算上了生成的部分?因为工具返回的拼接如果太长,即便截断历史,单次生成的最大长度也会撑爆显存。
我大概猜到问题出在哪了,vLLM的OfflineBatch本身是为单次生成设计的,它内部的KV cache和block管理在多次连续调用时不会自动清理之前的序列状态,除非你显式调一下llm.free()或者重置一下采样参数。我之前跑多轮对话也踩过这个坑,后来发现把每次调用包在with torch.no_grad()里,同时手动清空一下cache_config的gpu_memory_utilization占比,或者干脆给vLLM传个max_num_seqs=1强制限制并发,能缓解一些。但说实话,你这种Agent循环场景,vLLM的优化方向主要是高吞吐批处理,对这种长尾交互式的反复单序列推理其实不太友好,我更建议你用HF原生pipeline加上torch.compile,虽然慢一点但显存曲线是稳定的,因为每次前向传播的KV cache会随梯度释放。另外你history截断只保留最近几轮还不够,要检查一下是不是工具调用的返回结果里带了特殊token或者隐藏状态没被清理,有时候是拼回对话历史的tensor在计算图里被保留了。你可以试着在每个回合结束后跑一下gc.collect()和torch.cuda.empty_cache(),看显存能不能回落到初始值,如果能回落那就说明是缓存堆积,如果不能那就是某个张量被引用了。我自己的做法是干脆给每个循环回合单独起一个子进程跑推理,结果通过队列传回来,虽然笨但绝不出问题,代价就是多几十毫秒的启动时间。如果你找到vLLM官方的解决方案记得回来踢我一脚,我也想知道它到底支不支持这种动态序列释放。
大概率是vLLM的prefix cache在作祟,OfflineBatch建议手动调enable_prefix_caching=False试试。
我最近也踩过类似的坑,vLLM的cache确实会在长循环里累积,尤其OfflineBatch每次调用如果没显式清掉之前的序列,显存就慢慢涨上去了。你可以试试在每轮循环后调一下vllm的clear_cache或者重新初始化一个LLM实例,虽然笨但能稳住。另外别急着换HF,vLLM在长Agent场景下性能优势还是很明显的,实在不行就开--max-num-seqs限制并发看看。还得确认下是不是工具返回的文本太长,history截断只截了消息数但没限制单条内容大小,也会偷偷吃显存。
我也踩过类似的坑,vLLM的cache确实会在长循环里累积,尤其是OfflineBatch模式下如果每次传入的序列结构变化不大,旧KV cache不会主动释放。你可以试试显式调用free_cache接口,或者把每个回合的请求拆成独立batch,别复用同一个序列对象。另外,截断history的时候记得把对应的token id也同步清理,不然隐性长度还在涨。HF原生pipeline在显存控制上更直观,但速度会慢不少,建议先排查vLLM的版本和参数,比如enable_prefix_caching关掉可能有效。