1#!/usr/bin/env python32"""Convert a checkpoint's expert tensors into the TITAN expert-contiguous layout: one 4 KiB-aligned blob per (layer, expert)3holding its tensors back to back, so that the SSD tier reads an expert with a single pread. Dense tensors stay in the4original safetensors. Output: <out_dir>/experts.titan (blob) + <out_dir>/experts.titan.json (index). Never modifies the source.5"""6import argparse, hashlib, json, os, sys, time7sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))8from titan.registry import load_spec910ALIGN = 40961112def main():13 ap = argparse.ArgumentParser()14 ap.add_argument("model_dir")15 ap.add_argument("--out", default=None, help="default: ~/.cache/titan/<sha1 of realpath>/")16 a = ap.parse_args()17 spec = load_spec(a.model_dir)18 out = a.out or os.path.expanduser(f"~/.cache/titan/{hashlib.sha1(os.path.realpath(a.model_dir).encode()).hexdigest()[:12]}")19 os.makedirs(out, exist_ok=True)20 blob = os.path.join(out, "experts.titan")21 index = {"source": os.path.realpath(a.model_dir), "model_type": spec.model_type, "align": ALIGN, "experts": {}}22 fds = {}23 t0 = time.perf_counter(); written = 024 with open(blob, "wb") as w:25 pos = 026 for l in range(spec.num_layers):27 for e in range(spec.num_experts):28 es = spec.experts[(l, e)]29 parts = {}; off = 030 chunks = []31 for name, loc in es.tensors.items(): # fixed order: as in the registry (gate, up, down × weight, scales, biases)32 fd = fds.get(loc.file) or fds.setdefault(loc.file, os.open(loc.file, os.O_RDONLY))33 b = os.pread(fd, loc.nbytes, loc.offset)34 assert len(b) == loc.nbytes35 parts[name] = {"off": off, "n": loc.nbytes, "dtype": loc.dtype, "shape": list(loc.shape)}36 chunks.append(b); off += loc.nbytes37 data = b"".join(chunks)38 pad = (-len(data)) % ALIGN39 w.write(data); w.write(b"\0" * pad)40 index["experts"][f"{l},{e}"] = {"offset": pos, "nbytes": len(data), "tensors": parts}41 pos += len(data) + pad; written += len(data)42 print(f"layer {l} done, {pos/1e9:.2f} GB", file=sys.stderr)43 index["blob_bytes"] = pos44 json.dump(index, open(blob + ".json", "w"))45 print(json.dumps({"out": out, "experts": len(index["experts"]), "bytes": written, "blob_bytes": pos, "seconds": time.perf_counter() - t0}))4647if __name__ == "__main__":48 main()49