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

这问题我踩过差不多的坑，最后发现八成不是框架的锅，是你某个地方不小心把计算图或者中间变量存进list了，哪怕grad关掉，只要有引用在显存就下不来。empty_cache只是把缓存还给分配器，不代表释放给系统，你试着在循环里监控一下每个step前后的memory_allocated，看看是不是真的在涨。子进程隔离确实能根治，但开销太大，一般跑个几十轮真没必要，先拿pytorch的memory_snapshot分析一下是谁在占着吧。另外cuda graph如果用了，记得每轮要重新capture，不然旧节点会一直留着。

子进程隔离有点重了，先试试把每轮的输出detach后强制清缓存，大概率是缓存碎片化问题。

说实话你这个问题我太有共鸣了，之前跑一个类似的多轮推理agent也差点被显存搞到怀疑人生。我先说结论：大概率不是你代码逻辑漏删tensor，而是PyTorch的缓存分配器在长循环里确实有“内存碎片化”的问题，它不会把释放的块还给CUDA，而是留在自己的池子里复用，但每次step的tensor尺寸和shape一变化，旧的块就浪费了，新的又得重新申请，看起来就像显存阶梯式上涨。empty_cache()只是把池子清空还给驱动，但如果你每个step里还有中间变量被autograd的graph引用着（哪怕你关了梯度，某些算子内部的buffer也会被缓存），那它根本不会释放。我建议你先用torch.cuda.memory_summary()打印一下，看看是“reserved memory”在涨还是“allocated memory”在涨，前者是分配器缓存，后者才是真泄漏。如果确认是缓存问题，可以试试在循环外层固定输入shape、复用同一个tensor buffer，或者干脆把推理部分包成一个函数，每轮强制用torch.no_grad()并且调用gc.collect()配合empty_cache()，但说实话效果有限。你提到子进程隔离，那确实是终极方案，比如用multiprocessing或ray，每轮推理结束整个进程杀掉，显存完全释放，代价是通信和模型加载开销大，但稳定。我自己的做法是先用memory_summary定位，如果只是缓存问题就调大PYTORCH_CUDA_ALLOC_CONF的max_split_size_mb，把碎片化压下去，如果还不行就直接子进程，别跟框架的allocator较劲，太浪费时间。另外建议你查一下是不是有CUDA graph被隐式捕获了，有些库（比如transformers的编译模式）会在内部做graph捕捉，如果不手动释放，那才是真正的坑。总之先诊断再动手，别盲目换方案。

这问题我太熟了，之前搞长链推理的时候也撞过一模一样的情况。torch.cuda.empty_cache()其实只是把缓存块标记成可用，并不会真的把显存还给驱动，所以那玩意儿在循环里基本等于心理安慰。你查一下是不是有Variable或者Tensor的._base属性还挂着，有时候切片或者view操作会偷偷保留对原始大tensor的引用，del掉表面变量根本没释放底层存储。另外，如果用了torch.compile或者任何带autograd.Function的自定义算子，那graph的缓存确实会随着step数累积，建议每N步调一次torch._dynamo.reset()试试。子进程隔离确实是最暴力也最彻底的解法，但代价是每个step都要重新加载模型权重，如果模型不大倒还好，模型一大那IO开销也够呛。我更建议你先用nvidia-smi看下是cudaMalloc的碎片化增长还是真的peak usage在涨，如果是前者，用PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True启动能解决一大半。顺便问下你中间状态是不是用了python list存，如果是的话那list本身倒还好，但list里每个元素如果是不定长的序列，PyTorch的缓存分配器会为每次新shape重新分配池，旧的池又不会合并，这才是楼梯效应的元凶。

大概率不是框架的锅，你查查是不是tool返回的中间结果被隐式存进graph了，试试把每步输出detach后转numpy落盘。

我碰到过一模一样的，最后查出来是PyTorch的缓存分配器在搞鬼，它确实不会主动把显存还给驱动，empty_cache只是清空未使用的缓存块，但那些块还被分配器占着。你可以试试给每个step的推理包一个with torch.no_grad()，然后手动调一下allocator的set_per_process_memory_fraction，看看能不能把峰值压住。另外检查下是不是有tensor的grad_fn链没断干净，有时候del了但引用还在，用gc.collect()强制回收一下。子进程隔离太麻烦，性能损耗也大，我建议先跑个profiler看看是哪个op在累积内存，别急着换架构。

大概率还是缓存分配器在搞鬼，试试看每轮固定输入shape能不能缓解。子进程隔离有点重，先排查下是不是tokenize长度变化导致的碎片化。

大概率是你代码里某个list或dict在循环里悄悄累积了graph，试试每步结束后把相关变量显式置None再gc.collect()。

我遇到过类似情况，最后是关掉grad然后每轮调一下torch.cuda.synchronize()解决的，你可以先查查有没有隐藏的引用没断。

我之前跑长流程也遇到过一模一样的情况，empty_cache和del都用了，最后还是得靠限制PyTorch的缓存池上限才压住。不过你这个阶梯式上涨大概率不是缓存分配器的问题，更像是有计算图或者hook在某个地方被隐式保留，建议你每个step结束用torch.profiler看一眼内存分配栈，比瞎猜靠谱。子进程隔离确实能根治，但代价是通信开销大，如果每轮状态小可以试试，不然还是先查查是不是某些自定义层或工具返回值里带了autograd节点。

这问题我踩过一模一样的坑，大概率不是框架的锅，是你代码里某个循环变量悄悄被graph记住了。我之前跑强化学习也这样，后来发现是loss.backward()里产生的临时节点没清干净，建议你试试把推理部分包在torch.no_grad()里，然后每轮强制del掉中间变量后调用gc.collect()，再empty_cache，显存曲线能平缓很多。子进程隔离倒是终极方案，但开销太大，先排查下是不是某个自定义模块的forward里有隐藏的requires_grad状态没关掉。

我之前也踩过这个坑，后来发现是Python引用没断干净，比如藏在list或者dict里的中间变量，del只是删了名字，对象还被引用着。建议用gc.collect()配合torch.cuda.memory_summary()看看是谁在占着。另外autograd的hook如果注册了没remove，确实会一直挂着graph不放。子进程隔离虽然能解决，但通信开销挺大的，先试试手动断引用加gc吧。

我之前也踩过这个坑，最后发现是autograd的图没断干净。你虽然关了梯度，但Agent里如果有些tensor参与了requires_grad=True的运算，哪怕只是记录，计算图也会一直挂着不放。建议在每个step结束时手动把中间变量detach一下，或者用with torch.no_grad()把整个推理包起来。另外检查下有没有全局list或者dict在偷偷存tensor引用，那种最容易被忽略。