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

transformers/quantizers/quantizer_mxfp4.py

@@ -43,21 +43,17 @@ class Mxfp4HfQuantizer(HfQuantizer):
     FP4 quantization using fbgemm kernels
     """
 
-    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
         self.triton_kernels_hub = None
 
     def _lazy_import_kernels(self):
         """Lazy import and initialize kernels only when needed"""
         if self.triton_kernels_hub is None:
             try:
-                from kernels import get_kernel
+                from ..integrations.hub_kernels import get_kernel
 
                 self.triton_kernels_hub = get_kernel("kernels-community/triton_kernels")
             except ImportError:
@@ -74,7 +70,7 @@ class Mxfp4HfQuantizer(HfQuantizer):
         if self.quantization_config.dequantize:
             return
 
-        if not (torch.cuda.is_available() or torch.xpu.is_available()):
+        if not torch.cuda.is_available() and not torch.xpu.is_available():
             if self.pre_quantized:
                 logger.warning_once(
                     "Using MXFP4 quantized models requires a GPU, we will default to dequantizing the model to bf16"
@@ -131,185 +127,39 @@ class Mxfp4HfQuantizer(HfQuantizer):
                 "You have loaded an FP4 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 device_map = 'xpu'. "
             )
-        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 FP4 model with a device_map that contains a CPU or disk device."
                     "This is not supported when the model is quantized on the fly. "
                     "Please use a quantized checkpoint or remove the CPU or disk device from the device_map."
                 )
 
-    def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
-        if dtype is None:
-            dtype = torch.bfloat16
-            logger.info(
-                "Overriding dtype=%s with `dtype=torch.bfloat16` due to "
-                "requirements of `fbgemm-gpu` to enable model loading in fp4. "
-                "Pass your own dtype to specify the dtype of the remaining non-linear layers or pass"
-                " dtype=torch.bfloat16 to remove this warning.",
-                dtype,
-            )
-        return dtype
-
     def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
         from ..integrations import Mxfp4GptOssExperts
-        from ..models.gpt_oss.modeling_gpt_oss import GptOssExperts
 
-        if self.pre_quantized:
-            return False
-        # if we are dequantizing, the model doesn't have scales, and blocks only params like gate_up_proj and down_proj so we need to handle this case differently
-        if self.quantization_config.dequantize and ("blocks" in param_name or "scales" in param_name):
-            module, tensor_name = get_module_from_name(model, param_name[: -len("_blocks")])
-        else:
-            module, tensor_name = get_module_from_name(model, param_name)
-        if isinstance(module, Mxfp4GptOssExperts) or (
-            isinstance(module, GptOssExperts) and self.quantization_config.dequantize
-        ):
+        module, tensor_name = get_module_from_name(model, param_name)
+        if isinstance(module, Mxfp4GptOssExperts):
             if tensor_name in ["down_proj_bias", "gate_up_proj_bias"]:
                 return False
             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 (
-            Mxfp4GptOssExperts,
-            dequantize,
-            load_and_swizzle_mxfp4,
-            quantize_to_mxfp4,
-            swizzle_mxfp4,
-        )
-        from ..models.gpt_oss.modeling_gpt_oss import GptOssExperts
-
-        if not self.pre_quantized:
-            triton_kernels_hub = self._lazy_import_kernels()
-            module, _ = get_module_from_name(model, param_name)
-            with torch.device(target_device):
-                if isinstance(module, Mxfp4GptOssExperts):
-                    triton_weight_tensor, weight_scale = quantize_to_mxfp4(param_value, triton_kernels_hub)
-                    PrecisionConfig, FlexCtx, InFlexData = (
-                        triton_kernels_hub.matmul_ogs.PrecisionConfig,
-                        triton_kernels_hub.matmul_ogs.FlexCtx,
-                        triton_kernels_hub.matmul_ogs.InFlexData,
-                    )
-                    triton_weight_tensor, weight_scale = swizzle_mxfp4(
-                        triton_weight_tensor, weight_scale, triton_kernels_hub
-                    )
-
-                    proj = "gate_up_proj" if "gate_up_proj" in param_name else "down_proj"
-                    setattr(module, proj, triton_weight_tensor)
-                    setattr(
-                        module,
-                        f"{proj}_precision_config",
-                        PrecisionConfig(weight_scale=weight_scale, flex_ctx=FlexCtx(rhs_data=InFlexData())),
-                    )
-
-                    delattr(module, f"{proj}_blocks")
-                    delattr(module, f"{proj}_scales")
-
-        # The params going here are either gate_up_proj_blocks, or down_proj_blocks, or gate_up_proj_scales, or down_proj_scales
-        else:
-            #  This is when loading a quantized model (blocks and scales exist)
-            empty_param = kwargs.get("empty_param")
-            casting_dtype = kwargs.get("casting_dtype")
-            to_contiguous = kwargs.get("to_contiguous")
-            rank = kwargs.get("rank")
-            device_mesh = kwargs.get("device_mesh")
-            if ("blocks" in param_name or "scales" in param_name) and self.quantization_config.dequantize:
-                # blocks and scales have the same length that's why this works for both
-                module, _ = get_module_from_name(model, param_name[: -len("_blocks")])
-            else:
-                module, _ = get_module_from_name(model, param_name)
-
-            shard_kwargs = {
-                "empty_param": empty_param,
-                "casting_dtype": casting_dtype,
-                "to_contiguous": to_contiguous,
-                "rank": rank,
-                "device_mesh": device_mesh,
-                "model": model,
-            }
-
-            if isinstance(module, Mxfp4GptOssExperts) or (
-                isinstance(module, GptOssExperts) and self.quantization_config.dequantize
-            ):
-                if self.quantization_config.dequantize:
-                    # dq_param_name is the name of the parameter without the blocks or scales suffix, it's used in this case since we don't switch linears
-                    # so we only have the original param name
-                    dq_param_name = param_name[: -len("_blocks")]
-                    dequantize(module, param_name, param_value, target_device, dq_param_name, **shard_kwargs)
-                else:
-                    load_and_swizzle_mxfp4(
-                        module,
-                        param_name,
-                        param_value,
-                        target_device,
-                        self._lazy_import_kernels(),
-                        **shard_kwargs,
-                    )
-
     def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
-        # we are not really dequantizing, we are just removing everything related to quantization here
-        if self.quantization_config.dequantize:
-            self.remove_quantization_config(model)
         # clean cache due to triton ops
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
         elif torch.xpu.is_available():
             torch.xpu.empty_cache()
 
-    def update_expected_keys(self, model: "PreTrainedModel", expected_keys: list[str], checkpoint_keys: list[str]):
-        # Replace expected_keys for experts' gate_up_proj and down_proj with their _blocks and _scales variants
-        new_expected_keys = []
-        for key in expected_keys:
-            if key.endswith(".mlp.experts.gate_up_proj"):
-                base = key[: -len("gate_up_proj")]
-                new_expected_keys.append(base + "gate_up_proj_blocks")
-                new_expected_keys.append(base + "gate_up_proj_scales")
-            elif key.endswith(".mlp.experts.down_proj"):
-                base = key[: -len("down_proj")]
-                new_expected_keys.append(base + "down_proj_blocks")
-                new_expected_keys.append(base + "down_proj_scales")
-            elif not self.pre_quantized:
-                # in this case, we are quantizing the model so we need to update the keys as we changed the layers
-                if key.endswith(".mlp.experts.down_proj_blocks"):
-                    base = key[: -len("down_proj_blocks")]
-                    new_expected_keys.append(base + "down_proj")
-                elif key.endswith(".mlp.experts.gate_up_proj_blocks"):
-                    base = key[: -len("gate_up_proj_blocks")]
-                    new_expected_keys.append(base + "gate_up_proj")
-                elif key.endswith("scales"):
-                    # we remove it the scales as the checkpoint don't contain them
-                    continue
-                else:
-                    new_expected_keys.append(key)
-            else:
-                new_expected_keys.append(key)
-        return new_expected_keys
-
     def _process_model_before_weight_loading(
         self,
         model: "PreTrainedModel",
-        keep_in_fp32_modules: list[str] | None = None,
+        use_kernels: bool = False,
         **kwargs,
     ):
         from ..integrations import replace_with_mxfp4_linear
 
-        self.modules_to_not_convert = self.get_modules_to_not_convert(
-            model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules
-        )
-
-        use_kernels = kwargs.get("use_kernels", False)
         # if we are using kernels, we can't use the quantized model, since the forward pass is different and needs special handling
         if use_kernels:
             logger.warning_once(
@@ -318,30 +168,13 @@ class Mxfp4HfQuantizer(HfQuantizer):
             )
             self.quantization_config.dequantize = True
 
-        config = model.config
-        model = replace_with_mxfp4_linear(
-            model,
-            modules_to_not_convert=self.modules_to_not_convert,
-            quantization_config=self.quantization_config,
-            config=config,
+        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.config.quantization_config = self.quantization_config
-
-    def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> list[str]:
-        from ..integrations import Mxfp4GptOssExperts
-
-        not_missing_keys = []
-        for name, module in model.named_modules():
-            if isinstance(module, Mxfp4GptOssExperts):
-                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]
+        model = replace_with_mxfp4_linear(
+            model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config
+        )
 
     def update_tp_plan(self, config):
         if "GptOssConfig" in config.__class__.__name__:
@@ -369,20 +202,7 @@ class Mxfp4HfQuantizer(HfQuantizer):
                 )
         return config
 
-    def get_param_name(self, param_name: str) -> str:
-        if self.quantization_config.dequantize:
-            if "_blocks" in param_name:
-                return param_name.replace("_blocks", "")
-            elif "_scales" in param_name:
-                return param_name.replace("_scales", "")
-        elif not self.pre_quantized:
-            if param_name.endswith("gate_up_proj"):
-                return param_name.replace("gate_up_proj", "gate_up_proj_blocks")
-            if param_name.endswith("down_proj"):
-                return param_name.replace("down_proj", "down_proj_blocks")
-        return param_name
-
-    def get_state_dict_and_metadata(self, model, safe_serialization: bool = False):
+    def get_state_dict_and_metadata(self, model):
         from ..integrations import Mxfp4GptOssExperts
 
         state_dict = model.state_dict()
@@ -421,7 +241,7 @@ class Mxfp4HfQuantizer(HfQuantizer):
         metadata = {}
         return state_dict, metadata
 
-    def is_serializable(self, safe_serialization=None):
+    def is_serializable(self):
         return True
 
     @property

transformers/quantizers/quantizer_quanto.py

@@ -40,23 +40,17 @@ class QuantoHfQuantizer(HfQuantizer):
     Quantizer for the quanto library
     """
 
-    required_packages = ["quanto", "accelerate"]
-    requires_parameters_quantization = True
     requires_calibration = False
 
     def __init__(self, quantization_config: QuantoConfig, **kwargs):
         super().__init__(quantization_config, **kwargs)
-        self.post_init()
-
-    def post_init(self):
-        r"""
-        Safety checker
-        """
-        if self.quantization_config.activations is not None and not self.pre_quantized:
-            raise ValueError(
-                "We don't support quantizing the activations with transformers library."
-                "Use quanto library for more complex use cases such as activations quantization, calibration and quantization aware training."
-            )
+        map_to_param_size = {
+            "int8": 1,
+            "float8": 1,
+            "int4": 0.5,
+            "int2": 0.25,
+        }
+        self.quantized_param_size = map_to_param_size.get(self.quantization_config.weights, None)
 
     def validate_environment(self, *args, **kwargs):
         if not is_optimum_quanto_available():
@@ -67,42 +61,22 @@ class QuantoHfQuantizer(HfQuantizer):
             raise ImportError(
                 "Loading an optimum-quanto quantized model requires accelerate library (`pip install accelerate`)"
             )
-
-    def update_device_map(self, device_map):
-        if device_map is None:
-            device_map = {"": "cpu"}
-            logger.info(
-                "The device_map was not initialized. "
-                "Setting device_map to {'':'cpu'}. "
-                "If you want to use the model for inference, please set device_map ='auto'"
+        device_map = kwargs.get("device_map")
+        if isinstance(device_map, dict):
+            if len(device_map) > 1 and "cpu" in device_map.values() or "disk" in device_map.values():
+                raise ValueError(
+                    "You are attempting to load an model with a device_map that contains a CPU or disk device."
+                    "This is not supported with quanto when the model is quantized on the fly. "
+                    "Please remove the CPU or disk device from the device_map."
+                )
+        if self.quantization_config.activations is not None:
+            raise ValueError(
+                "We don't support quantizing the activations with transformers library."
+                "Use quanto library for more complex use cases such as activations quantization, calibration and quantization aware training."
             )
-        return device_map
-
-    def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
-        if dtype is None:
-            logger.info("You did not specify `dtype` in `from_pretrained`. Setting it to `torch.float32`.")
-            dtype = torch.float32
-        return dtype
-
-    def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> list[str]:
-        if is_optimum_quanto_available():
-            from optimum.quanto import QModuleMixin
-
-        not_missing_keys = []
-        for name, module in model.named_modules():
-            if isinstance(module, QModuleMixin):
-                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 param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
-        if is_optimum_quanto_available():
-            from optimum.quanto import QModuleMixin
+        from optimum.quanto import QModuleMixin
 
         module, tensor_name = get_module_from_name(model, param_name)
         # We only quantize the weights and the bias is not quantized.
@@ -116,50 +90,32 @@ class QuantoHfQuantizer(HfQuantizer):
         max_memory = {key: val * 0.90 for key, val in max_memory.items()}
         return max_memory
 
-    def create_quantized_param(
-        self,
-        model: "PreTrainedModel",
-        param_value: "torch.Tensor",
-        param_name: str,
-        target_device: "torch.device",
-        **kwargs,
-    ):
-        from ..modeling_utils import _load_parameter_into_model
-
-        _load_parameter_into_model(model, param_name, param_value.to(target_device))
-        module, _ = get_module_from_name(model, param_name)
-        module.freeze()
-        module.weight.requires_grad = False
-
-    def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype":
-        from accelerate.utils import CustomDtype
-
-        mapping = {
-            "int8": torch.int8,
-            "float8": CustomDtype.FP8,
-            "int4": CustomDtype.INT4,
-            "int2": CustomDtype.INT2,
-        }
-        target_dtype = mapping[self.quantization_config.weights]
-        return target_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) and self.quantized_param_size is not None:
+            return self.quantized_param_size
 
-    def _process_model_before_weight_loading(
-        self, model: "PreTrainedModel", keep_in_fp32_modules: list[str] | None = None, **kwargs
-    ):
+        return super().param_element_size(model, param_name, param)
+
+    def _process_model_before_weight_loading(self, model: "PreTrainedModel", **kwargs):
         from ..integrations import replace_with_quanto_layers
 
         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_quanto_layers(
+        model = replace_with_quanto_layers(
             model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config
         )
-        model.config.quantization_config = self.quantization_config
 
     @property
     def is_trainable(self) -> bool:
         return True
 
-    def is_serializable(self, safe_serialization=None):
+    def is_serializable(self):
         return False
+
+    def get_quantize_ops(self):
+        from ..integrations.quanto import QuantoQuantize
+
+        return QuantoQuantize(self)

transformers/quantizers/quantizer_quark.py

@@ -45,12 +45,6 @@ class QuarkHfQuantizer(HfQuantizer):
     """
 
     requires_calibration = True  # On-the-fly quantization with quark is not supported for now.
-    required_packages = ["quark"]
-
-    # Checkpoints are expected to be already quantized when loading a quark model. However, as some keys from
-    # the checkpoint might mismatch the model parameters keys, we use the `create_quantized_param` method
-    # to load the checkpoints, remapping the keys.
-    requires_parameters_quantization = True
 
     def __init__(self, quantization_config, **kwargs):
         super().__init__(quantization_config, **kwargs)
@@ -78,19 +72,44 @@ class QuarkHfQuantizer(HfQuantizer):
     def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
         return True
 
-    def create_quantized_param(self, model, param, param_name, param_device, **kwargs):
-        from ..modeling_utils import _load_parameter_into_model
-
-        postfix = param_name.split(".")[-1]
-
-        if postfix in CHECKPOINT_KEYS:
-            param_name = param_name.replace(postfix, CHECKPOINT_KEYS[postfix])
-
-        _load_parameter_into_model(model, param_name, param.to(param_device))
-
-    def is_serializable(self, safe_serialization=None):
+    def is_serializable(self):
         return False
 
     @property
     def is_trainable(self):
         return False
+
+    def get_weight_conversions(self):
+        from ..core_model_loading import WeightConverter
+        from ..integrations.quark import QuarkDeserialize
+        # In Quark, quantization is managed through a QParamsLinear module, which holds
+        # separate quantizers for the weights, inputs, and biases (e.g. weight_quantizer
+        # input_quantizer, bias_quantizer, etc.).
+        #
+        # When you call `module.state_dict()`, Quark automatically renames the quantizer
+        # parameters — for example, `input_quantizer.scale` becomes `input_scale` — and
+        # saves them directly at the parent module level.
+        #
+        # This means we cannot simply rename keys like `weight_scale` back to
+        # `weight_quantizer.scale` when loading the state_dict.
+        # Otherwise, the `missing_keys` list would still expect keys such as
+        # `weight_scale`, `bias_scale`, etc.
+        #
+        # To fix this, we keep the expected state_dict keys (like `weight_scale`,
+        # `bias_scale`, etc.) unchanged, and during the conversion step, we explicitly
+        # assign their values into the corresponding quantizer attributes
+        # (`weight_quantizer.scale`, `input_quantizer.scale`, and so on).
+
+        # You can notice here that in target_patterns we use the same key as the source_patterns,
+        # this is because we just want to collect the tensors, and we will rename them later in the convert function.
+        # We cannot rename directly or else the missing_keys list will not be able to find the tensors.
+        converters = []
+        for key in CHECKPOINT_KEYS.keys():
+            converters.append(
+                WeightConverter(
+                    source_patterns=[key],
+                    target_patterns=key,
+                    operations=[QuarkDeserialize(self)],
+                )
+            )
+        return converters

transformers/quantizers/quantizer_spqr.py

@@ -39,7 +39,6 @@ class SpQRHfQuantizer(HfQuantizer):
 
     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 torch.cuda.is_available():
@@ -52,36 +51,29 @@ class SpQRHfQuantizer(HfQuantizer):
             raise ImportError("Using `spqr` quantization requires SpQR: `pip install spqr_quant[gpu]`")
 
     def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
-        if dtype is None:
-            dtype = torch.float16
-            logger.info("Assuming SpQR inference on GPU and loading the model in `torch.float16`.")
-        elif dtype != torch.float16:
+        if dtype != torch.float16:
             raise ValueError(
-                "You cannot use any type other than torch.float16 for SpQR. Please either leave it None or set it to"
-                "torch.float16 explicitly."
+                "You cannot use any type other than torch.float16 for SpQR. Please set it totorch.float16 explicitly."
             )
         return dtype
 
     def _process_model_before_weight_loading(
         self,
         model: "PreTrainedModel",
-        keep_in_fp32_modules: list[str] | None = None,
         **kwargs,
     ):
         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_spqr_linear(
             model,
             quantization_config=self.quantization_config,
             modules_to_not_convert=self.modules_to_not_convert,
         )
-        model.config.quantization_config = self.quantization_config
 
     @property
     def is_trainable(self):
         return False
 
-    def is_serializable(self, safe_serialization=None):
+    def is_serializable(self):
         return True