From ec0aadf91afaddc628d57e0266ab96218975f65d Mon Sep 17 00:00:00 2001 From: Louie Date: Tue, 14 Jul 2026 23:59:49 +1000 Subject: [PATCH] olmoe: int8 container support (fixes SIGBUS), NEON int8 dot kernels (2.7x on ARM), macOS RSS fix, converter fallback (#187) - load_expert_w: detect I8/U8 tensors and read them raw with their .qs scales. Before this, pointing SNAP at a convert_olmoe.py container crashed with SIGBUS (st_read_f32 walks I8 data as 2-byte elements). Container misses now cost half the I/O and skip quantize_rows entirely: 2.08 -> 3.69 tok/s on the miss-heavy 16-slot ref.json run (M5). Raw bf16 checkpoints keep working. - matmul_q: Q8_0-style integer path on ARM (per-16 activation blocks, sdot on dotprod CPUs, vmull fallback). Same IDOT family as glm.c, same semantics: IDOT=0 stays byte-exact vs the oracle (12/12); default integer path can flip an argmax tie (documented in glm.c README). End-to-end on M5: 4.5 -> 12 tok/s warm-cache decode. - rss_gb: ru_maxrss is bytes on macOS, KB on Linux. RSS lines were reading '2029 GB' on Macs. - convert_olmoe.py: snapshot_download(local_files_only=True) raises LocalEntryNotFoundError when the repo is not cached; the download fallback was unreachable. Measured on M5 MacBook (10 cores, 24 GB, macOS 26.5), OLMoE-1B-7B-0125: ref.json greedy 12/12 with IDOT=0 on both container and raw paths. Co-authored-by: x Co-authored-by: Claude Fable 5 --- c/olmoe.c | 59 ++++++++++++++++++++++++++++++++++++++-- c/tools/convert_olmoe.py | 8 ++++-- 2 files changed, 62 insertions(+), 5 deletions(-) diff --git a/c/olmoe.c b/c/olmoe.c index 882e195..1c5ffaa 100644 --- a/c/olmoe.c +++ b/c/olmoe.c @@ -51,7 +51,11 @@ typedef struct { /* ---------- utility ---------- */ static double now_s(void) { struct timespec t; clock_gettime(CLOCK_MONOTONIC, &t); return t.tv_sec + t.tv_nsec*1e-9; } -static double rss_gb(void) { struct rusage r; getrusage(RUSAGE_SELF, &r); return r.ru_maxrss / (1024.0*1024.0); } +#if defined(__APPLE__) +static double rss_gb(void) { struct rusage r; getrusage(RUSAGE_SELF, &r); return r.ru_maxrss / (1024.0*1024.0*1024.0); } /* macOS: byte */ +#else +static double rss_gb(void) { struct rusage r; getrusage(RUSAGE_SELF, &r); return r.ru_maxrss / (1024.0*1024.0); } /* Linux: KB */ +#endif static float *falloc(int64_t n) { float *p = malloc(n*sizeof(float)); if(!p){fprintf(stderr,"OOM %ld\n",(long)n);exit(1);} return p; } /* y[S,O] = x[S,I] @ W^T, W e' [O,I] row-major */ @@ -69,8 +73,47 @@ static void matmul(float *y, const float *x, const float *W, int S, int I, int O } /* y[1,O] = x[1,I] @ W^T con W quantizzato: q[O,I] int8 + scala per riga. - * W[o,i] ~= q[o,i]*scale[o] -> y[o] = scale[o] * sum_i x[i]*q[o,i]. */ + * W[o,i] ~= q[o,i]*scale[o] -> y[o] = scale[o] * sum_i x[i]*q[o,i]. + * Su ARM: attivazione quantizzata Q8_0 (scala per blocco di 16) + dot int8 + * NEON (sdot dove c'e' dotprod) — stessa famiglia IDOT di glm.c, IDOT=0 per + * la via scalare byte-esatta. Misurato 2.7x end-to-end su M5. */ +#if defined(__ARM_NEON) +#include +static inline int32_t dot_i8_16(const int8_t *a, const int8_t *b) { + int32x4_t acc = vdupq_n_s32(0); + int8x16_t va = vld1q_s8(a), vb = vld1q_s8(b); +#if defined(__ARM_FEATURE_DOTPROD) + acc = vdotq_s32(acc, va, vb); +#else + acc = vpadalq_s16(acc, vmull_s8(vget_low_s8(va), vget_low_s8(vb))); + acc = vpadalq_s16(acc, vmull_s8(vget_high_s8(va), vget_high_s8(vb))); +#endif + return vaddvq_s32(acc); +} +#endif static void matmul_q(float *y, const float *x, const int8_t *q, const float *scale, int I, int O) { +#if defined(__ARM_NEON) + static int idot = -1; + if (idot < 0) { const char *e = getenv("IDOT"); idot = !(e && *e == '0'); } + if (idot && I % 16 == 0 && I <= 4096) { + int nb = I / 16; int8_t xi[4096]; float xs[256]; + for (int b = 0; b < nb; b++) { + const float *xb = x + b*16; + float am = 0.f; for (int i = 0; i < 16; i++) { float a = fabsf(xb[i]); if (a > am) am = a; } + float s = am/127.f; if (s < 1e-12f) s = 1e-12f; + xs[b] = s; float inv = 1.f/s; + for (int i = 0; i < 16; i++) xi[b*16+i] = (int8_t)lrintf(xb[i]*inv); + } + #pragma omp parallel for schedule(static) + for (int o = 0; o < O; o++) { + const int8_t *w = q + (int64_t)o * I; + float acc = 0.f; + for (int b = 0; b < nb; b++) acc += xs[b]*(float)dot_i8_16(xi+b*16, w+b*16); + y[o] = acc * scale[o]; + } + return; + } +#endif #pragma omp parallel for schedule(static) for (int o = 0; o < O; o++) { const int8_t *w = q + (int64_t)o * I; @@ -171,8 +214,18 @@ static void model_init(Model *m, const char *snap, int cap, int bits) { m->dense_load_s = now_s() - t0; } -/* legge un weight dal disco (streaming) e lo quantizza in q[O,I]+scale[O] */ +/* 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; + } st_read_f32(&m->S, name, tmp, 1); /* pread + fadvise DONTNEED */ quantize_rows(tmp, q, scale, O, I, m->quant_bits); } diff --git a/c/tools/convert_olmoe.py b/c/tools/convert_olmoe.py index 60d8744..dd45806 100644 --- a/c/tools/convert_olmoe.py +++ b/c/tools/convert_olmoe.py @@ -55,9 +55,13 @@ def main(): if args.repo: from huggingface_hub import snapshot_download + from huggingface_hub.errors import LocalEntryNotFoundError print(f"Downloading {args.repo}...") - src_dir = snapshot_download(args.repo, local_files_only=True, max_workers=4) - if not any(Path(src_dir).glob("*.safetensors")): + 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: