Revert "[NVFP4] Use observers to generate global weight scales " (#1507)

Reverts #1504

For some reason, automerge merged in this PR with just one approval

Author: dsikka

examples/quantization_w4a16_fp4/llama3_example.py

@@ -19,14 +19,6 @@
 # Apply quantization.
 oneshot(model=model, recipe=recipe)
 
-print("\n\n")
-print("========== SAMPLE GENERATION ==============")
-input_ids = tokenizer("Hello my name is", return_tensors="pt").input_ids.to("cuda")
-output = model.generate(input_ids, max_new_tokens=100)
-print(tokenizer.decode(output[0]))
-print("==========================================\n\n")
-
-
 # Save to disk in compressed-tensors format.
 SAVE_DIR = MODEL_ID.split("/")[1] + "-NVFP4A16"
 model.save_pretrained(SAVE_DIR, save_compressed=True)

examples/quantization_w4a4_fp4/llama3_example.py

@@ -67,14 +67,6 @@ def tokenize(sample):
     num_calibration_samples=NUM_CALIBRATION_SAMPLES,
 )
 
-print("\n\n")
-print("========== SAMPLE GENERATION ==============")
-input_ids = tokenizer("Hello my name is", return_tensors="pt").input_ids.to("cuda")
-output = model.generate(input_ids, max_new_tokens=100)
-print(tokenizer.decode(output[0]))
-print("==========================================\n\n")
-
-
 # Save to disk in compressed-tensors format.
 SAVE_DIR = MODEL_ID.split("/")[1] + "-NVFP4"
 model.save_pretrained(SAVE_DIR, save_compressed=True)

src/llmcompressor/modifiers/quantization/calibration.py

@@ -29,7 +29,6 @@
     "freeze_module_quantization",
     "apply_calibration_status",
     "reset_quantization_status",
-    "update_weight_global_scale",
 ]
 
 
@@ -67,13 +66,7 @@ def initialize_observer(
         module.register_module(f"{base_name}_observer", observer)
 
 
-def call_observer(
-    module: Module,
-    base_name: str,
-    value: Optional[torch.Tensor] = None,
-    should_calculate_gparam: bool = False,
-    should_calculate_qparams: bool = True,
-):
+def call_observer(module: Module, base_name: str, value: Optional[torch.Tensor] = None):
     """
     Call a module's attached input/weight/output observer using a provided value.
     Update the module's scale and zp using the observer's return values.
@@ -87,51 +80,54 @@ def call_observer(
         if base_name == "weight":
             value = module.weight
             g_idx = getattr(module, "weight_g_idx", None)
+            global_scale = getattr(module, f"{base_name}_global_scale", None)
         elif value is not None:
             g_idx = None
+            global_scale = None
         else:
             raise ValueError(
                 "Must provide a value to observe if not using weight observer"
             )
 
+        quantization_scheme = getattr(module, "quantization_scheme", None)
+        arg_name = "weights" if base_name == "weight" else f"{base_name}_activations"
+        quant_args = getattr(quantization_scheme, arg_name, None)
+
+        # We always calculate quantizaton parameters by default and no global parameters
+        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;
+        # should be moved here to unify global scale calculations
+        if (
+            quant_args.strategy == QuantizationStrategy.TENSOR_GROUP
+            and base_name == "input"
+        ):
+            should_calculate_gparam = True
+            should_calculate_qparams = False
+
         observer = getattr(module, f"{base_name}_observer")
+        observer_outputs = observer(
+            value,
+            g_idx=g_idx,
+            global_scale=global_scale,
+            should_calculate_gparam=should_calculate_gparam,
+        )
 
         if should_calculate_gparam:
-            global_scale = observer(
-                value,
-                should_calculate_gparam=True,
+            updated_global_scale = observer_outputs
+            update_parameter_data(
+                module, updated_global_scale, f"{base_name}_global_scale"
             )
-            update_parameter_data(module, global_scale, f"{base_name}_global_scale")
-        else:
-            global_scale = getattr(module, f"{base_name}_global_scale", None)
 
         if should_calculate_qparams:
-            updated_scale, updated_zero_point = observer(
-                value, g_idx=g_idx, global_scale=global_scale
-            )
+            # update scale and zero point
+            updated_scale, updated_zero_point = observer_outputs
             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_global_scale(module: Module):
-    if getattr_chain(module, "quantization_scheme.weights", None) is None:
-        return
-
-    if (
-        getattr_chain(module, "quantization_scheme.weights.strategy", None)
-        != QuantizationStrategy.TENSOR_GROUP
-    ):
-        return
-
-    call_observer(
-        module,
-        base_name="weight",
-        should_calculate_gparam=True,
-        should_calculate_qparams=False,
-    )
-    module.weight_observer.reset()
-
-
 def update_weight_zp_scale(module: Module):
     """
     marks a layer as ready for calibration which activates observers
@@ -169,24 +165,10 @@ def calibrate_activations(module: Module, value: torch.Tensor, base_name: str):
     if value.numel() == 0:
         return
 
-    quantization_scheme = getattr(module, "quantization_scheme", None)
-    quantization_args = getattr(quantization_scheme, f"{base_name}_activations", None)
-
-    calculate_qparams = True
-    calculate_gparam = False
-
-    if quantization_args is not None:
-        if quantization_args.dynamic in (True, DynamicType.LOCAL):
-            calculate_qparams = False
-        if quantization_args.strategy == QuantizationStrategy.TENSOR_GROUP:
-            calculate_gparam = True
-
     call_observer(
         module=module,
         base_name=base_name,
         value=value,
-        should_calculate_gparam=calculate_gparam,
-        should_calculate_qparams=calculate_qparams,
     )
 
 

src/llmcompressor/modifiers/quantization/quantization/base.py

@@ -2,12 +2,8 @@
 
 from llmcompressor.core import Event, EventType, State
 from llmcompressor.modifiers import Modifier
-from llmcompressor.modifiers.quantization.calibration import (
-    update_weight_global_scale,
-    update_weight_zp_scale,
-)
+from llmcompressor.modifiers.quantization.calibration import update_weight_zp_scale
 from llmcompressor.modifiers.quantization.quantization.mixin import QuantizationMixin
-from llmcompressor.modifiers.utils import update_fused_layer_weight_global_scales
 
 __all__ = ["QuantizationModifier"]
 
@@ -70,14 +66,7 @@ def on_start(self, state: State, event: Event, **kwargs):
         QuantizationMixin.start_calibration(self, state.model)
 
         modules = list(state.model.modules())
-        # TODO: this step can be combined with update_weight_zp_scale
-        # once update_fused_layer_weight_global_scales is removed
-        # and not required by vLLM
-        for module in tqdm.tqdm(modules):
-            update_weight_global_scale(module)
-
         for module in tqdm.tqdm(modules, desc="Calibrating weights"):
-            update_fused_layer_weight_global_scales(module)
             update_weight_zp_scale(module)
 
     def on_event(self, state: State, event: Event, **kwargs):

src/llmcompressor/modifiers/utils/__init__.py

@@ -1,4 +1,3 @@
 # flake8: noqa
 
 from .constants import *
-from .helpers import *

src/llmcompressor/modifiers/utils/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", "generate_gparam"]
 
 
 def get_observer_token_count(module: torch.nn.Module) -> Counter:
@@ -20,3 +26,29 @@ def get_observer_token_count(module: torch.nn.Module) -> Counter:
                 module._num_observed_tokens
             )
     return token_counts
+
+
+# TODO: we have a similar function in ct already
+# consolidate when adding weight global scale
+# generation
+def generate_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/helpers.py

@@ -2,10 +2,11 @@
 
 import torch
 from compressed_tensors.quantization.quant_args import QuantizationArgs
-from compressed_tensors.quantization.utils import calculate_qparams, generate_gparam
+from compressed_tensors.quantization.utils import calculate_qparams
 from compressed_tensors.utils import deprecated
 
 from llmcompressor.observers.base import Observer
+from llmcompressor.observers.helpers import generate_gparam
 
 __all__ = ["MinMaxObserver", "MovingAverageMinMaxObserver"]