| 151 | return x |
| 152 | |
| 153 | class SteInt2AsymQuantizer(nn.Module): |
| 154 | def __init__(self, q_group_size=64): |
| 155 | super().__init__() |
| 156 | self.q_group_size = q_group_size |
| 157 | self.bit = 2 |
| 158 | def forward(self, x): |
| 159 | org_w_shape = x.shape |
| 160 | |
| 161 | if self.q_group_size > 0: |
| 162 | assert org_w_shape[-1] % self.q_group_size == 0 |
| 163 | x = x.reshape(-1, self.q_group_size) |
| 164 | assert x.dim() == 2 |
| 165 | |
| 166 | max_val = x.amax(dim=1, keepdim=True) |
| 167 | min_val = x.amin(dim=1, keepdim=True) |
| 168 | max_int = 2 ** self.bit - 1 |
| 169 | min_int = 0 |
| 170 | scales = (max_val - min_val).clamp(min=1e-5) / max_int |
| 171 | zeros = (-torch.round(min_val / scales)).clamp_(min_int, max_int) |
| 172 | |
| 173 | assert torch.isnan(scales).sum() == 0 |
| 174 | assert torch.isnan(x).sum() == 0 |
| 175 | |
| 176 | x = (torch.clamp(Round.apply(x / scales) + |
| 177 | zeros, min_int, max_int) - zeros) * scales |
| 178 | assert torch.isnan(x).sum() == 0 |
| 179 | |
| 180 | x = x.reshape(org_w_shape) |
| 181 | |
| 182 | return x |
| 183 | |
| 184 | class SteN2F3Quantizer(nn.Module): |
| 185 | def __init__(self, q_group_size=128): |
no outgoing calls
no test coverage detected