Support skip_zeros in the XPU Triton optimizer kernels

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

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,
     )
 

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)