用PyTorch写Agent循环时显存越跑越高，是框架问题还是我代码问题？

缓存分配器背锅居多，试试torch.cuda.set_per_process_memory_fraction限制下峰值，比子进程省事。

子进程隔离太重了，先查查是不是tokenizer或tool返回的list没清干净，我上次就是这问题。

我之前也踩过这个坑，后来发现多半不是PyTorch的锅，是你每次step里可能隐式创建了计算图，比如某些操作没包在torch.no_grad()里，或者调用了会保留中间变量的函数。你可以试试在循环里跑一下torch.cuda.memory_summary()，看看具体是哪块分配的，大概率能定位到是某个固定操作在累积。子进程隔离太重了，不太建议，先排查一下是不是某个工具返回值里的tensor被隐式引用了，我上次就是被一个返回dict的装饰器坑了。

我之前跑长流程也遇到过一模一样的，最后发现不是你的问题，PyTorch的缓存分配器确实有这毛病，显存曲线会周期性跳增。empty_cache只是把缓存块还给分配器，但底层显存不一定还给驱动，所以治标不治本。建议你试试把每个step的输入输出都显式detach到CPU再存，或者干脆用pinned memory做中转，能明显缓解。至于子进程隔离，那确实是最彻底的方案，但代价是序列化开销，看你能不能接受。另外你确认一下有没有用torch.inference_mode替代no_grad，这个对某些算子能省不少临时显存。

说实话这锅大概率得PyTorch背，它的缓存策略在长循环里就是会这样，我之前跑RL的agent也踩过。你试过设置PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb吗？调小这个参数能减少碎片化，有时候能拖慢显存上涨速度。不过真要根治，还是建议每N轮手动把当前的激活和中间变量清一遍，再调一次empty_cache，别指望每轮都del就完事。子进程隔离我试过，稳是稳，但跨进程传模型状态特别麻烦，非必要别碰。

显存涨一般是缓存分配器的问题，试试每轮固定用同一批变量名，或者干脆把中间结果挪到CPU上。

我之前跑长链路的RLHF也遇到过一模一样的现象，最后定位到是PyTorch的缓存分配器在作祟，它确实不会主动把显存还回去，尤其是有小张量频繁申请释放的时候。你试试在每次step结束前把中间结果统一detach到CPU再置空，或者调一下PYTORCH_CUDA_ALLOC_CONF的max_split_size_mb参数，有时候能立竿见影。至于子进程隔离，我之前试过，虽然能解决OOM但通信开销太大，如果不是特别极端的场景不建议优先考虑。另外你确认过是不是某些自定义的Function或者hook在累积计算图吗，我之前就是被一个自定义autograd的retain_graph坑过。

说实话你这情况我太熟了，之前搞长链推理的时候也被这个楼梯曲线折磨过。torch.cuda.empty_cache()这玩意儿其实只清空缓存池里没用的块，但PyTorch的分配器默认会保留已申请过的显存，所以它治标不治本。你试试在循环外面包一个torch.inference_mode()，比关梯度更彻底，能省掉不少autograd的中间节点开销。另外你说的CUDA graph hook问题，如果用了torch.compile或者自定义autograd.Function，确实可能有context残留，检查下有没有把保存的tensor引用在step里清掉。不过我更怀疑是LLM推理库内部的问题，比如transformers的past_key_values或者beam search的缓存没释放，这个跟PyTorch关系不大。子进程隔离是个笨但有效的路子，但每轮启动模型太慢，性价比不高。我建议你先用nvidia-smi dmon实时看下每步的显存分配来源，区分是模型权重还是激活值，再决定优化方向。最后提醒下，如果用了vLLM或者TGI这类服务化推理，它们的显存管理是独立于PyTorch的，那才是真正的大头。

我之前也踩过这个坑，基本可以确定不是框架的锅，PyTorch的缓存分配器确实会占着显存不还，但更大概率是你某个step里不小心把中间结果保存到了self或者闭包里，查一下是不是有list在累积。empty_cache()只是清空未使用的缓存块，如果引用没断它根本不会释放。子进程隔离有点重了，建议先试试把每个step的输入输出都显式detach到CPU，或者用torch.cuda.set_per_process_memory_fraction限制一下，看看曲线是不是平了。顺便问下你用的是不是多卡或者有别的库在偷存东西，之前我遇到过一次是tokenizer的padding把整个batch都扩到最大长度了。

说实话这问题我太熟了，之前跑长链路的ReAct agent也撞过一模一样的墙。torch.cuda.empty_cache()这玩意儿其实只清空未占用的缓存块，并不强制把显存还给驱动，所以你看曲线还是锯齿状往上爬。我后来用pytorch的memory_stats()接口去打了快照，发现是CUDA caching allocator在后台保留了一堆碎片块，每个step的中间激活值大小不一样，分配器就不断开辟新块，旧块又因为尺寸不匹配没法复用。你那句“每步tensor都del了”其实del只是减引用，真正释放要等垃圾回收触发，建议在循环里显式调用gc.collect()试试，虽然听着玄学，但有时候真能压下去几十MB。另外检查一下是不是有Variable或者tensor的grad_fn还在图里挂着，即便你关了梯度，某些算子比如in-place操作还是会在autograd图里留node。子进程隔离倒是终极方案，但代价是IPC序列化开销大，慢得离谱，我后来是改用固定shape的buffer池，把工具返回的结果都pad到同一维度，再用view复用显存，曲线就平了。你不如先跑个一二十轮，每轮打印torch.cuda.memory_allocated()和memory_reserved()，看看差值是不是在持续扩大，如果allocated稳定但reserved涨，那就是分配器碎片问题，直接设PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True启动，能立竿见影。

我之前跑长序列推理也碰到过一模一样的情况，大概率不是框架bug，就是PyTorch的缓存分配器在搞鬼。你单纯empty_cache其实只是把空闲块标记一下，并没有真正还给驱动，建议用torch.cuda.set_per_process_memory_fraction限制一下上限，或者每跑完一步就查一下memory_allocated和memory_reserved的差值。至于子进程隔离，重的话成本太高，轻量场景可以试试把推理封装成单独的函数，用torch.inference_mode加环境变量PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True，这个组合我实测能压住那部分隐性增长。

我之前也踩过这个坑，大概率不是PyTorch本身的问题，而是缓存分配器把显存吃住了。你试过用torch.cuda.memory_summary()看看具体是哪块在涨吗？我后来是靠每轮强制调一下gc.collect()加上把不用的中间变量显式赋值None，才勉强稳住的。

另外你提到子进程隔离，其实如果只是长循环，不用那么重，试试把每个step的输入输出都挪到CPU上，只在真正要推理的时候再传回GPU，内存换显存，曲线会平滑很多。还有个细节，如果你用了任何带缓存的库（比如某些tokenizer），也可能悄悄占显存，得排查一下。

反正我最后是发现有个自定义的tool返回值没转成CPU，一直挂在GPU上，删了引用也没用，得手动调tensor.detach().cpu()才行。你可以先看看是不是也有这种“漏网之鱼”。

我之前也踩过这个坑，大概率不是你代码的问题，就是PyTorch缓存分配器在搞鬼，它申请了显存之后宁可留着复用也不还给驱动。你试试在循环里定期调一下torch.cuda.reset_peak_memory_stats()看曲线，再配合max_split_size_mb那个环境变量调调碎片，能缓解不少。子进程隔离倒是能彻底解决，但跨进程传数据那开销也挺烦的，除非你的step本身就很重，否则不太值当。另外检查下是不是有哪个中间结果被不小心塞进了某个list或者缓存里没清，这比显存碎片更隐蔽。

我之前也踩过这个坑，后来发现大概率不是框架bug，而是PyTorch的缓存分配器把显存块留着复用了，empty_cache只是清空未使用的缓存，并不会还给系统。你试试在每个step结束后把中间变量统一放进一个list再整体置空，或者用torch.cuda.set_per_process_memory_fraction限制上限，看曲线会不会平缓一些。另外如果你用了transformers的generate，它内部可能保留past_key_values，记得显式传past_key_values=None。子进程隔离倒是能彻底解决，但开销太大，我建议先排查一下是不是某个库在偷偷累积计算图，比如vmap或者functional_call这类操作。

大概率不是框架的锅，你试试在循环里加torch.cuda.synchronize()看下峰值，可能是有隐藏的list在累积计算图。

子进程隔离太笨重了，先查查是不是tokenizer或工具返回的tensor被隐式引用了，用tracemalloc能定位。

我之前也踩过这个坑，折腾了好久才反应过来大概率不是代码逻辑的问题。你试过empty_cache和del，但显存还在涨，基本就是PyTorch的缓存分配器在作祟，它会把释放的块留在池子里复用，但多步推理里每步的tensor形状可能略有差异，导致缓存碎片化，池子越撑越大。官方其实有个做法是torch.cuda.set_per_process_memory_fraction限制上限，或者干脆用torch.cuda.memory_stats去打印一下缓存池的reserved和allocated曲线，看看是不是reserved一直在涨而allocated没涨，这样能确认是不是分配器的问题。

至于autograd的hook，你既然关了梯度，理论上不该有残留，但保险起见可以把推理包在torch.no_grad()里，并且每轮结束调一下gc.collect()，有时候del只是减引用，循环里的闭包或者异常栈还挂着对象。子进程隔离确实能彻底解决，但代价是IPC拷贝开销很大，而且你状态都在GPU上，多进程反而更吃显存。

我自己的做法是给每个step固定输入长度，让tensor形状不变，这样缓存复用率高很多，实测跑了200轮都没再涨。还有个取巧的办法是每10轮手动调用torch.cuda.reset_peak_memory_stats()看峰值，但这不是释放。你如果真想查根源，可以开一下TORCH_CUDA_ALLOC_CONF=expandable_segments:True，这玩意儿能减少碎片，不过得看你的CUDA版本支持不支持。先试试形状固定加gc.collect，再不行再考虑子进程，别一上来就上重武器。

大概率是你代码里某个引用没释放，试试每轮显存快照对比下tensor地址，能揪出真凶。

我碰到过类似问题，最后发现是tokenizer的padding缓存没清，换子进程最省心。

我之前也踩过这个坑，大概率不是你代码的锅，就是PyTorch缓存分配器在长循环里只增不减，empty_cache只是清空未占用的块，但分配器持有的内存不会还给驱动。建议你先用torch.cuda.memory_summary()看看是不是缓存池在涨，如果是的话可以试试给每个step单独用torch.inference_mode()包起来，顺便检查下有没有不小心把中间变量挂到计算图上了。子进程隔离确实能根治，但代价是传输开销大，如果不想改架构，可以先设个PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128这种参数把碎片控制住，实测能缓解不少。

大概率还是缓存分配器的问题，不是代码逻辑错，试试用torch.cuda.memory_stats看下峰值是哪个op在涨。

实在不行就换vLLM或者把推理丢子进程，省心很多。

说实话我觉得大概率还是你代码的问题，但也不是说框架就完全没责任。PyTorch的缓存分配器确实会保留显存不还给系统，但那个“阶梯式上涨”通常不是缓存残留，而是某个你忽略的引用链在持续累积。比如你把中间结果塞进了某个list或者dict里没清，或者每个step都新建了计算图节点，哪怕关掉梯度，某些操作比如torch.no_grad()作用域外的张量运算还是会留下autograd元数据。我之前遇到过一个类似的情况，最后发现是给LLM的tokenizer输入里带了input_ids的.to(device)操作，每次生成新tensor但旧引用被模型的past_key_values缓存悄悄挂住了。

empty_cache()本来就只能释放空闲块，不能解决真正泄漏的对象。你试试用torch.cuda.memory_summary()看每个step的allocated和reserved的差值，如果allocated也在涨，那就不是缓存问题。另外子进程隔离是个粗暴但有效的方案，不过每轮都spawn进程开销太大，不如考虑用torch.compile或者干脆把长期Agent状态挪到CPU，只在推理瞬间把当前step的数据搬到GPU——反正LLM的输入输出都是小tensor，状态没必要全放显存。至于CUDA graph，你要是没手动捕获过，那基本可以排除，但autograd hook确实有可能在循环里被重复注册，建议检查一下有没有对同一个tensor调用.register_hook。最稳妥的做法是写个最小复现脚本，固定随机种子跑50轮，看显存曲线是否还涨，如果涨就把每步的torch.cuda.max_memory_allocated()打出来定位到具体操作。

这问题我上周刚踩完坑，大概率不是PyTorch的锅，而是你某个循环里悄悄保留了计算图或者缓存了中间变量。建议你试试在每轮step结束后强制gc.collect()，然后打印torch.cuda.memory_summary()看看到底是哪个op在涨，比瞎猜靠谱。另外如果用了functools.partial或者闭包传LLM调用，小心它把上一轮的输出隐式引用住了。子进程隔离是个笨但有效的方案，但代价是速度慢一半，建议先把代码里的retain_graph和no_grad范围检查一遍再说。

大概率还是缓存分配器的问题，试试每步固定shape加torch.cuda.set_per_process_memory_fraction限制上限。

我之前也踩过这坑，最后用pynvml监控发现是碎块太多，建议把中间结果全挪到CPU上。