1"""Model and expert registry.23Parses a Hugging Face / MLX MoE checkpoint directory (config.json + safetensors) into an explicit4description of layers, experts, routing semantics and the *byte ranges* of every expert tensor,5so that the expert store can read a single expert from disk with `pread` without touching the rest6of the file. Only formats that we have verified are accepted (see `SUPPORTED`).7"""8from __future__ import annotations910import json11import os12import struct13from dataclasses import dataclass, field14from typing import Dict, List, Tuple1516DTYPE_ITEMSIZE = {"F32": 4, "F16": 2, "BF16": 2, "U32": 4, "I32": 4, "U8": 1, "I8": 1, "U16": 2}1718# model_type -> (expert tensor key template, matrices, routing description)19SUPPORTED = {20 "olmoe": {21 "expert_key": "model.layers.{l}.mlp.experts.{e}.{m}_proj.{s}",22 "stacked_key": "model.layers.{l}.mlp.switch_mlp.{m}_proj.{s}",23 "matrices": ("gate", "up", "down"),24 "router_key": "model.layers.{l}.mlp.gate.{s}",25 "routing": "softmax_then_topk", # softmax over all experts, then top-k, weights = softmax probs (renormalised iff norm_topk_prob)26 "shared_experts": 0,27 },28 "qwen3_moe": {29 "expert_key": "model.layers.{l}.mlp.experts.{e}.{m}_proj.{s}",30 "stacked_key": "model.layers.{l}.mlp.switch_mlp.{m}_proj.{s}",31 "matrices": ("gate", "up", "down"),32 "router_key": "model.layers.{l}.mlp.gate.{s}",33 "routing": "softmax_then_topk",34 "shared_experts": 0,35 },36}373839@dataclass(frozen=True)40class TensorLoc:41 key: str42 dtype: str43 shape: Tuple[int, ...]44 file: str45 offset: int # absolute byte offset in file46 nbytes: int474849@dataclass50class ExpertSpec:51 layer: int52 expert: int53 tensors: Dict[str, TensorLoc] # e.g. "gate.weight" -> loc5455 @property56 def nbytes(self) -> int:57 return sum(t.nbytes for t in self.tensors.values())585960@dataclass61class ModelSpec:62 path: str63 model_type: str64 num_layers: int65 num_experts: int66 top_k: int67 norm_topk_prob: bool68 hidden_size: int69 intermediate_size: int70 quantization: Dict[str, int]71 routing: str72 shared_experts: int73 experts: Dict[Tuple[int, int], ExpertSpec] = field(default_factory=dict)74 dense: Dict[str, TensorLoc] = field(default_factory=dict)75 layout: str = "per_expert" # or "stacked"7677 @property78 def bytes_per_expert(self) -> int:79 return next(iter(self.experts.values())).nbytes8081 @property82 def expert_bytes_total(self) -> int:83 return sum(e.nbytes for e in self.experts.values())8485 @property86 def dense_bytes_total(self) -> int:87 return sum(t.nbytes for t in self.dense.values())8889 def summary(self) -> dict:90 return {91 "path": self.path,92 "model_type": self.model_type,93 "num_layers": self.num_layers,94 "num_experts": self.num_experts,95 "top_k": self.top_k,96 "norm_topk_prob": self.norm_topk_prob,97 "hidden_size": self.hidden_size,98 "intermediate_size": self.intermediate_size,99 "quantization": self.quantization,100 "routing": self.routing,101 "shared_experts": self.shared_experts,102 "layout": self.layout,103 "bytes_per_expert": self.bytes_per_expert,104 "expert_bytes_total": self.expert_bytes_total,105 "dense_bytes_total": self.dense_bytes_total,106 "experts_per_token": self.num_layers * self.top_k,107 "bytes_per_token_if_uncached": self.num_layers * self.top_k * self.bytes_per_expert,108 }109110111def read_safetensors_header(path: str) -> Dict[str, TensorLoc]:112 """Parse the safetensors header and return absolute byte locations for every tensor."""113 locs = {}114 with open(path, "rb") as f:115 (n,) = struct.unpack("<Q", f.read(8))116 header = json.loads(f.read(n))117 base = 8 + n118 for key, meta in header.items():119 if key == "__metadata__":120 continue121 a, b = meta["data_offsets"]122 locs[key] = TensorLoc(key, meta["dtype"], tuple(meta["shape"]), path, base + a, b - a)123 return locs124125126def load_spec(model_dir: str) -> ModelSpec:127 model_dir = os.path.realpath(model_dir)128 with open(os.path.join(model_dir, "config.json")) as f:129 cfg = json.load(f)130 mt = cfg["model_type"]131 if mt not in SUPPORTED:132 raise ValueError(f"model_type {mt!r} is not in the verified set {sorted(SUPPORTED)}")133 fmt = SUPPORTED[mt]134 locs: Dict[str, TensorLoc] = {}135 for fn in sorted(os.listdir(model_dir)):136 if fn.endswith(".safetensors"):137 locs.update(read_safetensors_header(os.path.join(model_dir, fn)))138 q = cfg.get("quantization") or cfg.get("quantization_config") or {}139 spec = ModelSpec(140 path=model_dir, model_type=mt, num_layers=cfg["num_hidden_layers"], num_experts=cfg["num_experts"],141 top_k=cfg["num_experts_per_tok"], norm_topk_prob=bool(cfg.get("norm_topk_prob", False)),142 hidden_size=cfg["hidden_size"], intermediate_size=cfg["intermediate_size"],143 quantization={"group_size": q.get("group_size"), "bits": q.get("bits")}, routing=fmt["routing"],144 shared_experts=fmt["shared_experts"],145 )146 per_expert_keys = set()147 for l in range(spec.num_layers):148 for e in range(spec.num_experts):149 tensors = {}150 for m in fmt["matrices"]:151 for s in ("weight", "scales", "biases"):152 k = fmt["expert_key"].format(l=l, e=e, m=m, s=s)153 if k in locs:154 tensors[f"{m}.{s}"] = locs[k]155 per_expert_keys.add(k)156 if not tensors:157 # stacked layout: switch_mlp.<m>_proj.<s> has shape [E, ...]; expert e is the contiguous slice e of the first axis158 for m in fmt["matrices"]:159 for s in ("weight", "scales", "biases"):160 k = fmt["stacked_key"].format(l=l, m=m, s=s)161 if k in locs:162 loc = locs[k]163 if loc.shape[0] != spec.num_experts:164 raise ValueError(f"{k}: first axis {loc.shape[0]} != num_experts {spec.num_experts}")165 per = loc.nbytes // spec.num_experts166 tensors[f"{m}.{s}"] = TensorLoc(f"{k}[{e}]", loc.dtype, tuple(loc.shape[1:]), loc.file, loc.offset + e * per, per)167 per_expert_keys.add(k)168 spec.layout = "stacked"169 if not tensors:170 raise ValueError(f"no expert tensors found for layer {l} expert {e} (neither per-expert nor stacked keys)")171 spec.experts[(l, e)] = ExpertSpec(l, e, tensors)172 spec.dense = {k: v for k, v in locs.items() if k not in per_expert_keys}173 return spec174175176if __name__ == "__main__":177 import sys178 s = load_spec(sys.argv[1])179 print(json.dumps(s.summary(), indent=2))180