Repository: huggingface/transformers
Pull request: [Quantization] Fixing mxfp4 saving using reverse_op

# What does this PR do?

Fixes mxfp4 quantization by using a reverse_op during saving

Changed files:
filename: src/transformers/integrations/mxfp4.py
patch: @@ -125,22 +125,24 @@ def convert(
         input_dict: dict[str, torch.Tensor],
         model: torch.nn.Module | None = None,
         full_layer_name: str | None = None,
-        missing_keys=None,
+        missing_keys: list[str] | None = None,
         **kwargs,
     ) -> dict[str, torch.Tensor]:
-        if "_blocks" in input_dict.keys():
-            if isinstance(input_dict["_blocks"], list):
-                blocks = input_dict["_blocks"][0]
+        param_data = {}
+        proj = "gate_up_proj" if "gate_up_proj" in full_layer_name else "down_proj"
+        if f"{proj}_blocks" in input_dict.keys():
+            if isinstance(input_dict[f"{proj}_blocks"], list):
+                param_data[f"{proj}_blocks"] = input_dict[f"{proj}_blocks"][0]
             else:
-                blocks = input_dict["_blocks"]
-        if "_scales" in input_dict.keys():
-            if isinstance(input_dict["_scales"], list):
-                scales = input_dict["_scales"][0]
+                param_data[f"{proj}_blocks"] = input_dict[f"{proj}_blocks"]
+        if f"{proj}_scales" in input_dict.keys():
+            if isinstance(input_dict[f"{proj}_scales"], list):
+                param_data[f"{proj}_scales"] = input_dict[f"{proj}_scales"][0]
             else:
-                scales = input_dict["_scales"]
+                param_data[f"{proj}_scales"] = input_dict[f"{proj}_scales"]
 
         # Here we are dequantizing the weights
-        dequantized = dequantize_convertops(blocks, scales)
+        dequantized = dequantize_convertops(param_data[f"{proj}_blocks"], param_data[f"{proj}_scales"])
         return {full_layer_name: dequantized}
 
 
@@ -157,26 +159,27 @@ def convert(
         **kwargs,
     ) -> dict[str, torch.Tensor]:
         param_data = {}
-        if "_blocks" in input_dict.keys():
-            if isinstance(input_dict["_blocks"], list):
-                param_data["_blocks"] = input_dict["_blocks"][0]
+        proj = "gate_up_proj" if "gate_up_proj" in full_layer_name else "down_proj"
+
+        if f"{proj}_blocks" in input_dict.keys():
+            if isinstance(input_dict[f"{proj}_blocks"], list):
+                param_data[f"{proj}_blocks"] = input_dict[f"{proj}_blocks"][0]
             else:
-                param_data["_blocks"] = input_dict["_blocks"]
-        if "_scales" in input_dict.keys():
-            if isinstance(input_dict["_scales"], list):
-                param_data["_scales"] = input_dict["_scales"][0]
+                param_data[f"{proj}_blocks"] = input_dict[f"{proj}_blocks"]
+        if f"{proj}_scales" in input_dict.keys():
+            if isinstance(input_dict[f"{proj}_scales"], list):
+                param_data[f"{proj}_scales"] = input_dict[f"{proj}_scales"][0]
             else:
-                param_data["_scales"] = input_dict["_scales"]
+                param_data[f"{proj}_scales"] = input_dict[f"{proj}_scales"]
 
         # Eagerly set tensors on the module and perform swizzle
         module, _ = get_module_from_name(model, full_layer_name)
-        proj = "gate_up_proj" if "gate_up_proj" in full_layer_name else "down_proj"
         swizzle_mxfp4_convertops(
-            param_data["_blocks"],
-            param_data["_scales"],
+            param_data[f"{proj}_blocks"],
+            param_data[f"{proj}_scales"],
             module,
             proj,
-            param_data["_blocks"].device,
+            param_data[f"{proj}_blocks"].device,
             triton_kernels_hub,
         )
         missing_keys.discard(f"{full_layer_name}")
@@ -186,6 +189,61 @@ def convert(
         # We don't use set_param_for_module since it expects mainly a torch.nn.Parameter or a safetensors pointer
         return {}
 
+    @property
+    def reverse_op(self) -> ConversionOps:
+        return Mxfp4ReverseDeserialize(self.hf_quantizer)
+
+
+class Mxfp4ReverseDeserialize(ConversionOps):
+    def __init__(self, hf_quantizer):
+        self.hf_quantizer = hf_quantizer
+
+    def convert(
+        self,
+        input_dict: dict[str, torch.Tensor],
+        model: torch.nn.Module | None = None,
+        full_layer_name: str | None = None,
+        missing_keys: list[str] | None = None,
+        **kwargs,
+    ) -> dict[str, torch.Tensor]:
+        num_local_experts = getattr(model.config, "num_local_experts", 32)
+        hidden_size = getattr(model.config, "hidden_size", 2880)
+
+        proj = "gate_up_proj" if "gate_up_proj" in full_layer_name else "down_proj"
+
+        name = full_layer_name.rsplit("_", 1)[0]
+        module, _ = get_module_from_name(model, full_layer_name)
+        state_dict = {}
+        if isinstance(module, Mxfp4GptOssExperts):
+            if "bias" in full_layer_name:
+                name = full_layer_name.replace("_blocks", "")
+                state_dict[name] = getattr(module, proj + "_bias")
+                return state_dict
+            if "gate_up_proj" in full_layer_name:
+                state_dict[f"{name}_blocks"] = (
+                    module.gate_up_proj.storage.layout.unswizzle_data(module.gate_up_proj.storage.data)
+                    .transpose(-1, -2)
+                    .reshape(num_local_experts, -1, 90, 16)
+                )
+                state_dict[f"{name}_scales"] = (
+                    module.gate_up_proj_precision_config.weight_scale.storage.layout.unswizzle_data(
+                        module.gate_up_proj_precision_config.weight_scale.storage.data
+                    ).transpose(-1, -2)
+                )
+            else:
+                state_dict[f"{name}_blocks"] = (
+                    module.down_proj.storage.layout.unswizzle_data(module.down_proj.storage.data)
+                    .transpose(-1, -2)
+                    .reshape(num_local_experts, hidden_size, 90, -1)
+                )
+                state_dict[f"{name}_scales"] = (
+                    module.down_proj_precision_config.weight_scale.storage.layout.unswizzle_data(
+                        module.down_proj_precision_config.weight_scale.storage.data
+                    ).transpose(-1, -2)
+                )
+
+        return state_dict
+
 
 # Copied from GPT_OSS repo and vllm
 def quantize_to_mxfp4(w, triton_kernels_hub):

filename: src/transformers/quantizers/quantizer_mxfp4.py
patch: @@ -94,9 +94,11 @@ def validate_environment(self, *args, **kwargs):
             is_device_supported_mxfp4 = compute_capability >= (7, 5)
             kernels_available = is_triton_available("3.4.0") and is_kernels_available()
         elif device.type == "cpu":
-            # CPU support mxfp4 in kernels
             is_device_supported_mxfp4 = True
             kernels_available = is_triton_available("3.5.0") and is_kernels_available()
+        else:
+            is_device_supported_mxfp4 = False
+            kernels_available = False
 
         if self.pre_quantized:
             # On unsupported GPUs or without kernels, we will dequantize the model to bf16
@@ -217,37 +219,33 @@ def get_state_dict_and_metadata(self, model):
         from ..integrations import Mxfp4GptOssExperts
 
         state_dict = model.state_dict()
-
-        # Get num_local_experts from model config
         num_local_experts = getattr(model.config, "num_local_experts", 32)
         hidden_size = getattr(model.config, "hidden_size", 2880)
 
         for name, module in model.named_modules():
-            if (
+            if not (
                 isinstance(module, Mxfp4GptOssExperts)
                 and hasattr(module, "gate_up_proj")
                 and hasattr(module, "down_proj")
             ):
-                state_dict[f"{name}.gate_up_proj_blocks"] = (
-                    module.gate_up_proj.storage.layout.unswizzle_data(module.gate_up_proj.storage.data)
-                    .transpose(-1, -2)
-                    .reshape(num_local_experts, -1, 90, 16)
-                )
-                state_dict[f"{name}.gate_up_proj_scales"] = (
-                    module.gate_up_proj_precision_config.weight_scale.storage.layout.unswizzle_data(
-                        module.gate_up_proj_precision_config.weight_scale.storage.data
-                    ).transpose(-1, -2)
-                )
-                state_dict[f"{name}.down_proj_blocks"] = (
-                    module.down_proj.storage.layout.unswizzle_data(module.down_proj.storage.data)
-                    .transpose(-1, -2)
-                    .reshape(num_local_experts, hidden_size, 90, -1)
-                )
-                state_dict[f"{name}.down_proj_scales"] = (
-                    module.down_proj_precision_config.weight_scale.storage.layout.unswizzle_data(
-                        module.down_proj_precision_config.weight_scale.storage.data
-                    ).transpose(-1, -2)
-                )
+                continue
+
+            for proj in ("gate_up_proj", "down_proj"):
+                triton_tensor = getattr(module, proj)
+                precision_config = getattr(module, f"{proj}_precision_config")
+
+                blocks = triton_tensor.storage.layout.unswizzle_data(triton_tensor.storage.data).transpose(-1, -2)
+                if proj == "gate_up_proj":
+                    blocks = blocks.reshape(num_local_experts, -1, 90, 16)
+                else:
+                    blocks = blocks.reshape(num_local_experts, hidden_size, 90, -1)
+
+                scales = precision_config.weight_scale.storage.layout.unswizzle_data(
+                    precision_config.weight_scale.storage.data
+                ).transpose(-1, -2)
+
+                state_dict[f"{name}.{proj}_blocks"] = blocks
+                state_dict[f"{name}.{proj}_scales"] = scales
 
         metadata = {}
         return state_dict, metadata
@@ -270,21 +268,29 @@ def get_quantize_ops(self):
     def get_weight_conversions(self):
         from ..integrations.mxfp4 import Mxfp4Dequantize, Mxfp4Deserialize
 
-        if self.pre_quantized:
-            if self.quantization_config.dequantize:
-                return [
-                    WeightConverter(
-                        source_patterns=["_blocks", "_scales"],
-                        target_patterns="",
-                        operations=[Mxfp4Dequantize(self)],
-                    )
-                ]
-            else:
-                return [
-                    WeightConverter(
-                        source_patterns=["_blocks", "_scales"],
-                        target_patterns="",
-                        operations=[Mxfp4Deserialize(self)],
-                    )
-                ]
-        return []
+        if self.pre_quantized and self.quantization_config.dequantize:
+            return [
+                WeightConverter(
+                    source_patterns=["down_proj_blocks", "down_proj_scales"],
+                    target_patterns=r"down_proj$",
+                    operations=[Mxfp4Dequantize(self)],
+                ),
+                WeightConverter(
+                    source_patterns=["gate_up_proj_blocks", "gate_up_proj_scales"],
+                    target_patterns=["gate_up_proj$"],
+                    operations=[Mxfp4Dequantize(self)],
+                ),
+            ]
+
+        return [
+            WeightConverter(
+                source_patterns=["gate_up_proj_blocks", "gate_up_proj_scales"],
+                target_patterns=r"gate_up_proj$",
+                operations=[Mxfp4Deserialize(self)],
+            ),
+            WeightConverter(
+                source_patterns=["down_proj_blocks", "down_proj_scales"],
+                target_patterns=r"down_proj$",
+                operations=[Mxfp4Deserialize(self)],
+            ),
+        ]

filename: tests/quantization/mxfp4/test_mxfp4.py
patch: @@ -60,6 +60,7 @@ def _patch_no_accelerator():
         stack.enter_context(patch("torch.cuda.is_available", return_value=False))
         if hasattr(torch, "xpu"):
             stack.enter_context(patch("torch.xpu.is_available", return_value=False))
+        stack.enter_context(patch("torch.accelerator.current_accelerator", return_value=None))
         yield