commit f8a6ea770cf7b94040abf92f4b97dcc965946f0f
author Sergey Lebedev <sergeyle@nvidia.com> Thu Nov 09 22:53:36 2023 +0000
committer PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> Thu Nov 09 22:53:40 2023 +0000
tree 00f88d99b8853052971e07f4b1f38ba04a256407
parent c7e0fa49b64694a2cdd1646833af527495555f2b

[UCC] Fix input tensor in scatter (#112246)

Input tensor is valid only for root rank. Fixes https://github.com/openucx/ucc/issues/859

Pull Request resolved: https://github.com/pytorch/pytorch/pull/112246
Approved by: https://github.com/Aidyn-A, https://github.com/Fuzzkatt, https://github.com/kwen2501

torch/csrc/distributed/c10d/ProcessGroupUCC.cpp

diff --git a/torch/csrc/distributed/c10d/ProcessGroupUCC.cpp b/torch/csrc/distributed/c10d/ProcessGroupUCC.cpp
index 35a15d7..cd6c4aa 100644
--- a/torch/csrc/distributed/c10d/ProcessGroupUCC.cpp
+++ b/torch/csrc/distributed/c10d/ProcessGroupUCC.cpp
@@ -1512,7 +1512,7 @@
       coll,
       std::unique_ptr<WorkData>(data),
       tensor.device(),
-      inputTensors[0],
+      (getRank() == opts.rootRank) ? inputTensors[0] : outputTensors,
       outputTensors,
       "ucc:scatter");
 }