Does pytorch xla work with gradient checkpointing?

[david-alexander-white]

❓ Questions and Help

I've got a decently big model that uses gradient checkpointing (using torch.utils.checkpoint), seems to max out mem usage at 4GB, and fits inside of my 8 GB of GPU memory with no problem.

However, when I try to run it on a TPU using XLA, it uses more than 8GM of memory and gets an OOM.

Reading the xla API guide, it seems to me that the lazy evaluation of xla tensors might actually stop the checkpointing, because xla might compile down what it sees as wasted computation (can explain this more if necessary).

However, I also understand there could be other issues, like something having to do with alignment to 128, which I also don't understand.

Just hoping to get some guidance to aid my further debugging.

[dlibenzi]

If the model is fully fused, I don't think checkpointing will be saving you anything with our Lazy Tensor design.
Any repeated graph will be squashed by XLA CSE pass.

[david-alexander-white]

Yeah, that’s kind of what I figured. Do you have any recommendations for how to handle this? I could try to switch to tensorflow, or use GPU, but wondering if it’s possible with xla — something like forcing an evaluation after the forward pass and then after each checkpointed segment on the backward pass. Or maybe somehow indicating that the optimizer shouldn’t squash this particular repetition. I’d be open to contributing on a fix to the library with some help, if that’s even a possibility.

…

On Thu, Jan 30, 2020 at 10:32 AM Davide Libenzi ***@***.***> wrote: If the model is fully fused, I don't think checkpointing will be saving you anything with our Lazy Tensor design. Any repeated graph will be squashed by XLA CSE pass. — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <#1571?email_source=notifications&email_token=AOHFHWIB6D4E7NOFSP4VPR3RAMMKZA5CNFSM4KNZR6G2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOEKMBCMI#issuecomment-580391217>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AOHFHWISEQVYXCCQXCCE5E3RAMMKZANCNFSM4KNZR6GQ> .

[dlibenzi]

Did you try to remove checkpointing with pytorch/xla?

[david-alexander-white]

Yeah, I ran it originally without checkpointing and got OOM, which is why I added it. That run had different hyperparams though, and could have been larger. Are you implying XLA should be able to handle this issue in its own and my use of checkpoint is making it worse? If not, im actually kind of excited about my proposal from the last thread and will probably try it out tomorrow.

…

On Thu, Jan 30, 2020 at 6:07 PM Davide Libenzi ***@***.***> wrote: Did you try to remove checkpointing with pytorch/xla? — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <#1571?email_source=notifications&email_token=AOHFHWJERVHT4BEN3TLLQ4TRAOBU5A5CNFSM4KNZR6G2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOEKNHWRQ#issuecomment-580549446>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AOHFHWPJKYGXJLBOBX4QWADRAOBU5ANCNFSM4KNZR6GQ> .

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[jysohn23]

Just tried out a simple MLP with gradient checkpointing and verified it seems to work as intended. The first fwd pass is run with torch.no_grad so not storing the activations is decided by PyTorch and XLA just takes that graph where fwd is run twice per bwd. Looks like the HBM footprint is reduced with a simple snippet like this one: https://gist.github.com/jysohn23/193afe7e8828aeeed0ea334900a2d3bb

Dumping the HLOs for the two cases:

• grad_ckpt_false_graph_0002.hlo.txt
• grad_ckpt_true_graph_0002.hlo.txt

Running those HLO graphs internally and profiling it also shows correctly that the activation tensors don't reside on HBM.

cc: @myleott