support gpt-oss mxfp4 directly loading (#1401)

Signed-off-by: Xin He <xin3.he@intel.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: xin3he

auto_round/modeling/fused_moe/moe_experts_interface.py

@@ -458,6 +458,9 @@ def _unfuse_single_projection(
     return linears
 
 
+_logged_memory_before_replacement = False  # To ensure we only log memory once before replacements
+
+
 def _unfuse_experts_weights_inplace(
     module: nn.Module,
     check_decorator: bool = True,
@@ -494,8 +497,12 @@ def _unfuse_experts_weights_inplace(
         logger.debug(f"Skipping unfuse for {module.__class__.__name__}: does not support @use_experts_implementation")
         return False
 
-    memory_monitor.update()
-    memory_monitor.log_summary("Before applying custom replacements")
+    global _logged_memory_before_replacement
+    if not _logged_memory_before_replacement:
+        _logged_memory_before_replacement = True
+        memory_monitor.update()
+        memory_monitor.log_summary("Before applying custom replacements")
+
     # Get first projection to determine num_experts and layout
     first_proj_name = next(iter(detected_projections))
     first_param = getattr(module, first_proj_name)
@@ -594,9 +601,6 @@ def _unfuse_experts_weights_inplace(
     # Install compact repr to collapse identical expert containers in print output
     _install_compact_expert_repr(module)
 
-    memory_monitor.update()
-    memory_monitor.log_summary("After applying custom replacements")
-
     return True
 
 
@@ -645,6 +649,9 @@ def prepare_model_for_moe_quantization(model: nn.Module, implementation: str = L
             model.config._experts_implementation = impl_to_set
             logger.debug(f"Set model.config._experts_implementation = '{impl_to_set}'")
 
+        memory_monitor.update()
+        memory_monitor.log_summary("After applying custom replacements")
+
     return unfused_modules
 
 

auto_round/modeling/fused_moe/replace_modules.py

@@ -101,7 +101,8 @@ def _log_first_moe_block(model: torch.nn.Module, label: str) -> None:
     for name, module in model.named_modules():
         if name.endswith(".experts"):
             logger.info(f"Experts ({label}) [{name}] ({module.__class__.__name__}):\n{module}")
-            return
+            return True
+    return False
 
 
 @dump_mem_usage("Materializing model", log_level="debug")
@@ -289,15 +290,24 @@ def apply_replacements(
         The model with modules replaced.
     """
     _import_required_replacements(model)
+    _raw_expert_is_logged = False
 
     # Custom replacements first
     if is_custom_model(model):
-        _log_first_moe_block(model, "before replacement")
+
+        if not _raw_expert_is_logged:
+            _raw_expert_is_logged = _log_first_moe_block(model, "before replacement")
+
         _apply_custom_replacements(model)
-        _log_first_moe_block(model, "after replacement")
+
     if auto_detect_moe and is_transformers_version_greater_or_equal_5():
-        _log_first_moe_block(model, "before replacement")
+
+        if not _raw_expert_is_logged:
+            _raw_expert_is_logged = _log_first_moe_block(model, "before replacement")
+
         _handle_moe_modules(model)
+
+    if _raw_expert_is_logged:
         _log_first_moe_block(model, "after replacement")
 
     return model

auto_round/special_model_handler.py

@@ -274,7 +274,6 @@ def get_predefined_ignore_layers(model: torch.nn.Module) -> list[str]:
     if not layers and is_moe_model_via_config(config):
         for name, _ in model.named_modules():
             if name.endswith(".gate"):
-                print(name)
                 layers.append(name)
 
     return list(dict.fromkeys(layers))

auto_round/utils/model.py

@@ -248,6 +248,40 @@ def _check_accelerate_version():
         )
 
 
+_MXFP4_SUPPORTED_MODEL_TYPES = {"gpt_oss"}
+
+
+def _is_mxfp4_model(model_path, trust_remote_code=True):
+    """Check if a model is an MXFP4 quantized model supported for direct loading.
+
+    Only checks when transformers >= 5.0.0. Returns False immediately for older versions,
+    adding zero overhead to non-MXFP4 model loading.
+    """
+    if version.parse(transformers.__version__) < version.parse("5.0.0"):
+        return False
+    from transformers import AutoConfig
+
+    try:  # in case of config loading failure for new models
+        config = AutoConfig.from_pretrained(model_path, trust_remote_code=trust_remote_code)
+    except:
+        return False
+
+    model_type = getattr(config, "model_type", "")
+    if model_type not in _MXFP4_SUPPORTED_MODEL_TYPES:
+        return False
+
+    quant_config = getattr(config, "quantization_config", None)
+    if quant_config is None:
+        return False
+
+    quant_method = (
+        quant_config.get("quant_method", "")
+        if isinstance(quant_config, dict)
+        else getattr(quant_config, "quant_method", "")
+    )
+    return quant_method == "mxfp4" and model_type in _MXFP4_SUPPORTED_MODEL_TYPES
+
+
 def llm_load_model(
     pretrained_model_name_or_path: str,
     platform: str = "hf",
@@ -284,6 +318,18 @@ def llm_load_model(
     if device_str is not None and "hpu" in device_str:
         torch_dtype = torch.bfloat16
 
+    is_mxfp4 = _is_mxfp4_model(pretrained_model_name_or_path, trust_remote_code=trust_remote_code)
+    load_kwargs = {
+        "torch_dtype": torch_dtype,
+        "trust_remote_code": trust_remote_code,
+        "device_map": "auto" if use_auto_mapping else None,
+    }
+    if is_mxfp4:
+        from transformers import Mxfp4Config
+
+        load_kwargs["quantization_config"] = Mxfp4Config(dequantized=True)
+        logger.info("Detected MXFP4 quantized model, using Mxfp4Config(dequantized=True) for loading.")
+
     is_glm = bool(re.search("chatglm", pretrained_model_name_or_path.lower()))
 
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, trust_remote_code=trust_remote_code)
@@ -295,40 +341,22 @@ def llm_load_model(
     if is_hpex_available():
         # For loading FP8 model on HPU
         with fake_cuda_for_hpu(), fake_triton_for_hpu(), override_cuda_device_capability():
-            model = model_cls.from_pretrained(
-                pretrained_model_name_or_path,
-                torch_dtype=torch_dtype,
-                trust_remote_code=trust_remote_code,
-                device_map="auto" if use_auto_mapping else None,
-            )
+            model = model_cls.from_pretrained(pretrained_model_name_or_path, **load_kwargs)
     else:
         try:
-            model = model_cls.from_pretrained(
-                pretrained_model_name_or_path,
-                torch_dtype=torch_dtype,
-                trust_remote_code=trust_remote_code,
-                device_map="auto" if use_auto_mapping else None,
-            )
+            model = model_cls.from_pretrained(pretrained_model_name_or_path, **load_kwargs)
         except ValueError as e:
             if "FP8 quantized" in str(e):
                 with override_cuda_device_capability():
-                    model = model_cls.from_pretrained(
-                        pretrained_model_name_or_path,
-                        torch_dtype=torch_dtype,
-                        trust_remote_code=trust_remote_code,
-                        device_map="auto" if use_auto_mapping else None,
-                    )
+                    model = model_cls.from_pretrained(pretrained_model_name_or_path, **load_kwargs)
                 logger.warning("the support for fp8 model as input is experimental, please use with caution.")
             else:
                 raise
 
         except OSError as e:
             logger.warning(f"fail to load {pretrained_model_name_or_path}, set trust_remote_code to False and retry.")
             model = model_cls.from_pretrained(
-                pretrained_model_name_or_path,
-                torch_dtype=torch_dtype,
-                trust_remote_code=False,
-                device_map="auto" if use_auto_mapping else None,
+                pretrained_model_name_or_path, **{**load_kwargs, "trust_remote_code": False}
             )
 
     model = model.eval()

test/test_cpu/models/test_moe_model.py

@@ -10,7 +10,7 @@
 
 from ...helpers import get_model_path, transformers_version
 
-gpt_oss_name_or_path = get_model_path("unsloth/gpt-oss-20b-BF16")
+gpt_oss_name_or_path = get_model_path("openai/gpt-oss-20b")
 llama4_name_or_path = get_model_path("meta-llama/Llama-4-Scout-17B-16E-Instruct")
 qwen3_vl_moe_name_or_path = get_model_path("Qwen/Qwen3-VL-30B-A3B-Instruct")
 # local path for debug

test/test_cuda/models/test_moe_model.py

@@ -14,7 +14,7 @@
 @pytest.fixture
 def setup_gpt_oss():
     """Fixture to set up the GPT-OSS model and tokenizer."""
-    model_name = "/models/gpt-oss-20b-BF16"
+    model_name = "openai/gpt-oss-20b"
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
     config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
     config.num_hidden_layers = 1  # Reduce layers for testing