Proposal custom quantizer

Author: pablomlago

src/brevitas/utils/python_utils.py

@@ -4,6 +4,14 @@
 from contextlib import contextmanager
 from enum import Enum
 import functools
+from typing import Callable
+from typing import Dict
+from typing import Generic
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import TypeVar
+from typing import Union
 
 
 class AutoName(str, Enum):
@@ -64,3 +72,43 @@ def run(*args, **kwargs):
         return function(*args, **kwargs)
 
     return run
+
+
+T = TypeVar("T")
+
+
+class Registry(Generic[T]):
+
+    def __init__(self, registry_name: Optional[str] = None) -> None:
+        self._registry_name = registry_name
+        self._registry: Dict[str, T] = {}
+
+    @property
+    def registry_name(self) -> str:
+        return "registry" if self._registry_name is None else self._registry_name
+
+    def register(self, names: Union[str, List[str]]) -> Callable[[T], T]:
+        if isinstance(names, str):
+            names = [names]
+
+        def decorator(value: T) -> T:
+            # Allow registering the same value to multiple keys
+            for name in names:
+                if name in self._registry:
+                    raise ValueError(f"'{name}' is already registered in {self.registry_name}.")
+                self._registry[name] = value
+            return value
+
+        return decorator
+
+    def get_registered_keys(self) -> Iterable[str]:
+        return self._registry.keys()
+
+    def get(self, name: str) -> T:
+        try:
+            return self._registry[name]
+        except KeyError:
+            available = ", ".join(sorted(self._registry)) or "<empty>"
+            raise ValueError(
+                f"'{name}' not found in {self.registry_name}. The available values are: {available}"
+            )

src/brevitas_examples/common/generative/quantize.py

@@ -468,7 +468,14 @@ def quant_format_from_string(quant_format):
             linear_input_quant = linear_input_quant.let(
                 **{
                     'group_dim': -1, 'group_size': input_group_size})
-    return linear_input_quant, weight_quant, input_quant, q_scaled_quant, k_transposed_quant, v_quant, attn_output_weights_quant
+    return {
+        'linear_input_quant': linear_input_quant,
+        'weight_quant': weight_quant,
+        'input_quant': input_quant,
+        'q_scaled_quant': q_scaled_quant,
+        'k_transposed_quant': k_transposed_quant,
+        'v_quant': v_quant,
+        'attn_output_weights_quant': attn_output_weights_quant}
 
 
 def generate_quant_maps(

src/brevitas_examples/common/generative/quantizers.py

@@ -3,6 +3,12 @@
 # SPDX-License-Identifier: BSD-3-Clause
 """
 
+from typing import ClassVar
+from typing import Dict
+from typing import Optional
+from typing import Type
+from typing import TypeVar
+
 from torch import nn
 
 from brevitas.core.function_wrapper.ops_ste import FloorSte
@@ -39,8 +45,10 @@
 from brevitas.quant.scaled_int import Int8ActPerTensorFloat
 from brevitas.quant.scaled_int import Int8WeightPerChannelFloat
 from brevitas.quant.scaled_int import Int8WeightPerChannelFloatHQO
+from brevitas.quant.scaled_int import Int8WeightPerTensorFloat
 from brevitas.quant.shifted_scaled_int import ShiftedUint8ActPerTensorFloat
 from brevitas.quant.shifted_scaled_int import ShiftedUint8WeightPerChannelFloat
+from brevitas.utils.python_utils import Registry
 
 from .quant_blocks import *
 
@@ -218,3 +226,31 @@ class FP8e4m3FNUZDynamicActPerRowFloat(Fp8e4m3FNUZActPerTensorFloat):
 
 class Fp8e4m3WeightPerChannelFloatMSE(MSESymmetricScale, Fp8e4m3WeightPerChannelFloat):
     pass
+
+
+# TODO: Subject to change
+class BaseQuantizer:
+    weight_quant: ClassVar[Optional[ExtendedInjector]] = None  # type: ignore
+    linear_input_quant: ClassVar[Optional[ExtendedInjector]] = None  # type: ignore
+    input_quant: ClassVar[Optional[ExtendedInjector]] = None  # type: ignore
+    q_scaled_quant: ClassVar[Optional[ExtendedInjector]] = None  # type: ignore
+    k_transposed_quant: ClassVar[Optional[ExtendedInjector]] = None  # type: ignore
+    v_quant: ClassVar[Optional[ExtendedInjector]] = None  # type: ignore
+    attn_output_weights_quant: ClassVar[Optional[ExtendedInjector]] = None  # type: ignore
+
+    @classmethod
+    def override_quantizers_dict(
+            cls: "BaseQuantizer",
+            quantizers_dict: Dict[str, Optional[ExtendedInjector]]):  # type: ignore
+        for key in quantizers_dict:
+            if hasattr(cls, key) and (value := getattr(cls, key)) is not None:
+                quantizers_dict[key] = value
+        return quantizers_dict
+
+
+CUSTOM_QUANTIZERS_REGISTRY = Registry[Type[BaseQuantizer]](registry_name="CustomQuantizersRegistry")
+
+
+@CUSTOM_QUANTIZERS_REGISTRY.register("custom_quant")
+class CustomQuantizerExample(BaseQuantizer):
+    weight_quant: ClassVar[Optional[ExtendedInjector]] = Int8WeightPerTensorFloat  # type: ignore

src/brevitas_examples/llm/llm_args.py

@@ -20,6 +20,13 @@ def create_args_parser() -> ArgumentParser:
         type=str,
         default="facebook/opt-125m",
         help='HF model name. Default: facebook/opt-125m.')
+    parser.add_argument(
+        '--custom-quantizer',
+        type=str,
+        default=None,
+        help=
+        'Override the quantization list with custom user defined quantizers. This must be a .py file with a list of seven quantizers. Default: None.'
+    )
     parser.add_argument(
         '--dtype',
         type=str,

src/brevitas_examples/llm/main.py

@@ -34,6 +34,7 @@
 from brevitas_examples.common.accelerate_utils.accelerate import update_internal_dict
 from brevitas_examples.common.generative.quantize import generate_quant_maps
 from brevitas_examples.common.generative.quantize import generate_quantizers
+from brevitas_examples.common.generative.quantizers import CUSTOM_QUANTIZERS_REGISTRY
 from brevitas_examples.common.parse_utils import override_defaults
 from brevitas_examples.common.parse_utils import parse_args
 from brevitas_examples.llm.gguf_export.export import save_quantized_as_gguf
@@ -411,7 +412,7 @@ def quantize_llm(args, extra_args=None):
                     'zero_point_affine_rescaling_init': args.weight_quant_rescaling_init}}
         if args.weight_narrow_range:
             weight_kwargs = {**weight_kwargs, **{'narrow_range': args.weight_narrow_range}}
-        linear_input_quant, weight_quant, input_quant, q_scaled_quant, k_transposed_quant, v_quant, attn_output_weights_quant = generate_quantizers(
+        quantizers_dict = generate_quantizers(
             weight_bit_width=args.weight_bit_width,
             weight_param_method=args.weight_param_method,
             weight_scale_precision=args.weight_scale_precision,
@@ -444,17 +445,11 @@ def quantize_llm(args, extra_args=None):
             quant_attn_mode='sdpa',
             scaling_min_val=args.scaling_min_val,
             weight_kwargs=weight_kwargs)
+        if args.custom_quantizer is not None:
+            custom_quantizer = CUSTOM_QUANTIZERS_REGISTRY.get(args.custom_quantizer)
+            quantizers_dict = custom_quantizer.override_quantizers_dict(quantizers_dict)
         layer_map = generate_quant_maps(
-            linear_input_quant=linear_input_quant,
-            weight_quant=weight_quant,
-            input_quant=input_quant,
-            q_scaled_quant=q_scaled_quant,
-            k_transposed_quant=k_transposed_quant,
-            v_quant=v_quant,
-            attn_output_weights_quant=attn_output_weights_quant,
-            dtype=dtype,
-            device=device,
-            quantize_embedding=False)
+            **quantizers_dict, dtype=dtype, device=device, quantize_embedding=False)
         if not args.quantize_last_layer:
             # Dynamo tracing changes the name of the modules, thus we need this workaround to pick
             # up the last module.