Repository: huggingface/transformers
Pull request: Fixed MXFP4 model storage issue

# What does this PR do?

The current **MXFP4** model throws an `TypeError` during storage. Fix this issue.

For example, running the test `RUN_SLOW=1 pytest -rA test_mxfp4.py::Mxfp4ModelTest::test_save_mxfp4` will result in the following error:

```
test_mxfp4.py::Mxfp4ModelTest::test_save_mxfp4 FAILED                                                                                                              [100%]

================================================================================ FAILURES ================================================================================
_____________________________________________________________________ Mxfp4ModelTest.test_save_mxfp4 _____________________________________________________________________

self = <mxfp4.test_mxfp4.Mxfp4ModelTest testMethod=test_save_mxfp4>

    def test_save_mxfp4(self):
        """Test saving quantized OpenAI MoE model with device_map"""
    
        model = GptOssForCausalLM.from_pretrained(
            self.model_name,
            torch_dtype=torch.bfloat16,
            device_map="auto",
        )
        tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        with tempfile.TemporaryDirectory() as tmp:
            # Save the model in mxfp4 format
>           model.save_pretrained(tmp)

test_mxfp4.py:456: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = GptOssForCausalLM(
  (model): GptOssModel(
    (embed_tokens): Embedding(201088, 2880, padding_idx=199999)
    (layers...
    (rotary_emb): GptOssRotaryEmbedding()
  )
  (lm_head): Linear(in_features=2880, out_features=201088, bias=False)
)
save_directory = '/tmp/tmptut51kyp', is_main_process = True, state_dict = None, save_function = <function save at 0x7fef7bd020e0>, push_to_hub = False
max_shard_size = '5GB', safe_serialization = True, variant = None, token = None, save_peft_format = True, kwargs = {}, use_auth_token = None
ignore_metadata_errors = False, _hf_peft_config_loaded = False, hf_quantizer = <transformers.quantizers.quantizer_mxfp4.Mxfp4HfQuantizer object at 0x7fee27aafbe0>
quantization_serializable = True

    def save_pretrained(
        self,
        save_directory: Union[str, os.PathLike],
        is_main_process: bool = True,
        state_dict: Optional[dict] = None,
        save_function: Callable = torch.save,
        push_to_hub: bool = False,
        max_shard_size: Union[int, str] = "5GB",
        safe_serialization: bool = True,
        variant: Optional[str] = None,
        token: Optional[Union[str, bool]] = None,
        save_peft_format: bool = True,
        **kwargs,
    ):
        """
        Save a model and its configuration file to a directory, so that it can be re-loaded using the
        [`~PreTrainedModel.from_pretrained`] class method.
    
        Arguments:
            save_directory (`str` or `os.PathLike`):
                Directory to which to save. Will be created if it doesn't exist.
            is_main_process (`bool`, *optional*, defaults to `True`):
                Whether the process calling this is the main process or not. Useful when in distributed training like
                TPUs and need to call this function on all processes. In this case, set `is_main_process=True` only on
                the main process to avoid race conditions.
            state_dict (nested dictionary of `torch.Tensor`):
                The state dictionary of the model to save. Will default to `self.state_dict()`, but can be used to only
                save parts of the model or if special precautions need to be taken when recovering the state dictionary
                of a model (like when using model parallelism).
            save_function (`Callable`):
                The function to use to save the state dictionary. Useful on distributed training like TPUs when one
                need to replace `torch.save` by another method.
            push_to_hub (`bool`, *optional*, defaults to `False`):
                Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
                repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
                namespace).
            max_shard_size (`int` or `str`, *optional*, defaults to `"5GB"`):
                The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size
                lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5MB"`).
                We default it to 5GB in order for models to be able to run easily on free-tier google colab instances
                without CPU OOM issues.
    
                <Tip warning={true}>
    
                If a single weight of the model is bigger than `max_shard_size`, it will be in its own checkpoint shard
                which will be bigger than `max_shard_size`.
    
                </Tip>
    
            safe_serialization (`bool`, *optional*, defaults to `True`):
                Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
            variant (`str`, *optional*):
                If specified, weights are saved in the format pytorch_model.<variant>.bin.
            token (`str` or `bool`, *optional*):
                The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use
                the token generated when running `hf auth login` (stored in `~/.huggingface`).
            save_peft_format (`bool`, *optional*, defaults to `True`):
                For backward compatibility with PEFT library, in case adapter weights are attached to the model, all
                keys of the state dict of adapters needs to be prepended with `base_model.model`. Advanced users can
                disable this behaviours by setting `save_peft_format` to `False`.
            kwargs (`dict[str, Any]`, *optional*):
                Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
        """
        use_auth_token = kwargs.pop("use_auth_token", None)
        ignore_metadata_errors = kwargs.pop("ignore_metadata_errors", False)
    
        if use_auth_token is not None:
            warnings.warn(
                "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
                FutureWarning,
            )
            if token is not None:
                raise ValueError(
                    "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
                )
            token = use_auth_token
    
        if token is not None:
            kwargs["token"] = token
    
        _hf_peft_config_loaded = getattr(self, "_hf_peft_config_loaded", False)
    
        hf_quantizer = getattr(self, "hf_quantizer", None)
        quantization_serializable = (
            hf_quantizer is not None
            and isinstance(hf_quantizer, HfQuantizer)
            and hf_quantizer.is_serializable(safe_serialization=safe_serialization)
        )
    
        if hf_quantizer is not None and not _hf_peft_config_loaded and not quantization_serializable:
            raise ValueError(
                f"The model is quantized with {hf_quantizer.quantization_config.quant_method} and is not serializable - check out the warnings from"
                " the logger on the traceback to understand the reason why the quantized model is not serializable."
            )
    
        if "save_config" in kwargs:
            warnings.warn(
                "`save_config` is deprecated and will be removed in v5 of Transformers. Use `is_main_process` instead."
            )
            is_main_process = kwargs.pop("save_config")
        if safe_serialization and not is_safetensors_available():
            raise ImportError("`safe_serialization` requires the `safetensors library: `pip install safetensors`.")
    
        # we need to check against tp_size, not tp_plan, as tp_plan is substituted to the class one
        if self._tp_size is not None and not is_huggingface_hub_greater_or_equal("0.31.4"):
            raise ImportError(
                "Saving a model with tensor parallelism requires `huggingface_hub` version 0.31.4 or higher."
            )
    
        if os.path.isfile(save_directory):
            logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
            return
    
        os.makedirs(save_directory, exist_ok=True)
    
        if push_to_hub:
            commit_message = kwargs.pop("commit_message", None)
            repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
            create_pr = kwargs.pop("create_pr", False)
            repo_id = self._create_repo(repo_id, **kwargs)
            files_timestamps = self._get_files_timestamps(save_directory)
    
        metadata = {}
        if hf_quantizer is not None:
>           state_dict, metadata = hf_quantizer.get_state_dict_and_metadata(self, safe_serialization)
E           TypeError: Mxfp4HfQuantizer.get_state_dict_and_metadata() takes 2 positional arguments but 3 were given

../../../src/transformers/modeling_utils.py:3918: TypeError
-------------------------------------------------------------------------- Captured stderr call --------------------------------------------------------------------------
`torch_dtype` is deprecated! Use `dtype` instead!
Fetching 41 files: 100%|██████████| 41/41 [00:00<00:00, 11247.72it/s]
Fetching 41 files: 100%|██████████| 41/41 [00:00<00:00, 10037.15it/s]
Loading checkpoint shards: 100%|██████████| 3/3 [00:07<00:00,  2.38s/it]
======================================================================== short test summary info =========================================================================
FAILED test_mxfp4.py::Mxfp4ModelTest::test_save_mxfp4 - TypeError: Mxfp4HfQuantizer.get_state_dict_and_metadata() takes 2 positional arguments but 3 were given
```

Changed files:
filename: src/transformers/quantizers/quantizer_mxfp4.py
patch: @@ -379,7 +379,7 @@ def update_param_name(self, param_name: str) -> str:
                 return param_name.replace("down_proj", "down_proj_blocks")
         return param_name
 
-    def get_state_dict_and_metadata(self, model):
+    def get_state_dict_and_metadata(self, model, safe_serialization: bool = False):
         from ..integrations import Mxfp4GptOssExperts
 
         state_dict = model.state_dict()