diff --git a/c/tools/convert_fp8_to_int4.py b/c/tools/convert_fp8_to_int4.py index b8e186e7..382125eb 100644 --- a/c/tools/convert_fp8_to_int4.py +++ b/c/tools/convert_fp8_to_int4.py @@ -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 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"]: @@ -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): @@ -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): @@ -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):