File size: 6,170 Bytes
e9f8aad 9d3ec6c e9f8aad 4abd289 e9f8aad 9d3ec6c e9f8aad 4abd289 e9f8aad 4abd289 e9f8aad 4abd289 e9f8aad 4abd289 e9f8aad 4abd289 e9f8aad | 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 | """One process per Core ML chunk (workaround: several long-context chunks in one process crash the Neural Engine runtime
on the second request on this Mac, while each chunk alone, or all on the GPU, runs fine). The parent runs the host side
(embeddings, masks, pointer head) and passes tensors along the chain through pipes.
WorkerPipeline(export_dir, chunk_dir, L, units).probs(enc)"""
import glob
import os
import multiprocessing as mp
import re
import shutil
import sys
import numpy as np
import torch
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) # d1a_coreml.py sits next to this file
BUCKETS = tuple(int(x) for x in os.environ.get("D1A_COREML_BUCKETS", "512").split(",")) # each length is its own ANE program; fewer keeps the ANE within its memory
def compiled_chunks(chunk_dir):
"""-> [(first layer, end layer, .mlmodelc, its .mlpackage)]. A released bundle's coreml/L512_<a>-<b>.mlpackage are
compiled once, each as a one-function package (L512), into coreml/.compiled/; a build tree's compiled/*.mlmodelc
(multifunction, one function per length) are used as they are."""
import coremltools as ct
out = []
for f in sorted(glob.glob(f"{chunk_dir}/*.mlmodelc")):
a, b = map(int, re.search(r"(\d+)-(\d+)\.mlmodelc$", f).groups())
out.append((a, b, f, f.replace("compiled/", "buckets/").replace(".mlmodelc", ".mlpackage")))
if out:
return out
os.makedirs(f"{chunk_dir}/.compiled", exist_ok=True)
for src in sorted(glob.glob(f"{chunk_dir}/L*_*.mlpackage")):
L, a, b = map(int, re.search(r"L(\d+)_(\d+)-(\d+)\.mlpackage$", src).groups())
dst = f"{chunk_dir}/.compiled/{a:02d}-{b:02d}.mlmodelc"
if not os.path.exists(dst):
d = ct.utils.MultiFunctionDescriptor(); d.add_function(src, src_function_name="main", target_function_name=f"L{L}")
d.default_function_name = f"L{L}"; tmp = f"{chunk_dir}/.compiled/{a:02d}-{b:02d}.mlpackage"
ct.utils.save_multifunction(d, tmp); shutil.move(ct.models.utils.compile_model(tmp), dst); shutil.rmtree(tmp)
out.append((a, b, dst, src))
return out
def _worker(path, spec_path, units, conn):
import coremltools as ct
cu = getattr(ct.ComputeUnit, units)
# a precompiled .mlmodelc (compiled once by coremltools.models.utils.compile_model): loading each length function
# from it never recompiles, where MLModel(.mlpackage) compiles the whole package again for every function
models = {L: ct.models.CompiledMLModel(path, compute_units=cu, function_name=f"L{L}") for L in BUCKETS}
spec = ct.utils.load_spec(spec_path)
fns = {f.name: f for f in spec.description.functions}
fn = fns.get(f"L{BUCKETS[0]}") or spec.description # a one-function package keeps its inputs in the description
conn.send(([i.name for i in fn.input], [o.name for o in fn.output]))
while True:
msg = conn.recv()
if msg is None:
break
L, feed = msg
import time as _t
t0 = _t.perf_counter()
out = models[L].predict(feed)
t1 = _t.perf_counter()
res = {k: np.array(v, dtype=np.float16 if v.dtype == np.float32 else v.dtype, copy=True) for k, v in out.items()}
res["_predict_ms"] = np.array((t1 - t0) * 1000)
conn.send(res)
class WorkerPipeline:
def __init__(self, export_dir, chunk_dir, L=None, units="CPU_AND_NE"):
from d1a_coreml import Export, HostEmbed
self.ex = Export(export_dir)
self.host = HostEmbed(self.ex)
self.hq, self.hk = self.ex.head()
self.T, self.L = self.ex.d1a["temperature"], L
ctx = mp.get_context("spawn")
self.chunks = []
for a, b, f, spec_path in compiled_chunks(chunk_dir):
parent, child = ctx.Pipe()
p = ctx.Process(target=_worker, args=(f, spec_path, units, child), daemon=True)
p.start()
self.chunks.append((a, b, p, parent))
self.io = getattr(self, "io", []) + [parent.recv()] # one at a time: concurrent first loads race in Core ML's compiler
def close(self):
for _, _, p, conn in self.chunks:
conn.send(None); p.join(5)
def hidden(self, enc):
from d1a_coreml import masks_for
self.last_predict_ms = []
n = len(enc["ids"])
L = next(b for b in BUCKETS if b >= n) # the smallest compiled length that fits
ids = list(enc["ids"]) + [self.ex.d1a.get("pad_id", 0)] * (L - n)
emb, ple = self.host(ids)
mf, ms, pos = masks_for(enc, self.ex.tc.sliding_window, L=L)
picks = list(enc["decide_idx"]) + [o for oo in enc["opt_idx"] for o in oo]
pk = np.zeros(128, np.int32); pk[: len(picks)] = picks
common = {"pos": pos.numpy().astype(np.int32), "mask_full": mf.clamp(min=-1e4).numpy().astype(np.float16),
"mask_slide": ms.clamp(min=-1e4).numpy().astype(np.float16)}
state = {"emb": emb.numpy().astype(np.float16), "ple_tok": ple.numpy().astype(np.float16), "picks": pk}
for (a, b, p, conn), (ins, outs) in zip(self.chunks, self.io):
conn.send((L, {k: (common[k] if k in common else state[k]) for k in ins}))
out = conn.recv()
self.last_predict_ms = getattr(self, "last_predict_ms", []) + [float(out.pop("_predict_ms"))]
if "hidden" in out:
return torch.tensor(out["hidden"][: len(picks)].astype(np.float32))
state["h_in"] = out["h"]
if "ple_out" in out:
state["ple"] = out["ple_out"]
for k, v in out.items():
if k.endswith("_out") and k != "ple_out":
state[k[: -len("_out")]] = v
raise RuntimeError("no chunk returned hidden")
def probs(self, enc):
h = self.hidden(enc)
Q = len(enc["decide_idx"])
out, j = [], Q
with torch.no_grad():
for qi, oo in enumerate(enc["opt_idx"]):
z = (self.hk(h[j: j + len(oo)]) @ self.hq(h[qi])) / (self.hq.out_features ** 0.5)
out.append(torch.softmax(z / self.T, -1).numpy()); j += len(oo)
return out
|