Absmean weight quantisation → {-1, 0, +1}. scale = mean(|W|) + ε (per weight-matrix, scalar) W_q = round_clamp(W / scale) via STE Returns (W_quantized_float, scale) where W_quantized_float still lives in fp32/bf16 so normal matmul works during training.
(w: torch.Tensor)
| 30 | # ────────────────────────────────────────────────────────────────────────────── |
| 31 | |
| 32 | def _quantize_weights_ternary(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 33 | """ |
| 34 | Absmean weight quantisation → {-1, 0, +1}. |
| 35 | |
| 36 | scale = mean(|W|) + ε (per weight-matrix, scalar) |
| 37 | W_q = round_clamp(W / scale) via STE |
| 38 | |
| 39 | Returns (W_quantized_float, scale) where W_quantized_float still lives in |
| 40 | fp32/bf16 so normal matmul works during training. |
| 41 | """ |
| 42 | scale = w.abs().mean().clamp(min=1e-8) |
| 43 | # STE: forward = round, backward = identity |
| 44 | w_scaled = w / scale |
| 45 | w_q = (w_scaled.round().clamp(-1, 1) - w_scaled).detach() + w_scaled |
| 46 | return w_q, scale |
| 47 | |
| 48 | |
| 49 | def _quantize_activations_int8(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |