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

transformers/quantizers/quantizer_fbgemm_fp8.py

@@ -19,21 +19,14 @@ from .base import HfQuantizer
 if TYPE_CHECKING:
     from ..modeling_utils import PreTrainedModel
 
-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 ..utils import is_accelerate_available, is_fbgemm_gpu_available, is_torch_available, logging
 from .quantizers_utils import get_module_from_name
 
 
 if is_torch_available():
     import torch
 
+
 logger = logging.get_logger(__name__)
 
 
@@ -42,41 +35,54 @@ 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_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():
+        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():
             raise ImportError(
-                "Loading an FP8 fbgemm quantized model on CUDA requires fbgemm-gpu library"
+                "Using fbgemm fp8 quantization 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 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)"
-                )
+
+        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)"
+            )
 
         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/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'. "
+                "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'. "
             )
-        elif isinstance(device_map, dict):
-            if not self.pre_quantized and ("cpu" in device_map.values() or "disk" in device_map.values()):
+        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())
+            ):
                 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. "
@@ -84,11 +90,19 @@ class FbgemmFp8HfQuantizer(HfQuantizer):
                 )
 
     def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
-        if dtype != torch.bfloat16:
-            logger.warning_once(
-                f"Setting dtype to {dtype}, but only bfloat16 is supported right now. Overwriting torch_dtype to bfloat16."
-            )
+        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"
+            )
         return dtype
 
     def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
@@ -108,25 +122,116 @@ 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, model._keep_in_fp32_modules
+            model, self.quantization_config.modules_to_not_convert, 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,
-            tp_plan=model._tp_plan,
+            config=config,
+            tp_plan=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 = {
@@ -174,14 +279,9 @@ class FbgemmFp8HfQuantizer(HfQuantizer):
 
         return config
 
-    def is_serializable(self):
+    def is_serializable(self, safe_serialization=None):
         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,19 +20,25 @@ 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 self.quantization_config.dequantize:
-            return
-
-        if not torch.cuda.is_available() and not is_torch_xpu_available():
+        if (not (torch.cuda.is_available() or is_torch_xpu_available())) 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,13 +52,10 @@ class FineGrainedFP8HfQuantizer(HfQuantizer):
             compute_capability = torch.cuda.get_device_capability()
             major, minor = compute_capability
             if (major < 8) or (major == 8 and minor < 9):
-                logger.warning_once(
+                raise ValueError(
                     "FP8 quantized models is only supported on GPUs with compute capability >= 8.9 (e.g 4090/H100)"
-                    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"
+                    f", actual = `{major}.{minor}`"
                 )
-                self.quantization_config.dequantize = True
-                return
 
         device_map = kwargs.get("device_map")
         if device_map is None:
@@ -61,12 +64,11 @@ 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 isinstance(device_map, dict):
+        elif device_map is not None:
             if (
                 not self.pre_quantized
-                and len(device_map) > 1
-                and "cpu" in device_map.values()
-                or "disk" in device_map.values()
+                and isinstance(device_map, dict)
+                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."
@@ -74,6 +76,76 @@ 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
 
@@ -85,34 +157,46 @@ 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, model._keep_in_fp32_modules
+            model, self.quantization_config.modules_to_not_convert, 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 independent -> figure out where to put
+    # gathers! But it should be model independant -> figure out where to put
     # the gather and that's it.
     def update_tp_plan(self, config):
         if "Qwen3" in config.__class__.__name__:
@@ -139,13 +223,17 @@ class FineGrainedFP8HfQuantizer(HfQuantizer):
 
         return config
 
-    def is_serializable(self):
+    def is_serializable(self, safe_serialization=None):
         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
 
@@ -158,9 +246,8 @@ 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", "activation_scale"],
+                    source_patterns=["weight$", "weight_scale_inv"],
                     target_patterns="weight",
                     operations=[Fp8Dequantize(self)],
                 )

transformers/quantizers/quantizer_fp_quant.py

@@ -36,10 +36,13 @@ 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():
@@ -65,35 +68,66 @@ 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):
-            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."
-                )
+        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."
+            )
 
     def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
-        if dtype != torch.bfloat16:
-            logger.warning_once(
-                f"Setting dtype to {dtype}, but only bfloat16 is supported right now. Overwriting torch_dtype to bfloat16."
-            )
+        if dtype is None:
+            logger.info("`dtype` is None. Setting `dtype=torch.bfloat16` for qutlass compatibility.")
             dtype = torch.bfloat16
-        return dtype
+        elif dtype != torch.bfloat16:
+            raise ValueError(f"Invalid `dtype` {dtype}. fp_quant quantization only supports `dtype=torch.bfloat16`.")
 
-    def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
-        from fp_quant import FPQuantLinear
+        return dtype
 
-        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 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 _process_model_before_weight_loading(
         self,
@@ -108,6 +142,20 @@ 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):
@@ -118,33 +166,15 @@ class FPQuantHfQuantizer(HfQuantizer):
             )
         return trainable
 
-    def is_serializable(self):
+    def is_serializable(self, safe_serialization=None):
         return True
 
-    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 []
+    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