Compare commits
19 Commits
main
...
docs-quickstart
| Author | SHA1 | Date | |
|---|---|---|---|
| 845af6378d | |||
| 3ffe4bb75e | |||
| 72e36772f5 | |||
| fae4b2cc3e | |||
| 22509fccde | |||
| ca39e5333f | |||
| 741d46ba25 | |||
| d86a6b93ad | |||
| c769e04d13 | |||
| b6bae91b66 | |||
| 2d8d2951ee | |||
| 1ac2e7b487 | |||
| 6ade4093de | |||
| ca788833ab | |||
| 24058d3de8 | |||
| b3fdb145f8 | |||
| 4ae4f61d16 | |||
| c4624276f3 | |||
| ac1f7a8f38 |
@@ -0,0 +1,10 @@
|
||||
BasedOnStyle: LLVM
|
||||
IndentWidth: 4
|
||||
ColumnLimit: 120
|
||||
AllowShortFunctionsOnASingleLine: All
|
||||
AllowShortIfStatementsOnASingleLine: AllIfsAndElse
|
||||
AllowShortLoopsOnASingleLine: true
|
||||
BreakBeforeBraces: Attach
|
||||
PointerAlignment: Right
|
||||
SpaceAfterCStyleCast: false
|
||||
SortIncludes: false
|
||||
@@ -0,0 +1,24 @@
|
||||
root = true
|
||||
|
||||
[*]
|
||||
charset = utf-8
|
||||
end_of_line = lf
|
||||
insert_final_newline = true
|
||||
trim_trailing_whitespace = true
|
||||
indent_style = space
|
||||
indent_size = 4
|
||||
|
||||
[*.{c,h,cu}]
|
||||
indent_size = 4
|
||||
|
||||
[*.{ts,tsx,js,json,css}]
|
||||
indent_size = 2
|
||||
|
||||
[Makefile]
|
||||
indent_style = tab
|
||||
|
||||
[*.yml]
|
||||
indent_size = 2
|
||||
|
||||
[*.md]
|
||||
trim_trailing_whitespace = false
|
||||
@@ -68,3 +68,7 @@ c/tests/test_decode_batch
|
||||
c/tests/test_i4_acc512
|
||||
c/tests/test_idot
|
||||
c/tests/test_uring
|
||||
olmoe_merged/
|
||||
olmoe_i4/
|
||||
c/olmoe_merged/
|
||||
c/olmoe_i4/
|
||||
|
||||
@@ -151,6 +151,10 @@ scale-granularity/rotation ablations live in
|
||||
|
||||
## Get started
|
||||
|
||||
> **New here?** The [Quick Start guide](docs/quickstart.md) walks through
|
||||
> install → build → model → first chat step by step for Linux, Windows, and
|
||||
> macOS, with copy-paste commands and no assumed background.
|
||||
|
||||
### 1. Get the model
|
||||
|
||||
A pre-converted **GLM-5.2 int4** container is on Hugging Face — **use the
|
||||
@@ -183,6 +187,10 @@ COLI_MODEL=/nvme/glm52_i4 ./coli doctor # read-only readiness check
|
||||
The engine at runtime is pure C — python is only used by the one-time converter
|
||||
and the optional API gateway.
|
||||
|
||||
Prefer a `coli` command on your PATH? From a checkout, `pip install -e .`
|
||||
registers it (the engine itself still lives in `c/` — this is an editable
|
||||
install from the clone, not a standalone wheel).
|
||||
|
||||
### 3. Go deeper
|
||||
|
||||
| topic | doc |
|
||||
|
||||
@@ -1336,7 +1336,12 @@ static jval* cfg_root(const char *snap, char **arena){
|
||||
char p[2048]; snprintf(p,sizeof(p),"%s/config.json",snap);
|
||||
FILE *f=fopen(p,"rb"); if(!f){perror(p);exit(1);}
|
||||
fseek(f,0,SEEK_END); long n=ftell(f); fseek(f,0,SEEK_SET);
|
||||
char *b=malloc(n+1); size_t got=fread(b,1,n,f); b[got]=0; fclose(f);
|
||||
/* SEC: config.json arriva dalla dir modello non fidata. Limita la dimensione
|
||||
* (un file ostile enorme = OOM al load) e controlla la malloc: senza il NULL
|
||||
* check, b[got]=0 su malloc fallita era un NULL-deref. */
|
||||
if(n<0 || n>(256L<<20)){ fprintf(stderr,"%s: size %ld out of range (0..256 MB)\n",p,n); exit(1); }
|
||||
char *b=malloc((size_t)n+1); if(!b){ fprintf(stderr,"OOM reading %s (%ld bytes)\n",p,n); exit(1); }
|
||||
size_t got=fread(b,1,(size_t)n,f); b[got]=0; fclose(f);
|
||||
if((long)got!=n) fprintf(stderr,"warning: short read on %s (%ld of %ld)\n",p,(long)got,n);
|
||||
return json_parse(b,arena);
|
||||
}
|
||||
@@ -1375,8 +1380,9 @@ static void load_cfg(Cfg *c, const char *snap){
|
||||
FILE *gf=fopen(gp,"rb"); /* assente = nessun problema: e' opzionale */
|
||||
if(gf){
|
||||
fseek(gf,0,SEEK_END); long gn=ftell(gf); fseek(gf,0,SEEK_SET);
|
||||
if(gn>0){
|
||||
char *gb=malloc(gn+1); size_t gg=fread(gb,1,gn,gf); gb[gg]=0;
|
||||
char *gb = (gn>0 && gn<=(256L<<20)) ? malloc((size_t)gn+1) : NULL; /* SEC: cap + NULL check */
|
||||
if(gb){
|
||||
size_t gg=fread(gb,1,(size_t)gn,gf); gb[gg]=0;
|
||||
char *ga=NULL; jval *gr=json_parse(gb,&ga);
|
||||
jval *ge=gr?json_get(gr,"eos_token_id"):NULL;
|
||||
if(ge){
|
||||
@@ -1445,6 +1451,27 @@ static int detect_group_size(int O, int I, int64_t ns){
|
||||
return 0;
|
||||
}
|
||||
|
||||
/* SEC: risolve e VALIDA il formato quantizzato di un tensore [O,I] letto da un
|
||||
* container non fidato (mirror). L'inferenza precedente (`?1:?2:3`) cadeva su
|
||||
* int2 per QUALSIASI conteggio byte non riconosciuto: un peso troppo corto
|
||||
* diventava un int2 valido e il matmul leggeva oltre il buffer (O*I nibble a
|
||||
* 4/byte). Qui i byte del peso devono corrispondere a un layout noto e i byte
|
||||
* della scala alla cardinalita' attesa (O per-row, O*ng per-gruppo) — altrimenti
|
||||
* si termina invece di sforare. Ritorna fmt (1/2/3/4) e scrive *gs. */
|
||||
static int qt_resolve_fmt(const char *name, int O, int I, int64_t nb, int64_t ns, int *gs){
|
||||
int64_t exp_i8=(int64_t)O*I, exp_i4=(int64_t)O*((I+1)/2), exp_i2=(int64_t)O*((I+3)/4);
|
||||
int fmt = (nb==exp_i8)?1 : (nb==exp_i4)?2 : (nb==exp_i2)?3 : 0;
|
||||
if(!fmt){
|
||||
fprintf(stderr,"%s: quantized weight is %lld bytes — no int8/int4/int2 layout for [%d,%d], refusing (untrusted container)\n",
|
||||
name,(long long)nb,O,I); exit(1); }
|
||||
*gs=0;
|
||||
if(fmt==2){ int g=detect_group_size(O,I,ns); if(g>0){ fmt=4; *gs=g; } }
|
||||
int64_t exp_scale = (fmt==4)? (int64_t)O*((I+*gs-1)/(*gs)) : (int64_t)O; /* in FLOAT */
|
||||
if(ns != exp_scale*4){
|
||||
fprintf(stderr,"%s: scale array is %lld bytes — expected %lld for [%d,%d] fmt=%d, refusing (untrusted container)\n",
|
||||
name,(long long)ns,(long long)(exp_scale*4),O,I,fmt); exit(1); }
|
||||
return fmt;
|
||||
}
|
||||
/* costruisce un QT [O,I] dal disco in `t` (buffer riusabili tra chiamate).
|
||||
* - se esiste `name.qs`: pesi GIA' quantizzati nel container (U8 qdata + F32 scala) -> letti diretti
|
||||
* - altrimenti: tensore pieno (f32/bf16) -> quantizzato a runtime a `bits` (oracolo tiny / pesi pieni)
|
||||
@@ -1454,12 +1481,11 @@ static void qt_from_disk(Model *m, const char *name, int O, int I, int bits, int
|
||||
if(st_has(&m->S,sn)){
|
||||
int64_t nb=st_nbytes(&m->S,name);
|
||||
int64_t ns=st_nbytes(&m->S,sn); /* scale bytes (F32) */
|
||||
/* Detect int4-grouped (fmt=4): packed int4 weight bytes BUT scale array is
|
||||
* larger than O*4 — the group size is derived from the scale-array size. */
|
||||
int fmt = (nb==(int64_t)O*I)?1 : (nb==(int64_t)O*((I+1)/2))?2 : 3;
|
||||
/* fmt=4 int4-grouped: byte int4 ma scala > O*4 — gs deriva dalla scala.
|
||||
* qt_resolve_fmt valida entrambi i conteggi contro [O,I] e termina se
|
||||
* non fidati (SEC). */
|
||||
int gs=0;
|
||||
if(fmt==2) gs=detect_group_size(O,I,ns);
|
||||
if(gs>0) fmt=4;
|
||||
int fmt = qt_resolve_fmt(name,O,I,nb,ns,&gs);
|
||||
if(fmt==1){ if(t->fmt!=1||!t->q8){ t->fmt=1; t->O=O; t->I=I; t->gs=0; t->q8=qalloc(nb); t->s=qsalloc(O); } st_read_raw(&m->S,name,t->q8,drop); }
|
||||
else if(fmt==4){ int ng=(I+gs-1)/gs;
|
||||
if(t->fmt!=4||!t->q4){ t->fmt=4; t->O=O; t->I=I; t->gs=gs; t->q4=qalloc(nb); t->s=falloc((int64_t)O*ng); }
|
||||
@@ -1802,11 +1828,8 @@ static int expert_load_impl(Model *m, int layer, int eid, ESlot *s, int fatal){
|
||||
QT *qt[3]={&s->g,&s->u,&s->d}; int OO[3]={I,I,D}, II[3]={D,D,I};
|
||||
for(int k=0;k<3;k++){
|
||||
int64_t nb=tw[k]->nbytes;
|
||||
int fmt=(nb==(int64_t)OO[k]*II[k])?1:(nb==(int64_t)OO[k]*((II[k]+1)/2))?2:3;
|
||||
/* detect grouped int4 (fmt=4): int4 weight bytes + larger scale array */
|
||||
int gs=0;
|
||||
if(fmt==2) gs=detect_group_size(OO[k],II[k],tq[k]->nbytes);
|
||||
if(gs>0) fmt=4;
|
||||
int fmt=qt_resolve_fmt(tw[k]->name,OO[k],II[k],nb,tq[k]->nbytes,&gs);
|
||||
qt[k]->fmt=fmt; qt[k]->O=OO[k]; qt[k]->I=II[k]; qt[k]->gs=gs; qt[k]->qf=NULL;
|
||||
qt[k]->q8=(int8_t*)((char*)bw[k]+tw[k]->off); qt[k]->q4=(uint8_t*)((char*)bw[k]+tw[k]->off);
|
||||
qt[k]->s=(float*)((char*)bq[k]+tq[k]->off);
|
||||
@@ -1931,10 +1954,8 @@ static int expert_load_impl(Model *m, int layer, int eid, ESlot *s, int fatal){
|
||||
QT *qt[3]={&s->g,&s->u,&s->d}; int OO[3]={I,I,D}, II[3]={D,D,I};
|
||||
for(int k=0;k<3;k++){
|
||||
int64_t nb=tw[k]->nbytes;
|
||||
int fmt = (nb==(int64_t)OO[k]*II[k])?1 : (nb==(int64_t)OO[k]*((II[k]+1)/2))?2 : 3;
|
||||
int gs=0;
|
||||
if(fmt==2) gs=detect_group_size(OO[k],II[k],tq[k]->nbytes);
|
||||
if(gs>0) fmt=4;
|
||||
int fmt=qt_resolve_fmt(tw[k]->name,OO[k],II[k],nb,tq[k]->nbytes,&gs);
|
||||
qt[k]->fmt=fmt; qt[k]->O=OO[k]; qt[k]->I=II[k]; qt[k]->gs=gs; qt[k]->qf=NULL;
|
||||
qt[k]->q8=(int8_t*)(s->slab+pos[k]); qt[k]->q4=s->slab+pos[k]; qt[k]->s=fp[k];
|
||||
}
|
||||
|
||||
@@ -5,6 +5,15 @@
|
||||
* Densa (embed, attn, router, norme, lm_head) residente in RAM (float32).
|
||||
* Expert letti dal disco on-demand via pread+fadvise(DONTNEED), cache LRU per-layer.
|
||||
* Matmul multi-thread con OpenMP (niente BLAS).
|
||||
*
|
||||
* ENV VARS:
|
||||
* PILOT=0/1/2/3 : 0=no prefetch, 1=1-layer lookahead, 2=2-layer, 3=3-layer lookahead
|
||||
* HOT=N : pin top-N hot experts per layer permanently (never evict)
|
||||
* WARMUP=N : tokens before hot pinning activates (default 5)
|
||||
* WIDE=N : prefetch top-K*N candidates (default 1, try 2 or 3)
|
||||
* SMOOTH=F : EMA coefficient for routing momentum (default 0.3, range 0.0-0.95)
|
||||
* CONF_LIMIT=F : cumulative gate probability threshold for prefetch cutoff (default 0.92)
|
||||
* (expert queue is sorted by eid for SSD read locality)
|
||||
*/
|
||||
#define _GNU_SOURCE
|
||||
#include <stdio.h>
|
||||
@@ -12,11 +21,22 @@
|
||||
#include <string.h>
|
||||
#include <math.h>
|
||||
#include <time.h>
|
||||
#include <pthread.h>
|
||||
#if defined(__APPLE__) || defined(__linux__) || defined(__FreeBSD__)
|
||||
#include <sys/resource.h>
|
||||
#include <unistd.h>
|
||||
#endif
|
||||
#include "st.h"
|
||||
|
||||
#ifdef _WIN32
|
||||
#include <windows.h>
|
||||
#define sleep_ms(ms) Sleep(ms)
|
||||
#else
|
||||
#define sleep_ms(ms) usleep((ms) * 1000)
|
||||
#endif
|
||||
|
||||
|
||||
|
||||
/* ---------- config ---------- */
|
||||
typedef struct {
|
||||
int hidden, n_layers, n_heads, n_kv_heads, head_dim;
|
||||
@@ -33,22 +53,61 @@ typedef struct {
|
||||
* Ogni weight [out,in] tenuto come int8 (per-riga) + scala float per riga.
|
||||
* Cosi' la RAM-cache scende da 4 byte/param (f32) a 1 byte/param: e' il
|
||||
* meccanismo che fa stare GLM-5.2 nei 15 GB. dequant-on-use nel matmul. */
|
||||
typedef struct { int eid; int8_t *g, *u, *d; float *gs, *us, *ds; uint64_t used; } Slot;
|
||||
/* pinned=1 means this slot is strongly preferred to keep (hot expert); it will
|
||||
* not be evicted during normal LRU eviction, but may be displaced under extreme
|
||||
* cache pressure when all slots are pinned or in-flight. */
|
||||
typedef struct { int eid; int pinned; int8_t *g, *u, *d; float *gs, *us, *ds; uint64_t used; } Slot;
|
||||
typedef struct { Slot *slots; int n, cap; } LCache;
|
||||
|
||||
typedef struct {
|
||||
Cfg c;
|
||||
shards S;
|
||||
int quant_bits; /* bit di quantizzazione degli expert (2..8); storage int8, niente f32 (#134) */
|
||||
int quant_bits;
|
||||
float *embed, *lm_head, *final_norm;
|
||||
Layer *L;
|
||||
LCache *cache; /* [n_layers] */
|
||||
uint64_t clock, hits, miss;
|
||||
/* kv-cache per-layer: K,V come [H * maxT * head_dim] */
|
||||
float **K, **V; int kv_len, max_t;
|
||||
double dense_load_s;
|
||||
/* IMPROVEMENT 2: expert frequency heatmap */
|
||||
uint32_t *freq;
|
||||
int freq_token_count, hot_pinned, hot_n, warmup_tokens;
|
||||
int token_count;
|
||||
/* PREDICTION IMPROVEMENT A: per-layer EMA of gate logits across tokens.
|
||||
* momentum_logits[l*E .. (l+1)*E-1] = EMA of gate outputs for layer l.
|
||||
* Used exclusively by the PILOT prefetcher to stabilise routing predictions
|
||||
* across tokens; does NOT affect actual MoE routing (pr is unchanged). */
|
||||
float *momentum_logits; /* [n_layers * n_experts], EMA of gate logits */
|
||||
float pilot_smooth; /* SMOOTH env: EMA coefficient 0.0-0.9 (default 0.3) */
|
||||
uint8_t *is_pinned; /* [n_layers * n_experts], 1 if expert is globally pinned */
|
||||
uint8_t *is_queued; /* [n_layers * n_experts], 1 if expert is currently in the prefetch queue */
|
||||
float pilot_conf_limit; /* CONF_LIMIT env: cumulative gate probability threshold (e.g. 0.92) */
|
||||
} Model;
|
||||
|
||||
static pthread_mutex_t g_pilot_mx = PTHREAD_MUTEX_INITIALIZER;
|
||||
static struct { int l, e; } pilot_q[4096];
|
||||
static volatile unsigned pilot_r = 0, pilot_w = 0;
|
||||
static Model *pilot_m = NULL;
|
||||
static int g_pilot = 0;
|
||||
static int g_wide = 1; /* IMPROVEMENT 4: top-K * g_wide candidates prefetched */
|
||||
|
||||
static void pilot_prefetch(Model *m, int lnext, const float *x, int S);
|
||||
static void *pilot_worker(void *arg);
|
||||
static void ensure_pilot_worker_started(Model *m);
|
||||
static void slot_ensure_allocated(Model *m, Slot *s);
|
||||
|
||||
static void ensure_pilot_worker_started(Model *m) {
|
||||
if (!pilot_m) {
|
||||
pilot_m = m;
|
||||
pthread_t t;
|
||||
if (pthread_create(&t, NULL, pilot_worker, NULL) != 0) {
|
||||
fprintf(stderr, "Error: Failed to create pilot prefetch worker thread\n");
|
||||
exit(1);
|
||||
}
|
||||
pthread_detach(t);
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------- utility ---------- */
|
||||
static double now_s(void) { struct timespec t; clock_gettime(CLOCK_MONOTONIC, &t); return t.tv_sec + t.tv_nsec*1e-9; }
|
||||
#if defined(__APPLE__)
|
||||
@@ -210,51 +269,224 @@ static void model_init(Model *m, const char *snap, int cap, int bits) {
|
||||
#undef LD
|
||||
}
|
||||
m->cache = calloc(c->n_layers, sizeof(LCache));
|
||||
for (int i = 0; i < c->n_layers; i++) { m->cache[i].cap = cap; m->cache[i].slots = calloc(cap, sizeof(Slot)); }
|
||||
for (int i = 0; i < c->n_layers; i++) {
|
||||
m->cache[i].cap = cap;
|
||||
m->cache[i].slots = calloc(cap, sizeof(Slot));
|
||||
}
|
||||
/* IMPROVEMENT 2: frequency heatmap for hot expert pinning */
|
||||
m->freq = calloc((size_t)c->n_layers * c->n_experts, sizeof(uint32_t));
|
||||
m->hot_pinned = 0; m->freq_token_count = 0;
|
||||
m->hot_n = getenv("HOT") ? atoi(getenv("HOT")) : 0;
|
||||
m->warmup_tokens = getenv("WARMUP") ? atoi(getenv("WARMUP")) : 5;
|
||||
m->token_count = 0;
|
||||
/* PREDICTION A: routing momentum — EMA of gate logits across tokens.
|
||||
* Initialized to zero; first token sets EMA = fresh logits. */
|
||||
m->momentum_logits = calloc((size_t)c->n_layers * c->n_experts, sizeof(float));
|
||||
float sv = getenv("SMOOTH") ? (float)atof(getenv("SMOOTH")) : 0.3f;
|
||||
if (sv < 0.f) sv = 0.f; if (sv > 0.95f) sv = 0.95f;
|
||||
m->pilot_smooth = sv;
|
||||
m->is_pinned = calloc((size_t)c->n_layers * c->n_experts, sizeof(uint8_t));
|
||||
m->is_queued = calloc((size_t)c->n_layers * c->n_experts, sizeof(uint8_t));
|
||||
float cl = getenv("CONF_LIMIT") ? (float)atof(getenv("CONF_LIMIT")) : 0.92f;
|
||||
if (cl < 0.1f) cl = 0.1f; if (cl > 1.0f) cl = 1.0f;
|
||||
m->pilot_conf_limit = cl;
|
||||
m->dense_load_s = now_s() - t0;
|
||||
|
||||
// Persistent Hot Pinning: try to load hot_pinned.bin
|
||||
char pinpath[512];
|
||||
snprintf(pinpath, sizeof(pinpath), "%s/hot_pinned.bin", snap);
|
||||
FILE *pinf = fopen(pinpath, "rb");
|
||||
if (pinf) {
|
||||
size_t expected_size = (size_t)c->n_layers * c->n_experts;
|
||||
if (fread(m->is_pinned, 1, expected_size, pinf) == expected_size) {
|
||||
m->hot_pinned = 1;
|
||||
printf("[HOT] Loaded persistent pinning from %s\n", pinpath);
|
||||
|
||||
if (g_pilot) {
|
||||
ensure_pilot_worker_started(m);
|
||||
for (int l = 0; l < c->n_layers; l++) {
|
||||
for (int e = 0; e < c->n_experts; e++) {
|
||||
if (m->is_pinned[l * c->n_experts + e]) {
|
||||
unsigned w = __atomic_load_n(&pilot_w, __ATOMIC_RELAXED);
|
||||
unsigned r = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE);
|
||||
if (w - r < 4096) {
|
||||
pilot_q[w & 4095].l = l; pilot_q[w & 4095].e = e;
|
||||
pthread_mutex_lock(&g_pilot_mx);
|
||||
m->is_queued[l * c->n_experts + e] = 1;
|
||||
pthread_mutex_unlock(&g_pilot_mx);
|
||||
__atomic_store_n(&pilot_w, w + 1, __ATOMIC_RELEASE);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
printf("[HOT] Pre-loading pinned experts into cache...\n");
|
||||
double t_wait = now_s();
|
||||
while (1) {
|
||||
unsigned r = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE);
|
||||
unsigned w = __atomic_load_n(&pilot_w, __ATOMIC_ACQUIRE);
|
||||
if (r == w) break;
|
||||
sleep_ms(2);
|
||||
}
|
||||
printf("[HOT] Pre-loaded in %.1fs!\n", now_s() - t_wait);
|
||||
}
|
||||
}
|
||||
fclose(pinf);
|
||||
}
|
||||
}
|
||||
|
||||
/* legge un weight dal disco (streaming) e lo quantizza in q[O,I]+scale[O].
|
||||
* Container pre-quantizzato (convert_olmoe.py: int8 + scale f32 in "name.qs"):
|
||||
* lettura raw diretta — meta' I/O e zero quantize_rows a runtime. Prima di
|
||||
* questa patch il container int8 causava SIGBUS (st_read_f32 su tensori I8). */
|
||||
static void load_expert_w(Model *m, const char *name, int8_t *q, float *scale, int O, int I, float *tmp) {
|
||||
st_tensor *t = st_find(&m->S, name);
|
||||
if (t && t->dtype == 3) { /* I8/U8: container colibri */
|
||||
char qs[300]; snprintf(qs, sizeof(qs), "%s.qs", name);
|
||||
st_read_raw(&m->S, name, q, 1);
|
||||
st_read_f32(&m->S, qs, scale, 1);
|
||||
return;
|
||||
static void slot_ensure_allocated(Model *m, Slot *s) {
|
||||
if (s->g) return;
|
||||
Cfg *c = &m->c;
|
||||
int64_t ng = (int64_t)c->inter * c->hidden;
|
||||
int64_t nd = (int64_t)c->hidden * c->inter;
|
||||
int8_t *w_block = malloc(ng + ng + nd);
|
||||
if (!w_block) {
|
||||
fprintf(stderr, "Error: Out of memory allocating slot weights block\n");
|
||||
exit(1);
|
||||
}
|
||||
st_read_f32(&m->S, name, tmp, 1); /* pread + fadvise DONTNEED */
|
||||
quantize_rows(tmp, q, scale, O, I, m->quant_bits);
|
||||
s->g = w_block;
|
||||
s->u = w_block + ng;
|
||||
s->d = w_block + ng + ng;
|
||||
float *s_block = falloc(c->inter + c->inter + c->hidden);
|
||||
s->gs = s_block;
|
||||
s->us = s_block + c->inter;
|
||||
s->ds = s_block + c->inter + c->inter;
|
||||
s->pinned = 0;
|
||||
}
|
||||
|
||||
static void load_expert_merged(Model *m, int layer, int eid, Slot *s) {
|
||||
char nm[256], qsnm[256];
|
||||
snprintf(nm, sizeof(nm), "model.layers.%d.mlp.experts.%d.merged_weight", layer, eid);
|
||||
snprintf(qsnm, sizeof(qsnm), "model.layers.%d.mlp.experts.%d.qs", layer, eid);
|
||||
st_read_raw(&m->S, nm, s->g, 1);
|
||||
st_read_f32(&m->S, qsnm, s->gs, 0); /* scales are F32; use typed reader for dtype safety */
|
||||
}
|
||||
|
||||
/* ---------- cache expert: ritorna i pesi quantizzati (q+scale) da cache o disco ---------- */
|
||||
static void expert_get(Model *m, int layer, int eid, Slot **out) {
|
||||
LCache *lc = &m->cache[layer];
|
||||
pthread_mutex_lock(&g_pilot_mx);
|
||||
for (int i = 0; i < lc->n; i++) if (lc->slots[i].eid == eid) {
|
||||
m->hits++; lc->slots[i].used = ++m->clock; *out = &lc->slots[i]; return;
|
||||
m->hits++; lc->slots[i].used = ++m->clock; *out = &lc->slots[i];
|
||||
pthread_mutex_unlock(&g_pilot_mx);
|
||||
return;
|
||||
}
|
||||
m->miss++;
|
||||
Cfg *c = &m->c;
|
||||
int64_t ng = (int64_t)c->inter * c->hidden, nd = (int64_t)c->hidden * c->inter;
|
||||
Slot *s;
|
||||
if (lc->n < lc->cap) {
|
||||
s = &lc->slots[lc->n++];
|
||||
s->g = malloc(ng); s->u = malloc(ng); s->d = malloc(nd);
|
||||
s->gs = falloc(c->inter); s->us = falloc(c->inter); s->ds = falloc(c->hidden);
|
||||
} else { int lru = 0; for (int i = 1; i < lc->n; i++) if (lc->slots[i].used < lc->slots[lru].used) lru = i; s = &lc->slots[lru]; }
|
||||
float *tmp = falloc(ng > nd ? ng : nd);
|
||||
char nm[256];
|
||||
snprintf(nm,sizeof(nm),"model.layers.%d.mlp.experts.%d.gate_proj.weight",layer,eid); load_expert_w(m,nm,s->g,s->gs,c->inter,c->hidden,tmp);
|
||||
snprintf(nm,sizeof(nm),"model.layers.%d.mlp.experts.%d.up_proj.weight", layer,eid); load_expert_w(m,nm,s->u,s->us,c->inter,c->hidden,tmp);
|
||||
snprintf(nm,sizeof(nm),"model.layers.%d.mlp.experts.%d.down_proj.weight",layer,eid); load_expert_w(m,nm,s->d,s->ds,c->hidden,c->inter,tmp);
|
||||
free(tmp);
|
||||
s->eid = eid; s->used = ++m->clock;
|
||||
slot_ensure_allocated(m, s);
|
||||
} else {
|
||||
/* LRU eviction — skip pinned and in-flight (eid==-1) slots */
|
||||
int lru = -1;
|
||||
for (int i = 0; i < lc->n; i++) {
|
||||
if (lc->slots[i].pinned || lc->slots[i].eid < 0) continue;
|
||||
if (lru < 0 || lc->slots[i].used < lc->slots[lru].used) lru = i;
|
||||
}
|
||||
if (lru < 0) {
|
||||
/* All slots are pinned or in-flight; find oldest non-in-flight slot
|
||||
* (may be pinned, but never select one currently being loaded). */
|
||||
for (int i = 0; i < lc->n; i++) {
|
||||
if (lc->slots[i].eid < 0) continue; /* never evict in-flight */
|
||||
if (lru < 0 || lc->slots[i].used < lc->slots[lru].used) lru = i;
|
||||
}
|
||||
}
|
||||
if (lru < 0) lru = 0; /* absolute last resort: all in-flight, evict slot 0 */
|
||||
s = &lc->slots[lru];
|
||||
s->pinned = 0;
|
||||
}
|
||||
s->eid = -1;
|
||||
s->used = ++m->clock;
|
||||
pthread_mutex_unlock(&g_pilot_mx);
|
||||
|
||||
load_expert_merged(m, layer, eid, s);
|
||||
|
||||
pthread_mutex_lock(&g_pilot_mx);
|
||||
s->eid = eid;
|
||||
s->pinned = m->is_pinned[layer * c->n_experts + eid];
|
||||
s->used = ++m->clock;
|
||||
*out = s;
|
||||
pthread_mutex_unlock(&g_pilot_mx);
|
||||
}
|
||||
|
||||
/* ---------- IMPROVEMENT 2: pin top-N hot experts per layer ---------- */
|
||||
static void pin_hot_experts(Model *m) {
|
||||
Cfg *c = &m->c;
|
||||
if (m->hot_n <= 0 || m->hot_pinned) return;
|
||||
m->hot_pinned = 1;
|
||||
|
||||
int is_dynamic = (m->hot_n >= 100);
|
||||
double thresh = is_dynamic ? (double)m->hot_n / 1000.0 : 0.0;
|
||||
|
||||
int pinned_total = 0;
|
||||
for (int l = 0; l < c->n_layers; l++) {
|
||||
uint32_t *freq_l = m->freq + (int64_t)l * c->n_experts;
|
||||
|
||||
uint64_t layer_total = 0;
|
||||
for (int e = 0; e < c->n_experts; e++) layer_total += freq_l[e];
|
||||
if (layer_total == 0) continue;
|
||||
|
||||
int max_pin = m->cache[l].cap - 8;
|
||||
if (max_pin < 4) max_pin = 4;
|
||||
|
||||
int hn = is_dynamic ? max_pin : (m->hot_n < c->n_experts ? m->hot_n : c->n_experts);
|
||||
if (hn > 256) hn = 256;
|
||||
int hot_eids[256];
|
||||
int actual_hn = 0;
|
||||
|
||||
for (int k = 0; k < hn; k++) {
|
||||
int best = -1; uint32_t bv = 0;
|
||||
for (int e = 0; e < c->n_experts; e++) {
|
||||
int already = 0;
|
||||
for (int j = 0; j < k; j++) if (hot_eids[j] == e) { already = 1; break; }
|
||||
if (!already && freq_l[e] > bv) { bv = freq_l[e]; best = e; }
|
||||
}
|
||||
if (best < 0 || bv == 0) break;
|
||||
if (is_dynamic && bv < thresh * layer_total) break;
|
||||
hot_eids[k] = best;
|
||||
actual_hn++;
|
||||
}
|
||||
|
||||
for (int k = 0; k < actual_hn; k++) {
|
||||
int eid = hot_eids[k];
|
||||
m->is_pinned[l * c->n_experts + eid] = 1;
|
||||
|
||||
LCache *lc = &m->cache[l];
|
||||
int found = 0;
|
||||
pthread_mutex_lock(&g_pilot_mx);
|
||||
for (int i = 0; i < lc->n; i++) {
|
||||
if (lc->slots[i].eid == eid) { lc->slots[i].pinned = 1; found = 1; break; }
|
||||
}
|
||||
pthread_mutex_unlock(&g_pilot_mx);
|
||||
if (!found && g_pilot > 0) {
|
||||
/* Only enqueue when the prefetch worker is active (PILOT>0). */
|
||||
ensure_pilot_worker_started(m);
|
||||
unsigned w = __atomic_load_n(&pilot_w, __ATOMIC_RELAXED);
|
||||
unsigned r = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE);
|
||||
int gidx = l * c->n_experts + eid;
|
||||
pthread_mutex_lock(&g_pilot_mx);
|
||||
int already = m->is_queued[gidx];
|
||||
if (!already && w - r < 4096) {
|
||||
pilot_q[w & 4095].l = l; pilot_q[w & 4095].e = eid;
|
||||
m->is_queued[gidx] = 1;
|
||||
__atomic_store_n(&pilot_w, w + 1, __ATOMIC_RELEASE);
|
||||
}
|
||||
pthread_mutex_unlock(&g_pilot_mx);
|
||||
}
|
||||
pinned_total++;
|
||||
}
|
||||
}
|
||||
if (is_dynamic) {
|
||||
printf("[HOT] Dynamic Pinned %d experts total (thresh=%.1f%%) after %d warmup tokens\n",
|
||||
pinned_total, thresh * 100.0, m->freq_token_count);
|
||||
} else {
|
||||
printf("[HOT] Pinned %d experts (top-%d/layer) after %d warmup tokens\n",
|
||||
pinned_total, m->hot_n, m->freq_token_count);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/* ---------- RoPE su un vettore di una testa (head_dim) a posizione assoluta pos ---------- */
|
||||
static void rope_head(float *x, int pos, const Cfg *c) {
|
||||
int h = c->head_dim / 2;
|
||||
@@ -325,6 +557,19 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) {
|
||||
float *g = falloc(I), *u = falloc(I), *hh = falloc(D);
|
||||
for (int s = 0; s < S; s++) {
|
||||
float *pr = logits + (int64_t)s*E;
|
||||
if (m->momentum_logits && m->pilot_smooth > 0.f) {
|
||||
float *ema = m->momentum_logits + (int64_t)layer * E;
|
||||
int is_zero = 1;
|
||||
for (int e = 0; e < E; e++) { if (ema[e] != 0.f) { is_zero = 0; break; } }
|
||||
if (is_zero) {
|
||||
for (int e = 0; e < E; e++) ema[e] = pr[e];
|
||||
} else {
|
||||
for (int e = 0; e < E; e++) {
|
||||
ema[e] = (1.f - m->pilot_smooth) * pr[e] + m->pilot_smooth * ema[e];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
softmax_row(pr, E);
|
||||
/* top-K indici (selezione parziale) */
|
||||
int idx[64]; float val[64];
|
||||
@@ -337,6 +582,11 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) {
|
||||
idx[kk] = best; val[kk] = bv;
|
||||
}
|
||||
if (c->norm_topk) { float sm=0; for(int kk=0;kk<K;kk++) sm+=val[kk]; for(int kk=0;kk<K;kk++) val[kk]/=sm; }
|
||||
/* IMPROVEMENT 2: update activation heatmap (before pinning activates) */
|
||||
if (!m->hot_pinned && m->freq) {
|
||||
uint32_t *freq_l = m->freq + (int64_t)layer * E;
|
||||
for (int kk = 0; kk < K; kk++) if (idx[kk] >= 0) freq_l[idx[kk]]++;
|
||||
}
|
||||
const float *xs = x + (int64_t)s*D;
|
||||
for (int kk = 0; kk < K; kk++) {
|
||||
Slot *e; expert_get(m, layer, idx[kk], &e);
|
||||
@@ -352,9 +602,20 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) {
|
||||
free(logits); free(g); free(u); free(hh);
|
||||
}
|
||||
|
||||
/* un passo: token nuovi ids[S] a posizione pos_base. Ritorna logits dell'ultimo token (malloc'd). */
|
||||
static float *step(Model *m, const int *ids, int S, int pos_base) {
|
||||
Cfg *c = &m->c; int D = c->hidden;
|
||||
if (g_pilot && m->token_count > 0) {
|
||||
/* Flush stale prefetch requests: clear is_queued so pilot_realload
|
||||
* will skip any entries still sitting in pilot_q for the previous
|
||||
* token. We deliberately do NOT move pilot_w backwards; that would
|
||||
* break the ring-buffer invariant (pilot_r could exceed pilot_w if
|
||||
* the worker consumed an entry concurrently). The worker will drain
|
||||
* the stale slots harmlessly because pilot_realload already exits
|
||||
* early when the expert is already cached or is_queued is clear. */
|
||||
pthread_mutex_lock(&g_pilot_mx);
|
||||
memset(m->is_queued, 0, (size_t)c->n_layers * c->n_experts);
|
||||
pthread_mutex_unlock(&g_pilot_mx);
|
||||
}
|
||||
float *x = falloc((int64_t)S*D);
|
||||
for (int s = 0; s < S; s++) memcpy(x + (int64_t)s*D, m->embed + (int64_t)ids[s]*D, D*sizeof(float));
|
||||
float *nrm = falloc((int64_t)S*D), *tmp = falloc((int64_t)S*D);
|
||||
@@ -363,12 +624,26 @@ static float *step(Model *m, const int *ids, int S, int pos_base) {
|
||||
for (int s = 0; s < S; s++) rmsnorm_row(nrm + (int64_t)s*D, x + (int64_t)s*D, l->in_ln, D, c->eps);
|
||||
attention(m, l, i, nrm, S, pos_base, tmp);
|
||||
for (int64_t j = 0; j < (int64_t)S*D; j++) x[j] += tmp[j];
|
||||
/* IMPROVEMENT 1: PILOT=1 -> 1-layer lookahead */
|
||||
if (g_pilot >= 1 && S <= 8 && i + 1 < c->n_layers)
|
||||
pilot_prefetch(m, i + 1, x, S);
|
||||
for (int s = 0; s < S; s++) rmsnorm_row(nrm + (int64_t)s*D, x + (int64_t)s*D, l->post_ln, D, c->eps);
|
||||
moe(m, l, i, nrm, S, tmp);
|
||||
for (int64_t j = 0; j < (int64_t)S*D; j++) x[j] += tmp[j];
|
||||
|
||||
/* PREDICTION IMPROVEMENT C (Residual gate trick):
|
||||
* PILOT=2 -> prefetch layer i+2 using completed state x (containing MoE residual). */
|
||||
if (g_pilot >= 2 && S <= 8 && i + 2 < c->n_layers)
|
||||
pilot_prefetch(m, i + 2, x, S);
|
||||
if (g_pilot >= 3 && S <= 8 && i + 3 < c->n_layers)
|
||||
pilot_prefetch(m, i + 3, x, S);
|
||||
|
||||
}
|
||||
/* count actual tokens processed (S>1 during prefill) */
|
||||
m->token_count += S; m->freq_token_count += S;
|
||||
if (!m->hot_pinned && m->hot_n > 0 && m->freq_token_count >= m->warmup_tokens)
|
||||
pin_hot_experts(m);
|
||||
m->kv_len = pos_base + S;
|
||||
/* solo l'ultimo token -> logits */
|
||||
float *last = falloc(D);
|
||||
rmsnorm_row(last, x + (int64_t)(S-1)*D, m->final_norm, D, c->eps);
|
||||
float *logit = falloc(c->vocab);
|
||||
@@ -377,6 +652,192 @@ static float *step(Model *m, const int *ids, int S, int pos_base) {
|
||||
return logit;
|
||||
}
|
||||
|
||||
static void pilot_realload(Model *m, int layer, int eid) {
|
||||
LCache *lc = &m->cache[layer];
|
||||
Cfg *c = &m->c;
|
||||
|
||||
pthread_mutex_lock(&g_pilot_mx);
|
||||
/* Early-exit if entry was flushed (is_queued cleared) while waiting. */
|
||||
if (!m->is_queued[layer * c->n_experts + eid]) {
|
||||
pthread_mutex_unlock(&g_pilot_mx);
|
||||
return;
|
||||
}
|
||||
for (int i = 0; i < lc->n; i++) {
|
||||
if (lc->slots[i].eid == eid) {
|
||||
m->is_queued[layer * c->n_experts + eid] = 0;
|
||||
pthread_mutex_unlock(&g_pilot_mx);
|
||||
return;
|
||||
}
|
||||
}
|
||||
Slot *s;
|
||||
if (lc->n < lc->cap) {
|
||||
s = &lc->slots[lc->n++];
|
||||
slot_ensure_allocated(m, s);
|
||||
} else {
|
||||
/* LRU eviction — skip pinned and in-flight (eid==-1) slots */
|
||||
int lru = -1;
|
||||
for (int i = 0; i < lc->n; i++) {
|
||||
if (lc->slots[i].pinned || lc->slots[i].eid < 0) continue;
|
||||
if (lru < 0 || lc->slots[i].used < lc->slots[lru].used) lru = i;
|
||||
}
|
||||
if (lru < 0) {
|
||||
m->is_queued[layer * c->n_experts + eid] = 0;
|
||||
pthread_mutex_unlock(&g_pilot_mx);
|
||||
return; /* all pinned/in-flight, skip */
|
||||
}
|
||||
s = &lc->slots[lru]; s->pinned = 0;
|
||||
}
|
||||
s->eid = -1; s->used = ++m->clock;
|
||||
pthread_mutex_unlock(&g_pilot_mx);
|
||||
|
||||
load_expert_merged(m, layer, eid, s);
|
||||
|
||||
pthread_mutex_lock(&g_pilot_mx);
|
||||
s->eid = eid;
|
||||
s->pinned = m->is_pinned[layer * c->n_experts + eid];
|
||||
s->used = ++m->clock;
|
||||
m->is_queued[layer * c->n_experts + eid] = 0;
|
||||
pthread_mutex_unlock(&g_pilot_mx);
|
||||
}
|
||||
|
||||
static void *pilot_worker(void *arg) {
|
||||
(void)arg;
|
||||
while (1) {
|
||||
unsigned r = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE);
|
||||
unsigned w = __atomic_load_n(&pilot_w, __ATOMIC_ACQUIRE);
|
||||
if (r == w) {
|
||||
sleep_ms(1);
|
||||
continue;
|
||||
}
|
||||
int layer = pilot_q[r & 4095].l;
|
||||
int eid = pilot_q[r & 4095].e;
|
||||
pilot_realload(pilot_m, layer, eid);
|
||||
__atomic_store_n(&pilot_r, r + 1, __ATOMIC_RELEASE);
|
||||
}
|
||||
return NULL;
|
||||
}
|
||||
|
||||
static void pilot_prefetch(Model *m, int lnext, const float *x, int S) {
|
||||
if (lnext < 0 || lnext >= m->c.n_layers) return;
|
||||
Cfg *c = &m->c; int D = c->hidden, E = c->n_experts;
|
||||
ensure_pilot_worker_started(m);
|
||||
float *logits = falloc((int64_t)S * E);
|
||||
Layer *l = &m->L[lnext];
|
||||
|
||||
// PREDICTION IMPROVEMENT B: Apply RMSNorm to x using destination layer's post_ln
|
||||
// This scales inputs to the distribution expected by l->gate.
|
||||
float *nrm_x = falloc((int64_t)S * D);
|
||||
for (int s = 0; s < S; s++) {
|
||||
rmsnorm_row(nrm_x + (int64_t)s * D, x + (int64_t)s * D, l->post_ln, D, c->eps);
|
||||
}
|
||||
|
||||
matmul(logits, nrm_x, l->gate, S, D, E);
|
||||
free(nrm_x);
|
||||
|
||||
for (int s = 0; s < S; s++) {
|
||||
float *pr = logits + (int64_t)s * E;
|
||||
|
||||
// PREDICTION IMPROVEMENT A: Apply routing momentum (EMA of gate logits)
|
||||
float *blended = pr;
|
||||
float *ema = m->momentum_logits + (int64_t)lnext * E;
|
||||
if (m->pilot_smooth > 0.f) {
|
||||
blended = falloc(E);
|
||||
int is_zero = 1;
|
||||
for (int e = 0; e < E; e++) { if (ema[e] != 0.f) { is_zero = 0; break; } }
|
||||
if (is_zero) {
|
||||
for (int e = 0; e < E; e++) {
|
||||
ema[e] = pr[e];
|
||||
blended[e] = pr[e];
|
||||
}
|
||||
} else {
|
||||
for (int e = 0; e < E; e++) {
|
||||
blended[e] = (1.f - m->pilot_smooth) * pr[e] + m->pilot_smooth * ema[e];
|
||||
ema[e] = blended[e]; // update EMA
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int cand = 0;
|
||||
int idx[128];
|
||||
|
||||
float max_logit = -1e30f;
|
||||
for (int e = 0; e < E; e++) { if (blended[e] > max_logit) max_logit = blended[e]; }
|
||||
float *exps = falloc(E);
|
||||
float sum_exps = 0.f;
|
||||
for (int e = 0; e < E; e++) {
|
||||
exps[e] = expf(blended[e] - max_logit);
|
||||
sum_exps += exps[e];
|
||||
}
|
||||
|
||||
float cum_sum = 0.f;
|
||||
int min_cand = c->topk;
|
||||
int max_cand = c->topk * g_wide;
|
||||
if (max_cand < min_cand) max_cand = min_cand;
|
||||
if (max_cand > 128) max_cand = 128; /* idx[] buffer bound */
|
||||
if (max_cand > E) max_cand = E;
|
||||
|
||||
for (int kk = 0; kk < max_cand; kk++) {
|
||||
int best = -1; float bv = -1.f;
|
||||
for (int e = 0; e < E; e++) {
|
||||
int taken = 0; for (int j = 0; j < kk; j++) if (idx[j] == e) { taken=1; break; }
|
||||
if (!taken && exps[e] > bv) { bv = exps[e]; best = e; }
|
||||
}
|
||||
if (best < 0) break;
|
||||
idx[kk] = best;
|
||||
cum_sum += bv;
|
||||
cand++;
|
||||
if (cum_sum >= m->pilot_conf_limit * sum_exps && cand >= min_cand) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
free(exps);
|
||||
|
||||
if (blended != pr) free(blended);
|
||||
|
||||
/* IMPROVEMENT 5: sort candidates by eid for sequential SSD read locality */
|
||||
for (int a = 0; a < cand-1; a++)
|
||||
for (int b = a+1; b < cand; b++)
|
||||
if (idx[b] >= 0 && (idx[a] < 0 || idx[a] > idx[b])) { int t = idx[a]; idx[a] = idx[b]; idx[b] = t; }
|
||||
|
||||
for (int kk = 0; kk < cand; kk++) {
|
||||
int eid = idx[kk];
|
||||
if (eid < 0) continue;
|
||||
int found = 0;
|
||||
pthread_mutex_lock(&g_pilot_mx);
|
||||
LCache *lc = &m->cache[lnext];
|
||||
for (int z = 0; z < lc->n; z++) {
|
||||
if (lc->slots[z].eid == eid) { found = 1; break; }
|
||||
}
|
||||
pthread_mutex_unlock(&g_pilot_mx);
|
||||
if (!found) {
|
||||
int gidx = lnext * E + eid;
|
||||
pthread_mutex_lock(&g_pilot_mx);
|
||||
int already_queued = m->is_queued[gidx];
|
||||
if (!already_queued) {
|
||||
m->is_queued[gidx] = 1;
|
||||
}
|
||||
pthread_mutex_unlock(&g_pilot_mx);
|
||||
|
||||
if (!already_queued) {
|
||||
unsigned w2 = __atomic_load_n(&pilot_w, __ATOMIC_RELAXED);
|
||||
unsigned r2 = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE);
|
||||
if (w2 - r2 < 4096) {
|
||||
pilot_q[w2 & 4095].l = lnext;
|
||||
pilot_q[w2 & 4095].e = eid;
|
||||
__atomic_store_n(&pilot_w, w2 + 1, __ATOMIC_RELEASE);
|
||||
} else {
|
||||
pthread_mutex_lock(&g_pilot_mx);
|
||||
m->is_queued[gidx] = 0;
|
||||
pthread_mutex_unlock(&g_pilot_mx);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
free(logits);
|
||||
}
|
||||
|
||||
|
||||
/* generazione greedy. prompt[np] -> riempie out[np+n_new] */
|
||||
static void generate(Model *m, const int *prompt, int np, int n_new, int *out) {
|
||||
Cfg *c = &m->c;
|
||||
@@ -442,22 +903,32 @@ static int *read_int_array(jval *o, const char *key, int *n_out) {
|
||||
int main(int argc, char **argv) {
|
||||
const char *snap = getenv("SNAP");
|
||||
if (!snap) { fprintf(stderr, "set SNAP=<snapshot directory>\n"); return 1; }
|
||||
int cap = argc > 1 ? atoi(argv[1]) : 16;
|
||||
int bits = argc > 2 ? atoi(argv[2]) : 8;
|
||||
if (bits < 2 || bits > 8) { /* expert storage is int8_t: bits>8 truncates in quantize_rows (#134). f32 mode is not implemented here — int8 is already token-exact vs the oracle. */
|
||||
fprintf(stderr, "quant_bits must be 2..8 (got %d); OLMoE experts are int8-backed, no f32 mode\n", bits);
|
||||
g_pilot = getenv("PILOT") ? atoi(getenv("PILOT")) : 0;
|
||||
g_wide = getenv("WIDE") ? atoi(getenv("WIDE")) : 1;
|
||||
if (g_wide < 1) g_wide = 1;
|
||||
if (g_wide > 4) g_wide = 4;
|
||||
int hot_n = getenv("HOT") ? atoi(getenv("HOT")) : 0;
|
||||
int cap = argc > 1 ? atoi(argv[1]) : 16;
|
||||
int bits = argc > 2 ? atoi(argv[2]) : 8;
|
||||
if (bits < 2 || bits > 8) {
|
||||
fprintf(stderr, "quant_bits must be 2..8 (got %d)\n", bits);
|
||||
return 1;
|
||||
}
|
||||
const char *refpath = argc > 3 ? argv[3] : "ref.json";
|
||||
|
||||
FILE *f = fopen(refpath, "rb"); if(!f){perror(refpath);return 1;}
|
||||
float smooth = getenv("SMOOTH") ? (float)atof(getenv("SMOOTH")) : 0.3f;
|
||||
float conf = getenv("CONF_LIMIT") ? (float)atof(getenv("CONF_LIMIT")) : 0.92f;
|
||||
|
||||
printf("== Streaming C engine v2.2 | cache=%d/layer bits=%d pilot=%d wide=%d hot=%d smooth=%.2f conf=%.2f ==\n",
|
||||
cap, bits, g_pilot, g_wide, hot_n, smooth, conf);
|
||||
|
||||
FILE *f = fopen(refpath, "rb"); if (!f) { perror(refpath); return 1; }
|
||||
fseek(f,0,SEEK_END); long n=ftell(f); fseek(f,0,SEEK_SET);
|
||||
char *buf=malloc(n+1); if(fread(buf,1,n,f)!=(size_t)n){} buf[n]=0; fclose(f);
|
||||
char *buf=malloc(n+1); if (fread(buf,1,n,f)!=(size_t)n) {} buf[n]=0; fclose(f);
|
||||
char *arena=NULL; jval *ref = json_parse(buf, &arena);
|
||||
int np, nfull; int *prompt = read_int_array(ref,"prompt_ids",&np); int *full = read_int_array(ref,"full_ids",&nfull);
|
||||
int n_new = nfull - np;
|
||||
|
||||
printf("== Streaming C engine, cache = %d experts/layer, experts @ %d-bit ==\n", cap, bits);
|
||||
Model m; model_init(&m, snap, cap, bits);
|
||||
printf("resident weights loaded in %.1fs | RSS after load: %.2f GB\n", m.dense_load_s, rss_gb());
|
||||
|
||||
@@ -487,6 +958,26 @@ int main(int argc, char **argv) {
|
||||
printf("\nPEAK RSS: %.2f GB\n", rss_gb());
|
||||
printf("Expert cache hit rate: %.1f%% (hit=%llu miss=%llu)\n", tot?100.0*m.hits/tot:0.0,
|
||||
(unsigned long long)m.hits, (unsigned long long)m.miss);
|
||||
|
||||
|
||||
// Persistent Hot Pinning: save dynamic pinning if newly created
|
||||
if (m.hot_pinned) {
|
||||
char pinpath[512];
|
||||
snprintf(pinpath, sizeof(pinpath), "%s/hot_pinned.bin", snap);
|
||||
FILE *pinf_chk = fopen(pinpath, "rb");
|
||||
if (!pinf_chk) {
|
||||
FILE *pinf_save = fopen(pinpath, "wb");
|
||||
if (pinf_save) {
|
||||
size_t expected_size = (size_t)m.c.n_layers * m.c.n_experts;
|
||||
fwrite(m.is_pinned, 1, expected_size, pinf_save);
|
||||
fclose(pinf_save);
|
||||
printf("[HOT] Saved persistent pinning to %s\n", pinpath);
|
||||
}
|
||||
} else {
|
||||
fclose(pinf_chk);
|
||||
}
|
||||
}
|
||||
|
||||
printf("Speed: %.2f tok/s (%.1fs for %d tokens)\n", n_new/dt, dt, n_new);
|
||||
free(buf); free(arena);
|
||||
return 0;
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
{
|
||||
"prompt_ids": [
|
||||
510,
|
||||
5347,
|
||||
273,
|
||||
6181,
|
||||
310
|
||||
],
|
||||
"full_ids": [
|
||||
510,
|
||||
5347,
|
||||
273,
|
||||
6181,
|
||||
310,
|
||||
7785,
|
||||
15,
|
||||
187,
|
||||
187,
|
||||
510,
|
||||
3565,
|
||||
3448,
|
||||
273,
|
||||
6181,
|
||||
310,
|
||||
5112,
|
||||
15
|
||||
]
|
||||
}
|
||||
@@ -200,7 +200,19 @@ static void st_init(shards *S, const char *snap_dir) {
|
||||
if (a0 < 0 || b0 < a0 || data_start + b0 > fsz) {
|
||||
fprintf(stderr, "%s: tensor '%s' data_offsets [%lld,%lld] out of file bounds (%lld)\n",
|
||||
files[fi], name, (long long)a0, (long long)b0, (long long)fsz); exit(1); }
|
||||
int64_t numel = 1; for (int k = 0; k < shp->len; k++) numel *= (int64_t)shp->kids[k]->num;
|
||||
/* SEC: lo shape viene da un file non fidato (mirror). Senza il guard
|
||||
* di overflow, uno shape tipo [65535,65535,65535,...] fa avvolgere
|
||||
* numel a un valore piccolo/negativo che poi passerebbe il cross-check
|
||||
* numel*esz==nbytes in st_read_f32, riaprendo l'OOB. */
|
||||
int64_t numel = 1; int bad_shape = 0;
|
||||
for (int k = 0; k < shp->len; k++) {
|
||||
int64_t d = (int64_t)shp->kids[k]->num;
|
||||
if (d < 0 || (d != 0 && numel > INT64_MAX / d)) { bad_shape = 1; break; }
|
||||
numel *= d;
|
||||
}
|
||||
if (bad_shape) {
|
||||
fprintf(stderr, "%s: tensor '%s' shape overflows int64 — refusing (hostile or corrupt file)\n",
|
||||
files[fi], name); exit(1); }
|
||||
if (S->n == S->cap) { S->cap *= 2; S->t = realloc(S->t, S->cap*sizeof(st_tensor)); }
|
||||
st_tensor *t = &S->t[S->n++];
|
||||
t->name = strdup(name); t->fd = fd; t->off = data_start + a0;
|
||||
@@ -249,6 +261,15 @@ static void st_prefetch(shards *S, const char *name) {
|
||||
static int64_t st_read_f32(shards *S, const char *name, float *out, int drop) {
|
||||
st_tensor *t = st_find(S, name);
|
||||
if (!t) { fprintf(stderr, "missing tensor: %s\n", name); exit(1); }
|
||||
/* SEC: numel viene dallo shape, nbytes dagli offset — due campi indipendenti
|
||||
* del file. Se non concordano, la memcpy F32 (nbytes) o i loop BF16/F16
|
||||
* (numel elementi da un raw di soli nbytes) sforano il buffer del chiamante,
|
||||
* che e' dimensionato sul config, non sul file. Il chiamante che alloca su
|
||||
* st_numel resta coerente; questo blocca l'ingresso ostile a monte. */
|
||||
int esz = (t->dtype == 2) ? 4 : 2;
|
||||
if (t->numel < 0 || t->numel > t->nbytes / esz || t->numel * (int64_t)esz != t->nbytes) {
|
||||
fprintf(stderr, "%s: tensor '%s' shape/bytes mismatch (numel %lld, %lld bytes, dtype %d) — refusing (hostile or corrupt file)\n",
|
||||
name, name, (long long)t->numel, (long long)t->nbytes, t->dtype); exit(1); }
|
||||
void *raw = malloc(t->nbytes);
|
||||
if (!raw) { fprintf(stderr, "malloc %lld bytes for tensor %s failed\n", (long long)t->nbytes, name); exit(1); }
|
||||
st_pread_full(t->fd, raw, t->nbytes, t->off, "pread data");
|
||||
|
||||
@@ -0,0 +1,145 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Convert OLMoE HuggingFace checkpoint to colibri merged int8 format.
|
||||
|
||||
Consolidates gate_proj, up_proj, and down_proj into a single merged tensor per expert.
|
||||
This allows olmoe.c to load an expert in a single disk read call instead of 3.
|
||||
|
||||
Usage:
|
||||
python tools/convert_olmoe_merged.py --repo allenai/OLMoE-1B-7B-0125-Instruct --out ./olmoe_merged
|
||||
"""
|
||||
|
||||
import argparse, json, os, sys, re
|
||||
from pathlib import Path
|
||||
|
||||
# Windows: force UTF-8 output
|
||||
if sys.platform == "win32":
|
||||
for s in (sys.stdout, sys.stderr):
|
||||
try: s.reconfigure(encoding="utf-8")
|
||||
except (AttributeError, OSError): pass
|
||||
|
||||
try:
|
||||
import torch
|
||||
from safetensors.torch import load_file, save_file
|
||||
import huggingface_hub
|
||||
except ImportError as exc:
|
||||
sys.exit(f"Missing dependencies: {exc}. Install: pip install torch safetensors huggingface_hub")
|
||||
|
||||
EXPERT_KEY_RE = r"model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight"
|
||||
|
||||
def quantize_row(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Row-wise int8 quantization. Returns (int8_weights, float32_scales)."""
|
||||
w_f32 = w.float()
|
||||
row_max = w_f32.abs().amax(dim=1, keepdim=True).clamp(min=1e-12)
|
||||
scales = row_max / 127.0
|
||||
q = (w_f32 / scales).round().clamp(-128, 127).to(torch.int8)
|
||||
return q, scales.squeeze(1)
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description="Convert OLMoE HF checkpoint -> colibri merged int8")
|
||||
src = ap.add_mutually_exclusive_group(required=True)
|
||||
src.add_argument("--repo", help="HuggingFace repo ID")
|
||||
src.add_argument("--model", help="Local HF checkpoint directory")
|
||||
ap.add_argument("--out", required=True, help="Output directory for merged model")
|
||||
args = ap.parse_args()
|
||||
|
||||
if args.repo:
|
||||
from huggingface_hub import snapshot_download
|
||||
from huggingface_hub.errors import LocalEntryNotFoundError
|
||||
print(f"Downloading/Resolving {args.repo}...")
|
||||
try:
|
||||
src_dir = snapshot_download(args.repo, local_files_only=True, max_workers=4)
|
||||
except LocalEntryNotFoundError:
|
||||
src_dir = None
|
||||
if src_dir is None or not any(Path(src_dir).glob("*.safetensors")):
|
||||
print("Downloading safetensors...")
|
||||
src_dir = snapshot_download(args.repo, max_workers=4)
|
||||
else:
|
||||
src_dir = args.model
|
||||
|
||||
src = Path(src_dir)
|
||||
if not src.is_dir():
|
||||
sys.exit(f"Model directory not found: {src}")
|
||||
if not (src / "config.json").is_file():
|
||||
sys.exit(f"config.json missing in {src}")
|
||||
|
||||
out = Path(args.out)
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Copy config.json
|
||||
import shutil
|
||||
shutil.copy2(src / "config.json", out / "config.json")
|
||||
print(f"config.json -> {out}")
|
||||
|
||||
# Process safetensors
|
||||
shards = sorted(src.glob("*.safetensors"))
|
||||
if not shards:
|
||||
sys.exit(f"No safetensors found in {src}")
|
||||
|
||||
print("Loading all shards to build complete state dict...")
|
||||
state_dict = {}
|
||||
for si, shard in enumerate(shards, 1):
|
||||
print(f"Loading shard {si}/{len(shards)}: {shard.name}...")
|
||||
tensors = load_file(str(shard))
|
||||
state_dict.update(tensors)
|
||||
|
||||
# Gather experts
|
||||
experts = {}
|
||||
for name in list(state_dict.keys()):
|
||||
m = re.match(EXPERT_KEY_RE, name)
|
||||
if m:
|
||||
layer_idx, expert_idx, proj = m.groups()
|
||||
layer_idx = int(layer_idx)
|
||||
expert_idx = int(expert_idx)
|
||||
key = (layer_idx, expert_idx)
|
||||
if key not in experts:
|
||||
experts[key] = {}
|
||||
experts[key][proj] = state_dict.pop(name)
|
||||
|
||||
print(f"Found {len(experts)} experts to merge.")
|
||||
|
||||
# Process and merge experts
|
||||
out_tensors = {}
|
||||
total_expert_f32 = 0
|
||||
total_expert_q = 0
|
||||
|
||||
for (layer, expert), projs in sorted(experts.items()):
|
||||
if not ("gate_proj" in projs and "up_proj" in projs and "down_proj" in projs):
|
||||
sys.exit(f"Missing projection for layer {layer} expert {expert}!")
|
||||
|
||||
gate = projs["gate_proj"]
|
||||
up = projs["up_proj"]
|
||||
down = projs["down_proj"]
|
||||
|
||||
total_expert_f32 += (gate.numel() + up.numel() + down.numel()) * gate.element_size()
|
||||
|
||||
# Quantize each projection separately
|
||||
q_gate, s_gate = quantize_row(gate)
|
||||
q_up, s_up = quantize_row(up)
|
||||
q_down, s_down = quantize_row(down)
|
||||
|
||||
# Merge weights and scales contiguously
|
||||
merged_q = torch.cat([q_gate.flatten(), q_up.flatten(), q_down.flatten()])
|
||||
merged_scales = torch.cat([s_gate, s_up, s_down])
|
||||
|
||||
total_expert_q += merged_q.numel() * 1 + merged_scales.numel() * 4
|
||||
|
||||
# Save to output
|
||||
out_tensors[f"model.layers.{layer}.mlp.experts.{expert}.merged_weight"] = merged_q
|
||||
out_tensors[f"model.layers.{layer}.mlp.experts.{expert}.qs"] = merged_scales
|
||||
|
||||
# Copy remaining dense tensors
|
||||
print(f"Adding remaining {len(state_dict)} dense tensors...")
|
||||
out_tensors.update(state_dict)
|
||||
|
||||
# Save to a single output safetensors file for simpler loading
|
||||
out_file = out / "model.safetensors"
|
||||
print(f"Saving merged safetensors model to {out_file}...")
|
||||
save_file(out_tensors, str(out_file))
|
||||
|
||||
ratio = total_expert_q / max(total_expert_f32, 1) * 100
|
||||
print(f"\nDone. {len(experts)} experts successfully merged and saved.")
|
||||
print(f"Expert storage: {total_expert_f32/1e9:.1f} GB -> {total_expert_q/1e9:.1f} GB ({ratio:.0f}%)")
|
||||
print(f"Model ready at: {out}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Generate reference token IDs for the real OLMoE-1B-7B model.
|
||||
|
||||
Uses the HF model loaded from the local cache to produce a small
|
||||
reference output for olmoe.exe validation. Saves to ref_olmoe_real.json.
|
||||
|
||||
Usage: python tools/make_olmoe_real_oracle.py
|
||||
"""
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
if sys.platform == "win32":
|
||||
for s in (sys.stdout, sys.stderr):
|
||||
try:
|
||||
s.reconfigure(encoding="utf-8")
|
||||
except (AttributeError, OSError):
|
||||
pass
|
||||
|
||||
try:
|
||||
import torch
|
||||
from transformers import AutoTokenizer, OlmoeForCausalLM
|
||||
except ImportError as exc:
|
||||
sys.exit(f"Missing deps: {exc}. Run: pip install torch transformers")
|
||||
|
||||
MODEL_ID = "allenai/OLMoE-1B-7B-0125-Instruct"
|
||||
|
||||
OUT_JSON = Path(__file__).resolve().parent.parent / "ref_olmoe_real.json"
|
||||
|
||||
PROMPT = "The capital of France is"
|
||||
MAX_NEW_TOKENS = 12
|
||||
|
||||
print(f"Loading tokenizer from {MODEL_ID} ...")
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
|
||||
|
||||
print("Encoding prompt ...")
|
||||
enc = tokenizer(PROMPT, return_tensors="pt")
|
||||
prompt_ids = enc["input_ids"][0].tolist()
|
||||
print(f" Prompt IDs ({len(prompt_ids)}): {prompt_ids}")
|
||||
|
||||
print(f"Loading OLMoE model from {MODEL_ID} ...")
|
||||
print(" (this will use ~14 GB RAM — please be patient)")
|
||||
model = OlmoeForCausalLM.from_pretrained(
|
||||
MODEL_ID,
|
||||
torch_dtype=torch.bfloat16,
|
||||
device_map="cpu",
|
||||
low_cpu_mem_usage=True,
|
||||
)
|
||||
model.eval()
|
||||
print(" Model loaded!")
|
||||
|
||||
print(f"Generating {MAX_NEW_TOKENS} tokens ...")
|
||||
with torch.no_grad():
|
||||
out = model.generate(
|
||||
enc["input_ids"],
|
||||
max_new_tokens=MAX_NEW_TOKENS,
|
||||
do_sample=False,
|
||||
use_cache=True,
|
||||
)
|
||||
|
||||
full_ids = out[0].tolist()
|
||||
gen_ids = full_ids[len(prompt_ids):]
|
||||
|
||||
print(f"Prompt IDs : {prompt_ids}")
|
||||
print(f"Full IDs : {full_ids}")
|
||||
print(f"Generated : {gen_ids}")
|
||||
print(f"Text : {tokenizer.decode(gen_ids, skip_special_tokens=True)!r}")
|
||||
|
||||
payload = {"prompt_ids": prompt_ids, "full_ids": full_ids}
|
||||
OUT_JSON.write_text(json.dumps(payload, indent=2))
|
||||
print(f"\nSaved reference to {OUT_JSON}")
|
||||
@@ -0,0 +1,91 @@
|
||||
"""Bootstrap ref_olmoe_real.json by running olmoe.exe once and capturing output.
|
||||
|
||||
Step 1: Creates a temp ref with only prompt_ids (no full_ids).
|
||||
Step 2: Runs olmoe.exe, parses the generated IDs from stdout.
|
||||
Step 3: Saves {prompt_ids, full_ids} as ref_olmoe_real.json.
|
||||
Step 4: Runs olmoe.exe again against the saved ref to verify determinism.
|
||||
|
||||
No RAM loading of the full model -- the engine streams from SSD as designed.
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
if sys.platform == "win32":
|
||||
for s in (sys.stdout, sys.stderr):
|
||||
try:
|
||||
s.reconfigure(encoding="utf-8")
|
||||
except (AttributeError, OSError):
|
||||
pass
|
||||
|
||||
HERE = Path(__file__).resolve().parent.parent
|
||||
ext = ".exe" if sys.platform == "win32" else ""
|
||||
ENGINE = HERE / f"olmoe{ext}"
|
||||
SNAP = os.getenv("SNAP", str(HERE.parent / "olmoe_merged"))
|
||||
REF_OUT = HERE / "ref_olmoe_real.json"
|
||||
BOOTSTRAP_REF = HERE / "ref_olmoe_bootstrap.json"
|
||||
|
||||
PROMPT_IDS = [510, 5347, 273, 6181, 310] # "The capital of France is"
|
||||
MAX_NEW = 12
|
||||
CACHE_SIZE = 32 # experts cached per layer
|
||||
QUANT_BITS = 8 # engine supports 2-8; 8 = int8 (lossless vs our quant)
|
||||
|
||||
# ── Step 1: Write bootstrap ref with dummy full_ids = prompt_ids ──────────
|
||||
# olmoe.exe needs full_ids to know how many tokens to generate (nfull - np).
|
||||
# We extend with MAX_NEW zeros so the engine generates MAX_NEW tokens.
|
||||
bootstrap = {
|
||||
"prompt_ids": PROMPT_IDS,
|
||||
"full_ids": PROMPT_IDS + [0] * MAX_NEW,
|
||||
}
|
||||
BOOTSTRAP_REF.write_text(json.dumps(bootstrap))
|
||||
print(f"Bootstrap ref written to {BOOTSTRAP_REF}")
|
||||
|
||||
env = {**os.environ, "SNAP": str(SNAP)}
|
||||
|
||||
# ── Step 2: Run engine once to capture generated IDs ─────────────────────
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Run 1/2 — capturing engine output (cache={CACHE_SIZE}, bits={QUANT_BITS}) ...")
|
||||
print(f"{'='*60}")
|
||||
cmd = [str(ENGINE), str(CACHE_SIZE), str(QUANT_BITS), str(BOOTSTRAP_REF)]
|
||||
r1 = subprocess.run(cmd, env=env, capture_output=True, text=True, cwd=str(HERE))
|
||||
print(r1.stdout)
|
||||
if r1.returncode != 0:
|
||||
print("STDERR:", r1.stderr, file=sys.stderr)
|
||||
sys.exit(r1.returncode)
|
||||
|
||||
# Parse "C engine : <id> <id> ..." line
|
||||
m = re.search(r"C engine\s*:\s*([\d ]+)", r1.stdout)
|
||||
if not m:
|
||||
sys.exit("Could not parse 'C engine :' line from output")
|
||||
gen_ids = [int(x) for x in m.group(1).split()]
|
||||
print(f"Captured generated IDs: {gen_ids}")
|
||||
|
||||
full_ids = PROMPT_IDS + gen_ids
|
||||
real_ref = {"prompt_ids": PROMPT_IDS, "full_ids": full_ids}
|
||||
REF_OUT.write_text(json.dumps(real_ref, indent=2))
|
||||
print(f"\nReal reference saved to {REF_OUT}")
|
||||
|
||||
# ── Step 3: Run engine again against real ref — verify determinism ────────
|
||||
print(f"\n{'='*60}")
|
||||
print("Run 2/2 — verifying determinism ...")
|
||||
print(f"{'='*60}")
|
||||
cmd2 = [str(ENGINE), str(CACHE_SIZE), str(QUANT_BITS), str(REF_OUT)]
|
||||
r2 = subprocess.run(cmd2, env=env, capture_output=True, text=True, cwd=str(HERE))
|
||||
print(r2.stdout)
|
||||
if r2.returncode != 0:
|
||||
print("STDERR:", r2.stderr, file=sys.stderr)
|
||||
sys.exit(r2.returncode)
|
||||
|
||||
if "Matching tokens: 12/12" in r2.stdout or f"Matching tokens: {MAX_NEW}/{MAX_NEW}" in r2.stdout:
|
||||
print("✓ Engine is DETERMINISTIC — same output on both runs!")
|
||||
else:
|
||||
m2 = re.search(r"Matching tokens: (\d+)/(\d+)", r2.stdout)
|
||||
if m2:
|
||||
print(f"⚠ Partial match: {m2.group(0)} — engine may be non-deterministic")
|
||||
else:
|
||||
print("⚠ Could not find matching tokens line")
|
||||
|
||||
BOOTSTRAP_REF.unlink(missing_ok=True)
|
||||
@@ -0,0 +1,5 @@
|
||||
"""colibrì — tiny engine, immense model."""
|
||||
|
||||
from colibri._version import __version__
|
||||
|
||||
__all__ = ["__version__"]
|
||||
@@ -0,0 +1,23 @@
|
||||
"""Version accessor for the pip package.
|
||||
|
||||
The single source of truth is c/version.py (#394): coli --version and the
|
||||
GitHub Release workflow read it, so the pip metadata must read the SAME file
|
||||
instead of carrying a second literal that would drift on the first bump.
|
||||
|
||||
From a checkout (the supported install: `pip install -e .`) the file is read
|
||||
directly. From an installed wheel c/ is not on disk, so fall back to the
|
||||
package metadata that setuptools baked at build time from that same file.
|
||||
"""
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
_ns = {}
|
||||
exec((Path(__file__).resolve().parent.parent / "c" / "version.py").read_text(), _ns)
|
||||
__version__ = _ns["__version__"]
|
||||
except OSError:
|
||||
from importlib.metadata import PackageNotFoundError, version
|
||||
|
||||
try:
|
||||
__version__ = version("colibri-engine")
|
||||
except PackageNotFoundError:
|
||||
__version__ = "0.0.0+unknown"
|
||||
@@ -0,0 +1,30 @@
|
||||
"""Entry point for `coli` when installed via pip.
|
||||
|
||||
Delegates to the original c/coli script which handles all subcommands.
|
||||
This wrapper exists so `pip install colibri-engine` creates a `coli` console
|
||||
script that works without the user having to add c/ to PATH manually.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import runpy
|
||||
|
||||
|
||||
def main():
|
||||
here = os.path.dirname(os.path.abspath(__file__))
|
||||
engine_dir = os.path.join(os.path.dirname(here), "c")
|
||||
coli_script = os.path.join(engine_dir, "coli")
|
||||
|
||||
if not os.path.exists(coli_script):
|
||||
sys.exit(
|
||||
"colibri engine directory not found.\n"
|
||||
"Install from source: git clone + pip install -e ."
|
||||
)
|
||||
|
||||
sys.path.insert(0, engine_dir)
|
||||
sys.argv[0] = coli_script
|
||||
runpy.run_path(coli_script, run_name="__main__")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,164 @@
|
||||
# Quick Start — from zero to a running model
|
||||
|
||||
A step-by-step guide for first-time users on **Linux**, **Windows**, and **macOS**.
|
||||
No prior experience with C, CUDA, or model conversion is assumed. If you get
|
||||
stuck, `./coli doctor` (below) tells you exactly what's missing.
|
||||
|
||||
> **What you're setting up:** colibrì runs a very large Mixture-of-Experts model
|
||||
> (e.g. GLM-5.2, 744B parameters) on a normal machine by streaming the model's
|
||||
> experts from disk instead of needing them all in RAM. The engine is a single
|
||||
> C program; Python is only used once, to prepare the model files.
|
||||
|
||||
---
|
||||
|
||||
## 0. What you need first (prerequisites)
|
||||
|
||||
| | Minimum | Recommended |
|
||||
|---|---|---|
|
||||
| **RAM** | ~16 GB | 24 GB+ |
|
||||
| **Free disk** | ~380 GB for the int4 model | a fast NVMe SSD (streaming speed = your token speed) |
|
||||
| **OS** | Linux, Windows 10/11, or macOS | any |
|
||||
| **Tools** | a C compiler + `make` + `git` + `python3` | — |
|
||||
|
||||
You do **not** need a GPU. A GPU only helps if you have one; the engine runs
|
||||
CPU-only by default.
|
||||
|
||||
---
|
||||
|
||||
## 1. Install the build tools
|
||||
|
||||
### Linux (Ubuntu / Debian)
|
||||
|
||||
```bash
|
||||
sudo apt update
|
||||
sudo apt install -y build-essential git python3
|
||||
```
|
||||
|
||||
`build-essential` gives you `gcc`, `make`, and OpenMP (libgomp) — everything the
|
||||
engine needs.
|
||||
|
||||
### Windows
|
||||
|
||||
You have two options.
|
||||
|
||||
**Option A — download a prebuilt binary (no compiler needed).**
|
||||
Grab the latest `colibri-<version>-windows-x86_64.zip` from the
|
||||
[Releases page](https://github.com/JustVugg/colibri/releases), unzip it, and
|
||||
skip to [step 3](#3-get-the-model). Python 3 (from
|
||||
[python.org](https://www.python.org/downloads/)) is still needed if you want to
|
||||
convert a model yourself.
|
||||
|
||||
**Option B — build from source with MSYS2.**
|
||||
Install [MSYS2](https://www.msys2.org/), open the **UCRT64** shell, and run:
|
||||
|
||||
```bash
|
||||
pacman -S --needed mingw-w64-ucrt-x86_64-gcc make git python
|
||||
```
|
||||
|
||||
### macOS
|
||||
|
||||
```bash
|
||||
xcode-select --install # C compiler (clang)
|
||||
brew install libomp git python # OpenMP for multithreading
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 2. Get the code and build the engine
|
||||
|
||||
```bash
|
||||
git clone https://github.com/JustVugg/colibri.git
|
||||
cd colibri/c
|
||||
./setup.sh
|
||||
```
|
||||
|
||||
`setup.sh` checks your compiler and OpenMP, builds the engine, and runs a tiny
|
||||
self-test. When it prints:
|
||||
|
||||
```
|
||||
engine self-test: 32/32 (expected 32/32)
|
||||
```
|
||||
|
||||
the engine is working correctly. (On Windows Option A you already have the
|
||||
binary — you can skip this step.)
|
||||
|
||||
---
|
||||
|
||||
## 3. Get the model
|
||||
|
||||
You have two paths.
|
||||
|
||||
### Easiest — download a ready-made int4 container
|
||||
|
||||
A pre-converted **GLM-5.2 int4** model is on Hugging Face. **Use the version
|
||||
with the int8 MTP heads** (the plain int4 heads disable speculative decoding —
|
||||
see [#8](https://github.com/JustVugg/colibri/issues/8)):
|
||||
|
||||
**https://huggingface.co/mateogrgic/GLM-5.2-colibri-int4-with-int8-mtp**
|
||||
|
||||
Download it into a folder on a fast disk, e.g. `/nvme/glm52_i4` (Linux/macOS) or
|
||||
`D:\glm52_i4` (Windows). It is about **372 GB**, so make sure you have the space.
|
||||
|
||||
### Or convert it yourself from the FP8 source
|
||||
|
||||
One resumable command downloads and converts the model shard by shard, so it
|
||||
never needs the full ~756 GB on disk at once:
|
||||
|
||||
```bash
|
||||
./coli convert --model /nvme/glm52_i4
|
||||
```
|
||||
|
||||
This step uses Python and runs only once. Safe to interrupt and re-run — it
|
||||
resumes where it left off.
|
||||
|
||||
---
|
||||
|
||||
## 4. Run it
|
||||
|
||||
Point `COLI_MODEL` at the folder from step 3 and start chatting:
|
||||
|
||||
```bash
|
||||
# Linux / macOS
|
||||
COLI_MODEL=/nvme/glm52_i4 ./coli chat
|
||||
|
||||
# Windows (UCRT64 shell)
|
||||
COLI_MODEL=/d/glm52_i4 ./coli chat
|
||||
```
|
||||
|
||||
Useful first commands:
|
||||
|
||||
```bash
|
||||
COLI_MODEL=/nvme/glm52_i4 ./coli doctor # read-only check: is everything ready?
|
||||
COLI_MODEL=/nvme/glm52_i4 ./coli plan # shows where the model will live (RAM/disk/GPU)
|
||||
COLI_MODEL=/nvme/glm52_i4 ./coli chat --topp 0.85 # faster: reads less from disk, same quality
|
||||
```
|
||||
|
||||
> **Tip:** `--topp 0.85` is worth adding on a disk-bound machine — it reads
|
||||
> fewer expert bytes per token with no quality loss, which directly means more
|
||||
> tokens per second.
|
||||
|
||||
---
|
||||
|
||||
## 5. What to expect
|
||||
|
||||
- **First launch loads the resident weights** (~10 GB) — this takes a moment.
|
||||
- **Speed depends on your disk.** The experts stream from storage, so a fast
|
||||
NVMe SSD is the single biggest factor in tokens/second. On a slow or shared
|
||||
disk, generation can be well under 1 token/second — that's expected, and it's
|
||||
the honest cost of running a 744B model on a small machine.
|
||||
- **It's still the full model.** Placement only changes speed, never the model's
|
||||
answers or precision.
|
||||
|
||||
If something doesn't work, run `./coli doctor` — it reports exactly what's
|
||||
missing (compiler, model files, permissions) and how to fix it.
|
||||
|
||||
---
|
||||
|
||||
## Where to go next
|
||||
|
||||
| Topic | Doc |
|
||||
|---|---|
|
||||
| Windows native build (and CUDA DLL) | [docs/windows.md](windows.md) |
|
||||
| Tuning: cache, prefetch, speculation | [docs/tuning.md](tuning.md) |
|
||||
| OpenAI-compatible API + web dashboard | [docs/api.md](api.md) |
|
||||
| Every environment variable | [docs/ENVIRONMENT.md](ENVIRONMENT.md) |
|
||||
@@ -0,0 +1,53 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=68.0"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "colibri-engine"
|
||||
dynamic = ["version"]
|
||||
description = "Tiny engine, immense model — run GLM-5.2 (744B MoE) locally"
|
||||
readme = "README.md"
|
||||
license = "Apache-2.0"
|
||||
requires-python = ">=3.10"
|
||||
authors = [
|
||||
{name = "JustVugg"},
|
||||
]
|
||||
classifiers = [
|
||||
"Development Status :: 4 - Beta",
|
||||
"Environment :: Console",
|
||||
"Intended Audience :: Science/Research",
|
||||
"Operating System :: POSIX :: Linux",
|
||||
"Operating System :: MacOS",
|
||||
"Operating System :: Microsoft :: Windows",
|
||||
"Programming Language :: Python :: 3",
|
||||
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
convert = [
|
||||
"numpy",
|
||||
"huggingface_hub",
|
||||
]
|
||||
oracle = [
|
||||
"torch>=2.0",
|
||||
"transformers>=4.40",
|
||||
"safetensors",
|
||||
]
|
||||
bench = [
|
||||
"tokenizers",
|
||||
"datasets",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
coli = "colibri.cli:main"
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://github.com/JustVugg/colibri"
|
||||
Issues = "https://github.com/JustVugg/colibri/issues"
|
||||
|
||||
[tool.setuptools.dynamic]
|
||||
version = {attr = "colibri._version.__version__"}
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["."]
|
||||
include = ["colibri*"]
|
||||
Reference in New Issue
Block a user