dsv41_fp4_stream.py

#!/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(w_e4m3: torch.Tensor, scale_e8m0: torch.Tensor):
    """(e4m3 [rows,256], e8m0 [rows,8]) -> (uint8 [rows,128], e8m0 [rows,8])

    Keeps the source's own block-32 / E8M0 layout, so the table is stored exactly the
    way DeepSeek stores an expert weight.  Scales are recomputed from the dequantized
    values (a power of two >= amax/6) rather than reused, so no block clips.
    """
    rows, cols = w_e4m3.shape
    nblk = cols // SRC_BLOCK
    assert scale_e8m0.shape == (rows, nblk)

    v = w_e4m3.float().view(rows, nblk, SRC_BLOCK) * scale_e8m0.float().unsqueeze(-1)
    amax = v.abs().amax(dim=-1, keepdim=True)
    # power-of-two scale, rounded UP so amax/scale <= 6 and nothing clips
    exp = torch.ceil(torch.log2((amax / FP4_MAX).clamp(min=1e-38)))
    exp = exp.clamp(min=-127, max=127)
    new_scale = torch.pow(torch.tensor(2.0), exp)
    codes = quantize_e2m1_codes(v / new_scale)
    packed = pack_e2m1(codes.view(rows, cols))
    return packed, new_scale.squeeze(-1).to(torch.float8_e8m0fnu), v, codes, new_scale


def engram_cosine(v, codes, new_scale, rows, nblk):
    """Mean cosine similarity per row between original and reconstructed values."""
    rec = FP4_TABLE[codes.long()].view(rows, nblk, SRC_BLOCK) * new_scale
    a = v.flatten(1); b = rec.flatten(1)
    num = (a * b).sum(1)
    den = a.norm(dim=1) * b.norm(dim=1)
    ok = den > 0
    return float((num[ok] / den[ok]).mean()) if ok.any() else 1.0


# ------------------------------------------------------------------ shard loop

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"
                w = f.get_tensor(name); s = f.get_tensor(sname)
                packed, new_scale, v, codes, ns = quantize_engram_to_fp4(w, s)
                out[base + ".weight"] = packed
                out[base + ".scale"] = new_scale
                stats["engram"] += 1
                stats["engram_cos"].append(
                    engram_cosine(v[:4096], codes[:4096], ns[:4096], min(4096, v.shape[0]), v.shape[1]))
                del v, codes, ns
            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")
    args = ap.parse_args()

    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": []}
    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={sum(stats['engram_cos'])/len(stats['engram_cos']):.5f}" if stats["engram_cos"] 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"]:
        print(f"[engram] mean cosine {sum(stats['engram_cos'])/len(stats['engram_cos']):.6f}")


if __name__ == "__main__":
    main()