diff --git a/auto_round/compressors/model_free.py b/auto_round/compressors/model_free.py index b6ad9839bf..5c4c7b05a9 100644 --- a/auto_round/compressors/model_free.py +++ b/auto_round/compressors/model_free.py @@ -40,7 +40,12 @@ * Preset names: ``MXFP4``, ``MXFP8``. * ``data_type="mx_fp"``, ``group_size=32``, ``bits in {4, 8}``. -Schemes that require special packing (FP8, NVFP4, GGUF, INT8_W8A8, +**NVFP4 E5M3** (saved in fake format): + +* Preset name: ``NVFP4_E5M3``. +* ``data_type="fp4_v2"``, ``group_size=16``, with high-precision QDQ weights. + +Schemes that require special packing (FP8, standard NVFP4, GGUF, INT8_W8A8, BF16, FPW8A16, ...) are **not** supported in model-free mode and will raise ``ValueError``. Use the standard AutoRound flow for those. @@ -49,6 +54,9 @@ * **INT schemes** → ``auto_round:auto_gptq`` packing format, ``quant_method="auto-round"``. * **MXFP schemes** → ``mxfp4-pack-quantized`` or ``mxfp8-quantized`` format, ``quant_method="compressed-tensors"``, compatible with vLLM / llm-compressor. +* **NVFP4_E5M3** → AutoRound format with packed ``.weight_packed`` and + ``.weight_scale`` tensors; use ``format="fake"`` explicitly for high-precision + QDQ ``.weight`` tensors. Usage (CLI) ----------- @@ -97,15 +105,16 @@ import shutil import sys import time -from concurrent.futures import ProcessPoolExecutor, as_completed +from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, ThreadPoolExecutor, as_completed, wait from dataclasses import asdict, fields +from functools import lru_cache from typing import Any, Callable, Optional, Union import torch from auto_round import envs from auto_round.compressors.config_resolution import thaw_mapping -from auto_round.compressors.utils import is_mx_fp +from auto_round.compressors.utils import is_mx_fp, is_nv_fp from auto_round.logger import logger from auto_round.schemes import PRESET_SCHEMES, QuantizationScheme, preset_name_to_scheme from auto_round.utils.common import AUDIO_MM_KEYS, VISION_MM_KEYS, compress_layer_names, to_standard_regex @@ -138,6 +147,7 @@ "W8A16", "MXFP4", "MXFP8", + "NVFP4_E5M3", "BF16", ) @@ -148,6 +158,8 @@ # Allowed ``bits`` values for MXFP weight quantization. _SUPPORTED_MXFP_BITS: tuple[int, ...] = (4, 8) +_NVFP4_E5M3_DATA_TYPE = "fp4_v2" + # Multimodal keywords kept in full precision by default. _NONTEXT_KEYWORDS: tuple[str, ...] = VISION_MM_KEYS + AUDIO_MM_KEYS @@ -565,6 +577,62 @@ def _quantize_weight_mxfp( } +def _quantize_weight_nvfp4_e5m3( + weight: torch.Tensor, + layer_name: str, + group_size: int = 16, + device: str = "cpu", +) -> dict[str, torch.Tensor]: + """Fake-quantize a 2D weight tensor to NVFP4 E5M3 and return its high-precision QDQ weight.""" + from auto_round.data_type.nvfp import fp4_v2 + + out_features, in_features = weight.shape + if group_size != 16: + raise ValueError(f"NVFP4_E5M3 requires group_size=16, got {group_size} for layer '{layer_name}'.") + if in_features % group_size != 0: + raise ValueError( + f"in_features={in_features} for layer '{layer_name}' is not divisible " + f"by NVFP4_E5M3 group_size={group_size}; cannot quantize." + ) + + weight_dev = weight.to(device) + qdq_weight, _, _ = fp4_v2(weight_dev, bits=4, group_size=group_size) + return {f"{layer_name}.weight": qdq_weight.to(dtype=weight.dtype, device="cpu")} + + +def _pack_weight_nvfp4_e5m3( + weight: torch.Tensor, + layer_name: str, + group_size: int = 16, + device: str = "cpu", +) -> dict[str, torch.Tensor]: + """Pack FP4 E2M1 weights with unsigned E5M3 block scales.""" + from auto_round.data_type.nvfp import fp4_v2 + from auto_round.export.export_to_autoround.qlinear_fp import QuantLinear + + out_features, in_features = weight.shape + if group_size != 16 or in_features % group_size != 0: + raise ValueError( + f"NVFP4_E5M3 requires in_features divisible by group_size=16, got {in_features} for '{layer_name}'." + ) + weight_dev = weight.to(device) + _, scale, _ = fp4_v2(weight_dev, bits=4, group_size=group_size) + # fp4_v2 may return a flattened per-group scale layout (e.g. [N, 1]); + # normalize to [out_features, in_features // group_size] before packing + # so serialized .weight_scale keeps the expected 2D shape. + scale = scale.reshape(out_features, in_features // group_size).to(torch.float32) + linear = torch.nn.Linear(in_features, out_features, bias=False, device=device, dtype=weight.dtype) + linear.weight = torch.nn.Parameter(weight_dev, requires_grad=False) + qlayer = QuantLinear( + 4, group_size, in_features, out_features, False, data_type="fp4_v2", act_bits=4, act_data_type="fp4_v2" + ) + qlayer.pack(linear, scale, device=device) + return { + f"{layer_name}.weight_packed": qlayer.weight_packed.to("cpu"), + f"{layer_name}.weight_scale": qlayer.weight_scale.to("cpu"), + } + + def _quantize_single_tensor( tensor_name: str, tensor: torch.Tensor, @@ -621,6 +689,24 @@ def _quantize_single_tensor( logger.warning(f"Failed to MXFP-quantize {layer_name}: {e}. Keeping original weight.") return layer_name, {tensor_name: tensor}, None, layer_name + # ---- NVFP4 E5M3 fake-quantization path ---- + if data_type == _NVFP4_E5M3_DATA_TYPE: + try: + quantize_e5m3 = ( + _quantize_weight_nvfp4_e5m3 if scheme.get("_output_format") == "fake" else _pack_weight_nvfp4_e5m3 + ) + out = quantize_e5m3( + weight=tensor, + layer_name=layer_name, + group_size=group_size, + device=device, + ) + logger.debug(f"Quantized (NVFP4_E5M3): {layer_name} (bits=4, group_size={group_size})") + return layer_name, out, layer_name, None + except Exception as e: + logger.warning(f"Failed to NVFP4_E5M3-quantize {layer_name}: {e}. Keeping original weight.") + return layer_name, {tensor_name: tensor}, None, layer_name + # ---- Integer WOQ path ---- try: qweight, qzeros, scales = quantize_func( @@ -660,12 +746,134 @@ def _collect_mxfp_source_entries(raw_tensors: dict[str, torch.Tensor]) -> list[t entries.append((layer_name, name, scale_key, 8)) elif name.endswith(".weight_packed") and tensor.dtype in (torch.int8, torch.uint8): layer_name = name[: -len(".weight_packed")] + # NVFP4 packed sources also use `.weight_packed` + `.weight_scale`, but + # are accompanied by global-scale tensors. Skip those here so they can + # be handled by the NVFP4 passthrough path instead of MXFP dequant. + if ( + f"{layer_name}.weight_global_scale" in raw_tensors + or f"{layer_name}.input_global_scale" in raw_tensors + or f"{layer_name}.weight_scale_2" in raw_tensors + or f"{layer_name}.input_scale" in raw_tensors + ): + continue scale_key = f"{layer_name}.weight_scale" if scale_key in raw_tensors: entries.append((layer_name, name, scale_key, 4)) return entries +def _normalize_nvfp4_source_tensors( + raw_tensors: dict[str, torch.Tensor], + shard_name: str | None = None, +) -> tuple[dict[str, torch.Tensor], list[str]]: + """Normalize legacy NVFP4 source naming to llm-compressor naming. + + Legacy checkpoints may store NVFP4 tensors as: + - ``.weight`` (packed U8) + - ``.weight_scale`` + - ``.weight_scale_2`` (reciprocal global scale) + - ``.input_scale`` (reciprocal global scale) + + For model-free passthrough and llm-compressor compatibility, convert to: + - ``.weight_packed`` + - ``.weight_scale`` + - ``.weight_global_scale`` + - ``.input_global_scale`` + """ + converted_layers: list[str] = [] + candidates: list[str] = [] + for name, tensor in list(raw_tensors.items()): + if not name.endswith(".weight"): + continue + layer_name = name[: -len(".weight")] + if tensor.dtype not in (torch.uint8, torch.int8): + continue + if f"{layer_name}.weight_scale" not in raw_tensors: + continue + has_legacy_global = f"{layer_name}.weight_scale_2" in raw_tensors or f"{layer_name}.input_scale" in raw_tensors + has_new_packed = f"{layer_name}.weight_packed" in raw_tensors + if has_legacy_global or has_new_packed: + candidates.append(layer_name) + + if not candidates: + return raw_tensors, converted_layers + + for layer_name in candidates: + weight_key = f"{layer_name}.weight" + weight_packed_key = f"{layer_name}.weight_packed" + weight_scale_2_key = f"{layer_name}.weight_scale_2" + input_scale_key = f"{layer_name}.input_scale" + weight_global_scale_key = f"{layer_name}.weight_global_scale" + input_global_scale_key = f"{layer_name}.input_global_scale" + + if weight_packed_key not in raw_tensors and weight_key in raw_tensors: + raw_tensors[weight_packed_key] = raw_tensors.pop(weight_key).view(torch.uint8).contiguous() + + if weight_scale_2_key in raw_tensors and weight_global_scale_key not in raw_tensors: + raw_tensors[weight_global_scale_key] = (1.0 / raw_tensors.pop(weight_scale_2_key).float()).to(torch.float32) + elif weight_scale_2_key in raw_tensors: + raw_tensors.pop(weight_scale_2_key) + + if input_scale_key in raw_tensors and input_global_scale_key not in raw_tensors: + raw_tensors[input_global_scale_key] = (1.0 / raw_tensors.pop(input_scale_key).float()).to(torch.float32) + elif input_scale_key in raw_tensors: + raw_tensors.pop(input_scale_key) + + converted_layers.append(layer_name) + + if converted_layers: + shard_prefix = f"[{shard_name}] " if shard_name else "" + logger.info(f"{shard_prefix}Normalized {len(converted_layers)} legacy NVFP4 layer(s) to llm-compressor naming.") + return raw_tensors, converted_layers + + +def _handle_nvfp4_source_tensors( + raw_tensors: dict[str, torch.Tensor], + matcher: "_PatternMatcher", +) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor], list[str]]: + """Passthrough NVFP4 source tensors when target scheme for the layer is NVFP4. + + This keeps already-quantized NVFP4 layers intact in model-free mode while + allowing other layers to be quantized normally. + """ + passthrough_tensors: dict[str, torch.Tensor] = {} + passthrough_layers: list[str] = [] + + for name, tensor in list(raw_tensors.items()): + if not name.endswith(".weight_packed") or tensor.dtype not in (torch.uint8, torch.int8): + continue + + layer_name = name[: -len(".weight_packed")] + scale_key = f"{layer_name}.weight_scale" + if scale_key not in raw_tensors: + continue + + scheme = matcher.resolve_scheme(f"{layer_name}.weight") + if scheme is None: + continue + scheme_bits = scheme.get("bits") + scheme_data_type = (scheme.get("data_type") or "").lower() + if not (scheme_bits == 4 and (is_nv_fp(scheme_data_type) or scheme_data_type == _NVFP4_E5M3_DATA_TYPE)): + continue + + keys_to_move = [name, scale_key] + weight_global_scale_key = f"{layer_name}.weight_global_scale" + input_global_scale_key = f"{layer_name}.input_global_scale" + if weight_global_scale_key in raw_tensors: + keys_to_move.append(weight_global_scale_key) + if input_global_scale_key in raw_tensors: + keys_to_move.append(input_global_scale_key) + + for key in keys_to_move: + passthrough_tensors[key] = raw_tensors.pop(key).to("cpu") + passthrough_layers.append(layer_name) + + if passthrough_layers: + logger.info(f"Handling NVFP4 source tensor(s): {len(passthrough_layers)} passthrough layer(s).") + + return raw_tensors, passthrough_tensors, passthrough_layers + + def _is_out_of_memory_error(exc: Exception) -> bool: if isinstance(exc, torch.OutOfMemoryError): return True @@ -701,6 +909,99 @@ def _dequantize_with_device_fallback( return on_cpu() +@lru_cache(maxsize=32) +def _load_weight_map_from_index(index_path: str) -> dict[str, str]: + """Load weight_map from an index file with a small process-local cache.""" + with open(index_path) as f: + index = json.load(f) + weight_map = index.get("weight_map", {}) + return weight_map if isinstance(weight_map, dict) else {} + + +def _hydrate_missing_fp8_scales_from_index( + raw_tensors: dict[str, torch.Tensor], + shard_path: str, + *, + shard_name: str | None = None, +) -> dict[str, torch.Tensor]: + """Populate missing ``.weight_scale_inv`` tensors from sibling shards. + + Some checkpoints shard FP8 weight and its corresponding scale tensor into + different ``.safetensors`` files. Model-free processing is shard-local, so + this helper hydrates missing ``.weight_scale_inv`` tensors by looking + up ``weight_map`` in ``*.safetensors.index.json`` and loading only the + needed tensors from referenced shards. + """ + if not shard_path.endswith(".safetensors"): + return raw_tensors + + weight_to_scale: dict[str, str] = {} + for name, tensor in raw_tensors.items(): + if not name.endswith(".weight"): + continue + if tensor.dtype != torch.float8_e4m3fn and tensor.element_size() != 1: + continue + scale_inv_name = name.replace(".weight", ".weight_scale_inv") + if scale_inv_name not in raw_tensors: + weight_to_scale[name] = scale_inv_name + + if not weight_to_scale: + return raw_tensors + + shard_dir = os.path.dirname(shard_path) + index_path = os.path.join(shard_dir, "model.safetensors.index.json") + if not os.path.exists(index_path): + candidates = sorted( + os.path.join(shard_dir, f) for f in os.listdir(shard_dir) if f.endswith(".safetensors.index.json") + ) + if not candidates: + return raw_tensors + index_path = candidates[0] + + try: + weight_map = _load_weight_map_from_index(index_path) + except Exception: + return raw_tensors + + current_shard = os.path.basename(shard_path) + scales_by_shard: dict[str, list[str]] = {} + for scale_name in weight_to_scale.values(): + target_shard = weight_map.get(scale_name) + if not target_shard or target_shard == current_shard: + continue + scales_by_shard.setdefault(target_shard, []).append(scale_name) + + if not scales_by_shard: + return raw_tensors + + from safetensors import safe_open + + hydrated = 0 + for target_shard, scale_names in scales_by_shard.items(): + target_path = os.path.join(shard_dir, target_shard) + if not os.path.exists(target_path): + continue + try: + with safe_open(target_path, framework="pt", device="cpu") as sf: + for scale_name in scale_names: + if scale_name in raw_tensors: + continue + try: + raw_tensors[scale_name] = sf.get_tensor(scale_name) + hydrated += 1 + except Exception: + # Tensor may be absent in this shard; skip lazily. + continue + except Exception: + continue + + if hydrated: + shard_prefix = f"[{shard_name}] " if shard_name else "" + logger.info(f"{shard_prefix}Hydrated {hydrated} FP8 scale tensor(s) from sibling shard(s) using index mapping.") + + return raw_tensors + + def _dequant_mxfp_tensors( raw_tensors: dict[str, torch.Tensor], device: str = "cpu", @@ -854,6 +1155,7 @@ def _dequant_fp8_tensors( block_size: list | None = None, device: str = "cpu", shard_name: str | None = None, + shard_path: str | None = None, ) -> dict[str, torch.Tensor]: """Dequantize DeepSeek-V3-style FP8 weight tensors to bfloat16. @@ -867,6 +1169,9 @@ def _dequant_fp8_tensors( """ from auto_round.utils.weight_handler import _dequant_fp8_linear_weight + if shard_path: + raw_tensors = _hydrate_missing_fp8_scales_from_index(raw_tensors, shard_path, shard_name=shard_name) + quant_entries: list[tuple[str, str]] = [] for name, tensor in raw_tensors.items(): if not name.endswith(".weight"): @@ -922,6 +1227,7 @@ def _process_shard( matcher: "_PatternMatcher | None" = None, fp8_block_size: list | None = None, model_type: str | None = None, + source_quantization_config: dict | None = None, enable_torch_compile: bool = False, ) -> tuple[dict[str, torch.Tensor], list[str], list[str]]: """Quantize eligible weights in a single safetensors shard. @@ -996,7 +1302,14 @@ def _process_shard( preserved_tensors[key] = raw_tensors.pop(key) # 1) model-type-specific preprocessing (format conversion only) - raw_tensors, source_state = _preprocess_model_type_source_tensors(raw_tensors, model_type=model_type) + raw_tensors, source_state = _preprocess_model_type_source_tensors( + raw_tensors, + model_type=model_type, + quantization_config=source_quantization_config, + ) + + # 1.5) normalize legacy NVFP4 names to llm-compressor naming. + raw_tensors, _converted_nvfp4_layers = _normalize_nvfp4_source_tensors(raw_tensors, shard_name=shard_name) # 2) generic MXFP handling for both preprocessed and normal source models raw_tensors, passthrough_tensors, passthrough_layers = _handle_mxfp_source_tensors( @@ -1009,11 +1322,20 @@ def _process_shard( output_tensors.update(passthrough_tensors) quantized_layers.extend(passthrough_layers) + # 3) NVFP4 passthrough for layers already stored in packed format. + raw_tensors, nvfp_passthrough_tensors, nvfp_passthrough_layers = _handle_nvfp4_source_tensors( + raw_tensors, + matcher, + ) + output_tensors.update(nvfp_passthrough_tensors) + quantized_layers.extend(nvfp_passthrough_layers) + raw_tensors = _dequant_fp8_tensors( raw_tensors, block_size=fp8_block_size, device=device, shard_name=shard_name, + shard_path=shard_path, ) raw_tensors.update(preserved_tensors) @@ -1042,6 +1364,51 @@ def _process_shard( # --------------------------------------------------------------------------- +def _get_llm_compressor_metadata() -> dict[str, str]: + """Return AutoRound provenance for model-free llm-compressor output.""" + # Keep metadata deterministic and installation-source agnostic. + return {"provider": "auto-round"} + + +def _build_nvfp4_e5m3_quantization_config(ignored_layers: list[str]) -> dict: + """Build compressed-tensors metadata for NVFP4 E5M3 without global scales.""" + from auto_round.export.export_to_llmcompressor.config import initialize_nvfp4_e5m3_quantization + + qconfig = initialize_nvfp4_e5m3_quantization(ignore=ignored_layers) + qconfig.update(_get_llm_compressor_metadata()) + return qconfig + + +def _get_mxfp_group_scheme_and_format( + group_bits: int, + group_data_type: str, + ignore: list[str], +): + """Return ``(QuantizationScheme, format_str)`` for a single (bits, data_type) group. + + Handles both MXFP (mx_fp) and NVFP4_E5M3 (fp4_v2) groups. + """ + from auto_round.export.export_to_llmcompressor.config import ( + initialize_nvfp4_e5m3_quantization, + initialize_quantization, + ) + + if group_data_type == _NVFP4_E5M3_DATA_TYPE: + # NVFP4_E5M3 is not a compressed-tensors preset; use its dedicated builder. + nvfp4_dict = initialize_nvfp4_e5m3_quantization(ignore=ignore) + # Extract the QuantizationScheme object from the dict config. + from compressed_tensors.quantization import QuantizationScheme + + raw_scheme = nvfp4_dict["config_groups"]["group_0"] + scheme_obj = QuantizationScheme.model_validate(raw_scheme) + return scheme_obj, "nvfp4-e5m3-pack-quantized" + else: + scheme_name = "MXFP4" if group_bits == 4 else "MXFP8" + fmt = "mxfp4-pack-quantized" if group_bits == 4 else "mxfp8-quantized" + tmp_qconfig = initialize_quantization(scheme=scheme_name, ignore=ignore) + return tmp_qconfig.config_groups["group_0"], fmt + + def _build_mxfp_quantization_config( default_scheme: dict, quantized_layers: list[str], @@ -1051,12 +1418,13 @@ def _build_mxfp_quantization_config( """Build a compressed-tensors / llm-compressor style quantization_config dict for MXFP4 / MXFP8 model-free output, including mixed-precision cases. - When *layer_config* contains layers that override the default bits (e.g. - some layers are MXFP8 while the default is MXFP4), the function creates - one ``config_group`` per distinct bit-width. Override groups list their - layers explicitly; the default-bits group uses ``targets=["Linear"]`` as a - catch-all. The top-level ``"format"`` is set to ``"mixed-precision"`` - when more than one group is produced. + When *layer_config* contains layers that override the default bits or + data_type (e.g. some layers are MXFP8 while the default is MXFP4, or + some layers use NVFP4_E5M3 while the default is MXFP8), the function + creates one ``config_group`` per distinct ``(bits, data_type)`` pair. + Override groups list their layers explicitly; the default group uses + ``targets=["Linear"]`` as a catch-all. The top-level ``"format"`` is + set to ``"mixed-precision"`` when more than one group is produced. Mirrors the per-group format produced by :mod:`auto_round.export.export_to_llmcompressor.export_to_fp`. @@ -1066,9 +1434,10 @@ def _build_mxfp_quantization_config( ) bits = default_scheme.get("bits", 4) + default_data_type = (default_scheme.get("data_type") or "mx_fp").lower() is_fp_default = (bits or 0) >= 16 # BF16/FP16 full-precision default - if not is_fp_default and bits not in _SUPPORTED_MXFP_BITS: + if not is_fp_default and bits not in _SUPPORTED_MXFP_BITS and default_data_type != _NVFP4_E5M3_DATA_TYPE: raise ValueError(f"Unsupported MXFP bits={bits} for model-free output.") # Default ignore list: any layer present in ignored_layers (deduped) that @@ -1077,8 +1446,9 @@ def _build_mxfp_quantization_config( quant_set = set(quantized_layers) ignore = [n for n in ignore if n not in quant_set] - # Resolve each quantized layer's effective bits using layer_config overrides. - scheme_groups: dict[int, list[str]] = {} # bits -> [layer_names] + # Resolve each quantized layer's effective (bits, data_type) using layer_config overrides. + # Key: (bits, data_type) to distinguish e.g. MXFP4 from NVFP4_E5M3 (both 4-bit). + scheme_groups: dict[tuple[int, str], list[str]] = {} # (bits, data_type) -> [layer_names] if layer_config: temp_matcher = _PatternMatcher( ignore_patterns=[], @@ -1088,51 +1458,65 @@ def _build_mxfp_quantization_config( for layer in quantized_layers: scheme = temp_matcher.resolve_scheme(f"{layer}.weight") layer_bits = scheme.get("bits", bits) if scheme is not None else bits - scheme_groups.setdefault(layer_bits, []).append(layer) + layer_dt = ( + (scheme.get("data_type") or default_data_type).lower() if scheme is not None else default_data_type + ) + scheme_groups.setdefault((layer_bits, layer_dt), []).append(layer) else: if not is_fp_default: - scheme_groups[bits] = list(quantized_layers) + scheme_groups[(bits, default_data_type)] = list(quantized_layers) # else: BF16 default with no layer_config → no MXFP layers; scheme_groups stays {} if len(scheme_groups) <= 1: - # Single scheme — use the actual MXFP bits from the group, not the - # default bits (which may be 16 for a BF16 default scheme). - actual_bits = next(iter(scheme_groups.keys())) if scheme_groups else bits - if actual_bits not in _SUPPORTED_MXFP_BITS: + # Single scheme — use the actual (bits, data_type) from the group. + if scheme_groups: + actual_bits, actual_dt = next(iter(scheme_groups.keys())) + else: + actual_bits, actual_dt = bits, default_data_type + if actual_dt != _NVFP4_E5M3_DATA_TYPE and actual_bits not in _SUPPORTED_MXFP_BITS: raise ValueError(f"Unsupported MXFP bits={actual_bits} for model-free output.") + group_scheme, fmt = _get_mxfp_group_scheme_and_format(actual_bits, actual_dt, ignore) + if actual_dt == _NVFP4_E5M3_DATA_TYPE: + from auto_round.export.export_to_llmcompressor.config import initialize_nvfp4_e5m3_quantization + + qconfig = initialize_nvfp4_e5m3_quantization(ignore=ignore) + if is_fp_default and scheme_groups: + qconfig["config_groups"]["group_0"]["targets"] = list(quantized_layers) + qconfig["format"] = fmt + qconfig.update(_get_llm_compressor_metadata()) + return qconfig + from auto_round.export.export_to_llmcompressor.config import initialize_quantization as _init_q + scheme_name = "MXFP4" if actual_bits == 4 else "MXFP8" - fmt = "mxfp4-pack-quantized" if actual_bits == 4 else "mxfp8-quantized" - qconfig = initialize_quantization(scheme=scheme_name, ignore=ignore) + qconfig = _init_q(scheme=scheme_name, ignore=ignore) if is_fp_default and scheme_groups: targets = list(quantized_layers) qconfig.config_groups["group_0"].targets = targets qconfig = qconfig.to_dict() qconfig["format"] = fmt - qconfig["provider"] = "auto-round" + qconfig.update(_get_llm_compressor_metadata()) return qconfig - # Mixed MXFP: build one config_group per distinct bit-width. - # Override groups (non-default bits) come first, default group last, + # Mixed precision: build one config_group per distinct (bits, data_type). + # Override groups (non-default key) come first, default group last, # ordered by descending bit-width within each partition so that the # higher-precision group gets the lower group index. + default_key = (bits, default_data_type) override_items = sorted( - [(b, layers) for b, layers in scheme_groups.items() if b != bits], - key=lambda x: x[0], + [(key, layers) for key, layers in scheme_groups.items() if key != default_key], + key=lambda x: x[0][0], reverse=True, ) - default_item = (bits, scheme_groups[bits]) if bits in scheme_groups else None + default_item = (default_key, scheme_groups[default_key]) if default_key in scheme_groups else None ordered = override_items + ([default_item] if default_item else []) config_groups: dict = {} group_formats: dict[str, str] = {} - for idx, (group_bits, layer_names) in enumerate(ordered): + for idx, ((group_bits, group_dt), layer_names) in enumerate(ordered): group_name = f"group_{idx}" - scheme_name = "MXFP4" if group_bits == 4 else "MXFP8" - fmt = "mxfp4-pack-quantized" if group_bits == 4 else "mxfp8-quantized" - is_default_group = group_bits == bits + is_default_group = (group_bits, group_dt) == default_key targets = ["Linear"] if is_default_group else layer_names - tmp_qconfig = initialize_quantization(scheme=scheme_name, ignore=ignore) - group_scheme = tmp_qconfig.config_groups["group_0"] + group_scheme, fmt = _get_mxfp_group_scheme_and_format(group_bits, group_dt, ignore) group_scheme.targets = targets config_groups[group_name] = group_scheme group_formats[group_name] = fmt @@ -1142,7 +1526,7 @@ def _build_mxfp_quantization_config( full_dict["format"] = "mixed-precision" for group_name, fmt in group_formats.items(): full_dict["config_groups"][group_name]["format"] = fmt - full_dict["provider"] = "auto-round" + full_dict.update(_get_llm_compressor_metadata()) return full_dict @@ -1381,6 +1765,8 @@ def _build_quantization_config( data_type = (default_scheme.get("data_type") or "int").lower() default_bits = default_scheme.get("bits", 4) is_fp_default = (default_bits or 0) >= 16 and not is_mx_fp(data_type) + if data_type == _NVFP4_E5M3_DATA_TYPE and format == "llm_compressor": + return _build_nvfp4_e5m3_quantization_config(ignored_layers) if is_mx_fp(data_type) or (is_fp_default and _layer_config_has_mxfp(layer_config)): if format in ("auto_round", "auto_round:auto_gptq"): return _build_mxfp_autoround_quantization_config( @@ -1411,7 +1797,10 @@ def _build_quantization_config( scheme_keys = [f.name for f in fields(QuantizationScheme)] # vllm only support auto_round:auto_gptq, but transformers cannot load it correctly when sym=False. # So we keep auto_round for asymmetric quantization to maintain compatibility with both. - packing_format = "auto_round:auto_gptq" if default_scheme.get("sym", True) else "auto_round" + if data_type == _NVFP4_E5M3_DATA_TYPE: + packing_format = "auto_round:fake" if format == "fake" else "auto_round:llm_compressor_nvfp4_e5m3" + else: + packing_format = "auto_round:auto_gptq" if default_scheme.get("sym", True) else "auto_round" qconfig = { "quant_method": "auto-round", @@ -1425,6 +1814,12 @@ def _build_quantization_config( "autoround_version": __version__, } + if data_type == _NVFP4_E5M3_DATA_TYPE: + for act_key in ("act_bits", "act_data_type", "act_group_size", "act_sym", "act_dynamic"): + value = default_scheme.get(act_key) + if value is not None: + qconfig[act_key] = value + if block_name_to_quantize: qconfig["block_name_to_quantize"] = block_name_to_quantize @@ -1453,6 +1848,8 @@ def _build_quantization_config( if non_linear_re.search(layer_name): continue extra_config[layer_name] = {"bits": 16, "data_type": "float"} + if data_type == _NVFP4_E5M3_DATA_TYPE: + extra_config[layer_name].update({"act_bits": 16, "act_data_type": "float"}) quantized_layer_set = set(quantized_layers) if "lm_head" in quantized_layer_set and "lm_head" not in extra_config: @@ -1549,6 +1946,7 @@ def _process_single_shard_task( ignore_patterns: list[str], fp8_block_size: list | None, model_type: str | None, + source_quantization_config: dict | None = None, quant_output_dir: str, total_shards: int, enable_torch_compile: bool = False, @@ -1568,6 +1966,48 @@ def _process_single_shard_task( if shard_path is None or not os.path.exists(shard_path): return shard_idx, shard_name, None, None, None, None, None + return _quantize_local_shard_task( + shard_idx, + shard_name, + shard_path=shard_path, + device=device, + default_scheme=default_scheme, + layer_config=layer_config, + ignore_patterns=ignore_patterns, + fp8_block_size=fp8_block_size, + model_type=model_type, + source_quantization_config=source_quantization_config, + quant_output_dir=quant_output_dir, + total_shards=total_shards, + enable_torch_compile=enable_torch_compile, + cleanup_source_shard=is_streaming, + ) + + +def _quantize_local_shard_task( + shard_idx: int, + shard_name: str, + *, + shard_path: str, + device: str, + default_scheme: dict, + layer_config: dict, + ignore_patterns: list[str], + fp8_block_size: list | None, + model_type: str | None, + source_quantization_config: dict | None, + quant_output_dir: str, + total_shards: int, + enable_torch_compile: bool = False, + cleanup_source_shard: bool = False, +) -> tuple[int, str, str | None, str | None, list[str] | None, list[str] | None, list[str] | None]: + """Quantize one already-downloaded shard and write the output shard. + + Returns lightweight metadata only so IPC does not transfer tensor storages. + """ + if shard_path is None or not os.path.exists(shard_path): + return shard_idx, shard_name, None, None, None, None, None + output_tensors, quantized, ignored = _process_shard( shard_path=shard_path, shard_name=shard_name, @@ -1577,6 +2017,7 @@ def _process_single_shard_task( device=device, fp8_block_size=fp8_block_size, model_type=model_type, + source_quantization_config=source_quantization_config, enable_torch_compile=enable_torch_compile, ) @@ -1591,13 +2032,12 @@ def _process_single_shard_task( tensor_names = list(local_weight_map.keys()) clear_memory() - if is_streaming: + if cleanup_source_shard: try: os.remove(shard_path) except OSError: pass - # Return only lightweight metadata to avoid IPC transfer of tensor storages. return shard_idx, shard_name, shard_path, out_shard_name, tensor_names, quantized, ignored @@ -1720,6 +2160,21 @@ def _validate_supported_scheme( ) return + if data_type == _NVFP4_E5M3_DATA_TYPE: + if bits != 4 or scheme_obj.group_size != 16 or act_bits != 4: + raise ValueError( + f"Model-free NVFP4_E5M3 requires bits=4, group_size=16, and act_bits=4, " + f"but '{scheme_input}' requests bits={bits}, group_size={scheme_obj.group_size}, " + f"act_bits={act_bits}." + ) + if (scheme_obj.act_data_type or "").lower() != _NVFP4_E5M3_DATA_TYPE or scheme_obj.act_group_size != 16: + raise ValueError( + f"Model-free NVFP4_E5M3 requires act_data_type='fp4_v2' and act_group_size=16, " + f"but '{scheme_input}' requests act_data_type='{scheme_obj.act_data_type}', " + f"act_group_size={scheme_obj.act_group_size}." + ) + return + if act_bits < 16: raise ValueError( f"Model-free mode only supports weight-only quantization (WOQ) schemes " @@ -1962,6 +2417,7 @@ class _ModelFreeCompressorCore: """ SUPPORTED_FORMATS: tuple[str, ...] = ( + "fake", "auto_round", "auto_round:auto_gptq", "llm_compressor", @@ -1987,7 +2443,7 @@ def __init__( self.scheme_input = scheme self.layer_config_input = layer_config self.ignore_layers_input = ignore_layers - self.format = format + self.format = format or "auto_round" self.device = device self.quant_lm_head = quant_lm_head self.quant_nontext_module = quant_nontext_module @@ -2043,6 +2499,7 @@ def _parse_scheme(self) -> None: _validate_supported_scheme(self.scheme_obj, self.scheme_input) ds = asdict(self.scheme_obj) self.default_scheme = {k: v for k, v in ds.items() if v is not None} + self.default_scheme["_output_format"] = self.format def _parse_layer_config(self) -> None: lc = copy.deepcopy(self.layer_config_input) if self.layer_config_input else {} @@ -2253,6 +2710,10 @@ def _quant_output_dir(self) -> str: # ------------------------------------------------------------------- def _process_all_shards(self) -> None: + if self.is_streaming: + self._process_all_shards_streaming_pipeline() + return + try: from tqdm import tqdm as _tqdm except ImportError: @@ -2284,6 +2745,7 @@ def _process_all_shards(self) -> None: ignore_patterns=self.ignore_patterns, fp8_block_size=self.fp8_block_size, model_type=self.model_type, + source_quantization_config=self.config.get("quantization_config", {}), enable_torch_compile=self.enable_torch_compile, quant_output_dir=self._quant_output_dir, total_shards=len(self.shard_names), @@ -2297,36 +2759,8 @@ def _process_all_shards(self) -> None: ) for future in shard_iter: - shard_idx, shard_name, shard_path, out_shard_name, tensor_names, quantized, ignored = future.result() - - if ( - shard_path is None - or out_shard_name is None - or tensor_names is None - or quantized is None - or ignored is None - ): - logger.warning(f"Shard not found: {shard_name}, skipping") - continue - - memory_monitor.update() - clear_memory() - if len(self.shard_names) > 1: - logger.info(f"Memory usage: {memory_monitor.get_summary()}") - - compressed_quantized = compress_layer_names(quantized) - compressed_ignored = compress_layer_names(ignored) - logger.info( - f"Shard {shard_idx + 1}/{len(self.shard_names)} ({shard_name}):\n" - f" Quantized layers ({len(quantized)}): {compressed_quantized}\n" - f" Ignored layers ({len(ignored)}): {compressed_ignored}" - ) - - self.all_quantized_layers.extend(quantized) - self.all_ignored_layers.extend(ignored) - - for tensor_name in tensor_names: - self.output_weight_map[tensor_name] = out_shard_name + result = future.result() + self._merge_shard_task_result(result) except KeyboardInterrupt: logger.warning("Interrupted by user; terminating model-free shard worker processes.") _force_cleanup_process_pool(pool) @@ -2337,6 +2771,152 @@ def _process_all_shards(self) -> None: finally: _force_cleanup_process_pool(pool) + def _merge_shard_task_result( + self, + result: tuple[int, str, str | None, str | None, list[str] | None, list[str] | None, list[str] | None], + ) -> None: + """Merge one shard-task result into global stats and weight map.""" + shard_idx, shard_name, shard_path, out_shard_name, tensor_names, quantized, ignored = result + if shard_path is None or out_shard_name is None or tensor_names is None or quantized is None or ignored is None: + logger.warning(f"Shard not found: {shard_name}, skipping") + return + + memory_monitor.update() + clear_memory() + if len(self.shard_names) > 1: + logger.info(f"Memory usage: {memory_monitor.get_summary()}") + + compressed_quantized = compress_layer_names(quantized) + compressed_ignored = compress_layer_names(ignored) + logger.info( + f"Shard {shard_idx + 1}/{len(self.shard_names)} ({shard_name}):\n" + f" Quantized layers ({len(quantized)}): {compressed_quantized}\n" + f" Ignored layers ({len(ignored)}): {compressed_ignored}" + ) + + self.all_quantized_layers.extend(quantized) + self.all_ignored_layers.extend(ignored) + for tensor_name in tensor_names: + self.output_weight_map[tensor_name] = out_shard_name + + def _process_all_shards_streaming_pipeline(self) -> None: + """Streaming-mode shard pipeline with dedicated downloader and quant workers. + + Design: + - one downloader worker serializes network bandwidth usage; + - N quant workers consume ready shards independently; + - shards are assigned for quantization as soon as download completes. + """ + try: + from tqdm import tqdm as _tqdm + except ImportError: + _tqdm = None + + if not self.shard_names: + return + + os.makedirs(self._quant_output_dir, exist_ok=True) + + worker_count = max(1, min(self.shard_parallelism, len(self.shard_names))) + prefetch_depth = max(2, worker_count) + total_shards = len(self.shard_names) + + download_pool: ThreadPoolExecutor | None = None + quant_pool: ProcessPoolExecutor | None = None + download_futures: dict = {} + quant_futures = set() + next_download_idx = 0 + completed_quant = 0 + + def _submit_next_download() -> bool: + nonlocal next_download_idx + if next_download_idx >= total_shards: + return False + shard_idx = next_download_idx + shard_name = self.shard_names[shard_idx] + future = download_pool.submit( + _prefetch_shard, + self.model_name_or_path, + shard_name, + self.work_dir, + self.source_dir, + self.is_streaming, + ) + download_futures[future] = (shard_idx, shard_name) + next_download_idx += 1 + return True + + try: + download_pool = ThreadPoolExecutor(max_workers=1) + quant_pool = ProcessPoolExecutor(max_workers=worker_count, mp_context=mp.get_context("spawn")) + + for _ in range(min(prefetch_depth, total_shards)): + _submit_next_download() + + progress = _tqdm(total=total_shards, desc="Processing shards", unit="shard") if _tqdm else None + + while completed_quant < total_shards: + wait_set = set(download_futures.keys()) | set(quant_futures) + if not wait_set: + break + + done, _ = wait(wait_set, return_when=FIRST_COMPLETED) + for future in done: + if future in download_futures: + shard_idx, shard_name = download_futures.pop(future) + shard_path = future.result() + if shard_path is None or not os.path.exists(shard_path): + logger.warning(f"Prefetch failed for shard {shard_name}, skipping") + completed_quant += 1 + if progress is not None: + progress.update(1) + else: + qf = quant_pool.submit( + _quantize_local_shard_task, + shard_idx, + shard_name, + shard_path=shard_path, + device=self.device, + default_scheme=self.default_scheme, + layer_config=self.layer_config, + ignore_patterns=self.ignore_patterns, + fp8_block_size=self.fp8_block_size, + model_type=self.model_type, + source_quantization_config=self.config.get("quantization_config", {}), + quant_output_dir=self._quant_output_dir, + total_shards=total_shards, + enable_torch_compile=self.enable_torch_compile, + cleanup_source_shard=True, + ) + quant_futures.add(qf) + + while len(download_futures) < prefetch_depth and _submit_next_download(): + pass + elif future in quant_futures: + quant_futures.remove(future) + result = future.result() + self._merge_shard_task_result(result) + completed_quant += 1 + if progress is not None: + progress.update(1) + + if progress is not None: + progress.close() + except KeyboardInterrupt: + logger.warning("Interrupted by user; terminating model-free shard workers.") + _force_cleanup_process_pool(quant_pool) + raise + except Exception: + _force_cleanup_process_pool(quant_pool) + raise + finally: + _force_cleanup_process_pool(quant_pool) + if download_pool is not None: + try: + download_pool.shutdown(wait=False, cancel_futures=True) + except Exception: + pass + # ------------------------------------------------------------------- # Output # ------------------------------------------------------------------- @@ -2354,6 +2934,7 @@ def _write_config_files(self) -> None: break block_name_to_quantize = ",".join(dict.fromkeys(block_prefixes)) or None + os.makedirs(self._quant_output_dir, exist_ok=True) quantization_config = _build_quantization_config( default_scheme=self.default_scheme, layer_config=self.layer_config, @@ -2365,7 +2946,6 @@ def _write_config_files(self) -> None: ) self.config["quantization_config"] = quantization_config - os.makedirs(self._quant_output_dir, exist_ok=True) with open(os.path.join(self._quant_output_dir, "config.json"), "w") as f: json.dump(self.config, f, indent=2) @@ -2476,6 +3056,8 @@ def run(self) -> str: if is_mx_fp(data_type): bits = self.default_scheme.get("bits", 4) packing_format = "mxfp4-pack-quantized" if bits == 4 else "mxfp8-quantized" + elif data_type == _NVFP4_E5M3_DATA_TYPE: + packing_format = "fake" if self.format == "fake" else "auto_round:llm_compressor_nvfp4_e5m3" else: packing_format = "auto_round:auto_gptq" @@ -2867,6 +3449,8 @@ def quantize_and_save( normalized_scheme is not None and is_mx_fp((normalized_scheme.data_type or "").lower()) ) or self._auto_scheme_family == "mx_fp": _accepted_formats = {"llm_compressor", "auto_round", "auto_round:auto_gptq"} + elif normalized_scheme is not None and (normalized_scheme.data_type or "").lower() == _NVFP4_E5M3_DATA_TYPE: + _accepted_formats = {"fake", "llm_compressor", "auto_round", "auto_round:auto_gptq"} elif _is_full_precision_default(self.scheme_input) and _layer_config_has_mxfp(self.layer_config_input): # BF16 default with MXFP layer_config overrides. _accepted_formats = {"llm_compressor", "auto_round", "auto_round:auto_gptq"} @@ -2921,8 +3505,11 @@ def _expand_e8m0_block_scale( Because every fine MXFP group lies entirely inside a single coarse block, the expansion is a pure ``repeat_interleave`` along both axes (no - interpolation). The returned tensor is ``uint8`` (raw E8M0 bytes), matching - the ``U8`` dtype used by llm-compressor ``weight_scale`` tensors. + interpolation). For DeepSeek variants that store *tail* blocks using + ceil-based 128x128 tiling (for example rows ``5 -> 576``), we expand with + 128 repeats and then slice the tail. The returned tensor is ``uint8`` (raw + E8M0 bytes), matching the ``U8`` dtype used by llm-compressor + ``weight_scale`` tensors. """ scale = scale.view(torch.uint8) if scale.dim() != 2: @@ -2932,23 +3519,41 @@ def _expand_e8m0_block_scale( target_cols = in_features // group_size rows, cols = scale.shape - if target_rows % rows != 0 or target_cols % cols != 0: - raise ValueError( - f"Cannot expand E8M0 block scale {tuple(scale.shape)} to " - f"({target_rows}, {target_cols}); shapes are not divisible." - ) - - if target_rows != rows: - scale = scale.repeat_interleave(target_rows // rows, dim=0) - if target_cols != cols: - scale = scale.repeat_interleave(target_cols // cols, dim=1) - return scale.contiguous() + # Standard path: exact divisibility between coarse and target shapes. + if target_rows % rows == 0 and target_cols % cols == 0: + if target_rows != rows: + scale = scale.repeat_interleave(target_rows // rows, dim=0) + if target_cols != cols: + scale = scale.repeat_interleave(target_cols // cols, dim=1) + return scale.contiguous() + + # DeepSeek FP8/UE8M0 path: coarse scales are laid out in ceil(./128) blocks. + # This handles tail blocks like rows=5 for out_features=576. + coarse_block = 128 + expected_rows = (out_features + coarse_block - 1) // coarse_block + expected_cols = (in_features + coarse_block - 1) // coarse_block + if rows == expected_rows and cols == expected_cols: + if coarse_block % group_size != 0: + raise ValueError( + f"Cannot expand DeepSeek E8M0 block scale with group_size={group_size}; " + f"{coarse_block} is not divisible by group_size." + ) + groups_per_block_col = coarse_block // group_size + scale = scale.repeat_interleave(coarse_block, dim=0)[:target_rows] + scale = scale.repeat_interleave(groups_per_block_col, dim=1)[:, :target_cols] + return scale.contiguous() + + raise ValueError( + f"Cannot expand E8M0 block scale {tuple(scale.shape)} to " + f"({target_rows}, {target_cols}); unsupported coarse/block layout." + ) def _preprocess_model_type_source_tensors( raw_tensors: dict[str, torch.Tensor], model_type: str | None, group_size: int = 32, + quantization_config: dict | None = None, ) -> tuple[dict[str, torch.Tensor], dict[str, int]]: """Apply model-type-specific source tensor normalization. @@ -2961,15 +3566,28 @@ def _preprocess_model_type_source_tensors( ``(raw_tensors, source_state)`` where ``source_state[layer]`` is the source MXFP bits (4 or 8) for model-type preprocessed layers. """ - if (model_type or "").lower() != "deepseek_v4": + model_type = (model_type or "").lower() + quantization_config = quantization_config or {} + is_deepseek_v4 = model_type == "deepseek_v4" + is_deepseek_v32_ue8m0 = ( + model_type == "deepseek_v32" + and quantization_config.get("quant_method") == "fp8" + and str(quantization_config.get("fmt", "")).lower() == "e4m3" + and str(quantization_config.get("scale_fmt", "")).lower() == "ue8m0" + ) + if not is_deepseek_v4 and not is_deepseek_v32_ue8m0: return raw_tensors, {} entries: list[tuple[str, str, bool]] = [] # (weight_name, scale_name, is_fp8) for name, tensor in raw_tensors.items(): if not name.endswith(".weight"): continue - scale_name = name[: -len(".weight")] + ".scale" - if scale_name not in raw_tensors: + layer_name = name[: -len(".weight")] + scale_candidates = [f"{layer_name}.scale"] + if is_deepseek_v32_ue8m0: + scale_candidates.extend((f"{layer_name}.weight_scale", f"{layer_name}.weight_scale_inv")) + scale_name = next((candidate for candidate in scale_candidates if candidate in raw_tensors), None) + if scale_name is None: continue if tensor.dtype == torch.float8_e4m3fn: entries.append((name, scale_name, True)) @@ -2992,6 +3610,21 @@ def _preprocess_model_type_source_tensors( weight_key = f"{layer_name}.weight" source_state[layer_name] = 8 n_fp8 += 1 + + # DeepSeek V32 UE8M0: the weight_scale_inv is stored in float32 + # but only the 8-bit exponent field is meaningful (UE8M0 encoding). + # Extract the biased exponent from each fp32 element as a uint8 byte + # so that _expand_e8m0_block_scale receives the expected U8 E8M0 tensor. + # float32 layout: sign(1) | exponent(8) | mantissa(23) + # → uint8 E8M0 = (view_as_int32 >> 23) & 0xFF + if scale.dtype == torch.float32: + sanitized_scale_name = ".".join("" if part.isdigit() else part for part in scale_name.split(".")) + logger.warning_once( + f"[{model_type}] Scale tensor '{sanitized_scale_name}' has dtype float32 with UE8M0 encoding " + f"(only the 8-bit exponent is significant). " + f"Extracting uint8 E8M0 exponent bytes from fp32 representation." + ) + scale = ((scale.view(torch.int32) >> 23) & 0xFF).to(torch.uint8) else: out_features = weight.shape[0] in_features = weight.shape[1] * 2 @@ -3005,7 +3638,7 @@ def _preprocess_model_type_source_tensors( raw_tensors[f"{layer_name}.weight_scale"] = weight_scale logger.info( - "Applied model_type preprocessing for deepseek_v4: " + f"Applied model_type preprocessing for {model_type}: " f"{n_fp8} MXFP8 layer(s), {n_fp4} MXFP4 layer(s) converted to llm-compressor naming." ) return raw_tensors, source_state diff --git a/auto_round/data_type/nvfp.py b/auto_round/data_type/nvfp.py index bab5b54105..582d652f31 100644 --- a/auto_round/data_type/nvfp.py +++ b/auto_round/data_type/nvfp.py @@ -206,10 +206,15 @@ def ref_fp4_quant(x, global_scale, block_size=16, v=0, max_scale=1.0): scale = global_scale * (vec_max * get_reciprocal(FLOAT4_E2M1_MAX)) scale = torch.clip(scale, 0, FLOAT8_UE5M3_MAX) scale = cast_to_ue5m3_ste(scale).to(torch.float32) - output_scale = get_reciprocal(scale * get_reciprocal(global_scale)) - scaled_x = x.to(torch.float32) * output_scale + v + dequant_scale = scale * get_reciprocal(global_scale) + scaled_x = torch.where( + dequant_scale == 0, + torch.zeros_like(x, dtype=torch.float32), + x.to(torch.float32) / dequant_scale, + ) + scaled_x = scaled_x + v clipped_x = torch.clamp(scaled_x, -6.0, 6.0) - return (cast_to_fp4(clipped_x) * get_reciprocal(output_scale)).reshape(m, n), scale + return (cast_to_fp4(clipped_x) * dequant_scale).reshape(m, n), scale @register_dtype("fp4_v2_with_global_scale") diff --git a/auto_round/envs.py b/auto_round/envs.py index b72b14ae39..c535dcdf2e 100644 --- a/auto_round/envs.py +++ b/auto_round/envs.py @@ -28,6 +28,7 @@ AR_AUTO_SCHEME_BATCH_SIZE: Optional[int] = None AR_AUTO_SCHEME_CACHE: Optional[str] = None AR_ENABLE_AUTO_SCHEME_PARALLEL: bool = True + AR_NVFP4_E5M3_CACHE_HP_WEIGHT: bool = False AR_DISK_STREAM_MODEL: bool = False AR_RESUME_DIR: Optional[str] = None @@ -101,6 +102,13 @@ def _get_optional_positive_int_env(name: str) -> Optional[int]: # set it to 0 when workers could exhaust host RAM or device memory. "AR_ENABLE_AUTO_SCHEME_PARALLEL": lambda: os.getenv("AR_ENABLE_AUTO_SCHEME_PARALLEL", "1").lower() in ("1", "true", "yes"), + # Controls whether NVFP4 E5M3 quant linear caches a dequantized high- + # precision weight after the first forward instead of dequantizing on + # every call. When enabled, the packed weight buffers are released after + # the cache is materialized, trading lower runtime overhead for higher + # steady-state memory usage. + "AR_NVFP4_E5M3_CACHE_HP_WEIGHT": lambda: os.getenv("AR_NVFP4_E5M3_CACHE_HP_WEIGHT", "0").lower() + in ("1", "true", "yes", "on"), # When set, the model is built as a meta-device skeleton and streamed # block-by-block from disk during quantization instead of being fully # materialized on CPU RAM up front. diff --git a/auto_round/experimental/qmodules/__init__.py b/auto_round/experimental/qmodules/__init__.py index 4f5ff8e829..e638f95bb0 100644 --- a/auto_round/experimental/qmodules/__init__.py +++ b/auto_round/experimental/qmodules/__init__.py @@ -19,5 +19,7 @@ MXINT4QuantLinear, ) +from auto_round.experimental.qmodules.fake import FakeActQuantLinear from auto_round.experimental.qmodules.nvfp4 import NVFP4QuantLinear +from auto_round.experimental.qmodules.nvfp4_e5m3 import CuteNVFP4E5M3QuantLinear, NVFP4E5M3QuantLinear from auto_round.experimental.qmodules.fp8_static import WeightFP8ActFP8StaticQuantLinear diff --git a/auto_round/experimental/qmodules/fake.py b/auto_round/experimental/qmodules/fake.py new file mode 100644 index 0000000000..0089a2555d --- /dev/null +++ b/auto_round/experimental/qmodules/fake.py @@ -0,0 +1,87 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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 Optional + +import torch + +from auto_round.data_type.utils import get_quant_func +from auto_round.experimental.qmodules.base import QModuleBase +from auto_round.schemes import QuantizationScheme + +__all__ = ["FakeActQuantLinear"] + + +class FakeActQuantLinear(QModuleBase): + """Linear with high-precision QDQ weights and runtime activation QDQ.""" + + def __init__( + self, + in_features: int, + out_features: int, + config: QuantizationScheme, + weight: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + dtype: torch.dtype = torch.bfloat16, + ): + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.config = config + if weight is None: + weight = torch.empty((out_features, in_features), dtype=dtype) + self.weight = torch.nn.Parameter(weight, requires_grad=False) + if bias is None: + self.register_parameter("bias", None) + else: + self.bias = torch.nn.Parameter(bias, requires_grad=False) + + @classmethod + def from_original(cls, config: QuantizationScheme, original_layer: torch.nn.Linear): + return cls( + in_features=original_layer.in_features, + out_features=original_layer.out_features, + config=config, + weight=original_layer.weight, + bias=original_layer.bias, + dtype=original_layer.weight.dtype, + ) + + @classmethod + def get_min_capability(cls) -> int: + return 0 + + def process_weights_after_loading(self, layer: torch.nn.Module): + return + + def post_init(self): + return + + def qdq_input(self, activation: torch.Tensor) -> torch.Tensor: + quant_func, _ = get_quant_func( + dtype=self.config.act_data_type, + bits=self.config.act_bits, + sym=self.config.act_sym, + ) + qdq_activation, _, _ = quant_func( + tensor=activation, + bits=self.config.act_bits, + group_size=self.config.act_group_size, + ) + return qdq_activation.to(activation.dtype) + + @torch.inference_mode() + def forward(self, activation: torch.Tensor) -> torch.Tensor: + qdq_activation = self.qdq_input(activation) + return torch.nn.functional.linear(qdq_activation, self.weight.to(qdq_activation.dtype), self.bias) diff --git a/auto_round/experimental/qmodules/nvfp4_e5m3.py b/auto_round/experimental/qmodules/nvfp4_e5m3.py new file mode 100644 index 0000000000..cec785c38b --- /dev/null +++ b/auto_round/experimental/qmodules/nvfp4_e5m3.py @@ -0,0 +1,161 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +import os +from typing import Optional, Union + +import torch + +from auto_round.data_type.nvfp import e5m3_to_float_tensor, fp4_v2 +from auto_round.experimental.qmodules.base import QModuleBase +from auto_round.experimental.qmodules.fp4_utils import unpack_fp4_from_uint8 +from auto_round.logger import logger +from auto_round.schemes import QuantizationScheme +from auto_round_extension.cuda.cute_nvfp4_e5m3 import ( + try_cute_fp4_v2_qdq, + try_cute_nvfp4_e5m3_linear, + try_cute_nvfp4_e5m3_weight_dq, +) + +__all__ = ["CuteNVFP4E5M3QuantLinear", "NVFP4E5M3QuantLinear"] + +_CACHE_WEIGHT_ENV = "AR_NVFP4_E5M3_CACHE_HP_WEIGHT" + + +def _resolve_cache_weight(cache_weight: Optional[bool], default: bool) -> bool: + if cache_weight is not None: + return cache_weight + value = os.getenv(_CACHE_WEIGHT_ENV) + if value is None: + return default + return value.strip().lower() not in {"0", "false", "no", "off"} + + +class NVFP4E5M3QuantLinear(QModuleBase): + """FP4 E2M1 weights and activations with unsigned E5M3 block scales.""" + + SUPPORTED_COMPUTE_DTYPE = [torch.bfloat16, torch.float16, torch.float32] + DEFAULT_CACHE_WEIGHT = False + + def __init__( + self, + in_features: int, + out_features: int, + config: QuantizationScheme, + weight: Optional[torch.Tensor] = None, + weight_scale: Optional[torch.Tensor] = None, + bias: Union[torch.Tensor, bool, None] = None, + dtype=torch.bfloat16, + cache_weight: Optional[bool] = None, + ): + super().__init__() + assert dtype in self.SUPPORTED_COMPUTE_DTYPE + assert config.group_size == 16 and config.act_group_size == 16 + self.in_features = in_features + self.out_features = out_features + self.group_size = config.group_size + self.config = config + self.dtype = dtype + self.cache_weight = _resolve_cache_weight(cache_weight, self.DEFAULT_CACHE_WEIGHT) + self._cached_weight = None + + packed_weight = torch.zeros((out_features, in_features // 2), dtype=torch.uint8) if weight is None else weight + self.register_buffer("weight_packed", packed_weight) + scale = ( + torch.empty((out_features, in_features // self.group_size), dtype=torch.uint8) + if weight_scale is None + else weight_scale + ) + self.register_buffer("weight_scale", scale) + + if bias is not None: + if isinstance(bias, bool): + bias = torch.zeros((out_features,), dtype=dtype) + self.bias = torch.nn.Parameter(bias, requires_grad=False) + else: + self.register_parameter("bias", None) + + @classmethod + def get_min_capability(cls) -> int: + logger.warning_once("NVFP4 E5M3 quantization uses reference PyTorch inference and may be slow.") + return 0 + + def dequant_weight_online(self) -> torch.Tensor: + unpacked = unpack_fp4_from_uint8(self.weight_packed, self.out_features, self.in_features, dtype=self.dtype).to( + torch.float32 + ) + scale = e5m3_to_float_tensor(self.weight_scale).reshape(-1, 1) + return (unpacked.reshape(-1, self.group_size) * scale).reshape(self.out_features, self.in_features) + + @property + def weight(self) -> torch.Tensor: + if self._cached_weight is None: + self._cached_weight = self.dequant_weight_online() + if self.cache_weight: + self.weight_packed = None + self.weight_scale = None + return self._cached_weight + + def clear_weight_cache(self) -> None: + if self.weight_packed is None: + raise RuntimeError("Cannot clear the cached weight after quantized weight buffers have been released.") + self._cached_weight = None + + def qdq_input(self, activation: torch.Tensor) -> torch.Tensor: + original_dtype = activation.dtype + qdq_activation, _, _ = fp4_v2( + activation.to(torch.float32), bits=self.config.act_bits, group_size=self.config.act_group_size + ) + return qdq_activation.to(original_dtype) + + @torch.inference_mode() + def forward(self, input: torch.Tensor) -> torch.Tensor: + qdq_input = self.qdq_input(input) + weight = self.weight if self.cache_weight else self.dequant_weight_online() + return torch.nn.functional.linear(qdq_input, weight.to(qdq_input.dtype), self.bias) + + @classmethod + def from_original(cls, config: QuantizationScheme, original_layer: torch.nn.Linear): + return cls( + in_features=original_layer.in_features, + out_features=original_layer.out_features, + config=config, + bias=original_layer.bias, + dtype=original_layer.weight.dtype, + ) + + +class CuteNVFP4E5M3QuantLinear(NVFP4E5M3QuantLinear): + """NVFP4 E5M3 linear that dispatches activation QDQ and GEMM to CuTe.""" + + def dequant_weight_online(self) -> torch.Tensor: + cute_weight = try_cute_nvfp4_e5m3_weight_dq(self.weight_packed, self.weight_scale, self.dtype) + if cute_weight is not None: + return cute_weight + return super().dequant_weight_online() + + def qdq_input(self, activation: torch.Tensor) -> torch.Tensor: + cute_qdq_activation = try_cute_fp4_v2_qdq(activation, self.config.act_group_size) + if cute_qdq_activation is not None: + return cute_qdq_activation + return super().qdq_input(activation) + + @torch.inference_mode() + def forward(self, input: torch.Tensor) -> torch.Tensor: + if self.cache_weight: + return super().forward(input) + fused_output = try_cute_nvfp4_e5m3_linear(input, self.weight_packed, self.weight_scale, self.bias) + if fused_output is not None: + return fused_output + return super().forward(input) diff --git a/auto_round/export/export_to_autoround/qlinear_fp.py b/auto_round/export/export_to_autoround/qlinear_fp.py index 338d624a7a..c3e0fb38bd 100644 --- a/auto_round/export/export_to_autoround/qlinear_fp.py +++ b/auto_round/export/export_to_autoround/qlinear_fp.py @@ -36,7 +36,7 @@ import auto_round.envs as envs from auto_round.compressors.utils import BackendDataType, is_mx_fp, is_nv_fp from auto_round.data_type.mxfp import FP32_EXPONENT_BIAS, FP32_MIN_NORMAL -from auto_round.data_type.nvfp import cast_to_fp4, get_reciprocal +from auto_round.data_type.nvfp import cast_to_fp4, float_to_e5m3_frexp, get_reciprocal from auto_round.data_type.utils import reshape_pad_tensor_by_group_size, revert_tensor_by_pad from auto_round.utils import get_packing_device, logger @@ -73,6 +73,7 @@ def __init__( raise NotImplementedError("Only 4,8 bits are supported.") self.is_mx = is_mx_fp(data_type) self.is_nv = is_nv_fp(data_type) + self.is_nvfp4_e5m3 = data_type == "fp4_v2" if self.is_mx: if group_size != 32: raise NotImplementedError(f"Only group_size 32 are supported for {BackendDataType.MX_FP} data type.") @@ -80,7 +81,7 @@ def __init__( raise NotImplementedError( f"in_feature must be divisible by {group_size} for {BackendDataType.MX_FP} data type." ) - if self.is_nv: + if self.is_nv or self.is_nvfp4_e5m3: if group_size % 16 != 0: raise NotImplementedError(f"Only group_size 16 are supported for {BackendDataType.NV_FP} data type.") if infeatures % group_size != 0: @@ -159,11 +160,17 @@ def pack(self, linear, scales, zeros=None, g_idx=None, global_scale=None, input_ ) scaled_tensor.clamp_(-6.0, 6.0) scaled_tensor = cast_to_fp4(scaled_tensor) + elif self.is_nvfp4_e5m3: + scaled_tensor = tensor / scales.reshape(tensor.shape[0], -1) + scaled_tensor.clamp_(-6.0, 6.0) + scaled_tensor = cast_to_fp4(scaled_tensor) else: scaled_tensor = tensor / (2 ** scales.reshape(tensor.shape[0], -1)) scaled_tensor = revert_tensor_by_pad(scaled_tensor, orig_shape=orig_shape, pad_len=pad_len) if self.is_mx: final_scale = (scales + E8M0_EXPONENT_BIAS).clamp(0, E8M0_EXPONENT_NAN_VAL).to(torch.uint8) + elif self.is_nvfp4_e5m3: + final_scale = float_to_e5m3_frexp(scales.to(torch.float32)) else: final_scale = scales.to(torch.float8_e4m3fn) diff --git a/auto_round/export/export_to_llmcompressor/config.py b/auto_round/export/export_to_llmcompressor/config.py index 62017ec736..1f5090e6c2 100644 --- a/auto_round/export/export_to_llmcompressor/config.py +++ b/auto_round/export/export_to_llmcompressor/config.py @@ -102,3 +102,38 @@ def initialize_quantization(scheme, targets=["Linear"], config_groups=None, kv_c quantization_status=QuantizationStatus.COMPRESSED, ignore=ignore, ) + + +def initialize_nvfp4_e5m3_quantization(ignore=None): + """Build stable compressed-tensors metadata for global-scale-free NVFP4 E5M3.""" + + def quant_args(dynamic): + return { + "actorder": None, + "block_structure": None, + "dynamic": dynamic, + "group_size": 16, + "num_bits": 4, + "observer": "minmax", + "observer_kwargs": {}, + "strategy": "tensor_group", + "symmetric": True, + "type": "float", + } + + return { + "config_groups": { + "group_0": { + "input_activations": quant_args("local"), + "output_activations": None, + "targets": ["Linear"], + "weights": quant_args(False), + } + }, + "format": "nvfp4-e5m3-pack-quantized", + "global_compression_ratio": None, + "ignore": list(dict.fromkeys(ignore or ["lm_head"])), + "kv_cache_scheme": None, + "quant_method": "compressed-tensors", + "quantization_status": "compressed", + } diff --git a/auto_round/export/export_to_llmcompressor/export_to_fp.py b/auto_round/export/export_to_llmcompressor/export_to_fp.py index d43a0962e8..250a3b73ac 100644 --- a/auto_round/export/export_to_llmcompressor/export_to_fp.py +++ b/auto_round/export/export_to_llmcompressor/export_to_fp.py @@ -140,6 +140,8 @@ def _get_scheme(bits, data_type): return "MXFP4" if bits == 4 else "MXFP8" if is_nv_fp(data_type): return "NVFP4" + if data_type == "fp4_v2": + return "NVFP4_E5M3" return None @@ -149,6 +151,8 @@ def _get_group_format(bits, data_type): return "mxfp4-pack-quantized" if bits == 4 else "mxfp8-quantized" if is_nv_fp(data_type): return "nvfp4-pack-quantized" + if data_type == "fp4_v2": + return "nvfp4-e5m3-pack-quantized" return "float-quantized" @@ -339,6 +343,11 @@ def save_quantized_as_fp( static_kv_dtype=serialization_dict.get("static_kv_dtype", None), static_attention_dtype=serialization_dict.get("static_attention_dtype", None), ) + elif data_type == "fp4_v2": + from auto_round.export.export_to_llmcompressor.config import initialize_nvfp4_e5m3_quantization + + quantization_config = initialize_nvfp4_e5m3_quantization(ignore=ignore) + quantization_config["provider"] = "auto-round" else: scheme = _get_scheme(bits, data_type) if scheme is None: diff --git a/auto_round/export/formats/backends/autoround.py b/auto_round/export/formats/backends/autoround.py index 21ab8daaae..05a02771db 100644 --- a/auto_round/export/formats/backends/autoround.py +++ b/auto_round/export/formats/backends/autoround.py @@ -41,6 +41,7 @@ class AutoRoundFormat(OutputFormat): "MXFP4", "MXFP8", "NVFP4", + "NVFP4_E5M3", "FPW8A16", "W2A16G64", "W2A16G32", @@ -67,7 +68,7 @@ def __init__(self, format: str, scheme: QuantizationScheme, ctx: Any): ) if enable_awq: self.backend = AutoAWQFormat("auto_round:auto_awq", scheme, ctx) - elif scheme.is_nv_fp() or scheme.is_mx_fp(): + elif scheme.is_nv_fp() or scheme.is_mx_fp() or scheme.data_type == BackendDataType.NVFP4_E5M3.value: self.backend = AutoRoundFormat(scheme.data_type, scheme, ctx) elif scheme.is_mx_int() and scheme.bits == 4: # only add mx_int4 now self.backend = AutoRoundFormat(scheme.data_type, scheme, ctx) @@ -87,7 +88,9 @@ def __init__(self, format: str, scheme: QuantizationScheme, ctx: Any): elif not format.startswith("auto_round"): if format == "mlx": self.backend = MLXFormat("mlx", scheme, ctx) - elif format.upper() not in list(BackendDataType.__members__.keys()): + elif format.upper() not in list(BackendDataType.__members__.keys()) and format not in { + BackendDataType.NVFP4_E5M3.value + }: raise KeyError(f"Unsupported backend format auto_round:{format}, please check") else: self.output_format = f"auto_round:{format}" @@ -131,7 +134,11 @@ def pack_layer(self, layer_name, model, device=None, **kwargs): backend = self.get_backend_name() - if self.output_format in [ + if self.output_format == f"auto_round:{BackendDataType.NVFP4_E5M3.value}": + from auto_round.export.export_to_llmcompressor.export_to_fp import pack_layer + + pack_func = pack_layer + elif self.output_format in [ f"auto_round:{BackendDataType.NV_FP.value}", f"auto_round:{BackendDataType.MX_FP.value}", f"auto_round:{BackendDataType.MX_FP_RCEIL.value}", @@ -185,7 +192,12 @@ def save_quantized( **kwargs, ) backend = self.get_backend_name() - if re.search(f"{BackendDataType.MX_FP.value}|{BackendDataType.NV_FP.value}", backend): + if backend == f"auto_round:{BackendDataType.NVFP4_E5M3.value}": + from auto_round.export.export_to_autoround.export_to_nvfp_mx import save_quantized_as_fp + + backend = "auto_round:llm_compressor_nvfp4_e5m3" + export_func = save_quantized_as_fp + elif re.search(f"{BackendDataType.MX_FP.value}|{BackendDataType.NV_FP.value}", backend): from auto_round.export.export_to_autoround.export_to_nvfp_mx import save_quantized_as_fp backend = "auto_round:llm_compressor" diff --git a/auto_round/export/formats/backends/fake.py b/auto_round/export/formats/backends/fake.py index 6c1e9a7030..c57fbe8fe4 100644 --- a/auto_round/export/formats/backends/fake.py +++ b/auto_round/export/formats/backends/fake.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy from typing import Any, Callable, Union import torch @@ -22,6 +23,16 @@ from auto_round.utils import copy_python_files_from_model_cache, unsupported_meta_device +def _serialize_quantization_config_value(value): + if isinstance(value, torch.dtype): + return str(value) + if isinstance(value, dict): + return {key: _serialize_quantization_config_value(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_serialize_quantization_config_value(item) for item in value] + return value + + @OutputFormat.register("fake") class FakeFormat(OutputFormat): support_schemes = None @@ -47,7 +58,44 @@ def save_quantized( serialization_dict: dict = None, **kwargs, ): - if not unsupported_meta_device(model): + has_meta_device = unsupported_meta_device(model) + if not inplace and not has_meta_device: + model = copy.deepcopy(model.to("cpu")) + + from auto_round.utils.model import set_module + from auto_round.wrapper import WrapperLinear, WrapperWALayer + + if not has_meta_device: + wrapped_modules = [ + (name, module) + for name, module in model.named_modules() + if name and isinstance(module, (WrapperLinear, WrapperWALayer)) + ] + wrapped_names = {name for name, _ in wrapped_modules} + for name, module in wrapped_modules: + if any(name.startswith(f"{parent}.") for parent in wrapped_names if parent != name): + continue + orig_layer = module.orig_layer + while isinstance(orig_layer, (WrapperLinear, WrapperWALayer)): + orig_layer = orig_layer.orig_layer + for attr_name in ("act_min_scale", "act_max_scale", "act_scale"): + orig_layer._parameters.pop(attr_name, None) + orig_layer._buffers.pop(attr_name, None) + if hasattr(orig_layer, attr_name): + delattr(orig_layer, attr_name) + set_module(model, name, orig_layer.to("cpu")) + + quantization_config = _serialize_quantization_config_value(dict(serialization_dict or {})) + quantization_config["quant_method"] = "auto-round" + quantization_config["packing_format"] = "auto_round:fake" + quantization_config["block_name_to_quantize"] = quantization_config.pop("to_quant_block_names", None) + from auto_round.export.utils import filter_quantization_config + + filter_quantization_config(quantization_config) + if hasattr(model, "config") and model.config is not None: + model.config.quantization_config = quantization_config + + if not has_meta_device: model = model.to("cpu") model.save_pretrained(output_dir) elif hasattr(model, "config") and model.config is not None: diff --git a/auto_round/export/formats/backends/llm_compressor.py b/auto_round/export/formats/backends/llm_compressor.py index 52d18037be..d1b41476b2 100644 --- a/auto_round/export/formats/backends/llm_compressor.py +++ b/auto_round/export/formats/backends/llm_compressor.py @@ -29,6 +29,7 @@ class LLMCompressorFormat(OutputFormat): "MXFP4", "MXFP8", "NVFP4", + "NVFP4_E5M3", "FPW8A16", "FP8_STATIC", "INT8", @@ -49,7 +50,12 @@ def __init__(self, format: str, scheme: QuantizationScheme, ctx: Any): if re.search("^(auto_round:)?llm_compressor", format): self.output_format = format self.backend = None - if scheme.is_nv_fp() or scheme.is_mx_fp(): + if scheme.data_type == "fp4_v2": + from auto_round.export.export_to_llmcompressor import check_compressed_tensors_supported + + check_compressed_tensors_supported(raise_error=True) + self.output_format = "llm_compressor:fp4_v2" + elif scheme.is_nv_fp() or scheme.is_mx_fp(): from auto_round.export.export_to_llmcompressor import check_compressed_tensors_supported check_compressed_tensors_supported(raise_error=True) @@ -135,6 +141,8 @@ def check_and_reset_format( ) if scheme.act_bits <= 8 and (not scheme.is_act_standard_fp() or scheme.act_dynamic): + if scheme.act_data_type == "fp4_v2": + return None, scheme, layer_config, quant_block_list if (scheme.is_act_nv_fp() and "static_gs" in scheme.act_data_type) or scheme.is_act_mx_fp(): return None, scheme, layer_config, quant_block_list elif scheme.is_dynamic_afp8() and scheme.is_block_wfp8(): @@ -159,7 +167,7 @@ def check_and_reset_format( def pack_layer(self, layer_name, model, device=None, **kwargs): if self.backend is not None: return self.backend.pack_layer(layer_name, model, device=device, **kwargs) - if re.search(f"{BackendDataType.MX_FP.value}|{BackendDataType.NV_FP.value}", self.output_format): + if re.search(f"{BackendDataType.MX_FP.value}|{BackendDataType.NV_FP.value}|fp4_v2", self.output_format): from auto_round.export.export_to_llmcompressor.export_to_fp import pack_layer return pack_layer(layer_name, model, device=device) @@ -195,7 +203,7 @@ def save_quantized( **kwargs, ) -> torch.nn.Module: backend = self.get_backend_name() - if re.search(f"{BackendDataType.MX_FP.value}|{BackendDataType.NV_FP.value}", backend): + if re.search(f"{BackendDataType.MX_FP.value}|{BackendDataType.NV_FP.value}|fp4_v2", backend): from auto_round.export.export_to_llmcompressor.export_to_fp import save_quantized_as_fp export_func = save_quantized_as_fp diff --git a/auto_round/export/formats/base.py b/auto_round/export/formats/base.py index 80636f99b2..b1b59cf3a7 100644 --- a/auto_round/export/formats/base.py +++ b/auto_round/export/formats/base.py @@ -54,6 +54,7 @@ class BackendDataType(str, Enum): MXFP8 = "mxfp8" MXFP4 = "mxfp4" NVFP4 = "nvfp4" + NVFP4_E5M3 = "fp4_v2" FP8 = "fp8" MX_FP = "mx_fp" NV_FP = "nv_fp" diff --git a/auto_round/inference/backend.py b/auto_round/inference/backend.py index ade2ca1ce7..00b5bcae46 100644 --- a/auto_round/inference/backend.py +++ b/auto_round/inference/backend.py @@ -194,10 +194,25 @@ def fp8_static_scheme_checker( return config == FP8_STATIC +def cute_nvfp4_e5m3_checker(in_feature: int, out_feature: int, config: QuantizationScheme) -> bool: + """Require CuTe DSL before selecting the optional NVFP4 E5M3 CUDA path.""" + del out_feature + if not in_feature_checker_group_size(in_feature, 0, config): + return False + try: + from auto_round_extension.cuda.cute_nvfp4_e5m3 import is_cute_dsl_available + + return torch.cuda.is_available() and is_cute_dsl_available() and torch.cuda.get_device_capability()[0] >= 8 + except (ImportError, RuntimeError): + return False + + GPTQ_FORMAT = ["auto_round:auto_gptq"] # zp+-1 GPTQ_FORMAT_NO_ZP = ["auto_round", "auto_round:gptqmodel"] AWQ_FORMAT = ["auto_round:auto_awq"] LLM_COMPRESSOR_FORMAT = ["auto_round:llm_compressor"] +FAKE_FORMAT = ["auto_round:fake"] +NVFP4_E5M3_LLM_COMPRESSOR_FORMAT = ["auto_round:llm_compressor_nvfp4_e5m3"] WOQ_DEFAULT_ACT_BITS = [None, 16, 32] # CPU backends that target Intel/x86 (ark / auto_round_kernel) cannot @@ -333,6 +348,24 @@ def fp8_static_scheme_checker( # NVFP4 +BackendInfos["auto_round:fake"] = BackendInfo( + device=["xpu", "cuda", "cpu"], + packing_format=FAKE_FORMAT, + sym=[True], + compute_dtype=["float32", "float16", "bfloat16"], + data_type=["fp4_v2"], + group_size=[16], + bits=[4], + act_bits=[4], + act_group_size=[16], + act_sym=[True], + act_data_type=["fp4_v2"], + priority=0, + checkers=[mxfp_nvfp_feature_checker], + alias=["auto_round", "torch"], + requirements=["auto-round>0.12.0"], +) + BackendInfos["auto_round:torch_nvfp4"] = BackendInfo( device=["xpu", "cuda", "cpu"], packing_format=LLM_COMPRESSOR_FORMAT, @@ -352,6 +385,44 @@ def fp8_static_scheme_checker( requirements=["auto-round>0.7.0"], ) +BackendInfos["auto_round:torch_nvfp4_e5m3"] = BackendInfo( + device=["xpu", "cuda", "cpu"], + packing_format=NVFP4_E5M3_LLM_COMPRESSOR_FORMAT, + sym=[True], + compute_dtype=["float32", "float16", "bfloat16"], + data_type=["fp4_v2"], + group_size=[16], + bits=[4], + act_bits=[4], + act_group_size=[16], + act_sym=[True], + act_data_type=["fp4_v2"], + act_dynamic=[True], + priority=3, + checkers=[mxfp_nvfp_feature_checker], + alias=["auto_round", "torch"], + requirements=["auto-round>0.12.0"], +) + +BackendInfos["auto_round:cute_nvfp4_e5m3"] = BackendInfo( + device=["cuda"], + packing_format=NVFP4_E5M3_LLM_COMPRESSOR_FORMAT, + sym=[True], + compute_dtype=["float32", "float16", "bfloat16"], + data_type=["fp4_v2"], + group_size=[16], + bits=[4], + act_bits=[4], + act_group_size=[16], + act_sym=[True], + act_data_type=["fp4_v2"], + act_dynamic=[True], + priority=6, + checkers=[cute_nvfp4_e5m3_checker], + alias=["cute_nvfp4_e5m3"], + requirements=["auto-round>0.12.0"], +) + BackendInfos["auto_round:tritonv2"] = BackendInfo( device=["cuda", "xpu"], data_type=["int"], @@ -779,8 +850,14 @@ def dynamic_import_inference_linear(backend, config, packing_format=None): return ar_qmodules.MXINT4QuantLinear if "torch_mxfp4" in backend: return ar_qmodules.MXFP4QuantLinear + if backend == "auto_round:cute_nvfp4_e5m3": + return ar_qmodules.CuteNVFP4E5M3QuantLinear + if backend == "auto_round:torch_nvfp4_e5m3": + return ar_qmodules.NVFP4E5M3QuantLinear if "torch_nvfp4" in backend: return ar_qmodules.NVFP4QuantLinear + if "auto_round:fake" in backend: + return ar_qmodules.FakeActQuantLinear if "auto_round_kernel" in backend or "ark" in backend: try: diff --git a/auto_round/inference/convert_model.py b/auto_round/inference/convert_model.py index b04c1bbaa1..665c072f1f 100644 --- a/auto_round/inference/convert_model.py +++ b/auto_round/inference/convert_model.py @@ -13,6 +13,7 @@ # limitations under the License. import os import re +from types import SimpleNamespace from typing import Union import torch @@ -70,7 +71,7 @@ def skip_not_convert_modules(model, quantization_config, layer_names, layer_conf if modules_to_not_convert: for layer_name in layer_names: if any([re.search(re.compile(n), layer_name) for n in modules_to_not_convert]): - layer_configs[layer_name] = {"bits": 16} + layer_configs[layer_name] = {"bits": 16, "act_bits": 16} return layer_configs @@ -564,6 +565,9 @@ def _create_quant_layer(layer, layer_backend, config, in_features, out_features, or BackendDataType.MXFP4.value in layer_backend or BackendDataType.NVFP4.value in layer_backend or BackendDataType.MXINT4.value in layer_backend + or layer_backend == "auto_round:torch_nvfp4_e5m3" + or layer_backend == "auto_round:cute_nvfp4_e5m3" + or layer_backend == "auto_round:fake" ): return QuantLinear.from_original(config, layer) @@ -832,6 +836,30 @@ def convert_hf_model(model: nn.Module, target_device: str = "cpu") -> tuple[nn.M if is_transformers_version_greater_or_equal_5(): disable_moe_conversion_mapping(model) quantization_config = model.config.quantization_config + config_format = ( + quantization_config.get("format") + if isinstance(quantization_config, dict) + else getattr(quantization_config, "format", None) + ) + if config_format == "nvfp4-e5m3-pack-quantized": + config_dict = quantization_config if isinstance(quantization_config, dict) else quantization_config.to_dict() + ignored = config_dict.get("ignore") or [] + # TODO: For experimental purpose, will delete when it's not necessary. + quantization_config = SimpleNamespace( + quant_method="auto-round", + packing_format="auto_round:llm_compressor_nvfp4_e5m3", + bits=4, + group_size=16, + sym=True, + data_type="fp4_v2", + act_bits=4, + act_group_size=16, + act_sym=True, + act_data_type="fp4_v2", + act_dynamic=True, + extra_config={name: {"bits": 16, "act_bits": 16} for name in ignored}, + ) + model.config.quantization_config = quantization_config # Check desc_act + static_groups if getattr(quantization_config, "desc_act", False): @@ -872,6 +900,12 @@ def convert_hf_model(model: nn.Module, target_device: str = "cpu") -> tuple[nn.M # Replace layers with quantized versions layer_configs = get_layer_config(model, quantization_config) used_backends = _replace_by_quant_layers(model, layer_configs, backend, target_device, packing_format) + logger.info( + "Inference backend selection: requested=%s, packing_format=%s, selected=%s", + backend, + packing_format, + ", ".join(used_backends), + ) # Apply rotation hooks (hadamard, spinquant, quarot, etc.) via unified dispatch. _has_rotation = getattr(quantization_config, "rotation_config", None) or getattr( diff --git a/auto_round/schemes.py b/auto_round/schemes.py index 31da6d95c4..ea2a05f813 100644 --- a/auto_round/schemes.py +++ b/auto_round/schemes.py @@ -706,6 +706,18 @@ def parse_scheme( } ) +NVFP4_E5M3 = QuantizationScheme.from_dict( + { + "bits": 4, + "group_size": 16, + "data_type": "fp4_v2", + "act_bits": 4, + "act_data_type": "fp4_v2", + "act_group_size": 16, + "act_sym": True, + } +) + FPW8A16 = QuantizationScheme.from_dict( { "bits": 8, @@ -802,6 +814,7 @@ def parse_scheme( "MXFP8": MXFP8, "MXFP8_RCEIL": MXFP8_RCEIL, "NVFP4": NVFP4, + "NVFP4_E5M3": NVFP4_E5M3, "FPW8A16": FPW8A16, "W2A16G64": W2A16G64, "W2A16G32": W2A16G32, diff --git a/auto_round/utils/common.py b/auto_round/utils/common.py index 4af0291a2c..764ab2eaf9 100644 --- a/auto_round/utils/common.py +++ b/auto_round/utils/common.py @@ -353,12 +353,50 @@ def monkey_patch_transformers(): _patch_tensor_get_dtype_for_prequantized_loading() _patch_default_rope_init() _patch_rotary_embedding_init_for_legacy_remote_code() + _patch_nvfp4_e5m3_compressed_tensors_quantizer() if parsed_version >= version.parse("4.56.0"): _patch_classmethod_kwargs(transformers.AutoModelForCausalLM, "from_pretrained", torch_dtype="dtype") else: _patch_classmethod_kwargs(transformers.AutoModelForCausalLM, "from_pretrained", dtype="torch_dtype") +def _patch_nvfp4_e5m3_compressed_tensors_quantizer(): + """Route the global-scale-free E5M3 format through AutoRound qmodules.""" + try: + from transformers.quantizers.quantizer_compressed_tensors import CompressedTensorsHfQuantizer + except ImportError: + return + + original_before = CompressedTensorsHfQuantizer._process_model_before_weight_loading + if getattr(original_before, "_auto_round_nvfp4_e5m3_patch", False): + return + original_after = CompressedTensorsHfQuantizer._process_model_after_weight_loading + + def is_nvfp4_e5m3(self): + compression_config = getattr(getattr(self, "compressor", None), "quantization_config", None) + return getattr(compression_config, "format", None) == "nvfp4-e5m3-pack-quantized" + + def patched_before(self, model, **kwargs): + if not is_nvfp4_e5m3(self): + return original_before(self, model, **kwargs) + from auto_round.inference.convert_model import convert_hf_model + + model.config.quantization_config = self.compressor.quantization_config + target_device = "cuda" if torch.cuda.is_available() else "cpu" + model, _ = convert_hf_model(model, target_device=target_device) + self.run_compressed = True + return model + + def patched_after(self, model, **kwargs): + if is_nvfp4_e5m3(self): + return model + return original_after(self, model, **kwargs) + + patched_before._auto_round_nvfp4_e5m3_patch = True + CompressedTensorsHfQuantizer._process_model_before_weight_loading = patched_before + CompressedTensorsHfQuantizer._process_model_after_weight_loading = patched_after + + @lru_cache(None) def monkey_patch(): monkey_patch_transformers() diff --git a/auto_round_extension/cuda/cute_nvfp4_e5m3.py b/auto_round_extension/cuda/cute_nvfp4_e5m3.py new file mode 100644 index 0000000000..2a7c0f994f --- /dev/null +++ b/auto_round_extension/cuda/cute_nvfp4_e5m3.py @@ -0,0 +1,270 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +"""Optional CuTe DSL dispatch for NVFP4 E5M3 activation QDQ.""" + +from functools import lru_cache +from importlib.util import find_spec +from typing import Optional + +import torch + +from auto_round.logger import logger + +_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32) +_GROUP_SIZE = 16 +_THREADS_PER_BLOCK = 256 +_FAILED_KERNEL_KEYS = set() +_FAILED_WEIGHT_DQ_KERNEL_KEYS = set() + + +@lru_cache(maxsize=1) +def is_cute_dsl_available() -> bool: + """Return whether the optional NVIDIA CuTe DSL package is installed.""" + return find_spec("cutlass") is not None + + +def can_use_cute_fp4_v2_qdq(activation: torch.Tensor, group_size: int) -> bool: + """Check whether an activation can use the CuTe QDQ kernel.""" + if not is_cute_dsl_available() or not activation.is_cuda: + return False + if group_size != _GROUP_SIZE or activation.dtype not in _SUPPORTED_DTYPES: + return False + if not activation.is_contiguous() or activation.shape[-1] % group_size: + return False + return torch.cuda.get_device_capability(activation.device)[0] >= 8 + + +def _make_qdq_kernel(): + import cutlass + import cutlass.cute as cute + from cutlass._mlir_helpers import math + + @cute.kernel + def qdq_kernel(input_tensor: cute.Tensor, output_tensor: cute.Tensor): + thread_idx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + group_idx = block_idx * _THREADS_PER_BLOCK + thread_idx + + if group_idx < input_tensor.shape[0]: + amax = cutlass.Float32(0.0) + for column_idx in range(_GROUP_SIZE): + value = cutlass.Float32(input_tensor[group_idx, column_idx]) + amax = cute.arch.fmax(amax, math.abs(value)) + + # This reproduces the E5M3 scale domain used by fp4_v2 without + # materializing scale or intermediate FP32 tensors in global memory. + scale = amax / cutlass.Float32(6.0) + if scale > cutlass.Float32(0.0): + exponent = math.floor(math.log2(scale)) + cutlass.Float32(1.0) + mantissa = scale / math.exp2(exponent) + mantissa_bits = math.roundeven((mantissa - cutlass.Float32(0.5)) * cutlass.Float32(16.0)) + mantissa_bits = cute.arch.fmin(cute.arch.fmax(mantissa_bits, 0.0), 7.0) + scale = (cutlass.Float32(1.0) + mantissa_bits / cutlass.Float32(8.0)) * math.exp2( + exponent - cutlass.Float32(1.0) + ) + + for column_idx in range(_GROUP_SIZE): + value = cutlass.Float32(input_tensor[group_idx, column_idx]) + scaled = value / scale if scale > cutlass.Float32(0.0) else cutlass.Float32(0.0) + magnitude = math.abs(scaled) + quantized = cutlass.Float32(0.0) + if magnitude < cutlass.Float32(2.0): + quantized = math.roundeven(magnitude * cutlass.Float32(2.0)) / cutlass.Float32(2.0) + elif magnitude < cutlass.Float32(4.0): + quantized = math.roundeven(magnitude) + else: + quantized = cutlass.Float32(2.0) * math.roundeven(magnitude / cutlass.Float32(2.0)) + quantized = cute.arch.fmin(cute.arch.fmax(quantized, 0.0), 6.0) + if scaled < cutlass.Float32(0.0): + quantized = -quantized + output_tensor[group_idx, column_idx] = (quantized * scale).to(output_tensor.element_type) + + @cute.jit + def launch_qdq(input_tensor: cute.Tensor, output_tensor: cute.Tensor): + groups = input_tensor.shape[0] + qdq_kernel(input_tensor, output_tensor).launch( + grid=((groups + _THREADS_PER_BLOCK - 1) // _THREADS_PER_BLOCK, 1, 1), + block=(_THREADS_PER_BLOCK, 1, 1), + ) + + return launch_qdq + + +def _make_weight_dq_kernel(): + import cutlass + import cutlass.cute as cute + from cutlass._mlir_helpers import math + + @cute.kernel + def weight_dq_kernel(weight_packed: cute.Tensor, weight_scale: cute.Tensor, output_tensor: cute.Tensor): + thread_idx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + group_idx = block_idx * _THREADS_PER_BLOCK + thread_idx + + if group_idx < output_tensor.shape[0]: + scale_byte = cutlass.Float32(weight_scale[group_idx, 0]) + exponent = math.floor(scale_byte / cutlass.Float32(8.0)) + mantissa = scale_byte - exponent * cutlass.Float32(8.0) + scale = (cutlass.Float32(1.0) + mantissa / cutlass.Float32(8.0)) * math.exp2( + exponent - cutlass.Float32(15.0) + ) + if exponent == cutlass.Float32(0.0): + scale = mantissa * math.exp2(cutlass.Float32(-17.0)) + + for column_idx in range(_GROUP_SIZE): + packed = cutlass.Float32(weight_packed[group_idx, column_idx // 2]) + packed_high = math.floor(packed / cutlass.Float32(16.0)) + nibble = packed - packed_high * cutlass.Float32(16.0) + if column_idx % 2: + nibble = packed_high + + sign = cutlass.Float32(1.0) + if nibble >= cutlass.Float32(8.0): + sign = cutlass.Float32(-1.0) + nibble = nibble - cutlass.Float32(8.0) + + value = cutlass.Float32(0.0) + if nibble == cutlass.Float32(1.0): + value = cutlass.Float32(0.5) + elif nibble == cutlass.Float32(2.0): + value = cutlass.Float32(1.0) + elif nibble == cutlass.Float32(3.0): + value = cutlass.Float32(1.5) + elif nibble == cutlass.Float32(4.0): + value = cutlass.Float32(2.0) + elif nibble == cutlass.Float32(5.0): + value = cutlass.Float32(3.0) + elif nibble == cutlass.Float32(6.0): + value = cutlass.Float32(4.0) + elif nibble == cutlass.Float32(7.0): + value = cutlass.Float32(6.0) + + output_tensor[group_idx, column_idx] = (sign * value * scale).to(output_tensor.element_type) + + @cute.jit + def launch_weight_dq(weight_packed: cute.Tensor, weight_scale: cute.Tensor, output_tensor: cute.Tensor): + groups = output_tensor.shape[0] + weight_dq_kernel(weight_packed, weight_scale, output_tensor).launch( + grid=((groups + _THREADS_PER_BLOCK - 1) // _THREADS_PER_BLOCK, 1, 1), + block=(_THREADS_PER_BLOCK, 1, 1), + ) + + return launch_weight_dq + + +@lru_cache(maxsize=None) +def _get_compiled_qdq_kernel(device_index: int, dtype: torch.dtype): + import cutlass.cute as cute + from cutlass.cute.runtime import from_dlpack + + input_tensor = torch.empty((1, _GROUP_SIZE), device=f"cuda:{device_index}", dtype=dtype) + output_tensor = torch.empty_like(input_tensor) + return cute.compile( + _make_qdq_kernel(), + from_dlpack(input_tensor).mark_layout_dynamic(), + from_dlpack(output_tensor).mark_layout_dynamic(), + ) + + +@lru_cache(maxsize=None) +def _get_compiled_weight_dq_kernel(device_index: int, dtype: torch.dtype): + import cutlass.cute as cute + from cutlass.cute.runtime import from_dlpack + + weight_packed = torch.empty((2, _GROUP_SIZE // 2), device=f"cuda:{device_index}", dtype=torch.uint8) + weight_scale = torch.empty((2, 1), device=f"cuda:{device_index}", dtype=torch.uint8) + output_tensor = torch.empty((2, _GROUP_SIZE), device=f"cuda:{device_index}", dtype=dtype) + return cute.compile( + _make_weight_dq_kernel(), + from_dlpack(weight_packed).mark_layout_dynamic(), + from_dlpack(weight_scale).mark_layout_dynamic(), + from_dlpack(output_tensor).mark_layout_dynamic(), + ) + + +def try_cute_fp4_v2_qdq(activation: torch.Tensor, group_size: int) -> Optional[torch.Tensor]: + """Run a CuTe DSL group-size-16 FP4 QDQ kernel when eligible.""" + if not can_use_cute_fp4_v2_qdq(activation, group_size): + return None + if torch.cuda.is_current_stream_capturing(): + return None + + kernel_key = (activation.device.index, activation.dtype) + if kernel_key in _FAILED_KERNEL_KEYS: + return None + + try: + groups = activation.numel() // group_size + output = torch.empty_like(activation) + compiled_qdq = _get_compiled_qdq_kernel(*kernel_key) + compiled_qdq(activation.reshape(groups, group_size), output.reshape(groups, group_size)) + return output + except Exception as error: + _FAILED_KERNEL_KEYS.add(kernel_key) + logger.warning_once("CuTe NVFP4 E5M3 QDQ failed; falling back to PyTorch reference: %s", error) + return None + + +def try_cute_nvfp4_e5m3_weight_dq( + weight_packed: torch.Tensor, weight_scale: torch.Tensor, dtype: torch.dtype +) -> Optional[torch.Tensor]: + """Dequantize packed FP4 E5M3 weights with CuTe.""" + if ( + not is_cute_dsl_available() + or not weight_packed.is_cuda + or weight_packed.dtype != torch.uint8 + or weight_scale.dtype != torch.uint8 + or weight_packed.ndim != 2 + or weight_scale.shape != (weight_packed.shape[0], weight_packed.shape[1] // (_GROUP_SIZE // 2)) + or dtype not in _SUPPORTED_DTYPES + or torch.cuda.get_device_capability(weight_packed.device)[0] < 8 + or torch.cuda.is_current_stream_capturing() + ): + return None + + kernel_key = (weight_packed.device.index, dtype) + if kernel_key in _FAILED_WEIGHT_DQ_KERNEL_KEYS: + return None + + try: + out_features = weight_packed.shape[0] + in_features = weight_packed.shape[1] * 2 + groups = weight_scale.numel() + packed_groups = weight_packed.reshape(groups, _GROUP_SIZE // 2) + scale_groups = weight_scale.reshape(groups, 1) + output = torch.empty((groups, _GROUP_SIZE), device=weight_packed.device, dtype=dtype) + compiled_weight_dq = _get_compiled_weight_dq_kernel(*kernel_key) + compiled_weight_dq(packed_groups, scale_groups, output) + return output.reshape(out_features, in_features) + except Exception as error: + _FAILED_WEIGHT_DQ_KERNEL_KEYS.add(kernel_key) + logger.warning_once("CuTe NVFP4 E5M3 weight DQ failed; falling back to PyTorch reference: %s", error) + return None + + +def try_cute_nvfp4_e5m3_linear( + activation: torch.Tensor, + weight_packed: torch.Tensor, + weight_scale: torch.Tensor, + bias: Optional[torch.Tensor], +) -> Optional[torch.Tensor]: + """Reserved second-stage dispatch point for fused QDQ, unpack, and GEMM. + + Returning ``None`` keeps the reference Linear path active until the packed-weight + mainloop has been validated against the existing E5M3 checkpoint format. + """ + del activation, weight_packed, weight_scale, bias + fused_output: Optional[torch.Tensor] = None + return fused_output diff --git a/docs/environments.md b/docs/environments.md index a61492314e..c5839ca457 100644 --- a/docs/environments.md +++ b/docs/environments.md @@ -161,6 +161,16 @@ export AR_AUTO_SCHEME_CACHE=/path/to/auto_scheme_cache export AR_ENABLE_AUTO_SCHEME_PARALLEL=0 ``` +### AR_NVFP4_E5M3_CACHE_HP_WEIGHT +- **Description**: Controls whether `NVFP4E5M3QuantLinear` caches a dequantized high-precision weight after the first forward pass, instead of dequantizing the packed FP4 weight on every call. +- **Default**: `False` (equivalent to `"0"`) +- **Valid Values**: `"1"`, `"true"`, `"yes"`, `"on"` (case-insensitive) enable caching; any other value disables caching +- **Usage**: Enable this when repeated inference throughput matters more than memory footprint. The current implementation releases `weight_packed` and `weight_scale` after materializing the cached high-precision weight, so steady-state memory usage increases and the cache cannot be cleared back to packed storage. + +```bash +export AR_NVFP4_E5M3_CACHE_HP_WEIGHT=1 +``` + ### AR_DISK_STREAM_MODEL - **Description**: When enabled, `AutoRound(model=, ...)` builds the model as a meta-device skeleton instead of fully materializing the checkpoint on CPU RAM up front, and streams each decoder block's real weights from the checkpoint's safetensors shards on demand -- materializing right before a block is used (calibration, tuning, or `AutoScheme` sensitivity scoring) and freeing it back to meta right after. This keeps peak CPU RAM roughly flat regardless of checkpoint size, instead of proportional to it. Non-block parameters (embeddings, `lm_head`, final norm) are still loaded up front, since they are typically small. Text-model AutoScheme scoring also supports combining this with parallel scoring, which is enabled by default; each worker streams its own block copy. - **Default**: `False` diff --git a/docs/environments_CN.md b/docs/environments_CN.md index 0272dbd1fe..75019764ea 100644 --- a/docs/environments_CN.md +++ b/docs/environments_CN.md @@ -161,6 +161,16 @@ export AR_AUTO_SCHEME_CACHE=/path/to/auto_scheme_cache export AR_ENABLE_AUTO_SCHEME_PARALLEL=0 ``` +### AR_NVFP4_E5M3_CACHE_HP_WEIGHT +- **描述**:控制 `NVFP4E5M3QuantLinear` 是否在首次前向后缓存解量化得到的高精度权重,而不是每次调用都从打包的 FP4 权重重新解量化。 +- **默认值**:`False`(等价于 `"0"`) +- **有效值**:`"1"`、`"true"`、`"yes"`、`"on"`(不区分大小写)表示启用缓存;其他值表示禁用缓存 +- **用途**:当重复推理吞吐比内存占用更重要时可启用。当前实现会在缓存高精度权重后释放 `weight_packed` 和 `weight_scale`,因此稳态内存占用会增大,且之后无法再切回打包存储。 + +```bash +export AR_NVFP4_E5M3_CACHE_HP_WEIGHT=1 +``` + ### AR_DISK_STREAM_MODEL - **描述**:启用后,`AutoRound(model=, ...)` 会将模型构建为 meta 设备骨架,而不是先把整个 checkpoint 完全加载到 CPU 内存;随后按需从 checkpoint 的 safetensors 分片中流式加载每个解码器块的真实权重——在该块被使用前(校准、调优或 `AutoScheme` 敏感度评分)才实体化,用完后立即释放回 meta。这样峰值 CPU 内存基本保持平稳,而不会随 checkpoint 大小成比例增长。非块参数(embedding、`lm_head`、最终归一化层)体积通常较小,仍会一次性加载。文本模型的 AutoScheme 评分也支持与默认启用的并行评分组合使用;每个 worker 会流式加载自己的 block 副本。 - **默认值**:`False` diff --git a/test/unit/test_cpu/core/test_format_decoupling.py b/test/unit/test_cpu/core/test_format_decoupling.py index b9dac1280c..d38dad7238 100644 --- a/test/unit/test_cpu/core/test_format_decoupling.py +++ b/test/unit/test_cpu/core/test_format_decoupling.py @@ -32,7 +32,9 @@ from auto_round.algorithms.quantization.config import QuantizationConfig from auto_round.compressors.base import BaseCompressor from auto_round.compressors.config_resolution import FormatResolution, ResolvedScheme +from auto_round.experimental import qmodules as ar_qmodules from auto_round.export.formats import resolve_formats +from auto_round.inference.backend import BackendInfos, dynamic_import_inference_linear, get_layer_backend from auto_round.schemes import QuantizationScheme, parse_scheme @@ -50,6 +52,7 @@ def _resolved_scheme(scheme_name): ("FP8_STATIC", "llm_compressor"): ["llm_compressor:fp8_static"], ("MXFP4", "auto_round"): ["auto_round:mx_fp"], ("NVFP4", "auto_round"): ["auto_round:nv_fp"], + ("NVFP4_E5M3", "auto_round"): ["auto_round:fp4_v2"], } @@ -67,6 +70,30 @@ def test_format_selection_baseline(): assert results == EXPECTED_FORMAT_SELECTION_BASELINE +def test_nvfp4_e5m3_autoround_uses_llm_compressor_packing_and_torch_fallback(): + scheme = _resolved_scheme("NVFP4_E5M3") + resolution = resolve_formats( + ResolvedScheme.from_scheme(scheme), format="auto_round", model=nn.Sequential(nn.Linear(16, 16)) + ) + + assert resolution.formats[0].get_backend_name() == "auto_round:fp4_v2" + assert BackendInfos["auto_round:torch_nvfp4_e5m3"].packing_format == ["auto_round:llm_compressor_nvfp4_e5m3"] + assert BackendInfos["auto_round:cute_nvfp4_e5m3"].priority > BackendInfos["auto_round:torch_nvfp4_e5m3"].priority + assert dynamic_import_inference_linear("auto_round:torch_nvfp4_e5m3", scheme) is ar_qmodules.NVFP4E5M3QuantLinear + assert dynamic_import_inference_linear("auto_round:cute_nvfp4_e5m3", scheme) is ar_qmodules.CuteNVFP4E5M3QuantLinear + assert ( + get_layer_backend( + "cpu", + "auto", + "auto_round:llm_compressor_nvfp4_e5m3", + scheme, + in_features=16, + out_features=16, + ) + == "auto_round:torch_nvfp4_e5m3" + ) + + def test_gguf_correction_propagates_to_scheme_and_quantize_config(monkeypatch): # W4A16's plain int-woq fields don't match every field gguf_args_check enforces # for a GGUF target; after resolving "gguf:q4_k_m" the scheme-bearing objects diff --git a/test/unit/test_cpu/export/test_fake_format.py b/test/unit/test_cpu/export/test_fake_format.py new file mode 100644 index 0000000000..b4b377ec6b --- /dev/null +++ b/test/unit/test_cpu/export/test_fake_format.py @@ -0,0 +1,188 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +import json +import os +from types import SimpleNamespace + +import torch +from transformers import AutoModelForCausalLM, OPTConfig, OPTForCausalLM + +from auto_round.experimental.qmodules.fake import FakeActQuantLinear +from auto_round.formats import FakeFormat +from auto_round.inference.convert_model import convert_hf_model +from auto_round.schemes import PRESET_SCHEMES +from auto_round.wrapper import WrapperWALayer + + +class _WrappedLinear(WrapperWALayer): + def __init__(self, linear): + torch.nn.Module.__init__(self) + self.orig_layer = linear + self.register_buffer("act_max_scale", torch.ones(1)) + + def forward(self, inputs): + return self.orig_layer(inputs) + + +class _SaveableModel(torch.nn.Module): + def __init__(self): + super().__init__() + linear = torch.nn.Linear(4, 3) + linear.register_parameter("act_max_scale", torch.nn.Parameter(torch.ones(1))) + self.linear = _WrappedLinear(linear) + self.config = SimpleNamespace() + + def save_pretrained(self, output_dir): + os.makedirs(output_dir, exist_ok=True) + torch.save(self.state_dict(), os.path.join(output_dir, "pytorch_model.bin")) + with open(os.path.join(output_dir, "config.json"), "w") as config_file: + json.dump({"quantization_config": self.config.quantization_config}, config_file) + + +def test_fake_format_unwraps_quantized_layers_before_save(tmp_path): + model = _SaveableModel() + expected_weight = model.linear.orig_layer.weight.detach().clone() + output_dir = str(tmp_path / "fake_model") + + FakeFormat("fake", PRESET_SCHEMES["NVFP4_E5M3"], SimpleNamespace(mllm=False)).save_quantized( + output_dir=output_dir, + model=model, + inplace=False, + serialization_dict={ + "bits": 4, + "group_size": 16, + "sym": True, + "data_type": "fp4_v2", + "act_bits": 4, + "act_group_size": 16, + "act_sym": True, + "act_data_type": "fp4_v2", + "supported_types": [torch.nn.Linear], + }, + ) + + state_dict = torch.load(os.path.join(output_dir, "pytorch_model.bin"), weights_only=True) + assert set(state_dict) == {"linear.weight", "linear.bias"} + assert torch.equal(state_dict["linear.weight"], expected_weight) + assert hasattr(model.linear, "orig_layer") + with open(os.path.join(output_dir, "config.json")) as config_file: + quantization_config = json.load(config_file)["quantization_config"] + assert "supported_types" not in quantization_config + assert quantization_config["packing_format"] == "auto_round:fake" + + +class _TinyLoadModel(torch.nn.Module): + def __init__(self, quantization_config): + super().__init__() + self.block = torch.nn.Module() + self.block.linear = torch.nn.Linear(16, 4) + self.lm_head = torch.nn.Linear(16, 4) + self.config = SimpleNamespace(quantization_config=quantization_config) + + +def test_fake_config_replaces_linear_and_qdq_activation_on_load(): + quantization_config = SimpleNamespace( + bits=4, + group_size=16, + sym=True, + data_type="fp4_v2", + act_bits=4, + act_group_size=16, + act_sym=True, + act_data_type="fp4_v2", + act_dynamic=True, + quant_method="auto-round", + packing_format="auto_round:fake", + block_name_to_quantize="block", + backend="auto", + extra_config={}, + modules_to_not_convert=[], + ) + model = _TinyLoadModel(quantization_config) + original_weight = model.block.linear.weight.detach().clone() + activation = torch.randn(2, 3, 16) + + model, used_backends = convert_hf_model(model, target_device="cpu") + + assert used_backends == ["auto_round:fake"] + assert isinstance(model.block.linear, FakeActQuantLinear) + assert torch.equal(model.block.linear.weight, original_weight) + qdq_activation = model.block.linear.qdq_input(activation) + assert not torch.equal(qdq_activation, activation) + expected = torch.nn.functional.linear(qdq_activation, model.block.linear.weight, model.block.linear.bias) + assert torch.equal(model.block.linear(activation), expected) + + +def test_fake_config_keeps_modules_to_not_convert_in_full_precision(): + quantization_config = SimpleNamespace( + bits=4, + group_size=16, + sym=True, + data_type="fp4_v2", + act_bits=4, + act_group_size=16, + act_sym=True, + act_data_type="fp4_v2", + act_dynamic=True, + quant_method="auto-round", + packing_format="auto_round:fake", + block_name_to_quantize="", + backend="auto", + extra_config={}, + modules_to_not_convert=["lm_head"], + ) + model = _TinyLoadModel(quantization_config) + original_lm_head = model.lm_head + + model, used_backends = convert_hf_model(model, target_device="cpu") + + assert used_backends == ["auto_round:fake"] + assert isinstance(model.block.linear, FakeActQuantLinear) + assert model.lm_head is original_lm_head + + +def test_transformers_load_replaces_fake_linear(tmp_path): + config = OPTConfig( + vocab_size=32, + hidden_size=16, + ffn_dim=32, + num_hidden_layers=1, + num_attention_heads=2, + max_position_embeddings=32, + word_embed_proj_dim=16, + ) + config.quantization_config = { + "bits": 4, + "group_size": 16, + "sym": True, + "data_type": "fp4_v2", + "act_bits": 4, + "act_group_size": 16, + "act_sym": True, + "act_data_type": "fp4_v2", + "act_dynamic": True, + "quant_method": "auto-round", + "packing_format": "auto_round:fake", + "block_name_to_quantize": "model.decoder.layers", + } + model_dir = str(tmp_path / "fake_opt") + OPTForCausalLM(config).save_pretrained(model_dir) + + loaded_model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="cpu") + + q_proj = loaded_model.model.decoder.layers[0].self_attn.q_proj + assert isinstance(q_proj, FakeActQuantLinear) + activation = torch.randn(1, 2, 16) + assert not torch.equal(q_proj.qdq_input(activation), activation) diff --git a/test/unit/test_cpu/quantization/test_model_free.py b/test/unit/test_cpu/quantization/test_model_free.py index 46ea0d35e9..4a02c718c4 100644 --- a/test/unit/test_cpu/quantization/test_model_free.py +++ b/test/unit/test_cpu/quantization/test_model_free.py @@ -67,6 +67,7 @@ def fail_run(): _process_shard, _process_single_shard_task, _quantize_weight_mxfp, + _quantize_weight_nvfp4_e5m3, _validate_auto_scheme_options, get_predefined_ignore_layers_from_config, is_model_free_supported_scheme, @@ -306,6 +307,127 @@ def test_quantizes_eligible_weights(self, tmp_path): assert "layer.fc1" in quantized assert "layer.fc1.qweight" in output and "layer.fc1.bias" in output + +def test_nvfp4_e5m3_model_free_fake_quantization(): + weight = torch.randn(8, 32) + output = _quantize_weight_nvfp4_e5m3(weight, "layer.fc", group_size=16) + + assert set(output) == {"layer.fc.weight"} + assert output["layer.fc.weight"].shape == weight.shape + assert output["layer.fc.weight"].dtype == weight.dtype + assert not torch.equal(output["layer.fc.weight"], weight) + assert is_model_free_supported_scheme("NVFP4_E5M3") + assert not is_model_free_supported_scheme("UNVFP4") + assert not is_model_free_supported_scheme("NVFP4+") + + +def test_nvfp4_e5m3_model_free_end_to_end(tmp_path): + tensors = { + "model.layers.0.self_attn.q_proj.weight": torch.randn(32, 32), + "lm_head.weight": torch.randn(64, 32), + } + model_dir = _make_model_dir(tmp_path, _LLAMA_CFG, tensors) + output_dir = str(tmp_path / "output") + os.makedirs(output_dir) + with open(os.path.join(output_dir, "quantization_config.json"), "w") as f: + json.dump({"stale": True}, f) + + compressor = _ModelFreeCompressorCore(model_name_or_path=model_dir, output_dir=output_dir, scheme="NVFP4_E5M3") + compressor.run() + + output_keys = _read_output_keys(output_dir) + assert "model.layers.0.self_attn.q_proj.weight" not in output_keys + assert "model.layers.0.self_attn.q_proj.weight_packed" in output_keys + assert "model.layers.0.self_attn.q_proj.weight_scale" in output_keys + assert "lm_head.weight" in output_keys + assert compressor.format == "auto_round" + quantization_config = _read_qconfig(output_dir) + assert quantization_config["packing_format"] == "auto_round:llm_compressor_nvfp4_e5m3" + assert quantization_config["data_type"] == "fp4_v2" + assert quantization_config["act_bits"] == 4 + assert quantization_config["act_data_type"] == "fp4_v2" + assert quantization_config["act_group_size"] == 16 + assert quantization_config["act_sym"] is True + assert quantization_config["extra_config"]["lm_head"] == { + "bits": 16, + "data_type": "float", + "act_bits": 16, + "act_data_type": "float", + } + assert os.path.exists(os.path.join(output_dir, "quantization_config.json")) + + +def test_nvfp4_e5m3_model_free_llm_compressor(tmp_path): + tensors = { + "model.layers.0.self_attn.q_proj.weight": torch.randn(32, 32), + "lm_head.weight": torch.randn(64, 32), + } + model_dir = _make_model_dir(tmp_path, _LLAMA_CFG, tensors) + output_dir = str(tmp_path / "output") + + compressor = _ModelFreeCompressorCore( + model_name_or_path=model_dir, + output_dir=output_dir, + scheme="NVFP4_E5M3", + format="llm_compressor", + ) + compressor.run() + + output_keys = _read_output_keys(output_dir) + prefix = "model.layers.0.self_attn.q_proj" + assert f"{prefix}.weight_packed" in output_keys + assert f"{prefix}.weight_scale" in output_keys + assert f"{prefix}.weight" not in output_keys + assert f"{prefix}.weight_global_scale" not in output_keys + assert f"{prefix}.input_global_scale" not in output_keys + quantization_config = _read_qconfig(output_dir) + group = quantization_config["config_groups"]["group_0"] + assert quantization_config["format"] == "nvfp4-e5m3-pack-quantized" + assert quantization_config["quant_method"] == "compressed-tensors" + assert quantization_config["provider"] == "auto-round" + assert group["weights"]["group_size"] == 16 + assert group["input_activations"]["dynamic"] == "local" + + +def test_model_free_legacy_nvfp4_is_normalized_and_passthrough(tmp_path): + prefix = "model.layers.0.mlp.down_proj" + tensors = { + f"{prefix}.weight": torch.randint(0, 256, (32, 64), dtype=torch.uint8), + f"{prefix}.weight_scale": torch.randint(0, 256, (32, 4), dtype=torch.uint8), + f"{prefix}.weight_scale_2": torch.tensor([2.0], dtype=torch.float32), + f"{prefix}.input_scale": torch.tensor([4.0], dtype=torch.float32), + "model.layers.0.self_attn.q_proj.weight": torch.randn(32, 32), + } + shard_path = str(tmp_path / "shard.safetensors") + save_file(tensors, shard_path) + + layer_config = { + prefix: { + "bits": 4, + "group_size": 16, + "sym": True, + "data_type": "nv_fp", + } + } + output, quantized, ignored = _process_shard(shard_path, _DEFAULT_SCHEME, layer_config, []) + + # Legacy naming should be normalized to llm-compressor-style global-scale keys. + assert f"{prefix}.weight_packed" in output + assert f"{prefix}.weight_scale" in output + assert f"{prefix}.weight_global_scale" in output + assert f"{prefix}.input_global_scale" in output + assert f"{prefix}.weight" not in output + assert f"{prefix}.weight_scale_2" not in output + assert f"{prefix}.input_scale" not in output + assert torch.allclose(output[f"{prefix}.weight_global_scale"], torch.tensor([0.5], dtype=torch.float32)) + assert torch.allclose(output[f"{prefix}.input_global_scale"], torch.tensor([0.25], dtype=torch.float32)) + + # The NVFP4 layer is treated as already quantized (passthrough) while + # other Linear layers in the shard are still quantized by model-free RTN. + assert prefix in quantized + assert "model.layers.0.self_attn.q_proj" in quantized + assert prefix not in ignored + def test_ignores_and_skips(self, tmp_path): shard_path = str(tmp_path / "shard.safetensors") save_file( @@ -414,6 +536,41 @@ def test_ignored_layer_preserves_original_fp8(self, tmp_path): # non-ignored layer should be quantized normally assert "layer" in quantized + def test_dequant_fp8_hydrates_scale_from_sibling_shard(self, tmp_path): + """When scale_inv is sharded separately, dequant should hydrate it via index.""" + shard_dir = tmp_path / "source" + shard_dir.mkdir(parents=True, exist_ok=True) + + weight_name = "model.layers.0.mlp.experts.1.gate_proj.weight" + scale_name = "model.layers.0.mlp.experts.1.gate_proj.weight_scale_inv" + + shard_a = shard_dir / "model-00001-of-00002.safetensors" + shard_b = shard_dir / "model-00002-of-00002.safetensors" + save_file({weight_name: torch.randn(2048, 7168, dtype=torch.bfloat16).to(torch.float8_e4m3fn)}, str(shard_a)) + save_file({scale_name: torch.ones(16, 56, dtype=torch.float32)}, str(shard_b)) + + with open(shard_dir / "model.safetensors.index.json", "w") as f: + json.dump( + { + "metadata": {"total_size": 0}, + "weight_map": { + weight_name: shard_a.name, + scale_name: shard_b.name, + }, + }, + f, + ) + + output, quantized, _ = _process_shard( + str(shard_a), + _DEFAULT_SCHEME, + {}, + [], + fp8_block_size=[128, 128], + ) + assert "model.layers.0.mlp.experts.1.gate_proj" in quantized + assert "model.layers.0.mlp.experts.1.gate_proj.qweight" in output + # =========================================================================== # Quantization config builder diff --git a/test/unit/test_cpu/quantization/test_nvfp4_quant_linear.py b/test/unit/test_cpu/quantization/test_nvfp4_quant_linear.py index 860ef0ef00..4789f593ab 100644 --- a/test/unit/test_cpu/quantization/test_nvfp4_quant_linear.py +++ b/test/unit/test_cpu/quantization/test_nvfp4_quant_linear.py @@ -1,10 +1,15 @@ +from types import SimpleNamespace +from unittest.mock import patch + import pytest import torch +from transformers.quantizers.auto import AutoHfQuantizer -from auto_round.data_type.nvfp import calculate_gparam +from auto_round.data_type.nvfp import calculate_gparam, fp4_v2 from auto_round.data_type.utils import get_quant_func from auto_round.experimental import qmodules as ar_qmodules from auto_round.export.export_to_autoround.qlinear_fp import QuantLinear as _FPLinear +from auto_round.export.export_to_llmcompressor.config import initialize_nvfp4_e5m3_quantization from auto_round.export.formats import BackendDataType from auto_round.schemes import PRESET_SCHEMES @@ -35,6 +40,137 @@ def test_calculate_gparam_with_float8_input(): assert torch.isfinite(global_scale) +def test_nvfp4_e5m3_compressed_tensors_loading_uses_no_global_scales(): + class TinyModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(model_type="tiny") + self.model = torch.nn.Module() + self.model.layers = torch.nn.ModuleList([torch.nn.ModuleDict({"fc": torch.nn.Linear(16, 8)})]) + + quantizer = AutoHfQuantizer.from_config(initialize_nvfp4_e5m3_quantization([])) + model = quantizer._process_model_before_weight_loading(TinyModel()) + layer = model.model.layers[0]["fc"] + + assert isinstance(layer, ar_qmodules.NVFP4E5M3QuantLinear) + assert set(layer.state_dict()) == {"bias", "weight_packed", "weight_scale"} + assert quantizer._process_model_after_weight_loading(model) is model + + +def test_nvfp4_e5m3_qdq_input_uses_reference_fallback_on_cpu(): + config = PRESET_SCHEMES["NVFP4_E5M3"] + layer = ar_qmodules.NVFP4E5M3QuantLinear(16, 8, config, dtype=torch.float32) + activation = torch.randn(2, 3, 16) + + expected, _, _ = fp4_v2(activation, bits=config.act_bits, group_size=config.act_group_size) + + assert torch.equal(layer.qdq_input(activation), expected) + + +def test_nvfp4_e5m3_forward_does_not_cache_dequantized_weight_by_default(): + config = PRESET_SCHEMES["NVFP4_E5M3"] + layer = ar_qmodules.NVFP4E5M3QuantLinear(16, 8, config, dtype=torch.float32) + activation = torch.randn(2, 16) + dequantized_weight = torch.randn(8, 16) + + with patch.object(layer, "dequant_weight_online", return_value=dequantized_weight) as dequant_weight_online: + layer(activation) + layer(activation) + + assert layer._cached_weight is None + assert dequant_weight_online.call_count == 2 + + +def test_nvfp4_e5m3_forward_caches_dequantized_weight_when_enabled(): + config = PRESET_SCHEMES["NVFP4_E5M3"] + layer = ar_qmodules.NVFP4E5M3QuantLinear(16, 8, config, dtype=torch.float32, cache_weight=True) + activation = torch.randn(2, 16) + dequantized_weight = torch.randn(8, 16) + + with patch.object(layer, "dequant_weight_online", return_value=dequantized_weight) as dequant_weight_online: + layer(activation) + layer(activation) + + dequant_weight_online.assert_called_once_with() + assert layer.weight_packed is None + assert layer.weight_scale is None + + +def test_nvfp4_e5m3_cannot_clear_released_quantized_weight_buffers(): + config = PRESET_SCHEMES["NVFP4_E5M3"] + layer = ar_qmodules.NVFP4E5M3QuantLinear(16, 8, config, dtype=torch.float32, cache_weight=True) + + layer(torch.randn(2, 16)) + + with pytest.raises(RuntimeError, match="quantized weight buffers have been released"): + layer.clear_weight_cache() + + +def test_cute_nvfp4_e5m3_does_not_cache_dequantized_weight_by_default(monkeypatch): + monkeypatch.delenv("AR_NVFP4_E5M3_CACHE_HP_WEIGHT", raising=False) + config = PRESET_SCHEMES["NVFP4_E5M3"] + layer = ar_qmodules.CuteNVFP4E5M3QuantLinear(16, 8, config, dtype=torch.float32) + activation = torch.randn(2, 16) + dequantized_weight = torch.randn(8, 16) + + with patch.object(layer, "dequant_weight_online", return_value=dequantized_weight) as dequant_weight_online: + layer(activation) + layer(activation) + + assert not layer.cache_weight + assert dequant_weight_online.call_count == 2 + + +def test_nvfp4_e5m3_cache_weight_environment_override(monkeypatch): + config = PRESET_SCHEMES["NVFP4_E5M3"] + monkeypatch.setenv("AR_NVFP4_E5M3_CACHE_HP_WEIGHT", "1") + assert ar_qmodules.NVFP4E5M3QuantLinear(16, 8, config, dtype=torch.float32).cache_weight + assert ar_qmodules.CuteNVFP4E5M3QuantLinear(16, 8, config, dtype=torch.float32).cache_weight + monkeypatch.setenv("AR_NVFP4_E5M3_CACHE_HP_WEIGHT", "0") + assert not ar_qmodules.CuteNVFP4E5M3QuantLinear(16, 8, config, dtype=torch.float32).cache_weight + + +def test_nvfp4_e5m3_torch_forward_does_not_call_cute(): + config = PRESET_SCHEMES["NVFP4_E5M3"] + layer = ar_qmodules.NVFP4E5M3QuantLinear(16, 8, config, dtype=torch.float32) + activation = torch.randn(2, 16) + + with patch( + "auto_round.experimental.qmodules.nvfp4_e5m3.try_cute_nvfp4_e5m3_linear", + ) as cute_linear: + output = layer(activation) + + assert output.shape == (2, 8) + cute_linear.assert_not_called() + + +def test_nvfp4_e5m3_cute_forward_uses_fused_output_when_available(): + config = PRESET_SCHEMES["NVFP4_E5M3"] + layer = ar_qmodules.CuteNVFP4E5M3QuantLinear(16, 8, config, dtype=torch.float32) + activation = torch.randn(2, 16) + fused_output = torch.randn(2, 8) + + with patch("auto_round.experimental.qmodules.nvfp4_e5m3.try_cute_nvfp4_e5m3_linear", return_value=fused_output): + assert layer(activation) is fused_output + + +def test_cute_nvfp4_e5m3_uses_cached_weight_when_enabled(): + config = PRESET_SCHEMES["NVFP4_E5M3"] + layer = ar_qmodules.CuteNVFP4E5M3QuantLinear(16, 8, config, dtype=torch.float32, cache_weight=True) + activation = torch.randn(2, 16) + dequantized_weight = torch.randn(8, 16) + + with ( + patch.object(layer, "dequant_weight_online", return_value=dequantized_weight) as dequant_weight_online, + patch("auto_round.experimental.qmodules.nvfp4_e5m3.try_cute_nvfp4_e5m3_linear") as cute_linear, + ): + layer(activation) + layer(activation) + + dequant_weight_online.assert_called_once_with() + cute_linear.assert_not_called() + + @pytest.mark.parametrize("scheme", [BackendDataType.NVFP4.value]) @torch.inference_mode() def test_nvfp4_quantlinear_from_original_and_forward(scheme): diff --git a/test/unit/test_cuda/quantization/test_nvfp4_e5m3_cute.py b/test/unit/test_cuda/quantization/test_nvfp4_e5m3_cute.py new file mode 100644 index 0000000000..f68290586a --- /dev/null +++ b/test/unit/test_cuda/quantization/test_nvfp4_e5m3_cute.py @@ -0,0 +1,66 @@ +import pytest +import torch + +from auto_round.data_type.nvfp import cast_to_fp4, e5m3_to_float_tensor, fp4_v2 +from auto_round.experimental.qmodules.fp4_utils import unpack_fp4_from_uint8 +from auto_round_extension.cuda.cute_nvfp4_e5m3 import ( + can_use_cute_fp4_v2_qdq, + try_cute_fp4_v2_qdq, + try_cute_nvfp4_e5m3_weight_dq, +) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize( + ("dtype", "atol"), + [(torch.float32, 1e-6), (torch.float16, 2e-3), (torch.bfloat16, 2e-2)], +) +def test_cute_fp4_v2_qdq_matches_reference(dtype, atol): + if torch.cuda.get_device_capability()[0] < 8: + pytest.skip("CuTe QDQ requires SM80 or newer") + + activation = torch.randn(32, 16, device="cuda", dtype=dtype) + if not can_use_cute_fp4_v2_qdq(activation, 16): + pytest.skip("CuTe DSL is not available") + + expected, _, _ = fp4_v2(activation.float(), bits=4, group_size=16) + actual = try_cute_fp4_v2_qdq(activation, 16) + + assert actual is not None + torch.testing.assert_close(actual, expected.to(dtype), rtol=0, atol=atol) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_cute_nvfp4_e5m3_weight_dq_matches_reference_for_multiple_groups(dtype): + if torch.cuda.get_device_capability()[0] < 8: + pytest.skip("CuTe weight DQ requires SM80 or newer") + + torch.manual_seed(42) + weight_packed = torch.randint(0, 256, (37, 64), device="cuda", dtype=torch.uint8) + weight_scale = torch.randint(0, 120, (37, 8), device="cuda", dtype=torch.uint8) + actual = try_cute_nvfp4_e5m3_weight_dq(weight_packed, weight_scale, dtype) + if actual is None: + pytest.skip("CuTe DSL is not available") + + unpacked = unpack_fp4_from_uint8(weight_packed, 37, 128, dtype=dtype).float() + expected = (unpacked.reshape(-1, 16) * e5m3_to_float_tensor(weight_scale).reshape(-1, 1)).reshape(37, 128).to(dtype) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_fp4_v2_uses_direct_scale_at_rounding_boundary(): + activation = torch.tensor([[1.25] + [1.0] * 15], dtype=torch.float32) + + actual, scale, _ = fp4_v2(activation, bits=4, group_size=16) + expected = cast_to_fp4(activation / scale) * scale + + torch.testing.assert_close(actual, expected) + + +def test_fp4_v2_direct_scale_handles_zero_group(): + activation = torch.zeros(1, 16, dtype=torch.float32) + + actual, scale, _ = fp4_v2(activation, bits=4, group_size=16) + + assert torch.equal(actual, activation) + assert torch.equal(scale, torch.zeros_like(scale))