最近在折腾基于Llama 3.1本地部署一个简单的AI Agent,用来做信息检索+总结。单用户测试还好,但一上到多轮对话(上下文积累到4k tokens以上)再加两三个并发请求,显存直接爆掉OOM。我试过vLLM和TGI,但感觉配置起来参数好复杂,什么max_num_batched_tokens、调度策略,看得有点懵。想问下有没有经验的同学,在不换显卡(目前是RTX 4090 24G)的前提下,怎么优化推理时的显存分配?或者有没有轻量的Agent框架能自动管理上下文长度和并发?先谢过各位前辈了。
部署开源大模型做Agent,遇到并发推理+长上下文时OOM怎么破?
全部回复
共 177 条说实话你这情况我太懂了,vLLM那堆参数确实劝退。我后来直接用llama.cpp的server模式,配合--parallel参数控制并发,再开个--ctx-size限制单条上下文,24G跑4k tokens加两三个请求基本稳。另外可以试试给Agent加个简单的滑动窗口,像LangChain的ConversationSummaryBufferMemory那样,把老对话压缩成摘要,比单纯调推理参数省心多了。
试试把max_model_len调成8k,配vLLM的continuous batching,24G够用。另外LangChain里加个自动裁剪历史消息的callback,比换框架省事。
4090 24G跑长上下文加并发确实挺头疼的,我之前也被vLLM那堆参数折磨过。后来发现其实不用硬啃调度策略,把max_model_len限制在你实际需要的4-8k,然后开paged attention,配合continuous batching,显存能省不少。另外你可以试试把Agent拆成同步阻塞模式,让并发请求排队而不是同时挤进显存,体感上比硬扛OOM强。轻量框架的话,LangChain那个conversation token buffer加自带的请求池也可以,但别指望它自动优化显存,核心还是得手动卡上下文长度。
说实话你这情况我太熟了,4090 24G跑Agent并发就是卡在显存墙和KV Cache上。vLLM那些参数看着唬人,但核心思路其实就一条:把max_num_batched_tokens调小,强制它做preemption,牺牲一点吞吐换显存稳定。另一个坑是长上下文,Llama 3.1的attention机制对KV Cache消耗太狠,4k tokens看似不大,但多轮并发时每路请求都在吃独立缓存,建议你在Agent框架层面加个滑动窗口,比如只保留最近2k tokens的摘要,别让原始对话无限堆。至于轻量框架,可以看看LlamaIndex或者Haystack的Agent模块,它们自带上下文裁剪策略,虽然灵活度不如自己手写,但省心很多。还有个偏方是开flash attention,能省不少内存,但你得确认TGI或vLLM版本支持。最后问一句,你试过把max_num_seqs限制到2或者1吗?有时候单路请求慢点,但总比OOM崩掉强。
显存不够就砍历史,用滑动窗口或摘要压缩上下文,4090跑4k还OOM多半是vLLM参数没调对。
4090跑这个确实容易爆,我之前也踩过坑。vLLM里把gpu_memory_utilization调到0.9左右,再把max_model_len设成4096别让它默认吃满,能省不少。长上下文那块建议上个滑动窗口或者定期摘要,别傻傻全塞进去。并发的话可以试试限制max_num_seqs,两三个并发其实不用开太大batch,反而稳。
4k上下文加两三个并发就爆24G,大概率是KV cache没管住。vLLM里把gpu_memory_utilization拉到0.9、max_model_len卡到实际需要的长度、开enable_prefix_caching,基本能救回来不少。Agent框架那层其实更该自己控上下文,比如每轮做摘要压缩或者滑动窗口,别把整段历史都塞进去。
我之前也踩过这个坑,4090跑长上下文并发确实很容易炸。你可以试试把vLLM的gpu_memory_utilization调到0.85左右,再把max_model_len卡在4096,别让它按默认的8k去预分配KV cache,能省不少。另外Agent那边最好加个滑动窗口或者摘要压缩,不然多轮对话越滚越大,再好的调度也扛不住。框架的话可以看看LangGraph配合自己写的上下文裁剪,轻量又可控,比指望框架全自动靠谱。
我之前也踩过这个坑,4090跑4k上下文加并发确实容易炸。vLLM里把gpu_memory_utilization调到0.85左右,再限制max_num_seqs别超过4,能缓解不少。另外可以试试开enable_prefix_caching,多轮对话里重复的system prompt能省不少显存。Agent框架的话可以看看Langroid或者CrewAI,它们对上下文裁剪有些默认策略,不用自己从头写。
我之前也踩过这个坑,4090跑Llama 3.1做多轮确实容易炸。你可以试试把vLLM的gpu_memory_utilization调到0.9,再把max_model_len降到4096以下,别让它默认吃满。另外enforce_eager打开能省点显存,代价是吞吐掉一些,但稳定性好很多。Agent框架这块可以看看LangGraph或者自己写个简单的滑动窗口截断,比堆框架更可控。
24G跑Llama 3.1加并发确实紧,我之前也踩过这坑。后来把gpu_memory_utilization压到0.85左右,再把max_model_len卡到4096,OOM基本就少了,你那些调度参数一开始不用全懂,先默认跑通再说。上下文那块可以考虑滑动窗口或者只保留最近几轮加摘要,不然多轮一积累很快就顶满。Agent框架的话LangGraph配个简单的内存裁剪就挺轻量的,别上太重的。
vLLM把gpu_memory_utilization调到0.85,再开enable_prefix_caching,4k上下文两三个并发应该能扛住。
4090跑4k上下文加两三个并发确实容易顶到天花板,可以试试把gpu_memory_utilization压到0.85左右,再开enable_prefix_caching,多轮对话里重复的system prompt能省不少。max_num_batched_tokens别设太大,2048到4096之间先试,调度用默认的就行,调狠了反而容易碎。轻量框架的话可以看下LangGraph配个滑动窗口或摘要压缩,上下文别硬堆满,超了就自动裁剪,比纯靠推理引擎扛要稳。
试试KV cache量化加PagedAttention,vLLM里开enable_prefix_caching能省不少显存。
我之前4090也踩过这个坑,vLLM里把gpu_memory_utilization压到0.85左右、max_model_len别按模型上限开,按实际需要设到6k,再限制max_num_seqs并发数,基本就不炸了。长上下文那块可以试试开启chunked prefill,prefill阶段显存峰值能降不少,代价是首token稍慢。Agent框架我推荐看看LangGraph配合自带的trim_messages,或者用mem0做外部记忆,别把所有历史都塞进prompt里。
这个OOM基本是KV cache吃满导致的,4k上下文下Llama 3.1 8B每个并发大概要占1G多显存,三个并发叠加权重和激活就顶不住了。你可以试试把vLLM的gpu_memory_utilization压到0.85左右,再手动设max_model_len到4096而不是默认的模型最大长度,这样KV cache的预留空间能省不少。max_num_batched_tokens调小到2048左右会降低吞吐但显存峰值明显下降,调度策略用默认的就行,别折腾priority。另外FP8 KV cache量化在40系卡上是支持的,能直接把KV cache砍一半,精度损失做检索总结基本感知不到。Agent框架这边可以看下LangGraph配合滑动窗口加摘要压缩,或者直接用mem0这类做记忆管理的,能自动控制塞进prompt的历史长度。真要省心就上SGLang,它的RadixAttention对多轮对话的KV复用比vLLM更狠,同样显存能扛更多并发。
有没有更详细的教程推荐?