From c4624276f37b1d0ac38b464f3d995d1627cf1d22 Mon Sep 17 00:00:00 2001 From: Egon Ruiter Date: Wed, 15 Jul 2026 18:04:25 +0200 Subject: [PATCH] feat(io): implement Consolidated Expert I/O to reduce expert disk reads by 3x and accelerate prefetching --- .gitignore | 4 + c/olmoe.c | 64 ++++++-------- c/tools/convert_olmoe_merged.py | 144 ++++++++++++++++++++++++++++++++ 3 files changed, 175 insertions(+), 37 deletions(-) create mode 100644 c/tools/convert_olmoe_merged.py diff --git a/.gitignore b/.gitignore index dba6a3f..c06902f 100644 --- a/.gitignore +++ b/.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/ diff --git a/c/olmoe.c b/c/olmoe.c index b0f9e8c..199ccaf 100644 --- a/c/olmoe.c +++ b/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; diff --git a/c/tools/convert_olmoe_merged.py b/c/tools/convert_olmoe_merged.py new file mode 100644 index 0000000..b39b123 --- /dev/null +++ b/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()