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

allreduce mean

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-25 08:52:20 -03:00
parent 7a3b0f276b
commit 5283c1bc06
2 changed files with 20 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

13
norch/csrc/distributed.cpp
View file

@ -56,6 +56,19 @@ void allreduce_sum_tensor(Tensor* tensor) {
cudaStreamDestroy(stream);
}
void allreduce_mean_tensor(Tensor* tensor) {
cudaStream_t stream;
cudaStreamCreate(&stream);
// Perform NCCL AllReduce operation to calculate the mean of all tensors across all processes
NCCL_CHECK(ncclAllReduce(tensor->data, tensor->data, tensor->size, ncclFloat, ncclSum, nccl_comm, stream));
tensor_div_scalar(tensor, tensor->data);
cudaStreamSynchronize(stream);
cudaStreamDestroy(stream);
}
void end_process_group() {
MPI_CHECK(MPI_Finalize());
NCCL_CHECK(ncclCommDestroy(nccl_comm));

7
norch/distributed/distributed.py
View file

@ -22,3 +22,10 @@ def allreduce_sum_tensor(tensor):
Tensor._C.allreduce_sum_tensor.restype = None
Tensor._C.allreduce_sum_tensor(tensor.tensor)
def allreduce_mean_tensor(tensor):
Tensor._C.allreduce_mean_tensor.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.allreduce_mean_tensor.restype = None
Tensor._C.allreduce_mean_tensor(tensor.tensor)