Repository navigation
Support skip_zeros in the XPU Triton optimizer kernels - #2097
Open
jiqing-feng wants to merge 1 commit into
Open
jiqing-feng wants to merge 1 commit into
jiqing-feng wants to merge 1 commit into
Conversation
`skip_zeros` leaves a parameter and its optimizer states completely untouched
wherever the gradient is exactly zero. It is useful for sparse/masked gradients
(embeddings, MoE experts that saw no tokens) where a plain optimizer would still
decay those rows towards zero via weight decay and the decaying momentum state.
The XPU/Triton backend raised `NotImplementedError` for it, so any caller that
set the flag hard-failed on XPU while silently working on CUDA.
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 for the default `skip_zeros=False` the
branch is compiled away and the generated kernel is unchanged. An A/B run of
all 14 optimizers x {fp32, fp16} against the unpatched build is bit-identical
(except LARS, whose `atomic_add` norm reduction is non-deterministic on its own)
and the measured step time is within noise at 4M and 64M elements.
Two deliberate differences from the CUDA kernels, which only implement the flag
partially and inconsistently:
* The zero test uses the raw gradient, before weight decay is folded in. CUDA's
32-bit 1-state kernel tests the post-decay gradient, which silently disables
`skip_zeros` for every non-zero parameter as soon as `weight_decay > 0`.
* AdEMAMix 32-bit and the 8-bit 2-state blockwise kernel honour the flag too;
CUDA ignores it there.
The CPU backend silently ignored the flag and produced wrong results; it now
raises `NotImplementedError` like the `default` backend until it is implemented.
jiqing-feng
force-pushed
the
xpu-triton-skip-zeros
branch
from
October 8, 2026 02:50
6284d52 to
688ace5
Compare
Contributor
Author
|
Hi @matthewdouglas . Would you please review this PR? Thanks! |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
skip_zerosis supported by the CUDA optimizer kernels but raisesNotImplementedErroron 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.whereright before the stores.SKIP_ZEROSis atl.constexpr, so the defaultskip_zeros=Falsecompiles 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
NotImplementedErrorlike thedefaultbackend.Tests
Adds
test_skip_zerostotests/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 theskip_zeros=Falseresult.Validation on Intel BMG
Correctness: 20 new tests pass on XPU. The rest of
tests/test_optim.pyis unchanged (same 30 pre-existingtest_optimizer32bitpaged 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_addnorm reduction is already non-deterministic run-to-run without this patch.No performance impact.