diff --git a/.gitignore b/.gitignore index b71fba4..667044d 100644 --- a/.gitignore +++ b/.gitignore @@ -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 diff --git a/README.md b/README.md index 19f42e0..663d3ab 100644 --- a/README.md +++ b/README.md @@ -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ì" diff --git a/engine.py b/engine.py deleted file mode 100644 index 31b41fb..0000000 --- a/engine.py +++ /dev/null @@ -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("= 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 diff --git a/profile_run.py b/profile_run.py deleted file mode 100644 index b821850..0000000 --- a/profile_run.py +++ /dev/null @@ -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}") diff --git a/quant_test.py b/quant_test.py deleted file mode 100644 index 2a6f7f8..0000000 --- a/quant_test.py +++ /dev/null @@ -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}") diff --git a/run.py b/run.py deleted file mode 100644 index c4acda0..0000000 --- a/run.py +++ /dev/null @@ -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)") diff --git a/s0_costmodel.py b/s0_costmodel.py deleted file mode 100644 index 160ef76..0000000 --- a/s0_costmodel.py +++ /dev/null @@ -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) diff --git a/s1_trace_hitrate.py b/s1_trace_hitrate.py deleted file mode 100644 index 11c884b..0000000 --- a/s1_trace_hitrate.py +++ /dev/null @@ -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).") diff --git a/s2_research.py b/s2_research.py deleted file mode 100644 index 00bb90a..0000000 --- a/s2_research.py +++ /dev/null @@ -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%).") diff --git a/sweep.py b/sweep.py deleted file mode 100644 index d088b7a..0000000 --- a/sweep.py +++ /dev/null @@ -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}") diff --git a/validate_ref.py b/validate_ref.py deleted file mode 100644 index 9c26bb9..0000000 --- a/validate_ref.py +++ /dev/null @@ -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)