fix allreduce distributed not working after 2nd iteration

norch/distributed/distributed.py

@@ -10,10 +10,10 @@ def init_process_group(rank, world_size, backend='nccl'):
     Tensor._C.init_process_group(rank, world_size)
 
 def get_rank():
-    return os.getenv('OMPI_COMM_WORLD_RANK', 0)
+    return int(os.getenv('OMPI_COMM_WORLD_RANK', 0))
 
 def get_world_size():
-    return os.getenv('OMPI_COMM_WORLD_SIZE', 1)
+    return int(os.getenv('OMPI_COMM_WORLD_SIZE', 1))
 
 def broadcast_tensor(tensor, src=0):
 

norch/nn/parallel.py

@@ -23,14 +23,17 @@ class DistributedDataParallel(Module):
         for _, _, parameter in self.parameters():
             dist.broadcast_tensor(parameter)
     
+    @staticmethod
     def allreduce_grads_hook(grad):
         """
         Everytime a gradient is assign to some value, it calculates mean of this gradient among all devices
         """
+        avg_grad = grad
         if isinstance(grad, norch.Tensor):
             dist.allreduce_sum_tensor(grad) 
-            grad /= dist.get_world_size()
-        return grad
+            avg_grad = grad / dist.get_world_size()
+            avg_grad = 0*grad + 5
+        return avg_grad
     
     def register_grads_hooks(self):
         """
@@ -38,7 +41,6 @@ class DistributedDataParallel(Module):
         """
         for _, _, parameter in self.parameters():
             parameter.register_hook(self.allreduce_grads_hook)
-
     
 
     

norch/optim/optimizers/sgd.py

@@ -18,6 +18,7 @@ class SGD(Optimizer):
             velocity = self.momentum * velocity - self.lr * parameter.grad
 
             updated_parameter = parameter + velocity
+            updated_parameter.hooks = parameter.hooks.copy()
 
             setattr(module, name, updated_parameter)
 

train.py

@@ -70,16 +70,14 @@ def main():
             loss = criterion(outputs, target)
             
             optimizer.zero_grad()
+            print(f"GRADIENT BEFORE: {model.module.fc2.bias.grad}")
             
-            print("#####################\n\nantes backward")
-            print(model.module.fc1.bias.grad)
             loss.backward()
-            print(model.module.fc1.bias.grad)
-            print("\n\n\n###############\n\npós backward")
+            print(f"GRADIENT AFTER: {model.module.fc2.bias.grad}")
+            print("\n\n")
 
 
             optimizer.step()
-            print("@@")
             break
 
         break