248 lines
11 KiB
Python
248 lines
11 KiB
Python
"""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 ctypes
|
|
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
|
|
self.malloc_trim = getattr(ctypes.CDLL(None), 'malloc_trim', None)
|
|
|
|
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
|
|
own_swap_kib = int(status.get('VmSwap', '0 kB').split()[0])
|
|
if self.swap_pressure_windows >= 3 and (available < 384*1024 or own_swap_kib > 128*1024):
|
|
raise MemoryError(f'Sustained host swap pressure plus constrained memory: available={available} KiB, process_swap={own_swap_kib} KiB')
|
|
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()
|
|
if self.malloc_trim is not None:
|
|
self.malloc_trim(0) # Release only this process's freed temporary heap pages.
|
|
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)
|