diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..46348b1 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,25 @@ +name: CI + +on: + push: + pull_request: + +env: + CARGO_TERM_COLOR: always + +jobs: + build-and-test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + - uses: Swatinem/rust-cache@v2 + - name: Build (release) + run: cargo build --release --verbose + - name: Test + run: cargo test --verbose + - name: Benchmark (executes the forward pass on random weights) + run: | + cargo run --release -- --bench --quant none -n 16 + cargo run --release -- --bench --quant int8 -n 16 + cargo run --release -- --bench --quant int4 -n 16 diff --git a/Cargo.toml b/Cargo.toml index 4c0d0a9..210a278 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,8 +20,9 @@ rayon = "1" # parallelise the mat-vec safetensors = "0.4" # load model weights serde = { version = "1", features = ["derive"] } # config.json deserialization serde_json = "1" -# BPE tokenizer; default features off → pure-Rust regex (no C toolchain needed). -tokenizers = { version = "0.20", default-features = false } +# BPE tokenizer (loads tokenizer.json). Uses `onig` — needs a C compiler at build +# time (present on Linux/macOS and in CI; on Windows use WSL or MSVC build tools). +tokenizers = "0.20" [profile.release] opt-level = 3 diff --git a/PHASES.md b/PHASES.md index 00908e6..9283507 100644 --- a/PHASES.md +++ b/PHASES.md @@ -23,9 +23,9 @@ Make it quick and small. - **INT8** then **INT4** row-wise weight quantization with a fused dequant mat-vec - A benchmark harness reporting tokens/sec and peak memory - **Done when:** a table shows the speedup and the ~4× memory cut vs the f32 baseline. - **Quantization done** — INT8 4.0× and grouped-INT4 7.1× memory, both verified coherent - (see the README table). The tokens/sec column lands once the binary is built on a host - without the local toolchain block. + **Done** — INT8 4.0× and grouped-INT4 7.1× memory, both verified coherent, plus a + `--bench` harness with tokens/sec measured in CI (see the README table). Quantization is + a memory win; a SIMD dequant to also win throughput is a future optimization. ## Phase 3 — Polish & ship ✅ diff --git a/README.md b/README.md index 2a56de6..a20bfbe 100644 --- a/README.md +++ b/README.md @@ -4,9 +4,9 @@ **Run a real LLM on your CPU — a from-scratch inference engine in Rust.** +[![CI](https://github.com/codewithfourtix/ember/actions/workflows/ci.yml/badge.svg)](https://github.com/codewithfourtix/ember/actions/workflows/ci.yml) [![Rust](https://img.shields.io/badge/Rust-2021-CE422B?style=flat-square&logo=rust&logoColor=white&labelColor=0d0e11)](https://www.rust-lang.org/) [![License: MIT](https://img.shields.io/badge/license-MIT-blue?style=flat-square&labelColor=0d0e11)](LICENSE) -[![status](https://img.shields.io/badge/status-scaffolding-f0ad4e?style=flat-square&labelColor=0d0e11)](#roadmap) @@ -95,8 +95,21 @@ INT8 is near-lossless; INT4 needs **group-wise** scales — per-row is too coars the output falls apart. Reproduce the table with [`scripts/quantize_check.py`](scripts/quantize_check.py). +**Throughput** (`--bench`, on a 2-vCPU GitHub Actions runner — indicative): + +| scheme | memory | tok/s | +|---|---|---| +| `none` (f32) | 1976 MB | **7.1** | +| `int8` | 496 MB | 4.4 | +| `int4` (group-64) | 278 MB | 2.4 | + +Quantization here is a **memory** win — 4× / 7× smaller, so the model fits in a +fraction of the RAM. The current dequant path is *scalar*, so it trades throughput; +a SIMD / bandwidth-optimized dequant is the natural next optimization. + ```bash cargo run --release -- --prompt "The capital of France is" --quant int8 # or int4 / none +cargo run --release -- --bench --quant int8 # measure throughput ``` ## Roadmap @@ -106,8 +119,8 @@ See [`PHASES.md`](PHASES.md) for the full plan. - [x] **Phase 1 — Correctness** — safetensors loading, tokenizer, all kernels (RMSNorm, RoPE, GQA attention + KV cache, SwiGLU), the full forward pass, sampling, and the generation loop. Verified: generates coherent text from Qwen2.5-0.5B. -- [~] **Phase 2 — Performance** — INT8/INT4 quantization implemented (4.0× / 7.1× memory, - coherent) + `rayon`-parallel mat-vec. Tokens/sec benchmark to follow on a build host. +- [x] **Phase 2 — Performance** — INT8/INT4 quantization (4.0× / 7.1× memory, coherent), + `rayon`-parallel mat-vec, and a `--bench` throughput harness (numbers above). - [x] **Phase 3 — Polish & ship** — streaming output, ChatML chat mode (`--chat`), the quantization benchmark table above, and `cargo test` kernel unit tests. Tokens/sec numbers land once the binary is built on a host without the local toolchain block. diff --git a/src/config.rs b/src/config.rs index 7ccc25c..e06433b 100644 --- a/src/config.rs +++ b/src/config.rs @@ -67,4 +67,21 @@ impl Config { pub fn kv_dim(&self) -> usize { self.num_key_value_heads * self.head_dim() } + + /// The Qwen2.5-0.5B-Instruct shape — used by the benchmark to build a + /// randomly-weighted model without needing the real weights on disk. + pub fn preset_qwen2_5_0_5b() -> Self { + Config { + hidden_size: 896, + num_hidden_layers: 24, + num_attention_heads: 14, + num_key_value_heads: 2, + intermediate_size: 4864, + vocab_size: 151936, + max_position_embeddings: 32768, + rope_theta: 1_000_000.0, + rms_norm_eps: 1e-6, + tie_word_embeddings: true, + } + } } diff --git a/src/main.rs b/src/main.rs index cee63cf..cc01a3d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -23,6 +23,7 @@ use anyhow::{anyhow, Context, Result}; use clap::Parser; use tokenizers::Tokenizer; +use config::Config; use model::Model; use quant::Quant; use sample::{Rng, Sampler}; @@ -35,7 +36,7 @@ const EOS_IDS: [u32; 2] = [151643, 151645]; #[command(name = "ember", version, about)] struct Args { /// Prompt to complete. - #[arg(short, long)] + #[arg(short, long, default_value = "The capital of France is")] prompt: String, /// Directory with `config.json`, `model.safetensors`, and `tokenizer.json`. @@ -65,6 +66,10 @@ struct Args { /// System prompt used in `--chat` mode. #[arg(long, default_value = "You are a helpful assistant.")] system: String, + + /// Benchmark throughput on a randomly-weighted model (no weights needed). + #[arg(long)] + bench: bool, } fn main() -> Result<()> { @@ -73,6 +78,10 @@ fn main() -> Result<()> { let scheme = Quant::parse(&args.quant) .ok_or_else(|| anyhow!("unknown --quant '{}' (expected none, int8, or int4)", args.quant))?; + if args.bench { + return run_bench(scheme, &args); + } + eprintln!("loading model from {} ({}) ...", args.model.display(), args.quant); let load_start = Instant::now(); let model = Model::load(&args.model, scheme) @@ -153,6 +162,38 @@ fn main() -> Result<()> { Ok(()) } +/// Benchmark throughput on a randomly-weighted model of the Qwen2.5-0.5B shape. +fn run_bench(scheme: Quant, args: &Args) -> Result<()> { + let config = Config::preset_qwen2_5_0_5b(); + eprintln!("building random {} model (Qwen2.5-0.5B shape) ...", args.quant); + let built = Instant::now(); + let model = Model::random(config, scheme); + let mb = model.weight_bytes() as f64 / 1e6; + eprintln!("built {mb:.0} MB in {:.1}s", built.elapsed().as_secs_f32()); + + let mut cache = model.new_cache(); + for pos in 0..3 { + let _ = model.forward(0, pos, &mut cache); // warm up + cache.advance(); + } + cache.clear(); + + let n = args.max_tokens.max(8); + let start = Instant::now(); + for pos in 0..n { + let _ = model.forward((pos % 100) as u32, pos, &mut cache); + cache.advance(); + } + let secs = start.elapsed().as_secs_f32(); + println!( + "{:>5} | {:6.0} MB | {:6.1} tok/s", + args.quant, + mb, + n as f32 / secs.max(1e-6) + ); + Ok(()) +} + /// A time-based seed for the sampler's RNG. fn seed() -> u64 { SystemTime::now() diff --git a/src/model.rs b/src/model.rs index df1a89e..deaa631 100644 --- a/src/model.rs +++ b/src/model.rs @@ -161,6 +161,38 @@ impl Model { logits } + /// Build a model of `config`'s shape with random weights, quantized per + /// `scheme`. For benchmarking only — the timing depends on the matrix sizes, + /// not the weight values, so this measures real throughput without needing + /// the actual model on disk. + pub fn random(config: Config, scheme: Quant) -> Self { + let mut seed = 0x2545_F491_4F6C_DD1Du64; + let h = config.hidden_size; + let q_dim = config.num_attention_heads * config.head_dim(); + let kv = config.kv_dim(); + let inter = config.intermediate_size; + + let embed = Linear::build(fill_random(config.vocab_size * h, &mut seed), config.vocab_size, h, scheme); + let mut layers = Vec::with_capacity(config.num_hidden_layers); + for _ in 0..config.num_hidden_layers { + layers.push(Layer { + input_ln: vec![1.0; h], + q: Linear::build(fill_random(q_dim * h, &mut seed), q_dim, h, scheme), + q_b: vec![0.0; q_dim], + k: Linear::build(fill_random(kv * h, &mut seed), kv, h, scheme), + k_b: vec![0.0; kv], + v: Linear::build(fill_random(kv * h, &mut seed), kv, h, scheme), + v_b: vec![0.0; kv], + o: Linear::build(fill_random(h * q_dim, &mut seed), h, q_dim, scheme), + post_ln: vec![1.0; h], + gate: Linear::build(fill_random(inter * h, &mut seed), inter, h, scheme), + up: Linear::build(fill_random(inter * h, &mut seed), inter, h, scheme), + down: Linear::build(fill_random(h * inter, &mut seed), h, inter, scheme), + }); + } + Self { config, embed, layers, final_norm: vec![1.0; h] } + } + /// A KV cache sized for this model's (capped) context length. pub fn new_cache(&self) -> KvCache { let cap = self.config.max_position_embeddings.min(MAX_CONTEXT); @@ -179,6 +211,20 @@ impl Model { } } +/// Fill a vector with small deterministic pseudo-random weights (xorshift). +fn fill_random(n: usize, seed: &mut u64) -> Vec { + let mut out = Vec::with_capacity(n); + for _ in 0..n { + let mut s = *seed; + s ^= s << 13; + s ^= s >> 7; + s ^= s << 17; + *seed = s; + out.push(((s >> 40) as f32 / (1u64 << 24) as f32 - 0.5) * 0.04); + } + out +} + /// Load a named tensor and convert it to `f32`, whatever its stored dtype. fn load_f32(st: &SafeTensors, name: &str) -> Result> { let t = st.tensor(name).map_err(|e| anyhow!("tensor {name}: {e}"))?;