File size: 6,002 Bytes
e9c8366 82adbbb | 1 2 3 4 5 6 7 8 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | """STE (Straight-Through Estimator) for quantization-aware training.
Forward: quantize (round) — non-differentiable
Backward: identity — gradient passes through as if round() was identity()
This allows training with fake quantization: the model "sees" quantized weights
in forward, but gradients flow to the latent fp32/fp16 weights in backward.
"""
import torch
class STEQuantize(torch.autograd.Function):
"""Straight-Through Estimator for symmetric quantization.
Forward: weight_int = clamp(round(weight / scale), -n_levels, n_levels)
For asymmetric (int8/int4): clamp(min=-n_levels, max=n_levels-1)
For symmetric (ternary {-1,0,+1}): clamp(min=-n_levels, max=n_levels)
Backward: d(weight_int)/d(weight) = 1 (identity, STE)
The clamp range depends on whether the quantization is symmetric (includes
both -n_levels and +n_levels, like ternary {-1,0,+1}) or asymmetric (int8
[-127,127], int4 [-7,7]). Use symmetric=True for ternary, False for int8/int4.
"""
@staticmethod
def forward(ctx, weight: torch.Tensor, scale: torch.Tensor, n_levels: int,
symmetric: bool = False) -> torch.Tensor:
"""Quantize weight and dequantize back (fake quantization).
Args:
weight: float tensor (latent weight)
scale: float tensor (per-channel or per-group scale)
n_levels: max quantized value (127 for int8, 7 for int4, 1 for ternary)
symmetric: if True, clamp to [-n_levels, n_levels] (ternary {-1,0,+1}).
if False, clamp to [-n_levels, n_levels-1] (int8 [-127,127]).
Returns:
Dequantized weight (fake quantized) — same shape as input.
"""
# Quantize
if symmetric:
q = torch.clamp(
torch.round(weight / scale),
min=-n_levels,
max=n_levels,
)
else:
q = torch.clamp(
torch.round(weight / scale),
min=-n_levels,
max=n_levels - 1,
)
# Dequantize back (fake quantization)
return q * scale
@staticmethod
def backward(ctx, grad_output: torch.Tensor):
# STE: gradient passes through as identity
# grad for weight = grad_output, grad for scale = None, grad for n_levels = None
return grad_output, None, None, None
def fake_quantize(weight: torch.Tensor, scale: torch.Tensor, n_levels: int,
symmetric: bool = False) -> torch.Tensor:
"""Apply fake quantization with STE for backward.
In forward: simulates quantization error.
In backward: gradient flows to weight (latent) as identity.
Args:
weight: latent weight (fp32 or fp16) — nn.Parameter
scale: quantization scale (per-channel or per-group)
n_levels: 127 for int8, 7 for int4, 1 for ternary
symmetric: True for ternary {-1,0,+1} (clamp to [-n, n]).
False for int8/int4 (clamp to [-n, n-1]).
Returns:
Fake-quantized weight (same dtype as input).
"""
return STEQuantize.apply(weight, scale, n_levels, symmetric)
class STECodebook(torch.autograd.Function):
"""STE for codebook quantization (learnable codebooks).
Forward: indices = argmin(|x - codebook|); deq = codebook[indices]
Backward: gradient flows to x (latent weight) as identity (STE),
AND to codebook (learnable) as identity for the selected entries.
This allows training both the latent weight AND the codebook via gradient
descent. The argmin (non-differentiable) is bypassed by STE.
"""
@staticmethod
def forward(ctx, weight: torch.Tensor, codebook: torch.Tensor,
indices: torch.Tensor) -> torch.Tensor:
"""Codebook lookup with precomputed indices.
Args:
weight: latent weight (for STE identity backward).
codebook: codebook tensor [K] or [K, dim].
indices: precomputed argmin indices (same shape as weight, or [..., dim]).
Returns:
Dequantized weight: codebook[indices], same shape as weight.
"""
ctx.save_for_backward(weight, codebook, indices)
if codebook.dim() == 1:
return codebook.to(weight.device)[indices.long()].to(torch.float32)
else:
# VQ: codebook [K, dim], indices [...]
return codebook.to(weight.device)[indices.long()].to(torch.float32)
@staticmethod
def backward(ctx, grad_output: torch.Tensor):
weight, codebook, indices = ctx.saved_tensors
# STE: grad to weight = grad_output (identity).
grad_weight = grad_output
# Grad to codebook: scatter-add gradient to selected entries.
grad_codebook = torch.zeros_like(codebook)
if codebook.dim() == 1:
grad_codebook.scatter_add_(0, indices.long().flatten(),
grad_output.flatten())
else:
# VQ: codebook [K, dim], indices [...]
grad_codebook.scatter_add_(0, indices.long().reshape(-1, 1).expand(-1, codebook.shape[1]),
grad_output.reshape(-1, codebook.shape[1]))
return grad_weight, grad_codebook, None
def fake_codebook_quantize(weight: torch.Tensor, codebook: torch.Tensor,
indices: torch.Tensor) -> torch.Tensor:
"""Apply codebook fake quantization with STE for backward.
Forward: codebook[indices] (precomputed argmin).
Backward: gradient to weight (identity STE) + gradient to codebook (scatter-add).
Args:
weight: latent weight (nn.Parameter, for STE identity).
codebook: codebook tensor [K] (for learnable codebook gradient).
indices: precomputed argmin indices.
Returns:
Fake-quantized weight (codebook[indices]), differentiable to both
weight and codebook.
"""
return STECodebook.apply(weight, codebook, indices) |