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

transformers/quantizers/quantizer_fbgemm_fp8.py

@@ -19,14 +19,21 @@ from .base import HfQuantizer
 if TYPE_CHECKING:
     from ..modeling_utils import PreTrainedModel
 
-from ..utils import is_accelerate_available, is_fbgemm_gpu_available, is_torch_available, logging
+from ..utils import (
+    is_accelerate_available,
+    is_fbgemm_gpu_available,
+    is_kernels_available,
+    is_torch_available,
+    is_torch_cuda_available,
+    is_torch_xpu_available,
+    logging,
+)
 from .quantizers_utils import get_module_from_name
 
 
 if is_torch_available():
     import torch
 
-
 logger = logging.get_logger(__name__)
 
 
@@ -35,54 +42,41 @@ class FbgemmFp8HfQuantizer(HfQuantizer):
     FP8 quantization using fbgemm kernels
     """
 
-    requires_parameters_quantization = True
     requires_calibration = False
 
-    required_packages = ["fbgemm-gpu", "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_torch_available():
-            raise ImportError(
-                "Using fbgemm fp8 quantization requires torch >= 2.1.0"
-                "Please install the latest version of torch ( pip install --upgrade torch )"
-            )
-        if not is_fbgemm_gpu_available():
+        if not is_torch_cuda_available() and not is_torch_xpu_available():
+            raise ImportError("Using fbgemm fp8 quantization requires a GPU or XPU")
+        if is_torch_xpu_available() and not is_kernels_available():
+            raise ImportError("Using FP8 fbgemm on XPU requires kernels (`pip install kernels`)")
+        if is_torch_cuda_available() and not is_fbgemm_gpu_available():
             raise ImportError(
-                "Using fbgemm fp8 quantization requires fbgemm-gpu library"
+                "Loading an FP8 fbgemm quantized model on CUDA requires fbgemm-gpu library"
                 "Please install the latest version of fbgemm-gpu library by following : https://pytorch.org/FBGEMM/fbgemm_gpu-development/InstallationInstructions.html#fbgemm-gpu-install-libraries"
             )
-
         if not is_accelerate_available():
             raise ImportError(
                 "Loading an FP8 quantized model requires accelerate (`pip install --upgrade accelerate`)"
             )
-
-        if not torch.cuda.is_available():
-            raise RuntimeError("Using FP8 quantized models with fbgemm kernels requires a GPU")
-
-        compute_capability = torch.cuda.get_device_capability()
-        major, minor = compute_capability
-        if major < 9:
-            raise ValueError(
-                "FP8 quantized models is only supported on GPUs with compute capability >= 9.0 (e.g H100)"
-            )
+        if is_torch_cuda_available():
+            compute_capability = torch.cuda.get_device_capability()
+            major, _ = compute_capability
+            if major < 9:
+                raise ValueError(
+                    "FP8 quantized models is only supported on GPUs with compute capability >= 9.0 (e.g H100)"
+                )
 
         device_map = kwargs.get("device_map")
         if device_map is None:
             logger.warning_once(
-                "You have loaded an FP8 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. To remove this warning, pass device_map = 'cuda'. "
+                "You have loaded an FP8 model on CPU and have a CUDA/XPU device available, make sure to set "
+                "your model on a GPU/XPU device in order to run your model. To remove this warning, pass device_map = 'cuda' or 'xpu' or 'auto'. "
             )
-        elif device_map is not None:
-            if (
-                not self.pre_quantized
-                and isinstance(device_map, dict)
-                and ("cpu" in device_map.values() or "disk" in device_map.values())
-            ):
+        elif isinstance(device_map, dict):
+            if not self.pre_quantized and ("cpu" in device_map.values() or "disk" in device_map.values()):
                 raise ValueError(
                     "You are attempting to load an FP8 model with a device_map that contains a CPU or disk device."
                     "This is not supported when the model is quantized on the fly. "
@@ -90,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:
@@ -122,116 +108,25 @@ class FbgemmFp8HfQuantizer(HfQuantizer):
                 return True
         return False
 
-    def create_quantized_param(
-        self,
-        model: "PreTrainedModel",
-        param_value: "torch.Tensor",
-        param_name: str,
-        target_device: "torch.device",
-        **kwargs,
-    ):
-        from ..integrations import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts
-
-        module, tensor_name = get_module_from_name(model, param_name)
-
-        # Sanity checks
-        if isinstance(module, FbgemmFp8Linear):
-            if self.pre_quantized or tensor_name == "bias":
-                if tensor_name == "weight" and param_value.dtype != torch.float8_e4m3fn:
-                    raise ValueError("Expect quantized weights but got an unquantized weight")
-            else:
-                if tensor_name == "weight_scale":
-                    raise ValueError("Expect unquantized weights but got a quantized weight_scale")
-        if isinstance(module, FbgemmFp8Llama4TextExperts):
-            if not (self.pre_quantized or tensor_name == "bias"):
-                if tensor_name == "gate_up_proj_scale" or tensor_name == "down_proj_scale":
-                    raise ValueError("Expect unquantized weights but got a quantized weight_scale")
-
-        if isinstance(module, FbgemmFp8Llama4TextExperts):
-            if tensor_name == "gate_up_proj":
-                # Process each expert separately
-                # Transpose the second and third dimension
-                transposed_param = param_value.transpose(1, 2)
-
-                # Reshape to 2D for quantization
-                original_shape = transposed_param.shape
-                flattened_param = transposed_param.reshape(-1, original_shape[-1])
-
-                # Quantize using per row instead of per column
-                new_value_flat, weight_scale_flat = torch.ops.fbgemm.quantize_fp8_per_row(flattened_param)
-
-                # Reshape back to original dimensions
-                new_value = new_value_flat.reshape(original_shape)
-                new_value = new_value.transpose(1, 2)
-                weight_scale = weight_scale_flat.reshape(original_shape[0], 1, original_shape[1])
-            elif tensor_name == "down_proj":
-                # Process each expert separately
-                # Transpose the weights for proper quantization
-                transposed_param = param_value.transpose(1, 2)
-
-                # Reshape to 2D for quantization
-                original_shape = transposed_param.shape
-                flattened_param = transposed_param.reshape(-1, original_shape[-1])
-
-                # Quantize using per column
-                new_value_flat, weight_scale_flat = torch.ops.fbgemm.quantize_fp8_per_row(flattened_param)
-
-                # Reshape back to original dimensions
-                new_value = new_value_flat.reshape(original_shape)
-                new_value = new_value.transpose(1, 2)
-                weight_scale = weight_scale_flat.reshape(original_shape[0], original_shape[1], 1)
-
-            module._parameters[f"{tensor_name}_scale"] = torch.nn.Parameter(weight_scale.to(target_device))
-        else:
-            new_value, weight_scale = torch.ops.fbgemm.quantize_fp8_per_row(param_value)
-            module._parameters[f"{tensor_name}_scale"] = torch.nn.Parameter(
-                weight_scale.view(weight_scale.shape[0], 1).to(target_device)
-            )
-
-        module._parameters[tensor_name] = torch.nn.Parameter(new_value.to(target_device))
-
-        del param_name
-
     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
 
-        tp_plan = model._tp_plan
         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
         )
 
-        config = model.config
         model = replace_with_fbgemm_fp8_linear(
             model,
             modules_to_not_convert=self.modules_to_not_convert,
             quantization_config=self.quantization_config,
             pre_quantized=self.pre_quantized,
-            config=config,
-            tp_plan=tp_plan,
+            tp_plan=model._tp_plan,
         )
 
-        model.config.quantization_config = self.quantization_config
-
-    def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> list[str]:
-        from ..integrations import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts
-
-        not_missing_keys = []
-        for name, module in model.named_modules():
-            if isinstance(module, (FbgemmFp8Linear, FbgemmFp8Llama4TextExperts)):
-                for missing in missing_keys:
-                    if (
-                        (name in missing or name in f"{prefix}.{missing}")
-                        and not missing.endswith(".weight")
-                        and not missing.endswith(".bias")
-                    ):
-                        not_missing_keys.append(missing)
-        return [k for k in missing_keys if k not in not_missing_keys]
-
     def update_tp_plan(self, config):
         if "Llama4" in config.__class__.__name__:
             text_plan = {
@@ -279,9 +174,14 @@ class FbgemmFp8HfQuantizer(HfQuantizer):
 
         return config
 
-    def is_serializable(self, safe_serialization=None):
+    def is_serializable(self):
         return True
 
     @property
     def is_trainable(self) -> bool:
         return False
+
+    def get_quantize_ops(self):
+        from ..integrations.fbgemm_fp8 import FbgemmFp8Quantize
+
+        return FbgemmFp8Quantize(self)

transformers/quantizers/quantizer_finegrained_fp8.py

@@ -20,25 +20,19 @@ class FineGrainedFP8HfQuantizer(HfQuantizer):
     Supports both e4m3fn formats based on platform.
     """
 
-    requires_parameters_quantization = True
     requires_calibration = False
-    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_torch_available():
-            raise ImportError(
-                "Using fp8 quantization requires torch >= 2.1.0"
-                "Please install the latest version of torch ( pip install --upgrade torch )"
-            )
-
         if not is_accelerate_available():
             raise ImportError("Loading an FP8 quantized model requires accelerate (`pip install accelerate`)")
 
-        if (not (torch.cuda.is_available() or is_torch_xpu_available())) and not self.quantization_config.dequantize:
+        if self.quantization_config.dequantize:
+            return
+
+        if not torch.cuda.is_available() and not is_torch_xpu_available():
             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"
@@ -52,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:
@@ -64,11 +61,12 @@ class FineGrainedFP8HfQuantizer(HfQuantizer):
                 "your model on a GPU or XPU device in order to run your model. To remove this warning, "
                 "pass device_map = 'cuda' or 'xpu'. "
             )
-        elif device_map is not None:
+        elif isinstance(device_map, dict):
             if (
                 not self.pre_quantized
-                and isinstance(device_map, dict)
-                and ("cpu" in device_map.values() or "disk" in device_map.values())
+                and len(device_map) > 1
+                and "cpu" in device_map.values()
+                or "disk" in device_map.values()
             ):
                 raise ValueError(
                     "You are attempting to load an FP8 model with a device_map that contains a cpu/disk device."
@@ -76,76 +74,6 @@ class FineGrainedFP8HfQuantizer(HfQuantizer):
                     "Please use a quantized checkpoint or remove the cpu/disk device from the device_map."
                 )
 
-    def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
-        if dtype is None:
-            logger.info("Setting dtype to torch.float32 as no dtype was specified in from_pretrained")
-            dtype = torch.float32
-        return dtype
-
-    # TODO: make this into a `ConversionType` ops -> potentially requires all weights on all ranks
-    # depending on the layer type (moe -> no if ep)
-    def create_quantized_param(
-        self,
-        model: "PreTrainedModel",
-        param_value: "torch.Tensor",
-        param_name: str,
-        target_device: "torch.device",
-        **kwargs,
-    ):
-        from ..integrations.finegrained_fp8 import FP8Linear
-        from ..modeling_utils import _load_parameter_into_model
-
-        # Sanity checks
-        module, tensor_name = get_module_from_name(model, param_name)
-        if isinstance(module, FP8Linear):
-            if self.pre_quantized or tensor_name == "bias":
-                if tensor_name == "weight" and param_value.dtype != torch.float8_e4m3fn:
-                    raise ValueError("Expect quantized weights but got an unquantized weight")
-            else:
-                return
-                # if tensor_name == "weight_scale_inv":
-                #     raise ValueError("Expect unquantized weights but got a quantized weight_scale")
-
-        param_value = param_value.to(target_device)
-
-        # Get FP8 min/max values
-        fp8_min = torch.finfo(torch.float8_e4m3fn).min
-        fp8_max = torch.finfo(torch.float8_e4m3fn).max
-
-        block_size_m, block_size_n = self.quantization_config.weight_block_size
-
-        rows, cols = param_value.shape[-2:]
-
-        if rows % block_size_m != 0 or cols % block_size_n != 0:
-            raise ValueError(
-                f"Matrix dimensions ({rows}, {cols}) must be divisible by block sizes ({block_size_m}, {block_size_n})"
-            )
-        param_value_orig_shape = param_value.shape
-
-        param_value = param_value.reshape(
-            -1, rows // block_size_m, block_size_m, cols // block_size_n, block_size_n
-        ).permute(0, 1, 3, 2, 4)
-
-        # Calculate scaling factor for each block
-        max_abs = torch.amax(torch.abs(param_value), dim=(-1, -2))
-        scale = fp8_max / max_abs
-        scale_orig_shape = scale.shape
-        scale = scale.unsqueeze(-1).unsqueeze(-1)
-
-        # Quantize the weights
-        quantized_param = torch.clamp(param_value * scale, min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
-
-        quantized_param = quantized_param.permute(0, 1, 3, 2, 4)
-        # Reshape back to matrix shape
-        quantized_param = quantized_param.reshape(param_value_orig_shape)
-
-        # Reshape scale to match the number of blocks
-        scale = scale.reshape(scale_orig_shape).squeeze().reciprocal()
-
-        # Load into the model
-        _load_parameter_into_model(model, param_name, quantized_param)
-        _load_parameter_into_model(model, param_name.rsplit(".", 1)[0] + ".weight_scale_inv", scale)
-
     def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
         from ..integrations.finegrained_fp8 import FP8Expert, FP8Linear
 
@@ -157,46 +85,34 @@ 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
 
-        # takes 2 fucking seconds
         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
         )
 
-        # while this one is 81ms :)
         model = replace_with_fp8_linear(
             model,
             modules_to_not_convert=self.modules_to_not_convert,
             quantization_config=self.quantization_config,
+            pre_quantized=self.pre_quantized,
         )
 
-        model.config.quantization_config = self.quantization_config
-
-    def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> list[str]:
-        from ..integrations import FP8Linear
-
-        not_missing_keys = []
-        for name, module in model.named_modules():
-            if isinstance(module, FP8Linear):
-                for missing in missing_keys:
-                    if (
-                        (name in missing or name in f"{prefix}.{missing}")
-                        and not missing.endswith(".weight")
-                        and not missing.endswith(".bias")
-                    ):
-                        not_missing_keys.append(missing)
-        return [k for k in missing_keys if k not in not_missing_keys]
-
     # 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__:
@@ -223,17 +139,13 @@ class FineGrainedFP8HfQuantizer(HfQuantizer):
 
         return config
 
-    def is_serializable(self, safe_serialization=None):
+    def is_serializable(self):
         return True
 
     @property
     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
 
@@ -246,8 +158,9 @@ class FineGrainedFP8HfQuantizer(HfQuantizer):
         if self.pre_quantized and self.quantization_config.dequantize:
             return [
                 # either use the dollar sign, or permute the source patterns to start matching against the scales first
+                # We also collect the activation scales, they will not be used
                 WeightConverter(
-                    source_patterns=["weight$", "weight_scale_inv"],
+                    source_patterns=["weight$", "weight_scale_inv", "activation_scale"],
                     target_patterns="weight",
                     operations=[Fp8Dequantize(self)],
                 )

transformers/quantizers/quantizer_fp_quant.py

@@ -36,13 +36,10 @@ class FPQuantHfQuantizer(HfQuantizer):
     """
 
     requires_calibration = False
-    requires_parameters_quantization = True
     is_qat_trainable = True
-    required_packages = ["fp_quant"]
 
     def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs):
         super().__init__(quantization_config, **kwargs)
-        self.quantization_config = quantization_config
 
     def validate_environment(self, device_map, **kwargs):
         if not torch.cuda.is_available() and not is_torch_xpu_available():
@@ -68,66 +65,35 @@ class FPQuantHfQuantizer(HfQuantizer):
                 "You are attempting to load a FPQuant model without setting device_map."
                 " Please set device_map comprised of 'cuda' devices."
             )
-        elif (
-            isinstance(device_map, dict)
-            and ("cpu" in device_map.values() or "disk" in device_map.values())
-            and not self.quantization_config.pseudoquantization
-        ):
-            raise ValueError(
-                "You are attempting to load a FPQuant 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."
-            )
+        elif isinstance(device_map, dict):
+            if (
+                not self.quantization_config.pseudoquantization
+                and len(device_map) > 1
+                and "cpu" in device_map.values()
+                or "disk" in device_map.values()
+            ):
+                raise ValueError(
+                    "You are attempting to load a FPQuant model with a device_map that contains a CPU or disk device."
+                    " This is not supported. Please remove the CPU or disk device from the device_map."
+                )
 
     def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
-        if dtype is None:
-            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 create_quantized_param(
-        self,
-        model: "PreTrainedModel",
-        param_value: "torch.Tensor",
-        param_name: str,
-        target_device: "torch.device",
-        **kwargs,
-    ):
-        module, _ = get_module_from_name(model, param_name)
-
-        if target_device == "cpu" and param_name.endswith("weight"):
-            # Works agains hard-coded missing key dispatch to CPU
-            return
-
-        # The module holds either:
-        #  * `weight` when `store_master_weights=True`
-        #  * `qweight` and `scales` when `store_master_weights=False` and `pseudoquantization=False`
-        #  * `dqweight` when `store_master_weights=False` and `pseudoquantization=True`
-
-        if param_name.endswith(".qweight"):
-            # Loading a real quantized checkpoint without master weights
-            module.qweight = torch.nn.Parameter(
-                param_value.to(target_device),
-                requires_grad=False,
-            )
-            module.weight = None
-            module.dqweight = None
-            return
-
-        if param_name.endswith(".dqweight"):
-            # Loading a pseudo-quantized checkpoint without master weights
-            module.dqweight = torch.nn.Parameter(param_value.to(target_device))
-            module.weight = None
-            module.qweight = None
-            module.scales = None
-            return
-
-        # Loading master weights or an unquantized checkpoint
-        module.weight = torch.nn.Parameter(param_value.to(target_device))
-        # Let pre-forward handle the quantization and set None where necessary
-        module.pre_forward()
+    def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
+        from fp_quant import FPQuantLinear
+
+        module, tensor_name = get_module_from_name(model, param_name)
+        if isinstance(module, FPQuantLinear) and tensor_name in ["weight", "qweight", "dqweight"]:
+            # Only quantize weights of FPQuantLinear modules that are not already quantized
+            return True
+        else:
+            return False
 
     def _process_model_before_weight_loading(
         self,
@@ -142,20 +108,6 @@ class FPQuantHfQuantizer(HfQuantizer):
             model,
             fp_quant_linear_config=adapt_fp_quant_config(self.quantization_config),
         )
-        model.config.quantization_config = self.quantization_config
-
-    def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> list[str]:
-        from fp_quant import FPQuantLinear
-
-        fp_quant_names = {name for name, module in model.named_modules() if isinstance(module, FPQuantLinear)}
-
-        def should_exclude(key: str) -> bool:
-            if key.endswith(".weight") or key.endswith(".bias"):
-                return False
-            full_key = f"{prefix}.{key}"
-            return any(name in key or name in full_key for name in fp_quant_names)
-
-        return [key for key in missing_keys if not should_exclude(key)]
 
     @property
     def is_trainable(self, model: Optional["PreTrainedModel"] = None):
@@ -166,15 +118,33 @@ class FPQuantHfQuantizer(HfQuantizer):
             )
         return trainable
 
-    def is_serializable(self, safe_serialization=None):
+    def is_serializable(self):
         return True
 
-    def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
-        from fp_quant import FPQuantLinear
-
-        module, tensor_name = get_module_from_name(model, param_name)
-        if isinstance(module, FPQuantLinear) and tensor_name in ["weight", "qweight", "dqweight"]:
-            # Only quantize weights of FPQuantLinear modules that are not already quantized
-            return True
-        else:
-            return False
+    def get_quantize_ops(self):
+        from ..integrations.fp_quant import FpQuantQuantize
+
+        return FpQuantQuantize(self)
+
+    def get_weight_conversions(self):
+        from ..core_model_loading import WeightConverter
+        from ..integrations.fp_quant import FpQuantDeserialize
+
+        if self.pre_quantized:
+            if self.quantization_config.pseudoquantization:
+                return [
+                    WeightConverter(
+                        source_patterns=[".dqweight"],
+                        target_patterns=".dqweight",
+                        operations=[FpQuantDeserialize(self)],
+                    ),
+                ]
+            else:
+                return [
+                    WeightConverter(
+                        source_patterns=[".qweight"],
+                        target_patterns=".qweight",
+                        operations=[FpQuantDeserialize(self)],
+                    ),
+                ]
+        return []