Feat (proxy): scale computation delegated to bias proxy

src/brevitas/nn/mixin/base.py

@@ -147,14 +147,9 @@ def state_dict(self, destination=None, prefix='', keep_vars=False):
 
  @staticmethod
  def gate_params_fwd(gate, quant_input):
- acc_scale = None
  quant_weight_ih = gate.input_weight()
  quant_weight_hh = gate.hidden_weight()
- if isinstance(quant_input, QuantTensor) and isinstance(quant_weight_ih, QuantTensor):
- acc_scale_shape = compute_channel_view_shape(quant_input.value, channel_dim=1)
- acc_scale = quant_weight_ih.scale.view(acc_scale_shape)
- acc_scale = acc_scale * quant_input.scale.view(acc_scale_shape)
- quant_bias = gate.bias_quant(gate.bias, acc_scale)
+ quant_bias = gate.bias_quant(gate.bias, quant_input, quant_weight_ih)
  return quant_weight_ih, quant_weight_hh, quant_bias
 
  def reset_parameters(self) -> None:

src/brevitas/nn/quant_layer.py

@@ -118,14 +118,6 @@ def inner_forward_impl(self, x: Tensor, quant_weight: Tensor, quant_bias: Option
  def max_acc_bit_width(self, input_bit_width: Tensor, quant_weight_bit_width: Tensor):
  pass
 
- def quant_output_scale_impl(
- self, inp: Tensor, quant_input_scale: Tensor, quant_weight_scale: Tensor):
- channel_dim = -1 if isinstance(self, torch.nn.Linear) else 1
- output_scale_shape = compute_channel_view_shape(inp, channel_dim=channel_dim)
- output_scale = quant_weight_scale.view(output_scale_shape)
- output_scale = output_scale * quant_input_scale.view(output_scale_shape)
- return output_scale
-
  @property
  def requires_export_handler(self):
  return (
@@ -150,7 +142,6 @@ def forward_impl(self, inp: Union[Tensor, QuantTensor]) -> Union[Tensor, QuantTe
  return out
 
  quant_input = self.input_quant(inp)
-
  quant_weight = self.quant_weight(quant_input)
 
  compute_output_quant_tensor = isinstance(quant_input, QuantTensor) and isinstance(
@@ -159,12 +150,8 @@ def forward_impl(self, inp: Union[Tensor, QuantTensor]) -> Union[Tensor, QuantTe
  self.output_quant.is_quant_enabled) and self.return_quant_tensor:
  raise RuntimeError("QuantLayer is not correctly configured")
 
- output_scale = None
- if isinstance(quant_input, QuantTensor) and isinstance(quant_weight, QuantTensor):
- output_scale = self.quant_output_scale_impl(inp, quant_input.scale, quant_weight.scale)
-
  if self.bias is not None:
- quant_bias = self.bias_quant(self.bias, output_scale)
+ quant_bias = self.bias_quant(self.bias, quant_input, quant_weight)
  else:
  quant_bias = None
  output_tensor = self.inner_forward_impl(quant_input, quant_weight, quant_bias)

src/brevitas/proxy/parameter_quant.py

@@ -17,6 +17,7 @@
 from brevitas.inject import BaseInjector as Injector
 from brevitas.quant_tensor import QuantTensor
 from brevitas.utils.quant_utils import _CachedIO
+from brevitas.utils.torch_utils import compute_channel_view_shape
 
 from .quant_proxy import QuantProxyFromInjector
 from .quant_proxy import QuantProxyProtocol
@@ -234,10 +235,38 @@ def bit_width(self):
  bit_width = self.__call__(self.tracked_parameter_list[0], zhs).bit_width
  return bit_width
 
- def forward(self,
- x: Tensor,
- input_scale: Optional[Tensor] = None) -> Union[Tensor, QuantTensor]:
+ def quant_output_scale_impl(
+ self, input: QuantTensor, weight: QuantTensor, module: torch.nn.Module) -> Tensor:
+ channel_dim = -1 if isinstance(module, torch.nn.Linear) else 1
+ output_scale_shape = compute_channel_view_shape(input, channel_dim=channel_dim)
+ output_scale = weight.scale.view(output_scale_shape)
+ output_scale = output_scale * input.scale.view(output_scale_shape)
+ return output_scale
+
+ def compute_bias_scale(
+ self,
+ input: Optional[Union[Tensor, QuantTensor]],
+ weight: Optional[Union[Tensor, QuantTensor]]) -> Optional[Tensor]:
+ if not self.requires_input_scale:
+ return None
+ if not isinstance(input, QuantTensor) or not isinstance(weight, QuantTensor):
+ return None
+ if len(self.tracked_module_list) > 1:
+ if not all(
+ [type[self.tracked_module_list[0]] == type[x] for x in self.tracked_module_list]):
+ raise RuntimeError(
+ "Bias quantizer shared across different type of layers with external scale is not supported."
+ )
+ scale = self.quant_output_scale_impl(input, weight, self.tracked_module_list[0])
+ return scale
+
+ def forward(
+ self,
+ x: Tensor,
+ input: Optional[Union[Tensor, QuantTensor]] = None,
+ weight: Optional[Union[Tensor, QuantTensor]] = None) -> Union[Tensor, QuantTensor]:
  out = x
+ input_scale = self.compute_bias_scale(input, weight)
  if self.is_quant_enabled:
  impl = self.export_handler if self.export_mode else self.tensor_quant
  if self.requires_input_scale and input_scale is None: