Stop an older unsloth_zoo turning a fully masked microbatch into NaN - #10892

Problem

unsloth_zoo's chunked fused cross entropy computes

divisor = n_items if n_items is not None else (labels != ignore_index).sum()

A microbatch with no trainable labels at all and no caller-supplied n_items therefore divides by zero. Loss and every gradient come back NaN, the optimizer state is poisoned, and the rest of the run is lost with no error raised.

Both halves really do co-occur:

• n_items is legitimately None. unsloth_zoo's own _unsloth_get_batch_samples clears num_items_in_batch when nothing downstream will divide by it (tests/test_loss_normalization_contract.py documents this).
• A sample truncated before its assistant turn is fully masked, so a small per_device_train_batch_size can make one an entire microbatch.

Reproduced end to end on Qwen3-0.6B with every 4th sample fully masked at batch_size=1: loss 12.32, then NaN for the remainder of the run.

Why here

The real fix is in unsloth_zoo (unslothai/unsloth-zoo#616). This exists for version skew -- a current unsloth against an older pinned unsloth_zoo -- and no-ops the moment zoo carries the fix. Detection is structural rather than a version compare, because the fix is not in a numbered release yet and dev builds and backports share release metadata.

Why it patches the autograd Function

Generated forwards under unsloth_compiled_cache/, and fused_losses/forward_adapter.py, do from unsloth_zoo.loss_utils import unsloth_fused_ce_loss, binding the function object into their own globals. Rebinding that name afterwards would never reach them.

That function resolves UnslothFusedLoss from its own module globals at call time, and torch resolves cls.forward at .apply() time, so a class level patch reaches already-imported callers with no cache invalidation. Verified on both routes into the kernel:

Both also pass with UNSLOTH_COMPILE_OVERWRITE=0.

The replacement preserves the original signature, because apply_autograd_function builds .apply()'s positional argument list from inspect.signature(forward), and clears the cached mapping afterwards. It short-circuits only the broken case and delegates everything else, so it does not reimplement chunking, the compile probe, or the backward contract. Emptiness is decided on the effective labels, after the causal shift and after the attention mask is folded in, since a batch can be non-empty before shifting and empty after.

Testing

Against a vulnerable unsloth_zoo 2026.9.3, 13/13:

• NaN reproduced before, finite 0.0 with zero finite gradients after
• dense and sparse batches, with and without n_items, still bit-match an unchunked reference
• the mask argument path, a trainable bias, and shift_labels=False
• idempotent; uninstall() restores the original NaN, proving the patch is what fixed it
• no-op verified against a zoo that already carries the fix

Dependency matrix, all green:

compileall, ruff --select E9,F63,F7,F82 clean; the new function ast-parses on Python 3.9 through 3.13.

Not affected, and not touched: GRPO uses its own autograd Function (UnslothEfficientGRPO), so vLLM / fast_inference runs do not reach this kernel, and unsloth_zoo/mlx/ has no reference to fused_losses at all.