[FPQuant] MXFP8 and MXFP4 backwards support (#41897)

* FP-Quant backwards

* fp-quant v0.3.0 docker

* availability version bump

* fp_quant==0.3.1

* fp_quant v0.3.2

Author: BlackSamorez

docker/transformers-quantization-latest-gpu/Dockerfile

@@ -81,7 +81,7 @@ RUN python3 -m pip uninstall -y flash-attn
 RUN cd transformers && python3 setup.py develop
 
 # Add fp-quant for quantization testing
-RUN python3 -m pip install --no-cache-dir "fp-quant>=0.2.0"
+RUN python3 -m pip install --no-cache-dir "fp-quant>=0.3.2"
 
 # Low usage or incompatible lib, will enable later on
 

src/transformers/integrations/fp_quant.py

@@ -35,6 +35,10 @@ def adapt_fp_quant_config(config: FPQuantConfig):
 
     if config.backward_dtype == "bf16":
         backward_dtype = FPQuantDtype.BF16
+    elif config.backward_dtype == "mxfp8":
+        backward_dtype = FPQuantDtype.MXFP8
+    elif config.backward_dtype == "mxfp4":
+        backward_dtype = FPQuantDtype.MXFP4
     else:
         raise ValueError(f"Unsupported backward dtype: {config.backward_dtype}")
 

src/transformers/utils/import_utils.py

@@ -973,13 +973,13 @@ def is_quark_available() -> bool:
 @lru_cache
 def is_fp_quant_available():
     is_available, fp_quant_version = _is_package_available("fp_quant", return_version=True)
-    return is_available and version.parse(fp_quant_version) >= version.parse("0.2.0")
+    return is_available and version.parse(fp_quant_version) >= version.parse("0.3.2")
 
 
 @lru_cache
 def is_qutlass_available():
     is_available, qutlass_version = _is_package_available("qutlass", return_version=True)
-    return is_available and version.parse(qutlass_version) >= version.parse("0.1.0")
+    return is_available and version.parse(qutlass_version) >= version.parse("0.2.0")
 
 
 @lru_cache

src/transformers/utils/quantization_config.py

@@ -1601,8 +1601,12 @@ def post_init(self):
         else:
             raise ValueError("Only 'mxfp4' and 'nvfp4' are supported for forward_dtype for now.")
 
-        if self.backward_dtype != "bf16":
-            raise ValueError("Only 'bf16' is supported for backward_dtype for now.")
+        if self.backward_dtype not in ["bf16", "mxfp8", "mxfp4"]:
+            raise ValueError("Only 'bf16', 'mxfp8' and 'mxfp4' are supported for backward_dtype for now.")
+
+        if self.backward_dtype != "bf16" and self.forward_dtype != "mxfp4":
+            raise ValueError("Only 'mxfp4' forward is compatible with non-bf16 backwards for now.")
+
         if self.transform_init not in ["hadamard", "identity", "gsr"]:
             raise ValueError("Only 'hadamard', 'identity' and 'gsr' are supported for transform_init.")
 

tests/quantization/fp_quant_integration/test_fp_quant.py

@@ -163,6 +163,13 @@ def getQuantizationConfig(cls):
         return FPQuantConfig(forward_dtype="mxfp4", pseudoquantization=False)
 
 
+@require_qutlass
+class FPQuantNVFP4Test(FPQuantBaseTest):
+    @classmethod
+    def getQuantizationConfig(cls):
+        return FPQuantConfig(forward_dtype="nvfp4", pseudoquantization=False)
+
+
 @require_qutlass
 class FPQuantMXFP4GS128Test(FPQuantBaseTest):
     @classmethod