最近在搭一个多步骤推理的Agent,每个step要调用好几次LLM和工具,中间状态都放GPU上。我发现只要跑超过30轮,显存就稳定上涨,最后直接OOM。我已经试过torch.cuda.empty_cache(),也确认过每个step的tensor都del了,梯度也关掉了。但显存曲线还是像楼梯一样一步一步往上走。有人说是PyTorch的缓存分配器不会主动释放,也有人说是CUDA graph或者autograd的hook没清干净。想问问大家:这种长循环场景下,有没有什么官方的best practice?还是说我应该干脆每轮把推理放到子进程里隔离?有点迷茫,求指条明路。
用PyTorch写Agent循环时显存越跑越高,是框架问题还是我代码问题?
全部回复
共 59 条说实话这问题我太有共鸣了,之前跑一个类似的长流程推理也差点被显存逼疯。你试的那几个方法我也都试过,empty_cache确实只是把缓存标记成可复用,但PyTorch的caching allocator并不会真的还给驱动,所以曲线看着就是只涨不跌。我后来查了一圈,发现如果每个step里用了torch.compile或者开启了CUDA graph,那部分缓存是根本不会释放的,必须显式调用torch.cuda.graph_pool或重置graph对象才行。至于autograd的hook,如果你用了任何带requires_grad的中间量,哪怕最后没backward,那些保存的tensor也可能被图结构引用着,建议所有推理tensor统一用torch.no_grad()包一层,并且检查一下有没有哪个工具函数偷偷把参数转成float32。子进程隔离倒是能彻底解决,但跨进程传tensor的开销和IPC的序列化成本在30轮循环里可能比显存还更拖慢速度,我觉得可以先试试在循环外创建一次CUDA context,然后每N步做一次torch.cuda.reset_peak_memory_stats()和手动触发gc.collect(),同时把模型输出都转成numpy再存,别让张量留在GPU上。最后提醒一下,检查一下是不是有某个库(比如transformers)默认开启了gradient checkpointing,那种隐式缓存特别容易漏。
我之前也踩过这个坑,大概率不是你代码逻辑的问题,就是PyTorch缓存分配器在长循环里“贪”着内存不还。你试试把每个step的输入输出都显式detach()到CPU再处理,或者干脆用torch.no_grad()包住整个推理段,能少很多隐性graph保留。另外,子进程隔离确实是个粗暴但有效的解法,就是IPC开销有点大,如果step之间状态依赖不重的话可以优先考虑。你查过NVIDIA的PyTorch官方issue没?我记得有个关于CUDA caching allocator碎片化的讨论,里面有人建议定期调一下PYTORCH_CUDA_ALLOC_CONF的max_split_size_mb参数。
大概率是你代码里有东西在偷偷累积计算图,查查循环里是否把loss或中间变量存进了list。
试试每轮显式 detach 输出再加 torch.cuda.synchronize(),不行就子进程隔离,省心。
这大概率不是框架的锅,PyTorch的缓存分配器本来就只会按需增长,不会主动把显存还给CUDA,所以你看到“楼梯式”上涨其实是缓存池在膨胀,不是实际泄漏。我之前跑长循环也踩过这个坑,后来用torch.cuda.reset_peak_memory_stats()配合每N步手动调一下cuda.memory_snapshot()看哪块没释放,发现是某些中间变量被闭包引用着没真正删掉。子进程隔离太重量级了,建议先试试把所有张量显式移回CPU或者干脆用torch.inference_mode()包住推理循环,能省不少事。你检查过是不是有多个step共用同一个计算图导致autograd保留图了吗?
我印象里torch的缓存分配器确实是个坑,它跟cudaMalloc不一样,用完的显存不会还给驱动,导致监控看着像内存泄漏。你试过用torch.cuda.memory_summary()看下是不是真的被缓存占着吗?另外检查下是不是有Variable被存在list里了,有时候自己没注意就引用住了。子进程隔离太麻烦了,不如先试试把每个step的输入输出都转成CPU上的python对象,只在必要的时候才往GPU搬张量。我之前遇到类似问题,最后发现是tokenizer的padding side导致的中间变量没释放,你可以往这个方向查查。
这问题我熟,之前跑长流程推理也踩过。大部分情况真不是代码逻辑漏了tensor,就是PyTorch的缓存分配器在作祟,它吃进去的显存不会自动吐出来,empty_cache只是清空未用块,治标不治本。你可以试试在每个step末尾加个torch.cuda.reset_peak_memory_stats()看下实际峰值,如果确认是缓存问题,可以调PYTORCH_CUDA_ALLOC_CONF的max_split_size_mb参数,或者干脆用garbage_collect + empty_cache组合拳。另外检查下有没有用到torch.no_grad()包住整个step,有时候中间变量被autograd图引用着没释放。子进程隔离倒是最终方案,但开销不小,建议先排查清楚再上。
这问题我折腾过,大概率不是框架的锅,是你代码里某个引用没释放干净。重点查一下每次step里有没有把中间结果append到list或者dict里存着,还有工具返回的response有没有被隐式转换成tensor。empty_cache只是把缓存还给分配器,不是清内存,真正要盯的是Python对象的生命周期。子进程隔离是最后手段,成本太高,建议先试试把推理包在with torch.no_grad()里,再看看tokenizer的padding有没有动态变化导致缓存key无限增长。
这个问题我踩过很久,最后发现大概率是你代码的问题,但根源不在tensor没释放,而是PyTorch的缓存分配器在长循环里会保留各种size的block,即便你del了,它也只是把显存标记成free,并不会真正还给CUDA,特别是每次LLM调用生成的中间tensor尺寸又不完全一样,缓存碎片化就会越堆越多。empty_cache其实只能清空空闲块,如果每次step里有新的显存分配请求触发了cudaMalloc,那缓存池还是会慢慢膨胀。你可以试试给每个step的推理包一个torch.inference_mode(),再把LLM的forward里所有不需要梯度的参数都显式detach,另外检查一下是不是有tokenizer或者工具返回值里的list被隐式转成tensor留在图上。至于子进程隔离,那是最后大招,进程间通信开销很大,多步骤agent延迟会翻倍,不太划算。还有个偏方:在循环末尾手动调一下torch.cuda.memory_reserved()看看峰值,如果持续增长,就说明是分配器碎片问题,可以定期把最大缓存块清掉,或者直接换vLLM这类推理框架来管理KV cache。另外确认下你是不是用了torch.compile,它的graph缓存有时候会保留额外状态,关掉试试看。
这问题我太熟了,之前跑RLHF的rollout也这样,del和empty_cache都治标不治本。你查下是不是tokenizer或者某些自定义module里保存了graph,或者用了torch.compile,那个会累积额外内存。子进程隔离肯定能解决,但开销大,我后来用pytorch的memory_stats接口逐轮打印,发现是某个常量tensor被反复注册到autograd里了,关掉requires_grad才彻底好。
这问题我踩过一模一样的坑,大概率不是你代码逻辑的问题,就是PyTorch缓存分配器在搞鬼。empty_cache只是把缓存标记为可用,并不会真正还给CUDA,尤其你这种循环里反复创建小tensor的场景,显存碎片化会越来越严重。建议先试试把每步的输入输出都统一成固定shape,减少分配粒度,再看下能不能用inference_mode替代no_grad,那个连autograd的元数据都不生成。如果还是涨,那子进程隔离反而是最省心的方案,就是跨进程传tensor稍微有点开销,但长任务稳定优先。
这问题我踩过一模一样的坑,最后发现真不是代码泄漏,就是PyTorch缓存分配器在搞鬼。你试试把torch.cuda.set_per_process_memory_fraction设个上限,或者干脆用PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True启动,能缓解不少。另外如果每个step的中间变量生命周期完全独立,可以考虑用torch.no_grad()包住整个循环,顺便把推理和工具调用拆到不同stream上,能压掉不少碎片化显存。子进程隔离确实一劳永逸,但通信开销大,先试试配置调优再说。
我之前也踩过这个坑,大概率不是框架的锅,是PyTorch缓存分配器在搞鬼——它确实不会把显存还给驱动,而是留在自己的池子里复用。你光调empty_cache其实没用,它只能清空未使用的缓存块,但如果你每个step里有微小显存碎片,池子就会越撑越大。建议先试试torch.cuda.set_per_process_memory_fraction硬性限制,或者把推理包在with torch.no_grad()里,再检查一下有没有隐性的Tensor保留在工具返回值里,比如某些库会悄悄缓存历史结果。子进程隔离倒是终极方案,但开销不小,可以先看看用torch.inference_mode()加上禁用gradient checkpointing能不能缓解。
这问题我也踩过,大概率不是框架bug,是缓存分配器的事,试试pytorch的malloc:async配合max_split_size_mb调参。
子进程隔离治标不治本,建议查下是不是tokenizer或者工具返回的list没清干净,我之前就是栽在hidden state引用上。
我之前也踩过这个坑,折腾了很久。你试过empty_cache没用挺正常的,它只是把缓存块标记为空闲,并不会还给驱动,显存占用看起来还是满的。我后来发现真正有问题的是torch的缓存分配器在长循环里会产生大量碎片,尤其是你每次step的tensor尺寸不一样的时候,它宁可保留旧块也不去复用。你可以试试在循环外固定好所有中间tensor的最大shape,或者用torch.cuda.set_per_process_memory_fraction限制个上限,至少能提前爆而不是悄悄涨。至于autograd hook,如果你用model.eval()还关不掉,那多半是某些自定义层的buffer在累积,比如用了inplace操作或者把list append到module的属性上。子进程隔离是个笨办法但确实干净,我见过有人用multiprocessing每轮起新进程,显存曲线直接变直线,就是IPC开销大,速度慢一半。另一个思路是查一下是不是LLM的tokenize结果里带了可变长sequence,导致attention mask的shape每次不同,这样缓存分配器就疯狂找新块。你可以打一下torch.cuda.memory_summary()看看是哪个op在分配,基本能定位到具体层。我个人觉得框架问题占八成,PyTorch对这类“高频小请求”场景确实优化得不够,但也不是没救,把推理尽量batch化或者用torch.compile试试,有时候图优化能帮你省掉不少中间保留。
这种楼梯式上涨大概率不是框架bug,我以前跑长流程也踩过。PyTorch的缓存分配器确实会预留显存,但你del了tensor之后它不一定立刻还给驱动,建议先试试torch.cuda.set_per_process_memory_fraction限制上限,观察曲线是不是变得平稳。另外你提到调LLM和工具,如果中间有Python对象在隐式持有graph或tensor引用,用gc.collect()强制回收一下可能比empty_cache管用。子进程隔离是最后手段,但每轮启动开销很大,多step场景不一定划算。
我遇到过类似情况,最后发现是某个tool返回的tensor虽然del了,但它的storage被autograd的saved_variables间接引用着。你确认一下每个step的loss或者中间变量是不是还被某个list存着,实在不行就显式把requires_grad置False再清一次。empty_cache本身只清缓存块,不清Python对象,别指望它能解决一切。子进程方案代价太高,不如试试把每次推理的输入输出都转成numpy,彻底切断GPU上的生命周期。
这问题八成出在CUDA caching allocator的碎片化上,30轮以后显存涨不代表泄漏,可能只是分配策略导致的“虚高”。你可以开一下torch.cuda.memory_summary()看看到底是reserved还是allocated在涨,如果reserved稳定但allocated上去
我之前跑长链条推理也踩过这个坑,最后发现八成是你代码里某个地方悄悄引用了上一轮的输出,比如把logits存进list做历史记录,或者某个hook没解绑。PyTorch的缓存分配器确实不会马上还显存给驱动,但它的峰值是平的,不会无限涨,你这种情况更像是某些tensor被意外保留在计算图里了,试试在循环末尾加torch.cuda.synchronize()再打印显存看看。子进程隔离倒是能根治,但代价是慢很多,建议先用tracemalloc或者pytorch的memory_snapshot定位一下到底是哪一行在累积,别急着上多进程。
另外你提到梯度关了,但如果是调LLM的库(比如transformers),它内部可能自己开了grad,或者keep_last_hidden_state之类的参数默认在累积。可以检查一下是不是每个step都在往同一个list里append输出,然后那个list又被下一个step的输入引用着,这种引用链会让旧数据一直活着。我自己的经验是把中间结果全移到cpu,或者干脆用dict覆盖而不是list追加,显存曲线立刻平了。
我之前也踩过这个坑,最后定位到是PyTorch的缓存分配器在长循环里不会自动把显存块还给驱动,empty_cache只是清空未占用的缓存块,但碎片化还是会让峰值越来越高。建议你试试把每个step的输入输出都强制转移到CPU再回来,或者干脆用torch.no_grad()包住整个循环体,autograd的hook有时候确实会残留。子进程隔离有点重了,但如果你用多进程推理,记得用mp.set_start_method('spawn')避免fork带来的CUDA句柄复制,否则显存翻倍更头疼。
遇到过长循环显存阶梯式上涨,大概率不是你的锅,PyTorch的缓存分配器确实有这毛病,尤其频繁创建小tensor时。我建议先看看是不是某个中间结果被隐式保留,比如list里存了loss或logits,哪怕del了引用计数没归零。另外试试torch.cuda.set_per_process_memory_fraction限制上限,能提前暴露泄漏点,比裸奔OOM好排查。子进程隔离是终极方案,但进程间通信和模型加载开销太大,除非你step间完全无状态,否则不推荐。我最后是靠把所有中间计算挪到CPU(只留模型权重在GPU)解决的,慢一点但稳得一批。
大概率是你代码问题,试试每个step结束把中间变量都显式赋None再gc.collect(),比empty_cache管用。
这问题我太有同感了,之前调一个类似的agent也卡在这。其实大概率不是框架的锅,torch的缓存分配器确实不会主动把显存还给驱动,但会复用,所以如果曲线是稳定阶梯状,多半是有tensor在偷偷逃逸到新的计算图上。你关了gradient但可能没关inference_mode,或者某个工具返回值里带了非叶子tensor的引用,每次循环都在动态构图。建议你先用torch.cuda.memory_summary()看看是tensor还是cache占的,如果是cache,可以试试torch.cuda.set_per_process_memory_fraction限制上限,或者直接torch.cuda.memory_reserved()对比一下。另外hook清没干净这个说法很可疑,你可以在每次step结束强制gc.collect()配合empty_cache,但注意别放在循环热路径上,会很慢。子进程隔离是终极方案但代价太大,我后来是改成把中间状态全部塞进CPU,只在LLM前向那几百毫秒搬回GPU,瞬间就稳了。还有个小坑,有些算子比如scatter或index_put会隐式创建graph,你最好在循环外统一建一个静态graph,或者干脆用torch.compile把推理部分包起来。总之先别急着怀疑框架,用memory_summary定位一下再动手。