[dynamo][moco] Disallow_in_graph distributed APIs

[anijain2305]

Stack from ghstack (oldest at bottom):

• [dynamo] hasattr on TensorVariable #100199
• -> [dynamo][moco] Disallow_in_graph distributed APIs #100071

cc @soumith @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire

[pytorch-bot]

🔗 Helpful Links

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

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

❗ 1 Active SEVs

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

• download.pytorch.org flaky

❌ 3 New Failures

As of commit 16687a2:

NEW FAILURES - The following jobs have failed:

• cuda11.8-py3.10-gcc7-sm86 / test (inductor_torchbench_dynamic, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh)
• cuda11.8-py3.10-gcc7-sm86 / test (inductor_torchbench, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh)
• linux-focal-cpu-py3.8-gcc7-inductor / test (inductor_torchbench_cpu_accuracy, 1, 1, linux.4xlarge) (gh)

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

[anijain2305]

@wconstab Wondering if you know of a better way here, instead of just graph breaking on them?

[wconstab]

Wondering if you know of a better way here, instead of just graph breaking on them?

No, we have to graph-break on them.

The missing piece is how to ensure we graph-break on the full set of them. I thought I stamped a PR (maybe from @yanboliang) a while back that skipped a large chunk of torch.distributed- looks like that did not land or else I can't find remnants of it.

We should disable for all of the collectives besides the new 'traceable/functional' ones (in _functional_collectives.py).

[yanboliang]

The missing piece is how to ensure we graph-break on the full set of them. I thought I stamped a PR (maybe from @yanboliang) a while back that skipped a large chunk of torch.distributed- looks like that did not land or else I can't find remnants of it.

Yes, that PR has been reverted multiple times, I finally use regex matching to skip torchrec.distribued(make it support torch.package), but only be used internally behind is_fbcode. I think we should skip these files right now, not only because they are collectives ops, but also they have memory leak issue due to their implementation.

[anijain2305]

@wconstab Is this too hacky? Problem is there is no easy way to tell what's a torch "op" and whats not.

[awgu]

Maybe cc @H-Huang @kwen2501 -- do we expect a function to use the exception_handler decorator if and only if it is a c10d op?

[wconstab]

hah yea its kinda hacky.

But otoh i think it's reasonable to have a set of dynamo_unsupported_distributed_c10d_ops and explicitly add them all. We should just get folks like @H-Huang to agree on the set of ops to initially tag.

[H-Huang]

I am missing the context here. Is the idea just to exclude all the operations defined in distributed_c10d.py?

[anijain2305]

Yes

[anijain2305]

Or specifically, the operators that do not work well with TorchDynamo tracing and we want to graph break (i.e. fallback to eager) on them

[anijain2305]

@H-Huang ping

[H-Huang]

I see, then would it be possible to just get all operations of the module through something like:
https://stackoverflow.com/questions/139180/how-to-list-all-functions-in-a-module

Then we don't have to create this list in distributed_c10d and worry about keeping it up to date

[anijain2305]

@H-Huang Yes, so we tried dir(module) and __dict__ etc but then it gives a large number of functions and many of them not really "operators".

Another way I tried earlier in a commit - 8e6dee3
was to record the ops in exception_handler. The assumption was that anyone adding a new op will decorate it withe exception handler. But, that looked little hacky as well.

[H-Huang]

I see. Thanks for clarifying. This looks like the route to go then. Should init_process_group be included in this list? This is not a conventional "operator"

[anijain2305]

Thanks for pointing that out. Removed it from the list.

[H-Huang]

LGTM

[anijain2305]

@pytorchbot merge

[pytorchmergebot]

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]

Merge failed

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

• pull / linux-focal-py3-clang7-android-ndk-r19c-gradle-custom-build-single-full-jit / build-and-test (default, 1, 1, linux.2xlarge)

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

[anijain2305]

@pytorchbot merge -f "unrelated CI error"

[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