[jiqing-feng]

Motivation

skip_zeros is supported by the CUDA optimizer kernels but raises NotImplementedError on XPU. This PR implements it so the XPU backend matches CUDA.

Implementation

Triton has no per-lane early exit, so the pre-update values are kept in registers and selected back in with tl.where right before the stores. SKIP_ZEROS is a tl.constexpr, so the default skip_zeros=False compiles the branch away entirely.

Covers all four kernels: 32-bit 1-state and 2-state, 8-bit blockwise 1-state and 2-state.

The zero test is applied to the raw gradient, before weight decay is folded in, so the flag stays meaningful when weight_decay > 0.

The CPU backend silently ignored the flag and produced incorrect results; it now raises NotImplementedError like the default backend.

Tests

Adds test_skip_zeros to tests/test_optim.py, parameterized over 10 optimizers x {fp32, fp16} x available devices. It asserts that elements with a zero gradient are bit-identical to their pre-step values, and that all other elements are bit-identical to the skip_zeros=False result.

Validation on Intel BMG

Correctness: 20 new tests pass on XPU. The rest of tests/test_optim.py is unchanged (same 30 pre-existing test_optimizer32bit paged failures before and after).

No change to the default path: an A/B run of 14 optimizers x {fp32, fp16} x 5 steps against the unpatched build produces bit-identical parameters and states. The only exception is LARS, whose atomic_add norm reduction is already non-deterministic run-to-run without this patch.

No performance impact.

[jiqing-feng]

Hi @matthewdouglas . Would you please review this PR? Thanks!

[github-actions]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[matthewdouglas]

Thanks for the PR. It seems there's test failures on CUDA, can we please correct this?

[jiqing-feng]

Fixed. The CUDA failures came from the new test_skip_zeros, not from the Triton kernels. CUDA only partially honors skip_zeros: the AdEMAMix 32-bit and 8-bit 2-state kernels ignore it, and the 32-bit 1-state kernels check for zero after weight decay is folded into the gradient, so with weight_decay > 0 nothing is skipped. The test now runs only on XPU, where the Triton kernels implement it fully. Making the CUDA kernels consistent would be a separate PR; happy to do that if you want it.