memory leakage problem when loss.backward()

Hi.

How can i solve memory leakage problem on loss.backward? Actually my code is a bit complex, which makes it hard to provide whole.

The main memory burden comes from here. Exactly after the execution of loss.backward(). The memory increases for every iteration, which leads to OOM.

loss = bce_extended(logits, y).sum()
with backpack(BatchGrad()):
  if real_sample:
      loss.backward(inputs = list(model.parameters()))

I also tried with disable(): which prevent the memory leak problem. However, it cannot be implemented with with backpack(BatchGrad()): when i want to get the per-sample gradient.

[f-dangel]

Hi,

just wanted to bring it up because I saw there is a .backward(..., create_graph=True) in your code: There’s a memory leak when using full_backward_hooks with create_graph=True in PyTorch (#82528). You could try installing PyTorch with the fix (#82788) to see if that’s causing the memory leak.

[f-dangel]

Sadly I don’t have any data yet how BackPACK compares to functorch in terms of runtime, e.g. for computing individual gradients.

But I think you will be able to port your existing code to functorch with relatively few changes to try if it works fast enough for your needs.