Metrics return sums (num_correct_tokens, entropy_sum) instead of means

00e29bd verified 2 months ago

3.85 kB

---
library_name: kernels
license: apache-2.0
tags:
- kernel
- triton
- cross-entropy
---
# fused-linear-ce

A Triton **fused linear cross-entropy** kernel: it computes the cross-entropy of
`hidden @ weight.T (+ bias)` against `targets` **without ever materialising the `(N, V)` logit
matrix**. At a large vocabulary this turns the single most expensive step of LM training from a
tens-of-gigabytes operation into a flat, small one.

## Usage

```python
# pip install -U kernels
from kernels import get_kernel
flce = get_kernel("qgallouedec/fused-linear-ce", version=1, trust_remote_code=True)
loss = flce.fused_linear_cross_entropy(hidden, lm_head.weight, labels)
```

Full signature:

```python
fused_linear_cross_entropy(
    hidden,             # (N, H)   float32 / bfloat16 / float16
    weight,             # (V, H)   same dtype as hidden
    targets,            # (N,)     int64
    bias=None,          # (V,) or None
    ignore_index=-100,
    logit_scale=1.0,
    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**
to `return_metrics=False`.

## Algorithm

Implements the method described in
[*Cut Your Losses in Large-Vocabulary Language Models*](https://huggingface.co/papers/2411.09009)
(Wijmans et al., 2024).

The forward pass tiles over the vocabulary. For each tile it computes the logit tile with `tl.dot`
into an fp32 accumulator, applies bias / `logit_scale` / softcap, and merges the tile into a
running log-sum-exp using the online-softmax rescaling that flash-attention uses. Nothing of size
`(N, V)` is ever written. The backward pass **recomputes** each logit tile rather than storing it,
forms `p - onehot`, and accumulates `grad_hidden = (p - onehot) @ weight` and
`grad_weight = (p - onehot).T @ hidden`. Recomputation is cheaper than the memory traffic it
avoids.

## Clean-room statement

This is a **clean-room implementation**, written from the published algorithm and from the
standard online-softmax formulation. It contains **no code derived from `apple/ml-cross-entropy`**
or from any other implementation of the method. The upstream package was used only as a
black-box numerical and performance oracle. Licensed **Apache-2.0**.

## Backends

`build/torch-cuda` only. The sources are pure Triton and carry no CUDA-specific intrinsics, but
ROCm and XPU builds are not shipped because they were not tested.

	---
	library_name: kernels
	license: apache-2.0
	tags:
	- kernel
	- triton
	- cross-entropy
	---

	# fused-linear-ce

	A Triton fused linear cross-entropy kernel: it computes the cross-entropy of
	`hidden @ weight.T (+ bias)` against `targets` **without ever materialising the `(N, V)` logit
	matrix**. At a large vocabulary this turns the single most expensive step of LM training from a
	tens-of-gigabytes operation into a flat, small one.

	## Usage

	```python
	# pip install -U kernels
	from kernels import get_kernel

	flce = get_kernel("qgallouedec/fused-linear-ce", version=1, trust_remote_code=True)

	loss = flce.fused_linear_cross_entropy(hidden, lm_head.weight, labels)
	```

	Full signature:

	```python
	fused_linear_cross_entropy(
	hidden, # (N, H) float32 / bfloat16 / float16
	weight, # (V, H) same dtype as hidden
	targets, # (N,) int64
	bias=None, # (V,) or None
	ignore_index=-100,
	logit_scale=1.0,
	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
	to `return_metrics=False`.

	## Algorithm

	Implements the method described in
	[Cut Your Losses in Large-Vocabulary Language Models](https://huggingface.co/papers/2411.09009)
	(Wijmans et al., 2024).

	The forward pass tiles over the vocabulary. For each tile it computes the logit tile with `tl.dot`
	into an fp32 accumulator, applies bias / `logit_scale` / softcap, and merges the tile into a
	running log-sum-exp using the online-softmax rescaling that flash-attention uses. Nothing of size
	`(N, V)` is ever written. The backward pass recomputes each logit tile rather than storing it,
	forms `p - onehot`, and accumulates `grad_hidden = (p - onehot) @ weight` and
	`grad_weight = (p - onehot).T @ hidden`. Recomputation is cheaper than the memory traffic it
	avoids.

	## Clean-room statement

	This is a clean-room implementation, written from the published algorithm and from the
	standard online-softmax formulation. It contains no code derived from `apple/ml-cross-entropy`
	or from any other implementation of the method. The upstream package was used only as a
	black-box numerical and performance oracle. Licensed Apache-2.0.

	## Backends

	`build/torch-cuda` only. The sources are pure Triton and carry no CUDA-specific intrinsics, but
	ROCm and XPU builds are not shipped because they were not tested.