Skip to content

Quantization

quantization

Symmetric per-tensor quantization for training and compression.

uniform_quantize 

uniform_quantize(model: Module, quant_max: int) -> tuple[Module, dict[str, dict[str, Tensor]]]

Fake-quantize model parameters in place and return their integer codes.

Each parameter w becomes clamp(round(s * w), -quant_max, quant_max) / s, with s = quant_max / max(abs(w)), or s = 1 for an all-zero tensor. Complex parameters are unsupported.

Parameters:

  • model (Module) –

    Module to quantize in place; pass a copy to keep the original weights.

  • quant_max (int) –

    Largest integer level, less than 2**63; codes use the smallest signed integer dtype that holds it.

Returns:

  • tuple[Module, dict[str, dict[str, Tensor]]] –

    The same model and a dictionary mapping parameter names to {"quant": codes, "scale": s}, both detached CPU tensors.

Examples:

>>> model, codes = uniform_quantize(nn.Linear(2, 3), quant_max=127)
>>> codes["weight"]["quant"].dtype
torch.int8
>>> restored = uniform_dequantize(nn.Linear(2, 3), codes)
>>> restored.weight.equal(model.weight)
True
Source code in inro/quantization.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
@torch.no_grad()
def uniform_quantize(
    model: nn.Module, quant_max: int
) -> tuple[nn.Module, dict[str, dict[str, Tensor]]]:
    """Fake-quantize model parameters in place and return their integer codes.

    Each parameter ``w`` becomes ``clamp(round(s * w), -quant_max, quant_max) / s``,
    with ``s = quant_max / max(abs(w))``, or ``s = 1`` for an all-zero tensor. Complex
    parameters are unsupported.

    Args:
        model: Module to quantize in place; pass a copy to keep the original weights.
        quant_max: Largest integer level, less than ``2**63``; codes use the smallest
            signed integer dtype that holds it.

    Returns:
        The same model and a dictionary mapping parameter names to
            ``{"quant": codes, "scale": s}``, both detached CPU tensors.

    Examples:
        >>> model, codes = uniform_quantize(nn.Linear(2, 3), quant_max=127)
        >>> codes["weight"]["quant"].dtype
        torch.int8
        >>> restored = uniform_dequantize(nn.Linear(2, 3), codes)
        >>> restored.weight.equal(model.weight)
        True
    """
    for dtype in (torch.int8, torch.int16, torch.int32, torch.int64):
        if quant_max <= torch.iinfo(dtype).max:
            break
    else:
        raise ValueError(f"quant_max must be less than 2**63, got {quant_max}")

    model_dict: dict[str, dict[str, Tensor]] = {}
    for name, param in model.named_parameters():
        peak = param.abs().max()
        scale = quant_max / peak if peak > 0 else peak.new_ones(())
        codes = torch.clamp(torch.round(scale * param), -quant_max, quant_max)
        param.copy_(codes / scale)

        model_dict[name] = {
            "quant": codes.cpu().to(dtype),
            "scale": scale.cpu(),
        }

    return model, model_dict

uniform_dequantize 

uniform_dequantize(model: Module, model_dict: dict[str, dict[str, Tensor]]) -> Module

Load parameters from uniform_quantize codes in place.

Parameters:

  • model (Module) –

    Module to overwrite in place.

  • model_dict (dict[str, dict[str, Tensor]]) –

    Codes and scales of every parameter of model, by name.

Returns:

  • Module –

    The same model, each parameter being quant.float() / scale.float().

Source code in inro/quantization.py 
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
@torch.no_grad()
def uniform_dequantize(
    model: nn.Module, model_dict: dict[str, dict[str, Tensor]]
) -> nn.Module:
    """Load parameters from :func:`uniform_quantize` codes in place.

    Args:
        model: Module to overwrite in place.
        model_dict: Codes and scales of every parameter of ``model``, by name.

    Returns:
        The same model, each parameter being ``quant.float() / scale.float()``.
    """
    for name, param in model.named_parameters():
        codes = model_dict[name]["quant"].float()
        scale = model_dict[name]["scale"].float()
        param.copy_(codes / scale)

    return model