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

transformers/quantizers/quantizer_hqq.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from collections import defaultdict
 from typing import TYPE_CHECKING
 
 from ..integrations import prepare_for_hqq_linear
@@ -49,10 +48,7 @@ class HqqHfQuantizer(HfQuantizer):
     nn.Linear modules are first tagged with quant_config in _process_model_before_weight_loading().
     """
 
-    use_keep_in_fp32_modules = False
-    requires_parameters_quantization = True
     requires_calibration = False
-    required_packages = ["hqq"]
 
     def __init__(self, quantization_config, **kwargs):
         if not is_hqq_available():
@@ -83,73 +79,67 @@ class HqqHfQuantizer(HfQuantizer):
             else:
                 self.using_multi_gpu = len(set(device_map.values())) > 1
 
-    def update_missing_keys(
-        self, model: "PreTrainedModel", missing_keys: list[str], prefix: str, **kwargs
-    ) -> list[str]:
-        if self.pre_quantized:
-            return [key for key in missing_keys if ("weight" not in key)]
-        else:
-            return missing_keys
-
-    # Adds missing keys for HQQLinear modules that are loaded but the model with initialized with torch.nn.Linear
-    def update_expected_keys(
-        self, model: "PreTrainedModel", expected_keys: list[str], loaded_keys: list[str]
-    ) -> list[str]:
-        if not self.pre_quantized:
-            return expected_keys
-
-        # Collects all quantizable (linear) layers
-        def _find_hqq_quantizable_layers(model, layers):
-            for name, module in model.named_children():
-                if isinstance(module, (torch.nn.Linear)):
-                    layers.add(module.name)
-                _find_hqq_quantizable_layers(module, layers)
-
-        new_keys = set(expected_keys)
-
-        # Name modules
-        for name, module in model.named_modules():
-            module.name = name
-
-        # valid modules are Linear layers that have HQQLinear state_dict. We ignore skip_modules and any layers with Linear state_dict() params
-        _valid_modules = set()
-        _find_hqq_quantizable_layers(model, _valid_modules)
-
-        # Remove skipped modules
-        _skipped_modules = set()
-        for _module in _valid_modules:
-            for _skip_module in model.config.quantization_config["skip_modules"]:
-                if _skip_module in _module:
-                    _skipped_modules.add(_module)
-        _valid_modules -= _skipped_modules
-
-        # Append new expected layers based on _ref_keys
-        _ref_keys = HQQLinear(
-            linear_layer=None,
-            quant_config=None,
-            compute_dtype=torch.float16,
-            device="cpu",
-            del_orig=False,
-        ).state_dict_keys() - {"bias"}
-
-        # Clean-up
-        _rm_keys = set()
-        for key in new_keys:
-            if any(_module in key for _module in _valid_modules):
-                _rm_keys.add(key)
-        new_keys -= _rm_keys
-        # At this point, new_keys contains all the keys of the layers that are NOT HQQLinear or torch.nn.Linear
-
-        # Re-populate Linear/HQQLinear
-        for _module in _valid_modules:
-            if _module + ".weight" in loaded_keys:
-                new_keys.add(_module + ".weight")
-            else:
-                new_keys.update({_module + "." + _ref_key for _ref_key in _ref_keys})
-            if _module + ".bias" in loaded_keys:
-                new_keys.add(_module + ".bias")
-
-        return list(new_keys)
+    # TODO: to remove
+    # Kept here in case we see some interest in adding support for it
+    # # Adds missing keys for HQQLinear modules that are loaded but the model with initialized with torch.nn.Linear
+    # def update_expected_keys(
+    #     self, model: "PreTrainedModel", expected_keys: list[str], loaded_keys: list[str]
+    # ) -> list[str]:
+    #     if not self.pre_quantized:
+    #         return expected_keys
+
+    #     # Collects all quantizable (linear) layers
+    #     def _find_hqq_quantizable_layers(model, layers):
+    #         for name, module in model.named_children():
+    #             if isinstance(module, (torch.nn.Linear)):
+    #                 layers.add(module.name)
+    #             _find_hqq_quantizable_layers(module, layers)
+
+    #     new_keys = set(expected_keys)
+
+    #     # Name modules
+    #     for name, module in model.named_modules():
+    #         module.name = name
+
+    #     # valid modules are Linear layers that have HQQLinear state_dict. We ignore skip_modules and any layers with Linear state_dict() params
+    #     _valid_modules = set()
+    #     _find_hqq_quantizable_layers(model, _valid_modules)
+
+    #     # Remove skipped modules
+    #     _skipped_modules = set()
+    #     for _module in _valid_modules:
+    #         for _skip_module in model.config.quantization_config["skip_modules"]:
+    #             if _skip_module in _module:
+    #                 _skipped_modules.add(_module)
+    #     _valid_modules -= _skipped_modules
+
+    #     # Append new expected layers based on _ref_keys
+    #     _ref_keys = HQQLinear(
+    #         linear_layer=None,
+    #         quant_config=None,
+    #         compute_dtype=torch.float16,
+    #         device="cpu",
+    #         del_orig=False,
+    #     ).state_dict_keys() - {"bias"}
+
+    #     # Clean-up
+    #     _rm_keys = set()
+    #     for key in new_keys:
+    #         if any(_module in key for _module in _valid_modules):
+    #             _rm_keys.add(key)
+    #     new_keys -= _rm_keys
+    #     # At this point, new_keys contains all the keys of the layers that are NOT HQQLinear or torch.nn.Linear
+
+    #     # Re-populate Linear/HQQLinear
+    #     for _module in _valid_modules:
+    #         if _module + ".weight" in loaded_keys:
+    #             new_keys.add(_module + ".weight")
+    #         else:
+    #             new_keys.update({_module + "." + _ref_key for _ref_key in _ref_keys})
+    #         if _module + ".bias" in loaded_keys:
+    #             new_keys.add(_module + ".bias")
+
+    #     return list(new_keys)
 
     def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
         module, _ = get_module_from_name(model, param_name)
@@ -157,87 +147,88 @@ class HqqHfQuantizer(HfQuantizer):
         # `create_quantized_param`, even when `self.is_quantized == True`
         return isinstance(module, torch.nn.Linear)
 
-    def create_quantized_param(
-        self,
-        model: "PreTrainedModel",
-        param_value: "torch.Tensor",
-        param_name: str,
-        target_device: "torch.device",
-        **kwargs,
-    ):
-        module, tensor_name = get_module_from_name(model, param_name)
-        module_name = param_name.rsplit(".", 1)[0]
-        parent_module, node = get_module_from_name(model, module_name)
-
-        quant_config = model.config.quantization_config["quant_config"]
-        skip_modules = model.config.quantization_config["skip_modules"]
-
-        # In this case we do not quantize this layer (it's explicitly skipped) -> simply load param
-        if any(skip_module in module.name for skip_module in skip_modules):
-            module.load_state_dict(
-                {tensor_name: param_value.to(device=target_device, dtype=self.dtype)}, strict=False, assign=True
-            )
-            return
-
-        # We need this hack as the model is not pre-prepared as an empty skeleton on meta device
-        if self.pre_quantized:
-            # Save them for later
-            if not hasattr(self, "hqq_params"):
-                self.hqq_params = defaultdict(dict)
-            self.hqq_params[module_name].update({tensor_name: param_value})
-            hqq_params = self.hqq_params[module_name]
-
-            # If they are all present and saved, make it a HQQLinear layer! (we cannot do it param after param because
-            # hqq does not support it...)
-            if all(k in hqq_params for k in self.hqq_keys) and ("bias" in hqq_params or module.bias is None):
-                hqq_layer = HQQLinear(
-                    linear_layer=None,
-                    quant_config=None,
-                    compute_dtype=self.dtype,
-                    device=target_device,
-                    del_orig=False,
-                )
-                hqq_layer.load_state_dict(hqq_params)
-
-                if hqq_layer.bias is not None and isinstance(hqq_layer.bias, torch.Tensor):
-                    hqq_layer.bias = torch.nn.Parameter(hqq_layer.bias)
-                if self.using_multi_gpu:
-                    hqq_layer = self._patch_layer_for_multigpu(hqq_layer)
-
-                setattr(parent_module, node, hqq_layer)
-                del self.hqq_params[module_name], module
-            return
-
-        # Load param in the module (without caring about device or dtype, it will be changed later)
-        module.load_state_dict({tensor_name: param_value}, strict=False, assign=True)
-
-        # If both the weight and bias have already been loaded, time to quantize!
-        module_is_ready = module.weight.device.type != "meta" and (
-            module.bias is None or module.bias.device.type != "meta"
-        )
-
-        if module_is_ready:
-            module_tag = ".".join(module.name.split(".")[-2:])
-            if "weight_quant_params" in quant_config:
-                module_quant_config = quant_config
-            elif module_tag in quant_config:
-                module_quant_config = quant_config[module_tag]
-
-            hqq_layer = HQQLinear(
-                module,
-                quant_config=module_quant_config,
-                compute_dtype=self.dtype,
-                device=target_device,
-                del_orig=True,
-            )
-
-            if hqq_layer.bias is not None and isinstance(hqq_layer.bias, torch.Tensor):
-                hqq_layer.bias = torch.nn.Parameter(hqq_layer.bias)
-
-            if self.using_multi_gpu:
-                hqq_layer = self._patch_layer_for_multigpu(hqq_layer)
-
-            setattr(parent_module, node, hqq_layer)
+    # TODO: to remove
+    # def create_quantized_param(
+    #     self,
+    #     model: "PreTrainedModel",
+    #     param_value: "torch.Tensor",
+    #     param_name: str,
+    #     target_device: "torch.device",
+    #     **kwargs,
+    # ):
+    #     module, tensor_name = get_module_from_name(model, param_name)
+    #     module_name = param_name.rsplit(".", 1)[0]
+    #     parent_module, node = get_module_from_name(model, module_name)
+
+    #     quant_config = model.config.quantization_config["quant_config"]
+    #     skip_modules = model.config.quantization_config["skip_modules"]
+
+    #     # In this case we do not quantize this layer (it's explicitly skipped) -> simply load param
+    #     if any(skip_module in module.name for skip_module in skip_modules):
+    #         module.load_state_dict(
+    #             {tensor_name: param_value.to(device=target_device, dtype=self.dtype)}, strict=False, assign=True
+    #         )
+    #         return
+
+    #     # We need this hack as the model is not pre-prepared as an empty skeleton on meta device
+    #     if self.pre_quantized:
+    #         # Save them for later
+    #         if not hasattr(self, "hqq_params"):
+    #             self.hqq_params = defaultdict(dict)
+    #         self.hqq_params[module_name].update({tensor_name: param_value})
+    #         hqq_params = self.hqq_params[module_name]
+
+    #         # If they are all present and saved, make it a HQQLinear layer! (we cannot do it param after param because
+    #         # hqq does not support it...)
+    #         if all(k in hqq_params for k in self.hqq_keys) and ("bias" in hqq_params or module.bias is None):
+    #             hqq_layer = HQQLinear(
+    #                 linear_layer=None,
+    #                 quant_config=None,
+    #                 compute_dtype=self.dtype,
+    #                 device=target_device,
+    #                 del_orig=False,
+    #             )
+    #             hqq_layer.load_state_dict(hqq_params)
+
+    #             if hqq_layer.bias is not None and isinstance(hqq_layer.bias, torch.Tensor):
+    #                 hqq_layer.bias = torch.nn.Parameter(hqq_layer.bias)
+    #             if self.using_multi_gpu:
+    #                 hqq_layer = self._patch_layer_for_multigpu(hqq_layer)
+
+    #             setattr(parent_module, node, hqq_layer)
+    #             del self.hqq_params[module_name], module
+    #         return
+
+    #     # Load param in the module (without caring about device or dtype, it will be changed later)
+    #     module.load_state_dict({tensor_name: param_value}, strict=False, assign=True)
+
+    #     # If both the weight and bias have already been loaded, time to quantize!
+    #     module_is_ready = module.weight.device.type != "meta" and (
+    #         module.bias is None or module.bias.device.type != "meta"
+    #     )
+
+    #     if module_is_ready:
+    #         module_tag = ".".join(module.name.split(".")[-2:])
+    #         if "weight_quant_params" in quant_config:
+    #             module_quant_config = quant_config
+    #         elif module_tag in quant_config:
+    #             module_quant_config = quant_config[module_tag]
+
+    #         hqq_layer = HQQLinear(
+    #             module,
+    #             quant_config=module_quant_config,
+    #             compute_dtype=self.dtype,
+    #             device=target_device,
+    #             del_orig=True,
+    #         )
+
+    #         if hqq_layer.bias is not None and isinstance(hqq_layer.bias, torch.Tensor):
+    #             hqq_layer.bias = torch.nn.Parameter(hqq_layer.bias)
+
+    #         if self.using_multi_gpu:
+    #             hqq_layer = self._patch_layer_for_multigpu(hqq_layer)
+
+    #         setattr(parent_module, node, hqq_layer)
 
     def _patch_layer_for_multigpu(self, hqq_layer):
         def forward_with_device(self, x):
@@ -263,7 +254,7 @@ class HqqHfQuantizer(HfQuantizer):
         model.is_hqq_serializable = self.is_serializable()
         return model
 
-    def is_serializable(self, safe_serialization=None):
+    def is_serializable(self):
         return True
 
     @property

transformers/quantizers/quantizer_mxfp4.py

@@ -43,14 +43,10 @@ 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):
@@ -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,12 +127,8 @@ 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. "
@@ -157,159 +149,30 @@ class Mxfp4HfQuantizer(HfQuantizer):
 
     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 +181,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, 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__:
@@ -382,7 +228,7 @@ class Mxfp4HfQuantizer(HfQuantizer):
                 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 +267,7 @@ class Mxfp4HfQuantizer(HfQuantizer):
         metadata = {}
         return state_dict, metadata
 
-    def is_serializable(self, safe_serialization=None):
+    def is_serializable(self):
         return True
 
     @property