| |
| """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, struct, sys, time |
| import torch |
| from safetensors import safe_open |
| from safetensors.torch import save_file |
|
|
| |
| |
| 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 |
| NVFP4_BLOCK = 16 |
|
|
| 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$') |
|
|
|
|
| |
|
|
| 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].to(x.device) |
| |
| mid = (grid[1:] + grid[:-1]) / 2 |
| code = torch.bucketize(mag, mid) |
| sign = (x < 0).to(torch.uint8) << 3 |
| return (code.to(torch.uint8) | sign) |
|
|
|
|
| |
|
|
| 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() |
| |
| ws_target = s.repeat_interleave(SRC_BLOCK // NVFP4_BLOCK, dim=1) |
|
|
| |
| |
| |
| 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() |
| |
| 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) |
| |
| weight = packed.view(torch.uint8).clone() |
| return weight, ws, gs.clone(), n_clamped, n_blocks |
|
|
|
|
| |
|
|
| 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") |
|
|
| |
| 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 |
|
|
|
|
| |
|
|
|
|
| |
|
|
| ST_DTYPE = {torch.uint8: ("U8", 1), torch.int8: ("I8", 1), |
| torch.float8_e4m3fn: ("F8_E4M3", 1), torch.float8_e8m0fnu: ("F8_E8M0", 1), |
| torch.bfloat16: ("BF16", 2), torch.float16: ("F16", 2), |
| torch.float32: ("F32", 4), torch.int32: ("I32", 4), torch.int64: ("I64", 8)} |
|
|
|
|
| def write_engram_shard(src_path, dst_path, device="cpu", chunk_rows=1 << 20): |
| """Quantize and write an Engram shard without ever holding the table in RAM. |
| |
| The obvious implementation -- build the packed tensor, hand the dict to save_file -- |
| needs ~47 GiB resident for the weights alone, and on a GB10 the GPU allocator draws on |
| that same pool. It got OOM-killed at 83%. So we emit the safetensors container by hand: |
| the format is an 8-byte little-endian header length, that many bytes of JSON naming |
| each tensor's dtype/shape/byte-range, then the raw buffers back to back. Every offset |
| is known before a single value is computed, so the packed table can be streamed to the |
| file a chunk at a time. |
| |
| Scales are the one thing still buffered (~3 GB for 384M rows x 8) because they are |
| produced by the same pass that produces the packed rows but live elsewhere in the file. |
| |
| Peak memory is one chunk plus the scale buffer. |
| """ |
| stats = {"engram": 0, "pass": 0, "cos_sum": 0.0, "cos_n": 0} |
| with safe_open(src_path, framework="pt") as f: |
| keys = list(f.keys()) |
| emb = [k for k in keys if ENGRAM_RE.match(k)] |
| consumed = set() |
| plan = [] |
| for name in emb: |
| base = name[:-len(".weight")] |
| sname = base + ".scale" |
| rows, cols = f.get_slice(name).get_shape() |
| nblk = cols // SRC_BLOCK |
| plan.append((base + ".weight", "U8", [rows, cols // 2], "stream", (name, sname))) |
| plan.append((base + ".scale", "F8_E8M0", [rows, nblk], "scale", base)) |
| consumed.add(name); consumed.add(sname) |
| for name in sorted(keys): |
| if name in consumed: |
| continue |
| plan.append((name, None, list(f.get_slice(name).get_shape()), "copy", name)) |
|
|
| |
| resolved = [] |
| for name, dts, shape, kind, srcref in plan: |
| if kind == "copy": |
| dts = ST_DTYPE[f.get_tensor(name).dtype][0] |
| resolved.append((name, dts, shape, kind, srcref)) |
|
|
| header, off = {}, 0 |
| for name, dts, shape, kind, srcref in resolved: |
| n = 1 |
| for d in shape: |
| n *= d |
| nbytes = n * {"U8": 1, "I8": 1, "F8_E4M3": 1, "F8_E8M0": 1, |
| "BF16": 2, "F16": 2, "F32": 4, "I32": 4, "I64": 8}[dts] |
| header[name] = {"dtype": dts, "shape": shape, "data_offsets": [off, off + nbytes]} |
| off += nbytes |
| header["__metadata__"] = {"format": "pt"} |
| blob = json.dumps(header).encode("utf-8") |
| pad = (-len(blob)) % 8 |
| blob += b" " * pad |
|
|
| scale_buf = {} |
| with open(dst_path, "wb", buffering=1 << 22) as out: |
| out.write(struct.pack("<Q", len(blob))) |
| out.write(blob) |
| for name, dts, shape, kind, srcref in resolved: |
| if kind == "stream": |
| wname, sname = srcref |
| base = name[:-len(".weight")] |
| rows, cols = f.get_slice(wname).get_shape() |
| nblk = cols // SRC_BLOCK |
| sb = torch.empty((rows, nblk), dtype=torch.float8_e8m0fnu) |
| wsl, ssl = f.get_slice(wname), f.get_slice(sname) |
| t0 = time.time() |
| 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) |
| ns = torch.pow(torch.tensor(2.0, device=v.device), exp) |
| codes = quantize_e2m1_codes(v / ns) |
| out.write(pack_e2m1(codes.view(n, cols)).to("cpu") |
| .flatten().view(torch.uint8).numpy().tobytes()) |
| sb[lo:hi] = ns.squeeze(-1).to(torch.float8_e8m0fnu).to("cpu") |
|
|
| k = min(4096, n) |
| a = v[:k].flatten(1) |
| b = (FP4_TABLE.to(v.device)[codes[:k].long()].view(k, nblk, SRC_BLOCK) |
| * ns[:k]).flatten(1) |
| num = (a * b).sum(1); den = a.norm(dim=1) * b.norm(dim=1) |
| ok = den > 0 |
| if ok.any(): |
| stats["cos_sum"] += float((num[ok] / den[ok]).sum()) |
| stats["cos_n"] += int(ok.sum()) |
| del v, codes, ns, amax, exp, a, b, num, den, ok |
| print(f" {base}: {hi}/{rows} ({100*hi/rows:.1f}%) " |
| f"cos={stats['cos_sum']/max(stats['cos_n'],1):.6f} " |
| f"[{time.time()-t0:.0f}s]", flush=True) |
| scale_buf[base + ".scale"] = sb |
| stats["engram"] += 1 |
| elif kind == "scale": |
| out.write(scale_buf.pop(srcref + ".scale") |
| .flatten().view(torch.uint8).numpy().tobytes()) |
| else: |
| t = f.get_tensor(srcref).contiguous() |
| |
| |
| out.write(t.flatten().view(torch.uint8).numpy().tobytes()) |
| stats["pass"] += 1 |
| return stats |
|
|
|
|
| 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 |
| 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 |
| has_engram = any(ENGRAM_RE.match(k) for k, v in index["weight_map"].items() if v == s) |
| if has_engram and do_g: |
| st = write_engram_shard(sp, os.path.join(args.out, s), device=DEVICE) |
| stats["engram"] += st["engram"]; stats["pass"] += st["pass"] |
| stats["engram_cos_sum"] += st["cos_sum"]; stats["engram_cos_n"] += st["cos_n"] |
| print(f" {s} [streamed engram shard] engram={st['engram']} pass={st['pass']} " |
| f"cos={st['cos_sum']/max(st['cos_n'],1):.6f} [{time.time()-t0:.0f}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() |
|
|
