Add sync_cache_decision_cross_ranks (#193043) (#193043)
Summary:
This diff introduces `sync_cache_decision_cross_ranks`. This mechanism make sure that all ranks have the same aot_autograd cache hit/miss results. This is important because during compilation process, there are some steps running collectives (e.g., `sync_decision_cross_ranks`), if some rank hit aot_autograd cache while other ranks do not. Then only ranks missed the cache will run the collective, which leads to a hang. With this flag on, we ensure unanimous decisions across all ranks.
Another change in this PR is that we introduced a global `_COMPILE_SYNC_PG` whose backend is gloo. Both `sync_cache_decision_cross_ranks` and `_sync_decision_cross_ranks` are using this pg. It's better than previous implementation with NCCL pg which is likely interleave with an in flight NCCL op from training execution.
Test Plan:
f1125112932
It's working:
{F1993930398}
Differential Revision: D114838651
Pull Request resolved: https://github.com/pytorch/pytorch/pull/193043
Approved by: https://github.com/bobrenjc93, https://github.com/aorenste S
Shuai Yang committed
866415f3f31162fc75ea8a7fb9dfa48d8a247a35
Parent: 9bc1ff8
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com>
on 8/24/2026, 10:22:24 PM