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)