Metrics return sums (num_correct_tokens, entropy_sum) instead of means

Browse files

Files changed (2) hide show

README.md +19 -7
build/torch-cuda/fused_linear_ce.py +26 -17

README.md CHANGED Viewed

@@ -38,23 +38,35 @@ fused_linear_cross_entropy(
     softcap=None,
     reduction="mean",   # "mean" | "sum" | "none"
     return_metrics=False,
-) -> loss  |  (loss, {"mean_token_accuracy": ..., "entropy": ...})
 ```
 `hidden`, `weight` and `bias` all receive gradients.
 ### Metrics for free
-Trainers that switch to a fused loss normally lose their `mean_token_accuracy` and `entropy`
-logging, because there are no logits left to compute them from. This kernel carries both through
-the same vocabulary sweep it already performs — the argmax rides along with the running row max,
-and the entropy term `sum_j p_j * logit_j` is accumulated with the same online rescaling as the
 log-sum-exp:
 ```python
 loss, metrics = fused_linear_cross_entropy(hidden, weight, targets, return_metrics=True)
-metrics["mean_token_accuracy"]  # scalar, over non-ignored positions
-metrics["entropy"]              # scalar, mean softmax entropy in nats
 ```
 Both are forward-only. `return_metrics=True` leaves the loss and the gradients **bit-identical**

     softcap=None,
     reduction="mean",   # "mean" | "sum" | "none"
     return_metrics=False,
+) -> loss  |  (loss, {"num_correct_tokens": ..., "entropy_sum": ...})
 ```
 `hidden`, `weight` and `bias` all receive gradients.
 ### Metrics for free
+Trainers that switch to a fused loss normally lose their token-accuracy and entropy logging,
+because there are no logits left to compute them from. This kernel carries both through the same
+vocabulary sweep it already performs — the argmax rides along with the running row max, and the
+entropy term `sum_j p_j * logit_j` is accumulated with the same online rescaling as the
 log-sum-exp:
 ```python
 loss, metrics = fused_linear_cross_entropy(hidden, weight, targets, return_metrics=True)
+metrics["num_correct_tokens"]  # int64 scalar, count of argmax == target
+metrics["entropy_sum"]         # float32 scalar, summed softmax entropy in nats
+```
+**Both are sums over non-ignored positions, not means.** Sums are the more general contract: a
+per-rank mean cannot be recombined across distributed ranks holding different token counts, so a
+distributed trainer must gather sums before dividing, and the accuracy count stays an exact
+integer instead of a float that callers have to multiply back up. A single-process caller
+recovers the means by dividing by a count it already has:
+```python
+n = (targets != -100).sum()
+mean_token_accuracy = metrics["num_correct_tokens"] / n
+mean_entropy = metrics["entropy_sum"] / n
 ```
 Both are forward-only. `return_metrics=True` leaves the loss and the gradients **bit-identical**

build/torch-cuda/fused_linear_ce.py CHANGED Viewed

@@ -369,24 +369,24 @@ class _FusedLinearCrossEntropy(torch.autograd.Function):
         if return_metrics:
             t = (t_p * alpha).sum(0)
-            entropy = torch.where(valid, lse - t / s, torch.zeros((), dtype=torch.float32, device=device))
-            entropy = entropy.sum() / n_valid
             argmax = a_p.gather(0, gsrc.unsqueeze(0)).squeeze(0)
-            accuracy = ((argmax == targets) & valid).sum().float() / n_valid
         else:
-            entropy = None
-            accuracy = None
         ctx.save_for_backward(hidden, weight, bias, targets, lse, valid, n_valid)
         ctx.ignore_index = ignore_index
         ctx.logit_scale = logit_scale
         ctx.softcap = softcap
         ctx.reduction = reduction
-        ctx.mark_non_differentiable(*[x for x in (entropy, accuracy) if x is not None])
-        return loss, entropy, accuracy
     @staticmethod
-    def backward(ctx, dloss, dentropy, daccuracy):
         hidden, weight, bias, targets, lse, valid, n_valid = ctx.saved_tensors
         N, H = hidden.shape
         V = weight.shape[0]
@@ -502,22 +502,31 @@ def fused_linear_cross_entropy(
             One of `"mean"` (sum over non-ignored positions divided by their count), `"sum"` or
             `"none"` (per-position losses, `0` at ignored positions).
         return_metrics (`bool`, *optional*, defaults to `False`):
-            Whether to also return `mean_token_accuracy` and `entropy`, computed in the same pass
-            as the loss. Enabling this does not change the loss or the gradients.
     Returns:
         `torch.Tensor` when `return_metrics=False`, otherwise a `tuple` of the loss and a `dict`
-        with keys:
-            - `mean_token_accuracy` (`torch.Tensor`):
-                Scalar, fraction of non-ignored positions whose argmax logit is the target.
-            - `entropy` (`torch.Tensor`):
-                Scalar, mean softmax entropy (in nats) over non-ignored positions.
     Examples:
     ```python
     >>> loss = fused_linear_cross_entropy(hidden, lm_head.weight, labels)
     >>> loss, metrics = fused_linear_cross_entropy(hidden, lm_head.weight, labels, return_metrics=True)
     ```
     """
     if hidden.ndim != 2:
@@ -535,9 +544,9 @@ def fused_linear_cross_entropy(
     if bias is not None:
         bias = bias.contiguous()
-    loss, entropy, accuracy = _FusedLinearCrossEntropy.apply(
         hidden, weight, targets, bias, ignore_index, logit_scale, softcap, reduction, return_metrics
     )
     if return_metrics:
-        return loss, {"mean_token_accuracy": accuracy, "entropy": entropy}
     return loss

         if return_metrics:
             t = (t_p * alpha).sum(0)
+            per_token_entropy = torch.where(valid, lse - t / s, torch.zeros((), dtype=torch.float32, device=device))
+            entropy_sum = per_token_entropy.sum()
             argmax = a_p.gather(0, gsrc.unsqueeze(0)).squeeze(0)
+            num_correct_tokens = ((argmax == targets) & valid).sum()
         else:
+            entropy_sum = None
+            num_correct_tokens = None
         ctx.save_for_backward(hidden, weight, bias, targets, lse, valid, n_valid)
         ctx.ignore_index = ignore_index
         ctx.logit_scale = logit_scale
         ctx.softcap = softcap
         ctx.reduction = reduction
+        ctx.mark_non_differentiable(*[x for x in (entropy_sum, num_correct_tokens) if x is not None])
+        return loss, entropy_sum, num_correct_tokens
     @staticmethod
+    def backward(ctx, dloss, dentropy_sum, dnum_correct_tokens):
         hidden, weight, bias, targets, lse, valid, n_valid = ctx.saved_tensors
         N, H = hidden.shape
         V = weight.shape[0]
             One of `"mean"` (sum over non-ignored positions divided by their count), `"sum"` or
             `"none"` (per-position losses, `0` at ignored positions).
         return_metrics (`bool`, *optional*, defaults to `False`):
+            Whether to also return `num_correct_tokens` and `entropy_sum`, computed in the same
+            pass as the loss. Both are **sums over non-ignored positions**, not means: a mean
+            cannot be recombined across distributed ranks holding different token counts, and the
+            accuracy count is exact as an integer. Divide by the number of non-ignored positions
+            to recover the per-token means. Enabling this does not change the loss or the
+            gradients.
     Returns:
         `torch.Tensor` when `return_metrics=False`, otherwise a `tuple` of the loss and a `dict`
+        whose values are both **sums over non-ignored positions**:
+            - `num_correct_tokens` (`torch.Tensor`):
+                Scalar `int64`, the number of non-ignored positions whose argmax logit is the
+                target.
+            - `entropy_sum` (`torch.Tensor`):
+                Scalar `float32`, the summed softmax entropy (in nats) over non-ignored
+                positions.
     Examples:
     ```python
     >>> loss = fused_linear_cross_entropy(hidden, lm_head.weight, labels)
     >>> loss, metrics = fused_linear_cross_entropy(hidden, lm_head.weight, labels, return_metrics=True)
+    >>> n = (labels != -100).sum()
+    >>> mean_token_accuracy = metrics["num_correct_tokens"] / n
+    >>> mean_entropy = metrics["entropy_sum"] / n
     ```
     """
     if hidden.ndim != 2:
     if bias is not None:
         bias = bias.contiguous()
+    loss, entropy_sum, num_correct_tokens = _FusedLinearCrossEntropy.apply(
         hidden, weight, targets, bias, ignore_index, logit_scale, softcap, reduction, return_metrics
     )
     if return_metrics:
+        return loss, {"num_correct_tokens": num_correct_tokens, "entropy_sum": entropy_sum}
     return loss