What is the best way to use torchlens to figure out what activations are being stored for the backward pass?

[vedantroy]

I'd like to understand what things my model is caching during the forward pass for the backward pass. Ideally sorted, so I can see the biggest memory-hogs.

Do you have a good understanding of how I can do this using torchlens?

I generated a model graph using the relevant method in torchlens, but I don't think that is the thing I need.

[johnmarktaylor91]

Interesting suggestion! TorchLens has a way of storing the gradients during a backwards pass, but currently doesn't have the feature you describe. It would not be impossible, I think, and it would fit well with Torchlens' other profiling features. My thought is that after the forward pass is finished, TorchLens could recursively crawl back from the outputs via the grad_fn of each tensor, and each time a grad_fn has a stored tensor (for later computing the gradients), add that storage amount to some field associated with that tensor. To save on runtime, I would probably make this an optional toggle, defaulting to disabled (since I suspect this might be a niche feature).

If this is a priority for you I can add it to the top of the list. I suspect it wouldn't take more than a few days of work.

[vedantroy]

No obligations, but ofc would be appreciated :)

For some reason, I assumed the activations were equivalent to what was being stored for the backwards pass. But looks like that's not the case?

[johnmarktaylor91]

TorchLens can save all activations if you tell it to, but not all of these activations are naturally saved by autograd during the forward pass--they're only saved if needed for computing the gradient later. So, the TorchLens saved activations will be a superset of the activations saved by autograd. Is that what you were asking, wanted to make sure I understand.

[vedantroy]

Yeah, that's what I was asking. I guess torchviz (https://github.com/szagoruyko/pytorchviz), does this -- but would be nice if the saved tensors from the backward could be mapped to op / layer names in the forward.

[johnmarktaylor91]

On second thought, there might be an easier way to implement your suggestion then. Really all that's needed is a checker to determine whether a particular layer executes a function that stores the activations during the forward pass. Then it would be simple to tally how much memory is used up by these tensors. Only question is how to implement this...