Final card and build scripts: measured numbers, Engram caveat, pruning analysis

finalize_repo.py

@@ -0,0 +1,95 @@
+#!/usr/bin/env python3
+"""Write config.json + copy the auxiliary files for the re-packed checkpoint.
+
+The weight re-packer only emits safetensors and an index. This adds everything else a
+consumer needs, and -- importantly -- writes a quantization_config that DESCRIBES what we
+actually did, because no released engine reads `deepseek_v41` yet and we are therefore
+defining the convention rather than matching one.
+
+Changes vs upstream config.json:
+  expert_dtype        "fp4"  -> "nvfp4"     experts are now block-16 E2M1 + E4M3 scale
+                                            + FP32 per-tensor global scale
+  expert_block_size   (new)  16
+  expert_scale_fmt    (new)  "e4m3"
+  engram_dtype        (new)  "fp4"          Engram tables are now E2M1
+  engram_block_size   (new)  32             keeping the source's own block/scale layout
+  engram_scale_fmt    (new)  "ue8m0"
+Everything else (dense fp8 block 32x32 / ue8m0) is untouched and left exactly as shipped.
+"""
+import glob, json, os, shutil, struct, sys
+
+src, out = sys.argv[1], sys.argv[2]
+
+
+def rebuild_index(out_dir):
+    """Rebuild model.safetensors.index.json by reading every output shard's header.
+
+    The re-packer is run in more than one pass (experts, then the two 94.5 GiB Engram
+    shards), and each pass writes an index covering only its own shards -- so the last
+    pass would otherwise clobber the rest. Reading the headers back is also a real check
+    that every shard on disk is a valid safetensors file.
+    """
+    DT = {"I8": 1, "U8": 1, "F8_E4M3": 1, "F8_E8M0": 1, "BF16": 2, "F16": 2,
+          "F32": 4, "F64": 8, "I32": 4, "I64": 8, "BOOL": 1}
+    wmap, total = {}, 0
+    shards = sorted(glob.glob(os.path.join(out_dir, "*.safetensors")))
+    for path in shards:
+        with open(path, "rb") as fh:
+            n = struct.unpack("<Q", fh.read(8))[0]
+            hdr = json.loads(fh.read(n))
+        base = os.path.basename(path)
+        for k, v in hdr.items():
+            if k == "__metadata__":
+                continue
+            assert k not in wmap, f"tensor {k} appears in two shards"
+            wmap[k] = base
+            cnt = 1
+            for d in v["shape"]:
+                cnt *= d
+            total += cnt * DT[v["dtype"]]
+    return {"metadata": {"total_size": total}, "weight_map": wmap}, len(shards), total
+
+
+idx, n_shards, total = rebuild_index(out)
+json.dump(idx, open(os.path.join(out, "model.safetensors.index.json"), "w"), indent=1)
+print(f"rebuilt index from {n_shards} shards: {len(idx['weight_map']):,} tensors, "
+      f"{total/2**30:.1f} GiB")
+
+cfg = json.load(open(os.path.join(src, "config.json")))
+q = dict(cfg.get("quantization_config", {}))
+assert q.get("expert_dtype") == "fp4", f"unexpected upstream quantization_config: {q}"
+q.update({
+    "expert_dtype": "nvfp4",
+    "expert_block_size": 16,
+    "expert_scale_fmt": "e4m3",
+    "expert_global_scale": True,
+    "engram_dtype": "fp4",
+    "engram_block_size": 32,
+    "engram_scale_fmt": "ue8m0",
+    "repacked_by": "LibertAI/dsv41_fp4_stream.py",
+})
+cfg["quantization_config"] = q
+json.dump(cfg, open(os.path.join(out, "config.json"), "w"), indent=2)
+print("wrote config.json:", json.dumps(q))
+
+FILES = ["tokenizer.json", "tokenizer_config.json", "LICENSE",
+         "DeepSeek_V41_Tech_Report.pdf"]
+DIRS = ["encoding", "inference"]
+for fn in FILES:
+    p = os.path.join(src, fn)
+    if os.path.exists(p):
+        shutil.copy2(p, os.path.join(out, fn)); print("  copied", fn)
+for d in DIRS:
+    p = os.path.join(src, d)
+    if os.path.isdir(p):
+        shutil.copytree(p, os.path.join(out, d), dirs_exist_ok=True); print("  copied", d + "/")
+
+# sanity: index must cover every tensor the source had, minus the scale tensors we folded
+si = json.load(open(os.path.join(src, "model.safetensors.index.json")))["weight_map"]
+oi = json.load(open(os.path.join(out, "model.safetensors.index.json")))["weight_map"]
+print(f"\nsource tensors {len(si):,} -> output tensors {len(oi):,}")
+missing = [k for k in si if k not in oi and not k.endswith(".scale")]
+extra = [k for k in oi if k not in si and not k.endswith((".weight_scale", ".weight_scale_2"))]
+print(f"missing (non-scale): {len(missing)}  {missing[:3]}")
+print(f"unexpected new:      {len(extra)}  {extra[:3]}")
+assert not missing, "output is missing tensors the source had"