File size: 12,456 Bytes
e1acc15 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 | """Identity-checked loader for the DSpark drafter tensors that lane S2 copied to x86host.
Source of truth: ``TENSORS.json`` written by S2 next to the extracted safetensors files. Every entry records the
tensor's name, dtype, shape, the destination file and byte range, and the sha256 of the *source* byte range on the
Spark (``/srv/dsv41/legal-v2a``, hardlinked from HeNryous EXL3 2.75 bpw @ 86f85253). Loading re-hashes each
destination byte range and refuses the checkpoint on any difference, on a missing or unexpected tensor, or on a
header that disagrees with the manifest. Nothing is converted here: tensors come back as zero-copy CPU views of the
mapped files, in the checkpoint's own dtypes.
The tensor set is derived from config.json (three ``mtp.{0,1,2}`` stages, 128 drafter experts, the target's embedding
and EXL3 head), not from the manifest, so a manifest that lost or gained tensors is caught too.
"""
from __future__ import annotations
import hashlib
import json
import mmap
import os
import struct
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
import torch
DTYPES = {
"BF16": torch.bfloat16,
"F16": torch.float16,
"F32": torch.float32,
"I16": torch.int16,
"I32": torch.int32,
"F8_E4M3": torch.float8_e4m3fn,
"F8_E8M0": torch.float8_e8m0fnu,
}
EXL3_PARTS = ("trellis", "suh", "svh", "mul1")
class IdentityError(RuntimeError):
"""The files on disk are not the tensors the manifest promises."""
@dataclass(frozen=True)
class DrafterShape:
"""The numbers of config.json the worker depends on (checked, never assumed)."""
hidden: int
vocab: int
n_heads: int
head_dim: int
rope_dim: int
q_lora: int
o_lora: int
o_groups: int
inter: int
n_experts: int
top_k: int
n_stages: int
n_target_layers: int
window: int
block: int
noise_token: int
markov_rank: int
hc_mult: int
hc_iters: int
hc_eps: float
rms_eps: float
swiglu_limit: float
route_scale: float
rope_theta: float
taps: tuple[int, ...]
@classmethod
def from_config(cls, cfg: dict) -> "DrafterShape":
t = cfg["text_config"]
n_layers = int(t["num_hidden_layers"])
ratios = t["compress_ratios"]
n_stages = int(t["num_nextn_predict_layers"])
if [int(r) for r in ratios[n_layers:n_layers + n_stages]] != [0] * n_stages:
raise IdentityError(f"drafter stages are not pure sliding-window layers: {ratios[n_layers:]}")
if t.get("scoring_func") != "sqrtsoftplus" or t.get("topk_method") != "noaux_tc" or not t.get("norm_topk_prob"):
raise IdentityError("routing is not sqrtsoftplus / noaux_tc / normalized top-k")
if t.get("hidden_act") != "silu":
raise IdentityError(f"hidden_act {t.get('hidden_act')!r}")
return cls(
hidden=int(t["hidden_size"]), vocab=int(t["vocab_size"]), n_heads=int(t["num_attention_heads"]),
head_dim=int(t["head_dim"]), rope_dim=int(t["qk_rope_head_dim"]), q_lora=int(t["q_lora_rank"]),
o_lora=int(t["o_lora_rank"]), o_groups=int(t["o_groups"]), inter=int(t["moe_intermediate_size"]),
n_experts=int(t["dspark_n_routed_experts"]), top_k=int(t["dspark_num_experts_per_tok"]),
n_stages=n_stages, n_target_layers=n_layers, window=int(t["sliding_window"]),
block=int(t["dspark_block_size"]), noise_token=int(t["dspark_noise_token_id"]),
markov_rank=int(t["dspark_markov_rank"]), hc_mult=int(t["hc_mult"]), hc_iters=int(t["hc_sinkhorn_iters"]),
hc_eps=float(t["hc_eps"]), rms_eps=float(t["rms_norm_eps"]), swiglu_limit=float(t["swiglu_limit"]),
route_scale=float(t["routed_scaling_factor"]), rope_theta=float(t["rope_theta"]),
taps=tuple(int(x) for x in t["dspark_target_layer_ids"]),
)
def required_names(s: DrafterShape) -> list[str]:
"""Every tensor the worker reads (plus ffn.gate.bias_vl, checked though text drafts never use it)."""
names = ["embed.weight"] + [f"head.{p}" for p in EXL3_PARTS]
for i in range(s.n_stages):
pre = f"mtp.{i}."
for lin in ("attn.wq_a", "attn.wkv", "attn.wq_b", "attn.wo_b",
"ffn.shared_experts.w1", "ffn.shared_experts.w2", "ffn.shared_experts.w3"):
names += [f"{pre}{lin}.{p}" for p in EXL3_PARTS]
names += [f"{pre}attn.wo_a.weight", f"{pre}attn.wo_a.scale", f"{pre}attn.attn_sink",
f"{pre}attn.q_norm.weight", f"{pre}attn.kv_norm.weight", f"{pre}attn_norm.weight",
f"{pre}ffn_norm.weight", f"{pre}ffn.gate.weight", f"{pre}ffn.gate.bias", f"{pre}ffn.gate.bias_vl"]
names += [f"{pre}hc_{s_}_{k}" for s_ in ("attn", "ffn") for k in ("fn", "base", "scale")]
for e in range(s.n_experts):
for w in ("w1", "w2", "w3"):
names += [f"{pre}ffn.experts.{e}.{w}.{p}" for p in EXL3_PARTS]
names += [f"mtp.0.main_proj.{p}" for p in EXL3_PARTS] + ["mtp.0.main_norm.weight"]
last = f"mtp.{s.n_stages - 1}."
names += [f"{last}norm.weight", f"{last}markov_head.embed.weight", f"{last}markov_head.head.weight",
f"{last}confidence_head.proj.weight"]
return names
def _read_header(path: str) -> tuple[int, dict]:
with open(path, "rb") as f:
n = struct.unpack("<Q", f.read(8))[0]
if n > 64 << 20:
raise IdentityError(f"{path}: implausible safetensors header length {n}")
return 8 + n, json.loads(f.read(n))
class Checkpoint:
"""The verified drafter tensors of one S2 delivery directory. ``get(name)`` returns a CPU view (no copy)."""
def __init__(self, root: str, verify: bool = True, threads: int = 16):
self.root = root
with open(os.path.join(root, "config.json")) as f:
self.config = json.load(f)
with open(os.path.join(root, "TENSORS.json")) as f:
self.manifest = json.load(f)
with open(os.path.join(root, "exl3_moe_plan.json")) as f:
self.plan_json = json.load(f)
self.shape = DrafterShape.from_config(self.config)
self._check_quant_config()
entries = {e["name"]: e for e in self.manifest["tensors"]}
need = required_names(self.shape)
missing = [n for n in need if n not in entries]
if missing:
raise IdentityError(f"{len(missing)} drafter tensors missing from TENSORS.json, e.g. {missing[:5]}")
# norm.weight (the target's final norm) is delivered too; the drafter uses mtp.2.norm, so it is ignored
extra = sorted(set(entries) - set(need) - {"norm.weight"})
if extra:
raise IdentityError(f"unexpected tensors in TENSORS.json: {extra[:5]} ({len(extra)})")
self.entries = {n: entries[n] for n in need}
self._maps: dict[str, mmap.mmap] = {}
self._headers: dict[str, tuple[int, dict]] = {}
from .identity import IdentityMismatch, resolve_inside
for e in self.entries.values():
fn = e["destination_file"]
if fn not in self._maps:
try: # never open a file outside the delivery (S5R fix 1)
path = resolve_inside(root, fn)
except IdentityMismatch as err:
raise IdentityError(f"{e['name']}: {err}") from err
self._headers[fn] = _read_header(path)
with open(path, "rb") as f:
self._maps[fn] = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
for name, e in self.entries.items():
self._check_header(name, e)
self.verified = False
if verify:
self.verify(threads)
def _check_quant_config(self) -> None:
q = self.config.get("quantization_config", {})
x = q.get("exl3_moe", {})
if q.get("quant_method") != "deepseek_v41_exl3" or x.get("exllamav3_version") != "1.6.0":
raise IdentityError(f"quantization_config is not deepseek_v41_exl3 / exllamav3 1.6.0: {q.get('quant_method')}")
if x.get("head") != "exl3":
raise IdentityError(f"head is {x.get('head')!r}, the worker implements the EXL3 head only")
if self.manifest.get("codebook") != "mul1" or self.plan_json.get("codebook") != "mul1":
raise IdentityError("codebook is not mul1")
self.head_k = int(x["head_k"])
self.drafter_k = {k: int(v) for k, v in x["dense"]["drafter"].items()}
def _check_header(self, name: str, e: dict) -> None:
fn = e["destination_file"]
hlen, hdr = self._headers[fn]
h = hdr.get(name)
if h is None:
raise IdentityError(f"{name}: not in {fn}'s header")
if h["dtype"] != e["dtype"] or list(h["shape"]) != list(e["shape"]):
raise IdentityError(f"{name}: header {h['dtype']} {h['shape']} vs manifest {e['dtype']} {e['shape']}")
a, b = h["data_offsets"]
if [hlen + a, hlen + b] != list(e["destination_byte_range"]):
raise IdentityError(f"{name}: header range {[hlen + a, hlen + b]} vs manifest {e['destination_byte_range']}")
if e["dtype"] not in DTYPES:
raise IdentityError(f"{name}: dtype {e['dtype']} not supported")
numel = 1
for d in e["shape"]:
numel *= int(d)
if numel * DTYPES[e["dtype"]].itemsize != b - a:
raise IdentityError(f"{name}: {b - a} bytes for {e['shape']} {e['dtype']}")
sa, sb = e["source_byte_range"]
if sb - sa != b - a:
raise IdentityError(f"{name}: source range {sb - sa} B vs destination {b - a} B")
def _bytes(self, name: str) -> memoryview:
e = self.entries[name]
a, b = e["destination_byte_range"]
return memoryview(self._maps[e["destination_file"]])[a:b]
def verify(self, threads: int = 16) -> dict:
"""sha256 of every destination byte range == the manifest's source-range sha256. Raises on any mismatch."""
def one(name: str) -> tuple[str, bool]:
return name, hashlib.sha256(self._bytes(name)).hexdigest() == self.entries[name]["source_byte_range_sha256"]
# big tensors first so the pool stays busy
order = sorted(self.entries, key=lambda n: -int(self.entries[n]["resident_packed_bytes"]))
with ThreadPoolExecutor(threads) as ex:
bad = [n for n, ok in ex.map(one, order) if not ok]
if bad:
raise IdentityError(f"sha256 mismatch for {len(bad)} tensors, e.g. {bad[:5]}")
self.verified = True
return {"tensors": len(order), "bytes": sum(int(self.entries[n]["resident_packed_bytes"]) for n in order)}
def get(self, name: str) -> torch.Tensor:
"""Zero-copy CPU view of the tensor bytes in the mapped file (read-only memory; copy before writing)."""
if not self.verified:
raise IdentityError("checkpoint not verified")
e = self.entries[name]
buf = self._bytes(name)
dt = DTYPES[e["dtype"]]
if len(buf) == 0:
return torch.empty(e["shape"], dtype=dt)
import warnings
with warnings.catch_warnings():
warnings.simplefilter("ignore") # torch warns that the mmap is not writable; nothing writes it
t = torch.frombuffer(buf, dtype=torch.uint8)
return t.view(dt).view(*e["shape"]) if e["shape"] else t.view(dt).reshape(())
def exl3(self, prefix: str) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
"""(trellis int16 [in/16, out/16, 16K], suh fp16 [in], svh fp16 [out], K) of an EXL3 linear; mul1 checked."""
trellis, suh, svh, mul1 = (self.get(f"{prefix}.{p}") for p in EXL3_PARTS)
words = trellis.shape[-1]
if trellis.dim() != 3 or words % 16:
raise IdentityError(f"{prefix}.trellis shape {tuple(trellis.shape)}")
if int(mul1.reshape(-1)[0].item()) & 0xFFFFFFFF != 0x83DCD12D:
raise IdentityError(f"{prefix}.mul1 marker {int(mul1.reshape(-1)[0].item()) & 0xFFFFFFFF:#x} is not mul1")
if suh.shape[0] != trellis.shape[0] * 16 or svh.shape[0] != trellis.shape[1] * 16:
raise IdentityError(f"{prefix}: suh/svh {tuple(suh.shape)}/{tuple(svh.shape)} vs trellis {tuple(trellis.shape)}")
return trellis, suh, svh, words // 16
def close(self) -> None:
for m in self._maps.values():
m.close()
self._maps.clear()
|