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

[BUG] FIx bug with f16 overfloat #482

Closed
wants to merge 2 commits into from
Closed

[BUG] FIx bug with f16 overfloat #482

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b8b7ae5
fix quntizatiion
vadiklyutiy Dec 21, 2024
8cba3a5
lint
vadiklyutiy Dec 21, 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
14 changes: 9 additions & 5 deletions python/hidet/graph/ops/quant/symmetric.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from typing import Union, List
from hidet import ir
from hidet.ir.type import DataType
from hidet.ir.dtypes import int32
from hidet.ir.dtypes import f32
from hidet.ir.expr import cast, if_then_else
from hidet.ir.compute.primitives import TensorNode, compute
from hidet.ir import primitives as prim
Expand All @@ -27,19 +27,23 @@ def __init__(self, w: TensorNode, quant_type: DataType, dims: Union[int, List[in
if not isinstance(dims, (list, tuple)):
dims = [dims]

# bf16 can't hold int16.max_value. Should convert to f32 first.
# For another types pair is similar.
# For cases when float type can hold quant_type.max_value we can skip this step but leave it for simplicity
wm = compute(
name='abs', shape=w.shape, fcompute=lambda *indices: if_then_else(w[indices] >= 0, w[indices], -w[indices])
name='abs',
shape=w.shape,
fcompute=lambda *indices: if_then_else(w[indices] >= 0, cast(w[indices], f32), cast(-w[indices], f32)),
)

scale = cops.reduce(wm, dims, keep_dim=False, reduce_type='max')
scale = compute(
name='scaling', shape=scale.shape, fcompute=lambda *indices: scale[indices] / quant_type.max_value
)

def scale_weight(*indices):
scale_indices = [indices[i] for i in range(len(indices)) if not i in dims]
# Have to cast to int32 first because there are several ways convert bf16 to int8
cast_to_int = cast(prim.round(w[indices] / scale[scale_indices]), int32)
return cast(cast_to_int, quant_type)
return cast(prim.round(w[indices] / scale[scale_indices]), quant_type)

wq = compute(name='quantize', shape=w.shape, fcompute=scale_weight)
super().__init__(
Expand Down
Loading