Files
colibri-strix/s2_research.py
T
JustVugg 1ae22a6135 colibrì: pure-C GLM-5.2 (744B MoE) engine with disk-streamed experts
Engine (c/glm.c): MLA attention with compressed KV, sigmoid noaux_tc router,
int8/int4/int2 quant kernels (AVX2), per-layer LRU expert cache + pinned
hot-store, batch-union MoE, native MTP speculative decoding (lossless),
multi-stop + official chat template, RAM auto-budget from MemAvailable.
Tokenizer: byte-level BPE in C. Tooling: coli CLI, disk-safe FP8→int4
converter, tiny-random oracle validation (TF 32/32, greedy 20/20).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-05 20:57:25 +02:00

111 lines
5.0 KiB
Python

"""
RICERCA - Il cardine del metodo: lo SKEW degli expert e' sfruttabile?
Se pochi expert "caldi" coprono gran parte delle attivazioni, allora la strategia
giusta per un modello che NON entra in RAM e':
- PIN dei caldi (residenti per sempre in RAM, profilati offline)
- STREAM dei freddi dal disco
invece di una LRU dinamica (che su RAM piccola va in pressione, l'abbiamo visto).
Test onesto: determino il "set caldo" dalla PRIMA meta' dei token, e misuro la
copertura sulla SECONDA meta' (mai vista). Confronto PIN-caldi statico vs LRU a parita' di K.
"""
import json, glob, collections, time
import torch
MODEL = "allenai/OLMoE-1B-7B-0924"
N_EXP, TOPK, N_LAYERS = 64, 8, 16
# testo piu' lungo e vario per statistiche decenti
PROMPTS = [
"The Roman Empire was one of the largest empires in history. At its height under "
"Trajan, it covered five million square kilometres and held seventy million people, "
"about a fifth of the world's population at the time. The empire's longevity and vast "
"extent ensured a lasting influence on language, religion, architecture, philosophy, law "
"and forms of government across the territory it once governed. ",
"Photosynthesis is a biological process used by plants, algae and some bacteria to "
"convert light energy into chemical energy stored in glucose. It occurs in the chloroplasts, "
"specifically using the green pigment chlorophyll. The process consumes carbon dioxide and "
"water and releases oxygen as a by-product, sustaining most life on Earth. ",
"def fibonacci(n):\n a, b = 0, 1\n result = []\n for _ in range(n):\n "
"result.append(a)\n a, b = b, a + b\n return result\n\n"
"class Stack:\n def __init__(self):\n self.items = []\n def push(self, x):\n"
" self.items.append(x)\n def pop(self):\n return self.items.pop()\n",
"L'economia mondiale nel ventunesimo secolo e' caratterizzata da una crescente "
"globalizzazione, dall'integrazione dei mercati finanziari e dalla rapida diffusione "
"delle tecnologie digitali. Le banche centrali giocano un ruolo cruciale nel mantenere "
"la stabilita' dei prezzi attraverso la politica monetaria. ",
]
def collect():
from transformers import AutoModelForCausalLM, AutoTokenizer
print("carico modello...", flush=True)
tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16,
low_cpu_mem_usage=True).eval()
trace = [[] for _ in range(N_LAYERS)]
for p in PROMPTS:
ids = tok(p, return_tensors="pt").input_ids
with torch.no_grad():
out = model(ids, output_router_logits=True, use_cache=False)
for li, rl in enumerate(out.router_logits):
for row in rl.topk(TOPK, -1).indices.tolist():
trace[li].append(tuple(row))
print(f" +{ids.shape[1]} token", flush=True)
return trace
def lru_hit(seq, K):
c = collections.OrderedDict(); hit = tot = 0
for experts in seq:
for e in experts:
tot += 1
if e in c: hit += 1; c.move_to_end(e)
else:
c[e] = 1
if len(c) > K: c.popitem(last=False)
return hit / tot
def static_hot_hit(train, test, K):
"""Set caldo = K piu' frequenti nel train; copertura misurata sul test."""
freq = collections.Counter(e for experts in train for e in experts)
hot = set(e for e, _ in freq.most_common(K))
hit = tot = 0
for experts in test:
for e in experts:
tot += 1
if e in hot: hit += 1
return hit / tot
if __name__ == "__main__":
trace = collect()
ntok = len(trace[0])
print(f"\nToken totali: {ntok} x {N_LAYERS} layer = {ntok*N_LAYERS*TOPK} accessi expert\n")
# skew: distribuzione di frequenza (media sui layer), e curva di copertura top-K
print("COPERTURA del set caldo (statico, profilato su prima meta', testato su seconda):")
print(f"{'K':>4} {'RAM':>7} {'pin-caldo':>10} {'LRU':>8} (uniforme=K/64)")
for K in (8, 12, 16, 24, 32):
cov_static, cov_lru = [], []
for li in range(N_LAYERS):
seq = trace[li]; h = len(seq) // 2
cov_static.append(static_hot_hit(seq[:h], seq[h:], K))
cov_lru.append(lru_hit(seq, K))
cs = sum(cov_static)/N_LAYERS; cl = sum(cov_lru)/N_LAYERS
ram = K * N_LAYERS * 12.6 / 1024
print(f"{K:>4} {ram:>5.1f}GB {cs*100:>9.1f}% {cl*100:>7.1f}% {K/64*100:>4.0f}%")
# quanto e' skewata la distribuzione? entropia normalizzata e top-8 share
import math
shares = []
for li in range(N_LAYERS):
freq = collections.Counter(e for ex in trace[li] for e in ex)
tot = sum(freq.values())
top8 = sum(c for _, c in freq.most_common(8)) / tot
shares.append(top8)
print(f"\nSkew: gli 8 expert piu' caldi (su 64) coprono in media "
f"{sum(shares)/len(shares)*100:.1f}% delle attivazioni (uniforme sarebbe 12.5%).")