README: try-it-on-better-hardware guide + predictions + support section; drop stage-0 research scripts from repo
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
JustVugg
+11
@@ -21,3 +21,14 @@ c/bench/
|
||||
*.safetensors
|
||||
stats*.txt
|
||||
*.log
|
||||
|
||||
# script di ricerca stadio-0 (esperimenti locali, non parte del motore)
|
||||
/engine.py
|
||||
/run.py
|
||||
/validate_ref.py
|
||||
/s0_costmodel.py
|
||||
/s1_trace_hitrate.py
|
||||
/s2_research.py
|
||||
/quant_test.py
|
||||
/profile_run.py
|
||||
/sweep.py
|
||||
|
||||
@@ -60,6 +60,52 @@ COLI_MODEL=/path/to/glm52_i4 ./coli chat
|
||||
|
||||
Useful knobs (env or flags): `--topp 0.7` adaptive expert top-p (30–40% less disk), `--ngen N` max tokens, `STATS=f`/`PIN=f PIN_GB=g` record expert usage then pin the hottest in RAM, `THINK=1` enable GLM-5.2's reasoning block, `DRAFT=n` MTP draft depth, `TF=1` teacher-forcing validation.
|
||||
|
||||
## Got a better machine? Try it — here's what to expect
|
||||
|
||||
colibrì was built on deliberately humble hardware (12 cores, 25 GB RAM, NVMe behind a WSL2 VHDX that caps random reads at ~1 GB/s). **Every one of those constraints is a knob your machine can turn up.** The engine needs: Linux (or WSL2), gcc with OpenMP, AVX2, ≥16 GB RAM, and the ~370 GB int4 model on a local NVMe (ext4 — never a network/9p mount).
|
||||
|
||||
**How to test it, in order:**
|
||||
|
||||
```bash
|
||||
cd c && ./setup.sh # build + architecture self-test (expects 32/32)
|
||||
|
||||
# 1) measure YOUR disk the way the engine uses it (parallel 19 MB random reads):
|
||||
gcc -O2 -fopenmp iobench.c -o iobench
|
||||
./iobench /path/to/glm52_i4/out-00069.safetensors 19 64 8 0 # buffered, 8 threads
|
||||
./iobench /path/to/glm52_i4/out-00069.safetensors 19 64 8 1 # O_DIRECT
|
||||
|
||||
# 2) chat; watch the per-turn stats line (tok/s, expert hit-rate, RSS):
|
||||
COLI_MODEL=/path/to/glm52_i4 ./coli chat
|
||||
|
||||
# 3) record expert usage, then pin the hottest experts in your spare RAM:
|
||||
STATS=stats.txt ./coli chat
|
||||
PIN=stats.txt PIN_GB=20 ./coli chat # scale PIN_GB to your free RAM
|
||||
|
||||
# 4) quality benchmarks (MMLU/HellaSwag/ARC):
|
||||
./coli bench
|
||||
```
|
||||
|
||||
**Back-of-envelope predictions** (decode is disk-bound: a cold token costs ~11.4 GB of expert reads; MTP speculation roughly halves the effective cost; RAM turns cold reads into free cache hits):
|
||||
|
||||
| machine | expected |
|
||||
|---|---|
|
||||
| this dev box (WSL2 VHDX, ~1 GB/s, 25 GB RAM) | ~0.05–0.1 tok/s cold — proven baseline |
|
||||
| native Linux, PCIe4 NVMe (~3–5 GB/s random), 32 GB | ~0.5–1 tok/s |
|
||||
| PCIe5 NVMe or 2×NVMe RAID0 (~8–12 GB/s), 64 GB (PIN ~40 GB of hot experts) | ~2–4 tok/s |
|
||||
| 128–256 GB RAM workstation (hot expert set mostly cached) | ~5–15 tok/s, matmul-bound — interactive |
|
||||
|
||||
These are estimates, not measurements — if you run colibrì on serious hardware, **please open an issue with your numbers**: real datapoints from better machines are exactly what this project needs next.
|
||||
|
||||
## Supporting the project
|
||||
|
||||
colibrì is a one-person project, written and tested entirely on a 12-core laptop with 25 GB of RAM — the numbers above are the ceiling of what I can measure at home. If this project is useful or interesting to you and you'd like to support its development (better test hardware translates *directly* into a faster engine for everyone: real NVMe scaling data, bigger pinned caches, int2/int3 quality sweeps on real benchmarks), you can:
|
||||
|
||||
- ⭐ star the repo and share it;
|
||||
- 🐛 open issues with benchmark numbers from your hardware;
|
||||
- 💬 reach out via GitHub issues if you'd like to sponsor development or donate hardware.
|
||||
|
||||
Every contribution, from a datapoint to a disk, moves the ceiling.
|
||||
|
||||
## Repo layout
|
||||
|
||||
```
|
||||
@@ -67,10 +113,10 @@ c/glm.c the engine (GLM-5.2 forward, streaming MoE, MTP, serve mode)
|
||||
c/st.h safetensors reader: pread + fadvise, no mmap (RSS stays flat)
|
||||
c/tok.h byte-level BPE tokenizer in C
|
||||
c/coli CLI: chat / run / bench / convert / info
|
||||
c/iobench.c parallel disk microbenchmark (measures what the engine feels)
|
||||
c/convert_fp8_to_int4.py disk-safe FP8 → int4 converter
|
||||
c/make_glm_oracle.py tiny-random oracle generator for validation
|
||||
c/olmoe.c stage-A engine (OLMoE), first validation target
|
||||
*.py research scripts (cost model, trace analysis, py engine)
|
||||
```
|
||||
|
||||
## Why "colibrì"
|
||||
|
||||
@@ -1,219 +0,0 @@
|
||||
"""
|
||||
Motore di inferenza MoE con EXPERT-STREAMING dal disco.
|
||||
|
||||
Idea (quella dell'utente, resa reale):
|
||||
- la parte DENSA (embedding, attenzione, router, norme, lm_head) sta in RAM;
|
||||
- gli EXPERT stanno su disco in un file safetensors mappato in memoria (mmap)
|
||||
e vengono letti SOLO quando un token li attiva;
|
||||
- una cache LRU tiene in RAM gli expert "caldi" -> meno letture da disco.
|
||||
|
||||
Cosi' un modello che NON entra in RAM gira lo stesso: in RAM ci tieni solo
|
||||
densa + cache, il resto vive sul disco. Validato qui su OLMoE-1B-7B.
|
||||
|
||||
NB: scritto per OLMoE (Llama-style con QK-norm). I punti specifici del modello
|
||||
(routing, norme) sono isolati cosi' che lo stesso scheletro valga per GLM/DeepSeek.
|
||||
"""
|
||||
import os, json, glob, struct, time, mmap, collections
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
ST_DTYPE = {"BF16": torch.bfloat16, "F16": torch.float16, "F32": torch.float32}
|
||||
|
||||
|
||||
class Shards:
|
||||
"""Indicizza i tensori in piu' file safetensors e li legge via mmap on-demand."""
|
||||
def __init__(self, snap_dir):
|
||||
self.index = {} # name -> (shard_path, abs_offset, nbytes, torch_dtype, shape)
|
||||
self.mm = {} # shard_path -> mmap
|
||||
for shard in sorted(glob.glob(os.path.join(snap_dir, "model-*.safetensors"))):
|
||||
with open(shard, "rb") as f:
|
||||
hlen = struct.unpack("<Q", f.read(8))[0]
|
||||
header = json.loads(f.read(hlen))
|
||||
data_start = 8 + hlen
|
||||
for name, meta in header.items():
|
||||
if name == "__metadata__":
|
||||
continue
|
||||
a, b = meta["data_offsets"]
|
||||
self.index[name] = (shard, data_start + a, b - a,
|
||||
ST_DTYPE[meta["dtype"]], tuple(meta["shape"]))
|
||||
fd = open(shard, "rb")
|
||||
self.mm[shard] = mmap.mmap(fd.fileno(), 0, prot=mmap.PROT_READ)
|
||||
|
||||
def read(self, name):
|
||||
"""Legge un tensore dal disco (mmap) e ne fa una copia RESIDENTE in RAM."""
|
||||
shard, off, nbytes, dt, shape = self.index[name]
|
||||
mv = memoryview(self.mm[shard])[off:off + nbytes]
|
||||
return torch.frombuffer(mv, dtype=dt).reshape(shape).clone()
|
||||
|
||||
def has(self, name):
|
||||
return name in self.index
|
||||
|
||||
|
||||
def quant_dequant(w, bits):
|
||||
"""Quantizzazione simmetrica per-riga a `bits` bit, poi dequantizza in bf16.
|
||||
Simula numericamente cosa darebbe un expert salvato a `bits` bit sul disco."""
|
||||
if bits >= 16:
|
||||
return w
|
||||
qmax = (1 << (bits - 1)) - 1 # int8->127, int4->7, int3->3, int2->1
|
||||
wf = w.float()
|
||||
scale = wf.abs().amax(dim=1, keepdim=True) / qmax
|
||||
scale = scale.clamp_min(1e-8)
|
||||
wq = torch.round(wf / scale).clamp(-qmax - 1, qmax)
|
||||
return (wq * scale).to(torch.bfloat16)
|
||||
|
||||
|
||||
class ExpertCache:
|
||||
"""Cache LRU degli expert. capacity = quanti expert teniamo residenti PER LAYER."""
|
||||
def __init__(self, shards, n_layers, capacity, quant_bits=16):
|
||||
self.shards = shards
|
||||
self.cap = capacity
|
||||
self.quant_bits = quant_bits
|
||||
self.caches = [collections.OrderedDict() for _ in range(n_layers)]
|
||||
self.hits = 0
|
||||
self.miss = 0
|
||||
|
||||
def get(self, layer, eid):
|
||||
"""Ritorna (gate_w, up_w, down_w) dell'expert, da cache o da disco."""
|
||||
c = self.caches[layer]
|
||||
if eid in c:
|
||||
self.hits += 1
|
||||
c.move_to_end(eid)
|
||||
return c[eid]
|
||||
self.miss += 1
|
||||
p = f"model.layers.{layer}.mlp.experts.{eid}."
|
||||
# tengo gli expert in bf16 (niente .float(): -24% tempo, -50% RAM, piu' fedele al riferimento)
|
||||
b = self.quant_bits
|
||||
w = (quant_dequant(self.shards.read(p + "gate_proj.weight"), b),
|
||||
quant_dequant(self.shards.read(p + "up_proj.weight"), b),
|
||||
quant_dequant(self.shards.read(p + "down_proj.weight"), b))
|
||||
c[eid] = w
|
||||
if len(c) > self.cap:
|
||||
c.popitem(last=False)
|
||||
return w
|
||||
|
||||
def hitrate(self):
|
||||
t = self.hits + self.miss
|
||||
return self.hits / t if t else 0.0
|
||||
|
||||
|
||||
def rmsnorm(x, w, eps=1e-5):
|
||||
x = x.float()
|
||||
x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
|
||||
return x * w.float()
|
||||
|
||||
|
||||
def rotate_half(x):
|
||||
h = x.shape[-1] // 2
|
||||
return torch.cat((-x[..., h:], x[..., :h]), dim=-1)
|
||||
|
||||
|
||||
class OlmoeStreaming:
|
||||
def __init__(self, snap_dir, expert_cap=16, quant_bits=16):
|
||||
self.cfg = json.load(open(os.path.join(snap_dir, "config.json")))
|
||||
self.shards = Shards(snap_dir)
|
||||
c = self.cfg
|
||||
self.L = c["num_hidden_layers"]
|
||||
self.H = c["num_attention_heads"]
|
||||
self.hd = c["hidden_size"] // self.H
|
||||
self.topk = c["num_experts_per_tok"]
|
||||
self.eps = c.get("rms_norm_eps", 1e-5)
|
||||
self.norm_topk = c.get("norm_topk_prob", False)
|
||||
theta = c.get("rope_theta", 10000.0)
|
||||
self.inv_freq = 1.0 / (theta ** (torch.arange(0, self.hd, 2).float() / self.hd))
|
||||
self.cache = ExpertCache(self.shards, self.L, expert_cap, quant_bits)
|
||||
|
||||
# --- parte DENSA: residente in RAM (float32) ---
|
||||
t = time.time()
|
||||
self.embed = self.shards.read("model.embed_tokens.weight").float()
|
||||
self.lm_head = self.shards.read("lm_head.weight").float()
|
||||
self.final_norm = self.shards.read("model.norm.weight").float()
|
||||
self.layers = []
|
||||
for i in range(self.L):
|
||||
p = f"model.layers.{i}."
|
||||
self.layers.append({
|
||||
"in_ln": self.shards.read(p + "input_layernorm.weight").float(),
|
||||
"post_ln":self.shards.read(p + "post_attention_layernorm.weight").float(),
|
||||
"q": self.shards.read(p + "self_attn.q_proj.weight").float(),
|
||||
"k": self.shards.read(p + "self_attn.k_proj.weight").float(),
|
||||
"v": self.shards.read(p + "self_attn.v_proj.weight").float(),
|
||||
"o": self.shards.read(p + "self_attn.o_proj.weight").float(),
|
||||
"qn": self.shards.read(p + "self_attn.q_norm.weight").float(),
|
||||
"kn": self.shards.read(p + "self_attn.k_norm.weight").float(),
|
||||
"gate": self.shards.read(p + "mlp.gate.weight").float(),
|
||||
})
|
||||
self.dense_load_s = time.time() - t
|
||||
|
||||
def _rope(self, x, pos):
|
||||
# x: (heads, seq, hd) ; pos: (seq,)
|
||||
freqs = torch.outer(pos.float(), self.inv_freq) # (seq, hd/2)
|
||||
emb = torch.cat((freqs, freqs), dim=-1) # (seq, hd)
|
||||
cos, sin = emb.cos(), emb.sin()
|
||||
return x * cos + rotate_half(x) * sin
|
||||
|
||||
def _attn(self, lw, x, pos, layer, kv):
|
||||
"""Attenzione con KV-cache. x = SOLO i token nuovi (S in prefill, 1 in decode).
|
||||
pos = posizioni assolute dei token nuovi. kv = lista per-layer dei (k,v) passati."""
|
||||
S = x.shape[0]
|
||||
q = rmsnorm(x @ lw["q"].T, lw["qn"], self.eps).view(S, self.H, self.hd).transpose(0, 1)
|
||||
k = rmsnorm(x @ lw["k"].T, lw["kn"], self.eps).view(S, self.H, self.hd).transpose(0, 1)
|
||||
v = (x @ lw["v"].T).view(S, self.H, self.hd).transpose(0, 1)
|
||||
q = self._rope(q, pos); k = self._rope(k, pos)
|
||||
if kv is not None and kv[layer] is not None: # concateno il passato
|
||||
pk, pv = kv[layer]
|
||||
k = torch.cat([pk, k], dim=1); v = torch.cat([pv, v], dim=1)
|
||||
if kv is not None:
|
||||
kv[layer] = (k, v)
|
||||
Tk = k.shape[1] # lunghezza totale (passato+nuovi)
|
||||
scores = (q @ k.transpose(-1, -2)) / (self.hd ** 0.5) # (H,S,Tk)
|
||||
# mask causale: query a posizione assoluta pos[i] vede key j<=pos[i]
|
||||
kpos = torch.arange(Tk)
|
||||
mask = torch.where(kpos[None, :] > pos[:, None], float("-inf"), 0.0) # -inf dove vietato
|
||||
a = F.softmax(scores + mask, dim=-1)
|
||||
out = (a @ v).transpose(0, 1).reshape(S, self.H * self.hd)
|
||||
return out @ lw["o"].T
|
||||
|
||||
def _moe(self, layer, lw, x):
|
||||
S = x.shape[0]
|
||||
logits = x @ lw["gate"].T # (S,64)
|
||||
probs = F.softmax(logits.float(), dim=-1)
|
||||
w, idx = torch.topk(probs, self.topk, dim=-1) # (S,topk)
|
||||
if self.norm_topk:
|
||||
w = w / w.sum(-1, keepdim=True)
|
||||
out = torch.zeros_like(x)
|
||||
# raggruppo per expert: per ogni expert davvero usato, processo i suoi token
|
||||
for eid in torch.unique(idx).tolist():
|
||||
sel = (idx == eid) # (S,topk) bool
|
||||
rows = sel.any(dim=-1).nonzero(as_tuple=True)[0]
|
||||
if rows.numel() == 0:
|
||||
continue
|
||||
gw, uw, dw = self.cache.get(layer, eid) # <-- streaming dal disco (bf16)
|
||||
xe = x[rows].to(torch.bfloat16) # calcolo expert in bf16
|
||||
h = (F.silu(xe @ gw.T) * (xe @ uw.T)) @ dw.T
|
||||
weight = (w[rows] * sel[rows].float()).sum(-1, keepdim=True)
|
||||
out[rows] += weight * h.float()
|
||||
return out
|
||||
|
||||
@torch.no_grad()
|
||||
def _step(self, ids_new, pos, kv):
|
||||
"""Un passo del modello sui token nuovi. Ritorna logit dell'ultimo token."""
|
||||
x = self.embed[torch.tensor(ids_new)] # (S,hidden)
|
||||
for i, lw in enumerate(self.layers):
|
||||
x = x + self._attn(lw, rmsnorm(x, lw["in_ln"], self.eps), pos, i, kv)
|
||||
x = x + self._moe(i, lw, rmsnorm(x, lw["post_ln"], self.eps))
|
||||
x = rmsnorm(x, self.final_norm, self.eps)
|
||||
return (x @ self.lm_head.T)[-1]
|
||||
|
||||
@torch.no_grad()
|
||||
def generate(self, token_ids, n_new, greedy=True):
|
||||
kv = [None] * self.L
|
||||
ids = list(token_ids)
|
||||
# PREFILL: tutti i token del prompt in un colpo, riempie la kv-cache
|
||||
logit = self._step(ids, torch.arange(len(ids)), kv)
|
||||
for s in range(n_new):
|
||||
nxt = int(torch.argmax(logit)) if greedy else int(torch.multinomial(F.softmax(logit, -1), 1))
|
||||
ids.append(nxt)
|
||||
if s == n_new - 1:
|
||||
break
|
||||
# DECODE: un solo token nuovo, usa la kv-cache (qui la cache expert torna a funzionare)
|
||||
logit = self._step([nxt], torch.tensor([len(ids) - 1]), kv)
|
||||
return ids
|
||||
@@ -1,18 +0,0 @@
|
||||
"""Profila dove va il tempo: lettura expert dal disco vs attenzione vs moe vs matmul."""
|
||||
import cProfile, pstats, io, glob, json
|
||||
from engine import OlmoeStreaming
|
||||
|
||||
snap = glob.glob("/home/vincenzo/.cache/huggingface/hub/models--allenai--OLMoE-1B-7B-0924/snapshots/*")[0]
|
||||
ref = json.load(open("ref.json"))
|
||||
m = OlmoeStreaming(snap, expert_cap=16)
|
||||
|
||||
pr = cProfile.Profile()
|
||||
pr.enable()
|
||||
m.generate(ref["prompt_ids"], 8, greedy=True) # 8 token bastano per il profilo
|
||||
pr.disable()
|
||||
|
||||
s = io.StringIO()
|
||||
ps = pstats.Stats(pr, stream=s).sort_stats("tottime")
|
||||
ps.print_stats(15)
|
||||
print(s.getvalue())
|
||||
print(f"Hit-rate: {m.cache.hitrate()*100:.1f}% hit={m.cache.hits} miss={m.cache.miss}")
|
||||
@@ -1,22 +0,0 @@
|
||||
"""La domanda che conta: a quanti bit l'output degli expert REGGE ancora?
|
||||
Quantizzo solo gli expert (la parte densa resta bf16) e confronto col riferimento."""
|
||||
import json, glob
|
||||
from engine import OlmoeStreaming
|
||||
|
||||
snap = glob.glob("/home/vincenzo/.cache/huggingface/hub/models--allenai--OLMoE-1B-7B-0924/snapshots/*")[0]
|
||||
ref = json.load(open("ref.json"))
|
||||
exp = ref["full_ids"][len(ref["prompt_ids"]):]
|
||||
n_new = len(exp)
|
||||
|
||||
print(f"{'bit':>4} {'MB/expert':>10} {'match':>7} testo")
|
||||
for bits in (16, 8, 4, 3, 2):
|
||||
m = OlmoeStreaming(snap, expert_cap=64, quant_bits=bits) # cap64: isola l'effetto quant
|
||||
out = m.generate(ref["prompt_ids"], n_new, greedy=True)
|
||||
gen = out[len(ref["prompt_ids"]):]
|
||||
match = sum(a == b for a, b in zip(gen, exp))
|
||||
mb = 6.29 * bits / 8 / 1.0 # ~6.29M param/expert * bit / 8 -> MB
|
||||
# decode testo per vedere se e' ancora sensato
|
||||
from transformers import AutoTokenizer
|
||||
tok = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924")
|
||||
txt = tok.decode(gen)
|
||||
print(f"{bits:>4} {mb:>9.1f}M {match:>4}/{n_new:<2} {txt!r}")
|
||||
@@ -1,31 +0,0 @@
|
||||
"""Lancia il motore streaming, confronta con il riferimento, misura RAM/hit-rate/velocita'."""
|
||||
import json, time, glob, sys, resource
|
||||
from engine import OlmoeStreaming
|
||||
|
||||
CAP = int(sys.argv[1]) if len(sys.argv) > 1 else 16 # expert residenti per layer (su 64)
|
||||
snap = glob.glob("/home/vincenzo/.cache/huggingface/hub/models--allenai--OLMoE-1B-7B-0924/snapshots/*")[0]
|
||||
ref = json.load(open("run_ref.json" if False else "ref.json"))
|
||||
|
||||
def rss_gb():
|
||||
return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / (1024**2) # KB->GB su Linux
|
||||
|
||||
print(f"== Motore streaming, cache = {CAP} expert/layer su 64 ==")
|
||||
t = time.time()
|
||||
m = OlmoeStreaming(snap, expert_cap=CAP)
|
||||
print(f"densa caricata in {m.dense_load_s:.1f}s | RSS dopo load densa: {rss_gb():.2f} GB")
|
||||
|
||||
n_new = len(ref["full_ids"]) - len(ref["prompt_ids"])
|
||||
t = time.time()
|
||||
out = m.generate(ref["prompt_ids"], n_new, greedy=True)
|
||||
dt = time.time() - t
|
||||
|
||||
# confronto
|
||||
gen = out[len(ref["prompt_ids"]):]
|
||||
exp = ref["full_ids"][len(ref["prompt_ids"]):]
|
||||
match = sum(a == b for a, b in zip(gen, exp))
|
||||
print(f"\nRiferimento (transformers): {exp}")
|
||||
print(f"Motore streaming : {gen}")
|
||||
print(f"Token coincidenti: {match}/{len(exp)}")
|
||||
print(f"\nRSS PICCO: {rss_gb():.2f} GB (modello completo in bf16 = ~13 GB)")
|
||||
print(f"Hit-rate cache expert: {m.cache.hitrate()*100:.1f}% (hit={m.cache.hits} miss={m.cache.miss})")
|
||||
print(f"Velocita': {n_new/dt:.2f} tok/s ({dt:.1f}s per {n_new} token, no kv-cache)")
|
||||
@@ -1,97 +0,0 @@
|
||||
"""
|
||||
Stadio 0 - Cost model + benchmark del disco per lo streaming degli expert MoE.
|
||||
|
||||
Domanda: e' FISICAMENTE possibile fare streaming degli expert da disco
|
||||
e generare a velocita' usabile, su QUESTA macchina?
|
||||
|
||||
Due numeri che servono:
|
||||
1. Banda effettiva del disco in lettura RANDOM, a blocchi grossi quanto un expert.
|
||||
2. Quanti byte/token dobbiamo leggere -> da cui il tetto di token/sec.
|
||||
|
||||
Nessun modello richiesto. Gira in secondi.
|
||||
"""
|
||||
import os, sys, time, mmap, random, argparse
|
||||
|
||||
MB = 1024 * 1024
|
||||
GB = 1024 * MB
|
||||
|
||||
|
||||
def bench_disk(path_dir, expert_mb=12.0, total_mb=2048, n_reads=200):
|
||||
"""Crea un file, poi misura lettura sequenziale e random a chunk = un expert."""
|
||||
os.makedirs(path_dir, exist_ok=True)
|
||||
fpath = os.path.join(path_dir, "_bench.bin")
|
||||
chunk = int(expert_mb * MB)
|
||||
total = int(total_mb * MB)
|
||||
total = (total // chunk) * chunk
|
||||
n_chunks = total // chunk
|
||||
|
||||
# scrittura
|
||||
t = time.time()
|
||||
with open(fpath, "wb") as f:
|
||||
buf = os.urandom(chunk)
|
||||
for _ in range(n_chunks):
|
||||
f.write(buf)
|
||||
f.flush(); os.fsync(f.fileno())
|
||||
write_bw = total / (time.time() - t) / GB
|
||||
|
||||
# prova a buttare via la page cache (best effort, serve permessi su Linux nativo)
|
||||
try:
|
||||
os.system("sync")
|
||||
with open("/proc/sys/vm/drop_caches", "w") as c:
|
||||
c.write("3")
|
||||
except Exception:
|
||||
pass # su /mnt/c o senza root non si puo': il numero sara' ottimistico
|
||||
|
||||
f = open(fpath, "rb")
|
||||
mm = mmap.mmap(f.fileno(), 0, prot=mmap.PROT_READ)
|
||||
|
||||
# random reads a chunk di un expert
|
||||
idxs = [random.randrange(n_chunks) for _ in range(n_reads)]
|
||||
s = 0
|
||||
t = time.time()
|
||||
for i in idxs:
|
||||
off = i * chunk
|
||||
s += mm[off] # tocca la prima pagina
|
||||
s += mm[off + chunk - 1] # e l'ultima -> forza il caricamento del range
|
||||
_ = bytes(mm[off:off + chunk]) # legge davvero l'intero expert
|
||||
rand_bw = (n_reads * chunk) / (time.time() - t) / GB
|
||||
|
||||
mm.close(); f.close()
|
||||
os.remove(fpath)
|
||||
return write_bw, rand_bw
|
||||
|
||||
|
||||
def cost_model(name, n_layers, n_active, expert_mb, disk_bw_gbs, ram_resident_gb):
|
||||
"""Stampa il tetto di token/sec in funzione dell'hit-rate della cache."""
|
||||
bytes_cold = n_layers * n_active * expert_mb / 1024 # GB letti per token se 0 cache
|
||||
print(f"\n--- {name} ---")
|
||||
print(f" layer={n_layers} expert_attivi/layer={n_active} expert={expert_mb:.1f} MB")
|
||||
print(f" parte densa residente in RAM stimata: ~{ram_resident_gb:.1f} GB")
|
||||
print(f" byte da streammare per token (cache fredda): {bytes_cold*1024:.0f} MB")
|
||||
print(f" tetto token/sec @ banda {disk_bw_gbs:.2f} GB/s, al variare dell'hit-rate cache:")
|
||||
for hit in (0.0, 0.5, 0.8, 0.9, 0.95, 0.99):
|
||||
eff = bytes_cold * (1 - hit)
|
||||
tps = disk_bw_gbs / eff if eff > 0 else float("inf")
|
||||
print(f" hit {hit*100:5.1f}% -> {tps:6.2f} tok/s")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--dir", default=".", help="cartella su cui benchmarkare il disco")
|
||||
ap.add_argument("--expert-mb", type=float, default=12.0)
|
||||
args = ap.parse_args()
|
||||
|
||||
print(f"Benchmark disco su: {os.path.abspath(args.dir)} (chunk={args.expert_mb} MB)")
|
||||
wbw, rbw = bench_disk(args.dir, expert_mb=args.expert_mb)
|
||||
print(f" scrittura seq : {wbw:.2f} GB/s")
|
||||
print(f" lettura random: {rbw:.2f} GB/s <-- numero che conta per lo streaming")
|
||||
|
||||
# Scenari. expert_mb a Q4 ~ (hidden*inter*3)*0.5B.
|
||||
# OLMoE 1B-7B: 16 layer, 8 attivi, hidden 2048 inter 1024 -> ~3 MB Q4
|
||||
cost_model("OLMoE 1B-7B (piccolo, lo useremo allo Stadio 1)",
|
||||
n_layers=16, n_active=8, expert_mb=3.0,
|
||||
disk_bw_gbs=rbw, ram_resident_gb=1.0)
|
||||
# DeepSeek-V3/V4 class: ~60 layer MoE, 8 attivi, expert ~6 MB Q2
|
||||
cost_model("DeepSeek/GLM class @ Q2 (il sogno finale)",
|
||||
n_layers=60, n_active=8, expert_mb=6.0,
|
||||
disk_bw_gbs=rbw, ram_resident_gb=10.0)
|
||||
@@ -1,126 +0,0 @@
|
||||
"""
|
||||
Pilastro 2 - La domanda che decide tutto:
|
||||
quanto e' alto l'hit-rate di una cache di expert su un router MoE VERO?
|
||||
|
||||
Facciamo girare OLMoE-1B-7B su testo reale, registriamo per ogni (layer, posizione)
|
||||
quali 8 expert su 64 vengono attivati, e poi SIMULIAMO diverse politiche di cache
|
||||
al variare della capacita' K (quanti expert/layer teniamo residenti in RAM).
|
||||
|
||||
Output: hit-rate per policy e per K -> mappato su token/sec col cost model.
|
||||
"""
|
||||
import sys, time, collections
|
||||
import torch
|
||||
|
||||
MODEL = "allenai/OLMoE-1B-7B-0924"
|
||||
TOPK = 8
|
||||
N_EXPERTS = 64
|
||||
N_LAYERS = 16
|
||||
EXPERT_MB = 12.6 # bf16
|
||||
|
||||
# Banda misurata allo Stadio 0 (random read, ext4). Aggiorna se rifai il bench.
|
||||
DISK_BW_GBS = 7.33
|
||||
|
||||
PROMPTS = [
|
||||
"The history of the Roman Empire spans over a thousand years, from the founding "
|
||||
"of the city to the fall of Constantinople. Its legacy in law, language and "
|
||||
"engineering still shapes the modern world.",
|
||||
"def quicksort(arr):\n if len(arr) <= 1:\n return arr\n pivot = arr[len(arr)//2]\n"
|
||||
" left = [x for x in arr if x < pivot]\n middle = [x for x in arr if x == pivot]\n"
|
||||
" right = [x for x in arr if x > pivot]\n return quicksort(left) + middle + quicksort(right)",
|
||||
"In quantum mechanics, the wave function encodes the probability amplitude of a "
|
||||
"particle's state. Measurement collapses this superposition into a definite outcome, "
|
||||
"a process that remains philosophically contested.",
|
||||
"La politica monetaria della banca centrale influenza i tassi di interesse, "
|
||||
"l'inflazione e l'occupazione. Alzare i tassi raffredda la domanda ma rischia "
|
||||
"di rallentare la crescita economica.",
|
||||
]
|
||||
|
||||
|
||||
def collect_trace():
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
print("Carico il modello (bf16, CPU)...", flush=True)
|
||||
t = time.time()
|
||||
tok = AutoTokenizer.from_pretrained(MODEL)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
MODEL, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True)
|
||||
model.eval()
|
||||
print(f" caricato in {time.time()-t:.0f}s", flush=True)
|
||||
|
||||
# trace[layer] = lista (in ordine di token) di tuple di 8 expert id
|
||||
trace = [[] for _ in range(N_LAYERS)]
|
||||
for pi, p in enumerate(PROMPTS):
|
||||
ids = tok(p, return_tensors="pt").input_ids
|
||||
with torch.no_grad():
|
||||
out = model(ids, output_router_logits=True, use_cache=False)
|
||||
# out.router_logits: tupla di N_LAYERS tensori (n_token, n_experts)
|
||||
for li, rl in enumerate(out.router_logits):
|
||||
topk = rl.topk(TOPK, dim=-1).indices # (n_token, 8)
|
||||
for row in topk.tolist():
|
||||
trace[li].append(tuple(sorted(row)))
|
||||
print(f" prompt {pi}: {ids.shape[1]} token", flush=True)
|
||||
return trace
|
||||
|
||||
|
||||
def simulate(trace, K, policy="lru"):
|
||||
"""Cache per-layer di capacita' K. Ritorna hit-rate globale sugli accessi a expert."""
|
||||
hits = total = 0
|
||||
for li in range(N_LAYERS):
|
||||
if policy == "lru":
|
||||
cache = collections.OrderedDict()
|
||||
elif policy == "lfu":
|
||||
cache = {} # eid -> freq
|
||||
freq = collections.Counter()
|
||||
for experts in trace[li]:
|
||||
for e in experts:
|
||||
total += 1
|
||||
if policy == "lru":
|
||||
if e in cache:
|
||||
hits += 1
|
||||
cache.move_to_end(e)
|
||||
else:
|
||||
cache[e] = 1
|
||||
if len(cache) > K:
|
||||
cache.popitem(last=False)
|
||||
elif policy == "lfu":
|
||||
freq[e] += 1
|
||||
if e in cache:
|
||||
hits += 1
|
||||
else:
|
||||
if len(cache) >= K:
|
||||
victim = min(cache, key=lambda x: freq[x])
|
||||
del cache[victim]
|
||||
cache[e] = 1
|
||||
return hits / total if total else 0.0
|
||||
|
||||
|
||||
def consecutive_reuse(trace):
|
||||
"""Frazione di expert al token t gia' attivi al token t-1 (stesso layer)."""
|
||||
same = tot = 0
|
||||
for li in range(N_LAYERS):
|
||||
seq = trace[li]
|
||||
for t in range(1, len(seq)):
|
||||
prev = set(seq[t-1]); cur = set(seq[t])
|
||||
same += len(prev & cur); tot += len(cur)
|
||||
return same / tot if tot else 0.0
|
||||
|
||||
|
||||
def tok_per_sec(hitrate):
|
||||
bytes_cold_gb = N_LAYERS * TOPK * EXPERT_MB / 1024
|
||||
eff = bytes_cold_gb * (1 - hitrate)
|
||||
return DISK_BW_GBS / eff if eff > 0 else float("inf")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
trace = collect_trace()
|
||||
ntok = sum(len(trace[0]) for _ in [0])
|
||||
print(f"\nToken totali tracciati: {len(trace[0])} x {N_LAYERS} layer")
|
||||
print(f"Riuso consecutivo (expert in comune t vs t-1): {consecutive_reuse(trace)*100:.1f}%")
|
||||
|
||||
print("\nHit-rate cache per-layer al variare di K (expert residenti su 64):")
|
||||
print(f"{'K':>4} {'RAM/GB':>7} {'LRU':>8} {'LFU':>8} {'tok/s@LRU':>10}")
|
||||
for K in (8, 12, 16, 24, 32, 48):
|
||||
ram = K * N_LAYERS * EXPERT_MB / 1024
|
||||
hl = simulate(trace, K, "lru")
|
||||
hf = simulate(trace, K, "lfu")
|
||||
print(f"{K:>4} {ram:>6.1f}G {hl*100:>7.1f}% {hf*100:>7.1f}% {tok_per_sec(hl):>9.1f}")
|
||||
print("\nNota: K=8 e' il minimo teorico (8 attivi/token). K=64 = tutto in RAM (no streaming).")
|
||||
-110
@@ -1,110 +0,0 @@
|
||||
"""
|
||||
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%).")
|
||||
@@ -1,20 +0,0 @@
|
||||
"""Sweep della cache: per ogni capacita' misura correttezza, hit-rate, RAM cache, velocita'."""
|
||||
import json, time, glob
|
||||
from engine import OlmoeStreaming
|
||||
|
||||
snap = glob.glob("/home/vincenzo/.cache/huggingface/hub/models--allenai--OLMoE-1B-7B-0924/snapshots/*")[0]
|
||||
ref = json.load(open("ref.json"))
|
||||
exp = ref["full_ids"][len(ref["prompt_ids"]):]
|
||||
n_new = len(exp)
|
||||
EXPERT_MB_BF16 = 12.6
|
||||
|
||||
print(f"{'cap':>4} {'RAMcache':>9} {'match':>6} {'hit%':>6} {'tok/s':>7} {'sec':>6}")
|
||||
for cap in (16, 32, 48, 64):
|
||||
m = OlmoeStreaming(snap, expert_cap=cap)
|
||||
t = time.time()
|
||||
out = m.generate(ref["prompt_ids"], n_new, greedy=True)
|
||||
dt = time.time() - t
|
||||
gen = out[len(ref["prompt_ids"]):]
|
||||
match = sum(a == b for a, b in zip(gen, exp))
|
||||
ram = cap * m.L * EXPERT_MB_BF16 / 1024
|
||||
print(f"{cap:>4} {ram:>7.1f}GB {match:>3}/{n_new:<2} {m.cache.hitrate()*100:>5.1f}% {n_new/dt:>7.2f} {dt:>6.1f}")
|
||||
@@ -1,17 +0,0 @@
|
||||
"""Genera il riferimento con transformers (greedy) e lo salva. Va lanciato da solo (usa ~13GB)."""
|
||||
import json, torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
PROMPT = "The capital of France is"
|
||||
N = 12
|
||||
tok = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924")
|
||||
ids = tok(PROMPT, return_tensors="pt").input_ids
|
||||
model = AutoModelForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0924",
|
||||
torch_dtype=torch.bfloat16, low_cpu_mem_usage=True).eval()
|
||||
with torch.no_grad():
|
||||
out = model.generate(ids, max_new_tokens=N, do_sample=False)
|
||||
full = out[0].tolist()
|
||||
json.dump({"prompt": PROMPT, "prompt_ids": ids[0].tolist(), "full_ids": full,
|
||||
"text": tok.decode(full)}, open("ref.json", "w"))
|
||||
print("RIFERIMENTO salvato:", repr(tok.decode(full)))
|
||||
print("full_ids:", full)
|
||||
Reference in New Issue
Block a user