diff --git a/docker/transformers-quantization-latest-gpu/Dockerfile b/docker/transformers-quantization-latest-gpu/Dockerfile index 18f3927da16c..9a63b7690524 100755 --- a/docker/transformers-quantization-latest-gpu/Dockerfile +++ b/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 diff --git a/src/transformers/integrations/fp_quant.py b/src/transformers/integrations/fp_quant.py index 0ac441e36f93..ccf933796165 100644 --- a/src/transformers/integrations/fp_quant.py +++ b/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}") diff --git a/src/transformers/utils/import_utils.py b/src/transformers/utils/import_utils.py index 7d939ed34e6f..2a740f4939bd 100644 --- a/src/transformers/utils/import_utils.py +++ b/src/transformers/utils/import_utils.py @@ -971,13 +971,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 diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index 12b4da8927d9..8ce450a60b9b 100644 --- a/src/transformers/utils/quantization_config.py +++ b/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.") diff --git a/tests/quantization/fp_quant_integration/test_fp_quant.py b/tests/quantization/fp_quant_integration/test_fp_quant.py index ce73f0d80a63..6ae0e7f23deb 100644 --- a/tests/quantization/fp_quant_integration/test_fp_quant.py +++ b/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