[pipelining] Fix more leaks and check leaks in tests #136584

wconstab · 2024-09-24T23:46:37Z

Stack from ghstack (oldest at bottom):

Fix two more leaks of the same variety as #136507 (see that PR desc and attached gdoc for debug details).

This time, also add a test-time check that helped to discover new leaks and ensure we won't accidently regress.

Adds check_tensor_leak util which internally asserts no tensors are being kept alive by other objects involved in py ref cycles.

Uses objgraph for a nice debug utility when a leak is found.

Credit to @H-Huang for pointing out objdump and helping debug the 'param_group["intermediates"]` leak.

I manually confirmed that all 3 of the leaks identified/fixed so far are caught by the unit test and checker.

Sample output, if I re-introduce a leak by commenting out del param_group["intermediates"] in _backward.py,
and run python test/distributed/pipelining/test_schedule_multiproc.py -k test_schedule_with_native_zero_bubble:

warnings.warn(
/data/users/whc/pytorch/torch/testing/_internal/common_utils.py:5341: UserWarning: 34 tensors were found in the garbage. Did you introduce a reference cycle?
warnings.warn(
/data/users/whc/pytorch/torch/testing/_internal/common_utils.py:5347: UserWarning: Dumping first 1 objgraphs of leaked tensors rendered to png
Graph written to /tmp/objgraph-ztz642h3.dot (19 nodes)
Graph viewer (xdot) not found, generating a png instead
Image generated as /tmp/objgraph-ztz642h3.png

rendering of /tmp/objgraph-ztz642h3.png:

cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @d4l3k @c-p-i-o

[ghstack-poisoned]

pytorch-bot · 2024-09-24T23:46:41Z

🔗 Helpful Links

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

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit 38426f0 with merge base 9992084 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

pull / before-test / llm-retrieval (gh) (matched llm-retrieval rule in flaky-rules.json)
The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.

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

wconstab · 2024-09-25T00:32:21Z

torch/testing/_internal/common_utils.py

+
+        gc.collect()
+        garbage_tensors = sum(int(isinstance(g, torch.Tensor)) for g in gc.garbage)
+        if garbage_tensors > 0:


oops, when i refactored this to a util, i dropped the self.assertFalse part so the tests won't actually fail anymore. need to fix that...

kwen2501

Thanks for the great tool!

H-Huang

Great addition!

H-Huang · 2024-09-25T16:04:38Z

torch/testing/_internal/common_utils.py

+        gc.set_debug(gc.DEBUG_SAVEALL)
+
+        # run the user code, after cleaning any existing refcycles
+        yield


For testing mentioned below or for better observability you could yield a Result class or something similar that has attributes like has_leak or leaked_tensors and then in the context manager use it like

with check_leaked_tensors() as result: ... code assert not result.has_leak print(result.leaked_tensors)

oh, i worked on the fix before i saw your comment. i ended up just yielding the garbage_objs list directly.

[ghstack-poisoned]

H-Huang · 2024-09-25T17:36:57Z

torch/distributed/pipelining/_backward.py

+        # `hook` -> cell -> param_group -> intermediates -> `hook`
+        # becuase we install the hook function onto each of the intermediate autograd nodes.
+        # We need to keep intermediates alive up until backward_weight, but we can free it now.
+        # del param_group["intermediates"]


accidentally commented this?

Thanks! fixed.

[ghstack-poisoned]

wconstab · 2024-09-25T17:53:21Z

@pytorchbot merge

pytorchmergebot · 2024-09-25T17:55:17Z

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

pytorchmergebot · 2024-09-25T18:54:41Z

Merge failed

Reason: 13 jobs have failed, first few of them are: pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 2, 5, lf.linux.4xlarge.nvidia.gpu), pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 1, 5, lf.linux.g5.4xlarge.nvidia.gpu), pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 2, 5, lf.linux.g5.4xlarge.nvidia.gpu), pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 3, 5, lf.linux.g5.4xlarge.nvidia.gpu), pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 5, 5, lf.linux.g5.4xlarge.nvidia.gpu)

Details for Dev Infra team

Raised by workflow job

[ghstack-poisoned]

Fix two more leaks of the same variety as #136507 (see that PR desc and attached gdoc for debug details). This time, also add a test-time check that helped to discover new leaks and ensure we won't accidently regress. Adds `check_tensor_leak` util which internally asserts no tensors are being kept alive by other objects involved in py ref cycles. Uses objgraph for a nice debug utility when a leak is found. Credit to H-Huang for pointing out objdump and helping debug the 'param_group["intermediates"]` leak. I manually confirmed that all 3 of the leaks identified/fixed so far are caught by the unit test and checker. Sample output, if I re-introduce a leak by commenting out `del param_group["intermediates"]` in _backward.py, and run `python test/distributed/pipelining/test_schedule_multiproc.py -k test_schedule_with_native_zero_bubble`: ``` warnings.warn( /data/users/whc/pytorch/torch/testing/_internal/common_utils.py:5341: UserWarning: 34 tensors were found in the garbage. Did you introduce a reference cycle? warnings.warn( /data/users/whc/pytorch/torch/testing/_internal/common_utils.py:5347: UserWarning: Dumping first 1 objgraphs of leaked tensors rendered to png Graph written to /tmp/objgraph-ztz642h3.dot (19 nodes) Graph viewer (xdot) not found, generating a png instead Image generated as /tmp/objgraph-ztz642h3.png ``` rendering of ` /tmp/objgraph-ztz642h3.png`: <img width="1671" alt="image" src="https://github.com/user-attachments/assets/9098ff29-224c-4533-935b-83c210ac2e22"> cc XilunWu H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 d4l3k c-p-i-o [ghstack-poisoned]

H-Huang · 2024-09-25T21:10:24Z

@pytorchbot merge

pytorchmergebot · 2024-09-25T21:12:07Z

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

pytorchmergebot · 2024-09-25T21:17:21Z

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

pull / before-test / llm-retrieval

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

H-Huang · 2024-09-26T00:02:16Z

@pytorchbot merge -i

pytorchmergebot · 2024-09-26T00:03:58Z

Merge started

Your change will be merged while ignoring the following 1 checks: pull / before-test / llm-retrieval

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

Pull Request resolved: #136678 Approved by: https://github.com/wconstab, https://github.com/kwen2501 ghstack dependencies: #136507, #136584

Fix two more leaks of the same variety as #136507 (see that PR desc and attached gdoc for debug details). This time, also add a test-time check that helped to discover new leaks and ensure we won't accidently regress. Adds `check_tensor_leak` util which internally asserts no tensors are being kept alive by other objects involved in py ref cycles. Uses objgraph for a nice debug utility when a leak is found. Credit to H-Huang for pointing out objdump and helping debug the 'param_group["intermediates"]` leak. Sample output, if I comment out `del param_group["intermediates"]` in _backward.py, and run `python test/distributed/pipelining/test_schedule_multiproc.py -k test_schedule_with_native_zero_bubble`: ``` warnings.warn( /data/users/whc/pytorch/torch/testing/_internal/common_utils.py:5341: UserWarning: 34 tensors were found in the garbage. Did you introduce a reference cycle? warnings.warn( /data/users/whc/pytorch/torch/testing/_internal/common_utils.py:5347: UserWarning: Dumping first 1 objgraphs of leaked tensors rendered to png Graph written to /tmp/objgraph-ztz642h3.dot (19 nodes) Graph viewer (xdot) not found, generating a png instead Image generated as /tmp/objgraph-ztz642h3.png ``` ghstack-source-id: 1132d0f Pull Request resolved: pytorch/pytorch#136584

Update

0671997

[ghstack-poisoned]

wconstab requested a review from a team as a code owner September 24, 2024 23:46

pytorch-bot bot added the oncall: distributed Add this issue/PR to distributed oncall triage queue label Sep 24, 2024

wconstab added the release notes: distributed (pipeline) release notes category label Sep 24, 2024

wconstab requested review from H-Huang and kwen2501 September 24, 2024 23:55

wconstab commented Sep 25, 2024

View reviewed changes

kwen2501 approved these changes Sep 25, 2024

View reviewed changes

H-Huang approved these changes Sep 25, 2024

View reviewed changes

Update

3dd8bb1

[ghstack-poisoned]

H-Huang reviewed Sep 25, 2024

View reviewed changes

Update

47fdca3

[ghstack-poisoned]

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

pytorchmergebot added the merging label Sep 25, 2024

pytorchmergebot removed the merging label Sep 25, 2024

Update

39e6605

[ghstack-poisoned]

H-Huang mentioned this pull request Sep 25, 2024

[PP] Remove modifications to autograd nodes in ZB #136678

Closed

pytorchmergebot added the merging label Sep 25, 2024

pytorchmergebot removed the merging label Sep 25, 2024

pytorchmergebot added the merging label Sep 26, 2024

pytorchmergebot added the Merged label Sep 26, 2024

pytorchmergebot closed this in 141cae2 Sep 26, 2024

pytorchmergebot removed the merging label Sep 26, 2024

pytorchmergebot pushed a commit that referenced this pull request Sep 27, 2024

[PP] Remove modifications to autograd nodes in ZB (#136678)

ad6c70b

Pull Request resolved: #136678 Approved by: https://github.com/wconstab, https://github.com/kwen2501 ghstack dependencies: #136507, #136584

github-actions bot deleted the gh/wconstab/337/head branch October 27, 2024 02:11

[pipelining] Fix more leaks and check leaks in tests #136584

[pipelining] Fix more leaks and check leaks in tests #136584

Uh oh!

Conversation

wconstab commented Sep 24, 2024 • edited by H-Huang Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Sep 24, 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/136584

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

wconstab Sep 25, 2024

Choose a reason for hiding this comment

Uh oh!

kwen2501 left a comment

Choose a reason for hiding this comment

Uh oh!

H-Huang left a comment

Choose a reason for hiding this comment

Uh oh!

H-Huang Sep 25, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wconstab Sep 25, 2024

Choose a reason for hiding this comment

Uh oh!

H-Huang Sep 25, 2024

Choose a reason for hiding this comment

Uh oh!

wconstab Sep 25, 2024

Choose a reason for hiding this comment

Uh oh!

wconstab commented Sep 25, 2024

Uh oh!

pytorchmergebot commented Sep 25, 2024

Merge started

Uh oh!

pytorchmergebot commented Sep 25, 2024

Merge failed

Uh oh!

H-Huang commented Sep 25, 2024

Uh oh!

pytorchmergebot commented Sep 25, 2024

Merge started

Uh oh!

pytorchmergebot commented Sep 25, 2024

Merge failed

Uh oh!

H-Huang commented Sep 26, 2024

Uh oh!

pytorchmergebot commented Sep 26, 2024

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

wconstab commented Sep 24, 2024 •

edited by H-Huang

Loading

pytorch-bot bot commented Sep 24, 2024 •

edited

Loading

H-Huang Sep 25, 2024 •

edited

Loading