diff --git a/fasterai/_modidx.py b/fasterai/_modidx.py index 9a396b6..26acd33 100644 --- a/fasterai/_modidx.py +++ b/fasterai/_modidx.py @@ -220,6 +220,34 @@ 'fasterai/export/onnx_exporter.py'), 'fasterai.export.onnx_exporter.verify_onnx': ( 'export/onnx_exporter.html#verify_onnx', 'fasterai/export/onnx_exporter.py')}, + 'fasterai.huggingface.huggingface': { 'fasterai.huggingface.huggingface.HFSparsifyCallback': ( 'huggingface/huggingface.html#hfsparsifycallback', + 'fasterai/huggingface/huggingface.py'), + 'fasterai.huggingface.huggingface.HFSparsifyCallback.__init__': ( 'huggingface/huggingface.html#hfsparsifycallback.__init__', + 'fasterai/huggingface/huggingface.py'), + 'fasterai.huggingface.huggingface.HFSparsifyCallback._sparsity_value': ( 'huggingface/huggingface.html#hfsparsifycallback._sparsity_value', + 'fasterai/huggingface/huggingface.py'), + 'fasterai.huggingface.huggingface.HFSparsifyCallback.on_epoch_end': ( 'huggingface/huggingface.html#hfsparsifycallback.on_epoch_end', + 'fasterai/huggingface/huggingface.py'), + 'fasterai.huggingface.huggingface.HFSparsifyCallback.on_log': ( 'huggingface/huggingface.html#hfsparsifycallback.on_log', + 'fasterai/huggingface/huggingface.py'), + 'fasterai.huggingface.huggingface.HFSparsifyCallback.on_optimizer_step': ( 'huggingface/huggingface.html#hfsparsifycallback.on_optimizer_step', + 'fasterai/huggingface/huggingface.py'), + 'fasterai.huggingface.huggingface.HFSparsifyCallback.on_step_begin': ( 'huggingface/huggingface.html#hfsparsifycallback.on_step_begin', + 'fasterai/huggingface/huggingface.py'), + 'fasterai.huggingface.huggingface.HFSparsifyCallback.on_train_begin': ( 'huggingface/huggingface.html#hfsparsifycallback.on_train_begin', + 'fasterai/huggingface/huggingface.py'), + 'fasterai.huggingface.huggingface.HFSparsifyCallback.on_train_end': ( 'huggingface/huggingface.html#hfsparsifycallback.on_train_end', + 'fasterai/huggingface/huggingface.py'), + 'fasterai.huggingface.huggingface._has_transformers': ( 'huggingface/huggingface.html#_has_transformers', + 'fasterai/huggingface/huggingface.py'), + 'fasterai.huggingface.huggingface._load_model': ( 'huggingface/huggingface.html#_load_model', + 'fasterai/huggingface/huggingface.py'), + 'fasterai.huggingface.huggingface._require_transformers': ( 'huggingface/huggingface.html#_require_transformers', + 'fasterai/huggingface/huggingface.py'), + 'fasterai.huggingface.huggingface._save_compressed': ( 'huggingface/huggingface.html#_save_compressed', + 'fasterai/huggingface/huggingface.py'), + 'fasterai.huggingface.huggingface.sparsify_model': ( 'huggingface/huggingface.html#sparsify_model', + 'fasterai/huggingface/huggingface.py')}, 'fasterai.misc.all': {}, 'fasterai.misc.bn_folding': { 'fasterai.misc.bn_folding.BN_Folder': ( 'misc/bn_folding.html#bn_folder', 'fasterai/misc/bn_folding.py'), @@ -305,6 +333,42 @@ 'fasterai/prune/pruner.py'), 'fasterai.prune.pruner.Pruner.restore_attention_layers': ( 'prune/pruner.html#pruner.restore_attention_layers', 'fasterai/prune/pruner.py')}, + 'fasterai.quantize.adaround': { 'fasterai.quantize.adaround._calibrate_activations': ( 'quantize/adaround.html#_calibrate_activations', + 'fasterai/quantize/adaround.py'), + 'fasterai.quantize.adaround._capture_inputs': ( 'quantize/adaround.html#_capture_inputs', + 'fasterai/quantize/adaround.py'), + 'fasterai.quantize.adaround._fake_quant_act': ( 'quantize/adaround.html#_fake_quant_act', + 'fasterai/quantize/adaround.py'), + 'fasterai.quantize.adaround._hard_quant_from_V': ( 'quantize/adaround.html#_hard_quant_from_v', + 'fasterai/quantize/adaround.py'), + 'fasterai.quantize.adaround._init_V': ( 'quantize/adaround.html#_init_v', + 'fasterai/quantize/adaround.py'), + 'fasterai.quantize.adaround._iter_target_layers': ( 'quantize/adaround.html#_iter_target_layers', + 'fasterai/quantize/adaround.py'), + 'fasterai.quantize.adaround._layer_fwd': ( 'quantize/adaround.html#_layer_fwd', + 'fasterai/quantize/adaround.py'), + 'fasterai.quantize.adaround._normalize_calibration': ( 'quantize/adaround.html#_normalize_calibration', + 'fasterai/quantize/adaround.py'), + 'fasterai.quantize.adaround._pack_int4': ( 'quantize/adaround.html#_pack_int4', + 'fasterai/quantize/adaround.py'), + 'fasterai.quantize.adaround._rect_sigmoid': ( 'quantize/adaround.html#_rect_sigmoid', + 'fasterai/quantize/adaround.py'), + 'fasterai.quantize.adaround._set_child': ( 'quantize/adaround.html#_set_child', + 'fasterai/quantize/adaround.py'), + 'fasterai.quantize.adaround._unpack_int4': ( 'quantize/adaround.html#_unpack_int4', + 'fasterai/quantize/adaround.py'), + 'fasterai.quantize.adaround.adaround_layer': ( 'quantize/adaround.html#adaround_layer', + 'fasterai/quantize/adaround.py'), + 'fasterai.quantize.adaround.adaround_quantize': ( 'quantize/adaround.html#adaround_quantize', + 'fasterai/quantize/adaround.py'), + 'fasterai.quantize.adaround.adaround_to_onnx': ( 'quantize/adaround.html#adaround_to_onnx', + 'fasterai/quantize/adaround.py'), + 'fasterai.quantize.adaround.fold_bn': ( 'quantize/adaround.html#fold_bn', + 'fasterai/quantize/adaround.py'), + 'fasterai.quantize.adaround.rtn_quant': ( 'quantize/adaround.html#rtn_quant', + 'fasterai/quantize/adaround.py'), + 'fasterai.quantize.adaround.weight_scale': ( 'quantize/adaround.html#weight_scale', + 'fasterai/quantize/adaround.py')}, 'fasterai.quantize.all': {}, 'fasterai.quantize.quantize_callback': { 'fasterai.quantize.quantize_callback.QuantizeCallback': ( 'quantize/quantize_callback.html#quantizecallback', 'fasterai/quantize/quantize_callback.py'), diff --git a/fasterai/quantize/adaround.py b/fasterai/quantize/adaround.py new file mode 100644 index 0000000..3eb847b --- /dev/null +++ b/fasterai/quantize/adaround.py @@ -0,0 +1,400 @@ +# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/quantize/adaround.ipynb. + +# %% ../../nbs/quantize/adaround.ipynb #1746d882 +from __future__ import annotations +import copy +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.utils.fusion import fuse_conv_bn_eval + +# %% auto #0 +__all__ = ['fold_bn', 'weight_scale', 'rtn_quant', 'adaround_layer', 'adaround_quantize', 'adaround_to_onnx'] + +# %% ../../nbs/quantize/adaround.ipynb #80a30f53 +def _set_child(parent, name, new): + "Replace child module `name` of `parent` (handles Sequential integer indices)." + if name.isdigit(): parent[int(name)] = new + else: setattr(parent, name, new) + +def fold_bn(model: nn.Module # model containing Conv2d→BatchNorm2d pairs + ) -> nn.Module: # same model with each Conv-BN pair fused in-place + "Fuse each `Conv2d`→`BatchNorm2d` pair into the conv (eval mode) and replace the BN with Identity." + for mod in model.modules(): + kids = list(mod.named_children()) + for i in range(len(kids) - 1): + (n1, c1), (n2, c2) = kids[i], kids[i + 1] + if isinstance(c1, nn.Conv2d) and isinstance(c2, nn.BatchNorm2d): + _set_child(mod, n1, fuse_conv_bn_eval(c1, c2)) + _set_child(mod, n2, nn.Identity()) + return model + +# %% ../../nbs/quantize/adaround.ipynb #c9f0e214 +def weight_scale(W: torch.Tensor, # weight tensor [out_channels, ...] + w_bits: int = 4, # bit-width (signed) + granularity: str = 'channel', # 'channel' (per-output-channel) or 'tensor' + ): # (scale, qmin, qmax) — dequant scale and signed clip range + "Signed weight quantization `scale` and clip range `[qmin, qmax]` on the ONNX/ORT grid." + qmax = 2 ** (w_bits - 1) - 1 # INT4 -> 7, INT8 -> 127 + qmin = -(2 ** (w_bits - 1)) # INT4 -> -8, INT8 -> -128 (asymmetric signed grid) + denom = 2 ** (w_bits - 1) - 0.5 # INT4 -> 7.5 (matches onnxruntime per-channel int4) + if granularity == 'channel': + r = W.detach().abs().reshape(W.shape[0], -1).amax(1).clamp_min(1e-12) + scale = (r / denom).reshape([-1] + [1] * (W.dim() - 1)) + else: + scale = W.detach().abs().amax().clamp_min(1e-12) / denom + return scale, qmin, qmax + +def rtn_quant(W: torch.Tensor, # weight tensor to fake-quantize + w_bits: int = 4, # bit-width + granularity: str = 'channel', # 'channel' or 'tensor' + ) -> torch.Tensor: # dequantized (fake-quant) weights, same shape + "Round-to-nearest (RTN) signed fake-quant — the AdaRound baseline (== AdaRound at init)." + scale, qmin, qmax = weight_scale(W, w_bits, granularity) + floorWs = torch.floor(W / scale) + q = floorWs + ((W / scale - floorWs) >= 0.5).float() # round half up (matches the AdaRound hard round) + return torch.clamp(q, qmin, qmax) * scale + +# %% ../../nbs/quantize/adaround.ipynb #d948dfff +# AdaRound rectified-sigmoid soft rounding (Nagel et al., 2020): +# h(V) = clamp(sigmoid(V)*(zeta-gamma)+gamma, 0, 1) +# W_q(V) = s * clamp(floor(W/s) + h(V), qmin, qmax) +_ZETA, _GAMMA = 1.1, -0.1 + +def _rect_sigmoid(V: torch.Tensor) -> torch.Tensor: + "Rectified sigmoid squashing `V` into a soft rounding decision in [0, 1]." + return torch.clamp(torch.sigmoid(V) * (_ZETA - _GAMMA) + _GAMMA, 0, 1) + +def _init_V(frac: torch.Tensor) -> torch.Tensor: + "Initialise `V` so that `h(V) == frac(W/s)` — makes AdaRound identical to RTN at init." + s0 = ((frac - _GAMMA) / (_ZETA - _GAMMA)).clamp(1e-6, 1 - 1e-6) + return torch.log(s0 / (1 - s0)) + +def _hard_quant_from_V(V, floorWs, scale, qmin, qmax): + "Deployed weights: hard-round the learned soft rounding `h(V)` onto the signed grid." + h_hard = (_rect_sigmoid(V) >= 0.5).float() + return scale * torch.clamp(floorWs + h_hard, qmin, qmax) + +def _layer_fwd(mod, x, w, b): + "Forward one Conv2d/Linear layer with an explicit weight/bias (for reconstruction)." + if isinstance(mod, nn.Conv2d): + return F.conv2d(x, w, b, mod.stride, mod.padding, mod.dilation, mod.groups) + return F.linear(x, w, b) + +# %% ../../nbs/quantize/adaround.ipynb #ab024063 +def adaround_layer(layer: nn.Module, # Conv2d/Linear whose rounding is optimized + x_cal: torch.Tensor, # cached FP inputs to this layer [N, ...] + w_bits: int = 4, # weight bit-width + *, + granularity: str = 'channel', # 'channel' or 'tensor' + iters: int = 2000, # optimization iterations + lr: float = 1e-2, # Adam learning rate + lam: float = 0.01, # rounding-regularizer weight + beta_start: float = 20.0, # anneal start (encourages soft) + beta_end: float = 2.0, # anneal end (encourages binary) + warmup: float = 0.2, # fraction of iters before regularizer turns on + batch_size: int = 32, # minibatch of cached inputs per step + device: str | torch.device | None = None, # compute device (default: layer's) + ): # (Wq_hard, info) — deployed fake-quant weights + diagnostics + "Learn per-weight soft rounding for one layer, minimising output reconstruction ‖W_q(V)x − Wx‖² + regularizer." + device = torch.device(device) if device is not None else layer.weight.device + W = layer.weight.detach().float().to(device) + scale, qmin, qmax = weight_scale(W, w_bits, granularity) + floorWs = torch.floor(W / scale) + frac = (W / scale) - floorWs # in [0, 1) + V = _init_V(frac).clone().requires_grad_(True) # init so h(V) == frac (== RTN) + opt = torch.optim.Adam([V], lr=lr) + bias = layer.bias.detach().float().to(device) if layer.bias is not None else None + N = x_cal.shape[0] + warm_iters = int(warmup * iters) + for it in range(iters): + idx = torch.randint(0, N, (min(batch_size, N),)) + xb = x_cal[idx].to(device, dtype=torch.float32) + h = _rect_sigmoid(V) + Wq = scale * torch.clamp(floorWs + h, qmin, qmax) + out_q = _layer_fwd(layer, xb, Wq, bias) + with torch.no_grad(): + out_fp = _layer_fwd(layer, xb, W, bias) + rec = (out_q - out_fp).pow(2).flatten(1).sum(1).mean() # per-sample SSE, averaged + if it < warm_iters: + beta, lam_t = beta_start, 0.0 # warmup: reconstruction only + else: + rel = (it - warm_iters) / max(1, iters - warm_iters) + beta = beta_end + (beta_start - beta_end) * max(0.0, 1 - rel) + lam_t = lam + round_loss = lam_t * (1 - (2 * h - 1).abs().pow(beta)).sum() # push h -> {0, 1} + loss = rec + round_loss + opt.zero_grad(); loss.backward(); opt.step() + with torch.no_grad(): + h = _rect_sigmoid(V) + not_binary = ((h > 0.02) & (h < 0.98)).float().mean().item() # unresolved fraction + Wq_hard = _hard_quant_from_V(V, floorWs, scale, qmin, qmax) + return Wq_hard, {'not_binary_frac': not_binary, 'w_bits': w_bits, 'granularity': granularity} + +# %% ../../nbs/quantize/adaround.ipynb #10fa2568 +def _iter_target_layers(model: nn.Module, # model to scan + layer_type = (nn.Conv2d, nn.Linear), # module type(s) to target + ): + "Iterate over modules of `layer_type` (the AdaRound analogue of `Sparsifier._iter_layers`)." + for m in model.modules(): + if isinstance(m, layer_type): + yield m + +def _normalize_calibration(calibration_dl, # DataLoaders | DataLoader | list of batches + n_batches: int | None = None, # cap on number of batches + ) -> list: # list of input tensors (xb) + "Normalise the three accepted calibration inputs to a list of input tensors." + dl = calibration_dl + if hasattr(dl, 'valid') and hasattr(dl, 'train'): # fastai DataLoaders -> use validation dl + dl = dl.valid + batches = [] + for b in dl: + xb = b[0] if isinstance(b, (list, tuple)) else b # (xb, yb) batch or a bare input tensor + batches.append(xb) + if n_batches is not None and len(batches) >= n_batches: + break + return batches + +@torch.no_grad() +def _capture_inputs(model, # (pristine) FP model + layers, # target layers to hook + batches, # list of input tensors + device, # compute device + store_dtype = torch.float16, # cache dtype (saves memory) + ) -> dict: # {layer: cached inputs [N, ...] on CPU} + "Cache the FP input tensor to every target layer in a single forward pass." + cache = {l: [] for l in layers} + def make_hook(layer): + def hook(mod, inp): + cache[layer].append(inp[0].detach().to(store_dtype).cpu()) + return hook + handles = [l.register_forward_pre_hook(make_hook(l)) for l in layers] + model.eval() + for xb in batches: + model(xb.to(device)) + for h in handles: h.remove() + return {l: torch.cat(v) for l, v in cache.items()} + +# %% ../../nbs/quantize/adaround.ipynb #aeba40b3 +def _fake_quant_act(x, scale, qmax): + "Per-tensor symmetric fake-quant of an activation tensor." + return torch.clamp(torch.round(x / scale), -qmax, qmax) * scale + +@torch.no_grad() +def _calibrate_activations(model, # weight-quantized model + layers, # target layers + batches, # calibration input tensors + act_bits: int, # activation bit-width + device, # compute device + n_batches: int = 8, # batches used for statistics + ): + "Calibrate MSE-optimal per-tensor activation scales and install persistent fake-quant pre-hooks." + qmax = 2 ** (act_bits - 1) - 1 + samples = {l: [] for l in layers} + def obs(mod, inp): + v = inp[0].detach().flatten() + if v.numel() > 20000: + v = v[torch.randint(0, v.numel(), (20000,), device=v.device)] + samples[mod].append(v.float().cpu()) + handles = [l.register_forward_pre_hook(obs) for l in layers] + model.eval() + for i, xb in enumerate(batches): + if i >= n_batches: break + model(xb.to(device)) + for h in handles: h.remove() + alphas = [1.0, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7, 0.65, 0.6, 0.55, 0.5, 0.4, 0.3] + def make_hook(scale, qmax): + def hook(mod, inp): + return (_fake_quant_act(inp[0], scale, qmax),) + tuple(inp[1:]) + return hook + for l in layers: + a = torch.cat(samples[l]); mx = a.abs().max().item() + best_mse, best_clip = float('inf'), mx + for al in alphas: # search clip minimising quant MSE + clip = al * mx + if clip <= 0: continue + scale = clip / qmax + q = torch.clamp(torch.round(a / scale), -qmax, qmax) * scale + mse = ((q - a) ** 2).mean().item() + if mse < best_mse: best_mse, best_clip = mse, clip + l.register_forward_pre_hook(make_hook(max(best_clip, 1e-12) / qmax, qmax)) + +# %% ../../nbs/quantize/adaround.ipynb #094c67f0 +def adaround_quantize( + model: nn.Module, # model whose weights -> INT-`w_bits` with learned rounding + calibration_dl, # fastai DataLoaders | torch DataLoader | list of (xb, yb)/xb batches -- REQUIRED + w_bits: int = 4, # weight bit-width + granularity: str = 'channel', # 'channel' (per-output-channel) or 'tensor' + act_bits: int | None = None, # activation bit-width (None -> weight-only) + iters: int = 2000, # AdaRound optimization iters per layer + lr: float = 1e-2, # Adam learning rate + fold_batchnorm: bool = True, # fuse Conv->BN before quantizing + layer_type = (nn.Conv2d, nn.Linear), # module type(s) to quantize + *, + lam: float = 0.01, # AdaRound rounding-regularizer weight + batch_size: int = 32, # minibatch of cached inputs per AdaRound step + n_calib_batches: int | None = None, # cap on calibration batches used + device: str | torch.device | None = None, # compute device (default: model's) + verbose: bool = False, # print per-layer progress +) -> nn.Module: # AdaRound fake-quantized model + """Post-training INT weight quantization with learned rounding (AdaRound, Nagel et al. 2020). + + Signed per-output-channel (or per-tensor) fake-quant whose rounding is *learned* per layer by + minimising output reconstruction error, recovering most of the INT4 round-to-nearest cliff for + redundant CNNs (ResNet-class ~ INT8; depthwise-heavy nets stay architecture-dependently lossy). + + Returns an accuracy-faithful **fake-quantized** model (weights quantized then de-quantized + in-place). It also stashes the exact per-layer quant params on `model._adaround_qparams` so + `adaround_to_onnx` can write a *faithful* INT4 QDQ ONNX (torchao is the other deployment path). In + the dimensional `Quantizer` this is `criteria='adaround'`.""" + if calibration_dl is None: + raise ValueError("adaround_quantize requires calibration data -- pass a DataLoaders, a " + "DataLoader, or a list of input batches as `calibration_dl` (got None).") + if granularity not in ('channel', 'tensor'): + raise ValueError(f"granularity must be 'channel' or 'tensor', got {granularity!r}.") + if not 2 <= w_bits <= 8: + raise ValueError(f"w_bits must be between 2 and 8, got {w_bits}.") + if act_bits is not None and not 2 <= act_bits <= 8: + raise ValueError(f"act_bits must be between 2 and 8 (or None), got {act_bits}.") + + qmodel = copy.deepcopy(model).eval() + device = torch.device(device) if device is not None else next(qmodel.parameters()).device + qmodel = qmodel.to(device) + if fold_batchnorm: fold_bn(qmodel) + + batches = _normalize_calibration(calibration_dl, n_calib_batches) + if len(batches) == 0: + raise ValueError("calibration_dl yielded no batches.") + + layers = list(_iter_target_layers(qmodel, layer_type)) + layer_names = {m: n for n, m in qmodel.named_modules()} # module -> dotted name (matches ONNX inits) + caches = _capture_inputs(qmodel, layers, batches, device) # FP inputs (parallel AdaRound) + + # stash exact per-layer quant params for faithful INT4 ONNX export (adaround_to_onnx). + # the scale is *not* recoverable post-hoc: AdaRound bakes W <- scale*int, so capture it here. + qmodel._adaround_qparams = {} # {name: (per-channel scale, qmin, qmax)} + for i, layer in enumerate(layers): + scale, qmin, qmax = weight_scale(layer.weight.detach().float(), w_bits, granularity) + qmodel._adaround_qparams[layer_names[layer]] = (scale.detach().cpu(), int(qmin), int(qmax)) + Wq, info = adaround_layer(layer, caches[layer], w_bits, granularity=granularity, + iters=iters, lr=lr, lam=lam, batch_size=batch_size, device=device) + layer.weight.data = Wq.to(layer.weight.dtype) + if verbose: + print(f"[{i+1}/{len(layers)}] {type(layer).__name__} " + f"not_binary={info['not_binary_frac']:.3f}") + + if act_bits is not None: + _calibrate_activations(qmodel, layers, batches, act_bits, device) + return qmodel + +# %% ../../nbs/quantize/adaround.ipynb #b5556811 +def _pack_int4(flat: np.ndarray, # 1-D array of signed ints in [-8, 7] + ) -> bytes: # packed two-nibbles-per-byte (onnxruntime INT4 raw_data layout) + "Pack signed int4 values two-per-byte (low nibble first), matching onnxruntime's INT4 initializers." + v = (flat.astype(np.int32) & 0xF).astype(np.uint8) + n = v.size; pad = (-n) % 2 + if pad: v = np.concatenate([v, np.zeros(pad, np.uint8)]) + v = v.reshape(-1, 2) + return (v[:, 0] | (v[:, 1] << 4)).astype(np.uint8).tobytes() + +def _unpack_int4(raw: bytes, # packed int4 bytes (two nibbles per byte) + n: int, # number of int4 values to recover + ) -> np.ndarray: # 1-D int8 array of the signed int4 values in [-8, 7] + "Inverse of `_pack_int4` (sign-extends each 4-bit nibble)." + b = np.frombuffer(raw, dtype=np.uint8) + out = np.empty(b.size * 2, dtype=np.int16) + out[0::2] = b & 0xF + out[1::2] = (b >> 4) & 0xF + out = out[:n] + return np.where(out >= 8, out - 16, out).astype(np.int8) + +def adaround_to_onnx( + model: nn.Module, # a model returned by `adaround_quantize` (carries `_adaround_qparams`) + sample: torch.Tensor, # one example input (batched or [C,H,W]) shaping the ONNX graph + path: str, # output path for the INT4 QDQ ONNX file + *, + input_name: str = 'image', # ONNX graph input name + calibration_dl = None, # optional calibration data for activation quant (else synthesized from `sample`) +) -> str: # `path` to the written INT4 QDQ ONNX + """Export an AdaRound model to a *faithful* INT4 QDQ ONNX. + + Standard ONNX exporters recalibrate the weight scale inside `quantize_static`, discarding + AdaRound's learned up/down rounding (a ~4-7 pt accuracy drop at INT4). This helper instead + writes AdaRound's **exact** per-channel `scale` and int4 weights into the QDQ initializers, so + the deployed ONNX carries the *learned* rounding and matches the fake-quant accuracy. + + The graph is produced in three steps: (1) export the fake-quant model to FP32 ONNX (legacy + exporter, dynamic batch); (2) run onnxruntime `quantize_static` (QDQ, per-channel QInt4 weights, + QInt8 activations) to get the correct QDQ *structure* with ORT's own (wrong) scales; (3) overwrite + the weight `_scale` and int4 `_quantized` initializers with AdaRound's exact params from + `model._adaround_qparams`. Requires `onnx` and `onnxruntime`.""" + qparams = getattr(model, '_adaround_qparams', None) + if qparams is None: + raise ValueError("model has no `_adaround_qparams` -- call adaround_quantize first " + "(the exact per-layer scales cannot be recovered from the baked weights).") + try: + import onnx + from onnx import numpy_helper + from onnxruntime.quantization import (quantize_static, QuantType, QuantFormat, + CalibrationDataReader, quant_pre_process, + CalibrationMethod) + except ImportError as e: + raise ImportError("adaround_to_onnx requires onnx and onnxruntime. " + "Install with: pip install onnx onnxruntime") from e + import os, tempfile + + orig_device = next(model.parameters()).device + model.eval().cpu() + dummy = sample.detach().cpu().float() + if dummy.dim() == 3: dummy = dummy.unsqueeze(0) + + # AdaRound-baked weights keyed by module name (== _adaround_qparams keys == ONNX init prefixes) + ada_w = {n: m.weight.detach().float().cpu() + for n, m in model.named_modules() if isinstance(m, (nn.Conv2d, nn.Linear))} + + # calibration images for the activation-quant statistics -> list of float32 arrays + if calibration_dl is not None: + calib = [b.detach().cpu().float().numpy() for b in _normalize_calibration(calibration_dl)] + else: # synthesize a few perturbed copies of `sample` so MinMax sees a plausible activation range + calib = [dummy.numpy()] + [(dummy + 0.1 * torch.randn_like(dummy)).numpy() for _ in range(3)] + + class _DR(CalibrationDataReader): + def __init__(self): self._it = iter([{input_name: a} for a in calib]) + def get_next(self): return next(self._it, None) + + with tempfile.TemporaryDirectory() as td: + fp, pre = os.path.join(td, 'fp32.onnx'), os.path.join(td, 'pre.onnx') + # (1) FP32 ONNX -- dynamo=False (the dynamo exporter breaks ORT's quant_pre_process) + torch.onnx.export(model, dummy, fp, input_names=[input_name], output_names=['output'], + dynamic_axes={input_name: {0: 'batch'}, 'output': {0: 'batch'}}, + opset_version=18, dynamo=False) + # (2) standard ORT INT4 QDQ -- correct graph, ORT's (wrong) scales/ints + quant_pre_process(fp, pre) + quantize_static(pre, path, _DR(), quant_format=QuantFormat.QDQ, per_channel=True, + weight_type=QuantType.QInt4, activation_type=QuantType.QInt8, + op_types_to_quantize=['Conv', 'Gemm'], + calibrate_method=CalibrationMethod.MinMax) + model.to(orig_device) + + # (3) overwrite the weight initializers with AdaRound's EXACT int4 + scale + m = onnx.load(path) + inits = {i.name: i for i in m.graph.initializer} + for nm in list(inits): + if not nm.endswith('.weight_quantized'): continue + base = nm[:-len('_quantized')] # '.weight' + mod = base[:-len('.weight')] # '' + if mod not in qparams: continue + scale, qmin, qmax = qparams[mod] + scale = scale.detach().cpu() + inits[base + '_scale'].CopyFrom( + numpy_helper.from_array(scale.reshape(-1).float().numpy(), base + '_scale')) + int_ada = torch.round(ada_w[mod].double() / scale.double()).clamp(qmin, qmax).to(torch.int8).numpy() + t = inits[nm]; t.ClearField('int32_data'); t.raw_data = _pack_int4(int_ada.reshape(-1)) + # drop the value_info that duplicates graph IO (an artifact ORT can leave behind) + io = {o.name for o in m.graph.output} | {i.name for i in m.graph.input} + keep = [vi for vi in m.graph.value_info if vi.name not in io] + del m.graph.value_info[:]; m.graph.value_info.extend(keep) + onnx.save(m, path) + return path + diff --git a/nbs/quantize/DIMENSIONAL_DESIGN.md b/nbs/quantize/DIMENSIONAL_DESIGN.md new file mode 100644 index 0000000..25bf9bb --- /dev/null +++ b/nbs/quantize/DIMENSIONAL_DESIGN.md @@ -0,0 +1,103 @@ +# Dimensional `Quantizer` — design north-star + +**Status:** design spec (target architecture). AdaRound is the first concrete piece built toward it. + +## Principle + +fasterai's pruning API is a small set of **composable, orthogonal dimensions** — `Sparsifier(granularity, context, criteria, schedule)`. Every pruning technique is a *value on an axis*, not a bespoke function. This composability (one grammar, sensitivity-driven) is the moat, more than any single technique. + +Quantization has the **same dimensional structure**. This spec defines a `Quantizer` that mirrors `Sparsifier`, so a user who knows one knows the other, and every quant technique — existing (`quantize_mixed`, SmoothQuant) and new (AdaRound, GPTQ) — is a composable value, not a separate feature. + +## The dimensions (4 shared with pruning + 1 new) + +| dimension | Pruning (`Sparsifier`) | Quantization (`Quantizer`) | +|---|---|---| +| **granularity** — scope a scale is shared over | weight / vector / kernel / filter | `tensor` / `channel` / `group` / `block` | +| **context** — allocation across layers | local / global | `uniform` / `mixed` (per-layer bits, sensitivity-driven) | +| **criteria** — how values are decided | large_final / movement / gradient / Wanda | `minmax` / `percentile` / `mse` / `adaround` / `gptq` | +| **schedule** — when it is applied | one_shot / iterative / agp / cos | `None` (PTQ) / a `Schedule` (progressive QAT) | +| **precision** — how aggressive *(new; no pruning analog)* | — | `int8` / `int4` / `int16` / `fp16` / `bf16`, or `{layer: dtype}` for mixed | + +## Target API + +```python +class Quantizer: + "Reduce numerical precision along composable dimensions — the quantization analogue of Sparsifier." + def __init__(self, + model: nn.Module, + granularity: str = 'channel', # 'tensor'|'channel'|'group'|'block' + context: str = 'uniform', # 'uniform' | 'mixed' (per-layer, sensitivity-driven) + criteria: QuantCriteria = minmax, # minmax|percentile|mse|adaround|gptq + schedule: Schedule | None = None, # None -> PTQ one-shot; Schedule -> progressive QAT + dtype: str | dict[str, str] = 'int8', # 'int8'|'int4'|'int16'|'fp16'|'bf16' OR {layer: dtype} for mixed + *, + act_dtype: str | None = 'int8', # activation precision (None -> weight-only, e.g. W4A16) + symmetric: bool = True, # symmetric vs asymmetric zero-point + dynamic: bool = False, # activations: dynamic (per-inference) vs static (calibrated) + group_size: int = 128, # block width for 'group'/'block' granularity (INT4) + layer_type: type = nn.Conv2d, # which modules to quantize (like Sparsifier) + ): ... + + def quantize_model(self, calibration_dl=None) -> nn.Module: + "PTQ: calibrate then apply. For QAT, use QuantizeCallback(schedule=...)." +``` + +`dtype` is **a value or a per-layer dict** — exactly like `Sparsifier.sparsity`. That single choice makes uniform vs mixed-precision fall out of one API. + +## `QuantCriteria` — the calibration strategy (mirrors pruning `Criteria`) + +```python +class QuantCriteria: + "How to choose (scale, zero_point) for a tensor — the quant analogue of a pruning Criteria." + needs_data: bool # needs calibration_dl? (adaround/gptq/percentile -> yes) + def params(self, w, n_bits, granularity, symmetric, calib=None) -> tuple[Tensor, Tensor]: ... + +minmax = QuantCriteria(...) # min/max range — closed-form, data-free for weights +percentile = lambda p=99.9: ... # clip outliers at the p-th percentile +mse = QuantCriteria(...) # scale minimising quantization MSE +adaround = QuantCriteria(...) # learned per-weight rounding (needs calib) — the "make INT4 safe" criterion +gptq = QuantCriteria(...) # OBC / Hessian-guided reconstruction (needs calib) +``` + +## Every technique is a config, not a feature + +| technique | how it is expressed | +|---|---| +| per-channel INT8 (near-lossless default) | `granularity='channel', dtype='int8'` | +| per-tensor (naive baseline) | `granularity='tensor'` | +| per-group INT4 weight-only | `granularity='group', dtype='int4', act_dtype=None, group_size=128` | +| **AdaRound** | `criteria=adaround` | +| **GPTQ** | `criteria=gptq` | +| percentile / MSE calibration | `criteria=percentile()` / `criteria=mse` | +| **mixed-precision** (was `quantize_mixed()`) | `context='mixed'` -> per-layer `dtype` dict from sensitivity | +| W4A16 | `dtype='int4', act_dtype='fp16'` | +| dynamic quant | `dynamic=True` | +| data-free | pass a *synthetic* `calibration_dl` (BN-stat-generated) — orthogonal to the axes | +| **SmoothQuant** | a pre-quant **transform** (like `BN_Folder`), not a Quantizer dim | +| ~~CLE / Bias-Correction~~ | pre-quant transform — **validated redundant with per-channel; not integrated** | + +## Mechanism / decision split (same as pruning) + +- **fasterai `Quantizer` = mechanism.** Applies quantization at a given granularity/dtype/criteria; accepts `dtype` as a **per-layer dict** when `context='mixed'`. It does not decide the allocation. +- **fasterrecipes = decision.** For `context='mixed'`, it calls sensitivity analysis to produce the per-layer bit-width dict, then hands it to the Quantizer — exactly how `Sparsifier` takes a sparsity dict that fasterrecipes computes. + +**The same sensitivity engine feeds both pruning `layer_targets` and quantization bit-width** -> one allocation brain, two techniques -> cross-technique co-optimization (allocate ratio *and* bits per layer from one sensitivity pass). This is the differentiator neither Pruna nor Embedl has. + +## Validated evidence behind the design + +- **per-channel INT8 is already near-lossless** (−0.4 pt on MobileNetV2). "Safe INT8" is solved by per-channel granularity; it needs no help. +- **AdaRound makes INT4 viable — but architecture-dependently.** ResNet-class: INT4+AdaRound ≈ INT8 (~1 pt). MobileNet-class: still ~57 pt below INT8 (depthwise convs are the INT4-killer). => INT4 is the "up to ~8× smaller at ~1 pt" lever *for redundant CNNs*, and `context='mixed'` (route depthwise to INT8) is how efficient nets are handled. +- **INT4 is a size lever, not a CPU-speed one** (our own benchmark: W4A32 is ~4× *slower* on CPU, no fast kernel). Market INT4 as size/memory. +- **CLE / Bias-Correction: validated redundant** with per-channel (within noise) and a −24 pt ReLU6 landmine — *not* integrated. + +## Migration path (incremental, non-breaking) + +The current `Quantizer` is backend-oriented (`backend`, `method`). We do **not** big-bang refactor it. Instead: +1. Add `AdaRound` as a self-contained criterion + a thin apply path (this branch) — shaped to become `criteria='adaround'`. +2. Revive mixed-precision (`quantize_mixed`) on current master as `context='mixed'`, fed by sensitivity. +3. Add SmoothQuant as a pre-quant transform. +4. Unify under the dimensional `Quantizer` once the pieces exist. + +## Honest constraint (encoded, not hidden) + +The API can *express* combinations that don't deploy fast (per-group INT4 on CPU, mixed-precision on a runtime that ignores it). That is intentional: `quantize_model` gives faithful **accuracy** (fake-quant), and `benchmark_on_device` (fasterrecipes) says which combo actually **pays off** on the target. Express anything; verify what ships. diff --git a/nbs/quantize/adaround.ipynb b/nbs/quantize/adaround.ipynb new file mode 100644 index 0000000..cf157d3 --- /dev/null +++ b/nbs/quantize/adaround.ipynb @@ -0,0 +1,883 @@ +{ + "cells": [ + { + "cell_type": "raw", + "id": "711e8d73", + "metadata": {}, + "source": [ + "---\n", + "description: Post-training INT4 weight quantization with learned rounding (AdaRound)\n", + "output-file: adaround.html\n", + "title: AdaRound\n", + "skip_showdoc: true\n", + "---" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e5db50e4", + "metadata": {}, + "outputs": [], + "source": [ + "#| default_exp quantize.adaround" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1746d882", + "metadata": {}, + "outputs": [], + "source": [ + "#| export\n", + "from __future__ import annotations\n", + "import copy\n", + "import numpy as np\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.nn.functional as F\n", + "from torch.nn.utils.fusion import fuse_conv_bn_eval" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "447d6abe", + "metadata": {}, + "outputs": [], + "source": [ + "#| include: false\n", + "from nbdev.showdoc import *" + ] + }, + { + "cell_type": "markdown", + "id": "41fe7f45", + "metadata": {}, + "source": [ + "## Overview\n", + "\n", + "`adaround_quantize` applies **AdaRound** (Nagel et al., 2020, *\"Up or Down? Adaptive Rounding for\n", + "Post-Training Quantization\"*) — the learned-rounding criterion that makes **INT4 weight\n", + "quantization** viable for redundant CNNs.\n", + "\n", + "Naive round-to-nearest (RTN) rounds every weight to its closest grid point independently. At INT4\n", + "that greedy choice is *jointly* far from optimal: small per-weight rounding errors accumulate\n", + "through the layer and collapse accuracy (the \"INT4 cliff\"). AdaRound instead **learns**, per layer,\n", + "whether each weight should round *up* or *down* — minimising the layer's output reconstruction error\n", + "`‖W_q x − W x‖²` on a little calibration data. It recovers most of the INT4 cliff on ResNet-class\n", + "networks (≈ INT8, ~1 pt) while leaving depthwise-heavy nets (MobileNet) architecture-dependently\n", + "lossy — so INT4 is a *size* lever for redundant CNNs, not a universal one.\n", + "\n", + "### How it maps to the dimensional `Quantizer`\n", + "\n", + "This is the first concrete criterion built toward the dimensional `Quantizer`\n", + "(see `DIMENSIONAL_DESIGN.md`). In that grammar it is `criteria='adaround'`:\n", + "\n", + "| dimension | value here |\n", + "|---|---|\n", + "| **granularity** | `'channel'` (per-output-channel) or `'tensor'` |\n", + "| **criteria** | `adaround` — learned rounding (this notebook) |\n", + "| **precision** | `w_bits` weights, optional `act_bits` activations |\n", + "| **schedule** | `None` — one-shot PTQ |\n", + "\n", + "The returned model is an **accuracy-faithful fake-quantized** model (weights are quantized then\n", + "de-quantized in-place). Real INT4-kernel export (torchao / ONNX-QDQ) is the deployment follow-up." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "80a30f53", + "metadata": {}, + "outputs": [], + "source": [ + "#| export\n", + "def _set_child(parent, name, new):\n", + " \"Replace child module `name` of `parent` (handles Sequential integer indices).\"\n", + " if name.isdigit(): parent[int(name)] = new\n", + " else: setattr(parent, name, new)\n", + "\n", + "def fold_bn(model: nn.Module # model containing Conv2d→BatchNorm2d pairs\n", + " ) -> nn.Module: # same model with each Conv-BN pair fused in-place\n", + " \"Fuse each `Conv2d`→`BatchNorm2d` pair into the conv (eval mode) and replace the BN with Identity.\"\n", + " for mod in model.modules():\n", + " kids = list(mod.named_children())\n", + " for i in range(len(kids) - 1):\n", + " (n1, c1), (n2, c2) = kids[i], kids[i + 1]\n", + " if isinstance(c1, nn.Conv2d) and isinstance(c2, nn.BatchNorm2d):\n", + " _set_child(mod, n1, fuse_conv_bn_eval(c1, c2))\n", + " _set_child(mod, n2, nn.Identity())\n", + " return model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c9f0e214", + "metadata": {}, + "outputs": [], + "source": [ + "#| export\n", + "def weight_scale(W: torch.Tensor, # weight tensor [out_channels, ...]\n", + " w_bits: int = 4, # bit-width (signed)\n", + " granularity: str = 'channel', # 'channel' (per-output-channel) or 'tensor'\n", + " ): # (scale, qmin, qmax) — dequant scale and signed clip range\n", + " \"Signed weight quantization `scale` and clip range `[qmin, qmax]` on the ONNX/ORT grid.\"\n", + " qmax = 2 ** (w_bits - 1) - 1 # INT4 -> 7, INT8 -> 127\n", + " qmin = -(2 ** (w_bits - 1)) # INT4 -> -8, INT8 -> -128 (asymmetric signed grid)\n", + " denom = 2 ** (w_bits - 1) - 0.5 # INT4 -> 7.5 (matches onnxruntime per-channel int4)\n", + " if granularity == 'channel':\n", + " r = W.detach().abs().reshape(W.shape[0], -1).amax(1).clamp_min(1e-12)\n", + " scale = (r / denom).reshape([-1] + [1] * (W.dim() - 1))\n", + " else:\n", + " scale = W.detach().abs().amax().clamp_min(1e-12) / denom\n", + " return scale, qmin, qmax\n", + "\n", + "def rtn_quant(W: torch.Tensor, # weight tensor to fake-quantize\n", + " w_bits: int = 4, # bit-width\n", + " granularity: str = 'channel', # 'channel' or 'tensor'\n", + " ) -> torch.Tensor: # dequantized (fake-quant) weights, same shape\n", + " \"Round-to-nearest (RTN) signed fake-quant — the AdaRound baseline (== AdaRound at init).\"\n", + " scale, qmin, qmax = weight_scale(W, w_bits, granularity)\n", + " floorWs = torch.floor(W / scale)\n", + " q = floorWs + ((W / scale - floorWs) >= 0.5).float() # round half up (matches the AdaRound hard round)\n", + " return torch.clamp(q, qmin, qmax) * scale" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d948dfff", + "metadata": {}, + "outputs": [], + "source": [ + "#| export\n", + "# AdaRound rectified-sigmoid soft rounding (Nagel et al., 2020):\n", + "# h(V) = clamp(sigmoid(V)*(zeta-gamma)+gamma, 0, 1)\n", + "# W_q(V) = s * clamp(floor(W/s) + h(V), qmin, qmax)\n", + "_ZETA, _GAMMA = 1.1, -0.1\n", + "\n", + "def _rect_sigmoid(V: torch.Tensor) -> torch.Tensor:\n", + " \"Rectified sigmoid squashing `V` into a soft rounding decision in [0, 1].\"\n", + " return torch.clamp(torch.sigmoid(V) * (_ZETA - _GAMMA) + _GAMMA, 0, 1)\n", + "\n", + "def _init_V(frac: torch.Tensor) -> torch.Tensor:\n", + " \"Initialise `V` so that `h(V) == frac(W/s)` — makes AdaRound identical to RTN at init.\"\n", + " s0 = ((frac - _GAMMA) / (_ZETA - _GAMMA)).clamp(1e-6, 1 - 1e-6)\n", + " return torch.log(s0 / (1 - s0))\n", + "\n", + "def _hard_quant_from_V(V, floorWs, scale, qmin, qmax):\n", + " \"Deployed weights: hard-round the learned soft rounding `h(V)` onto the signed grid.\"\n", + " h_hard = (_rect_sigmoid(V) >= 0.5).float()\n", + " return scale * torch.clamp(floorWs + h_hard, qmin, qmax)\n", + "\n", + "def _layer_fwd(mod, x, w, b):\n", + " \"Forward one Conv2d/Linear layer with an explicit weight/bias (for reconstruction).\"\n", + " if isinstance(mod, nn.Conv2d):\n", + " return F.conv2d(x, w, b, mod.stride, mod.padding, mod.dilation, mod.groups)\n", + " return F.linear(x, w, b)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ab024063", + "metadata": {}, + "outputs": [], + "source": [ + "#| export\n", + "def adaround_layer(layer: nn.Module, # Conv2d/Linear whose rounding is optimized\n", + " x_cal: torch.Tensor, # cached FP inputs to this layer [N, ...]\n", + " w_bits: int = 4, # weight bit-width\n", + " *,\n", + " granularity: str = 'channel', # 'channel' or 'tensor'\n", + " iters: int = 2000, # optimization iterations\n", + " lr: float = 1e-2, # Adam learning rate\n", + " lam: float = 0.01, # rounding-regularizer weight\n", + " beta_start: float = 20.0, # anneal start (encourages soft)\n", + " beta_end: float = 2.0, # anneal end (encourages binary)\n", + " warmup: float = 0.2, # fraction of iters before regularizer turns on\n", + " batch_size: int = 32, # minibatch of cached inputs per step\n", + " device: str | torch.device | None = None, # compute device (default: layer's)\n", + " ): # (Wq_hard, info) — deployed fake-quant weights + diagnostics\n", + " \"Learn per-weight soft rounding for one layer, minimising output reconstruction ‖W_q(V)x − Wx‖² + regularizer.\"\n", + " device = torch.device(device) if device is not None else layer.weight.device\n", + " W = layer.weight.detach().float().to(device)\n", + " scale, qmin, qmax = weight_scale(W, w_bits, granularity)\n", + " floorWs = torch.floor(W / scale)\n", + " frac = (W / scale) - floorWs # in [0, 1)\n", + " V = _init_V(frac).clone().requires_grad_(True) # init so h(V) == frac (== RTN)\n", + " opt = torch.optim.Adam([V], lr=lr)\n", + " bias = layer.bias.detach().float().to(device) if layer.bias is not None else None\n", + " N = x_cal.shape[0]\n", + " warm_iters = int(warmup * iters)\n", + " for it in range(iters):\n", + " idx = torch.randint(0, N, (min(batch_size, N),))\n", + " xb = x_cal[idx].to(device, dtype=torch.float32)\n", + " h = _rect_sigmoid(V)\n", + " Wq = scale * torch.clamp(floorWs + h, qmin, qmax)\n", + " out_q = _layer_fwd(layer, xb, Wq, bias)\n", + " with torch.no_grad():\n", + " out_fp = _layer_fwd(layer, xb, W, bias)\n", + " rec = (out_q - out_fp).pow(2).flatten(1).sum(1).mean() # per-sample SSE, averaged\n", + " if it < warm_iters:\n", + " beta, lam_t = beta_start, 0.0 # warmup: reconstruction only\n", + " else:\n", + " rel = (it - warm_iters) / max(1, iters - warm_iters)\n", + " beta = beta_end + (beta_start - beta_end) * max(0.0, 1 - rel)\n", + " lam_t = lam\n", + " round_loss = lam_t * (1 - (2 * h - 1).abs().pow(beta)).sum() # push h -> {0, 1}\n", + " loss = rec + round_loss\n", + " opt.zero_grad(); loss.backward(); opt.step()\n", + " with torch.no_grad():\n", + " h = _rect_sigmoid(V)\n", + " not_binary = ((h > 0.02) & (h < 0.98)).float().mean().item() # unresolved fraction\n", + " Wq_hard = _hard_quant_from_V(V, floorWs, scale, qmin, qmax)\n", + " return Wq_hard, {'not_binary_frac': not_binary, 'w_bits': w_bits, 'granularity': granularity}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10fa2568", + "metadata": {}, + "outputs": [], + "source": [ + "#| export\n", + "def _iter_target_layers(model: nn.Module, # model to scan\n", + " layer_type = (nn.Conv2d, nn.Linear), # module type(s) to target\n", + " ):\n", + " \"Iterate over modules of `layer_type` (the AdaRound analogue of `Sparsifier._iter_layers`).\"\n", + " for m in model.modules():\n", + " if isinstance(m, layer_type):\n", + " yield m\n", + "\n", + "def _normalize_calibration(calibration_dl, # DataLoaders | DataLoader | list of batches\n", + " n_batches: int | None = None, # cap on number of batches\n", + " ) -> list: # list of input tensors (xb)\n", + " \"Normalise the three accepted calibration inputs to a list of input tensors.\"\n", + " dl = calibration_dl\n", + " if hasattr(dl, 'valid') and hasattr(dl, 'train'): # fastai DataLoaders -> use validation dl\n", + " dl = dl.valid\n", + " batches = []\n", + " for b in dl:\n", + " xb = b[0] if isinstance(b, (list, tuple)) else b # (xb, yb) batch or a bare input tensor\n", + " batches.append(xb)\n", + " if n_batches is not None and len(batches) >= n_batches:\n", + " break\n", + " return batches\n", + "\n", + "@torch.no_grad()\n", + "def _capture_inputs(model, # (pristine) FP model\n", + " layers, # target layers to hook\n", + " batches, # list of input tensors\n", + " device, # compute device\n", + " store_dtype = torch.float16, # cache dtype (saves memory)\n", + " ) -> dict: # {layer: cached inputs [N, ...] on CPU}\n", + " \"Cache the FP input tensor to every target layer in a single forward pass.\"\n", + " cache = {l: [] for l in layers}\n", + " def make_hook(layer):\n", + " def hook(mod, inp):\n", + " cache[layer].append(inp[0].detach().to(store_dtype).cpu())\n", + " return hook\n", + " handles = [l.register_forward_pre_hook(make_hook(l)) for l in layers]\n", + " model.eval()\n", + " for xb in batches:\n", + " model(xb.to(device))\n", + " for h in handles: h.remove()\n", + " return {l: torch.cat(v) for l, v in cache.items()}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aeba40b3", + "metadata": {}, + "outputs": [], + "source": [ + "#| export\n", + "def _fake_quant_act(x, scale, qmax):\n", + " \"Per-tensor symmetric fake-quant of an activation tensor.\"\n", + " return torch.clamp(torch.round(x / scale), -qmax, qmax) * scale\n", + "\n", + "@torch.no_grad()\n", + "def _calibrate_activations(model, # weight-quantized model\n", + " layers, # target layers\n", + " batches, # calibration input tensors\n", + " act_bits: int, # activation bit-width\n", + " device, # compute device\n", + " n_batches: int = 8, # batches used for statistics\n", + " ):\n", + " \"Calibrate MSE-optimal per-tensor activation scales and install persistent fake-quant pre-hooks.\"\n", + " qmax = 2 ** (act_bits - 1) - 1\n", + " samples = {l: [] for l in layers}\n", + " def obs(mod, inp):\n", + " v = inp[0].detach().flatten()\n", + " if v.numel() > 20000:\n", + " v = v[torch.randint(0, v.numel(), (20000,), device=v.device)]\n", + " samples[mod].append(v.float().cpu())\n", + " handles = [l.register_forward_pre_hook(obs) for l in layers]\n", + " model.eval()\n", + " for i, xb in enumerate(batches):\n", + " if i >= n_batches: break\n", + " model(xb.to(device))\n", + " for h in handles: h.remove()\n", + " alphas = [1.0, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7, 0.65, 0.6, 0.55, 0.5, 0.4, 0.3]\n", + " def make_hook(scale, qmax):\n", + " def hook(mod, inp):\n", + " return (_fake_quant_act(inp[0], scale, qmax),) + tuple(inp[1:])\n", + " return hook\n", + " for l in layers:\n", + " a = torch.cat(samples[l]); mx = a.abs().max().item()\n", + " best_mse, best_clip = float('inf'), mx\n", + " for al in alphas: # search clip minimising quant MSE\n", + " clip = al * mx\n", + " if clip <= 0: continue\n", + " scale = clip / qmax\n", + " q = torch.clamp(torch.round(a / scale), -qmax, qmax) * scale\n", + " mse = ((q - a) ** 2).mean().item()\n", + " if mse < best_mse: best_mse, best_clip = mse, clip\n", + " l.register_forward_pre_hook(make_hook(max(best_clip, 1e-12) / qmax, qmax))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "094c67f0", + "metadata": {}, + "outputs": [], + "source": [ + "#| export\n", + "def adaround_quantize(\n", + " model: nn.Module, # model whose weights -> INT-`w_bits` with learned rounding\n", + " calibration_dl, # fastai DataLoaders | torch DataLoader | list of (xb, yb)/xb batches -- REQUIRED\n", + " w_bits: int = 4, # weight bit-width\n", + " granularity: str = 'channel', # 'channel' (per-output-channel) or 'tensor'\n", + " act_bits: int | None = None, # activation bit-width (None -> weight-only)\n", + " iters: int = 2000, # AdaRound optimization iters per layer\n", + " lr: float = 1e-2, # Adam learning rate\n", + " fold_batchnorm: bool = True, # fuse Conv->BN before quantizing\n", + " layer_type = (nn.Conv2d, nn.Linear), # module type(s) to quantize\n", + " *,\n", + " lam: float = 0.01, # AdaRound rounding-regularizer weight\n", + " batch_size: int = 32, # minibatch of cached inputs per AdaRound step\n", + " n_calib_batches: int | None = None, # cap on calibration batches used\n", + " device: str | torch.device | None = None, # compute device (default: model's)\n", + " verbose: bool = False, # print per-layer progress\n", + ") -> nn.Module: # AdaRound fake-quantized model\n", + " \"\"\"Post-training INT weight quantization with learned rounding (AdaRound, Nagel et al. 2020).\n", + "\n", + " Signed per-output-channel (or per-tensor) fake-quant whose rounding is *learned* per layer by\n", + " minimising output reconstruction error, recovering most of the INT4 round-to-nearest cliff for\n", + " redundant CNNs (ResNet-class ~ INT8; depthwise-heavy nets stay architecture-dependently lossy).\n", + "\n", + " Returns an accuracy-faithful **fake-quantized** model (weights quantized then de-quantized\n", + " in-place). It also stashes the exact per-layer quant params on `model._adaround_qparams` so\n", + " `adaround_to_onnx` can write a *faithful* INT4 QDQ ONNX (torchao is the other deployment path). In\n", + " the dimensional `Quantizer` this is `criteria='adaround'`.\"\"\"\n", + " if calibration_dl is None:\n", + " raise ValueError(\"adaround_quantize requires calibration data -- pass a DataLoaders, a \"\n", + " \"DataLoader, or a list of input batches as `calibration_dl` (got None).\")\n", + " if granularity not in ('channel', 'tensor'):\n", + " raise ValueError(f\"granularity must be 'channel' or 'tensor', got {granularity!r}.\")\n", + " if not 2 <= w_bits <= 8:\n", + " raise ValueError(f\"w_bits must be between 2 and 8, got {w_bits}.\")\n", + " if act_bits is not None and not 2 <= act_bits <= 8:\n", + " raise ValueError(f\"act_bits must be between 2 and 8 (or None), got {act_bits}.\")\n", + "\n", + " qmodel = copy.deepcopy(model).eval()\n", + " device = torch.device(device) if device is not None else next(qmodel.parameters()).device\n", + " qmodel = qmodel.to(device)\n", + " if fold_batchnorm: fold_bn(qmodel)\n", + "\n", + " batches = _normalize_calibration(calibration_dl, n_calib_batches)\n", + " if len(batches) == 0:\n", + " raise ValueError(\"calibration_dl yielded no batches.\")\n", + "\n", + " layers = list(_iter_target_layers(qmodel, layer_type))\n", + " layer_names = {m: n for n, m in qmodel.named_modules()} # module -> dotted name (matches ONNX inits)\n", + " caches = _capture_inputs(qmodel, layers, batches, device) # FP inputs (parallel AdaRound)\n", + "\n", + " # stash exact per-layer quant params for faithful INT4 ONNX export (adaround_to_onnx).\n", + " # the scale is *not* recoverable post-hoc: AdaRound bakes W <- scale*int, so capture it here.\n", + " qmodel._adaround_qparams = {} # {name: (per-channel scale, qmin, qmax)}\n", + " for i, layer in enumerate(layers):\n", + " scale, qmin, qmax = weight_scale(layer.weight.detach().float(), w_bits, granularity)\n", + " qmodel._adaround_qparams[layer_names[layer]] = (scale.detach().cpu(), int(qmin), int(qmax))\n", + " Wq, info = adaround_layer(layer, caches[layer], w_bits, granularity=granularity,\n", + " iters=iters, lr=lr, lam=lam, batch_size=batch_size, device=device)\n", + " layer.weight.data = Wq.to(layer.weight.dtype)\n", + " if verbose:\n", + " print(f\"[{i+1}/{len(layers)}] {type(layer).__name__} \"\n", + " f\"not_binary={info['not_binary_frac']:.3f}\")\n", + "\n", + " if act_bits is not None:\n", + " _calibrate_activations(qmodel, layers, batches, act_bits, device)\n", + " return qmodel" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b5556811", + "metadata": {}, + "outputs": [], + "source": [ + "#| export\n", + "def _pack_int4(flat: np.ndarray, # 1-D array of signed ints in [-8, 7]\n", + " ) -> bytes: # packed two-nibbles-per-byte (onnxruntime INT4 raw_data layout)\n", + " \"Pack signed int4 values two-per-byte (low nibble first), matching onnxruntime's INT4 initializers.\"\n", + " v = (flat.astype(np.int32) & 0xF).astype(np.uint8)\n", + " n = v.size; pad = (-n) % 2\n", + " if pad: v = np.concatenate([v, np.zeros(pad, np.uint8)])\n", + " v = v.reshape(-1, 2)\n", + " return (v[:, 0] | (v[:, 1] << 4)).astype(np.uint8).tobytes()\n", + "\n", + "def _unpack_int4(raw: bytes, # packed int4 bytes (two nibbles per byte)\n", + " n: int, # number of int4 values to recover\n", + " ) -> np.ndarray: # 1-D int8 array of the signed int4 values in [-8, 7]\n", + " \"Inverse of `_pack_int4` (sign-extends each 4-bit nibble).\"\n", + " b = np.frombuffer(raw, dtype=np.uint8)\n", + " out = np.empty(b.size * 2, dtype=np.int16)\n", + " out[0::2] = b & 0xF\n", + " out[1::2] = (b >> 4) & 0xF\n", + " out = out[:n]\n", + " return np.where(out >= 8, out - 16, out).astype(np.int8)\n", + "\n", + "def adaround_to_onnx(\n", + " model: nn.Module, # a model returned by `adaround_quantize` (carries `_adaround_qparams`)\n", + " sample: torch.Tensor, # one example input (batched or [C,H,W]) shaping the ONNX graph\n", + " path: str, # output path for the INT4 QDQ ONNX file\n", + " *,\n", + " input_name: str = 'image', # ONNX graph input name\n", + " calibration_dl = None, # optional calibration data for activation quant (else synthesized from `sample`)\n", + ") -> str: # `path` to the written INT4 QDQ ONNX\n", + " \"\"\"Export an AdaRound model to a *faithful* INT4 QDQ ONNX.\n", + "\n", + " Standard ONNX exporters recalibrate the weight scale inside `quantize_static`, discarding\n", + " AdaRound's learned up/down rounding (a ~4-7 pt accuracy drop at INT4). This helper instead\n", + " writes AdaRound's **exact** per-channel `scale` and int4 weights into the QDQ initializers, so\n", + " the deployed ONNX carries the *learned* rounding and matches the fake-quant accuracy.\n", + "\n", + " The graph is produced in three steps: (1) export the fake-quant model to FP32 ONNX (legacy\n", + " exporter, dynamic batch); (2) run onnxruntime `quantize_static` (QDQ, per-channel QInt4 weights,\n", + " QInt8 activations) to get the correct QDQ *structure* with ORT's own (wrong) scales; (3) overwrite\n", + " the weight `_scale` and int4 `_quantized` initializers with AdaRound's exact params from\n", + " `model._adaround_qparams`. Requires `onnx` and `onnxruntime`.\"\"\"\n", + " qparams = getattr(model, '_adaround_qparams', None)\n", + " if qparams is None:\n", + " raise ValueError(\"model has no `_adaround_qparams` -- call adaround_quantize first \"\n", + " \"(the exact per-layer scales cannot be recovered from the baked weights).\")\n", + " try:\n", + " import onnx\n", + " from onnx import numpy_helper\n", + " from onnxruntime.quantization import (quantize_static, QuantType, QuantFormat,\n", + " CalibrationDataReader, quant_pre_process,\n", + " CalibrationMethod)\n", + " except ImportError as e:\n", + " raise ImportError(\"adaround_to_onnx requires onnx and onnxruntime. \"\n", + " \"Install with: pip install onnx onnxruntime\") from e\n", + " import os, tempfile\n", + "\n", + " orig_device = next(model.parameters()).device\n", + " model.eval().cpu()\n", + " dummy = sample.detach().cpu().float()\n", + " if dummy.dim() == 3: dummy = dummy.unsqueeze(0)\n", + "\n", + " # AdaRound-baked weights keyed by module name (== _adaround_qparams keys == ONNX init prefixes)\n", + " ada_w = {n: m.weight.detach().float().cpu()\n", + " for n, m in model.named_modules() if isinstance(m, (nn.Conv2d, nn.Linear))}\n", + "\n", + " # calibration images for the activation-quant statistics -> list of float32 arrays\n", + " if calibration_dl is not None:\n", + " calib = [b.detach().cpu().float().numpy() for b in _normalize_calibration(calibration_dl)]\n", + " else: # synthesize a few perturbed copies of `sample` so MinMax sees a plausible activation range\n", + " calib = [dummy.numpy()] + [(dummy + 0.1 * torch.randn_like(dummy)).numpy() for _ in range(3)]\n", + "\n", + " class _DR(CalibrationDataReader):\n", + " def __init__(self): self._it = iter([{input_name: a} for a in calib])\n", + " def get_next(self): return next(self._it, None)\n", + "\n", + " with tempfile.TemporaryDirectory() as td:\n", + " fp, pre = os.path.join(td, 'fp32.onnx'), os.path.join(td, 'pre.onnx')\n", + " # (1) FP32 ONNX -- dynamo=False (the dynamo exporter breaks ORT's quant_pre_process)\n", + " torch.onnx.export(model, dummy, fp, input_names=[input_name], output_names=['output'],\n", + " dynamic_axes={input_name: {0: 'batch'}, 'output': {0: 'batch'}},\n", + " opset_version=18, dynamo=False)\n", + " # (2) standard ORT INT4 QDQ -- correct graph, ORT's (wrong) scales/ints\n", + " quant_pre_process(fp, pre)\n", + " quantize_static(pre, path, _DR(), quant_format=QuantFormat.QDQ, per_channel=True,\n", + " weight_type=QuantType.QInt4, activation_type=QuantType.QInt8,\n", + " op_types_to_quantize=['Conv', 'Gemm'],\n", + " calibrate_method=CalibrationMethod.MinMax)\n", + " model.to(orig_device)\n", + "\n", + " # (3) overwrite the weight initializers with AdaRound's EXACT int4 + scale\n", + " m = onnx.load(path)\n", + " inits = {i.name: i for i in m.graph.initializer}\n", + " for nm in list(inits):\n", + " if not nm.endswith('.weight_quantized'): continue\n", + " base = nm[:-len('_quantized')] # '.weight'\n", + " mod = base[:-len('.weight')] # ''\n", + " if mod not in qparams: continue\n", + " scale, qmin, qmax = qparams[mod]\n", + " scale = scale.detach().cpu()\n", + " inits[base + '_scale'].CopyFrom(\n", + " numpy_helper.from_array(scale.reshape(-1).float().numpy(), base + '_scale'))\n", + " int_ada = torch.round(ada_w[mod].double() / scale.double()).clamp(qmin, qmax).to(torch.int8).numpy()\n", + " t = inits[nm]; t.ClearField('int32_data'); t.raw_data = _pack_int4(int_ada.reshape(-1))\n", + " # drop the value_info that duplicates graph IO (an artifact ORT can leave behind)\n", + " io = {o.name for o in m.graph.output} | {i.name for i in m.graph.input}\n", + " keep = [vi for vi in m.graph.value_info if vi.name not in io]\n", + " del m.graph.value_info[:]; m.graph.value_info.extend(keep)\n", + " onnx.save(m, path)\n", + " return path\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fe6706a3", + "metadata": {}, + "outputs": [], + "source": [ + "show_doc(adaround_quantize)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "54785994", + "metadata": {}, + "outputs": [], + "source": [ + "show_doc(adaround_to_onnx)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "98cc9c0b", + "metadata": {}, + "outputs": [], + "source": [ + "show_doc(rtn_quant)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "61397d6d", + "metadata": {}, + "outputs": [], + "source": [ + "show_doc(adaround_layer)" + ] + }, + { + "cell_type": "markdown", + "id": "026a6fb1", + "metadata": {}, + "source": [ + "---\n", + "\n", + "## Usage\n", + "\n", + "```python\n", + "from fasterai.quantize.adaround import adaround_quantize\n", + "\n", + "# Weight-only INT4 with learned rounding, per-output-channel scales\n", + "qmodel = adaround_quantize(\n", + " model,\n", + " calibration_dl=dls.valid, # a fastai DataLoaders, a torch DataLoader, or a list of batches\n", + " w_bits=4,\n", + " granularity='channel',\n", + " iters=2000,\n", + ")\n", + "\n", + "# W4A8 (also fake-quantize activations to INT8)\n", + "qmodel = adaround_quantize(model, dls.valid, w_bits=4, act_bits=8)\n", + "```\n", + "\n", + "`calibration_dl` accepts three shapes, normalised internally:\n", + "\n", + "- a **fastai `DataLoaders`** (its `.valid` loader is used),\n", + "- a **torch `DataLoader`** yielding `(xb, yb)` batches, or\n", + "- a plain **list of batches** — each a `(xb, yb)` tuple or a bare input tensor." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c53a8251", + "metadata": {}, + "outputs": [], + "source": "#| hide\nfrom fastcore.test import *\ntorch.manual_seed(0)\n\ndef _tiny_net():\n \"Small 2-conv + 1-linear net used across the fast tests.\"\n return nn.Sequential(\n nn.Conv2d(3, 8, 3, padding=1), nn.ReLU(),\n nn.Conv2d(8, 16, 3, padding=1), nn.AdaptiveAvgPool2d(1),\n nn.Flatten(), nn.Linear(16, 4),\n ).eval()\n\n_calib = [torch.randn(4, 3, 16, 16) for _ in range(3)] # list-of-batches calibration\n\n# 1) SANITY GATE (correctness): at init (iters=0) AdaRound weights == RTN weights,\n# because V is initialised so h(V) == frac(W/s). If this fails, the port is wrong.\n_net = _tiny_net()\n_q0 = adaround_quantize(copy.deepcopy(_net), _calib, w_bits=4, iters=0,\n fold_batchnorm=False, device='cpu')\n_orig = [m for m in _net.modules() if isinstance(m, (nn.Conv2d, nn.Linear))]\n_q0l = [m for m in _q0.modules() if isinstance(m, (nn.Conv2d, nn.Linear))]\ntest_eq(len(_q0l), 3) # 2 conv + 1 linear\nfor o, q in zip(_orig, _q0l):\n test_close(q.weight, rtn_quant(o.weight.detach(), 4, 'channel'), eps=1e-5)" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "46e8bf57", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# 2) RECONSTRUCTION IMPROVES: after a short AdaRound run the per-layer reconstruction\n", + "# error ‖W_q x − W x‖² is <= the RTN reconstruction error.\n", + "torch.manual_seed(0)\n", + "_conv = nn.Conv2d(4, 8, 3, padding=1).eval()\n", + "_x = torch.randn(32, 4, 8, 8)\n", + "_W = _conv.weight.detach()\n", + "with torch.no_grad():\n", + " _out_fp = F.conv2d(_x, _W, _conv.bias, _conv.stride, _conv.padding)\n", + " _Wq_rtn = rtn_quant(_W, 4, 'channel')\n", + " _rec_rtn = (F.conv2d(_x, _Wq_rtn, _conv.bias, _conv.stride, _conv.padding) - _out_fp).pow(2).sum().item()\n", + "_Wq_ada, _info = adaround_layer(_conv, _x, 4, iters=400, batch_size=16, device='cpu')\n", + "with torch.no_grad():\n", + " _rec_ada = (F.conv2d(_x, _Wq_ada, _conv.bias, _conv.stride, _conv.padding) - _out_fp).pow(2).sum().item()\n", + "assert _rec_ada <= _rec_rtn + 1e-6, f\"AdaRound rec {_rec_ada:.4g} should be <= RTN rec {_rec_rtn:.4g}\"\n", + "# init V so h(V)=frac -> diagnostics dict is well-formed\n", + "test_eq(set(_info), {'not_binary_frac', 'w_bits', 'granularity'})" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5759c02d", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# 4) FAILURE MODES / USER-SIDE INPUTS\n", + "_net = _tiny_net()\n", + "# Missing calibration data -> clear ValueError mentioning calibration\n", + "with ExceptionExpected(ValueError, regex='calibration'):\n", + " adaround_quantize(_net, None, w_bits=4)\n", + "# Unsupported granularity -> clear error\n", + "with ExceptionExpected(ValueError, regex='granularity'):\n", + " adaround_quantize(_net, _calib, granularity='block')\n", + "# Unsupported w_bits -> clear error\n", + "with ExceptionExpected(ValueError, regex='w_bits'):\n", + " adaround_quantize(_net, _calib, w_bits=1)\n", + "# Unsupported act_bits -> clear error\n", + "with ExceptionExpected(ValueError, regex='act_bits'):\n", + " adaround_quantize(_net, _calib, w_bits=4, act_bits=1)\n", + "\n", + "# A plain LIST OF BATCHES works end-to-end (cross-function consistency, not only a DataLoaders)\n", + "_qm = adaround_quantize(copy.deepcopy(_net), _calib, w_bits=4, iters=5,\n", + " fold_batchnorm=False, device='cpu')\n", + "_o = _qm(torch.randn(2, 3, 16, 16))\n", + "test_eq(_o.shape, (2, 4))\n", + "assert torch.isfinite(_o).all()\n", + "# Weight-only leaves the model callable; act_bits installs activation fake-quant hooks\n", + "_qm_a = adaround_quantize(copy.deepcopy(_net), _calib, w_bits=4, act_bits=8, iters=5,\n", + " fold_batchnorm=False, device='cpu')\n", + "assert torch.isfinite(_qm_a(torch.randn(2, 3, 16, 16))).all()\n", + "# 'tensor' granularity is a valid alternative to 'channel'\n", + "_qm_t = adaround_quantize(copy.deepcopy(_net), _calib, w_bits=4, granularity='tensor',\n", + " iters=5, fold_batchnorm=False, device='cpu')\n", + "assert torch.isfinite(_qm_t(torch.randn(2, 3, 16, 16))).all()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "be3600ea", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# adaround_quantize stashes _adaround_qparams for the faithful ONNX export\n", + "torch.manual_seed(0)\n", + "_qm = adaround_quantize(copy.deepcopy(_tiny_net()), _calib, w_bits=4, iters=5,\n", + " fold_batchnorm=False, device='cpu')\n", + "assert hasattr(_qm, '_adaround_qparams')\n", + "_named = {n: m for n, m in _qm.named_modules() if isinstance(m, (nn.Conv2d, nn.Linear))}\n", + "test_eq(set(_qm._adaround_qparams), set(_named)) # one entry per quantized layer\n", + "for _n, (_s, _qmin, _qmax) in _qm._adaround_qparams.items():\n", + " test_eq((_qmin, _qmax), (-8, 7)) # signed INT4 grid\n", + " test_eq(_s.reshape(-1).numel(), _named[_n].weight.shape[0]) # scale is per-output-channel\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a0987625", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# int4 pack/unpack round-trips across the full signed nibble range [-8, 7]\n", + "_ints = np.array([-8, -7, -1, 0, 1, 6, 7, -8, 7, 0, -3, 4, 5], dtype=np.int8)\n", + "test_eq(_unpack_int4(_pack_int4(_ints), _ints.size).tolist(), _ints.tolist())\n", + "_r = np.random.randint(-8, 8, size=257).astype(np.int8) # odd length exercises the pad branch\n", + "test_eq(_unpack_int4(_pack_int4(_r), _r.size).tolist(), _r.tolist())\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "34ff577a", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# adaround_to_onnx on a plain model (no _adaround_qparams) -> clear ValueError, checked before any onnx import\n", + "with ExceptionExpected(ValueError, regex='adaround_quantize'):\n", + " adaround_to_onnx(_tiny_net(), torch.randn(1, 3, 16, 16), '/tmp/_should_not_write.onnx')\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "608d0f31", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "#| slow\n", + "# 3) BEHAVIORAL RECOVERY (real model, real data): INT4+AdaRound top-1 > INT4-RTN top-1.\n", + "# Small Imagenette slice, ResNet-18 pretrained on ImageNet (1000-way argmax).\n", + "import os, torchvision\n", + "from torchvision import transforms\n", + "\n", + "_DATA = '/home/nathan/.fastai/data/imagenette2-320'\n", + "# Imagenette WNID (alphabetical = ImageFolder idx 0..9) -> ImageNet class index\n", + "_WNID2IN = {'n01440764': 0, 'n02102040': 217, 'n02979186': 482, 'n03000684': 491,\n", + " 'n03028079': 497, 'n03394916': 566, 'n03417042': 569, 'n03425413': 571,\n", + " 'n03445777': 574, 'n03888257': 701}\n", + "_dev = 'cuda' if torch.cuda.is_available() else 'cpu'\n", + "\n", + "_tfm = transforms.Compose([\n", + " transforms.Resize(232, interpolation=transforms.InterpolationMode.BILINEAR),\n", + " transforms.CenterCrop(224), transforms.ToTensor(),\n", + " transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])\n", + "_ds = torchvision.datasets.ImageFolder(os.path.join(_DATA, 'val'), transform=_tfm)\n", + "_i2w = {v: k for k, v in _ds.class_to_idx.items()}\n", + "_remap = {i: _WNID2IN[_i2w[i]] for i in range(len(_i2w))}\n", + "_dl = torch.utils.data.DataLoader(_ds, batch_size=64, shuffle=True, num_workers=8)\n", + "_xs, _ys = [], []\n", + "for _xb, _yb in _dl:\n", + " _xs.append(_xb); _ys.append(torch.tensor([_remap[int(t)] for t in _yb]))\n", + " if sum(x.shape[0] for x in _xs) >= 384: break\n", + "_X = torch.cat(_xs)[:384]; _Y = torch.cat(_ys)[:384]\n", + "_calib_rn = [_X[i:i+32] for i in range(0, 128, 32)] # 128 calibration images\n", + "_Xe, _Ye = _X[128:], _Y[128:] # 256 eval images\n", + "\n", + "_rn = torchvision.models.resnet18(weights='IMAGENET1K_V1').eval().to(_dev)\n", + "fold_bn(_rn)\n", + "\n", + "@torch.no_grad()\n", + "def _top1(m):\n", + " m.eval(); ok = 0\n", + " for i in range(0, _Xe.shape[0], 64):\n", + " ok += (m(_Xe[i:i+64].to(_dev)).argmax(1).cpu() == _Ye[i:i+64]).sum().item()\n", + " return 100.0 * ok / _Xe.shape[0]\n", + "\n", + "_acc_fp = _top1(_rn)\n", + "\n", + "_rtn = copy.deepcopy(_rn) # INT4 RTN baseline (weight-only)\n", + "for _l in _iter_target_layers(_rtn):\n", + " _l.weight.data = rtn_quant(_l.weight.detach().float(), 4, 'channel').to(_l.weight.dtype)\n", + "_acc_rtn = _top1(_rtn)\n", + "\n", + "_ada = adaround_quantize(_rn, _calib_rn, w_bits=4, iters=1000, device=_dev) # INT4 + AdaRound\n", + "_acc_ada = _top1(_ada)\n", + "\n", + "print(f\"ResNet-18 Imagenette top-1 | FP={_acc_fp:.1f} INT4-RTN={_acc_rtn:.1f} INT4-AdaRound={_acc_ada:.1f}\")\n", + "assert _acc_ada > _acc_rtn, f\"AdaRound {_acc_ada:.1f} should recover over RTN {_acc_rtn:.1f}\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "75512f27", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "#| slow\n", + "# FAITHFULNESS (the point): the deployed INT4 ONNX matches the fake-quant top-1, i.e. AdaRound's\n", + "# learned rounding survives export. Reuses `_ada`/`_Xe`/`_Ye`/`_acc_ada` from the cell above.\n", + "try:\n", + " import onnx, onnxruntime as ort\n", + " from onnx import numpy_helper\n", + " _have_onnx = True\n", + "except ImportError:\n", + " _have_onnx = False\n", + "\n", + "if _have_onnx:\n", + " _onnx_int4 = '/tmp/fasterai_adaround_int4.onnx'\n", + " adaround_to_onnx(_ada, _Xe[:1], _onnx_int4, input_name='image', calibration_dl=_calib_rn)\n", + "\n", + " # (a) deployed ONNX top-1 via onnxruntime on the SAME eval set as the fake-quant top-1\n", + " _sess = ort.InferenceSession(_onnx_int4, providers=['CPUExecutionProvider'])\n", + " _iname = _sess.get_inputs()[0].name\n", + " _ok = 0\n", + " for _i in range(0, _Xe.shape[0], 64):\n", + " _out = _sess.run(None, {_iname: _Xe[_i:_i+64].cpu().numpy().astype(np.float32)})[0]\n", + " _ok += int((_out.argmax(1) == _Ye[_i:_i+64].numpy()).sum())\n", + " _acc_onnx = 100.0 * _ok / _Xe.shape[0]\n", + "\n", + " # (b) FP32 ONNX export of the same graph -> the INT4 file must be much smaller (genuine int4 storage)\n", + " _onnx_fp32 = '/tmp/fasterai_adaround_fp32.onnx'\n", + " torch.onnx.export(copy.deepcopy(_ada).eval().cpu(), _Xe[:1].cpu(), _onnx_fp32,\n", + " input_names=['image'], output_names=['output'],\n", + " dynamic_axes={'image': {0: 'b'}, 'output': {0: 'b'}}, opset_version=18, dynamo=False)\n", + " _sz_int4, _sz_fp32 = os.path.getsize(_onnx_int4), os.path.getsize(_onnx_fp32)\n", + "\n", + " # (c) the ONNX initializers carry AdaRound's EXACT weights (dequant == baked fake-quant weights)\n", + " _mo = onnx.load(_onnx_int4); _inits = {i.name: i for i in _mo.graph.initializer}\n", + " _adaw = {n: m.weight.detach().float().cpu() for n, m in _ada.named_modules()\n", + " if isinstance(m, (nn.Conv2d, nn.Linear))}\n", + " _maxerr = 0.0\n", + " for _nm in _inits:\n", + " if not _nm.endswith('.weight_quantized'): continue\n", + " _b = _nm[:-len('_quantized')]; _mod = _b[:-len('.weight')]\n", + " if _mod not in _adaw: continue\n", + " _q = numpy_helper.to_array(_inits[_nm]).astype(np.float64)\n", + " _sc = numpy_helper.to_array(_inits[_b + '_scale']).astype(np.float64)\n", + " _shp = _q.shape\n", + " _w = (_q.reshape(_shp[0], -1) * _sc.reshape(-1, 1)).reshape(_shp)\n", + " _maxerr = max(_maxerr, float(np.abs(_w - _adaw[_mod].double().numpy()).max()))\n", + "\n", + " print(f\"AdaRound INT4 faithful export | fake-quant={_acc_ada:.2f} deployed-ONNX={_acc_onnx:.2f} \"\n", + " f\"delta={_acc_ada-_acc_onnx:+.2f} pt\")\n", + " print(f\" onnx size: int4={_sz_int4/1e6:.2f} MB fp32={_sz_fp32/1e6:.2f} MB \"\n", + " f\"({_sz_fp32/_sz_int4:.1f}x) | max|ONNX dequant - AdaRound baked|={_maxerr:.2e}\")\n", + " assert abs(_acc_ada - _acc_onnx) <= 1.5, \\\n", + " f\"deployed ONNX top-1 {_acc_onnx:.2f} should be within 1.5 pt of fake-quant {_acc_ada:.2f}\"\n", + " assert _sz_int4 < 0.7 * _sz_fp32, \\\n", + " f\"INT4 ONNX ({_sz_int4} B) should be much smaller than FP32 ONNX ({_sz_fp32} B)\"\n", + " assert _maxerr < 1e-5, f\"faithful ONNX weights drifted from AdaRound baked weights ({_maxerr})\"\n" + ] + }, + { + "cell_type": "markdown", + "id": "fd6f8075", + "metadata": {}, + "source": [ + "---\n", + "\n", + "## See Also\n", + "\n", + "- [Dimensional Quantizer design](DIMENSIONAL_DESIGN.md) — the north-star API. AdaRound is\n", + " `criteria='adaround'` in that grammar (composable with `granularity`, `precision`, `context`).\n", + "- [Quantizer](quantizer.html) — backend-oriented INT8 PTQ/QAT (torch.ao / torchao).\n", + "- [QuantizeCallback](quantize_callback.html) — apply quantization during fastai training.\n", + "- [BN Folding](../misc/bn_folding.html) — the Conv→BN fusion AdaRound runs first." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "python3", + "language": "python", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}