fix import distributed nn

norch/csrc/distributed.cpp

@@ -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);
 }

norch/csrc/distributed.h

@@ -24,6 +24,7 @@ extern "C" {
   void init_process_group(int rank, int world_size);
   void broadcast_tensor(Tensor* tensor);
   void allreduce_sum_tensor(Tensor* tensor);
+  void allreduce_mean_tensor(Tensor* tensor);
 }
 
 #endif /* DISTRIBUTED_H */

norch/nn/__init__.py

@@ -1,4 +1,5 @@
 from .modules import *
 from .activation import *
 from .loss import *
-from .functional import *
+from .functional import *
+from .distributed import *

norch/nn/distributed.py

@@ -14,7 +14,7 @@ class DistributedDataParallel(Module):
     def backward(self):
         self.module.backward()
         for module, name, parameter in self.parameters():
-            dist.allreduce_sum_tensor(parameter)
+            dist.allreduce_mean_tensor(parameter)