functorch.grad support for autograd.Function

[zou3519]

Stack from ghstack (oldest at bottom):

• [pyfunctorch] Generate a more meaningful name for _SingleLevelAutogradFunction #90418
• Fix circular import in torch/autograd/function.py #90415
• functorch.jvp support for autograd.Function #90077
• Move private forward grad mode helpers to torch.autograd.forward_ad #90240
• autograd.Function supports vmap staticmethod #90037
• -> functorch.grad support for autograd.Function #89860
• Add setup_context staticmethod to autograd.Function #89859
• Add feature flag for the autograd.Function extension #89858

Happy to split this PR more if it helps.

This PR adds functorch.grad support for autograd.Function. There's a lot
going on; here is the high level picture and there are more details as
comments in the code.

Mechanism (PyOperator)

• Somehow, autograd.Function needs to dispatch with functorch. This is
  necessary because every layer of functorch needs to see the
  autograd.Function; grad layers need to preserve the backward pass.
• The mechanism for this is via PyOperator. If functorch transforms are
  active, then we wrap the autograd.Function in a custom_function_call
  PyOperator where we are able to define various rules for functorch
  transforms.
• custom_function_call has a rule for the functorch grad transform.

autograd.Function changes

• I needed to make some changes to autograd.Function to make this work.
• First, this PR splits autograd.Function into a _SingleLevelFunction
  (that works with a single level of functorch transform) and
  autograd.Function (which works with multiple levels). This is necessary
  because functorch's grad rule needs some way of specifying a backward
  pass for that level only.
• This PR changes autograd.Function's apply to eitehr call
  custom_function_call (if functorch is active) or super().apply (if
  functorch isn't active).

Testing

• Most of this PR is just testing. It creates an autograd.Function
  OpInfo database that then gets passed to the functorch grad-based tests
  (grad, vjp, vjpvjp).
• Since functorch transform tests are autogenerated from OpInfo tests,
  this is the easiest way to test various autograd.Function with
  functorch.

Future

• jvp and vmap support coming next
• better error message (functorch only supports autograd.Function that
  have the optional setup_context staticmethod)
• documentation to come when we remove the feature flag

cc @mlazos @soumith @voznesenskym @yanboliang @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @chunyuan-w @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• macos-m1-12 nodes are down

✅ No Failures

As of commit 3f953d1:
💚 Looks good so far! There are no failures yet. 💚

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

[zou3519]

Naming question: I refer to autograd.Function as "AutogradFunction" sometimes, and sometimes as "CustomFunction". What do you guys want me to call it?

[soulitzer]

Should the instance of CustomFunctionPyOperator be private? I see it being called as part of "apply', but is there a case when someone would someone call it directly?

[zou3519]

Everything inside torch/_functorch/autograd_function.py is private, including custom_function_call. Users should not be calling it directly.

Developers should interact with custom_function_call directly, instead of CustomFunctionPyOperator if that is the question. custom_function_call is an instance of a PyOperator, CustomFunctionPyOperator is just a class.

[zou3519]

I would prefer to handle this in a follow-up:

• due to the feature flag this is not exercised in the fast-path
• It might involve a refactor of torch/_ops.py

[zou3519]

There appear to be some issues with dynamo, but other than that, this should be ready for review

[soulitzer]

Pretty cool!

[soulitzer]

Should the instance of CustomFunctionPyOperator be private? I see it being called as part of "apply', but is there a case when someone would someone call it directly?

[soulitzer]

When does this happen? Why didn't we need it before?

[zou3519]

This is necessary to support functorch.vjp. functorch.vjp constructs a graph using GradWrapperTensors which then all become dead tensor wrappers.

In the C++ the functorch grad transform always unwraps dead tensor wrappers (

pytorch/aten/src/ATen/functorch/ADInterpreters.cpp

Lines 60 to 68 in a6caa9c

	// if is a grad transform, and the operation is in-place, and the mutated
	// argument is not currently wrapped in a TensorWrapper, then we need to
	// error out otherwise the result is silently incorrect
	checkForInvalidMutationOnCaptures(op, stack, current_level);
	// materialize live GradWrappers
	auto maybeTransformGradWrappers = [&](const Tensor& tensor) {
	return materializeGradWrappers(tensor, current_level);
	};

, it happens inside checkForInvalidMutationOnCaptures).

In Python, we want to be consistent with that behavior (

pytorch/torch/_functorch/pyfunctorch.py

Lines 107 to 110 in a6caa9c

	def process(self, op, args, kwargs):
	kernel = op.functorch_table[TransformType.Grad]
	args, kwargs = self.lift(args, kwargs)
	return kernel(self, *args, **kwargs)

). pyfunctorch today just materializes the wrappers but does not actually unwrap dead wrappers, so this PR fixes it.

Previously we didn't need it because I didn't actually test functorch.vjp; this PR exercises the functorch vjp tests and so we hit this case.

[zou3519]

Should the instance of CustomFunctionPyOperator be private? I see it being called as part of "apply', but is there a case when someone would someone call it directly?

I can't comment on this directly for some reason, but here goes:

Everything in the torch/_functorch/autograd_function.py file is private. Users should not be calling custom_function_call or CustomFunctionPyOperator directly. autograd.Function invokes custom_function_call as an implementation detail.

EDIT: I commented directly on the thread.

[pytorch-bot]

❌ 🤖 pytorchbot command failed:

@pytorchbot merge: error: argument -f/--force: expected one argument

usage: @pytorchbot merge [-g | -f MESSAGE | -l] [-r [{viable/strict,master}]]

Try @pytorchbot --help for more info.

[zou3519]

@pytorchbot merge -f "all test failures look unrelated (they didn't fail on the topmost PR)

[pytorch-bot]

❌ 🤖 pytorchbot command failed:

Got EOF while in a quoted string```
Try `@pytorchbot --help` for more info.

[zou3519]

@pytorchbot merge -f "all test failures look unrelated (they didn't fail on the topmost PR)"

[pytorchmergebot]

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[soulitzer]

I think you forgot to define f for this one

[zou3519]

Lol, looks like the expected failure masked this. Thanks for catching.