Fix bias rounding (#137)

Author: MaximGorkem

izer/checkpoint.py

@@ -183,8 +183,8 @@ def load(
             if bias_name in checkpoint_state and seq not in no_bias:
                 wb = checkpoint_state[wb_name].numpy().astype(np.int64) \
                      if wb_name in checkpoint_state else 8
-                w = checkpoint_state[bias_name].numpy(). \
-                    astype(np.int64) // 2**(wb - 1)
+                w = (checkpoint_state[bias_name] // 2**(wb - 1)).numpy(). \
+                    astype(np.int64)
 
                 if np.all(w == 0):
                     wprint(f'All bias values for `{bias_name}` are zero.')

izer/quantize.py

@@ -221,6 +221,7 @@ def get_max_bit_shift(t, shift_quantile, return_bit_shift=False):
                 bias = bias.add(.5).floor().clamp(min=-(2**(clamp_bits+tc.dev.ACTIVATION_BITS-2)),
                                                   max=2**(clamp_bits+tc.dev.ACTIVATION_BITS-2)-1)
 
+                bias = (bias // 128) * 128
                 # Store modified bias back into model
                 new_checkpoint_state[bias_name] = bias