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 <x@Mac.fritz.box>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Louie
2026-07-14 23:59:49 +10:00
committed by GitHub
parent 32629eae7f
commit ec0aadf91a
2 changed files with 62 additions and 5 deletions
+56 -3
View File
@@ -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 <arm_neon.h>
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);
}
+6 -2
View File
@@ -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: