File size: 4,854 Bytes
e9c8366 | 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 | """BitNet 1.58 — ternary quantization primitives (weights {-1, 0, +1}).
Pure primitive module: tensor-level quantization functions only. No
nn.Module wrapper classes — the unified QuantizedModule (base.py) is the
single place for quantized layer wrappers.
BitNet b1.58: each weight is quantized to one of three values: {-1, 0, +1}.
A single per-tensor (or per-channel) scale factor reconstructs the weight:
W ~= ternary_weight * scale
Quantization (data-free, from weights only):
scale = mean(abs(W)) # per-tensor (BitNet original)
scale = mean(abs(W), dim=1) # per-channel (better for uneven layers)
ternary = round(W / scale) # clamped to {-1, 0, +1}
W_approx = ternary * scale
STE (Straight-Through Estimator):
Forward: ternary = round(W_latent / scale) # non-differentiable
Backward: grad flows to W_latent as identity # STE bypass
Inference: dequant = ternary.to(float) * scale, then normal matmul/conv.
Storage: ternary weights packed 2 values per int8 byte (2 bits each + padding),
scale is float32 per-tensor or per-channel.
"""
import torch
from agiws_neural_quant.training.ste import STEQuantize
# ---------------------------------------------------------------------------
# Ternary quantization core (data-free, from weight tensor)
# ---------------------------------------------------------------------------
def ternarize_tensor(
w: torch.Tensor,
scale_mode: str = "per-channel",
) -> tuple[torch.Tensor, torch.Tensor]:
"""Quantize a weight tensor to ternary {-1, 0, +1} + scale.
Args:
w: float weight tensor. For Linear: [out, in]. For Conv3d: [out, in, kT, kH, kW].
scale_mode: 'per-tensor' (single scale) or 'per-channel' (one scale per output channel).
Returns:
(ternary, scale) where:
ternary: int8 tensor with values {-1, 0, +1}, same shape as w
scale: float32 tensor — scalar (per-tensor) or [out] (per-channel)
"""
w = w.detach().float()
if scale_mode == "per-tensor":
scale = w.abs().mean().clamp(min=1e-8)
ternary = torch.clamp(torch.round(w / scale), min=-1, max=1).to(torch.int8)
return ternary, scale.reshape(1)
elif scale_mode == "per-channel":
# Per output channel (dim 0 for both Linear [out,in] and Conv [out,in,kT,kH,kW])
reduce_dims = tuple(d for d in range(1, w.dim()))
scale = w.abs().mean(dim=reduce_dims).clamp(min=1e-8) # [out]
# Reshape scale for broadcasting
reshape = [1] * w.dim()
reshape[0] = w.shape[0]
ternary = torch.clamp(torch.round(w / scale.reshape(reshape)), min=-1, max=1).to(torch.int8)
return ternary, scale
else:
raise ValueError(f"scale_mode must be 'per-tensor' or 'per-channel', got {scale_mode!r}")
def ternary_dequantize(ternary: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
"""Reconstruct float weight from ternary + scale.
Args:
ternary: int8 tensor with values {-1, 0, +1}
scale: scalar tensor (per-tensor) or [out] tensor (per-channel)
Returns:
float32 dequantized weight, same shape as ternary
"""
t = ternary.to(torch.float32)
if scale.dim() == 0 or scale.numel() == 1:
return t * scale.to(torch.float32)
else:
# Per-channel: reshape scale for broadcasting over [out, ...]
reshape = [1] * ternary.dim()
reshape[0] = ternary.shape[0]
return t * scale.to(torch.float32).reshape(reshape)
def fake_ternarize(w: torch.Tensor, scale_mode: str = "per-channel") -> torch.Tensor:
"""STE fake ternarization for training (QAT / block-wise distillation).
Forward: ternarize {-1,0,+1} and dequantize back (simulates quantization error).
Backward: gradient flows to w as identity (STE).
Args:
w: latent float weight (nn.Parameter, requires_grad=True)
scale_mode: 'per-tensor' or 'per-channel'
Returns:
Fake-quantized weight (float, same shape as w), gradient-connected to w via STE.
"""
w = w.float()
if scale_mode == "per-tensor":
scale = w.abs().mean().clamp(min=1e-8).detach()
# STE: symmetric=True for ternary {-1,0,+1} (clamp to [-1, +1], not [-1, 0])
return STEQuantize.apply(w, scale.unsqueeze(0), 1, True)
elif scale_mode == "per-channel":
reduce_dims = tuple(d for d in range(1, w.dim()))
scale = w.abs().mean(dim=reduce_dims).clamp(min=1e-8).detach()
reshape = [1] * w.dim()
reshape[0] = w.shape[0]
# STE: symmetric=True for ternary {-1,0,+1}
return STEQuantize.apply(w, scale.reshape(reshape), 1, True)
else:
raise ValueError(f"scale_mode must be 'per-tensor' or 'per-channel', got {scale_mode!r}")
__all__ = [
"ternarize_tensor",
"ternary_dequantize",
"fake_ternarize",
] |