Zayden-Jung/exp-model-parameter-similarity
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement bounded streaming CPU functional reuse pilot

Browse Source
This commit is contained in:
Codex Bot
2026-11-09 04:40:33 +08:00
parent d4012592b6
commit 573cbedd98
15 changed files with 940 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
functional/model.py
+242
View File
@@ -0,0 +1,242 @@
"""Streaming FP32 text-only Qwen3.5 reference, original BF16 weights.
Mathematical semantics checked against Transformers Qwen3.5, commit
fd3cf0140de03f7619ac95b5cfed05b711ff9e22 (Apache-2.0, Qwen Team/HuggingFace).
No cached decode, padding, images, generation, fused kernels or autograd.
Full-sequence gated delta recurrence replaces the algebraically equivalent
chunk algorithm; both are independently compared in remote tests.
"""
import json
import time
from pathlib import Path
import numpy as np
from pilot.tensors import load_tensors, decode
PREFIX = 'model.language_model.'
def sigmoid(x):
return np.exp(-np.logaddexp(np.float32(0), -x))
def silu(x):
return x * sigmoid(x)
def norm(x, weight, eps=1e-6, centered=True):
return (x / np.sqrt(np.mean(x*x, axis=-1, keepdims=True) + eps)
* (1 + weight if centered else weight)).astype(np.float32)
def rope(x, theta=10000000., fraction=.25):
# Text-only: all three MRoPE position axes coincide.
dim = int(x.shape[-1] * fraction)
inv = np.float32(1) / np.float32(theta) ** (np.arange(0, dim, 2, dtype=np.float32) / dim)
freq = np.arange(x.shape[1], dtype=np.float32)[:, None] * inv[None, :]
angle = np.concatenate([freq, freq], axis=-1)[None, :, None, :]
a = x[..., :dim]
rotated = np.concatenate([-a[..., dim//2:], a[..., :dim//2]], axis=-1)
return np.concatenate([a*np.cos(angle) + rotated*np.sin(angle), x[..., dim:]], axis=-1)
def delta_rule(q, k, v, g, beta):
q = q / np.sqrt(np.sum(q*q, axis=-1, keepdims=True) + 1e-6)
k = k / np.sqrt(np.sum(k*k, axis=-1, keepdims=True) + 1e-6)
q = q / np.float32(np.sqrt(q.shape[-1]))
state = np.zeros((q.shape[0], q.shape[2], q.shape[3], v.shape[3]), dtype=np.float32)
out = np.empty_like(v)
for t in range(q.shape[1]):
state *= np.exp(g[:, t, :, None, None])
predicted = np.sum(state * k[:, t, :, :, None], axis=-2)
delta = (v[:, t] - predicted) * beta[:, t, :, None]
state += k[:, t, :, :, None] * delta[:, :, None, :]
out[:, t] = np.sum(state * q[:, t, :, :, None], axis=-2)
return out, state
def causal_conv(x, kernel):
out = np.zeros_like(x)
width = kernel.shape[-1]
for j in range(width):
lag = width - 1 - j
if lag < x.shape[1]:
out[:, lag:] += x[:, :x.shape[1]-lag] * kernel[:, j]
return silu(out)
def attention(q, k, v):
repeats = q.shape[2] // k.shape[2]
k, v = np.repeat(k, repeats, axis=2), np.repeat(v, repeats, axis=2)
scores = np.einsum('bthd,bshd->bhts', q, k) / np.float32(np.sqrt(q.shape[-1]))
mask = np.triu(np.ones(scores.shape[-2:], dtype=bool), 1)
scores = np.where(mask, -np.inf, scores)
exp = np.exp(scores - scores.max(axis=-1, keepdims=True))
p = exp / exp.sum(axis=-1, keepdims=True)
return np.einsum('bhts,bshd->bthd', p, v), p
class StreamingLM:
def __init__(self, weights, config, seconds=7200, enforce_memory=True):
self.tensors = load_tensors(Path(weights))
self.config = config.get('text_config', config)
self.deadline = time.monotonic() + seconds
self.enforce_memory = enforce_memory
self.eps = self.config['rms_norm_eps']
self.norm_replacement = None
self.capture_qk = False
self.qk_inputs = {}
self.swap_observation = None
self.swap_pressure_windows = 0
def guard(self):
if time.monotonic() > self.deadline:
raise TimeoutError('CPU computation deadline reached')
if not self.enforce_memory:
return
status = dict(line.split(':', 1) for line in Path('/proc/self/status').read_text().splitlines() if ':' in line)
rss_kib = int(status['VmRSS'].split()[0])
available = int(next(line.split()[1] for line in Path('/proc/meminfo').read_text().splitlines() if line.startswith('MemAvailable:')))
now = time.monotonic()
if self.swap_observation is None or now-self.swap_observation[0] >= 5:
counters = dict(line.split() for line in Path('/proc/vmstat').read_text().splitlines())
count = int(counters['pswpin'])+int(counters['pswpout'])
if self.swap_observation is not None:
elapsed = now-self.swap_observation[0]
rate = (count-self.swap_observation[1])*4096/elapsed
self.swap_pressure_windows = self.swap_pressure_windows+1 if rate > 4*1024**2 else 0
if self.swap_pressure_windows >= 3:
raise MemoryError('Sustained host swap IO above 4 MiB/s in three observation windows')
self.swap_observation = (now,count)
if rss_kib > 512*1024 or available < 256*1024:
raise MemoryError(f'Soft memory stop: RSS {rss_kib} KiB, available {available} KiB')
def small(self, name):
t = self.tensors[PREFIX + name]
if t.numel > 1048576:
raise ValueError('Use blockwise linear for large tensors')
x = t.read().reshape(t.shape)
t.evict()
return x
def linear(self, x, name, rows=256):
t = self.tensors[PREFIX + name + '.weight']
outdim, indim = t.shape
flat = x.reshape(-1, indim)
result = np.empty((flat.shape[0], outdim), dtype=np.float32)
for start in range(0, outdim, rows):
self.guard()
end = min(start + rows, outdim)
w = t.read(slice(start*indim, end*indim)).reshape(end-start, indim)
result[:, start:end] = flat @ w.T
del w
t.evict()
return result.reshape(*x.shape[:-1], outdim)
def embedding(self, ids):
t = self.tensors[PREFIX + 'embed_tokens.weight']
raw = t.raw.reshape(t.shape)
result = decode(raw[ids], t.dtype)
t.evict()
return result
def normalization(self, x, name):
if getattr(self, 'norm_replacement', None) and name == self.norm_replacement[0]:
name = self.norm_replacement[1]
return norm(x, self.small(name + '.weight'), self.eps)
def mlp(self, h, layer):
p = f'layers.{layer}.mlp.'
gate = silu(self.linear(h, p+'gate_proj'))
up = self.linear(h, p+'up_proj')
return self.linear(gate*up, p+'down_proj')
def full_attention(self, h, layer, shared_p=None):
p = f'layers.{layer}.self_attn.'
c = self.config
heads, kvheads, dim = c['num_attention_heads'], c['num_key_value_heads'], c['head_dim']
# Q and output gate are interleaved per head, not globally split.
mixed = self.linear(h, p+'q_proj').reshape(*h.shape[:-1], heads, 2*dim)
q, gate = mixed[..., :dim], mixed[..., dim:]
k = self.linear(h, p+'k_proj').reshape(*h.shape[:-1], kvheads, dim)
v = self.linear(h, p+'v_proj').reshape(*h.shape[:-1], kvheads, dim)
if getattr(self, 'capture_qk', False):
self.qk_inputs[layer] = (q.copy(), k.copy())
q = self.normalization(q, p+'q_norm')
k = self.normalization(k, p+'k_norm')
rp = c['rope_parameters']
q, k = (rope(x, rp['rope_theta'], rp['partial_rotary_factor']) for x in (q,k))
out, probs = attention(q,k,v)
if shared_p is not None:
out = np.einsum('bhts,bshd->bthd', shared_p, np.repeat(v, heads//kvheads, axis=2))
out = (out*sigmoid(gate)).reshape(*h.shape[:-1], -1)
return self.linear(out, p+'o_proj'), probs
def linear_attention(self, h, layer):
p = f'layers.{layer}.linear_attn.'
c = self.config
kh, vh = c['linear_num_key_heads'], c['linear_num_value_heads']
kd, vd = c['linear_key_head_dim'], c['linear_value_head_dim']
qkv = self.linear(h, p+'in_proj_qkv')
qkv = causal_conv(qkv, self.small(p+'conv1d.weight')[:,0,:])
q, k, v = np.split(qkv, [kh*kd, 2*kh*kd], axis=-1)
q = q.reshape(*h.shape[:-1], kh, kd)
k = k.reshape(*h.shape[:-1], kh, kd)
v = v.reshape(*h.shape[:-1], vh, vd)
if vh != kh:
q, k = np.repeat(q, vh//kh, axis=2), np.repeat(k, vh//kh, axis=2)
z = self.linear(h, p+'in_proj_z').reshape(*h.shape[:-1], vh, vd)
beta = sigmoid(self.linear(h, p+'in_proj_b'))
a = self.linear(h, p+'in_proj_a')
g = -np.exp(self.small(p+'A_log')) * np.logaddexp(np.float32(0), a+self.small(p+'dt_bias'))
out, _ = delta_rule(q,k,v,g,beta)
out = norm(out, self.small(p+'norm.weight'), self.eps, centered=False)*silu(z)
return self.linear(out.reshape(*h.shape[:-1], -1), p+'out_proj')
def forward(self, ids, capture=None, replacement=None, progress=None):
h = self.embedding(ids)
saved = {}
capture = set(capture or [])
for layer, kind in enumerate(self.config['layer_types']):
self.guard()
raw = h
n = self.normalization(h, f'layers.{layer}.input_layernorm')
if ('attn', layer) in capture:
saved[f'attn_{layer}'] = n.copy()
donor = replacement[2] if replacement and replacement[:2] == ('attn',layer) else layer
if kind == 'full_attention':
out, _ = self.full_attention(n,donor)
else:
out = self.linear_attention(n,donor)
h = raw + out
n = self.normalization(h, f'layers.{layer}.post_attention_layernorm')
if ('mlp',layer) in capture:
saved[f'mlp_{layer}'] = n.copy()
donor = replacement[2] if replacement and replacement[:2] == ('mlp',layer) else layer
h = h + self.mlp(n,donor)
if not np.isfinite(h).all():
raise FloatingPointError(f'Nonfinite hidden states at layer {layer}')
if progress:
progress(layer)
return self.normalization(h,'norm'), saved
def nll(self, h, ids, rows=512):
# Tied output embedding; online logsumexp avoids full vocabulary logits.
x = h[:,:-1].reshape(-1,h.shape[-1])
targets = ids[:,1:].reshape(-1)
t = self.tensors[PREFIX+'embed_tokens.weight']
mx = np.full(x.shape[0],-np.inf,dtype=np.float64)
total = np.zeros(x.shape[0],dtype=np.float64)
correct = np.empty(x.shape[0],dtype=np.float64)
for start in range(0,t.shape[0],rows):
self.guard()
end=min(start+rows,t.shape[0])
w=t.read(slice(start*t.shape[1],end*t.shape[1])).reshape(end-start,t.shape[1])
scores=x@w.T
m=np.maximum(mx,scores.max(axis=-1))
total=total*np.exp(mx-m)+np.exp(scores-m[:,None]).sum(axis=-1)
mx=m
which=(targets>=start)&(targets<end)
correct[which]=scores[np.flatnonzero(which),targets[which]-start]
del w,scores
t.evict()
return (mx+np.log(total)-correct).reshape(ids.shape[0],-1)