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

fix import distributed nn

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-25 08:57:35 -03:00
parent 5283c1bc06
commit 1e0c408ffb
4 changed files with 6 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

4
norch/csrc/distributed.cpp
View file

@ -63,8 +63,8 @@ void allreduce_mean_tensor(Tensor* tensor) {
// 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);
tensor_div_scalar_cuda(tensor, world_size, tensor->data);
cudaStreamSynchronize(stream);
cudaStreamDestroy(stream);
}