What is z-loss and when would you add it?

Z-loss is an auxiliary term defined as α(log Z(x))^2 added to the cross-entropy loss. It forces the output softmax normalizer toward one, preventing logit scale drift that causes numerical instability in bfloat16. It is typically applied from step one with α = 10^-4. It costs nothing in compute and does not change target predictions, only the logit scale.