kempnerforge.training.loss

Loss function registry for KempnerForge.

Registers loss functions and provides build_loss_fn() to compose them with config-driven options (chunk size, z-loss). Follows the same builder pattern as build_optimizer.

Functions

build_loss_fn(config)	Build a composed loss function from training config.
chunked_cross_entropy_loss(logits, labels[, ...])	Cross-entropy computed in token-dimension chunks.
cross_entropy_loss(logits, labels[, fp32])	Standard cross-entropy loss for language modeling.
z_loss(logits, weight)	Logit magnitude regularizer (PaLM / Gemini).

kempnerforge.training.loss.cross_entropy_loss(logits, labels, fp32=False)

Standard cross-entropy loss for language modeling.

Uses ignore_index=-100 so packed-sequence boundary tokens (labeled -100 by the dataset) are excluded from the loss. When no -100 labels are present, this has zero overhead. A batch with every label -100 returns an exact float32 zero that stays on the autograd graph of logits.

With fp32=True the loss and its gradient are computed in fp32 even when the logits are bf16 (train.fp32_cross_entropy).

Parameters:

• logits (torch.Tensor)

• labels (torch.Tensor)

• fp32 (bool)

Return type:

torch.Tensor

kempnerforge.training.loss.chunked_cross_entropy_loss(logits, labels, chunk_size=4096, fp32=False)

Cross-entropy computed in token-dimension chunks.

Chunks along the token dimension and uses PyTorch’s fused CE kernel per chunk, avoiding an explicit float32 materialization of the full logit tensor. For Llama-3 7B (vocab=128K, batch=4, seq=4096), the manual logsumexp path would create a ~8 GB float32 copy; this implementation avoids that entirely.

Uses ignore_index=-100 so packed-sequence boundary tokens are excluded. When no -100 labels are present, this has zero overhead. A batch with every label -100 returns an exact float32 zero that stays on the autograd graph of logits.

Note: the input logit tensor (B*S, V) is still fully materialized by the model’s output head before reaching this function. For deeper savings (never materializing the full logit tensor), the output projection itself must be chunked in the model forward pass — a future enhancement.

Parameters:

• logits (torch.Tensor) – (batch, seq, vocab) or (tokens, vocab).

• labels (torch.Tensor) – (batch, seq) or (tokens,).

• chunk_size (int) – Number of tokens per chunk.

• fp32 (bool) – Compute the loss and its gradient in fp32 from bf16 logits, still without an fp32 copy of the full logits (train.fp32_cross_entropy).

Return type:

torch.Tensor

kempnerforge.training.loss.z_loss(logits, weight)

Logit magnitude regularizer (PaLM / Gemini).

Penalizes large logit magnitudes to prevent logit drift that causes NaN/divergence in long training runs. Negligible compute cost.

Formula: weight * mean(logsumexp(logits, dim=-1) ** 2)

Parameters:

• logits (torch.Tensor) – Model output logits, shape (batch, seq, vocab) or (tokens, vocab).

• weight (float) – Regularization weight (PaLM uses 1e-4).

Returns:

Scalar z-loss term to add to the main loss; zero on the graph for logits with no rows (a packed micro-batch without text).

Return type:

torch.Tensor

kempnerforge.training.loss.build_loss_fn(config)

Build a composed loss function from training config.

Follows the build_optimizer pattern: config in, callable out. Binds chunk_size for chunked CE and fp32 when train.fp32_cross_entropy is set, and composes z-loss, so the caller gets a clean (logits, labels) -> Tensor interface.

Return type:

Callable[[torch.Tensor, torch.Tensor], torch.Tensor]