Merge branch 'main' into egor/8bit_opt2

Author: Egor-Krivov

bitsandbytes/_ops.py

@@ -352,7 +352,7 @@ def _(
 
 torch.library.define(
     "bitsandbytes::optimizer_update_32bit",
-    "(str optimizer_name, Tensor g, Tensor p, Tensor state1, Tensor! state2, Tensor! unorm_vec, float max_unorm, float param_norm, float beta1, float beta2, float beta3, float alpha, float eps, float weight_decay, int step, float lr, float gnorm_scale, bool skip_zeros) -> ()",
+    "(str optimizer_name, Tensor(a0!) g, Tensor(a1!) p, Tensor(a2!) state1, Tensor(a3!)? state2, Tensor(a4!)? unorm_vec, float max_unorm, float param_norm, float beta1, float beta2, float beta3, float alpha, float eps, float weight_decay, int step, float lr, float gnorm_scale, bool skip_zeros=False) -> ()",
 )
 
 
@@ -395,7 +395,7 @@ def _(
 
 torch.library.define(
     "bitsandbytes::optimizer_update_8bit_blockwise",
-    "(str optimizer_name, Tensor g, Tensor p, Tensor state1, Tensor! state2, float beta1, float beta2, float beta3, float alpha, float eps, int step, float lr, Tensor qmap1, Tensor! qmap2, Tensor absmax1, Tensor! absmax2, float weight_decay, float gnorm_scale, bool skip_zeros) -> ()",
+    "(str optimizer_name, Tensor(a0!) g, Tensor(a1!) p, Tensor(a2!) state1, Tensor(a3!)? state2, float beta1, float beta2, float beta3, float alpha, float eps, int step, float lr, Tensor(a4!) qmap1, Tensor(a5!)? qmap2, Tensor(a6!) absmax1, Tensor(a7!)? absmax2, float weight_decay, float gnorm_scale, bool skip_zeros=False) -> ()",
 )
 
 
@@ -417,8 +417,8 @@ def _(
     qmap2: Optional[torch.Tensor],
     absmax1: torch.Tensor,
     absmax2: Optional[torch.Tensor],
-    weight_decay: float = 0.0,
-    gnorm_scale: float = 1.0,
+    weight_decay: float,
+    gnorm_scale: float,
     skip_zeros=False,
 ) -> None:
     torch._check(

bitsandbytes/backends/cuda/ops.py

@@ -686,8 +686,8 @@ def _optimizer_update_8bit_blockwise_impl(
     qmap2: Optional[torch.Tensor],
     absmax1: torch.Tensor,
     absmax2: Optional[torch.Tensor],
-    weight_decay: float = 0.0,
-    gnorm_scale: float = 1.0,
+    weight_decay: float,
+    gnorm_scale: float,
     skip_zeros=False,
 ) -> None:
     # torch._check(

tests/helpers.py

@@ -21,7 +21,8 @@
 def get_available_devices(no_cpu=False):
     if "BNB_TEST_DEVICE" in os.environ:
         # If the environment variable is set, use it directly.
-        return [d for d in os.environ["BNB_TEST_DEVICE"] if d.lower() != "cpu"]
+        device = os.environ["BNB_TEST_DEVICE"]
+        return [] if no_cpu and device == "cpu" else [device]
 
     devices = [] if HIP_ENVIRONMENT else ["cpu"] if not no_cpu else []
 

tests/test_optim.py

@@ -170,6 +170,7 @@ def rm_path(path):
 @pytest.mark.parametrize("dim1", [1024], ids=id_formatter("dim1"))
 @pytest.mark.parametrize("dim2", [32, 1024, 4097, 1], ids=id_formatter("dim2"))
 @pytest.mark.parametrize("device", get_available_devices(no_cpu=True), ids=id_formatter("device"))
+@pytest.mark.skipif(not get_available_devices(no_cpu=True), reason="No device")
 def test_optimizer32bit(dim1, dim2, gtype, optim_name, device):
     if optim_name.startswith("paged_") and sys.platform == "win32":
         pytest.skip("Paged optimizers can have issues on Windows.")
@@ -250,6 +251,7 @@ def test_optimizer32bit(dim1, dim2, gtype, optim_name, device):
 @pytest.mark.parametrize("dim2", [32, 1024, 4097], ids=id_formatter("dim2"))
 @pytest.mark.parametrize("gtype", [torch.float32, torch.float16], ids=describe_dtype)
 @pytest.mark.parametrize("device", get_available_devices(no_cpu=True))
+@pytest.mark.skipif(not get_available_devices(no_cpu=True), reason="No device")
 def test_global_config(dim1, dim2, gtype, device):
     if dim1 == 1 and dim2 == 1:
         return
@@ -306,6 +308,7 @@ def test_global_config(dim1, dim2, gtype, device):
 @pytest.mark.parametrize("dim2", [32, 1024, 4097], ids=id_formatter("dim2"))
 @pytest.mark.parametrize("dim1", [1024], ids=id_formatter("dim1"))
 @pytest.mark.parametrize("device", get_available_devices(no_cpu=True))
+@pytest.mark.skipif(not get_available_devices(no_cpu=True), reason="No device")
 def test_optimizer8bit(dim1, dim2, gtype, optim_name, device):
     torch.set_printoptions(precision=6)