File size: 7,772 Bytes
78a1f66 | 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 | """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()
class D3RowPart(nn.Module):
"""Layers [start, stop) of D3RowEmbeds, for checkpoints too large to trace or load as one graph (d3-flash 9B).
hidden [1, L, H] in; the last part ends with the gather at the last real token and the final norm ([1, H]),
the others return hidden [1, L, H] for the next part."""
def __init__(self, cfg: TextConfig, seq_len: int, start: int, stop: int, chunk: int = 64):
super().__init__()
self.cfg, self.L, self.start, self.stop = cfg, seq_len, start, stop
self.last = stop == cfg.num_layers
self.layers = nn.ModuleList([DecoderLayer(cfg, i, seq_len, chunk) for i in range(start, stop)])
if self.last:
self.norm = RMSNorm(cfg.hidden_size, cfg.eps)
def forward(self, hidden, cos, sin, key_valid, last_onehot):
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])) if self.last else x
def load_row_part(root: str, seq_len: int, start: int, stop: int, chunk: int = 64, dtype=torch.float16) -> D3RowPart:
"""Only this part's weights are read (layer i -> layers[i - start])."""
cfg = load_text_config(root)
prefix = "language_model."
state = {}
from safetensors import safe_open
for path in sorted(__import__("glob").glob(f"{root}/model*.safetensors")):
with safe_open(path, "pt") as f:
for k in f.keys():
if not k.startswith(prefix):
continue
name = k[len(prefix):]
if name.startswith("layers."):
i = int(name.split(".")[1])
if start <= i < stop:
state[f"layers.{i - start}." + name.split(".", 2)[2]] = f.get_tensor(k).to(dtype)
elif name.startswith("norm.") and stop == cfg.num_layers:
state[name] = f.get_tensor(k).to(dtype)
model = D3RowPart(cfg, seq_len, start, stop, 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()
|