diff --git a/training/tests/unit/test_rl_metrics_aliases.py b/training/tests/unit/test_rl_metrics_aliases.py index d286bb9e..b54ce2d6 100644 --- a/training/tests/unit/test_rl_metrics_aliases.py +++ b/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) diff --git a/training/utils/logging.py b/training/utils/logging.py index 746eb77e..a4a6a5fa 100644 --- a/training/utils/logging.py +++ b/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 diff --git a/training/utils/rl/metrics.py b/training/utils/rl/metrics.py index 7a7d3ab1..aa9aec72 100644 --- a/training/utils/rl/metrics.py +++ b/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] = []