Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fix (quant/float): restore fix for log(0) #968

Merged
merged 2 commits into from
Jun 11, 2024
Merged

Fix (quant/float): restore fix for log(0) #968

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
db7a6f5
Fix (quant/float): restore fix for log(0)
Giuseppe5 Jun 10, 2024
02fd910
correct dtype
Giuseppe5 Jun 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/brevitas/core/quant/float.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@


class FloatQuant(brevitas.jit.ScriptModule):
__constants__ = ['signed']
__constants__ = ['signed', 'eps']

def __init__(
self,
Expand Down Expand Up @@ -71,7 +71,7 @@ def quantize(self, x: torch.Tensor):
scale = scaling_impl_value / float_scaling_impl_value
scaled_x = x / scale
internal_scale = float_internal_scale(
scaled_x, self.mantissa_bit_width(), self.fp_internal_scale_min())
scaled_x, self.mantissa_bit_width(), self.fp_internal_scale_min(), self.eps)
val_fp_quant = internal_scale * self.float_to_int_impl(scaled_x / internal_scale)
return val_fp_quant, scale

Expand Down
10 changes: 8 additions & 2 deletions src/brevitas/quant_tensor/float_quant_tensor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,10 @@ def training(self):
def saturating(self):
return self.saturating_t.item()

@property
def eps(self):
return torch.finfo(self.scale.dtype).tiny

def __torch_function__(self, func, types, args=(), kwargs=None):
if kwargs is None:
kwargs = {}
Expand All @@ -99,7 +103,8 @@ def _pre_round_float_value(self):
scale = self.scale.type(torch.float32)
minifloat_value = value / scale
fp_internal_scale = 1. - self.exponent_bias - self.mantissa_bit_width
int_scale = float_internal_scale(self.value, self.mantissa_bit_width, fp_internal_scale)
int_scale = float_internal_scale(
self.value, self.mantissa_bit_width, fp_internal_scale, self.eps)
minifloat_value = minifloat_value / int_scale
return minifloat_value

Expand Down Expand Up @@ -135,7 +140,8 @@ def minifloat(self, float_datatype=True):

if self.is_valid:
fp_internal_scale = 1. - self.exponent_bias - self.mantissa_bit_width
int_scale = float_internal_scale(self.value, self.mantissa_bit_width, fp_internal_scale)
int_scale = float_internal_scale(
self.value, self.mantissa_bit_width, fp_internal_scale, self.eps)
float_value = torch.round(self._pre_round_float_value) * int_scale
return float_value.type(self.scale.dtype)
else:
Expand Down
8 changes: 5 additions & 3 deletions src/brevitas/utils/torch_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,10 +93,12 @@ def compute_channel_view_shape(tensor: torch.Tensor, channel_dim: int):

@brevitas.jit.script
def float_internal_scale(
x: torch.Tensor, mantissa_bit_width: torch.Tensor,
fp_internal_scale_min: torch.Tensor) -> torch.Tensor:
x: torch.Tensor,
mantissa_bit_width: torch.Tensor,
fp_internal_scale_min: torch.Tensor,
eps: float) -> torch.Tensor:

internal_scale = floor_ste(torch.log2(torch.abs(x))) - mantissa_bit_width
internal_scale = floor_ste(torch.log2(torch.abs(x) + eps)) - mantissa_bit_width
internal_scale = torch.clamp_min(internal_scale, fp_internal_scale_min)
internal_scale = torch.exp2(internal_scale)
return internal_scale
3 changes: 2 additions & 1 deletion tests/brevitas/core/test_float_quant.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,8 +193,9 @@ def test_inner_scale(inp, minifloat_format, scale):
max_value = max_val if max_available_float is None else torch.min(
max_value, max_available_float)
# call internal scale
eps = torch.finfo(inp.dtype).tiny
internal_scale = float_internal_scale(
scaled_inp, float_quant.mantissa_bit_width(), float_quant.fp_internal_scale_min())
scaled_inp, float_quant.mantissa_bit_width(), float_quant.fp_internal_scale_min(), eps)
val_fp_quant = internal_scale * float_quant.float_to_int_impl(scaled_inp / internal_scale)
if signed:
val_fp_quant = torch.clip(val_fp_quant, -1. * max_val, max_val)
Expand Down
Loading