[custom_op] explicit autograd API #101824

zou3519 · 2023-05-18T21:11:40Z

Stack from ghstack:

This PR adds an explicit API for registering a backward formula for a
CustomOp. In the end state, we will likely have this explicit API and a
magic API (which is sugar on top of an explicit API), since different
parties of users prefer different ones.

Concretely, to define a backward formula for a CustomOp:

a user must provide us a "save for backward" function that accepts
(inputs, output) and returns exactly what they want saved for backward
a user must provide us a "backward" function that accepts
(ctx, saved, *grads) and returns us the grad_inputs. The grad_inputs
are returned as a dict mapping str to a gradient.
Please see the changes in custom_op_db.py for examples of the API.

There are a number of pieces to this PR and I'm happy to split it if it
helps. They are:

The actual APIs for specifying the two functions
(impl_save_for_backward, impl_backward)
The autograd kernel: we take the functions the user give us and
construct an autograd.Function object that we then register to
the Autograd dispatch key
Indirection for the autograd kernel. We add a layer of indirection so
that one can swap out the autograd kernel. This is necessary because by
default, we register an "autograd not implemented" kernel as the
Autograd implementation but then swap it for the actual kernel when the
user provides it.

Test Plan:

We apply this API to give backward formulas for things in
custom_op_db. We then hook up custom_op_db to the Autograd OpInfo tests.
Various tests in test_python_dispatch.py to check error cases.

This PR adds an explicit API for registering a backward formula for a CustomOp. In the end state, we will likely have this explicit API and a magic API (which is sugar on top of an explicit API), since different parties of users prefer different ones. Concretely, to define a backward formula for a CustomOp: - a user must provide us a "save for backward" function that accepts (inputs, output) and returns exactly what they want saved for backward - a user must provide us a "backward" function that accepts (ctx, saved, *grads) and returns us the grad_inputs. The grad_inputs are returned as a dict mapping str to a gradient. Please see the changes in custom_op_db.py for examples of the API. There are a number of pieces to this PR and I'm happy to split it if it helps. They are: - The actual APIs for specifying the two functions (impl_save_for_backward, impl_backward) - The autograd kernel: we take the functions the user give us and construct an autograd.Function object that we then register to the Autograd dispatch key - Indirection for the autograd kernel. We add a layer of indirection so that one can swap out the autograd kernel. This is necessary because by default, we register an "autograd not implemented" kernel as the Autograd implementation but then swap it for the actual kernel when the user provides it. Test Plan: - We apply this API to give backward formulas for things in custom_op_db. We then hook up custom_op_db to the Autograd OpInfo tests. - Various tests in test_python_dispatch.py to check error cases. [ghstack-poisoned]

pytorch-bot · 2023-05-18T21:11:43Z

🔗 Helpful Links

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

📄 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 the bot commands wiki or our office hours

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

✅ No Failures

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

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

This PR adds an explicit API for registering a backward formula for a CustomOp. In the end state, we will likely have this explicit API and a magic API (which is sugar on top of an explicit API), since different parties of users prefer different ones. Concretely, to define a backward formula for a CustomOp: - a user must provide us a "save for backward" function that accepts (inputs, output) and returns exactly what they want saved for backward - a user must provide us a "backward" function that accepts (ctx, saved, *grads) and returns us the grad_inputs. The grad_inputs are returned as a dict mapping str to a gradient. Please see the changes in custom_op_db.py for examples of the API. There are a number of pieces to this PR and I'm happy to split it if it helps. They are: - The actual APIs for specifying the two functions (impl_save_for_backward, impl_backward) - The autograd kernel: we take the functions the user give us and construct an autograd.Function object that we then register to the Autograd dispatch key - Indirection for the autograd kernel. We add a layer of indirection so that one can swap out the autograd kernel. This is necessary because by default, we register an "autograd not implemented" kernel as the Autograd implementation but then swap it for the actual kernel when the user provides it. Test Plan: - We apply this API to give backward formulas for things in custom_op_db. We then hook up custom_op_db to the Autograd OpInfo tests. - Various tests in test_python_dispatch.py to check error cases. ghstack-source-id: 55ff699 Pull Request resolved: #101824

torch/_custom_op/autograd.py

ezyang

LGTM. But perhaps you want to have one of the autograd experts take a look over.

ezyang

LGTM. But perhaps you want to have one of the autograd experts take a look over.

torch/_custom_op/autograd.py

albanD · 2023-05-22T10:21:38Z

torch/_custom_op/autograd.py

+            return grad_inputs_dict_to_flat_tuple(grad_inputs_dict, args_info)
+
+        generated_cls = gen_autograd_function(
+            forward_op._opname + 'CustomOp', forward, backward)


nit: lower case vs camel case?

torch/_custom_op/autograd.py

torch/_custom_op/impl.py

albanD · 2023-05-22T11:32:45Z

torch/testing/_internal/custom_op_db.py

+
+@numpy_mul.impl_backward()
+def numpy_mul_backward(ctx, saved, grad_out):
+    grad_x = grad_out * saved['y'] if saved['x_requires_grad'] else None


The ctx should be telling you what to compute gradients for really.
This is fine as a test but we should not recommend users do this. We should just expose needs_input_grad on the ctx.

Agree, I didn't add needs_input_grad in this PR because it was getting long.

I can add it as a follow-up, unless you want to see it in this PR?

No need to add it here as long as this is not in any user-facing example!

albanD · 2023-05-22T11:34:52Z

torch/_custom_op/autograd.py

+            save_for_backward_fn_inputs = namedtuple_args(schema, args)
+            to_save = save_for_backward_fn(save_for_backward_fn_inputs, output)
+
+            save_pytree_for_backward(ctx, (to_save, args_info))


nit we should be able to skip this when grad mode is disabled?

It is a little complicated to do this because autograd.Function's forward sets grad_mode to False. This means that we would need to shepherd the knowledge of if grad mode is disabled or not from somewhere else.

Since this is a nit I'm just going to add this to the wishlist.

torch/_custom_op/autograd.py

test/test_python_dispatch.py

torch/_custom_op/autograd.py

test/test_python_dispatch.py

torch/_custom_op/autograd.py

zou3519

.

This PR adds an explicit API for registering a backward formula for a CustomOp. In the end state, we will likely have this explicit API and a magic API (which is sugar on top of an explicit API), since different parties of users prefer different ones. Concretely, to define a backward formula for a CustomOp: - a user must provide us a "save for backward" function that accepts (inputs, output) and returns exactly what they want saved for backward - a user must provide us a "backward" function that accepts (ctx, saved, *grads) and returns us the grad_inputs. The grad_inputs are returned as a dict mapping str to a gradient. Please see the changes in custom_op_db.py for examples of the API. There are a number of pieces to this PR and I'm happy to split it if it helps. They are: - The actual APIs for specifying the two functions (impl_save_for_backward, impl_backward) - The autograd kernel: we take the functions the user give us and construct an autograd.Function object that we then register to the Autograd dispatch key - Indirection for the autograd kernel. We add a layer of indirection so that one can swap out the autograd kernel. This is necessary because by default, we register an "autograd not implemented" kernel as the Autograd implementation but then swap it for the actual kernel when the user provides it. Test Plan: - We apply this API to give backward formulas for things in custom_op_db. We then hook up custom_op_db to the Autograd OpInfo tests. - Various tests in test_python_dispatch.py to check error cases. [ghstack-poisoned]

This PR adds an explicit API for registering a backward formula for a CustomOp. In the end state, we will likely have this explicit API and a magic API (which is sugar on top of an explicit API), since different parties of users prefer different ones. Concretely, to define a backward formula for a CustomOp: - a user must provide us a "save for backward" function that accepts (inputs, output) and returns exactly what they want saved for backward - a user must provide us a "backward" function that accepts (ctx, saved, *grads) and returns us the grad_inputs. The grad_inputs are returned as a dict mapping str to a gradient. Please see the changes in custom_op_db.py for examples of the API. There are a number of pieces to this PR and I'm happy to split it if it helps. They are: - The actual APIs for specifying the two functions (impl_save_for_backward, impl_backward) - The autograd kernel: we take the functions the user give us and construct an autograd.Function object that we then register to the Autograd dispatch key - Indirection for the autograd kernel. We add a layer of indirection so that one can swap out the autograd kernel. This is necessary because by default, we register an "autograd not implemented" kernel as the Autograd implementation but then swap it for the actual kernel when the user provides it. Test Plan: - We apply this API to give backward formulas for things in custom_op_db. We then hook up custom_op_db to the Autograd OpInfo tests. - Various tests in test_python_dispatch.py to check error cases. ghstack-source-id: 553a1ca Pull Request resolved: #101824

This PR adds an explicit API for registering a backward formula for a CustomOp. In the end state, we will likely have this explicit API and a magic API (which is sugar on top of an explicit API), since different parties of users prefer different ones. Concretely, to define a backward formula for a CustomOp: - a user must provide us a "save for backward" function that accepts (inputs, output) and returns exactly what they want saved for backward - a user must provide us a "backward" function that accepts (ctx, saved, *grads) and returns us the grad_inputs. The grad_inputs are returned as a dict mapping str to a gradient. Please see the changes in custom_op_db.py for examples of the API. There are a number of pieces to this PR and I'm happy to split it if it helps. They are: - The actual APIs for specifying the two functions (impl_save_for_backward, impl_backward) - The autograd kernel: we take the functions the user give us and construct an autograd.Function object that we then register to the Autograd dispatch key - Indirection for the autograd kernel. We add a layer of indirection so that one can swap out the autograd kernel. This is necessary because by default, we register an "autograd not implemented" kernel as the Autograd implementation but then swap it for the actual kernel when the user provides it. Test Plan: - We apply this API to give backward formulas for things in custom_op_db. We then hook up custom_op_db to the Autograd OpInfo tests. - Various tests in test_python_dispatch.py to check error cases. [ghstack-poisoned]

This PR adds an explicit API for registering a backward formula for a CustomOp. In the end state, we will likely have this explicit API and a magic API (which is sugar on top of an explicit API), since different parties of users prefer different ones. Concretely, to define a backward formula for a CustomOp: - a user must provide us a "save for backward" function that accepts (inputs, output) and returns exactly what they want saved for backward - a user must provide us a "backward" function that accepts (ctx, saved, *grads) and returns us the grad_inputs. The grad_inputs are returned as a dict mapping str to a gradient. Please see the changes in custom_op_db.py for examples of the API. There are a number of pieces to this PR and I'm happy to split it if it helps. They are: - The actual APIs for specifying the two functions (impl_save_for_backward, impl_backward) - The autograd kernel: we take the functions the user give us and construct an autograd.Function object that we then register to the Autograd dispatch key - Indirection for the autograd kernel. We add a layer of indirection so that one can swap out the autograd kernel. This is necessary because by default, we register an "autograd not implemented" kernel as the Autograd implementation but then swap it for the actual kernel when the user provides it. Test Plan: - We apply this API to give backward formulas for things in custom_op_db. We then hook up custom_op_db to the Autograd OpInfo tests. - Various tests in test_python_dispatch.py to check error cases. ghstack-source-id: c40cade Pull Request resolved: #101824

zou3519 · 2023-05-23T18:28:50Z

@pytorchbot merge

pytorchmergebot · 2023-05-23T18:31:21Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

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

zou3519 · 2023-05-23T18:33:55Z

torch/_custom_op/impl.py

+        were declared to be Tensors in the CustomOp definition must be accounted
+        for in the dict. The gradient may be a Tensor or None.
+
+        TODO(rzou): Add example when this PR is closer to landing.


Sorry, I remembered this after I typed the merge command. Will do as a follow-up.

vadimkantorov · 2023-06-27T15:01:34Z

Regarding the syntax, could custom_op be combined somehow with forward impl?
currently:

        @custom_op(f'{TestCustomOp.test_ns}::foo')
        def foo(x: torch.Tensor) -> torch.Tensor:
            ...

        @foo.impl(['cpu', 'cuda'])
        def foo_impl(x):
            return x.sin()

        @foo.impl_backward()
        def foo_backward(ctx, saved, grad):
            return grad * saved.cos()

proposed:

        @custom_op_with_forward(f'{TestCustomOp.test_ns}::foo', device_types = ['cpu', 'cuda'])
        def foo(x : torch.Tensor) -> torch.Tensor:
            return x.sin()

        @foo.impl_backward()
        def foo_backward(ctx, saved, grad):
            return grad * saved.cos()

zou3519 requested review from mruberry, ngimel and soulitzer as code owners May 18, 2023 21:11

This was referenced May 18, 2023

[custom_op] Add a test for symints #101822

Closed

[custom_op] Create a new torch._custom_op namespace #101823

Closed

zou3519 requested review from albanD, bdhirsh and ezyang and removed request for mruberry and ngimel May 19, 2023 14:02

ezyang reviewed May 21, 2023

View reviewed changes

torch/_custom_op/autograd.py Outdated Show resolved Hide resolved

ezyang approved these changes May 21, 2023

View reviewed changes

albanD reviewed May 22, 2023

View reviewed changes

soulitzer reviewed May 22, 2023

View reviewed changes

zou3519 commented May 22, 2023

View reviewed changes

zou3519 added ciflow/trunk Trigger trunk jobs on your pull request release notes: composability release notes category labels May 23, 2023

pytorchmergebot added the merging label May 23, 2023

pytorchmergebot added Merged and removed merging labels May 23, 2023

pytorchmergebot closed this in 723f111 May 23, 2023

zou3519 commented May 23, 2023

View reviewed changes

facebook-github-bot deleted the gh/zou3519/663/head branch June 8, 2023 19:36

[custom_op] explicit autograd API #101824

[custom_op] explicit autograd API #101824

Uh oh!

Conversation

zou3519 commented May 18, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented May 18, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

✅ No Failures

Uh oh!

Uh oh!

ezyang left a comment

Choose a reason for hiding this comment

Uh oh!

ezyang left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zou3519 left a comment

Choose a reason for hiding this comment

Uh oh!

zou3519 commented May 23, 2023

Uh oh!

pytorchmergebot commented May 23, 2023

Merge started

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vadimkantorov commented Jun 27, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

zou3519 commented May 18, 2023 •

edited

Loading

pytorch-bot bot commented May 18, 2023 •

edited

Loading