Open JackCaoG opened 4 months ago
This sounds exciting!
Awesome! One thing we may need to handle is autocast state with gradient checkpointing - the upstream restores state using device modules (e.g. torch.cuda
or torch.cpu
), and it fetches the device module using getattr(torch, device)
which won't work out-of-the-box for us.
We can probably just extend the _get_device_module
logic in the upstream to support torch_xla
.
@jonb377 we should update the checkpoint code to use the new device-generic API for amp here https://github.com/pytorch/pytorch/pull/124479
cc @tengyifei
🚀 Feature
Today PyTorch/XLA ask user to use its own version of the gradient checkpointing in https://github.com/pytorch/xla/blob/d1235858628417ed7abc0d61e6e9be50df3e1a87/torch_xla/utils/checkpoint.py#L145-L146 We should extend upstream's api instead of asking user to use our version.
Motivation
Upstream gradient checkpointing doesn't work because XLA's CSE(common subexpression elimation) pass will undo the gradient checkpointing. More details in https://github.com/pytorch/xla/issues/5766#issuecomment-1792913756 . As a result I copied the upstream checkpointing and add a
optimization_barrier_
on inputs of backward recompute. This is bad becausePitch
I chatted with @soulitzer , there is a way to pass context manager to extend the gradient checkpointing behavior. @soulitzer even went ahead and wrote a draft
What we should verified is that
optimization_barrier_
should only be applied on the input of the backward recompute, not to all the backwards. I think we should take above code and play with it and verified if we can use this approach to extend gradient checkpointing.Also on top of the
optimization_barrier_
we also do some pytorch/xla rng seed state management in https://github.com/pytorch/xla/blob/d1235858628417ed7abc0d61e6e9be50df3e1a87/torch_xla/utils/checkpoint.py#L151-L155 We should think about how to handle this part in the extension as well.cc @alanwaketan @jonb377 @albanD