expand observers to calculate gparams, add example for activations

Author: dsikka

examples/quantization_w4a4_fp4/llama3_example.py

@@ -0,0 +1,74 @@
+from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from llmcompressor import oneshot
+from llmcompressor.modifiers.quantization import QuantizationModifier
+
+MODEL_ID = "meta-llama/Llama-3.3-70B-Instruct"
+# MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+
+# Load model.
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID, device_map="auto", torch_dtype="auto"
+)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+
+DATASET_ID = "HuggingFaceH4/ultrachat_200k"
+DATASET_SPLIT = "train_sft"
+
+# Select number of samples. 512 samples is a good place to start.
+# Increasing the number of samples can improve accuracy.
+NUM_CALIBRATION_SAMPLES = 20
+MAX_SEQUENCE_LENGTH = 2048
+
+# Load dataset and preprocess.
+ds = load_dataset(DATASET_ID, split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]")
+ds = ds.shuffle(seed=42)
+
+
+def preprocess(example):
+    return {
+        "text": tokenizer.apply_chat_template(
+            example["messages"],
+            tokenize=False,
+        )
+    }
+
+
+ds = ds.map(preprocess)
+
+
+# Tokenize inputs.
+def tokenize(sample):
+    return tokenizer(
+        sample["text"],
+        padding=False,
+        max_length=MAX_SEQUENCE_LENGTH,
+        truncation=True,
+        add_special_tokens=False,
+    )
+
+
+ds = ds.map(tokenize, remove_columns=ds.column_names)
+
+# Configure the quantization algorithm and scheme.
+# In this case, we:
+#   * quantize the weights to fp4 with per group 16 via ptq
+#   * calibrate a global_scale for activations, which will be used to
+#       quantize activations to fp4 on the fly
+recipe = QuantizationModifier(targets="Linear", scheme="NVFP4", ignore=["lm_head"])
+
+# Apply quantization.
+oneshot(
+    model=model,
+    dataset=ds,
+    recipe=recipe,
+    max_seq_length=MAX_SEQUENCE_LENGTH,
+    num_calibration_samples=NUM_CALIBRATION_SAMPLES,
+)
+
+# Save to disk in compressed-tensors format.
+SAVE_DIR = MODEL_ID.split("/")[1] + "-NVFP4-v4"
+model.save_pretrained(SAVE_DIR, save_compressed=True)
+tokenizer.save_pretrained(SAVE_DIR)

src/llmcompressor/modifiers/quantization/calibration.py

@@ -5,6 +5,7 @@
     KVCacheScaleType,
     QuantizationScheme,
     QuantizationStatus,
+    QuantizationStrategy,
 )
 from compressed_tensors.quantization.lifecycle.forward import forward_quantize
 from compressed_tensors.quantization.utils import is_kv_cache_quant_scheme
@@ -84,14 +85,48 @@ def call_observer(module: Module, base_name: str, value: Optional[torch.Tensor]
                 "Must provide a value to observe if not using weight observer"
             )
 
+        quantization_scheme = getattr(module, "quantization_scheme", None)
+        should_calculate_gparam = False
+        should_calculate_qparams = True
+
+        # TODO: will update to be the case for both weight and input in a follow-up
+        # weight global calculate is currently done in ct right now; s
+        # should be moved here to unify global scale calculations
+        if (
+            quantization_scheme.strategy == QuantizationStrategy.TENSOR_GROUP
+            and base_name == "input"
+        ):
+            should_calculate_gparam = True
+            should_calculate_qparams = False
+
         observer = getattr(module, f"{base_name}_observer")
-        updated_scale, updated_zero_point = observer(
-            value, g_idx=g_idx, global_scale=global_scale
+        observer_outputs = observer(
+            value,
+            g_idx=g_idx,
+            global_scale=global_scale,
+            should_calculate_gparam=should_calculate_gparam,
+            should_calculate_qparams=should_calculate_qparams,
         )
 
-        # update scale and zero point
-        update_parameter_data(module, updated_scale, f"{base_name}_scale")
-        update_parameter_data(module, updated_zero_point, f"{base_name}_zero_point")
+        if should_calculate_qparams:
+            if should_calculate_gparam:
+                updated_scale, updated_zero_point, updated_global_scale = (
+                    observer_outputs
+                )
+            else:
+                updated_scale, updated_zero_point = observer_outputs
+        else:
+            updated_global_scale = observer_outputs
+
+        if should_calculate_gparam:
+            update_parameter_data(
+                module, updated_global_scale, f"{base_name}_global_scale"
+            )
+
+        if should_calculate_qparams:
+            # update scale and zero point
+            update_parameter_data(module, updated_scale, f"{base_name}_scale")
+            update_parameter_data(module, updated_zero_point, f"{base_name}_zero_point")
 
 
 def update_weight_zp_scale(module: Module):

src/llmcompressor/modifiers/quantization/quantization/mixin.py

@@ -2,6 +2,7 @@
 
 import torch
 from compressed_tensors.quantization import (
+    DynamicType,
     QuantizationArgs,
     QuantizationConfig,
     QuantizationScheme,
@@ -212,7 +213,10 @@ def _initialize_observers(self, module: torch.nn.Module):
             return
 
         scheme: QuantizationScheme = module.quantization_scheme
-        input = scheme.input_activations and not scheme.input_activations.dynamic
+        input = scheme.input_activations and scheme.input_activations.dynamic in (
+            False,
+            DynamicType.LOCAL,
+        )
         weight = scheme.weights is not None
         output = scheme.output_activations and not scheme.output_activations.dynamic
         is_attention = is_attention_module(module)
@@ -241,7 +245,10 @@ def _initialize_hooks(self, model: torch.nn.Module) -> Set[RemovableHandle]:
                 continue
 
             scheme: QuantizationScheme = module.quantization_scheme
-            input = scheme.input_activations and not scheme.input_activations.dynamic
+            input = scheme.input_activations and scheme.input_activations.dynamic in (
+                False,
+                DynamicType.LOCAL,
+            )
             output = scheme.output_activations and not scheme.output_activations.dynamic
             is_attention = is_attention_module(module)
 

src/llmcompressor/observers/base.py

@@ -73,11 +73,14 @@ def post_calculate_qparams(self) -> None:
         Run any logic specific to its observers after running calculate_qparams
         """
 
+    # TODO: use a different name?
     def get_qparams(
         self,
         observed: Optional[Tensor] = None,
         g_idx: Optional[Tensor] = None,
         global_scale: Optional[Tensor] = None,
+        should_calculate_gparam: bool = False,
+        should_calculate_qparams: bool = True,
     ) -> Tuple[FloatTensor, IntTensor]:
         """
         Convenience function to wrap overwritten calculate_qparams
@@ -101,6 +104,14 @@ def get_qparams(
                 QuantizationStrategy.TENSOR_GROUP,
                 QuantizationStrategy.GROUP,
             ):
+                # Global params are for the entire tensor
+                if should_calculate_gparam:
+                    return self.calculate_qparams(
+                        observed,
+                        should_calculate_gparam=True,
+                        should_calculate_qparams=False,
+                    )
+
                 rows = observed.shape[0]
                 columns = observed.shape[1]
                 num_groups = int(ceil(columns / group_size))
@@ -137,7 +148,7 @@ def get_qparams(
                         observed[:, start:end],
                         0,
                         tensor_id=group_index,
-                        global_scale=global_scale
+                        global_scale=global_scale,
                     )
 
                     self._scale[:, group_index] = scale.squeeze(1)

src/llmcompressor/observers/helpers.py

@@ -1,8 +1,14 @@
 from collections import Counter
+from typing import Optional
 
 import torch
+from compressed_tensors.quantization.quant_args import (
+    FP4_E2M1_DATA,
+    FP8_E4M3_DATA,
+    FloatArgs,
+)
 
-__all__ = ["get_observer_token_count"]
+__all__ = ["get_observer_token_count", "calculate_gparam"]
 
 
 def get_observer_token_count(module: torch.nn.Module) -> Counter:
@@ -20,3 +26,26 @@ def get_observer_token_count(module: torch.nn.Module) -> Counter:
                 module._num_observed_tokens
             )
     return token_counts
+
+
+def calculate_gparam(
+    updated_min_val: torch.Tensor,
+    updated_max_val: torch.Tensor,
+    scale_data: Optional[FloatArgs] = FP8_E4M3_DATA,
+    quant_data: Optional[FloatArgs] = FP4_E2M1_DATA,
+    dtype: Optional[torch.dtype] = torch.float32,
+):
+    """
+    Generate a global scale for an entire tensor (input_tensor).
+    Goal of the scale is to ensure that the quantization (local) scale
+    falls into the approproiate dtype range.
+
+    E.g. for NVFP4, group (local) scales are in dtype FP8. The global_scale
+    attempts to use the entire FP8 dtype range while mapping a per-group max
+    to the FP4 max.
+    """
+    min_vals = torch.min(updated_min_val, torch.zeros_like(updated_min_val))
+    max_vals = torch.max(updated_max_val, torch.zeros_like(updated_max_val))
+    max_val_pos = torch.max(torch.abs(min_vals), torch.abs(max_vals))
+    global_scale = scale_data.max * quant_data.max / max_val_pos
+    return global_scale.to(dtype)

src/llmcompressor/observers/min_max.py

@@ -6,6 +6,7 @@
 from compressed_tensors.utils import deprecated
 
 from llmcompressor.observers.base import Observer
+from llmcompressor.observers.helpers import calculate_gparam
 
 __all__ = ["MinMaxObserver", "MovingAverageMinMaxObserver"]
 
@@ -35,6 +36,8 @@ def calculate_qparams(
         reduce_dims: Optional[Tuple[int]] = None,
         tensor_id: Optional[Any] = None,
         global_scale: Optional[torch.Tensor] = None,
+        should_calculate_gparam: bool = False,
+        should_calculate_qparams: bool = True,
     ) -> Tuple[torch.FloatTensor, torch.IntTensor]:
         """
         Updates the observed min and max using a moving average smoothed by the
@@ -83,13 +86,24 @@ def calculate_qparams(
         self.min_val[tensor_id] = updated_min_val
         self.max_val[tensor_id] = updated_max_val
 
-        return calculate_qparams(
+        if should_calculate_gparam:
+            global_scale = calculate_gparam(
+                updated_min_val=updated_max_val, updated_max_val=updated_max_val
+            )
+            if not should_calculate_qparams:
+                return global_scale
+
+        scale, zero_point = calculate_qparams(
             min_vals=updated_min_val,
             max_vals=updated_max_val,
             quantization_args=self.quantization_args,
             global_scale=global_scale,
         )
 
+        if should_calculate_gparam:
+            return scale, zero_point, global_scale
+        return scale, zero_point
+
     def get_qparams_along_dim(
         self,
         observed: torch.Tensor,

src/llmcompressor/observers/mse.py

@@ -6,6 +6,7 @@
 from torch import FloatTensor, IntTensor, Tensor
 
 from llmcompressor.observers.base import Observer
+from llmcompressor.observers.helpers import calculate_gparam
 
 __all__ = ["MovingAverageMSEObserver"]
 
@@ -115,6 +116,8 @@ def calculate_qparams(
         reduce_dims: Optional[Tuple[int]] = None,
         tensor_id: Optional[Any] = None,
         global_scale: Optional[torch.Tensor] = None,
+        should_calculate_gparam: bool = False,
+        should_calculate_qparams: bool = True,
     ) -> Tuple[FloatTensor, IntTensor]:
         """
         Updates the mse-clipped min and max values of the observed tensor using
@@ -149,13 +152,24 @@ def calculate_qparams(
         self.min_val[tensor_id] = updated_min_val
         self.max_val[tensor_id] = updated_max_val
 
-        return calculate_qparams(
+        if should_calculate_gparam:
+            global_scale = calculate_gparam(
+                updated_min_val=updated_max_val, updated_max_val=updated_max_val
+            )
+            if not should_calculate_qparams:
+                return global_scale
+
+        scale, zero_point = calculate_qparams(
             min_vals=updated_min_val,
             max_vals=updated_max_val,
             quantization_args=self.quantization_args,
             global_scale=global_scale,
         )
 
+        if should_calculate_gparam:
+            return scale, zero_point, global_scale
+        return scale, zero_point
+
     def get_qparams_along_dim(
         self,
         observed,

src/llmcompressor/transformers/compression/quantization_format.py

@@ -61,13 +61,13 @@ def infer_quantization_format(
         )
         is_weight_only = len(input_args) == 0 and len(weight_args) > 0
 
-        if is_weight_only:  # w4a16 and w8a16
-            if (
-                weight_args[0].num_bits == 4
-                and weight_args[0].type == QuantizationType.FLOAT.value
-            ):
-                return CompressionFormat.nvfp4_pack_quantized
+        if (
+            weight_args[0].num_bits == 4
+            and weight_args[0].type == QuantizationType.FLOAT.value
+        ):
+            return CompressionFormat.nvfp4_pack_quantized
 
+        if is_weight_only:  # w4a16 and w8a16
             is_valid_pack = all(
                 weight_arg.num_bits in [4, 8]
                 and weight_arg.type == QuantizationType.INT.value