diff --git a/docs/user_guide/QUANTIZATION.md b/docs/user_guide/QUANTIZATION.md index e0ddae8a..46e93386 100644 --- a/docs/user_guide/QUANTIZATION.md +++ b/docs/user_guide/QUANTIZATION.md @@ -16,6 +16,8 @@ Quantization is a powerful technique to reduce the memory footprint and computat |int8_per_tensor|quantize weights and activations to int8 (dynamic quantization) with tensorwise method.|>=sm80, Ampere or newer| |int8_weight_only|quantize only weights to int8, keep activations in full precision|>=sm80, Ampere or newer| |int4_weight_only|quantize only weights to int4, keep activations in full precision|>=sm90, Hopper or newer, TMA required| +|nvfp4|TorchAO dynamic activation and NVFP4 weight quantization. The current integration uses Triton activation quantization and dynamic per-tensor scaling.|>=sm100, Blackwell or newer| +|nvfp4_weight_only|TorchAO NVFP4 weight-only quantization with dynamic per-tensor weight scaling; activations remain in their original precision.|>=sm100, Blackwell or newer| |svdq_int4_r{32...}|post-training SVDQuant (W4A4) with calibration and checkpoint serialization for users who want the higher-accuracy PTQ workflow.|>=sm80, Ampere or newer, excluded Hopper (NO INT4 MMA)| |svdq_nvfp4_r{32...}|post-training SVDQuant (W4A4) with NVFP4 packed weights/activations, calibration, and checkpoint serialization. Only svdq_kwargs["runtime_kernel"]="v1" is currently supported.|>=sm120, Blackwell or newer| |svdq_int4_r{32...}_dq|quantize weights and activations to int4 with SVDQuant dynamic quantization (W4A4) without any calibration.|>=sm80, Ampere or newer, excluded Hopper (NO INT4 MMA)| @@ -289,13 +291,42 @@ INT4 quantization can provide even better memory reduction compared to FP8 or IN Please note that users should also install mslk kernel library to enable INT8/INT4 quantization features. The int4_weight_only w4a16 compute kennel requires architectures >= sm90 (Hopper or newer, TMA required). For older architectures, users can use int8_weight_only quantization for better compatibility. ```bash -# stable: mslk (change cu130 to cu129 if using CUDA 12.9), required torch>=2.11.0 -uv pip install torch==2.11.0 mslk --index-url https://download.pytorch.org/whl/cu130 --upgrade -# nightly: mslk (change cu130 to cu129 if using CUDA 12.9), required torch>=2.11.0 -uv pip install --pre torch mslk --index-url https://download.pytorch.org/whl/nightly/cu130 --upgrade +# stable: mslk (change cu13x to cu129 if using CUDA 12.9), torch>=2.11.0 +# mslk version matching: torch==2.11.x mslk==1.1.0; torch==2.12.x mslk==1.2.0 +pip install torch==2.11.0 mslk==1.1.0 --index-url https://download.pytorch.org/whl/cu130 +pip install torch==2.12.0 mslk==1.2.0 --index-url https://download.pytorch.org/whl/cu132 +# nightly: mslk (change cu13x to cu129 if using CUDA 12.9), required torch>=2.11.0 +pip install --pre torch mslk --index-url https://download.pytorch.org/whl/nightly/cu132 ``` + In the case of distributed inference (context parallelism or tensor parallelism), we recommend users to use float8 quantization to avoid potential compatibility issues. + +## TorchAO NVFP4 Quantization + +nvfp4 uses NVFP4DynamicActivationNVFP4WeightConfig from TorchAO's MX formats prototype. Cache-DiT currently fixes both use_triton_kernel=True and use_dynamic_per_tensor_scale=True. It is an online dynamic quantization path and is separate from the svdq_nvfp4_* SVDQuant path. (Required mslk kernels library) + +For example: + +```python +import cache_dit +from cache_dit import QuantizeConfig + +cache_dit.enable_cache( + pipe, cache_config=..., quantize_config=QuantizeConfig(quant_type="nvfp4"), +) +``` + +The current implementation requires Blackwell or newer GPUs (SM100+) and skips Linear layers whose two weight dimensions are not divisible by 16 or whose output dimensions are too small for the NVFP4 kernel. The supported CLI shortcut is: + +```bash +python3 -m cache_dit.generate flux --nvfp4 +python3 -m cache_dit.generate flux --nvfp4 --compile +python3 -m cache_dit.generate flux --nvfp4-weight-only +``` + +The model must use bfloat16 Linear weights. The Triton activation kernel can fall back internally when a runtime shape does not satisfy its kernel constraints. + ## SVDQuant (W4A4) PTQ Cache-DiT provides a native SVDQuant PTQ workflow for W4A4 quantization (with high performance W4A4 GEMM kernels and an easy-to-use PTQ interface). The public API is intentionally small: build a QuantizeConfig, quantize with cache_dit.quantize(...), then reload with cache_dit.load(...). Cache-DiT now supports both INT4 and NVFP4 SVDQuant PTQ flows. We highly recommend using native SVDQuant support in Cache-DiT for W4A4 quantization, as it can provide high performance and better usability compared to other third-party low-bit quantization libraries. @@ -824,7 +855,7 @@ Conversion complete. Load the quantized model with: **Experimental weight-only smooth strategy:** ```bash -# INT4 +# SVDQ INT4 cache-dit-convert \ --model-path black-forest-labs/FLUX.2-klein-4B \ --save-dir ./FLUX.2-klein-4B-svdq \ @@ -832,7 +863,7 @@ cache-dit-convert \ --svdq-smooth-strategy weight \ --svdq-calibrate-precision medium -# NVFP4 +# SVDQ NVFP4 cache-dit-convert \ --model-path black-forest-labs/FLUX.2-klein-4B \ --save-dir ./FLUX.2-klein-4B-svdq-nvfp4 \ @@ -844,7 +875,7 @@ cache-dit-convert \ **With v2 runtime kernel and verbose logging:** ```bash -# INT4 only, NVFP4 DQ is not supported with v2 runtime kernel for now. +# SVDQ INT4 only, SVDQ NVFP4 DQ is not supported with v2 runtime kernel for now. cache-dit-convert \ --model-path /path/to/FLUX.1-dev \ --save-dir ./FLUX.1-dev-svdq \ diff --git a/pyproject.toml b/pyproject.toml index afe2e45b..7d3b4713 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,7 @@ ray = [ ] quantization = [ - "torchao>=0.14.1", + "torchao>=0.17.0", "bitsandbytes>=0.48.1", "safetensors>=0.5.3", ] diff --git a/src/cache_dit/_utils/utils.py b/src/cache_dit/_utils/utils.py index 1501391c..130d1601 100644 --- a/src/cache_dit/_utils/utils.py +++ b/src/cache_dit/_utils/utils.py @@ -425,6 +425,8 @@ def get_args(parse: bool = True, ) -> argparse.ArgumentParser | argparse.Namespa "float8_per_tensor", "float8_per_block", "float8_weight_only", + "nvfp4", + "nvfp4_weight_only", "int8_per_row", "int8_per_tensor", "int8_weight_only", @@ -516,6 +518,18 @@ def get_args(parse: bool = True, ) -> argparse.ArgumentParser | argparse.Namespa default=False, help="Enable int4 weight-only quantization for transformer", ) + parser.add_argument( + "--nvfp4", + action="store_true", + default=False, + help="Enable TorchAO NVFP4 dynamic activation and weight quantization for transformer", + ) + parser.add_argument( + "--nvfp4-weight-only", + action="store_true", + default=False, + help="Enable TorchAO NVFP4 weight-only quantization for transformer", + ) parser.add_argument( "--svdq-int4-r32-dq", "--svdq-r32", @@ -1183,38 +1197,36 @@ def get_base_args(parse: bool = True) -> argparse.Namespace | argparse.ArgumentP def maybe_postprocess_args(args: argparse.Namespace) -> argparse.Namespace: # Force enable quantization if quantize_type is specified - if args.float8_per_row: - args.quantize_type = "float8_per_row" - elif args.float8_per_tensor: - args.quantize_type = "float8_per_tensor" - elif args.float8_per_block: - args.quantize_type = "float8_per_block" - elif args.float8_weight_only: - args.quantize_type = "float8_weight_only" - elif args.int8_per_row: - args.quantize_type = "int8_per_row" - elif args.int8_per_tensor: - args.quantize_type = "int8_per_tensor" - elif args.int8_weight_only: - args.quantize_type = "int8_weight_only" - elif args.int4_weight_only: - args.quantize_type = "int4_weight_only" - elif args.svdq_int4_r32_dq: - args.quantize_type = "svdq_int4_r32_dq" - elif args.svdq_int4_r64_dq: - args.quantize_type = "svdq_int4_r64_dq" - elif args.svdq_int4_r128_dq: - args.quantize_type = "svdq_int4_r128_dq" - elif args.svdq_int4_r256_dq: - args.quantize_type = "svdq_int4_r256_dq" - elif args.svdq_nvfp4_r32_dq: - args.quantize_type = "svdq_nvfp4_r32_dq" - elif args.svdq_nvfp4_r64_dq: - args.quantize_type = "svdq_nvfp4_r64_dq" - elif args.svdq_nvfp4_r128_dq: - args.quantize_type = "svdq_nvfp4_r128_dq" - elif args.svdq_nvfp4_r256_dq: - args.quantize_type = "svdq_nvfp4_r256_dq" + quantize_shortcuts = [ + ("float8_per_row", args.float8_per_row), + ("float8_per_tensor", args.float8_per_tensor), + ("float8_per_block", args.float8_per_block), + ("float8_weight_only", args.float8_weight_only), + ("int8_per_row", args.int8_per_row), + ("int8_per_tensor", args.int8_per_tensor), + ("int8_weight_only", args.int8_weight_only), + ("int4_weight_only", args.int4_weight_only), + ("nvfp4", args.nvfp4), + ("nvfp4_weight_only", args.nvfp4_weight_only), + ("svdq_int4_r32_dq", args.svdq_int4_r32_dq), + ("svdq_int4_r64_dq", args.svdq_int4_r64_dq), + ("svdq_int4_r128_dq", args.svdq_int4_r128_dq), + ("svdq_int4_r256_dq", args.svdq_int4_r256_dq), + ("svdq_nvfp4_r32_dq", args.svdq_nvfp4_r32_dq), + ("svdq_nvfp4_r64_dq", args.svdq_nvfp4_r64_dq), + ("svdq_nvfp4_r128_dq", args.svdq_nvfp4_r128_dq), + ("svdq_nvfp4_r256_dq", args.svdq_nvfp4_r256_dq), + ] + selected_shortcuts = [quant_type for quant_type, enabled in quantize_shortcuts if enabled] + if len(selected_shortcuts) > 1: + raise ValueError(f"Quantization shortcuts are mutually exclusive, got {selected_shortcuts}.") + if selected_shortcuts: + selected_quant_type = selected_shortcuts[0] + if args.quantize_type is not None and args.quantize_type != selected_quant_type: + raise ValueError( + f"--quantize-type {args.quantize_type} conflicts with --{selected_quant_type.replace('_', '-')}." + ) + args.quantize_type = selected_quant_type if args.quantize_type is not None: args.quantize = True diff --git a/src/cache_dit/quantization/config.py b/src/cache_dit/quantization/config.py index 54dc2baf..148aff8b 100644 --- a/src/cache_dit/quantization/config.py +++ b/src/cache_dit/quantization/config.py @@ -333,7 +333,7 @@ class QuantizeConfig: backend: str | QuantizeBackend = QuantizeBackend.AUTO # Quantization type, currently support "float8_weight_only" and "float8_per_row", # "float8_per_tensor", "float8_per_block", "int8_per_row", "int8_per_tensor", - # "int8_weight_only", "int4_weight_only", etc. + # "int8_weight_only", "int4_weight_only", "nvfp4", "nvfp4_weight_only", etc. quant_type: str = "float8_per_row" # The layers specified in this variable will be excluded from quantization, # even if they are in the repeated blocks or not filtered out by filter_fn. diff --git a/src/cache_dit/quantization/torchao/quantize_ao.py b/src/cache_dit/quantization/torchao/quantize_ao.py index 7a50b316..1f470c89 100644 --- a/src/cache_dit/quantization/torchao/quantize_ao.py +++ b/src/cache_dit/quantization/torchao/quantize_ao.py @@ -18,6 +18,8 @@ "float8_per_tensor", "float8_per_block", "float8_weight_only", + "nvfp4", + "nvfp4_weight_only", "int8_per_row", "int8_per_tensor", "int8_weight_only", @@ -29,6 +31,8 @@ "float8_per_row", "float8_per_block", "float8_weight_only", + "nvfp4", + "nvfp4_weight_only", "int8_per_tensor", "int8_per_row", "int8_weight_only", @@ -198,6 +202,12 @@ def __post_init__(self): 9, ), "FP8 requires Ada or newer GPUs (>=sm89), but got " + str( current_platform.get_device_capability()) + if self.is_nvfp4(): + assert current_platform.get_device_capability() >= ( + 10, + 0, + ), "NVFP4 requires Blackwell or newer GPUs (>=sm100), but got " + str( + current_platform.get_device_capability()) @staticmethod def from_config( @@ -411,6 +421,12 @@ def is_float8_per_block(self) -> bool: def is_float8_weight_only(self) -> bool: return self.quant_type == "float8_weight_only" + def is_nvfp4(self) -> bool: + return self.quant_type in ("nvfp4", "nvfp4_weight_only") + + def is_nvfp4_weight_only(self) -> bool: + return self.quant_type == "nvfp4_weight_only" + def required_fallback(self) -> bool: # Currently, only support float8 per-tensor fallback for rowwise layers if # regional quantiztion is enabled. Not support fallback for int8/int4/weight-only @@ -485,6 +501,7 @@ def _get_torchao_config(quant_type: str, **kwargs) -> AOBaseConfig: Float8DynamicActivationFloat8WeightConfig, PerRow, ) + from torchao.quantization.quantize_.common import KernelPreference quant_config = Float8DynamicActivationFloat8WeightConfig( weight_dtype=kwargs.get( @@ -496,6 +513,7 @@ def _get_torchao_config(quant_type: str, **kwargs) -> AOBaseConfig: torch.float8_e4m3fn, ), granularity=(PerRow(), PerRow()), + kernel_preference=KernelPreference.TORCH, ) elif quant_type == "float8_per_tensor": from torchao.quantization import ( @@ -550,6 +568,20 @@ def _get_torchao_config(quant_type: str, **kwargs) -> AOBaseConfig: torch.float8_e4m3fn, ), ) + elif quant_type == "nvfp4": + from torchao.prototype.mx_formats.inference_workflow import ( + NVFP4DynamicActivationNVFP4WeightConfig, ) + + quant_config = NVFP4DynamicActivationNVFP4WeightConfig( + use_triton_kernel=True, + use_dynamic_per_tensor_scale=True, + ) + + elif quant_type == "nvfp4_weight_only": + from torchao.prototype.mx_formats.inference_workflow import NVFP4WeightOnlyConfig + + quant_config = NVFP4WeightOnlyConfig(use_dynamic_per_tensor_scale=True) + elif quant_type == "int8_per_row": from torchao.quantization import ( Int8DynamicActivationInt8WeightConfig, @@ -658,6 +690,20 @@ def _is_curr_plan_allow_to_quantize(name: str) -> bool: # precision plan logger.debug(skip_reason) return False + if curr_quant_type in ("nvfp4", "nvfp4_weight_only"): + weight_shape = tuple(m.weight.shape) + if weight_shape[-2] % 16 != 0 or weight_shape[-1] % 16 != 0: + skip_reason = _skip_reason(f"weight_shape{weight_shape}%16!=0") + quant_ctx.skipped_map[curr_quant_type].append(skip_reason) + logger.debug(skip_reason) + return False + if (curr_quant_type == "nvfp4" + and (weight_shape[-2] <= 64 or (weight_shape[-1] <= 1024 and weight_shape[-2] <= 1024))): + skip_reason = _skip_reason(f"weight_shape{weight_shape} is too small") + quant_ctx.skipped_map[curr_quant_type].append(skip_reason) + logger.debug(skip_reason) + return False + # check blockwise fp8 support for linear layers, if not supported, # skip quantization for that layer. if curr_quant_type in [ diff --git a/tests/quantization/test_torchao_nvfp4.py b/tests/quantization/test_torchao_nvfp4.py new file mode 100644 index 00000000..bdadac52 --- /dev/null +++ b/tests/quantization/test_torchao_nvfp4.py @@ -0,0 +1,133 @@ +import importlib.util + +import pytest +import torch + +from cache_dit._utils.utils import get_args, maybe_postprocess_args +from cache_dit.quantization import QuantizeConfig, quantize +from cache_dit.quantization.torchao.quantize_ao import _get_torchao_config + +torchao = pytest.importorskip("torchao") +from torchao.quantization.quantize_.common import KernelPreference + +nvfp4_workflow = pytest.importorskip("torchao.prototype.mx_formats.inference_workflow") +NVFP4Tensor = pytest.importorskip("torchao.prototype.mx_formats.nvfp4_tensor").NVFP4Tensor + + +def test_nvfp4_config_and_cli() -> None: + float8_config = _get_torchao_config("float8_per_row") + assert float8_config.kernel_preference is KernelPreference.TORCH + + config = QuantizeConfig(quant_type="nvfp4") + assert config.strify() == "nvfp4" + assert config.backend.value == "TORCHAO" + + torchao_config = _get_torchao_config("nvfp4") + assert isinstance(torchao_config, nvfp4_workflow.NVFP4DynamicActivationNVFP4WeightConfig) + assert torchao_config.use_triton_kernel is True + assert torchao_config.use_dynamic_per_tensor_scale is True + + weight_only_config = QuantizeConfig(quant_type="nvfp4_weight_only") + assert weight_only_config.strify() == "nvfp4_weight_only" + weight_only_torchao_config = _get_torchao_config("nvfp4_weight_only") + assert isinstance(weight_only_torchao_config, nvfp4_workflow.NVFP4WeightOnlyConfig) + assert weight_only_torchao_config.use_dynamic_per_tensor_scale is True + + parser = get_args(parse=False) + args = maybe_postprocess_args(parser.parse_args(["--nvfp4"])) + assert args.quantize is True + assert args.quantize_type == "nvfp4" + + args = maybe_postprocess_args(parser.parse_args(["--quantize-type", "nvfp4"])) + assert args.quantize is True + assert args.quantize_type == "nvfp4" + + args = maybe_postprocess_args(parser.parse_args(["--nvfp4-weight-only"])) + assert args.quantize is True + assert args.quantize_type == "nvfp4_weight_only" + + with pytest.raises(ValueError, match="mutually exclusive"): + maybe_postprocess_args(parser.parse_args(["--nvfp4", "--svdq-nvfp4-r128-dq"])) + + +def _require_nvfp4_device() -> None: + if not torch.cuda.is_available(): + pytest.skip("TorchAO NVFP4 tests require CUDA.") + if torch.cuda.get_device_capability() < (10, 0): + pytest.skip("TorchAO dynamic NVFP4 requires SM100 or newer.") + if importlib.util.find_spec("mslk") is None: + pytest.skip("TorchAO dynamic NVFP4 tests require MSLK.") + + +def test_nvfp4_quantizes_supported_linears_and_skips_unsupported_linears() -> None: + _require_nvfp4_device() + + class ToyModel(torch.nn.Module): + + def __init__(self) -> None: + super().__init__() + self.large = torch.nn.Linear(2048, 2048, bias=False, dtype=torch.bfloat16) + self.small = torch.nn.Linear(1024, 64, bias=False, dtype=torch.bfloat16) + self.unaligned = torch.nn.Linear(1025, 2048, bias=False, dtype=torch.bfloat16) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + return self.large(input) + + model = ToyModel().cuda() + quantize( + model, + QuantizeConfig(quant_type="nvfp4", regional_quantize=False), + ) + + assert isinstance(model.large.weight, NVFP4Tensor) + assert not isinstance(model.small.weight, NVFP4Tensor) + assert not isinstance(model.unaligned.weight, NVFP4Tensor) + assert model.large.weight.act_quant_kwargs.use_triton_kernel is True + assert model.large.weight.act_quant_kwargs.use_dynamic_per_tensor_scale is True + + output = model(torch.randn(128, 2048, device="cuda", dtype=torch.bfloat16)) + assert output.shape == (128, 2048) + assert torch.isfinite(output).all() + + +def test_nvfp4_compiles_and_runs() -> None: + _require_nvfp4_device() + + model = torch.nn.Linear(2048, 2048, bias=False, dtype=torch.bfloat16, device="cuda") + quantize( + model, + QuantizeConfig(quant_type="nvfp4", regional_quantize=False), + ) + compiled_model = torch.compile(model, fullgraph=True) + output = compiled_model(torch.randn(128, 2048, device="cuda", dtype=torch.bfloat16)) + assert output.shape == (128, 2048) + assert torch.isfinite(output).all() + + +def test_nvfp4_weight_only_quantizes_and_compiles() -> None: + _require_nvfp4_device() + + class ToyModel(torch.nn.Module): + + def __init__(self) -> None: + super().__init__() + self.large = torch.nn.Linear(2048, 2048, bias=False, dtype=torch.bfloat16) + self.unaligned = torch.nn.Linear(1025, 2048, bias=False, dtype=torch.bfloat16) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + return self.large(input) + + model = ToyModel().cuda() + quantize( + model, + QuantizeConfig(quant_type="nvfp4_weight_only", regional_quantize=False), + ) + + assert isinstance(model.large.weight, NVFP4Tensor) + assert model.large.weight.act_quant_kwargs is None + assert not isinstance(model.unaligned.weight, NVFP4Tensor) + + compiled_model = torch.compile(model, fullgraph=True) + output = compiled_model(torch.randn(128, 2048, device="cuda", dtype=torch.bfloat16)) + assert output.shape == (128, 2048) + assert torch.isfinite(output).all()