Minor refactoring

Author: pablomlago

src/brevitas/utils/python_utils.py

@@ -12,6 +12,7 @@
 from typing import Optional
 from typing import TypeVar
 from typing import Union
+import warnings
 
 
 class AutoName(str, Enum):
@@ -83,9 +84,16 @@ def __init__(self, registry_name: Optional[str] = None) -> None:
         self._registry_name = registry_name
         self._registry: Dict[str, T] = {}
 
+    @staticmethod
+    def register(
+        registry: "Registry[T]",
+        names: Union[str, List[str]],
+    ) -> Callable[[T], T]:
+        return registry.register(names)
+
     @property
     def registry_name(self) -> str:
-        return "registry" if self._registry_name is None else self._registry_name
+        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):
@@ -95,7 +103,9 @@ 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}.")
+                    warnings.warn(
+                        f"'{name}' is already registered in {self.registry_name}. Overwriting the existing value."
+                    )
                 self._registry[name] = value
             return value
 

src/brevitas_examples/common/generative/quantizers.py

@@ -7,7 +7,7 @@
 from typing import Dict
 from typing import Optional
 from typing import Type
-from typing import TypeVar
+from typing import TypeAlias
 
 from torch import nn
 
@@ -52,6 +52,34 @@
 
 from .quant_blocks import *
 
+# Prevents Pylance from raising "Variable not allowed in type expression" error in every type hint in BaseQuantizer
+QuantInjector: TypeAlias = ExtendedInjector  # type: ignore
+
+
+class BaseQuantizer:
+    weight_quant: ClassVar[Optional[QuantInjector]] = None
+    linear_input_quant: ClassVar[Optional[QuantInjector]] = None
+    input_quant: ClassVar[Optional[QuantInjector]] = None
+    q_scaled_quant: ClassVar[Optional[QuantInjector]] = None
+    k_transposed_quant: ClassVar[Optional[QuantInjector]] = None
+    v_quant: ClassVar[Optional[QuantInjector]] = None
+    attn_output_weights_quant: ClassVar[Optional[QuantInjector]] = None
+
+    @classmethod
+    def override_quantizers_dict(
+            cls: "BaseQuantizer",
+            quantizers_dict: Dict[str,
+                                  Optional[QuantInjector]]) -> Dict[str, Optional[QuantInjector]]:
+        # Overrides the quantizers in the input dictionary
+        for key in quantizers_dict:
+            if (value := getattr(cls, key)) is not None:
+                quantizers_dict[key] = value
+        return quantizers_dict
+
+
+# Registry for custom quantizers
+CUSTOM_QUANTIZERS_REGISTRY = Registry[Type[BaseQuantizer]](registry_name="CustomQuantizersRegistry")
+
 
 class DynamicActProxyMixin(ExtendedInjector):
     proxy_class = DynamicActQuantProxyFromInjector
@@ -226,31 +254,3 @@ 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/stable_diffusion/main.py

@@ -530,7 +530,7 @@ def sdpa_zp_stats_type():
             input_kwargs=input_kwargs)
 
         layer_map = generate_quant_maps(
-            *quantizers, dtype=dtype, device=args.device, quantize_embedding=False)
+            **quantizers, dtype=dtype, device=args.device, quantize_embedding=False)
 
         linear_qkwargs = layer_map[torch.nn.Linear][1]
         linear_qkwargs[

tests/brevitas/utils/test_python_utils.py

@@ -3,7 +3,10 @@
 
 from enum import auto
 
+import pytest
+
 from brevitas.utils.python_utils import AutoName
+from brevitas.utils.python_utils import Registry
 
 
 class TestEnum(AutoName):
@@ -38,3 +41,77 @@ def test_eq_enum():
 
 def test_neq_enum():
     assert TestEnum.FIRST != TestEnum.SECOND
+
+
+class TestRegistry:
+
+    def test_register_single_name(self):
+        r = Registry()
+
+        @r.register("k")
+        class Dummy:
+            pass
+
+        dummy = r.get("k")
+
+        assert dummy is Dummy
+        assert len(r.get_registered_keys()) == 1
+        assert next(iter(r.get_registered_keys())) == "k"
+
+    def test_register_multiple_names(self):
+        r = Registry()
+
+        @r.register(["k1", "k2"])
+        class Dummy:
+            pass
+
+        dummy1 = r.get("k1")
+        dummy2 = r.get("k2")
+
+        assert dummy1 is Dummy
+        assert dummy2 is Dummy
+        assert len(r.get_registered_keys()) == 2
+        assert set(r.get_registered_keys()) == {"k1", "k2"}
+
+    def test_register_single_name_static(self):
+        r = Registry()
+
+        @Registry.register(r, "k")
+        class Dummy:
+            pass
+
+        dummy = r.get("k")
+
+        assert dummy is Dummy
+        assert len(r.get_registered_keys()) == 1
+        assert next(iter(r.get_registered_keys())) == "k"
+
+    def test_register_duplicate_raises_warning(self):
+        r = Registry("TestRegistry")
+
+        r.register("dup")("k")
+        # Patch warnings.warn and check that it is called with the expected message
+        with pytest.warns(UserWarning) as record:
+            r.register("dup")("k")
+        msg = str(record[0].message)
+        assert "'dup' is already registered in TestRegistry. Overwriting the existing value." == msg
+
+    def test_get_missing_empty_raises_valueerror(self):
+        r = Registry("TestRegistry")
+
+        with pytest.raises(ValueError) as excinfo:
+            r.get("missing")
+
+        msg = str(excinfo.value)
+        assert msg == "'missing' not found in TestRegistry. The available values are: <empty>"
+
+    def test_get_missing_raises_valueerror(self):
+        r = Registry("TestRegistry")
+
+        r.register("k1")("v1")
+        r.register("k2")("v2")
+        with pytest.raises(ValueError) as excinfo:
+            r.get("missing")
+
+        msg = str(excinfo.value)
+        assert msg == "'missing' not found in TestRegistry. The available values are: k1, k2"

tests/brevitas_examples/test_llm_cases.py

@@ -27,7 +27,7 @@ class LLMRunCases:
             "mistral",  #"mixtral",
             "opt",],
     )
-    def case_small_models_with_ppl(self, run_dict, default_run_args, request):
+    def case_small_models_run(self, run_dict, default_run_args, request):
         yield process_args_and_metrics(default_run_args, run_dict)
 
     # yapf: disable
@@ -93,6 +93,26 @@ def case_small_models_toggle_args(self, run_dict, default_run_args, request):
             pytest.skip(reason=f'MSE as weight_param_method requires JIT to be disabled')
         yield process_args_and_metrics(default_run_args, run_dict)
 
+    @pytest_cases.parametrize(
+        "run_dict",
+        [
+            {
+                "model": "hf-internal-testing/tiny-random-LlamaForCausalLM",
+                "custom_quantizer": "example_int8_weight_quant"},],
+        ids=[
+            "llama",]
+    )
+    def case_small_models_custom_quantizer(self, run_dict, default_run_args, request):
+        from brevitas.quant.scaled_int import Int8WeightPerTensorFloat
+        from brevitas.utils.python_utils import Registry
+        from brevitas_examples.common.generative.quantizers import BaseQuantizer
+        from brevitas_examples.common.generative.quantizers import CUSTOM_QUANTIZERS_REGISTRY
+        @Registry.register(CUSTOM_QUANTIZERS_REGISTRY, "example_int8_weight_quant")
+        class ExampleInt8WeightQuantizer(BaseQuantizer):
+            weight_quant = Int8WeightPerTensorFloat
+        yield process_args_and_metrics(default_run_args, run_dict)
+
+
 class LLMPerplexityCases:
 
     METRICS = ["float_ppl", "quant_ppl"]