add metrics

Author: acisseJZhong

torchtitan/components/loss.py

@@ -372,12 +372,24 @@ def __init__(self, config: Config, *, compile_config: CompileConfig | None = Non
 def compute_logprobs(
     logits: torch.Tensor,
     labels: torch.Tensor,
-) -> torch.Tensor:
-    """Compute per-position logprobs from logits and labels.
+    *,
+    return_entropy: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    """Per-position logprobs from logits and labels, optionally with entropy.
 
     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.
+
+    When ``return_entropy`` is set, also returns per-token Shannon entropy
+    ``H(p) = logsumexp(logits) - sum(softmax(logits) * logits)``, shape
+    ``[batch, seq_len]``. Both share the single vocab gather + fp32 upcast.
+    Entropy 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.
+
+    Returns ``logprobs`` when ``return_entropy`` is False, else
+    ``(logprobs, entropy)``.
     """
     if isinstance(logits, DTensor):
         # TODO: pass `grad_placements=[Replicate(), ...]` to make the autograd
@@ -406,13 +418,22 @@ def compute_logprobs(
             dst=spmd.I,
         )
 
+    # Single bf16->fp32 upcast, reused by both logprobs and (optionally) entropy.
+    logits = logits.float()
     B, L, V = logits.shape
-    return -F.cross_entropy(
-        logits.float().reshape(B * L, V),
+    logprobs = -F.cross_entropy(
+        logits.reshape(B * L, V),
         labels.reshape(B * L),
         reduction="none",
         ignore_index=IGNORE_INDEX,
     ).reshape(B, L)
+    if not return_entropy:
+        return logprobs
+    with torch.no_grad():
+        entropy = torch.logsumexp(logits, dim=-1) - (
+            torch.softmax(logits, dim=-1) * logits
+        ).sum(dim=-1)
+    return logprobs, entropy
 
 
 class GradAccumulator:

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

@@ -78,7 +78,11 @@ 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, computed under no_grad inside).
+        trainer_logprobs, token_entropy = compute_logprobs(
+            logits, labels, return_entropy=True
+        )
         # 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 +127,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