From c42d3345993dd6058a80d92f4be15b08cab1bf02 Mon Sep 17 00:00:00 2001
From: lucasdelimanogueira <lucasln96@hotmail.com>
Date: Mon, 3 Jun 2024 15:14:42 -0300
Subject: [PATCH] fix allreduce distributed not working after 2nd iteration

---
 norch/__pycache__/tensor.cpython-38.pyc       | Bin 20576 -> 20576 bytes
 norch/distributed/distributed.py              |   4 ++--
 norch/nn/__pycache__/parameter.cpython-38.pyc | Bin 736 -> 736 bytes
 norch/nn/parallel.py                          |   8 +++++---
 .../optimizers/__pycache__/sgd.cpython-38.pyc | Bin 1243 -> 1274 bytes
 norch/optim/optimizers/sgd.py                 |   1 +
 train.py                                      |   8 +++-----
 7 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc
index 8505635cae1e65dbd99dee53f7aa7be84ac97f3f..4b1058ad59bd3ce7565d925312dbd43e2454c801 100644
GIT binary patch
delta 22
ccmaE`fbqcsM($8vUM>b8_{SQzkvlX108qjPNdN!<

delta 22
ccmaE`fbqcsM($8vUM>b82tOaYkvlX108r=#YXATM

diff --git a/norch/distributed/distributed.py b/norch/distributed/distributed.py
index 98bc2a0..d534ec4 100644
--- a/norch/distributed/distributed.py
+++ b/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):
 
diff --git a/norch/nn/__pycache__/parameter.cpython-38.pyc b/norch/nn/__pycache__/parameter.cpython-38.pyc
index 14bb0692963118abaefee14811bc7562f61be0e2..f85766ad4dc99ef0d5031c9469f4ec7d80a3ca68 100644
GIT binary patch
delta 20
acmaFB`hb-?l$V!_0SLab#ckxi$^-y5$OSt9

delta 20
acmaFB`hb-?l$V!_0SJDciQUM3l?eblg9Zcu

diff --git a/norch/nn/parallel.py b/norch/nn/parallel.py
index c021537..8d6fef9 100644
--- a/norch/nn/parallel.py
+++ b/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)
-
     
 
     
diff --git a/norch/optim/optimizers/__pycache__/sgd.cpython-38.pyc b/norch/optim/optimizers/__pycache__/sgd.cpython-38.pyc
index 7336b8ecea49832c6630b61684847c7482efdc92..0ee226c0b3d139a685b589b2150df2735cacdb97 100644
GIT binary patch
delta 178
zcmcc3`HPb`l$V!_0SNxG#ig}w<lV+3&;#U^Fw`()F*Y;0FvQkOzRcw0QNx<WzJOyP
zLk)X8dkJR^Qw?(sdowd5Lk$~HA#)0&Bts3O4MPp<0<MJ&AW;Zi$6Cmh!Vt`$$>dkW
z4m4Jid$K>XJ`Za~etvfGEtcf`g38IY%-IrRAQv!7Fmf^RF^Ms8F!C_*09hQ2lm9Tg
RFe*=WWU*(An%u||1pv%QCvpG)

delta 147
zcmeyxd7G0rl$V!_0SI(9N2lp+<lV+3PzU6dFw`()F*Y;0FvR9fzRcugRl;7wRKr}u
z-ptI%P{Rh4WKLm}WT;`ZVW?qUz_E}4BnqJmxl$N{88n&vir9e$X>v}EW!9fOlQ~O7
l2xJYT1S1y{ACnjp2NMsFo-EDc!l*bofyJKj*5m~&Q2=Or9y9;|

diff --git a/norch/optim/optimizers/sgd.py b/norch/optim/optimizers/sgd.py
index 1d785c7..62375f6 100644
--- a/norch/optim/optimizers/sgd.py
+++ b/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)
 
diff --git a/train.py b/train.py
index 4eb30c6..666355c 100644
--- a/train.py
+++ b/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