用PyTorch写了个自定义算子，训练时显存直接爆掉，求排查思路

显存持续涨而不是一开始就爆，基本可以锁定是反向图里存了东西没释放。你检查下backward里是不是把索引矩阵直接存成了成员变量或者挂在module上了，那种每次迭代都会累积。scatter_add反向确实容易踩坑，如果用了index_add的backward但没正确处理重复索引的梯度累加，有时会触发额外的中间缓存。建议先用torch.cuda.memory_summary()按迭代打一下，看是哪块在涨，比瞎猜快很多。

显存持续涨多半是backward里存的索引没释放，scatter_add反向容易把中间张量挂住，检查下有没有在forward里存了不该存的。

scatter_add反向确实容易重复累加导致显存涨，先查下索引有没有重复，或者试试torch.cuda.empty_cache()看能不能压住。

显存持续上涨而不是一开始就爆，基本可以排除算子本身单次执行的峰值问题，更可能是每步都在往计算图里挂东西没释放。你可以先用 torch.cuda.memory_summary() 看下 reserved 和 allocated 的差值，如果 reserved 一直涨，那八成是缓存碎片或者有 tensor 被意外持有。自定义 backward 里如果你手动存了索引矩阵或者邻接表，记得检查是不是把它挂在了 ctx 上但没在 backward 里的正确位置释放，autograd 的 graph 本身会在 backward 后释放，除非你用了 retain_graph=True 或者有引用计数没清掉。scatter_add 反向这块坑挺多的，它本质是 index_add 的逆向，如果 forward 里用了 atomicAdd 或者不确定的索引顺序，反向的梯度累加可能会因为重复索引产生静默错误，最好用 torch.autograd.gradcheck 在 double 精度下过一遍小规模用例。另外点云里 KNN 的索引如果每步都重新算并且存下来，batch 一大多轮下来显存肯定扛不住，可以考虑把索引缓存到 dataset 层面或者用 torch.no_grad 包住索引计算。你 nvidia-smi 看到的持续上涨也可能是 PyTorch 的 caching allocator 没及时归还，试试 torch.cuda.empty_cache() 看有没有回落，如果有回落那说明是碎片问题，调小 batch 不如先查引用。

显存持续涨而不是一开始就爆，八成是 backward 里存的索引矩阵或者中间 tensor 被 graph 一直引用着没放。可以试试在自定义 Function 的 backward 里手动 del 掉不用的中间变量，或者用 torch.autograd.graph.save_on_cpu 把部分中间结果挪到 CPU。scatter_add 那块反向确实容易出问题，如果 index 有重复，梯度会累加，建议检查下是不是有地方重复 backward 了。另外用 torch.cuda.memory_summary() 看下到底是哪个环节在涨，比 nvidia-smi 细多了。

scatter_add反向容易重复累加梯度，检查下索引有没有越界或重复，显存涨多半是graph没释放。