aparler/PyNorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

test broadcast and all reduce

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-24 18:31:59 -03:00
parent 55db13836b
commit 94b3582a46
5 changed files with 17 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

6
norch/csrc/distributed.cpp
View file

@ -46,17 +46,13 @@ void broadcast_tensor(Tensor* tensor) {
cudaStreamDestroy(stream);
}
void allreduce_mean_tensor(Tensor* tensor) {
void allreduce_sum_tensor(Tensor* tensor) {
cudaStream_t stream;
cudaStreamCreate(&stream);
// Perform NCCL AllReduce operation to calculate the sum of all tensors across all processes
NCCL_CHECK(ncclAllReduce(tensor->data, tensor->data, tensor->size, ncclFloat, ncclSum, nccl_comm, stream));
// Divide the sum by the world size to get the mean
float scale = 1.0f / world_size;
scalar_mul_tensor_cuda(tensor, scale, tensor->data);
cudaStreamSynchronize(stream);
cudaStreamDestroy(stream);
}