Update save_on_cpu and checkpointing to work with functorch wrapped tensors

[soulitzer]

Stack from ghstack (oldest at bottom):

• -> Update save_on_cpu and checkpointing to work with functorch wrapped tensors #89166
• Improve hooks ordering behavior #85849
• Prevent module full_backward_hook from erroring in double backward #88357

Design doc: https://docs.google.com/document/d/1OX5__xKsZP-natgEnsRrK4gfD0wSN9WS6j4kekfn-IA/edit

This approach saves the inner-most tensor. As we unwrap, we also append to a list of rewrap functions that capture the necessary information to restore the metadata on the original wrappers. This PR tries to do most things in Python, but there are probably some APIs that could exist (or maybe already exist) that could simplify this PR.

• This PR does very weird things to stash autograd metadata:
  • The rewrap function needs to capture autograd metadata so that the higher order graphs don't get disconnected, we reuse TensorWrapper to do this, but in a way that is careful not to save the original TensorWrapper's data
  • During packing, we do a clone on the original TensorWrapper, then replace the value_ with an empty tensor, so this new dataless TensorWrapper gets captured instead by rewrap fn
  • During unpacking, when we run the rewrap fn, we just replace the value_ with the value we desire (this could either be the
    recomputed value or value that was previously offloaded)
  • The API exposed to replace value_ is set_data!
• There doesn't seem to be a reliable way to uniquely identify a tensor since id() gets reused, using data_ptr helps but it is
  also not enough sometimes. In this PR, I'm also using the first element of the Tensor to get a test to pass.

Unanswered questions:

• Why did we need to enable grad mode while packing (where was it disabled)

Other prototypes:

• [doesn't work] Update save_on_cpu to handle saving the same possibly wrapped tensor multiple times #89159 (alternate approach that saves the outer-most tensor instead and unwraps the necessary number of layers during unpack - the issue is that we cannot tell when we are saving the outer-most tensor)
• [WIP] dedupe tensors saved by save_on_cpu when running grad(grad #88976 (same approach as this PR, but in cpp, unfinished)

TODO:

• verify in tests that we are actually saving the correct amount of tensors
• try a non-zero bdim
• make that assert more precise

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/89166

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 7 Failures

As of commit a02b774:

The following jobs have failed:

• lintrunner
• linux-bionic-cuda11.6-py3.10-gcc7 / test (distributed, 3, 3, linux.8xlarge.nvidia.gpu)
• linux-bionic-cuda11.6-py3.10-gcc7 / test (functorch, 1, 1, linux.4xlarge.nvidia.gpu)
• linux-bionic-py3.7-clang9 / test (functorch, 1, 1, linux.2xlarge)
• linux-focal-py3.7-clang7-asan / test (functorch, 1, 1, linux.2xlarge)
• linux-focal-py3.7-gcc7 / test (functorch, 1, 1, linux.2xlarge)
• win-vs2019-cpu-py3 / test (functorch, 1, 1, windows.4xlarge)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[soulitzer]

inner -> unwrapped

[soulitzer]

Not sure this is true actually, if we consider the case;

grad(
   forward
        grad(
               forward  (tensor packed here)
               backward (or here)
   backward             (but unpacked here)

Lifting shouldn't happen, but I haven't seen the assert triggered.

[soulitzer]

@zou3519 In case you started to read this, I plan to rewrite the checkpointing logic actually, which actually may turn out to be a lot simpler than save_on_cpu (which I had mostly been thinking about before this).

From Alban's comments offline:

• the checkpoint application may be able to be simplified (unlike save_on_cpu, we don't really need to worry about recomputing both grad(t) and t), this should get rid of most of the unwrapping + wrapping logic.
• need to consider the case where we checkpoint a function that calls backward inside (just need to early stop when when the size of the storage reaches the size of what we expect to be saved) or else it would infinite loop
• agrees that reenabling functorch modes if backward is done later should be done, and should be straightforward

[zou3519]

need to consider the case where we checkpoint a function that calls backward inside (just need to early stop when when the size of the storage reaches the size of what we expect to be saved) or else it would infinite loop

For this case, let's say we're doing

def f(x):
  y = x.sin()
  z = y.sin()
  return z

def g(x):
  return checkpoint(grad_1(f), x)

grad_0(g)(x)

grad_1 executes something like:

y = x.sin()
z = y.sin()
# --- 
recomputed_y = x.sin()
gy = recomputed_y.cos()
gx = gy * x.cos()

If we "early stop checkpointing", what happens in grad_0? Do gy, recomputed_y get recomputed, or are they saved?

[github-actions]

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.