Merge pull request #362 from EgonRuiter/prefetcher-v3
feat(prefetch): async expert prefetcher v3.2 for the OLMoE testbed Testbed-only scope: olmoe.c + tools/oracle files; glm.c untouched (the production engine already ships the equivalent techniques: coalesced slab preads, PILOT lookahead, persistent PIN). Trivial .gitignore conflict resolved keeping both sides. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
JustVugg
@@ -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)
|
||||
Reference in New Issue
Block a user