Fix (core): po2 for float quantization (#1033)

src/brevitas/core/quant/float.py

@@ -68,10 +68,6 @@ def __init__(
 
  @brevitas.jit.script_method
  def quantize(self, x: torch.Tensor, scale: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
- if self.float_scaling_impl is not None:
- float_scaling_impl_value = self.float_scaling_impl(
- self.exponent_bit_width(), self.mantissa_bit_width(), self.exponent_bias())
- scale = scale / float_scaling_impl_value
  x = self.input_view_impl(x)
  scaled_x = x / scale
  internal_scale = float_internal_scale(
@@ -85,7 +81,12 @@ def dequantize(self, y, scale):
 
  @brevitas.jit.script_method
  def forward(self, x):
- scale = self.scaling_impl(x)
+ if self.float_scaling_impl is not None:
+ float_scaling_impl_value = self.float_scaling_impl(
+ self.exponent_bit_width(), self.mantissa_bit_width(), self.exponent_bias())
+ else:
+ float_scaling_impl_value = None
+ scale = self.scaling_impl(x, float_scaling_impl_value)
  if self.observer_only:
  y = x
  saturating, inf_values, nan_values = self.float_clamp_impl.saturating, self.float_clamp_impl.inf_values, self.float_clamp_impl.nan_values

src/brevitas/core/quant/int.py

@@ -150,9 +150,8 @@ def __init__(
  @brevitas.jit.script_method
  def forward(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
  bit_width = self.msb_clamp_bit_width_impl()
- threshold = self.scaling_impl(x)
  int_threshold = self.int_scaling_impl(bit_width)
- scale = threshold / int_threshold
+ scale = self.scaling_impl(x, int_threshold)
  zero_point = self.zero_point_impl(x, scale, bit_width)
  if self.observer_only:
  y = x
@@ -189,8 +188,7 @@ def forward(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Te
  pre_threshold = self.pre_scaling_impl(x)
  pre_scale = pre_threshold / int_threshold
  pre_zero_point = self.pre_zero_point_impl(x, pre_scale, bit_width)
- threshold = self.scaling_impl(x)
- scale = threshold / int_threshold
+ scale = self.scaling_impl(x, int_threshold)
  zero_point = self.zero_point_impl(x, scale, bit_width)
  if self.observer_only:
  y = x
@@ -258,8 +256,7 @@ def forward(self, x: Tensor, input_bit_width: Tensor,
  pre_threshold = self.pre_scaling_impl(x, input_bit_width, input_is_signed)
  pre_scale = pre_threshold / int_threshold
  pre_zero_point = self.pre_zero_point_impl(x, pre_scale, bit_width)
- threshold = self.scaling_impl(x)
- scale = threshold / int_threshold
+ scale = self.scaling_impl(x, int_threshold)
  zero_point = self.zero_point_impl(x, scale, bit_width)
  if self.observer_only:
  y = x

src/brevitas/core/restrict_val.py

@@ -36,7 +36,7 @@ def __init__(self, scaling_min_val: Optional[float], restrict_value_impl: Option
  self.restrict_value_impl = Identity()
 
  @brevitas.jit.script_method
- def forward(self, x: torch.Tensor):
+ def forward(self, x: Tensor):
  x = self.restrict_value_impl(x)
  x = self.clamp_min_ste(x)
  return x
@@ -52,7 +52,7 @@ def __init__(self, restrict_value_impl: Optional[Module]):
  self.restrict_value_impl = Identity()
 
  @brevitas.jit.script_method
- def forward(self, x: torch.Tensor):
+ def forward(self, x: Tensor):
  x = self.restrict_value_impl(x)
  return x
 
@@ -68,7 +68,7 @@ def __init__(self, scaling_min_val: Optional[float]):
  self.min_val = scaling_min_val
 
  @brevitas.jit.script_method
- def forward(self, x: torch.Tensor):
+ def forward(self, x: Tensor):
  x = self.clamp_min_ste(x)
  return x
 
@@ -90,8 +90,11 @@ def restrict_init_module(self):
  def restrict_init_inplace_module(self):
  return Identity()
 
+ def combine_scale_threshold(self, x: Tensor, threshold: Tensor) -> Tensor:
+ return x / threshold
+
  @brevitas.jit.script_method
- def forward(self, x: torch.Tensor) -> Tensor:
+ def forward(self, x: Tensor) -> Tensor:
  return x
 
 
@@ -104,7 +107,7 @@ def __init__(self):
  def restrict_init_float(self, x: float):
  return math.log2(x)
 
- def restrict_init_tensor(self, x: torch.Tensor):
+ def restrict_init_tensor(self, x: Tensor):
  return torch.log2(x)
 
  def restrict_init_module(self):
@@ -113,8 +116,11 @@ def restrict_init_module(self):
  def restrict_init_inplace_module(self):
  return InplaceLogTwo()
 
+ def combine_scale_threshold(self, x: Tensor, threshold: Tensor) -> Tensor:
+ return x - threshold
+
  @brevitas.jit.script_method
- def forward(self, x: torch.Tensor):
+ def forward(self, x: Tensor):
  x = self.power_of_two(x)
  return x
 
@@ -128,7 +134,7 @@ def __init__(self, restrict_value_float_to_int_impl: Module = RoundSte()):
  def restrict_init_float(self, x: float):
  return x
 
- def restrict_init_tensor(self, x: torch.Tensor):
+ def restrict_init_tensor(self, x: Tensor):
  return x
 
  def restrict_init_module(self):
@@ -137,8 +143,11 @@ def restrict_init_module(self):
  def restrict_init_inplace_module(self):
  return Identity()
 
+ def combine_scale_threshold(self, x: Tensor, threshold: Tensor) -> Tensor:
+ return x / threshold
+
  @brevitas.jit.script_method
- def forward(self, x: torch.Tensor):
+ def forward(self, x: Tensor):
  x = self.float_to_int_impl(x)
  return x
 
@@ -153,7 +162,7 @@ def __init__(self, restrict_value_float_to_int_impl: Module = RoundSte()):
  def restrict_init_float(self, x: float):
  return math.log2(x)
 
- def restrict_init_tensor(self, x: torch.Tensor):
+ def restrict_init_tensor(self, x: Tensor):
  return torch.log2(x)
 
  def restrict_init_module(self):
@@ -162,8 +171,11 @@ def restrict_init_module(self):
  def restrict_init_inplace_module(self):
  return InplaceLogTwo()
 
+ def combine_scale_threshold(self, x: Tensor, threshold: Tensor) -> Tensor:
+ return x - threshold
+
  @brevitas.jit.script_method
- def forward(self, x: torch.Tensor):
+ def forward(self, x: Tensor):
  x = self.float_to_int_impl(x)
  x = self.power_of_two(x)
  return x

src/brevitas/core/scaling/runtime.py

@@ -11,6 +11,7 @@
 import brevitas.config as config
 from brevitas.core.function_wrapper import Identity
 from brevitas.core.restrict_val import _RestrictClampValue
+from brevitas.core.restrict_val import FloatRestrictValue
 from brevitas.core.stats import _ParameterListStats
 from brevitas.core.stats import _RuntimeStats
 from brevitas.core.stats import DEFAULT_MOMENTUM
@@ -27,8 +28,8 @@ def __init__(
  scaling_stats_input_view_shape_impl: Module,
  scaling_stats_input_concat_dim: int,
  tracked_parameter_list: List[torch.nn.Parameter],
- restrict_scaling_impl: Module,
  scaling_shape: Tuple[int, ...],
+ restrict_scaling_impl: Module = FloatRestrictValue(),
  affine_rescaling: bool = False,
  affine_shift_scale: bool = False,
  scaling_min_val: Optional[float] = None,
@@ -51,9 +52,12 @@ def __init__(
  device)
 
  @brevitas.jit.script_method
- def forward(self, ignored: torch.Tensor) -> torch.Tensor:
+ def forward(
+ self, ignored: torch.Tensor, threshold: Optional[torch.Tensor] = None) -> torch.Tensor:
  stats = self.parameter_list_stats()
- return self.stats_scaling_impl(stats)
+ if threshold is None:
+ threshold = torch.ones(1).type_as(stats)
+ return self.stats_scaling_impl(stats, threshold)
 
 
 class _StatsScaling(brevitas.jit.ScriptModule):
@@ -78,10 +82,16 @@ def __init__(
  self.affine_rescaling = Identity()
  self.restrict_clamp_scaling = _RestrictClampValue(scaling_min_val, restrict_scaling_impl)
  self.restrict_scaling_pre = restrict_scaling_impl.restrict_init_module()
+ self.restrict_scaling_impl = restrict_scaling_impl
 
  @brevitas.jit.script_method
- def forward(self, stats: torch.Tensor) -> torch.Tensor:
+ def forward(
+ self, stats: torch.Tensor, threshold: Optional[torch.Tensor] = None) -> torch.Tensor:
+ if threshold is None:
+ threshold = torch.ones(1).type_as(stats)
+ threshold = self.restrict_scaling_pre(threshold)
  stats = self.restrict_scaling_pre(stats)
+ stats = self.restrict_scaling_impl.combine_scale_threshold(stats, threshold)
  stats = self.affine_rescaling(stats)
  stats = self.restrict_clamp_scaling(stats)
  return stats
@@ -93,10 +103,10 @@ def __init__(
  self,
  scaling_stats_impl: Module,
  scaling_stats_input_view_shape_impl: Module,
- restrict_scaling_impl: Module,
  scaling_shape: Tuple[int, ...],
  affine_rescaling: bool = False,
  affine_shift_scale: bool = False,
+ restrict_scaling_impl: Module = FloatRestrictValue(),
  scaling_stats_momentum: float = DEFAULT_MOMENTUM,
  scaling_min_val: Optional[float] = None,
  dtype: Optional[torch.dtype] = None,
@@ -120,9 +130,9 @@ def __init__(
  device)
 
  @brevitas.jit.script_method
- def forward(self, x: torch.Tensor):
+ def forward(self, x: torch.Tensor, threshold: Optional[torch.Tensor] = None) -> torch.Tensor:
  stats = self.runtime_stats(x)
- return self.stats_scaling_impl(stats)
+ return self.stats_scaling_impl(stats, threshold)
 
 
 class _AffineRescaling(brevitas.jit.ScriptModule):
@@ -163,13 +173,13 @@ def _load_from_state_dict(
 class RuntimeDynamicGroupStatsScaling(brevitas.jit.ScriptModule):
 
  def __init__(
-  self,
-  group_size: int,
-  group_dim: int,
-  input_view_impl: torch.nn.Module,
-  scaling_stats_impl: torch.nn.Module,
-  scaling_min_val: Optional[float],
-  restrict_scaling_impl: Optional[torch.nn.Module]) -> None:
+ self,
+ group_size: int,
+ group_dim: int,
+ input_view_impl: Module,
+ scaling_stats_impl: Module,
+ scaling_min_val: Optional[float],
+ restrict_scaling_impl: Module = FloatRestrictValue()) -> None:
  super(RuntimeDynamicGroupStatsScaling, self).__init__()
  self.group_size = group_size
  self.group_dim = group_dim
@@ -179,9 +189,14 @@ def __init__(
  self.restrict_clamp_scaling = _RestrictClampValue(scaling_min_val, restrict_scaling_impl)
 
  @brevitas.jit.script_method
- def forward(self, stats_input) -> torch.Tensor:
+ def forward(
+ self,
+ stats_input: torch.Tensor,
+ threshold: Optional[torch.Tensor] = None) -> torch.Tensor:
+ if threshold is None:
+ threshold = torch.ones(1).type_as(stats_input)
  stats_input_reshaped = self.input_view_impl(stats_input)
- out = self.scaling_stats_impl(stats_input_reshaped)
+ out = self.scaling_stats_impl(stats_input_reshaped) / threshold
  # Scaling min val
  out = self.restrict_clamp_scaling(out)
  return out