[dynamo] compiled_autograd support for post_acc_grad hooks #112326

voznesenskym · 2023-10-28T20:20:20Z

Stack from ghstack (oldest at bottom):

cc @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng

…_grad hooks [ghstack-poisoned]

pytorch-bot · 2023-10-28T20:20:23Z

🔗 Helpful Links

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

📄 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 f2d3f04 with merge base a7a0955 ():
💚 Looks good so far! There are no failures yet. 💚

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

voznesenskym · 2023-10-28T20:21:35Z

torch/csrc/autograd/functions/accumulate_grad.cpp

-
-  return variable_list();
+  // A little hack - this is only here for the purpose of hooks. It will get cleared.
+  return variable_list({fake_variable_copy});


@jansel surely there's a better way to smuggle stuff out? This does get a fake tensor into the hook (if you set a breakpoint in def post_acc_grad_hook you can see it having the correct value, with grad.

Don't smuggle it out. Just call into a virtual method on tensor_post_acc_grad_hooks().

Eager returns nothing here, compiled autograd should do the same.

voznesenskym · 2023-10-28T20:21:56Z

torch/csrc/dynamo/python_compiled_autograd.cpp

+      if(typeid(*call.node) == typeid(torch::autograd::AccumulateGrad)) {
+        // The return of AccumulateGrad should be [], but we hack it to make hooks work.
+        // This restores it correctly.
+        outputs = variable_list();
+      }


@jansel see comment above. I dislike this.

voznesenskym · 2023-10-28T20:22:27Z

torch/_dynamo/compiled_autograd.py

+        assert len(inputs) == 1
+        hook = self.hooks_proxy[hook_id]


@jansel this is an odd one - if we don't unpack the argument, and leave it to pass in inputs, the hook is invoked with a list, which is incorrect.

jansel · 2023-10-29T02:21:46Z

torch/_dynamo/compiled_autograd.py

+        with disable_proxy_modes_tracing():
+            inputs = maybe_clone(inputs[0])
+            self.bind_tensors_to_proxies([inputs], proxies)
+        return inputs


Does this type of hook actually return anything? The call to it in eager just ignores the result.

jansel · 2023-10-29T02:22:28Z

torch/csrc/autograd/functions/accumulate_grad.cpp

-
-  return variable_list();
+  // A little hack - this is only here for the purpose of hooks. It will get cleared.
+  return variable_list({fake_variable_copy});


Don't smuggle it out. Just call into a virtual method on tensor_post_acc_grad_hooks().

Eager returns nothing here, compiled autograd should do the same.

jansel · 2023-10-29T02:24:38Z

torch/csrc/dynamo/python_compiled_autograd.cpp


      SwapSavedVariables saved(compiler_call, state);
      variable_list outputs = call.node->apply_with_saved(inputs, saved);
+      if (!call.post_acc_grad_hooks.empty()) {


Don't do this here, in eager it is called in accumulate_grad.cpp. We should do the same.

Whats the right way to invoke the python fn here?

Im doing this outside because I couldn't figure out the plumbing. I also tried to do it as an inductor op, but I dont think that supports hook signatures.

Just follow the same pattern as the other hooks.

I thought that's what I did here. post and pre hook are called above and below. No other hook is called from a AccumulateGrad or from within a node from what I saw? I don't understand how to call a python hook from within the accumulate_grad implementation without plumbing in a bunch of python notions into it?

Even if within apply_with_saved within AccumulateGrad we grab the hook and call it, you still need to pass a py_compiler fromthe_autograd_compiler, right? I didn't want to start mucking around with the apply_with_saved interface. Would you mind just giving me a little more detail on how you would do it here?

Like, we could skip the py stuff and just do

auto& hook = tensor_post_acc_grad_hooks(); if (hook != nullptr) { (*hook)(variable); }

But then it will bypass all the proxy and fake tensor stuff, which I don't think is what we want? Unless you want me to invoke the above with the fake tensor between before/after...?

Well, I guess this works, but not sure if this is correct. I'll PR what I have and we can discuss it there.

…or post_acc_grad hooks" cc penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]

ghstack-source-id: d7014b1 Pull Request resolved: #112326 [dynamo][wip][not working yet] compiled_autograd support for post_acc_grad hooks

…or post_acc_grad hooks" cc penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]

jansel · 2023-10-30T17:52:48Z

torch/csrc/autograd/functions/accumulate_grad.cpp

+  auto& hook = tensor_post_acc_grad_hooks();
+  if (hook != nullptr) {
+    (*hook)(variable_copy);
+  }


This seems wrong. It will just trace through the hook (which might not be tracable).

Need something like:

hook->apply_with_saved(variable_copy, saved)

Then inside the handler for apply_with_saved the hooks is lifted to an input.

Asked offline, I do not understand.

torch/csrc/dynamo/compiled_autograd.h

cc penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]

voznesenskym · 2023-10-30T22:22:20Z

test/dynamo/test_hooks.py

-                def hook(p):
-                    p.add_(p.grad)
+                def hook(input_t):
+                    input_t.mul_(2)


cc penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]

jansel

I'd expect this to work at this point, though see minor comment.

torch/csrc/autograd/python_hook.h

jansel · 2023-10-31T00:21:25Z

torch/csrc/dynamo/compiled_autograd.h


  AutogradCompilerCall& compiler;
  TraceState& state;
+  PyObject* py_compiler;


Add comment about borrowed(?) ownership.

yes, borrowed.

voznesenskym · 2023-10-31T01:31:10Z

I'd expect this to work at this point, though see minor comment.

yes, it works.

cc penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]

voznesenskym · 2023-10-31T22:49:59Z

@pytorchbot merge

pytorchmergebot · 2023-10-31T22:52:36Z

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

…12326) Pull Request resolved: pytorch#112326 Approved by: https://github.com/jansel ghstack dependencies: pytorch#112325

[dynamo][wip][not working yet] compiled_autograd support for post_acc…

54df4f4

…_grad hooks [ghstack-poisoned]

voznesenskym requested review from albanD and soulitzer as code owners October 28, 2023 20:20

This was referenced Oct 28, 2023

[fsdp] Replace acc_grad hooking with register_post_accumulate_grad_hook on flat_param #112184

Closed

[dynamo] Add support for register_post_accumulate_grad_hook #112325

Closed

github-actions bot requested review from SherlockNoMad, antoniojkim, bdhirsh, ezyang, miladm and wconstab October 28, 2023 20:20

github-actions bot added module: dynamo ciflow/inductor labels Oct 28, 2023

voznesenskym commented Oct 28, 2023

View reviewed changes

jansel requested changes Oct 29, 2023

View reviewed changes

Update on "[dynamo][wip][not working yet] compiled_autograd support f…

dcdb48a

…or post_acc_grad hooks" cc penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]

voznesenskym added a commit that referenced this pull request Oct 29, 2023

[dynamo] compiled_autograd support for post_acc_grad hooks

ac36f79

ghstack-source-id: d7014b1 Pull Request resolved: #112326 [dynamo][wip][not working yet] compiled_autograd support for post_acc_grad hooks

voznesenskym requested a review from jansel October 29, 2023 18:24

Update on "[dynamo][wip][not working yet] compiled_autograd support f…

f47318a

…or post_acc_grad hooks" cc penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]

voznesenskym changed the title ~~[dynamo][wip][not working yet] compiled_autograd support for post_acc_grad hooks~~ [dynamo] compiled_autograd support for post_acc_grad hooks Oct 30, 2023

jansel requested changes Oct 30, 2023

View reviewed changes

Update on "[dynamo] compiled_autograd support for post_acc_grad hooks"

50340c2

cc penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]

voznesenskym commented Oct 30, 2023

View reviewed changes

test/dynamo/test_hooks.py

def hook(p):

p.add_(p.grad)

def hook(input_t):

input_t.mul_(2)

Copy link

Collaborator Author

voznesenskym Oct 30, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#112446

Update on "[dynamo] compiled_autograd support for post_acc_grad hooks"

123fb93

cc penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]

voznesenskym requested a review from jansel October 30, 2023 22:25

jansel requested changes Oct 31, 2023

View reviewed changes

Update on "[dynamo] compiled_autograd support for post_acc_grad hooks"

f2d3f04

cc penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]

voznesenskym requested a review from jansel October 31, 2023 01:52

jansel approved these changes Oct 31, 2023

View reviewed changes

voznesenskym added ciflow/trunk Trigger trunk jobs on your pull request topic: not user facing topic category labels Oct 31, 2023

pytorchmergebot added the merging label Oct 31, 2023

pytorchmergebot added the Merged label Oct 31, 2023

pytorchmergebot removed the merging label Oct 31, 2023

pytorchmergebot closed this in 0f4d290 Oct 31, 2023

facebook-github-bot deleted the gh/voznesenskym/258/head branch November 4, 2023 14:26

xuhancn pushed a commit to xuhancn/pytorch that referenced this pull request Nov 7, 2023

[dynamo] compiled_autograd support for post_acc_grad hooks (pytorch#1…

1ab522e

…12326) Pull Request resolved: pytorch#112326 Approved by: https://github.com/jansel ghstack dependencies: pytorch#112325

[dynamo] compiled_autograd support for post_acc_grad hooks #112326

[dynamo] compiled_autograd support for post_acc_grad hooks #112326

Uh oh!

Conversation

voznesenskym commented Oct 28, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Oct 28, 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/112326

✅ No Failures

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!

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!

Choose a reason for hiding this comment

Uh oh!

jansel left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

voznesenskym commented Oct 31, 2023

Uh oh!

voznesenskym commented Oct 31, 2023

Uh oh!

pytorchmergebot commented Oct 31, 2023

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

voznesenskym commented Oct 28, 2023 •

edited

Loading

pytorch-bot bot commented Oct 28, 2023 •

edited

Loading