File size: 5,467 Bytes
f217779 | 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 | """d3-lite (Qwen3.5-0.8B, code readout v1) as one Core ML row graph.
One question per call, as upstream: the row is right-padded to the bucket length L. The 6 full-attention layers are
noncausal (decision_config `noncausal_full_attention`): they see every real token and mask only the padding keys. The
Gated DeltaNet layers are causal scans, so padding after the last real token cannot reach it. Output is the final-normed
hidden state at the last real token; the 255 x 1024 FP32 readout runs on the host, as upstream.
"""
import json
import numpy as np
import torch
from safetensors.torch import load_file
from torch import nn
from qwen35_export import DecoderLayer, RMSNorm, TextConfig, rope_cos_sin
MASK_VALUE = -1e4
class D3Row(nn.Module):
def __init__(self, cfg: TextConfig, seq_len: int, chunk: int = 64):
super().__init__()
self.cfg, self.L = cfg, seq_len
self.embed_tokens = nn.Embedding(1, cfg.hidden_size) # replaced on load
self.layers = nn.ModuleList([DecoderLayer(cfg, i, seq_len, chunk) for i in range(cfg.num_layers)])
self.norm = RMSNorm(cfg.hidden_size, cfg.eps)
def forward(self, input_ids, cos, sin, key_valid, last_onehot):
"""input_ids [1, L] int32; cos/sin [L, R]; key_valid [L] (1 real, 0 pad); last_onehot [1, L] -> hidden [1, H]."""
mask = ((1.0 - key_valid) * MASK_VALUE)[None, None, None, :]
x = self.embed_tokens(input_ids)
for layer in self.layers:
x = layer(x, cos, sin, mask)
return self.norm(torch.matmul(last_onehot, x[0]))
class D3RowEmbeds(nn.Module):
"""D3Row for image requests: the host gathers token embeddings and splices the vision tower's tokens in at the
image-pad positions, and builds cos/sin from the 3-D M-RoPE positions (text and image tokens differ)."""
def __init__(self, cfg: TextConfig, seq_len: int, chunk: int = 64):
super().__init__()
self.cfg, self.L = cfg, seq_len
self.layers = nn.ModuleList([DecoderLayer(cfg, i, seq_len, chunk) for i in range(cfg.num_layers)])
self.norm = RMSNorm(cfg.hidden_size, cfg.eps)
def forward(self, hidden, cos, sin, key_valid, last_onehot):
"""hidden [1, L, H]; cos/sin [L, R]; key_valid [L]; last_onehot [1, L] -> final hidden [1, H]."""
mask = ((1.0 - key_valid) * MASK_VALUE)[None, None, None, :]
x = hidden
for layer in self.layers:
x = layer(x, cos, sin, mask)
return self.norm(torch.matmul(last_onehot, x[0]))
def load_row_embeds(root: str, seq_len: int, chunk: int = 64, dtype=torch.float32) -> tuple[D3RowEmbeds, torch.Tensor]:
"""(graph, embedding table [V, H]). dtype=float16 halves peak memory for tracing the larger checkpoints (d3-mini:
fp32 weights plus their copy do not fit in 24 GB); weights are assigned in place, not copied."""
cfg = load_text_config(root)
prefix = "language_model."
state = {}
from safetensors import safe_open
files = sorted(__import__("glob").glob(f"{root}/model*.safetensors"))
for path in files:
with safe_open(path, "pt") as f:
for k in f.keys():
if k.startswith(prefix):
state[k[len(prefix):]] = f.get_tensor(k).to(dtype)
embed = state.pop("embed_tokens.weight")
model = D3RowEmbeds(cfg, seq_len, chunk).to(dtype)
missing, unexpected = model.load_state_dict(state, strict=False, assign=True)
if missing or unexpected:
raise RuntimeError(f"weight mismatch: missing={missing[:8]} unexpected={unexpected[:8]}")
return model.eval(), embed
def load_text_config(root: str) -> TextConfig:
return TextConfig(json.load(open(f"{root}/config.json"))["text_config"])
def load_row(root: str, seq_len: int, chunk: int = 64) -> D3Row:
cfg = load_text_config(root)
weights = load_file(f"{root}/model.safetensors")
prefix = "language_model."
state = {k[len(prefix):]: v.float() for k, v in weights.items() if k.startswith(prefix)}
model = D3Row(cfg, seq_len, chunk)
model.embed_tokens = nn.Embedding.from_pretrained(state.pop("embed_tokens.weight"), freeze=True)
missing, unexpected = model.load_state_dict(state, strict=False)
missing = [m for m in missing if not m.startswith("embed_tokens")]
if missing or unexpected:
raise RuntimeError(f"weight mismatch: missing={missing[:8]} unexpected={unexpected[:8]}")
return model.eval()
def load_readout(root: str) -> np.ndarray:
return load_file(f"{root}/readout.safetensors")["weight"].float().numpy()
def row_inputs(cfg: TextConfig, ids, L: int, pad_id: int) -> dict[str, np.ndarray]:
n = len(ids)
if not 1 <= n <= L:
raise ValueError(f"{n} tokens do not fit L{L}")
input_ids = np.full((1, L), pad_id, dtype=np.int32)
input_ids[0, :n] = ids
pos = torch.arange(L)
pos[n:] = n - 1
cos, sin = rope_cos_sin(cfg, pos[None].expand(3, -1))
valid = np.zeros(L, dtype=np.float32)
valid[:n] = 1
last = np.zeros((1, L), dtype=np.float32)
last[0, n - 1] = 1
return {"input_ids": input_ids, "cos": cos.numpy(), "sin": sin.numpy(), "key_valid": valid, "last_onehot": last}
def code_probabilities(hidden: np.ndarray, readout: np.ndarray, count: int, temperature: float = 1.0) -> np.ndarray:
logits = hidden.astype(np.float64).reshape(-1) @ readout[:count].astype(np.float64).T / temperature
logits -= logits.max()
p = np.exp(logits)
return p / p.sum()
|