From b51d1cbe02367a94c60fae3d06d3824e7b58d986 Mon Sep 17 00:00:00 2001 From: sigma-andex <77549848+sigma-andex@users.noreply.github.com> Date: Thu, 21 Dec 2023 13:12:34 +0000 Subject: [PATCH] First integration attempt --- examples/remix/app/routes/_index.tsx | 3 +- packages/ask-react/src/client.ts | 4 +- packages/llama/Cargo.lock | 297 ++------------ packages/llama/Cargo.toml | 4 +- packages/llama/build-lib.sh | 1 + packages/llama/package.json | 4 +- packages/llama/src/bin/m2.rs | 171 ++++++++ packages/llama/src/bin/worker.rs | 5 - packages/llama/src/lib.rs | 2 +- packages/llama/src/main.rs_ | 548 ++++++++++++++++++++++++++ packages/llama/src/quantized_llama.rs | 509 ++++++++++++++++++++++++ packages/llama/src/worker.rs | 111 ------ 12 files changed, 1279 insertions(+), 380 deletions(-) mode change 100644 => 100755 packages/llama/build-lib.sh create mode 100644 packages/llama/src/bin/m2.rs delete mode 100644 packages/llama/src/bin/worker.rs create mode 100644 packages/llama/src/main.rs_ create mode 100644 packages/llama/src/quantized_llama.rs diff --git a/examples/remix/app/routes/_index.tsx b/examples/remix/app/routes/_index.tsx index db59292..4623fb0 100644 --- a/examples/remix/app/routes/_index.tsx +++ b/examples/remix/app/routes/_index.tsx @@ -15,8 +15,7 @@ export default function Index() { }); return (
-

@text.yoga/ask

-

remix example

+

@text.yoga/ask

{ }; }, []); - const weightsURL = `http://localhost:5173/model.bin`; + const weightsURL = `http://localhost:5173/model2.gguf`; const modelID = "stories15M"; - const tokenizerURL = `http://localhost:5173/tokenizer.json`; + const tokenizerURL = `http://localhost:5173/tokenizer2.json`; const [questionValue, setQuestionValue] = useState(""); diff --git a/packages/llama/Cargo.lock b/packages/llama/Cargo.lock index e750954..a9ff590 100644 --- a/packages/llama/Cargo.lock +++ b/packages/llama/Cargo.lock @@ -98,12 +98,6 @@ version = "1.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" -[[package]] -name = "boolinator" -version = "2.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfa8873f51c92e232f9bac4065cddef41b714152812bfc5f7672ba16d6ef8cd9" - [[package]] name = "bumpalo" version = "3.14.0" @@ -382,20 +376,6 @@ dependencies = [ "percent-encoding", ] -[[package]] -name = "futures" -version = "0.3.29" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da0290714b38af9b4a7b094b8a37086d1b4e61f2df9122c3cad2577669145335" -dependencies = [ - "futures-channel", - "futures-core", - "futures-io", - "futures-sink", - "futures-task", - "futures-util", -] - [[package]] name = "futures-channel" version = "0.3.29" @@ -403,7 +383,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ff4dd66668b557604244583e3e1e1eada8c5c2e96a6d0d6653ede395b78bbacb" dependencies = [ "futures-core", - "futures-sink", ] [[package]] @@ -412,53 +391,12 @@ version = "0.3.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eb1d22c66e66d9d72e1758f0bd7d4fd0bee04cad842ee34587d68c07e45d088c" -[[package]] -name = "futures-io" -version = "0.3.29" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8bf34a163b5c4c52d0478a4d757da8fb65cabef42ba90515efee0f6f9fa45aaa" - -[[package]] -name = "futures-macro" -version = "0.3.29" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53b153fd91e4b0147f4aced87be237c98248656bb01050b96bf3ee89220a8ddb" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.41", -] - [[package]] name = "futures-sink" version = "0.3.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e36d3378ee38c2a36ad710c5d30c2911d752cb941c00c72dbabfb786a7970817" -[[package]] -name = "futures-task" -version = "0.3.29" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "efd193069b0ddadc69c46389b740bbccdd97203899b48d09c5f7969591d6bae2" - -[[package]] -name = "futures-util" -version = "0.3.29" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a19526d624e703a3179b3d322efec918b6246ea0fa51d41124525f00f1cc8104" -dependencies = [ - "futures-channel", - "futures-core", - "futures-io", - "futures-macro", - "futures-sink", - "futures-task", - "memchr", - "pin-project-lite", - "pin-utils", - "slab", -] - [[package]] name = "gemm" version = "0.16.15" @@ -611,7 +549,7 @@ dependencies = [ "gloo-storage", "gloo-timers", "gloo-utils", - "gloo-worker 0.2.1", + "gloo-worker", ] [[package]] @@ -744,23 +682,6 @@ dependencies = [ "web-sys", ] -[[package]] -name = "gloo-worker" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09110b5555bcafe508cee0fb94308af9aac7a85f980d3c88b270d117c6c6911d" -dependencies = [ - "anymap2", - "bincode", - "gloo-console", - "gloo-utils", - "js-sys", - "serde", - "slab", - "wasm-bindgen", - "web-sys", -] - [[package]] name = "gloo-worker" version = "0.2.1" @@ -792,12 +713,6 @@ dependencies = [ "rand_distr", ] -[[package]] -name = "hashbrown" -version = "0.12.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" - [[package]] name = "hermit-abi" version = "0.3.3" @@ -821,25 +736,6 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" -[[package]] -name = "implicit-clone" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfd6201e7c30ccb24773cac7efa6fec1e06189d414b7439ce756a481c8bfbf53" -dependencies = [ - "indexmap", -] - -[[package]] -name = "indexmap" -version = "1.9.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99" -dependencies = [ - "autocfg", - "hashbrown", -] - [[package]] name = "itertools" version = "0.11.0" @@ -1058,63 +954,12 @@ version = "0.2.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8afb450f006bf6385ca15ef45d71d2288452bc3683ce2e2cacc0d18e4be60b58" -[[package]] -name = "pin-utils" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" - -[[package]] -name = "pinned" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a829027bd95e54cfe13e3e258a1ae7b645960553fb82b75ff852c29688ee595b" -dependencies = [ - "futures", - "rustversion", - "thiserror", -] - [[package]] name = "ppv-lite86" version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" -[[package]] -name = "prettyplease" -version = "0.1.25" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c8646e95016a7a6c4adea95bafa8a16baab64b583356217f2c85db4a39d9a86" -dependencies = [ - "proc-macro2", - "syn 1.0.109", -] - -[[package]] -name = "proc-macro-error" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c" -dependencies = [ - "proc-macro-error-attr", - "proc-macro2", - "quote", - "syn 1.0.109", - "version_check", -] - -[[package]] -name = "proc-macro-error-attr" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869" -dependencies = [ - "proc-macro2", - "quote", - "version_check", -] - [[package]] name = "proc-macro2" version = "1.0.70" @@ -1124,23 +969,6 @@ dependencies = [ "unicode-ident", ] -[[package]] -name = "prokio" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "03b55e106e5791fa5a13abd13c85d6127312e8e09098059ca2bc9b03ca4cf488" -dependencies = [ - "futures", - "gloo", - "num_cpus", - "once_cell", - "pin-project", - "pinned", - "tokio", - "tokio-stream", - "wasm-bindgen-futures", -] - [[package]] name = "pulp" version = "0.18.6" @@ -1295,12 +1123,6 @@ version = "0.1.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d626bb9dae77e28219937af045c257c28bfd3f69333c512553507f5f9798cb76" -[[package]] -name = "rustversion" -version = "1.0.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ffc183a10b4478d04cbbbfc96d0873219d962dd5accaff2ffbd4ceb7df837f4" - [[package]] name = "ryu" version = "1.0.16" @@ -1387,12 +1209,12 @@ dependencies = [ ] [[package]] -name = "slab" -version = "0.4.9" +name = "sharded-slab" +version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f92a496fb766b417c996b9c5e57daf2f7ad3b0bebe1ccfca4856390e3d3bb67" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" dependencies = [ - "autocfg", + "lazy_static", ] [[package]] @@ -1478,12 +1300,12 @@ dependencies = [ "serde", "serde_json", "tokenizers", + "tracing-wasm", "wasm-bindgen", "wasm-bindgen-futures", "wasm-logger", "web-sys", - "yew", - "yew-agent", + "web-time", ] [[package]] @@ -1506,6 +1328,16 @@ dependencies = [ "syn 2.0.41", ] +[[package]] +name = "thread_local" +version = "1.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fdd6f064ccff2d6567adcb3873ca630700f00b5ad3f060c25b5dcfd9a4ce152" +dependencies = [ + "cfg-if", + "once_cell", +] + [[package]] name = "tokenizers" version = "0.15.0" @@ -1537,27 +1369,6 @@ dependencies = [ "unicode_categories", ] -[[package]] -name = "tokio" -version = "1.35.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "841d45b238a16291a4e1584e61820b8ae57d696cc5015c459c229ccc6990cc1c" -dependencies = [ - "backtrace", - "pin-project-lite", -] - -[[package]] -name = "tokio-stream" -version = "0.1.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "397c988d37662c7dda6d2208364a706264bf3d6138b11d436cbac0ad38832842" -dependencies = [ - "futures-core", - "pin-project-lite", - "tokio", -] - [[package]] name = "tracing" version = "0.1.40" @@ -1589,6 +1400,28 @@ dependencies = [ "once_cell", ] +[[package]] +name = "tracing-subscriber" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b" +dependencies = [ + "sharded-slab", + "thread_local", + "tracing-core", +] + +[[package]] +name = "tracing-wasm" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4575c663a174420fa2d78f4108ff68f65bf2fbb7dd89f33749b6e826b3626e07" +dependencies = [ + "tracing", + "tracing-subscriber", + "wasm-bindgen", +] + [[package]] name = "unicode-ident" version = "1.0.12" @@ -1622,12 +1455,6 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" -[[package]] -name = "version_check" -version = "0.9.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f" - [[package]] name = "wasi" version = "0.11.0+wasi-snapshot-preview1" @@ -1731,53 +1558,13 @@ dependencies = [ ] [[package]] -name = "yew" -version = "0.20.0" +name = "web-time" +version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5dbecfe44343b70cc2932c3eb445425969ae21754a8ab3a0966981c1cf7af1cc" +checksum = "57099a701fb3a8043f993e8228dc24229c7b942e2b009a1b962e54489ba1d3bf" dependencies = [ - "console_error_panic_hook", - "futures", - "gloo", - "implicit-clone", - "indexmap", "js-sys", - "prokio", - "rustversion", - "serde", - "slab", - "thiserror", - "tokio", - "tracing", "wasm-bindgen", - "wasm-bindgen-futures", - "web-sys", - "yew-macro", -] - -[[package]] -name = "yew-agent" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b06f7c5ed97fff22816bb00d3d82ebc0fc1119d7bbb9e07e62c0d2853f51920a" -dependencies = [ - "gloo-worker 0.1.2", - "yew", -] - -[[package]] -name = "yew-macro" -version = "0.20.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b64c253c1d401f1ea868ca9988db63958cfa15a69f739101f338d6f05eea8301" -dependencies = [ - "boolinator", - "once_cell", - "prettyplease", - "proc-macro-error", - "proc-macro2", - "quote", - "syn 1.0.109", ] [[package]] diff --git a/packages/llama/Cargo.toml b/packages/llama/Cargo.toml index 0c8ac84..99b1fbe 100644 --- a/packages/llama/Cargo.toml +++ b/packages/llama/Cargo.toml @@ -32,11 +32,11 @@ rand = "0.8.5" serde = {version = "1.0.193", features = ["derive"]} serde_json = "1.0.108" tokenizers = { version = "0.15.0", features = ["unstable_wasm"], default-features = false } +tracing-wasm = "0.2.1" wasm-bindgen = "0.2.87" wasm-bindgen-futures = "0.4.37" wasm-logger = "0.2" -yew = {version = "0.20.0", features = ["csr"]} -yew-agent = "0.2.0" +web-time = { version = "0.2.3" } [dependencies.web-sys] features = [ diff --git a/packages/llama/build-lib.sh b/packages/llama/build-lib.sh old mode 100644 new mode 100755 index e3ce367..c478956 --- a/packages/llama/build-lib.sh +++ b/packages/llama/build-lib.sh @@ -1,2 +1,3 @@ cargo build --target wasm32-unknown-unknown --release wasm-bindgen ./target/wasm32-unknown-unknown/release/m.wasm --out-dir dist --target web +wasm-bindgen ./target/wasm32-unknown-unknown/release/m2.wasm --out-dir dist --target web diff --git a/packages/llama/package.json b/packages/llama/package.json index ad82f51..759e7f7 100644 --- a/packages/llama/package.json +++ b/packages/llama/package.json @@ -7,8 +7,8 @@ ], "exports": { ".": { - "import": "./dist/m.js", - "types": "./dist/m.d.ts" + "import": "./dist/m2.js", + "types": "./dist/m2.d.ts" } }, "scripts": { diff --git a/packages/llama/src/bin/m2.rs b/packages/llama/src/bin/m2.rs new file mode 100644 index 0000000..7ed6ec8 --- /dev/null +++ b/packages/llama/src/bin/m2.rs @@ -0,0 +1,171 @@ +use std::io::Cursor; + +use candle::{quantized::gguf_file, Device, Tensor}; +use candle_transformers::{ + generation::LogitsProcessor, models::quantized_llama::ModelWeights as M, +}; + +use text_yoga_ai::quantized_llama::ModelWeights; +use tokenizers::Tokenizer; +// use text_yoga_ai::worker::{Model as M, ModelData}; +use wasm_bindgen::prelude::*; +use web_time as time; + +#[wasm_bindgen] +pub struct Model { + inner: ModelWeights, + logits_processor: LogitsProcessor, + tokens: Vec, + repeat_penalty: f32, + tokenizer: Tokenizer, +} + +impl Model { + fn process(&mut self, tokens: &[u32]) -> candle::Result { + const REPEAT_LAST_N: usize = 64; + let dev = Device::Cpu; + let input = Tensor::new(tokens, &dev)?.unsqueeze(0)?; + let logits = self.inner.forward(&input, tokens.len())?; + let logits = logits.squeeze(0)?; + let logits = if self.repeat_penalty == 1. || tokens.is_empty() { + logits + } else { + let start_at = self.tokens.len().saturating_sub(REPEAT_LAST_N); + candle_transformers::utils::apply_repeat_penalty( + &logits, + self.repeat_penalty, + &self.tokens[start_at..], + )? + }; + + let next_token = self.logits_processor.sample(&logits)?; + self.tokens.push(next_token); + let text = match self.tokenizer.id_to_token(next_token) { + Some(text) => text.replace('▁', " ").replace("<0x0A>", "\n"), + None => "".to_string(), + }; + Ok(text) + } +} + +#[wasm_bindgen] +impl Model { + #[wasm_bindgen(constructor)] + pub fn new(weights: Vec, tokenizer: Vec) -> Result { + // let model = M::load(ModelData { + // tokenizer, + // model: weights, + // }); + // let logits_processor = LogitsProcessor::new(299792458, None, None); + // match model { + // Ok(inner) => Ok(Self { + // inner, + // logits_processor, + // tokens: vec![], + // repeat_penalty: 1., + // }), + // Err(e) => Err(JsError::new(&e.to_string())), + // } + let seed = 299792458; + let temperature: Option = Some(0.0); + let top_p: Option = None; + let repeat_penalty: f32 = 1.; + let start = time::Instant::now(); + let mut cursor = Cursor::new(&weights); + let mut cursor2 = Cursor::new(&weights); + let mut model: ModelWeights = { + let model = gguf_file::Content::read(&mut cursor)?; + let mut total_size_in_bytes = 0; + for (_, tensor) in model.tensor_infos.iter() { + let elem_count = tensor.shape.elem_count(); + total_size_in_bytes += + elem_count * tensor.ggml_dtype.type_size() / tensor.ggml_dtype.blck_size(); + } + println!( + "loaded {:?} tensors ({}) in {:.2}s", + model.tensor_infos.len(), + &format_size(total_size_in_bytes), + start.elapsed().as_secs_f32(), + ); + ModelWeights::from_gguf(model, &mut cursor2)? + }; + println!("model built"); + + let tokenizer = Tokenizer::from_bytes(&tokenizer) + .map_err(|msg| JsError::new("Failed to load tokenizer."))?; + let mut logits_processor = LogitsProcessor::new(seed, temperature, top_p); + Ok(Self { + inner: model, + logits_processor, + repeat_penalty, + tokens: vec![], + tokenizer, + }) + } + + #[wasm_bindgen] + pub fn get_seq_len(&mut self) -> usize { + // self.inner.config.seq_len + 100 + } + + #[wasm_bindgen] + pub fn init_with_prompt( + &mut self, + prompt: String, + temp: f64, + top_p: f64, + repeat_penalty: f32, + seed: u64, + ) -> Result { + // First reset the cache. + // { + // let mut cache = self.inner.cache.kvs.lock().unwrap(); + // for elem in cache.iter_mut() { + // *elem = None + // } + // } + let temp = if temp <= 0. { None } else { Some(temp) }; + let top_p = if top_p <= 0. || top_p >= 1. { + None + } else { + Some(top_p) + }; + self.logits_processor = LogitsProcessor::new(seed, temp, top_p); + self.repeat_penalty = repeat_penalty; + self.tokens.clear(); + let tokens = self + .tokenizer + .encode(prompt, true) + .map_err(|m| JsError::new(&m.to_string()))? + .get_ids() + .to_vec(); + let text = self + .process(&tokens) + .map_err(|m| JsError::new(&m.to_string()))?; + Ok(text) + } + + #[wasm_bindgen] + pub fn next_token(&mut self) -> Result { + let last_token = *self.tokens.last().unwrap(); + let text = self + .process(&[last_token]) + .map_err(|m| JsError::new(&m.to_string()))?; + Ok(text) + } +} + +fn format_size(size_in_bytes: usize) -> String { + if size_in_bytes < 1_000 { + format!("{}B", size_in_bytes) + } else if size_in_bytes < 1_000_000 { + format!("{:.2}KB", size_in_bytes as f64 / 1e3) + } else if size_in_bytes < 1_000_000_000 { + format!("{:.2}MB", size_in_bytes as f64 / 1e6) + } else { + format!("{:.2}GB", size_in_bytes as f64 / 1e9) + } +} + +fn main() {} diff --git a/packages/llama/src/bin/worker.rs b/packages/llama/src/bin/worker.rs deleted file mode 100644 index dda1dc0..0000000 --- a/packages/llama/src/bin/worker.rs +++ /dev/null @@ -1,5 +0,0 @@ -use yew_agent::PublicWorker; -fn main() { - console_error_panic_hook::set_once(); - text_yoga_ai::Worker::register(); -} diff --git a/packages/llama/src/lib.rs b/packages/llama/src/lib.rs index b5fbe3d..311f73c 100644 --- a/packages/llama/src/lib.rs +++ b/packages/llama/src/lib.rs @@ -1,3 +1,3 @@ pub mod model; +pub mod quantized_llama; pub mod worker; -pub use worker::Worker; diff --git a/packages/llama/src/main.rs_ b/packages/llama/src/main.rs_ new file mode 100644 index 0000000..df758b4 --- /dev/null +++ b/packages/llama/src/main.rs_ @@ -0,0 +1,548 @@ +#[cfg(feature = "mkl")] +extern crate intel_mkl_src; + +#[cfg(feature = "accelerate")] +extern crate accelerate_src; + +use clap::{Parser, ValueEnum}; +use std::io::Write; +use tokenizers::Tokenizer; + +use candle::quantized::{ggml_file, gguf_file}; +use candle::{Device, Tensor}; +use candle_transformers::generation::LogitsProcessor; + +use candle_examples::token_output_stream::TokenOutputStream; +use candle_transformers::models::quantized_llama as model; +use model::ModelWeights; + +const DEFAULT_PROMPT: &str = "My favorite theorem is "; + +#[derive(Debug)] +enum Prompt { + Interactive, + Chat, + One(String), +} + +#[derive(Clone, Debug, Copy, PartialEq, Eq, ValueEnum)] +enum Which { + #[value(name = "7b")] + L7b, + #[value(name = "13b")] + L13b, + #[value(name = "70b")] + L70b, + #[value(name = "7b-chat")] + L7bChat, + #[value(name = "13b-chat")] + L13bChat, + #[value(name = "70b-chat")] + L70bChat, + #[value(name = "7b-code")] + L7bCode, + #[value(name = "13b-code")] + L13bCode, + #[value(name = "32b-code")] + L34bCode, + #[value(name = "7b-leo")] + Leo7b, + #[value(name = "13b-leo")] + Leo13b, + #[value(name = "7b-mistral")] + Mistral7b, + #[value(name = "7b-mistral-instruct")] + Mistral7bInstruct, + #[value(name = "7b-zephyr-a")] + Zephyr7bAlpha, + #[value(name = "7b-zephyr-b")] + Zephyr7bBeta, + #[value(name = "7b-open-chat-3.5")] + OpenChat35, + #[value(name = "7b-starling-a")] + Starling7bAlpha, + #[value(name = "mixtral")] + Mixtral, + #[value(name = "mixtral-instruct")] + MixtralInstruct, +} + +impl Which { + fn is_mistral(&self) -> bool { + match self { + Self::L7b + | Self::L13b + | Self::L70b + | Self::L7bChat + | Self::L13bChat + | Self::L70bChat + | Self::L7bCode + | Self::L13bCode + | Self::L34bCode + | Self::Leo7b + | Self::Leo13b => false, + // Zephyr and OpenChat are fine tuned versions of mistral and should be treated in the + // same way. Starling is a fine tuned version of OpenChat. + Self::OpenChat35 + | Self::Starling7bAlpha + | Self::Zephyr7bAlpha + | Self::Zephyr7bBeta + | Self::Mixtral + | Self::MixtralInstruct + | Self::Mistral7b + | Self::Mistral7bInstruct => true, + } + } + + fn is_zephyr(&self) -> bool { + match self { + Self::L7b + | Self::L13b + | Self::L70b + | Self::L7bChat + | Self::L13bChat + | Self::L70bChat + | Self::L7bCode + | Self::L13bCode + | Self::L34bCode + | Self::Leo7b + | Self::Leo13b + | Self::Mixtral + | Self::MixtralInstruct + | Self::Mistral7b + | Self::Mistral7bInstruct + | Self::OpenChat35 + | Self::Starling7bAlpha => false, + Self::Zephyr7bAlpha | Self::Zephyr7bBeta => true, + } + } + + fn is_open_chat(&self) -> bool { + match self { + Self::L7b + | Self::L13b + | Self::L70b + | Self::L7bChat + | Self::L13bChat + | Self::L70bChat + | Self::L7bCode + | Self::L13bCode + | Self::L34bCode + | Self::Leo7b + | Self::Leo13b + | Self::Mixtral + | Self::MixtralInstruct + | Self::Mistral7b + | Self::Mistral7bInstruct + | Self::Zephyr7bAlpha + | Self::Zephyr7bBeta => false, + Self::OpenChat35 | Self::Starling7bAlpha => true, + } + } + + fn tokenizer_repo(&self) -> &'static str { + match self { + Which::L7b + | Which::L13b + | Which::L70b + | Which::L7bChat + | Which::L13bChat + | Which::L70bChat + | Which::L7bCode + | Which::L13bCode + | Which::L34bCode => "hf-internal-testing/llama-tokenizer", + Which::Leo7b => "LeoLM/leo-hessianai-7b", + Which::Leo13b => "LeoLM/leo-hessianai-13b", + Which::Mixtral => "mistralai/Mixtral-8x7B-v0.1", + Which::MixtralInstruct => "mistralai/Mixtral-8x7B-Instruct-v0.1", + Which::Mistral7b + | Which::Mistral7bInstruct + | Which::Zephyr7bAlpha + | Which::Zephyr7bBeta => "mistralai/Mistral-7B-v0.1", + Which::OpenChat35 => "openchat/openchat_3.5", + Which::Starling7bAlpha => "berkeley-nest/Starling-LM-7B-alpha", + } + } +} + +#[derive(Parser, Debug)] +#[command(author, version, about, long_about = None)] +struct Args { + /// GGML file to load, typically a .bin file generated by the quantize command from llama.cpp + #[arg(long)] + model: Option, + + /// The initial prompt, use 'interactive' for entering multiple prompts in an interactive way + /// and 'chat' for an interactive model where history of previous prompts and generated tokens + /// is preserved. + #[arg(long)] + prompt: Option, + + /// The length of the sample to generate (in tokens). + #[arg(short = 'n', long, default_value_t = 1000)] + sample_len: usize, + + /// The tokenizer config in json format. + #[arg(long)] + tokenizer: Option, + + /// The temperature used to generate samples, use 0 for greedy sampling. + #[arg(long, default_value_t = 0.8)] + temperature: f64, + + /// Nucleus sampling probability cutoff. + #[arg(long)] + top_p: Option, + + /// The seed to use when generating random samples. + #[arg(long, default_value_t = 299792458)] + seed: u64, + + /// Enable tracing (generates a trace-timestamp.json file). + #[arg(long)] + tracing: bool, + + /// Display the token for the specified prompt. + #[arg(long)] + verbose_prompt: bool, + + /// Penalty to be applied for repeating tokens, 1. means no penalty. + #[arg(long, default_value_t = 1.1)] + repeat_penalty: f32, + + /// The context size to consider for the repeat penalty. + #[arg(long, default_value_t = 64)] + repeat_last_n: usize, + + /// The model size to use. + #[arg(long, default_value = "7b")] + which: Which, + + /// Group-Query Attention, use 8 for the 70B version of LLaMAv2. + #[arg(long)] + gqa: Option, +} + +impl Args { + fn tokenizer(&self) -> anyhow::Result { + let tokenizer_path = match &self.tokenizer { + Some(config) => std::path::PathBuf::from(config), + None => { + let api = hf_hub::api::sync::Api::new()?; + let repo = self.which.tokenizer_repo(); + let api = api.model(repo.to_string()); + api.get("tokenizer.json")? + } + }; + Tokenizer::from_file(tokenizer_path).map_err(anyhow::Error::msg) + } + + fn model(&self) -> anyhow::Result { + let model_path = match &self.model { + Some(config) => std::path::PathBuf::from(config), + None => { + let (repo, filename) = match self.which { + Which::L7b => ("TheBloke/Llama-2-7B-GGML", "llama-2-7b.ggmlv3.q4_0.bin"), + Which::L13b => ("TheBloke/Llama-2-13B-GGML", "llama-2-13b.ggmlv3.q4_0.bin"), + Which::L70b => ("TheBloke/Llama-2-70B-GGML", "llama-2-70b.ggmlv3.q4_0.bin"), + Which::L7bChat => ( + "TheBloke/Llama-2-7B-Chat-GGML", + "llama-2-7b-chat.ggmlv3.q4_0.bin", + ), + Which::L13bChat => ( + "TheBloke/Llama-2-13B-Chat-GGML", + "llama-2-13b-chat.ggmlv3.q4_0.bin", + ), + Which::L70bChat => ( + "TheBloke/Llama-2-70B-Chat-GGML", + "llama-2-70b-chat.ggmlv3.q4_0.bin", + ), + Which::L7bCode => ("TheBloke/CodeLlama-7B-GGUF", "codellama-7b.Q8_0.gguf"), + Which::L13bCode => ("TheBloke/CodeLlama-13B-GGUF", "codellama-13b.Q8_0.gguf"), + Which::L34bCode => ("TheBloke/CodeLlama-34B-GGUF", "codellama-34b.Q8_0.gguf"), + Which::Leo7b => ( + "TheBloke/leo-hessianai-7B-GGUF", + "leo-hessianai-7b.Q4_K_M.gguf", + ), + Which::Leo13b => ( + "TheBloke/leo-hessianai-13B-GGUF", + "leo-hessianai-13b.Q4_K_M.gguf", + ), + Which::Mixtral => ( + "TheBloke/Mixtral-8x7B-v0.1-GGUF", + "mixtral-8x7b-v0.1.Q4_K_M.gguf", + ), + Which::MixtralInstruct => ( + "TheBloke/Mixtral-8x7B-Instruct-v0.1-GGUF", + "mixtral-8x7b-instruct-v0.1.Q4_K_M.gguf", + ), + Which::Mistral7b => ( + "TheBloke/Mistral-7B-v0.1-GGUF", + "mistral-7b-v0.1.Q4_K_S.gguf", + ), + Which::Mistral7bInstruct => ( + "TheBloke/Mistral-7B-Instruct-v0.1-GGUF", + "mistral-7b-instruct-v0.1.Q4_K_S.gguf", + ), + Which::Zephyr7bAlpha => ( + "TheBloke/zephyr-7B-alpha-GGUF", + "zephyr-7b-alpha.Q4_K_M.gguf", + ), + Which::Zephyr7bBeta => { + ("TheBloke/zephyr-7B-beta-GGUF", "zephyr-7b-beta.Q4_K_M.gguf") + } + Which::OpenChat35 => ("TheBloke/openchat_3.5-GGUF", "openchat_3.5.Q4_K_M.gguf"), + Which::Starling7bAlpha => ( + "TheBloke/Starling-LM-7B-alpha-GGUF", + "starling-lm-7b-alpha.Q4_K_M.gguf", + ), + }; + let api = hf_hub::api::sync::Api::new()?; + let api = api.model(repo.to_string()); + api.get(filename)? + } + }; + Ok(model_path) + } +} + +fn format_size(size_in_bytes: usize) -> String { + if size_in_bytes < 1_000 { + format!("{}B", size_in_bytes) + } else if size_in_bytes < 1_000_000 { + format!("{:.2}KB", size_in_bytes as f64 / 1e3) + } else if size_in_bytes < 1_000_000_000 { + format!("{:.2}MB", size_in_bytes as f64 / 1e6) + } else { + format!("{:.2}GB", size_in_bytes as f64 / 1e9) + } +} + +fn main() -> anyhow::Result<()> { + use tracing_chrome::ChromeLayerBuilder; + use tracing_subscriber::prelude::*; + + let args = Args::parse(); + let temperature = if args.temperature == 0. { + None + } else { + Some(args.temperature) + }; + let _guard = if args.tracing { + let (chrome_layer, guard) = ChromeLayerBuilder::new().build(); + tracing_subscriber::registry().with(chrome_layer).init(); + Some(guard) + } else { + None + }; + + println!( + "avx: {}, neon: {}, simd128: {}, f16c: {}", + candle::utils::with_avx(), + candle::utils::with_neon(), + candle::utils::with_simd128(), + candle::utils::with_f16c() + ); + println!( + "temp: {:.2} repeat-penalty: {:.2} repeat-last-n: {}", + args.temperature, args.repeat_penalty, args.repeat_last_n + ); + + let model_path = args.model()?; + let mut file = std::fs::File::open(&model_path)?; + let start = std::time::Instant::now(); + + let mut model = match model_path.extension().and_then(|v| v.to_str()) { + Some("gguf") => { + let model = gguf_file::Content::read(&mut file)?; + let mut total_size_in_bytes = 0; + for (_, tensor) in model.tensor_infos.iter() { + let elem_count = tensor.shape.elem_count(); + total_size_in_bytes += + elem_count * tensor.ggml_dtype.type_size() / tensor.ggml_dtype.blck_size(); + } + println!( + "loaded {:?} tensors ({}) in {:.2}s", + model.tensor_infos.len(), + &format_size(total_size_in_bytes), + start.elapsed().as_secs_f32(), + ); + ModelWeights::from_gguf(model, &mut file)? + } + Some("ggml" | "bin") | Some(_) | None => { + let model = ggml_file::Content::read(&mut file)?; + let mut total_size_in_bytes = 0; + for (_, tensor) in model.tensors.iter() { + let elem_count = tensor.shape().elem_count(); + total_size_in_bytes += + elem_count * tensor.dtype().type_size() / tensor.dtype().blck_size(); + } + println!( + "loaded {:?} tensors ({}) in {:.2}s", + model.tensors.len(), + &format_size(total_size_in_bytes), + start.elapsed().as_secs_f32(), + ); + println!("params: {:?}", model.hparams); + let default_gqa = match args.which { + Which::L7b + | Which::L13b + | Which::L7bChat + | Which::L13bChat + | Which::L7bCode + | Which::L13bCode + | Which::L34bCode + | Which::Leo7b + | Which::Leo13b => 1, + Which::Mixtral + | Which::MixtralInstruct + | Which::Mistral7b + | Which::Mistral7bInstruct + | Which::Zephyr7bAlpha + | Which::Zephyr7bBeta + | Which::L70b + | Which::L70bChat + | Which::OpenChat35 + | Which::Starling7bAlpha => 8, + }; + ModelWeights::from_ggml(model, args.gqa.unwrap_or(default_gqa))? + } + }; + println!("model built"); + + let tokenizer = args.tokenizer()?; + let mut tos = TokenOutputStream::new(tokenizer); + let prompt = match args.prompt.as_deref() { + Some("chat") => Prompt::Chat, + Some("interactive") => Prompt::Interactive, + Some(s) => Prompt::One(s.to_string()), + None => Prompt::One(DEFAULT_PROMPT.to_string()), + }; + + let mut pre_prompt_tokens = vec![]; + for prompt_index in 0.. { + let prompt_str = match &prompt { + Prompt::One(prompt) => prompt.clone(), + Prompt::Interactive | Prompt::Chat => { + let is_interactive = matches!(prompt, Prompt::Interactive); + print!("> "); + std::io::stdout().flush()?; + let mut prompt = String::new(); + std::io::stdin().read_line(&mut prompt)?; + if prompt.ends_with('\n') { + prompt.pop(); + if prompt.ends_with('\r') { + prompt.pop(); + } + } + if args.which.is_open_chat() { + format!("GPT4 Correct User: {prompt}<|end_of_turn|>GPT4 Correct Assistant:") + } else if args.which.is_zephyr() { + if prompt_index == 0 || is_interactive { + format!("<|system|>\n\n<|user|>\n{prompt}\n<|assistant|>",) + } else { + format!("<|user|>\n{prompt}\n<|assistant|>") + } + } else if args.which.is_mistral() { + format!("[INST] {prompt} [/INST]") + } else { + prompt + } + } + }; + print!("{}", &prompt_str); + let tokens = tos + .tokenizer() + .encode(prompt_str, true) + .map_err(anyhow::Error::msg)?; + if args.verbose_prompt { + for (token, id) in tokens.get_tokens().iter().zip(tokens.get_ids().iter()) { + let token = token.replace('▁', " ").replace("<0x0A>", "\n"); + println!("{id:7} -> '{token}'"); + } + } + + let prompt_tokens = [&pre_prompt_tokens, tokens.get_ids()].concat(); + let to_sample = args.sample_len.saturating_sub(1); + let prompt_tokens = if prompt_tokens.len() + to_sample > model::MAX_SEQ_LEN - 10 { + let to_remove = prompt_tokens.len() + to_sample + 10 - model::MAX_SEQ_LEN; + prompt_tokens[prompt_tokens.len().saturating_sub(to_remove)..].to_vec() + } else { + prompt_tokens + }; + let mut all_tokens = vec![]; + let mut logits_processor = LogitsProcessor::new(args.seed, temperature, args.top_p); + + let start_prompt_processing = std::time::Instant::now(); + let mut next_token = { + let input = Tensor::new(prompt_tokens.as_slice(), &Device::Cpu)?.unsqueeze(0)?; + let logits = model.forward(&input, 0)?; + let logits = logits.squeeze(0)?; + logits_processor.sample(&logits)? + }; + let prompt_dt = start_prompt_processing.elapsed(); + all_tokens.push(next_token); + if let Some(t) = tos.next_token(next_token)? { + print!("{t}"); + std::io::stdout().flush()?; + } + + let eos_token = if args.which.is_open_chat() { + "<|end_of_turn|>" + } else { + "" + }; + let eos_token = *tos.tokenizer().get_vocab(true).get(eos_token).unwrap(); + let start_post_prompt = std::time::Instant::now(); + let mut sampled = 0; + for index in 0..to_sample { + let input = Tensor::new(&[next_token], &Device::Cpu)?.unsqueeze(0)?; + let logits = model.forward(&input, prompt_tokens.len() + index)?; + let logits = logits.squeeze(0)?; + let logits = if args.repeat_penalty == 1. { + logits + } else { + let start_at = all_tokens.len().saturating_sub(args.repeat_last_n); + candle_transformers::utils::apply_repeat_penalty( + &logits, + args.repeat_penalty, + &all_tokens[start_at..], + )? + }; + next_token = logits_processor.sample(&logits)?; + all_tokens.push(next_token); + if let Some(t) = tos.next_token(next_token)? { + print!("{t}"); + std::io::stdout().flush()?; + } + sampled += 1; + if next_token == eos_token { + break; + }; + } + if let Some(rest) = tos.decode_rest().map_err(candle::Error::msg)? { + print!("{rest}"); + } + std::io::stdout().flush()?; + let dt = start_post_prompt.elapsed(); + println!( + "\n\n{:4} prompt tokens processed: {:.2} token/s", + prompt_tokens.len(), + prompt_tokens.len() as f64 / prompt_dt.as_secs_f64(), + ); + println!( + "{sampled:4} tokens generated: {:.2} token/s", + sampled as f64 / dt.as_secs_f64(), + ); + + match prompt { + Prompt::One(_) => break, + Prompt::Interactive => {} + Prompt::Chat => { + pre_prompt_tokens = [prompt_tokens.as_slice(), all_tokens.as_slice()].concat() + } + } + } + + Ok(()) +} diff --git a/packages/llama/src/quantized_llama.rs b/packages/llama/src/quantized_llama.rs new file mode 100644 index 0000000..4ef22e3 --- /dev/null +++ b/packages/llama/src/quantized_llama.rs @@ -0,0 +1,509 @@ +use std::collections::HashMap; + +use candle::quantized::QTensor; +use candle::quantized::{ggml_file, gguf_file}; +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::{Embedding, Module}; + +pub const MAX_SEQ_LEN: usize = 4096; + +#[derive(Debug, Clone)] +struct RmsNorm { + inner: candle_nn::LayerNorm, + // span: tracing::Span, +} + +impl RmsNorm { + fn new(scale: QTensor, eps: f32) -> Result { + // let span = tracing::span!(tracing::Level::TRACE, "rms-norm"); + let scale = scale.dequantize(&Device::Cpu)?; + let inner = candle_nn::LayerNorm::rms_norm(scale, eps as f64); + Ok(Self { inner + // , span + }) + } + + fn forward(&self, x: &Tensor) -> Result { + // let _enter = self.span.enter(); + self.inner.forward(x) + } +} + +// QMatMul wrapper adding some tracing. +#[derive(Debug, Clone)] +struct QMatMul { + inner: candle::quantized::QMatMul, + // span: tracing::Span, +} + +impl QMatMul { + fn from_qtensor(qtensor: QTensor) -> Result { + let inner = candle::quantized::QMatMul::from_qtensor(qtensor)?; + // let span = tracing::span!(tracing::Level::TRACE, "qmatmul"); + Ok(Self { + inner, + // span + }) + } + + fn forward(&self, xs: &Tensor) -> Result { + // let _enter = self.span.enter(); + self.inner.forward(xs) + } +} + +#[derive(Debug, Clone)] +struct Mlp { + feed_forward_w1: QMatMul, + feed_forward_w2: QMatMul, + feed_forward_w3: QMatMul, +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let w1 = self.feed_forward_w1.forward(xs)?; + let w3 = self.feed_forward_w3.forward(xs)?; + self.feed_forward_w2 + .forward(&(candle_nn::ops::silu(&w1)? * w3)?) + } +} + +#[derive(Debug, Clone)] +enum MlpOrMoe { + Mlp(Mlp), + MoE { + n_expert_used: usize, + feed_forward_gate_inp: QMatMul, + experts: Vec, + }, +} + +impl Module for MlpOrMoe { + fn forward(&self, xs: &Tensor) -> Result { + match self { + Self::MoE { + feed_forward_gate_inp, + experts, + n_expert_used, + } => { + let (b_size, seq_len, hidden_dim) = xs.dims3()?; + let xs = xs.reshape(((), hidden_dim))?; + let router_logits = feed_forward_gate_inp.forward(&xs)?; + let routing_weights = candle_nn::ops::softmax_last_dim(&router_logits)?; + + // In order to extract topk, we extract the data from the tensor and manipulate it + // directly. Maybe we will want to use some custom ops instead at some point. + let routing_weights = routing_weights.to_dtype(DType::F32)?.to_vec2::()?; + + // routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + // top_x contains the row indexes to evaluate for each expert. + let mut top_x = vec![vec![]; experts.len()]; + let mut selected_rws = vec![vec![]; experts.len()]; + for (row_idx, rw) in routing_weights.iter().enumerate() { + let mut dst = (0..rw.len() as u32).collect::>(); + dst.sort_by(|&i, &j| rw[j as usize].total_cmp(&rw[i as usize])); + let mut sum_routing_weights = 0f32; + for &expert_idx in dst.iter().take(*n_expert_used) { + let expert_idx = expert_idx as usize; + let routing_weight = rw[expert_idx]; + sum_routing_weights += routing_weight; + top_x[expert_idx].push(row_idx as u32); + } + for &expert_idx in dst.iter().take(*n_expert_used) { + let expert_idx = expert_idx as usize; + let routing_weight = rw[expert_idx]; + selected_rws[expert_idx].push(routing_weight / sum_routing_weights) + } + } + + // routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + // expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) + + let mut ys = xs.zeros_like()?; + for (expert_idx, expert_layer) in experts.iter().enumerate() { + let top_x = &top_x[expert_idx]; + if top_x.is_empty() { + continue; + } + let top_x = Tensor::new(top_x.as_slice(), xs.device())?; + let selected_rws = + Tensor::new(selected_rws[expert_idx].as_slice(), xs.device())? + .reshape(((), 1))?; + // Index the correct hidden states and compute the expert hidden state for + // the current expert. We need to make sure to multiply the output hidden + // states by `routing_weights` on the corresponding tokens (top-1 and top-2) + let current_state = xs.index_select(&top_x, 0)?.reshape(((), hidden_dim))?; + // current_hidden_states = expert_layer(current_state, routing_weights[top_x_list, idx_list, None]) + let current_hidden_states = expert_layer.forward(¤t_state)?; + let current_hidden_states = + current_hidden_states.broadcast_mul(&selected_rws)?; + ys = ys.index_add(&top_x, ¤t_hidden_states, 0)?; + } + + let ys = ys.reshape((b_size, seq_len, hidden_dim))?; + Ok(ys) + } + Self::Mlp(mlp) => mlp.forward(xs), + } + } +} + +#[derive(Debug, Clone)] +struct LayerWeights { + attention_wq: QMatMul, + attention_wk: QMatMul, + attention_wv: QMatMul, + attention_wo: QMatMul, + attention_norm: RmsNorm, + mlp_or_moe: MlpOrMoe, + ffn_norm: RmsNorm, + n_head: usize, + n_kv_head: usize, + head_dim: usize, + cos: Tensor, + sin: Tensor, + kv_cache: Option<(Tensor, Tensor)>, + // span_attn: tracing::Span, + // span_rot: tracing::Span, + // span_mlp: tracing::Span, +} + +fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: f32) -> Result { + let shape = mask.shape(); + let on_true = Tensor::new(on_true, on_false.device())?.broadcast_as(shape.dims())?; + let m = mask.where_cond(&on_true, on_false)?; + Ok(m) +} + +impl LayerWeights { + fn apply_rotary_emb(&self, x: &Tensor, index_pos: usize) -> Result { + // let _enter = self.span_rot.enter(); + let (b_sz, n_head, seq_len, n_embd) = x.dims4()?; + let cos = self + .cos + .narrow(0, index_pos, seq_len)? + .reshape((seq_len, n_embd / 2, 1))?; + let sin = self + .sin + .narrow(0, index_pos, seq_len)? + .reshape((seq_len, n_embd / 2, 1))?; + let cos = cos.broadcast_as((b_sz, 1, seq_len, n_embd / 2, 1))?; + let sin = sin.broadcast_as((b_sz, 1, seq_len, n_embd / 2, 1))?; + // This mimics the llama.cpp behavior. + // https://github.com/ggerganov/llama.cpp/blob/1f0bccb27929e261744c979bc75114955da49e98/ggml.c#L12104-L12105 + // The x0 and x1 value are interleaved on the n_embd (= head_dim) dimension. + // The resulting y0 and y1 are also interleaved with: + // y0 = x0*cos - x1*sin + // y1 = x0*sin + x1*cos + let x = x.reshape((b_sz, n_head, seq_len, n_embd / 2, 2))?; + let x0 = x.narrow(D::Minus1, 0, 1)?; + let x1 = x.narrow(D::Minus1, 1, 1)?; + let y0 = (x0.broadcast_mul(&cos)? - x1.broadcast_mul(&sin)?)?; + let y1 = (x0.broadcast_mul(&sin)? + x1.broadcast_mul(&cos)?)?; + let rope = Tensor::cat(&[y0, y1], D::Minus1)?; + let rope = rope.flatten_from(D::Minus2)?; + Ok(rope) + } + + fn forward_attn(&mut self, x: &Tensor, mask: &Tensor, index_pos: usize) -> Result { + // let _enter = self.span_attn.enter(); + let (b_sz, seq_len, n_embd) = x.dims3()?; + let q = self.attention_wq.forward(x)?; + let k = self.attention_wk.forward(x)?; + let v = self.attention_wv.forward(x)?; + + let q = q + .reshape((b_sz, seq_len, self.n_head, self.head_dim))? + .transpose(1, 2)?; + let k = k + .reshape((b_sz, seq_len, self.n_kv_head, self.head_dim))? + .transpose(1, 2)?; + let v = v + .reshape((b_sz, seq_len, self.n_kv_head, self.head_dim))? + .transpose(1, 2)?; + + let q = self.apply_rotary_emb(&q, index_pos)?; + let k = self.apply_rotary_emb(&k, index_pos)?; + + let (k, v) = match &self.kv_cache { + None => (k, v), + Some((k_cache, v_cache)) => { + if index_pos == 0 { + (k, v) + } else { + let k = Tensor::cat(&[k_cache, &k], 2)?.contiguous()?; + let v = Tensor::cat(&[v_cache, &v], 2)?.contiguous()?; + (k, v) + } + } + }; + self.kv_cache = Some((k.clone(), v.clone())); + + // Support for MQA, useful for 70B models. + let k = self.repeat_kv(k)?; + let v = self.repeat_kv(v)?; + + let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?; + let mask = mask.broadcast_as(att.shape())?; + let att = masked_fill(&att, &mask, f32::NEG_INFINITY)?; + let att = candle_nn::ops::softmax_last_dim(&att)?; + // Convert to contiguous as matmul doesn't support strided vs for now. + let y = att.matmul(&v.contiguous()?)?; + let y = y.transpose(1, 2)?.reshape(&[b_sz, seq_len, n_embd])?; + let y = self.attention_wo.forward(&y)?; + Ok(y) + } + + fn repeat_kv(&self, x: Tensor) -> Result { + let n_rep = self.n_head / self.n_kv_head; + if n_rep == 1 { + Ok(x) + } else { + let (b_sz, n_kv_head, seq_len, head_dim) = x.dims4()?; + let x = x + .unsqueeze(2)? + .expand((b_sz, n_kv_head, n_rep, seq_len, head_dim))? + .reshape((b_sz, n_kv_head * n_rep, seq_len, head_dim))?; + Ok(x) + } + } +} + +#[derive(Debug, Clone)] +pub struct ModelWeights { + tok_embeddings: Embedding, + layers: Vec, + norm: RmsNorm, + output: QMatMul, + masks: HashMap, + // span: tracing::Span, + // span_output: tracing::Span, +} + +fn precomput_freqs_cis(head_dim: usize, freq_base: f32) -> Result<(Tensor, Tensor)> { + let theta: Vec<_> = (0..head_dim) + .step_by(2) + .map(|i| 1f32 / freq_base.powf(i as f32 / head_dim as f32)) + .collect(); + let theta = Tensor::new(theta.as_slice(), &Device::Cpu)?; + let idx_theta = Tensor::arange(0, MAX_SEQ_LEN as u32, &Device::Cpu)? + .to_dtype(DType::F32)? + .reshape((MAX_SEQ_LEN, 1))? + .matmul(&theta.reshape((1, theta.elem_count()))?)?; + let cos = idx_theta.cos()?; + let sin = idx_theta.sin()?; + Ok((cos, sin)) +} + +impl ModelWeights { + pub fn from_ggml(mut ct: ggml_file::Content, gqa: usize) -> Result { + let cpu = &Device::Cpu; + let head_dim = (ct.hparams.n_embd / ct.hparams.n_head) as usize; + let (cos, sin) = precomput_freqs_cis(head_dim, 10000.)?; + let tok_embeddings = ct.remove("tok_embeddings.weight")?; + let tok_embeddings = tok_embeddings.dequantize(cpu)?; + let norm = RmsNorm::new(ct.remove("norm.weight")?, 1e-5)?; + let output = ct.remove("output.weight")?; + let mut layers = Vec::with_capacity(ct.hparams.n_layer as usize); + for layer_idx in 0..ct.hparams.n_layer { + let prefix = format!("layers.{layer_idx}"); + let attention_wq = ct.remove(&format!("{prefix}.attention.wq.weight"))?; + let attention_wk = ct.remove(&format!("{prefix}.attention.wk.weight"))?; + let attention_wv = ct.remove(&format!("{prefix}.attention.wv.weight"))?; + let attention_wo = ct.remove(&format!("{prefix}.attention.wo.weight"))?; + let mlp_or_moe = { + let feed_forward_w1 = ct.remove(&format!("{prefix}.feed_forward.w1.weight"))?; + let feed_forward_w2 = ct.remove(&format!("{prefix}.feed_forward.w2.weight"))?; + let feed_forward_w3 = ct.remove(&format!("{prefix}.feed_forward.w3.weight"))?; + MlpOrMoe::Mlp(Mlp { + feed_forward_w1: QMatMul::from_qtensor(feed_forward_w1)?, + feed_forward_w2: QMatMul::from_qtensor(feed_forward_w2)?, + feed_forward_w3: QMatMul::from_qtensor(feed_forward_w3)?, + }) + }; + let attention_norm = ct.remove(&format!("{prefix}.attention_norm.weight"))?; + let ffn_norm = ct.remove(&format!("{prefix}.ffn_norm.weight"))?; + // let span_attn = tracing::span!(tracing::Level::TRACE, "attn"); + // let span_rot = tracing::span!(tracing::Level::TRACE, "attn-rot"); + // let span_mlp = tracing::span!(tracing::Level::TRACE, "attn-mlp"); + layers.push(LayerWeights { + attention_wq: QMatMul::from_qtensor(attention_wq)?, + attention_wk: QMatMul::from_qtensor(attention_wk)?, + attention_wv: QMatMul::from_qtensor(attention_wv)?, + attention_wo: QMatMul::from_qtensor(attention_wo)?, + attention_norm: RmsNorm::new(attention_norm, 1e-5)?, + mlp_or_moe, + ffn_norm: RmsNorm::new(ffn_norm, 1e-5)?, + n_head: ct.hparams.n_head as usize, + n_kv_head: ct.hparams.n_head as usize / gqa, + head_dim: (ct.hparams.n_embd / ct.hparams.n_head) as usize, + cos: cos.clone(), + sin: sin.clone(), + kv_cache: None, + // span_attn, + // span_rot, + // span_mlp, + }) + } + // let span = tracing::span!(tracing::Level::TRACE, "model"); + // let span_output = tracing::span!(tracing::Level::TRACE, "output"); + Ok(Self { + tok_embeddings: Embedding::new(tok_embeddings, ct.hparams.n_embd as usize), + layers, + norm, + output: QMatMul::from_qtensor(output)?, + masks: HashMap::new(), + // span, + // span_output, + }) + } + + pub fn from_gguf( + ct: gguf_file::Content, + reader: &mut R, + ) -> Result { + let cpu = &Device::Cpu; + let md_get = |s: &str| match ct.metadata.get(s) { + None => candle::bail!("cannot find {s} in metadata"), + Some(v) => Ok(v), + }; + + // Parameter extraction from metadata. + let n_expert = md_get("llama.expert_count") + .and_then(|v| v.to_u32()) + .unwrap_or(0) as usize; + let n_expert_used = md_get("llama.expert_used_count") + .and_then(|v| v.to_u32()) + .unwrap_or(0) as usize; + let head_count = md_get("llama.attention.head_count")?.to_u32()? as usize; + let head_count_kv = md_get("llama.attention.head_count_kv")?.to_u32()? as usize; + let block_count = md_get("llama.block_count")?.to_u32()? as usize; + let embedding_length = md_get("llama.embedding_length")?.to_u32()? as usize; + let rope_dim = md_get("llama.rope.dimension_count")?.to_u32()? as usize; + // Strangely this value is generally 1e-6 in GGUF file but used to be 1e-5 by default. + let rms_norm_eps = md_get("llama.attention.layer_norm_rms_epsilon")?.to_f32()?; + + let rope_freq_base = md_get("llama.rope.freq_base") + .and_then(|m| m.to_f32()) + .unwrap_or(10000f32); + let (cos, sin) = precomput_freqs_cis(rope_dim, rope_freq_base)?; + + let tok_embeddings = ct.tensor(reader, "token_embd.weight")?; + let tok_embeddings = tok_embeddings.dequantize(cpu)?; + let norm = RmsNorm::new(ct.tensor(reader, "output_norm.weight")?, rms_norm_eps)?; + let output = ct.tensor(reader, "output.weight")?; + let mut layers = Vec::with_capacity(block_count); + for layer_idx in 0..block_count { + let prefix = format!("blk.{layer_idx}"); + let attention_wq = ct.tensor(reader, &format!("{prefix}.attn_q.weight"))?; + let attention_wk = ct.tensor(reader, &format!("{prefix}.attn_k.weight"))?; + let attention_wv = ct.tensor(reader, &format!("{prefix}.attn_v.weight"))?; + let attention_wo = ct.tensor(reader, &format!("{prefix}.attn_output.weight"))?; + let mlp_or_moe = if n_expert <= 1 { + let feed_forward_w1 = ct.tensor(reader, &format!("{prefix}.ffn_gate.weight"))?; + let feed_forward_w2 = ct.tensor(reader, &format!("{prefix}.ffn_down.weight"))?; + let feed_forward_w3 = ct.tensor(reader, &format!("{prefix}.ffn_up.weight"))?; + MlpOrMoe::Mlp(Mlp { + feed_forward_w1: QMatMul::from_qtensor(feed_forward_w1)?, + feed_forward_w2: QMatMul::from_qtensor(feed_forward_w2)?, + feed_forward_w3: QMatMul::from_qtensor(feed_forward_w3)?, + }) + } else { + let feed_forward_gate_inp = + ct.tensor(reader, &format!("{prefix}.ffn_gate_inp.weight"))?; + let mut experts = Vec::with_capacity(n_expert); + for i in 0..n_expert { + let feed_forward_w1 = + ct.tensor(reader, &format!("{prefix}.ffn_gate.{i}.weight"))?; + let feed_forward_w2 = + ct.tensor(reader, &format!("{prefix}.ffn_down.{i}.weight"))?; + let feed_forward_w3 = + ct.tensor(reader, &format!("{prefix}.ffn_up.{i}.weight"))?; + experts.push(Mlp { + feed_forward_w1: QMatMul::from_qtensor(feed_forward_w1)?, + feed_forward_w2: QMatMul::from_qtensor(feed_forward_w2)?, + feed_forward_w3: QMatMul::from_qtensor(feed_forward_w3)?, + }) + } + MlpOrMoe::MoE { + n_expert_used, + feed_forward_gate_inp: QMatMul::from_qtensor(feed_forward_gate_inp)?, + experts, + } + }; + let attention_norm = ct.tensor(reader, &format!("{prefix}.attn_norm.weight"))?; + let ffn_norm = ct.tensor(reader, &format!("{prefix}.ffn_norm.weight"))?; + // let span_attn = tracing::span!(tracing::Level::TRACE, "attn"); + // let span_rot = tracing::span!(tracing::Level::TRACE, "attn-rot"); + // let span_mlp = tracing::span!(tracing::Level::TRACE, "attn-mlp"); + layers.push(LayerWeights { + attention_wq: QMatMul::from_qtensor(attention_wq)?, + attention_wk: QMatMul::from_qtensor(attention_wk)?, + attention_wv: QMatMul::from_qtensor(attention_wv)?, + attention_wo: QMatMul::from_qtensor(attention_wo)?, + attention_norm: RmsNorm::new(attention_norm, rms_norm_eps)?, + mlp_or_moe, + ffn_norm: RmsNorm::new(ffn_norm, rms_norm_eps)?, + n_head: head_count, + n_kv_head: head_count_kv, + head_dim: embedding_length / head_count, + cos: cos.clone(), + sin: sin.clone(), + kv_cache: None, + // span_attn, + // span_rot, + // span_mlp, + }) + } + // let span = tracing::span!(tracing::Level::TRACE, "model"); + // let span_output = tracing::span!(tracing::Level::TRACE, "output"); + Ok(Self { + tok_embeddings: Embedding::new(tok_embeddings, embedding_length), + layers, + norm, + output: QMatMul::from_qtensor(output)?, + masks: HashMap::new(), + // span, + // span_output, + }) + } + + fn mask(&mut self, t: usize) -> Result { + if let Some(mask) = self.masks.get(&t) { + Ok(mask.clone()) + } else { + let mask: Vec<_> = (0..t) + .flat_map(|i| (0..t).map(move |j| u8::from(j > i))) + .collect(); + let mask = Tensor::from_slice(&mask, (t, t), &Device::Cpu)?; + self.masks.insert(t, mask.clone()); + Ok(mask) + } + } + + pub fn forward(&mut self, x: &Tensor, index_pos: usize) -> Result { + let (_b_sz, seq_len) = x.dims2()?; + let mask = self.mask(seq_len)?; + // let _enter = self.span.enter(); + let mut layer_in = self.tok_embeddings.forward(x)?; + for layer in self.layers.iter_mut() { + let x = layer_in; + let residual = &x; + let x = layer.attention_norm.forward(&x)?; + let attn = layer.forward_attn(&x, &mask, index_pos)?; + let x = (attn + residual)?; + + // MLP + // let _enter = layer.span_mlp.enter(); + let residual = &x; + let x = layer.ffn_norm.forward(&x)?; + let x = layer.mlp_or_moe.forward(&x)?; + let x = (x + residual)?; + layer_in = x + } + let x = self.norm.forward(&layer_in)?; + let x = x.i((.., seq_len - 1, ..))?; + // let _enter = self.span_output.enter(); + self.output.forward(&x) + } +} diff --git a/packages/llama/src/worker.rs b/packages/llama/src/worker.rs index 79dd2f3..a34cb96 100644 --- a/packages/llama/src/worker.rs +++ b/packages/llama/src/worker.rs @@ -2,11 +2,9 @@ use crate::model::{Cache, Config, Llama}; use byteorder::{LittleEndian, ReadBytesExt}; use candle::{DType, Device, IndexOp, Result, Shape, Tensor}; use candle_nn::VarBuilder; -use candle_transformers::generation::LogitsProcessor; use serde::{Deserialize, Serialize}; use tokenizers::Tokenizer; use wasm_bindgen::prelude::*; -use yew_agent::{HandlerId, Public, WorkerLink}; #[wasm_bindgen] extern "C" { @@ -56,59 +54,6 @@ pub struct Model { pub tokenizer: Tokenizer, } -impl Model { - fn run( - &self, - link: &WorkerLink, - id: HandlerId, - temp: f64, - top_p: f64, - prompt: String, - ) -> Result<()> { - let dev = Device::Cpu; - let temp = if temp <= 0. { None } else { Some(temp) }; - let top_p = if top_p <= 0. || top_p >= 1.0 { - None - } else { - Some(top_p) - }; - console_log!("temp: {temp:?} top_p: {top_p:?} prompt: {prompt}"); - let mut logits_processor = LogitsProcessor::new(299792458, temp, top_p); - let mut index_pos = 0; - let mut tokens = self - .tokenizer - .encode(prompt.to_string(), true) - .map_err(|m| candle::Error::Msg(m.to_string()))? - .get_ids() - .to_vec(); - link.respond(id, Ok(WorkerOutput::Generated(prompt))); - - for index in 0.. { - if tokens.len() >= self.config.seq_len { - break; - } - let context_size = if self.cache.use_kv_cache && index > 0 { - 1 - } else { - tokens.len() - }; - let ctxt = &tokens[tokens.len().saturating_sub(context_size)..]; - let input = Tensor::new(ctxt, &dev)?.unsqueeze(0)?; - let logits = self.llama.forward(&input, index_pos)?; - let logits = logits.squeeze(0)?; - index_pos += ctxt.len(); - - let next_token = logits_processor.sample(&logits)?; - tokens.push(next_token); - if let Some(text) = self.tokenizer.id_to_token(next_token) { - let text = text.replace('▁', " ").replace("<0x0A>", "\n"); - link.respond(id, Ok(WorkerOutput::Generated(text))); - } - } - Ok(()) - } -} - impl Config { fn from_reader(r: &mut R) -> Result { let dim = read_i32(r)? as usize; @@ -266,11 +211,6 @@ impl Model { } } -pub struct Worker { - link: WorkerLink, - model: Option, -} - #[derive(Serialize, Deserialize)] pub enum WorkerInput { ModelData(ModelData), @@ -283,54 +223,3 @@ pub enum WorkerOutput { GenerationDone(std::result::Result<(), String>), WeightsLoaded, } - -impl yew_agent::Worker for Worker { - type Input = WorkerInput; - type Message = (); - type Output = std::result::Result; - type Reach = Public; - - fn create(link: WorkerLink) -> Self { - Self { link, model: None } - } - - fn update(&mut self, _msg: Self::Message) { - // no messaging - } - - fn handle_input(&mut self, msg: Self::Input, id: HandlerId) { - let output = match msg { - WorkerInput::ModelData(md) => match Model::load(md) { - Ok(model) => { - self.model = Some(model); - Ok(WorkerOutput::WeightsLoaded) - } - Err(err) => Err(format!("model creation error {err:?}")), - }, - WorkerInput::Run(temp, top_p, prompt) => match &mut self.model { - None => Err("model has not been set yet".to_string()), - Some(model) => { - { - let mut cache = model.cache.kvs.lock().unwrap(); - for elem in cache.iter_mut() { - *elem = None - } - } - let result = model - .run(&self.link, id, temp, top_p, prompt) - .map_err(|e| e.to_string()); - Ok(WorkerOutput::GenerationDone(result)) - } - }, - }; - self.link.respond(id, output); - } - - fn name_of_resource() -> &'static str { - "worker.js" - } - - fn resource_path_is_relative() -> bool { - true - } -}