FP-Quant support

docker/transformers-quantization-latest-gpu/Dockerfile

@@ -78,6 +78,9 @@ RUN git clone https://github.com/NetEase-FuXi/EETQ.git && cd EETQ/ && git submod
 # RUN python3 -m pip install --no-cache-dir flute-kernel==0.4.1
 # RUN python3 -m pip install --no-cache-dir git+https://github.com/Dao-AILab/fast-hadamard-transform.git
 
+# Add fp-quant for quantization testing
+RUN python3 -m pip install --no-cache-dir "fp-quant>=0.1.6"
+
 # Add compressed-tensors for quantization testing
 RUN python3 -m pip install --no-cache-dir compressed-tensors
 

docs/source/en/_toctree.yml

@@ -179,6 +179,8 @@
     title: FBGEMM
   - local: quantization/finegrained_fp8
     title: Fine-grained FP8
+  - local: quantization/fp_quant
+    title: FP-Quant
   - local: gguf
     title: GGUF
   - local: quantization/gptq

docs/source/en/main_classes/quantization.md

@@ -93,6 +93,10 @@ Learn how to quantize models in the [Quantization](../quantization) guide.
 
 [[autodoc]] QuarkConfig
 
+## FPQuantConfig
+
+[[autodoc]] FPQuantConfig
+
 ## AutoRoundConfig
 
 [[autodoc]] AutoRoundConfig

docs/source/en/quantization/fp_quant.md

@@ -0,0 +1,66 @@
+<!--Copyright 2025 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+
+-->
+
+# FP-Quant
+
+[FP-Quant](https://github.com/IST-DASLab/FP-Quant) is a family of quantization algorithms tailored for the Blackwell generation of Nvidia GPUs. The goal is to allow for efficient post-training quantization (PTQ) and quantization-aware trainin (QAT) of LLMs in the [MXFP4 and NVFP4 data-types](https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf).
+
+Currently, only PTQ with MXFP4 is supported. Models can either be quantized on the fly with `quantization_config=FPQuantConfig()`:
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer, FPQuantConfig
+import torch
+
+model = AutoModelForCausalLM.from_pretrained(
+    "qwen/Qwen3-8B",
+    quantization_config=FPQuantConfig(),
+    device_map="cuda",
+    torch_dtype=torch.bfloat16,
+)
+```
+
+or pre-processed with GPTQ for better quality (see [FP Format Quantization Harness](https://github.com/IST-DASLab/FP-Quant)).
+
+A **Blackwell-generation GPU is required** to run the kernels. Runtime support for FP-Quant is implemented through the [QuTLASS](https://github.com/IST-DASLab/qutlass) library and a lightweight PyTorch interface lib [`fp_quant`](https://github.com/IST-DASLab/FP-Quant/tree/master/inference_lib). We recommend installing the former **from source** and the latter with  `pip install fp_quant`.
+
+Users **without a Blackwell-generation GPU** , can use the method with `quantization_config=FPQuantConfig(pseudoquant=True)` without having to install [QuTLASS](https://github.com/IST-DASLab/qutlass). This would provide no speedups but would fully emulate the effect of quantization.
+
+> [!TIP]
+> Find models pre-quantized with FP-Quant in the official ISTA-DASLab [collection](https://huggingface.co/collections/ISTA-DASLab/fp-quant-6877c186103a21d3a02568ee).
+
+## torch.compile
+
+FP-Quant is fully compatible with [torch.compile](https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html).
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, FPQuantConfig
+
+model = AutoModelForCausalLM.from_pretrained(
+    "qwen/Qwen3-8B",
+    quantization_config=FPQuantConfig(),
+    device_map="cuda",
+    torch_dtype=torch.bfloat16,
+)
+
+model.forward = torch.compile(model.forward, mode="max-autotune", fullgraph=True)
+```
+
+## Speedups
+
+FP-Quant currently performs best for very large batch size processing.
+
+See [QuTLASS README](https://github.com/IST-DASLab/qutlass/blob/main/README.md) for speedups.

docs/source/en/quantization/overview.md

@@ -30,6 +30,7 @@ Use the Space below to help you pick a quantization method depending on your har
 | [bitsandbytes](./bitsandbytes)            | 🟢                   | 🟡 |     🟢     | 🟡 | 🔴                    | 🟡 | 🟢 | 4/8          | 🟢               | 🟢                          | 🟢                      | https://github.com/bitsandbytes-foundation/bitsandbytes |
 | [compressed-tensors](./compressed_tensors) | 🔴                   | 🟢              |     🟢     | 🟢        | 🔴                                 | 🔴              | 🔴              | 1/8          | 🟢               | 🟢                          | 🟢                      | https://github.com/neuralmagic/compressed-tensors |
 | [EETQ](./eetq)                            | 🟢                   | 🔴              | 🟢        | 🔴        | 🔴                                 | 🔴              | ?               | 8            | 🟢               | 🟢                          | 🟢                      | https://github.com/NetEase-FuXi/EETQ        |
+| [FP-Quant](./fp_quant)                          | 🟢                   | 🔴              | 🟢        | 🔴        | 🔴                                 | 🔴              | 🟢              | 4           | 🔴               | 🟢                          | 🟢                      | https://github.com/IST-DASLab/FP-Quant      |
 | [GGUF / GGML (llama.cpp)](../gguf)        | 🟢                   | 🟢              | 🟢        | 🔴        | 🟢                                 | 🔴              | 🔴              | 1/8          | 🔴               | [See Notes](../gguf)     | [See Notes](../gguf) | https://github.com/ggerganov/llama.cpp      |
 | [GPTQModel](./gptq)                       | 🔴                   | 🟢 | 🟢        | 🟢        | 🟢                                 | 🟢 | 🔴              | 2/3/4/8      | 🟢               | 🟢                          | 🟢                      | https://github.com/ModelCloud/GPTQModel        |
 | [AutoGPTQ](./gptq)                        | 🔴                   | 🔴              | 🟢        | 🟢        | 🔴                                 | 🔴              | 🔴              | 2/3/4/8      | 🟢               | 🟢                          | 🟢                      | https://github.com/AutoGPTQ/AutoGPTQ        |

src/transformers/__init__.py

@@ -275,6 +275,7 @@
         "HqqConfig",
         "QuantoConfig",
         "QuarkConfig",
+        "FPQuantConfig",
         "SpQRConfig",
         "TorchAoConfig",
         "VptqConfig",
@@ -961,6 +962,7 @@
         EetqConfig,
         FbgemmFp8Config,
         FineGrainedFP8Config,
+        FPQuantConfig,
         GPTQConfig,
         HiggsConfig,
         HqqConfig,

src/transformers/integrations/fp_quant.py

@@ -0,0 +1,47 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"FP-Quant integration file"
+
+from ..utils import (
+    is_fp_quant_available,
+)
+
+
+if is_fp_quant_available():
+    from fp_quant import FPQuantConfig as FPQuantLinearConfig
+    from fp_quant import FPQuantDtype
+
+from transformers.utils.quantization_config import FPQuantConfig
+
+
+def adapt_fp_quant_config(config: FPQuantConfig):
+    if config.forward_dtype == "mxfp4":
+        forward_dtype = FPQuantDtype.MXFP4
+    else:
+        raise ValueError(f"Unsupported forward dtype: {config.forward_dtype}")
+
+    if config.backward_dtype == "bf16":
+        backward_dtype = FPQuantDtype.BF16
+    else:
+        raise ValueError(f"Unsupported backward dtype: {config.backward_dtype}")
+
+    return FPQuantLinearConfig(
+        forward_dtype=forward_dtype,
+        forward_method=config.forward_method,
+        backward_dtype=backward_dtype,
+        store_master_weights=config.store_master_weights,
+        hadamard_group_size=config.hadamard_group_size,
+        pseudoquantization=config.pseudoquantization,
+        modules_to_not_convert=config.modules_to_not_convert,
+    )

src/transformers/quantizers/auto.py

@@ -27,6 +27,7 @@
     EetqConfig,
     FbgemmFp8Config,
     FineGrainedFP8Config,
+    FPQuantConfig,
     GPTQConfig,
     HiggsConfig,
     HqqConfig,
@@ -49,6 +50,7 @@
 from .quantizer_eetq import EetqHfQuantizer
 from .quantizer_fbgemm_fp8 import FbgemmFp8HfQuantizer
 from .quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer
+from .quantizer_fp_quant import FPQuantHfQuantizer
 from .quantizer_gptq import GptqHfQuantizer
 from .quantizer_higgs import HiggsHfQuantizer
 from .quantizer_hqq import HqqHfQuantizer
@@ -67,6 +69,7 @@
     "aqlm": AqlmHfQuantizer,
     "quanto": QuantoHfQuantizer,
     "quark": QuarkHfQuantizer,
+    "fp_quant": FPQuantHfQuantizer,
     "eetq": EetqHfQuantizer,
     "higgs": HiggsHfQuantizer,
     "hqq": HqqHfQuantizer,
@@ -89,6 +92,7 @@
     "aqlm": AqlmConfig,
     "quanto": QuantoConfig,
     "quark": QuarkConfig,
+    "fp_quant": FPQuantConfig,
     "hqq": HqqConfig,
     "compressed-tensors": CompressedTensorsConfig,
     "fbgemm_fp8": FbgemmFp8Config,

src/transformers/quantizers/quantizer_fp_quant.py

@@ -0,0 +1,183 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING, Any, Optional
+
+from .base import HfQuantizer
+from .quantizers_utils import get_module_from_name
+
+
+if TYPE_CHECKING:
+    from ..modeling_utils import PreTrainedModel
+
+from ..utils import is_fp_quant_available, is_qutlass_available, is_torch_available, logging
+from ..utils.quantization_config import QuantizationConfigMixin
+
+
+if is_torch_available():
+    import torch
+
+logger = logging.get_logger(__name__)
+
+
+class FPQuantHfQuantizer(HfQuantizer):
+    """
+    Quantizer for the FP-Quant method. Enables the loading of prequantized models and in-flight quantization of full-precision models.
+    """
+
+    requires_calibration = False
+    requires_parameters_quantization = True
+    is_qat_trainable = False
+    required_packages = ["fp_quant"]
+
+    def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs):
+        super().__init__(quantization_config, **kwargs)
+        self.quantization_config = quantization_config
+
+    def validate_environment(self, device_map, **kwargs):
+        if not torch.cuda.is_available():
+            raise NotImplementedError(
+                "FPQuant quantization is only supported on GPU. Please use a different quantizer."
+            )
+
+        if not is_qutlass_available() and not self.quantization_config.pseudoquantization:
+            raise ImportError(
+                "Using `fp_quant` with real quantization requires a **Blackwell GPU** and qutlass: `git clone https://github.com/IST-DASLab/qutlass.git && cd qutlass && pip install --no-build-isolation .`. You can use `FPQuantConfig(pseudoquantization=True, ...)` to use Triton-based pseudo-quantization. It doesn't provide any speedups but emulates the quantization behavior of the real quantization."
+            )
+
+        if self.quantization_config.pseudoquantization:
+            logger.warning(
+                "Using pseudo-quantization for FP-Quant. This doesn't provide any speedups but emulates the quantization behavior of the real quantization."
+            )
+
+        if not is_fp_quant_available():
+            raise ImportError("Using `fp_quant` quantization requires fp_quant: `pip install fp_quant`")
+
+        if device_map is None:
+            raise ValueError(
+                "You are attempting to load a FPQuant model without setting device_map."
+                " Please set device_map comprised of 'cuda' devices."
+            )
+        elif isinstance(device_map, dict) and ("cpu" in device_map.values() or "disk" in device_map.values()):
+            raise ValueError(
+                "You are attempting to load a FPQuant model with a device_map that contains a CPU or disk device."
+                " This is not supported. Please remove the CPU or disk device from the device_map."
+            )
+
+    def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
+        if torch_dtype is None:
+            logger.info("`torch_dtype` is None. Setting `torch_dtype=torch.bfloat16` for qutlass compatibility.")
+            torch_dtype = torch.bfloat16
+        elif torch_dtype != torch.bfloat16:
+            raise ValueError(
+                f"Invalid `torch_dtype` {torch_dtype}. fp_quant quantization only supports `torch_dtype=torch.bfloat16`."
+            )
+
+        return torch_dtype
+
+    def create_quantized_param(
+        self,
+        model: "PreTrainedModel",
+        param_value: "torch.Tensor",
+        param_name: str,
+        target_device: "torch.device",
+        state_dict: dict[str, Any],
+        unexpected_keys: Optional[list[str]] = None,
+    ):
+        module, _ = get_module_from_name(model, param_name)
+
+        # The module holds either:
+        #  * `weight` when `store_master_weights=True`
+        #  * `qweight` and `scales` when `store_master_weights=False` and `pseudoquantization=False`
+        #  * `dqweight` when `store_master_weights=False` and `pseudoquantization=True`
+
+        if param_name.endswith(".qweight"):
+            # Loading a real quantized checkpoint without master weights
+            module.qweight = torch.nn.Parameter(
+                param_value.to(target_device),
+                requires_grad=False,
+            )
+            module.weight = None
+            module.dqweight = None
+            return
+
+        if param_name.endswith(".dqweight"):
+            # Loading a pseudo-quantized checkpoint without master weights
+            module.dqweight = torch.nn.Parameter(param_value.to(target_device))
+            module.weight = None
+            module.qweight = None
+            module.scales = None
+            return
+
+        # Loading master weights or an unquantized checkpoint
+        module.weight = torch.nn.Parameter(param_value.to(target_device))
+        # Let pre-forward handle the quantization and set None where necessary
+        module.pre_forward()
+
+        if unexpected_keys is not None and param_name in unexpected_keys:
+            unexpected_keys.remove(param_name)
+
+    def _process_model_before_weight_loading(
+        self,
+        model: "PreTrainedModel",
+        **kwargs,
+    ):
+        from fp_quant import replace_with_fp_quant_linear
+
+        from ..integrations.fp_quant import adapt_fp_quant_config
+
+        replace_with_fp_quant_linear(
+            model,
+            fp_quant_linear_config=adapt_fp_quant_config(self.quantization_config),
+        )
+        model.config.quantization_config = self.quantization_config
+
+    def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
+        return model
+
+    def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> list[str]:
+        from fp_quant import FPQuantLinear
+
+        fp_quant_names = {name for name, module in model.named_modules() if isinstance(module, FPQuantLinear)}
+
+        def should_exclude(key: str) -> bool:
+            if key.endswith(".weight") or key.endswith(".bias"):
+                return False
+            full_key = f"{prefix}.{key}"
+            return any(name in key or name in full_key for name in fp_quant_names)
+
+        return [key for key in missing_keys if not should_exclude(key)]
+
+    @property
+    def is_trainable(self, model: Optional["PreTrainedModel"] = None):
+        return False
+
+    def is_serializable(self, safe_serialization=None):
+        return True
+
+    def check_quantized_param(
+        self,
+        model: "PreTrainedModel",
+        param_value: "torch.Tensor",
+        param_name: str,
+        state_dict: dict[str, Any],
+        **kwargs,
+    ) -> bool:
+        from fp_quant import FPQuantLinear
+
+        module, tensor_name = get_module_from_name(model, param_name)
+        if isinstance(module, FPQuantLinear) and tensor_name in ["weight", "qweight", "dqweight"]:
+            # Only quantize weights of FPQuantLinear modules that are not already quantized
+            return True
+        else:
+            return False