From 688ace5f4bfee12d64467caf638254b05932e08c Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Thu, 8 Oct 2026 02:37:19 +0000 Subject: [PATCH] Support `skip_zeros` in the XPU Triton optimizer kernels `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. --- bitsandbytes/backends/cpu/ops.py | 6 ++ bitsandbytes/backends/triton/kernels_optim.py | 66 ++++++++++++++--- tests/test_optim.py | 70 +++++++++++++++++++ 3 files changed, 134 insertions(+), 8 deletions(-) diff --git a/bitsandbytes/backends/cpu/ops.py b/bitsandbytes/backends/cpu/ops.py index 44fb5dceb..8fadfe8b1 100755 --- a/bitsandbytes/backends/cpu/ops.py +++ b/bitsandbytes/backends/cpu/ops.py @@ -362,6 +362,9 @@ def _optimizer_update_32bit_cpu( gnorm_scale: float, skip_zeros: bool = False, ) -> None: + if skip_zeros: + raise NotImplementedError("skip_zeros is not supported yet") + g_float = g.float() * gnorm_scale p_float = p.data.float() @@ -487,6 +490,9 @@ def _optimizer_update_8bit_blockwise_cpu( gnorm_scale: float, skip_zeros: bool = False, ) -> None: + if skip_zeros: + raise NotImplementedError("skip_zeros is not supported yet") + blocksize = 256 # Dequantize states diff --git a/bitsandbytes/backends/triton/kernels_optim.py b/bitsandbytes/backends/triton/kernels_optim.py index f7eb2e213..d73e869af 100644 --- a/bitsandbytes/backends/triton/kernels_optim.py +++ b/bitsandbytes/backends/triton/kernels_optim.py @@ -168,7 +168,7 @@ def _optimizer_update_2state_32bit_triton_kernel( beta2_step, lr, gnorm_scale: tl.constexpr, - skip_zeros, + SKIP_ZEROS: tl.constexpr, n_elements, OPTIMIZER_ID: tl.constexpr, BLOCK_SIZE: tl.constexpr, @@ -187,9 +187,17 @@ def _optimizer_update_2state_32bit_triton_kernel( if OPTIMIZER_ID == 5: # ADEMAMIX s3_vals = tl.load(state1_ptr + n_elements + offsets, mask=mask, other=0.0) + s3_prev = s3_vals g_vals = gnorm_scale * g_vals + # skip_zeros: Triton has no per-lane early exit, so keep the pre-update values and + # select them back in before storing. Compiled away when SKIP_ZEROS is False. + nonzero_grad = g_vals != 0.0 + p_prev = p_vals + s1_prev = s1_vals + s2_prev = s2_vals + update_scale = 1.0 if max_unorm > 0.0: current_unorm = tl.sqrt(tl.load(unorm_ptr)) @@ -225,6 +233,13 @@ def _optimizer_update_2state_32bit_triton_kernel( adaptive_term = (tl.sqrt(s2_vals) / correction2) + eps p_vals = p_vals - lr * (mixed_momentum / adaptive_term) + if SKIP_ZEROS: + p_vals = tl.where(nonzero_grad, p_vals, p_prev) + s1_vals = tl.where(nonzero_grad, s1_vals, s1_prev) + s2_vals = tl.where(nonzero_grad, s2_vals, s2_prev) + if OPTIMIZER_ID == 5: # ADEMAMIX + s3_vals = tl.where(nonzero_grad, s3_vals, s3_prev) + tl.store(p_ptr + offsets, p_vals, mask=mask) tl.store(state1_ptr + offsets, s1_vals, mask=mask) tl.store(state2_ptr + offsets, s2_vals, mask=mask) @@ -253,7 +268,7 @@ def _optimizer_update_1state_32bit_triton_kernel( beta2_step, lr, gnorm_scale: tl.constexpr, - skip_zeros, + SKIP_ZEROS: tl.constexpr, n_elements, OPTIMIZER_ID: tl.constexpr, BLOCK_SIZE: tl.constexpr, @@ -270,6 +285,12 @@ def _optimizer_update_1state_32bit_triton_kernel( s1_vals = tl.load(state1_ptr + offsets, mask=mask, other=0.0) g_vals = gnorm_scale * g_vals + + # skip_zeros: test the raw gradient, before weight decay is folded in below. + nonzero_grad = g_vals != 0.0 + p_prev = p_vals + s1_prev = s1_vals + # Coupled (L2) weight decay: fold wd into the gradient. This is correct for # MOMENTUM/RMSPROP/ADAGRAD, but NOT for LION (id 4), which uses *decoupled* # (AdamW-style) weight decay applied to the param directly (see the LION branch @@ -317,6 +338,10 @@ def _optimizer_update_1state_32bit_triton_kernel( update_val = lr * g_vals / (tl.sqrt(s1_vals) + eps) p_vals = p_vals - update_val + if SKIP_ZEROS: + p_vals = tl.where(nonzero_grad, p_vals, p_prev) + s1_vals = tl.where(nonzero_grad, s1_vals, s1_prev) + tl.store(p_ptr + offsets, p_vals, mask=mask) tl.store(state1_ptr + offsets, s1_vals, mask=mask) @@ -380,9 +405,6 @@ def optimizer_update_32bit_impl( """ 32-bit optimizer implemented by Triton """ - if skip_zeros: - raise NotImplementedError("skip_zeros is not supported on XPU yet") - BLOCK_SIZE = 256 N_PER_TH = 1 # Number of blocks processed per thread. grid = (triton.cdiv(p.numel(), BLOCK_SIZE * N_PER_TH),) @@ -900,6 +922,7 @@ def _optimizer_update_1state_8bit_blockwise_triton_kernel( BLOCK_SIZE_N: tl.constexpr, N_PER_TH: tl.constexpr, OPTIMIZER_ID: tl.constexpr, + SKIP_ZEROS: tl.constexpr, ): """ Triton kernel for 8-bit optimizers that use one momentum state. @@ -916,6 +939,11 @@ def _optimizer_update_1state_8bit_blockwise_triton_kernel( p = tl.load(p_ptr + offsets, mask=mask, other=0.0).to(tl.float32) s1 = dequant_8bit_blockwise_kernel_util(state1_ptr, offsets, qmap1_ptr, absmax1_ptr, mask, BLOCK_SIZE_N) + # skip_zeros: test the raw gradient, before weight decay is folded in below. + nonzero_grad = g != 0.0 + p_prev = p + s1_prev = s1 + # 3. Optimizer-specific updates # LION if weight_decay > 0.0 and OPTIMIZER_ID == 2: @@ -950,6 +978,10 @@ def _optimizer_update_1state_8bit_blockwise_triton_kernel( s1 = s1 * beta2 + (1.0 - beta2) * g # 4. Store updated parameter and requantized state + if SKIP_ZEROS: + p = tl.where(nonzero_grad, p, p_prev) + s1 = tl.where(nonzero_grad, s1, s1_prev) + tl.store(p_ptr + offsets, p.to(p_ptr.dtype.element_ty), mask=mask) s1_codes, new_absmax1 = quantize_8bit_blockwise_kernel_util(s1, qmap1_ptr, 256, BLOCK_SIZE_N, N_PER_TH) tl.store(state1_ptr + offsets, s1_codes, mask=mask) @@ -985,6 +1017,7 @@ def _optimizer_update_2state_8bit_blockwise_triton_kernel( BLOCK_SIZE_N: tl.constexpr, N_PER_TH: tl.constexpr, OPTIMIZER_ID: tl.constexpr, + SKIP_ZEROS: tl.constexpr, ): """ Triton kernel for 8-bit optimizers that use two momentum states. @@ -1000,10 +1033,15 @@ def _optimizer_update_2state_8bit_blockwise_triton_kernel( g = tl.load(g_ptr + offsets, mask=mask, other=0.0).to(tl.float32) * gnorm_scale p = tl.load(p_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + nonzero_grad = g != 0.0 + p_prev = p + # 3. Optimizer-specific updates if OPTIMIZER_ID == 3: # ADAM s1 = dequant_8bit_blockwise_kernel_util(state1_ptr, offsets, qmap1_ptr, absmax1_ptr, mask, BLOCK_SIZE_N) s2 = dequant_8bit_blockwise_kernel_util(state2_ptr, offsets, qmap2_ptr, absmax2_ptr, mask, BLOCK_SIZE_N) + s1_prev = s1 + s2_prev = s2 s1 = s1 * beta1 + (1.0 - beta1) * g s2 = s2 * beta2 + (1.0 - beta2) * g * g @@ -1021,6 +1059,11 @@ def _optimizer_update_2state_8bit_blockwise_triton_kernel( denom = tl.sqrt(s2) / tl.sqrt(bias_correction2) + eps p -= (lr / bias_correction1) * (s1 / denom) + if SKIP_ZEROS: + p = tl.where(nonzero_grad, p, p_prev) + s1 = tl.where(nonzero_grad, s1, s1_prev) + s2 = tl.where(nonzero_grad, s2, s2_prev) + # Store updated parameter tl.store(p_ptr + offsets, p.to(p_ptr.dtype.element_ty), mask=mask) @@ -1045,6 +1088,9 @@ def _optimizer_update_2state_8bit_blockwise_triton_kernel( BLOCK_SIZE_N, ) nu = dequant_8bit_blockwise_kernel_util(state2_ptr, offsets, qmap2_ptr, absmax2_ptr, mask, BLOCK_SIZE_N) + m1_prev = m1 + m2_prev = m2 + nu_prev = nu m1 = m1 * beta1 + (1.0 - beta1) * g m2 = m2 * beta3 + (1.0 - beta3) * g @@ -1064,6 +1110,12 @@ def _optimizer_update_2state_8bit_blockwise_triton_kernel( p -= lr * update + if SKIP_ZEROS: + p = tl.where(nonzero_grad, p, p_prev) + m1 = tl.where(nonzero_grad, m1, m1_prev) + m2 = tl.where(nonzero_grad, m2, m2_prev) + nu = tl.where(nonzero_grad, nu, nu_prev) + # Store updated parameter tl.store(p_ptr + offsets, p.to(p_ptr.dtype.element_ty), mask=mask) @@ -1117,9 +1169,6 @@ def optimizer_update_8bit_blockwise_impl( gnorm_scale: float = 1.0, skip_zeros=False, ) -> None: - if skip_zeros: - raise NotImplementedError("skip_zeros is not supported on XPU yet") - if optimizer_name == "ademamix": # Handle AdEMAMIX's stacked state tensors if state1.dim() < 2 or state1.shape[0] != 2: @@ -1166,6 +1215,7 @@ def optimizer_update_8bit_blockwise_impl( BLOCK_SIZE_N=BLOCK_SIZE, N_PER_TH=N_PER_TH, OPTIMIZER_ID=optimizer_id, + SKIP_ZEROS=skip_zeros, num_warps=2, ) diff --git a/tests/test_optim.py b/tests/test_optim.py index 29736311d..cfe6f3650 100644 --- a/tests/test_optim.py +++ b/tests/test_optim.py @@ -292,6 +292,76 @@ def test_lion32bit_weight_decay(dim1, dim2, gtype, device): p2.copy_(p1.data) +skip_zeros_optimizers = { + "adam": lambda pxx: bnb.optim.Adam(pxx, lr=0.01, weight_decay=0.1), + "adam8bit_blockwise": lambda pxx: bnb.optim.Adam8bit(pxx, lr=0.01, weight_decay=0.1, min_8bit_size=0), + "momentum": lambda pxx: bnb.optim.SGD(pxx, lr=0.01, momentum=0.9, weight_decay=0.1), + "momentum8bit_blockwise": lambda pxx: bnb.optim.SGD8bit( + pxx, lr=0.01, momentum=0.9, weight_decay=0.1, min_8bit_size=0 + ), + "rmsprop": lambda pxx: bnb.optim.RMSprop(pxx, lr=0.01, alpha=0.9, weight_decay=0.1), + "rmsprop8bit_blockwise": lambda pxx: bnb.optim.RMSprop8bit( + pxx, lr=0.01, alpha=0.9, weight_decay=0.1, min_8bit_size=0 + ), + "lion": lambda pxx: bnb.optim.Lion(pxx, lr=0.01, weight_decay=0.1), + "lion8bit_blockwise": lambda pxx: bnb.optim.Lion8bit(pxx, lr=0.01, weight_decay=0.1, min_8bit_size=0), + "ademamix": lambda pxx: bnb.optim.AdEMAMix(pxx, lr=0.01, weight_decay=0.1), + "ademamix8bit_blockwise": lambda pxx: bnb.optim.AdEMAMix8bit(pxx, lr=0.01, weight_decay=0.1, min_8bit_size=0), +} + + +@pytest.mark.parametrize("optim_name", list(skip_zeros_optimizers), ids=id_formatter("opt")) +@pytest.mark.parametrize("gtype", [torch.float32, torch.float16], ids=describe_dtype) +@pytest.mark.parametrize("device", get_available_devices(), ids=id_formatter("device")) +@pytest.mark.skipif(not get_available_devices(), reason="No device") +def test_skip_zeros(optim_name, gtype, device): + """With `skip_zeros=True`, elements whose gradient is exactly 0 must be left untouched + (no weight decay, no state update), while all other elements get exactly the same + update as with `skip_zeros=False`. + """ + dim1, dim2 = 1024, 32 + + # Generate everything up-front so that both runs below see bit-identical inputs. + torch.manual_seed(1337) + p_init = (torch.randn(dim1, dim2, dtype=gtype) * 0.1).to(device) + warmup_grads = [(torch.randn(dim1, dim2, dtype=gtype) * 0.01).to(device) for _ in range(k)] + zeros = (torch.rand(dim1, dim2) < 0.5).to(device) + g = (torch.randn(dim1, dim2, dtype=gtype) * 0.01).to(device) + g[zeros] = 0.0 + + def run(skip_zeros): + mng = bnb.optim.GlobalOptimManager.get_instance() + mng.initialize() + p = p_init.clone() + if skip_zeros: + mng.override_config(p, "skip_zeros", True) + optimizer = skip_zeros_optimizers[optim_name]([p]) + + # Warm up with dense gradients so the states are populated and non-trivial. + for warmup_grad in warmup_grads: + p.grad = warmup_grad.clone() + optimizer.step() + + before = p.clone() + p.grad = g.clone() + optimizer.step() + return before, p.clone() + + try: + before_ref, p_ref = run(skip_zeros=False) + before_skip, p_skip = run(skip_zeros=True) + except NotImplementedError as e: + pytest.skip(f"skip_zeros is not supported on {device}: {e}") + + assert torch.equal(before_ref, before_skip), "warm-up is not deterministic" + + # Guard against a vacuous test: without skip_zeros the zero-gradient elements do move. + assert not torch.equal(p_ref[zeros], before_ref[zeros]), "test is vacuous" + + torch.testing.assert_close(p_skip[zeros], before_skip[zeros], rtol=0, atol=0) + torch.testing.assert_close(p_skip[~zeros], p_ref[~zeros], rtol=0, atol=0) + + @pytest.mark.parametrize("dim1", [1024], ids=id_formatter("dim1")) @pytest.mark.parametrize("dim2", [32, 1024, 4097], ids=id_formatter("dim2")) @pytest.mark.parametrize("gtype", [torch.float32, torch.float16], ids=describe_dtype)