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

converter: per-tensor-type mixed-precision control

Browse Source 
Split the resident weight classification into 5 sub-types so each can
get different precision:
  sh   = shared expert (highest sensitivity, fires every token)
  o    = o_proj (reconstructs output, biggest attn tensor)
  kvb  = kv_b_proj (reconstructs KV cache on every decode)
  attn = q_a/q_b/kv_a (other attention projections)
  dmlp = dense MLP (first 3 layers)

New args: --shared-bits, --o-bits, --kvb-bits, --attn-bits, --dmlp-bits
Each defaults to ebits (backward compat). When set, the converter applies
that precision to just that tensor type.

Research-backed plan: put the 3 compounding tensors (shared expert, o_proj,
kv_b_proj) at int8 and everything else at grouped int4. Extra RAM cost:
only +5.3 GB (those tensors are small vs the 372 GB expert pool on disk).
This commit is contained in:
woolcoxm
2026-07-14 20:32:14 -04:00
parent 3fd47b7bbd
commit e141db047d
1 changed files with 52 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/tools/convert_fp8_to_int4.py
+52 -8
View File
@@ -116,7 +116,7 @@ def layer_idx(name):
def classify(name, n_layers, keep_mtp=False, keep_idx=False):
if name.endswith("_scale_inv"): return "consumed" # FP8 base: gestito col suo peso
# NVFP4 (modelopt): i sidecar delle scale sono consumati insieme al loro .weight U8.
# NVFP4 (modelopt): i sidecar delle scale sono consumati insieme al loro U8 .weight.
# EN: NVFP4 (modelopt): scale sidecars are consumed together with their U8 .weight.
if name.endswith((".weight_scale", ".weight_scale_2", ".input_scale")): return "consumed"
li = layer_idx(name)
@@ -137,7 +137,20 @@ def classify(name, n_layers, keep_mtp=False, keep_idx=False):
if name.endswith("norm.weight") or name == "model.norm.weight": return "f32"
if name in ("model.embed_tokens.weight", "lm_head.weight"): return "io"
if ".mlp.experts." in name and name.endswith(".weight"): return "x" # expert ROUTED (streaming)
if name.endswith(".weight"): return "q" # attn/dense-mlp/shared (residente)
# Split resident weights by type for mixed-precision control:
# "sh" = shared expert (fires on every token, highest sensitivity)
# "o" = o_proj attention (reconstructs output, biggest attn tensor)
# "kvb" = kv_b_proj (reconstructs KV cache on every decode step)
# "attn" = other attention projections (q_a, q_b, kv_a)
# "dmlp" = dense MLP (first 3 layers)
if "shared_experts" in name: return "sh"
if name.endswith("o_proj.weight"): return "o"
if name.endswith("kv_b_proj.weight"): return "kvb"
if any(name.endswith(k) for k in ("q_a_proj.weight", "q_b_proj.weight",
"kv_a_proj_with_mqa.weight")): return "attn"
if any(name.endswith(k) for k in ("mlp.gate_proj.weight", "mlp.up_proj.weight",
"mlp.down_proj.weight")): return "dmlp"
if name.endswith(".weight"): return "q" # fallback: other resident weights
return "f32"
# ---------- dequant NVFP4 (modelopt) di UN tensore expert -> f32 [O,I] ----------
@@ -202,7 +215,7 @@ def dequant(f, name, keys):
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, group_size=0):
keep_mtp=False, keep_idx=False, group_size=0, bits_map=None):
from safetensors import safe_open
with safe_open(path, framework="pt") as f:
keys = set(f.keys())
@@ -213,7 +226,15 @@ def convert_shard(path, out_dict, n_layers, ebits, io_bits, xbits,
if kind == "f32":
out_dict[name] = w.astype(np.float32)
else:
bits = io_bits if kind == "io" else xbits if kind == "x" else ebits
# Resolve bits for this tensor type: use bits_map override if provided,
# otherwise fall back to the classic ebits/xbits/io_bits scheme.
if bits_map and kind in bits_map:
bits = bits_map[kind]
else:
bits = io_bits if kind == "io" else xbits if kind == "x" else ebits
# Any unknown kind that fell through classify as "q"
if bits_map and kind not in bits_map and kind not in ("io", "x", "sh", "o", "kvb", "attn", "dmlp"):
bits = ebits
if w.ndim != 2: # es. bias 1D non previsto come 'q' -> tienilo f32
out_dict[name] = w.astype(np.float32); continue
if group_size > 0 and bits <= 4:
@@ -234,6 +255,18 @@ 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
# Mixed-precision: per-tensor-type bit overrides. Default = ebits (all same).
# Set these higher to protect sensitive tensors from quantization error.
ap.add_argument("--shared-bits", type=int, default=None,
help="bits for shared expert (fires on every token, highest sensitivity). Default=ebits")
ap.add_argument("--o-bits", type=int, default=None,
help="bits for o_proj attention (reconstructs output, biggest attn tensor). Default=ebits")
ap.add_argument("--kvb-bits", type=int, default=None,
help="bits for kv_b_proj (reconstructs KV cache on every decode). Default=ebits")
ap.add_argument("--attn-bits", type=int, default=None,
help="bits for other attention projections (q_a, q_b, kv_a). Default=ebits")
ap.add_argument("--dmlp-bits", type=int, default=None,
help="bits for dense MLP (first 3 layers). 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)
@@ -255,6 +288,17 @@ def main():
a.ebits = 8 if (a.mtp or a.indexer) else 4
if a.xbits is None: a.xbits = a.ebits
# Build per-type bits map. If a type-specific arg is set, use it; otherwise the
# converter falls back to ebits for that type.
bits_map = {}
if a.shared_bits is not None: bits_map["sh"] = a.shared_bits
if a.o_bits is not None: bits_map["o"] = a.o_bits
if a.kvb_bits is not None: bits_map["kvb"] = a.kvb_bits
if a.attn_bits is not None: bits_map["attn"] = a.attn_bits
if a.dmlp_bits is not None: bits_map["dmlp"] = a.dmlp_bits
if bits_map:
print(f"[MIXED] precision map: " + ", ".join(f"{k}={v}bit" for k,v in sorted(bits_map.items())))
if a.selftest_nvfp4:
import torch
# 1) LUT e2m1: i 16 codici devono decodificare esattamente ai valori attesi.
@@ -336,7 +380,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, group_size=a.group_size)
out = {}; convert_shard(sp, out, a.n_layers, a.ebits, a.io_bits, a.xbits, group_size=a.group_size, bits_map=bits_map)
save_file(out, os.path.join(a.outdir, f"out-{i:05d}.safetensors"))
# copia config + tokenizer
for fn in ["config.json"]:
@@ -579,7 +623,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, group_size=a.group_size)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_mtp=True, group_size=a.group_size, bits_map=bits_map)
save_file(out, outp)
os.remove(p)
for blob in glob.glob(os.path.join(tmp, "**", "*"), recursive=True):
@@ -599,7 +643,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, group_size=a.group_size)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_idx=True, group_size=a.group_size, bits_map=bits_map)
if out: save_file(out, outp)
os.remove(p)
for blob in glob.glob(os.path.join(tmp, "**", "*"), recursive=True):
@@ -613,7 +657,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, group_size=a.group_size)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, group_size=a.group_size, bits_map=bits_map)
save_file(out, outp)
os.remove(p) # <-- cancella subito lo shard fp8
for blob in glob.glob(os.path.join(tmp, "**", "*"), recursive=True):