"""Phase B: dump MTP training triples from the trunk over the corpus. For each corpus record, renders the chat template over (messages + completion), runs the quantized trunk once, and stores per position t: hidden[t] post-norm_f trunk hidden (bf16) -- the MTP head input token[t+1] next token id -- fused with hidden[t] token[t+2] hard target topk_logits[t+1] trunk's top-K logits at t+1 -- soft target The soft target is the point: the trunk's distribution at position t+1 IS P(t+2 | prefix, t+1) — exactly what the head should predict given (hidden[t], token[t+1]). Training with KL to this aligns the head to the trunk's conditionals regardless of who authored the text; the corpus only controls state coverage. Loss = CE(hard) + lambda * KL(soft), lambda ~1. Storage: one .npz per corpus shard chunk in triples/: hiddens (N, H) bf16 saved as uint16 view tokens (N+2,) uint32 (positions align: hidden[i] pairs tokens[i+1], tokens[i+2]) topk_ids (N, K) uint32; topk_lp (N, K) float16 (K=64, log-softmaxed) doc_bounds: (n_docs, 2) start/end into N (docs are independent sequences) Only completion-region positions are dumped (prompt positions carry off-distribution states for the head's serving regime, where it drafts during generation). A position budget per doc caps degenerate long docs. MEMORY: loads the full quantized model (~81 GB) via mlx_lm. Run while the omlx server is idle or stopped if headroom is tight; peak ~90 GB + activations. Usage: python3 dump_triples.py [--corpus-glob 'corpus/*.jsonl'] [--out triples] [--topk 64] [--max-docs N] [--resume] --resume continues after whatever triples-*.npz already exist (numbering and doc position carry on from them). Without it, a non-empty --out is an error rather than a silent overwrite. """ import argparse import glob as globmod import json import time from pathlib import Path import mlx.core as mx import numpy as np MODEL_DIR = "/Users/david/AI/kenosistron3-oq5e-mtp" ROOT = Path(__file__).parent CHUNK_TOKENS = 4096 # prefill chunk (fits the SSD kernel fast path) FLUSH_POSITIONS = 200_000 # positions per output .npz (~1.6 GB hiddens) def load_model(): # omlx's oq loader handles the oQ5e quantization; mlx_lm.load handles the # rest of the architecture. Try mlx_lm first; fall back to omlx. try: from mlx_lm import load model, tokenizer = load(MODEL_DIR) return model, tokenizer except Exception as e: raise SystemExit( f"mlx_lm.load failed ({e}); if this is an oQ-format issue, run " "through omlx's loader instead (omlx.utils.model_loading)." ) def render(tokenizer, messages, completion): """Token ids for the full exchange + index where the completion starts.""" prompt_ids = tokenizer.apply_chat_template( messages, add_generation_prompt=True, tokenize=True ) comp_ids = tokenizer.encode(completion, add_special_tokens=False) return list(prompt_ids) + list(comp_ids), len(prompt_ids) def trunk_hiddens(model, ids): """Post-norm_f hiddens + logits for a full sequence, chunked prefill.""" cache = model.make_cache() hiddens = [] logits = [] x = mx.array(ids)[None] for s0 in range(0, x.shape[1], CHUNK_TOKENS): chunk = x[:, s0 : s0 + CHUNK_TOKENS] lg, hd = model(chunk, cache=cache, return_hidden=True) mx.eval(lg, hd) hiddens.append(hd[0]) logits.append(lg[0]) return mx.concatenate(hiddens), mx.concatenate(logits) def main(): ap = argparse.ArgumentParser() ap.add_argument("--corpus-glob", default=str(ROOT / "corpus" / "*.jsonl")) ap.add_argument("--out", default=str(ROOT / "triples")) ap.add_argument("--topk", type=int, default=64) ap.add_argument("--max-docs", type=int, default=None) ap.add_argument("--max-positions-per-doc", type=int, default=1600) ap.add_argument( "--resume", action="store_true", help="continue after the existing triples-*.npz instead of overwriting", ) args = ap.parse_args() out = Path(args.out) out.mkdir(exist_ok=True) # The MTP-side patches must be active so return_hidden exists. from omlx.patches.mlx_lm_mtp import apply_mlx_lm_mtp_patch apply_mlx_lm_mtp_patch() model, tokenizer = load_model() docs = [] for f in sorted(globmod.glob(args.corpus_glob)): for line in open(f): docs.append(json.loads(line)) if args.max_docs: docs = docs[: args.max_docs] print(f"{len(docs)} corpus docs") existing = sorted(out.glob("triples-*.npz")) start_di, docs_done, shard_start = 0, 0, 0 if existing: if not args.resume: raise SystemExit( f"{len(existing)} shard(s) already in {out}; pass --resume to " "continue after them (re-running without it overwrites from 0000)." ) shard_start = len(existing) for f in existing: with np.load(f) as d: docs_done += d["doc_bounds"].shape[0] # Replay the filters below to find where those docs ended. They depend # only on render() and lengths -- not the trunk -- so the fast-forward # is exact and costs no model forwards. seen, start_di = 0, len(docs) for di, doc in enumerate(docs): if seen >= docs_done: start_di = di break ids, comp_start = render(tokenizer, doc["messages"], doc["completion"]) if len(ids) - comp_start < 8: continue lo = max(comp_start - 1, 0) hi = min(len(ids) - 2, lo + args.max_positions_per_doc) if hi <= lo: continue seen += 1 if seen < docs_done: raise SystemExit( f"corpus yields only {seen} usable docs but shards hold " f"{docs_done}; corpus-glob likely differs from the first run." ) print( f"resume: {docs_done} docs in {shard_start} shard(s), " f"restarting at corpus index {start_di}", flush=True, ) buf_h, buf_tok, buf_ids, buf_lp, bounds = [], [], [], [], [] n_pos, shard_i, t0 = 0, shard_start, time.time() def flush(): nonlocal buf_h, buf_tok, buf_ids, buf_lp, bounds, shard_i if not buf_h: return np.savez( out / f"triples-{shard_i:04d}.npz", hiddens=np.concatenate(buf_h), tokens=np.concatenate(buf_tok), topk_ids=np.concatenate(buf_ids), topk_lp=np.concatenate(buf_lp), doc_bounds=np.array(bounds, dtype=np.int64), ) print(f" wrote triples-{shard_i:04d}.npz") shard_i += 1 buf_h, buf_tok, buf_ids, buf_lp, bounds = [], [], [], [], [] written = docs_done for di, doc in enumerate(docs): if di < start_di: continue ids, comp_start = render(tokenizer, doc["messages"], doc["completion"]) if len(ids) - comp_start < 8: continue h, lg = trunk_hiddens(model, ids) # positions t in the completion region with t+2 in range lo = max(comp_start - 1, 0) hi = min(len(ids) - 2, lo + args.max_positions_per_doc) if hi <= lo: continue hs = h[lo:hi] # (n, H) lg1 = lg[lo + 1 : hi + 1].astype(mx.float32) # trunk logits at t+1 lp1 = lg1 - mx.logsumexp(lg1, axis=-1, keepdims=True) kidx = mx.argpartition(-lp1, kth=args.topk - 1, axis=-1)[:, : args.topk] klp = mx.take_along_axis(lp1, kidx, axis=-1) mx.eval(hs, kidx, klp) start = sum(x.shape[0] for x in buf_h) buf_h.append(np.array(hs.astype(mx.bfloat16).view(mx.uint16))) buf_tok.append(np.array(ids[lo : hi + 2], dtype=np.uint32)) buf_ids.append(np.array(kidx.astype(mx.uint32))) buf_lp.append(np.array(klp.astype(mx.float16))) bounds.append((start, start + (hi - lo))) n_pos += hi - lo written += 1 if n_pos >= FLUSH_POSITIONS: flush() n_pos = 0 if written % 50 == 0: dt = time.time() - t0 print(f"[{dt/60:5.1f}m] {written}/{len(docs)} docs", flush=True) flush() print(f"DONE: {written} docs in {(time.time()-t0)/60:.1f} min -> {out}") if __name__ == "__main__": main()