Hierarchical scales

.github/workflows/base.yml.template

@@ -20,7 +20,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

.github/workflows/base_reduced.yml.template

@@ -22,7 +22,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

.github/workflows/develop_install.yml

@@ -30,7 +30,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

.github/workflows/end_to_end.yml

@@ -32,7 +32,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

.github/workflows/examples_llm_pytest.yml

@@ -34,7 +34,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

.github/workflows/examples_pytest.yml

@@ -34,7 +34,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

.github/workflows/finn_integration.yml

@@ -30,7 +30,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

.github/workflows/notebook.yml

@@ -32,7 +32,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

.github/workflows/ort_integration.yml

@@ -30,7 +30,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

.github/workflows/pytest.yml

@@ -34,7 +34,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

.github/workflows/reduced_develop_install.yml

@@ -32,7 +32,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

.github/workflows/reduced_end_to_end.yml

@@ -34,7 +34,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

.github/workflows/reduced_examples_llm_pytest.yml

@@ -33,7 +33,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

.github/workflows/reduced_examples_pytest.yml

@@ -33,7 +33,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

.github/workflows/reduced_finn_integration.yml

@@ -32,7 +32,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

.github/workflows/reduced_notebook.yml

@@ -34,7 +34,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

.github/workflows/reduced_ort_integration.yml

@@ -32,7 +32,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

.github/workflows/reduced_pytest.yml

@@ -33,7 +33,7 @@ jobs:
     steps:
 
     - name: Checkout repo
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: Set up Python
       uses: actions/setup-python@v2

src/brevitas/config.py

@@ -25,3 +25,4 @@ def env_to_bool(name, default):
 _FULL_STATE_DICT = False
 _IS_INSIDE_QUANT_LAYER = None
 _ONGOING_EXPORT = None
+_RETROCOMPATIBLE_SCALING = False

src/brevitas/core/quant/float.py

@@ -67,12 +67,13 @@ def __init__(
 
     @brevitas.jit.script_method
     def quantize(self, x: torch.Tensor):
-        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())
-            scale = scale / float_scaling_impl_value
+        else:
+            float_scaling_impl_value = None
+        scale = self.scaling_impl(x, float_scaling_impl_value)
         x = self.input_view_impl(x)
         scaled_x = x / scale
         internal_scale = float_internal_scale(

src/brevitas/core/quant/int.py

@@ -8,6 +8,7 @@
 from torch.nn import Module
 
 import brevitas
+from brevitas.core.function_wrapper.misc import Identity
 from brevitas.core.quant.delay import DelayWrapper
 from brevitas.core.utils import StatelessBuffer
 from brevitas.function.ops_ste import round_ste
@@ -138,20 +139,24 @@ def __init__(
             scaling_impl: Module,
             int_scaling_impl: Module,
             zero_point_impl: Module,
-            bit_width_impl: Module):
+            bit_width_impl: Module,
+            scaling_int_quant: Optional[Module] = None):
         super(RescalingIntQuant, self).__init__()
         self.int_quant = int_quant
         self.scaling_impl = scaling_impl
         self.int_scaling_impl = int_scaling_impl
         self.zero_point_impl = zero_point_impl
         self.msb_clamp_bit_width_impl = bit_width_impl
+        if scaling_int_quant is None:
+            self.scaling_int_quant = Identity()
+        else:
+            self.scaling_int_quant = scaling_int_quant
 
     @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)
         y = self.int_quant(scale, zero_point, bit_width, x)
         return y, scale, zero_point, bit_width
@@ -184,8 +189,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)
         y = self.decoupled_int_quant(pre_scale, pre_zero_point, scale, zero_point, bit_width, x)
         return y, scale, zero_point, bit_width, pre_scale, pre_zero_point
@@ -250,8 +254,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)
         y = self.decoupled_int_quant(pre_scale, pre_zero_point, scale, zero_point, bit_width, x)
         return y, scale, zero_point, bit_width, pre_scale, pre_zero_point

src/brevitas/core/restrict_val.py

@@ -90,6 +90,9 @@ def restrict_init_module(self):
     def restrict_init_inplace_module(self):
         return Identity()
 
+    def retrocompatibility_op(self, x):
+        return x
+
     @brevitas.jit.script_method
     def forward(self, x: torch.Tensor) -> Tensor:
         return x
@@ -113,6 +116,9 @@ def restrict_init_module(self):
     def restrict_init_inplace_module(self):
         return InplaceLogTwo()
 
+    def retrocompatibility_op(self, x):
+        return self.power_of_two(x)
+
     @brevitas.jit.script_method
     def forward(self, x: torch.Tensor):
         x = self.power_of_two(x)
@@ -137,6 +143,9 @@ def restrict_init_module(self):
     def restrict_init_inplace_module(self):
         return Identity()
 
+    def retrocompatibility_op(self, x):
+        return x
+
     @brevitas.jit.script_method
     def forward(self, x: torch.Tensor):
         x = self.float_to_int_impl(x)
@@ -162,8 +171,38 @@ def restrict_init_module(self):
     def restrict_init_inplace_module(self):
         return InplaceLogTwo()
 
+    def retrocompatibility_op(self, x):
+        return self.power_of_two(x)
+
     @brevitas.jit.script_method
     def forward(self, x: torch.Tensor):
         x = self.float_to_int_impl(x)
         x = self.power_of_two(x)
         return x
+
+
+class QuantRestrictValue(brevitas.jit.ScriptModule):
+
+    def __init__(self, restrict_value_float_to_int_impl: Module):
+        super(QuantRestrictValue, self).__init__()
+        self.float_to_int_impl = restrict_value_float_to_int_impl
+
+    def restrict_init_float(self, x: float):
+        return Identity()
+
+    def restrict_init_tensor(self, x: torch.Tensor):
+        return Identity()
+
+    def restrict_init_module(self):
+        return Identity()
+
+    def restrict_init_inplace_module(self):
+        return Identity()
+
+    def retrocompatibility_op(self, x):
+        return Identity()
+
+    @brevitas.jit.script_method
+    def forward(self, x: torch.Tensor):
+        o, *_ = self.float_to_int_impl(x)
+        return o

src/brevitas/core/scaling/runtime.py

@@ -50,10 +50,11 @@ def __init__(
             dtype,
             device)
 
-    @brevitas.jit.script_method
-    def forward(self, ignored: torch.Tensor) -> torch.Tensor:
-        stats = self.parameter_list_stats()
-        return self.stats_scaling_impl(stats)
+    def forward(self, x: torch.Tensor, threshold: Optional[torch.Tensor] = None) -> torch.Tensor:
+        stats = self.parameter_list_stats(x)
+        if threshold is None:
+            threshold = torch.ones(1).type_as(stats)
+        return self.stats_scaling_impl(stats, threshold)
 
 
 class _StatsScaling(brevitas.jit.ScriptModule):
@@ -80,8 +81,11 @@ def __init__(
         self.restrict_scaling_pre = restrict_scaling_impl.restrict_init_module()
 
     @brevitas.jit.script_method
-    def forward(self, stats: torch.Tensor) -> torch.Tensor:
-        stats = self.restrict_scaling_pre(stats)
+    def forward(
+            self, stats: torch.Tensor, threshold: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if threshold is None:
+            threshold = torch.ones(1).type_as(stats)
+        stats = self.restrict_scaling_pre(stats / threshold)
         stats = self.affine_rescaling(stats)
         stats = self.restrict_clamp_scaling(stats)
         return stats
@@ -120,9 +124,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):
@@ -179,9 +183,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