File size: 3,984 Bytes
5e32a69 | 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 | """d3 text requests on Core ML (Python host), for the checkpoints whose text graph takes embeddings (d3-mini).
Upstream's own request preparation runs on a weightless (meta-device) d3, so prompts and token ids are exactly
upstream's; Core ML runs the text graph per question (host embedding gather, text positions on all three M-RoPE axes);
FP32 readout on the host.
m = D3CoreMLText("<d3 checkpoint>", "<coreml dir>", text_buckets=[256, 640])
m.system_one(state=..., questions={...})"""
import sys
import time
from pathlib import Path
import coremltools as ct
import numpy as np
import torch
HERE = Path(__file__).parent
class D3CoreMLText:
def __init__(self, root, coreml_dir, text_buckets=None, compute_units=ct.ComputeUnit.CPU_AND_GPU):
root, cdir = Path(root), Path(coreml_dir)
sys.path.insert(0, str(root))
sys.path.insert(0, str(HERE.parent))
import d3_runtime as up
from d3_graph import load_text_config
from qwen35_export import rope_cos_sin
from safetensors.torch import load_file
self.up, self.rope_cos_sin = up, rope_cos_sin
self.prep = up.D3.from_pretrained(str(root), device="meta", verify="none") # no weights: prompts only
self.temperature = self.prep.temperature
self.tcfg = load_text_config(str(root))
self.readout = load_file(str(root / "readout.safetensors"))["weight"].float().numpy()
self.embed = np.memmap(cdir / "embeddings.f16", dtype=np.float16, mode="r").reshape(-1, self.tcfg.hidden_size)
self.text = {int(p.stem[4:]): ct.models.MLModel(str(p), compute_units=compute_units)
for p in sorted(cdir.glob("d3_E*.mlpackage")) if not text_buckets or int(p.stem[4:]) in text_buckets}
self.pad = self.prep.pad_id
self.model_name = root.name + "-coreml"
self.timing = []
def warm(self):
for L, model in self.text.items():
for _ in range(2):
self._final(list(range(10)), L)
def _final(self, ids, L):
n = len(ids)
x = np.zeros((1, L, self.tcfg.hidden_size), dtype=np.float16)
x[0, :n] = self.embed[np.asarray(ids)]
x[0, n:] = self.embed[self.pad]
pos = torch.arange(L).clamp(max=n - 1)[None].expand(3, -1)
cos, sin = self.rope_cos_sin(self.tcfg, pos)
valid = np.zeros(L, dtype=np.float16); valid[:n] = 1
last = np.zeros((1, L), dtype=np.float16); last[0, n - 1] = 1
return self.text[L].predict({"hidden": x, "cos": cos.numpy().astype(np.float16),
"sin": sin.numpy().astype(np.float16), "key_valid": valid,
"last_onehot": last})["final"].reshape(-1)
def system_one(self, *, state, questions):
prepared = self.prep.prepare(state, questions)
answers, used, self.timing = {}, 0, []
for key in prepared.keys:
if key in prepared.errors:
answers[key] = prepared.errors[key]
continue
ids, question = prepared.sequences[key], prepared.questions[key]
L = next((b for b in sorted(self.text) if b >= len(ids)), None)
if L is None:
answers[key] = {"type": question.kind, "error": "max_length_exceeded",
"message": f"{len(ids)} tokens exceed the largest Core ML bucket {max(self.text)}"}
continue
t = time.perf_counter()
h = self._final(ids, L)
self.timing.append((len(ids), L, round(time.perf_counter() - t, 3)))
logits = self.readout[:len(question.keys)].astype(np.float64) @ h.astype(np.float64) / self.temperature
p = np.exp(logits - logits.max())
answers[key] = self.up.product_answer(question, (p / p.sum()).tolist())
used += len(ids)
return {"model": self.model_name, "answers": answers, "usage": {"input_tokens": used, "output_tokens": 0}}
|