fix tracking variance error (IBM#27)

Signed-off-by: Mayank Mishra <[email protected]>

dolomite_engine/train_utils.py

@@ -133,8 +133,10 @@ def train_step(
 
 def all_reduce_metrics_tracker(metrics_tracker: MetricsTrackingDict) -> MetricsTrackingDict:
     tensor = [metrics_tracker[key] for key in metrics_tracker]
-    tensor = torch.stack(tensor)
-    torch.distributed.all_reduce(tensor.cpu(), group=ProcessGroupManager.get_data_parallel_group())
+    tensor = torch.stack(tensor) / ProcessGroupManager.get_data_parallel_world_size()
+    tensor = tensor.cpu()
+    # gloo op doesn't support averaging so we do sum and divide by world size above
+    torch.distributed.all_reduce(tensor, group=ProcessGroupManager.get_data_parallel_group())
     tensor = tensor.tolist()
 
     for i, key in enumerate(metrics_tracker):

dolomite_engine/utils/loss_dict.py

@@ -15,6 +15,8 @@ def __add__(self, x: MetricsTrackingDict | dict | float | int) -> MetricsTrackin
         elif isinstance(x, (int, float)):
             for key in self.data:
                 self.data[key] += x
+        else:
+            raise ValueError()
 
         return self
 
@@ -28,6 +30,8 @@ def __truediv__(self, x: MetricsTrackingDict | dict | float | int) -> MetricsTra
         elif isinstance(x, (int, float)):
             for key in self.data:
                 self.data[key] /= x
+        else:
+            raise ValueError()
 
         return self