[NCCL] Don't override waitUntilInitialized's setting of comm->initialized_ (#137210)

pytorchbot · eqy · web-flow · commit 0b1b609ed738 · 2024-10-02T14:51:50.000-07:00
[NCCL] Don't override `waitUntilInitialized`'s setting of `comm->initialized_` (#136155) #133630 sets `initialized_` to `true` which causes previous wait codepaths to skip necessary waits, see also ##136151 CC @shuqiangzhang @wconstab Pull Request resolved: #136155 Approved by: https://github.com/fduwjj, https://github.com/kwen2501, https://github.com/c-p-i-o, https://github.com/shuqiangzhang (cherry picked from commit e3aa5e2) Co-authored-by: eqy <eddiey@nvidia.com>
diff --git a/torch/csrc/distributed/c10d/NCCLUtils.cpp b/torch/csrc/distributed/c10d/NCCLUtils.cpp
@@ -84,7 +84,9 @@ std::shared_ptr<NCCLComm> NCCLComm::split(
       std::nullopt);
   ++source->ncclCommSplitCounter_;
   comm->rank_ = rank;
-  comm->initialized_ = true;
+  if (!nccl_use_nonblocking()) {
+    comm->initialized_ = true;
+  }
   return comm;
 }
 #endif
