Feat: remove quant metadata from quantlayer

notebooks/03_anatomy_of_a_quantizer.ipynb

@@ -775,7 +775,7 @@
  "cell_type": "markdown",
  "metadata": {},
  "source": [
- "Indeed we can verify that `quant_weight_scale()` is equal to `weight.abs().max()`:"
+ "Indeed we can verify that `weight_quant.scale()` is equal to `weight.abs().max()`:"
  ]
  },
  {
@@ -792,7 +792,7 @@
  }
  ],
  "source": [
- "assert_with_message((param_from_max_quant_conv.quant_weight_scale() == param_from_max_quant_conv.weight.abs().max()).item())"
+ "assert_with_message((param_from_max_quant_conv.weight_quant.scale() == param_from_max_quant_conv.weight.abs().max()).item())"
  ]
  },
  {
@@ -1024,7 +1024,7 @@
  }
  ],
  "source": [
- "assert_with_message((quant_conv1.quant_weight_scale() == quant_conv2.quant_weight_scale()).item())"
+ "assert_with_message((quant_conv1.weight_quant.scale() == quant_conv2.weight_quant.scale()).item())"
  ]
  },
  {
@@ -1059,9 +1059,9 @@
  " return module.weight.abs().mean()\n",
  " \n",
  "quant_conv1 = QuantConv2d(3, 2, (3, 3), weight_quant=SharedParamFromMeanWeightQuantizer)\n",
- "old_quant_conv1_scale = quant_conv1.quant_weight_scale()\n",
+ "old_quant_conv1_scale = quant_conv1.weight_quant.scale()\n",
  "quant_conv2 = QuantConv2d(3, 2, (3, 3), weight_quant=quant_conv1.weight_quant)\n",
- "new_quant_conv1_scale = quant_conv1.quant_weight_scale()\n",
+ "new_quant_conv1_scale = quant_conv1.weight_quant.scale()\n",
  "\n",
  "assert_with_message(not (old_quant_conv1_scale == new_quant_conv1_scale).item())"
  ]
@@ -1080,7 +1080,7 @@
  }
  ],
  "source": [
- "assert_with_message((new_quant_conv1_scale == quant_conv2.quant_weight_scale()).item())"
+ "assert_with_message((new_quant_conv1_scale == quant_conv2.weight_quant.scale()).item())"
  ]
  },
  {
@@ -1134,7 +1134,7 @@
  "quant_conv_w_init = QuantConv2d(3, 2, (3, 3), weight_quant=ParamFromMaxWeightQuantizer)\n",
  "torch.nn.init.uniform_(quant_conv_w_init.weight)\n",
  "\n",
- "assert_with_message(not (quant_conv_w_init.weight.abs().max() == quant_conv_w_init.quant_weight_scale()).item())"
+ "assert_with_message(not (quant_conv_w_init.weight.abs().max() == quant_conv_w_init.weight_quant.scale()).item())"
  ]
  },
  {
@@ -1160,7 +1160,7 @@
  "source": [
  "quant_conv_w_init.weight_quant.init_tensor_quant()\n",
  "\n",
- "assert_with_message((quant_conv_w_init.weight.abs().max() == quant_conv_w_init.quant_weight_scale()).item())"
+ "assert_with_message((quant_conv_w_init.weight.abs().max() == quant_conv_w_init.weight_quant.scale()).item())"
  ]
  },
  {

src/brevitas/export/common/handler/base.py

@@ -127,24 +127,3 @@ def zero_point_with_dtype(cls, signed, bit_width, zero_point):
  return zero_point.type(torch.int8)
  else:
  return zero_point.type(torch.int32)
-
- @classmethod
- def quant_input_zero_point(cls, module):
- signed = module.is_quant_input_signed
- zero_point = module.quant_input_zero_point()
- bit_width = module.quant_input_bit_width()
- return cls.zero_point_with_dtype(signed, bit_width, zero_point)
-
- @classmethod
- def quant_weight_zero_point(cls, module):
- signed = module.is_quant_weight_signed
- zero_point = module.quant_weight_zero_point()
- bit_width = module.quant_weight_bit_width()
- return cls.zero_point_with_dtype(signed, bit_width, zero_point)
-
- @classmethod
- def quant_output_zero_point(cls, module):
- signed = module.is_quant_output_signed
- zero_point = module.quant_output_zero_point()
- bit_width = module.quant_output_bit_width()
- return cls.zero_point_with_dtype(signed, bit_width, zero_point)

src/brevitas/export/manager.py

@@ -14,13 +14,13 @@
 from torch.nn import Module
 
 from brevitas import config
-from brevitas.nn.mixin.base import _CachedIO
 from brevitas.nn.mixin.base import QuantLayerMixin
 from brevitas.nn.mixin.base import QuantRecurrentLayerMixin
 from brevitas.proxy.quant_proxy import QuantProxyProtocol
 from brevitas.quant_tensor import QuantTensor
 from brevitas.utils.jit_utils import clear_class_registry
 from brevitas.utils.python_utils import patch
+from brevitas.utils.quant_utils import _CachedIO
 
 
 class _JitTraceExportWrapper(nn.Module):
@@ -64,18 +64,11 @@ def _override_bias_caching_mode(m: Module, enabled: bool):
  m.cache_inference_quant_bias = enabled
 
 
-def _override_inp_caching_mode(m: Module, enabled: bool):
- if hasattr(m, 'cache_inference_quant_inp'):
- if not hasattr(m, "cache_inference_quant_inp_backup"):
- m.cache_inference_quant_inp_backup = m.cache_inference_quant_inp
- m.cache_inference_quant_inp = enabled
-
-
-def _override_out_caching_mode(m: Module, enabled: bool):
- if hasattr(m, 'cache_inference_quant_out'):
- if not hasattr(m, "cache_inference_quant_out_backup"):
- m.cache_inference_quant_out_backup = m.cache_inference_quant_out
- m.cache_inference_quant_out = enabled
+def _override_act_caching_mode(m: Module, enabled: bool):
+ if hasattr(m, 'cache_inference_quant_act'):
+ if not hasattr(m, "cache_inference_quant_act_backup"):
+ m.cache_inference_quant_act_backup = m.cache_inference_quant_act
+ m.cache_inference_quant_act = enabled
 
 
 def _restore_quant_metadata_caching_mode(m: Module):
@@ -90,16 +83,10 @@ def _restore_bias_caching_mode(m: Module):
  del m.cache_inference_quant_bias_backup
 
 
-def _restore_inp_caching_mode(m: Module):
- if hasattr(m, "cache_inference_quant_inp_backup"):
- m.cache_inference_quant_inp = m.cache_inference_quant_inp_backup
- del m.cache_inference_quant_inp_backup
-
-
-def _restore_out_caching_mode(m: Module):
- if hasattr(m, "cache_inference_quant_out_backup"):
- m.cache_inference_quant_out = m.cache_inference_quant_out_backup
- del m.cache_inference_quant_out_backup
+def _restore_act_caching_mode(m: Module):
+ if hasattr(m, "cache_inference_quant_act_backup"):
+ m.cache_inference_quant_act = m.cache_inference_quant_act_backup
+ del m.cache_inference_quant_act_backup
 
 
 def _set_recurrent_layer_export_mode(model: Module, enabled: bool):
@@ -202,14 +189,12 @@ def _cache_inp_out(cls, module, *args, **kwargs):
  # force enable caching
  module.apply(lambda m: _override_quant_metadata_caching_mode(m, enabled=True))
  module.apply(lambda m: _override_bias_caching_mode(m, enabled=True))
- module.apply(lambda m: _override_inp_caching_mode(m, enabled=True))
- module.apply(lambda m: _override_out_caching_mode(m, enabled=True))
+ module.apply(lambda m: _override_act_caching_mode(m, enabled=True))
  _ = module.forward(*args, **kwargs)
  # Restore previous caching properties
  module.apply(lambda m: _restore_quant_metadata_caching_mode(m))
  module.apply(lambda m: _restore_bias_caching_mode(m))
- module.apply(lambda m: _restore_inp_caching_mode(m))
- module.apply(lambda m: _restore_out_caching_mode(m))
+ module.apply(lambda m: _restore_act_caching_mode(m))
 
  @classmethod
  def jit_inference_trace(

src/brevitas/export/onnx/manager.py

@@ -24,8 +24,8 @@
 from brevitas import torch_version
 from brevitas.quant_tensor import QuantTensor
 
-from ..manager import _override_inp_caching_mode
-from ..manager import _restore_inp_caching_mode
+from ..manager import _override_act_caching_mode
+from ..manager import _restore_act_caching_mode
 from ..manager import BaseManager
 from ..manager import ExportContext
 
@@ -120,7 +120,7 @@ def export_onnx(
  # enable export mode, this triggers collecting export values into handlers
  cls.set_export_mode(module, enabled=True)
  # temporarily disable input caching to avoid collectives empty debug values
- module.apply(lambda m: _override_inp_caching_mode(m, enabled=False))
+ module.apply(lambda m: _override_act_caching_mode(m, enabled=False))
  # perform export pass
  if export_path is not None:
  export_target = export_path
@@ -130,7 +130,7 @@ def export_onnx(
  torch.onnx.export(module, args, export_target, **onnx_export_kwargs)
 
  # restore the model to previous properties
- module.apply(lambda m: _restore_inp_caching_mode(m))
+ module.apply(lambda m: _restore_act_caching_mode(m))
  cls.set_export_mode(module, enabled=False)
  module.train(training_state)
 

src/brevitas/graph/gpfq.py

@@ -236,7 +236,7 @@ def update_batch(self, module, input, current_layer):
  raise StopFwdException
 
  def single_layer_update(self):
- assert not self.layer.weight_quant_requires_quant_input, "Error: GPFQ does not support weight quantizers that require quantized inputs."
+ assert not self.layer.weight_quant.requires_quant_input, "Error: GPFQ does not support weight quantizers that require quantized inputs."
  weight = self.layer.weight.data
  dev = weight.device
  dtype = weight.dtype
@@ -360,7 +360,7 @@ def single_layer_update(self):
  input_is_signed = self.quant_input.signed
  T = get_upper_bound_on_l1_norm(
  torch.tensor(self.accumulator_bit_width), input_bit_width, input_is_signed)
- s = self.layer.quant_weight_scale()
+ s = self.layer.weight_quant.scale()
  if s.ndim > 1:
  s = s.view(self.groups, -1) # [Groups, OC/Groups]
 

src/brevitas/graph/gptq.py

@@ -194,7 +194,7 @@ def update_batch(self, module, input, current_layer):
  raise StopFwdException
 
  def single_layer_update(self, percdamp=.01):
- assert not self.layer.weight_quant_requires_quant_input, "Error: GPTQ does not support weight quantizers that require quantized inputs."
+ assert not self.layer.weight_quant.requires_quant_input, "Error: GPTQ does not support weight quantizers that require quantized inputs."
  if hasattr(self.layer, 'allocate_params'):
  self.layer.allocate_params(self.layer)
  weight = self.layer.weight.data

src/brevitas/graph/quantize.py

@@ -237,10 +237,10 @@ def align_input_quant(
  # If it is a QuantIdentity already, simply modify tensor_quant or the scaling implementations
  # based on whether we need to align the sign or not
  if isinstance(module, qnn.QuantIdentity):
- if align_sign or module.is_quant_act_signed == shared_quant_identity.is_quant_act_signed:
+ if align_sign or module.input_quant.is_signed == shared_quant_identity.input_quant.is_signed:
  return shared_quant_identity
  else:
- assert not module.is_quant_act_signed and shared_quant_identity.is_quant_act_signed
+ assert not module.input_quant.is_signed and shared_quant_identity.input_quant.is_signed
  quant_module_class, quant_module_kwargs = quant_identity_map['unsigned']
  return (
  quant_module_class,

src/brevitas/graph/quantize_impl.py

@@ -80,8 +80,8 @@ def are_inputs_unsigned(model, node, is_unsigned_list, quant_act_map, unsigned_a
  elif isinstance(inp_module, tuple(SIGN_PRESERVING_MODULES)):
  are_inputs_unsigned(
  model, inp_node, is_unsigned_list, quant_act_map, unsigned_act_tuple)
- elif hasattr(inp_module, 'is_quant_act_signed'):
- is_unsigned_list.append(not inp_module.is_quant_act_signed)
+ elif hasattr(inp_module, 'input_quant'):
+ is_unsigned_list.append(not inp_module.input_quant.is_signed)
  else:
  is_unsigned_list.append(False)
  elif inp_node.op == 'call_function':

src/brevitas/nn/mixin/act.py

@@ -4,13 +4,11 @@
 from abc import ABCMeta
 from abc import abstractmethod
 from typing import Optional, Type, Union
-from warnings import warn
 
 from torch.nn import Module
 
 from brevitas.inject import ExtendedInjector
 from brevitas.inject import Injector
-from brevitas.proxy.runtime_quant import ActQuantProxyFromInjector
 from brevitas.proxy.runtime_quant import ActQuantProxyProtocol
 from brevitas.quant import NoneActQuant
 
@@ -34,31 +32,6 @@ def __init__(self, input_quant: Optional[ActQuantType], **kwargs):
  input_passthrough_act=True,
  **kwargs)
 
- @property
- def is_input_quant_enabled(self):
- return self.input_quant.is_quant_enabled
-
- @property
- def is_quant_input_narrow_range(self): # TODO make abstract once narrow range can be cached
- return self.input_quant.is_narrow_range
-
- @property
- @abstractmethod
- def is_quant_input_signed(self):
- pass
-
- @abstractmethod
- def quant_input_scale(self):
- pass
-
- @abstractmethod
- def quant_input_zero_point(self):
- pass
-
- @abstractmethod
- def quant_input_bit_width(self):
- pass
-
 
 class QuantOutputMixin(QuantProxyMixin):
  __metaclass__ = ABCMeta
@@ -75,31 +48,6 @@ def __init__(self, output_quant: Optional[ActQuantType], **kwargs):
  output_passthrough_act=True,
  **kwargs)
 
- @property
- def is_output_quant_enabled(self):
- return self.output_quant.is_quant_enabled
-
- @property
- def is_quant_output_narrow_range(self): # TODO make abstract once narrow range can be cached
- return self.output_quant.is_narrow_range
-
- @property
- @abstractmethod
- def is_quant_output_signed(self):
- pass
-
- @abstractmethod
- def quant_output_scale(self):
- pass
-
- @abstractmethod
- def quant_output_zero_point(self):
- pass
-
- @abstractmethod
- def quant_output_bit_width(self):
- pass
-
 
 class QuantNonLinearActMixin(QuantProxyMixin):
  __metaclass__ = ABCMeta
@@ -124,28 +72,3 @@ def __init__(
  none_quant_injector=NoneActQuant,
  **prefixed_kwargs,
  **kwargs)
-
- @property
- def is_act_quant_enabled(self):
- return self.act_quant.is_quant_enabled
-
- @property
- def is_quant_act_narrow_range(self): # TODO make abstract once narrow range can be cached
- return self.act_quant.is_narrow_range
-
- @property
- @abstractmethod
- def is_quant_act_signed(self):
- pass
-
- @abstractmethod
- def quant_act_scale(self):
- pass
-
- @abstractmethod
- def quant_act_zero_point(self):
- pass
-
- @abstractmethod
- def quant_act_bit_width(self):
- pass