feat(io): implement Consolidated Expert I/O to reduce expert disk reads by 3x and accelerate prefetching

.gitignore

@@ -67,3 +67,7 @@ c/bench/
 c/tests/test_decode_batch
 c/tests/test_i4_acc512
 c/tests/test_idot
+olmoe_merged/
+olmoe_i4/
+c/olmoe_merged/
+c/olmoe_i4/

c/olmoe.c

@@ -271,20 +271,29 @@ 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].
- * 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);
+
+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);
+    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_raw(&m->S, qsnm, s->gs, 1);
 }
 
 /* ---------- cache expert: ritorna i pesi quantizzati (q+scale) da cache o disco ---------- */
@@ -298,13 +307,10 @@ static void expert_get(Model *m, int layer, int eid, Slot **out) {
     }
     m->miss++; lc->layer_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);
-        s->pinned = 0;
+        slot_ensure_allocated(m, s);
     } else {
         /* IMPROVEMENT 2: LRU eviction — never evict pinned experts */
         int lru = -1;
@@ -320,12 +326,7 @@ static void expert_get(Model *m, int layer, int eid, Slot **out) {
     s->used = ++m->clock;
     pthread_mutex_unlock(&g_pilot_mx);
 
-    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);
+    load_expert_merged(m, layer, eid, s);
 
     pthread_mutex_lock(&g_pilot_mx);
     s->eid = eid;
@@ -551,7 +552,6 @@ static float *step(Model *m, const int *ids, int S, int pos_base) {
 static void pilot_realload(Model *m, int layer, int eid) {
     LCache *lc = &m->cache[layer];
     Cfg *c = &m->c;
-    int64_t ng = (int64_t)c->inter * c->hidden, nd = (int64_t)c->hidden * c->inter;
 
     pthread_mutex_lock(&g_pilot_mx);
     for (int i = 0; i < lc->n; i++) {
@@ -560,9 +560,7 @@ static void pilot_realload(Model *m, int layer, int eid) {
     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);
-        s->pinned = 0;
+        slot_ensure_allocated(m, s);
     } else {
         /* IMPROVEMENT 2: never evict pinned experts */
         int lru = -1;
@@ -576,15 +574,7 @@ static void pilot_realload(Model *m, int layer, int eid) {
     s->eid = -1; s->used = ++m->clock;
     pthread_mutex_unlock(&g_pilot_mx);
 
-    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);
+    load_expert_merged(m, layer, eid, s);
 
     pthread_mutex_lock(&g_pilot_mx);
     s->eid = eid;

c/tools/convert_olmoe_merged.py

@@ -0,0 +1,144 @@
+#!/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
+except ImportError as exc:
+    sys.exit(f"Missing dependencies: {exc}. Install: pip install torch safetensors")
+
+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()