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

transformers/quantizers/quantizer_bnb_4bit.py

@@ -11,6 +11,7 @@
 # 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
@@ -37,20 +38,43 @@ 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
+    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
     """
 
+    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(
@@ -66,9 +90,17 @@ class Bnb4BitHfQuantizer(HfQuantizer):
         validate_bnb_backend_availability(raise_exception=True)
 
         device_map = kwargs.get("device_map")
-        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):
+        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():
                 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 "
@@ -78,25 +110,115 @@ class Bnb4BitHfQuantizer(HfQuantizer):
                     "for more details. "
                 )
 
-    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 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
 
-        return super().param_element_size(model, param_name, param)
+    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)]
 
     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():
@@ -116,23 +238,33 @@ 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, model._keep_in_fp32_modules
+            model, self.quantization_config.llm_int8_skip_modules, keep_in_fp32_modules
         )
 
-        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)
+        # 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)
         model = replace_with_bnb_linear(
             model,
             modules_to_not_convert=self.modules_to_not_convert,
@@ -140,22 +272,27 @@ 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):
+    def is_serializable(self, safe_serialization=None):
         return True
 
     @property
     def is_trainable(self) -> bool:
         return True
 
-    def _dequantize(self, model, dtype=None):
+    def _dequantize(self, model):
         from ..integrations import dequantize_and_replace
 
-        model = dequantize_and_replace(model, quantization_config=self.quantization_config, dtype=dtype)
+        model = dequantize_and_replace(
+            model, self.modules_to_not_convert, quantization_config=self.quantization_config
+        )
         return model
 
     def get_quantize_ops(self):

transformers/quantizers/quantizer_bnb_8bit.py

@@ -25,8 +25,6 @@ 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,
 )
@@ -37,20 +35,34 @@ 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
+    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
     """
 
+    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(
@@ -66,9 +78,17 @@ class Bnb8BitHfQuantizer(HfQuantizer):
         validate_bnb_backend_availability(raise_exception=True)
 
         device_map = kwargs.get("device_map")
-        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):
+        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():
                 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 "
@@ -83,14 +103,23 @@ 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:
@@ -102,12 +131,14 @@ class Bnb8BitHfQuantizer(HfQuantizer):
             )
         return device_map
 
-    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 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_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
         import bitsandbytes as bnb
@@ -115,6 +146,47 @@ 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()
@@ -124,18 +196,28 @@ 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, model._keep_in_fp32_modules
+            model, self.quantization_config.llm_int8_skip_modules, keep_in_fp32_modules
         )
 
-        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)
+        # 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)
 
         model = replace_with_bnb_linear(
             model,
@@ -144,17 +226,21 @@ class Bnb8BitHfQuantizer(HfQuantizer):
             pre_quantized=self.pre_quantized,
         )
 
-    def is_serializable(self):
+        model.config.quantization_config = self.quantization_config
+
+    def is_serializable(self, safe_serialization=None):
         return True
 
     @property
     def is_trainable(self) -> bool:
         return True
 
-    def _dequantize(self, model, dtype=None):
+    def _dequantize(self, model):
         from ..integrations import dequantize_and_replace
 
-        model = dequantize_and_replace(model, quantization_config=self.quantization_config, dtype=dtype)
+        model = dequantize_and_replace(
+            model, self.modules_to_not_convert, quantization_config=self.quantization_config
+        )
         return model
 
     def get_quantize_ops(self):

transformers/quantizers/quantizer_compressed_tensors.py

@@ -31,6 +31,7 @@ class CompressedTensorsHfQuantizer(HfQuantizer):
     """
 
     requires_calibration = True
+    required_packages = ["compressed_tensors"]
 
     def __init__(self, quantization_config: CompressedTensorsConfig, **kwargs):
         super().__init__(quantization_config, **kwargs)
@@ -57,9 +58,15 @@ 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 != torch.float16:
+        if dtype is None:
+            logger.info("Loading model using torch.float16 for compressed-tensors quantization")
+            dtype = torch.float16
+        elif dtype != torch.float16:
             logger.info("We suggest you to set `dtype=torch.float16` for better efficiency with compressed_tensors.")
         return dtype
 
@@ -106,6 +113,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) -> bool:
+    def is_serializable(self, safe_serialization=None) -> 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_kernels_available, is_torch_available, logging
+from ..utils import is_accelerate_available, is_eetq_available, is_torch_available, logging
 from .quantizers_utils import get_module_from_name
 
 
@@ -32,17 +32,40 @@ logger = logging.get_logger(__name__)
 
 class EetqHfQuantizer(HfQuantizer):
     """
-    8-bit quantization from EETQ quantization method
+    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
     """
 
+    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_kernels_available():
-            raise ImportError("Loading an EETQ quantized model requires kernels (`pip install kernels`)")
+        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_accelerate_available():
             raise ImportError("Loading an EETQ quantized model requires accelerate (`pip install accelerate`)")
@@ -56,20 +79,29 @@ 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 isinstance(device_map, dict):
-            if len(device_map) > 1 and "cpu" in device_map.values() or "disk" in device_map.values():
+        elif device_map is not None:
+            if isinstance(device_map, dict) 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 != torch.float16:
+        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:
             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 ..integrations.eetq import EetqLinear
+        from eetq import EetqLinear
 
         module, tensor_name = get_module_from_name(model, param_name)
 
@@ -80,29 +112,55 @@ 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, model._keep_in_fp32_modules
+            model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules
         )
 
         model = replace_with_eetq_linear(
-            model, modules_to_not_convert=self.modules_to_not_convert, pre_quantized=self.pre_quantized
+            model,
+            modules_to_not_convert=self.modules_to_not_convert,
+            quantization_config=self.quantization_config,
+            pre_quantized=self.pre_quantized,
         )
 
-    def is_serializable(self):
+        model.config.quantization_config = self.quantization_config
+
+    def is_serializable(self, safe_serialization=None):
         return True
 
     @property
     def is_trainable(self) -> bool:
         return True
-
-    def get_quantize_ops(self):
-        from ..integrations.eetq import EetqQuantize
-
-        return EetqQuantize(self)