File size: 9,328 Bytes
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 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | """D1A (Gemma 4 text) -> Core ML pilot (jonpol01/d1a#224): rebuild a D1A MLX export as a torch module, split into
- host side: the two embedding tables (token embedding, per-layer embedding: 2.75B of E2B's 4.6B parameters) read
row by row from the export's safetensors (dequantized per request, never held whole), and
- Core: everything that computes (per-layer projection, 35 decoder layers, final norm, pointer head), which is what
gets converted to Core ML.
Inputs follow D1A's packed encoding (d1a.core.encoding.encode): one record, block-causal mask, positions restarting per
question after the state, readout at each question's <decide> and </opt> tokens."""
import json
import math
from pathlib import Path
import numpy as np
import torch
from safetensors import safe_open
from torch import nn
from transformers import AutoConfig
from transformers.models.gemma4 import modeling_gemma4 as G
PREFIX = "language_model.model."
def dequant(packed, scales, biases, bits, group=64):
"""MLX affine dequantization: packed uint32 [out, in*bits/32] (lowest bits first) -> float32 [out, in]."""
p = packed.view(torch.int32).to(torch.int64) & 0xFFFFFFFF
per = 32 // bits
shifts = torch.arange(per, dtype=torch.int64) * bits
q = ((p[..., None] >> shifts) & ((1 << bits) - 1)).reshape(p.shape[0], -1).to(torch.float32)
s, b = scales.to(torch.float32), biases.to(torch.float32)
return (q.view(q.shape[0], -1, group) * s[..., None] + b[..., None]).reshape(q.shape[0], -1)
class _Shards:
"""model.safetensors or model-0000N-of-0000M.safetensors (via model.safetensors.index.json), read like one safe_open."""
def __init__(self, path):
idx = Path(path) / "model.safetensors.index.json"
if idx.exists():
wm = json.loads(idx.read_text())["weight_map"]
files = {f: safe_open(str(Path(path) / f), "pt") for f in sorted(set(wm.values()))}
self.where = {k: files[f] for k, f in wm.items()}
else: # a released Core ML bundle ships only the embedding tables, as embeddings.safetensors
name = "model.safetensors" if (Path(path) / "model.safetensors").exists() else "embeddings.safetensors"
f = safe_open(str(Path(path) / name), "pt")
self.where = {k: f for k in f.keys()}
def keys(self):
return self.where.keys()
def get_tensor(self, k):
return self.where[k].get_tensor(k)
def get_slice(self, k):
return self.where[k].get_slice(k)
class Export:
"""A D1A MLX export directory (config.json, d1a_config.json, model.safetensors, head.safetensors)."""
def __init__(self, path):
self.path = Path(path)
self.cfg = AutoConfig.from_pretrained(self.path)
self.tc = self.cfg.text_config
self.tc._attn_implementation = "eager"
self.d1a = json.loads((self.path / "d1a_config.json").read_text())
q = self.d1a["quantization"]
self.bits = {"default": (q["bits"], q["group_size"]), "embed_tokens_per_layer": (q[PREFIX.replace("model.", "") + "model.embed_tokens_per_layer"]["bits"] if False else q["language_model.model.embed_tokens_per_layer"]["bits"], q["language_model.model.embed_tokens_per_layer"]["group_size"])}
self.f = _Shards(self.path)
def bits_of(self, name):
return self.bits["embed_tokens_per_layer"] if name.endswith("embed_tokens_per_layer") else self.bits["default"]
def tensor(self, name):
"""A float32 tensor for `name` (without PREFIX), dequantized if it is stored quantized."""
full = PREFIX + name
keys = set(self.f.keys())
if full + ".weight" in keys and full + ".scales" in keys:
bits, group = self.bits_of(name)
return dequant(self.f.get_tensor(full + ".weight"), self.f.get_tensor(full + ".scales"), self.f.get_tensor(full + ".biases"), bits, group)
return self.f.get_tensor(full).to(torch.float32)
def text_model(self, dtype=torch.float32):
"""The Gemma 4 text trunk with every non-embedding weight loaded (embedding tables stay empty: host side)."""
with torch.device("meta"):
m = G.Gemma4TextModel(self.tc)
m = m.to_empty(device="cpu")
state = {}
keys = set(self.f.keys())
for name, p in m.named_parameters():
if name.startswith("embed_tokens"):
continue
base = name[: -len(".weight")] if name.endswith(".weight") else name
if PREFIX + base + ".scales" in keys:
state[name] = self.tensor(base)
else:
state[name] = self.f.get_tensor(PREFIX + name).to(torch.float32)
assert state[name].shape == p.shape, (name, state[name].shape, p.shape)
for name, b in m.named_buffers():
if PREFIX + name in keys:
state[name] = self.f.get_tensor(PREFIX + name).to(torch.float32)
missing, unexpected = m.load_state_dict(state, strict=False)
missing = [k for k in missing if not k.startswith("embed_tokens")]
assert not missing and not unexpected, (missing, unexpected)
# buffers transformers computes itself (rotary inv_freq, embed_scale): re-run their init on CPU
m.rotary_emb = G.Gemma4TextRotaryEmbedding(self.tc)
return m.to(dtype).eval()
def head(self):
h = safe_open(str(self.path / "head.safetensors"), "pt")
dp, d = h.get_tensor("q.weight").shape # pointer width, hidden size (E2B 1536, E4B 2560)
q, k = nn.Linear(d, dp), nn.Linear(d, dp)
q.weight.data, q.bias.data = h.get_tensor("q.weight").float(), h.get_tensor("q.bias").float()
k.weight.data, k.bias.data = h.get_tensor("k.weight").float(), h.get_tensor("k.bias").float()
return q, k
class HostEmbed:
"""Row lookups of the two embedding tables, dequantized per request: emb = E[ids] * sqrt(hidden) and
ple = E_pl[ids] * sqrt(ple_dim), reshaped [L, layers, ple_dim] (transformers' get_per_layer_inputs)."""
def __init__(self, ex):
self.ex, self.tc = ex, ex.tc
self.f = ex.f
def _rows(self, name, ids):
full = PREFIX + name
bits, group = self.ex.bits_of(name)
w, s, b = (self.f.get_slice(full + suf) for suf in (".weight", ".scales", ".biases"))
uniq, inv = np.unique(np.asarray(ids), return_inverse=True)
rows = [dequant(w[int(i): int(i) + 1], s[int(i): int(i) + 1], b[int(i): int(i) + 1], bits, group) for i in uniq]
return torch.cat(rows)[torch.as_tensor(inv)]
def __call__(self, ids):
emb = self._rows("embed_tokens", ids) * (self.tc.hidden_size ** 0.5)
ple = self._rows("embed_tokens_per_layer", ids) * (self.tc.hidden_size_per_layer_input ** 0.5)
return emb[None], ple.view(1, len(ids), self.tc.num_hidden_layers, self.tc.hidden_size_per_layer_input)
def masks_for(enc, window, L=None, dtype=torch.float32):
"""D1A's packed masks (d1a.backends.torch.branch_masks) for one record, padded to L: additive [1, 1, L, L]."""
seg, pos = enc["seg"], enc["pos"]
n = len(seg)
L = L or n
s = torch.full((L,), -1)
s[:n] = torch.tensor(seg)
p = torch.zeros(L, dtype=torch.long)
p[:n] = torch.tensor(pos)
causal = torch.tril(torch.ones(L, L, dtype=torch.bool))
same = (s[None, :] == s[:, None]) | (s[None, :] == 0)
allow = (causal & same & (s != -1)[None, :]) | torch.eye(L, dtype=torch.bool)
neg = torch.finfo(dtype).min
full = torch.zeros(L, L, dtype=dtype).masked_fill(~allow, neg)
far = (p[:, None] - p[None, :] >= window) & ~torch.eye(L, dtype=torch.bool)
slide = full.masked_fill(far, neg)
return full[None, None], slide[None, None], p[None]
class Core(nn.Module):
"""Everything that computes: per-layer projection, decoder layers, final norm, pointer head.
forward(emb [1,L,d], ple_tok [1,L,layers,256], pos [1,L] int, mask_full, mask_slide [1,1,L,L],
decide [Q] int, opts [Q,K] int, opt_valid [Q,K] float (1 used, 0 padding)) -> logits [Q,K]
(raw, temperature not applied; padded options get a large negative)."""
def __init__(self, text, head_q, head_k, silu_tanh=True):
super().__init__()
self.text, self.hq, self.hk = text, head_q, head_k
self.types = text.config.layer_types
self.scale = 1 / math.sqrt(head_q.out_features)
def forward(self, emb, ple_tok, pos, mask_full, mask_slide, decide, opts, opt_valid):
t = self.text
ple = t.project_per_layer_inputs(emb, ple_tok)
rope = {lt: t.rotary_emb(emb, pos, lt) for lt in t.unique_layer_types}
masks = {"full_attention": mask_full, "sliding_attention": mask_slide}
from collections import UserDict
kv = UserDict()
h = emb
for i, layer in enumerate(t.layers):
h = layer(h, ple[:, :, i, :], shared_kv_states=kv, position_embeddings=rope[self.types[i]],
attention_mask=masks[self.types[i]], position_ids=pos)
h = t.norm(h)[0] # [L, d]
q = self.hq(h[decide]) # [Q, dp]
k = self.hk(h[opts.reshape(-1)]).view(*opts.shape, -1) # [Q, K, dp]
z = (k * q[:, None, :]).sum(-1) * self.scale
return z + (opt_valid - 1.0) * 1e4
|