Use task submitter TLS in gloo working threads (#142184)

dvrogozh · pytorchmergebot · commit 5872a8c6b00a · 2024-12-06T17:03:17.000Z
Fixes: #86830 CC: @albanD Pull Request resolved: #142184 Approved by: https://github.com/albanD
diff --git a/test/distributed/test_c10d_gloo.py b/test/distributed/test_c10d_gloo.py
@@ -1188,6 +1188,11 @@ def test_allgather_noncontiguous_input(self):
         # Take a column of 2D tensor, such that memory is not dense
         self._test_allgather_basics(lambda t: t.expand(2, 2).contiguous()[:, 0])
 
+    @requires_gloo()
+    def test_allgather_inference_mode(self):
+        with torch.inference_mode():
+            self._test_allgather_basics(lambda t: t.clone())
+
     def _test_allgather_stress(self, inputs, fn):
         store = c10d.FileStore(self.file_name, self.world_size)
         pg = self._create_process_group_gloo(
diff --git a/torch/csrc/distributed/c10d/ProcessGroupGloo.cpp b/torch/csrc/distributed/c10d/ProcessGroupGloo.cpp
@@ -426,6 +426,7 @@ void ProcessGroupGloo::AsyncWork::execute(
     work->recordFunctionBeforeCallback_();
   }
   try {
+    at::ThreadLocalStateGuard g(work->getTLS());
     work->run();
   } catch (...) {
     work->finishWorkGlooError(std::current_exception());
diff --git a/torch/csrc/distributed/c10d/ProcessGroupGloo.hpp b/torch/csrc/distributed/c10d/ProcessGroupGloo.hpp
@@ -22,6 +22,8 @@
 #include <torch/csrc/distributed/c10d/Types.hpp>
 #include <torch/csrc/distributed/c10d/Utils.hpp>
 
+#include <ATen/ThreadLocalState.h>
+
 namespace c10d {
 
 constexpr const char* GLOO_BACKEND_NAME = "gloo";
@@ -73,6 +75,10 @@ class TORCH_API ProcessGroupGloo : public Backend {
     c10::intrusive_ptr<c10::ivalue::Future> getFuture() override;
     uint64_t getSequencenumber() const override;
 
+    inline at::ThreadLocalState getTLS() const {
+      return tls_;
+    }
+
    protected:
     friend class ProcessGroupGloo;
 
@@ -87,6 +93,7 @@ class TORCH_API ProcessGroupGloo : public Backend {
     c10::intrusive_ptr<at::ivalue::Future> future_;
     std::function<void()> recordFunctionBeforeCallback_;
     const uint64_t seq_;
+    at::ThreadLocalState tls_;
   };
 
   // Wrap c10d store as Gloo store

Original file line number	Diff line number	Diff line change
`@@ -426,6 +426,7 @@ void ProcessGroupGloo::AsyncWork::execute(`
`426`	`426`	`work->recordFunctionBeforeCallback_();`
`427`	`427`	`}`
`428`	`428`	`try {`
	`429`	`+ at::ThreadLocalStateGuard g(work->getTLS());`
`429`	`430`	`work->run();`
`430`	`431`	`} catch (...) {`
`431`	`432`	`work->finishWorkGlooError(std::current_exception());`