Feat (core): use runtime parameter for scale

src/brevitas/core/scaling/runtime.py

@@ -51,8 +51,8 @@ def __init__(
  device)
 
  @brevitas.jit.script_method
- def forward(self, ignored: torch.Tensor) -> torch.Tensor:
- stats = self.parameter_list_stats()
+ def forward(self, x: Optional[torch.Tensor]) -> torch.Tensor:
+ stats = self.parameter_list_stats(x)
  return self.stats_scaling_impl(stats)
 
 

src/brevitas/core/scaling/standalone.py

@@ -201,12 +201,12 @@ def __init__(
  self.value = Parameter(torch.full(scaling_shape, 1.0, dtype=dtype, device=device))
 
  @brevitas.jit.script_method
- def forward(self, ignored: torch.Tensor) -> torch.Tensor:
+ def forward(self, x: Optional[torch.Tensor]) -> torch.Tensor:
  if self.init_done:
  value = abs_binary_sign_grad(self.stats_scaling_impl.restrict_clamp_scaling(self.value))
  return value
  else:
- stats = self.parameter_list_stats()
+ stats = self.parameter_list_stats(x)
  # workaround to avoid find_ununsed_parameter=True in DDP
  stats = stats + 0. * self.value
  if self.local_loss_mode:

src/brevitas/core/stats/stats_wrapper.py

@@ -13,6 +13,7 @@
 from brevitas.core.utils import inplace_tensor_mul
 
 from .view_wrapper import _ViewCatParameterWrapper
+from .view_wrapper import _ViewParameter
 from .view_wrapper import _ViewParameterWrapper
 
 DEFAULT_MOMENTUM = 0.1
@@ -96,8 +97,12 @@ def __init__(
  super(_ParameterListStats, self).__init__()
 
  self.stats_input_concat_dim = stats_input_concat_dim
- self.first_tracked_param = _ViewParameterWrapper(
- tracked_parameter_list[0], stats_input_view_shape_impl)
+ if len(tracked_parameter_list) >= 1:
+ self.first_tracked_param = _ViewParameterWrapper(
+ tracked_parameter_list[0], stats_input_view_shape_impl)
+ else:
+ self.first_tracked_param = _ViewParameter(stats_input_view_shape_impl)
+
  if len(tracked_parameter_list) > 1:
  extra_list = [
  _ViewCatParameterWrapper(
@@ -109,10 +114,12 @@ def __init__(
  self.stats = _Stats(stats_impl, stats_output_shape)
 
  @brevitas.jit.script_method
- def forward(self) -> torch.Tensor:
- stats_input = self.first_tracked_param()
+ def forward(self, x: Optional[torch.Tensor] = None) -> torch.Tensor:
  if self.extra_tracked_params_list is not None:
+ stats_input = self.first_tracked_param(None)
  for extra_tracked_param in self.extra_tracked_params_list:
  stats_input = extra_tracked_param(stats_input)
+ else:
+ stats_input = self.first_tracked_param(x)
  out = self.stats(stats_input)
  return out

src/brevitas/core/stats/view_wrapper.py

@@ -1,6 +1,8 @@
 # Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
 # SPDX-License-Identifier: BSD-3-Clause
 
+from typing import Optional
+
 import torch
 from torch import Tensor
 from torch.nn import Module
@@ -19,8 +21,12 @@ def __init__(self, parameter: Parameter, view_shape_impl: Module) -> None:
  self.view_shape_impl = view_shape_impl
 
  @brevitas.jit.script_method
- def forward(self) -> Tensor:
- return self.view_shape_impl(self.parameter)
+ def forward(self, x: Optional[Tensor]) -> Tensor:
+ if x is not None:
+ parameter = x
+ else:
+ parameter = self.parameter
+ return self.view_shape_impl(parameter)
 
  def _load_from_state_dict(
  self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
@@ -39,6 +45,17 @@ def state_dict(self, destination=None, prefix='', keep_vars=False):
  return output_dict
 
 
+class _ViewParameter(brevitas.jit.ScriptModule):
+
+ def __init__(self, view_shape_impl: Module) -> None:
+ super(_ViewParameter, self).__init__()
+ self.view_shape_impl = view_shape_impl
+
+ @brevitas.jit.script_method
+ def forward(self, x: Tensor) -> Tensor:
+ return self.view_shape_impl(x)
+
+
 class _ViewCatParameterWrapper(brevitas.jit.ScriptModule):
  __constants__ = ['cat_dim']
 

src/brevitas/core/zero_point.py

@@ -82,7 +82,7 @@ def __init__(
 
  @brevitas.jit.script_method
  def forward(self, x: Tensor, scale: Tensor, bit_width: Tensor) -> torch.Tensor:
- stats = self.parameter_list_stats()
+ stats = self.parameter_list_stats(x)
  return self.scale_shift_zero_point(-stats, scale, bit_width)
 
 
@@ -266,7 +266,7 @@ def forward(self, x: Tensor, scale: Tensor, bit_width: Tensor) -> torch.Tensor:
  value = self.scale_shift_zero_point(value, scale, bit_width)
  return value
  else:
- stats = self.parameter_list_stats()
+ stats = self.parameter_list_stats(x)
  # workaround to avoid find_ununsed_parameter=True in DDP
  stats = stats + 0. * self.value
  if self.local_loss_mode:

src/brevitas_examples/llm/main.py

@@ -268,6 +268,8 @@ def main(args):
 
  model = offload_model(model)
 
+ model(**calibration_loader[0])
+
  if args.act_calibration:
  print("Apply act calibration...")
  apply_calibration(model, calibration_loader)