add TORCH_FORCE_SYNCHRONOUS_COLLECTIVES to force functional collectives to be synchronous

[bdhirsh]

One way that OOM's can show up with using torch.compile with distributed code is if our compile fails to issue a wait_tensor properly on a functional collective.

This PR adds an env var that forces all functional collectives to be synchronous at runtime (making issuing wait_tensor at runtime a no-op, so if we fail to issue wait_tensors properly at runtime then we will not leak memory).

If we see a mem leak / OOM when using compile with distributed code, this flag should at least quickly tell us if the mem leak is coming from not waiting on a collective, vs. something else entirely.

cc @yifuwang / @wconstab, wdyt of this env var?

Stack from ghstack (oldest at bottom):

• -> add TORCH_FORCE_SYNCHRONOUS_COLLECTIVES to force functional collectives to be synchronous #128331
• dont prune unused symint graphargs from inner subclass tensors #128045
• inductor: avoiding moving constructor to cuda when it would cause h2d sync in index_put_ fallback #130338

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang @d4l3k

[pytorch-bot]

🔗 Helpful Links

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

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

• Too many GPU instances requested - any *.nvidia.gpu and ephemeral instances should experience significant queue times

✅ No Failures

As of commit 2b4a770 with merge base 81df076 ():
💚 Looks good so far! There are no failures yet. 💚

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

[ezyang]

Doesn't have to be in this PR, but you should think about where this should be doc'ed and how people who need to know about it can find out about it.

[bdhirsh]

yep, very fair (@drisspg - per that BE project you were looking into, do we have a nice way for users to find TORCH_* env vars already? Otherwise maybe if/once we agree that this env var is reasonable, I can look for a place in the docs to mention it)

[Skylion007]

Use std::getenv. Also what if the user wants to explicitly disable it?

[bdhirsh]

I'll probably just switch to use the same API's to read envvars as the other distributed code: https://github.com/pytorch/pytorch/blob/main/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L780

[wconstab]

I like the idea of adding the env, but not sure how to name it.

most of the distributed envs are named TORCH_NCCL_* and @shuqiang recently doc'd them all in this page, you should probably add an entry there and also in compile docs?

The 2 confusing points regarding naming and behavior are

• its not NCCL specific, so probably TORCH_ is better
• the env only affects functional collectives, which might confuse users of non-functional collectives

TORCH_FUNCOL_SYNC?

[bdhirsh]

The 2 confusing points regarding naming and behavior are

@wconstab thanks for the feedback. The point about most env vars including NCCL in their name but this not being nccl-specific is fair. TORCH_FUNCOL_SYNC sounds reasonable to me, I can change it to that.

[yifuwang]

This looks great! Thanks for adding it.

Curious what are the most common sources of unwaited collectives based on your observation. It is more due to graph break boundary handling, or users calling the _c10d_functional ops without calling wait? For the latter, we can add a pass to lint/fix the program.