Paolify/RVC_4
0
1# Copyright (c) Facebook, Inc. and its affiliates.2# All rights reserved.3#4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7"""8Differentiable quantizer based on scaled noise injection.9"""10from dataclasses import dataclass11import math12import typing as tp13 14import torch15 16from .base import BaseQuantizer17from .uniform import uniform_quantize, uniform_unquantize18from .utils import simple_repr19 20 21class DiffQuantizer(BaseQuantizer):22 @dataclass23 class _QuantizedParam(BaseQuantizer._QuantizedParam):24 logit: torch.nn.Parameter25 26 def __init__(self, model: torch.nn.Module, min_size: float = 0.01, float16: bool = False,27 group_size: int = 1, min_bits: float = 2, max_bits: float = 15,28 param="bits", noise="gaussian",29 init_bits: float = 8, extra_bits: float = 0, suffix: str = "_diffq",30 exclude: tp.List[str] = [], detect_bound: bool = True):31 """32 Differentiable quantizer based on scaled noise injection.33 For every parameter `p` in the model, this introduces a number of bits parameter34 `b` with the same dimensions (when group_size = 1).35 Before each forward, `p` is replaced by `p + U`36 with U uniform iid noise with range [-d/2, d/2], with `d` the uniform quantization37 step for `b` bits.38 This noise approximates the quantization noise in a differentiable manner, both39 with respect to the unquantized parameter `p` and the number of bits `b`.40 41 At eveluation (as detected with `model.eval()`), the model is replaced42 by its true quantized version, and restored when going back to training.43 44 When doing actual quantization (for serialization, or evaluation),45 the number of bits is rounded to the nearest integer, and needs to be stored along.46 This will cost a few bits per dimension. To reduce this cost, one can use `group_size`,47 which will use a single noise level for multiple weight entries.48 49 You can use the `DiffQuantizer.model_size` method to get a differentiable estimate of the50 model size in MB. You can then use this estimate as a penalty in your training loss.51 52 Args:53 model (torch.nn.Module): model to quantize54 min_size (float): minimum size in MB of a parameter to be quantized.55 float16 (bool): if a layer is smaller than min_size, should we still do float16?56 group_size (int): weight entries are groupped together to reduce the number57 of noise scales to store. This should divide the size of all parameters58 bigger than min_size.59 min_bits (float): minimal number of bits.60 max_bits (float): maximal number of bits.61 init_bits (float): initial number of bits.62 extra_bits (float): extra bits to add for actual quantization (before roundoff).63 suffix (str): suffix used for the name of the extra noise scale parameters.64 exclude (list[str]): list of patterns used to match parameters to exclude.65 For instance `['bias']` to exclude all bias terms.66 detect_bound (bool): if True, will detect bound parameters and reuse67 the same quantized tensor for both, as well as the same number of bits.68 69 ..Warning::70 You must call `model.training()` and `model.eval()` for `DiffQuantizer` work properly.71 72 """73 self.group_size = group_size74 self.min_bits = min_bits75 self.max_bits = max_bits76 self.init_bits = init_bits77 self.extra_bits = extra_bits78 self.suffix = suffix79 self.param = param80 self.noise = noise81 assert noise in ["gaussian", "uniform"]82 self._optimizer_setup = False83 84 self._min_noise = 1 / (2 ** self.max_bits - 1)85 self._max_noise = 1 / (2 ** self.min_bits - 1)86 87 assert group_size >= 088 assert min_bits < init_bits < max_bits, \89 "init_bits must be between min_bits and max_bits excluded3"90 91 for name, _ in model.named_parameters():92 if name.endswith(suffix):93 raise RuntimeError("The model already has some noise scales parameters, "94 "maybe you used twice a DiffQuantizer on the same model?.")95 96 super().__init__(model, min_size, float16, exclude, detect_bound)97 98 def _get_bits(self, logit: torch.Tensor):99 if self.param == "noise":100 return torch.log2(1 + 1 / self._get_noise_scale(logit))101 else:102 t = torch.sigmoid(logit)103 return self.max_bits * t + (1 - t) * self.min_bits104 105 def _get_noise_scale(self, logit: torch.Tensor):106 if self.param == "noise":107 t = torch.sigmoid(logit)108 return torch.exp(t * math.log(self._min_noise) + (1 - t) * math.log(self._max_noise))109 else:110 return 1 / (2 ** self._get_bits(logit) - 1)111 112 def _register_param(self, name, param, module, other):113 if other is not None:114 return self.__class__._QuantizedParam(115 name=name, param=param, module=module, logit=other.logit, other=other)116 assert self.group_size == 0 or param.numel() % self.group_size == 0117 # we want the initial number of bits to be init_bits.118 if self.param == "noise":119 noise_scale = 1 / (2 ** self.init_bits - 1)120 t = (math.log(noise_scale) - math.log(self._max_noise)) / (121 math.log(self._min_noise) - math.log(self._max_noise))122 else:123 t = (self.init_bits - self.min_bits) / (self.max_bits - self.min_bits)124 assert 0 < t < 1125 logit = torch.logit(torch.tensor(float(t)))126 assert abs(self._get_bits(logit) - self.init_bits) < 1e-5127 if self.group_size > 0:128 nparam = param.numel() // self.group_size129 else:130 nparam = 1131 logit = torch.nn.Parameter(132 torch.full(133 (nparam,),134 logit,135 device=param.device))136 module.register_parameter(name + self.suffix, logit)137 return self.__class__._QuantizedParam(138 name=name, param=param, module=module, logit=logit, other=None)139 140 def clear_optimizer(self, optimizer: torch.optim.Optimizer):141 params = [qp.logit for qp in self._qparams]142 143 for group in optimizer.param_groups:144 new_params = []145 for q in list(group["params"]):146 matched = False147 for p in params:148 if p is q:149 matched = True150 if not matched:151 new_params.append(q)152 group["params"][:] = new_params153 154 def setup_optimizer(self, optimizer: torch.optim.Optimizer,155 lr: float = 1e-3, **kwargs):156 """157 Setup the optimizer to tune the number of bits. In particular, this will deactivate158 weight decay for the bits parameters.159 160 Args:161 optimizer (torch.Optimizer): optimizer to use.162 lr (float): specific learning rate for the bits parameters. 1e-3163 is perfect for Adam.,w164 kwargs (dict): overrides for other optimization parameters for the bits.165 """166 assert not self._optimizer_setup167 self._optimizer_setup = True168 169 params = [qp.logit for qp in self._qparams]170 171 for group in optimizer.param_groups:172 for q in list(group["params"]):173 for p in params:174 if p is q:175 raise RuntimeError("You should create the optimizer "176 "before the quantizer!")177 178 group = {"params": params, "lr": lr, "weight_decay": 0}179 group.update(kwargs)180 optimizer.add_param_group(group)181 182 def no_optimizer(self):183 """184 Call this if you do not want to use an optimizer.185 """186 self._optimizer_setup = True187 188 def check_unused(self):189 for qparam in self._qparams:190 if qparam.other is not None:191 continue192 grad = qparam.param.grad193 if grad is None or (grad == 0).all():194 if qparam.logit.grad is not None:195 qparam.logit.grad.data.zero_()196 197 def model_size(self, exact=False):198 """199 Differentiable estimate of the model size.200 The size is returned in MB.201 202 If `exact` is True, then the output is no longer differentiable but203 reflect exactly an achievable size, even without compression,204 i.e.same as returned by `naive_model_size()`.205 """206 total = super().model_size()207 subtotal = 0208 for qparam in self._qparams:209 # only count the first appearance of a Parameter210 if qparam.other is not None:211 continue212 bits = self.extra_bits + self._get_bits(qparam.logit)213 if exact:214 bits = bits.round().clamp(1, 15)215 if self.group_size == 0:216 group_size = qparam.param.numel()217 else:218 group_size = self.group_size219 subtotal += group_size * bits.sum()220 subtotal += 2 * 32 # param scale221 222 # Number of bits to represent each number of bits223 bits_bits = math.ceil(math.log2(1 + (bits.max().round().item() - self.min_bits)))224 subtotal += 8 # 8 bits for bits_bits225 subtotal += bits_bits * bits.numel()226 227 subtotal /= 2 ** 20 * 8 # bits -> MegaBytes228 return total + subtotal229 230 def true_model_size(self):231 """232 Naive model size without zlib compression.233 """234 return self.model_size(exact=True).item()235 236 def _pre_forward_train(self):237 if not self._optimizer_setup:238 raise RuntimeError("You must call `setup_optimizer()` on your optimizer "239 "before starting training.")240 for qparam in self._qparams:241 if qparam.other is not None:242 noisy = qparam.other.module._parameters[qparam.other.name]243 else:244 bits = self._get_bits(qparam.logit)[:, None]245 if self.group_size == 0:246 p_flat = qparam.param.view(-1)247 else:248 p_flat = qparam.param.view(-1, self.group_size)249 scale = p_flat.max() - p_flat.min()250 unit = 1 / (2**bits - 1)251 if self.noise == "uniform":252 noise_source = (torch.rand_like(p_flat) - 0.5)253 elif self.noise == "gaussian":254 noise_source = torch.randn_like(p_flat) / 2255 noise = scale * unit * noise_source256 noisy = p_flat + noise257 # We bypass the checks by PyTorch on parameters being leafs258 qparam.module._parameters[qparam.name] = noisy.view_as(qparam.param)259 return True260 261 def _post_forward_train(self):262 for qparam in self._qparams:263 qparam.module._parameters[qparam.name] = qparam.param264 return True265 266 def _quantize_param(self, qparam: _QuantizedParam) -> tp.Any:267 bits = self.extra_bits + self._get_bits(qparam.logit)268 bits = bits.round().clamp(1, 15)[:, None].byte()269 if self.group_size == 0:270 p = qparam.param.data.view(-1)271 else:272 p = qparam.param.data.view(-1, self.group_size)273 levels, scales = uniform_quantize(p, bits)274 return levels, scales, bits275 276 def _unquantize_param(self, qparam: _QuantizedParam, quantized: tp.Any) -> torch.Tensor:277 levels, param_scale, bits = quantized278 return uniform_unquantize(levels, param_scale, bits).view_as(qparam.param.data)279 280 def detach(self):281 super().detach()282 for qparam in self._qparams:283 delattr(qparam.module, qparam.name + self.suffix)284 285 def __repr__(self):286 return simple_repr(self)287