| 182 | return x |
| 183 | |
| 184 | class SteN2F3Quantizer(nn.Module): |
| 185 | def __init__(self, q_group_size=128): |
| 186 | super().__init__() |
| 187 | self.q_group_size = q_group_size |
| 188 | |
| 189 | def forward(self, x): |
| 190 | org_w_shape = x.shape |
| 191 | |
| 192 | # reshape to groupsize |
| 193 | if self.q_group_size > 0: |
| 194 | assert org_w_shape[-1] % self.q_group_size == 0 |
| 195 | qx = x.reshape(-1, self.q_group_size) |
| 196 | elif self.q_group_size == -1: |
| 197 | qx = x.reshape(-1, x.shape[-1]) |
| 198 | assert qx.dim() == 2 |
| 199 | |
| 200 | # Get the Min Max |
| 201 | max_val = qx.amax(dim=1, keepdim=True) |
| 202 | min_val = qx.amin(dim=1, keepdim=True) |
| 203 | |
| 204 | |
| 205 | scale_pos = torch.abs(max_val) |
| 206 | scale_neg = torch.abs(min_val) |
| 207 | |
| 208 | dev = qx.device |
| 209 | x_pos = torch.zeros_like(qx) |
| 210 | x_neg = torch.zeros_like(qx) |
| 211 | x_pos = torch.where(qx >= 0, qx, x_pos) |
| 212 | x_neg = torch.where(qx < 0, qx, x_neg) |
| 213 | q_pos = x_pos / scale_pos |
| 214 | q_neg = x_neg / scale_neg |
| 215 | |
| 216 | q_pos, q_neg = self.round_pass(q_pos, q_neg, dev) |
| 217 | |
| 218 | qx = q_pos * scale_pos + q_neg * scale_neg |
| 219 | |
| 220 | qx = qx.reshape(org_w_shape) |
| 221 | |
| 222 | return qx |
| 223 | |
| 224 | def round_n2f3(self, q_pos, q_neg, dev): |
| 225 | q_pos = torch.where(q_pos >= 0.8114928305149078, torch.tensor(1.0).to(dev), q_pos) |
| 226 | q_pos = torch.where((q_pos < 0.8114928305149078) & (q_pos >= 0.5024898052215576), torch.tensor(0.6229856610298157).to(dev), q_pos) |
| 227 | q_pos = torch.where((q_pos < 0.5024898052215576) & (q_pos >= 0.2826657369732857), torch.tensor(0.3819939494132996).to(dev), q_pos) |
| 228 | q_pos = torch.where((q_pos < 0.2826657369732857) & (q_pos >= 0.0916687622666359), torch.tensor(0.1833375245332718).to(dev), q_pos) |
| 229 | q_pos = torch.where(q_pos < 0.0916687622666359, torch.tensor(0).to(dev), q_pos) |
| 230 | |
| 231 | q_neg = torch.where(q_neg >= -0.1234657019376755, torch.tensor(0).to(dev), q_neg) |
| 232 | q_neg = torch.where((q_neg < -0.1234657019376755) & (q_neg >= -0.39097706973552704), torch.tensor(-0.2469314038753510).to(dev), q_neg) |
| 233 | q_neg = torch.where((q_neg < -0.39097706973552704) & (q_neg >= -0.7675113677978516), torch.tensor(-0.5350227355957031).to(dev), q_neg) |
| 234 | q_neg = torch.where(q_neg < -0.7675113677978516, torch.tensor(-1.0).to(dev), q_neg) |
| 235 | |
| 236 | return q_pos, q_neg |
| 237 | |
| 238 | def round_pass(self, q_pos, q_neg, dev): |
| 239 | y_grad_pos, y_grad_neg = q_pos, q_neg |
| 240 | y_pos, y_neg = self.round_n2f3(q_pos, q_neg, dev) |
| 241 |
no outgoing calls