fix(metrics): align WandB axes and weighted minibatch stats

training/tests/unit/test_rl_metrics_aliases.py

@@ -110,6 +110,22 @@ def test_averages_across_minibatches(self):
         assert metrics["train/ppo_clip_frac"] == (0.0 + 0.1 + 0.3) / 3
         assert metrics["train/ppo_ratio_mean"] == (1.0 + 1.05 + 1.2) / 3
 
+    def test_uses_weights_when_minibatches_have_different_sizes(self):
+        fwd_bwds = [
+            self._fake_fwd_bwd(ppo_clip_frac=0.0, ppo_ratio_mean=1.00),
+            self._fake_fwd_bwd(ppo_clip_frac=0.5, ppo_ratio_mean=1.20),
+        ]
+        metrics = compute_step_metrics(
+            prompt_groups=[],
+            fwd_bwd_results=fwd_bwds,
+            optim_result=None,
+            n_accum=len(fwd_bwds),
+            timing_metrics={},
+            fwd_bwd_weights=[3, 1],
+        )
+        assert metrics["train/ppo_clip_frac"] == 0.125
+        assert metrics["train/ppo_ratio_mean"] == 1.05
+
     def test_k1_matches_single_fwd_bwd_behavior(self):
         """K=1: report the single minibatch's metrics directly (pre-PR behavior)."""
         only = self._fake_fwd_bwd(ppo_clip_frac=0.42, ppo_ratio_mean=1.07)

training/utils/logging.py

@@ -33,12 +33,23 @@ def setup_wandb(wb: WandBConfig, config: dict[str, Any]) -> bool:
 
         wandb.init(entity=wb.entity, project=wb.project, name=wb.run_name, config=config)
         if wandb.run is not None:
+            # Align metric families with the step axis they naturally track.
+            # Recipes that sample and train in one loop emit both train/step
+            # and rollout/step in the same metrics dict, so WandB can plot
+            # train and rollout metrics on their own axes.
             wandb.define_metric("train/step")
+            wandb.define_metric("rollout/step")
+            wandb.define_metric("eval/step")
             wandb.define_metric("train/*", step_metric="train/step")
-            wandb.define_metric("perf/*", step_metric="train/step")
-            wandb.define_metric("rollout/*", step_metric="train/step")
+            wandb.define_metric("rollout/*", step_metric="rollout/step")
+            wandb.define_metric("perf/*", step_metric="rollout/step")
+            wandb.define_metric("multi_turn/*", step_metric="rollout/step")
+            wandb.define_metric("passrate/*", step_metric="rollout/step")
+            wandb.define_metric("eval/*", step_metric="eval/step")
             wandb.define_metric("batch/*", step_metric="train/step")
             wandb.define_metric("infra/*", step_metric="train/step")
+            wandb.define_metric("version/*", step_metric="train/step")
+            wandb.define_metric("async/*", step_metric="train/step")
             logger.info("WandB: %s", wandb.run.url)
         return True
     except ImportError:
@@ -73,12 +84,23 @@ def compute_pass_at_k(
 
 
 def wandb_log(metrics: dict[str, Any], step: int) -> None:
-    """Log metrics to WandB if available."""
+    """Log metrics to WandB if available.
+
+    Pass the metrics dict directly so WandB resolves each metric family
+    from the step metrics declared in ``setup_wandb``.  ``step`` is written
+    under both train and rollout axes when the caller did not include them,
+    which preserves single-loop recipe behavior without forcing all metrics
+    onto the same x-axis.
+    """
     try:
         import wandb
 
-        if wandb.run is not None:
-            wandb.log(metrics, step=step, commit=True)
+        if wandb.run is None:
+            return
+        out = dict(metrics)
+        out.setdefault("train/step", step)
+        out.setdefault("rollout/step", step)
+        wandb.log(out, commit=True)
     except ImportError:
         pass
 

training/utils/rl/metrics.py

@@ -102,6 +102,7 @@ def compute_step_metrics(
     timing_metrics: dict[str, Any],
     loop_stats: dict | None = None,
     completions_per_prompt: int = 1,
+    fwd_bwd_weights: Sequence[float] | None = None,
 ) -> dict[str, Any]:
     """Compute all per-step wandb metrics from prompt groups and remote results.
 
@@ -119,22 +120,25 @@ def compute_step_metrics(
             if k not in _SKIP_REMOTE_KEYS:
                 metrics[f"train/{k}"] = v
 
-    # Average fwd_bwd metrics across inner minibatches. With ppo_n_minibatches=1
-    # this reduces to the pre-PR behavior (one result, mean == that result). With
-    # K>1 the last minibatch alone is misleading — early minibatches haven't
-    # drifted yet so their ppo_clip_frac is ~0, while later ones clip more; the
-    # headline claim of this PR is that clipping fires, which the mean reports
-    # honestly rather than only showing the most-clipped minibatch.
+    # Average fwd_bwd metrics across inner minibatches. ppo_n_minibatches uses
+    # ceil(n / num_minibatches), so the last minibatch is often smaller than
+    # the rest. Callers that know per-minibatch weights can pass them to avoid
+    # biasing train metrics toward a small tail minibatch.
     if fwd_bwd_results:
+        if fwd_bwd_weights is not None and len(fwd_bwd_weights) == len(fwd_bwd_results):
+            weights = [float(w) for w in fwd_bwd_weights]
+        else:
+            weights = [1.0] * len(fwd_bwd_results)
+        total_weight = sum(weights) or float(len(fwd_bwd_results))
+
         accum: dict[str, float] = {}
-        for result in fwd_bwd_results:
+        for result, weight in zip(fwd_bwd_results, weights):
             for k, v in result.metrics.items():
                 if k in _SKIP_REMOTE_KEYS:
                     continue
-                accum[k] = accum.get(k, 0.0) + v
-        n = len(fwd_bwd_results)
+                accum[k] = accum.get(k, 0.0) + float(v) * weight
         for k, v in accum.items():
-            metrics[f"train/{k}"] = v / n
+            metrics[f"train/{k}"] = v / total_weight
 
     all_rewards: list[float] = []
     all_comp_lens: list[int] = []