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

transformers/quantizers/base.py

@@ -12,17 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from abc import ABC, abstractmethod
-from copy import deepcopy
 from typing import TYPE_CHECKING, Any
 
-from ..utils import is_accelerate_available, is_torch_available, logging
+from ..utils import is_torch_available, logging
 from ..utils.quantization_config import QuantizationConfigMixin, QuantizationMethod
 from .quantizers_utils import get_module_from_name
 
 
-if is_accelerate_available():
-    from accelerate.utils import find_tied_parameters
-
 if TYPE_CHECKING:
     from ..modeling_utils import PreTrainedModel
 
@@ -45,50 +41,31 @@ def _assign_original_dtype(module, original_dtype):
         _assign_original_dtype(child, original_dtype)
 
 
-def get_keys_to_not_convert(model):
+def get_keys_to_not_convert(model) -> list:
     r"""
-    An utility function to get the key of the module to keep in full precision if any For example for CausalLM modules
-    we may want to keep the lm_head in full precision for numerical stability reasons. For other architectures, we want
-    to keep the tied weights of the model. The function will return a list of the keys of the modules to not convert in
-    int8.
-
-    Parameters:
-    model (`torch.nn.Module`):
-        Input model
+    Function to automatically detect keys to not convert for usage like quantization. For example for CausalLM modules
+    we may want to keep the lm_head in full precision for numerical stability reasons.
     """
-    # Create a copy of the model and tie the weights, then
-    # check if it contains tied weights
-    tied_model = deepcopy(model)  # this has 0 cost since it is done inside `init_empty_weights` context manager`
-    tied_model.tie_weights()
-
-    tied_params = find_tied_parameters(tied_model)
-    tied_keys = sum(tied_params, [])
-    has_tied_params = len(tied_keys) > 0
-
-    # If there is not tied weights, we want to keep the lm_head（output_embedding) in full precision
-    if not has_tied_params:
-        output_emb = model.get_output_embeddings()
-        if output_emb is not None:
-            list_last_module = [name for name, module in model.named_modules() if id(module) == id(output_emb)]
-            return list_last_module
-
-    # otherwise, no tied weights, no output embedding defined, simply keep the last module in full precision
-    list_modules = list(model.named_parameters())
-    list_last_module = [list_modules[-1][0]]
-    # add last module together with tied weights
-    intersection = set(list_last_module) - set(tied_keys)
-    list_untouched = list(set(tied_keys)) + list(intersection)
-
-    # remove ".weight" from the keys
-    names_to_remove = [".weight", ".bias"]
-    filtered_module_names = []
-    for name in list_untouched:
-        for name_to_remove in names_to_remove:
-            if name_to_remove in name:
-                name = name.replace(name_to_remove, "")
-        filtered_module_names.append(name)
-
-    return filtered_module_names
+    # remove tied weights
+    tied_keys = set()
+    if len(model.all_tied_weights_keys) > 0:
+        tied_keys = set(model.all_tied_weights_keys.values()) | set(model.all_tied_weights_keys.keys())
+
+    # remove last module
+    last_module_key = {list(model.named_parameters())[-1][0]}
+
+    # remove output emb
+    output_emb_module = model.get_output_embeddings()
+    output_emb_keys = {
+        name
+        for name, module in model.named_modules()
+        if output_emb_module is not None and id(module) == id(output_emb_module)
+    }
+    modules_to_not_convert = tied_keys | last_module_key | output_emb_keys
+
+    modules_to_not_convert = list({k.removesuffix(".weight") for k in modules_to_not_convert})
+
+    return list(modules_to_not_convert)
 
 
 class HfQuantizer(ABC):
@@ -100,26 +77,14 @@ class HfQuantizer(ABC):
     Attributes
         quantization_config (`transformers.utils.quantization_config.QuantizationConfigMixin`):
             The quantization config that defines the quantization parameters of your model that you want to quantize.
-        modules_to_not_convert (`list[str]`, *optional*):
-            The list of module names to not convert when quantizing the model.
-        required_packages (`list[str]`, *optional*):
-            The list of required pip packages to install prior to using the quantizer
         requires_calibration (`bool`):
             Whether the quantization method requires to calibrate the model before using it.
-        requires_parameters_quantization (`bool`):
-            Whether the quantization method requires to create a new Parameter. For example, for bitsandbytes, it is
-            required to create a new xxxParameter in order to properly quantize the model.
     """
 
     requires_calibration = False
-    required_packages = None
-    requires_parameters_quantization = False
 
     def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs):
         self.quantization_config = quantization_config
-
-        # -- Handle extra kwargs below --
-        self.modules_to_not_convert = kwargs.pop("modules_to_not_convert", [])
         self.pre_quantized = kwargs.pop("pre_quantized", True)
 
         if not self.pre_quantized and self.requires_calibration:
@@ -182,53 +147,16 @@ class HfQuantizer(ABC):
                 return mapping[custom_dtype]
         return param.element_size()
 
-    def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> list[str]:
-        """
-        Override this method if you want to adjust the `missing_keys`.
-
-        Args:
-            missing_keys (`list[str]`, *optional*):
-                The list of missing keys in the checkpoint compared to the state dict of the model
-        """
-        return missing_keys
-
-    def update_expected_keys(self, model, expected_keys: list[str], loaded_keys: list[str]) -> list[str]:
-        """
-        Override this method if you want to adjust the `update_expected_keys`.
-
-        Args:
-            expected_keys (`list[str]`, *optional*):
-                The list of the expected keys in the initialized model.
-            loaded_keys (`list[str]`, *optional*):
-                The list of the loaded keys in the checkpoint.
-        """
-        return expected_keys
-
-    def update_unexpected_keys(self, model, unexpected_keys: list[str]) -> list[str]:
-        return unexpected_keys
-
     def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]:
         """adjust max_memory argument for infer_auto_device_map() if extra memory is needed for quantization"""
         return max_memory
 
     def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
         """
-        Check whether a given param needs quantization as defined by `create_quantized_param`.
+        Check whether a given param needs to be quantized.
         """
         return False
 
-    def create_quantized_param(self, *args, **kwargs):
-        """
-        Take needed components from state_dict (those from which `param_needs_quantization` is True) and create
-        quantized param.
-        It usually also load the new param directly in the `model`.
-        Note: only applicable if requires_parameters_quantization == True.
-        """
-        if not self.requires_parameters_quantization:
-            raise AttributeError(
-                f"`.create_quantized_param()` method is not supported by quantizer class {self.__class__.__name__}."
-            )
-
     def validate_environment(self, *args, **kwargs):
         """
         This method is used to potentially check for potential conflicts with arguments that are
@@ -288,6 +216,11 @@ class HfQuantizer(ABC):
             kwargs (`dict`, *optional*):
                 The keyword arguments that are passed along `_process_model_after_weight_loading`.
         """
+        model.config.quantization_config = self.quantization_config
+
+        if self.pre_quantized and getattr(self.quantization_config, "dequantize", False):
+            self.remove_quantization_config(model)
+
         return self._process_model_after_weight_loading(model, **kwargs)
 
     def remove_quantization_config(self, model):
@@ -310,13 +243,7 @@ class HfQuantizer(ABC):
         Note not all quantization schemes support this.
         """
         model = self._dequantize(model)
-
-        # Delete quantizer and quantization config
-        del model.hf_quantizer
-        del model.config.quantization_config
-        del model.config._pre_quantization_dtype
-        del model.quantization_method
-        model.is_quantized = False
+        self.remove_quantization_config(model)
 
         return model
 
@@ -360,6 +287,8 @@ class HfQuantizer(ABC):
         if keep_in_fp32_modules is not None:
             modules_to_not_convert.extend(keep_in_fp32_modules)
 
+        modules_to_not_convert = list(set(modules_to_not_convert))
+
         return modules_to_not_convert
 
     @property
@@ -372,16 +301,12 @@ class HfQuantizer(ABC):
         """Flag indicating whether the quantized model can be compiled"""
         return False
 
-    def get_state_dict_and_metadata(self, model, safe_serialization=False):
+    def get_state_dict_and_metadata(self, model):
         """Get state dict and metadata. Useful when we need to modify a bit the state dict due to quantization"""
         return None, {}
 
-    def update_state_dict_with_metadata(self, state_dict, metadata):
-        """Update state dict with metadata. Default behaviour returns state_dict"""
-        return state_dict
-
     @abstractmethod
-    def is_serializable(self, safe_serialization=None): ...
+    def is_serializable(self): ...
 
     @property
     @abstractmethod

transformers/quantizers/quantizer_aqlm.py

@@ -39,12 +39,9 @@ class AqlmHfQuantizer(HfQuantizer):
     """
 
     requires_calibration = True
-    required_packages = ["aqlm"]
-    optimum_quantizer = None
 
     def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs):
         super().__init__(quantization_config, **kwargs)
-        self.quantization_config = quantization_config
 
     def validate_environment(self, *args, **kwargs):
         if not is_accelerate_available():
@@ -77,7 +74,6 @@ class AqlmHfQuantizer(HfQuantizer):
             quantization_config=self.quantization_config,
             linear_weights_not_to_quantize=self.quantization_config.linear_weights_not_to_quantize,
         )
-        model.config.quantization_config = self.quantization_config
 
     @property
     def is_trainable(self) -> bool:
@@ -90,5 +86,5 @@ class AqlmHfQuantizer(HfQuantizer):
             )
             return False
 
-    def is_serializable(self, safe_serialization=None):
+    def is_serializable(self):
         return True

transformers/quantizers/quantizer_auto_round.py

@@ -36,7 +36,6 @@ class AutoRoundQuantizer(HfQuantizer):
 
     # AutoRound requires data calibration - we support only inference
     requires_calibration = True
-    required_packages = ["auto_round"]
 
     def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs):
         super().__init__(quantization_config, **kwargs)
@@ -76,6 +75,6 @@ class AutoRoundQuantizer(HfQuantizer):
     def is_trainable(self) -> bool:
         return False
 
-    def is_serializable(self, safe_serialization=None):
+    def is_serializable(self):
         ## for gptq/awq models, the quantization config will be changed
         return True

transformers/quantizers/quantizer_awq.py

@@ -22,8 +22,8 @@ from .base import HfQuantizer
 if TYPE_CHECKING:
     from ..modeling_utils import PreTrainedModel
 
-from ..utils import is_accelerate_available, is_auto_awq_available, is_torch_available, logging
-from ..utils.quantization_config import AWQLinearVersion
+from ..utils import is_accelerate_available, is_gptqmodel_available, is_torch_available, logging
+from ..utils.quantization_config import AwqBackend
 
 
 if is_torch_available():
@@ -40,60 +40,18 @@ class AwqQuantizer(HfQuantizer):
     # AWQ requires data calibration - we support only inference
     requires_calibration = True
 
-    required_packages = ["awq", "accelerate"]
-
     def __init__(self, quantization_config, **kwargs):
         super().__init__(quantization_config, **kwargs)
 
-    def validate_environment(self, device_map, **kwargs):
-        if not is_auto_awq_available():
-            raise ImportError("Loading an AWQ quantized model requires auto-awq library (`pip install autoawq`)")
+    def validate_environment(self, **kwargs):
+        if not is_gptqmodel_available():
+            raise ImportError(
+                "Loading an AWQ quantized model requires gptqmodel. Please install it with `pip install gptqmodel`"
+            )
 
         if not is_accelerate_available():
             raise ImportError("Loading an AWQ quantized model requires accelerate (`pip install accelerate`)")
 
-        if (
-            self.quantization_config.version == AWQLinearVersion.GEMM
-            and not torch.cuda.is_available()
-            and not torch.xpu.is_available()
-        ):
-            logger.warning_once("No CUDA or XPU found, consider switching to the IPEX version for CPU-only execution.")
-            self.quantization_config.version = AWQLinearVersion.IPEX
-
-        if self.quantization_config.version == AWQLinearVersion.IPEX:
-            if version.parse(importlib.metadata.version("autoawq")) < version.parse("0.2.6"):
-                raise RuntimeError(
-                    "To use IPEX backend, you need autoawq>0.2.6. Please install the latest version or from source."
-                )
-            if device_map is None:
-                logger.warning_once(
-                    "You have loaded an AWQ model without setting device_map, please set 'cpu' or 'xpu' or 'auto'"
-                )
-            elif isinstance(device_map, dict) and "disk" in device_map.values():
-                raise ValueError(
-                    "You are attempting to load an IPEX version AWQ model with a device_map that contains disk device."
-                    " This is not supported. Please make sure only cpu and xpu in the device_map."
-                )
-        else:
-            if not torch.cuda.is_available() and not torch.xpu.is_available():
-                raise RuntimeError(
-                    "GPU is required to run AWQ quantized model. You can use IPEX version AWQ if you have an Intel CPU"
-                )
-
-            if device_map is None:
-                logger.warning_once(
-                    "You have loaded an AWQ model on CPU and have a CUDA/XPU 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 any(
-                    forbidden in device_map.values() for forbidden in ("cpu", torch.device("cpu"), "disk")
-                ):
-                    raise ValueError(
-                        "You are attempting to load an AWQ 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):
         if dtype is None:
             dtype = torch.float16
@@ -116,42 +74,22 @@ class AwqQuantizer(HfQuantizer):
             model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules, add_default_skips=True
         )
 
-        model, has_been_replaced = replace_with_awq_linear(
-            model, quantization_config=self.quantization_config, modules_to_not_convert=self.modules_to_not_convert
+        model = replace_with_awq_linear(
+            model,
+            quantization_config=self.quantization_config,
+            modules_to_not_convert=self.modules_to_not_convert,
+            device_map=kwargs.get("device_map"),
         )
 
         model = replace_quantization_scales(model, model.config.model_type)
 
-        if not has_been_replaced:
-            logger.warning(
-                "You are loading an AWQ model but no linear modules were found in your model."
-                " Please double check your model architecture, or submit an issue on github if you think this is a bug."
-            )
-
     def _process_model_after_weight_loading(self, model, **kwargs):
-        if self.quantization_config.do_fuse:
-            from ..integrations import fuse_awq_modules
+        from gptqmodel.utils.model import hf_gptqmodel_post_init
 
-            model = fuse_awq_modules(model, self.quantization_config)
-            model._awq_is_fused = True  # TODO: consider storing this flag in model.config instead
-
-        if self.quantization_config.version == AWQLinearVersion.EXLLAMA:
-            from ..integrations import post_init_awq_exllama_modules
-
-            model = post_init_awq_exllama_modules(model, self.quantization_config.exllama_config)
-
-        if self.quantization_config.version == AWQLinearVersion.IPEX:
-            from ..integrations import post_init_awq_ipex_modules
-
-            model = post_init_awq_ipex_modules(model)
-
-    def is_serializable(self, safe_serialization=None):
-        # AWQ through auto-awq has been always serializable, except if the model is fused.
-        if self.quantization_config.do_fuse:
-            logger.warning("You cannot save an AWQ model that uses fused modules!")
-            return False
+        hf_gptqmodel_post_init(model, use_act_order=self.quantization_config.desc_act)
 
-        if self.quantization_config.version == AWQLinearVersion.EXLLAMA:
+    def is_serializable(self):
+        if self.quantization_config.backend in [AwqBackend.EXLLAMA_V1, AwqBackend.EXLLAMA_V2]:
             logger.warning("You cannot save an AWQ model that uses Exllama backend!")
             return False
 
@@ -159,6 +97,4 @@ class AwqQuantizer(HfQuantizer):
 
     @property
     def is_trainable(self):
-        # AWQ supports PEFT fine-tuning from version 0.2.0
-        MIN_AWQ_VERSION_FOR_PEFT = "0.2.0"
-        return version.parse(importlib.metadata.version("autoawq")) >= version.parse(MIN_AWQ_VERSION_FOR_PEFT)
+        return version.parse(importlib.metadata.version("gptqmodel")) >= version.parse("5.0.0")

transformers/quantizers/quantizer_bitnet.py

@@ -37,14 +37,10 @@ class BitNetHfQuantizer(HfQuantizer):
     Check out the paper introducing this method: https://huggingface.co/papers/2402.17764
     """
 
-    requires_parameters_quantization = False
     requires_calibration = True
 
-    required_packages = ["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_accelerate_available():
@@ -62,8 +58,8 @@ class BitNetHfQuantizer(HfQuantizer):
                 "You have loaded a BitNet 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 a BitNet 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."
@@ -85,7 +81,6 @@ class BitNetHfQuantizer(HfQuantizer):
             model,
             modules_to_not_convert=self.modules_to_not_convert,
             quantization_config=self.quantization_config,
-            pre_quantized=self.pre_quantized,
         )
 
     def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]:
@@ -96,7 +91,7 @@ class BitNetHfQuantizer(HfQuantizer):
         target_dtype = torch.int8
         return target_dtype
 
-    def is_serializable(self, safe_serialization=None):
+    def is_serializable(self):
         return True
 
     @property

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,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 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}",
@@ -90,17 +75,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 "
@@ -117,13 +94,11 @@ class Bnb4BitHfQuantizer(HfQuantizer):
             logger.info("target_dtype {target_dtype} is replaced by `CustomDtype.INT4` for 4-bit BnB quantization")
         return CustomDtype.INT4
 
-    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
+        # TODO: maybe remove
+        # # 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)
@@ -142,71 +117,13 @@ class Bnb4BitHfQuantizer(HfQuantizer):
                 )
         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":
+        # TODO: remove ? is it still true ? we will move to dtype = "auto" so it will likely be either fp16 or bf16
         if dtype is None:
             # We force the `dtype` to be float16, this is a requirement from `bitsandbytes`
             logger.info(
@@ -238,7 +155,6 @@ 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",
@@ -248,23 +164,15 @@ class Bnb4BitHfQuantizer(HfQuantizer):
     ):
         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
         )
 
-        # 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,15 +180,12 @@ 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
@@ -290,9 +195,7 @@ class Bnb4BitHfQuantizer(HfQuantizer):
     def _dequantize(self, model):
         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)
         return model
 
     def get_quantize_ops(self):