Fix mean mismatch metrics

Author: YanhuiDua

xtuner/v1/rl/base/rollout_is.py

@@ -467,7 +467,7 @@ def compute_mismatch_metrics(
     - When Speed Kills Stability: https://yingru.notion.site/When-Speed-Kills-Stability-271211a558b7808d8b12d403fd15edda
     """
     metrics = {}
-
+    metrics["valid"] = response_mask.any().item()
     # 1. Training policy perplexity (always available)
     # Formula: exp(-1/|T| * Σ log π_training(y_t|y_<t))
     # where |T| is the number of tokens generated by the model
@@ -529,18 +529,32 @@ def compute_mismatch_metrics(
 
 def merge_rollout_is_metrics(rollout_is_metrics: list[dict[str, float]], device="cuda") -> dict[str, float]:
     metrics = {}
-    for key in rollout_is_metrics[0].keys():
-        all_values = [m[key] for m in rollout_is_metrics]
+    keys = [k for k in rollout_is_metrics[0].keys() if k != "mismatch/valid"]
+
+    for key in keys:
+        values = []
+        valids = []
+        for m in rollout_is_metrics:
+            is_valid = m.get("mismatch/valid", True)
+            valids.append(float(is_valid))
+            values.append(m[key] if is_valid else 0.0)  # set to 0.0 if invalid
+        value_tensor = torch.tensor(values, dtype=torch.float32, device=device)
+        valid_tensor = torch.tensor(valids, dtype=torch.float32, device=device)
+
+        # Aggregate across all processes
         if "max" in key:
-            max_value = torch.tensor(all_values).max().to(torch.float32).to(device)
+            max_value = value_tensor.max()
             dist.all_reduce(max_value, op=dist.ReduceOp.MAX)
             metrics[key] = max_value.item()
         elif "min" in key:
-            min_value = torch.tensor(all_values).min().to(torch.float32).to(device)
+            min_value = value_tensor.min()
             dist.all_reduce(min_value, op=dist.ReduceOp.MIN)
             metrics[key] = min_value.item()
         else:
-            mean_value = torch.tensor(all_values).mean().to(torch.float32).to(device)
-            dist.all_reduce(mean_value, op=dist.ReduceOp.AVG)
-            metrics[key] = mean_value.item()
+            sum_value = value_tensor.sum()
+            count_value = valid_tensor.sum()
+            dist.all_reduce(sum_value, op=dist.ReduceOp.SUM)
+            dist.all_reduce(count_value, op=dist.ReduceOp.SUM)
+            metrics[key] = sum_value.item() / count_value.item() if count_value.item() > 0 else 0.0
+
     return metrics

xtuner/v1/rl/base/worker.py

@@ -477,7 +477,9 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int):
             rollout_is_metrics = merge_rollout_is_metrics(all_rollout_is_metrics, DEVICE)
             if len(rollout_is_metrics) > 0:
                 logger_msg += f"\n rollout importance sampling metrics:\n{json.dumps(rollout_is_metrics, indent=4)}"
-        self.logger.info(logger_msg)
+
+        if self.rank == 0:
+            self.logger.info(logger_msg)
 
         if self._has_ref:
             # ref logprobs are inplaced updated in compute_actor_logprobs