url/colibri
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

experiment: group-scaled int4 (fmt=4) — one scale per 128 elements, not per row

Browse Source 
Root cause of gibberish output: the int4 quantization uses one F32 scale per
output row (2048 scales for a 2048x6144 matrix). The FP8 source has 128x128
block scales — 48x finer. This destroys reasoning while keeping surface fluency.

Changes:
- Converter: quant_int4_grouped() with --group-size 128 arg. Same nibble
  packing, but one scale per group of 128 elements along the input dim.
- Engine QT struct: added 'gs' field (group size, 0=per-row backward compat)
- Engine qt_from_disk: auto-detects fmt=4 when scale array is O*ceil(I/128)
  elements instead of O. Old per-row models (fmt=2) work unchanged.
- Engine matmul_i4_grouped(): AVX2 kernel that applies per-group scales.
  Accumulator resets at each group boundary: dot(x[grp],w[grp]) * scale[grp].
- Engine matmul_qt_ex: dispatches to grouped kernel for fmt=4 (always exact,
  no IDOT approximation since the point is quality)
- Engine expert_load: both mmap and slab+pread paths detect fmt=4 from
  scale array size and set gs=128
- qt_bytes: fmt=4 reports correct memory including group scales

Backward compatible: existing per-row int4 models work unchanged.
The fused gate+up pair path (matmul_i4_pair) falls back to separate
matmul_qt calls for fmt=4 — minor perf cost, correctness preserved.
This commit is contained in:
woolcoxm
2026-07-14 18:57:35 -04:00
parent c71e5d1cf8
commit 498ab0c20e
2 changed files with 122 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/tools/convert_fp8_to_int4.py
+45 -7
View File
@@ -51,6 +51,38 @@ def quant_int4(w, bits): # -> (qbytes U8 [O*ceil(I/2)], s
out[:, :v1.shape[1]] |= (v1 << 4)
return out.reshape(-1), s[:, 0].astype(np.float32)
def quant_int4_grouped(w, bits, gs=128):
"""Group-scaled int4: one scale per group of `gs` elements along the input dim.
Drastically reduces quantization error vs per-row scaling — matches the FP8
source's 128x128 block-scale granularity. Output layout:
qbytes: same packed nibble format as quant_int4
scales: f32 [O * ngroups] where ngroups = ceil(I/gs), laid out as
s[o * ngroups + g] = scale for row o, group g.
The engine detects this format (fmt=4) by checking the .qs array size."""
O, I = w.shape
qmax = (1 << (bits - 1)) - 1
ngroups = (I + gs - 1) // gs
# pad I to a multiple of gs for clean reshape, then trim
Ipad = ngroups * gs
wpad = np.zeros((O, Ipad), np.float32)
wpad[:, :I] = w
wr = wpad.reshape(O, ngroups, gs) # [O, ngroups, gs]
amax = np.abs(wr).max(axis=2, keepdims=True) # [O, ngroups, 1]
s = np.maximum(amax / qmax, 1e-8) # [O, ngroups, 1]
q = np.clip(np.rint(wr / s), -8, qmax).astype(np.int32) # [O, ngroups, gs]
q = q.reshape(O, Ipad)[:, :I] # trim padding -> [O, I]
# pack nibbles (identical to quant_int4)
rb = (I + 1) // 2
out = np.zeros((O, rb), np.uint8)
v0 = (q[:, 0::2] + 8).astype(np.uint8)
out[:, :v0.shape[1]] = v0
if I > 1:
v1 = (q[:, 1::2] + 8).astype(np.uint8)
out[:, :v1.shape[1]] |= (v1 << 4)
# scales: flatten [O, ngroups] -> [O * ngroups]
s_flat = s[:, :, 0].astype(np.float32).reshape(-1)
return out.reshape(-1), s_flat
def quant_int2(w, bits): # -> (qbytes U8 [O*ceil(I/4)], scale f32 [O]); 4/byte
O, I = w.shape
qmax = (1 << (bits - 1)) - 1 # bits=2 -> qmax=1, valori [-2,1]
@@ -169,7 +201,8 @@ def dequant(f, name, keys):
return (w * sc).numpy()
return f.get_tensor(name).to(torch.float32).numpy()
def convert_shard(path, out_dict, n_layers, ebits, io_bits, xbits, keep_mtp=False, keep_idx=False):
def convert_shard(path, out_dict, n_layers, ebits, io_bits, xbits,
keep_mtp=False, keep_idx=False, group_size=0):
from safetensors import safe_open
with safe_open(path, framework="pt") as f:
keys = set(f.keys())
@@ -183,8 +216,11 @@ def convert_shard(path, out_dict, n_layers, ebits, io_bits, xbits, keep_mtp=Fals
bits = io_bits if kind == "io" else xbits if kind == "x" else ebits
if w.ndim != 2: # es. bias 1D non previsto come 'q' -> tienilo f32
out_dict[name] = w.astype(np.float32); continue
q, s = (quant_int2(w, bits) if bits <= 2 else
quant_int4(w, bits) if bits <= 4 else quant_int8(w, bits))
if group_size > 0 and bits <= 4:
q, s = quant_int4_grouped(w, bits, group_size)
else:
q, s = (quant_int2(w, bits) if bits <= 2 else
quant_int4(w, bits) if bits <= 4 else quant_int8(w, bits))
out_dict[name] = q
out_dict[name + ".qs"] = s
@@ -198,6 +234,8 @@ def main():
ap.add_argument("--ebits", type=int, default=None) # bit residenti (default 4; 8 per --mtp/--indexer)
ap.add_argument("--io-bits", type=int, default=8) # bit di embed/lm_head
ap.add_argument("--xbits", type=int, default=None) # bit degli expert ROUTED (streaming); default=ebits
ap.add_argument("--group-size", type=int, default=0, # 0 = per-row (backward compat); 128 = group-scaled
help="group size for int4 scales: 0=per-row (default), 128=one scale per 128 elements (much better quality)")
ap.add_argument("--n-layers", type=int, default=78)
ap.add_argument("--min-free-gb", type=float, default=20.0)
ap.add_argument("--selftest", action="store_true")
@@ -298,7 +336,7 @@ def main():
shards = sorted(glob.glob(os.path.join(a.indir, "*.safetensors")))
from safetensors.numpy import save_file
for i, sp in enumerate(shards):
out = {}; convert_shard(sp, out, a.n_layers, a.ebits, a.io_bits, a.xbits)
out = {}; convert_shard(sp, out, a.n_layers, a.ebits, a.io_bits, a.xbits, group_size=a.group_size)
save_file(out, os.path.join(a.outdir, f"out-{i:05d}.safetensors"))
# copia config + tokenizer
for fn in ["config.json"]:
@@ -541,7 +579,7 @@ def main():
if os.path.exists(outp): print(f"[MTP] {outp} already done"); continue
print(f"[MTP {i+1}/{len(mtp_shards)}] downloading {sh}...", flush=True)
p = download_retry(a.repo, sh, tmp)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_mtp=True)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_mtp=True, group_size=a.group_size)
save_file(out, outp)
os.remove(p)
for blob in glob.glob(os.path.join(tmp, "**", "*"), recursive=True):
@@ -561,7 +599,7 @@ def main():
if os.path.exists(outp): continue # gia' fatto -> ripartibile
print(f"[IDX {i+1}/{len(idx_shards)}] downloading {sh}...", flush=True)
p = download_retry(a.repo, sh, tmp)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_idx=True)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_idx=True, group_size=a.group_size)
if out: save_file(out, outp)
os.remove(p)
for blob in glob.glob(os.path.join(tmp, "**", "*"), recursive=True):
@@ -575,7 +613,7 @@ def main():
if os.path.exists(outp): continue # gia' fatto -> ripartibile
print(f"[{i+1}/{len(shards)}] downloading {sh} ({free_gb(a.outdir):.0f} GB free)...", flush=True)
p = download_retry(a.repo, sh, tmp)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, group_size=a.group_size)
save_file(out, outp)
os.remove(p) # <-- cancella subito lo shard fp8
for blob in glob.glob(os.path.join(tmp, "**", "*"), recursive=True):