#!/usr/bin/env python3 """Shard-streaming FP4 re-packer for deepseek-ai/DeepSeek-V4.1-Flash (DeepseekV41ForCausalLM / deepseek_v41, 552B backbone + 196B Engram, multimodal, CED). WHY THIS IS NOT THE USUAL QUANTIZER ----------------------------------- Every prior LibertAI NVFP4 release (Qwen3.6/3.8, Gemma-4, Hy3, GLM-5.3, GLM-5.3-Flash) started from a BF16 checkpoint and took the routed experts to 4 bits for a ~70% cut. DeepSeek-V4.1-Flash ships FOUR-BIT ALREADY. From its own config.json: "quant_method": "fp8", "weight_block_size": [32,32], "scale_fmt": "ue8m0", "expert_dtype": "fp4" Measured over all 48 shards / 96,085 tensors (475.2 GiB): routed experts 268.9 GiB 56.6% I8 packed E2M1 + F8_E8M0 scales, block 1x32 Engram tables 189.1 GiB 39.8% F8_E4M3 + F8_E8M0 scales, block 1x32 MTP (3 layers) 7.4 GiB 1.6% experts likewise E2M1 attention/dense 4.9 GiB 1.0% F8_E4M3 block 32x32 embed/head/vision 4.0 GiB 0.8% BF16 So there is no "-70%" to be had. This script does the two things that ARE left: 1. --experts nvfp4 MXFP4 (block-32, E8M0) -> NVFP4 (block-16, E4M3 + f32 global). This is a LOSSLESS FORMAT TRANSCODE, not a re-quantization, and it is lossless for a specific reason: NVFP4's 16-element blocks are a strict refinement of the source's 32-element blocks, and the source values already sit on the E2M1 grid. Both NVFP4 half-blocks of a source block therefore share one scale s, so if we emit the nibbles UNCHANGED and force weight_scale * weight_scale_2 == s exactly, every reconstructed value is bit-identical to the source. Since E8M0 scales are powers of two and E4M3 represents 2^k exactly for k in [-9, 8], that identity holds whenever s/global lands in E4M3's exponent window -- checked per tensor, and the global scale is chosen to centre the window. Costs +0.25 bit/weight (4.5 vs 4.25), i.e. the experts GROW 268.9 -> ~284.6 GiB. Ship it only for engines whose Blackwell MoE kernels want NVFP4 and have no MX path. 2. --engram fp4 F8_E4M3 -> E2M1, keeping the source's own block-32 E8M0 scale layout (i.e. stored exactly the way DeepSeek stores its experts, so any kernel that can dequant an expert can dequant the table). This IS a real quantization and the only real size lever in the checkpoint: 189.1 -> ~97 GiB. Engram is a hashed n-gram GATHER, not a GEMM -- ~48 rows x 256 B ~= 12 KB per token -- so it is cheap to dequant on lookup and NVMe-friendly if offloaded. Never holds more than one shard in RAM. CPU only by default. Source tensor naming (flat, no "model." prefix, "ffn" not "mlp"): layers.{L}.ffn.experts.{E}.w{1,2,3}.weight int8 [out, in/2] layers.{L}.ffn.experts.{E}.w{1,2,3}.scale e8m0 [out, in/32] layers.{L}.engram.embed.weight e4m3 [rows, 256] layers.{L}.engram.embed.scale e8m0 [rows, 8] mtp.{M}.ffn.experts.{E}... same as backbone """ import argparse, json, os, re, shutil, sys, time import torch from safetensors import safe_open from safetensors.torch import save_file # E2M1, index = raw 4-bit code. Taken verbatim from the checkpoint's own # inference/convert.py FP4_TABLE -- not reconstructed from the spec. FP4_TABLE = torch.tensor( [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], dtype=torch.float32) FP4_MAX = 6.0 SRC_BLOCK = 32 # DeepSeek's block, both for experts and Engram NVFP4_BLOCK = 16 # NVFP4's block EXPERT_RE = re.compile(r'^(?:layers|mtp)\.\d+\.(?:\.?\d+\.)?.*ffn\.experts\.\d+\.w[123]\.weight$') ENGRAM_RE = re.compile(r'^layers\.\d+\.engram\.embed\.weight$') # ---------------------------------------------------------------- E2M1 helpers def unpack_e2m1(packed: torch.Tensor) -> torch.Tensor: """int8/uint8 [out, in/2] -> float32 [out, in] of E2M1 grid values. Nibble order is LOW FIRST -- convert.py does torch.stack([FP4_TABLE[low], FP4_TABLE[high]], dim=-1).flatten() so byte b holds element 2i in its low nibble and 2i+1 in its high nibble. """ u = packed.view(torch.uint8) low = (u & 0x0F).long() high = ((u >> 4) & 0x0F).long() return torch.stack([FP4_TABLE[low], FP4_TABLE[high]], dim=-1).flatten(-2) def pack_e2m1(codes: torch.Tensor) -> torch.Tensor: """uint8 codes [out, in] (values 0..15) -> packed uint8 [out, in/2], low nibble first.""" assert codes.shape[-1] % 2 == 0 c = codes.view(*codes.shape[:-1], -1, 2) return (c[..., 0] | (c[..., 1] << 4)).to(torch.uint8) def quantize_e2m1_codes(x: torch.Tensor) -> torch.Tensor: """float [..., N] already divided by its block scale -> nearest E2M1 code (uint8). Round-to-nearest on the magnitude grid {0,.5,1,1.5,2,3,4,6}; ties go to the even code, matching what the FP4 hardware conversion does. """ mag = x.abs().clamp(max=FP4_MAX) grid = FP4_TABLE[:8] # 8 magnitudes, ascending # midpoints between consecutive grid points: .25 .75 1.25 1.75 2.5 3.5 5.0 mid = (grid[1:] + grid[:-1]) / 2 code = torch.bucketize(mag, mid) # 0..7 sign = (x < 0).to(torch.uint8) << 3 return (code.to(torch.uint8) | sign) # ------------------------------------------------------- experts: MXFP4->NVFP4 def transcode_expert_to_nvfp4(packed: torch.Tensor, scale_e8m0: torch.Tensor): """(int8 [out,in/2], e8m0 [out,in/32]) -> (uint8 [out,in/2], e4m3 [out,in/16], f32 []) Lossless when every block scale lands in E4M3's exact power-of-two window. Returns (weight, weight_scale, weight_scale_2, n_clamped, n_blocks). """ out_dim, half = packed.shape in_dim = half * 2 assert scale_e8m0.shape == (out_dim, in_dim // SRC_BLOCK), \ f"scale {tuple(scale_e8m0.shape)} vs weight {tuple(packed.shape)}" s = scale_e8m0.float() # exact powers of two # each source block-32 covers two NVFP4 block-16s -> repeat along the block axis ws_target = s.repeat_interleave(SRC_BLOCK // NVFP4_BLOCK, dim=1) # [out, in/16] # Choose a per-tensor global scale that centres the block scales in E4M3's # exactly-representable power-of-two window [2^-9, 2^8]. Work in log2 so the # global stays a power of two and the division below is exact. finite = ws_target[ws_target > 0] if finite.numel() == 0: gexp = 0 else: lo = torch.log2(finite.min()).item() hi = torch.log2(finite.max()).item() # centre of the available window is 2^-0.5; aim the scale midpoint there gexp = int(round((lo + hi) / 2 + 0.5)) gs = torch.tensor(2.0, dtype=torch.float32) ** gexp ratio = ws_target / gs n_blocks = ratio.numel() lo_lim, hi_lim = 2.0 ** -9, 2.0 ** 8 clamped_mask = (ratio > 0) & ((ratio < lo_lim) | (ratio > hi_lim)) n_clamped = int(clamped_mask.sum()) ratio = ratio.clamp(min=lo_lim, max=hi_lim) ws = ratio.to(torch.float8_e4m3fn) # nibbles pass through untouched -- this is what makes it lossless weight = packed.view(torch.uint8).clone() return weight, ws, gs.clone(), n_clamped, n_blocks # ------------------------------------------------------------ engram: fp8->fp4 def quantize_engram_to_fp4(f, name, sname, rows, cols, chunk_rows=1 << 20, device="cpu"): """Streamed E4M3 -> E2M1 for one Engram table, read in row chunks. Each table is a SINGLE ~94.5 GiB tensor of ~384M rows x 256, so it must never be materialised whole: dequantising it to float32 would want ~393 GB. We preallocate the packed output (~47 GiB) and fill it a chunk at a time via safetensors slices. Keeps the source's block-32 / E8M0 layout, so a table row ends up stored exactly the way DeepSeek stores an expert weight. Scales are recomputed as a power of two >= amax/6 rather than reused, so no block clips. Returns (packed uint8 [rows, cols/2], scale e8m0 [rows, cols/32], cos_sum, cos_n). """ nblk = cols // SRC_BLOCK packed = torch.empty((rows, cols // 2), dtype=torch.uint8) scales = torch.empty((rows, nblk), dtype=torch.float8_e8m0fnu) wsl = f.get_slice(name) ssl = f.get_slice(sname) cos_sum, cos_n = 0.0, 0 for lo in range(0, rows, chunk_rows): hi = min(lo + chunk_rows, rows) n = hi - lo v = (wsl[lo:hi].to(device).float().view(n, nblk, SRC_BLOCK) * ssl[lo:hi].to(device).float().unsqueeze(-1)) amax = v.abs().amax(dim=-1, keepdim=True) exp = torch.ceil(torch.log2((amax / FP4_MAX).clamp(min=1e-38))).clamp(-127, 127) new_scale = torch.pow(torch.tensor(2.0, device=v.device), exp) codes = quantize_e2m1_codes(v / new_scale) packed[lo:hi] = pack_e2m1(codes.view(n, cols)).to("cpu") scales[lo:hi] = new_scale.squeeze(-1).to(torch.float8_e8m0fnu).to("cpu") # cosine on a subsample of each chunk, so the figure spans the whole table k = min(4096, n) a = v[:k].flatten(1) b = (FP4_TABLE.to(v.device)[codes[:k].long()].view(k, nblk, SRC_BLOCK) * new_scale[:k]).flatten(1) num = (a * b).sum(1); den = a.norm(dim=1) * b.norm(dim=1) ok = den > 0 if ok.any(): cos_sum += float((num[ok] / den[ok]).sum()); cos_n += int(ok.sum()) del v, codes, new_scale, amax, exp, a, b, num, den, ok print(f" engram {name}: rows {hi}/{rows} ({100*hi/rows:.1f}%) " f"cos={cos_sum/max(cos_n,1):.6f}", flush=True) return packed, scales, cos_sum, cos_n # ------------------------------------------------------------------ shard loop DEVICE = "cpu" def process_shard(src, dst, do_experts, do_engram, stats): out = {} with safe_open(src, framework="pt") as f: keys = list(f.keys()) for name in keys: if do_experts and EXPERT_RE.match(name): base = name[:-len(".weight")] sname = base + ".scale" if sname not in keys: out[name] = f.get_tensor(name); stats["pass"] += 1; continue w, ws, ws2, nc, nb = transcode_expert_to_nvfp4( f.get_tensor(name), f.get_tensor(sname)) out[base + ".weight"] = w out[base + ".weight_scale"] = ws out[base + ".weight_scale_2"] = ws2 stats["experts"] += 1 stats["clamped"] += nc stats["blocks"] += nb elif name.endswith(".scale") and do_experts and EXPERT_RE.match(name[:-len(".scale")] + ".weight"): continue # consumed above elif do_engram and ENGRAM_RE.match(name): base = name[:-len(".weight")] sname = base + ".scale" shp = f.get_slice(name).get_shape() packed, new_scale, cs, cn = quantize_engram_to_fp4( f, name, sname, shp[0], shp[1], device=DEVICE) out[base + ".weight"] = packed out[base + ".scale"] = new_scale stats["engram"] += 1 stats["engram_cos_sum"] += cs stats["engram_cos_n"] += cn elif name.endswith(".scale") and do_engram and ENGRAM_RE.match(name[:-len(".scale")] + ".weight"): continue else: out[name] = f.get_tensor(name); stats["pass"] += 1 save_file(out, dst, metadata={"format": "pt"}) wmap, total = {}, 0 for k, v in out.items(): wmap[k] = os.path.basename(dst) total += v.numel() * v.element_size() return wmap, total def main(): ap = argparse.ArgumentParser() ap.add_argument("--src", required=True) ap.add_argument("--out", required=True) ap.add_argument("--experts", choices=["nvfp4", "keep"], default="nvfp4") ap.add_argument("--engram", choices=["fp4", "keep"], default="fp4") ap.add_argument("--only-shards", type=int, default=0) ap.add_argument("--shard-list", default=None, help="comma-separated shard filenames") ap.add_argument("--device", default="cpu", help="device for the Engram loop ('cuda' is ~10x on 98G elements/table)") args = ap.parse_args() global DEVICE DEVICE = args.device if DEVICE != "cpu": print(f"[engram device] {DEVICE}", flush=True) index = json.load(open(os.path.join(args.src, "model.safetensors.index.json"))) shards = sorted(set(index["weight_map"].values())) if args.shard_list: want = set(args.shard_list.split(",")) shards = [s for s in shards if s in want] if args.only_shards: shards = shards[:args.only_shards] os.makedirs(args.out, exist_ok=True) do_e = args.experts == "nvfp4" do_g = args.engram == "fp4" stats = {"experts": 0, "engram": 0, "pass": 0, "clamped": 0, "blocks": 0, "engram_cos_sum": 0.0, "engram_cos_n": 0} full_map, total, t0 = {}, 0, time.time() for s in shards: sp = os.path.join(args.src, s) if not os.path.exists(sp): print(f" SKIP (absent) {s}", flush=True); continue wmap, tb = process_shard(sp, os.path.join(args.out, s), do_e, do_g, stats) full_map.update(wmap); total += tb cos = (f" engram_cos={stats['engram_cos_sum']/stats['engram_cos_n']:.6f}" if stats["engram_cos_n"] else "") print(f" {s} experts={stats['experts']} engram={stats['engram']} " f"pass={stats['pass']} clamped={stats['clamped']}/{stats['blocks']}{cos} " f"[{time.time()-t0:.0f}s]", flush=True) json.dump({"metadata": {"total_size": total}, "weight_map": full_map}, open(os.path.join(args.out, "model.safetensors.index.json"), "w"), indent=1) print(f"[DONE] {total/2**30:.1f} GiB experts={stats['experts']} engram={stats['engram']} " f"pass={stats['pass']} lossless_blocks={stats['blocks']-stats['clamped']}/{stats['blocks']}") if stats["engram_cos_n"]: print(f"[engram] mean cosine {stats['engram_cos_sum']/stats['engram_cos_n']:.6f} " f"over {stats['engram_cos_n']} sampled rows") if __name__ == "__main__": main()