Questions about using the reduce_scatter api in nccl.h. How should the input and output parameters in the api be defined?

sushe · April 9, 2025, 7:32am

Questions about using the reduce_scatter api in nccl.h. How should the input and output parameters in the api be defined?
Check the code: On line 752, const auto count = inputs[0].numel() / len;
In ncclReduceScatter, this count is required to be recv_count. They don’t seem to match.

github.com/pytorch/pytorch

torch/csrc/cuda/nccl.cpp

main


      
                  data_type,
                  to_nccl_red_op(op),
                  to_nccl_comm(comm),
                  stream));
            }
          #else
            TORCH_CHECK(false, "PyTorch built without NCCL support");
          #endif
          }
          
          void reduce_scatter(
              const std::vector<at::Tensor>& inputs,
              std::vector<at::Tensor>& outputs,
              int32_t op,
              const stream_list& streams,
              const comm_list& user_comms) {
          #ifdef USE_NCCL
            using namespace torch::cuda::nccl::detail;
            const auto len = inputs.size();
            check_inputs(inputs, outputs, 1, len);

neoncube · April 9, 2025, 10:42am

It’s tricky. I found this tutorial to be extremely helpful when I was learning scatter(): Understand torch.scatter_(). First, note that scatter_() is an… | by Yu Yang | Medium

sushe · April 9, 2025, 12:03pm

Thank you for your reply. My question is about the NCCL communication library, and it may be different from yours.

neoncube · April 9, 2025, 3:18pm

Yes, you’re right, sorry