cut-cross-entropy

Cross-entropy over a large vocabulary without ever materializing the logits, loadable through kernels. The reference baseline is F.cross_entropy on materialized fp32 logits, matched below 1e-5 relative on the loss with gradients at bf16 rounding. Companion to det-cross-entropy, which fixes the same reduction for reproducibility rather than memory. After Wijmans et al., Cut Your Losses in Large-Vocabulary Language Models (2024).

The loss of a language model is a softmax over E @ C.T, and the usual route forms the [N, V] logit tensor: at 65,536 tokens against a 128,256-token vocabulary that is 16 GB in bf16, held by autograd along with the log-softmax until backward. Nothing in the loss needs it to exist. This kernel blocks the vocabulary into panels, folds each panel into a running online-softmax state, and rebuilds panels during backward, so the training step that could not fit on the card at all simply runs, in gigabytes instead of tens of gigabytes.

Forward and backward at vocabulary 128,256, measured live: the materialized path climbs 16 to 63 GB, crawls at 14 s once past the card, and at 65,536 tokens raises CUDA out of memory; the blocked path runs every row in 1.2 to 3.4 GB, agreeing with the reference wherever both run.

Usage

import torch
from kernels import get_kernel

cce = get_kernel("phanerozoic/cut-cross-entropy", version=1, trust_remote_code=True)

hidden = model.model(input_ids).last_hidden_state      # [B, T, D] bf16
loss = cce.linear_cross_entropy(
    hidden.reshape(-1, hidden.shape[-1]),              # [N, D]
    model.lm_head.weight,                              # [V, D]
    labels.reshape(-1),                                # [N] int64
)
loss.backward()

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark. The classifier is passed at call time, so its gradient lands in lm_head.weight.grad as usual.

API

Method

The forward needs only three per-token quantities: the row maximum, the row sum of exponentials, and the target logit; the backward needs only the softmax, a fixed function of that state rebuildable one block at a time. Both passes block the vocabulary into BV-column panels. Each panel's logits are one bf16 tensor-core GEMM (fp32 accumulate); the forward folds the panel into a running (max, sumexp) by the online-softmax rescale and captures the target logit as its column passes. The backward recomputes the panel, rebuilds p = exp(x - lse), subtracts the one-hot, and contracts into dE and dC; a panel owns a disjoint vocabulary slice, so its dC block is a single beta = 0 write with no atomics. Resident cost is the [N, BV] panel rather than [N, V].

Measured

Peak allocation and wall-clock for a full forward and backward against F.cross_entropy on materialized logits, bf16, D = 2048 unless noted:

Two regimes. While the materialized path fits it is 1.3x to 1.6x faster, because this kernel recomputes each logit panel in the backward instead of storing it. At 65,536 tokens the materialized path wants 48.6 GB on a 48 GB card and takes 39.5 seconds where memory oversubscription exists, or fails outright where it does not; the blocked path uses 3.1 GB and 2.0 seconds. The crossover is not a tuning parameter but the point where [N, V] stops fitting.

Correctness

• Loss: maximum relative deviation below 1e-5 across shapes to (16, 2048, 128256) and all three reductions; most cases bitwise.
• Gradients: dE and dC agree with autograd through the materialized path to 2.6e-3 relative worst case (bf16 gradient rounding).
• Masking: ignore_index rows contribute nothing and receive exactly zero gradient; a fully masked batch gives zero loss and finite zero gradients.
• A ragged vocabulary with a target inside the short final panel matches to 1e-4; repeated forward passes are bitwise identical.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+ and bf16 tensor cores.
• E [N, D] and C [V, D] bf16, targets [N] int64.
• The [V, D] weight gradient is unavoidable; this kernel removes the [N, V] term. Small vocabularies leave little to save.
• Slower than the materialized path whenever that path fits comfortably; use this where the logit tensor is the constraint.

References

Wijmans, Huval, Hertzberg, Koltun, Krähenbühl, "Cut Your Losses in Large-Vocabulary Language Models" (2024); Milakov and Gimelshein, "Online normalizer calculation for softmax" (2018).

License

Apache-2.0.