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

transformers/quantizers/base.py

@@ -31,16 +31,6 @@ else:
 logger = logging.get_logger(__file__)
 
 
-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 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
@@ -68,6 +58,14 @@ def get_keys_to_not_convert(model) -> list:
     return list(modules_to_not_convert)
 
 
+def _assign_is_quantized(model):
+    from ..modeling_utils import PreTrainedModel
+
+    for module in model.modules():
+        if isinstance(module, PreTrainedModel):
+            module.config._is_quantized = True
+
+
 class HfQuantizer(ABC):
     """
     Abstract class of the HuggingFace quantizer. Supports for now quantizing HF transformers models for inference and/or quantization.
@@ -118,33 +116,7 @@ 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 adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]:
@@ -176,7 +148,7 @@ class HfQuantizer(ABC):
     def _process_model_before_weight_loading(self, model, **kwargs):
         return model
 
-    def preprocess_model(self, model: "PreTrainedModel", config, dtype=None, checkpoint_files=None, **kwargs):
+    def preprocess_model(self, model: "PreTrainedModel", dtype=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
@@ -194,14 +166,6 @@ 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
 
@@ -220,6 +184,8 @@ class HfQuantizer(ABC):
 
         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)
 
@@ -231,34 +197,25 @@ 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):
+    def dequantize(self, model, dtype=None):
         """
         Potentially dequantize the model to retrieve the original model, with some loss in accuracy / performance.
         Note not all quantization schemes support this.
         """
-        model = self._dequantize(model)
+        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)
 
         return model
 
-    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):
+    def _dequantize(self, model, dtype=None):
         raise NotImplementedError(
             f"{self.quantization_config.quant_method} has no implementation of `dequantize`, please raise an issue on GitHub."
         )
@@ -313,15 +270,13 @@ class HfQuantizer(ABC):
     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 init_empty_weights():
+                with torch.device("meta"):
                     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,13 +23,10 @@ if TYPE_CHECKING:
     from ..modeling_utils import PreTrainedModel
 
 from ..integrations import replace_with_aqlm_linear
-from ..utils import is_accelerate_available, is_aqlm_available, is_torch_available, logging
+from ..utils import is_accelerate_available, is_aqlm_available, logging
 from ..utils.quantization_config import QuantizationConfigMixin
 
 
-if is_torch_available():
-    import torch
-
 logger = logging.get_logger(__name__)
 
 
@@ -50,20 +47,6 @@ 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",

transformers/quantizers/quantizer_auto_round.py

@@ -19,13 +19,10 @@ from .base import HfQuantizer
 if TYPE_CHECKING:
     from ..modeling_utils import PreTrainedModel
 
-from ..utils import is_auto_round_available, is_torch_available, logging
+from ..utils import is_auto_round_available, logging
 from ..utils.quantization_config import QuantizationConfigMixin
 
 
-if is_torch_available():
-    import torch
-
 logger = logging.get_logger(__name__)
 
 
@@ -47,12 +44,6 @@ 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.")

transformers/quantizers/quantizer_awq.py

@@ -53,10 +53,7 @@ class AwqQuantizer(HfQuantizer):
             raise ImportError("Loading an AWQ quantized model requires accelerate (`pip install accelerate`)")
 
     def update_dtype(self, dtype):
-        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()):
+        if 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`."
             )
@@ -65,13 +62,11 @@ 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", keep_in_fp32_modules: list[str] | None = None, **kwargs
-    ):
+    def _process_model_before_weight_loading(self, model: "PreTrainedModel", **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, keep_in_fp32_modules, add_default_skips=True
+            model, self.quantization_config.modules_to_not_convert, model._keep_in_fp32_modules, add_default_skips=True
         )
 
         model = replace_with_awq_linear(

transformers/quantizers/quantizer_bitnet.py

@@ -68,13 +68,12 @@ 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, keep_in_fp32_modules
+            model, self.quantization_config.modules_to_not_convert, model._keep_in_fp32_modules
         )
 
         model = replace_with_bitnet_linear(
@@ -87,10 +86,6 @@ class BitNetHfQuantizer(HfQuantizer):
         max_memory = {key: val * 0.90 for key, val in max_memory.items()}
         return max_memory
 
-    def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype":
-        target_dtype = torch.int8
-        return target_dtype
-
     def is_serializable(self):
         return True
 

transformers/quantizers/quantizer_bnb_4bit.py

@@ -51,15 +51,6 @@ class Bnb4BitHfQuantizer(HfQuantizer):
     def __init__(self, quantization_config, **kwargs):
         super().__init__(quantization_config, **kwargs)
 
-        # 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(
@@ -87,55 +78,25 @@ class Bnb4BitHfQuantizer(HfQuantizer):
                     "for more details. "
                 )
 
-    def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype":
-        from accelerate.utils import CustomDtype
+    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
 
-        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 param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
         import bitsandbytes as bnb
 
-        # 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)
         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 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
 
-    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(
-                "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():
@@ -159,13 +120,12 @@ class Bnb4BitHfQuantizer(HfQuantizer):
         self,
         model: "PreTrainedModel",
         device_map,
-        keep_in_fp32_modules: list[str] | None = None,
         **kwargs,
     ):
         from ..integrations import replace_with_bnb_linear
 
         self.modules_to_not_convert = self.get_modules_to_not_convert(
-            model, self.quantization_config.llm_int8_skip_modules, keep_in_fp32_modules
+            model, self.quantization_config.llm_int8_skip_modules, model._keep_in_fp32_modules
         )
 
         if self.quantization_config.llm_int8_enable_fp32_cpu_offload:
@@ -192,10 +152,10 @@ class Bnb4BitHfQuantizer(HfQuantizer):
     def is_trainable(self) -> bool:
         return True
 
-    def _dequantize(self, model):
+    def _dequantize(self, model, dtype=None):
         from ..integrations import dequantize_and_replace
 
-        model = dequantize_and_replace(model, quantization_config=self.quantization_config)
+        model = dequantize_and_replace(model, quantization_config=self.quantization_config, dtype=dtype)
         return model
 
     def get_quantize_ops(self):

transformers/quantizers/quantizer_bnb_8bit.py

@@ -83,19 +83,6 @@ 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():
@@ -115,8 +102,12 @@ class Bnb8BitHfQuantizer(HfQuantizer):
             )
         return device_map
 
-    def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype":
-        return torch.int8
+    def param_element_size(self, model: "PreTrainedModel", param_name: str, param: "torch.Tensor") -> float:
+        "Return the element size (in bytes) for `param_name`."
+        if self.param_needs_quantization(model, param_name):
+            # 8-bit
+            return 1
+        return super().param_element_size(model, param_name, param)
 
     def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
         import bitsandbytes as bnb
@@ -133,13 +124,12 @@ class Bnb8BitHfQuantizer(HfQuantizer):
         self,
         model: "PreTrainedModel",
         device_map,
-        keep_in_fp32_modules: list[str] | None = None,
         **kwargs,
     ):
         from ..integrations import replace_with_bnb_linear
 
         self.modules_to_not_convert = self.get_modules_to_not_convert(
-            model, self.quantization_config.llm_int8_skip_modules, keep_in_fp32_modules
+            model, self.quantization_config.llm_int8_skip_modules, model._keep_in_fp32_modules
         )
 
         if self.quantization_config.llm_int8_enable_fp32_cpu_offload:
@@ -161,10 +151,10 @@ class Bnb8BitHfQuantizer(HfQuantizer):
     def is_trainable(self) -> bool:
         return True
 
-    def _dequantize(self, model):
+    def _dequantize(self, model, dtype=None):
         from ..integrations import dequantize_and_replace
 
-        model = dequantize_and_replace(model, quantization_config=self.quantization_config)
+        model = dequantize_and_replace(model, quantization_config=self.quantization_config, dtype=dtype)
         return model
 
     def get_quantize_ops(self):

transformers/quantizers/quantizer_compressed_tensors.py

@@ -59,10 +59,7 @@ class CompressedTensorsHfQuantizer(HfQuantizer):
             )
 
     def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
-        if dtype is None:
-            logger.info("Loading model using torch.float16 for compressed-tensors quantization")
-            dtype = torch.float16
-        elif dtype != torch.float16:
+        if dtype != torch.float16:
             logger.info("We suggest you to set `dtype=torch.float16` for better efficiency with compressed_tensors.")
         return dtype
 

transformers/quantizers/quantizer_eetq.py

@@ -64,16 +64,7 @@ class EetqHfQuantizer(HfQuantizer):
                 )
 
     def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
-        if dtype is None:
-            dtype = torch.float16
-            logger.info(
-                "Overriding dtype=%s with `dtype=torch.float16` due to "
-                "requirements of `eetq` to enable model loading in 8-bit. "
-                "Pass your own dtype to specify the dtype of the remaining non-linear layers or pass"
-                " dtype=torch.float16 to remove this warning.",
-                dtype,
-            )
-        elif dtype != torch.float16:
+        if dtype != torch.float16:
             logger.info("We suggest you to set `dtype=torch.float16` for better efficiency with EETQ.")
         return dtype
 
@@ -92,13 +83,12 @@ class EetqHfQuantizer(HfQuantizer):
     def _process_model_before_weight_loading(
         self,
         model: "PreTrainedModel",
-        keep_in_fp32_modules: list[str] | None = None,
         **kwargs,
     ):
         from ..integrations import replace_with_eetq_linear
 
         self.modules_to_not_convert = self.get_modules_to_not_convert(
-            model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules
+            model, self.quantization_config.modules_to_not_convert, model._keep_in_fp32_modules
         )
 
         model = replace_with_eetq_linear(

transformers/quantizers/quantizer_fbgemm_fp8.py

@@ -84,19 +84,11 @@ class FbgemmFp8HfQuantizer(HfQuantizer):
                 )
 
     def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
-        if dtype is None:
-            dtype = torch.bfloat16
-            logger.info(
-                "Overriding dtype=%s with `dtype=torch.bloat16` due to "
-                "requirements of `fbgemm-gpu` to enable model loading in fp8. "
-                "Pass your own dtype to specify the dtype of the remaining non-linear layers or pass"
-                " dtype=torch.bfloat16 to remove this warning.",
-                dtype,
-            )
-        elif dtype == torch.float16:
-            raise ValueError(
-                "You cannot use FP8 with dtype=torch.float16. We recommend you passing dtype=torch.bfloat16"
+        if dtype != torch.bfloat16:
+            logger.warning_once(
+                f"Setting dtype to {dtype}, but only bfloat16 is supported right now. Overwriting torch_dtype to bfloat16."
             )
+            dtype = torch.bfloat16
         return dtype
 
     def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
@@ -119,13 +111,12 @@ class FbgemmFp8HfQuantizer(HfQuantizer):
     def _process_model_before_weight_loading(
         self,
         model: "PreTrainedModel",
-        keep_in_fp32_modules: list[str] | None = None,
         **kwargs,
     ):
         from ..integrations import replace_with_fbgemm_fp8_linear
 
         self.modules_to_not_convert = self.get_modules_to_not_convert(
-            model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules
+            model, self.quantization_config.modules_to_not_convert, model._keep_in_fp32_modules
         )
 
         model = replace_with_fbgemm_fp8_linear(

transformers/quantizers/quantizer_finegrained_fp8.py

@@ -33,7 +33,7 @@ class FineGrainedFP8HfQuantizer(HfQuantizer):
             return
 
         if not torch.cuda.is_available() and not is_torch_xpu_available():
-            if self.pre_quantized and not self.quantization_config.dequantize:
+            if self.pre_quantized:
                 logger.warning_once(
                     "Using FP8 quantized models requires a GPU or XPU, we will default to dequantizing the model to bf16 since no GPU or XPU is available"
                 )
@@ -46,10 +46,13 @@ class FineGrainedFP8HfQuantizer(HfQuantizer):
             compute_capability = torch.cuda.get_device_capability()
             major, minor = compute_capability
             if (major < 8) or (major == 8 and minor < 9):
-                raise ValueError(
+                logger.warning_once(
                     "FP8 quantized models is only supported on GPUs with compute capability >= 8.9 (e.g 4090/H100)"
-                    f", actual = `{major}.{minor}`"
+                    f", actual = `{major}.{minor}`. We will default to dequantizing the model to bf16. Feel free "
+                    f"to use a different quantization method like bitsandbytes or torchao"
                 )
+                self.quantization_config.dequantize = True
+                return
 
         device_map = kwargs.get("device_map")
         if device_map is None:
@@ -82,16 +85,22 @@ class FineGrainedFP8HfQuantizer(HfQuantizer):
                 return True
         return False
 
+    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, this is neeed as when `pre_quantized`` is False, we don't set the dtype of the FP8Linear in order to correctly load the weights
+            return 1
+        return super().param_element_size(model, param_name, param)
+
     def _process_model_before_weight_loading(
         self,
         model: "PreTrainedModel",
-        keep_in_fp32_modules: list[str] | None = None,
         **kwargs,
     ):
         from ..integrations.finegrained_fp8 import replace_with_fp8_linear
 
         self.modules_to_not_convert = self.get_modules_to_not_convert(
-            model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules
+            model, self.quantization_config.modules_to_not_convert, model._keep_in_fp32_modules
         )
 
         model = replace_with_fp8_linear(
@@ -103,7 +112,7 @@ class FineGrainedFP8HfQuantizer(HfQuantizer):
 
     # NOTE: TP is applied before quantization so this is only to add hooks.
     # Quantization is incompatible with DTensors, so we have to anyway have
-    # gathers! But it should be model independant -> figure out where to put
+    # gathers! But it should be model independent -> figure out where to put
     # the gather and that's it.
     def update_tp_plan(self, config):
         if "Qwen3" in config.__class__.__name__:
@@ -137,10 +146,6 @@ class FineGrainedFP8HfQuantizer(HfQuantizer):
     def is_trainable(self) -> bool:
         return False
 
-    def get_accelerator_warm_up_factor(self):
-        # Pre-processing is done cleanly, so we can allocate everything here
-        return 2
-
     def get_quantize_ops(self):
         from ..integrations.finegrained_fp8 import Fp8Quantize
 

transformers/quantizers/quantizer_fp_quant.py

@@ -78,11 +78,11 @@ class FPQuantHfQuantizer(HfQuantizer):
                 )
 
     def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
-        if dtype is None:
-            logger.info("`dtype` is None. Setting `dtype=torch.bfloat16` for qutlass compatibility.")
+        if dtype != torch.bfloat16:
+            logger.warning_once(
+                f"Setting dtype to {dtype}, but only bfloat16 is supported right now. Overwriting torch_dtype to bfloat16."
+            )
             dtype = torch.bfloat16
-        elif dtype != torch.bfloat16:
-            raise ValueError(f"Invalid `dtype` {dtype}. fp_quant quantization only supports `dtype=torch.bfloat16`.")
         return dtype
 
     def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:

transformers/quantizers/quantizer_gptq.py

@@ -66,10 +66,7 @@ class GptqHfQuantizer(HfQuantizer):
             raise ImportError("The gptqmodel version should be >= 1.4.3, optimum version should >= 1.24.0")
 
     def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
-        if dtype is None:
-            dtype = torch.float16
-            logger.info("Loading the model in `torch.float16`. To overwrite it, set `dtype` manually.")
-        elif dtype != torch.float16:
+        if dtype != torch.float16:
             logger.info("We suggest you to set `dtype=torch.float16` for better efficiency with GPTQ.")
         return dtype
 

transformers/quantizers/quantizer_higgs.py

@@ -69,10 +69,7 @@ class HiggsHfQuantizer(HfQuantizer):
                 )
 
     def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
-        if dtype is None:
-            logger.info("`dtype` is None. Setting `dtype=torch.float16` for FLUTE compatibility.")
-            dtype = torch.float16
-        elif dtype != torch.float16 and dtype != torch.bfloat16:
+        if dtype != torch.float16 and dtype != torch.bfloat16:
             raise ValueError(
                 f"Invalid `dtype` {dtype}. HIGGS quantization only supports `dtype=torch.float16` or `dtype=torch.bfloat16`."
             )
@@ -116,13 +113,12 @@ class HiggsHfQuantizer(HfQuantizer):
     def _process_model_before_weight_loading(
         self,
         model: "PreTrainedModel",
-        keep_in_fp32_modules: list[str] | None = None,
         **kwargs,
     ):
         from ..integrations import replace_with_higgs_linear
 
         self.modules_to_not_convert = self.get_modules_to_not_convert(
-            model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules
+            model, self.quantization_config.modules_to_not_convert, model._keep_in_fp32_modules
         )
 
         replace_with_higgs_linear(