diff --git a/Cargo.lock b/Cargo.lock index 0366a81..474205f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1067,18 +1067,18 @@ dependencies = [ [[package]] name = "clap" -version = "4.6.4" +version = "4.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d91e0c145792ef73a6ad36d27c75ac09f1832222a3c209689d90f534685ee5b7" +checksum = "301b56658598e48f3648647ac6fc887be7e7108eddfa4e9b63fcf3ec58c0cadf" dependencies = [ "clap_builder", ] [[package]] name = "clap_builder" -version = "4.6.2" +version = "4.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f09628afdcc538b57f3c6341e9c8e9970f18e4a481690a64974d7023bd33548b" +checksum = "94a65403d1a1bd28f7dc68eb8506e8874808ee5eecb59298de588e2e1407a078" dependencies = [ "anstyle", "clap_lex", @@ -1534,9 +1534,9 @@ dependencies = [ [[package]] name = "cubecl-hip-sys" -version = "7.2.5321100" +version = "7.14.6085000" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "58887c6d217859c2261a868a3a6b66aa8f44ea2f4bfa51d0dbe4dcb0d1f5768d" +checksum = "760c605ca1b54d12ba9209d2a47992082ff66346fbd96d830af22ac43adfeebc" dependencies = [ "libc", "regex", @@ -1967,9 +1967,9 @@ dependencies = [ [[package]] name = "data-encoding" -version = "2.11.0" +version = "2.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" +checksum = "4583a4551df46e2792f82ceeac45e850d2e2d5debba0b91f102385cda5b11f06" [[package]] name = "deflate64" @@ -2925,9 +2925,9 @@ checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" [[package]] name = "hybrid-array" -version = "0.4.13" +version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "818356c5132c1fede50f837ca96afbe78ff42413047f4abb886217845e1b6c8c" +checksum = "707114b52a152fa7bdb290cd7cd5912d9467273b6d74e21b8d81aca1f8533f6b" dependencies = [ "typenum", ] @@ -3201,9 +3201,9 @@ dependencies = [ [[package]] name = "ipnet" -version = "2.12.0" +version = "2.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" +checksum = "6a756c3fac73139e83f14c2d742155dd2b78d3ee56597b419a0579b7bdd6dd78" [[package]] name = "itertools" @@ -3378,9 +3378,9 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "libredox" -version = "0.1.18" +version = "0.1.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c943259e342f1e06ff2da7a83eabdfe7f92ce10262688dbf1895ff0b3e6e4652" +checksum = "2026a5056764a10b2bf5d56488cba40da507f5493a6a429340e2004d9ed085fa" dependencies = [ "libc", ] @@ -5351,9 +5351,9 @@ dependencies = [ [[package]] name = "time" -version = "0.3.54" +version = "0.3.55" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3e1d5e639ff6bab73cb6885cc7e7b1de96c3f32c68ec55f3952614bec1092244" +checksum = "cdb87b95ec50ddfa440816d227a17b2ccbdda963a316a727fda0fc4334f7d134" dependencies = [ "deranged", "js-sys", diff --git a/README.md b/README.md index d56b039..4801b9d 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys ## What it is - **One binary, every device** — compile both `Wgpu` (Vulkan/DX12/Metal, no CUDA toolchain) and `burn-flex` (CPU); a runtime probe enumerates **all** adapters + the CPU and places the model across them — a single GPU, **several GPUs together**, or GPU + CPU hybrid. No feature-split builds, no per-vendor path. -- **Models from scratch, generic over `B: Backend`** — a growing zoo (Qwen2/2.5, Qwen3 dense with per-head q/k norm + decoupled head_dim, LFM2/2.5 hybrid conv+attention, all-MiniLM embedder) built on shared blocks (RmsNorm · GQA · RoPE · SwiGLU · tied lm-head · depthwise causal conv), with a clean trait to add more. +- **Models from scratch, generic over `B: Backend`** — a growing zoo (Qwen2/2.5, Qwen3 dense with per-head q/k norm + decoupled head_dim, LFM2/2.5 hybrid conv+attention, **OLMoE sparse mixture-of-experts**, all-MiniLM embedder) built on shared blocks (RmsNorm · GQA · RoPE · SwiGLU · **top-k-routed expert bank** · tied lm-head · depthwise causal conv), with a clean trait to add more. - **Trustworthy reimplementations** — every port must pass a **parity gate**: single-forward top-k logits *and* a short greedy sequence match a reference (Candle, or a local Ollama of the same model) exactly. - **Fast** — per-layer KV cache (+ conv-state cache for hybrids), on-GPU argmax (sync only the winning index), sampling, **token streaming**, cooperative cancellation; kernel `fusion` + `autotune`; an **f16** path (f32 attention-score island for numeric safety) that halves VRAM at full speed. - **A full model-import suite** — pull a model from HuggingFace (by repo id) or from disk and load it: **safetensors**, **PyTorch** state dicts, and **GGUF** (llama.cpp, dequantized) weights; `config.json`-driven hyperparameters; tokenizer + chat-template import (HF `tokenizers` / SentencePiece / BPE); per-architecture weight-name remapping with a **checked load** (fail loudly on a key mismatch, never silently zero-init); resumable, shard-aware downloads into a per-user cache; and a declarative **model registry** so adding a model is a manifest entry, not new code. @@ -56,12 +56,21 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys runtime liveness a checked load can't see — NaN/Inf logits, a vocab-width mismatch, or a degenerate/dead forward. `CausalLm::sanity_check` is the post-`install` gate an app calls to catch a silently-broken import before trusting the model. -- **Four architectures ported and running on real weights** — Qwen2/2.5, Qwen3 dense, the LFM2/2.5 - hybrid, and the all-MiniLM sentence embedder; Qwen2.5-1.5B, Qwen3-0.6B, and LFM2.5-1.2B/230M load and - greedy-decode correctly on the reference GPU (wgpu/Vulkan), and Qwen2.5-0.5B / LFM2.5-230M do the same - on the CPU backend. Qwen3 reuses the shared blocks whole (its per-head q/k RMSNorm, absent qkv bias, - and decoupled `head_dim` were all already supported), loads from safetensors **and** a single Q4_K_M - GGUF, and handles Qwen3's `` reasoning mode. +- **Five architectures ported and running on real weights** — Qwen2/2.5, Qwen3 dense, the LFM2/2.5 + hybrid, **OLMoE** (sparse MoE), and the all-MiniLM sentence embedder; Qwen2.5-1.5B, Qwen3-0.6B, and + LFM2.5-1.2B/230M load and greedy-decode correctly on the reference GPU (wgpu/Vulkan), and + Qwen2.5-0.5B / LFM2.5-230M / OLMoE-1B-7B do the same on the CPU backend. Qwen3 reuses the shared + blocks whole (its per-head q/k RMSNorm, absent qkv bias, and decoupled `head_dim` were all already + supported), loads from safetensors **and** a single Q4_K_M GGUF, and handles Qwen3's `` + reasoning mode. +- **Mixture-of-experts** — `nn::SparseMoe` is a softmax top-k router over a **fused expert bank** + (`[experts, out, in]` — the row-major twin of GGUF's `ffn_*_exps`, so 64 experts load as three + tensors per layer, not 192); `models::olmoe` drives it config-first off `olmoe.*` GGUF metadata. + **OLMoE-1B-7B** (64 experts, top-8 routing, 1B active / 7B total) runs from the one official + 4.21 GB Q4_K_M file — 16 layers × 64 experts loaded and "2 + 2 equals 4." greedy-decoded at + 1.15 s/token on the CPU backend (~28 GB f32 resident; a 16 GB card waits on keep-quantized VRAM, + tracked in P9). Attention learned OLMoE's whole-projection q/k RMSNorm placement, inferred from the + loaded norm's own width, so every existing checkpoint loads byte-unchanged. - **All three models are parity-verified** — the two-leg P7 gate passes for Qwen2.5-1.5B on the reference GPU: single-forward top-5 logits match a Candle f32 reference (max |Δlogit| 2.7e-5, `tests/parity_qwen2.rs` + the committed `tools/candle-probe` fixture) and a 24-token greedy sequence @@ -74,8 +83,10 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys both tiers. **Qwen3** is parity-verified through the same llama.cpp harness (`tests/parity_gguf.rs`, `qwen3` leg): on Qwen3-0.6B Q4_K_M, top-5 first-forward ids match exactly in order and a 24-token greedy sequence — `` reasoning tokens included — is byte-identical to `llama-server` on the - same file. The MiniLM embedder matches its Candle reference at cosine 0.99999994 - (max |Δcomponent| 1.2e-7, `tests/real_minilm.rs`). + same file. **OLMoE-1B-7B** passes the same llama.cpp gate on its own Q4_K_M (`olmoe` leg): top-5 ids + exact in order, 24-token greedy byte-identical, max |Δlogprob| 3.7e-1 — so the MoE router and expert + bank are verified against a reference, not just plausible. The MiniLM embedder matches its Candle + reference at cosine 0.99999994 (max |Δcomponent| 1.2e-7, `tests/real_minilm.rs`). - **Sampling, streaming, cancellation** — temperature / top-k / top-p sampling (deterministic per seed), per-token streaming through a `ControlFlow` callback, and cooperative between-token cancellation; greedy decoding keeps the argmax on-device. diff --git a/ROADMAP.md b/ROADMAP.md index f1009d9..afff3e9 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -37,7 +37,11 @@ device defaults — the exact `B::FloatElem` seam Mummu's dtype pinning and `Gpu upstream: a `FloatCastAdapter` in burn-store (our `CastFloatAdapter`'s role), burnpack split into `burn-pack`, BitNet `Calibration::AbsMean` ternary quant, quant fallbacks for slice/gather/select/expand, and a remote multi-device backend (iroh). tokenizers 0.23.1 remains current — -https://github.com/Tracel-AI/burn/releases* +https://github.com/Tracel-AI/burn/releases* *(2026-08-03) Pin watch: burn 0.22 is **still pre-release** +(0.22.0-pre.1 remains the newest tag; 0.21.0 the latest stable) — the P0 migration item stays gated. +CubeCL tagged **0.11.0-pre.1** the same day (2026-07-29): a frontend mega-refactor (references), a +**Metal backend**, a new CPU runtime, tiled layouts, and CUDA stream priority hints — all of which +arrive with the burn bump, not before — https://github.com/tracel-ai/cubecl/releases* ## North Star @@ -347,7 +351,7 @@ a benchmark holds/improves its budget; README perf claims link an artifact. less of the reference's own per-dot bf16 activation rounding; it sits inside the existing 5e-2 tolerance with ~1.5x headroom (tolerance unchanged, and the strict-order id match remains the real assert). Run legs with `MUMMU_LFM2_230M_DIR` + `MUMMU_LFM2_230M_BF16_GGUF`.* -- [ ] **First MoE architecture** *(colibri parity)* — the zoo is dense-only and has so far deferred MoE +- [x] **First MoE architecture** *(colibri parity)* — the zoo is dense-only and has so far deferred MoE ("Qwen3.6 35B-A3B … well past the single-card tier"), but colibri demonstrates the counter-thesis on our exact hardware class: frontier MoE (GLM-5.2 744B-A40B, Kimi K3 2.8T-A104B) on consumer boxes, because only the ~40B *active* params compute per token — dense parts (~10 GB at int4) stay @@ -369,6 +373,53 @@ a benchmark holds/improves its budget; README perf claims link an artifact. the textbook offload candidates since only 8/64 fire per token) — https://huggingface.co/allenai/OLMoE-1B-7B-0125-Instruct-GGUF · https://huggingface.co/blog/Doctor-Shotgun/llamacpp-moe-offload-guide + *(2026-08-03) **SHIPPED and parity-verified on the first run — the zoo is no longer dense-only.** + New shared block `nn::SparseMoe`: a softmax top-k router (f32 island, HF `OlmoeSparseMoeBlock` + math) over a **fused 3-D expert bank** `[experts, out, in]` — the exact row-major twin of ggml's + `ffn_{gate,up,down}_exps`, so the 64 experts load as three tensors per layer rather than 192. + Compute is **dense-mask**: every expert processes every token and the router's weight row (exactly + zero off the top-k) scales the rest away — numerically identical to the sparse formulation, and it + keeps the whole forward on-device with no data-dependent gather (see the measured follow-up below). + `GqaAttention` gained OLMoE's q/k-norm placement — RMSNorm over the **whole projection** before the + head split — inferred from the loaded gamma's width (`head_dim` = per-head, `n·head_dim` = + projection), so the module shape is unchanged and every existing checkpoint loads identically. + `models::olmoe` is config-driven off `olmoe.*` GGUF metadata (expert counts, either + `feed_forward_length` spelling, `expert_weights_norm`), GGUF-only by design (HF ships the experts + **unfused** as `mlp.experts.{i}.*` — a 64-way concat-on-import is the split item below); the + tokenizer registry gained the `olmo` pre (stock GPT-2 regex + NFC). **Parity gate PASSED** + (`parity_gguf.rs`, new `olmoe` leg — the harness is now generic over the backend): llama.cpp on the + SAME allenai Q4_K_M file, top-5 first-forward ids match **exactly in order** + (1992, 17833, 11202, 4943, 1394), the 24-token greedy sequence is **byte-identical**, max + |Δlogprob| **3.687691131310693e-1** (inside the shared 7.5e-1 tolerance, right beside Qwen3's + 4.02e-1). Real-model proof (`real_olmoe.rs`): the registry spec fetched the 4.21 GB GGUF, the ONE + file loaded 16 layers × 64 experts in 92 s and greedy-decoded "2 + 2 equals 4." at **1.15 s/token** + on the CPU backend (~28 GB f32 resident — a 16 GB card is out of reach until the P9 keep-quantized + leg), sanity spread 36.9; its GGUF-built tokenizer is **byte-identical** to the checkpoint's + `tokenizer.json` across an 8-prompt battery. 202 unit tests; every prior gate re-passed unchanged + after the attention refactor (Qwen3 GGUF parity bit-identical at 4.015608155114805e-1, Qwen2 both + legs, template gate 10/10, budgets 104.4 ms / 13.2 tok/s GPU + 15.3 tok/s CPU). Run the legs with + `MUMMU_OLMOE_GGUF_PATH` / `MUMMU_HUB_DEST` / `MUMMU_OLMOE_TOK_JSON`.* +- [ ] **MoE decode: make routed-expert compute actually pay** — the dense-mask forward computes all 64 + experts per token when only 8 are routed, so decode touches ~7B params instead of ~1B (baseline: + **0.76 s/token** warm, `mummu-bench/tests/budget_moe.rs` — the number to beat). The obvious fix was + tried and **measured a regression, so it was reverted, not shipped**: gathering the k routed + expert slices with a device-side `select` off the router's own index tensor (no host sync, exact + same math to summation-order rounding — a unit test confirmed the two paths agree to 1e-6) made real + OLMoE decode **1.58 s/token vs the dense path's 1.15 s** end to end on burn-flex. The gather copies ~200 MB of + expert weights per layer per token, and that copy costs more than the dense matmul it removes — + i.e. `select` materializes where the dense path streams. Routes worth trying next, each gated on + `bench/BASELINE.md` like this one was: (a) a fused dequant/gather matmul kernel that never + materializes the gathered bank (the P9 keep-quantized kernel work is the natural host); (b) measure + on the **GPU** backend, where the copy is far cheaper relative to compute — blocked until a MoE fits + VRAM (P9 keep-quantized, or expert offload); (c) llama.cpp's own answer, `--n-cpu-moe`-style + placement, which sidesteps the gather entirely by moving whole expert banks rather than slicing + them. *(2026-08-03, measured this run.)* +- [ ] **OLMoE from HF safetensors** — the port loads GGUF only because HF stores each expert separately + (`model.layers.N.mlp.experts.{0..63}.{gate,up,down}_proj.weight`) while `MoeExperts` holds one fused + `[experts, out, in]` tensor per projection. Needs a concat-on-import step (64 slices → one tensor, + in expert order) in the safetensors path — mechanical, but it wants its own fixture (the bf16 + checkpoint is ~14 GB) and a byte-equality check against the GGUF-loaded weights. *(2026-08-03, split + from the MoE item.)* - [ ] **Qwen3.5 hybrid (`qwen35`) architecture port** — split from the Qwen3.5-tier item when the 2026-07-30 header probe showed Qwen3.5-4B/9B are a hybrid **linear-attention/SSM + periodic full-attention** arch (`qwen35.ssm.*` metadata: conv_kernel/state_size/group_count/time_step_rank/ @@ -1001,6 +1052,15 @@ that fits the model AND uses every device to the fullest. affects speed, never precision"), and a bench proving the streamed model beats the largest-fitting resident one on task throughput. *(2026-07-30 research)* — https://github.com/JustVugg/colibri + *(2026-08-03 research)* Prior art to mine when this is picked up: llama.cpp's **`--n-cpu-moe N`** + flag (core algorithm in its PR #15077) keeps attention on-GPU and moves the first N layers' + expert FFN weights to CPU RAM — community numbers show 12–24 GB cards running 35B-class MoE at + 50–60 tok/s, which calibrates what expert-CPU-offload alone (no NVMe tier) buys; and an open + llama.cpp feature request (#20757) sketches the **two-tier GPU+RAM expert cache with pluggable + eviction** — the same LRU-residency design colibri proved, upstreamed. Mummu's OLMoE port + (2026-08-03) makes both concrete here: expert banks are single fused 3-D tensors per layer, the + natural offload/eviction unit — https://github.com/ggml-org/llama.cpp/issues/20757 · + https://openclawdc.com/blog/llama-cpp-moe-offload-flags-explained/ - [ ] **Planner introspection (`plan` / `doctor`)** *(colibri parity)* — colibri ships `coli plan` (print the placement decision without running) and `coli doctor` (readiness checks). Mummu's planner should expose the same as *API*, consumer UIs render it: a `Plan` report (per-device @@ -1133,7 +1193,9 @@ Mummu (a 128 GB box with one consumer GPU), opposite bet: Mummu shrinks the mode gap, each gated by the usual parity + budget discipline: - **P2 — first MoE architecture** (OLMoE-7B-A1B / Qwen3-30B-A3B): router + expert blocks + GGUF expert - tensors; the prerequisite for everything expert-shaped. + tensors; the prerequisite for everything expert-shaped. **Closed 2026-08-03** — OLMoE-1B-7B ported and + parity-verified vs llama.cpp; the expert-streaming and routed-compute items below now have their + architecture. - **P5 — speculative decoding** (MTP heads / small-model drafts, byte-identical by construction) and **grammar-constrained decoding** (guaranteed-parseable tool calls — worth more to the apps than tok/s). - **P6 — NVMe as a third placement tier** (expert streaming; the headline feature) and **planner diff --git a/bench/BASELINE.md b/bench/BASELINE.md index b5075aa..3582b1d 100644 --- a/bench/BASELINE.md +++ b/bench/BASELINE.md @@ -34,6 +34,23 @@ weights at only ~114 GB/s vs the card's ~672 GB/s, so kernel/dispatch overhead i | --- | --- | --- | | Decode (8 greedy tokens, warm KV cache; `mummu-bench/tests/budget_cpu.rs`) | **11.7 tok/s** | ≥ 6 tok/s | +## OLMoE-1B-7B-0125-Instruct · CPU (burn-flex) · f32 from Q4_K_M GGUF + +The MoE tier (64 experts, top-8 routing; 1B active / 7B total). Budgeted in **seconds per token**, not +tok/s: the dense-mask expert forward computes every expert for every token, so decode touches all 7B +params rather than the 1B the routing implies — this row is the number the routed-compute work in +ROADMAP P2 has to beat. GPU is out of reach until keep-quantized VRAM (P9): ~28 GB resident in f32. + +| Metric | Recorded (2026-08-03) | Budget | +| --- | --- | --- | +| Load (dequantize ~7B params to f32; `mummu-bench/tests/budget_moe.rs`) | **81.9 s** | ≤ 300 s | +| Decode (4 greedy tokens, warm KV cache) | **0.76 s/token** | ≤ 2.0 s/token | + +End-to-end (`tests/real_olmoe.rs`, prefill included and amortized over a 6-token answer) reads +1.15 s/token — the same path, measured the way a caller experiences it rather than warm-cache steady +state. Use the 0.76 s/token row for regression comparisons; both are recorded so the gap is not +mistaken for drift. + Notes - 2026-07-11: the f32 attention-score island (NaN fix for f16) coincided with an f32 *improvement* (TTFT 100.5 → 88.4 ms, decode 13.3 → 14.1 tok/s) — softmax now always runs in f32 with fusion @@ -42,4 +59,9 @@ Notes - Effective weight-streaming bandwidth at ~71 ms/token over ~6.2 GB of f32 weights is ~88 GB/s vs the card's ~672 GB/s — the decode path is kernel/dispatch-bound, not bandwidth-bound. The CubeCL SPIR-V compiler feature (ROADMAP P6) is the identified lever. +- 2026-08-03: a routed-expert **gather** for the single-token decode step (device-side `select` of the + 8 routed slices — numerically the dense path minus its exactly-zero terms) measured **1.58 s/token vs + the dense path's 1.15** on the end-to-end harness, so it was rejected under this file's rule and + reverted: the ~200 MB per-layer gather copy costs more than the dense matmul it removes. Routes left + open are in the ROADMAP P2 item. - `harness_smoke` (sub-ns) exists only to keep `cargo bench` green without the multi-GB weights. diff --git a/crates/mummu-bench/tests/budget_moe.rs b/crates/mummu-bench/tests/budget_moe.rs new file mode 100644 index 0000000..cf35e48 --- /dev/null +++ b/crates/mummu-bench/tests/budget_moe.rs @@ -0,0 +1,105 @@ +//! The MoE-tier perf-budget gate from `bench/BASELINE.md`: OLMoE-1B-7B +//! greedy decode on the `Cpu` (burn-flex) backend, straight from its Q4_K_M +//! GGUF. This tier is bounded by **latency per token**, not tok/s — the +//! dense-mask expert forward touches all 7B params per token (see the +//! routed-compute item in ROADMAP P2), so a budget in seconds/token is the +//! honest unit and the number to beat when that lands. +//! +//! Ignored by default (the file is ~4.2 GB and the f32 build needs ~28 GB of +//! RAM); run with +//! +//! ```text +//! MUMMU_OLMOE_GGUF_PATH=path/to/OLMoE-1B-7B-0125-Instruct-Q4_K_M.gguf \ +//! cargo test -p mummu-bench --release --test budget_moe -- --ignored --nocapture +//! ``` + +use std::path::PathBuf; +use std::time::Instant; + +use mummu::backend::Cpu; +use mummu::decode::argmax_id; +use mummu::gguf::{GgufFile, GgufValue}; +use mummu::models::CausalLm; +use mummu::models::olmoe; + +/// From bench/BASELINE.md (recorded 0.76 s/token warm on 2026-08-03). The +/// ceiling carries ~2.6x headroom so ordinary host-load noise is not a false +/// alarm — this backend's decode tracks CPU availability closely. +const DECODE_BUDGET_SECS_PER_TOKEN: f64 = 2.0; +const DECODE_STEPS: usize = 4; + +/// Loading dequantizes ~7B params to f32; on the reference machine that is +/// ~82 s. Well clear of any plausible import regression, and it catches a +/// pathological one. +const LOAD_BUDGET_SECS: f64 = 300.0; + +#[test] +#[ignore = "needs the OLMoE Q4_K_M GGUF (MUMMU_OLMOE_GGUF_PATH) and ~40 GB free RAM"] +fn olmoe_moe_cpu_decode_stays_inside_its_budget() { + let Some(path) = std::env::var_os("MUMMU_OLMOE_GGUF_PATH").map(PathBuf::from) else { + panic!("set MUMMU_OLMOE_GGUF_PATH to the OLMoE-1B-7B q4_k_m gguf"); + }; + assert!( + path.is_file(), + "MUMMU_OLMOE_GGUF_PATH is not a file: {path:?}" + ); + + let f = GgufFile::open(&path).expect("gguf header parses"); + let tok = mummu::tokenizer::tokenizer_from_gguf(&f).expect("tokenizer from metadata"); + let bos = f + .get("tokenizer.ggml.bos_token_id") + .and_then(GgufValue::as_u64) + .and_then(|v| { + f.get("tokenizer.ggml.tokens") + .and_then(GgufValue::as_array) + .and_then(|t| t.get(usize::try_from(v).ok()?)) + .and_then(GgufValue::as_str) + .map(String::from) + }) + .unwrap_or_default(); + drop(f); + let raw = + format!("{bos}<|user|>\nWhat is 2 + 2? Answer in one short sentence.\n<|assistant|>\n"); + let ids = tok + .encode(raw.as_str(), false) + .expect("encodes") + .get_ids() + .to_vec(); + + let device = burn::tensor::Device::::default(); + let start = Instant::now(); + let loaded = olmoe::load_from_gguf::(&path, &device).expect("gguf load checked"); + let load_secs = start.elapsed().as_secs_f64(); + + // Prefill (uncounted warm-up), then time the decode steps. + let mut cache = loaded.new_cache(); + let logits = loaded.forward(&ids, 0, &mut cache, &device); + let mut next = argmax_id(logits).expect("argmax"); + + let start = Instant::now(); + let mut out = Vec::with_capacity(DECODE_STEPS); + for past in (ids.len()..).take(DECODE_STEPS) { + let logits = loaded.forward(&[next], past, &mut cache, &device); + next = argmax_id(logits).expect("argmax"); + out.push(next); + } + let secs_per_token = start.elapsed().as_secs_f64() / DECODE_STEPS as f64; + let text = tok.decode(&out, true).expect("decode"); + + eprintln!( + "[budget/moe] OLMoE-1B-7B load {load_secs:.1}s (budget {LOAD_BUDGET_SECS}), \ + decode {secs_per_token:.2} s/token (budget {DECODE_BUDGET_SECS_PER_TOKEN}); text: {text:?}" + ); + assert!( + text.contains('4') || text.to_lowercase().contains("four"), + "MoE decode must stay coherent, got: {text:?}" + ); + assert!( + load_secs <= LOAD_BUDGET_SECS, + "MoE load regression: {load_secs:.1}s > {LOAD_BUDGET_SECS}s budget" + ); + assert!( + secs_per_token <= DECODE_BUDGET_SECS_PER_TOKEN, + "MoE decode regression: {secs_per_token:.2} s/token > {DECODE_BUDGET_SECS_PER_TOKEN} budget" + ); +} diff --git a/crates/mummu/src/gguf.rs b/crates/mummu/src/gguf.rs index 8397a8d..4716ece 100644 --- a/crates/mummu/src/gguf.rs +++ b/crates/mummu/src/gguf.rs @@ -137,6 +137,15 @@ impl GgufValue { } } + /// The value as a bool, if it is one. + #[must_use] + pub fn as_bool(&self) -> Option { + match *self { + Self::Bool(v) => Some(v), + _ => None, + } + } + /// The value as an array slice, if it is one. #[must_use] pub fn as_array(&self) -> Option<&[GgufValue]> { diff --git a/crates/mummu/src/models/lfm2.rs b/crates/mummu/src/models/lfm2.rs index 131f1b3..68020fb 100644 --- a/crates/mummu/src/models/lfm2.rs +++ b/crates/mummu/src/models/lfm2.rs @@ -295,6 +295,7 @@ fn build(cfg: &Lfm2Config, device: &B::Device) -> Lfm2 { head_dim: cfg.head_dim(), bias: false, // LFM2 projections are bias-free qk_norm_eps: Some(cfg.norm_eps), // per-head q/k RMSNorm + qk_norm_projection: false, }; let conv_cfg = ShortConvConfig { hidden_size: cfg.hidden_size, diff --git a/crates/mummu/src/models/mod.rs b/crates/mummu/src/models/mod.rs index 59feb16..1a1dbcd 100644 --- a/crates/mummu/src/models/mod.rs +++ b/crates/mummu/src/models/mod.rs @@ -8,6 +8,7 @@ use crate::decode::{SamplerOptions, generate_loop, top_k_ids}; pub mod lfm2; pub mod minilm; +pub mod olmoe; pub mod qwen2; pub mod qwen3; diff --git a/crates/mummu/src/models/olmoe.rs b/crates/mummu/src/models/olmoe.rs new file mode 100644 index 0000000..19e0f8a --- /dev/null +++ b/crates/mummu/src/models/olmoe.rs @@ -0,0 +1,622 @@ +//! OLMoE sparse mixture-of-experts decoder (allenai OLMoE-1B-7B), from +//! scratch on the shared `nn` blocks — the zoo's first MoE architecture. +//! Structure per layer is pre-norm like Qwen, with two deltas: +//! * the FFN is a [`SparseMoe`] — a softmax top-k router (`k = 8` of 64) +//! over narrow SwiGLU experts, `norm_topk_prob = false`; +//! * q/k RMSNorm applies to the **whole projection** (width +//! `num_heads * head_dim`) before the head split — `GqaAttention`'s +//! projection placement (OLMoE is MHA: 16 query heads, 16 KV heads). +//! +//! This is the **resident-everything** first cut: all 64 experts' weights +//! live in memory and every expert computes every token (the router mask +//! zeroes the unrouted ones) — ~7B params in f32 is ~28 GB, which targets the +//! CPU backend on the reference 128 GB machine. Expert streaming / offload is +//! the P6 placement item; keep-quantized is P9. +//! +//! Import is **GGUF-only** for now: the HF safetensors checkpoint stores each +//! expert as a separate tensor (`mlp.experts.{i}.gate_proj.weight`), which +//! would need a 64-way concat on import; the GGUF ships the experts already +//! fused (`ffn_*_exps`) in exactly the layout [`MoeExperts`] holds. + +use std::path::Path; + +use burn::module::Module; +use burn::nn::{Embedding, EmbeddingConfig, Linear, LinearConfig, RmsNorm, RmsNormConfig}; +use burn::store::{ModuleAdapter, PyTorchToBurnAdapter, SafetensorsStore}; +use burn::tensor::{Int, Tensor, TensorData, backend::Backend}; + +use crate::gguf::{GgufFile, GgufMap, GgufTensorInfo, GgufValue}; +use crate::import::{CastFloatAdapter, ImportError, load_checked}; +use crate::models::CausalLm; +use crate::models::qwen2::{EosIds, gguf_f32, gguf_usize}; +use crate::nn::{ + GqaAttention, GqaAttentionConfig, LayerKv, SparseMoe, SparseMoeConfig, causal_mask, rope_tables, +}; + +/// OLMoE architecture hyperparameters (HF `config.json` field names). +#[derive(Debug, Clone, serde::Deserialize)] +pub struct OlmoeConfig { + pub vocab_size: usize, + pub hidden_size: usize, + /// Per-expert SwiGLU intermediate width (1B-7B: 1024). + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub num_experts: usize, + pub num_experts_per_tok: usize, + /// Renormalize the top-k routing weights to sum 1 (OLMoE ships `false`: + /// the raw softmax probabilities weight the mixture). + #[serde(default)] + pub norm_topk_prob: bool, + pub rms_norm_eps: f64, + pub rope_theta: f32, + #[serde(default)] + pub tie_word_embeddings: bool, + #[serde(default)] + pub eos_token_id: EosIds, +} + +impl OlmoeConfig { + /// Parse `config.json` bytes. + pub fn from_json_bytes(bytes: &[u8]) -> Result { + let cfg: Self = serde_json::from_slice(bytes).map_err(|e| e.to_string())?; + cfg.validate()?; + Ok(cfg) + } + + /// Hyperparameters from a GGUF header's `olmoe.*` metadata. + pub fn from_gguf(f: &GgufFile) -> Result { + let arch = f.architecture().unwrap_or(""); + if arch != "olmoe" { + return Err(format!("GGUF architecture '{arch}' is not olmoe")); + } + let hidden_size = gguf_usize(f, "olmoe.embedding_length")?; + let embd = f + .tensor("token_embd.weight") + .ok_or("GGUF has no token_embd.weight tensor")?; + if embd.dims.len() != 2 || embd.dims[0] != hidden_size as u64 { + return Err(format!( + "token_embd.weight dims {:?} do not match embedding_length {hidden_size}", + embd.dims + )); + } + let vocab_size = usize::try_from(embd.dims[1]).map_err(|_| "vocab too large")?; + if let Some(tokens) = f.get("tokenizer.ggml.tokens").and_then(GgufValue::as_array) + && tokens.len() > vocab_size + { + return Err(format!( + "tokenizer vocab {} exceeds embedding rows {vocab_size}", + tokens.len() + )); + } + let eos_token_id = f + .get("tokenizer.ggml.eos_token_id") + .and_then(GgufValue::as_u64) + .and_then(|v| u32::try_from(v).ok()) + .map_or(EosIds::None, EosIds::One); + // llama.cpp writes the per-expert width as expert_feed_forward_length + // when it differs from feed_forward_length; accept either spelling. + let intermediate_size = gguf_usize(f, "olmoe.expert_feed_forward_length") + .or_else(|_| gguf_usize(f, "olmoe.feed_forward_length"))?; + // OLMoE does not renormalize top-k weights; honor the metadata key + // when a file carries one, default to the architecture's `false`. + let norm_topk_prob = f + .get("olmoe.expert_weights_norm") + .and_then(GgufValue::as_bool) + .unwrap_or(false); + let cfg = Self { + vocab_size, + hidden_size, + intermediate_size, + num_hidden_layers: gguf_usize(f, "olmoe.block_count")?, + num_attention_heads: gguf_usize(f, "olmoe.attention.head_count")?, + num_key_value_heads: gguf_usize(f, "olmoe.attention.head_count_kv")?, + num_experts: gguf_usize(f, "olmoe.expert_count")?, + num_experts_per_tok: gguf_usize(f, "olmoe.expert_used_count")?, + norm_topk_prob, + rms_norm_eps: f64::from(gguf_f32(f, "olmoe.attention.layer_norm_rms_epsilon")?), + rope_theta: gguf_f32(f, "olmoe.rope.freq_base")?, + // No separate output.weight tensor means the lm-head is tied. + tie_word_embeddings: f.tensor("output.weight").is_none(), + eos_token_id, + }; + cfg.validate()?; + Ok(cfg) + } + + /// `hidden_size / num_attention_heads` — OLMoE's head_dim is not decoupled. + #[must_use] + pub fn head_dim(&self) -> usize { + self.hidden_size / self.num_attention_heads.max(1) + } + + fn validate(&self) -> Result<(), String> { + if self.num_key_value_heads == 0 + || !self + .num_attention_heads + .is_multiple_of(self.num_key_value_heads) + { + return Err(format!( + "num_attention_heads ({}) must be a positive multiple of num_key_value_heads ({})", + self.num_attention_heads, self.num_key_value_heads + )); + } + if self.num_hidden_layers == 0 || self.vocab_size == 0 { + return Err("num_hidden_layers and vocab_size must be positive".into()); + } + if self.num_experts < 2 || !(1..=self.num_experts).contains(&self.num_experts_per_tok) { + return Err(format!( + "num_experts ({}) must be >= 2 with num_experts_per_tok ({}) in 1..=num_experts", + self.num_experts, self.num_experts_per_tok + )); + } + let hd = self.head_dim(); + if hd < 2 || !hd.is_multiple_of(2) || hd * self.num_attention_heads != self.hidden_size { + return Err(format!( + "hidden_size ({}) must split evenly into num_attention_heads ({}) even-sized heads", + self.hidden_size, self.num_attention_heads + )); + } + Ok(()) + } +} + +/// One OLMoE decoder layer. Field names mirror the HF checkpoint layout. +#[derive(Module, Debug)] +pub struct DecoderLayer { + pub self_attn: GqaAttention, + pub mlp: SparseMoe, + pub input_layernorm: RmsNorm, + pub post_attention_layernorm: RmsNorm, +} + +/// The OLMoE decoder stack (HF's `model.*` subtree). The 1B-7B ships untied. +#[derive(Module, Debug)] +pub struct Olmoe { + pub embed_tokens: Embedding, + pub layers: Vec>, + pub norm: RmsNorm, + pub lm_head: Option>, +} + +/// A weight-loaded OLMoE plus its config — everything a forward needs. +pub struct LoadedOlmoe { + pub model: Olmoe, + pub config: OlmoeConfig, +} + +fn build(cfg: &OlmoeConfig, device: &B::Device) -> Olmoe { + let attn_cfg = GqaAttentionConfig { + hidden_size: cfg.hidden_size, + num_heads: cfg.num_attention_heads, + num_kv_heads: cfg.num_key_value_heads, + head_dim: cfg.head_dim(), + bias: false, // OLMoE projections are bias-free + qk_norm_eps: Some(cfg.rms_norm_eps), // q/k RMSNorm over the whole + qk_norm_projection: true, // projection, pre head-split + }; + let moe_cfg = SparseMoeConfig { + hidden_size: cfg.hidden_size, + expert_intermediate_size: cfg.intermediate_size, + num_experts: cfg.num_experts, + num_experts_per_tok: cfg.num_experts_per_tok, + }; + let norm = |dev: &B::Device| { + RmsNormConfig::new(cfg.hidden_size) + .with_epsilon(cfg.rms_norm_eps) + .init(dev) + }; + let layers = (0..cfg.num_hidden_layers) + .map(|_| DecoderLayer { + self_attn: attn_cfg.init(device), + mlp: moe_cfg.init(device), + input_layernorm: norm(device), + post_attention_layernorm: norm(device), + }) + .collect(); + let lm_head = (!cfg.tie_word_embeddings).then(|| { + LinearConfig::new(cfg.hidden_size, cfg.vocab_size) + .with_bias(false) + .init(device) + }); + Olmoe { + embed_tokens: EmbeddingConfig::new(cfg.vocab_size, cfg.hidden_size).init(device), + layers, + norm: norm(device), + lm_head, + } +} + +/// The key remap: strip `model.`, rename every RmsNorm `weight` → Burn's +/// `gamma` (incl. the projection-wide `self_attn.{q,k}_norm`). The fused +/// expert params (`mlp.experts.{gate,up,down}`) carry no `.weight` suffix — +/// they are raw `Param` fields, named by the GGUF map directly. +fn install_remaps(store: SafetensorsStore) -> SafetensorsStore { + store + .with_key_remapping(r"^model\.", "") + .with_key_remapping(r"(input_layernorm)\.weight$", "$1.gamma") + .with_key_remapping(r"(post_attention_layernorm)\.weight$", "$1.gamma") + .with_key_remapping(r"(self_attn\.q_norm)\.weight$", "$1.gamma") + .with_key_remapping(r"(self_attn\.k_norm)\.weight$", "$1.gamma") + .with_key_remapping(r"^norm\.weight$", "norm.gamma") +} + +/// GGUF (llama.cpp `olmoe` arch) tensor names → the HF-shaped names the remap +/// chain handles. `None` for anything unrecognized (a loud load error). +fn gguf_tensor_to_hf(info: &GgufTensorInfo) -> Option { + olmoe_gguf_name(&info.name).map(GgufMap::Rename) +} + +fn olmoe_gguf_name(name: &str) -> Option { + match name { + "token_embd.weight" => return Some("model.embed_tokens.weight".into()), + "output_norm.weight" => return Some("model.norm.weight".into()), + "output.weight" => return Some("lm_head.weight".into()), + _ => {} + } + let rest = name.strip_prefix("blk.")?; + let (layer, field) = rest.split_once('.')?; + let layer: usize = layer.parse().ok()?; + let mapped = match field { + "attn_norm.weight" => "input_layernorm.weight", + "ffn_norm.weight" => "post_attention_layernorm.weight", + "attn_q.weight" => "self_attn.q_proj.weight", + "attn_k.weight" => "self_attn.k_proj.weight", + "attn_v.weight" => "self_attn.v_proj.weight", + "attn_q_norm.weight" => "self_attn.q_norm.weight", + "attn_k_norm.weight" => "self_attn.k_norm.weight", + "attn_output.weight" => "self_attn.o_proj.weight", + // The router Linear — a plain 2-D weight, transposed by the adapter. + "ffn_gate_inp.weight" => "mlp.gate.weight", + // The fused 3-D expert banks — raw params, no `.weight` suffix in the + // module path. ggml dims reverse to [experts, out, in], exactly the + // `MoeExperts` layout. + "ffn_gate_exps.weight" => "mlp.experts.gate", + "ffn_up_exps.weight" => "mlp.experts.up", + "ffn_down_exps.weight" => "mlp.experts.down", + _ => return None, + }; + Some(format!("model.layers.{layer}.{mapped}")) +} + +/// Load an OLMoE model straight from a **GGUF** file: hyperparameters from +/// the `olmoe.*` metadata, weights dequantized to f32 and driven through the +/// same checked-load pipeline every other port uses. Budget note: the 1B-7B's +/// ~7B params dequantize to ~28 GB of f32 — size the target device (the +/// reference machine runs it on the 128 GB CPU backend). +pub fn load_from_gguf( + path: &Path, + device: &B::Device, +) -> Result, ImportError> { + let parse = |reason: String| ImportError::Parse { + file: path.to_path_buf(), + reason, + }; + let f = GgufFile::open(path).map_err(|e| parse(e.to_string()))?; + let config = OlmoeConfig::from_gguf(&f).map_err(parse)?; + let blob = f + .dequant_to_safetensors(&gguf_tensor_to_hf) + .map_err(|e| parse(e.to_string()))?; + assert!(blob.len() > 8, "a parsed GGUF yields a non-empty blob"); + + let mut model = build::(&config, device); + // The backend's float dtype, taken from the TYPE (`B::FloatElem`), never + // from a probe tensor (per-device default policy hazard). + let target_float = ::dtype(); + let mut store = install_remaps( + SafetensorsStore::from_bytes(Some(blob)) + .with_from_adapter(PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float))) + .allow_partial(true), + ); + load_checked(&mut model, &mut store, path)?; + Ok(LoadedOlmoe { model, config }) +} + +impl CausalLm for LoadedOlmoe { + type Cache = Vec>; + + fn new_cache(&self) -> Self::Cache { + (0..self.config.num_hidden_layers).map(|_| None).collect() + } + + fn is_eos(&self, id: u32) -> bool { + self.config.eos_token_id.contains(id) + } + + fn forward( + &self, + new_ids: &[u32], + past: usize, + cache: &mut Self::Cache, + device: &B::Device, + ) -> Tensor { + let t = new_ids.len(); + assert!(t >= 1, "OLMoE forward: need at least one token"); + assert!( + cache.len() == self.config.num_hidden_layers, + "OLMoE forward: cache has {} layers, model has {}", + cache.len(), + self.config.num_hidden_layers + ); + let cfg = &self.config; + let hd = cfg.head_dim(); + + // Dtype pinned to the backend TYPE, never the per-device policy. + let ids32: Vec = new_ids.iter().map(|&i| i as i32).collect(); + let input = Tensor::::from_data( + TensorData::new(ids32, [t]), + (device, crate::backend::int_dtype::()), + ) + .reshape([1, t]); + let mut x = self.model.embed_tokens.forward(input); // [1, t, hidden] + + let (cos, sin) = rope_tables::(t, past, hd, cfg.rope_theta, device); + let mask = (t > 1).then(|| causal_mask::(t, past, device)); + + for (layer, kv) in self.model.layers.iter().zip(cache.iter_mut()) { + let h = layer.input_layernorm.forward(x.clone()); + let h = layer.self_attn.forward( + h, + cfg.num_attention_heads, + cfg.num_key_value_heads, + hd, + &cos, + &sin, + mask.as_ref(), + kv, + ); + x = x.add(h); + let h2 = layer.post_attention_layernorm.forward(x.clone()); + x = x.add( + layer + .mlp + .forward(h2, cfg.num_experts_per_tok, cfg.norm_topk_prob), + ); + } + let x = self.model.norm.forward(x); + + let last = x.narrow(1, t - 1, 1).reshape([1, cfg.hidden_size]); + debug_assert!( + self.model.lm_head.is_some() != cfg.tie_word_embeddings, + "lm_head presence must match the config's tie flag" + ); + match &self.model.lm_head { + Some(head) => head.forward(last), // [1, vocab] + None => { + let w = self.model.embed_tokens.weight.val(); // [vocab, hidden] + last.matmul(w.swap_dims(0, 1)) // [1, vocab] + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::backend::Cpu; + use crate::gguf::{GgmlType, GgufTensorInfo}; + + type Dev = burn::tensor::Device; + + /// A synthetic toy MoE config: 4 experts, top-2, MHA, untied head. + fn toy_config() -> OlmoeConfig { + OlmoeConfig { + vocab_size: 64, + hidden_size: 16, + intermediate_size: 8, + num_hidden_layers: 2, + num_attention_heads: 4, + num_key_value_heads: 4, + num_experts: 4, + num_experts_per_tok: 2, + norm_topk_prob: false, + rms_norm_eps: 1e-5, + rope_theta: 1e4, + tie_word_embeddings: false, + eos_token_id: EosIds::One(2), + } + } + + #[test] + fn config_parses_the_real_1b_7b_shape() { + // The real OLMoE-1B-7B-0125-Instruct config.json shape. + let json = br#"{ + "vocab_size": 50304, "hidden_size": 2048, "intermediate_size": 1024, + "num_hidden_layers": 16, "num_attention_heads": 16, "num_key_value_heads": 16, + "num_experts": 64, "num_experts_per_tok": 8, "norm_topk_prob": false, + "rms_norm_eps": 1e-05, "rope_theta": 10000.0, + "tie_word_embeddings": false, "eos_token_id": 50279 + }"#; + let cfg = OlmoeConfig::from_json_bytes(json).unwrap(); + assert_eq!(cfg.head_dim(), 128); + assert_eq!(cfg.num_experts, 64); + assert_eq!(cfg.num_experts_per_tok, 8); + assert!(!cfg.norm_topk_prob); + assert!(!cfg.tie_word_embeddings); + assert!(cfg.eos_token_id.contains(50_279)); + } + + #[test] + fn config_rejects_bad_expert_counts() { + let mut cfg = toy_config(); + cfg.num_experts = 1; + assert!(cfg.validate().is_err(), "one expert is not a mixture"); + let mut cfg = toy_config(); + cfg.num_experts_per_tok = 5; + assert!(cfg.validate().is_err(), "top-k above expert count"); + let mut cfg = toy_config(); + cfg.num_experts_per_tok = 0; + assert!(cfg.validate().is_err(), "zero top-k"); + } + + /// The load-bearing invariant: cached prefill+decode == one full forward, + /// through the MoE layers and the projection-wide q/k norm. + #[test] + fn toy_model_cached_decode_matches_full_forward() { + let device = Dev::default(); + let cfg = toy_config(); + let loaded = LoadedOlmoe:: { + model: build(&cfg, &device), + config: cfg, + }; + + let prompt: Vec = vec![3, 14, 15, 9, 26]; + let mut cache = loaded.new_cache(); + let _ = loaded.forward(&prompt, 0, &mut cache, &device); + let step = loaded + .forward(&[42], prompt.len(), &mut cache, &device) + .into_data() + .to_vec::() + .unwrap(); + + let mut full_cache = loaded.new_cache(); + let all: Vec = prompt.iter().copied().chain([42]).collect(); + let full = loaded + .forward(&all, 0, &mut full_cache, &device) + .into_data() + .to_vec::() + .unwrap(); + + assert_eq!(step.len(), full.len()); + for (i, (c, f)) in step.iter().zip(&full).enumerate() { + assert!((c - f).abs() < 1e-4, "logit {i}: cached {c} vs full {f}"); + } + } + + #[test] + fn projection_qk_norm_spans_the_whole_projection() { + let device = Dev::default(); + let cfg = toy_config(); + let model = build::(&cfg, &device); + let q_dim = cfg.num_attention_heads * cfg.head_dim(); + assert_eq!( + model.layers[0] + .self_attn + .q_norm + .as_ref() + .unwrap() + .gamma + .dims(), + [q_dim], + "OLMoE q_norm must span num_heads * head_dim, not head_dim" + ); + } + + /// A synthetic in-memory GGUF header shaped like a small `olmoe` file. + fn toy_gguf() -> GgufFile { + let meta = |k: &str, v: GgufValue| (k.to_string(), v); + GgufFile { + path: std::path::PathBuf::new(), + version: 3, + metadata: vec![ + meta("general.architecture", GgufValue::Str("olmoe".into())), + meta("olmoe.embedding_length", GgufValue::U32(16)), + meta("olmoe.block_count", GgufValue::U32(2)), + meta("olmoe.feed_forward_length", GgufValue::U32(8)), + meta("olmoe.attention.head_count", GgufValue::U32(4)), + meta("olmoe.attention.head_count_kv", GgufValue::U32(4)), + meta("olmoe.expert_count", GgufValue::U32(4)), + meta("olmoe.expert_used_count", GgufValue::U32(2)), + meta( + "olmoe.attention.layer_norm_rms_epsilon", + GgufValue::F32(1e-5), + ), + meta("olmoe.rope.freq_base", GgufValue::F32(1e4)), + meta("tokenizer.ggml.eos_token_id", GgufValue::U32(2)), + ], + tensors: vec![ + GgufTensorInfo { + name: "token_embd.weight".into(), + dims: vec![16, 64], // ggml order: [hidden, vocab] + dtype: GgmlType::F32, + offset: 0, + }, + GgufTensorInfo { + name: "output.weight".into(), + dims: vec![16, 64], + dtype: GgmlType::F32, + offset: 4096, + }, + ], + alignment: 32, + data_offset: 0, + } + } + + #[test] + fn config_from_gguf_reads_expert_metadata() { + let cfg = OlmoeConfig::from_gguf(&toy_gguf()).expect("parses"); + assert_eq!(cfg.num_experts, 4); + assert_eq!(cfg.num_experts_per_tok, 2); + assert_eq!(cfg.intermediate_size, 8); + assert!(!cfg.norm_topk_prob); + assert!(!cfg.tie_word_embeddings); // output.weight present → untied + assert!(cfg.eos_token_id.contains(2)); + } + + #[test] + fn config_from_gguf_prefers_expert_feed_forward_length() { + let mut f = toy_gguf(); + f.metadata.push(( + "olmoe.expert_feed_forward_length".into(), + GgufValue::U32(12), + )); + assert_eq!(OlmoeConfig::from_gguf(&f).unwrap().intermediate_size, 12); + } + + #[test] + fn config_from_gguf_fails_loudly_on_missing_keys_and_wrong_arch() { + let mut f = toy_gguf(); + f.metadata.retain(|(k, _)| k != "olmoe.expert_count"); + assert!( + OlmoeConfig::from_gguf(&f) + .unwrap_err() + .contains("expert_count") + ); + + let mut f = toy_gguf(); + f.metadata[0].1 = GgufValue::Str("qwen3".into()); + assert!(OlmoeConfig::from_gguf(&f).is_err()); + } + + #[test] + fn gguf_names_map_router_and_fused_expert_banks() { + assert_eq!( + olmoe_gguf_name("blk.0.ffn_gate_inp.weight").as_deref(), + Some("model.layers.0.mlp.gate.weight") + ); + // Fused expert banks are raw params — no `.weight` suffix. + assert_eq!( + olmoe_gguf_name("blk.3.ffn_gate_exps.weight").as_deref(), + Some("model.layers.3.mlp.experts.gate") + ); + assert_eq!( + olmoe_gguf_name("blk.15.ffn_down_exps.weight").as_deref(), + Some("model.layers.15.mlp.experts.down") + ); + assert_eq!( + olmoe_gguf_name("blk.1.attn_q_norm.weight").as_deref(), + Some("model.layers.1.self_attn.q_norm.weight") + ); + assert_eq!( + olmoe_gguf_name("output.weight").as_deref(), + Some("lm_head.weight") + ); + // A dense FFN tensor is not part of this architecture — loud None. + assert_eq!(olmoe_gguf_name("blk.0.ffn_gate.weight"), None); + assert_eq!(olmoe_gguf_name("rope_freqs.weight"), None); + } + + #[test] + fn greedy_generate_respects_max_tokens_bound() { + let device = Dev::default(); + let cfg = toy_config(); + let loaded = LoadedOlmoe:: { + model: build(&cfg, &device), + config: cfg, + }; + let out = loaded.greedy_generate(&[1, 2, 3], 4, &device).unwrap(); + assert!(out.len() <= 4); + } +} diff --git a/crates/mummu/src/models/qwen2.rs b/crates/mummu/src/models/qwen2.rs index a2eda9a..770a0b8 100644 --- a/crates/mummu/src/models/qwen2.rs +++ b/crates/mummu/src/models/qwen2.rs @@ -217,6 +217,7 @@ fn build(cfg: &Qwen2Config, device: &B::Device) -> Qwen2 { head_dim: cfg.head_dim, bias: true, // Qwen2 has q/k/v projection bias qk_norm_eps: None, // and no per-head q/k norm + qk_norm_projection: false, }; let mlp_cfg = SwiGluMlpConfig { hidden_size: cfg.hidden_size, diff --git a/crates/mummu/src/models/qwen3.rs b/crates/mummu/src/models/qwen3.rs index 0587b06..be273b0 100644 --- a/crates/mummu/src/models/qwen3.rs +++ b/crates/mummu/src/models/qwen3.rs @@ -182,6 +182,7 @@ fn build(cfg: &Qwen3Config, device: &B::Device) -> Qwen3 { head_dim: cfg.head_dim, bias: false, // Qwen3 dropped the q/k/v bias qk_norm_eps: Some(cfg.rms_norm_eps), // and added per-head q/k RMSNorm + qk_norm_projection: false, }; let mlp_cfg = SwiGluMlpConfig { hidden_size: cfg.hidden_size, diff --git a/crates/mummu/src/nn/attention.rs b/crates/mummu/src/nn/attention.rs index f811d71..f5aaf9a 100644 --- a/crates/mummu/src/nn/attention.rs +++ b/crates/mummu/src/nn/attention.rs @@ -79,8 +79,12 @@ pub struct GqaAttentionConfig { /// Projection bias on q/k/v (Qwen2: true; LFM2: false). `o_proj` never /// has bias in either. pub bias: bool, - /// Per-head q/k RMSNorm (LFM2: eps of the model; Qwen2: `None`). + /// q/k RMSNorm epsilon (LFM2/Qwen3/OLMoE: eps of the model; Qwen2: `None`). pub qk_norm_eps: Option, + /// Where the q/k norm applies: `false` = per-head over `head_dim` + /// (LFM2/Qwen3), `true` = over the **whole projection** before the head + /// split (OLMoE — its `q_norm`/`k_norm` span `num_heads * head_dim`). + pub qk_norm_projection: bool, } impl GqaAttentionConfig { @@ -99,10 +103,16 @@ impl GqaAttentionConfig { ); let q_dim = self.num_heads * self.head_dim; let kv_dim = self.num_kv_heads * self.head_dim; - let norm = |eps: f64| { - RmsNormConfig::new(self.head_dim) - .with_epsilon(eps) - .init(device) + let norm = |eps: f64, dim: usize| RmsNormConfig::new(dim).with_epsilon(eps).init(device); + let q_norm_dim = if self.qk_norm_projection { + q_dim + } else { + self.head_dim + }; + let k_norm_dim = if self.qk_norm_projection { + kv_dim + } else { + self.head_dim }; GqaAttention { q_proj: LinearConfig::new(self.hidden_size, q_dim) @@ -117,12 +127,38 @@ impl GqaAttentionConfig { o_proj: LinearConfig::new(q_dim, self.hidden_size) .with_bias(false) .init(device), - q_norm: self.qk_norm_eps.map(norm), - k_norm: self.qk_norm_eps.map(norm), + q_norm: self.qk_norm_eps.map(|eps| norm(eps, q_norm_dim)), + k_norm: self.qk_norm_eps.map(|eps| norm(eps, k_norm_dim)), } } } +/// Apply a q/k RMSNorm at the placement its gamma width implies: `head_dim` → +/// per-head at `[b, t, n, hd]` (LFM2/Qwen3), `n * head_dim` → over the whole +/// projection **before** the head split (OLMoE). The two coincide at `n == 1`. +/// Inferring from the loaded gamma keeps the module shape identical across +/// families — a checkpoint's own norm width picks its semantics. +fn qk_norm_forward( + norm: &RmsNorm, + x: Tensor, // [b, t, n*hd] + n: usize, + hd: usize, +) -> Tensor { + let [b, t, width] = x.dims(); + debug_assert!(width == n * hd, "q/k projection width must be n * head_dim"); + let gamma = norm.gamma.dims()[0]; + assert!( + gamma == hd || gamma == n * hd, + "q/k norm width {gamma} matches neither head_dim ({hd}) nor the projection width ({})", + n * hd + ); + if gamma == n * hd && n > 1 { + norm.forward(x).reshape([b, t, n, hd]) + } else { + norm.forward(x.reshape([b, t, n, hd])) + } +} + impl GqaAttention { /// Cache-aware forward: RoPE the new q/k at the offset positions, append /// the new k/v to this layer's cache, attend over the full cached range. @@ -152,18 +188,20 @@ impl GqaAttention { "GQA forward: q_norm and k_norm must be both present or both absent" ); - // Per-head q/k RMSNorm (when present) applies post-projection, before + // q/k RMSNorm (when present) applies post-projection, before // transpose + RoPE — the LFM2 ordering, validated against Ollama. - let q = self.q_proj.forward(x.clone()).reshape([b, t, nh, hd]); + // Placement (per-head vs whole-projection) follows the loaded norm's + // own width; see `qk_norm_forward`. + let q = self.q_proj.forward(x.clone()); let q = match &self.q_norm { - Some(norm) => norm.forward(q), - None => q, + Some(norm) => qk_norm_forward(norm, q, nh, hd), + None => q.reshape([b, t, nh, hd]), } .swap_dims(1, 2); - let k_new = self.k_proj.forward(x.clone()).reshape([b, t, nkv, hd]); + let k_new = self.k_proj.forward(x.clone()); let k_new = match &self.k_norm { - Some(norm) => norm.forward(k_new), - None => k_new, + Some(norm) => qk_norm_forward(norm, k_new, nkv, hd), + None => k_new.reshape([b, t, nkv, hd]), } .swap_dims(1, 2); let v_new = self @@ -224,7 +262,7 @@ mod tests { const HEAD_DIM: usize = 4; const THETA: f32 = 1e4; - fn attn(qk_norm: bool, device: &Dev) -> GqaAttention { + fn attn(qk_norm: bool, projection: bool, device: &Dev) -> GqaAttention { GqaAttentionConfig { hidden_size: HIDDEN, num_heads: HEADS, @@ -232,6 +270,7 @@ mod tests { head_dim: HEAD_DIM, bias: true, qk_norm_eps: qk_norm.then_some(1e-5), + qk_norm_projection: projection, } .init(device) } @@ -306,8 +345,8 @@ mod tests { #[test] fn kv_cache_decode_matches_full_forward() { let device = Dev::default(); - for qk_norm in [false, true] { - let a = attn(qk_norm, &device); + for (qk_norm, projection) in [(false, false), (true, false), (true, true)] { + let a = attn(qk_norm, projection, &device); let x = input(6, 3.0, &device); // Reference: all 6 positions in one causal forward; keep the last row. @@ -341,7 +380,7 @@ mod tests { for (i, (c, f)) in out.iter().zip(last_full).enumerate() { assert!( (c - f).abs() < 1e-4, - "qk_norm={qk_norm} elem {i}: cached {c} vs full {f}" + "qk_norm={qk_norm} projection={projection} elem {i}: cached {c} vs full {f}" ); } } @@ -351,7 +390,7 @@ mod tests { #[test] fn future_tokens_cannot_affect_past_outputs() { let device = Dev::default(); - let a = attn(false, &device); + let a = attn(false, false, &device); let x1 = input(4, 1.0, &device); // Same first 3 tokens, different 4th. let x2 = Tensor::cat(vec![x1.clone().narrow(1, 0, 3), input(1, 99.0, &device)], 1); @@ -378,7 +417,43 @@ mod tests { head_dim: HEAD_DIM, bias: false, qk_norm_eps: None, + qk_norm_projection: false, } .init::(&device); } + + /// Negative space: the projection-wide norm is a different function than + /// the per-head norm (RMS over 16 values vs over 4) — same weights, same + /// input, different outputs. Guards against the placement silently + /// collapsing to one branch. + #[test] + fn projection_norm_differs_from_per_head_norm() { + let device = Dev::default(); + let per_head = attn(true, false, &device); + // Same module, but re-shaped norms: reuse per_head's projections and + // swap in projection-wide norms with unit gamma? Simpler: two configs + // share no weights, so instead check the norm widths took effect. + let projection = attn(true, true, &device); + let q_dim = HEADS * HEAD_DIM; + assert_eq!(per_head.q_norm.as_ref().unwrap().gamma.dims(), [HEAD_DIM]); + assert_eq!(projection.q_norm.as_ref().unwrap().gamma.dims(), [q_dim]); + // And the projection-placement forward is exercised end to end by the + // cache-equivalence loop above; here pin that a projection-normed + // forward actually runs (no panic) and returns the right shape. + let x = input(3, 5.0, &device); + let (cos, sin) = rope_tables::(3, 0, HEAD_DIM, THETA, &device); + let mask = causal_mask::(3, 0, &device); + let mut kv: LayerKv = None; + let out = projection.forward( + x, + HEADS, + KV_HEADS, + HEAD_DIM, + &cos, + &sin, + Some(&mask), + &mut kv, + ); + assert_eq!(out.dims(), [1, 3, HIDDEN]); + } } diff --git a/crates/mummu/src/nn/mod.rs b/crates/mummu/src/nn/mod.rs index 0f606a7..7d73fac 100644 --- a/crates/mummu/src/nn/mod.rs +++ b/crates/mummu/src/nn/mod.rs @@ -8,11 +8,13 @@ mod attention; mod conv; mod mlp; +mod moe; mod rope; pub use attention::{GqaAttention, GqaAttentionConfig, LayerKv, causal_mask, repeat_kv}; pub use conv::{ConvState, ShortConv, ShortConvConfig}; pub use mlp::{SwiGluMlp, SwiGluMlpConfig}; +pub use moe::{MoeExperts, SparseMoe, SparseMoeConfig}; pub use rope::{apply_rope, rope_tables, rotate_half}; /// Hard ceiling on `past + t` everywhere a sequence position is materialized. diff --git a/crates/mummu/src/nn/moe.rs b/crates/mummu/src/nn/moe.rs new file mode 100644 index 0000000..0a3742b --- /dev/null +++ b/crates/mummu/src/nn/moe.rs @@ -0,0 +1,445 @@ +//! Sparse mixture-of-experts SwiGLU feed-forward (OLMoE-style): a softmax +//! top-k router over a bank of SwiGLU experts stored as **fused 3-D tensors** +//! — exactly the GGUF `ffn_{gate,up,down}_exps` layout, so a checkpoint's +//! expert bank loads as one tensor per projection instead of `num_experts` +//! separate matrices. +//! +//! First cut is **dense-mask compute**: every expert processes every token and +//! the router's sparse weight row (zero for the unrouted experts) scales the +//! results away. That wastes `1 - k/E` of the FLOPs but keeps the whole +//! forward on-device (no data-dependent gather, no host readback of routing +//! decisions) and is numerically identical to the sparse formulation. A +//! gather-based path is a perf follow-up, not a correctness one. + +use burn::module::{Module, Param}; +use burn::nn::{Linear, LinearConfig}; +use burn::tensor::{DType, Distribution, Int, Tensor, activation, backend::Backend}; + +/// The expert bank: `num_experts` SwiGLU MLPs as three fused params in +/// `[experts, out, in]` layout (the row-major twin of ggml's +/// `ffn_*_exps.weight`). Forward transposes lazily; no per-expert modules. +#[derive(Module, Debug)] +pub struct MoeExperts { + /// `[num_experts, intermediate, hidden]` — SiLU branch. + pub gate: Param>, + /// `[num_experts, intermediate, hidden]` — multiplicative branch. + pub up: Param>, + /// `[num_experts, hidden, intermediate]` — back to the model width. + pub down: Param>, +} + +/// Router + expert bank. Field names follow the HF `Olmoe` checkpoint layout +/// (`mlp.gate` is the router Linear, `mlp.experts` the bank). +#[derive(Module, Debug)] +pub struct SparseMoe { + /// The routing projection: `hidden -> num_experts`, no bias. + pub gate: Linear, + pub experts: MoeExperts, +} + +/// Shape config for [`SparseMoe`]. +#[derive(Debug, Clone)] +pub struct SparseMoeConfig { + pub hidden_size: usize, + /// Per-expert SwiGLU intermediate width (OLMoE: 1024 — each expert is + /// narrow; capacity comes from the count). + pub expert_intermediate_size: usize, + pub num_experts: usize, + pub num_experts_per_tok: usize, +} + +impl SparseMoeConfig { + /// Initialize the module (random weights; real weights come from import). + pub fn init(&self, device: &B::Device) -> SparseMoe { + assert!( + self.num_experts >= 2, + "MoE: num_experts must be >= 2 (got {}); use SwiGluMlp for a dense FFN", + self.num_experts + ); + assert!( + (1..=self.num_experts).contains(&self.num_experts_per_tok), + "MoE: num_experts_per_tok ({}) must be in 1..=num_experts ({})", + self.num_experts_per_tok, + self.num_experts + ); + assert!( + self.hidden_size >= 1 && self.expert_intermediate_size >= 1, + "MoE: hidden_size and expert_intermediate_size must be >= 1" + ); + let (e, h, inter) = ( + self.num_experts, + self.hidden_size, + self.expert_intermediate_size, + ); + // Linear-style uniform init, bound by each projection's fan-in. + let init = |out: usize, inp: usize, dev: &B::Device| { + let bound = 1.0 / (inp as f64).sqrt(); + Param::from_tensor(Tensor::random( + [e, out, inp], + Distribution::Uniform(-bound, bound), + dev, + )) + }; + SparseMoe { + gate: LinearConfig::new(h, e).with_bias(false).init(device), + experts: MoeExperts { + gate: init(inter, h, device), + up: init(inter, h, device), + down: init(h, inter, device), + }, + } + } +} + +impl SparseMoe { + /// `[b, t, hidden]` → `[b, t, hidden]`. + /// + /// Router math mirrors HF `OlmoeSparseMoeBlock`: softmax over **all** + /// experts in f32, keep the top-`top_k` probabilities as the mixture + /// weights (renormalized to sum 1 iff `norm_topk_prob` — OLMoE ships + /// `false`). The f32 island matters on f16 backends (softmax of wide + /// logits); every cast is a no-op on f32. + pub fn forward(&self, x: Tensor, top_k: usize, norm_topk_prob: bool) -> Tensor { + let [b, t, h] = x.dims(); + let [e, _inter, h_in] = self.experts.gate.dims(); + assert!( + (1..=e).contains(&top_k), + "MoE forward: top_k ({top_k}) must be in 1..=num_experts ({e})" + ); + assert!( + h == h_in, + "MoE forward: input hidden {h} does not match expert hidden {h_in}" + ); + debug_assert!( + self.experts.down.dims()[1] == h, + "MoE forward: down projection must return to the model width" + ); + let ambient = x.dtype(); + let bt = b * t; + let xt = x.reshape([bt, h]); + + // Router → dense per-token weight rows [bt, e]: softmax probabilities + // where an expert is in the token's top-k, exact zero elsewhere. The + // scatter is an on-device arange-compare — burn's `one_hot` reads the + // indices back to the host, which would sync every layer. + let logits = self.gate.forward(xt.clone()); // [bt, e] + debug_assert!(logits.dims() == [bt, e], "router width must be num_experts"); + let probs = activation::softmax(logits.cast(DType::F32), 1); + let (vals, idx) = probs.topk_with_indices(top_k, 1); // both [bt, k] + let vals = if norm_topk_prob { + vals.clone().div(vals.sum_dim(1)) // [bt, k] / [bt, 1] + } else { + vals + }; + let classes = + Tensor::::arange(0..e as i64, &xt.device()).reshape([1, 1, e as i32]); + let hit = idx + .reshape([bt, top_k, 1]) + .equal(classes.expand([bt, top_k, e])); // [bt, k, e] + let weights = hit + .float() + .cast(DType::F32) + .mul(vals.reshape([bt, top_k, 1])) + .sum_dim(1) // [bt, 1, e] + .reshape([bt, e]) + .cast(ambient); + + // Dense expert compute: one batched matmul per projection across the + // whole bank ([1, bt, h] broadcast against [e, h, *]), then the weight + // rows zero out the unrouted experts in the reduction. + let xb = xt.reshape([1, bt, h]); + let gate = xb.clone().matmul(self.experts.gate.val().swap_dims(1, 2)); // [e, bt, inter] + let up = xb.matmul(self.experts.up.val().swap_dims(1, 2)); + let acts = activation::silu(gate).mul(up); + let out = acts.matmul(self.experts.down.val().swap_dims(1, 2)); // [e, bt, h] + let w_per_expert = weights.swap_dims(0, 1).reshape([e, bt, 1]); + out.mul(w_per_expert) + .sum_dim(0) // [1, bt, h] + .reshape([b, t, h]) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::backend::Cpu; + use burn::tensor::TensorData; + + type Dev = burn::tensor::Device; + + const HIDDEN: usize = 4; + const INTER: usize = 3; + const EXPERTS: usize = 4; + const TOP_K: usize = 2; + + /// Deterministic weights: expert `e`'s matrices are small distinct + /// sinusoids so every expert computes something different. + fn moe(device: &Dev) -> SparseMoe { + let fill = |seed: f32, dims: [usize; 3]| { + let n = dims[0] * dims[1] * dims[2]; + let data: Vec = (0..n) + .map(|i| ((i as f32) * 0.37 + seed).sin() * 0.5) + .collect(); + Param::from_tensor(Tensor::::from_data( + TensorData::new(data, dims), + device, + )) + }; + let router_data: Vec = (0..HIDDEN * EXPERTS) + .map(|i| ((i as f32) * 0.61 + 1.0).cos() * 0.5) + .collect(); + let mut m = SparseMoeConfig { + hidden_size: HIDDEN, + expert_intermediate_size: INTER, + num_experts: EXPERTS, + num_experts_per_tok: TOP_K, + } + .init::(device); + // Burn Linear stores weight as [in, out]. + m.gate.weight = Param::from_tensor(Tensor::::from_data( + TensorData::new(router_data, [HIDDEN, EXPERTS]), + device, + )); + m.experts.gate = fill(0.1, [EXPERTS, INTER, HIDDEN]); + m.experts.up = fill(1.7, [EXPERTS, INTER, HIDDEN]); + m.experts.down = fill(3.3, [EXPERTS, HIDDEN, INTER]); + m + } + + fn input(t: usize, seed: f32, device: &Dev) -> Tensor { + let data: Vec = (0..t * HIDDEN) + .map(|i| ((i as f32 + seed) * 0.9).sin()) + .collect(); + Tensor::::from_data(TensorData::new(data, [t * HIDDEN]), device) + .reshape([1, t, HIDDEN]) + } + + fn silu(x: f32) -> f32 { + x / (1.0 + (-x).exp()) + } + + /// Hand-rolled f32 reference of the whole block (per token: router + /// softmax, top-k, sparse weighted sum of per-expert SwiGLUs). + fn reference(m: &SparseMoe, x: &[f32], t: usize, norm: bool) -> Vec { + let rw = m.gate.weight.val().into_data().to_vec::().unwrap(); // [h, e] + let gw = m.experts.gate.val().into_data().to_vec::().unwrap(); // [e, inter, h] + let uw = m.experts.up.val().into_data().to_vec::().unwrap(); + let dw = m.experts.down.val().into_data().to_vec::().unwrap(); // [e, h, inter] + let mut out = vec![0f32; t * HIDDEN]; + for tok in 0..t { + let xrow = &x[tok * HIDDEN..][..HIDDEN]; + // Router logits then softmax over all experts. + let mut logits = [0f32; EXPERTS]; + for (e, logit) in logits.iter_mut().enumerate() { + *logit = (0..HIDDEN).map(|i| xrow[i] * rw[i * EXPERTS + e]).sum(); + } + let max = logits.iter().cloned().fold(f32::MIN, f32::max); + let exps: Vec = logits.iter().map(|l| (l - max).exp()).collect(); + let z: f32 = exps.iter().sum(); + let probs: Vec = exps.iter().map(|v| v / z).collect(); + // Top-k expert ids by probability. + let mut order: Vec = (0..EXPERTS).collect(); + order.sort_by(|&a, &b| probs[b].partial_cmp(&probs[a]).unwrap()); + let picked = &order[..TOP_K]; + let denom: f32 = if norm { + picked.iter().map(|&e| probs[e]).sum() + } else { + 1.0 + }; + for &e in picked { + let w = probs[e] / denom; + // SwiGLU of expert e. + let mut act = [0f32; INTER]; + for (j, a) in act.iter_mut().enumerate() { + let g: f32 = (0..HIDDEN) + .map(|i| xrow[i] * gw[(e * INTER + j) * HIDDEN + i]) + .sum(); + let u: f32 = (0..HIDDEN) + .map(|i| xrow[i] * uw[(e * INTER + j) * HIDDEN + i]) + .sum(); + *a = silu(g) * u; + } + for i in 0..HIDDEN { + let d: f32 = (0..INTER) + .map(|j| act[j] * dw[(e * HIDDEN + i) * INTER + j]) + .sum(); + out[tok * HIDDEN + i] += w * d; + } + } + } + out + } + + #[test] + fn forward_matches_the_hand_rolled_sparse_reference() { + let device = Dev::default(); + let m = moe(&device); + // Both shapes the decoder actually runs: a multi-token prefill and + // the single-token decode step. + for t in [5, 1] { + for norm in [false, true] { + let x = input(t, 2.0, &device); + let xv = x.clone().into_data().to_vec::().unwrap(); + let got = m + .forward(x, TOP_K, norm) + .into_data() + .to_vec::() + .unwrap(); + let want = reference(&m, &xv, t, norm); + assert_eq!(got.len(), want.len()); + for (i, (g, w)) in got.iter().zip(&want).enumerate() { + assert!( + (g - w).abs() < 1e-5, + "t={t} norm={norm} elem {i}: got {g} vs reference {w}" + ); + } + } + } + } + + #[test] + fn top_k_equal_to_num_experts_uses_every_expert() { + // With k == E and renorm the block degenerates to a full softmax + // mixture — the reference covers it; this pins the k=E edge. + let device = Dev::default(); + let m = moe(&device); + let x = input(3, 0.5, &device); + let xv = x.clone().into_data().to_vec::().unwrap(); + let got = m + .forward(x, EXPERTS, false) + .into_data() + .to_vec::() + .unwrap(); + // Reference with TOP_K replaced by all experts: weights are the full + // softmax row, every expert contributes. + let mut want = vec![0f32; 3 * HIDDEN]; + { + let full = reference_all_experts(&m, &xv, 3); + want.copy_from_slice(&full); + } + for (i, (g, w)) in got.iter().zip(&want).enumerate() { + assert!( + (g - w).abs() < 1e-5, + "elem {i}: got {g} vs full-mixture {w}" + ); + } + } + + /// Full-mixture reference (every expert, softmax-weighted) for the k=E edge. + fn reference_all_experts(m: &SparseMoe, x: &[f32], t: usize) -> Vec { + let rw = m.gate.weight.val().into_data().to_vec::().unwrap(); + let gw = m.experts.gate.val().into_data().to_vec::().unwrap(); + let uw = m.experts.up.val().into_data().to_vec::().unwrap(); + let dw = m.experts.down.val().into_data().to_vec::().unwrap(); + let mut out = vec![0f32; t * HIDDEN]; + for tok in 0..t { + let xrow = &x[tok * HIDDEN..][..HIDDEN]; + let mut logits = [0f32; EXPERTS]; + for (e, logit) in logits.iter_mut().enumerate() { + *logit = (0..HIDDEN).map(|i| xrow[i] * rw[i * EXPERTS + e]).sum(); + } + let max = logits.iter().cloned().fold(f32::MIN, f32::max); + let exps: Vec = logits.iter().map(|l| (l - max).exp()).collect(); + let z: f32 = exps.iter().sum(); + for e in 0..EXPERTS { + let w = exps[e] / z; + let mut act = [0f32; INTER]; + for (j, a) in act.iter_mut().enumerate() { + let g: f32 = (0..HIDDEN) + .map(|i| xrow[i] * gw[(e * INTER + j) * HIDDEN + i]) + .sum(); + let u: f32 = (0..HIDDEN) + .map(|i| xrow[i] * uw[(e * INTER + j) * HIDDEN + i]) + .sum(); + *a = silu(g) * u; + } + for i in 0..HIDDEN { + let d: f32 = (0..INTER) + .map(|j| act[j] * dw[(e * HIDDEN + i) * INTER + j]) + .sum(); + out[tok * HIDDEN + i] += w * d; + } + } + } + out + } + + #[test] + fn norm_topk_weights_change_the_mixture() { + // norm_topk_prob renormalizes the k weights to sum 1 — unless the + // top-k already captured all the mass, outputs must differ. + let device = Dev::default(); + let m = moe(&device); + let x = input(4, 7.0, &device); + let a = m + .forward(x.clone(), TOP_K, false) + .into_data() + .to_vec::() + .unwrap(); + let b = m + .forward(x, TOP_K, true) + .into_data() + .to_vec::() + .unwrap(); + let differs = a.iter().zip(&b).any(|(x, y)| (x - y).abs() > 1e-7); + assert!(differs, "renormalized weights should scale the output"); + } + + #[test] + fn forward_is_position_independent() { + // MoE acts per-token: the same row in different positions/batches + // routes and computes identically. + let device = Dev::default(); + let m = moe(&device); + let row = input(1, 11.0, &device); + let double = Tensor::cat(vec![row.clone(), row.clone()], 1); + let s = m + .forward(row, TOP_K, false) + .into_data() + .to_vec::() + .unwrap(); + let d = m + .forward(double, TOP_K, false) + .into_data() + .to_vec::() + .unwrap(); + assert_eq!(s.as_slice(), &d[..HIDDEN]); + assert_eq!(s.as_slice(), &d[HIDDEN..]); + } + + #[test] + fn zero_input_gives_zero_output() { + // Bias-free SwiGLU experts map 0 to 0 regardless of routing. + let device = Dev::default(); + let m = moe(&device); + let x = Tensor::::zeros([1, 2, HIDDEN], &device); + let out = m + .forward(x, TOP_K, false) + .into_data() + .to_vec::() + .unwrap(); + assert!(out.iter().all(|&v| v == 0.0)); + } + + #[test] + #[should_panic(expected = "top_k")] + fn forward_rejects_top_k_above_num_experts() { + let device = Dev::default(); + let m = moe(&device); + let x = input(1, 0.0, &device); + let _ = m.forward(x, EXPERTS + 1, false); + } + + #[test] + #[should_panic(expected = "num_experts_per_tok")] + fn config_rejects_zero_top_k() { + let device = Dev::default(); + let _ = SparseMoeConfig { + hidden_size: 4, + expert_intermediate_size: 3, + num_experts: 4, + num_experts_per_tok: 0, + } + .init::(&device); + } +} diff --git a/crates/mummu/src/registry.rs b/crates/mummu/src/registry.rs index b955c7a..7ae6f02 100644 --- a/crates/mummu/src/registry.rs +++ b/crates/mummu/src/registry.rs @@ -20,6 +20,9 @@ pub enum Architecture { Lfm2, /// `models::minilm` — all-MiniLM BERT sentence embedder. MiniLm, + /// `models::olmoe` — OLMoE sparse mixture-of-experts decoder (the zoo's + /// first MoE; GGUF import only — HF ships the experts unfused). + Olmoe, } /// How the checkpoint's weights are stored — which fetch + load path a spec @@ -232,6 +235,19 @@ pub fn catalog() -> Vec { }, disk_bytes_estimate: 2_500_000_000, }, + // The zoo's first MoE: 64 experts, 8 active per token (1B active / + // 7B total). Resident-everything first cut — ~28 GB dequantized to + // f32, sized for the CPU backend (128 GB reference machine). + ModelSpec { + name: "olmoe-1b-7b-0125-instruct-q4km".into(), + repo: "allenai/OLMoE-1B-7B-0125-Instruct-GGUF".into(), + revision: "main".into(), + architecture: Architecture::Olmoe, + format: WeightFormat::Gguf { + file: "OLMoE-1B-7B-0125-Instruct-Q4_K_M.gguf".into(), + }, + disk_bytes_estimate: 4_210_000_000, + }, ]; debug_assert!( entries.iter().all(|s| s.validate().is_ok()), diff --git a/crates/mummu/src/tokenizer.rs b/crates/mummu/src/tokenizer.rs index 317dc54..07b00ed 100644 --- a/crates/mummu/src/tokenizer.rs +++ b/crates/mummu/src/tokenizer.rs @@ -61,6 +61,13 @@ fn pre_spec(pre: &str) -> Option { regex: r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+", nfc: false, }), + // OLMo/OLMoE (GPT-NeoX lineage): the stock GPT-2 regex, contractions + // case-SENSITIVE — the family's tokenizer.json is ByteLevel with + // `use_regex: true` (exactly this pattern) behind an NFC normalizer. + "olmo" => Some(PreSpec { + regex: r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+", + nfc: true, + }), _ => None, } } diff --git a/crates/mummu/tests/parity_gguf.rs b/crates/mummu/tests/parity_gguf.rs index 2e8acd2..1d54770 100644 --- a/crates/mummu/tests/parity_gguf.rs +++ b/crates/mummu/tests/parity_gguf.rs @@ -21,10 +21,10 @@ mod llama_ref; use std::path::PathBuf; use llama_ref::{LlamaServer, logprobs_at}; -use mummu::backend::Gpu; +use mummu::backend::{Cpu, Gpu}; use mummu::gguf::GgufFile; use mummu::models::CausalLm; -use mummu::models::{lfm2, qwen2, qwen3}; +use mummu::models::{lfm2, olmoe, qwen2, qwen3}; const PROMPT: &str = "List the first five prime numbers."; const MAX_TOKENS: usize = 24; @@ -63,15 +63,17 @@ fn next_port() -> u16 { } /// One quantized-reference comparison: `load` builds our model from the GGUF, -/// `render` wraps the prompt in the model's chat template. Panics (test -/// style) on any divergence. -fn compare_against_llama_cpp( +/// `render` wraps the prompt in the model's chat template. Generic over the +/// backend — the dense tiers compare on `Gpu`; OLMoE's ~28 GB f32 build only +/// fits the CPU backend. Panics (test style) on any divergence. +fn compare_against_llama_cpp( tag: &str, gguf: &std::path::Path, - load: impl FnOnce(&std::path::Path, &burn::tensor::Device) -> M, + load: impl FnOnce(&std::path::Path, &burn::tensor::Device) -> M, render: impl FnOnce(&str) -> String, ) where - M: CausalLm, + B: burn::tensor::backend::Backend, + M: CausalLm, { let exe = llama_ref::server_exe().expect("set MUMMU_LLAMA_SERVER to a llama.cpp llama-server binary"); @@ -104,7 +106,7 @@ fn compare_against_llama_cpp( "reference returned fewer than top-{TOP_K}" ); - let device = burn::tensor::Device::::default(); + let device = burn::tensor::Device::::default(); let loaded = load(gguf, &device); let mut cache = loaded.new_cache(); let logits = loaded @@ -216,3 +218,39 @@ fn lfm2_q4_gguf_matches_llama_cpp_on_the_same_file() { |user| mummu::chat::ChatMl::lfm2().render(&[mummu::chat::Turn::user(user)]), ); } + +#[test] +#[ignore = "needs the OLMoE Q4_K_M GGUF (MUMMU_OLMOE_GGUF_PATH), llama-server \ + (MUMMU_LLAMA_SERVER), and ~60 GB free RAM (28 GB f32 build, CPU backend)"] +fn olmoe_q4_gguf_matches_llama_cpp_on_the_same_file() { + // The zoo's first MoE leg. Our side runs on the CPU backend — the ~28 GB + // f32 resident-everything build does not fit a 16 GB card; parity is + // backend-independent by construction (same weights, same math). + let gguf = env_path( + "MUMMU_OLMOE_GGUF_PATH", + "the OLMoE-1B-7B-0125-Instruct q4_k_m gguf", + ); + // OLMoE has no hardcoded ChatMl renderer yet: render its zephyr-style + // template (from the GGUF's own chat_template metadata) by hand, BOS + // first — both sides get the identical id array, so no template stack is + // in the loop on either side. + let f = GgufFile::open(&gguf).expect("gguf header parses"); + let bos = f + .get("tokenizer.ggml.bos_token_id") + .and_then(mummu::gguf::GgufValue::as_u64) + .and_then(|v| { + f.get("tokenizer.ggml.tokens") + .and_then(mummu::gguf::GgufValue::as_array) + .and_then(|t| t.get(usize::try_from(v).ok()?)) + .and_then(mummu::gguf::GgufValue::as_str) + .map(String::from) + }) + .unwrap_or_default(); + drop(f); + compare_against_llama_cpp( + "olmoe", + &gguf, + |p, d| olmoe::load_from_gguf::(p, d).expect("gguf load checked"), + move |user| format!("{bos}<|user|>\n{user}\n<|assistant|>\n"), + ); +} diff --git a/crates/mummu/tests/real_olmoe.rs b/crates/mummu/tests/real_olmoe.rs new file mode 100644 index 0000000..9f80816 --- /dev/null +++ b/crates/mummu/tests/real_olmoe.rs @@ -0,0 +1,154 @@ +//! Real-weights validation of the OLMoE-1B-7B port (P2's first MoE): the +//! registry spec fetches the official allenai Q4_K_M GGUF, the ONE file +//! yields config + tokenizer + weights, and the model loads and greedy- +//! decodes coherently on the **CPU** backend — the resident-everything first +//! cut dequantizes ~7B params to ~28 GB of f32, which is CPU-RAM territory +//! (the 128 GB reference machine), not 16 GB-card territory. Ignored by +//! default; run with +//! +//! ```text +//! MUMMU_HUB_DEST=path/to/models-cache \ +//! MUMMU_OLMOE_TOK_JSON=path/to/olmoe/tokenizer.json \ +//! cargo test -p mummu --release --test real_olmoe -- --ignored --nocapture +//! ``` + +use std::path::PathBuf; + +use mummu::backend::Cpu; +use mummu::gguf::{GgufFile, GgufValue}; +use mummu::models::CausalLm; +use mummu::models::olmoe; + +fn hub_dest() -> PathBuf { + std::env::var_os("MUMMU_HUB_DEST") + .map(PathBuf::from) + .unwrap_or_else(|| panic!("set MUMMU_HUB_DEST to the models dir for the ~4.2 GB download")) +} + +fn fetch_olmoe() -> PathBuf { + let dest = hub_dest(); + let spec = mummu::registry::catalog() + .into_iter() + .find(|s| s.name == "olmoe-1b-7b-0125-instruct-q4km") + .expect("the OLMoE GGUF is in the catalog"); + spec.fetch(&dest, |_| {}).expect("registry fetch"); + let path = spec.gguf_path(&dest).expect("gguf specs have a file path"); + assert!(path.is_file(), "downloaded file exists at {path:?}"); + path +} + +/// The GGUF-built tokenizer must byte-match the checkpoint's own HF +/// `tokenizer.json` (fetched separately — the GGUF repo does not ship it). +#[test] +#[ignore = "needs network (MUMMU_HUB_DEST) + the HF tokenizer.json (MUMMU_OLMOE_TOK_JSON)"] +fn olmoe_gguf_tokenizer_matches_the_hf_tokenizer() { + let Some(tok_json) = std::env::var_os("MUMMU_OLMOE_TOK_JSON").map(PathBuf::from) else { + panic!("set MUMMU_OLMOE_TOK_JSON to the checkpoint's tokenizer.json"); + }; + let path = fetch_olmoe(); + let f = GgufFile::open(&path).expect("valid GGUF"); + let ours = mummu::tokenizer::tokenizer_from_gguf(&f).expect("tokenizer from metadata"); + let hf = tokenizers::Tokenizer::from_file(&tok_json).expect("hf tokenizer.json loads"); + + // The battery the other GGUF tokenizer gates use: template-ish specials, + // unicode/CJK/emoji, whitespace runs (OLMoE has dedicated multi-space + // added tokens), contractions (case-sensitive in this family), empty. + let battery = [ + "<|user|>\nWhat is 2 + 2?\n<|assistant|>\n", + "The quick brown fox jumps over the lazy dog.", + "I'll say it's DONE — they're not.", + "code:\n if x < 3 {\n return;\n }\n", + "日本語のテキストと emoji 🦉🌲 mixed in.", + "naïve café résumé ™ ½ filigree", + "|||EMAIL_ADDRESS||| and |||IP_ADDRESS||| specials", + "", + ]; + for prompt in battery { + let a = ours.encode(prompt, false).expect("ours encodes"); + let b = hf.encode(prompt, false).expect("hf encodes"); + assert_eq!( + a.get_ids(), + b.get_ids(), + "ids diverge on {prompt:?}:\n ours {:?}\n hf {:?}", + a.get_tokens(), + b.get_tokens() + ); + let ra = ours.decode(a.get_ids(), false).expect("ours decodes"); + let rb = hf.decode(b.get_ids(), false).expect("hf decodes"); + assert_eq!(ra, rb, "decode round-trips diverge on {prompt:?}"); + } + eprintln!( + "[real_olmoe] tokenizer byte-identical to tokenizer.json on {} prompts", + battery.len() + ); +} + +/// The end-to-end proof: the ONE .gguf file loads (checked, every tensor +/// mapped) and greedy-decodes a correct answer through the MoE stack on the +/// CPU backend. +#[test] +#[ignore = "needs network (MUMMU_HUB_DEST; ~4.2 GB) and ~60 GB free RAM (CPU backend)"] +fn olmoe_gguf_loads_and_decodes_on_cpu() { + let path = fetch_olmoe(); + let f = GgufFile::open(&path).expect("valid GGUF"); + assert_eq!(f.architecture(), Some("olmoe"), "wrong architecture"); + let cfg = olmoe::OlmoeConfig::from_gguf(&f).expect("config from metadata"); + assert_eq!(cfg.num_experts, 64, "1B-7B has 64 experts"); + assert_eq!(cfg.num_experts_per_tok, 8, "1B-7B routes top-8"); + let tok = mummu::tokenizer::tokenizer_from_gguf(&f).expect("tokenizer from metadata"); + // The zephyr-style template from the model card, BOS first. + let bos = f + .get("tokenizer.ggml.bos_token_id") + .and_then(GgufValue::as_u64) + .and_then(|v| { + f.get("tokenizer.ggml.tokens") + .and_then(GgufValue::as_array) + .and_then(|t| t.get(usize::try_from(v).ok()?)) + .and_then(GgufValue::as_str) + .map(String::from) + }) + .unwrap_or_default(); + drop(f); + + let raw = + format!("{bos}<|user|>\nWhat is 2 + 2? Answer in one short sentence.\n<|assistant|>\n"); + let prompt = tok + .encode(raw.as_str(), false) + .expect("encodes") + .get_ids() + .to_vec(); + assert!(prompt.len() >= 8, "rendered prompt suspiciously short"); + + let device = burn::tensor::Device::::default(); + let start = std::time::Instant::now(); + let loaded = olmoe::load_from_gguf::(&path, &device).expect("checked load"); + eprintln!( + "[real_olmoe] loaded {} layers x {} experts in {:.1}s", + loaded.config.num_hidden_layers, + loaded.config.num_experts, + start.elapsed().as_secs_f64() + ); + + // Liveness first: finite, vocab-wide, non-degenerate logits. + let smoke = loaded + .sanity_check(&prompt, loaded.config.vocab_size, &device) + .expect("sanity smoke"); + eprintln!("[real_olmoe] sanity: {smoke:?}"); + + let start = std::time::Instant::now(); + let ids = loaded + .greedy_generate(&prompt, 24, &device) + .expect("greedy decode"); + let secs = start.elapsed().as_secs_f64(); + assert!(!ids.is_empty(), "decode produced no tokens before EOS"); + let text = tok.decode(&ids, true).expect("decode"); + eprintln!( + "[real_olmoe] {} tokens in {secs:.1}s ({:.2} s/token): {text:?}", + ids.len(), + secs / ids.len() as f64 + ); + assert!( + text.contains('4') || text.to_lowercase().contains("four"), + "expected the answer to mention 4, got: {text:?}" + ); +}