| 17 | |
| 18 | |
| 19 | class DifferentiableQuantizer: |
| 20 | def __init__(self, name: str, config: AttributeQuantizerConfig): |
| 21 | self.name = name |
| 22 | self.config = config |
| 23 | |
| 24 | def quantize(self, tensor: Tensor, step: int) -> QuantizeResult: |
| 25 | if not self.config.enabled or self.config.bitwidth is None: |
| 26 | return QuantizeResult(value=tensor, q_step=None, metadata={}) |
| 27 | |
| 28 | bitwidth = self._select_bitwidth(step) |
| 29 | if bitwidth is None: |
| 30 | return QuantizeResult(value=tensor, q_step=None, metadata={}) |
| 31 | |
| 32 | clamp_range = self.config.clamp_range |
| 33 | if clamp_range is None: |
| 34 | raise ValueError(f"Quantizer '{self.name}' requires clamp_range when enabled") |
| 35 | |
| 36 | lower, upper = clamp_range |
| 37 | if upper <= lower: |
| 38 | raise ValueError(f"Quantizer '{self.name}' has invalid clamp range {clamp_range}") |
| 39 | |
| 40 | q_step = self._create_q_step(tensor, lower, upper, bitwidth) |
| 41 | |
| 42 | mode = self.config.mode |
| 43 | if mode == "noise": |
| 44 | value = self._apply_noise_quantization(tensor, lower, upper, q_step) |
| 45 | elif mode == "round": |
| 46 | value = self._apply_round_quantization(tensor, lower, upper, bitwidth) |
| 47 | else: |
| 48 | raise ValueError(f"Unsupported quantizer mode '{mode}' for attribute '{self.name}'") |
| 49 | |
| 50 | return QuantizeResult( |
| 51 | value=value, |
| 52 | q_step=q_step, |
| 53 | metadata={"bitwidth": torch.tensor(bitwidth, device=tensor.device, dtype=tensor.dtype)}, |
| 54 | ) |
| 55 | |
| 56 | def _select_bitwidth(self, step: int) -> Optional[int]: |
| 57 | bitwidth = self.config.bitwidth |
| 58 | if bitwidth is None: |
| 59 | return None |
| 60 | warmup_steps = self.config.warmup_steps |
| 61 | if warmup_steps is not None and step < warmup_steps: |
| 62 | return self.config.warmup_bitwidth or bitwidth |
| 63 | return bitwidth |
| 64 | |
| 65 | def _create_q_step(self, tensor: Tensor, lower: float, upper: float, bitwidth: int) -> Tensor: |
| 66 | denom = (2 ** bitwidth) - 1 |
| 67 | step_value = (upper - lower) / denom |
| 68 | return torch.as_tensor(step_value, dtype=tensor.dtype, device=tensor.device) |
| 69 | |
| 70 | def _apply_noise_quantization(self, tensor: Tensor, lower: float, upper: float, q_step: Tensor) -> Tensor: |
| 71 | clamped = tensor.clamp(lower, upper) |
| 72 | noise = torch.empty_like(tensor).uniform_(-0.5, 0.5) |
| 73 | quantized = clamped + noise * q_step |
| 74 | return quantized.clamp(lower, upper) |
| 75 | |
| 76 | def _apply_round_quantization(self, tensor: Tensor, lower: float, upper: float, bitwidth: int) -> Tensor: |
no outgoing calls