Skip to content

Converter: per-tensor-type mixed precision (int8 for shared/o_proj/kv_b_proj, int4 rest) #243

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
woolcoxm wants to merge 1 commit into from
Closed

Converter: per-tensor-type mixed precision (int8 for shared/o_proj/kv_b_proj, int4 rest) #243

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 52 additions & 8 deletions c/tools/convert_fp8_to_int4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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] ----------
Expand Down Expand Up @@ -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())
Expand All @@ -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:
Expand All @@ -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)
Expand All @@ -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.
Expand Down Expand Up @@ -336,7 +380,7 @@ def get_slice(s, n): return None
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"]:
Expand Down Expand Up @@ -579,7 +623,7 @@ def _download_single(url, fn, out, part, expected):
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):
Expand All @@ -599,7 +643,7 @@ def _download_single(url, fn, out, part, expected):
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):
Expand All @@ -613,7 +657,7 @@ def _download_single(url, fn, out, part, expected):
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):
Expand Down