[dynamo] Properly prune dead cell local variables #136891

StrongerXi · 2024-09-27T19:45:10Z

This patch updates the prune_dead_locals logic to do slightly more aggressive pruning for cell local variables, in absence of side-effects, e.g., a cell variable can be pruned when its user function(s) will never be used again.

See added tests for examples; note that a few tests in test/dynamo/test_higher_order_ops.py also got updated because we are no longer returning the unnecessary graph output.

Fixes #127350, #124653

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @rec

pytorch-bot · 2024-09-27T19:45:14Z

🔗 Helpful Links

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

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

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

linux-foundation-easycla · 2024-09-27T19:45:14Z

The committers listed above are authorized under a signed CLA.

✅ login: StrongerXi / name: Xiangxi Guo (Ryan) (42de084)

StrongerXi · 2024-09-27T19:52:53Z

/easycla

StrongerXi · 2024-09-27T21:26:37Z

torch/_dynamo/symbolic_convert.py

+        # 2. In the fast path of `OutputGraph::compile_subgraph`, we populate
+        #    the tuple of cells _after_ emitting the `MAKE FUNCTION` bytecode,
+        #    via `STORE DEREF`; these `STORE DEREF` are generated partly based
+        #    on the current `symbolic_locals`.


I thought about another approach where we emit STORE_DEREF as part of NestedUserFunctionVariable::reconstruct, and get rid of the special handling of cellvars here; I might look into that more if I can't fix the existing test failures.

I'm starting to think there might be a better solution with this approach, by replacing uses of ClosureVariable with NewCellVariable. It feels related to #136814 as well.

StrongerXi · 2024-09-27T23:19:20Z

Fix side-effect related issues -- Dynamo must treat captured variables as live (therefore generate bytecode to update these cells), if any of their user functions escapes the lifetime of the function that defined the captured variables.

StrongerXi · 2024-09-28T01:03:51Z

Run SideEffects::prune_dead_object_new first, which might enable more cellvar pruning. This undoes the limiting effect from the last push.

jansel

This looks like a good change, but tests are failing

StrongerXi · 2024-09-30T15:44:11Z

torch/_dynamo/symbolic_convert.py

+            # Manually release the self-referential nested function, which
+            # captures `self.symbolic_locals` and affects parts of PT test/logic
+            # that are sensitive to when certain objects get released.
+            del visit


Attempt to fix failing tests; this is somewhat surprising, or maybe expected?

This is surprising to me - which tests are failing because of this? The memleak tests?

These ones, I assume self.symbolic_locals had a reference to x, so before we run the full gc.collect(), x won't be freed (tests do pass if we add a gc.collect() after del x).
https://github.com/pytorch/pytorch/blob/54e2b2a6d4f6b6d27637b941f35f254ef42edbad/test/dynamo/test_misc.py#L4753-L4797

Somewhat similar to the above, but more behind-the-scene logic
https://github.com/pytorch/pytorch/blob/54e2b2a6d4f6b6d27637b941f35f254ef42edbad/test/dynamo/test_misc.py#L10756-L10766

This one, which can't seem to be fixed by adding a gc.collect() in the test itself.
https://github.com/pytorch/pytorch/blob/54e2b2a6d4f6b6d27637b941f35f254ef42edbad/test/test_autograd.py#L5464-L5465

What exactly is the reference cycle being caused here?

Quite subtle:

Inside this frame, visit is a cell variable, because visit as a nested function references and captures itself. So you have visit.__closure__[0] being a cell that references visit. Example:

def foo(): def bar(): return bar() print(bar.__closure__[0].cell_contents is bar) foo() # prints True

This is annoying but I don't have any ideas.

Were you able to try any of the other methods that we discussed on Monday to prevent the refcycle without using del?

StrongerXi · 2024-09-30T18:26:55Z

torch/_dynamo/symbolic_convert.py

+            # NOTE: Don't trace from the cell locals which aren't explicitly
+            # read anymore; if they are indirectly used, they will be reached by
+            # other roots. These initially excluded cells are the ones that will
+            # hopefully be pruned.


Another fix.

StrongerXi · 2024-09-30T18:40:07Z

Rebase to fix a inductor failure which seems unrelated to my change.

StrongerXi · 2024-09-30T22:55:52Z

@jansel I just reproduced the failed test inductor/test_group_batch_fusion.py::TestGroupBatchFusion::test_pointwise_op_fusion on my gpu devserver, on latest main c07ebaf, I think it's not related to my patch? It also looks like a inductor-only test.

zou3519 · 2024-10-01T17:53:09Z

@StrongerXi if the test isn't related -- try to rebase? Otherwise you'll need to force merge this when it's been approved

zou3519 · 2024-10-01T17:55:10Z

test/dynamo/test_misc.py

+
+            # Returning `func` forces dynamo to output `x` in the compiled graph, so
+            # that we can store it as `func`'s closure.
+            return func, func()


the outputs are (x, x + x) ?

Hmm (func, x + x) I'd say?

Oh I meant, what are the 2 outputs in the compiled graph?

Ah, yep. Lemme document that.

facebook-github-bot · 2024-10-02T19:19:33Z

@StrongerXi has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

StrongerXi · 2024-10-02T23:27:09Z

Chatted with @yanboliang, I'll update the test, which itself isn't super important. The main concern here is this patch breaks brittle patterns other parts of the stack rely on, and their effect on our internal workloads. I'll also hold off "fixing" the fusion pass, for the same reason.

Another 1 line fix in test/inductor/test_group_batch_fusion.py, for similar reason as above.

facebook-github-bot · 2024-10-02T23:28:13Z

@StrongerXi has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

StrongerXi · 2024-10-03T16:11:12Z

All green, but I'm still waiting for some internal experiments to guard against potential regression.

StrongerXi · 2024-10-04T22:37:49Z

Need to investigate compile-time increase and qps drop for internal workflow.

As title.

StrongerXi · 2024-10-07T22:11:18Z

Still investigating. Update documentation about why we skip InlinedClosureVariable in VariableTracker visiting.

StrongerXi · 2024-10-09T17:50:16Z

Internal workflow passed (previous results turned out to be noise). Remove inaccurate documentation on why we skip InlinedClosureVariable in search for must-keep-alive cells, now that #137510 is about to land.

facebook-github-bot · 2024-10-09T17:54:12Z

@StrongerXi has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

This patch updates the `prune_dead_locals` logic to do slightly more aggressive pruning for cell local variables, in absence of side-effects, e.g., a cell variable can be pruned when its user function(s) will never be used again. See added tests for examples; note that a few tests in `test/dynamo/test_higher_order_ops.py` also got updated because we are no longer returning the unnecessary graph output. Fixes pytorch#127350, pytorch#124653

StrongerXi · 2024-10-10T00:15:41Z

Chatted with @yanboliang, I'll update the test, which itself isn't super important. The main concern here is this patch breaks brittle patterns other parts of the stack rely on, and their effect on our internal workloads. I'll also hold off "fixing" the fusion pass, for the same reason.

Another 1 line fix in test/inductor/test_group_batch_fusion.py, for similar reason as above.

Another 1 line fix of this sort, in a test that requires fbgemm.

facebook-github-bot · 2024-10-10T00:16:13Z

@StrongerXi has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

StrongerXi · 2024-10-10T15:54:18Z

All green externally and internally. Merging.

StrongerXi · 2024-10-10T15:54:37Z

All green externally and internally. Merging.

StrongerXi · 2024-10-10T15:55:09Z

@pytorchbot merge

pytorchmergebot · 2024-10-10T15:57:01Z

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

pytorch-bot bot added ciflow/inductor module: dynamo labels Sep 27, 2024

StrongerXi requested review from anijain2305, jansel and williamwen42 and removed request for williamwen42 September 27, 2024 19:45

StrongerXi linked an issue Sep 27, 2024 that may be closed by this pull request

dynamo creates unnecessary buffers #124653

Closed

StrongerXi added the topic: not user facing topic category label Sep 27, 2024

StrongerXi commented Sep 27, 2024

View reviewed changes

StrongerXi force-pushed the fix-closure-var branch 2 times, most recently from 37a7f44 to d75022e Compare September 27, 2024 23:15

StrongerXi requested a review from zou3519 September 27, 2024 23:23

StrongerXi force-pushed the fix-closure-var branch from d75022e to 08b1c20 Compare September 28, 2024 01:02

jansel requested changes Sep 28, 2024

View reviewed changes

StrongerXi force-pushed the fix-closure-var branch from 08b1c20 to b7e4ecb Compare September 30, 2024 15:43

StrongerXi commented Sep 30, 2024

View reviewed changes

StrongerXi force-pushed the fix-closure-var branch from b7e4ecb to 54e2b2a Compare September 30, 2024 18:26

StrongerXi commented Sep 30, 2024

View reviewed changes

StrongerXi force-pushed the fix-closure-var branch from 54e2b2a to 43a1d59 Compare September 30, 2024 18:39

StrongerXi requested a review from jansel September 30, 2024 22:55

zou3519 reviewed Oct 1, 2024

View reviewed changes

zou3519 approved these changes Oct 2, 2024

View reviewed changes

StrongerXi force-pushed the fix-closure-var branch from f290a99 to 7020b9c Compare October 2, 2024 23:26

StrongerXi added a commit that referenced this pull request Oct 7, 2024

[DO NOT LAND] PR to trigger benchmark CI for #136891

238b2dc

As title.

StrongerXi force-pushed the fix-closure-var branch from 7020b9c to c48469e Compare October 7, 2024 22:10

StrongerXi force-pushed the fix-closure-var branch from c48469e to 2dc6537 Compare October 9, 2024 17:46

StrongerXi force-pushed the fix-closure-var branch from 2dc6537 to 42de084 Compare October 10, 2024 00:13

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 10, 2024

pytorchmergebot added the merging label Oct 10, 2024

pytorchmergebot added the Merged label Oct 10, 2024

pytorchmergebot closed this in dd7c289 Oct 10, 2024

pytorchmergebot removed the merging label Oct 10, 2024

This was referenced Oct 14, 2024

[dynamo] Remove mutable_local.source and index on VariableTracker rather than MutableLocalBase #137905

Closed

dynamo creates unnecessary buffers #124653

Closed

[dynamo] Properly prune dead cell local variables #136891

[dynamo] Properly prune dead cell local variables #136891

Uh oh!

Conversation

StrongerXi commented Sep 27, 2024 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Sep 27, 2024 • 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/136891

✅ No Failures

Uh oh!

linux-foundation-easycla bot commented Sep 27, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

StrongerXi commented Sep 27, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

StrongerXi commented Sep 27, 2024

Uh oh!

StrongerXi commented Sep 28, 2024

Uh oh!

jansel left a comment

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!

StrongerXi commented Sep 30, 2024

Uh oh!

StrongerXi commented Sep 30, 2024

Uh oh!

zou3519 commented Oct 1, 2024

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!

facebook-github-bot commented Oct 2, 2024

Uh oh!

StrongerXi commented Oct 2, 2024

Uh oh!

facebook-github-bot commented Oct 2, 2024

Uh oh!

StrongerXi commented Oct 3, 2024

Uh oh!

StrongerXi commented Oct 4, 2024

Uh oh!

StrongerXi commented Oct 7, 2024

Uh oh!

StrongerXi commented Oct 9, 2024

Uh oh!

facebook-github-bot commented Oct 9, 2024

Uh oh!

StrongerXi commented Oct 10, 2024

Uh oh!

facebook-github-bot commented Oct 10, 2024

Uh oh!

StrongerXi commented Oct 10, 2024

StrongerXi commented Sep 27, 2024 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Sep 27, 2024 •

edited

Loading

linux-foundation-easycla bot commented Sep 27, 2024 •

edited

Loading