transformers 5.0.0rc0__py3-none-any.whl → 5.0.0rc2__py3-none-any.whl

transformers/quantizers/quantizer_bnb_4bit.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from collections import defaultdict
 from typing import TYPE_CHECKING
 
 from .base import HfQuantizer
@@ -38,43 +37,20 @@ if is_torch_available():
     import torch
 
     from ..core_model_loading import WeightConverter
-    from ..pytorch_utils import Conv1D
 
 logger = logging.get_logger(__name__)
 
 
 class Bnb4BitHfQuantizer(HfQuantizer):
     """
-    4-bit quantization from bitsandbytes quantization method:
-        before loading: converts transformer layers into Linear4bit during loading: load 16bit weight and pass to the
-        layer object after: quantizes individual weights in Linear4bit into 4bit at the first .cuda() call
-        saving:
-            from state dict, as usual; saves weights and `quant_state` components
-        loading:
-            need to locate `quant_state` components and pass to Param4bit constructor
+    4-bit quantization from bitsandbytes quantization method
     """
 
-    use_keep_in_fp32_modules = True
-    requires_parameters_quantization = True
     requires_calibration = False
 
-    required_packages = ["bitsandbytes", "accelerate"]
-
     def __init__(self, quantization_config, **kwargs):
         super().__init__(quantization_config, **kwargs)
 
-        if self.quantization_config.llm_int8_skip_modules is not None:
-            self.modules_to_not_convert = self.quantization_config.llm_int8_skip_modules
-
-        # This describes the additional items that are saved on the state dict (on the params themselves)
-        self.bnb_keys = [
-            f"quant_state.bitsandbytes__{self.quantization_config.bnb_4bit_quant_type}",
-            "absmax",
-            "quant_map",
-        ]
-        if self.quantization_config.bnb_4bit_use_double_quant:
-            self.bnb_keys.extend(["nested_absmax", "nested_quant_map"])
-
     def validate_environment(self, *args, **kwargs):
         if not is_accelerate_available():
             raise ImportError(
@@ -90,17 +66,9 @@ class Bnb4BitHfQuantizer(HfQuantizer):
         validate_bnb_backend_availability(raise_exception=True)
 
         device_map = kwargs.get("device_map")
-        if (
-            device_map is not None
-            and isinstance(device_map, dict)
-            and not self.quantization_config.llm_int8_enable_fp32_cpu_offload
-        ):
-            device_map_without_lm_head = {
-                key: device_map[key] for key in device_map if key not in self.modules_to_not_convert
-            }
-            if set(device_map.values()) == {"cpu"}:
-                pass
-            elif "cpu" in device_map_without_lm_head.values() or "disk" in device_map_without_lm_head.values():
+        if not self.quantization_config.llm_int8_enable_fp32_cpu_offload and isinstance(device_map, dict):
+            values = set(device_map.values())
+            if values != {"cpu"} and ("cpu" in values or "disk" in values):
                 raise ValueError(
                     "Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit the "
                     "quantized model. If you want to dispatch the model on the CPU or the disk while keeping these modules "
@@ -110,115 +78,25 @@ class Bnb4BitHfQuantizer(HfQuantizer):
                     "for more details. "
                 )
 
-    def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype":
-        from accelerate.utils import CustomDtype
-
-        if target_dtype != torch.int8:
-            logger.info("target_dtype {target_dtype} is replaced by `CustomDtype.INT4` for 4-bit BnB quantization")
-        return CustomDtype.INT4
+    def param_element_size(self, model: "PreTrainedModel", param_name: str, param: "torch.Tensor") -> float:
+        "Return the element size (in bytes) for `param_name`."
+        if self.param_needs_quantization(model, param_name):
+            # 4 bit
+            return 0.5
 
-    def update_unexpected_keys(self, model, unexpected_keys: list[str]) -> list[str]:
-        return [k for k in unexpected_keys if not any(k.endswith(x) for x in self.bnb_keys)]
+        return super().param_element_size(model, param_name, param)
 
     def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
         import bitsandbytes as bnb
 
-        # They are on the params themselves, so we cannot easily extract the module from the name
-        if any(param_name.endswith(x) for x in self.bnb_keys):
-            return True
         module, name = get_module_from_name(model, param_name)
         return isinstance(module, bnb.nn.Linear4bit) and name != "bias"
 
-    def get_param_name(self, param_name: str) -> str:
-        """
-        Get the right param_name in order to get the module associated with the param.
-        This is useful for quantized stats lile absmax or quant_map as we need to update the param_name to get the module as they are stored in ...weight.absmax.
-        """
-        if self.pre_quantized:
-            # We need to get the param name of quantized weights and not its components. Otherwise, we won't be able to get the nn.Module associated.
-            if any(param_name.endswith(x) for x in self.bnb_keys):
-                param_name = (
-                    param_name.rsplit(".", 1)[0] if "quant_state." not in param_name else param_name.rsplit(".", 2)[0]
-                )
-        return param_name
-
-    def create_quantized_param(
-        self,
-        model: "PreTrainedModel",
-        param_value: "torch.Tensor",
-        param_name: str,
-        target_device: "torch.device",
-        **kwargs,
-    ):
-        import bitsandbytes as bnb
-
-        full_name = param_name
-
-        # update param name to get the weights instead of the quantized stats
-        param_name = self.get_param_name(param_name)
-        module, tensor_name = get_module_from_name(model, param_name)
-
-        # `torch.Tensor.to(<int num>)` is not supported by `torch_npu` (see this [issue](https://github.com/Ascend/pytorch/issues/16)).
-        if isinstance(target_device, int) and is_torch_npu_available():
-            target_device = f"npu:{target_device}"
-
-        # construct `new_value` for the module._parameters[tensor_name]
-        if self.pre_quantized:
-            module_name = param_name.rsplit(".", 1)[0]
-            # Save the states for later quantization when they are all gathered
-            if not hasattr(self, "param_quant_stats"):
-                self.param_quant_stats = defaultdict(dict)
-            self.param_quant_stats[module_name].update({full_name: param_value})
-
-            # We are ready for quantization in this case (note, the +1 is for the weight itself)
-            if len(self.param_quant_stats[module_name]) == len(self.bnb_keys) + 1:
-                weight = self.param_quant_stats[module_name].pop(f"{module_name}.weight")
-                new_value = bnb.nn.Params4bit.from_prequantized(
-                    data=weight,
-                    quantized_stats=self.param_quant_stats[module_name],
-                    requires_grad=False,
-                    device=target_device,
-                    module=module,
-                )
-                # Set it
-                module._parameters[tensor_name] = new_value
-                # Delete the states
-                del self.param_quant_stats[module_name]
-        else:
-            new_value = param_value.to("cpu")
-            old_value = getattr(module, tensor_name)
-
-            # Support models using `Conv1D` in place of `nn.Linear` (e.g. openai-community/gpt2) by transposing the weight matrix prior to quantization.
-            # Since weights are saved in the correct "orientation", we skip transposing when loading.
-            if issubclass(module.source_cls, Conv1D):
-                new_value = new_value.T
-
-            kwargs = old_value.__dict__
-            kwargs.pop("_is_hf_initialized", None)
-            new_value = bnb.nn.Params4bit(new_value, requires_grad=False, **kwargs).to(target_device)
-
-            module._parameters[tensor_name] = new_value
-
-    # Copied from transformers.quantizers.quantizer_bnb_8bit.Bnb8BitHfQuantizer.adjust_max_memory
     def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]:
         # need more space for buffers that are created during quantization
         max_memory = {key: val * 0.90 for key, val in max_memory.items()}
         return max_memory
 
-    # Copied from transformers.quantizers.quantizer_bnb_8bit.Bnb8BitHfQuantizer.update_dtype
-    def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
-        if dtype is None:
-            # We force the `dtype` to be float16, this is a requirement from `bitsandbytes`
-            logger.info(
-                "Overriding dtype=%s with `dtype=torch.float16` due to "
-                "requirements of `bitsandbytes` to enable model loading in 8-bit or 4-bit. "
-                "Pass your own dtype to specify the dtype of the remaining non-linear layers or pass"
-                " dtype=torch.float16 to remove this warning.",
-                dtype,
-            )
-            dtype = torch.float16
-        return dtype
-
     def update_device_map(self, device_map):
         if device_map is None:
             if torch.cuda.is_available():
@@ -238,33 +116,23 @@ class Bnb4BitHfQuantizer(HfQuantizer):
             )
         return device_map
 
-    # Copied from transformers.quantizers.quantizer_bnb_8bit.Bnb8BitHfQuantizer._process_model_before_weight_loading
     def _process_model_before_weight_loading(
         self,
         model: "PreTrainedModel",
         device_map,
-        keep_in_fp32_modules: list[str] | None = None,
         **kwargs,
     ):
         from ..integrations import replace_with_bnb_linear
 
-        llm_int8_enable_fp32_cpu_offload = self.quantization_config.llm_int8_enable_fp32_cpu_offload
-
         self.modules_to_not_convert = self.get_modules_to_not_convert(
-            model, self.quantization_config.llm_int8_skip_modules, keep_in_fp32_modules
+            model, self.quantization_config.llm_int8_skip_modules, model._keep_in_fp32_modules
         )
 
-        # Extend `self.modules_to_not_convert` to keys that are supposed to be offloaded to `cpu` or `disk`
-        if isinstance(device_map, dict) and len(device_map.keys()) > 1:
-            keys_on_cpu = [key for key, value in device_map.items() if value in ["disk", "cpu"]]
+        if self.quantization_config.llm_int8_enable_fp32_cpu_offload:
+            if isinstance(device_map, dict):
+                keys_on_cpu = [key for key, value in device_map.items() if value in ["disk", "cpu"]]
+                self.modules_to_not_convert.extend(keys_on_cpu)
 
-            if len(keys_on_cpu) > 0 and not llm_int8_enable_fp32_cpu_offload:
-                raise ValueError(
-                    "If you want to offload some keys to `cpu` or `disk`, you need to set "
-                    "`llm_int8_enable_fp32_cpu_offload=True`. Note that these modules will not be "
-                    " converted to 8-bit but kept in 32-bit."
-                )
-            self.modules_to_not_convert.extend(keys_on_cpu)
         model = replace_with_bnb_linear(
             model,
             modules_to_not_convert=self.modules_to_not_convert,
@@ -272,27 +140,22 @@ class Bnb4BitHfQuantizer(HfQuantizer):
             pre_quantized=self.pre_quantized,
         )
 
-        model.config.quantization_config = self.quantization_config
-
-    # Copied from transformers.quantizers.quantizer_bnb_8bit.Bnb8BitHfQuantizer._process_model_after_weight_loading with 8bit->4bit
     def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
         model.is_loaded_in_4bit = True
         model.is_4bit_serializable = self.is_serializable()
         return model
 
-    def is_serializable(self, safe_serialization=None):
+    def is_serializable(self):
         return True
 
     @property
     def is_trainable(self) -> bool:
         return True
 
-    def _dequantize(self, model):
+    def _dequantize(self, model, dtype=None):
         from ..integrations import dequantize_and_replace
 
-        model = dequantize_and_replace(
-            model, self.modules_to_not_convert, quantization_config=self.quantization_config
-        )
+        model = dequantize_and_replace(model, quantization_config=self.quantization_config, dtype=dtype)
         return model
 
     def get_quantize_ops(self):

transformers/quantizers/quantizer_bnb_8bit.py

@@ -25,6 +25,8 @@ from ..utils import (
     is_accelerate_available,
     is_bitsandbytes_available,
     is_torch_available,
+    is_torch_hpu_available,
+    is_torch_npu_available,
     is_torch_xpu_available,
     logging,
 )
@@ -35,34 +37,20 @@ if is_torch_available():
     import torch
 
     from ..core_model_loading import WeightConverter
-    from ..pytorch_utils import Conv1D
 
 logger = logging.get_logger(__name__)
 
 
 class Bnb8BitHfQuantizer(HfQuantizer):
     """
-    8-bit quantization from bitsandbytes quantization method:
-        before loading: converts transformer layers into Linear8bitLt during loading: load 16bit weight and pass to the
-        layer object after: quantizes individual weights in Linear8bitLt into 8bit at fitst .cuda() call
-    saving:
-        from state dict, as usual; saves weights and 'SCB' component
-    loading:
-        need to locate SCB component and pass to the Linear8bitLt object
+    8-bit quantization from bitsandbytes quantization method
     """
 
-    use_keep_in_fp32_modules = True
-    requires_parameters_quantization = True
     requires_calibration = False
 
-    required_packages = ["bitsandbytes", "accelerate"]
-
     def __init__(self, quantization_config, **kwargs):
         super().__init__(quantization_config, **kwargs)
 
-        if self.quantization_config.llm_int8_skip_modules is not None:
-            self.modules_to_not_convert = self.quantization_config.llm_int8_skip_modules
-
     def validate_environment(self, *args, **kwargs):
         if not is_accelerate_available():
             raise ImportError(
@@ -78,17 +66,9 @@ class Bnb8BitHfQuantizer(HfQuantizer):
         validate_bnb_backend_availability(raise_exception=True)
 
         device_map = kwargs.get("device_map")
-        if (
-            device_map is not None
-            and isinstance(device_map, dict)
-            and not self.quantization_config.llm_int8_enable_fp32_cpu_offload
-        ):
-            device_map_without_lm_head = {
-                key: device_map[key] for key in device_map if key not in self.modules_to_not_convert
-            }
-            if set(device_map.values()) == {"cpu"}:
-                pass
-            elif "cpu" in device_map_without_lm_head.values() or "disk" in device_map_without_lm_head.values():
+        if not self.quantization_config.llm_int8_enable_fp32_cpu_offload and isinstance(device_map, dict):
+            values = set(device_map.values())
+            if values != {"cpu"} and ("cpu" in values or "disk" in values):
                 raise ValueError(
                     "Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit the "
                     "quantized model. If you want to dispatch the model on the CPU or the disk while keeping these modules "
@@ -103,23 +83,14 @@ class Bnb8BitHfQuantizer(HfQuantizer):
         max_memory = {key: val * 0.90 for key, val in max_memory.items()}
         return max_memory
 
-    def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
-        if dtype is None:
-            # We force the `dtype` to be float16, this is a requirement from `bitsandbytes`
-            logger.info(
-                "Overriding dtype=%s with `dtype=torch.float16` due to "
-                "requirements of `bitsandbytes` to enable model loading in 8-bit or 4-bit. "
-                "Pass your own dtype to specify the dtype of the remaining non-linear layers or pass"
-                " dtype=torch.float16 to remove this warning.",
-                dtype,
-            )
-            dtype = torch.float16
-        return dtype
-
     def update_device_map(self, device_map):
         if device_map is None:
             if torch.cuda.is_available():
                 device_map = {"": torch.cuda.current_device()}
+            elif is_torch_npu_available():
+                device_map = {"": f"npu:{torch.npu.current_device()}"}
+            elif is_torch_hpu_available():
+                device_map = {"": f"hpu:{torch.hpu.current_device()}"}
             elif is_torch_xpu_available():
                 device_map = {"": torch.xpu.current_device()}
             else:
@@ -131,14 +102,12 @@ class Bnb8BitHfQuantizer(HfQuantizer):
             )
         return device_map
 
-    def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype":
-        if target_dtype != torch.int8:
-            logger.info("target_dtype {target_dtype} is replaced by `torch.int8` for 8-bit BnB quantization")
-        return torch.int8
-
-    def update_unexpected_keys(self, model, unexpected_keys: list[str]) -> list[str]:
-        bnb_keys = ["SCB", "weight_format"]
-        return [k for k in unexpected_keys if not any(k.endswith(x) for x in bnb_keys)]
+    def param_element_size(self, model: "PreTrainedModel", param_name: str, param: "torch.Tensor") -> float:
+        "Return the element size (in bytes) for `param_name`."
+        if self.param_needs_quantization(model, param_name):
+            # 8-bit
+            return 1
+        return super().param_element_size(model, param_name, param)
 
     def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
         import bitsandbytes as bnb
@@ -146,47 +115,6 @@ class Bnb8BitHfQuantizer(HfQuantizer):
         module, name = get_module_from_name(model, param_name)
         return isinstance(module, bnb.nn.Linear8bitLt) and name != "bias"
 
-    def create_quantized_param(
-        self,
-        model: "PreTrainedModel",
-        param_value: "torch.Tensor",
-        param_name: str,
-        target_device: "torch.device",
-        **kwargs,
-    ):
-        import bitsandbytes as bnb
-
-        module, tensor_name = get_module_from_name(model, param_name)
-
-        if self.pre_quantized and not self.is_serializable():
-            raise ValueError(
-                "Detected int8 weights but the version of bitsandbytes is not compatible with int8 serialization. "
-                "Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`."
-            )
-        # Those 2 can only happen when self.pre_quantized == True
-        if tensor_name == "SCB":
-            setattr(module.weight, "SCB", param_value.to(target_device))
-            return
-        # It's not used, but it's getting serialized for BC reason...
-        elif tensor_name == "weight_format":
-            return
-
-        # Support models using `Conv1D` in place of `nn.Linear` (e.g. openai-community/gpt2) by transposing the weight matrix prior to quantization.
-        # Since weights are saved in the correct "orientation", we skip transposing when loading.
-        if issubclass(module.source_cls, Conv1D) and not self.pre_quantized:
-            param_value = param_value.T
-
-        old_value = getattr(module, tensor_name)
-        kwargs = old_value.__dict__
-        kwargs.pop("_is_hf_initialized", None)
-        # Need to pop SCB and reset it because of bnb internals that modifies its value when switching devices ...
-        SCB = kwargs.pop("SCB", None)
-        new_value = bnb.nn.Int8Params(param_value.to("cpu"), requires_grad=False, **kwargs).to(target_device)
-        if SCB is not None:
-            setattr(new_value, "SCB", SCB)
-        # Set it to the module
-        module._parameters[tensor_name] = new_value
-
     def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
         model.is_loaded_in_8bit = True
         model.is_8bit_serializable = self.is_serializable()
@@ -196,28 +124,18 @@ class Bnb8BitHfQuantizer(HfQuantizer):
         self,
         model: "PreTrainedModel",
         device_map,
-        keep_in_fp32_modules: list[str] | None = None,
         **kwargs,
     ):
         from ..integrations import replace_with_bnb_linear
 
-        llm_int8_enable_fp32_cpu_offload = self.quantization_config.llm_int8_enable_fp32_cpu_offload
-
         self.modules_to_not_convert = self.get_modules_to_not_convert(
-            model, self.quantization_config.llm_int8_skip_modules, keep_in_fp32_modules
+            model, self.quantization_config.llm_int8_skip_modules, model._keep_in_fp32_modules
         )
 
-        # Extend `self.modules_to_not_convert` to keys that are supposed to be offloaded to `cpu` or `disk`
-        if isinstance(device_map, dict) and len(device_map.keys()) > 1:
-            keys_on_cpu = [key for key, value in device_map.items() if value in ["disk", "cpu"]]
-
-            if len(keys_on_cpu) > 0 and not llm_int8_enable_fp32_cpu_offload:
-                raise ValueError(
-                    "If you want to offload some keys to `cpu` or `disk`, you need to set "
-                    "`llm_int8_enable_fp32_cpu_offload=True`. Note that these modules will not be "
-                    " converted to 8-bit but kept in 32-bit."
-                )
-            self.modules_to_not_convert.extend(keys_on_cpu)
+        if self.quantization_config.llm_int8_enable_fp32_cpu_offload:
+            if isinstance(device_map, dict):
+                keys_on_cpu = [key for key, value in device_map.items() if value in ["disk", "cpu"]]
+                self.modules_to_not_convert.extend(keys_on_cpu)
 
         model = replace_with_bnb_linear(
             model,
@@ -226,21 +144,17 @@ class Bnb8BitHfQuantizer(HfQuantizer):
             pre_quantized=self.pre_quantized,
         )
 
-        model.config.quantization_config = self.quantization_config
-
-    def is_serializable(self, safe_serialization=None):
+    def is_serializable(self):
         return True
 
     @property
     def is_trainable(self) -> bool:
         return True
 
-    def _dequantize(self, model):
+    def _dequantize(self, model, dtype=None):
         from ..integrations import dequantize_and_replace
 
-        model = dequantize_and_replace(
-            model, self.modules_to_not_convert, quantization_config=self.quantization_config
-        )
+        model = dequantize_and_replace(model, quantization_config=self.quantization_config, dtype=dtype)
         return model
 
     def get_quantize_ops(self):

transformers/quantizers/quantizer_compressed_tensors.py

@@ -31,7 +31,6 @@ class CompressedTensorsHfQuantizer(HfQuantizer):
     """
 
     requires_calibration = True
-    required_packages = ["compressed_tensors"]
 
     def __init__(self, quantization_config: CompressedTensorsConfig, **kwargs):
         super().__init__(quantization_config, **kwargs)
@@ -58,15 +57,9 @@ class CompressedTensorsHfQuantizer(HfQuantizer):
                 "Using `compressed_tensors` quantized models requires the compressed-tensors library: "
                 "`pip install compressed-tensors`"
             )
-        if not is_torch_available():
-            # torch already should be installed as part of compressed tensors
-            raise ImportError("torch is required for using compressed-tensors quantization")
 
     def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
-        if dtype is None:
-            logger.info("Loading model using torch.float16 for compressed-tensors quantization")
-            dtype = torch.float16
-        elif dtype != torch.float16:
+        if dtype != torch.float16:
             logger.info("We suggest you to set `dtype=torch.float16` for better efficiency with compressed_tensors.")
         return dtype
 
@@ -113,6 +106,6 @@ class CompressedTensorsHfQuantizer(HfQuantizer):
         # models need to be decompressed carry out qat
         return not self.run_compressed or not self.quantization_config.is_quantization_compressed
 
-    def is_serializable(self, safe_serialization=None) -> bool:
+    def is_serializable(self) -> bool:
         """Models quantized using compressed tensors can be saved to disk"""
         return True

transformers/quantizers/quantizer_eetq.py

@@ -19,7 +19,7 @@ from .base import HfQuantizer
 if TYPE_CHECKING:
     from ..modeling_utils import PreTrainedModel
 
-from ..utils import is_accelerate_available, is_eetq_available, is_torch_available, logging
+from ..utils import is_accelerate_available, is_kernels_available, is_torch_available, logging
 from .quantizers_utils import get_module_from_name
 
 
@@ -32,40 +32,17 @@ logger = logging.get_logger(__name__)
 
 class EetqHfQuantizer(HfQuantizer):
     """
-    8-bit quantization from EETQ quantization method:
-        before loading: converts transformer layers into W8A16Linear during loading: load 16bit weight and pass to the
-        layer object after: quantizes individual weights in Linear8bitLt into 8bit at first .cuda() call
+    8-bit quantization from EETQ quantization method
     """
 
-    requires_parameters_quantization = True
     requires_calibration = False
 
-    required_packages = ["eetq", "accelerate"]
-
     def __init__(self, quantization_config, **kwargs):
         super().__init__(quantization_config, **kwargs)
-        self.quantization_config = quantization_config
 
     def validate_environment(self, *args, **kwargs):
-        if not is_eetq_available():
-            raise ImportError(
-                "Using `eetq` 8-bit quantization requires eetq."
-                "Please install the latest version of eetq from : https://github.com/NetEase-FuXi/EETQ"
-            )
-
-        try:
-            import eetq  # noqa: F401
-        except ImportError as exc:
-            if "shard_checkpoint" in str(exc):
-                # EETQ 1.0.0 is currently broken with the latest transformers because it tries to import the removed
-                # shard_checkpoint function, see https://github.com/NetEase-FuXi/EETQ/issues/34.
-                # TODO: Update message once eetq releases a fix
-                raise ImportError(
-                    "You are using a version of EETQ that is incompatible with the current transformers version. "
-                    "Either downgrade transformers to <= v4.46.3 or, if available, upgrade EETQ to > v1.0.0."
-                ) from exc
-            else:
-                raise
+        if not is_kernels_available():
+            raise ImportError("Loading an EETQ quantized model requires kernels (`pip install kernels`)")
 
         if not is_accelerate_available():
             raise ImportError("Loading an EETQ quantized model requires accelerate (`pip install accelerate`)")
@@ -79,29 +56,20 @@ class EetqHfQuantizer(HfQuantizer):
                 "You have loaded an EETQ model on CPU and have a CUDA device available, make sure to set "
                 "your model on a GPU device in order to run your model."
             )
-        elif device_map is not None:
-            if isinstance(device_map, dict) and ("cpu" in device_map.values() or "disk" in device_map.values()):
+        elif isinstance(device_map, dict):
+            if len(device_map) > 1 and "cpu" in device_map.values() or "disk" in device_map.values():
                 raise ValueError(
                     "You are attempting to load an EETQ model with a device_map that contains a CPU or disk device."
                     " This is not supported. Please remove the CPU or disk device from the device_map."
                 )
 
     def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
-        if dtype is None:
-            dtype = torch.float16
-            logger.info(
-                "Overriding dtype=%s with `dtype=torch.float16` due to "
-                "requirements of `eetq` to enable model loading in 8-bit. "
-                "Pass your own dtype to specify the dtype of the remaining non-linear layers or pass"
-                " dtype=torch.float16 to remove this warning.",
-                dtype,
-            )
-        elif dtype != torch.float16:
+        if dtype != torch.float16:
             logger.info("We suggest you to set `dtype=torch.float16` for better efficiency with EETQ.")
         return dtype
 
     def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
-        from eetq import EetqLinear
+        from ..integrations.eetq import EetqLinear
 
         module, tensor_name = get_module_from_name(model, param_name)
 
@@ -112,55 +80,29 @@ class EetqHfQuantizer(HfQuantizer):
                 return True
         return False
 
-    def create_quantized_param(
-        self,
-        model: "PreTrainedModel",
-        param_value: "torch.Tensor",
-        param_name: str,
-        target_device: "torch.device",
-        **kwargs,
-    ):
-        from eetq import EetqLinear, quantize_and_preprocess_weights
-
-        module, tensor_name = get_module_from_name(model, param_name)
-        new_value, weight_scale = quantize_and_preprocess_weights(param_value)
-
-        # Samity check
-        if isinstance(module, EetqLinear):
-            if self.pre_quantized or tensor_name == "bias":
-                if tensor_name == "weight" and param_value.dtype != torch.int8:
-                    raise ValueError("Expect quantized weights but got an unquantized weight")
-            else:
-                if tensor_name == "weight_scale":
-                    raise ValueError("Expect unquantized weights but got a quantized weight_scale")
-
-        module._buffers[tensor_name] = new_value.to(target_device)
-        module.register("weight_scales", weight_scale.to(target_device))
-
     def _process_model_before_weight_loading(
         self,
         model: "PreTrainedModel",
-        keep_in_fp32_modules: list[str] | None = None,
         **kwargs,
     ):
         from ..integrations import replace_with_eetq_linear
 
         self.modules_to_not_convert = self.get_modules_to_not_convert(
-            model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules
+            model, self.quantization_config.modules_to_not_convert, model._keep_in_fp32_modules
         )
 
         model = replace_with_eetq_linear(
-            model,
-            modules_to_not_convert=self.modules_to_not_convert,
-            quantization_config=self.quantization_config,
-            pre_quantized=self.pre_quantized,
+            model, modules_to_not_convert=self.modules_to_not_convert, pre_quantized=self.pre_quantized
         )
 
-        model.config.quantization_config = self.quantization_config
-
-    def is_serializable(self, safe_serialization=None):
+    def is_serializable(self):
         return True
 
     @property
     def is_trainable(self) -> bool:
         return True
+
+    def get_quantize_ops(self):
+        from ..integrations.eetq import EetqQuantize
+
+        return EetqQuantize(self)