From 4b672fcb1d8eb36520ced48a3751b4115e152217 Mon Sep 17 00:00:00 2001 From: Zeel Date: Fri, 10 Jul 2026 11:38:33 -0400 Subject: [PATCH 1/3] feat: add per-head fp8 kv attention calibration Signed-off-by: Zeel --- auto_round/autoround.py | 2 + auto_round/cli/main.py | 4 + auto_round/cli/parser.py | 14 +++ auto_round/compressors/base.py | 22 +++- auto_round/compressors/entry.py | 2 + auto_round/context/compress.py | 4 + auto_round/experimental/attention.py | 42 +++++-- auto_round/experimental/kv_cache.py | 23 ++-- auto_round/experimental/utils.py | 52 ++++++++- .../export_to_llmcompressor/export_to_fp.py | 18 ++- .../export_to_static_fp.py | 36 +++--- test/test_cpu/export/test_export.py | 107 ++++++++++++++---- test/test_cpu/export/test_llmc_format.py | 57 ++++++++++ .../test_experimental_fp8_granularity.py | 21 ++++ 14 files changed, 341 insertions(+), 63 deletions(-) create mode 100644 test/test_cpu/test_experimental_fp8_granularity.py diff --git a/auto_round/autoround.py b/auto_round/autoround.py index 2b67300dec..dcfb84d2d2 100644 --- a/auto_round/autoround.py +++ b/auto_round/autoround.py @@ -59,6 +59,8 @@ "enable_deterministic_algorithms", "static_kv_dtype", "static_attention_dtype", + "static_kv_granularity", + "static_attention_granularity", "rotation_config", "processor", "image_processor", diff --git a/auto_round/cli/main.py b/auto_round/cli/main.py index 67d38c7ad8..823444dcf3 100644 --- a/auto_round/cli/main.py +++ b/auto_round/cli/main.py @@ -64,6 +64,10 @@ def _build_entry_base_kwargs(args, *, low_cpu_mem_usage, enable_torch_compile, l "layer_config": layer_config, "model_dtype": args.model_dtype, "trust_remote_code": not args.disable_trust_remote_code, + "static_kv_dtype": args.static_kv_dtype, + "static_kv_granularity": args.static_kv_granularity, + "static_attention_dtype": args.static_attention_dtype, + "static_attention_granularity": args.static_attention_granularity, } diff --git a/auto_round/cli/parser.py b/auto_round/cli/parser.py index ce68bf283a..fea70bc342 100644 --- a/auto_round/cli/parser.py +++ b/auto_round/cli/parser.py @@ -173,6 +173,13 @@ def build_quantize_parser(*, prog: str = "auto_round quantize") -> argparse.Argu choices=["fp8", "float8_e4m3fn"], help="Static KV-cache quantization data type.", ) + rt.add_argument( + "--static_kv_granularity", + default="tensor", + type=str, + choices=["tensor", "head"], + help="Static KV-cache FP8 calibration granularity.", + ) rt.add_argument( "--static_attention_dtype", default=None, @@ -180,6 +187,13 @@ def build_quantize_parser(*, prog: str = "auto_round quantize") -> argparse.Argu choices=["fp8", "float8_e4m3fn"], help="Static attention quantization data type.", ) + rt.add_argument( + "--static_attention_granularity", + default="tensor", + type=str, + choices=["tensor", "head"], + help="Static attention FP8 calibration granularity.", + ) # ---- Evaluation ---- ev = parser.add_argument_group("Evaluation Arguments") diff --git a/auto_round/compressors/base.py b/auto_round/compressors/base.py index f0060b54c8..4c34da0ca7 100644 --- a/auto_round/compressors/base.py +++ b/auto_round/compressors/base.py @@ -98,6 +98,8 @@ class SerializedCompressorConfig: supported_types: Optional[list[str]] = SUPPORTED_LAYER_TYPES static_attention_dtype: Optional[str] = None static_kv_dtype: Optional[str] = None + static_attention_granularity: Optional[str] = "tensor" + static_kv_granularity: Optional[str] = "tensor" super_bits: Optional[int] = None super_group_size: Optional[int] = None to_quant_block_names: Optional[list[str]] = None @@ -308,12 +310,18 @@ def __init__( if device is not None: logger.warning("`device` is deprecated, please use `device_map` instead") + from auto_round.experimental.utils import normalize_fp8_granularity + self.static_attention_dtype = kwargs.pop("static_attention_dtype", None) + self.static_attention_granularity = normalize_fp8_granularity( + kwargs.pop("static_attention_granularity", "tensor") + ) # Attention static dtype if self.static_attention_dtype is not None: logger.warning("The static attention dtype is experimental and currently has limited support.") # KV cache, this one does not affect tuning but will collect some infos during tuning self.static_kv_dtype = kwargs.pop("static_kv_dtype", None) + self.static_kv_granularity = normalize_fp8_granularity(kwargs.pop("static_kv_granularity", "tensor")) if self.static_kv_dtype is not None: logger.warning("The static kv is experimental and currently has limited support.") @@ -401,6 +409,8 @@ def __init__( formats=self.formats, static_kv_dtype=self.static_kv_dtype, static_attention_dtype=self.static_attention_dtype, + static_kv_granularity=self.static_kv_granularity, + static_attention_granularity=self.static_attention_granularity, ) self.shard_writer = None @@ -1605,13 +1615,21 @@ def quantize_and_save( if self.static_attention_dtype is not None: from auto_round.experimental.attention import attention_quant_ctx - with attention_quant_ctx(self.model_context.model, static_attention_dtype=self.static_attention_dtype): + with attention_quant_ctx( + self.model_context.model, + static_attention_dtype=self.static_attention_dtype, + static_attention_granularity=self.static_attention_granularity, + ): self.quantize() self.model_context.quantized = True elif self.static_kv_dtype is not None: from auto_round.experimental.kv_cache import kvcache_quant_context - with kvcache_quant_context(self.model_context.model, static_kv_dtype=self.static_kv_dtype): + with kvcache_quant_context( + self.model_context.model, + static_kv_dtype=self.static_kv_dtype, + static_kv_granularity=self.static_kv_granularity, + ): self.quantize() self.model_context.quantized = True else: diff --git a/auto_round/compressors/entry.py b/auto_round/compressors/entry.py index 9cdff97af8..6b28c08057 100644 --- a/auto_round/compressors/entry.py +++ b/auto_round/compressors/entry.py @@ -36,6 +36,8 @@ "enable_deterministic_algorithms", "static_kv_dtype", "static_attention_dtype", + "static_kv_granularity", + "static_attention_granularity", } _ENTRY_MLLM_KWARGS = {"processor", "image_processor", "template", "extra_data_dir", "quant_nontext_module"} _ENTRY_DIFFUSION_KWARGS = {"guidance_scale", "num_inference_steps", "generator_seed"} diff --git a/auto_round/context/compress.py b/auto_round/context/compress.py index b8dfc3070a..c4c97f74c6 100644 --- a/auto_round/context/compress.py +++ b/auto_round/context/compress.py @@ -40,6 +40,8 @@ def __init__( output_dir: str = "./compressed_models", static_kv_dtype: Optional[torch.dtype] = None, # TODO later this should be scheme wenhuach static_attention_dtype: Optional[torch.dtype] = None, + static_kv_granularity: str = "tensor", + static_attention_granularity: str = "tensor", **kwargs, ) -> None: super().__init__() @@ -58,6 +60,8 @@ def __init__( self.is_immediate_saving = is_immediate_saving self.static_kv_dtype = static_kv_dtype self.static_attention_dtype = static_attention_dtype + self.static_kv_granularity = static_kv_granularity + self.static_attention_granularity = static_attention_granularity def clear_memory(self, tensor=None): """Clear GPU/CPU memory only when ``low_gpu_mem_usage`` is enabled.""" diff --git a/auto_round/experimental/attention.py b/auto_round/experimental/attention.py index 0a2361423b..8459f0bd01 100644 --- a/auto_round/experimental/attention.py +++ b/auto_round/experimental/attention.py @@ -32,8 +32,9 @@ from auto_round.experimental.kv_cache import kvcache_quant_context from auto_round.experimental.utils import ( clean_model_parameters_and_buffers_, + fp8_qdq, is_attention_module, - per_tensor_fp8_qdq, + normalize_fp8_granularity, update_parameter_data, ) from auto_round.utils import logger @@ -69,9 +70,10 @@ class QuantizedAttentionImpl(torch.nn.Module): _original_impl = "sdpa" - def __init__(self, config: PretrainedConfig, attn_module: Module): + def __init__(self, config: PretrainedConfig, attn_module: Module, granularity: str = "tensor"): super().__init__() self.config = config + self.granularity = normalize_fp8_granularity(granularity) self.attn_module = ref(attn_module) # avoid circular references # register query max device = next(attn_module.parameters()).device @@ -89,14 +91,17 @@ def forward( *args, **kwargs, ): - cur_query_max = query.abs().max() + if self.granularity == "head": + cur_query_max = query.abs().amax(dim=(0, 2, 3)) + else: + cur_query_max = query.abs().max() query_max = torch.max( getattr(module, QUERY_MAX_NAME).data, cur_query_max.detach().to(getattr(module, QUERY_MAX_NAME).data.device), ) update_parameter_data(module, query_max, QUERY_MAX_NAME) - _, query_scale = per_tensor_fp8_qdq(query, tensor_max=query_max) - update_parameter_data(module, query_scale.squeeze(0).detach(), QUERY_SCALE_NAME) + _, query_scale = fp8_qdq(query, tensor_max=query_max, granularity=self.granularity) + update_parameter_data(module, query_scale.reshape(-1).detach(), QUERY_SCALE_NAME) # original attention return ALL_ATTENTION_FUNCTIONS[self._original_impl]( module, @@ -118,7 +123,7 @@ def _ct_hooked_attention(module: Module, *args, **kwargs): return ALL_ATTENTION_FUNCTIONS[_original_impl](module, *args, **kwargs) # pylint: disable=E0601 -def init_hooked_attention(module: Module, config): +def init_hooked_attention(module: Module, config, granularity: str = "tensor"): """ Initialize `QuantizedAttentionImpl` and `QuantizedKVCache` instances attached to attention @@ -127,7 +132,7 @@ def init_hooked_attention(module: Module, config): :param module: attention module to initialize with """ if not hasattr(module, ATTN_IMPL_ATTR_NAME): - module.register_module(ATTN_IMPL_ATTR_NAME, QuantizedAttentionImpl(config, module)) + module.register_module(ATTN_IMPL_ATTR_NAME, QuantizedAttentionImpl(config, module, granularity=granularity)) if config._attn_implementation != HOOKED_ATTENTION_NAME: # assumes only one model at a time global _original_impl @@ -139,10 +144,10 @@ def init_hooked_attention(module: Module, config): # initialize_hooked_kv_cache(model, module) -def prep_attention_module_for_calibration(module: torch.nn.Module, config): +def prep_attention_module_for_calibration(module: torch.nn.Module, config, granularity: str = "tensor"): if is_attention_module(module): logger.trace(f"Preparing attention module {module.__class__.__name__} for calibration") - init_hooked_attention(module, config) + init_hooked_attention(module, config, granularity=granularity) def clean_up_hooked_attention(module, model): @@ -155,12 +160,25 @@ def clean_up_hooked_attention(module, model): @contextlib.contextmanager -def attention_quant_ctx(model: PreTrainedModel, static_attention_dtype=torch.float8_e4m3fn): +def attention_quant_ctx( + model: PreTrainedModel, + static_attention_dtype=torch.float8_e4m3fn, + static_attention_granularity: str = "tensor", +): try: # Setup phase: Initialize hooked attention - prepare_fn = partial(prep_attention_module_for_calibration, config=model.config) + static_attention_granularity = normalize_fp8_granularity(static_attention_granularity) + prepare_fn = partial( + prep_attention_module_for_calibration, + config=model.config, + granularity=static_attention_granularity, + ) model.apply(prepare_fn) - with kvcache_quant_context(model, static_kv_dtype=static_attention_dtype): + with kvcache_quant_context( + model, + static_kv_dtype=static_attention_dtype, + static_kv_granularity=static_attention_granularity, + ): yield model finally: clean_fn = partial(clean_up_hooked_attention, model=model) diff --git a/auto_round/experimental/kv_cache.py b/auto_round/experimental/kv_cache.py index 1774315d72..ba35c408ea 100644 --- a/auto_round/experimental/kv_cache.py +++ b/auto_round/experimental/kv_cache.py @@ -25,9 +25,10 @@ from transformers.cache_utils import DynamicCache from auto_round.experimental.utils import ( + fp8_qdq, is_attention_module, + normalize_fp8_granularity, normalize_static_kv_dtype, - per_tensor_fp8_qdq, update_parameter_data, ) from auto_round.utils import logger @@ -106,9 +107,10 @@ def __new__(cls, *args, **kwargs): cls._instance = super(QuantizedKVParameterCache, cls).__new__(cls) return cls._instance - def __init__(self, dtype: torch.dtype = torch.float8_e4m3fn): + def __init__(self, dtype: torch.dtype = torch.float8_e4m3fn, granularity: str = "tensor"): assert dtype == torch.float8_e4m3fn, "Only fp8_e4m3fn is supported for now." + self.granularity = normalize_fp8_granularity(granularity) if not self._initialized: super().__init__() @@ -172,13 +174,13 @@ def _quant_dequant(self, tensor: torch.Tensor, kv_type: KVCacheScaleType, layer_ assert kv_type == KVCacheScaleType.VALUE scales = self.v_scales - qdq_tensor, scale = per_tensor_fp8_qdq(tensor) + qdq_tensor, scale = fp8_qdq(tensor, granularity=self.granularity) # Detach scale to prevent holding computation graph references - _pad_and_append_at_idx_(scales, layer_idx, scale.squeeze(0).detach()) + _pad_and_append_at_idx_(scales, layer_idx, scale.reshape(-1).detach()) return qdq_tensor -def initialize_quantized_kv_cache(module: torch.nn.Module, dtype=torch.float8_e4m3fn): +def initialize_quantized_kv_cache(module: torch.nn.Module, dtype=torch.float8_e4m3fn, granularity: str = "tensor"): """ Initialize a quantized kv_cache on a module (analogous to initializing an observer) """ @@ -189,7 +191,7 @@ def initialize_quantized_kv_cache(module: torch.nn.Module, dtype=torch.float8_e4 if isinstance(existing_kv_cache, QuantizedKVParameterCache): return - quantized_kv_cache = QuantizedKVParameterCache(dtype=dtype) + quantized_kv_cache = QuantizedKVParameterCache(dtype=dtype, granularity=granularity) setattr(module, "kv_cache", quantized_kv_cache) logger.debug(f"Initialized quantized kv_cache for {module.__class__.__name__} {getattr(module, 'layer_idx', None)}") init_scale = torch.tensor([0.0], device=next(module.parameters()).device) @@ -234,15 +236,20 @@ def prep_attention_module_for_calibration(module: torch.nn.Module): @contextlib.contextmanager -def kvcache_quant_context(model: torch.nn.Module, static_kv_dtype=torch.float8_e4m3fn): +def kvcache_quant_context( + model: torch.nn.Module, static_kv_dtype=torch.float8_e4m3fn, static_kv_granularity: str = "tensor" +): """Context manager for FP8 KV cache quantization operations.""" try: # Setup phase: Initialize KV cache for quantization static_kv_dtype = normalize_static_kv_dtype(static_kv_dtype) + static_kv_granularity = normalize_fp8_granularity(static_kv_granularity) if static_kv_dtype != torch.float8_e4m3fn: logger.warning(f"Ignoring static kv dtype {static_kv_dtype}, only fp8_e4m3fn is supported.") else: - initialize_fn = partial(initialize_quantized_kv_cache, dtype=static_kv_dtype) + initialize_fn = partial( + initialize_quantized_kv_cache, dtype=static_kv_dtype, granularity=static_kv_granularity + ) model.apply(initialize_fn) model.apply(prep_attention_module_for_calibration) diff --git a/auto_round/experimental/utils.py b/auto_round/experimental/utils.py index c9a5dd2aad..e3aba53896 100644 --- a/auto_round/experimental/utils.py +++ b/auto_round/experimental/utils.py @@ -23,6 +23,22 @@ SUPPORTED_QUANTIZATION_SCHEMES = ["MXFP8", "MXFP4", "NVFP4"] +FP8_GRANULARITY_TENSOR = "tensor" +FP8_GRANULARITY_HEAD = "head" + + +def normalize_fp8_granularity(granularity: str | None) -> str: + if granularity is None: + return FP8_GRANULARITY_TENSOR + granularity = granularity.lower() + if granularity in (FP8_GRANULARITY_TENSOR, FP8_GRANULARITY_HEAD): + return granularity + raise ValueError( + f"Invalid FP8 calibration granularity: {granularity}. " + f"Supported granularities are: {FP8_GRANULARITY_TENSOR}, {FP8_GRANULARITY_HEAD}." + ) + + def per_tensor_fp8_qdq( tensor: torch.Tensor, tensor_max: None | torch.Tensor = None ) -> tuple[torch.Tensor, torch.Tensor]: @@ -32,6 +48,37 @@ def per_tensor_fp8_qdq( return qdq_tensor, scale +def per_head_fp8_qdq(tensor: torch.Tensor, tensor_max: None | torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize/dequantize an attention tensor with one FP8 scale per attention head.""" + if tensor.dim() < 4: + raise ValueError(f"Per-head FP8 calibration expects [batch, heads, seq, head_dim], got {tuple(tensor.shape)}") + + info = torch.finfo(torch.float8_e4m3fn) + orig_dtype = tensor.dtype + if tensor.dtype == torch.float16: + tensor = tensor.to(torch.bfloat16) + + if tensor_max is None: + max_tensor = tensor.abs().amax(dim=(0, 2, 3)).to(torch.float32) + else: + max_tensor = tensor_max.to(tensor.device, dtype=torch.float32) + + scale = torch.clip(max_tensor / info.max, min=float(1.0 / (info.max * 512.0))) + qdq_scale = scale.view(1, -1, 1, 1) + fp8_res = torch.clip(tensor / qdq_scale, info.min, info.max) + fp8_res = fp8_res.to(torch.float8_e4m3fn) + return (fp8_res.to(tensor.dtype) * qdq_scale).to(orig_dtype), scale + + +def fp8_qdq( + tensor: torch.Tensor, tensor_max: None | torch.Tensor = None, granularity: str | None = None +) -> tuple[torch.Tensor, torch.Tensor]: + granularity = normalize_fp8_granularity(granularity) + if granularity == FP8_GRANULARITY_HEAD: + return per_head_fp8_qdq(tensor, tensor_max=tensor_max) + return per_tensor_fp8_qdq(tensor, tensor_max=tensor_max) + + @torch.compiler.disable def update_parameter_data(module: torch.nn.Module, new_val: torch.Tensor, name: str): """ @@ -41,7 +88,10 @@ def update_parameter_data(module: torch.nn.Module, new_val: torch.Tensor, name: if hasattr(module, name): param = getattr(module, name) if isinstance(param, torch.nn.Parameter): - param.data.copy_(new_val) + if param.shape == new_val.shape: + param.data.copy_(new_val) + else: + setattr(module, name, torch.nn.Parameter(new_val)) else: module.register_parameter(name, torch.nn.Parameter(new_val)) else: 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..c061095da2 100644 --- a/auto_round/export/export_to_llmcompressor/export_to_fp.py +++ b/auto_round/export/export_to_llmcompressor/export_to_fp.py @@ -161,6 +161,8 @@ def _build_mixed_fp_quantization_config( model, static_kv_dtype=None, static_attention_dtype=None, + static_kv_granularity="tensor", + static_attention_granularity="tensor", ): """Build a quantization config dict for mixed-precision scenarios. @@ -204,13 +206,16 @@ def _build_mixed_fp_quantization_config( use_fp8_attention = _use_fp8_attention(static_attention_dtype) if use_fp8_attention: - attention_config = _get_attention_config(model) + attention_config = _get_attention_config(model, static_attention_granularity) else: attention_config = None + kv_granularity = static_attention_granularity if use_fp8_attention else static_kv_granularity quantization_config = initialize_quantization( scheme=None, config_groups=config_groups, - kv_cache_scheme=_construct_kv_scheme() if (_use_fp8_kv(static_kv_dtype) or use_fp8_attention) else None, + kv_cache_scheme=( + _construct_kv_scheme(kv_granularity) if (_use_fp8_kv(static_kv_dtype) or use_fp8_attention) else None + ), ignore=ignore, ) quantization_config = quantization_config.to_dict() @@ -322,8 +327,11 @@ def save_quantized_as_fp( is_mixed = len(scheme_groups) > 1 use_fp8_attention = _use_fp8_attention(serialization_dict.get("static_attention_dtype", None)) + static_attention_granularity = serialization_dict.get("static_attention_granularity", "tensor") + static_kv_granularity = serialization_dict.get("static_kv_granularity", "tensor") + kv_granularity = static_attention_granularity if use_fp8_attention else static_kv_granularity kv_cache_scheme = ( - _construct_kv_scheme() + _construct_kv_scheme(kv_granularity) if (_use_fp8_kv(serialization_dict.get("static_kv_dtype", None)) or use_fp8_attention) else None ) @@ -338,6 +346,8 @@ def save_quantized_as_fp( model, static_kv_dtype=serialization_dict.get("static_kv_dtype", None), static_attention_dtype=serialization_dict.get("static_attention_dtype", None), + static_kv_granularity=static_kv_granularity, + static_attention_granularity=static_attention_granularity, ) else: scheme = _get_scheme(bits, data_type) @@ -351,7 +361,7 @@ def save_quantized_as_fp( ignore=ignore, ) if use_fp8_attention: - attention_config = _get_attention_config(model) + attention_config = _get_attention_config(model, static_attention_granularity) else: attention_config = None setattr(quantization_config, "format", format) diff --git a/auto_round/export/export_to_llmcompressor/export_to_static_fp.py b/auto_round/export/export_to_llmcompressor/export_to_static_fp.py index 82cf6d2a2f..a2cc52b079 100644 --- a/auto_round/export/export_to_llmcompressor/export_to_static_fp.py +++ b/auto_round/export/export_to_llmcompressor/export_to_static_fp.py @@ -62,25 +62,32 @@ def pack_layer(layer_name: str, model: torch.nn.Module, data_type: str, device: fp8_pack_layer(layer_name, model, data_type, device, unsqueeze=True) -def _construct_fp8_args(): +def _get_fp8_strategy(granularity: str = "tensor"): + from compressed_tensors.quantization import QuantizationStrategy # pylint: disable=E0401 + + if granularity == "head": + return QuantizationStrategy.ATTN_HEAD + return QuantizationStrategy.TENSOR + + +def _construct_fp8_args(granularity: str = "tensor"): from compressed_tensors.quantization import ( # pylint: disable=E0401 QuantizationArgs, - QuantizationStrategy, QuantizationType, ) return QuantizationArgs( num_bits=8, type=QuantizationType.FLOAT, - strategy=QuantizationStrategy.TENSOR, + strategy=_get_fp8_strategy(granularity), symmetric=True, dynamic=False, ) -def _construct_kv_scheme(): +def _construct_kv_scheme(granularity: str = "tensor"): """Construct the default KV cache quantization scheme for FP8_STATIC export.""" - default_kv_scheme = _construct_fp8_args() + default_kv_scheme = _construct_fp8_args(granularity) logger.warning_once( "Using default KV cache scheme: %s. " @@ -119,7 +126,7 @@ def _get_attention_targets(model: torch.nn.Module) -> list[str]: return attention_targets -def _get_attention_config(model: torch.nn.Module) -> dict | None: +def _get_attention_config(model: torch.nn.Module, granularity: str = "tensor") -> dict | None: """Return attention FP8 config as a standalone dict for the top-level ``attention_input_activations`` field. @@ -134,7 +141,7 @@ def _get_attention_config(model: torch.nn.Module) -> dict | None: return { "targets": attention_targets, - "input_activations": _construct_fp8_args().model_dump(), + "input_activations": _construct_fp8_args(granularity).model_dump(), } @@ -237,18 +244,15 @@ def save_quantized_as_static_fp( config_groups = {} scheme = QuantizationScheme(targets=targets, **scheme_args) config_groups["group_0"] = scheme + static_attention_granularity = serialization_dict.get("static_attention_granularity", "tensor") + static_kv_granularity = serialization_dict.get("static_kv_granularity", "tensor") use_fp8_attention = _use_fp8_attention(serialization_dict.get("static_attention_dtype", None)) - if use_fp8_attention: - attention_config = _get_attention_config(model) - else: - attention_config = None + use_fp8_kv = _use_fp8_kv(serialization_dict.get("static_kv_dtype", None)) + attention_config = _get_attention_config(model, static_attention_granularity) if use_fp8_attention else None + kv_granularity = static_attention_granularity if use_fp8_attention else static_kv_granularity quantization_config = QuantizationConfig( config_groups=config_groups, - kv_cache_scheme=( - _construct_kv_scheme() - if (_use_fp8_kv(serialization_dict.get("static_kv_dtype", None)) or use_fp8_attention) - else None - ), + kv_cache_scheme=_construct_kv_scheme(kv_granularity) if (use_fp8_kv or use_fp8_attention) else None, quantization_status=QuantizationStatus.COMPRESSED, ignore=ignore, ) diff --git a/test/test_cpu/export/test_export.py b/test/test_cpu/export/test_export.py index 3b6c411eb6..db63294f9f 100644 --- a/test/test_cpu/export/test_export.py +++ b/test/test_cpu/export/test_export.py @@ -224,11 +224,21 @@ def test_static_afp8_export(self, static_kv_dtype): ) quantized_model_path = self.save_dir _, quantized_model_path = autoround.quantize_and_save(output_dir=quantized_model_path, format="auto_round") - f = safe_open(os.path.join(quantized_model_path, "model.safetensors"), framework="pt") - assert "model.decoder.layers.8.self_attn.k_proj.input_scale" in f.keys() - assert "model.decoder.layers.8.self_attn.k_proj.weight_scale" in f.keys() - assert f.get_tensor("model.decoder.layers.5.self_attn.v_proj.input_scale").shape == torch.Size([1]) - assert f.get_tensor("model.decoder.layers.5.self_attn.v_proj.weight").dtype == torch.float8_e4m3fn + with safe_open(os.path.join(quantized_model_path, "model.safetensors"), framework="pt") as f: + assert "model.decoder.layers.8.self_attn.k_proj.input_scale" in f.keys() + assert "model.decoder.layers.8.self_attn.k_proj.weight_scale" in f.keys() + assert f.get_tensor("model.decoder.layers.5.self_attn.v_proj.input_scale").shape == torch.Size([1]) + assert f.get_tensor("model.decoder.layers.5.self_attn.v_proj.weight").dtype == torch.float8_e4m3fn + if static_kv_dtype == "fp8": + assert "model.decoder.layers.8.self_attn.k_scale" in f.keys() + assert "model.decoder.layers.8.self_attn.v_scale" in f.keys() + assert f.get_tensor("model.decoder.layers.5.self_attn.v_scale").shape == torch.Size([1]) + assert f.get_tensor("model.decoder.layers.5.self_attn.k_scale").shape == torch.Size([1]) + assert ( + f.get_tensor("model.decoder.layers.5.self_attn.k_scale").dtype == torch.float32 + or f.get_tensor("model.decoder.layers.5.self_attn.k_scale").dtype == torch.bfloat16 + ) + if static_kv_dtype is None: with torch.no_grad(): import transformers @@ -243,7 +253,10 @@ def test_static_afp8_export(self, static_kv_dtype): assert ( model.model.decoder.layers[0].self_attn.k_proj.__class__.__name__ == "WeightFP8ActFP8StaticQuantLinear" - ), f"Expected WeightFP8ActFP8StaticQuantLinear, got {model.model.decoder.layers[0].self_attn.k_proj.__class__.__name__}" + ), ( + "Expected WeightFP8ActFP8StaticQuantLinear, " + f"got {model.model.decoder.layers[0].self_attn.k_proj.__class__.__name__}" + ) tokenizer = transformers.AutoTokenizer.from_pretrained(quantized_model_path) prompt = "AI is " encode = tokenizer.encode(prompt, return_tensors="pt") @@ -257,16 +270,6 @@ def test_static_afp8_export(self, static_kv_dtype): print(f"Output: {output}") assert output is not None, "Output should not be None" - if static_kv_dtype == "fp8": - assert "model.decoder.layers.8.self_attn.k_scale" in f.keys() - assert "model.decoder.layers.8.self_attn.v_scale" in f.keys() - assert f.get_tensor("model.decoder.layers.5.self_attn.v_scale").shape == torch.Size([1]) - assert f.get_tensor("model.decoder.layers.5.self_attn.k_scale").shape == torch.Size([1]) - assert ( - f.get_tensor("model.decoder.layers.5.self_attn.k_scale").dtype == torch.float32 - or f.get_tensor("model.decoder.layers.5.self_attn.k_scale").dtype == torch.bfloat16 - ) - model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", trust_remote_code=True) autoround = AutoRound( model, @@ -285,11 +288,45 @@ def test_static_afp8_export(self, static_kv_dtype): quantized_model_path = self.save_dir _, quantized_model_path = autoround.quantize_and_save(output_dir=quantized_model_path, format="auto_round") + with safe_open(os.path.join(quantized_model_path, "model.safetensors"), framework="pt") as f: + assert "model.decoder.layers.8.self_attn.k_proj.input_scale" in f.keys() + assert "model.decoder.layers.8.self_attn.k_proj.weight_scale" in f.keys() + assert f.get_tensor("model.decoder.layers.5.self_attn.v_proj.input_scale").shape == torch.Size([1]) + assert f.get_tensor("model.decoder.layers.5.self_attn.v_proj.weight").dtype == torch.float8_e4m3fn + + def test_static_afp8_per_head_export(self): + import os + + from safetensors import safe_open + + model_name = self.model_name + model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", trust_remote_code=True) + autoround = AutoRound( + model, + self.tokenizer, + bits=8, + group_size=-1, + iters=0, + scheme="fp8_static", + nsamples=2, + seqlen=2, + static_kv_dtype="fp8", + static_kv_granularity="head", + ) + _, quantized_model_path = autoround.quantize_and_save(output_dir=self.save_dir, format="auto_round") f = safe_open(os.path.join(quantized_model_path, "model.safetensors"), framework="pt") - assert "model.decoder.layers.8.self_attn.k_proj.input_scale" in f.keys() - assert "model.decoder.layers.8.self_attn.k_proj.weight_scale" in f.keys() - assert f.get_tensor("model.decoder.layers.5.self_attn.v_proj.input_scale").shape == torch.Size([1]) - assert f.get_tensor("model.decoder.layers.5.self_attn.v_proj.weight").dtype == torch.float8_e4m3fn + assert f.get_tensor("model.decoder.layers.5.self_attn.k_scale").shape == torch.Size( + [model.config.num_attention_heads] + ) + assert f.get_tensor("model.decoder.layers.5.self_attn.v_scale").shape == torch.Size( + [model.config.num_attention_heads] + ) + + with open(os.path.join(quantized_model_path, "config.json")) as config_file: + config = json.load(config_file) + quantization_config = config["quantization_config"] + assert quantization_config["static_kv_dtype"] == "fp8" + assert quantization_config["static_kv_granularity"] == "head" def test_static_fp8_attn(self): import os @@ -321,6 +358,36 @@ def test_static_fp8_attn(self): assert f.get_tensor(weight_name).shape == torch.Size([1]) assert f.get_tensor(weight_name).dtype == torch.float32 or f.get_tensor(weight_name).dtype == torch.bfloat16 + def test_static_fp8_per_head_attn(self): + import os + + from safetensors import safe_open + + model_name = self.model_name + model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", trust_remote_code=True) + autoround = AutoRound( + model, + self.tokenizer, + iters=0, + nsamples=2, + seqlen=2, + scheme="FP8_STATIC", + static_attention_dtype="fp8", + static_attention_granularity="head", + ) + _, quantized_model_path = autoround.quantize_and_save(output_dir=self.save_dir, format="auto_round") + f = safe_open(os.path.join(quantized_model_path, "model.safetensors"), framework="pt") + for attr in ("k_scale", "v_scale", "q_scale"): + weight_name = f"model.decoder.layers.5.self_attn.{attr}" + assert weight_name in f.keys() + assert f.get_tensor(weight_name).shape == torch.Size([model.config.num_attention_heads]) + + with open(os.path.join(quantized_model_path, "config.json")) as config_file: + config = json.load(config_file) + quantization_config = config["quantization_config"] + assert quantization_config["static_attention_dtype"] == "fp8" + assert quantization_config["static_attention_granularity"] == "head" + def test_awq_lmhead_export(self, dataloader): bits, sym, group_size = 4, False, 128 model_name = get_model_path("microsoft/phi-4") diff --git a/test/test_cpu/export/test_llmc_format.py b/test/test_cpu/export/test_llmc_format.py index 8a6e6d95f2..e27d82d820 100644 --- a/test/test_cpu/export/test_llmc_format.py +++ b/test/test_cpu/export/test_llmc_format.py @@ -150,6 +150,35 @@ def test_mxfp8_llmcompressor_kv_config(self, tiny_opt_model_path, tmp_path): assert kv_cache_scheme["dynamic"] is False assert kv_cache_scheme["symmetric"] is True + def test_mxfp8_llmcompressor_per_head_kv_config(self, tiny_opt_model_path, tmp_path): + from safetensors import safe_open + + ar = AutoRound( + model=tiny_opt_model_path, + iters=0, + disable_opt_rtn=True, + scheme="mxfp8", + static_kv_dtype="fp8", + static_kv_granularity="head", + ) + compressed_model, quantized_model_path = ar.quantize_and_save(output_dir=tmp_path, format="llm_compressor") + + with open(os.path.join(quantized_model_path, "config.json")) as f: + config = json.load(f) + + kv_cache_scheme = config["quantization_config"]["kv_cache_scheme"] + assert kv_cache_scheme is not None + assert kv_cache_scheme["strategy"] == "attn_head" + + num_kv_heads = compressed_model.config.num_attention_heads + if hasattr(compressed_model.config, "num_key_value_heads"): + num_kv_heads = compressed_model.config.num_key_value_heads + with safe_open(os.path.join(quantized_model_path, "model.safetensors"), framework="pt") as f: + k_scale = f.get_tensor("model.decoder.layers.0.self_attn.k_scale") + v_scale = f.get_tensor("model.decoder.layers.0.self_attn.v_scale") + assert k_scale.shape == torch.Size([num_kv_heads]) + assert v_scale.shape == torch.Size([num_kv_heads]) + def test_mxfp8_llmcompressor_attention_config(self, tiny_opt_model_path, tmp_path): ar = AutoRound( model=tiny_opt_model_path, @@ -185,6 +214,34 @@ def test_mxfp8_llmcompressor_attention_config(self, tiny_opt_model_path, tmp_pat assert quantization_config["kv_cache_scheme"] is not None assert getattr(compressed_model.model.decoder.layers[0].self_attn, "q_scale", None) is not None + def test_mxfp8_llmcompressor_per_head_attention_config(self, tiny_opt_model_path, tmp_path): + from safetensors import safe_open + + ar = AutoRound( + model=tiny_opt_model_path, + iters=0, + disable_opt_rtn=True, + scheme="mxfp8", + static_attention_dtype="fp8", + static_attention_granularity="head", + ) + compressed_model, quantized_model_path = ar.quantize_and_save(output_dir=tmp_path, format="llm_compressor") + + with open(os.path.join(quantized_model_path, "config.json")) as f: + saved_config = json.load(f) + + attention_config = saved_config["quantization_config"]["attention_input_activations"] + assert attention_config is not None + assert attention_config["targets"] == [compressed_model.model.decoder.layers[0].self_attn.__class__.__name__] + assert attention_config["input_activations"]["strategy"] == "attn_head" + assert saved_config["quantization_config"]["kv_cache_scheme"]["strategy"] == "attn_head" + + with safe_open(os.path.join(quantized_model_path, "model.safetensors"), framework="pt") as f: + q_scale = f.get_tensor("model.decoder.layers.0.self_attn.q_scale") + k_scale = f.get_tensor("model.decoder.layers.0.self_attn.k_scale") + assert q_scale.shape == torch.Size([compressed_model.config.num_attention_heads]) + assert k_scale.ndim == 1 + def test_mixed_precision_llmcompressor_format(self, tiny_opt_model_path, tmp_path): scheme = AutoScheme( avg_bits=7, diff --git a/test/test_cpu/test_experimental_fp8_granularity.py b/test/test_cpu/test_experimental_fp8_granularity.py new file mode 100644 index 0000000000..c8ea96f198 --- /dev/null +++ b/test/test_cpu/test_experimental_fp8_granularity.py @@ -0,0 +1,21 @@ +import torch + +from auto_round.experimental.utils import fp8_qdq + + +def test_per_head_fp8_qdq_preserves_mha_head_axis(): + tensor = torch.randn(2, 4, 3, 8) + + qdq_tensor, scale = fp8_qdq(tensor, granularity="head") + + assert qdq_tensor.shape == tensor.shape + assert scale.shape == torch.Size([4]) + + +def test_per_head_fp8_qdq_preserves_gqa_kv_head_axis(): + tensor = torch.randn(2, 2, 3, 8) + + qdq_tensor, scale = fp8_qdq(tensor, granularity="head") + + assert qdq_tensor.shape == tensor.shape + assert scale.shape == torch.Size([2]) From e9766a51bdcb17550cb13b8ff4386be14f4f4e38 Mon Sep 17 00:00:00 2001 From: yiliu30 Date: Fri, 7 Aug 2026 07:33:26 +0000 Subject: [PATCH 2/3] fix: address fp8 review comments --- auto_round/experimental/utils.py | 12 +++++++++ test/unit/test_cpu/export/test_export.py | 27 ++++++++++--------- .../test_experimental_fp8_granularity.py | 8 ++++++ 3 files changed, 35 insertions(+), 12 deletions(-) diff --git a/auto_round/experimental/utils.py b/auto_round/experimental/utils.py index e3aba53896..3727bcf4c3 100644 --- a/auto_round/experimental/utils.py +++ b/auto_round/experimental/utils.py @@ -52,6 +52,11 @@ def per_head_fp8_qdq(tensor: torch.Tensor, tensor_max: None | torch.Tensor = Non """Quantize/dequantize an attention tensor with one FP8 scale per attention head.""" if tensor.dim() < 4: raise ValueError(f"Per-head FP8 calibration expects [batch, heads, seq, head_dim], got {tuple(tensor.shape)}") + if tensor_max is not None and tensor_max.shape != (tensor.shape[1],): + raise ValueError( + f"Per-head FP8 calibration expects tensor_max shape ({tensor.shape[1]},), " + f"got {tuple(tensor_max.shape)}" + ) info = torch.finfo(torch.float8_e4m3fn) orig_dtype = tensor.dtype @@ -91,6 +96,13 @@ def update_parameter_data(module: torch.nn.Module, new_val: torch.Tensor, name: if param.shape == new_val.shape: param.data.copy_(new_val) else: + logger.warning( + "Replacing parameter %s in module %s with a new shape: %s -> %s", + name, + module.__class__.__name__, + tuple(param.shape), + tuple(new_val.shape), + ) setattr(module, name, torch.nn.Parameter(new_val)) else: module.register_parameter(name, torch.nn.Parameter(new_val)) diff --git a/test/unit/test_cpu/export/test_export.py b/test/unit/test_cpu/export/test_export.py index f26e7d9514..13aad317f0 100644 --- a/test/unit/test_cpu/export/test_export.py +++ b/test/unit/test_cpu/export/test_export.py @@ -348,18 +348,21 @@ def test_static_fp8_attn(self): ) quantized_model_path = self.save_dir _, quantized_model_path = autoround.quantize_and_save(output_dir=quantized_model_path, format="auto_round") - f = safe_open(os.path.join(quantized_model_path, "model.safetensors"), framework="pt") - assert "model.decoder.layers.0.self_attn.k_proj.input_scale" in f.keys() - assert "model.decoder.layers.0.self_attn.k_proj.weight_scale" in f.keys() - assert f.get_tensor("model.decoder.layers.0.self_attn.v_proj.input_scale").shape == torch.Size([1]) - assert f.get_tensor("model.decoder.layers.0.self_attn.v_proj.weight").dtype == torch.float8_e4m3fn - check_attrs = ["k_scale", "v_scale", "q_scale"] - for attr in check_attrs: - weight_name = f"model.decoder.layers.0.self_attn.{attr}" - assert weight_name in f.keys() - assert f.get_tensor(weight_name).shape == torch.Size([1]) - assert f.get_tensor(weight_name).dtype == torch.float32 or f.get_tensor(weight_name).dtype == torch.bfloat16 - assert not any(key.endswith(".q_max") for key in f.keys()) + with safe_open(os.path.join(quantized_model_path, "model.safetensors"), framework="pt") as f: + assert "model.decoder.layers.0.self_attn.k_proj.input_scale" in f.keys() + assert "model.decoder.layers.0.self_attn.k_proj.weight_scale" in f.keys() + assert f.get_tensor("model.decoder.layers.0.self_attn.v_proj.input_scale").shape == torch.Size([1]) + assert f.get_tensor("model.decoder.layers.0.self_attn.v_proj.weight").dtype == torch.float8_e4m3fn + check_attrs = ["k_scale", "v_scale", "q_scale"] + for attr in check_attrs: + weight_name = f"model.decoder.layers.0.self_attn.{attr}" + assert weight_name in f.keys() + assert f.get_tensor(weight_name).shape == torch.Size([1]) + assert ( + f.get_tensor(weight_name).dtype == torch.float32 + or f.get_tensor(weight_name).dtype == torch.bfloat16 + ) + assert not any(key.endswith(".q_max") for key in f.keys()) def test_static_fp8_per_head_attn(self): import os diff --git a/test/unit/test_cpu/test_experimental_fp8_granularity.py b/test/unit/test_cpu/test_experimental_fp8_granularity.py index c8ea96f198..61bcfcf516 100644 --- a/test/unit/test_cpu/test_experimental_fp8_granularity.py +++ b/test/unit/test_cpu/test_experimental_fp8_granularity.py @@ -1,3 +1,4 @@ +import pytest import torch from auto_round.experimental.utils import fp8_qdq @@ -19,3 +20,10 @@ def test_per_head_fp8_qdq_preserves_gqa_kv_head_axis(): assert qdq_tensor.shape == tensor.shape assert scale.shape == torch.Size([2]) + + +def test_per_head_fp8_qdq_rejects_mismatched_tensor_max_shape(): + tensor = torch.randn(2, 4, 3, 8) + + with pytest.raises(ValueError, match=r"tensor_max shape \(4,\)"): + fp8_qdq(tensor, tensor_max=torch.ones(2), granularity="head") From 3a3aa43f0c735d8b1214e41f528cd02bc2e61e51 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 07:36:52 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- auto_round/experimental/utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/auto_round/experimental/utils.py b/auto_round/experimental/utils.py index 3727bcf4c3..2b106f1f10 100644 --- a/auto_round/experimental/utils.py +++ b/auto_round/experimental/utils.py @@ -54,8 +54,7 @@ def per_head_fp8_qdq(tensor: torch.Tensor, tensor_max: None | torch.Tensor = Non raise ValueError(f"Per-head FP8 calibration expects [batch, heads, seq, head_dim], got {tuple(tensor.shape)}") if tensor_max is not None and tensor_max.shape != (tensor.shape[1],): raise ValueError( - f"Per-head FP8 calibration expects tensor_max shape ({tensor.shape[1]},), " - f"got {tuple(tensor_max.shape)}" + f"Per-head FP8 calibration expects tensor_max shape ({tensor.shape[1]},), " f"got {tuple(tensor_max.shape)}" ) info = torch.finfo(torch.float8_e4m3fn)