commit 1ea224c2a4c70484a21954acf85ced387c31aaae
author mikey dagitses <pytorch@dagits.es> Tue May 09 15:29:08 2023 +0000
committer PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> Tue May 09 15:29:08 2023 +0000
tree c09abac72135dec81353e4ae49c8c35e46d1cc45
parent bc3108c2e295b5a2ece6c70881d1f0df6c1b3be6

make torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp data_ptr-correct (#100888)

make torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp data_ptr-correct

Test Plan: Rely on CI.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/100888
Approved by: https://github.com/ezyang

torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp

diff --git a/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp b/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp
index 3894b57..255a730 100644
--- a/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp
+++ b/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp
@@ -85,12 +85,15 @@
   ncclRedOp_t preMulSum;
   bool has_tensor = preMulSupplement->tensor_factor.defined();
   auto residence = has_tensor ? ncclScalarDevice : ncclScalarHostImmediate;
-  T* ptr_factor =
-      has_tensor ? preMulSupplement->tensor_factor.data_ptr<T>() : nullptr;
+  const T* ptr_factor = has_tensor
+      ? preMulSupplement->tensor_factor.const_data_ptr<T>()
+      : nullptr;
   T scalar_factor = T(preMulSupplement->double_factor);
   ncclRedOpCreatePreMulSum(
       &preMulSum,
-      has_tensor ? ptr_factor : &scalar_factor,
+      // https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/ops.html#ncclredopcreatepremulsum
+      // tells us that the scalar input is strictly a multiplier.
+      /*scalar=*/has_tensor ? const_cast<T*>(ptr_factor) : &scalar_factor,
       dataType,
       residence,
       comm);