File size: 4,393 Bytes
d8edfe6 | 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 | """v3.1b: next-token decoder functions n16/n32/n64/n96 -> mf_next/dec_mf_next_fp16.mlpackage (multifunction, shared weights).
Input ids int32 [1,T] = BOS + left text, RIGHT-padded with anything (causal model: pads never reach position `last`)
last int32 [1] = index of the last real token (= len(BOS + left ids) - 1)
Output lp fp32 [1,V] = log-softmax over the 16k vocab of the NEXT token after ids[0, last] (max-shifted logsumexp).
Same QDec body as the shipped t{T} functions (fp16, ane_kwargs, t16 single-block GDN, t>=32 chunked c16).
usage: build_next.py [--merge-test] (--merge-test: also merge with the shipped mf/parts/*fp16 into a throwaway package
to measure whether weights de-duplicate; prints sizes, then deletes it)"""
import sys, time, shutil, torch, numpy as np
import torch.nn.functional as F
from pathlib import Path
from qdec import QDec, ane_kwargs, FF
TS = (16, 32, 64, 96)
class QNext(QDec):
def __init__(self, T, **kw):
super().__init__(T, B=1, **kw)
self.register_buffer("pos", torch.arange(T, dtype=torch.int32)[None, :]) # [1,T]
def forward(self, ids, last):
x = F.embedding(ids, self.emb)
for m in self.layers:
h = m.n1(x)
x = x + (self.gdn(m, h) if m.kind == "linear_attention" else self.attn(m, h))
gu = m.gate_up(m.n2(x))
x = x + m.down(F.silu(gu[..., :FF]) * gu[..., FF:])
# pick row `last` BEFORE the final RMSNorm (row-wise, so exact): one-hot mask * x, sum over T
sel = (self.pos == last.to(torch.int32)[:, None]).to(x.dtype) # [1,T]
h = self.norm((sel[:, :, None] * x).sum(1)) # [1,640]
logits = h @ self.emb.t() # tied head [1,V]
mx = logits.max(-1, keepdim=True).values # max-shifted lse (ANE fp16, qdec.py)
return logits - (mx + torch.log(torch.exp(logits - mx).sum(-1, keepdim=True)))
if __name__ == "__main__":
import coremltools as ct
from coremltools.models.utils import MultiFunctionDescriptor, save_multifunction
out = Path("mf_next"); (out / "parts").mkdir(parents=True, exist_ok=True)
size = lambda p: sum(f.stat().st_size for f in Path(p).rglob("*") if f.is_file()) / 1e6
torch.set_num_threads(1)
parts = {}
for T in TS:
m = QNext(T, chunk=16 if T > 16 else None, **ane_kwargs()).eval()
ex = (torch.zeros((1, T), dtype=torch.long), torch.tensor([T - 1], dtype=torch.int32))
with torch.no_grad():
tr = torch.jit.trace(m, ex)
t0 = time.time()
mm = ct.convert(tr, inputs=[ct.TensorType("ids", shape=(1, T), dtype=np.int32), ct.TensorType("last", shape=(1,), dtype=np.int32)],
outputs=[ct.TensorType("lp", dtype=np.float32)], compute_precision=ct.precision.FLOAT16,
minimum_deployment_target=ct.target.macOS15, convert_to="mlprogram")
p = out / "parts" / f"dec_b1_n{T}_fp16.mlpackage"; shutil.rmtree(p, ignore_errors=True); mm.save(str(p)); parts[T] = p
print(f"OK part n{T} {time.time()-t0:.1f}s {size(p):.1f} MB", flush=True)
def merge(fns, dst, default):
d = MultiFunctionDescriptor()
for name, p in fns.items():
d.add_function(str(p), src_function_name="main", target_function_name=name)
d.default_function_name = default
shutil.rmtree(dst, ignore_errors=True); save_multifunction(d, str(dst))
return size(dst)
dst = out / "dec_mf_next_fp16.mlpackage"
print(f"OK multifunction n16..n96: {merge({f'n{T}': p for T, p in parts.items()}, dst, 'n16'):.1f} MB "
f"(parts total {sum(size(p) for p in parts.values()):.1f} MB) -> {dst}")
if "--merge-test" in sys.argv: # shipped t* parts + new n* parts in one package: do the weights de-duplicate?
fns = {f"t{T}": Path(f"mf/parts/dec_b3_t{T}_fp16.mlpackage") for T in TS} | {f"n{T}": p for T, p in parts.items()}
mt = out / "merge_test_fp16.mlpackage"
print(f"MERGE-TEST t16..t96 + n16..n96: {merge(fns, mt, 't16'):.1f} MB (shipped mf/dec_mf_fp16 {size('mf/dec_mf_fp16.mlpackage'):.1f} MB)")
for w in mt.rglob("*.bin"):
print(" weight file", w.relative_to(mt), f"{w.stat().st_size/1e6:.1f} MB")
shutil.rmtree(mt)
|