remove hook for bnb 4-bit (#3223)

* relax dispatch for bnb

* style

src/accelerate/big_modeling.py

@@ -37,6 +37,7 @@
     find_tied_parameters,
     get_balanced_memory,
     infer_auto_device_map,
+    is_bnb_available,
     is_mlu_available,
     is_musa_available,
     is_npu_available,
@@ -351,16 +352,19 @@ def dispatch_model(
     # Error early if the device map is incomplete.
     check_device_map(model, device_map)
 
-    # for backward compatibility
-    is_bnb_quantized = (
-        getattr(model, "is_quantized", False) or getattr(model, "is_loaded_in_8bit", False)
-    ) and getattr(model, "quantization_method", "bitsandbytes") == "bitsandbytes"
+    # We need to force hook for quantized model that can't be moved with to()
+    if getattr(model, "quantization_method", "bitsandbytes") == "bitsandbytes":
+        # since bnb 0.43.2, we can move 4-bit model
+        if getattr(model, "is_loaded_in_8bit", False) or (
+            getattr(model, "is_loaded_in_4bit", False) and not is_bnb_available(min_version="0.43.2")
+        ):
+            force_hooks = True
 
     # We attach hooks if the device_map has at least 2 different devices or if
     # force_hooks is set to `True`. Otherwise, the model in already loaded
     # in the unique device and the user can decide where to dispatch the model.
     # If the model is quantized, we always force-dispatch the model
-    if (len(set(device_map.values())) > 1) or is_bnb_quantized or force_hooks:
+    if (len(set(device_map.values())) > 1) or force_hooks:
         if main_device is None:
             if set(device_map.values()) == {"cpu"} or set(device_map.values()) == {"cpu", "disk"}:
                 main_device = "cpu"

src/accelerate/utils/imports.py

@@ -179,8 +179,13 @@ def is_8bit_bnb_available():
     return False
 
 
-def is_bnb_available():
-    return _is_package_available("bitsandbytes")
+def is_bnb_available(min_version=None):
+    package_exists = _is_package_available("bitsandbytes")
+    if package_exists and min_version is not None:
+        bnb_version = version.parse(importlib.metadata.version("bitsandbytes"))
+        return compare_versions(bnb_version, ">=", min_version)
+    else:
+        return package_exists
 
 
 def is_bitsandbytes_multi_backend_available():