All-reduce ntokens_seen before logging (#1509)

Currently, `ntokens_seen` is only locally logged. I think it is almost
always desirable to only track the global quantity (the only use case I
can see for per-device tracking is for debugging?).

Therefore, I propose to all-reduce `ntokens_seen` before logging.

Author: runame

torchtitan/distributed/utils.py

@@ -62,6 +62,16 @@ def dist_max(
     )
 
 
+def dist_sum(
+    x: torch.Tensor,
+    mesh: DeviceMesh,
+    extra_pg: dist.ProcessGroup | None = None,
+) -> float:
+    return _dist_reduce(
+        x, reduceOp=c10d.ReduceOp.SUM.name, mesh=mesh, extra_pg=extra_pg
+    )
+
+
 def dist_mean(
     x: torch.Tensor,
     mesh: DeviceMesh,

torchtitan/train.py

@@ -498,15 +498,23 @@ def train_step(
         if parallel_dims.dp_cp_enabled:
             loss = loss.detach()
             ft_pg = self.ft_manager.loss_sync_pg
-            global_avg_loss, global_max_loss = (
+            global_avg_loss, global_max_loss, global_ntokens_seen = (
                 dist_utils.dist_mean(loss, parallel_dims.world_mesh["dp_cp"], ft_pg),
                 dist_utils.dist_max(loss, parallel_dims.world_mesh["dp_cp"], ft_pg),
+                dist_utils.dist_sum(
+                    torch.tensor(
+                        self.ntokens_seen, dtype=torch.int64, device=self.device
+                    ),
+                    parallel_dims.world_mesh["dp_cp"],
+                    ft_pg,
+                ),
             )
         else:
             global_avg_loss = global_max_loss = loss.detach().item()
+            global_ntokens_seen = self.ntokens_seen
 
         extra_metrics = {
-            "n_tokens_seen": self.ntokens_seen,
+            "n_tokens_seen": global_ntokens_seen,
             "lr": lr,
         }
         self.metrics_processor.log(