Losses¶

Two loss functions are registered in kempnerforge/training/loss.py: cross_entropy and chunked_cross_entropy. z-loss is a regularizer that wraps either — not a separate registry entry.

`cross_entropy`¶

[train]
loss_fn = "cross_entropy"

Direct call to F.cross_entropy over flattened (batch·seq_len, vocab_size) logits and (batch·seq_len,) labels. ignore_index = -100 is hard-coded; labels with -100 are skipped in the reduction. This is how document-packed sequences mask out cross-document attention — tokens at packing boundaries get label -100 so the loss doesn’t cross document edges.

When every label is -100 it returns an exact fp32 zero that stays on the autograd graph of the logits, so there is no NaN from dividing by zero valid tokens and backward and FSDP’s collectives still run on that rank.

`chunked_cross_entropy`¶

[train]
loss_fn = "chunked_cross_entropy"
ce_chunk_size = 4096                    # 0 -> auto 4096

Identical output to cross_entropy, but computed in chunks along the flattened token dimension. Use this when the full logit tensor would be large enough to matter — a 7B Llama-3 with batch=4, seq=4096, vocab=128000 materializes an 8 GB float32 cross-entropy intermediate if computed end-to-end.

Mechanics:

The loop iterates for i in range(0, num_tokens, chunk_size) and accumulates per-chunk loss with reduction="sum".
At the end, divides by the number of non-ignored tokens across all chunks to recover the mean.
ignore_index = -100 is threaded through per chunk.

Important: chunking saves memory inside the loss kernel, not in the output projection. The (batch, seq_len, vocab_size) logit tensor that the model produces is still materialized — chunked_cross_entropy only avoids the second full-size tensor that F.cross_entropy internally allocates for the softmax.

ce_chunk_size = 0 in the config is a sentinel for “use the default 4096” — see build_loss_fn.

fp32 cross-entropy¶

[train]
loss_fn = "cross_entropy"               # or chunked_cross_entropy
fp32_cross_entropy = true               # default false

Under mixed_precision = "bf16" the output head returns bf16 logits, and both losses above compute the log-softmax, the loss and its gradient in bf16. Stacks that run the forward pass under torch.autocast compute cross-entropy in fp32 instead. fp32_cross_entropy = true does the same for either loss, without the fp32 copy of the logits (8.4 GB for batch=4, seq=4096, vocab=128256) that upcasting them would allocate, and without the fp32 log-probabilities autograd would keep for backward. It changes training numerics slightly. With the default false, both losses are unchanged.

Both losses then run one custom autograd function in fp32 (fp64 logits stay fp64):

Forward: takes the log-sum-exp and target logit in fp32 and keeps only the per-token log-sum-exp (one fp32 value per token).
Backward: recomputes the softmax from that log-sum-exp, subtracts the one-hot target, and writes the gradient in the logits’ dtype. Rows labeled -100 get a zero gradient.

cross_entropy does this over the whole batch at once; chunked_cross_entropy does it ce_chunk_size tokens at a time. On CUDA both steps are torch.compiled (once per process, as FlexAttention is) into fused kernels that read the bf16 logits directly. On CPU the same functions run eagerly.

Measured on one H100 for batch=4, seq=4096, vocab=128256 (bf16 logits of 3.9 GiB), loss forward and backward:

	time	memory beyond the logits	loss error vs fp64
`cross_entropy` (bf16)	12.4 ms	11.7 GiB	2.1e-3
`F.cross_entropy(logits.float())`	35.4 ms	23.5 GiB	4.3e-8
`chunked_cross_entropy`, `fp32_cross_entropy = true`, `ce_chunk_size = 4096`	8.8 ms	4.9 GiB	7.5e-8

The first call on each new shape compiles for a few seconds.

Limits:

If torch.compile is disabled or falls back to eager on CUDA, whole-batch cross_entropy holds several fp32 tensors the size of the logits at once in its backward; chunked_cross_entropy bounds them to one chunk. Use chunked_cross_entropy if compile is off.
The backward cannot itself be differentiated: with create_graph=True the gradient comes back detached, so a second backward through it raises.
The flag is rejected for any loss_fn other than these two.

z-loss (logit magnitude regularizer)¶

[train]
loss_fn = "cross_entropy"               # or chunked_cross_entropy
z_loss_weight = 1e-4                    # 0 -> disabled

z-loss penalizes large log-sum-exp values — equivalent to penalizing the softmax normalizer from running off. It is a PaLM / Gemma trick for keeping logits in a numerically sane range without clamping.

Formula:

z_loss = z_loss_weight * mean(logsumexp(logits)^2)

How it wires in: build_loss_fn wraps the base loss with a closure that adds z_loss(logits, z_weight) to the returned loss. The wrap happens once at setup; the per-step overhead is one logsumexp over the logit tensor.

Recommended weights from the literature:

Value	Source
`1e-4`	PaLM
`1e-4` – `2e-4`	Gemma
`0.0`	off; default

z_loss_weight = 0.0 makes the wrap a no-op — the closure still runs but z_loss() returns 0.0 immediately without computing logsumexp.

MoE auxiliary loss is separate¶

The MoE load-balancing loss is not part of this pipeline. It lives on the model (model.get_moe_aux_loss()) and is added to the base loss in the training loop:

loss = loss_fn(logits, labels)
if mc.is_moe:
    loss = loss + mc.moe_aux_loss_weight * model.get_moe_aux_loss()

See MoE for how get_moe_aux_loss() sums the router aux losses: one per MoE layer, and on a MoT model one per modality on each MoE layer.

Picking one¶

Situation	Loss	z-loss
Small model, tight compile, small vocab	`cross_entropy`	off
7B+ with large vocab (≥ 50K)	`chunked_cross_entropy`	optional
Muon / aggressive LRs / long runs	either	`1e-4`
PaLM / Gemma-style recipe	`chunked_cross_entropy`	`1e-4`

The shipped 7b_16gpu_muon.toml is the reference combination: Muon + chunked cross-entropy + z-loss — exercises all three paths together.