Hugging Face's logo

RegalInvader
/
Marco-Nano-Instruct-GGUF

Text Generation
GGUF
llama.cpp
pocketpal
on-device
Mixture of Experts
multilingual
quantization
distillation
1-bit
2-bit
3-bit
conversational
Model card Files
xet
 Community
Marco-Nano-Instruct-GGUF / code /src /lowbit /quant /ternary.py
RegalInvader's picture
RegalInvader
Release: Marco-Nano-Instruct 1-, 2- and 3-bit GGUF for phones
e5235fb verified
Raw History Blame Contribute Delete
28.3 kB
"""Ternary weight quantisers (spec §4): Q-A absmean, Q-B TWN, Q-C learnable (LSQ-style) scale.
Weights have shape [..., in]; groups of `group_size` run along the input (last) dimension, matching
llama.cpp's block layout. Stacked expert tensors [E, out, in] work unchanged.
The forward weight is s * q with q in {-1, 0, +1}. It is built by a custom autograd function rather
than `w + (w_hat - w).detach()`: in fp32 that expression is not always bit-equal to w_hat, and at
7.7B weights the difference flips ~1e5 bf16 values. Gradients to the latent are the identity (plain
STE) or zero where |w|/s > clip_threshold (clipped STE).
Scales are canonicalised to bf16 precision (DECISIONS.md P-03): the trainer's matmul sees
bf16(s * q) = bf16(s) * q, and bf16(s) is exactly representable in the fp16 export format, so the
fake-quant, packed and GGUF students share one set of weights by construction.
D-51 adds three N-bit expert formats, each trained on exactly the grid of a llama.cpp GGUF type so the export is
lossless (section "N-bit grids" below): int2 = Q2_0, int4 = Q4_0, int3 = Q3_K (its down projection Q4_0).
D-57 adds llama.cpp's non-linear 4-bit grid: iq4nl = IQ4_NL, iq4xs = IQ4_XS (its down projection IQ4_NL), and
int3nl = Q3_K with its down projection IQ4_NL.
"""
from __future__ import annotations
import math
from dataclasses import dataclass
import torch
METHODS = ("absmean", "twn", "lsq", "int2", "int3", "int4", # int*: D-51, see "N-bit grids"
"int3nl", "iq4nl", "iq4xs", "int2nl", # D-57
"codebook") # D-58: llama.cpp's codebook formats, frozen codes (`quant.codebook`)
STES = ("plain", "clipped")
TWN_DELTA = 0.7
@dataclass(frozen=True)
class QuantSpec:
method: str = "absmean" # absmean (Q-A) | twn (Q-B) | lsq (Q-C)
group_size: int = 128
ste: str = "plain" # plain | clipped
clip_threshold: float = 1.5
eps: float = 1e-8
def __post_init__(self) -> None:
if self.method not in METHODS:
raise ValueError(f"method {self.method!r} not in {METHODS}")
if self.ste not in STES:
raise ValueError(f"ste {self.ste!r} not in {STES}")
if self.group_size <= 0:
raise ValueError("group_size must be positive")
if self.method in GRIDS:
if self.group_size != GRIDS[self.method].group:
raise ValueError(f"{self.method} is GGUF {GRIDS[self.method].gguf}: group size must be "
f"{GRIDS[self.method].group}, not {self.group_size}")
if self.ste != "plain":
raise ValueError(f"{self.method}: only the plain STE is built")
def _grouped(w: torch.Tensor, group: int) -> torch.Tensor:
if w.shape[-1] % group:
raise ValueError(f"input dim {w.shape[-1]} not divisible by group size {group}")
return w.unflatten(-1, (-1, group))
def canonical_scale(s: torch.Tensor) -> torch.Tensor:
"""Round scales to bf16 precision (nearest, ties to even), returned as fp32 (exactly representable in fp16).
Done on the fp32 bits: torch.compile elides a `.to(bfloat16).to(float32)` round trip inside a fused kernel,
which left every scale unrounded (T2.3). Scales are positive and finite, so the bit trick is exact."""
b = s.to(torch.float32).view(torch.int32)
b = (b + 0x7FFF + ((b >> 16) & 1)) & -0x10000
return b.view(torch.float32)
def init_scale(w: torch.Tensor, spec: QuantSpec, importance: torch.Tensor | None = None) -> torch.Tensor:
"""Per-group scale [..., in/g] in fp32, canonicalised. Q-C initialises from Q-A. The N-bit methods return
their (signed) per-group scale, `nbit_scale`.
Group sums accumulate in fp64, as the reference does, so the scale does not depend on reduction order: under
torch.compile fp32 sums moved scales across bf16 rounding boundaries (0.085% of Marco's codes, T2.3)."""
if spec.method in GRIDS:
return nbit_scale(w, spec, importance)
wg = _grouped(w.detach().double(), spec.group_size).abs()
if spec.method == "twn":
delta = TWN_DELTA * wg.mean(-1, keepdim=True)
above = wg > delta
s = (wg * above).sum(-1) / above.sum(-1).clamp_min(1)
else:
s = wg.mean(-1)
return canonical_scale(s.float().clamp_min(spec.eps))
def ternary_codes(w: torch.Tensor, s: torch.Tensor, spec: QuantSpec) -> torch.Tensor:
"""int8 codes in {-1, 0, 1}, shape of w.
Round-half-even of w/s clamped to [-1, 1] is sign(w)·[|w| > s/2]. With s bf16-exact, s/2 is exact, so the
comparison is exact where an fp32 division is not (compiled kernels may divide differently).
The N-bit methods return their grid codes (`nbit_codes`)."""
if spec.method in GRIDS:
return nbit_codes(w, s, spec)
wg = _grouped(w.detach().float(), spec.group_size)
if spec.method == "twn":
a = wg.abs()
delta = TWN_DELTA * a.double().mean(-1, keepdim=True) # fp64, as in init_scale
q = torch.sign(wg) * (a > delta)
else:
q = torch.sign(wg) * (wg.abs() > 0.5 * s.unsqueeze(-1))
return q.flatten(-2).to(torch.int8)
class _TernarySTE(torch.autograd.Function):
"""Forward s*q exactly; backward identity (or clipped) to w, LSQ gradient to a learnable s."""
@staticmethod
def forward(ctx, w, s, spec: QuantSpec):
codes = ternary_codes(w, s, spec)
s_c = canonical_scale(s.detach().float())
w_hat = (_grouped(codes.to(torch.float32), spec.group_size) * s_c.unsqueeze(-1)).flatten(-2)
ctx.spec = spec
ctx.save_for_backward(w, s_c, codes) # int8: a quarter of fp32 codes (Marco: 0.27 vs 1.1 GB per layer)
return w_hat.to(w.dtype)
@staticmethod
def backward(ctx, g):
w, s_c, codes = ctx.saved_tensors
spec: QuantSpec = ctx.spec
gw = g
ratio = None
if spec.ste == "clipped" or (spec.method == "lsq" and ctx.needs_input_grad[1]):
ratio = _grouped(w.float(), spec.group_size) / s_c.unsqueeze(-1)
if spec.ste == "clipped":
gw = g * (ratio.abs() <= spec.clip_threshold).flatten(-2).to(g.dtype)
gs = None
if spec.method == "lsq" and ctx.needs_input_grad[1]:
# LSQ: d(s*q)/ds = q - w/s inside the rounding range, q outside; grad scale 1/sqrt(g*Qp), Qp = 1
qg = _grouped(codes.float(), spec.group_size)
inside = ratio.abs() <= 1.5
dws = torch.where(inside, qg - ratio, qg)
gs = (_grouped(g.float(), spec.group_size) * dws).sum(-1) / math.sqrt(spec.group_size)
return gw, gs, None
def fake_quant(w: torch.Tensor, spec: QuantSpec, scale: torch.Tensor | None = None,
importance: torch.Tensor | None = None) -> torch.Tensor:
"""Ternary (or D-51 N-bit) forward weight with STE. `scale` is required (a trained parameter) for lsq.
D-57 `importance` [..., in] (w without its output dim) weights the N-bit scale search (`_error_weight`)."""
if spec.method in GRIDS:
return _NBitSTE.apply(w, spec, importance)
if spec.method == "lsq":
if scale is None:
raise ValueError("lsq needs a learnable scale (initialise with init_scale)")
s = scale
else:
s = init_scale(w, spec)
return _TernarySTE.apply(w, s, spec)
def scale_for(w: torch.Tensor, spec: QuantSpec, scale: torch.Tensor | None = None) -> torch.Tensor:
"""The canonical scale actually used in the forward pass."""
return canonical_scale(scale.detach().float()) if spec.method == "lsq" else init_scale(w, spec)
# ---------------------------------------------------------------- packing
def pack_2bit(codes: torch.Tensor) -> torch.Tensor:
"""{-1,0,1} int8 [..., n] -> uint8 [..., n/4]; element i in bits 2*(i%4) of byte i//4, stored as q+1."""
if codes.shape[-1] % 4:
raise ValueError("last dim must be divisible by 4")
u = (codes.to(torch.int16) + 1).to(torch.uint8).unflatten(-1, (-1, 4))
return u[..., 0] | (u[..., 1] << 2) | (u[..., 2] << 4) | (u[..., 3] << 6)
def unpack_2bit(packed: torch.Tensor) -> torch.Tensor:
parts = [(packed >> (2 * i)) & 3 for i in range(4)]
return (torch.stack(parts, -1).flatten(-2).to(torch.int16) - 1).to(torch.int8)
_POW3 = (1, 3, 9, 27, 81)
def pack_trits(codes: torch.Tensor) -> torch.Tensor:
"""{-1,0,1} int8 [..., n] -> uint8 [..., ceil(n/5)]; byte = sum_i (q_i+1) * 3^i (5 trits/byte, 1.6 bits)."""
n = codes.shape[-1]
pad = (-n) % 5
t = (codes.to(torch.int16) + 1)
if pad:
t = torch.nn.functional.pad(t, (0, pad), value=1) # pad with code 0
t = t.unflatten(-1, (-1, 5))
pw = torch.tensor(_POW3, dtype=torch.int16, device=codes.device)
return (t * pw).sum(-1).to(torch.uint8)
def unpack_trits(packed: torch.Tensor, n: int) -> torch.Tensor:
v = packed.to(torch.int16)
out = []
for _ in range(5):
out.append(v % 3)
v = v // 3
t = torch.stack(out, -1).flatten(-2)[..., :n]
return (t - 1).to(torch.int8)
def export_packed(w: torch.Tensor, spec: QuantSpec, scale: torch.Tensor | None = None,
importance: torch.Tensor | None = None) -> dict:
"""Packed student tensor: 2-bit codes and fp16 scales (exactly the forward's bf16 scales). N-bit methods:
`export_packed_nbit` (with the forward's `importance`, so the codes are the trained ones)."""
if spec.method in GRIDS:
return export_packed_nbit(w, spec, importance)
s = scale_for(w, spec, scale)
return {"codes": pack_2bit(ternary_codes(w, s, spec)), "scales": s.to(torch.float16),
"group_size": spec.group_size, "shape": tuple(w.shape)}
def dequant_packed(p: dict, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
if p.get("method") in GRIDS:
return dequant_packed_nbit(p, dtype)
q = unpack_2bit(p["codes"]).to(torch.float32)
w = (_grouped(q, p["group_size"]) * p["scales"].to(torch.float32).unsqueeze(-1)).flatten(-2)
return w.to(dtype)
# ---------------------------------------------------------------- N-bit grids (D-51)
#
# Each N-bit format is the value grid of one llama.cpp GGUF type, so the GGUF export is a repack, not a quantisation:
# int2 = Q2_0: blocks of 64, fp16 d, value q·d with q in {-1, 0, 1, 2}.
# int4 = Q4_0: blocks of 32, fp16 d, value q·d with q in -8..7.
# int3 = Q3_K: super-blocks of 256 with an fp16 d; 16 sub-blocks of 16, each with a 6-bit scale m in -32..31;
# value d·m·q with q in -4..3 (ggml's dequantize_row_q3_K). Q3_K needs rows of a multiple of 256, and
# Marco's expert down projection has rows of 384 (moe_intermediate_size), so int3 trains its down
# projection on int4's grid and exports it as Q4_0 (`proj_spec`): int3 is a mixed 3/4-bit model.
#
# Scales are signed (both GGUF types store a signed d), so each group puts its extreme weight E (the largest |w|,
# the positive one on a tie) on the long end of its asymmetric grid: s = E / ext with ext = +2 (int2), -8 (int4),
# -4 (int3's sub-blocks), which for c = 1 is llama.cpp's own Q4_0 rule. A fixed search then tries s = c·E/ext for
# the few c in `cands` and keeps the one with the least squared error over the group (the first on a tie); int3
# instead tries m = round(|t|/d) + k for the offsets k in `cands` (below). On Marco's experts (scratch measurement
# on 6-9 expert tensors of layers 2, 14, 26, 2026-10-22) the search cuts the round-to-nearest error against c = 1
# (k = 0) alone by 19% (int2), 6% (int4) and 12% (int3), to within 1-4% of an oracle search (151 values of c; every
# m in 1..31).
#
# bf16 exactness (P-03): the trainer's matmul sees bf16(s·q), so s·q must be exact in bf16 (8 significant bits)
# for every code. int2's codes have one significant bit, so its scales keep bf16's 8; int4's need up to three
# (5, 7), so its scales keep 5. int3's d·m·q needs d·m·3 in 8 bits, and m alone takes up to 5, so d is a power of
# two: the smallest at or above max|t|/31 over the super-block (t = the sub-block's E/ext), and
# m = round(|t|/d) + k clamped to 1..31 (offsets, not fractions: a small sub-block's m is a few units, where
# fractions of it round to the same m). Every |scale| (int3: d) is at least 2^-14, fp16's smallest normal, so all are
# exact in fp16 and the GGUF dequantises to the trained weights bit for bit (`export.gguf.check_round_trip`).
#
# Determinism (T2.3): group extremes are exact; scale arithmetic runs in fp64 and is rounded on the fp32 bits; codes
# are round-half-even of w/s made exact by checking an fp32 guess against the half-way points (q ± 0.5)·s, which
# are exact (≤ 13 significant bits); the search sums squared errors in fp64. Compiled and eager runs agree
# (tests/test_nbit.py). The gradient is the plain STE.
@dataclass(frozen=True)
class Grid:
gguf: str
lo: int
hi: int
group: int
ext: int # the grid end a group's extreme weight lands on (its sign orients the scale)
sig: int # significant bits kept in a scale (int3: unused, d is a power of two)
cands: tuple # the search's fractions c of E/ext (int3, iq4xs: integer offsets k to m)
values: tuple | None = None # D-57 non-linear grids: the value of each code lo..hi (None: the code itself)
# D-57: IQ4_NL's 16 levels (ggml's kvalues_iq4nl), shared by IQ4_XS. Codes are indices 0..15 into this table.
IQ4NL_VALUES = (-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113)
GRIDS = {"int2": Grid("Q2_0", -1, 2, 64, 2, 8, (1.0, 0.9, 0.8, 0.7, 0.6)),
"int4": Grid("Q4_0", -8, 7, 32, -8, 5, (1.0, 15 / 16, 7 / 8)),
"int3": Grid("Q3_K", -4, 3, 16, -4, 0, (-3, -2, -1, 0, 1)),
"iq4nl": Grid("IQ4_NL", 0, 15, 32, -127, 11, (1.0, 0.96, 0.92, 0.88, 0.84, 0.8, 0.76), IQ4NL_VALUES),
"iq4xs": Grid("IQ4_XS", 0, 15, 32, -127, 11, (-1, 0, 1), IQ4NL_VALUES)}
GRIDS["int3nl"] = GRIDS["int3"] # D-57: Q3_K gate/up; only its down projection differs (`model.student.proj_spec`)
GRIDS["int2nl"] = GRIDS["int2"] # D-57: Q2_0 gate/up with an IQ4_NL down projection (`model.student.proj_spec`)
Q3K_METHODS = ("int3", "int3nl")
INT2_METHODS = ("int2", "int2nl")
Q3K_SUPER, Q3K_M_MAX = 256, 31
IQ4XS_SUPER, IQ4XS_M_MAX = 256, 31
SCALE_MIN = 2.0 ** -14 # fp16's smallest normal
_SLICE = 1 << 24 # eager calls on stacked tensors (export, eval) run in slices of about this many weights
def _sliced(fn, w: torch.Tensor, s: torch.Tensor | None = None):
"""fn(w) or fn(w, s) on slices of dim 0 when eager: the search's temporaries are several times its input, and a
Marco layer's gate_up is 182M weights. Groups never cross dim 0, so the result is the same."""
call = (lambda a, b: fn(a)) if s is None else fn
if torch.compiler.is_compiling() or w.dim() < 2 or w.numel() <= _SLICE:
return call(w, s)
n = max(1, _SLICE // (w.numel() // w.shape[0]))
parts = [call(w[i:i + n], None if s is None else s[i:i + n]) for i in range(0, w.shape[0], n)]
if isinstance(parts[0], tuple):
return tuple(torch.cat(p) for p in zip(*parts))
return torch.cat(parts)
def _round_sig(x: torch.Tensor, sig: int) -> torch.Tensor:
"""|x| rounded to `sig` significant bits (nearest, ties to even) on its fp32 bits, at least 2^-14; sign kept
(-0 counts as positive). For positive x ≥ 2^-14 and sig = 8 this is `canonical_scale`."""
a = x.float().abs().clamp_min(SCALE_MIN)
drop = 24 - sig
b = a.view(torch.int32)
b = (b + ((1 << (drop - 1)) - 1) + ((b >> drop) & 1)) & -(1 << drop)
a = b.view(torch.float32)
return torch.where(x < 0, -a, a)
def _extreme(wg: torch.Tensor) -> torch.Tensor:
"""Each group's largest-magnitude weight, signed; the positive one on a tie."""
mx, mn = wg.amax(-1), wg.amin(-1)
return torch.where(mx >= -mn, mx, mn)
def _round_grid(x: torch.Tensor, a: torch.Tensor, lo: int, hi: int) -> torch.Tensor:
"""round-half-even(x / a) clamped to [lo, hi], exactly, for a > 0 with few significant bits: the fp32 quotient's
guess is off by at most one, and comparing x with the exact half-way points (q ± 0.5)·a corrects it."""
q = torch.round(x / a).clamp(lo - 1, hi + 1)
up, dn = (q + 0.5) * a, (q - 0.5) * a
odd = torch.remainder(q, 2) == 1 # at a tie the even neighbour wins: q ± 1 is even when q is odd
q = q + ((x > up) | ((x == up) & odd)).to(x.dtype) - ((x < dn) | ((x == dn) & odd)).to(x.dtype)
return q.clamp(lo, hi)
def _nearest_index(x: torch.Tensor, values: tuple) -> torch.Tensor:
"""D-57: index of the nearest table value to x (the lower one at a midpoint), as float codes 0..len-1."""
v = torch.tensor(values, dtype=torch.float32, device=x.device)
mids = (v[1:] + v[:-1]) / 2 # half-integers, exact
return torch.bucketize(x.float(), mids).to(x.dtype) # = (x > mids).sum(-1), by binary search
def _code_values(q: torch.Tensor, g: Grid) -> torch.Tensor:
"""The grid value of each code: the code itself, or (D-57) its entry in the table."""
if g.values is None:
return q
return torch.tensor(g.values, dtype=torch.float32, device=q.device)[q.long()]
def _grid_codes(wg: torch.Tensor, s: torch.Tensor, g: Grid) -> torch.Tensor:
"""Codes of grouped weights [..., n, g] against signed scales [..., n]: round(w / s) on the grid (D-57: the
nearest table value; the quotient is not exact there, which only moves weights within a hair of a midpoint)."""
s = s.unsqueeze(-1)
x = torch.where(s < 0, -wg, wg)
if g.values is not None:
return _nearest_index(x / s.abs(), g.values)
return _round_grid(x, s.abs(), g.lo, g.hi)
def _search(wg: torch.Tensor, g: Grid, cands: list[torch.Tensor], unit: torch.Tensor | None = None,
wt: torch.Tensor | None = None) -> torch.Tensor:
"""The candidate (a scale, or int3's m with scale unit·m) whose codes leave the least squared error per group
(D-57: weighted by `wt`, broadcastable to wg, when given)."""
best = err = None
for c in cands:
s = c if unit is None else unit * c
q = _grid_codes(wg, s, g)
d2 = (wg - _code_values(q, g) * s.unsqueeze(-1)).double().square()
e = (d2 if wt is None else d2 * wt).sum(-1)
if best is None:
best, err = c, e
else:
take = e < err
best, err = torch.where(take, c, best), torch.where(take, e, err)
return best
def _error_weight(wg: torch.Tensor, imp: torch.Tensor | None, g: Grid) -> torch.Tensor | None:
"""D-57: llama.cpp's imatrix weighting of a group's squared errors, imp_j · sqrt(σ² + w_j²) with σ² = 2·mean(w²)
over the group (its quantize_row_iq4_nl_impl). `imp` [..., in] is the mean squared input activation of each
input column; w is [..., out, in], so imp broadcasts over the output rows. None without an importance."""
if imp is None or g.values is None: # only the non-linear IQ4 grids: Q2_0/Q4_0 keep their trained rounding
return None
ig = _grouped(imp.float().unsqueeze(-2), g.group)
sigma2 = 2 * wg.square().mean(-1, keepdim=True)
return (ig * (sigma2 + wg.square()).sqrt()).double()
def _flat_scale(w: torch.Tensor, g: Grid, imp: torch.Tensor | None = None) -> torch.Tensor:
wg = _grouped(w, g.group)
e = _extreme(wg).double()
wt = _error_weight(wg, imp, g)
s = _search(wg, g, [_round_sig(e * (c / g.ext), g.sig) for c in g.cands], wt=wt)
if g.values is None:
return s
# D-57 non-linear grids: one (weighted) least-squares refit of the scale to the chosen codes, kept where it lowers
# the error with its own nearest codes, as llama.cpp's IQ4 quantiser does (sumqx / sumq2)
v = _code_values(_grid_codes(wg, s, g), g).double()
ww = torch.ones_like(v) if wt is None else wt
num, den = (ww * wg.double() * v).sum(-1), (ww * v * v).sum(-1)
ls = torch.where(den > 0, num / den.clamp_min(1e-30), s.double())
return _search(wg, g, [s, _round_sig(ls, g.sig)], wt=wt)
def _q3k_factors(w: torch.Tensor, g: Grid) -> tuple[torch.Tensor, torch.Tensor]:
"""Q3_K's scales of w [..., C]: d [..., C/256] (fp32 powers of two) and m [..., C/16] (signed integers in fp32,
1 ≤ |m| ≤ 31)."""
if w.shape[-1] % Q3K_SUPER:
raise ValueError(f"int3 (Q3_K) needs rows of a multiple of {Q3K_SUPER}, not {w.shape[-1]}")
wg = _grouped(w, g.group)
t = _extreme(wg).double() / g.ext # exact: ext is a power of two
x = (_grouped(t.abs(), Q3K_SUPER // g.group).amax(-1) / Q3K_M_MAX).float().clamp_min(SCALE_MIN)
d = ((x.view(torch.int32) + 0x7FFFFF) & 0x7F800000).view(torch.float32) # the power of two at or above x
unit = d.unsqueeze(-1).expand(*d.shape, Q3K_SUPER // g.group).flatten(-2) # d per sub-block
sign = torch.where(t < 0, -1.0, 1.0).float()
m0 = (t.abs() / unit.double()).round() # exact quotient (d is a power of two), half to even
cands = [sign * (m0 + k).clamp(1, Q3K_M_MAX).float() for k in g.cands]
return d, _search(wg, g, cands, unit)
def _q3k_scale(d: torch.Tensor, m: torch.Tensor) -> torch.Tensor:
return d.unsqueeze(-1).expand(*d.shape, Q3K_SUPER // GRIDS["int3"].group).flatten(-2) * m
def _iq4xs_factors(w: torch.Tensor, g: Grid, imp: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]:
"""D-57 IQ4_XS's scales of w [..., C]: d [..., C/256] (fp16-exact fp32) and m [..., C/32] (signed integers in
fp32, 1 ≤ |m| ≤ 31; GGUF stores m + 32 in 6 bits). Each sub-block's own IQ4_NL search sets its target t; d puts
the largest |t| of the super-block on m = 31, and m = round(|t|/d) + k for k in `cands` keeps the least error.
d·m has at most 16 significant bits, so d·m·value is exact in fp32, as ggml computes it."""
if w.shape[-1] % IQ4XS_SUPER:
raise ValueError(f"iq4xs (IQ4_XS) needs rows of a multiple of {IQ4XS_SUPER}, not {w.shape[-1]}")
wg = _grouped(w, g.group)
t = _flat_scale(w, GRIDS["iq4nl"], imp).double()
d = _round_sig(_grouped(t.abs(), IQ4XS_SUPER // g.group).amax(-1) / IQ4XS_M_MAX, g.sig)
unit = d.unsqueeze(-1).expand(*d.shape, IQ4XS_SUPER // g.group).flatten(-2)
sign = torch.where(t < 0, -1.0, 1.0).float()
m0 = (t.abs() / unit.double()).round()
cands = [sign * (m0 + k).clamp(1, IQ4XS_M_MAX).float() for k in g.cands]
return d, _search(wg, g, cands, unit, wt=_error_weight(wg, imp, g))
def _iq4xs_scale(d: torch.Tensor, m: torch.Tensor) -> torch.Tensor:
return d.unsqueeze(-1).expand(*d.shape, IQ4XS_SUPER // GRIDS["iq4xs"].group).flatten(-2) * m
def nbit_scale(w: torch.Tensor, spec: QuantSpec, importance: torch.Tensor | None = None) -> torch.Tensor:
"""The forward's signed per-group scale [..., in/g] in fp32 (int3: d·m per sub-block of 16; iq4xs: d·m per 32).
D-57: `importance` [..., in] weights the search of the flat grids and IQ4_XS (Q3_K's search stays unweighted)."""
g = GRIDS[spec.method]
w = w.detach().float()
if spec.method in Q3K_METHODS:
return _q3k_scale(*_sliced(lambda x: _q3k_factors(x, g), w))
if spec.method == "iq4xs":
if importance is None:
return _iq4xs_scale(*_sliced(lambda x: _iq4xs_factors(x, g), w))
return _iq4xs_scale(*_sliced(lambda x, i: _iq4xs_factors(x, g, i), w, importance))
if importance is None:
return _sliced(lambda x: _flat_scale(x, g), w)
return _sliced(lambda x, i: _flat_scale(x, g, i), w, importance)
def nbit_codes(w: torch.Tensor, s: torch.Tensor, spec: QuantSpec) -> torch.Tensor:
"""int8 codes on the format's grid, shape of w, for the signed scales s [..., in/g]."""
g = GRIDS[spec.method]
return _sliced(lambda x, y: _grid_codes(_grouped(x, g.group), y, g).flatten(-2).to(torch.int8),
w.detach().float(), s)
class _NBitSTE(torch.autograd.Function):
"""D-51: forward s·q on the format's grid, exactly; backward the plain STE (the identity to the latent)."""
@staticmethod
def forward(ctx, w, spec: QuantSpec, importance=None):
s = nbit_scale(w, spec, importance)
q = nbit_codes(w, s, spec)
v = _code_values(q.to(torch.float32), GRIDS[spec.method])
w_hat = (_grouped(v, spec.group_size) * s.unsqueeze(-1)).flatten(-2)
return w_hat.to(w.dtype)
@staticmethod
def backward(ctx, g):
return g, None, None
def pack_4bit(u: torch.Tensor) -> torch.Tensor:
"""uint8 values 0..15 [..., n] -> uint8 [..., n/2]; element 2i in the low nibble of byte i, 2i+1 in the high."""
u = u.to(torch.uint8).unflatten(-1, (-1, 2))
return u[..., 0] | (u[..., 1] << 4)
def unpack_4bit(packed: torch.Tensor) -> torch.Tensor:
return torch.stack([packed & 15, packed >> 4], -1).flatten(-2)
def export_packed_nbit(w: torch.Tensor, spec: QuantSpec, importance: torch.Tensor | None = None) -> dict:
"""Packed N-bit tensor: codes (int2: 2-bit as `pack_2bit`, q+1; int3/int4: 4-bit, q - lo) and fp16 scales
exactly as the forward used them; int3 adds `sub_scales` (int8 m), its `scales` being Q3_K's d."""
g = GRIDS[spec.method]
w = w.detach().float()
out = {"method": spec.method, "group_size": g.group, "shape": tuple(w.shape)}
if spec.method in Q3K_METHODS:
d, m = _sliced(lambda x: _q3k_factors(x, g), w)
s = _q3k_scale(d, m)
out.update(scales=d.to(torch.float16), sub_scales=m.to(torch.int8))
elif spec.method == "iq4xs":
if importance is None:
d, m = _sliced(lambda x: _iq4xs_factors(x, g), w)
else:
d, m = _sliced(lambda x, i: _iq4xs_factors(x, g, i), w, importance)
s = _iq4xs_scale(d, m)
out.update(scales=d.to(torch.float16), sub_scales=m.to(torch.int8))
else:
s = nbit_scale(w, spec, importance)
out["scales"] = s.to(torch.float16)
q = nbit_codes(w, s, spec)
out["codes"] = (pack_2bit(q) if spec.method in INT2_METHODS
else pack_4bit((q.to(torch.int16) - g.lo).to(torch.uint8)))
return out
def dequant_packed_nbit(p: dict, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
g = GRIDS[p["method"]]
if p["method"] in INT2_METHODS:
q = unpack_2bit(p["codes"]).to(torch.float32)
else:
q = _code_values((unpack_4bit(p["codes"]).to(torch.int16) + g.lo).to(torch.float32), g)
if p["method"] in Q3K_METHODS:
s = _q3k_scale(p["scales"].to(torch.float32), p["sub_scales"].to(torch.float32))
elif p["method"] == "iq4xs":
s = _iq4xs_scale(p["scales"].to(torch.float32), p["sub_scales"].to(torch.float32))
else:
s = p["scales"].to(torch.float32)
return (_grouped(q, g.group) * s.unsqueeze(-1)).flatten(-2).to(dtype)
# ---------------------------------------------------------------- diagnostics
def zero_fraction(codes: torch.Tensor) -> float:
return (codes == 0).float().mean().item()
def flip_rate(old: torch.Tensor, new: torch.Tensor) -> float:
return (old != new).float().mean().item()
# ---------------------------------------------------------------- slow references (tests only)
def reference_codes_and_scale(w: torch.Tensor, spec: QuantSpec) -> tuple[torch.Tensor, torch.Tensor]:
"""Element-by-element reference for absmean and twn over a 2-D tensor."""
rows, cols = w.shape
g = spec.group_size
codes = torch.zeros(rows, cols, dtype=torch.int8)
scales = torch.zeros(rows, cols // g, dtype=torch.float32)
for r in range(rows):
for j in range(cols // g):
grp = [float(x) for x in w[r, j * g:(j + 1) * g]]
mean_abs = sum(abs(x) for x in grp) / g
if spec.method == "twn":
delta = TWN_DELTA * mean_abs
kept = [abs(x) for x in grp if abs(x) > delta]
s = sum(kept) / max(len(kept), 1)
q = [(1 if x > 0 else -1) if abs(x) > delta else 0 for x in grp]
else:
s = mean_abs
q = None
s = max(s, spec.eps)
s = float(torch.tensor(s, dtype=torch.float32).to(torch.bfloat16).float())
if q is None:
# round-half-even of x/s clamped to [-1, 1], in exact arithmetic (s/2 is exact)
q = [(1 if x > 0 else -1) if abs(x) > 0.5 * s else 0 for x in grp]
scales[r, j] = s
codes[r, j * g:(j + 1) * g] = torch.tensor(q, dtype=torch.int8)
return codes, scales