Repository: huggingface/transformers
Pull request: Enable mxfp4 model on CPU

CPU mxfp4 moe kernel already implemented. Fix the check to enable mxfp4 on CPU.

Waiting for CPU kernel merge and release.

Changed files:
filename: src/transformers/integrations/hub_kernels.py
patch: @@ -158,6 +158,12 @@ def use_kernel_func_from_hub(func_name: str):
                     layer_name="MegaBlocksMoeMLP",
                 )
             },
+            "cpu": {
+                Mode.INFERENCE: LayerRepository(
+                    repo_id="kernels-community/megablocks",
+                    layer_name="CPUMegaBlocksMoeMLP",
+                )
+            },
         },
         "FastGELU": {
             "cuda": {

filename: src/transformers/integrations/mxfp4.py
patch: @@ -524,8 +524,12 @@ def load_and_swizzle_mxfp4(module, param_name, param_value, target_device, trito
             blocks = blocks.reshape(local_experts, module.intermediate_size * 2, -1)
         else:
             blocks = blocks.reshape(local_experts, -1, module.intermediate_size // 2)
-        if getattr(target_device, "type", target_device) == "cpu":
-            target_device = torch.accelerator.current_accelerator().type if hasattr(torch, "accelerator") else "cuda"
+        if (
+            getattr(target_device, "type", target_device) == "cpu"
+            and hasattr(torch, "accelerator")
+            and torch.accelerator.current_accelerator() is not None
+        ):
+            target_device = torch.accelerator.current_accelerator().type
         blocks = blocks.to(target_device).contiguous()
         scales = scales.to(target_device).contiguous()
         with on_device(target_device):
@@ -564,8 +568,12 @@ def swizzle_mxfp4_convertops(blocks, scales, module, proj, target_device, triton
     )
 
     local_experts = blocks.size(0)
-    if getattr(target_device, "type", target_device) == "cpu":
-        target_device = torch.accelerator.current_accelerator().type if hasattr(torch, "accelerator") else "cuda"
+    if (
+        getattr(target_device, "type", target_device) == "cpu"
+        and hasattr(torch, "accelerator")
+        and torch.accelerator.current_accelerator() is not None
+    ):
+        target_device = torch.accelerator.current_accelerator().type
 
     blocks = blocks.to(target_device).contiguous()
     scales = scales.to(target_device).contiguous()
@@ -574,8 +582,6 @@ def swizzle_mxfp4_convertops(blocks, scales, module, proj, target_device, triton
         blocks = blocks.reshape(local_experts, module.intermediate_size * 2, -1)
     else:
         blocks = blocks.reshape(local_experts, -1, module.intermediate_size // 2)
-    if getattr(target_device, "type", target_device) == "cpu":
-        target_device = "cuda"
 
     with on_device(target_device):
         triton_weight_tensor, weight_scale = swizzle_mxfp4(

filename: src/transformers/quantizers/quantizer_mxfp4.py
patch: @@ -70,30 +70,37 @@ def validate_environment(self, *args, **kwargs):
         if self.quantization_config.dequantize:
             return
 
-        if not torch.cuda.is_available() and not torch.xpu.is_available():
+        if not is_accelerate_available():
+            raise ImportError("Using mxfp4 requires Accelerate: `pip install accelerate`")
+
+        device = torch.accelerator.current_accelerator() or torch.device("cpu")
+        if device.type not in ["cuda", "xpu", "cpu"]:
             if self.pre_quantized:
                 logger.warning_once(
-                    "Using MXFP4 quantized models requires a GPU, we will default to dequantizing the model to bf16"
+                    f"Using MXFP4 quantized models requires model on cuda/xpu/cpu, but found {device}, we will default to dequantizing the model to bf16. To use mxfp4, please disable the current accelerator."
                 )
                 self.quantization_config.dequantize = True
                 return
             else:
-                raise RuntimeError("Quantizing a model using MXFP4 requires a GPU")
-
-        if not is_accelerate_available():
-            raise ImportError("Using mxfp4 requires Accelerate: `pip install accelerate`")
+                raise RuntimeError(
+                    f"Quantizing a model using MXFP4 requires model on cuda/xpu/cpu, but found {device}. To use mxfp4, please disable the current accelerator."
+                )
 
         if torch.xpu.is_available():
-            gpu_is_supported = True
+            is_device_supported_mxfp4 = True
             kernels_available = is_triton_available("3.5.0") and is_kernels_available()
-        else:
+        elif torch.cuda.is_available():
             compute_capability = torch.cuda.get_device_capability()
-            gpu_is_supported = compute_capability >= (7, 5)
+            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()
 
         if self.pre_quantized:
             # On unsupported GPUs or without kernels, we will dequantize the model to bf16
-            if not gpu_is_supported:
+            if not is_device_supported_mxfp4:
                 logger.warning_once(
                     "MXFP4 quantization is only supported on GPUs with compute capability >= 7.5 (e.g T4, A100, L4, H100, or B200) or XPUs (e.g Intel® Data Center GPU Max Series) "
                     "We will default to dequantizing the model to bf16."
@@ -107,32 +114,27 @@ def validate_environment(self, *args, **kwargs):
                 )
                 self.quantization_config.dequantize = True
                 return
-        elif not gpu_is_supported:
+        elif not is_device_supported_mxfp4:
             # we can't quantize the model in this case so we raise an error
             raise ValueError(
-                "MXFP4 quantization is only supported on GPUs with compute capability >= 7.5 (e.g T4, A100, L4, H100, or B200) or XPUs (e.g Intel® Data Center GPU Max Series) "
+                "MXFP4 quantization is only supported on GPUs with compute capability >= 7.5 (e.g T4, A100, L4, H100, or B200) or XPUs (e.g Intel® Data Center GPU Max Series) or CPU"
             )
         elif not kernels_available:
             # we can't quantize the model in this case so we raise an error
             raise ValueError(
-                "MXFP4 quantization requires Triton and kernels installed: CUDA requires Triton >= 3.4.0, XPU requires Triton >= 3.5.0"
+                "MXFP4 quantization requires Triton and kernels installed: CUDA requires Triton >= 3.4.0, XPU/CPU requires Triton >= 3.5.0"
             )
 
         if not self.pre_quantized:
             self._lazy_import_kernels()
 
         device_map = kwargs.get("device_map")
-        if device_map is None:
-            logger.warning_once(
-                "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 isinstance(device_map, dict):
-            if not self.pre_quantized and ("cpu" in device_map.values() or "disk" in device_map.values()):
+        if device_map is not None and isinstance(device_map, dict):
+            if not self.pre_quantized and "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."
+                    "You are attempting to load an FP4 model with a device_map that contains a 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."
+                    "Please use a quantized checkpoint or remove the disk device from the device_map."
                 )
 
     def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
@@ -161,13 +163,22 @@ def _process_model_before_weight_loading(
         from ..integrations import replace_with_mxfp4_linear
 
         # 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:
+        # only CPU kernels can work with pre-quantized models
+        device = torch.accelerator.current_accelerator() or torch.device("cpu")
+        if use_kernels and device.type not in ["cpu"]:
             logger.warning_once(
                 "You are using full precision kernels, we will dequantize the model to bf16. "
                 "To use the quantized model with quantization kernels, please set use_kernels=False"
             )
             self.quantization_config.dequantize = True
 
+        if not use_kernels and device.type in ["cpu"]:
+            logger.warning_once(
+                "MXFP4 inference on CPU requires use_kernels=True, but use_kernels is disabled. "
+                "We will dequantize the model to bf16. To run MXFP4 natively on CPU, please set use_kernels=True."
+            )
+            self.quantization_config.dequantize = True
+
         self.modules_to_not_convert = self.get_modules_to_not_convert(
             model, self.quantization_config.modules_to_not_convert, model._keep_in_fp32_modules
         )

filename: tests/models/gpt_oss/test_modeling_gpt_oss.py
patch: @@ -35,7 +35,6 @@
     require_deterministic_for_xpu,
     require_kernels,
     require_torch,
-    require_torch_accelerator,
     require_torch_gpu,
     slow,
     torch_device,
@@ -231,7 +230,6 @@ def generate_config_key(quantized, model, kernels, attn_impl, mode):
 
 
 @slow
-@require_torch_accelerator
 class GptOssIntegrationTest(unittest.TestCase):
     input_text = [
         "Roses are red, violets",
@@ -252,8 +250,13 @@ def tearDown(self):
     # ------------------------
     # Non-distributed inference
     # ------------------------
-    @staticmethod
-    def load_and_forward(model_id, attn_implementation, input_text, mode="eval", **pretrained_kwargs):
+    def load_and_forward(self, model_id, attn_implementation, input_text, mode="eval", **pretrained_kwargs):
+        if torch_device == "cpu":
+            if attn_implementation == "kernels-community/vllm-flash-attn3":
+                self.skipTest("vllm-flash-attn3 is not supported on CPU.")
+            if pretrained_kwargs.get("kernels", False) and mode == "train":
+                self.skipTest("CPU kernels only support inference.")
+
         model = AutoModelForCausalLM.from_pretrained(
             model_id,
             dtype=torch.bfloat16,
@@ -361,6 +364,12 @@ def run_distributed_test(quantized, model, kernels, attn_impl, mode):
     @parameterized.expand(PARAMETERS)
     @require_deterministic_for_xpu
     def test_model_outputs(self, quantized, model, kernels, attn_impl, mode):
+        if torch_device == "cpu":
+            if attn_impl == "kernels-community/vllm-flash-attn3":
+                self.skipTest("vllm-flash-attn3 is not supported on CPU.")
+            if kernels and mode == "train":
+                self.skipTest("CPU kernels only support inference.")
+
         if torch_device == "xpu" and attn_impl == "kernels-community/vllm-flash-attn3":
             self.skipTest("flash attention 3 is not supported on XPU yet.")
 
@@ -426,6 +435,9 @@ def test_model_outputs(self, quantized, model, kernels, attn_impl, mode):
     # ------------------------
     @parameterized.expand(PARAMETERS)
     def test_model_outputs_distributed(self, quantized, model, kernels, attn_impl, mode):
+        if torch_device == "cpu":
+            self.skipTest("Skip TP on CPU until verified.")
+
         if torch_device == "xpu" and attn_impl == "kernels-community/vllm-flash-attn3":
             self.skipTest("flash attention 3 is not supported on XPU yet.")
 
@@ -436,6 +448,12 @@ def test_model_outputs_distributed(self, quantized, model, kernels, attn_impl, m
     # ------------------------
     @parameterized.expand(PARAMETERS)
     def test_training_step(self, quantized, model, kernels, attn_impl, mode):
+        if torch_device == "cpu":
+            if attn_impl == "kernels-community/vllm-flash-attn3":
+                self.skipTest("vllm-flash-attn3 is not supported on CPU.")
+            if kernels and mode == "train":
+                self.skipTest("CPU kernels only support inference.")
+
         if mode != "train":
             self.skipTest("This test is only for training mode.")
 

filename: tests/quantization/mxfp4/test_mxfp4.py
patch: @@ -41,6 +41,8 @@
     REQUIRE_TRITON_MXFP4 = require_triton(min_version="3.4.0")
 elif hasattr(torch, "xpu") and torch.xpu.is_available():
     REQUIRE_TRITON_MXFP4 = require_triton(min_version="3.5.0")
+elif torch_device == "cpu":
+    REQUIRE_TRITON_MXFP4 = require_triton(min_version="3.5.0")
 else:
     REQUIRE_TRITON_MXFP4 = unittest.skip("test requires CUDA or XPU")
 
@@ -129,9 +131,8 @@ def test_quantizer_validation_no_accelerator(self):
             config = Mxfp4Config()
             quantizer = Mxfp4HfQuantizer(config)
             quantizer.pre_quantized = False
-
-            with self.assertRaises(RuntimeError):
-                quantizer.validate_environment()
+            # CPU already supported MXFP4
+            quantizer.validate_environment()
 
     @require_torch_gpu
     def test_quantizer_validation_low_compute_capability(self):
@@ -193,8 +194,8 @@ def test_quantizer_validation_order_dequantize_before_accelerator_check(self):
             quantizer = Mxfp4HfQuantizer(config)
             quantizer.pre_quantized = False
 
-            with self.assertRaises(RuntimeError):
-                quantizer.validate_environment()
+            # CPU already supported MXFP4
+            quantizer.validate_environment()
 
     def test_quantizer_validation_missing_triton(self):
         """Test quantizer validation when triton is not available"""
@@ -278,7 +279,7 @@ def test_quantize_to_mxfp4(self):
         quantizer = Mxfp4HfQuantizer(config)
 
         # Create dummy weight tensor
-        device = "xpu" if (hasattr(torch, "xpu") and torch.xpu.is_available()) else "cuda"
+        device = torch_device
         w = torch.randn(32, 64, 128, dtype=torch.bfloat16, device=torch.device(device))
 
         quantized_w, w_scale = quantize_to_mxfp4(w, quantizer._lazy_import_kernels())
@@ -369,9 +370,8 @@ def test_model_device_map_validation(self):
         quantizer = Mxfp4HfQuantizer(config)
         quantizer.pre_quantized = False
 
-        # Test with CPU in device map (should raise error for non-pre-quantized)
-        with self.assertRaises(ValueError):
-            quantizer.validate_environment(device_map={"": "cpu"})
+        # Test with CPU in device map (CPU already support mxfp4)
+        quantizer.validate_environment(device_map={"": "cpu"})
 
     def test_memory_footprint_comparison(self):
         """Test memory footprint differences between quantized and unquantized models"""