Chunked Linear Cross-Entropy

Chunked Linear Cross-Entropy is the optimized cross-entropy specialization of the chunked linear loss strategy. It chunks the final projection, applies a fused cross-entropy-and-logits-gradient kernel to each temporary logits tile, and avoids materializing the full [tokens, vocab] logits tensor.

For the generic scheduling strategy and BYO loss API, see Strategies: Chunked Linear Loss.

Kernel design

The implementation chunks over the batch/token dimension. For each chunk, it:

1. Computes a temporary logits tile: logits_chunk = x_chunk @ W
2. Runs the fused cross-entropy-and-logits-gradient kernel on that tile
3. Adds the chunk’s loss into the total loss
4. Reuses the tile as dlogits, the gradient with respect to the logits
5. Immediately uses dlogits to accumulate the input and weight gradients

The standalone cross-entropy kernel computes the numerically stable row loss and the logits gradient together. The outer autograd backward step scales and returns the chunk-accumulated grad_input and grad_weight.

Usage

Functional API

import torch
from greyhound.nn.functional import chunked_linear_cross_entropy

# inputs: [num_tokens, hidden_dim], weight: [vocab_size, hidden_dim]
inputs = torch.randn(8192, 4096, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(32000, 4096, device="cuda", dtype=torch.bfloat16)
target = torch.randint(0, 32000, (8192,), device="cuda")

loss = chunked_linear_cross_entropy(inputs, weight, target)
loss.backward()

# With z-loss regularization
loss = chunked_linear_cross_entropy(inputs, weight, target, z_loss_multiplier=1e-4)

# With sum reduction and ignore_index
loss = chunked_linear_cross_entropy(
    inputs, weight, target,
    ignore_index=-100,
    reduction="sum",
)

# Opt into lower-memory/faster weight-dtype gradient accumulation
loss = chunked_linear_cross_entropy(
    inputs,
    weight,
    target,
    grad_weight_accum_dtype="weight",
)

greyhound.nn.functional.chunked_linear_cross_entropy

Benchmarks

Benchmarks use bfloat16 inputs and a batch size of 8192 tokens. The x-axis is the vocabulary size.

Full pass

Forward pass

These standalone cross-entropy plots provide the forward-only comparison point for the loss computation itself:

Backward pass

Full pass

Standalone cross-entropy

Full pass

Standalone cross-entropy