add metrics

Author: acisseJZhong

torchtitan/components/loss.py

@@ -365,15 +365,13 @@ def __init__(self, config: Config, *, compile_config: CompileConfig | None = Non
         self._maybe_compile(compile_config)
 
 
-def compute_logprobs(
-    logits: torch.Tensor,
-    labels: torch.Tensor,
-) -> torch.Tensor:
-    """Compute per-position logprobs from logits and labels.
+def _gather_vocab_logits(logits: torch.Tensor) -> torch.Tensor:
+    """Redistribute vocab-sharded TP logits to full-vocab local logits.
 
-    Output shape matches input: ``[batch, seq_len]``. Any DTensor placement
-    handling is centralized here so RL losses that call ``compute_logprobs`` do
-    not need to duplicate the vocab-gather logic.
+    Per-token quantities that need the full vocab distribution (logprobs,
+    entropy) share this gather so the DTensor / spmd_types handling lives in one
+    place. Returns a plain (local) tensor; a non-vocab-sharded input is returned
+    unchanged.
     """
     if isinstance(logits, DTensor):
         # TODO: pass `grad_placements=[Replicate(), ...]` to make the autograd
@@ -401,7 +399,17 @@ def compute_logprobs(
             src=spmd.S(-1),
             dst=spmd.I,
         )
+    return logits
+
 
+def _logprobs_from_full_vocab(
+    logits: torch.Tensor, labels: torch.Tensor
+) -> torch.Tensor:
+    """Per-token logprobs from already-gathered full-vocab ``logits`` [B, L, V].
+
+    Split out so ``compute_logprobs`` and ``compute_logprobs_and_entropy`` share
+    one logprob code path (and one gather at the caller).
+    """
     B, L, V = logits.shape
     return -F.cross_entropy(
         logits.float().reshape(B * L, V),
@@ -411,6 +419,44 @@ def compute_logprobs(
     ).reshape(B, L)
 
 
+def compute_logprobs(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+) -> torch.Tensor:
+    """Compute per-position logprobs from logits and labels.
+
+    Output shape matches input: ``[batch, seq_len]``. Vocab-sharded TP logits are
+    gathered via ``_gather_vocab_logits`` so callers do not duplicate that logic.
+    """
+    return _logprobs_from_full_vocab(_gather_vocab_logits(logits), labels)
+
+
+def compute_logprobs_and_entropy(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per-token logprobs and Shannon entropy from a single vocab gather.
+
+    Equivalent to ``compute_logprobs`` plus a separate entropy pass, but gathers
+    the (TP-sharded) vocab only once.
+
+    - logprobs use the same fused ``F.cross_entropy`` path as ``compute_logprobs``.
+    - entropy ``H(p) = logsumexp(logits) - sum(softmax(logits) * logits)`` is a
+      metric only, so it is computed under ``no_grad``: it never contributes
+      gradient and must not build an autograd graph over the [B, L, V] softmax.
+
+    Both outputs are ``[batch, seq_len]``.
+    """
+    logits = _gather_vocab_logits(logits)
+    logprobs = _logprobs_from_full_vocab(logits, labels)
+    with torch.no_grad():
+        logits_fp32 = logits.float()
+        entropy = torch.logsumexp(logits_fp32, dim=-1) - (
+            torch.softmax(logits_fp32, dim=-1) * logits_fp32
+        ).sum(dim=-1)
+    return logprobs, entropy
+
+
 class GradAccumulator:
     """Accumulates chunk gradients into a pre-allocated buffer.
 

torchtitan/experiments/rl/actors/trainer.py

@@ -438,7 +438,7 @@ async def optim_step(self) -> OptimStepOutput:
         if len(current_lrs) != 1:
             raise ValueError(
                 "RL metrics only support a single optimizer LR for "
-                f"train/lr; got {current_lrs}"
+                f"trainer/lr; got {current_lrs}"
             )
         current_lr = float(current_lrs[0])
 
@@ -466,9 +466,9 @@ async def optim_step(self) -> OptimStepOutput:
         return OptimStepOutput(
             policy_version=self.policy_version,
             metrics={
-                "train/grad_norm/mean": float(grad_norm.item()),
-                "train/lr": current_lr,
-                "train/policy_version": float(self.policy_version),
+                "trainer/grad_norm/mean": float(grad_norm.item()),
+                "trainer/lr": current_lr,
+                "trainer/policy_version": float(self.policy_version),
             },
         )
 

torchtitan/experiments/rl/losses/dapo.py

@@ -12,7 +12,7 @@
 
 import torch
 
-from torchtitan.components.loss import BaseLoss, compute_logprobs
+from torchtitan.components.loss import BaseLoss, compute_logprobs_and_entropy
 from torchtitan.config import CompileConfig
 
 # Clamp |log(pi_theta/pi_old)| before exp() so a large generator/trainer
@@ -78,7 +78,9 @@ def __call__(
             (loss, metrics) where loss is a scalar tensor and metrics is a dict of
             scalar tensors pre-normalized for SUM reduction across DP ranks.
         """
-        trainer_logprobs = compute_logprobs(logits, labels)
+        # One vocab gather feeds both: differentiable logprobs (for the loss) and
+        # detached per-token entropy (a metric, already under no_grad in the helper).
+        trainer_logprobs, token_entropy = compute_logprobs_and_entropy(logits, labels)
         # A non-finite generator logprob (notably under cudagraph) has no valid
         # old-policy reference, so DROP that token from the loss + denominator (cleaner
         # than nan->0, which trains it as if it were on-policy). `response_mask` keeps
@@ -123,13 +125,18 @@ def __call__(
                     (~torch.isfinite(generator_logprobs)).float() * response_mask
                 ).sum()
                 / loss_denominator,
+                # Mean per-token log-ratio (log p_trainer - log q_generator) over
+                # sampled tokens. This is the k1 Monte-Carlo estimate of -KL(q || p).
                 "bit_wise/logprob_diff/mean": diff_for_metrics.float().sum()
                 / loss_denominator,
                 "bit_wise/ratio_tokens_different/mean": (
                     (diff_for_metrics.abs() > 1e-6).float() * loss_mask
                 ).sum()
                 / loss_denominator,
                 "bit_wise/logprob_diff/max": diff_for_metrics.abs().max(),
+                # Mean trainer-policy entropy H(p) over response tokens.
+                "trainer/entropy/mean": (token_entropy * response_mask).sum()
+                / loss_denominator,
             }
 
         return loss, metrics

torchtitan/experiments/rl/observability/metrics/processor.py

@@ -64,8 +64,9 @@ class Config(Configurable.Config):
                 # --- learning signals (not perf, but you want them in the same glance) ---
                 "loss/mean",
                 "rollout_reward/_mean",
-                "train/grad_norm/mean",
-                "train/lr",
+                "trainer/entropy/mean",
+                "trainer/grad_norm/mean",
+                "trainer/lr",
             ]
         )
         """Regex search patterns selecting console keys for train log lines.

torchtitan/experiments/rl/tests/test_metrics.py

@@ -460,7 +460,7 @@ def test_console_prefix_renders_before_step(
         (-0.001234, "-0.0012"),
         (99.999, "100.00"),
         (5, "5.0"),
-        # Very small values fall back to scientific so train/lr ≈ 2e-6
+        # Very small values fall back to scientific so trainer/lr ≈ 2e-6
         # doesn't render as `0.00000`.
         (2e-6, "2.0e-06"),
         (1e-5, "0.00001"),  # boundary: 5 decimals still resolves the digit