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

transformers/quantizers/base.py

@@ -12,13 +12,17 @@
 # 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_torch_available, logging
+from ..utils import is_accelerate_available, 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
 
@@ -31,39 +35,60 @@ else:
 logger = logging.get_logger(__file__)
 
 
-def get_keys_to_not_convert(model) -> list:
-    r"""
-    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.
-    """
-    # 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)
+def _assign_original_dtype(module, original_dtype):
+    # not very nice in a recursive function but it avoids a circular import
+    from ..modeling_utils import PreTrainedModel
 
+    for child in module.children():
+        if isinstance(child, PreTrainedModel):
+            child.config._pre_quantization_dtype = original_dtype
+        _assign_original_dtype(child, original_dtype)
 
-def _assign_is_quantized(model):
-    from ..modeling_utils import PreTrainedModel
 
-    for module in model.modules():
-        if isinstance(module, PreTrainedModel):
-            module.config._is_quantized = True
+def get_keys_to_not_convert(model):
+    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
+    """
+    # 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
 
 
 class HfQuantizer(ABC):
@@ -75,14 +100,26 @@ 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:
@@ -116,19 +153,82 @@ class HfQuantizer(ABC):
         """
         return device_map
 
+    def adjust_target_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
+        """
+        Override this method if you want to adjust the `target_dtype` variable used in `from_pretrained`
+        to compute the device_map in case the device_map is a `str`. E.g. for bitsandbytes we force-set `target_dtype`
+        to `torch.int8` and for 4-bit we pass a custom enum `accelerate.CustomDtype.int4`.
+
+        Args:
+            dtype (`torch.dtype`, *optional*):
+                The dtype that is used to compute the device_map.
+        """
+        return dtype
+
     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):
+            from accelerate.utils import CustomDtype
+
+            mapping = {
+                torch.int8: 1,
+                CustomDtype.INT4: 0.5,
+                CustomDtype.FP8: 1,
+                CustomDtype.INT2: 0.25,
+            }
+            # The value passed is actually not used when the method is overridden
+            if (custom_dtype := self.adjust_target_dtype(torch.float16)) in mapping:
+                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 to be quantized.
+        Check whether a given param needs quantization as defined by `create_quantized_param`.
         """
         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
@@ -148,7 +248,7 @@ class HfQuantizer(ABC):
     def _process_model_before_weight_loading(self, model, **kwargs):
         return model
 
-    def preprocess_model(self, model: "PreTrainedModel", dtype=None, **kwargs):
+    def preprocess_model(self, model: "PreTrainedModel", config, dtype=None, checkpoint_files=None, **kwargs):
         """
         Setting model attributes and/or converting model before weights loading. At this point
         the model should be initialized on the meta device so you can freely manipulate the skeleton
@@ -166,6 +266,14 @@ class HfQuantizer(ABC):
             self._convert_model_for_quantization(model)
         self._process_model_before_weight_loading(model, **kwargs)
 
+        # We store the original dtype for quantized models as we cannot easily retrieve it
+        # once the weights have been quantized
+        # Note that once you have loaded a quantized model, you can't change its dtype so this will
+        # remain a single source of truth
+        original_dtype = dtype if dtype is not None else torch.get_default_dtype()
+        config._pre_quantization_dtype = original_dtype
+        _assign_original_dtype(model, original_dtype)
+
     def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
         return model
 
@@ -180,13 +288,6 @@ 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)
-        else:
-            _assign_is_quantized(model)
-
         return self._process_model_after_weight_loading(model, **kwargs)
 
     def remove_quantization_config(self, model):
@@ -197,25 +298,40 @@ class HfQuantizer(ABC):
             del model.hf_quantizer
         if hasattr(model.config, "quantization_config"):
             del model.config.quantization_config
+        if hasattr(model.config, "_pre_quantization_dtype"):
+            del model.config._pre_quantization_dtype
         if hasattr(model, "quantization_method"):
             del model.quantization_method
         model.is_quantized = False
 
-    def dequantize(self, model, dtype=None):
+    def dequantize(self, model):
         """
         Potentially dequantize the model to retrieve the original model, with some loss in accuracy / performance.
         Note not all quantization schemes support this.
         """
-        if dtype is None:
-            # using the same dtype we used to load the model. If we don't do that, we might have issues with modules we didn't quantize.
-            # or we need to upcast everything to the same dtype
-            dtype = model.config.dtype
-        model = self._dequantize(model, dtype=dtype)
-        self.remove_quantization_config(model)
+        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
 
         return model
 
-    def _dequantize(self, model, dtype=None):
+    def get_accelerator_warm_up_factor(self):
+        """
+        The factor to be used in `caching_allocator_warmup` to get the number of bytes to pre-allocate to warm up accelerator.
+        A factor of 2 means we allocate all bytes in the empty model (since we allocate in fp16), a factor of 4 means
+        we allocate half the memory of the weights residing in the empty model, etc...
+        """
+        # By default we return 4, i.e. half the model size (this corresponds to the case where the model is not
+        # really pre-processed, i.e. we do not have the info that weights are going to be 8 bits before actual
+        # weight loading)
+        return 4
+
+    def _dequantize(self, model):
         raise NotImplementedError(
             f"{self.quantization_config.quant_method} has no implementation of `dequantize`, please raise an issue on GitHub."
         )
@@ -244,8 +360,6 @@ 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
@@ -258,25 +372,31 @@ class HfQuantizer(ABC):
         """Flag indicating whether the quantized model can be compiled"""
         return False
 
-    def get_state_dict_and_metadata(self, model):
+    def get_state_dict_and_metadata(self, model, safe_serialization=False):
         """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): ...
+    def is_serializable(self, safe_serialization=None): ...
 
     @property
     @abstractmethod
     def is_trainable(self): ...
 
     def _convert_model_for_quantization(self, model):
+        from accelerate import init_empty_weights
+
         for name, module in model.named_modules():
             module_class_name = module.__class__.__name__
             if module_class_name in MODULES_TO_PATCH_FOR_QUANTIZATION and (
                 self.quantization_config.quant_method
                 in MODULES_TO_PATCH_FOR_QUANTIZATION[module_class_name]["quantization_methods"]
             ):
-                with torch.device("meta"):
+                with init_empty_weights():
                     parent_module, name = get_module_from_name(model, name)
                     parent_module._modules[name] = MODULES_TO_PATCH_FOR_QUANTIZATION[module_class_name]["module_name"](
                         model.config.get_text_config()

transformers/quantizers/quantizer_aqlm.py

@@ -23,10 +23,13 @@ if TYPE_CHECKING:
     from ..modeling_utils import PreTrainedModel
 
 from ..integrations import replace_with_aqlm_linear
-from ..utils import is_accelerate_available, is_aqlm_available, logging
+from ..utils import is_accelerate_available, is_aqlm_available, is_torch_available, logging
 from ..utils.quantization_config import QuantizationConfigMixin
 
 
+if is_torch_available():
+    import torch
+
 logger = logging.get_logger(__name__)
 
 
@@ -36,9 +39,12 @@ 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():
@@ -47,6 +53,20 @@ class AqlmHfQuantizer(HfQuantizer):
         if not is_aqlm_available():
             raise ImportError("Using `aqlm` quantization requires AQLM: `pip install aqlm[gpu,cpu]`")
 
+    def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
+        if dtype is None:
+            if torch.cuda.is_available():
+                dtype = torch.float16
+                logger.info(
+                    "CUDA available. Assuming AQLM inference on GPU and loading the model in `torch.float16`. To overwrite it, set `dtype` manually."
+                )
+            else:
+                dtype = torch.float32
+                logger.info(
+                    "CUDA is unavailable. Assuming AQLM inference on CPU and loading the model in `torch.float32`. To overwrite it, set `dtype` manually."
+                )
+        return dtype
+
     def _process_model_before_weight_loading(
         self,
         model: "PreTrainedModel",
@@ -57,6 +77,7 @@ 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:
@@ -69,5 +90,5 @@ class AqlmHfQuantizer(HfQuantizer):
             )
             return False
 
-    def is_serializable(self):
+    def is_serializable(self, safe_serialization=None):
         return True

transformers/quantizers/quantizer_auto_round.py

@@ -19,10 +19,13 @@ from .base import HfQuantizer
 if TYPE_CHECKING:
     from ..modeling_utils import PreTrainedModel
 
-from ..utils import is_auto_round_available, logging
+from ..utils import is_auto_round_available, is_torch_available, logging
 from ..utils.quantization_config import QuantizationConfigMixin
 
 
+if is_torch_available():
+    import torch
+
 logger = logging.get_logger(__name__)
 
 
@@ -33,6 +36,7 @@ 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)
@@ -44,6 +48,12 @@ class AutoRoundQuantizer(HfQuantizer):
                 "Loading an AutoRound quantized model requires auto-round library (`pip install 'auto-round>=0.5'`)"
             )
 
+    def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
+        if dtype is None:
+            dtype = torch.bfloat16
+            logger.info("Loading the model in `torch.bfloat16`. To overwrite it, set `dtype` manually.")
+        return dtype
+
     def _process_model_before_weight_loading(self, model: "PreTrainedModel", **kwargs):
         if model.__class__.main_input_name != "input_ids":
             logger.warning("AutoRound offers only limited support for models that are not strictly text-based.")
@@ -66,6 +76,6 @@ class AutoRoundQuantizer(HfQuantizer):
     def is_trainable(self) -> bool:
         return False
 
-    def is_serializable(self):
+    def is_serializable(self, safe_serialization=None):
         ## 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_gptqmodel_available, is_torch_available, logging
-from ..utils.quantization_config import AwqBackend
+from ..utils import is_accelerate_available, is_auto_awq_available, is_torch_available, logging
+from ..utils.quantization_config import AWQLinearVersion
 
 
 if is_torch_available():
@@ -40,20 +40,65 @@ 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, **kwargs):
-        if not is_gptqmodel_available():
-            raise ImportError(
-                "Loading an AWQ quantized model requires gptqmodel. Please install it with `pip install gptqmodel`"
-            )
+    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`)")
 
         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 == torch.bfloat16 and (torch.cuda.is_available() or torch.xpu.is_available()):
+        if dtype is None:
+            dtype = torch.float16
+            logger.info("Loading the model in `torch.float16`. To overwrite it, set `dtype` manually.")
+        elif dtype == torch.bfloat16 and (torch.cuda.is_available() or torch.xpu.is_available()):
             logger.warning(
                 "`torch.bfloat16` is not supported for AWQ CUDA/XPU kernels yet. Casting to `torch.float16`."
             )
@@ -62,29 +107,51 @@ class AwqQuantizer(HfQuantizer):
             logger.warning("We suggest you to set `dtype=torch.float16` for better efficiency on CUDA/XPU with AWQ.")
         return dtype
 
-    def _process_model_before_weight_loading(self, model: "PreTrainedModel", **kwargs):
+    def _process_model_before_weight_loading(
+        self, model: "PreTrainedModel", keep_in_fp32_modules: list[str] | None = None, **kwargs
+    ):
         from ..integrations import replace_quantization_scales, replace_with_awq_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, add_default_skips=True
+            model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules, add_default_skips=True
         )
 
-        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, has_been_replaced = replace_with_awq_linear(
+            model, quantization_config=self.quantization_config, modules_to_not_convert=self.modules_to_not_convert
         )
 
         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):
-        from gptqmodel.utils.model import hf_gptqmodel_post_init
+        if self.quantization_config.do_fuse:
+            from ..integrations import fuse_awq_modules
 
-        hf_gptqmodel_post_init(model, use_act_order=self.quantization_config.desc_act)
+            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
 
-    def is_serializable(self):
-        if self.quantization_config.backend in [AwqBackend.EXLLAMA_V1, AwqBackend.EXLLAMA_V2]:
+        if self.quantization_config.version == AWQLinearVersion.EXLLAMA:
             logger.warning("You cannot save an AWQ model that uses Exllama backend!")
             return False
 
@@ -92,4 +159,6 @@ class AwqQuantizer(HfQuantizer):
 
     @property
     def is_trainable(self):
-        return version.parse(importlib.metadata.version("gptqmodel")) >= version.parse("5.0.0")
+        # 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)

transformers/quantizers/quantizer_bitnet.py

@@ -37,10 +37,14 @@ 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():
@@ -58,8 +62,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 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 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."
@@ -68,25 +72,31 @@ class BitNetHfQuantizer(HfQuantizer):
     def _process_model_before_weight_loading(
         self,
         model: "PreTrainedModel",
+        keep_in_fp32_modules: list[str] | None = None,
         **kwargs,
     ):
         from ..integrations import replace_with_bitnet_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_bitnet_linear(
             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]:
         max_memory = {key: val * 0.90 for key, val in max_memory.items()}
         return max_memory
 
-    def is_serializable(self):
+    def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype":
+        target_dtype = torch.int8
+        return target_dtype
+
+    def is_serializable(self, safe_serialization=None):
         return True
 
     @property