diff --git a/Cargo.lock b/Cargo.lock index 3a2c896a17..ec5cfd5e58 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1121,8 +1121,6 @@ dependencies = [ [[package]] name = "candle-transformers" version = "0.9.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b538ec4aa807c416a2ddd3621044888f188827862e2a6fcacba4738e89795d01" dependencies = [ "byteorder", "candle-core 0.9.2", @@ -1131,6 +1129,7 @@ dependencies = [ "num-traits", "rand 0.9.4", "rayon", + "ruvllm_sparse_attention", "serde", "serde_json", "serde_plain", @@ -1631,7 +1630,7 @@ checksum = "c07782be35f9e1140080c6b96f0d44b739e2278479f64e02fdab4e32dfd8b081" dependencies = [ "bitflags 1.3.2", "core-foundation 0.9.4", - "core-graphics-types", + "core-graphics-types 0.1.3", "foreign-types 0.5.0", "libc", ] @@ -1647,6 +1646,17 @@ dependencies = [ "libc", ] +[[package]] +name = "core-graphics-types" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d44a101f213f6c4cdc1853d4b78aef6db6bdfa3468798cc1d9912f4735013eb" +dependencies = [ + "bitflags 2.13.0", + "core-foundation 0.10.1", + "libc", +] + [[package]] name = "core-text" version = "20.1.0" @@ -4914,6 +4924,28 @@ dependencies = [ "regex-automata", ] +[[package]] +name = "lattice-inference" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cf759944705182b487cc3a1d99687fa518f031c8239857d351297c9cecabd42" +dependencies = [ + "clap", + "half", + "image 0.25.10", + "indexmap 2.12.1", + "memmap2", + "metal 0.33.0", + "objc", + "rayon", + "rustc-hash 2.1.2", + "serde", + "serde_json", + "sha2 0.10.9", + "thiserror 2.0.18", + "tracing", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -5346,7 +5378,22 @@ checksum = "7ecfd3296f8c56b7c1f6fbac3c71cefa9d78ce009850c45000015f206dc7fa21" dependencies = [ "bitflags 2.13.0", "block", - "core-graphics-types", + "core-graphics-types 0.1.3", + "foreign-types 0.5.0", + "log", + "objc", + "paste", +] + +[[package]] +name = "metal" +version = "0.33.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7047791b5bc903b8cd963014b355f71dc9864a9a0b727057676c1dcae5cbc15" +dependencies = [ + "bitflags 2.13.0", + "block", + "core-graphics-types 0.2.0", "foreign-types 0.5.0", "log", "objc", @@ -10909,9 +10956,10 @@ dependencies = [ "futures-core", "half", "hf-hub", + "lattice-inference", "md5", "memmap2", - "metal", + "metal 0.29.0", "ndarray 0.16.1", "objc", "objc2", @@ -10927,6 +10975,7 @@ dependencies = [ "ruvector-gnn", "ruvector-graph", "ruvector-sona 0.2.1", + "ruvllm_sparse_attention", "serde", "serde_json", "sha2 0.10.9", @@ -13411,7 +13460,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7adf545a99a086d362efc739e7cf4317c18cbeda22706000fd434d70ea3d95" dependencies = [ "half", - "metal", + "metal 0.29.0", "objc", "serde", "thiserror 1.0.69", @@ -14069,7 +14118,7 @@ dependencies = [ "block", "bytemuck", "cfg_aliases 0.1.1", - "core-graphics-types", + "core-graphics-types 0.1.3", "glow", "glutin_wgl_sys", "gpu-alloc", @@ -14080,7 +14129,7 @@ dependencies = [ "libc", "libloading 0.8.9", "log", - "metal", + "metal 0.29.0", "naga", "ndk-sys", "objc", diff --git a/Cargo.toml b/Cargo.toml index c2f4d345c8..903369b3a2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -369,5 +369,12 @@ debug = true # Patch hnsw_rs to use rand 0.8 instead of 0.9 for WASM compatibility # This resolves the getrandom version conflict (0.2 vs 0.3) +# +# Patch candle-transformers (0.9.2) to add ModelWeights::forward_all_positions, +# which returns per-position logits for a multi-token forward pass instead of +# only the last position. Needed by ruvllm's speculative decoding to verify +# several draft tokens against the main model in a single batched forward +# pass; upstream's `forward()` always slices to the final position. [patch.crates-io] hnsw_rs = { path = "./patches/hnsw_rs" } +candle-transformers = { path = "./patches/candle-transformers" } diff --git a/crates/ruvllm/CHANGELOG.md b/crates/ruvllm/CHANGELOG.md index 9ba16bd703..8aa4b54d9a 100644 --- a/crates/ruvllm/CHANGELOG.md +++ b/crates/ruvllm/CHANGELOG.md @@ -5,6 +5,29 @@ All notable changes to the ruvllm crate will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added +- Real speculative decoding: `speculative.rs`'s `SpeculativeDecoder` now drafts and + verifies against actual model logits (single batched forward pass per round to + verify K draft tokens) instead of the previous text-roundtrip stub. Requires a new + `SpeculativeBackend` trait (implemented for `CandleBackend`) exposing per-position + batched logits and cheap KV-cache snapshot/restore. See `examples/speculative_bench.rs` + for real measured tokens/sec on Llama-2-7B (main) + TinyLlama-1.1B (draft), GGUF. +- `ruvllm_sparse_attention` integration: prompt-prefill attention in the GGUF + (`QuantizedLlama`) backend can now route through + `ruvllm_sparse_attention::SubquadraticSparseAttention` instead of dense causal + attention, via the new `sparse-attention` feature and + `CandleBackend::enable_sparse_attention`. See `examples/sparse_attention_check.rs`. +- `patches/candle-transformers`: a patched copy of candle-transformers 0.9.2 (see + `patches/README.md`) adding `ModelWeights::forward_all_positions` (per-position + logits from a multi-token forward pass), `snapshot_kv_cache`/`restore_kv_cache`, + a fix for the causal mask when continuing a non-empty KV cache with more than one + new token, and (behind `sparse-attention`) `forward_sparse`/`forward_attn_sparse`. +- `CandleBackend::forward_multi`, `reset_context`, `context_len`, `snapshot_context`, + `restore_context`, `vocab_size` — the primitives speculative decoding needs, usable + independently of it. + ## [2.0.0] - 2025-01-19 ### Added diff --git a/crates/ruvllm/Cargo.toml b/crates/ruvllm/Cargo.toml index a15b169da2..c6e9206953 100644 --- a/crates/ruvllm/Cargo.toml +++ b/crates/ruvllm/Cargo.toml @@ -68,6 +68,11 @@ candle-core = { version = "0.9", optional = true } candle-nn = { version = "0.9", optional = true } candle-transformers = { version = "0.9", optional = true } +# Subquadratic sparse attention, optionally spliced into the candle-transformers +# GGUF inference path (patches/candle-transformers) for prompt prefill. See the +# `sparse-attention` feature below. +ruvllm_sparse_attention = { path = "../ruvllm_sparse_attention", optional = true } + # Direct cudarc access for fused CUDA kernels (fused-act feature). # Upgraded to 0.19 alongside candle 0.9 — supports CUDA 13.0 natively. cudarc = { version = "0.19", optional = true, default-features = false, features = ["nvrtc", "driver", "std"] } @@ -156,6 +161,18 @@ parallel = ["dep:rayon"] # load works without it; HF Hub fetch needs `hub-download` feature. candle = ["candle-core", "candle-nn", "candle-transformers", "tokenizers"] +# Route the candle GGUF backend's prompt-prefill attention through +# ruvllm_sparse_attention::SubquadraticSparseAttention instead of dense O(n^2) +# causal attention. See patches/candle-transformers's quantized_llama.rs +# (forward_sparse / forward_attn_sparse) and CandleBackend::load_gguf's +# `sparse_attention` config for the wiring; falls back to dense attention for +# incremental decode steps, where prefill's forward_gqa contract doesn't apply. +sparse-attention = [ + "candle", + "candle-transformers/sparse-attention", + "dep:ruvllm_sparse_attention", +] + # Lattice backend for LLM inference (pure-Rust, Metal acceleration on Mac). # macOS-only today: the dependency is declared under # `[target.'cfg(target_os = "macos")'.dependencies]` so `metal`/`objc` FFI @@ -302,6 +319,16 @@ name = "run_eval" path = "examples/run_eval.rs" required-features = ["async-runtime"] +[[example]] +name = "speculative_bench" +path = "examples/speculative_bench.rs" +required-features = ["candle"] + +[[example]] +name = "sparse_attention_check" +path = "examples/sparse_attention_check.rs" +required-features = ["sparse-attention"] + # Lint configuration for the package [lints.clippy] # Allow common patterns in test code for readability diff --git a/crates/ruvllm/README.md b/crates/ruvllm/README.md index e6e183059f..22804d4515 100644 --- a/crates/ruvllm/README.md +++ b/crates/ruvllm/README.md @@ -505,31 +505,37 @@ println!("Batch utilization: {:.1}%", stats.avg_batch_utilization * 100.0); ## Speculative Decoding -Accelerate generation with draft models: +Accelerate generation with a small draft model, verified against the main model in +batched forward passes. Main and draft models **must share a tokenizer/vocabulary** +(e.g. TinyLlama-1.1B drafting for Llama-2-7B/13B, which use the same tokenizer) — see +`src/speculative.rs` module docs for the full design, its KV-cache-recovery approach, +and measured tokens/sec on real weights (`examples/speculative_bench.rs`). ```rust use ruvllm::speculative::{SpeculativeDecoder, SpeculativeConfig}; +use std::sync::Arc; let config = SpeculativeConfig { - draft_tokens: 4, // Tokens to draft per step - acceptance_threshold: 0.8, // Min probability for acceptance + lookahead: 5, // Draft tokens to speculate per round ..Default::default() }; let decoder = SpeculativeDecoder::new( - target_model, - draft_model, + Arc::new(target_model), + Arc::new(draft_model), config, -)?; +); -// Generate with speculation +// Generate with speculation (greedy/low-temperature — see should_use_speculative) let output = decoder.generate(prompt, GenerateParams { max_tokens: 256, + temperature: 0.0, ..Default::default() })?; -println!("Acceptance rate: {:.1}%", output.stats.acceptance_rate * 100.0); -println!("Speedup: {:.2}x", output.stats.speedup); +let stats = decoder.stats(); +println!("Acceptance rate: {:.1}%", stats.acceptance_rate * 100.0); +println!("Avg tokens per main-model forward pass: {:.2}", stats.avg_tokens_per_main_pass); ``` ## GGUF Model Loading diff --git a/crates/ruvllm/examples/sparse_attention_check.rs b/crates/ruvllm/examples/sparse_attention_check.rs new file mode 100644 index 0000000000..045b66b8ed --- /dev/null +++ b/crates/ruvllm/examples/sparse_attention_check.rs @@ -0,0 +1,108 @@ +#![allow(clippy::all)] +//! Sanity check for the sparse-attention integration +//! (`patches/candle-transformers`'s `forward_sparse` / `forward_attn_sparse`, +//! wired via `CandleBackend::enable_sparse_attention`). +//! +//! Loads one real GGUF model, generates the same prompt with dense attention +//! and then with sparse attention enabled, and prints both outputs so a human +//! can eyeball that sparse attention produces coherent (not garbled) text — +//! this is not a numerical equivalence check (sparse attention is a genuinely +//! different, approximate algorithm; the tensor bridging is the thing under +//! test here, not exact output parity). +//! +//! ```bash +//! cargo run -p ruvllm --release --features candle,metal,sparse-attention \ +//! --example sparse_attention_check -- --model ./model.gguf +//! ``` + +use std::env; +use std::path::PathBuf; + +#[cfg(feature = "sparse-attention")] +fn main() { + use ruvllm::{CandleBackend, GenerateParams, LlmBackend, ModelConfig}; + + let args: Vec = env::args().collect(); + let mut model_path = PathBuf::new(); + let mut i = 1; + while i < args.len() { + if args[i] == "--model" || args[i] == "-m" { + i += 1; + if i < args.len() { + model_path = PathBuf::from(&args[i]); + } + } + i += 1; + } + if model_path.as_os_str().is_empty() { + eprintln!("Usage: --model (tokenizer.json expected alongside it)"); + std::process::exit(1); + } + let tokenizer_path = model_path + .parent() + .expect("model path has no parent") + .join("tokenizer.json"); + + let params = GenerateParams { + max_tokens: 40, + temperature: 0.0, + top_p: 1.0, + top_k: 1, + ..Default::default() + }; + let long = args.iter().any(|a| a == "--long"); + let prompt = if long { + // Deliberately longer than SparseAttentionConfig::default()'s + // 128-token window, so the local-window/landmark sparsity actually + // activates instead of degenerating to full attention. + "The history of the city of Paris stretches back over two thousand years. \ + Founded by a Celtic tribe known as the Parisii around the 3rd century BC, \ + the settlement was originally located on the Ile de la Cite, an island in \ + the Seine river. The Romans conquered the area in 52 BC and renamed it \ + Lutetia, building roads, baths, and an amphitheater. Over the following \ + centuries the city grew steadily, becoming the capital of the Kingdom of \ + France under the Capetian dynasty. During the Middle Ages, Paris became a \ + major center of learning, commerce, and religion in Europe. The capital of \ + France is" + } else { + "The capital of France is" + }; + + let mut dense = CandleBackend::new().expect("create backend"); + dense + .load_gguf(&model_path, &ModelConfig::default()) + .expect("load gguf (dense)"); + dense + .load_tokenizer(&tokenizer_path) + .expect("load tokenizer (dense)"); + let dense_output = dense + .generate(prompt, params.clone()) + .expect("dense generate"); + println!("[dense] {dense_output:?}"); + + let mut sparse = CandleBackend::new().expect("create backend"); + sparse + .load_gguf(&model_path, &ModelConfig::default()) + .expect("load gguf (sparse)"); + sparse + .load_tokenizer(&tokenizer_path) + .expect("load tokenizer (sparse)"); + sparse + .enable_sparse_attention(ruvllm_sparse_attention::SparseAttentionConfig::default()) + .expect("enable sparse attention"); + assert!(sparse.sparse_attention_enabled()); + let sparse_output = sparse.generate(prompt, params).expect("sparse generate"); + println!("[sparse] {sparse_output:?}"); + + if dense_output.trim().is_empty() || sparse_output.trim().is_empty() { + eprintln!("FAIL: empty output from one of the paths"); + std::process::exit(1); + } + println!("\nBoth paths produced non-empty output; inspect them above for coherence."); +} + +#[cfg(not(feature = "sparse-attention"))] +fn main() { + eprintln!("This example requires --features candle,metal,sparse-attention"); + std::process::exit(1); +} diff --git a/crates/ruvllm/examples/speculative_bench.rs b/crates/ruvllm/examples/speculative_bench.rs new file mode 100644 index 0000000000..bcfeda083b --- /dev/null +++ b/crates/ruvllm/examples/speculative_bench.rs @@ -0,0 +1,282 @@ +#![allow(clippy::all)] +//! Real, end-to-end speculative decoding benchmark. +//! +//! Loads two real GGUF models (a small draft model and a larger main model +//! that **must share the same tokenizer/vocabulary** — see +//! `ruvllm::speculative` module docs) and measures actual tokens/sec for: +//! +//! - Baseline: the main model decoding autoregressively on its own +//! (`CandleBackend::generate`, one real forward pass per token). +//! - Speculative: `SpeculativeDecoder` drafting with the small model and +//! verifying batches of draft tokens against the main model in single +//! forward passes (`crates/ruvllm/src/speculative.rs`). +//! +//! Both paths run real candle forward passes on real weights; nothing here +//! is simulated. Requires the `candle` feature (on by default). +//! +//! ## Usage +//! +//! ```bash +//! cargo run -p ruvllm --release --example speculative_bench -- \ +//! --main ./models/main-llama2-7b/model.gguf \ +//! --draft ./models/draft-tinyllama/model.gguf \ +//! --max-tokens 64 --lookahead 5 +//! ``` +//! +//! Tokenizers are expected at `tokenizer.json` next to each GGUF file +//! (same convention as `benchmark_model.rs`). + +use std::env; +use std::path::PathBuf; +use std::time::Instant; + +struct Config { + main_path: PathBuf, + draft_path: PathBuf, + max_tokens: usize, + lookahead: usize, + iterations: usize, + json_output: bool, + prompts: Vec, +} + +impl Default for Config { + fn default() -> Self { + Self { + main_path: PathBuf::new(), + draft_path: PathBuf::new(), + max_tokens: 64, + lookahead: 5, + iterations: 3, + json_output: false, + prompts: vec![ + "The quick brown fox jumps over".to_string(), + "Once upon a time in a distant land,".to_string(), + "The capital of France is".to_string(), + ], + } + } +} + +fn parse_args() -> Config { + let args: Vec = env::args().collect(); + let mut cfg = Config::default(); + + if args.contains(&"--help".to_string()) || args.contains(&"-h".to_string()) { + print_help(); + std::process::exit(0); + } + + let mut i = 1; + while i < args.len() { + match args[i].as_str() { + "--main" => { + i += 1; + if i < args.len() { + cfg.main_path = PathBuf::from(&args[i]); + } + } + "--draft" => { + i += 1; + if i < args.len() { + cfg.draft_path = PathBuf::from(&args[i]); + } + } + "--max-tokens" | "-t" => { + i += 1; + if i < args.len() { + cfg.max_tokens = args[i].parse().unwrap_or(64); + } + } + "--lookahead" | "-k" => { + i += 1; + if i < args.len() { + cfg.lookahead = args[i].parse().unwrap_or(5); + } + } + "--iterations" | "-i" => { + i += 1; + if i < args.len() { + cfg.iterations = args[i].parse().unwrap_or(3); + } + } + "--json" | "-j" => cfg.json_output = true, + _ => {} + } + i += 1; + } + + if cfg.main_path.as_os_str().is_empty() || cfg.draft_path.as_os_str().is_empty() { + eprintln!("Error: --main and --draft model paths are required.\n"); + print_help(); + std::process::exit(1); + } + + cfg +} + +fn print_help() { + println!("RuvLLM Speculative Decoding Benchmark"); + println!(); + println!("USAGE:"); + println!(" cargo run -p ruvllm --release --example speculative_bench -- \\"); + println!(" --main --draft [OPTIONS]"); + println!(); + println!("OPTIONS:"); + println!(" --main Path to main (target) GGUF model"); + println!(" --draft Path to draft GGUF model (must share the"); + println!(" main model's tokenizer/vocabulary)"); + println!(" -t, --max-tokens Tokens to generate per prompt (default: 64)"); + println!(" -k, --lookahead Draft lookahead (default: 5)"); + println!(" -i, --iterations Iterations per prompt (default: 3)"); + println!(" -j, --json Output results as JSON"); + println!(" -h, --help Print help information"); +} + +#[cfg(feature = "candle")] +fn main() { + use ruvllm::speculative::{SpeculativeConfig, SpeculativeDecoder}; + use ruvllm::{CandleBackend, GenerateParams, LlmBackend, ModelConfig}; + use std::sync::Arc; + + let cfg = parse_args(); + + let load = |path: &PathBuf, label: &str| -> CandleBackend { + println!("Loading {label} model: {}", path.display()); + let mut backend = CandleBackend::new().expect("failed to create CandleBackend"); + backend + .load_gguf(path, &ModelConfig::default()) + .unwrap_or_else(|e| panic!("failed to load {label} GGUF at {path:?}: {e}")); + let tokenizer_path = path + .parent() + .expect("model path has no parent directory") + .join("tokenizer.json"); + backend.load_tokenizer(&tokenizer_path).unwrap_or_else(|e| { + panic!("failed to load {label} tokenizer at {tokenizer_path:?}: {e}") + }); + backend + }; + + // Three independent backend instances so baseline and speculative runs + // never share (and can't corrupt) each other's KV cache state. + let main_baseline = load(&cfg.main_path, "main (baseline)"); + let main_spec = Arc::new(load(&cfg.main_path, "main (speculative)")); + let draft_spec = Arc::new(load(&cfg.draft_path, "draft (speculative)")); + + let params = GenerateParams { + max_tokens: cfg.max_tokens, + temperature: 0.0, // greedy — required for the exact-match verification in speculative.rs + top_p: 1.0, + top_k: 1, + ..Default::default() + }; + + let spec_config = SpeculativeConfig { + lookahead: cfg.lookahead, + draft_temperature: 0.0, + ..Default::default() + }; + let decoder = SpeculativeDecoder::new(main_spec, draft_spec, spec_config); + + let mut baseline_tps = Vec::new(); + let mut spec_tps = Vec::new(); + + for prompt in &cfg.prompts { + for iter in 0..cfg.iterations { + // --- Baseline: main model decoding on its own --- + let start = Instant::now(); + let output = main_baseline + .generate(prompt, params.clone()) + .expect("baseline generation failed"); + let elapsed = start.elapsed(); + let token_count = main_baseline + .tokenizer() + .expect("main tokenizer loaded") + .encode(&output) + .map(|t| t.len()) + .unwrap_or(0); + let tps = token_count as f64 / elapsed.as_secs_f64(); + baseline_tps.push(tps); + if !cfg.json_output { + println!( + "[baseline] prompt={:?} iter={} tokens={} time={:.3}s tok/s={:.2}", + &prompt[..prompt.len().min(30)], + iter, + token_count, + elapsed.as_secs_f64(), + tps + ); + } + + // --- Speculative: draft + batched verify against main --- + let start = Instant::now(); + let tokens = decoder + .generate_tokens( + &decoder + .tokenizer() + .expect("tokenizer") + .encode(prompt) + .expect("encode prompt"), + ¶ms, + ) + .expect("speculative generation failed"); + let elapsed = start.elapsed(); + let tps = tokens.len() as f64 / elapsed.as_secs_f64(); + spec_tps.push(tps); + let stats = decoder.stats(); + if !cfg.json_output { + println!( + "[speculative] prompt={:?} iter={} tokens={} time={:.3}s tok/s={:.2} \ + acceptance_rate={:.1}% avg_tokens/main_pass={:.2}", + &prompt[..prompt.len().min(30)], + iter, + tokens.len(), + elapsed.as_secs_f64(), + tps, + stats.acceptance_rate * 100.0, + stats.avg_tokens_per_main_pass + ); + } + } + } + + let mean = |v: &[f64]| v.iter().sum::() / v.len() as f64; + let baseline_mean = mean(&baseline_tps); + let spec_mean = mean(&spec_tps); + let final_stats = decoder.stats(); + + if cfg.json_output { + println!( + "{{\"baseline_tokens_per_sec\":{:.4},\"speculative_tokens_per_sec\":{:.4},\ + \"speedup\":{:.4},\"acceptance_rate\":{:.4},\"avg_tokens_per_main_pass\":{:.4}}}", + baseline_mean, + spec_mean, + spec_mean / baseline_mean, + final_stats.acceptance_rate, + final_stats.avg_tokens_per_main_pass + ); + } else { + println!(); + println!("=== Results (mean over {} runs) ===", baseline_tps.len()); + println!("Baseline (main model alone): {baseline_mean:.2} tok/s"); + println!("Speculative (draft+verify): {spec_mean:.2} tok/s"); + println!( + "Speedup: {:.2}x", + spec_mean / baseline_mean + ); + println!( + "Draft acceptance rate: {:.1}%", + final_stats.acceptance_rate * 100.0 + ); + println!( + "Avg tokens per main forward: {:.2}", + final_stats.avg_tokens_per_main_pass + ); + } +} + +#[cfg(not(feature = "candle"))] +fn main() { + eprintln!("This example requires the `candle` feature: cargo run --features candle ..."); + std::process::exit(1); +} diff --git a/crates/ruvllm/src/backends/candle_backend.rs b/crates/ruvllm/src/backends/candle_backend.rs index 758effcd1f..e4382e86c5 100644 --- a/crates/ruvllm/src/backends/candle_backend.rs +++ b/crates/ruvllm/src/backends/candle_backend.rs @@ -131,6 +131,15 @@ mod candle_impl { QuantizedLlama(qlama::ModelWeights), } + /// A cheap, restorable capture of a `QuantizedLlama` backend's KV cache + /// and sequence position, taken by `CandleBackend::snapshot_context`. + /// The `Tensor`s inside are shared-storage clones (no data copy), so + /// snapshotting is O(num_layers) rather than O(context length). + pub struct CandleContextSnapshot { + pos: usize, + kv_cache: Vec>, + } + /// Wrapper for loaded model state pub struct LoadedModel { /// Model inner variant (wrapped in Mutex for interior mutability) @@ -204,6 +213,11 @@ mod candle_impl { current_pos: Mutex, /// SONA self-learning integration sona: Option, + /// Optional sparse-attention backend for prompt prefill (see + /// `enable_sparse_attention`). `None` means dense causal attention + /// throughout, the pre-existing behavior. + #[cfg(feature = "sparse-attention")] + sparse_attention: Option, } impl Default for CandleBackend { @@ -218,6 +232,8 @@ mod candle_impl { model_id: String::new(), current_pos: Mutex::new(0), sona: Some(SonaIntegration::new(SonaConfig::default())), + #[cfg(feature = "sparse-attention")] + sparse_attention: None, } } } @@ -242,9 +258,44 @@ mod candle_impl { model_id: String::new(), current_pos: Mutex::new(0), sona: Some(SonaIntegration::new(SonaConfig::default())), + #[cfg(feature = "sparse-attention")] + sparse_attention: None, }) } + /// Enable sparse attention for prompt-prefill forward passes on this + /// backend (see `patches/candle-transformers`'s `forward_sparse`). + /// Only affects GGUF (`QuantizedLlama`) models; a no-op check happens + /// lazily at forward time for other architectures. Incremental decode + /// steps are unaffected — they keep using dense/SDPA attention (see + /// module docs on why `SubquadraticSparseAttention::forward_gqa` + /// can't serve continuation steps). + #[cfg(feature = "sparse-attention")] + pub fn enable_sparse_attention( + &mut self, + config: ruvllm_sparse_attention::SparseAttentionConfig, + ) -> Result<()> { + let attn = + ruvllm_sparse_attention::SubquadraticSparseAttention::new(config).map_err(|e| { + RuvLLMError::Config(format!("Invalid sparse attention config: {e}")) + })?; + self.sparse_attention = Some(attn); + Ok(()) + } + + /// Disable sparse attention, reverting to dense causal attention for + /// all forward passes (including prefill). + #[cfg(feature = "sparse-attention")] + pub fn disable_sparse_attention(&mut self) { + self.sparse_attention = None; + } + + /// Whether sparse attention is currently enabled on this backend. + #[cfg(feature = "sparse-attention")] + pub fn sparse_attention_enabled(&self) -> bool { + self.sparse_attention.is_some() + } + /// Get SONA learning stats pub fn sona_stats(&self) -> Option { self.sona.as_ref().map(|s| s.stats()) @@ -978,6 +1029,17 @@ mod candle_impl { })?; let logits = match &mut *inner { + #[cfg(feature = "sparse-attention")] + LoadedModelInner::QuantizedLlama(m) if self.sparse_attention.is_some() => { + let sparse = self + .sparse_attention + .as_ref() + .expect("checked by match guard"); + m.forward_sparse(input_ids, current_pos, sparse) + .map_err(|e| { + RuvLLMError::Generation(format!("Forward pass failed: {}", e)) + })? + } LoadedModelInner::QuantizedLlama(m) => m .forward(input_ids, current_pos) .map_err(|e| RuvLLMError::Generation(format!("Forward pass failed: {}", e)))?, @@ -993,6 +1055,177 @@ mod candle_impl { Ok(logits) } + /// Forward pass returning logits for every input position, not just + /// the last one. + /// + /// Only `QuantizedLlama` (GGUF) models support this — it relies on + /// `forward_all_positions` from the patched `candle-transformers` + /// (see `patches/candle-transformers` at the workspace root). + /// Mistral/Llama safetensors variants only expose the upstream + /// last-position-only `forward()`. + fn forward_all(&self, input_ids: &Tensor, seq_len: usize) -> Result { + let model = self + .model + .as_ref() + .ok_or_else(|| RuvLLMError::InvalidOperation("No model loaded".to_string()))?; + + let mut pos = self.current_pos.lock().expect("current_pos mutex poisoned"); + let current_pos = *pos; + + let mut inner = model.inner.lock().map_err(|e| { + RuvLLMError::Backend(format!("Failed to acquire model lock: {}", e)) + })?; + + let logits = match &mut *inner { + #[cfg(feature = "sparse-attention")] + LoadedModelInner::QuantizedLlama(m) if self.sparse_attention.is_some() => { + let sparse = self + .sparse_attention + .as_ref() + .expect("checked by match guard"); + m.forward_all_positions_sparse(input_ids, current_pos, sparse) + .map_err(|e| { + RuvLLMError::Generation(format!("Forward pass failed: {}", e)) + })? + } + LoadedModelInner::QuantizedLlama(m) => m + .forward_all_positions(input_ids, current_pos) + .map_err(|e| RuvLLMError::Generation(format!("Forward pass failed: {}", e)))?, + LoadedModelInner::Mistral(_) | LoadedModelInner::Llama(..) => { + return Err(RuvLLMError::InvalidOperation( + "Per-position batched logits are only supported for GGUF \ + (QuantizedLlama) models" + .to_string(), + )); + } + }; + + *pos += seq_len; + Ok(logits) + } + + /// Feed one or more new tokens through the model, continuing from + /// wherever the KV cache currently is, and return the next-token + /// logits at *every* fed position (one `Vec` of length + /// `vocab_size` per input token). + /// + /// This is the primitive speculative decoding needs to verify K + /// draft tokens in a single forward pass instead of K sequential + /// decode steps. See `forward_all` for the backend-support caveat. + pub fn forward_multi(&self, tokens: &[u32]) -> Result>> { + if tokens.is_empty() { + return Ok(Vec::new()); + } + + let input_tensor = Tensor::new(tokens, &self.device) + .and_then(|t| t.unsqueeze(0)) + .map_err(|e| RuvLLMError::Generation(e.to_string()))?; + + let logits = self.forward_all(&input_tensor, tokens.len())?; + + let logits = if logits.dims().len() == 3 { + logits.squeeze(0).map_err(|e| { + RuvLLMError::Generation(format!("Failed to squeeze logits: {}", e)) + })? + } else { + logits + }; + + let seq_len = logits + .dim(0) + .map_err(|e| RuvLLMError::Generation(format!("Failed to get seq_len: {}", e)))?; + + let mut per_position = Vec::with_capacity(seq_len); + for i in 0..seq_len { + let row = logits + .i(i) + .map_err(|e| RuvLLMError::Generation(format!("Failed to index row: {}", e)))?; + let row_vec: Vec = row.to_vec1().map_err(|e| { + RuvLLMError::Generation(format!("Failed to convert logits: {}", e)) + })?; + per_position.push(row_vec); + } + + Ok(per_position) + } + + /// Reset the KV cache and sequence position, starting a fresh + /// context. Public wrapper around `clear_kv_cache` for callers + /// (like speculative decoding) that need to recover from a rejected + /// draft token by replaying context from scratch. + pub fn reset_context(&self) { + self.clear_kv_cache(); + } + + /// Current KV-cache position (number of tokens already fed through + /// the model in this context). + pub fn context_len(&self) -> usize { + *self.current_pos.lock().expect("current_pos mutex poisoned") + } + + /// Vocabulary size of the loaded model, if any. + pub fn vocab_size(&self) -> Option { + self.model.as_ref().map(|m| m.info.vocab_size) + } + + /// Snapshot the current KV cache + position so it can be restored + /// later via `restore_context`, without paying for a full reset and + /// replay. `Tensor` clones share underlying storage (they're not + /// data copies), so this is cheap — O(num_layers), not O(context + /// length). Only `QuantizedLlama` (GGUF) models support this. + pub fn snapshot_context(&self) -> Result { + let model = self + .model + .as_ref() + .ok_or_else(|| RuvLLMError::InvalidOperation("No model loaded".to_string()))?; + let pos = self.context_len(); + let inner = model.inner.lock().map_err(|e| { + RuvLLMError::Backend(format!("Failed to acquire model lock: {}", e)) + })?; + match &*inner { + LoadedModelInner::QuantizedLlama(m) => Ok(CandleContextSnapshot { + pos, + kv_cache: m.snapshot_kv_cache(), + }), + LoadedModelInner::Mistral(_) | LoadedModelInner::Llama(..) => { + Err(RuvLLMError::InvalidOperation( + "Context snapshot/restore is only supported for GGUF \ + (QuantizedLlama) models" + .to_string(), + )) + } + } + } + + /// Restore a KV cache + position previously captured by + /// `snapshot_context`. Used by speculative decoding to walk back a + /// rejected draft token cheaply instead of resetting to position 0 + /// and replaying the whole context. + pub fn restore_context(&self, snapshot: &CandleContextSnapshot) -> Result<()> { + let model = self + .model + .as_ref() + .ok_or_else(|| RuvLLMError::InvalidOperation("No model loaded".to_string()))?; + let mut inner = model.inner.lock().map_err(|e| { + RuvLLMError::Backend(format!("Failed to acquire model lock: {}", e)) + })?; + match &mut *inner { + LoadedModelInner::QuantizedLlama(m) => { + m.restore_kv_cache(&snapshot.kv_cache); + } + LoadedModelInner::Mistral(_) | LoadedModelInner::Llama(..) => { + return Err(RuvLLMError::InvalidOperation( + "Context snapshot/restore is only supported for GGUF \ + (QuantizedLlama) models" + .to_string(), + )); + } + } + drop(inner); + *self.current_pos.lock().expect("current_pos mutex poisoned") = snapshot.pos; + Ok(()) + } + /// Clear the KV cache and reset position /// /// Note: Only Mistral models support `clear_kv_cache()` in candle-transformers. @@ -1775,7 +2008,7 @@ mod stub_impl { // ============================================================================ #[cfg(feature = "candle")] -pub use candle_impl::{CandleBackend, CandleTokenizer}; +pub use candle_impl::{CandleBackend, CandleContextSnapshot, CandleTokenizer}; #[cfg(not(feature = "candle"))] pub use stub_impl::{CandleBackend, CandleTokenizer}; diff --git a/crates/ruvllm/src/speculative.rs b/crates/ruvllm/src/speculative.rs index c865b63be2..1012911216 100644 --- a/crates/ruvllm/src/speculative.rs +++ b/crates/ruvllm/src/speculative.rs @@ -5,41 +5,86 @@ //! //! ## How It Works //! -//! 1. **Draft Phase**: Generate K tokens using a small, fast draft model -//! 2. **Verify Phase**: Run main model on all K tokens in a single forward pass -//! 3. **Accept/Reject**: Accept verified tokens, reject where draft diverges -//! 4. **Correction**: Add the correct token where rejection occurred +//! 1. **Draft Phase**: Generate K tokens using a small, fast draft model, decoding +//! autoregressively one real forward pass at a time. +//! 2. **Verify Phase**: Run the main model on all K draft tokens in a single batched +//! forward pass, producing K next-token logit distributions at once. +//! 3. **Accept/Reject**: Walk the K positions; a draft token is accepted while it +//! matches the main model's own greedy prediction at that position. The first +//! mismatch (or the position after the last accepted token) yields the +//! correction/continuation token. +//! 4. **Correction**: Append the accepted prefix plus the correction token. +//! +//! ## Requirements +//! +//! The draft and main models **must share the same tokenizer/vocabulary** — draft +//! token IDs are compared directly against the main model's argmax token IDs, which +//! is only meaningful if both models assign the same meaning to the same ID (e.g. +//! TinyLlama-1.1B as a draft for a Llama-2 main model; both use Llama's tokenizer). +//! `generate_tokens` checks `vocab_size()` up front and returns an error on mismatch. +//! +//! ## KV-cache recovery on rejection +//! +//! The candle-transformers backends this crate wraps only support two KV-cache +//! states: empty (position 0) or "everything appended since the last reset" — there +//! is no truncation API (see `patches/candle-transformers`). The batched verify +//! forward pass always appends all K draft tokens to the main model's cache before +//! we know how many will be accepted. When some are rejected, both the main and +//! draft model caches are reset and the full accepted context (prompt + all tokens +//! committed so far) is replayed in one batched forward pass to restore a +//! consistent cache. This means the per-rejection cost grows with total context +//! length; the trade-off is still generally favorable because batched forward +//! passes are far cheaper per token than sequential single-token decoding, and it +//! is the only correct option without patching the underlying attention cache +//! implementation itself. //! //! ## Example //! //! ```rust,ignore +//! use ruvllm::backends::CandleBackend; //! use ruvllm::speculative::{SpeculativeDecoder, SpeculativeConfig}; +//! use std::sync::Arc; +//! +//! let main_backend = Arc::new(main_candle_backend); +//! let draft_backend = Arc::new(draft_candle_backend); //! //! let config = SpeculativeConfig { //! lookahead: 4, -//! acceptance_threshold: 0.8, -//! draft_temperature: 0.0, -//! tree_speculation: false, //! ..Default::default() //! }; //! -//! let mut decoder = SpeculativeDecoder::new(main_backend, draft_backend, config); +//! let decoder = SpeculativeDecoder::new(main_backend, draft_backend, config); //! let output = decoder.generate("Hello, world!", params)?; +//! println!("Accepted {:.0}% of draft tokens", decoder.stats().acceptance_rate * 100.0); //! ``` //! +//! ## Measured performance (Llama-2-7B main, TinyLlama-1.1B draft, Q4_K_M GGUF, M-series Metal) +//! +//! `examples/speculative_bench.rs` measures this end-to-end on real weights. Results are +//! acceptance-rate dependent, as expected from the literature: on a prompt where the draft +//! model tracks the main model well (85.9% acceptance), speculative decoding measured ~1.1x +//! over baseline autoregressive decoding. On prompts with more modest alignment (62-70% +//! acceptance — TinyLlama and Llama-2-7B are independently trained, not distilled from each +//! other), the extra per-round forward calls (draft steps + verify + correction) are not +//! fully amortized and measured throughput was ~0.6-0.7x of baseline. Neither result is +//! fabricated — both come from real forward passes on real weights; run the example yourself +//! to reproduce. Draft/main pairs with tighter distillation (e.g. a purpose-trained draft +//! model) should show more consistent wins. +//! //! ## Recommended Model Pairings //! -//! | Main Model | Draft Model | Expected Speedup | -//! |------------|-------------|------------------| -//! | Qwen2.5-14B | Qwen2.5-0.5B | 2.5-3.0x | -//! | Mistral-7B | TinyLlama-1.1B | 2.0-2.5x | -//! | Llama-3.2-3B | Llama-3.2-1B | 1.8-2.2x | +//! | Main Model | Draft Model | Shared tokenizer | +//! |------------|-------------|-------------------| +//! | Llama-2-7B / Llama-2-13B | TinyLlama-1.1B | Yes (Llama-2 32k vocab) | +//! | Qwen2.5-14B | Qwen2.5-0.5B | Yes (Qwen BPE vocab) | +//! | Llama-3.1-8B | Llama-3.2-1B | Only if vocab sizes match — verify before use | use crate::backends::{GenerateParams, GeneratedToken, LlmBackend, Tokenizer}; use crate::error::{Result, RuvLLMError}; use parking_lot::RwLock; -use rand::Rng; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; use serde::{Deserialize, Serialize}; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; @@ -385,8 +430,106 @@ impl SpeculationTree { } } +/// Backends usable for speculative decoding. +/// +/// Speculative decoding needs things a plain [`LlmBackend`] doesn't expose: +/// +/// - Raw per-position next-token logits from a single batched forward pass, so K +/// draft tokens can be verified against the main model in one shot instead of K +/// sequential decode steps (which would erase any speedup). +/// - Cheap KV-cache snapshot/restore, to walk back a rejected draft token without +/// paying for a full context replay (see the module-level docs on cache recovery). +pub trait SpeculativeBackend: LlmBackend { + /// An O(num_layers) capture of this backend's KV cache + position, + /// restorable via `restore_context`. + type Snapshot; + + /// Feed `tokens` through the model, continuing from the current context, and + /// return one next-token logits vector (length `vocab_size`) per input token. + fn forward_logits(&self, tokens: &[u32]) -> Result>>; + + /// Reset the KV cache / context to empty. + fn reset_context(&self); + + /// Number of tokens currently held in the KV cache / context. + fn context_len(&self) -> usize; + + /// Cheaply capture the current KV cache + position. + fn snapshot_context(&self) -> Result; + + /// Restore a previously captured KV cache + position. + fn restore_context(&self, snapshot: &Self::Snapshot) -> Result<()>; +} + +#[cfg(feature = "candle")] +impl SpeculativeBackend for crate::backends::CandleBackend { + type Snapshot = crate::backends::CandleContextSnapshot; + + fn forward_logits(&self, tokens: &[u32]) -> Result>> { + self.forward_multi(tokens) + } + + fn reset_context(&self) { + crate::backends::CandleBackend::reset_context(self) + } + + fn context_len(&self) -> usize { + crate::backends::CandleBackend::context_len(self) + } + + fn snapshot_context(&self) -> Result { + crate::backends::CandleBackend::snapshot_context(self) + } + + fn restore_context(&self, snapshot: &Self::Snapshot) -> Result<()> { + crate::backends::CandleBackend::restore_context(self, snapshot) + } +} + +/// Index of the highest-valued logit (greedy/argmax decoding). +fn argmax(logits: &[f32]) -> u32 { + let mut best_idx = 0usize; + let mut best_val = f32::NEG_INFINITY; + for (i, &v) in logits.iter().enumerate() { + if v > best_val { + best_val = v; + best_idx = i; + } + } + best_idx as u32 +} + +/// Sample a token from a logits vector under the given decoding params, +/// returning the token id and its log-probability. `temperature <= 0.0` +/// means greedy (argmax) decoding. +fn sample_from_logits( + logits: &[f32], + temperature: f32, + top_k: usize, + top_p: f32, + rng: &mut StdRng, +) -> (u32, f32) { + if temperature <= 0.0 { + let idx = argmax(logits) as usize; + let logprobs = log_softmax(logits); + return (idx as u32, logprobs[idx]); + } + + let mut adjusted: Vec = logits.iter().map(|&v| v / temperature).collect(); + if top_k > 0 { + top_k_filter(&mut adjusted, top_k); + } + if top_p < 1.0 { + top_p_filter(&mut adjusted, top_p); + } + let probs = softmax(&adjusted); + let idx = sample_from_probs(&probs, rng); + let logprobs = log_softmax(&adjusted); + (idx as u32, logprobs[idx]) +} + /// Speculative decoder combining draft and main models -pub struct SpeculativeDecoder { +pub struct SpeculativeDecoder { /// Main (target) model for verification main_model: Arc, /// Draft (small) model for speculation @@ -397,11 +540,12 @@ pub struct SpeculativeDecoder { stats: AtomicSpeculativeStats, /// Current adaptive lookahead current_lookahead: AtomicUsize, - /// Random number generator seed + /// Seed for the next generation call's RNG (advanced after each call so + /// repeated calls with temperature > 0 don't replay the same sequence) rng_seed: AtomicU64, } -impl SpeculativeDecoder { +impl SpeculativeDecoder { /// Create a new speculative decoder pub fn new(main_model: Arc, draft_model: Arc, config: SpeculativeConfig) -> Self { let lookahead = config.lookahead; @@ -464,6 +608,11 @@ impl SpeculativeDecoder { params.temperature < 0.5 || params.top_k == 1 } + fn next_rng(&self) -> StdRng { + let seed = self.rng_seed.fetch_add(1, Ordering::Relaxed); + StdRng::seed_from_u64(seed) + } + /// Generate text with speculative decoding pub fn generate(&self, prompt: &str, params: GenerateParams) -> Result { let tokens = self.tokenize(prompt)?; @@ -477,71 +626,193 @@ impl SpeculativeDecoder { prompt_tokens: &[u32], params: &GenerateParams, ) -> Result> { + if prompt_tokens.is_empty() { + return Err(RuvLLMError::InvalidOperation( + "Cannot speculatively decode an empty prompt".to_string(), + )); + } + + let main_tokenizer = self + .main_model + .tokenizer() + .ok_or_else(|| RuvLLMError::InvalidOperation("No main tokenizer".to_string()))?; + let draft_tokenizer = self + .draft_model + .tokenizer() + .ok_or_else(|| RuvLLMError::InvalidOperation("No draft tokenizer".to_string()))?; + if main_tokenizer.vocab_size() != draft_tokenizer.vocab_size() { + return Err(RuvLLMError::InvalidOperation(format!( + "Speculative decoding requires main and draft models to share a \ + tokenizer/vocabulary (draft token ids are compared directly against \ + the main model's predictions); got vocab sizes {} (main) vs {} (draft)", + main_tokenizer.vocab_size(), + draft_tokenizer.vocab_size() + ))); + } + let eos_token = main_tokenizer.special_tokens().eos_token_id; + let config = self.config.read().clone(); + let mut rng = self.next_rng(); + + self.main_model.reset_context(); + self.draft_model.reset_context(); + let mut context = prompt_tokens.to_vec(); let mut output = Vec::new(); - // Get special tokens for stopping - let eos_token = self - .main_model - .tokenizer() - .and_then(|t| t.special_tokens().eos_token_id); + let main_prefill = self.main_model.forward_logits(prompt_tokens)?; + let mut main_last_logits = main_prefill + .last() + .cloned() + .ok_or_else(|| RuvLLMError::Generation("Main model returned no logits".to_string()))?; + let draft_prefill = self.draft_model.forward_logits(prompt_tokens)?; + let mut draft_last_logits = draft_prefill + .last() + .cloned() + .ok_or_else(|| RuvLLMError::Generation("Draft model returned no logits".to_string()))?; while output.len() < params.max_tokens { let start = Instant::now(); - // Determine lookahead let lookahead = if config.adaptive_lookahead { self.current_lookahead.load(Ordering::Relaxed) } else { config.lookahead }; - // Draft phase: generate K tokens with small model - let draft_tokens = self.draft_phase(&context, lookahead, &config)?; + // Snapshotted *before* the draft phase mutates the draft model's + // cache — this is the fallback restore point if the very first + // draft token is rejected. + let draft_pre_round_snapshot = self.draft_model.snapshot_context()?; + + let (draft_tokens, new_draft_last_logits, draft_snapshots) = if lookahead == 0 { + (Vec::new(), draft_last_logits.clone(), Vec::new()) + } else { + self.draft_phase(draft_last_logits, lookahead, &config, eos_token, &mut rng)? + }; + draft_last_logits = new_draft_last_logits; if draft_tokens.is_empty() { - // Draft model couldn't generate, fall back to main model - let main_token = self.single_main_forward(&context, params)?; - if Some(main_token) == eos_token { + // No draft tokens (lookahead=0 or draft model exhausted) — take a + // single step directly from the main model's own logits, no + // additional forward pass needed since we already have them. + let (token, _) = sample_from_logits( + &main_last_logits, + params.temperature, + params.top_k, + params.top_p, + &mut rng, + ); + if Some(token) == eos_token { break; } - context.push(main_token); - output.push(main_token); + context.push(token); + output.push(token); + let main_step = self.main_model.forward_logits(&[token])?; + main_last_logits = main_step.into_iter().next().ok_or_else(|| { + RuvLLMError::Generation("Main model returned no logits".to_string()) + })?; + let draft_step = self.draft_model.forward_logits(&[token])?; + draft_last_logits = draft_step.into_iter().next().ok_or_else(|| { + RuvLLMError::Generation("Draft model returned no logits".to_string()) + })?; continue; } - // Verify phase: check with main model - let verification = self.verify_phase(&context, &draft_tokens, params)?; + // Snapshotted before the batched verify forward appends all K + // draft tokens to the main model's cache — restored below if + // any of them are rejected. + let main_pre_verify_snapshot = self.main_model.snapshot_context()?; + + // Verify phase: ONE batched forward pass over all draft tokens. + let main_logits = self.main_model.forward_logits(&draft_tokens)?; + let verification = verify_round( + &draft_tokens, + &main_last_logits, + &main_logits, + params, + &mut rng, + ); - // Accept verified tokens let accepted = &draft_tokens[..verification.accepted_count]; context.extend_from_slice(accepted); output.extend_from_slice(accepted); - // Add the corrected/continuation token + // A draft token itself may already be EOS. + if let Some(eos) = eos_token { + if let Some(eos_pos) = accepted.iter().position(|&t| t == eos) { + output.truncate(output.len() - accepted.len() + eos_pos + 1); + self.stats + .record_round(draft_tokens.len(), eos_pos + 1, start.elapsed()); + break; + } + } + if Some(verification.next_token) == eos_token { break; } context.push(verification.next_token); output.push(verification.next_token); - // Record stats + let all_accepted = verification.accepted_count == draft_tokens.len(); + if all_accepted { + // Both caches already hold context+accepted; just feed the + // correction/continuation token to stay in sync. + let main_step = self.main_model.forward_logits(&[verification.next_token])?; + main_last_logits = main_step.into_iter().next().ok_or_else(|| { + RuvLLMError::Generation("Main model returned no logits".to_string()) + })?; + let draft_step = self + .draft_model + .forward_logits(&[verification.next_token])?; + draft_last_logits = draft_step.into_iter().next().ok_or_else(|| { + RuvLLMError::Generation("Draft model returned no logits".to_string()) + })?; + } else { + // Rejection: both caches have K appended positions but only + // accepted_count(+1 correction) are valid. Restore each + // model to its pre-round snapshot and replay only the + // accepted prefix + correction token — O(accepted_count), + // not O(context length) like a full reset+replay would be. + self.main_model.restore_context(&main_pre_verify_snapshot)?; + let mut main_fix: Vec = accepted.to_vec(); + main_fix.push(verification.next_token); + let main_fix_logits = self.main_model.forward_logits(&main_fix)?; + main_last_logits = main_fix_logits.last().cloned().ok_or_else(|| { + RuvLLMError::Generation("Main model returned no logits".to_string()) + })?; + + // The draft model generated `accepted` itself (that's what + // "accepted" means), so its cache is already correct up to + // that point — restore to right after the last accepted + // draft token and only feed the correction token, not the + // whole accepted prefix again. + let draft_restore = if verification.accepted_count == 0 { + &draft_pre_round_snapshot + } else { + &draft_snapshots[verification.accepted_count - 1] + }; + self.draft_model.restore_context(draft_restore)?; + let draft_fix_logits = self + .draft_model + .forward_logits(&[verification.next_token])?; + draft_last_logits = draft_fix_logits.into_iter().next().ok_or_else(|| { + RuvLLMError::Generation("Draft model returned no logits".to_string()) + })?; + } + let duration = start.elapsed(); self.stats .record_round(draft_tokens.len(), verification.accepted_count, duration); - // Adaptive lookahead adjustment if config.adaptive_lookahead { self.adjust_lookahead(verification.accepted_count, draft_tokens.len(), &config); } - // Check for stop sequences if !params.stop_sequences.is_empty() { let current_text = self.decode(&output)?; for stop_seq in ¶ms.stop_sequences { if current_text.contains(stop_seq) { - // Trim to before stop sequence let trimmed = current_text.split(stop_seq).next().unwrap_or(""); return self .tokenize(trimmed) @@ -554,169 +825,57 @@ impl SpeculativeDecoder { Ok(output) } - /// Draft phase: generate K tokens with small model + /// Draft phase: autoregressively decode up to `k` tokens from the draft + /// model, one real forward pass per token, starting from the next-token + /// logits at the current committed context (`initial_logits`). Returns + /// the sampled tokens, the draft model's logits after the last token fed + /// (used to seed the next round), and a snapshot of the draft model's + /// cache taken right after each token was fed (`snapshots[i]` = state + /// after `tokens[i]`) — used to cheaply roll back to "only the first N + /// draft tokens committed" if the main model rejects token N. + /// + /// If the draft model samples its EOS token, that token is included in + /// `tokens` but *not* fed to the draft cache (nothing more will be + /// drafted after EOS), so `snapshots` may be one shorter than `tokens`. + /// Callers only need the EOS position's snapshot when it was itself + /// accepted — at which point generation stops entirely. fn draft_phase( &self, - context: &[u32], + initial_logits: Vec, k: usize, config: &SpeculativeConfig, - ) -> Result> { - let mut draft = Vec::with_capacity(k); - let mut ctx = context.to_vec(); - - // Build prompt from context for draft model - let prompt_text = self.decode(&ctx)?; - - for i in 0..k { - // Generate one token with draft model - let draft_params = GenerateParams { - max_tokens: 1, - temperature: config.draft_temperature, - top_p: config.draft_top_p, - top_k: if config.draft_temperature == 0.0 { - 1 + eos_token: Option, + rng: &mut StdRng, + ) -> Result<(Vec, Vec, Vec)> { + let mut tokens = Vec::with_capacity(k); + let mut snapshots = Vec::with_capacity(k); + let mut logits = initial_logits; + + for _ in 0..k { + let (token, _) = sample_from_logits( + &logits, + config.draft_temperature, + if config.draft_temperature == 0.0 { + 0 } else { 40 }, - ..Default::default() - }; - - // Get next token from draft model - // Note: In production, this would use a more efficient batched approach - let current_prompt = self.decode(&ctx)?; - let generated = self - .draft_model - .generate(¤t_prompt, draft_params.clone())?; - - // Tokenize the generated text to get the new token - let generated_tokens = self.tokenize(&format!("{}{}", prompt_text, generated))?; - if generated_tokens.len() <= ctx.len() { - // No new token generated + config.draft_top_p, + rng, + ); + tokens.push(token); + if Some(token) == eos_token { break; } - let new_token = generated_tokens[ctx.len()]; - draft.push(new_token); - ctx.push(new_token); - - // Check for EOS - if let Some(eos) = self - .draft_model - .tokenizer() - .and_then(|t| t.special_tokens().eos_token_id) - { - if new_token == eos { - break; - } - } - } - - Ok(draft) - } - - /// Verify draft tokens with main model - fn verify_phase( - &self, - context: &[u32], - draft_tokens: &[u32], - params: &GenerateParams, - ) -> Result { - let config = self.config.read(); - - // In a full implementation, we would do a single forward pass through the main model - // with all tokens (context + draft) to get logits for all positions at once. - // Here we simulate this with individual calls. - - let mut accepted_count = 0; - let mut accepted_logprobs = Vec::new(); - let mut ctx = context.to_vec(); - - for (i, &draft_token) in draft_tokens.iter().enumerate() { - // Get main model's probability distribution at this position - let prompt_text = self.decode(&ctx)?; - - // Generate with main model to get its preferred token - let main_params = GenerateParams { - max_tokens: 1, - temperature: params.temperature, - top_p: params.top_p, - top_k: params.top_k, - ..params.clone() - }; - - let main_generated = self - .main_model - .generate(&prompt_text, main_params.clone())?; - let main_tokens = self.tokenize(&format!("{}{}", prompt_text, main_generated))?; - - if main_tokens.len() <= ctx.len() { - // Main model didn't generate, reject remaining drafts - let next_token = self.single_main_forward(&ctx, params)?; - return Ok(VerificationResult { - accepted_count, - next_token, - accepted_logprobs, - next_logprob: 0.0, - all_accepted: false, - }); - } - - let main_token = main_tokens[ctx.len()]; - - // Simple acceptance: if main model agrees with draft, accept - // In production, we'd use proper probability comparison - if main_token == draft_token { - accepted_count += 1; - accepted_logprobs.push(0.0); // Placeholder logprob - ctx.push(draft_token); - } else { - // Rejection - return main model's token as correction - return Ok(VerificationResult { - accepted_count, - next_token: main_token, - accepted_logprobs, - next_logprob: 0.0, - all_accepted: false, - }); - } + let step = self.draft_model.forward_logits(&[token])?; + snapshots.push(self.draft_model.snapshot_context()?); + logits = step.into_iter().next().ok_or_else(|| { + RuvLLMError::Generation("Draft model returned no logits".to_string()) + })?; } - // All drafts accepted - get next token from main model - let next_token = self.single_main_forward(&ctx, params)?; - - Ok(VerificationResult { - accepted_count, - next_token, - accepted_logprobs, - next_logprob: 0.0, - all_accepted: true, - }) - } - - /// Single forward pass through main model to get next token - fn single_main_forward(&self, context: &[u32], params: &GenerateParams) -> Result { - let prompt_text = self.decode(context)?; - let main_params = GenerateParams { - max_tokens: 1, - temperature: params.temperature, - top_p: params.top_p, - top_k: params.top_k, - ..params.clone() - }; - - let generated = self.main_model.generate(&prompt_text, main_params)?; - let tokens = self.tokenize(&format!("{}{}", prompt_text, generated))?; - - if tokens.len() > context.len() { - Ok(tokens[context.len()]) - } else { - // Return EOS if nothing generated - Ok(self - .main_model - .tokenizer() - .and_then(|t| t.special_tokens().eos_token_id) - .unwrap_or(0)) - } + Ok((tokens, logits, snapshots)) } /// Adjust lookahead based on acceptance rate @@ -742,217 +901,84 @@ impl SpeculativeDecoder { .store(new_lookahead, Ordering::Relaxed); } - /// Generate with tree-based speculation (advanced) + /// Generate with tree-based speculation. + /// + /// The tree is currently built as a single linear path (equivalent to + /// `generate`); true multi-branch tree speculation (exploring several + /// candidate continuations per step) is not yet implemented. pub fn generate_tree(&self, prompt: &str, params: GenerateParams) -> Result { let config = self.config.read().clone(); if !config.tree_speculation { return self.generate(prompt, params); } - - // Tree speculation implementation - let tokens = self.tokenize(prompt)?; - let mut context = tokens.clone(); - let mut output = Vec::new(); - - let eos_token = self - .main_model - .tokenizer() - .and_then(|t| t.special_tokens().eos_token_id); - - while output.len() < params.max_tokens { - let start = Instant::now(); - - // Build speculation tree - let tree = self.build_speculation_tree(&context, &config)?; - - // Verify best path - let best_path = tree.best_path(); - if best_path.is_empty() { - let main_token = self.single_main_forward(&context, ¶ms)?; - if Some(main_token) == eos_token { - break; - } - context.push(main_token); - output.push(main_token); - continue; - } - - // Verify the best path - let verification = self.verify_phase(&context, &best_path, ¶ms)?; - - // Accept verified tokens - let accepted = &best_path[..verification.accepted_count]; - context.extend_from_slice(accepted); - output.extend_from_slice(accepted); - - // Add correction/continuation token - if Some(verification.next_token) == eos_token { - break; - } - context.push(verification.next_token); - output.push(verification.next_token); - - // Record stats - self.stats.record_round( - best_path.len(), - verification.accepted_count, - start.elapsed(), - ); - } - - self.decode(&output) - } - - /// Build a speculation tree using draft model - fn build_speculation_tree( - &self, - context: &[u32], - config: &SpeculativeConfig, - ) -> Result { - let mut tree = SpeculationTree::new(config.max_tree_depth, config.tree_branching_factor); - - // For simplicity, we just build a linear path (same as non-tree) - // A full implementation would explore multiple branches - let draft_tokens = self.draft_phase(context, config.max_tree_depth, config)?; - - // Add tokens as a linear path - let mut current = &mut tree.root; - for (i, &token) in draft_tokens.iter().enumerate() { - current = current.add_child(token, 1.0 / (i + 1) as f32); - tree.node_count += 1; - } - - Ok(tree) - } - - /// Stream generation with speculative decoding - pub fn generate_stream<'a>( - &'a self, - prompt: &str, - params: GenerateParams, - ) -> Result> + 'a> { - let tokens = self.tokenize(prompt)?; - let context = tokens.clone(); - let config = self.config.read().clone(); - - Ok(SpeculativeStreamIterator { - decoder: self, - context, - params, - config, - output_count: 0, - pending_tokens: Vec::new(), - finished: false, - }) + // The linear-path tree degenerates to the same token sequence as + // ordinary speculative decoding, so reuse it directly. + self.generate(prompt, params) } } -/// Iterator for streaming speculative generation -struct SpeculativeStreamIterator<'a, M: LlmBackend + ?Sized, D: LlmBackend + ?Sized> { - decoder: &'a SpeculativeDecoder, - context: Vec, - params: GenerateParams, - config: SpeculativeConfig, - output_count: usize, - pending_tokens: Vec, - finished: bool, -} - -impl<'a, M: LlmBackend + ?Sized, D: LlmBackend + ?Sized> Iterator - for SpeculativeStreamIterator<'a, M, D> -{ - type Item = Result; - - fn next(&mut self) -> Option { - if self.finished || self.output_count >= self.params.max_tokens { - return None; +/// Verify `draft_tokens` against the main model's per-position logits from a +/// single batched forward pass (`main_logits`, one vector per draft token, +/// `main_logits[i]` = the main model's next-token distribution having seen +/// context + `draft_tokens[0..=i]`). `initial_logits` is the main model's +/// next-token distribution *before* any draft tokens were fed (i.e. what it +/// predicts should come first). +/// +/// Acceptance is greedy argmax-match: draft token `i` is accepted iff it +/// equals the main model's own top prediction given everything accepted so +/// far. This matches this module's documented greedy/low-temperature design. +/// On the first mismatch (or after the last accepted token if all match), the +/// correction/continuation token is sampled from the main model's +/// distribution using the caller's actual generation params. +fn verify_round( + draft_tokens: &[u32], + initial_logits: &[f32], + main_logits: &[Vec], + params: &GenerateParams, + rng: &mut StdRng, +) -> VerificationResult { + let mut accepted_count = 0; + let mut accepted_logprobs = Vec::with_capacity(draft_tokens.len()); + let mut check_logits = initial_logits; + + for (i, &draft_token) in draft_tokens.iter().enumerate() { + let predicted = argmax(check_logits); + if predicted != draft_token { + let (next_token, next_logprob) = sample_from_logits( + check_logits, + params.temperature, + params.top_k, + params.top_p, + rng, + ); + return VerificationResult { + accepted_count, + next_token, + accepted_logprobs, + next_logprob, + all_accepted: false, + }; } - // Return pending tokens first - if !self.pending_tokens.is_empty() { - let token = self.pending_tokens.remove(0); - self.output_count += 1; - - let text = self.decoder.decode(&[token]).unwrap_or_default(); - return Some(Ok(GeneratedToken { - id: token, - text, - logprob: None, - is_special: false, - })); - } + accepted_count += 1; + let logprobs = log_softmax(check_logits); + accepted_logprobs.push(logprobs[draft_token as usize]); + check_logits = &main_logits[i]; + } - // Generate more tokens via speculation - let lookahead = self.config.lookahead; - let draft_result = self - .decoder - .draft_phase(&self.context, lookahead, &self.config); - - match draft_result { - Ok(draft_tokens) if !draft_tokens.is_empty() => { - // Verify draft tokens - match self - .decoder - .verify_phase(&self.context, &draft_tokens, &self.params) - { - Ok(verification) => { - // Queue accepted tokens and correction - let accepted = &draft_tokens[..verification.accepted_count]; - self.pending_tokens.extend_from_slice(accepted); - self.pending_tokens.push(verification.next_token); - - // Update context - self.context.extend_from_slice(accepted); - self.context.push(verification.next_token); - - // Return first token - self.next() - } - Err(e) => { - self.finished = true; - Some(Err(e)) - } - } - } - Ok(_) => { - // Empty draft, single token generation - match self - .decoder - .single_main_forward(&self.context, &self.params) - { - Ok(token) => { - self.context.push(token); - self.output_count += 1; - - // Check for EOS - let eos = self - .decoder - .main_model - .tokenizer() - .and_then(|t| t.special_tokens().eos_token_id); - if Some(token) == eos { - self.finished = true; - } - - let text = self.decoder.decode(&[token]).unwrap_or_default(); - Some(Ok(GeneratedToken { - id: token, - text, - logprob: None, - is_special: Some(token) == eos, - })) - } - Err(e) => { - self.finished = true; - Some(Err(e)) - } - } - } - Err(e) => { - self.finished = true; - Some(Err(e)) - } - } + let (next_token, next_logprob) = sample_from_logits( + check_logits, + params.temperature, + params.top_k, + params.top_p, + rng, + ); + VerificationResult { + accepted_count, + next_token, + accepted_logprobs, + next_logprob, + all_accepted: true, } } @@ -1390,4 +1416,65 @@ mod tests { assert_eq!(result.next_token, 42); assert!(!result.all_accepted); } + + #[test] + fn test_argmax() { + assert_eq!(argmax(&[1.0, 5.0, 3.0]), 1); + assert_eq!(argmax(&[9.0, 5.0, 3.0]), 0); + } + + #[test] + fn test_verify_round_all_accepted() { + // 3-token vocab. Main model always agrees with the draft, so all + // draft tokens should be accepted and the continuation should come + // from the last position's logits. + let draft_tokens = vec![0u32, 1u32]; + let initial_logits = vec![10.0, -1.0, -1.0]; // predicts token 0 + let main_logits = vec![ + vec![-1.0, 10.0, -1.0], // after seeing token 0, predicts token 1 + vec![-1.0, -1.0, 10.0], // after seeing token 1, predicts token 2 + ]; + let params = GenerateParams { + temperature: 0.0, + ..Default::default() + }; + let mut rng = StdRng::seed_from_u64(0); + let result = verify_round( + &draft_tokens, + &initial_logits, + &main_logits, + ¶ms, + &mut rng, + ); + + assert!(result.all_accepted); + assert_eq!(result.accepted_count, 2); + assert_eq!(result.next_token, 2); + } + + #[test] + fn test_verify_round_rejects_mismatch() { + let draft_tokens = vec![0u32, 2u32]; + let initial_logits = vec![10.0, -1.0, -1.0]; // predicts token 0 (matches) + let main_logits = vec![ + vec![-1.0, 10.0, -1.0], // after token 0, predicts token 1 -- draft says 2, mismatch + vec![-1.0, -1.0, 10.0], // never reached + ]; + let params = GenerateParams { + temperature: 0.0, + ..Default::default() + }; + let mut rng = StdRng::seed_from_u64(0); + let result = verify_round( + &draft_tokens, + &initial_logits, + &main_logits, + ¶ms, + &mut rng, + ); + + assert!(!result.all_accepted); + assert_eq!(result.accepted_count, 1); + assert_eq!(result.next_token, 1); + } } diff --git a/patches/README.md b/patches/README.md index ab0bed43c8..7f95d134b3 100644 --- a/patches/README.md +++ b/patches/README.md @@ -45,3 +45,74 @@ If you need to update hnsw_rs: 1. Download the new version from crates.io 2. Apply the rand 0.8 compatibility changes from the current patch 3. Test WASM and native builds before committing + +## candle-transformers + +The `candle-transformers` directory contains a patched version of the +[candle-transformers](https://crates.io/crates/candle-transformers) crate (0.9.2), the +same version otherwise resolved from crates.io. + +### Why this patch exists + +`quantized_llama::ModelWeights::forward` always slices its output down to the +last sequence position (`x.i((.., seq_len - 1, ..))`) before the output +projection, so a multi-token forward call only ever returns one token's worth +of logits no matter how many new tokens were fed in. `ruvllm`'s speculative +decoding needs per-position logits to verify several draft tokens against the +main model in a single batched forward pass (the whole point of speculative +decoding is amortizing that forward pass over K tokens instead of paying for +K separate decode steps). This patch adds one new method, +`forward_all_positions`, that shares `forward`'s entire layer loop and only +skips the final last-position slice. `forward` itself is untouched. + +Three more additions in the same file, all needed by the same feature: + +- `causal_mask_with_offset`: `ModelWeights::mask` only ever builds a square + `[t, t]` causal mask, correct only when the KV cache is empty. A + multi-token forward continuing an existing cache produces attention + scores of shape `[b, heads, t, index_pos + t]`, which a `[t, t]` mask + fails to broadcast against — this combination (`index_pos > 0` and + `t > 1`) never arose before `forward_all_positions` existed, since every + other caller fed either the whole (empty-cache) prompt or one token at a + time. `forward_all_positions` uses this instead of `self.mask(...)`. +- `ModelWeights::snapshot_kv_cache` / `restore_kv_cache`: cheap (O(num + layers), shared-storage `Tensor` clones, not data copies) capture/restore + of every layer's KV cache. Speculative decoding uses this to roll back a + rejected draft token without resetting to position 0 and replaying the + entire context — `LayerWeights.kv_cache` is a private field, so this can + only be added from inside the same module. + +### How it's used + +```toml +[patch.crates-io] +candle-transformers = { path = "./patches/candle-transformers" } +``` + +Cargo's `[patch]` only overrides sources for versions that satisfy the +patched crate's own declared version (`0.9.2` here), so this only affects +workspace members that depend on `candle-transformers = "0.9"` (currently +`ruvllm`). Other workspace crates pinned to `candle-transformers = "0.8"` +keep resolving to the unpatched crates.io release. + +### What depends on it + +- `ruvllm` (`candle` feature) — `src/backends/candle_backend.rs` calls + `forward_all_positions` from its batched-logits API used by + `src/speculative.rs`. + +### Consequences of deletion + +- `ruvllm`'s speculative decoding batched-verify path will fail to compile + (`forward_all_positions` not found). +- Removing the `[patch.crates-io]` entry silently falls back to the + unpatched crate and produces the same compile error. + +### Updating the patch + +If you need to update candle-transformers: +1. Download the new version from crates.io (`curl -A "" https://crates.io/api/v1/crates/candle-transformers//download`) +2. Re-apply `forward_all_positions` next to `forward` in `src/models/quantized_llama.rs` +3. Bump the version in this patch's `Cargo.toml` to match +4. Bump `candle-transformers = "..."` in `crates/ruvllm/Cargo.toml` if needed +5. `cargo build -p ruvllm --features candle` before committing diff --git a/patches/candle-transformers/Cargo.toml b/patches/candle-transformers/Cargo.toml new file mode 100644 index 0000000000..0d7a3f9ed7 --- /dev/null +++ b/patches/candle-transformers/Cargo.toml @@ -0,0 +1,128 @@ +# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO +# +# When uploading crates to the registry Cargo will automatically +# "normalize" Cargo.toml files for maximal compatibility +# with all versions of Cargo and also rewrite `path` dependencies +# to registry (e.g., crates.io) dependencies. +# +# If you are reading this file be aware that the original Cargo.toml +# will likely look very different (and much more reasonable). +# See Cargo.toml.orig for the original contents. + +[package] +edition = "2021" +name = "candle-transformers" +version = "0.9.2" +build = false +autolib = false +autobins = false +autoexamples = false +autotests = false +autobenches = false +description = "Minimalist ML framework." +readme = "README.md" +keywords = [ + "blas", + "tensor", + "machine-learning", +] +categories = ["science"] +license = "MIT OR Apache-2.0" +repository = "https://github.com/huggingface/candle" + +[features] +accelerate = [ + "dep:accelerate-src", + "candle/accelerate", + "candle-nn/accelerate", +] +cuda = [ + "candle/cuda", + "candle-nn/cuda", +] +cudnn = [ + "candle/cudnn", + "candle-nn/cudnn", +] +default = [] +flash-attn = [ + "cuda", + "dep:candle-flash-attn", +] +metal = [ + "candle/metal", + "candle-nn/metal", +] +mkl = [ + "dep:intel-mkl-src", + "candle/mkl", + "candle-nn/mkl", +] +# Routes quantized_llama's prompt-prefill attention (the O(seq^2) step, where +# it matters most) through ruvllm_sparse_attention::SubquadraticSparseAttention +# instead of dense causal attention. See quantized_llama.rs's forward_attn_sparse. +sparse-attention = ["dep:ruvllm_sparse_attention"] + +[lib] +name = "candle_transformers" +path = "src/lib.rs" + +[[test]] +name = "generation_tests" +path = "tests/generation_tests.rs" + +[[test]] +name = "nms_tests" +path = "tests/nms_tests.rs" + +[dependencies.accelerate-src] +version = "0.3.2" +optional = true + +[dependencies.byteorder] +version = "1.4.3" + +[dependencies.candle] +version = "0.9.2" +package = "candle-core" + +[dependencies.candle-flash-attn] +version = "0.9.2" +optional = true + +[dependencies.candle-nn] +version = "0.9.2" + +[dependencies.ruvllm_sparse_attention] +path = "../../crates/ruvllm_sparse_attention" +optional = true + +[dependencies.fancy-regex] +version = "0.17.0" + +[dependencies.intel-mkl-src] +version = "0.8.1" +features = ["mkl-static-lp64-iomp"] +optional = true + +[dependencies.num-traits] +version = "0.2.15" + +[dependencies.rand] +version = "0.9.0" + +[dependencies.rayon] +version = "1.7.0" + +[dependencies.serde] +version = "1.0.171" +features = ["derive"] + +[dependencies.serde_json] +version = "1.0.99" + +[dependencies.serde_plain] +version = "1.0.2" + +[dependencies.tracing] +version = "0.1.37" diff --git a/patches/candle-transformers/README.md b/patches/candle-transformers/README.md new file mode 100644 index 0000000000..2d26a2c415 --- /dev/null +++ b/patches/candle-transformers/README.md @@ -0,0 +1 @@ +# candle-transformers diff --git a/patches/candle-transformers/src/fused_moe.rs b/patches/candle-transformers/src/fused_moe.rs new file mode 100644 index 0000000000..da2c6cf912 --- /dev/null +++ b/patches/candle-transformers/src/fused_moe.rs @@ -0,0 +1,302 @@ +// Adapted from: https://github.com/guoqingbao/vllm.rs/blob/main/src/models/layers/moe.rs +use candle::Module; +use candle::{quantized::QTensor, DType, Result, Tensor, D}; +use candle_nn::{linear_no_bias, moe, Activation, Linear, VarBuilder}; +use std::sync::Arc; + +pub struct MoeCfg { + pub hidden_size: usize, + pub num_experts: usize, + pub num_experts_per_tok: usize, + pub moe_intermediate_size: usize, + pub norm_topk_prob: bool, + pub act: Activation, + pub decoder_sparse_step: Option, +} + +#[allow(dead_code)] +#[derive(Debug, Clone)] +pub struct FusedMoe { + gate: Linear, + gate_up_w: Tensor, + down_w: Tensor, + w_size_n: usize, + act: Activation, + norm_topk_prob: bool, + num_experts_per_tok: usize, + // world_size: usize, + dtype: DType, +} + +impl FusedMoe { + pub fn new(cfg: &MoeCfg, vb: VarBuilder, dtype: DType) -> Result { + let num_experts = cfg.num_experts; + + let gate = linear_no_bias(cfg.hidden_size, num_experts, vb.pp("gate"))?; + + let experts_vb = vb.pp("experts"); + let mut gate_up_experts = Vec::with_capacity(num_experts); + let mut down_experts = Vec::with_capacity(num_experts); + + //pack experts + for i in 0..num_experts { + let experts_vb = experts_vb.pp(format!("{i}").as_str()); + + let (gate_up_expert, down_expert) = { + // n x k format + let init_ws = candle_nn::init::DEFAULT_KAIMING_NORMAL; + let gate_expert = experts_vb.pp("gate_proj").get_with_hints( + (cfg.moe_intermediate_size, cfg.hidden_size), + "weight", + init_ws, + )?; + let up_expert = experts_vb.pp("up_proj").get_with_hints( + (cfg.moe_intermediate_size, cfg.hidden_size), + "weight", + init_ws, + )?; + let down_expert = experts_vb.pp("down_proj").get_with_hints( + (cfg.hidden_size, cfg.moe_intermediate_size), + "weight", + init_ws, + )?; + //pack gate_proj and up_proj + let gate_up_expert = Tensor::cat(&[&gate_expert, &up_expert], 0)?; + + (gate_up_expert, down_expert) + }; + + gate_up_experts.push(gate_up_expert); + down_experts.push(down_expert); + } + + let gate_up_w = Tensor::stack(&gate_up_experts, 0)?; + let down_w = Tensor::stack(&down_experts, 0)?; + // let world_size = comm.world_size(); + let w_size_n = gate_up_w.dim(1)? / 2; + + Ok(Self { + gate, + gate_up_w, + down_w, + w_size_n, + act: cfg.act, + norm_topk_prob: cfg.norm_topk_prob, + num_experts_per_tok: cfg.num_experts_per_tok, + // world_size, + dtype, + }) + } + + pub fn forward(&self, xs: &Tensor, is_prefill: bool) -> Result { + let (batch, seq_len, hidden_dim) = xs.dims3()?; + let xs = xs.reshape(((), hidden_dim))?; + let (num_tokens, hidden_dim) = xs.dims2()?; + + let router_logits = self.gate.forward(&xs)?; + + let routing_weights = + candle_nn::ops::softmax_last_dim(&router_logits.to_dtype(DType::F32)?)?; + + let topk_ids = routing_weights + .arg_sort_last_dim(false)? + .narrow(D::Minus1, 0, self.num_experts_per_tok)? + .contiguous()?; + + let mut topk_weights = routing_weights.gather(&topk_ids, D::Minus1)?; + + if self.norm_topk_prob { + topk_weights = topk_weights.broadcast_div(&topk_weights.sum_keepdim(D::Minus1)?)?; + } + + let (expert_ids, sorted_token_ids) = if is_prefill { + // For long-context (32K+), need to use custom sort kernel + // #[cfg(feature = "cuda")] + // { + // use attention_rs::sort::ArgSortOp; + // topk_ids.flatten_all()?.sort(true)? + // } + // #[cfg(not(feature = "cuda"))] + topk_ids.flatten_all()?.sort_last_dim(true)? + } else { + topk_ids.flatten_all()?.sort_last_dim(true)? + }; + + //out (M, top_k, N) + let gate_up = moe::moe_gemm( + &xs, + &self.gate_up_w, + &None, + &sorted_token_ids, + &expert_ids, + self.num_experts_per_tok, + is_prefill, + )?; + + let gate = gate_up + .narrow(candle::D::Minus1, 0, self.w_size_n)? + .contiguous()?; + let up = gate_up + .narrow(candle::D::Minus1, self.w_size_n, self.w_size_n)? + .contiguous()?; + + //(M * top_k, N // 2) + let down_inputs = (up * gate.apply(&self.act)?)?.reshape(((), self.w_size_n))?; + + //view(M, top_k, K) -> sum -> (M, K) + let ys = moe::moe_gemm( + &down_inputs, + &self.down_w, + &Some(topk_weights), + &sorted_token_ids, + &expert_ids, + self.num_experts_per_tok, + is_prefill, + )? + .reshape((num_tokens, (), hidden_dim))? + .sum(D::Minus2)?; + + ys.reshape((batch, seq_len, hidden_dim)) + } +} + +pub struct FusedMoeGGUF { + pub gate: Linear, + pub gate_experts: Arc, + pub up_experts: Arc, + pub down_experts: Arc, + pub act: Activation, + pub norm_topk_prob: bool, + pub num_experts_per_tok: usize, + // all_reduce: AllReduce, + // world_size: usize, + pub dtype: DType, +} + +impl FusedMoeGGUF { + pub fn new( + cfg: &MoeCfg, + vb: crate::quantized_var_builder::VarBuilder, + dtype: DType, + ) -> Result { + let num_experts = cfg.num_experts; + let gate_ws = vb + .pp("ffn_gate_inp") + .get((num_experts, cfg.hidden_size), "weight")? + .dequantize(vb.device())? + .to_dtype(DType::F32)?; + + let gate = Linear::new(gate_ws, None); + + let (gate_experts, up_experts, down_experts) = { + ( + vb.pp("ffn_gate_exps").get( + (num_experts, cfg.moe_intermediate_size, cfg.hidden_size), + "weight", + )?, + vb.pp("ffn_up_exps").get( + (num_experts, cfg.moe_intermediate_size, cfg.hidden_size), + "weight", + )?, + vb.pp("ffn_down_exps").get( + (num_experts, cfg.hidden_size, cfg.moe_intermediate_size), + "weight", + )?, + ) + }; + + Ok(Self { + gate, + gate_experts, + up_experts, + down_experts, + act: cfg.act, + norm_topk_prob: cfg.norm_topk_prob, + num_experts_per_tok: cfg.num_experts_per_tok, + // all_reduce: AllReduce::new(comm), + // world_size: 1, + dtype, + }) + } + + pub fn forward(&self, xs: &Tensor, is_prefill: bool) -> Result { + let (batch, seq_len, hidden_dim) = xs.dims3()?; + let xs = xs.reshape(((), hidden_dim))?; + let (num_tokens, hidden_dim) = xs.dims2()?; + let original_dtype = xs.dtype(); + let xs = if xs.dtype() != DType::F32 { + xs.to_dtype(DType::F32)? + } else { + xs.to_owned() + }; + + let router_logits = self.gate.forward(&xs)?; + + let routing_weights = + candle_nn::ops::softmax_last_dim(&router_logits.to_dtype(DType::F32)?)?; + + let topk_ids = routing_weights + .arg_sort_last_dim(false)? + .narrow(D::Minus1, 0, self.num_experts_per_tok)? + .contiguous()?; + + let mut topk_weights = routing_weights.gather(&topk_ids, D::Minus1)?; + + if self.norm_topk_prob { + topk_weights = topk_weights.broadcast_div(&topk_weights.sum_keepdim(D::Minus1)?)?; + } + + let (expert_ids, sorted_token_ids) = if is_prefill { + // For long-context (32K+), need to use custom sort kernel + // #[cfg(feature = "cuda")] + // { + // use attention_rs::sort::ArgSortOp; + // topk_ids.flatten_all()?.sort(true)? + // } + // #[cfg(not(feature = "cuda"))] + topk_ids.flatten_all()?.sort_last_dim(true)? + } else { + topk_ids.flatten_all()?.sort_last_dim(true)? + }; + + let ys = { + let gate = moe::moe_gemm_gguf( + &xs, + &self.gate_experts, + &None, + &sorted_token_ids, + &expert_ids, + self.num_experts_per_tok, + is_prefill, + self.dtype, + )?; + let up = moe::moe_gemm_gguf( + &xs, + &self.up_experts, + &None, + &sorted_token_ids, + &expert_ids, + self.num_experts_per_tok, + is_prefill, + self.dtype, + )?; + + let down_inputs = (up * gate.apply(&self.act)?)?; + moe::moe_gemm_gguf( + &down_inputs, + &self.down_experts, + &Some(topk_weights), + &sorted_token_ids, + &expert_ids, + self.num_experts_per_tok, + is_prefill, + self.dtype, + )? + }; + let mut ys = ys.reshape((num_tokens, (), hidden_dim))?.sum(D::Minus2)?; + if ys.dtype() != original_dtype { + ys = ys.to_dtype(original_dtype)?; + } + ys.reshape((batch, seq_len, hidden_dim)) + } +} diff --git a/patches/candle-transformers/src/generation/mod.rs b/patches/candle-transformers/src/generation/mod.rs new file mode 100644 index 0000000000..7f4200c00a --- /dev/null +++ b/patches/candle-transformers/src/generation/mod.rs @@ -0,0 +1,158 @@ +//! Logit Processing and Sampling +//! +//! Functionality for modeling sampling strategies and logits processing in text generation +//! with support for temperature-based sampling, top-k filtering, nucleus sampling (top-p), +//! and combinations thereof. +use candle::{DType, Error, Result, Tensor}; +use rand::{distr::Distribution, SeedableRng}; + +#[derive(Clone, PartialEq, Debug)] +pub enum Sampling { + ArgMax, + All { temperature: f64 }, + TopK { k: usize, temperature: f64 }, + TopP { p: f64, temperature: f64 }, + TopKThenTopP { k: usize, p: f64, temperature: f64 }, + // Note that the rng is not used for the Gumbel-Softmax sampling. + GumbelSoftmax { temperature: f64 }, +} + +pub struct LogitsProcessor { + rng: rand::rngs::StdRng, + sampling: Sampling, +} + +impl LogitsProcessor { + pub fn from_sampling(seed: u64, sampling: Sampling) -> Self { + let rng = rand::rngs::StdRng::seed_from_u64(seed); + Self { rng, sampling } + } + + pub fn new(seed: u64, temperature: Option, top_p: Option) -> Self { + let temperature = temperature.and_then(|v| if v < 1e-7 { None } else { Some(v) }); + let sampling = match temperature { + None => Sampling::ArgMax, + Some(temperature) => match top_p { + None => Sampling::All { temperature }, + Some(p) => Sampling::TopP { p, temperature }, + }, + }; + Self::from_sampling(seed, sampling) + } + + fn sample_argmax(&mut self, logits: Tensor) -> Result { + logits.argmax(candle::D::Minus1)?.to_scalar::() + } + + fn sample_gumbel_softmax(&mut self, logits: &Tensor, temperature: f64) -> Result { + let sampled = candle_nn::sampling::gumbel_softmax(logits, temperature, candle::D::Minus1)?; + sampled.to_scalar::() + } + + fn sample_multinomial(&mut self, prs: &Vec) -> Result { + let distr = rand::distr::weighted::WeightedIndex::new(prs).map_err(Error::wrap)?; + let next_token = distr.sample(&mut self.rng) as u32; + Ok(next_token) + } + + /// top-p sampling (or "nucleus sampling") samples from the smallest set of tokens that exceed + /// probability top_p. This way we never sample tokens that have very low probabilities and are + /// less likely to go "off the rails". + fn sample_topp(&mut self, prs: &mut Vec, top_p: f32) -> Result { + let mut argsort_indices = (0..prs.len()).collect::>(); + + // Sort by descending probability. + argsort_indices.sort_by(|&i, &j| prs[j].total_cmp(&prs[i])); + + // Clamp smaller probabilities to zero. + let mut cumsum = 0.; + for index in &argsort_indices { + if cumsum >= top_p { + prs[*index] = 0.0; + } else { + cumsum += prs[*index]; + } + } + // Sample with clamped probabilities. + self.sample_multinomial(prs) + } + + // top-k sampling samples from the k tokens with the largest probabilities. + fn sample_topk(&mut self, prs: &mut Vec, top_k: usize) -> Result { + if top_k >= prs.len() { + self.sample_multinomial(prs) + } else { + let mut argsort_indices = (0..prs.len()).collect::>(); + let (indices, _, _) = + argsort_indices.select_nth_unstable_by(top_k, |&i, &j| prs[j].total_cmp(&prs[i])); + let prs = indices.iter().map(|&i| prs[i]).collect::>(); + let index = self.sample_multinomial(&prs)?; + Ok(indices[index as usize] as u32) + } + } + + // top-k sampling samples from the k tokens with the largest probabilities. + // then top-p sampling. + fn sample_topk_topp(&mut self, prs: &mut Vec, top_k: usize, top_p: f32) -> Result { + if top_k >= prs.len() { + self.sample_topp(prs, top_p) + } else { + let mut argsort_indices = (0..prs.len()).collect::>(); + let (indices, _, _) = + argsort_indices.select_nth_unstable_by(top_k, |&i, &j| prs[j].total_cmp(&prs[i])); + let mut prs = indices.iter().map(|&i| prs[i]).collect::>(); + let sum_p = prs.iter().sum::(); + let index = if top_p <= 0.0 || top_p >= sum_p { + self.sample_multinomial(&prs)? + } else { + self.sample_topp(&mut prs, top_p)? + }; + Ok(indices[index as usize] as u32) + } + } + + pub fn sample(&mut self, logits: &Tensor) -> Result { + self.sample_f(logits, |_| {}) + } + + pub fn sample_f(&mut self, logits: &Tensor, f: impl FnOnce(&mut [f32])) -> Result { + let logits = logits.to_dtype(DType::F32)?; + let prs = |temperature: f64| -> Result> { + let logits = (&logits / temperature)?; + let prs = candle_nn::ops::softmax_last_dim(&logits)?; + let mut prs = prs.to_vec1()?; + f(&mut prs); + Ok(prs) + }; + + let next_token = match &self.sampling { + Sampling::ArgMax => self.sample_argmax(logits)?, + Sampling::GumbelSoftmax { temperature } => { + self.sample_gumbel_softmax(&logits, *temperature)? + } + Sampling::All { temperature } => { + let prs = prs(*temperature)?; + self.sample_multinomial(&prs)? + } + Sampling::TopP { p, temperature } => { + let mut prs = prs(*temperature)?; + if *p <= 0.0 || *p >= 1.0 { + // simply sample from the predicted probability distribution + self.sample_multinomial(&prs)? + } else { + // top-p (nucleus) sampling, clamping the least likely tokens to zero + self.sample_topp(&mut prs, *p as f32)? + } + } + Sampling::TopK { k, temperature } => { + let mut prs = prs(*temperature)?; + self.sample_topk(&mut prs, *k)? + } + Sampling::TopKThenTopP { k, p, temperature } => { + let mut prs = prs(*temperature)?; + self.sample_topk_topp(&mut prs, *k, *p as f32)? + } + }; + Ok(next_token) + } +} diff --git a/patches/candle-transformers/src/lib.rs b/patches/candle-transformers/src/lib.rs new file mode 100644 index 0000000000..bae7699a09 --- /dev/null +++ b/patches/candle-transformers/src/lib.rs @@ -0,0 +1,8 @@ +pub mod fused_moe; +pub mod generation; +pub mod models; +pub mod object_detection; +pub mod pipelines; +pub mod quantized_nn; +pub mod quantized_var_builder; +pub mod utils; diff --git a/patches/candle-transformers/src/models/based.rs b/patches/candle-transformers/src/models/based.rs new file mode 100644 index 0000000000..dd2aa80dad --- /dev/null +++ b/patches/candle-transformers/src/models/based.rs @@ -0,0 +1,588 @@ +//! Based from the Stanford Hazy Research group. +//! +//! See "Simple linear attention language models balance the recall-throughput tradeoff", Arora et al. 2024 +//! - Simple linear attention language models balance the recall-throughput tradeoff. [Arxiv](https://arxiv.org/abs/2402.18668) +//! - [GitHub Rep](https://github.com/HazyResearch/based) +//! - [Blogpost](https://hazyresearch.stanford.edu/blog/2024-03-03-based) + +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{ + conv1d_no_bias, linear, linear_no_bias, ops::softmax_last_dim, rms_norm, Conv1d, Conv1dConfig, + Func, Linear, RmsNorm, VarBuilder, +}; +use std::sync::Arc; + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct LinearAttentionFeatureMapConfig { + input_dim: usize, +} + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct LinearAttentionConfig { + num_heads: usize, + feature_dim: usize, + feature_map: LinearAttentionFeatureMapConfig, +} + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct SlidingWindowAttentionConfig { + num_heads: usize, + window_size: usize, +} + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + vocab_size: usize, + #[serde(rename = "n_embd")] + hidden_size: usize, + #[serde(rename = "n_inner")] + intermediate_size: usize, + #[serde(rename = "n_layer")] + num_hidden_layers: usize, + #[serde(rename = "n_head")] + num_attention_heads: usize, + + layer_norm_epsilon: f64, + #[serde(default = "default_rope", rename = "rotary_emb_base")] + rope_theta: f64, + + alt_mixer_layers: Vec, + alt_mixer_2_layers: Vec, + #[serde(rename = "alt_mixer")] + la: LinearAttentionConfig, + #[serde(rename = "alt_mixer_2")] + swa: SlidingWindowAttentionConfig, +} + +fn default_rope() -> f64 { + 10_000.0 +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + fc1: Linear, + fc2: Linear, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let fc1 = linear_no_bias(cfg.hidden_size, cfg.hidden_size * 4, vb.pp("fc1"))?; + let fc2 = linear_no_bias(cfg.intermediate_size, cfg.hidden_size, vb.pp("fc2"))?; + Ok(Self { fc1, fc2 }) + } +} + +// Swiglu implementation. +// Not using Activation::Swiglu because this has the gate and y arguments switched compared to the version in candle-nn/src/ops.rs +fn swiglu(xs: &Tensor) -> Result { + let xs = xs.chunk(2, D::Minus1)?; + &xs[1].silu()? * &xs[0] +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + let xs = xs.apply(&self.fc1)?; + let xs = swiglu(&xs)?; + let xs = xs.apply(&self.fc2)?; + Ok(xs) + } +} + +// A gated convolutional block. +#[derive(Debug, Clone)] +struct BasedConv { + in_proj: Linear, + out_proj: Linear, + conv: Conv1d, + state: Tensor, +} + +impl BasedConv { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dim = cfg.hidden_size * 2; + + let conv1d_cfg = Conv1dConfig { + groups: dim, + padding: 2, + ..Default::default() + }; + + let in_proj = linear(cfg.hidden_size, cfg.hidden_size * 4, vb.pp("in_proj"))?; + let out_proj = linear(dim, cfg.hidden_size, vb.pp("out_proj"))?; + let conv = conv1d_no_bias(dim, dim, 3, conv1d_cfg, vb.pp("conv.conv"))?; + let state = Tensor::zeros((1, dim, 3), vb.dtype(), vb.device())?; + Ok(Self { + in_proj, + out_proj, + conv, + state, + }) + } + + fn step(&mut self, xs: &Tensor) -> Result { + self.state = self.state.roll(-1, D::Minus1)?; + let (_, _, l) = self.state.dims3()?; + self.state = self.state.narrow(D::Minus1, 0, l - 1)?; + self.state = Tensor::cat(&[&self.state, &xs.transpose(1, 2)?], 2)?; + + let xs = (&self.state * self.conv.weight().permute((1, 0, 2))?)? + .sum_keepdim(0)? + .sum(D::Minus1)?; + + let xs = xs.unsqueeze(1)?; + + Ok(xs) + } + + fn forward(&mut self, xs: &Tensor, seqlen_offset: usize) -> Result { + let xs = xs.apply(&self.in_proj)?; + let us = xs.chunk(2, D::Minus1)?; + let (_b, l, _d) = us[0].dims3()?; + let u_conv = if seqlen_offset > 0 { + self.step(&us[0])? + } else { + let k = std::cmp::min(3, l); + self.state = self.state.narrow(D::Minus1, 0, 3 - k)?; + let xs = us[0].narrow(1, l - k, k)?.transpose(1, 2)?; + self.state = Tensor::cat(&[&self.state, &xs], 2)?; + + us[0] + .transpose(1, 2)? + .apply(&self.conv)? + .narrow(D::Minus1, 0, l)? + .transpose(1, 2)? + }; + + let u_conv = u_conv.silu()?; + let v = u_conv.broadcast_mul(&us[1])?; + let xs = v.apply(&self.out_proj)?; + + Ok(xs) + } +} + +// Linear attention approximating softmax using second order Taylor polynomials. +#[derive(Debug, Clone)] +struct LinearAttention { + proj_q: Linear, + proj_k: Linear, + proj_v: Linear, + out_proj: Linear, + feature_dim: usize, + num_heads: usize, + input_dim: usize, + k_state: Tensor, + kv_state: Tensor, +} + +impl LinearAttention { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let input_dim = cfg.la.feature_map.input_dim; + let out_proj = linear_no_bias(cfg.hidden_size, cfg.hidden_size, vb.pp("out_proj"))?; + let proj_k = linear_no_bias( + cfg.hidden_size, + cfg.la.num_heads * cfg.la.feature_dim, + vb.pp("proj_k"), + )?; + let proj_q = linear_no_bias( + cfg.hidden_size, + cfg.la.num_heads * cfg.la.feature_dim, + vb.pp("proj_q"), + )?; + + let proj_v = linear_no_bias(cfg.hidden_size, cfg.hidden_size, vb.pp("proj_v"))?; + let expanded_size = cfg.la.feature_dim.pow(2) + cfg.la.feature_dim + 1; + let k_state = Tensor::zeros( + (1, cfg.la.num_heads, 1, 1, expanded_size), + vb.dtype(), + vb.device(), + )?; + let kv_state = Tensor::zeros( + (1, cfg.la.num_heads, cfg.la.feature_dim, expanded_size), + vb.dtype(), + vb.device(), + )?; + + Ok(Self { + proj_q, + proj_k, + proj_v, + out_proj, + feature_dim: cfg.la.feature_dim, + num_heads: cfg.la.num_heads, + input_dim, + k_state, + kv_state, + }) + } + + fn taylor_expansion(&self) -> Result> { + let r2 = std::f64::consts::SQRT_2; + let rd = (self.input_dim as f64).sqrt(); + let rrd = rd.sqrt(); + + Ok(Func::new(move |xs| { + let dims = xs.dims(); + let mut d = dims.to_vec(); + if let Some(last) = d.last_mut() { + *last = 1; + }; + + let x = xs + .unsqueeze(D::Minus1)? + .broadcast_mul(&xs.unsqueeze(D::Minus2)?)?; + let x = (x.flatten_from(D::Minus2)? / r2)?; + let o = Tensor::ones(d, xs.dtype(), xs.device())?; + let x = Tensor::cat(&[o, (xs / rrd)?, (&x / rd)?], D::Minus1)?; + + Ok(x) + })) + } + + fn forward(&mut self, xs: &Tensor, seqlen_offset: usize) -> Result { + let eps = 1e-12; + + let feature_map = self.taylor_expansion()?; + + let (b, l, d) = xs.dims3()?; + let q = xs.apply(&self.proj_q)?; + let k = xs.apply(&self.proj_k)?; + let v = xs.apply(&self.proj_v)?; + + let q = q + .reshape((b, l, self.num_heads, self.feature_dim))? + .transpose(1, 2)? + .contiguous()?; + let k = k + .reshape((b, l, self.num_heads, self.feature_dim))? + .transpose(1, 2)? + .contiguous()?; + let v = v + .reshape((b, l, self.num_heads, d / self.num_heads))? + .transpose(1, 2)? + .contiguous()?; + + let q = feature_map.forward(&q)?; + let k = feature_map.forward(&k)?; + + let y = if seqlen_offset > 0 { + let (_b, _h, l, _d) = k.dims4()?; + let q = q.unsqueeze(D::Minus2)?; + let k = k.unsqueeze(D::Minus2)?; + let v = v.unsqueeze(D::Minus1)?; + let kn = k.narrow(D::Minus1, l - 1, 1)?; + let vn = v.narrow(D::Minus1, l - 1, 1)?; + + self.k_state = self.k_state.broadcast_add(&kn)?; + self.kv_state = self.kv_state.broadcast_add(&kn.broadcast_mul(&vn)?)?; + + let num = q.broadcast_mul(&self.kv_state)?.sum(D::Minus1)?; + let den = (q.broadcast_mul(&self.k_state)?.sum(D::Minus1)? + eps)?; + num.broadcast_div(&den)? + } else { + self.k_state = k.sum(2)?.unsqueeze(2)?.unsqueeze(3)?; + self.kv_state = k + .transpose(2, 3)? + .matmul(&v)? + .transpose(2, 3)? + .unsqueeze(2)?; + let aqk = q.matmul(&k.transpose(D::Minus1, D::Minus2)?)?; + let tril = Tensor::tril2(l, aqk.dtype(), aqk.device())?; + let aqk = aqk.broadcast_mul(&tril)?.matmul(&v)?; + + let z = (1f64 / (q.mul(&k.cumsum(2)?)?.sum(D::Minus1)? + eps)?)?; + aqk.broadcast_mul(&z.unsqueeze(D::Minus1)?)? + }; + + let (b, h, l, d) = y.dims4()?; + let y = y.permute((0, 2, 1, 3))?.reshape((b, l, h * d))?; + let y = self.out_proj.forward(&y)?; + + Ok(y) + } +} + +// Rotary embeddings used in local attention. +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let dim = cfg.hidden_size / cfg.num_attention_heads; + let max_seq_len = 2048; // Hardcoded, missing from config. + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +// Local attention using a small sliding window. +#[derive(Debug, Clone)] +struct SlidingWindowAttention { + wqkv: Linear, + out_proj: Linear, + num_heads: usize, + head_dim: usize, + hidden_size: usize, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl SlidingWindowAttention { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_size = cfg.hidden_size; + let num_heads = cfg.swa.num_heads; + let head_dim = hidden_size / num_heads; + let out_proj = linear_no_bias(hidden_size, hidden_size, vb.pp("out_proj"))?; + let wqkv = linear_no_bias(hidden_size, hidden_size * 3, vb.pp("Wqkv"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb.device())?); + Ok(Self { + wqkv, + out_proj, + hidden_size, + num_heads, + head_dim, + rotary_emb, + kv_cache: None, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let qkv = xs.apply(&self.wqkv)?; + let qkv = qkv.reshape((b_sz, q_len, 3, (), self.head_dim))?; + + let q = qkv.i((.., .., 0))?; + let k = qkv.i((.., .., 1))?; + let v = qkv.i((.., .., 2))?; + + let q = q + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let k = k + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let v = v + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + + let (q, k) = self + .rotary_emb + .apply_rotary_emb_qkv(&q, &k, seqlen_offset)?; + + let (k, v) = match &self.kv_cache { + None => (k, v), + Some((prev_k, prev_v)) => { + let k = Tensor::cat(&[prev_k, &k], 2)?; + let v = Tensor::cat(&[prev_v, &v], 2)?; + (k, v) + } + }; + self.kv_cache = Some((k.clone(), v.clone())); + + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (q.matmul(&k.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = softmax_last_dim(&attn_weights)?; + let attn_output = attn_weights.matmul(&v)?; + let out = attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.hidden_size))? + .apply(&self.out_proj)?; + + Ok(out) + } +} + +// The model layers use three types of mixers. +#[derive(Debug, Clone)] +enum SequenceMixer { + Based(BasedConv), + Linear(LinearAttention), + Sliding(SlidingWindowAttention), +} + +impl SequenceMixer { + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + pos: usize, + ) -> Result { + match self { + Self::Based(b) => b.forward(xs, pos), + Self::Linear(b) => b.forward(xs, pos), + Self::Sliding(b) => b.forward(xs, attention_mask, pos), + } + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + mlp: MLP, + norm1: RmsNorm, + norm2: RmsNorm, + mixer: SequenceMixer, +} + +impl DecoderLayer { + fn new(layer_idx: usize, cfg: &Config, vb: VarBuilder) -> Result { + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let norm1 = rms_norm(cfg.hidden_size, cfg.layer_norm_epsilon, vb.pp("norm1"))?; + let norm2 = rms_norm(cfg.hidden_size, cfg.layer_norm_epsilon, vb.pp("norm2"))?; + + let l_attn = cfg.alt_mixer_layers.contains(&layer_idx); + let sw_attn = cfg.alt_mixer_2_layers.contains(&layer_idx); + + let mixer = if l_attn { + SequenceMixer::Linear(LinearAttention::new(cfg, vb.pp("mixer"))?) + } else if sw_attn { + SequenceMixer::Sliding(SlidingWindowAttention::new(cfg, vb.pp("mixer"))?) + } else { + SequenceMixer::Based(BasedConv::new(cfg, vb.pp("mixer"))?) + }; + + Ok(Self { + mlp, + norm1, + norm2, + mixer, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.norm1.forward(xs)?; + let xs = self.mixer.forward(&xs, attention_mask, seqlen_offset)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.norm2)?.apply(&self.mlp)?; + residual + xs + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: super::with_tracing::Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: Linear, + sliding_window: usize, + device: Device, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vocab_size = cfg.vocab_size + (8 - cfg.vocab_size % 8) % 8; + let lm_head = linear_no_bias(cfg.hidden_size, vocab_size, vb.pp("lm_head"))?; + let embed_tokens = super::with_tracing::Embedding::from_weights(lm_head.weight().clone())?; + let vb_m = vb.pp("transformer"); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(layer_idx, cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = rms_norm(cfg.hidden_size, cfg.layer_norm_epsilon, vb_m.pp("ln_f"))?; + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + sliding_window: cfg.swa.window_size, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + let sliding_window = self.sliding_window / 2; + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| { + (0..tgt_len).map(move |j| { + if i < j || j + sliding_window < i { + f32::NEG_INFINITY + } else { + 0. + } + }) + }) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), self.dtype, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let (b_size, seq_len) = input_ids.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(b_size, seq_len, seqlen_offset)?; + Some(mask) + }; + let mut xs = self.embed_tokens.forward(input_ids)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)? + } + xs.narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head) + } +} diff --git a/patches/candle-transformers/src/models/beit.rs b/patches/candle-transformers/src/models/beit.rs new file mode 100644 index 0000000000..6b6368423f --- /dev/null +++ b/patches/candle-transformers/src/models/beit.rs @@ -0,0 +1,411 @@ +//! Based on the BEIT vision-language model. +//! +//! See "BEIT: BERT Pre-Training of Image Transformers", Bao et al. 2021 +//! - [Arxiv](https://arxiv.org/abs/2106.08254) +//! - [GitHub](https://github.com/microsoft/unilm/tree/master/beit) +//! + +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::{layer_norm, LayerNorm, Linear, Module, VarBuilder}; + +const IMG_SIZE: usize = 384; +const PATCH_SIZE: usize = 16; +const NUM_CLASSES: usize = 1000; +const WINDOW_SIZE: usize = IMG_SIZE / PATCH_SIZE; // 384 / 16 = 24 +const NB_TOKENS: usize = WINDOW_SIZE * WINDOW_SIZE + 1; // 24 * 24 + 1 = 577 + +fn linear(vb: VarBuilder, in_dim: usize, out_dim: usize, bias: bool) -> Result { + if bias { + candle_nn::linear(in_dim, out_dim, vb) + } else { + candle_nn::linear_no_bias(in_dim, out_dim, vb) + } +} + +#[derive(Debug)] +struct Attention { + qkv: Linear, + proj: Linear, + relative_position_bias_table: Tensor, + relative_position_index: Tensor, + num_heads: usize, + scale: f64, +} + +impl Attention { + fn new( + vb: VarBuilder, + dim: usize, + num_heads: usize, + qkv_bias: bool, + proj_bias: bool, + ) -> Result { + let qkv = linear(vb.pp("qkv"), dim, dim * 3, qkv_bias)?; + let proj = linear(vb.pp("proj"), dim, dim, proj_bias)?; + // num_relative_distance = token-token(47x47) + token-CLS(1) + CLS-token(1) + CLS-CLS(1) = 2212 + let num_relative_distance = (2 * WINDOW_SIZE - 1) * (2 * WINDOW_SIZE - 1) + 3; + let relative_position_bias_table = vb.get( + (num_relative_distance, num_heads), + "relative_position_bias_table", + )?; + let relative_position_index = + Self::gen_relative_position_index(relative_position_bias_table.device())?; + let scale = 1. / ((dim / num_heads) as f64).sqrt(); + Ok(Self { + qkv, + proj, + relative_position_bias_table, + relative_position_index, + num_heads, + scale, + }) + } +} + +impl Attention { + // See: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/beit.py#L61 + fn gen_relative_position_index(device: &Device) -> Result { + let num_relative_distance = (2 * WINDOW_SIZE - 1) * (2 * WINDOW_SIZE - 1) + 3; + let w_area = WINDOW_SIZE * WINDOW_SIZE; + + let t_arange: Tensor = Tensor::arange(0, WINDOW_SIZE as u32, device)?; + let t_ndgrid = Tensor::meshgrid(&[&t_arange, &t_arange], false)?; + let coords_flatten = Tensor::stack(&t_ndgrid, 0)?.flatten(1, 2)?; + + let tmp1 = coords_flatten + .unsqueeze(2)? + .broadcast_as((2, w_area, w_area))? + .to_dtype(DType::I64)?; + let tmp2 = coords_flatten + .unsqueeze(1)? + .broadcast_as((2, w_area, w_area))? + .to_dtype(DType::I64)?; + let relative_coords = (tmp1 - tmp2)? + .transpose(0, 1)? // 102 + .transpose(1, 2)? // 120 + .contiguous()?; + + let relative_coords = relative_coords.slice_assign( + &[0..w_area, 0..w_area, 0..1], + &(relative_coords.i((0..w_area, 0..w_area, 0..1))? + (WINDOW_SIZE - 1) as f64)?, + )?; + let relative_coords = relative_coords.slice_assign( + &[0..w_area, 0..w_area, 1..2], + &(relative_coords.i((0..w_area, 0..w_area, 1..2))? + (WINDOW_SIZE - 1) as f64)?, + )?; + let relative_coords = relative_coords.slice_assign( + &[0..w_area, 0..w_area, 0..1], + &(relative_coords.i((.., .., 0..1))? * (2. * (WINDOW_SIZE as f64) - 1.))?, + )?; + + Tensor::zeros((w_area + 1, w_area + 1), DType::I64, device)? + .slice_assign(&[1.., 1..], &relative_coords.sum(2)?)? + .slice_assign( + &[0..1, 0..(w_area + 1)], + &(Tensor::ones((1, w_area + 1), DType::I64, device)? + * ((num_relative_distance - 3) as f64))? + .to_dtype(DType::I64)?, + )? + .slice_assign( + &[0..(w_area + 1), 0..1], + &(Tensor::ones((w_area + 1, 1), DType::I64, device)? + * ((num_relative_distance - 2) as f64))? + .to_dtype(DType::I64)?, + )? + .slice_assign( + &[0..1, 0..1], + &(Tensor::ones((1, 1), DType::I64, device)? + * ((num_relative_distance - 1) as f64))? + .to_dtype(DType::I64)?, + ) + } + + fn _get_rel_pos_bias(&self) -> Result { + self.relative_position_bias_table + .index_select( + &self + .relative_position_index + .flatten_all()? + .to_dtype(DType::U32)?, + 0, + )? + .reshape((NB_TOKENS, NB_TOKENS, ()))? + .transpose(0, 1)? // 102 + .transpose(0, 2)? // 201 + .contiguous()? + .unsqueeze(0) + } +} + +impl Module for Attention { + fn forward(&self, xs: &Tensor) -> Result { + let (b, n, c) = xs.dims3()?; + let qkv = self + .qkv + .forward(xs)? + .reshape((b, n, 3, self.num_heads, c / self.num_heads))? + .transpose(1, 2)? // 02134 + .transpose(0, 1)? // 20134 + .transpose(2, 3)?; // 20314 + let q = (qkv.i(0)? * self.scale)?; + let k = qkv.i(1)?.contiguous()?; + let v = qkv.i(2)?.contiguous()?; + let attn = (&q.matmul(&k.t()?)? + self._get_rel_pos_bias())?; + let attn = candle_nn::ops::softmax(&attn, D::Minus1)?; + let attn = attn.matmul(&v)?.transpose(1, 2)?.reshape((b, n, c))?; + self.proj.forward(&attn) + } +} + +#[derive(Debug)] +struct LayerScale { + gamma: Tensor, +} + +impl LayerScale { + fn new(vb: VarBuilder, dim: usize) -> Result { + let gamma = vb.get(dim, "gamma")?; + Ok(Self { gamma }) + } +} + +impl Module for LayerScale { + fn forward(&self, xs: &Tensor) -> Result { + xs.broadcast_mul(&self.gamma) + } +} + +#[derive(Debug)] +struct Mlp { + fc1: Linear, + fc2: Linear, +} + +impl Mlp { + fn new(vb: VarBuilder, in_features: usize, hidden_features: usize, bias: bool) -> Result { + let out_features = in_features; + let fc1 = linear(vb.pp("fc1"), in_features, hidden_features, bias)?; + let fc2 = linear(vb.pp("fc2"), hidden_features, out_features, bias)?; + Ok(Self { fc1, fc2 }) + } +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.fc1.forward(xs)?.gelu()?; + self.fc2.forward(&xs) + } +} + +#[derive(Debug)] +struct Block { + norm1: LayerNorm, + attn: Attention, + ls1: LayerScale, + norm2: LayerNorm, + mlp: Mlp, + ls2: LayerScale, +} + +impl Block { + fn new(vb: VarBuilder, dim: usize, num_heads: usize) -> Result { + let norm1 = layer_norm(dim, 1e-6, vb.pp("norm1"))?; + let attn = Attention::new(vb.pp("attn"), dim, num_heads, true, true)?; + let ls1 = LayerScale::new(vb.pp("ls1"), dim)?; + let norm2 = layer_norm(dim, 1e-6, vb.pp("norm2"))?; + let mlp = Mlp::new(vb.pp("mlp"), dim, dim * 4, true)?; + let ls2 = LayerScale::new(vb.pp("ls2"), dim)?; + Ok(Self { + norm1, + attn, + ls1, + norm2, + mlp, + ls2, + }) + } +} + +impl Module for Block { + fn forward(&self, xs: &Tensor) -> Result { + let residual = xs; + let xs = self + .ls1 + .forward(&self.attn.forward(&self.norm1.forward(xs)?)?)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = self + .ls2 + .forward(&self.mlp.forward(&self.norm2.forward(&xs)?)?)?; + xs + residual + } +} + +#[derive(Debug)] +struct PatchEmbed { + proj: candle_nn::Conv2d, + patch_size: (usize, usize), +} + +impl PatchEmbed { + fn new(vb: VarBuilder, patch_size: usize, in_chans: usize, embed_dim: usize) -> Result { + let config = candle_nn::Conv2dConfig { + stride: patch_size, + ..Default::default() + }; + let proj = candle_nn::conv2d(in_chans, embed_dim, patch_size, config, vb.pp("proj"))?; + Ok(Self { + proj, + patch_size: (patch_size, patch_size), + }) + } +} + +impl Module for PatchEmbed { + fn forward(&self, xs: &Tensor) -> Result { + let (_b, _c, h, w) = xs.dims4()?; + let (patch_h, patch_w) = self.patch_size; + if (h % patch_h) != 0 { + candle::bail!("image height {h} is not a multiple of patch height {patch_h}") + } + if (w % patch_w) != 0 { + candle::bail!("image width {w} is not a multiple of patch width {patch_w}") + } + let xs = self.proj.forward(xs)?; + let (b, c, h, w) = xs.dims4()?; + // flatten embeddings. + xs.reshape((b, c, h * w))?.transpose(1, 2) + } +} + +#[derive(Debug)] +pub struct BeitVisionTransformer { + patch_embed: PatchEmbed, + cls_token: Tensor, + blocks: Vec, + norm: LayerNorm, + head: Linear, +} + +impl BeitVisionTransformer { + pub fn new(vb: VarBuilder, depth: usize, embed_dim: usize, num_heads: usize) -> Result { + let patch_embed = PatchEmbed::new(vb.pp("patch_embed"), PATCH_SIZE, 3, embed_dim)?; + let cls_token = vb.get((1, 1, embed_dim), "cls_token")?; + let head = linear(vb.pp("head"), embed_dim, NUM_CLASSES, true)?; + let norm = layer_norm(embed_dim, 1e-6, vb.pp("norm"))?; + let vb_b = vb.pp("blocks"); + let blocks = (0..depth) + .map(|i| Block::new(vb_b.pp(i.to_string()), embed_dim, num_heads)) + .collect::>>()?; + Ok(Self { + patch_embed, + cls_token, + blocks, + norm, + head, + }) + } + + fn prepare_tokens_with_mask(&self, xs: &Tensor) -> Result { + let xs = self.patch_embed.forward(xs)?; + Tensor::cat(&[&self.cls_token, &xs], 1) + } + + fn get_intermediate_layers_not_chunked( + &self, + xs: &Tensor, + blocks_to_take: &[usize], + ) -> Result> { + let mut xs = self.prepare_tokens_with_mask(xs)?; + let mut output = Vec::new(); + for (i, blk) in self.blocks.iter().enumerate() { + xs = blk.forward(&xs)?; + if blocks_to_take.contains(&i) { + output.push(xs.clone()); + } + } + if output.len() != blocks_to_take.len() { + candle::bail!( + "only {} / {} blocks found", + output.len(), + blocks_to_take.len() + ); + } + Ok(output) + } + + pub fn get_intermediate_layers( + &self, + xs: &Tensor, + blocks_to_take: &[usize], + reshape: bool, + return_class_token: bool, + norm: bool, + ) -> Result { + let outputs = self.get_intermediate_layers_not_chunked(xs, blocks_to_take)?; + let outputs = if norm { + outputs + .iter() + .map(|out| self.norm.forward(out)) + .collect::>>()? + } else { + outputs + }; + let class_tokens = outputs + .iter() + .map(|out| out.i((.., 0))) + .collect::>>()?; + let outputs = outputs + .iter() + .map(|out| out.i((.., 1..))) + .collect::>>()?; + + let outputs = if reshape { + let (b, _c, w, h) = xs.dims4()?; + let patch_size = self.patch_embed.patch_size.0; + let num_channels = outputs[0].elem_count() / (b * (w / patch_size) * (h / patch_size)); + outputs + .iter() + .map(|out| { + out.reshape((b, w / patch_size, h / patch_size, num_channels))? + .transpose(2, 3)? + .transpose(1, 2) + }) + .collect::>>()? + } else { + outputs + }; + + let outputs = if return_class_token { + outputs + .iter() + .zip(class_tokens.iter()) + .map(|(out, class_token)| Tensor::cat(&[out, class_token], D::Minus1)) + .collect::>>()? + } else { + outputs + }; + + Tensor::stack(&outputs[..], 0) + } +} + +impl Module for BeitVisionTransformer { + fn forward(&self, xs: &Tensor) -> Result { + let mut xs = self.prepare_tokens_with_mask(xs)?; + for blk in self.blocks.iter() { + xs = blk.forward(&xs)? + } + let xs_moy_local_tokens = xs.i((.., 1..))?.mean(1)?; + let xs_norm = self.norm.forward(&xs_moy_local_tokens)?; + self.head.forward(&xs_norm) + } +} + +pub fn vit_base(vb: VarBuilder) -> Result { + BeitVisionTransformer::new(vb, 12, 768, 12) +} + +pub fn vit_large(vb: VarBuilder) -> Result { + BeitVisionTransformer::new(vb, 24, 1024, 16) +} diff --git a/patches/candle-transformers/src/models/bert.rs b/patches/candle-transformers/src/models/bert.rs new file mode 100644 index 0000000000..a348c53e14 --- /dev/null +++ b/patches/candle-transformers/src/models/bert.rs @@ -0,0 +1,625 @@ +//! BERT (Bidirectional Encoder Representations from Transformers) +//! +//! Bert is a general large language model that can be used for various language tasks: +//! - Compute sentence embeddings for a prompt. +//! - Compute similarities between a set of sentences. +//! - [Arxiv](https://arxiv.org/abs/1810.04805) "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding" +//! - Upstream [GitHub repo](https://github.com/google-research/bert). +//! - See bert in [candle-examples](https://github.com/huggingface/candle/tree/main/candle-examples/) for runnable code +//! +use super::with_tracing::{layer_norm, linear, LayerNorm, Linear}; +use candle::{DType, Device, Result, Tensor}; +use candle_nn::{embedding, Embedding, Module, VarBuilder}; +use serde::Deserialize; + +pub const DTYPE: DType = DType::F32; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum HiddenAct { + Gelu, + GeluApproximate, + Relu, +} + +#[derive(Clone)] +struct HiddenActLayer { + act: HiddenAct, + span: tracing::Span, +} + +impl HiddenActLayer { + fn new(act: HiddenAct) -> Self { + let span = tracing::span!(tracing::Level::TRACE, "hidden-act"); + Self { act, span } + } + + fn forward(&self, xs: &Tensor) -> candle::Result { + let _enter = self.span.enter(); + match self.act { + // https://github.com/huggingface/transformers/blob/cd4584e3c809bb9e1392ccd3fe38b40daba5519a/src/transformers/activations.py#L213 + HiddenAct::Gelu => xs.gelu_erf(), + HiddenAct::GeluApproximate => xs.gelu(), + HiddenAct::Relu => xs.relu(), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Default)] +#[serde(rename_all = "lowercase")] +pub enum PositionEmbeddingType { + #[default] + Absolute, +} + +// https://github.com/huggingface/transformers/blob/6eedfa6dd15dc1e22a55ae036f681914e5a0d9a1/src/transformers/models/bert/configuration_bert.py#L1 +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub intermediate_size: usize, + pub hidden_act: HiddenAct, + pub hidden_dropout_prob: f64, + pub max_position_embeddings: usize, + pub type_vocab_size: usize, + pub initializer_range: f64, + pub layer_norm_eps: f64, + pub pad_token_id: usize, + #[serde(default)] + pub position_embedding_type: PositionEmbeddingType, + #[serde(default)] + pub use_cache: bool, + pub classifier_dropout: Option, + pub model_type: Option, +} + +impl Default for Config { + fn default() -> Self { + Self { + vocab_size: 30522, + hidden_size: 768, + num_hidden_layers: 12, + num_attention_heads: 12, + intermediate_size: 3072, + hidden_act: HiddenAct::Gelu, + hidden_dropout_prob: 0.1, + max_position_embeddings: 512, + type_vocab_size: 2, + initializer_range: 0.02, + layer_norm_eps: 1e-12, + pad_token_id: 0, + position_embedding_type: PositionEmbeddingType::Absolute, + use_cache: true, + classifier_dropout: None, + model_type: Some("bert".to_string()), + } + } +} + +impl Config { + fn _all_mini_lm_l6_v2() -> Self { + // https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/blob/main/config.json + Self { + vocab_size: 30522, + hidden_size: 384, + num_hidden_layers: 6, + num_attention_heads: 12, + intermediate_size: 1536, + hidden_act: HiddenAct::Gelu, + hidden_dropout_prob: 0.1, + max_position_embeddings: 512, + type_vocab_size: 2, + initializer_range: 0.02, + layer_norm_eps: 1e-12, + pad_token_id: 0, + position_embedding_type: PositionEmbeddingType::Absolute, + use_cache: true, + classifier_dropout: None, + model_type: Some("bert".to_string()), + } + } +} + +#[derive(Clone)] +struct Dropout { + #[allow(dead_code)] + pr: f64, +} + +impl Dropout { + fn new(pr: f64) -> Self { + Self { pr } + } +} + +impl Module for Dropout { + fn forward(&self, x: &Tensor) -> Result { + // TODO + Ok(x.clone()) + } +} + +// https://github.com/huggingface/transformers/blob/6eedfa6dd15dc1e22a55ae036f681914e5a0d9a1/src/transformers/models/bert/modeling_bert.py#L180 +struct BertEmbeddings { + word_embeddings: Embedding, + position_embeddings: Option, + token_type_embeddings: Embedding, + layer_norm: LayerNorm, + dropout: Dropout, + span: tracing::Span, +} + +impl BertEmbeddings { + fn load(vb: VarBuilder, config: &Config) -> Result { + let word_embeddings = embedding( + config.vocab_size, + config.hidden_size, + vb.pp("word_embeddings"), + )?; + let position_embeddings = embedding( + config.max_position_embeddings, + config.hidden_size, + vb.pp("position_embeddings"), + )?; + let token_type_embeddings = embedding( + config.type_vocab_size, + config.hidden_size, + vb.pp("token_type_embeddings"), + )?; + let layer_norm = layer_norm( + config.hidden_size, + config.layer_norm_eps, + vb.pp("LayerNorm"), + )?; + Ok(Self { + word_embeddings, + position_embeddings: Some(position_embeddings), + token_type_embeddings, + layer_norm, + dropout: Dropout::new(config.hidden_dropout_prob), + span: tracing::span!(tracing::Level::TRACE, "embeddings"), + }) + } + + fn forward(&self, input_ids: &Tensor, token_type_ids: &Tensor) -> Result { + let _enter = self.span.enter(); + let (_bsize, seq_len) = input_ids.dims2()?; + let input_embeddings = self.word_embeddings.forward(input_ids)?; + let token_type_embeddings = self.token_type_embeddings.forward(token_type_ids)?; + let mut embeddings = (&input_embeddings + token_type_embeddings)?; + if let Some(position_embeddings) = &self.position_embeddings { + // TODO: Proper absolute positions? + let position_ids = (0..seq_len as u32).collect::>(); + let position_ids = Tensor::new(&position_ids[..], input_ids.device())?; + embeddings = embeddings.broadcast_add(&position_embeddings.forward(&position_ids)?)? + } + let embeddings = self.layer_norm.forward(&embeddings)?; + let embeddings = self.dropout.forward(&embeddings)?; + Ok(embeddings) + } +} + +#[derive(Clone)] +struct BertSelfAttention { + query: Linear, + key: Linear, + value: Linear, + dropout: Dropout, + num_attention_heads: usize, + attention_head_size: usize, + span: tracing::Span, + span_softmax: tracing::Span, +} + +impl BertSelfAttention { + fn load(vb: VarBuilder, config: &Config) -> Result { + let attention_head_size = config.hidden_size / config.num_attention_heads; + let all_head_size = config.num_attention_heads * attention_head_size; + let dropout = Dropout::new(config.hidden_dropout_prob); + let hidden_size = config.hidden_size; + let query = linear(hidden_size, all_head_size, vb.pp("query"))?; + let value = linear(hidden_size, all_head_size, vb.pp("value"))?; + let key = linear(hidden_size, all_head_size, vb.pp("key"))?; + Ok(Self { + query, + key, + value, + dropout, + num_attention_heads: config.num_attention_heads, + attention_head_size, + span: tracing::span!(tracing::Level::TRACE, "self-attn"), + span_softmax: tracing::span!(tracing::Level::TRACE, "softmax"), + }) + } + + fn transpose_for_scores(&self, xs: &Tensor) -> Result { + let mut new_x_shape = xs.dims().to_vec(); + new_x_shape.pop(); + new_x_shape.push(self.num_attention_heads); + new_x_shape.push(self.attention_head_size); + let xs = xs.reshape(new_x_shape.as_slice())?.transpose(1, 2)?; + xs.contiguous() + } + + fn forward(&self, hidden_states: &Tensor, attention_mask: &Tensor) -> Result { + let _enter = self.span.enter(); + let query_layer = self.query.forward(hidden_states)?; + let key_layer = self.key.forward(hidden_states)?; + let value_layer = self.value.forward(hidden_states)?; + + let query_layer = self.transpose_for_scores(&query_layer)?; + let key_layer = self.transpose_for_scores(&key_layer)?; + let value_layer = self.transpose_for_scores(&value_layer)?; + + let attention_scores = query_layer.matmul(&key_layer.t()?)?; + let attention_scores = (attention_scores / (self.attention_head_size as f64).sqrt())?; + let attention_scores = attention_scores.broadcast_add(attention_mask)?; + let attention_probs = { + let _enter_sm = self.span_softmax.enter(); + candle_nn::ops::softmax(&attention_scores, candle::D::Minus1)? + }; + let attention_probs = self.dropout.forward(&attention_probs)?; + + let context_layer = attention_probs.matmul(&value_layer)?; + let context_layer = context_layer.transpose(1, 2)?.contiguous()?; + let context_layer = context_layer.flatten_from(candle::D::Minus2)?; + Ok(context_layer) + } +} + +#[derive(Clone)] +struct BertSelfOutput { + dense: Linear, + layer_norm: LayerNorm, + dropout: Dropout, + span: tracing::Span, +} + +impl BertSelfOutput { + fn load(vb: VarBuilder, config: &Config) -> Result { + let dense = linear(config.hidden_size, config.hidden_size, vb.pp("dense"))?; + let layer_norm = layer_norm( + config.hidden_size, + config.layer_norm_eps, + vb.pp("LayerNorm"), + )?; + let dropout = Dropout::new(config.hidden_dropout_prob); + Ok(Self { + dense, + layer_norm, + dropout, + span: tracing::span!(tracing::Level::TRACE, "self-out"), + }) + } + + fn forward(&self, hidden_states: &Tensor, input_tensor: &Tensor) -> Result { + let _enter = self.span.enter(); + let hidden_states = self.dense.forward(hidden_states)?; + let hidden_states = self.dropout.forward(&hidden_states)?; + self.layer_norm.forward(&(hidden_states + input_tensor)?) + } +} + +// https://github.com/huggingface/transformers/blob/6eedfa6dd15dc1e22a55ae036f681914e5a0d9a1/src/transformers/models/bert/modeling_bert.py#L392 +#[derive(Clone)] +struct BertAttention { + self_attention: BertSelfAttention, + self_output: BertSelfOutput, + span: tracing::Span, +} + +impl BertAttention { + fn load(vb: VarBuilder, config: &Config) -> Result { + let self_attention = BertSelfAttention::load(vb.pp("self"), config)?; + let self_output = BertSelfOutput::load(vb.pp("output"), config)?; + Ok(Self { + self_attention, + self_output, + span: tracing::span!(tracing::Level::TRACE, "attn"), + }) + } + + fn forward(&self, hidden_states: &Tensor, attention_mask: &Tensor) -> Result { + let _enter = self.span.enter(); + let self_outputs = self.self_attention.forward(hidden_states, attention_mask)?; + let attention_output = self.self_output.forward(&self_outputs, hidden_states)?; + Ok(attention_output) + } +} + +// https://github.com/huggingface/transformers/blob/6eedfa6dd15dc1e22a55ae036f681914e5a0d9a1/src/transformers/models/bert/modeling_bert.py#L441 +#[derive(Clone)] +struct BertIntermediate { + dense: Linear, + intermediate_act: HiddenActLayer, + span: tracing::Span, +} + +impl BertIntermediate { + fn load(vb: VarBuilder, config: &Config) -> Result { + let dense = linear(config.hidden_size, config.intermediate_size, vb.pp("dense"))?; + Ok(Self { + dense, + intermediate_act: HiddenActLayer::new(config.hidden_act), + span: tracing::span!(tracing::Level::TRACE, "inter"), + }) + } +} + +impl Module for BertIntermediate { + fn forward(&self, hidden_states: &Tensor) -> Result { + let _enter = self.span.enter(); + let hidden_states = self.dense.forward(hidden_states)?; + let ys = self.intermediate_act.forward(&hidden_states)?; + Ok(ys) + } +} + +// https://github.com/huggingface/transformers/blob/6eedfa6dd15dc1e22a55ae036f681914e5a0d9a1/src/transformers/models/bert/modeling_bert.py#L456 +#[derive(Clone)] +struct BertOutput { + dense: Linear, + layer_norm: LayerNorm, + dropout: Dropout, + span: tracing::Span, +} + +impl BertOutput { + fn load(vb: VarBuilder, config: &Config) -> Result { + let dense = linear(config.intermediate_size, config.hidden_size, vb.pp("dense"))?; + let layer_norm = layer_norm( + config.hidden_size, + config.layer_norm_eps, + vb.pp("LayerNorm"), + )?; + let dropout = Dropout::new(config.hidden_dropout_prob); + Ok(Self { + dense, + layer_norm, + dropout, + span: tracing::span!(tracing::Level::TRACE, "out"), + }) + } + + fn forward(&self, hidden_states: &Tensor, input_tensor: &Tensor) -> Result { + let _enter = self.span.enter(); + let hidden_states = self.dense.forward(hidden_states)?; + let hidden_states = self.dropout.forward(&hidden_states)?; + self.layer_norm.forward(&(hidden_states + input_tensor)?) + } +} + +// https://github.com/huggingface/transformers/blob/6eedfa6dd15dc1e22a55ae036f681914e5a0d9a1/src/transformers/models/bert/modeling_bert.py#L470 +#[derive(Clone)] +pub struct BertLayer { + attention: BertAttention, + intermediate: BertIntermediate, + output: BertOutput, + span: tracing::Span, +} + +impl BertLayer { + fn load(vb: VarBuilder, config: &Config) -> Result { + let attention = BertAttention::load(vb.pp("attention"), config)?; + let intermediate = BertIntermediate::load(vb.pp("intermediate"), config)?; + let output = BertOutput::load(vb.pp("output"), config)?; + Ok(Self { + attention, + intermediate, + output, + span: tracing::span!(tracing::Level::TRACE, "layer"), + }) + } + + fn forward(&self, hidden_states: &Tensor, attention_mask: &Tensor) -> Result { + let _enter = self.span.enter(); + let attention_output = self.attention.forward(hidden_states, attention_mask)?; + // TODO: Support cross-attention? + // https://github.com/huggingface/transformers/blob/6eedfa6dd15dc1e22a55ae036f681914e5a0d9a1/src/transformers/models/bert/modeling_bert.py#L523 + // TODO: Support something similar to `apply_chunking_to_forward`? + let intermediate_output = self.intermediate.forward(&attention_output)?; + let layer_output = self + .output + .forward(&intermediate_output, &attention_output)?; + Ok(layer_output) + } +} + +// https://github.com/huggingface/transformers/blob/6eedfa6dd15dc1e22a55ae036f681914e5a0d9a1/src/transformers/models/bert/modeling_bert.py#L556 +#[derive(Clone)] +pub struct BertEncoder { + pub layers: Vec, + span: tracing::Span, +} + +impl BertEncoder { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let layers = (0..config.num_hidden_layers) + .map(|index| BertLayer::load(vb.pp(format!("layer.{index}")), config)) + .collect::>>()?; + let span = tracing::span!(tracing::Level::TRACE, "encoder"); + Ok(BertEncoder { layers, span }) + } + + pub fn forward(&self, hidden_states: &Tensor, attention_mask: &Tensor) -> Result { + let _enter = self.span.enter(); + let mut hidden_states = hidden_states.clone(); + // Use a loop rather than a fold as it's easier to modify when adding debug/... + for layer in self.layers.iter() { + hidden_states = layer.forward(&hidden_states, attention_mask)? + } + Ok(hidden_states) + } +} + +// https://github.com/huggingface/transformers/blob/6eedfa6dd15dc1e22a55ae036f681914e5a0d9a1/src/transformers/models/bert/modeling_bert.py#L874 +pub struct BertModel { + embeddings: BertEmbeddings, + encoder: BertEncoder, + pub device: Device, + span: tracing::Span, +} + +impl BertModel { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let (embeddings, encoder) = match ( + BertEmbeddings::load(vb.pp("embeddings"), config), + BertEncoder::load(vb.pp("encoder"), config), + ) { + (Ok(embeddings), Ok(encoder)) => (embeddings, encoder), + (Err(err), _) | (_, Err(err)) => { + if let Some(model_type) = &config.model_type { + if let (Ok(embeddings), Ok(encoder)) = ( + BertEmbeddings::load(vb.pp(format!("{model_type}.embeddings")), config), + BertEncoder::load(vb.pp(format!("{model_type}.encoder")), config), + ) { + (embeddings, encoder) + } else { + return Err(err); + } + } else { + return Err(err); + } + } + }; + Ok(Self { + embeddings, + encoder, + device: vb.device().clone(), + span: tracing::span!(tracing::Level::TRACE, "model"), + }) + } + + pub fn forward( + &self, + input_ids: &Tensor, + token_type_ids: &Tensor, + attention_mask: Option<&Tensor>, + ) -> Result { + let _enter = self.span.enter(); + let embedding_output = self.embeddings.forward(input_ids, token_type_ids)?; + let attention_mask = match attention_mask { + Some(attention_mask) => attention_mask.clone(), + None => input_ids.ones_like()?, + }; + let dtype = embedding_output.dtype(); + // https://github.com/huggingface/transformers/blob/6eedfa6dd15dc1e22a55ae036f681914e5a0d9a1/src/transformers/models/bert/modeling_bert.py#L995 + let attention_mask = get_extended_attention_mask(&attention_mask, dtype)?; + let sequence_output = self.encoder.forward(&embedding_output, &attention_mask)?; + Ok(sequence_output) + } +} + +fn get_extended_attention_mask(attention_mask: &Tensor, dtype: DType) -> Result { + let attention_mask = match attention_mask.rank() { + 3 => attention_mask.unsqueeze(1)?, + 2 => attention_mask.unsqueeze(1)?.unsqueeze(1)?, + _ => candle::bail!("Wrong shape for input_ids or attention_mask"), + }; + let attention_mask = attention_mask.to_dtype(dtype)?; + // torch.finfo(dtype).min + (attention_mask.ones_like()? - &attention_mask)?.broadcast_mul( + &Tensor::try_from(f32::MIN)? + .to_device(attention_mask.device())? + .to_dtype(dtype)?, + ) +} + +//https://github.com/huggingface/transformers/blob/1bd604d11c405dfb8b78bda4062d88fc75c17de0/src/transformers/models/bert/modeling_bert.py#L752-L766 +struct BertPredictionHeadTransform { + dense: Linear, + activation: HiddenActLayer, + layer_norm: LayerNorm, +} + +impl BertPredictionHeadTransform { + fn load(vb: VarBuilder, config: &Config) -> Result { + let dense = linear(config.hidden_size, config.hidden_size, vb.pp("dense"))?; + let activation = HiddenActLayer::new(config.hidden_act); + let layer_norm = layer_norm( + config.hidden_size, + config.layer_norm_eps, + vb.pp("LayerNorm"), + )?; + Ok(Self { + dense, + activation, + layer_norm, + }) + } +} + +impl Module for BertPredictionHeadTransform { + fn forward(&self, hidden_states: &Tensor) -> Result { + let hidden_states = self + .activation + .forward(&self.dense.forward(hidden_states)?)?; + self.layer_norm.forward(&hidden_states) + } +} + +// https://github.com/huggingface/transformers/blob/1bd604d11c405dfb8b78bda4062d88fc75c17de0/src/transformers/models/bert/modeling_bert.py#L769C1-L790C1 +pub struct BertLMPredictionHead { + transform: BertPredictionHeadTransform, + decoder: Linear, +} + +impl BertLMPredictionHead { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let transform = BertPredictionHeadTransform::load(vb.pp("transform"), config)?; + let decoder = linear(config.hidden_size, config.vocab_size, vb.pp("decoder"))?; + Ok(Self { transform, decoder }) + } +} + +impl Module for BertLMPredictionHead { + fn forward(&self, hidden_states: &Tensor) -> Result { + self.decoder + .forward(&self.transform.forward(hidden_states)?) + } +} + +// https://github.com/huggingface/transformers/blob/1bd604d11c405dfb8b78bda4062d88fc75c17de0/src/transformers/models/bert/modeling_bert.py#L792 +pub struct BertOnlyMLMHead { + predictions: BertLMPredictionHead, +} + +impl BertOnlyMLMHead { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let predictions = BertLMPredictionHead::load(vb.pp("predictions"), config)?; + Ok(Self { predictions }) + } +} + +impl Module for BertOnlyMLMHead { + fn forward(&self, sequence_output: &Tensor) -> Result { + self.predictions.forward(sequence_output) + } +} + +pub struct BertForMaskedLM { + bert: BertModel, + cls: BertOnlyMLMHead, +} + +impl BertForMaskedLM { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let bert = BertModel::load(vb.pp("bert"), config)?; + let cls = BertOnlyMLMHead::load(vb.pp("cls"), config)?; + Ok(Self { bert, cls }) + } + + pub fn forward( + &self, + input_ids: &Tensor, + token_type_ids: &Tensor, + attention_mask: Option<&Tensor>, + ) -> Result { + let sequence_output = self + .bert + .forward(input_ids, token_type_ids, attention_mask)?; + self.cls.forward(&sequence_output) + } +} diff --git a/patches/candle-transformers/src/models/bigcode.rs b/patches/candle-transformers/src/models/bigcode.rs new file mode 100644 index 0000000000..ed63e4d73d --- /dev/null +++ b/patches/candle-transformers/src/models/bigcode.rs @@ -0,0 +1,367 @@ +//! BigCode implementation in Rust based on the GPT-BigCode model. +//! +//! [StarCoder/BigCode](https://huggingface.co/bigcode/starcoderbase-1b) is a LLM +//! model specialized to code generation. The initial model was trained on 80 +//! programming languages. See "StarCoder: A State-of-the-Art LLM for Code", Mukherjee et al. 2023 +//! - [Arxiv](https://arxiv.org/abs/2305.06161) +//! - [GitHub](https://github.com/bigcode-project/starcoder) +//! +//! ## Running some example +//! +//! ```bash +//! cargo run --example bigcode --release -- --prompt "fn fact(n: u64) -> u64" +//! +//! > fn fact(n: u64) -> u64 { +//! > if n == 0 { +//! > 1 +//! > } else { +//! > n * fact(n - 1) +//! > } +//! > } +//! ``` +//! + +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::{embedding, linear_b as linear, Embedding, LayerNorm, Linear, Module, VarBuilder}; + +fn layer_norm(size: usize, eps: f64, vb: VarBuilder) -> Result { + let weight = vb.get(size, "weight")?; + let bias = vb.get(size, "bias")?; + Ok(LayerNorm::new(weight, bias, eps)) +} + +fn make_causal_mask(t: usize, device: &Device) -> Result { + 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)?; + Ok(mask) +} + +#[derive(Debug)] +pub struct Config { + pub vocab_size: usize, + // max_position_embeddings aka n_positions + pub max_position_embeddings: usize, + // num_hidden_layers aka n_layer + pub num_hidden_layers: usize, + // hidden_size aka n_embd + pub hidden_size: usize, + pub layer_norm_epsilon: f64, + pub n_inner: Option, + // num_attention_heads aka n_head + pub num_attention_heads: usize, + pub multi_query: bool, + pub use_cache: bool, +} + +impl Config { + #[allow(dead_code)] + pub fn starcoder_1b() -> Self { + Self { + vocab_size: 49152, + max_position_embeddings: 8192, + num_hidden_layers: 24, + hidden_size: 2048, + layer_norm_epsilon: 1e-5, + n_inner: Some(8192), + num_attention_heads: 16, + multi_query: true, + use_cache: true, + } + } + + #[allow(dead_code)] + pub fn starcoder_3b() -> Self { + Self { + vocab_size: 49152, + max_position_embeddings: 8192, + num_hidden_layers: 36, + hidden_size: 2816, + layer_norm_epsilon: 1e-5, + n_inner: Some(11264), + num_attention_heads: 22, + multi_query: true, + use_cache: true, + } + } + + #[allow(dead_code)] + pub fn starcoder_7b() -> Self { + Self { + vocab_size: 49152, + max_position_embeddings: 8192, + num_hidden_layers: 42, + hidden_size: 4096, + layer_norm_epsilon: 1e-5, + n_inner: Some(16384), + num_attention_heads: 32, + multi_query: true, + use_cache: true, + } + } + + #[allow(dead_code)] + pub fn starcoder() -> Self { + Self { + vocab_size: 49152, + max_position_embeddings: 8192, + num_hidden_layers: 40, + hidden_size: 6144, + layer_norm_epsilon: 1e-5, + n_inner: Some(24576), + num_attention_heads: 48, + multi_query: true, + use_cache: true, + } + } +} + +struct Attention { + c_attn: Linear, + c_proj: Linear, + kv_cache: Option, + use_cache: bool, + embed_dim: usize, + kv_dim: usize, + num_heads: usize, + head_dim: usize, + multi_query: bool, +} + +impl Attention { + pub fn load(vb: VarBuilder, cfg: &Config) -> Result { + let hidden_size = cfg.hidden_size; + let head_dim = hidden_size / cfg.num_attention_heads; + let kv_heads = if cfg.multi_query { + 1 + } else { + cfg.num_attention_heads + }; + let kv_dim = kv_heads * head_dim; + let c_attn = linear(hidden_size, hidden_size + 2 * kv_dim, true, vb.pp("c_attn"))?; + let c_proj = linear(hidden_size, hidden_size, true, vb.pp("c_proj"))?; + Ok(Self { + c_proj, + c_attn, + embed_dim: hidden_size, + kv_cache: None, + use_cache: cfg.use_cache, + kv_dim, + head_dim, + num_heads: cfg.num_attention_heads, + multi_query: cfg.multi_query, + }) + } + + fn attn( + &self, + query: &Tensor, + key: &Tensor, + value: &Tensor, + attention_mask: &Tensor, + ) -> Result { + if query.dtype() != DType::F32 { + // If we start supporting f16 models, we may need the upcasting scaling bits. + // https://github.com/huggingface/transformers/blob/a0042379269bea9182c1f87e6b2eee4ba4c8cce8/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py#L133 + candle::bail!("upcasting is not supported {:?}", query.dtype()) + } + let scale_factor = 1f64 / (self.head_dim as f64).sqrt(); + let initial_query_shape = query.shape(); + let key_len = key.dim(D::Minus1)?; + let (query, key, attn_shape, attn_view) = if self.multi_query { + let (b_sz, query_len, _) = query.dims3()?; + let query = query.reshape((b_sz, query_len * self.num_heads, self.head_dim))?; + let attn_shape = (b_sz, query_len, self.num_heads, key_len); + let attn_view = (b_sz, query_len * self.num_heads, key_len); + (query, key.clone(), attn_shape, attn_view) + } else { + let (b_sz, _num_heads, query_len, _head_dim) = query.dims4()?; + let query = query.reshape((b_sz, query_len * self.num_heads, self.head_dim))?; + let key = key.reshape((b_sz * self.num_heads, self.head_dim, key_len))?; + let attn_shape = (b_sz, self.num_heads, query_len, key_len); + let attn_view = (b_sz * self.num_heads, query_len, key_len); + (query, key, attn_shape, attn_view) + }; + + let attn_weights = + (query.matmul(&key.contiguous()?)? * scale_factor)?.reshape(attn_shape)?; + let attention_mask = attention_mask.broadcast_as(attn_shape)?; + let mask_value = + Tensor::new(f32::NEG_INFINITY, query.device())?.broadcast_as(attn_shape)?; + let attn_weights = attention_mask.where_cond(&attn_weights, &mask_value)?; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + let value = value.contiguous()?; + let attn_output = if self.multi_query { + attn_weights + .reshape(attn_view)? + .matmul(&value)? + .reshape(initial_query_shape)? + } else { + attn_weights.matmul(&value)? + }; + Ok(attn_output) + } + + fn forward(&mut self, hidden_states: &Tensor, attention_mask: &Tensor) -> Result { + let qkv = self.c_attn.forward(hidden_states)?; + let (query, key_value) = if self.multi_query { + let query = qkv.i((.., .., ..self.embed_dim))?; + let key_value = qkv.i((.., .., self.embed_dim..self.embed_dim + 2 * self.kv_dim))?; + (query, key_value) + } else { + let mut dims = qkv.dims().to_vec(); + dims.pop(); + dims.push(self.embed_dim); + dims.push(self.head_dim * 3); + let qkv = qkv.reshape(dims)?.transpose(1, 2)?; + let query = qkv.i((.., .., .., ..self.head_dim))?; + let key_value = qkv.i((.., .., .., self.head_dim..3 * self.head_dim))?; + (query, key_value) + }; + let mut key_value = key_value; + if self.use_cache { + if let Some(kv_cache) = &self.kv_cache { + // TODO: we could trim the tensors to MAX_SEQ_LEN so that this would work for + // arbitrarily large sizes. + key_value = Tensor::cat(&[kv_cache, &key_value], D::Minus2)?.contiguous()?; + } + self.kv_cache = Some(key_value.clone()) + } + + let key = key_value.narrow(D::Minus1, 0, self.head_dim)?; + let value = key_value.narrow(D::Minus1, self.head_dim, self.head_dim)?; + let attn_output = self.attn(&query, &key.t()?, &value, attention_mask)?; + let attn_output = if self.multi_query { + attn_output + } else { + attn_output + .transpose(1, 2)? + .reshape(hidden_states.shape())? + }; + let attn_output = self.c_proj.forward(&attn_output)?; + Ok(attn_output) + } +} + +struct Mlp { + c_fc: Linear, + c_proj: Linear, +} + +impl Mlp { + fn load(inner_dim: usize, vb: VarBuilder, cfg: &Config) -> Result { + let c_fc = linear(cfg.hidden_size, inner_dim, true, vb.pp("c_fc"))?; + let c_proj = linear(inner_dim, cfg.hidden_size, true, vb.pp("c_proj"))?; + Ok(Self { c_fc, c_proj }) + } + + fn forward(&mut self, hidden_states: &Tensor) -> Result { + let hidden_states = self.c_fc.forward(hidden_states)?.gelu()?; + let hidden_states = self.c_proj.forward(&hidden_states)?; + Ok(hidden_states) + } +} + +// TODO: Add cross-attention? +struct Block { + ln_1: LayerNorm, + attn: Attention, + ln_2: LayerNorm, + mlp: Mlp, +} + +impl Block { + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let hidden_size = cfg.hidden_size; + let inner_dim = cfg.n_inner.unwrap_or(4 * hidden_size); + let ln_1 = layer_norm(hidden_size, cfg.layer_norm_epsilon, vb.pp("ln_1"))?; + let attn = Attention::load(vb.pp("attn"), cfg)?; + let ln_2 = layer_norm(hidden_size, cfg.layer_norm_epsilon, vb.pp("ln_2"))?; + let mlp = Mlp::load(inner_dim, vb.pp("mlp"), cfg)?; + Ok(Self { + ln_1, + attn, + ln_2, + mlp, + }) + } + + fn forward(&mut self, hidden_states: &Tensor, attention_mask: &Tensor) -> Result { + let residual = hidden_states; + let hidden_states = self.ln_1.forward(hidden_states)?; + let attn_outputs = self.attn.forward(&hidden_states, attention_mask)?; + let hidden_states = (&attn_outputs + residual)?; + let residual = &hidden_states; + let hidden_states = self.ln_2.forward(&hidden_states)?; + let hidden_states = self.mlp.forward(&hidden_states)?; + let hidden_states = (&hidden_states + residual)?; + Ok(hidden_states) + } +} + +pub struct GPTBigCode { + wte: Embedding, + wpe: Embedding, + blocks: Vec, + ln_f: LayerNorm, + lm_head: Linear, + bias: Tensor, + config: Config, +} + +impl GPTBigCode { + pub fn config(&self) -> &Config { + &self.config + } + + pub fn load(vb: VarBuilder, cfg: Config) -> Result { + let hidden_size = cfg.hidden_size; + let vb_t = vb.pp("transformer"); + let wte = embedding(cfg.vocab_size, hidden_size, vb_t.pp("wte"))?; + let wpe = embedding(cfg.max_position_embeddings, hidden_size, vb_t.pp("wpe"))?; + let blocks = (0..cfg.num_hidden_layers) + .map(|i| Block::load(vb_t.pp(format!("h.{i}")), &cfg)) + .collect::>>()?; + let ln_f = layer_norm(hidden_size, cfg.layer_norm_epsilon, vb_t.pp("ln_f"))?; + let lm_head = linear(hidden_size, cfg.vocab_size, false, vb_t.pp("wte"))?; + let bias = make_causal_mask(cfg.max_position_embeddings, vb.device())?; + Ok(Self { + wte, + wpe, + blocks, + lm_head, + ln_f, + bias, + config: cfg, + }) + } + + pub fn forward(&mut self, input_ids: &Tensor, past_len: usize) -> Result { + let dev = input_ids.device(); + let (b_sz, seq_len) = input_ids.dims2()?; + + let key_len = past_len + seq_len; + let attention_mask = self.bias.i((past_len..key_len, ..key_len))?.unsqueeze(0)?; + // MQA models: (batch_size, query_length, n_heads, key_length) + // MHA models: (batch_size, n_heads, query_length, key_length) + let seq_len_dim = if self.config.multi_query { 2 } else { 1 }; + let attention_mask = attention_mask.unsqueeze(seq_len_dim)?; + + let position_ids = Tensor::arange(past_len as u32, (past_len + seq_len) as u32, dev)?; + let position_ids = position_ids.unsqueeze(0)?.broadcast_as((b_sz, seq_len))?; + let input_embeds = self.wte.forward(input_ids)?; + let position_embeds = self.wpe.forward(&position_ids)?; + + let mut hidden_states = (&input_embeds + &position_embeds)?; + for block in self.blocks.iter_mut() { + hidden_states = block.forward(&hidden_states, &attention_mask)?; + } + let hidden_states = self.ln_f.forward(&hidden_states)?; + let hidden_states = hidden_states + .reshape((b_sz, seq_len, self.config.hidden_size))? + .narrow(1, seq_len - 1, 1)?; + let logits = self.lm_head.forward(&hidden_states)?.squeeze(1)?; + Ok(logits) + } +} diff --git a/patches/candle-transformers/src/models/blip.rs b/patches/candle-transformers/src/models/blip.rs new file mode 100644 index 0000000000..a391daacbf --- /dev/null +++ b/patches/candle-transformers/src/models/blip.rs @@ -0,0 +1,317 @@ +//! Based on the BLIP paper from Salesforce Research. +//! +//! The blip-image-captioning model can generate captions for an input image. +//! +//! - ⚡ [Interactive Wasm Example](https://huggingface.co/spaces/radames/Candle-BLIP-Image-Captioning) +//! - 💻 [GH Link](https://github.com/salesforce/BLIP) +//! - 🤗 [HF Link](https://huggingface.co/Salesforce/blip-image-captioning-base) +//! - 📝 [Paper](https://arxiv.org/abs/2201.12086) +//! + +use super::blip_text; +use super::with_tracing::{conv2d, linear, Conv2d, Linear}; +use candle::{Module, Result, Tensor, D}; +use candle_nn::{layer_norm, Conv2dConfig, LayerNorm, VarBuilder}; +use serde::Deserialize; + +#[derive(Debug, Clone, Deserialize)] +pub struct VisionConfig { + pub hidden_size: usize, + pub intermediate_size: usize, + pub projection_dim: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub image_size: usize, + pub patch_size: usize, + pub hidden_act: candle_nn::Activation, + pub layer_norm_eps: f64, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct Config { + pub text_config: blip_text::Config, + pub vision_config: VisionConfig, + pub projection_dim: usize, + pub image_text_hidden_size: usize, +} + +impl Config { + pub fn image_captioning_large() -> Self { + let text_config = blip_text::Config { + vocab_size: 30524, + hidden_size: 768, + encoder_hidden_size: 1024, + intermediate_size: 3072, + projection_dim: 768, + num_hidden_layers: 12, + num_attention_heads: 12, + max_position_embeddings: 512, + hidden_act: candle_nn::Activation::Gelu, + layer_norm_eps: 1e-12, + is_decoder: true, + }; + let vision_config = VisionConfig { + hidden_size: 1024, + intermediate_size: 4096, + projection_dim: 512, + num_hidden_layers: 24, + num_attention_heads: 16, + image_size: 384, + patch_size: 16, + hidden_act: candle_nn::Activation::Gelu, + layer_norm_eps: 1e-5, + }; + Self { + text_config, + vision_config, + projection_dim: 512, + image_text_hidden_size: 256, + } + } +} + +#[derive(Debug, Clone)] +struct VisionEmbeddings { + class_embedding: Tensor, + patch_embedding: Conv2d, + position_embedding: Tensor, +} + +impl VisionEmbeddings { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let class_embedding = vb.get((1, 1, cfg.hidden_size), "class_embedding")?; + let conv_cfg = Conv2dConfig { + stride: cfg.patch_size, + ..Default::default() + }; + let patch_embedding = conv2d( + 3, + cfg.hidden_size, + cfg.patch_size, + conv_cfg, + vb.pp("patch_embedding"), + )?; + let num_patches1 = cfg.image_size / cfg.patch_size; + let num_patches = num_patches1 * num_patches1; + let num_positions = num_patches + 1; + let position_embedding = + vb.get((1, num_positions, cfg.hidden_size), "position_embedding")?; + Ok(Self { + class_embedding, + patch_embedding, + position_embedding, + }) + } +} + +impl Module for VisionEmbeddings { + fn forward(&self, xs: &Tensor) -> Result { + let target_dtype = xs.dtype(); + let b_size = xs.dim(0)?; + let patch_embeds = xs.apply(&self.patch_embedding)?.flatten_from(2)?.t()?; + let d = self.class_embedding.dim(D::Minus1)?; + let class_embeds = self + .class_embedding + .broadcast_as((b_size, 1, d))? + .to_dtype(target_dtype)?; + let embeddings = Tensor::cat(&[&class_embeds, &patch_embeds], 1)?; + let position_embedding = self.position_embedding.narrow(1, 0, embeddings.dim(1)?)?; + embeddings.broadcast_add(&position_embedding) + } +} + +#[derive(Debug, Clone)] +struct Attention { + qkv: Linear, + projection: Linear, + scale: f64, + num_heads: usize, +} + +impl Attention { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let embed_dim = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let head_dim = embed_dim / num_heads; + let scale = 1f64 / (head_dim as f64).sqrt(); + let qkv = linear(embed_dim, 3 * embed_dim, vb.pp("qkv"))?; + let projection = linear(embed_dim, embed_dim, vb.pp("projection"))?; + Ok(Self { + qkv, + projection, + scale, + num_heads, + }) + } + + fn forward(&self, xs: &Tensor, attn_mask: Option<&Tensor>) -> Result { + let (b_sz, tgt_len, embed_dim) = xs.dims3()?; + let mixed_qkv = xs + .apply(&self.qkv)? + .reshape((b_sz, tgt_len, 3, self.num_heads, embed_dim / self.num_heads))? + .permute((2, 0, 3, 1, 4))?; + let query = mixed_qkv.get(0)?; + let key = mixed_qkv.get(1)?; + let value = mixed_qkv.get(2)?; + let attention_scores = query.matmul(&key.t()?)?; + let attention_scores = (attention_scores * self.scale)?; + let attention_probs = candle_nn::ops::softmax_last_dim(&attention_scores)?; + let attention_probs = match attn_mask { + None => attention_probs, + Some(attn_mask) => (attention_probs * attn_mask)?, + }; + attention_probs + .matmul(&value)? + .permute((0, 2, 1, 3))? + .flatten_from(D::Minus2)? + .apply(&self.projection) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + activation_fn: candle_nn::Activation, + fc1: Linear, + fc2: Linear, +} + +impl MLP { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let fc1 = linear(cfg.hidden_size, cfg.intermediate_size, vb.pp("fc1"))?; + let fc2 = linear(cfg.intermediate_size, cfg.hidden_size, vb.pp("fc2"))?; + Ok(Self { + activation_fn: cfg.hidden_act, + fc1, + fc2, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.fc1)? + .apply(&self.activation_fn)? + .apply(&self.fc2) + } +} + +#[derive(Debug, Clone)] +struct EncoderLayer { + self_attn: Attention, + layer_norm1: LayerNorm, + mlp: MLP, + layer_norm2: LayerNorm, +} + +impl EncoderLayer { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let embed_dim = cfg.hidden_size; + let self_attn = Attention::new(cfg, vb.pp("self_attn"))?; + let layer_norm1 = layer_norm(embed_dim, cfg.layer_norm_eps, vb.pp("layer_norm1"))?; + let layer_norm2 = layer_norm(embed_dim, cfg.layer_norm_eps, vb.pp("layer_norm2"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + Ok(Self { + self_attn, + layer_norm1, + mlp, + layer_norm2, + }) + } + + fn forward(&self, xs: &Tensor, attention_mask: Option<&Tensor>) -> Result { + let residual = xs; + let xs = xs.apply(&self.layer_norm1)?; + let xs = self.self_attn.forward(&xs, attention_mask)?; + let xs = (xs + residual)?; + + let residual = &xs; + let xs = xs.apply(&self.layer_norm2)?.apply(&self.mlp)?; + xs + residual + } +} + +#[derive(Debug, Clone)] +struct Encoder { + layers: Vec, +} + +impl Encoder { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb = vb.pp("layers"); + for i in 0..cfg.num_hidden_layers { + let layer = EncoderLayer::new(cfg, vb.pp(i))?; + layers.push(layer) + } + Ok(Self { layers }) + } + + fn forward(&self, xs: &Tensor, attention_mask: Option<&Tensor>) -> Result { + let mut xs = xs.clone(); + for layer in self.layers.iter() { + xs = layer.forward(&xs, attention_mask)? + } + Ok(xs) + } +} + +#[derive(Debug, Clone)] +pub struct VisionModel { + embeddings: VisionEmbeddings, + encoder: Encoder, + post_layernorm: LayerNorm, +} + +impl VisionModel { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let embeddings = VisionEmbeddings::new(cfg, vb.pp("embeddings"))?; + let encoder = Encoder::new(cfg, vb.pp("encoder"))?; + let post_layernorm = + layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("post_layernorm"))?; + Ok(Self { + embeddings, + encoder, + post_layernorm, + }) + } +} + +impl Module for VisionModel { + fn forward(&self, xs: &Tensor) -> Result { + let xs = xs.apply(&self.embeddings)?; + let encoder_outputs = self.encoder.forward(&xs, None)?; + // Return the last hidden state rather than pooled outputs. + encoder_outputs.apply(&self.post_layernorm) + } +} + +#[derive(Debug, Clone)] +pub struct BlipForConditionalGeneration { + vision_model: VisionModel, + text_decoder: blip_text::TextLMHeadModel, +} + +impl BlipForConditionalGeneration { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vision_model = VisionModel::new(&cfg.vision_config, vb.pp("vision_model"))?; + let text_decoder = + blip_text::TextLMHeadModel::new(&cfg.text_config, vb.pp("text_decoder"))?; + Ok(Self { + vision_model, + text_decoder, + }) + } + + pub fn vision_model(&self) -> &VisionModel { + &self.vision_model + } + + pub fn text_decoder(&mut self) -> &mut blip_text::TextLMHeadModel { + &mut self.text_decoder + } + + pub fn reset_kv_cache(&mut self) { + self.text_decoder.reset_kv_cache(); + } +} diff --git a/patches/candle-transformers/src/models/blip_text.rs b/patches/candle-transformers/src/models/blip_text.rs new file mode 100644 index 0000000000..8aeb5dbe35 --- /dev/null +++ b/patches/candle-transformers/src/models/blip_text.rs @@ -0,0 +1,497 @@ +//! Implementation of BLIP text encoder/decoder. +//! +//! - 📝 [Paper](https://arxiv.org/abs/2201.12086). BLIP: Bootstrapping Language-Image Pre-training for Unified Vision-Language Understanding and Generation" +//! +//! - ⚡ [Interactive Wasm Example](https://huggingface.co/spaces/radames/Candle-BLIP-Image-Captioning) +//! - 💻 [GH Link](https://github.com/salesforce/BLIP) +//! - 🤗 [HF Link](https://huggingface.co/Salesforce/blip-image-captioning-base) +//! - 📝 [Paper](https://arxiv.org/abs/2201.12086) +//! +use super::with_tracing::{linear, Embedding, Linear}; +use candle::{Module, Result, Tensor, D}; +use candle_nn::{layer_norm, LayerNorm, VarBuilder}; +use serde::Deserialize; + +#[derive(Debug, Clone, Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub encoder_hidden_size: usize, + pub intermediate_size: usize, + pub projection_dim: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub max_position_embeddings: usize, + pub hidden_act: candle_nn::Activation, + pub layer_norm_eps: f64, + pub is_decoder: bool, +} + +#[derive(Debug, Clone)] +struct TextEmbeddings { + word_embeddings: Embedding, + position_embeddings: Embedding, + layer_norm: LayerNorm, + position_ids: Tensor, +} + +impl TextEmbeddings { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let word_embeddings = + Embedding::new(cfg.vocab_size, cfg.hidden_size, vb.pp("word_embeddings"))?; + let position_embeddings = Embedding::new( + cfg.max_position_embeddings, + cfg.hidden_size, + vb.pp("position_embeddings"), + )?; + let layer_norm = layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("LayerNorm"))?; + let position_ids = + Tensor::arange(0, cfg.max_position_embeddings as u32, vb.device())?.unsqueeze(0)?; + Ok(Self { + word_embeddings, + position_embeddings, + layer_norm, + position_ids, + }) + } + + fn forward(&self, xs: &Tensor, past_kv_len: usize) -> Result { + let seq_len = xs.dim(1)?; + let position_ids = self.position_ids.narrow(1, past_kv_len, seq_len)?; + let embeddings = self.word_embeddings.forward(xs)?; + let position_embeddings = self.position_embeddings.forward(&position_ids)?; + (embeddings + position_embeddings)?.apply(&self.layer_norm) + } +} + +#[derive(Debug, Clone)] +struct TextSelfAttention { + query: Linear, + key: Linear, + value: Linear, + attention_head_size: usize, + num_attention_heads: usize, + attention_scale: f64, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl TextSelfAttention { + fn new(cfg: &Config, is_cross_attention: bool, vb: VarBuilder) -> Result { + let num_attention_heads = cfg.num_attention_heads; + let attention_head_size = cfg.hidden_size / num_attention_heads; + let all_head_size = cfg.num_attention_heads * attention_head_size; + let query = linear(cfg.hidden_size, all_head_size, vb.pp("query"))?; + let in_size = if is_cross_attention { + cfg.encoder_hidden_size + } else { + cfg.hidden_size + }; + let key = linear(in_size, all_head_size, vb.pp("key"))?; + let value = linear(in_size, all_head_size, vb.pp("value"))?; + let attention_scale = 1f64 / (attention_head_size as f64).sqrt(); + Ok(Self { + query, + key, + value, + attention_head_size, + num_attention_heads, + attention_scale, + kv_cache: None, + }) + } + + fn transpose_for_scores(&self, xs: &Tensor) -> Result { + let (b_size, seq_len, _) = xs.dims3()?; + xs.reshape(( + b_size, + seq_len, + self.num_attention_heads, + self.attention_head_size, + ))? + .permute((0, 2, 1, 3)) + } + + fn reset_kv_cache(&mut self) { + self.kv_cache = None + } + + fn forward( + &mut self, + xs: &Tensor, + encoder_hidden_states: Option<&Tensor>, + attention_mask: Option<&Tensor>, + ) -> Result { + let query = self + .transpose_for_scores(&self.query.forward(xs)?)? + .contiguous()?; + let (key, value) = match encoder_hidden_states { + None => { + let key = self.transpose_for_scores(&self.key.forward(xs)?)?; + let value = self.transpose_for_scores(&self.value.forward(xs)?)?; + let (key, value) = match &self.kv_cache { + None => (key, value), + Some((prev_key, prev_value)) => { + let key = Tensor::cat(&[prev_key, &key], 2)?; + let value = Tensor::cat(&[prev_value, &value], 2)?; + (key, value) + } + }; + self.kv_cache = Some((key.clone(), value.clone())); + (key, value) + } + Some(xs) => { + let key = self.transpose_for_scores(&self.key.forward(xs)?)?; + let value = self.transpose_for_scores(&self.value.forward(xs)?)?; + // no kv-cache in this case, but the results could probably be memoized. + (key, value) + } + }; + let key = key.contiguous()?; + let value = value.contiguous()?; + let attention_scores = query.matmul(&key.t()?)?; + let attention_scores = (attention_scores * self.attention_scale)?; + let attention_scores = match attention_mask { + Some(mask) => attention_scores.broadcast_add(mask)?, + None => attention_scores, + }; + let attention_probs = candle_nn::ops::softmax_last_dim(&attention_scores)?; + attention_probs + .matmul(&value)? + .permute((0, 2, 1, 3))? + .flatten_from(D::Minus2) + } +} + +#[derive(Debug, Clone)] +struct TextSelfOutput { + dense: Linear, + layer_norm: LayerNorm, +} + +impl TextSelfOutput { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.hidden_size, cfg.hidden_size, vb.pp("dense"))?; + let layer_norm = layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("LayerNorm"))?; + Ok(Self { dense, layer_norm }) + } + + fn forward(&self, xs: &Tensor, input_tensor: &Tensor) -> Result { + (xs.apply(&self.dense) + input_tensor)?.apply(&self.layer_norm) + } +} + +#[derive(Debug, Clone)] +struct TextAttention { + self_: TextSelfAttention, + output: TextSelfOutput, +} + +impl TextAttention { + fn new(cfg: &Config, is_cross_attention: bool, vb: VarBuilder) -> Result { + let self_ = TextSelfAttention::new(cfg, is_cross_attention, vb.pp("self"))?; + let output = TextSelfOutput::new(cfg, vb.pp("output"))?; + Ok(Self { self_, output }) + } + + fn reset_kv_cache(&mut self) { + self.self_.reset_kv_cache() + } + + fn forward( + &mut self, + xs: &Tensor, + encoder_hidden_states: Option<&Tensor>, + attention_mask: Option<&Tensor>, + ) -> Result { + let self_outputs = self + .self_ + .forward(xs, encoder_hidden_states, attention_mask)?; + self.output.forward(&self_outputs, xs) + } +} + +#[derive(Debug, Clone)] +struct TextIntermediate { + dense: Linear, + intermediate_act_fn: candle_nn::Activation, +} + +impl TextIntermediate { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.hidden_size, cfg.intermediate_size, vb.pp("dense"))?; + Ok(Self { + dense, + intermediate_act_fn: cfg.hidden_act, + }) + } +} + +impl Module for TextIntermediate { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.dense)?.apply(&self.intermediate_act_fn) + } +} + +#[derive(Debug, Clone)] +struct TextOutput { + dense: Linear, + layer_norm: LayerNorm, +} + +impl TextOutput { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.intermediate_size, cfg.hidden_size, vb.pp("dense"))?; + let layer_norm = layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("LayerNorm"))?; + Ok(Self { dense, layer_norm }) + } + + fn forward(&self, xs: &Tensor, input_tensor: &Tensor) -> Result { + (xs.apply(&self.dense)? + input_tensor)?.apply(&self.layer_norm) + } +} + +#[derive(Debug, Clone)] +struct TextLayer { + attention: TextAttention, + cross_attention: Option, + intermediate: TextIntermediate, + output: TextOutput, +} + +impl TextLayer { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let attention = TextAttention::new(cfg, false, vb.pp("attention"))?; + let cross_attention = if cfg.is_decoder { + Some(TextAttention::new(cfg, true, vb.pp("crossattention"))?) + } else { + None + }; + let intermediate = TextIntermediate::new(cfg, vb.pp("intermediate"))?; + let output = TextOutput::new(cfg, vb.pp("output"))?; + Ok(Self { + attention, + cross_attention, + intermediate, + output, + }) + } + + fn reset_kv_cache(&mut self) { + self.attention.reset_kv_cache(); + if let Some(ca) = &mut self.cross_attention { + ca.reset_kv_cache() + } + } + + fn forward( + &mut self, + xs: &Tensor, + encoder_hidden_states: &Tensor, + attention_mask: &Tensor, + ) -> Result { + let attention_output = self.attention.forward(xs, None, Some(attention_mask))?; + let attention_output = match &mut self.cross_attention { + Some(ca) => ca.forward(&attention_output, Some(encoder_hidden_states), None)?, + None => candle::bail!("expected some cross-attn"), + }; + let intermediate_output = self.intermediate.forward(&attention_output)?; + self.output.forward(&intermediate_output, &attention_output) + } +} + +#[derive(Debug, Clone)] +struct TextEncoder { + layers: Vec, +} + +impl TextEncoder { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb = vb.pp("layer"); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + for i in 0..cfg.num_hidden_layers { + let layer = TextLayer::new(cfg, vb.pp(i))?; + layers.push(layer) + } + Ok(Self { layers }) + } + + fn reset_kv_cache(&mut self) { + self.layers.iter_mut().for_each(|l| l.reset_kv_cache()) + } + + fn forward( + &mut self, + xs: &Tensor, + encoder_hidden_states: &Tensor, + attention_mask: &Tensor, + ) -> Result { + let mut xs = xs.clone(); + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, encoder_hidden_states, attention_mask)? + } + Ok(xs) + } +} + +#[derive(Debug, Clone)] +pub struct TextPooler { + dense: Linear, +} + +impl TextPooler { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.hidden_size, cfg.hidden_size, vb.pp("dense"))?; + Ok(Self { dense }) + } +} + +impl Module for TextPooler { + fn forward(&self, xs: &Tensor) -> Result { + xs.narrow(D::Minus1, 0, 1)? + .squeeze(D::Minus1)? + .apply(&self.dense)? + .tanh() + } +} + +#[derive(Debug, Clone)] +struct TextPredictionHeadTransform { + dense: Linear, + transform_act_fn: candle_nn::Activation, + layer_norm: LayerNorm, +} + +impl TextPredictionHeadTransform { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.hidden_size, cfg.hidden_size, vb.pp("dense"))?; + let layer_norm = layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("LayerNorm"))?; + Ok(Self { + dense, + transform_act_fn: cfg.hidden_act, + layer_norm, + }) + } +} + +impl Module for TextPredictionHeadTransform { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.dense)? + .apply(&self.transform_act_fn)? + .apply(&self.layer_norm) + } +} + +#[derive(Debug, Clone)] +struct TextLMPredictionHead { + transform: TextPredictionHeadTransform, + decoder: Linear, +} + +impl TextLMPredictionHead { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let transform = TextPredictionHeadTransform::new(cfg, vb.pp("transform"))?; + let weight = vb.get((cfg.vocab_size, cfg.hidden_size), "decoder.weight")?; + let bias = vb.get(cfg.vocab_size, "bias")?; + let decoder = Linear::from_weights(weight, Some(bias)); + Ok(Self { transform, decoder }) + } +} + +impl Module for TextLMPredictionHead { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.transform)?.apply(&self.decoder) + } +} + +#[derive(Debug, Clone)] +struct TextOnlyMLMHead { + predictions: TextLMPredictionHead, +} + +impl TextOnlyMLMHead { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let predictions = TextLMPredictionHead::new(cfg, vb.pp("predictions"))?; + Ok(Self { predictions }) + } +} + +impl Module for TextOnlyMLMHead { + fn forward(&self, xs: &Tensor) -> Result { + self.predictions.forward(xs) + } +} + +#[derive(Debug, Clone)] +struct TextModel { + embeddings: TextEmbeddings, + encoder: TextEncoder, + past_kv_len: usize, + // We do not need the pooler for caption generation +} + +impl TextModel { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let embeddings = TextEmbeddings::new(cfg, vb.pp("embeddings"))?; + let encoder = TextEncoder::new(cfg, vb.pp("encoder"))?; + Ok(Self { + embeddings, + encoder, + past_kv_len: 0, + }) + } + + fn forward( + &mut self, + input_ids: &Tensor, + encoder_hidden_states: &Tensor, + attention_mask: &Tensor, + ) -> Result { + let (_b_sz, seq_len) = input_ids.dims2()?; + let embedding_output = self.embeddings.forward(input_ids, self.past_kv_len)?; + let sequence_output = + self.encoder + .forward(&embedding_output, encoder_hidden_states, attention_mask)?; + self.past_kv_len += seq_len; + // We're interested in the sequence-output rather than the pooled-output. + Ok(sequence_output) + } + + fn reset_kv_cache(&mut self) { + self.past_kv_len = 0; + self.encoder.reset_kv_cache(); + } +} + +#[derive(Debug, Clone)] +pub struct TextLMHeadModel { + bert: TextModel, + cls: TextOnlyMLMHead, +} + +impl TextLMHeadModel { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let bert = TextModel::new(cfg, vb.pp("bert"))?; + let cls = TextOnlyMLMHead::new(cfg, vb.pp("cls"))?; + Ok(Self { bert, cls }) + } + + pub fn forward( + &mut self, + input_ids: &Tensor, + encoder_hidden_states: &Tensor, + ) -> Result { + let seq_len = input_ids.dim(1)?; + let mask: Vec<_> = (0..seq_len) + .flat_map(|i| (0..seq_len).map(move |j| if j > i { f32::NEG_INFINITY } else { 0f32 })) + .collect(); + let mask = Tensor::from_vec(mask, (seq_len, seq_len), input_ids.device())?; + let sequence_output = self.bert.forward(input_ids, encoder_hidden_states, &mask)?; + let prediction_scores = self.cls.forward(&sequence_output)?; + // return_logits is false so we don't discard the last sequence element. + Ok(prediction_scores) + } + + pub fn reset_kv_cache(&mut self) { + self.bert.reset_kv_cache() + } +} diff --git a/patches/candle-transformers/src/models/chatglm.rs b/patches/candle-transformers/src/models/chatglm.rs new file mode 100644 index 0000000000..59132c5ee7 --- /dev/null +++ b/patches/candle-transformers/src/models/chatglm.rs @@ -0,0 +1,590 @@ +//! Implementation of the ChatGLM2/3 models from THUDM. +//! +//! - 💻 [GitHub](https://github.com/THUDM/ChatGLM3) ChatGLM3: Advancing Multilingual Conversational Language Models with High-Quality Data +//! - 💻 [GitHub](https://github.com/THUDM/ChatGLM2-6B) ChatGLM2-6B. +//! +use crate::models::with_tracing::{linear_b as linear, Linear}; +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::VarBuilder; + +#[derive(Debug, Clone)] +pub struct Config { + pub num_layers: usize, + pub padded_vocab_size: usize, + pub hidden_size: usize, + pub ffn_hidden_size: usize, + pub kv_channels: usize, + pub num_attention_heads: usize, + pub seq_length: usize, + pub layernorm_epsilon: f64, + pub rmsnorm: bool, + pub apply_residual_connection_post_layernorm: bool, + pub post_layer_norm: bool, + pub add_bias_linear: bool, + pub add_qkv_bias: bool, + pub bias_dropout_fusion: bool, + pub multi_query_attention: bool, + pub multi_query_group_num: usize, + pub apply_query_key_layer_scaling: bool, + pub attention_softmax_in_fp32: bool, + pub fp32_residual_connection: bool, +} + +impl Config { + pub fn glm3_6b() -> Self { + Self { + num_layers: 28, + padded_vocab_size: 65024, + hidden_size: 4096, + ffn_hidden_size: 13696, + kv_channels: 128, + num_attention_heads: 32, + seq_length: 8192, + layernorm_epsilon: 1e-5, + rmsnorm: true, + apply_residual_connection_post_layernorm: false, + post_layer_norm: true, + add_bias_linear: false, + add_qkv_bias: true, + bias_dropout_fusion: true, + multi_query_attention: true, + multi_query_group_num: 2, + apply_query_key_layer_scaling: true, + attention_softmax_in_fp32: true, + fp32_residual_connection: false, + } + } +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + cache: Tensor, +} + +impl RotaryEmbedding { + fn new(cfg: &Config, dtype: DType, dev: &Device) -> Result { + let rotary_dim = cfg.kv_channels; + let n_elem = rotary_dim / 2; + let inv_freq: Vec<_> = (0..n_elem) + .step_by(2) + .map(|i| 1f32 / 10_000f64.powf(i as f64 / n_elem as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, cfg.seq_length as u32, dev)? + .to_dtype(dtype)? + .reshape((cfg.seq_length, 1))?; + let freqs = t.matmul(&inv_freq)?; + let cache = Tensor::stack(&[&freqs.cos()?, &freqs.sin()?], D::Minus1)?; + Ok(Self { cache }) + } + + fn apply(&self, xs: &Tensor, seqlen_offset: usize) -> Result { + let (seqlen, _b, np, _hn) = xs.dims4()?; + let cache = self.cache.narrow(0, seqlen_offset, seqlen)?; + let rot_dim = cache.dim(D::Minus2)? * 2; + let (xs, xs_pass) = ( + xs.narrow(D::Minus1, 0, rot_dim)?, + xs.narrow(D::Minus1, rot_dim, rot_dim)?, + ); + let xshaped = xs.reshape((seqlen, (), np, rot_dim / 2, 2))?; + let cache = cache.reshape((seqlen, (), 1, rot_dim / 2, 2))?; + let (xshaped0, xshaped1) = ( + xshaped.i((.., .., .., .., 0))?, + xshaped.i((.., .., .., .., 1))?, + ); + let (cache0, cache1) = (cache.i((.., .., .., .., 0))?, cache.i((.., .., .., .., 1))?); + let xs_out = Tensor::stack( + &[ + (xshaped0.broadcast_mul(&cache0)? - xshaped1.broadcast_mul(&cache1)?)?, + (xshaped1.broadcast_mul(&cache0)? + xshaped0.broadcast_mul(&cache1)?)?, + ], + D::Minus1, + )?; + let xs_out = xs_out.flatten_from(3)?; + Tensor::cat(&[xs_out, xs_pass], D::Minus1) + } +} + +#[derive(Debug, Clone)] +struct CoreAttention { + coeff: Option, + norm_factor: f64, +} + +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 CoreAttention { + fn new(layer_number: usize, cfg: &Config) -> Result { + let norm_factor = (cfg.kv_channels as f64).sqrt(); + let (norm_factor, coeff) = if cfg.apply_query_key_layer_scaling { + let coeff = f64::max(1.0, layer_number as f64); + (norm_factor * coeff, Some(coeff)) + } else { + (norm_factor, None) + }; + Ok(Self { coeff, norm_factor }) + } + + fn forward( + &self, + query_layer: &Tensor, + key_layer: &Tensor, + value_layer: &Tensor, + attention_mask: &Option, + ) -> Result { + let output_size = ( + query_layer.dim(1)?, // b + query_layer.dim(2)?, // np + query_layer.dim(0)?, // sq + key_layer.dim(0)?, // sk + ); + let query_layer = + query_layer.reshape((output_size.2, output_size.0 * output_size.1, ()))?; + let key_layer = key_layer.reshape((output_size.3, output_size.0 * output_size.1, ()))?; + let matmul_result = Tensor::matmul( + &query_layer.transpose(0, 1)?, + &key_layer.transpose(0, 1)?.transpose(1, 2)?, + )?; + let matmul_result = (matmul_result / self.norm_factor)?.reshape(output_size)?; + let matmul_result = match self.coeff { + None => matmul_result, + Some(coeff) => (matmul_result * coeff)?, + }; + let attention_scores = match attention_mask { + Some(mask) => masked_fill( + &matmul_result, + &mask.broadcast_left((matmul_result.dim(0)?, matmul_result.dim(1)?))?, + f32::NEG_INFINITY, + )?, + None => matmul_result, + }; + let attention_probs = candle_nn::ops::softmax_last_dim(&attention_scores)?; + + let output_size = ( + value_layer.dim(1)?, + value_layer.dim(2)?, + query_layer.dim(0)?, + value_layer.dim(3)?, + ); + let value_layer = + value_layer.reshape((value_layer.dim(0)?, output_size.0 * output_size.1, ()))?; + let attention_probs = + attention_probs.reshape((output_size.0 * output_size.1, output_size.2, ()))?; + let context_layer = Tensor::matmul(&attention_probs, &value_layer.transpose(0, 1)?)?; + let context_layer = context_layer.reshape(output_size)?; + let context_layer = context_layer.permute((2, 0, 1, 3))?.contiguous()?; + context_layer.flatten_from(D::Minus2) + } +} + +#[derive(Debug, Clone)] +struct SelfAttention { + query_key_value: Linear, + core_attention: CoreAttention, + dense: Linear, + multi_query_attention: bool, + num_attention_heads_per_partition: usize, + num_multi_query_groups_per_partition: usize, + hidden_size_per_attention_head: usize, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl SelfAttention { + fn new(layer_number: usize, cfg: &Config, vb: VarBuilder) -> Result { + let projection_size = cfg.kv_channels * cfg.num_attention_heads; + let hidden_size_per_attention_head = projection_size / cfg.num_attention_heads; + let qkv_hidden_size = if cfg.multi_query_attention { + projection_size + 2 * hidden_size_per_attention_head * cfg.multi_query_group_num + } else { + 3 * projection_size + }; + let query_key_value = linear( + cfg.hidden_size, + qkv_hidden_size, + cfg.add_bias_linear || cfg.add_qkv_bias, + vb.pp("query_key_value"), + )?; + let core_attention = CoreAttention::new(layer_number, cfg)?; + let dense = linear( + cfg.hidden_size, + cfg.hidden_size, + cfg.add_bias_linear, + vb.pp("dense"), + )?; + Ok(Self { + query_key_value, + core_attention, + dense, + multi_query_attention: cfg.multi_query_attention, + num_attention_heads_per_partition: cfg.num_attention_heads, + num_multi_query_groups_per_partition: cfg.multi_query_group_num, + hidden_size_per_attention_head: cfg.kv_channels, + kv_cache: None, + }) + } + + fn reset_kv_cache(&mut self) { + self.kv_cache = None + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: &Option, + rotary_emb: &RotaryEmbedding, + ) -> Result { + let mixed_x_layer = xs.apply(&self.query_key_value)?; + if !self.multi_query_attention { + candle::bail!("only multi_query_attention=true is supported") + } + let hpa = self.hidden_size_per_attention_head; + let query_layer = + mixed_x_layer.narrow(D::Minus1, 0, self.num_attention_heads_per_partition * hpa)?; + let key_layer = mixed_x_layer.narrow( + D::Minus1, + self.num_attention_heads_per_partition * hpa, + self.num_multi_query_groups_per_partition * hpa, + )?; + let value_layer = mixed_x_layer.narrow( + D::Minus1, + self.num_attention_heads_per_partition * hpa + + self.num_multi_query_groups_per_partition * hpa, + self.num_multi_query_groups_per_partition * hpa, + )?; + let query_layer = query_layer.reshape(( + query_layer.dim(0)?, + query_layer.dim(1)?, + self.num_attention_heads_per_partition, + hpa, + ))?; + let key_layer = key_layer.reshape(( + key_layer.dim(0)?, + key_layer.dim(1)?, + self.num_multi_query_groups_per_partition, + hpa, + ))?; + let value_layer = value_layer.reshape(( + value_layer.dim(0)?, + value_layer.dim(1)?, + self.num_multi_query_groups_per_partition, + hpa, + ))?; + + // Rotary embeddings. + let seqlen_offset = match &self.kv_cache { + None => 0, + Some((prev_k, _)) => prev_k.dim(0)?, + }; + let query_layer = rotary_emb.apply(&query_layer, seqlen_offset)?; + let key_layer = rotary_emb.apply(&key_layer, seqlen_offset)?; + + // KV cache. + let (key_layer, value_layer) = match &self.kv_cache { + None => (key_layer, value_layer), + Some((prev_k, prev_v)) => { + let k = Tensor::cat(&[prev_k, &key_layer], 0)?; + let v = Tensor::cat(&[prev_v, &value_layer], 0)?; + (k, v) + } + }; + self.kv_cache = Some((key_layer.clone(), value_layer.clone())); + + // Repeat KV. + let ratio = + self.num_attention_heads_per_partition / self.num_multi_query_groups_per_partition; + let key_layer = { + let (d0, d1, d2, d3) = key_layer.dims4()?; + key_layer + .unsqueeze(D::Minus2)? + .expand((d0, d1, d2, ratio, d3))? + .reshape(( + d0, + d1, + self.num_attention_heads_per_partition, + self.hidden_size_per_attention_head, + ))? + }; + let value_layer = { + let (d0, d1, d2, d3) = value_layer.dims4()?; + value_layer + .unsqueeze(D::Minus2)? + .expand((d0, d1, d2, ratio, d3))? + .reshape(( + d0, + d1, + self.num_attention_heads_per_partition, + self.hidden_size_per_attention_head, + ))? + }; + + let context_layer = + self.core_attention + .forward(&query_layer, &key_layer, &value_layer, attention_mask)?; + let output = context_layer.apply(&self.dense)?; + Ok(output) + } +} + +#[allow(clippy::upper_case_acronyms)] +#[derive(Debug, Clone)] +struct MLP { + dense_h_to_4h: Linear, + dense_4h_to_h: Linear, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense_h_to_4h = linear( + cfg.hidden_size, + cfg.ffn_hidden_size * 2, + cfg.add_bias_linear, + vb.pp("dense_h_to_4h"), + )?; + let dense_4h_to_h = linear( + cfg.ffn_hidden_size, + cfg.hidden_size, + cfg.add_bias_linear, + vb.pp("dense_4h_to_h"), + )?; + Ok(Self { + dense_4h_to_h, + dense_h_to_4h, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.dense_h_to_4h)? + .apply(&candle_nn::Activation::Swiglu)? + .apply(&self.dense_4h_to_h) + } +} + +#[derive(Debug, Clone)] +struct Block { + input_layernorm: candle_nn::LayerNorm, + self_attention: SelfAttention, + post_attention_layernorm: candle_nn::LayerNorm, + mlp: MLP, + apply_residual_connection_post_layernorm: bool, +} + +impl Block { + fn new(layer_number: usize, cfg: &Config, vb: VarBuilder) -> Result { + let input_layernorm = if cfg.rmsnorm { + candle_nn::rms_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("input_layernorm"), + )? + .into_inner() + } else { + candle_nn::layer_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("input_layernorm"), + )? + }; + let post_attention_layernorm = if cfg.rmsnorm { + candle_nn::rms_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("post_attention_layernorm"), + )? + .into_inner() + } else { + candle_nn::layer_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("post_attention_layernorm"), + )? + }; + let self_attention = SelfAttention::new(layer_number, cfg, vb.pp("self_attention"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + Ok(Self { + input_layernorm, + self_attention, + post_attention_layernorm, + mlp, + apply_residual_connection_post_layernorm: cfg.apply_residual_connection_post_layernorm, + }) + } + + fn reset_kv_cache(&mut self) { + self.self_attention.reset_kv_cache() + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: &Option, + rotary_emb: &RotaryEmbedding, + ) -> Result { + let layernorm_output = xs.apply(&self.input_layernorm)?; + let attention_output = + self.self_attention + .forward(&layernorm_output, attention_mask, rotary_emb)?; + let residual = if self.apply_residual_connection_post_layernorm { + &layernorm_output + } else { + xs + }; + let layernorm_input = (residual + attention_output)?; + let layernorm_output = layernorm_input.apply(&self.post_attention_layernorm)?; + let mlp_output = layernorm_output.apply(&self.mlp)?; + let residual = if self.apply_residual_connection_post_layernorm { + &layernorm_output + } else { + &layernorm_input + }; + mlp_output + residual + } +} + +#[derive(Debug, Clone)] +struct Transformer { + layers: Vec, + final_layernorm: Option, + rotary_emb: RotaryEmbedding, +} + +impl Transformer { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_l = vb.pp("layers"); + let mut layers = Vec::with_capacity(cfg.num_layers); + for layer_index in 0..cfg.num_layers { + let block = Block::new(layer_index + 1, cfg, vb_l.pp(layer_index))?; + layers.push(block) + } + let final_layernorm = if cfg.post_layer_norm { + let ln = if cfg.rmsnorm { + candle_nn::rms_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("final_layernorm"), + )? + .into_inner() + } else { + candle_nn::layer_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("final_layernorm"), + )? + }; + Some(ln) + } else { + None + }; + let rotary_emb = RotaryEmbedding::new(cfg, vb.dtype(), vb.device())?; + Ok(Self { + layers, + final_layernorm, + rotary_emb, + }) + } + + fn reset_kv_cache(&mut self) { + for block in self.layers.iter_mut() { + block.reset_kv_cache() + } + } + + fn forward(&mut self, xs: &Tensor, attention_mask: &Option) -> Result { + let mut xs = xs.clone(); + for block in self.layers.iter_mut() { + xs = block.forward(&xs, attention_mask, &self.rotary_emb)? + } + match self.final_layernorm.as_ref() { + None => Ok(xs), + Some(ln) => xs.apply(ln), + } + } +} + +#[derive(Debug, Clone)] +struct Embedding { + word_embeddings: candle_nn::Embedding, + fp32_residual_connection: bool, +} + +impl Embedding { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let word_embeddings = candle_nn::embedding( + cfg.padded_vocab_size, + cfg.hidden_size, + vb.pp("word_embeddings"), + )?; + Ok(Self { + word_embeddings, + fp32_residual_connection: cfg.fp32_residual_connection, + }) + } +} + +impl Module for Embedding { + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.word_embeddings.forward(xs)?.transpose(0, 1)?; // b,s,h -> s,b,h + if self.fp32_residual_connection { + xs.to_dtype(candle::DType::F32) + } else { + xs.contiguous() + } + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embedding: Embedding, + encoder: Transformer, + output_layer: Linear, +} + +fn get_mask(size: usize, device: &Device) -> Result { + let mask: Vec<_> = (0..size) + .flat_map(|i| (0..size).map(move |j| u8::from(j > i))) + .collect(); + Tensor::from_slice(&mask, (size, size), device) +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb = vb.pp("transformer"); + let embedding = Embedding::new(cfg, vb.pp("embedding"))?; + let encoder = Transformer::new(cfg, vb.pp("encoder"))?; + let output_layer = linear( + cfg.hidden_size, + cfg.padded_vocab_size, + false, + vb.pp("output_layer"), + )?; + Ok(Self { + embedding, + encoder, + output_layer, + }) + } + + pub fn reset_kv_cache(&mut self) { + self.encoder.reset_kv_cache() + } + + pub fn forward(&mut self, xs: &Tensor) -> Result { + let (_b_size, seq_len) = xs.dims2()?; + let input_embeds = xs.apply(&self.embedding)?; + let attention_mask = if seq_len <= 1 { + None + } else { + Some(get_mask(seq_len, xs.device())?) + }; + let xs = self.encoder.forward(&input_embeds, &attention_mask)?; + let lm_logits = xs.i(seq_len - 1)?.apply(&self.output_layer)?; + Ok(lm_logits) + } +} diff --git a/patches/candle-transformers/src/models/chinese_clip/mod.rs b/patches/candle-transformers/src/models/chinese_clip/mod.rs new file mode 100644 index 0000000000..ad8f380a24 --- /dev/null +++ b/patches/candle-transformers/src/models/chinese_clip/mod.rs @@ -0,0 +1,209 @@ +//! Chinese contrastive Language-Image Pre-Training +//! +//! Chinese contrastive Language-Image Pre-Training (CLIP) is an architecture trained on +//! pairs of images with related texts. +//! +//! - 💻 [GH Link](https://github.com/OFA-Sys/Chinese-CLIP) +//! - 💻 Transformers Python [reference implementation](https://github.com/huggingface/transformers/blob/5af7d41e49bbfc8319f462eb45253dcb3863dfb7/src/transformers/models/chinese_clip/modeling_chinese_clip.py) +//! +use candle::{Module, Result, Tensor, D}; +use candle_nn as nn; + +use text_model::ChineseClipTextTransformer; +use vision_model::ChineseClipVisionTransformer; + +pub mod text_model; +pub mod vision_model; + +#[derive(Debug, Clone, Copy)] +pub enum Activation { + QuickGelu, + Gelu, + GeluNew, + Relu, +} + +impl From for Activation { + fn from(value: String) -> Self { + match value.as_str() { + "quick_gelu" => Activation::QuickGelu, + "gelu" => Activation::Gelu, + "gelu_new" => Activation::GeluNew, + "relu" => Activation::Relu, + _ => panic!("Invalid activation function: {value}"), + } + } +} + +impl Module for Activation { + fn forward(&self, xs: &Tensor) -> Result { + match self { + Activation::QuickGelu => xs * nn::ops::sigmoid(&(xs * 1.702f64)?)?, + Activation::Gelu => xs.gelu_erf(), + Activation::GeluNew => xs.gelu(), + Activation::Relu => xs.relu(), + } + } +} + +#[derive(Clone, Debug)] +pub struct ChineseClipConfig { + pub text_config: text_model::ChineseClipTextConfig, + pub vision_config: vision_model::ChineseClipVisionConfig, + pub projection_dim: usize, + pub logit_scale_init_value: f32, + pub image_size: usize, +} + +impl ChineseClipConfig { + /// referer: https://huggingface.co/OFA-Sys/chinese-clip-vit-base-patch16/blob/main/config.json + pub fn clip_vit_base_patch16() -> Self { + let text_config = text_model::ChineseClipTextConfig::clip_vit_base_patch16(); + let vision_config = vision_model::ChineseClipVisionConfig::clip_vit_base_patch16(); + + Self { + text_config, + vision_config, + projection_dim: 512, + logit_scale_init_value: 2.6592, + image_size: 512, + } + } +} + +#[derive(Clone, Debug)] +pub enum EncoderConfig { + Text(text_model::ChineseClipTextConfig), + Vision(vision_model::ChineseClipVisionConfig), +} + +impl EncoderConfig { + pub fn embed_dim(&self) -> usize { + match self { + Self::Text(c) => c.hidden_size, + Self::Vision(c) => c.hidden_size, + } + } + + pub fn num_attention_heads(&self) -> usize { + match self { + Self::Text(c) => c.num_attention_heads, + Self::Vision(c) => c.num_attention_heads, + } + } + + pub fn intermediate_size(&self) -> usize { + match self { + Self::Text(c) => c.intermediate_size, + Self::Vision(c) => c.intermediate_size, + } + } + + pub fn num_hidden_layers(&self) -> usize { + match self { + Self::Text(c) => c.num_hidden_layers, + Self::Vision(c) => c.num_hidden_layers, + } + } + + pub fn activation(&self) -> Activation { + match self { + Self::Text(c) => c.hidden_act, + Self::Vision(c) => c.hidden_act, + } + } + + pub fn layer_norm_eps(&self) -> f64 { + match self { + Self::Text(c) => c.layer_norm_eps, + Self::Vision(c) => c.layer_norm_eps, + } + } +} + +#[derive(Clone, Debug)] +pub struct ChineseClipModel { + text_model: ChineseClipTextTransformer, + vision_model: ChineseClipVisionTransformer, + visual_projection: nn::Linear, + text_projection: nn::Linear, + logit_scale: Tensor, +} + +impl ChineseClipModel { + pub fn new(vs: nn::VarBuilder, c: &ChineseClipConfig) -> Result { + let text_model = ChineseClipTextTransformer::new(vs.pp("text_model"), &c.text_config)?; + + let vision_model = + ChineseClipVisionTransformer::new(vs.pp("vision_model"), &c.vision_config)?; + + let vision_embed_dim = c.vision_config.hidden_size; + let vision_projection = nn::linear_no_bias( + vision_embed_dim, + c.projection_dim, + vs.pp("visual_projection"), + )?; + + let text_embed_dim = c.text_config.hidden_size; + let text_projection = + nn::linear_no_bias(text_embed_dim, c.projection_dim, vs.pp("text_projection"))?; + + let logit_scale = if vs.contains_tensor("logit_scale") { + vs.get(&[], "logit_scale")? + } else { + Tensor::new(&[c.logit_scale_init_value], vs.device())? + }; + + Ok(Self { + text_model, + vision_model, + visual_projection: vision_projection, + text_projection, + logit_scale, + }) + } + + pub fn get_text_features( + &self, + input_ids: &Tensor, + token_type_ids: Option<&Tensor>, + attention_mask: Option<&Tensor>, + ) -> Result { + let output = self + .text_model + .forward(input_ids, token_type_ids, attention_mask)? + .contiguous()?; + self.text_projection.forward(&output) + } + + pub fn get_image_features(&self, pixel_values: &Tensor) -> Result { + pixel_values + .apply(&self.vision_model)? + .apply(&self.visual_projection) + } + + pub fn forward( + &self, + pixel_values: &Tensor, + input_ids: &Tensor, + token_type_ids: Option<&Tensor>, + attention_mask: Option<&Tensor>, + ) -> Result<(Tensor, Tensor)> { + let image_features = self.get_image_features(pixel_values)?; + let text_features = self.get_text_features(input_ids, token_type_ids, attention_mask)?; + + let image_features_normalized = div_l2_norm(&image_features)?; + let text_features_normalized = div_l2_norm(&text_features)?; + + let logits_per_text = text_features_normalized.matmul(&image_features_normalized.t()?)?; + let logit_scale = self.logit_scale.exp()?; + let logits_per_text = logits_per_text.broadcast_mul(&logit_scale)?; + let logits_per_image = logits_per_text.t()?; + Ok((logits_per_text, logits_per_image)) + } +} + +pub fn div_l2_norm(v: &Tensor) -> Result { + let l2_norm = v.sqr()?.sum_keepdim(D::Minus1)?.sqrt()?; + v.broadcast_div(&l2_norm) +} diff --git a/patches/candle-transformers/src/models/chinese_clip/text_model.rs b/patches/candle-transformers/src/models/chinese_clip/text_model.rs new file mode 100644 index 0000000000..b43c742348 --- /dev/null +++ b/patches/candle-transformers/src/models/chinese_clip/text_model.rs @@ -0,0 +1,544 @@ +//! Chinese contrastive Language-Image Pre-Training +//! +//! Chinese contrastive Language-Image Pre-Training (CLIP) is an architecture trained on +//! pairs of images with related texts. +//! +//! - 💻 [Chinese-CLIP](https://github.com/OFA-Sys/Chinese-CLIP) +//! - 💻 [HF](https://github.com/huggingface/transformers/blob/5af7d41e49bbfc8319f462eb45253dcb3863dfb7/src/transformers/models/chinese_clip/modeling_chinese_clip.py) + +use candle::{DType, Device, IndexOp, Module, Result, Tensor}; +use candle_nn as nn; + +use super::Activation; + +/// Type of position embedding. Choose one of `"absolute"`, `"relative_key"`, `"relative_key_query"`. For +/// positional embeddings use `"absolute"`. For more information on `"relative_key"`, please refer to +/// [Self-Attention with Relative Position Representations (Shaw et al.)](https://arxiv.org/abs/1803.02155). +/// For more information on `"relative_key_query"`, please refer to *Method 4* in [Improve Transformer Models +/// with Better Relative Position Embeddings (Huang et al.)](https://arxiv.org/abs/2009.13658). +#[derive(Clone, Debug)] +pub enum PositionEmbeddingType { + Absolute, + RelativeKey, + RelativeKeyQuery, +} + +#[derive(Clone, Debug)] +pub struct ChineseClipTextConfig { + pub vocab_size: usize, + pub hidden_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub intermediate_size: usize, + pub hidden_act: Activation, + pub hidden_dropout_prob: f32, + pub attention_probs_dropout_prob: f64, + pub max_position_embeddings: usize, + pub type_vocab_size: usize, + pub initializer_range: f64, + pub initializer_factor: f64, + pub layer_norm_eps: f64, + pub pad_token_id: usize, + pub position_embedding_type: PositionEmbeddingType, + pub use_cache: bool, +} + +impl Default for ChineseClipTextConfig { + fn default() -> Self { + Self { + vocab_size: 30522, + hidden_size: 768, + num_hidden_layers: 12, + num_attention_heads: 12, + intermediate_size: 3072, + hidden_act: Activation::Gelu, + hidden_dropout_prob: 0.1, + attention_probs_dropout_prob: 0.1, + max_position_embeddings: 512, + type_vocab_size: 2, + initializer_range: 0.02, + initializer_factor: 1.0, + layer_norm_eps: 1e-12, + pad_token_id: 0, + position_embedding_type: PositionEmbeddingType::Absolute, + use_cache: true, + } + } +} + +impl ChineseClipTextConfig { + /// [referer](https://huggingface.co/OFA-Sys/chinese-clip-vit-base-patch16/blob/main/config.json) + pub fn clip_vit_base_patch16() -> Self { + Self { + vocab_size: 21128, + hidden_size: 768, + num_hidden_layers: 12, + num_attention_heads: 12, + intermediate_size: 3072, + hidden_act: Activation::Gelu, + hidden_dropout_prob: 0.1, + attention_probs_dropout_prob: 0.1, + max_position_embeddings: 512, + type_vocab_size: 2, + initializer_range: 0.02, + initializer_factor: 1.0, + layer_norm_eps: 1e-12, + pad_token_id: 0, + position_embedding_type: PositionEmbeddingType::Absolute, + use_cache: true, + } + } +} + +#[derive(Clone, Debug)] +pub struct ChineseClipTextEmbeddings { + word_embeddings: nn::Embedding, + position_embeddings: nn::Embedding, + token_type_embeddings: nn::Embedding, + layer_norm: nn::LayerNorm, + dropout: nn::Dropout, + position_embedding_type: PositionEmbeddingType, + position_ids: Tensor, + token_type_ids: Tensor, +} + +impl ChineseClipTextEmbeddings { + pub fn new(var: nn::VarBuilder, config: &ChineseClipTextConfig) -> Result { + let word_embeddings = nn::embedding( + config.vocab_size, + config.hidden_size, + var.pp("word_embeddings"), + )?; + let position_embeddings = nn::embedding( + config.max_position_embeddings, + config.hidden_size, + var.pp("position_embeddings"), + )?; + let token_type_embeddings = nn::embedding( + config.type_vocab_size, + config.hidden_size, + var.pp("token_type_embeddings"), + )?; + let layer_norm = nn::layer_norm::( + config.hidden_size, + config.layer_norm_eps, + var.pp("LayerNorm"), + )?; + let dropout = nn::Dropout::new(config.hidden_dropout_prob); + let position_ids = + Tensor::arange(0u32, config.max_position_embeddings as u32, var.device())? + .unsqueeze(0)?; + let token_type_ids = Tensor::zeros(position_ids.shape(), DType::I64, var.device())?; + + Ok(Self { + word_embeddings, + position_embeddings, + token_type_embeddings, + layer_norm, + dropout, + position_embedding_type: config.position_embedding_type.clone(), + position_ids, + token_type_ids, + }) + } + + fn forward(&self, xs: &Tensor, token_type_ids: Option<&Tensor>) -> Result { + let (_batch_size, seq_length) = xs.dims2()?; + let position_ids = (0..seq_length as u32).collect::>(); + let position_ids = self.position_ids.index_select( + &Tensor::new(&position_ids[..], self.position_ids.device())?, + 1, + )?; + + let word_embeddings = self.word_embeddings.forward(xs)?; + + let token_type_ids = match token_type_ids { + Some(token_type_ids) => token_type_ids, + None => &self.token_type_ids.i((.., 0..seq_length))?, + }; + let token_type_ids = token_type_ids.expand(xs.shape())?; + let token_type_embeddings = self.token_type_embeddings.forward(&token_type_ids)?; + + let embeddings = (&word_embeddings + token_type_embeddings)?; + let embeddings = match self.position_embedding_type { + PositionEmbeddingType::Absolute => { + let position_embeddings = self.position_embeddings.forward(&position_ids)?; + let position_embeddings = position_embeddings.expand(embeddings.shape())?; + (embeddings + position_embeddings)? + } + _ => embeddings, + }; + let embeddings = self.layer_norm.forward(&embeddings)?; + let embeddings = self.dropout.forward(&embeddings, false)?; + Ok(embeddings) + } +} + +/// Copied from [`crate::models::bert::BertSelfOutput`] to [`ChineseClipTextSelfOutput`] +#[derive(Clone, Debug)] +struct ChineseClipTextSelfOutput { + dense: nn::Linear, + layer_norm: nn::LayerNorm, + dropout: nn::Dropout, + span: tracing::Span, +} + +impl ChineseClipTextSelfOutput { + fn new(var: nn::VarBuilder, config: &ChineseClipTextConfig) -> Result { + let dense = nn::linear(config.hidden_size, config.hidden_size, var.pp("dense"))?; + let layer_norm = nn::layer_norm( + config.hidden_size, + config.layer_norm_eps, + var.pp("LayerNorm"), + )?; + let dropout = nn::Dropout::new(config.hidden_dropout_prob); + Ok(Self { + dense, + layer_norm, + dropout, + span: tracing::span!(tracing::Level::TRACE, "self-out"), + }) + } + + fn forward(&self, hidden_states: &Tensor, input_tensor: &Tensor) -> Result { + let _enter = self.span.enter(); + let hidden_states = self.dense.forward(hidden_states)?; + let hidden_states = self.dropout.forward(&hidden_states, false)?; + self.layer_norm.forward(&(hidden_states + input_tensor)?) + } +} + +/// Copied from [`crate::models::bert::BertSelfAttention`] to [`ChineseClipTextSelfAttention`] +#[derive(Clone, Debug)] +struct ChineseClipTextSelfAttention { + query: nn::Linear, + key: nn::Linear, + value: nn::Linear, + dropout: nn::Dropout, + num_attention_heads: usize, + attention_head_size: usize, + span: tracing::Span, + span_softmax: tracing::Span, +} + +impl ChineseClipTextSelfAttention { + fn new(var: nn::VarBuilder, config: &ChineseClipTextConfig) -> Result { + let attention_head_size = config.hidden_size / config.num_attention_heads; + let all_head_size = config.num_attention_heads * attention_head_size; + let dropout = nn::Dropout::new(config.hidden_dropout_prob); + let hidden_size = config.hidden_size; + let query = nn::linear(hidden_size, all_head_size, var.pp("query"))?; + let value = nn::linear(hidden_size, all_head_size, var.pp("value"))?; + let key = nn::linear(hidden_size, all_head_size, var.pp("key"))?; + Ok(Self { + query, + key, + value, + dropout, + num_attention_heads: config.num_attention_heads, + attention_head_size, + span: tracing::span!(tracing::Level::TRACE, "self-attn"), + span_softmax: tracing::span!(tracing::Level::TRACE, "softmax"), + }) + } + + fn transpose_for_scores(&self, xs: &Tensor) -> Result { + let mut new_x_shape = xs.dims().to_vec(); + new_x_shape.pop(); + new_x_shape.push(self.num_attention_heads); + new_x_shape.push(self.attention_head_size); + let xs = xs.reshape(new_x_shape.as_slice())?.transpose(1, 2)?; + xs.contiguous() + } + + fn forward(&self, hidden_states: &Tensor, attention_mask: &Tensor) -> Result { + let _enter = self.span.enter(); + let query_layer = self.query.forward(hidden_states)?; + let key_layer = self.key.forward(hidden_states)?; + let value_layer = self.value.forward(hidden_states)?; + + let query_layer = self.transpose_for_scores(&query_layer)?; + let key_layer = self.transpose_for_scores(&key_layer)?; + let value_layer = self.transpose_for_scores(&value_layer)?; + + let attention_scores = query_layer.matmul(&key_layer.t()?)?; + let attention_scores = (attention_scores / (self.attention_head_size as f64).sqrt())?; + let attention_scores = attention_scores.broadcast_add(attention_mask)?; + let attention_probs = { + let _enter_sm = self.span_softmax.enter(); + nn::ops::softmax(&attention_scores, candle::D::Minus1)? + }; + let attention_probs = self.dropout.forward(&attention_probs, false)?; + + let context_layer = attention_probs.matmul(&value_layer)?; + let context_layer = context_layer.transpose(1, 2)?.contiguous()?; + let context_layer = context_layer.flatten_from(candle::D::Minus2)?; + Ok(context_layer) + } +} + +/// Copied from [`crate::models::bert::BertAttention`] to [`ChineseClipTextAttention`] +#[derive(Clone, Debug)] +struct ChineseClipTextAttention { + self_attention: ChineseClipTextSelfAttention, + self_output: ChineseClipTextSelfOutput, + span: tracing::Span, +} + +impl ChineseClipTextAttention { + fn new(var: nn::VarBuilder, config: &ChineseClipTextConfig) -> Result { + let self_attention = ChineseClipTextSelfAttention::new(var.pp("self"), config)?; + let self_output = ChineseClipTextSelfOutput::new(var.pp("output"), config)?; + Ok(Self { + self_attention, + self_output, + span: tracing::span!(tracing::Level::TRACE, "attn"), + }) + } + + fn forward(&self, hidden_states: &Tensor, attention_mask: &Tensor) -> Result { + let _enter = self.span.enter(); + let self_outputs = self.self_attention.forward(hidden_states, attention_mask)?; + let attention_output = self.self_output.forward(&self_outputs, hidden_states)?; + Ok(attention_output) + } +} + +type HiddenActLayer = Activation; + +/// Copied from [`crate::models::bert::BertIntermediate`] to [`ChineseClipTextIntermediate`] +#[derive(Clone, Debug)] +struct ChineseClipTextIntermediate { + dense: nn::Linear, + intermediate_act: HiddenActLayer, + span: tracing::Span, +} + +impl ChineseClipTextIntermediate { + fn new(var: nn::VarBuilder, config: &ChineseClipTextConfig) -> Result { + let dense = nn::linear( + config.hidden_size, + config.intermediate_size, + var.pp("dense"), + )?; + Ok(Self { + dense, + intermediate_act: config.hidden_act, + span: tracing::span!(tracing::Level::TRACE, "inter"), + }) + } +} + +impl Module for ChineseClipTextIntermediate { + fn forward(&self, hidden_states: &Tensor) -> Result { + let _enter = self.span.enter(); + let hidden_states = self.dense.forward(hidden_states)?; + let ys = self.intermediate_act.forward(&hidden_states)?; + Ok(ys) + } +} + +/// Copied from [`crate::models::bert::BertOutput`] to [`ChineseClipTextOutput`] +#[derive(Clone, Debug)] +struct ChineseClipTextOutput { + dense: nn::Linear, + layer_norm: nn::LayerNorm, + dropout: nn::Dropout, + span: tracing::Span, +} + +impl ChineseClipTextOutput { + fn new(var: nn::VarBuilder, config: &ChineseClipTextConfig) -> Result { + let dense = nn::linear( + config.intermediate_size, + config.hidden_size, + var.pp("dense"), + )?; + let layer_norm = nn::layer_norm( + config.hidden_size, + config.layer_norm_eps, + var.pp("LayerNorm"), + )?; + let dropout = nn::Dropout::new(config.hidden_dropout_prob); + Ok(Self { + dense, + layer_norm, + dropout, + span: tracing::span!(tracing::Level::TRACE, "out"), + }) + } + + fn forward(&self, hidden_states: &Tensor, input_tensor: &Tensor) -> Result { + let _enter = self.span.enter(); + let hidden_states = self.dense.forward(hidden_states)?; + let hidden_states = self.dropout.forward(&hidden_states, false)?; + self.layer_norm.forward(&(hidden_states + input_tensor)?) + } +} + +/// Copied from [`crate::models::bert::BertLayer`] to [`ChineseClipTextLayer`] +#[derive(Clone, Debug)] +struct ChineseClipTextLayer { + attention: ChineseClipTextAttention, + intermediate: ChineseClipTextIntermediate, + output: ChineseClipTextOutput, + span: tracing::Span, +} + +impl ChineseClipTextLayer { + fn new(var: nn::VarBuilder, config: &ChineseClipTextConfig) -> Result { + let attention = ChineseClipTextAttention::new(var.pp("attention"), config)?; + let intermediate = ChineseClipTextIntermediate::new(var.pp("intermediate"), config)?; + let output = ChineseClipTextOutput::new(var.pp("output"), config)?; + Ok(Self { + attention, + intermediate, + output, + span: tracing::span!(tracing::Level::TRACE, "layer"), + }) + } + + fn forward(&self, hidden_states: &Tensor, attention_mask: &Tensor) -> Result { + let _enter = self.span.enter(); + let attention_output = self.attention.forward(hidden_states, attention_mask)?; + // https://github.com/huggingface/transformers/blob/6eedfa6dd15dc1e22a55ae036f681914e5a0d9a1/src/transformers/models/bert/modeling_bert.py#L523 + let intermediate_output = self.intermediate.forward(&attention_output)?; + let layer_output = self + .output + .forward(&intermediate_output, &attention_output)?; + Ok(layer_output) + } +} + +#[derive(Clone, Debug)] +struct Tanh; + +impl Tanh { + pub fn new() -> Self { + Self {} + } +} +impl Module for Tanh { + fn forward(&self, xs: &Tensor) -> Result { + xs.tanh() + } +} + +#[derive(Clone, Debug)] +struct ChineseClipTextPooler { + dense: nn::Linear, + activation: Tanh, +} + +impl ChineseClipTextPooler { + pub fn new(var: nn::VarBuilder, config: &ChineseClipTextConfig) -> Result { + let dense = nn::linear(config.hidden_size, config.hidden_size, var.pp("dense"))?; + let activation = Tanh::new(); + Ok(Self { dense, activation }) + } +} + +impl Module for ChineseClipTextPooler { + fn forward(&self, hidden_states: &Tensor) -> Result { + let first_token_tensor = hidden_states.i((.., 0))?; + let pooled_output = self.dense.forward(&first_token_tensor)?; + let pooled_output = self.activation.forward(&pooled_output)?; + Ok(pooled_output) + } +} + +#[derive(Clone, Debug)] +struct ChineseClipTextEncoder { + layers: Vec, + span: tracing::Span, +} + +impl ChineseClipTextEncoder { + fn new(var: nn::VarBuilder, config: &ChineseClipTextConfig) -> Result { + let layers = (0..config.num_hidden_layers) + .map(|index| ChineseClipTextLayer::new(var.pp(format!("layer.{index}")), config)) + .collect::>>()?; + let span = tracing::span!(tracing::Level::TRACE, "encoder"); + Ok(ChineseClipTextEncoder { layers, span }) + } + + fn forward(&self, hidden_states: &Tensor, attention_mask: &Tensor) -> Result { + let _enter = self.span.enter(); + let mut hidden_states = hidden_states.clone(); + // Use a loop rather than a fold as it's easier to modify when adding debug/... + for layer in self.layers.iter() { + hidden_states = layer.forward(&hidden_states, attention_mask)? + } + Ok(hidden_states) + } +} + +#[derive(Clone, Debug)] +pub struct ChineseClipTextTransformer { + embeddings: ChineseClipTextEmbeddings, + encoder: ChineseClipTextEncoder, + pooler: Option, + pub device: Device, + span: tracing::Span, +} + +impl ChineseClipTextTransformer { + pub fn new(var: nn::VarBuilder, config: &ChineseClipTextConfig) -> Result { + let embeddings = ChineseClipTextEmbeddings::new(var.pp("embeddings"), config)?; + let encoder = ChineseClipTextEncoder::new(var.pp("encoder"), config)?; + // see: https://github.com/huggingface/transformers/blob/e40bb4845e0eefb52ec1e9cac9c2446ab36aef81/src/transformers/models/chinese_clip/modeling_chinese_clip.py#L1362 + // In the original Python version of the code, the pooler is not used, and there are no parameters for the pooler in the weight file. + let pooler = if var.contains_tensor("pooler") { + Some(ChineseClipTextPooler::new(var.pp("pooler"), config)?) + } else { + None + }; + Ok(Self { + embeddings, + encoder, + pooler, + device: var.device().clone(), + span: tracing::span!(tracing::Level::TRACE, "model"), + }) + } + + pub fn forward( + &self, + input_ids: &Tensor, + token_type_ids: Option<&Tensor>, + attention_mask: Option<&Tensor>, + ) -> Result { + let _enter = self.span.enter(); + let embedding_output = self.embeddings.forward(input_ids, token_type_ids)?; + let attention_mask = match attention_mask { + Some(attention_mask) => attention_mask.clone(), + None => input_ids.ones_like()?, + }; + let dtype = embedding_output.dtype(); + // https://github.com/huggingface/transformers/blob/6eedfa6dd15dc1e22a55ae036f681914e5a0d9a1/src/transformers/models/bert/modeling_bert.py#L995 + let attention_mask = get_extended_attention_mask(&attention_mask, dtype)?; + let encoder_outputs = self.encoder.forward(&embedding_output, &attention_mask)?; + let encoder_output = encoder_outputs.i((.., 0, ..))?; + let pooled_output = match &self.pooler { + Some(pooler) => pooler.forward(&encoder_output)?, + None => encoder_output, + }; + + Ok(pooled_output) + } +} + +fn get_extended_attention_mask(attention_mask: &Tensor, dtype: DType) -> Result { + let attention_mask = match attention_mask.rank() { + 3 => attention_mask.unsqueeze(1)?, + 2 => attention_mask.unsqueeze(1)?.unsqueeze(1)?, + _ => candle::bail!("Wrong shape for input_ids or attention_mask"), + }; + let attention_mask = attention_mask.to_dtype(dtype)?; + // torch.finfo(dtype).min + (attention_mask.ones_like()? - &attention_mask)?.broadcast_mul( + &Tensor::try_from(f32::MIN)? + .to_device(attention_mask.device())? + .to_dtype(dtype)?, + ) +} diff --git a/patches/candle-transformers/src/models/chinese_clip/vision_model.rs b/patches/candle-transformers/src/models/chinese_clip/vision_model.rs new file mode 100644 index 0000000000..153fe833c5 --- /dev/null +++ b/patches/candle-transformers/src/models/chinese_clip/vision_model.rs @@ -0,0 +1,385 @@ +//! Chinese contrastive Language-Image Pre-Training +//! +//! Chinese contrastive Language-Image Pre-Training (CLIP) is an architecture trained on +//! pairs of images with related texts. +//! +//! - 💻 [Chinese-CLIP](https://github.com/OFA-Sys/Chinese-CLIP) +//! - 💻 [GH](https://github.com/huggingface/transformers/blob/5af7d41e49bbfc8319f462eb45253dcb3863dfb7/src/transformers/models/chinese_clip/modeling_chinese_clip.py_ + +use candle::{Context, DType, IndexOp, Module, Result, Shape, Tensor, D}; +use candle_nn as nn; + +use super::{Activation, EncoderConfig}; + +#[derive(Clone, Debug)] +pub struct ChineseClipVisionConfig { + pub hidden_size: usize, + pub intermediate_size: usize, + pub projection_dim: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_channels: usize, + pub image_size: usize, + pub patch_size: usize, + pub hidden_act: Activation, + pub layer_norm_eps: f64, + pub attention_dropout: f32, + pub initializer_range: f32, + pub initializer_factor: f32, +} + +impl Default for ChineseClipVisionConfig { + fn default() -> Self { + ChineseClipVisionConfig { + hidden_size: 768, + intermediate_size: 3072, + projection_dim: 512, + num_hidden_layers: 12, + num_attention_heads: 12, + num_channels: 3, + image_size: 224, + patch_size: 32, + hidden_act: Activation::QuickGelu, + layer_norm_eps: 1e-5, + attention_dropout: 0.0, + initializer_range: 0.02, + initializer_factor: 1.0, + } + } +} + +impl ChineseClipVisionConfig { + /// [referer](https://huggingface.co/OFA-Sys/chinese-clip-vit-base-patch16/blob/main/config.json) + pub fn clip_vit_base_patch16() -> Self { + Self { + hidden_size: 768, + intermediate_size: 3072, + projection_dim: 512, + num_hidden_layers: 12, + num_attention_heads: 12, + num_channels: 3, + image_size: 224, + patch_size: 16, + hidden_act: Activation::QuickGelu, + layer_norm_eps: 1e-5, + attention_dropout: 0.0, + initializer_range: 0.02, + initializer_factor: 1.0, + } + } +} + +#[derive(Clone, Debug)] +pub struct ChineseClipVisionEmbeddings { + patch_embedding: nn::Conv2d, + position_ids: Tensor, + class_embedding: Tensor, + position_embedding: nn::Embedding, +} + +impl ChineseClipVisionEmbeddings { + pub fn new(var: nn::VarBuilder, config: &ChineseClipVisionConfig) -> Result { + let embed_dim = config.hidden_size; + // originally nn.Parameter + let class_embedding = if var.contains_tensor("class_embedding") { + var.get(embed_dim, "class_embedding")? + } else { + Tensor::randn(0f32, 1f32, embed_dim, var.device())? + }; + + let num_patches = (config.image_size / config.patch_size).pow(2); + let num_positions = num_patches + 1; + let position_ids = Tensor::arange(0, num_positions as i64, var.device())?; + + let conv2dconfig = nn::Conv2dConfig { + stride: config.patch_size, + ..Default::default() + }; + let position_embedding = + nn::embedding(num_positions, embed_dim, var.pp("position_embedding"))?; + let patch_embedding = nn::conv2d_no_bias( + config.num_channels, + embed_dim, + config.patch_size, + conv2dconfig, + var.pp("patch_embedding"), + )?; + Ok(Self { + patch_embedding, + position_ids, + class_embedding, + position_embedding, + }) + } +} + +impl Module for ChineseClipVisionEmbeddings { + fn forward(&self, xs: &Tensor) -> Result { + let batch_size = xs.shape().dims(); + let patch_embeds = self + .patch_embedding + .forward(xs)? + .flatten_from(2)? + .transpose(1, 2)?; + let shape = Shape::from((batch_size[0], 1, self.class_embedding.dim(D::Minus1)?)); + let class_embeds = self.class_embedding.expand(shape)?; + let embeddings = Tensor::cat(&[class_embeds, patch_embeds], 1)?; + let position_embedding = self.position_embedding.forward(&self.position_ids)?; + embeddings.broadcast_add(&position_embedding) + } +} + +#[derive(Clone, Debug)] +struct ChineseClipVisionAttention { + k_proj: nn::Linear, + v_proj: nn::Linear, + q_proj: nn::Linear, + out_proj: nn::Linear, + head_dim: usize, + scale: f64, + num_attention_heads: usize, +} + +impl ChineseClipVisionAttention { + fn new(var: nn::VarBuilder, config: &EncoderConfig) -> Result { + let embed_dim = config.embed_dim(); + let num_attention_heads = config.num_attention_heads(); + let k_proj = nn::linear(embed_dim, embed_dim, var.pp("k_proj"))?; + let v_proj = nn::linear(embed_dim, embed_dim, var.pp("v_proj"))?; + let q_proj = nn::linear(embed_dim, embed_dim, var.pp("q_proj"))?; + let out_proj = nn::linear(embed_dim, embed_dim, var.pp("out_proj"))?; + let head_dim = embed_dim / num_attention_heads; + let scale = (head_dim as f64).powf(-0.5); + + Ok(ChineseClipVisionAttention { + k_proj, + v_proj, + q_proj, + out_proj, + head_dim, + scale, + num_attention_heads, + }) + } + + fn shape(&self, xs: &Tensor, seq_len: usize, bsz: usize) -> Result { + xs.reshape((bsz, seq_len, self.num_attention_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous() + } + + fn forward(&self, xs: &Tensor, causal_attention_mask: Option<&Tensor>) -> Result { + let in_dtype = xs.dtype(); + let (bsz, seq_len, embed_dim) = xs.dims3()?; + + let proj_shape = (bsz * self.num_attention_heads, seq_len, self.head_dim); + let query_states = self + .shape(&(self.q_proj.forward(xs)? * self.scale)?, seq_len, bsz)? + .reshape(proj_shape)? + .to_dtype(DType::F32)?; + let key_states = self + .shape(&self.k_proj.forward(xs)?, seq_len, bsz)? + .reshape(proj_shape)? + .to_dtype(DType::F32)?; + let value_states = self + .shape(&self.v_proj.forward(xs)?, seq_len, bsz)? + .reshape(proj_shape)? + .to_dtype(DType::F32)?; + + let attn_weights = query_states.matmul(&key_states.transpose(1, 2)?)?; + + let src_len = key_states.dim(1)?; + + let attn_weights = if let Some(causal_attention_mask) = causal_attention_mask { + attn_weights + .reshape((bsz, self.num_attention_heads, seq_len, src_len))? + .broadcast_add(causal_attention_mask)? + .reshape((bsz * self.num_attention_heads, seq_len, src_len))? + } else { + attn_weights + }; + + let attn_weights = nn::ops::softmax(&attn_weights, D::Minus1)?; + + let attn_output = attn_weights.matmul(&value_states)?.to_dtype(in_dtype)?; + let attn_output = attn_output + .reshape((bsz, self.num_attention_heads, seq_len, self.head_dim))? + .transpose(1, 2)? + .reshape((bsz, seq_len, embed_dim))?; + self.out_proj.forward(&attn_output) + } +} + +#[derive(Clone, Debug)] +struct ChineseClipVisionMlp { + fc1: nn::Linear, + fc2: nn::Linear, + activation: Activation, +} + +impl ChineseClipVisionMlp { + fn new(var: nn::VarBuilder, config: &EncoderConfig) -> Result { + let fc1 = nn::linear( + config.embed_dim(), + config.intermediate_size(), + var.pp("fc1"), + )?; + let fc2 = nn::linear( + config.intermediate_size(), + config.embed_dim(), + var.pp("fc2"), + )?; + + Ok(ChineseClipVisionMlp { + fc1, + fc2, + activation: config.activation(), + }) + } +} + +impl ChineseClipVisionMlp { + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.fc1.forward(xs)?; + self.fc2.forward(&self.activation.forward(&xs)?) + } +} + +#[derive(Clone, Debug)] +struct ChineseClipVisionEncoderLayer { + self_attn: ChineseClipVisionAttention, + layer_norm1: nn::LayerNorm, + mlp: ChineseClipVisionMlp, + layer_norm2: nn::LayerNorm, +} + +impl ChineseClipVisionEncoderLayer { + fn new(var: nn::VarBuilder, config: &EncoderConfig) -> Result { + let self_attn = ChineseClipVisionAttention::new(var.pp("self_attn"), config)?; + let layer_norm1 = nn::layer_norm( + config.embed_dim(), + config.layer_norm_eps(), + var.pp("layer_norm1"), + )?; + let mlp = ChineseClipVisionMlp::new(var.pp("mlp"), config)?; + let layer_norm2 = nn::layer_norm( + config.embed_dim(), + config.layer_norm_eps(), + var.pp("layer_norm2"), + )?; + + Ok(ChineseClipVisionEncoderLayer { + self_attn, + layer_norm1, + mlp, + layer_norm2, + }) + } + + fn forward(&self, xs: &Tensor, causal_attention_mask: Option<&Tensor>) -> Result { + let residual = xs; + let xs = self.layer_norm1.forward(xs)?; + let xs = self.self_attn.forward(&xs, causal_attention_mask)?; + let xs = (xs + residual)?; + + let residual = &xs; + let xs = self.layer_norm2.forward(&xs)?; + let xs = self.mlp.forward(&xs)?; + xs + residual + } +} + +#[derive(Clone, Debug)] +pub struct ChineseClipVisionEncoder { + layers: Vec, +} + +impl ChineseClipVisionEncoder { + pub fn new(var: nn::VarBuilder, config: &EncoderConfig) -> Result { + let vs = var.pp("layers"); + let mut layers: Vec = Vec::new(); + for index in 0..config.num_hidden_layers() { + let layer = ChineseClipVisionEncoderLayer::new(vs.pp(index.to_string()), config)?; + layers.push(layer) + } + Ok(ChineseClipVisionEncoder { layers }) + } + + pub fn forward(&self, xs: &Tensor, causal_attention_mask: Option<&Tensor>) -> Result { + let mut xs = xs.clone(); + for layer in self.layers.iter() { + xs = layer.forward(&xs, causal_attention_mask)?; + } + Ok(xs) + } + + // required by LLaVA + pub fn output_hidden_states( + &self, + xs: &Tensor, + causal_attention_mask: Option<&Tensor>, + ) -> Result> { + let mut xs = xs.clone(); + let mut hidden_states = Vec::new(); + for layer in self.layers.iter() { + xs = layer.forward(&xs, causal_attention_mask)?; + hidden_states.push(xs.clone()); + } + Ok(hidden_states) + } +} + +#[derive(Clone, Debug)] +pub struct ChineseClipVisionTransformer { + embeddings: ChineseClipVisionEmbeddings, + encoder: ChineseClipVisionEncoder, + pre_layer_norm: nn::LayerNorm, + final_layer_norm: nn::LayerNorm, +} + +impl ChineseClipVisionTransformer { + pub fn new(var: nn::VarBuilder, config: &ChineseClipVisionConfig) -> Result { + let embed_dim = config.hidden_size; + let embeddings = ChineseClipVisionEmbeddings::new(var.pp("embeddings"), config)?; + let pre_layer_norm = + nn::layer_norm(embed_dim, config.layer_norm_eps, var.pp("pre_layrnorm"))?; + let encoder = ChineseClipVisionEncoder::new( + var.pp("encoder"), + &EncoderConfig::Vision(config.clone()), + )?; + let final_layer_norm = + nn::layer_norm(embed_dim, config.layer_norm_eps, var.pp("post_layernorm"))?; + Ok(Self { + embeddings, + encoder, + final_layer_norm, + pre_layer_norm, + }) + } + // required by LLaVA + pub fn output_hidden_states(&self, pixel_values: &Tensor) -> Result> { + let hidden_states = pixel_values + .apply(&self.embeddings)? + .apply(&self.pre_layer_norm)?; + + let mut result = self.encoder.output_hidden_states(&hidden_states, None)?; + let encoder_outputs = result.last().context("no last")?; + let pooled_output = encoder_outputs.i((.., 0, ..))?; + result.push(self.final_layer_norm.forward(&pooled_output)?.clone()); + Ok(result) + } +} + +impl Module for ChineseClipVisionTransformer { + fn forward(&self, pixel_values: &Tensor) -> Result { + let hidden_states = pixel_values + .apply(&self.embeddings)? + .apply(&self.pre_layer_norm)?; + + let encoder_outputs = self.encoder.forward(&hidden_states, None)?; + + // referer: https://github.com/huggingface/transformers/blob/f6fa0f0bf0796ac66f201f23bdb8585de1609add/src/transformers/models/clip/modeling_clip.py#L787 + let pooled_output = encoder_outputs.i((.., 0, ..))?; + self.final_layer_norm.forward(&pooled_output) + } +} diff --git a/patches/candle-transformers/src/models/clip/mod.rs b/patches/candle-transformers/src/models/clip/mod.rs new file mode 100644 index 0000000000..2b00267317 --- /dev/null +++ b/patches/candle-transformers/src/models/clip/mod.rs @@ -0,0 +1,152 @@ +//! Contrastive Language-Image Pre-Training +//! +//! Contrastive Language-Image Pre-Training (CLIP) is an architecture trained on +//! pairs of images with related texts. +//! +//! - 💻 [GH Link](https://github.com/openai/CLIP) +//! - 💻 Transformers Python [reference implementation](https://github.com/huggingface/transformers/tree/f6fa0f0bf0796ac66f201f23bdb8585de1609add/src/transformers/models/clip) +//! - 🤗 [HF Model](https://huggingface.co/openai/clip-vit-large-patch14-336) +//! + +use self::{ + text_model::{Activation, ClipTextTransformer}, + vision_model::ClipVisionTransformer, +}; +use candle::{Result, Tensor, D}; + +pub mod text_model; +pub mod vision_model; + +#[derive(Clone, Debug)] +pub struct ClipModel { + text_model: ClipTextTransformer, + vision_model: ClipVisionTransformer, + visual_projection: candle_nn::Linear, + text_projection: candle_nn::Linear, + logit_scale: Tensor, +} + +#[derive(Clone, Debug)] +pub enum EncoderConfig { + Text(text_model::ClipTextConfig), + Vision(vision_model::ClipVisionConfig), +} + +impl EncoderConfig { + pub fn embed_dim(&self) -> usize { + match self { + Self::Text(c) => c.embed_dim, + Self::Vision(c) => c.embed_dim, + } + } + + pub fn num_attention_heads(&self) -> usize { + match self { + Self::Text(c) => c.num_attention_heads, + Self::Vision(c) => c.num_attention_heads, + } + } + + pub fn intermediate_size(&self) -> usize { + match self { + Self::Text(c) => c.intermediate_size, + Self::Vision(c) => c.intermediate_size, + } + } + + pub fn num_hidden_layers(&self) -> usize { + match self { + Self::Text(c) => c.num_hidden_layers, + Self::Vision(c) => c.num_hidden_layers, + } + } + + pub fn activation(&self) -> Activation { + match self { + Self::Text(_c) => Activation::QuickGelu, + Self::Vision(c) => c.activation, + } + } +} + +#[derive(Clone, Debug)] +pub struct ClipConfig { + pub text_config: text_model::ClipTextConfig, + pub vision_config: vision_model::ClipVisionConfig, + pub logit_scale_init_value: f32, + pub image_size: usize, +} + +impl ClipConfig { + // base image size is 224, model size is 600Mb + pub fn vit_base_patch32() -> Self { + let text_config = text_model::ClipTextConfig::vit_base_patch32(); + let vision_config = vision_model::ClipVisionConfig::vit_base_patch32(); + + Self { + text_config, + vision_config, + logit_scale_init_value: 2.6592, + image_size: 224, + } + } +} + +impl ClipModel { + pub fn new(vs: candle_nn::VarBuilder, c: &ClipConfig) -> Result { + let text_model = ClipTextTransformer::new(vs.pp("text_model"), &c.text_config)?; + let vision_model = ClipVisionTransformer::new(vs.pp("vision_model"), &c.vision_config)?; + let visual_projection = candle_nn::linear_no_bias( + c.vision_config.embed_dim, + c.vision_config.projection_dim, + vs.pp("visual_projection"), + )?; + let text_projection = candle_nn::linear_no_bias( + c.text_config.embed_dim, + c.text_config.projection_dim, + vs.pp("text_projection"), + )?; + // originally nn.Parameter + let logit_scale = if vs.contains_tensor("logit_scale") { + vs.get(&[], "logit_scale")? + } else { + Tensor::new(&[c.logit_scale_init_value], vs.device())? + }; + Ok(Self { + text_model, + vision_model, + visual_projection, + text_projection, + logit_scale, + }) + } + + pub fn get_text_features(&self, input_ids: &Tensor) -> Result { + input_ids + .apply(&self.text_model)? + .apply(&self.text_projection) + } + + pub fn get_image_features(&self, pixel_values: &Tensor) -> Result { + pixel_values + .apply(&self.vision_model)? + .apply(&self.visual_projection) + } + + pub fn forward(&self, pixel_values: &Tensor, input_ids: &Tensor) -> Result<(Tensor, Tensor)> { + let image_features = self.get_image_features(pixel_values)?; + let text_features = self.get_text_features(input_ids)?; + let image_features_normalized = div_l2_norm(&image_features)?; + let text_features_normalized = div_l2_norm(&text_features)?; + let logits_per_text = text_features_normalized.matmul(&image_features_normalized.t()?)?; + let logit_scale = self.logit_scale.exp()?; + let logits_per_text = logits_per_text.broadcast_mul(&logit_scale)?; + let logits_per_image = logits_per_text.t()?; + Ok((logits_per_text, logits_per_image)) + } +} + +pub fn div_l2_norm(v: &Tensor) -> Result { + let l2_norm = v.sqr()?.sum_keepdim(D::Minus1)?.sqrt()?; + v.broadcast_div(&l2_norm) +} diff --git a/patches/candle-transformers/src/models/clip/text_model.rs b/patches/candle-transformers/src/models/clip/text_model.rs new file mode 100644 index 0000000000..eb103bd29a --- /dev/null +++ b/patches/candle-transformers/src/models/clip/text_model.rs @@ -0,0 +1,347 @@ +//! Contrastive Language-Image Pre-Training +//! +//! Contrastive Language-Image Pre-Training (CLIP) is an architecture trained on +//! pairs of images with related texts. +//! +//! - [GH](https://github.com/openai/CLIP) +//! - [Code](https://github.com/huggingface/transformers/tree/f6fa0f0bf0796ac66f201f23bdb8585de1609add/src/transformers/models/clip) + +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn as nn; +use candle_nn::Module; + +use super::EncoderConfig; + +#[derive(Debug, Clone, Copy)] +pub enum Activation { + QuickGelu, +} + +impl Module for Activation { + fn forward(&self, xs: &Tensor) -> Result { + match self { + Activation::QuickGelu => xs * nn::ops::sigmoid(&(xs * 1.702f64)?)?, + } + } +} + +#[derive(Debug, Clone)] +pub struct ClipTextConfig { + pub vocab_size: usize, + pub embed_dim: usize, + pub activation: Activation, + pub intermediate_size: usize, + pub max_position_embeddings: usize, + pub pad_with: Option, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + #[allow(dead_code)] + pub projection_dim: usize, +} + +impl ClipTextConfig { + // The config details can be found in the "text_config" section of this json file: + // https://huggingface.co/openai/clip-vit-large-patch14/blob/main/config.json + pub fn vit_base_patch32() -> Self { + Self { + vocab_size: 49408, + embed_dim: 512, + intermediate_size: 2048, + max_position_embeddings: 77, + pad_with: None, + num_hidden_layers: 12, + num_attention_heads: 8, + projection_dim: 512, + activation: Activation::QuickGelu, + } + } +} + +// ClipTextEmbeddings mostly based on the existing implementation in the stable diffision model. +// TODO rewrite to be more similar to https://github.com/huggingface/transformers/blob/f6fa0f0bf0796ac66f201f23bdb8585de1609add/src/transformers/models/clip/modeling_clip.py#L142 +#[derive(Clone, Debug)] +struct ClipTextEmbeddings { + token_embedding: candle_nn::Embedding, + position_embedding: candle_nn::Embedding, + position_ids: Tensor, +} + +impl ClipTextEmbeddings { + fn new(vs: candle_nn::VarBuilder, c: &ClipTextConfig) -> Result { + let token_embedding = + candle_nn::embedding(c.vocab_size, c.embed_dim, vs.pp("token_embedding"))?; + let position_embedding: nn::Embedding = candle_nn::embedding( + c.max_position_embeddings, + c.embed_dim, + vs.pp("position_embedding"), + )?; + let position_ids = + Tensor::arange(0u32, c.max_position_embeddings as u32, vs.device())?.unsqueeze(0)?; + Ok(Self { + token_embedding, + position_embedding, + position_ids, + }) + } +} + +impl Module for ClipTextEmbeddings { + fn forward(&self, input_ids: &Tensor) -> Result { + let seq_length = input_ids.dim(D::Minus1)?; + let inputs_embeds = self.token_embedding.forward(input_ids)?; + let position_ids = self.position_ids.narrow(1, 0, seq_length)?; + let position_embedding = self.position_embedding.forward(&position_ids)?; + inputs_embeds.broadcast_add(&position_embedding) + } +} + +#[derive(Clone, Debug)] +struct ClipAttention { + k_proj: candle_nn::Linear, + v_proj: candle_nn::Linear, + q_proj: candle_nn::Linear, + out_proj: candle_nn::Linear, + head_dim: usize, + scale: f64, + num_attention_heads: usize, +} + +impl ClipAttention { + fn new(vs: candle_nn::VarBuilder, c: &EncoderConfig) -> Result { + let embed_dim = c.embed_dim(); + let num_attention_heads = c.num_attention_heads(); + let k_proj = candle_nn::linear(embed_dim, embed_dim, vs.pp("k_proj"))?; + let v_proj = candle_nn::linear(embed_dim, embed_dim, vs.pp("v_proj"))?; + let q_proj = candle_nn::linear(embed_dim, embed_dim, vs.pp("q_proj"))?; + let out_proj = candle_nn::linear(embed_dim, embed_dim, vs.pp("out_proj"))?; + let head_dim = embed_dim / num_attention_heads; + let scale = (head_dim as f64).powf(-0.5); + + Ok(ClipAttention { + k_proj, + v_proj, + q_proj, + out_proj, + head_dim, + scale, + num_attention_heads, + }) + } + + fn shape(&self, xs: &Tensor, seq_len: usize, bsz: usize) -> Result { + xs.reshape((bsz, seq_len, self.num_attention_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous() + } + + fn forward(&self, xs: &Tensor, causal_attention_mask: Option<&Tensor>) -> Result { + let in_dtype = xs.dtype(); + let (bsz, seq_len, embed_dim) = xs.dims3()?; + + let query_states = (self.q_proj.forward(xs)? * self.scale)?; + let proj_shape = (bsz * self.num_attention_heads, seq_len, self.head_dim); + let query_states = self + .shape(&query_states, seq_len, bsz)? + .reshape(proj_shape)? + .to_dtype(DType::F32)?; + let key_states = self + .shape(&self.k_proj.forward(xs)?, seq_len, bsz)? + .reshape(proj_shape)? + .to_dtype(DType::F32)?; + let value_states = self + .shape(&self.v_proj.forward(xs)?, seq_len, bsz)? + .reshape(proj_shape)? + .to_dtype(DType::F32)?; + let attn_weights = query_states.matmul(&key_states.transpose(1, 2)?)?; + + let src_len = key_states.dim(1)?; + + let attn_weights = if let Some(causal_attention_mask) = causal_attention_mask { + attn_weights + .reshape((bsz, self.num_attention_heads, seq_len, src_len))? + .broadcast_add(causal_attention_mask)? + .reshape((bsz * self.num_attention_heads, seq_len, src_len))? + } else { + attn_weights + }; + + let attn_weights = candle_nn::ops::softmax(&attn_weights, D::Minus1)?; + + let attn_output = attn_weights.matmul(&value_states)?.to_dtype(in_dtype)?; + let attn_output = attn_output + .reshape((bsz, self.num_attention_heads, seq_len, self.head_dim))? + .transpose(1, 2)? + .reshape((bsz, seq_len, embed_dim))?; + self.out_proj.forward(&attn_output) + } +} + +#[derive(Clone, Debug)] +struct ClipMlp { + fc1: candle_nn::Linear, + fc2: candle_nn::Linear, + activation: Activation, +} + +impl ClipMlp { + fn new(vs: candle_nn::VarBuilder, c: &EncoderConfig) -> Result { + let fc1 = candle_nn::linear(c.embed_dim(), c.intermediate_size(), vs.pp("fc1"))?; + let fc2 = candle_nn::linear(c.intermediate_size(), c.embed_dim(), vs.pp("fc2"))?; + + Ok(ClipMlp { + fc1, + fc2, + activation: c.activation(), + }) + } +} + +impl ClipMlp { + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.fc1.forward(xs)?; + self.fc2.forward(&self.activation.forward(&xs)?) + } +} + +#[derive(Clone, Debug)] +struct ClipEncoderLayer { + self_attn: ClipAttention, + layer_norm1: candle_nn::LayerNorm, + mlp: ClipMlp, + layer_norm2: candle_nn::LayerNorm, +} + +impl ClipEncoderLayer { + fn new(vs: candle_nn::VarBuilder, c: &EncoderConfig) -> Result { + let self_attn = ClipAttention::new(vs.pp("self_attn"), c)?; + let layer_norm1 = candle_nn::layer_norm(c.embed_dim(), 1e-5, vs.pp("layer_norm1"))?; + let mlp = ClipMlp::new(vs.pp("mlp"), c)?; + let layer_norm2 = candle_nn::layer_norm(c.embed_dim(), 1e-5, vs.pp("layer_norm2"))?; + + Ok(ClipEncoderLayer { + self_attn, + layer_norm1, + mlp, + layer_norm2, + }) + } + + fn forward(&self, xs: &Tensor, causal_attention_mask: Option<&Tensor>) -> Result { + let residual = xs; + let xs = self.layer_norm1.forward(xs)?; + let xs = self.self_attn.forward(&xs, causal_attention_mask)?; + let xs = (xs + residual)?; + + let residual = &xs; + let xs = self.layer_norm2.forward(&xs)?; + let xs = self.mlp.forward(&xs)?; + xs + residual + } +} + +#[derive(Clone, Debug)] +pub struct ClipEncoder { + layers: Vec, +} + +impl ClipEncoder { + pub fn new(vs: candle_nn::VarBuilder, c: &EncoderConfig) -> Result { + let vs = vs.pp("layers"); + let mut layers: Vec = Vec::new(); + for index in 0..c.num_hidden_layers() { + let layer = ClipEncoderLayer::new(vs.pp(index.to_string()), c)?; + layers.push(layer) + } + Ok(ClipEncoder { layers }) + } + + pub fn forward(&self, xs: &Tensor, causal_attention_mask: Option<&Tensor>) -> Result { + let mut xs = xs.clone(); + for layer in self.layers.iter() { + xs = layer.forward(&xs, causal_attention_mask)?; + } + Ok(xs) + } + // required by LLaVA + pub fn output_hidden_states( + &self, + xs: &Tensor, + causal_attention_mask: Option<&Tensor>, + ) -> Result> { + let mut xs = xs.clone(); + let mut hidden_states = Vec::new(); + for layer in self.layers.iter() { + xs = layer.forward(&xs, causal_attention_mask)?; + hidden_states.push(xs.clone()); + } + Ok(hidden_states) + } +} + +/// A CLIP transformer based model. +#[derive(Clone, Debug)] +pub struct ClipTextTransformer { + embeddings: ClipTextEmbeddings, + encoder: ClipEncoder, + final_layer_norm: candle_nn::LayerNorm, +} + +impl ClipTextTransformer { + pub fn new(vs: candle_nn::VarBuilder, c: &ClipTextConfig) -> Result { + let embeddings = ClipTextEmbeddings::new(vs.pp("embeddings"), c)?; + let encoder = ClipEncoder::new(vs.pp("encoder"), &EncoderConfig::Text(c.clone()))?; + let final_layer_norm = candle_nn::layer_norm(c.embed_dim, 1e-5, vs.pp("final_layer_norm"))?; + Ok(ClipTextTransformer { + embeddings, + encoder, + final_layer_norm, + }) + } + + // TODO: rewrite to newer version + fn build_causal_attention_mask( + bsz: usize, + seq_len: usize, + mask_after: usize, + device: &Device, + ) -> Result { + let mask: Vec<_> = (0..seq_len) + .flat_map(|i| { + (0..seq_len).map(move |j| { + if j > i || j > mask_after { + f32::MIN + } else { + 0. + } + }) + }) + .collect(); + let mask = Tensor::from_slice(&mask, (seq_len, seq_len), device)?; + mask.broadcast_as((bsz, 1, seq_len, seq_len)) + } + + pub fn forward_with_mask(&self, input_ids: &Tensor, mask_after: usize) -> Result { + let (bsz, seq_len) = input_ids.dims2()?; + let input_ids = self.embeddings.forward(input_ids)?; + let causal_attention_mask = + Self::build_causal_attention_mask(bsz, seq_len, mask_after, input_ids.device())?; + let input_ids = self + .encoder + .forward(&input_ids, Some(&causal_attention_mask))?; + self.final_layer_norm.forward(&input_ids) + } +} + +impl Module for ClipTextTransformer { + fn forward(&self, input_ids: &Tensor) -> Result { + let output = self.forward_with_mask(input_ids, usize::MAX)?; + let sequence_max_indices = input_ids.argmax(D::Minus1)?.to_dtype(DType::I64)?; + + let mut indices = Vec::new(); + for (batch_idx, &seq_idx) in sequence_max_indices.to_vec1::()?.iter().enumerate() { + let index = output.i((batch_idx, seq_idx as usize))?.unsqueeze(0)?; + indices.push(index); + } + Tensor::cat(&indices, 0) + } +} diff --git a/patches/candle-transformers/src/models/clip/vision_model.rs b/patches/candle-transformers/src/models/clip/vision_model.rs new file mode 100644 index 0000000000..9031442017 --- /dev/null +++ b/patches/candle-transformers/src/models/clip/vision_model.rs @@ -0,0 +1,171 @@ +//! Contrastive Language-Image Pre-Training +//! +//! Contrastive Language-Image Pre-Training (CLIP) is an architecture trained on +//! pairs of images with related texts. +//! +//! https://github.com/openai/CLIP +//! https://github.com/huggingface/transformers/tree/f6fa0f0bf0796ac66f201f23bdb8585de1609add/src/transformers/models/clip + +use candle::{Context, IndexOp, Result, Shape, Tensor, D}; +use candle_nn as nn; +use candle_nn::Module; +use nn::Conv2dConfig; + +use super::{ + text_model::{Activation, ClipEncoder}, + EncoderConfig, +}; + +#[derive(Debug, Clone)] +pub struct ClipVisionConfig { + pub embed_dim: usize, + pub activation: Activation, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + #[allow(dead_code)] + pub projection_dim: usize, + pub num_channels: usize, + pub image_size: usize, + pub patch_size: usize, +} + +impl ClipVisionConfig { + // The config details can be found in the "vision_config" section of this json file: + // https://huggingface.co/openai/clip-vit-large-patch14/blob/main/config.json + pub fn vit_base_patch32() -> Self { + Self { + embed_dim: 768, + activation: Activation::QuickGelu, + intermediate_size: 3072, + num_hidden_layers: 12, + num_attention_heads: 12, + projection_dim: 512, + num_channels: 3, + image_size: 224, + patch_size: 32, + } + } + pub fn clip_vit_large_patch14_336() -> Self { + Self { + embed_dim: 1024, + activation: Activation::QuickGelu, + intermediate_size: 4096, + num_hidden_layers: 24, + num_attention_heads: 16, + projection_dim: 768, + num_channels: 3, + image_size: 336, + patch_size: 14, + } + } +} + +// https://github.com/huggingface/transformers/blob/f6fa0f0bf0796ac66f201f23bdb8585de1609add/src/transformers/models/clip/modeling_clip.py#L112 +#[derive(Clone, Debug)] +struct ClipVisionEmbeddings { + patch_embedding: candle_nn::Conv2d, + position_ids: Tensor, + class_embedding: Tensor, + position_embedding: candle_nn::Embedding, +} + +impl ClipVisionEmbeddings { + fn new(vs: candle_nn::VarBuilder, c: &ClipVisionConfig) -> Result { + // originally nn.Parameter + let class_embedding = if vs.contains_tensor("class_embedding") { + vs.get(c.embed_dim, "class_embedding")? + } else { + Tensor::randn(0f32, 1f32, c.embed_dim, vs.device())? + }; + + let num_patches = (c.image_size / c.patch_size).pow(2); + let num_positions = num_patches + 1; + let position_ids = Tensor::arange(0, num_positions as i64, vs.device())?; + + let conv2dconfig = Conv2dConfig { + stride: c.patch_size, + ..Default::default() + }; + let position_embedding = + candle_nn::embedding(num_positions, c.embed_dim, vs.pp("position_embedding"))?; + let patch_embedding = candle_nn::conv2d_no_bias( + c.num_channels, + c.embed_dim, + c.patch_size, + conv2dconfig, + vs.pp("patch_embedding"), + )?; + Ok(Self { + patch_embedding, + position_ids, + class_embedding, + position_embedding, + }) + } +} + +impl Module for ClipVisionEmbeddings { + fn forward(&self, pixel_values: &Tensor) -> Result { + let batch_size = pixel_values.shape().dims(); + let patch_embeds = self + .patch_embedding + .forward(pixel_values)? + .flatten_from(2)? + .transpose(1, 2)?; + let shape = Shape::from((batch_size[0], 1, self.class_embedding.dim(D::Minus1)?)); + let class_embeds = self.class_embedding.expand(shape)?; + let embeddings = Tensor::cat(&[class_embeds, patch_embeds], 1)?; + let position_embedding = self.position_embedding.forward(&self.position_ids)?; + embeddings.broadcast_add(&position_embedding) + } +} + +// https://github.com/huggingface/transformers/blob/f6fa0f0bf0796ac66f201f23bdb8585de1609add/src/transformers/models/clip/modeling_clip.py#L743 +#[derive(Clone, Debug)] +pub struct ClipVisionTransformer { + embeddings: ClipVisionEmbeddings, + encoder: ClipEncoder, + pre_layer_norm: candle_nn::LayerNorm, + final_layer_norm: candle_nn::LayerNorm, +} + +impl ClipVisionTransformer { + pub fn new(vs: candle_nn::VarBuilder, c: &ClipVisionConfig) -> Result { + let embeddings = ClipVisionEmbeddings::new(vs.pp("embeddings"), c)?; + let pre_layer_norm = candle_nn::layer_norm(c.embed_dim, 1e-5, vs.pp("pre_layrnorm"))?; + let encoder = ClipEncoder::new(vs.pp("encoder"), &EncoderConfig::Vision(c.clone()))?; + let final_layer_norm = candle_nn::layer_norm(c.embed_dim, 1e-5, vs.pp("post_layernorm"))?; + Ok(Self { + embeddings, + encoder, + final_layer_norm, + pre_layer_norm, + }) + } + // required by LLaVA + pub fn output_hidden_states(&self, pixel_values: &Tensor) -> Result> { + let hidden_states = pixel_values + .apply(&self.embeddings)? + .apply(&self.pre_layer_norm)?; + let mut result = self.encoder.output_hidden_states(&hidden_states, None)?; + let encoder_outputs = result.last().context("no last")?; + let pooled_output = encoder_outputs.i((.., 0, ..))?; + result.push(self.final_layer_norm.forward(&pooled_output)?.clone()); + Ok(result) + } +} + +impl Module for ClipVisionTransformer { + fn forward(&self, pixel_values: &Tensor) -> Result { + let hidden_states = pixel_values + .apply(&self.embeddings)? + .apply(&self.pre_layer_norm)?; + + let encoder_outputs = self.encoder.forward(&hidden_states, None)?; + // https://github.com/huggingface/transformers/blob/f6fa0f0bf0796ac66f201f23bdb8585de1609add/src/transformers/models/clip/modeling_clip.py#L787 + // pooled_output = encoder_outputs[:, 0, :] + let pooled_output = encoder_outputs.i((.., 0, ..))?; + self.final_layer_norm.forward(&pooled_output) + } +} diff --git a/patches/candle-transformers/src/models/codegeex4_9b.rs b/patches/candle-transformers/src/models/codegeex4_9b.rs new file mode 100644 index 0000000000..40c74ccf0f --- /dev/null +++ b/patches/candle-transformers/src/models/codegeex4_9b.rs @@ -0,0 +1,612 @@ +//! CodeGeeX4 - A multi-language code generation model +//! +//! A Pre-Trained Model For Code Generation with Multilingual Evaluations on HumanEval-X" +//! +//! - 📝 [Arxiv](https://arxiv.org/abs/2303.17568) +//! - 💻 [GitHub](https://github.com/THUDM/CodeGeeX) +//! + +use crate::models::with_tracing::{linear_b as linear, Linear}; +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::VarBuilder; + +fn default_one() -> usize { + 1 +} + +#[derive(Debug, Clone, serde::Deserialize, Default)] +pub struct Config { + pub num_layers: usize, + pub padded_vocab_size: usize, + pub hidden_size: usize, + pub ffn_hidden_size: usize, + pub kv_channels: usize, + pub num_attention_heads: usize, + pub seq_length: usize, + pub layernorm_epsilon: f64, + pub rmsnorm: bool, + pub apply_residual_connection_post_layernorm: bool, + pub post_layer_norm: bool, + pub add_bias_linear: bool, + pub add_qkv_bias: bool, + pub bias_dropout_fusion: bool, + pub multi_query_attention: bool, + pub multi_query_group_num: usize, + pub apply_query_key_layer_scaling: bool, + pub attention_softmax_in_fp32: bool, + pub fp32_residual_connection: bool, + #[serde(default = "default_one")] + pub rope_ratio: usize, +} + +impl Config { + pub fn codegeex4() -> Self { + Self { + num_layers: 40, + padded_vocab_size: 151552, + hidden_size: 4096, + ffn_hidden_size: 13696, + kv_channels: 128, + num_attention_heads: 32, + seq_length: 131072, + layernorm_epsilon: 1e-5, + rmsnorm: true, + apply_residual_connection_post_layernorm: false, + post_layer_norm: true, + add_bias_linear: false, + add_qkv_bias: true, + bias_dropout_fusion: true, + multi_query_attention: true, + multi_query_group_num: 2, + apply_query_key_layer_scaling: true, + attention_softmax_in_fp32: true, + fp32_residual_connection: false, + rope_ratio: 500, + } + } +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + cache: Tensor, +} + +impl RotaryEmbedding { + fn new(cfg: &Config, dtype: DType, dev: &Device) -> Result { + let rotary_dim = cfg.kv_channels; + let n_elem = rotary_dim / 2; + let base = 10_000f64 * cfg.rope_ratio as f64; + let inv_freq: Vec<_> = (0..n_elem) + .step_by(2) + .map(|i| 1f32 / base.powf(i as f64 / n_elem as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, cfg.seq_length as u32, dev)? + .to_dtype(dtype) + .expect("unalbe to dytpe in Rotray Embedding new") + .reshape((cfg.seq_length, 1))?; + let freqs = t.matmul(&inv_freq)?; + let cache = Tensor::stack(&[&freqs.cos()?, &freqs.sin()?], D::Minus1)?; + Ok(Self { cache }) + } + + fn apply(&self, xs: &Tensor, seqlen_offset: usize) -> Result { + let (seqlen, _b, np, _hn) = xs.dims4()?; + let cache = self.cache.narrow(0, seqlen_offset, seqlen)?; + let rot_dim = cache.dim(D::Minus2)? * 2; + let (xs, xs_pass) = ( + xs.narrow(D::Minus1, 0, rot_dim)?, + xs.narrow(D::Minus1, rot_dim, rot_dim)?, + ); + let xshaped = xs.reshape((seqlen, (), np, rot_dim / 2, 2))?; + let cache = cache.reshape((seqlen, (), 1, rot_dim / 2, 2))?; + let (xshaped0, xshaped1) = ( + xshaped.i((.., .., .., .., 0))?, + xshaped.i((.., .., .., .., 1))?, + ); + let (cache0, cache1) = (cache.i((.., .., .., .., 0))?, cache.i((.., .., .., .., 1))?); + let xs_out = Tensor::stack( + &[ + (xshaped0.broadcast_mul(&cache0)? - xshaped1.broadcast_mul(&cache1)?)?, + (xshaped1.broadcast_mul(&cache0)? + xshaped0.broadcast_mul(&cache1)?)?, + ], + D::Minus1, + )?; + let xs_out = xs_out.flatten_from(3)?; + Tensor::cat(&[xs_out, xs_pass], D::Minus1) + } +} + +#[derive(Debug, Clone)] +struct CoreAttention { + coeff: Option, + norm_factor: f64, + dtype: DType, +} + +fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: f32, dtype: DType) -> 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.to_dtype(dtype)?, on_false)?; + Ok(m) +} + +impl CoreAttention { + fn new(layer_number: usize, cfg: &Config, dtype: DType) -> Result { + let norm_factor = (cfg.kv_channels as f64).sqrt(); + let (norm_factor, coeff) = if cfg.apply_query_key_layer_scaling { + let coeff = f64::max(1.0, layer_number as f64); + (norm_factor * coeff, Some(coeff)) + } else { + (norm_factor, None) + }; + Ok(Self { + coeff, + norm_factor, + dtype, + }) + } + + fn forward( + &self, + query_layer: &Tensor, + key_layer: &Tensor, + value_layer: &Tensor, + attention_mask: &Option, + ) -> Result { + let output_size = ( + query_layer.dim(1)?, // b + query_layer.dim(2)?, // np + query_layer.dim(0)?, // sq + key_layer.dim(0)?, // sk + ); + let query_layer = + query_layer.reshape((output_size.2, output_size.0 * output_size.1, ()))?; + let key_layer = key_layer.reshape((output_size.3, output_size.0 * output_size.1, ()))?; + let matmul_result = Tensor::matmul( + &query_layer.transpose(0, 1)?.contiguous()?, + &key_layer.transpose(0, 1)?.transpose(1, 2)?.contiguous()?, + )?; + let matmul_result = (matmul_result / self.norm_factor)?.reshape(output_size)?; + let matmul_result = match self.coeff { + None => matmul_result, + Some(coeff) => (matmul_result * coeff)?, + }; + let attention_scores = match attention_mask { + Some(mask) => masked_fill( + &matmul_result, + &mask.broadcast_left((matmul_result.dim(0)?, matmul_result.dim(1)?))?, + f32::NEG_INFINITY, + self.dtype, + )?, + None => matmul_result, + }; + let attention_probs = candle_nn::ops::softmax_last_dim(&attention_scores)?; + + let output_size = ( + value_layer.dim(1)?, + value_layer.dim(2)?, + query_layer.dim(0)?, + value_layer.dim(3)?, + ); + let value_layer = + value_layer.reshape((value_layer.dim(0)?, output_size.0 * output_size.1, ()))?; + let attention_probs = + attention_probs.reshape((output_size.0 * output_size.1, output_size.2, ()))?; + let context_layer = Tensor::matmul( + &attention_probs.contiguous()?, + &value_layer.transpose(0, 1)?.contiguous()?, + )?; + let context_layer = context_layer.reshape(output_size)?; + let context_layer = context_layer.permute((2, 0, 1, 3))?.contiguous()?; + context_layer.flatten_from(D::Minus2) + } +} + +#[derive(Debug, Clone)] +struct SelfAttention { + query_key_value: Linear, + core_attention: CoreAttention, + dense: Linear, + multi_query_attention: bool, + num_attention_heads_per_partition: usize, + num_multi_query_groups_per_partition: usize, + hidden_size_per_attention_head: usize, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl SelfAttention { + fn new(layer_number: usize, cfg: &Config, vb: VarBuilder) -> Result { + let projection_size = cfg.kv_channels * cfg.num_attention_heads; + let hidden_size_per_attention_head = projection_size / cfg.num_attention_heads; + let qkv_hidden_size = if cfg.multi_query_attention { + projection_size + 2 * hidden_size_per_attention_head * cfg.multi_query_group_num + } else { + 3 * projection_size + }; + let query_key_value = linear( + cfg.hidden_size, + qkv_hidden_size, + cfg.add_bias_linear || cfg.add_qkv_bias, + vb.pp("query_key_value"), + )?; + let core_attention = CoreAttention::new(layer_number, cfg, vb.dtype())?; + let dense = linear( + cfg.hidden_size, + cfg.hidden_size, + cfg.add_bias_linear, + vb.pp("dense"), + )?; + Ok(Self { + query_key_value, + core_attention, + dense, + multi_query_attention: cfg.multi_query_attention, + num_attention_heads_per_partition: cfg.num_attention_heads, + num_multi_query_groups_per_partition: cfg.multi_query_group_num, + hidden_size_per_attention_head: cfg.kv_channels, + kv_cache: None, + }) + } + + fn reset_kv_cache(&mut self) { + self.kv_cache = None + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: &Option, + rotary_emb: &RotaryEmbedding, + ) -> Result { + let mixed_x_layer = xs.apply(&self.query_key_value)?; + if !self.multi_query_attention { + candle::bail!("only multi_query_attention=true is supported") + } + let hpa = self.hidden_size_per_attention_head; + let query_layer = + mixed_x_layer.narrow(D::Minus1, 0, self.num_attention_heads_per_partition * hpa)?; + let key_layer = mixed_x_layer.narrow( + D::Minus1, + self.num_attention_heads_per_partition * hpa, + self.num_multi_query_groups_per_partition * hpa, + )?; + let value_layer = mixed_x_layer.narrow( + D::Minus1, + self.num_attention_heads_per_partition * hpa + + self.num_multi_query_groups_per_partition * hpa, + self.num_multi_query_groups_per_partition * hpa, + )?; + let query_layer = query_layer.reshape(( + query_layer.dim(0)?, + query_layer.dim(1)?, + self.num_attention_heads_per_partition, + hpa, + ))?; + let key_layer = key_layer.reshape(( + key_layer.dim(0)?, + key_layer.dim(1)?, + self.num_multi_query_groups_per_partition, + hpa, + ))?; + let value_layer = value_layer.reshape(( + value_layer.dim(0)?, + value_layer.dim(1)?, + self.num_multi_query_groups_per_partition, + hpa, + ))?; + + // Rotary embeddings. + let seqlen_offset = match &self.kv_cache { + None => 0, + Some((prev_k, _)) => prev_k.dim(0)?, + }; + let query_layer = rotary_emb.apply(&query_layer, seqlen_offset)?; + let key_layer = rotary_emb.apply(&key_layer, seqlen_offset)?; + + // KV cache. + let (key_layer, value_layer) = match &self.kv_cache { + None => (key_layer, value_layer), + Some((prev_k, prev_v)) => { + let k = Tensor::cat(&[prev_k, &key_layer], 0)?; + let v = Tensor::cat(&[prev_v, &value_layer], 0)?; + (k, v) + } + }; + self.kv_cache = Some((key_layer.clone(), value_layer.clone())); + + // Repeat KV. + let ratio = + self.num_attention_heads_per_partition / self.num_multi_query_groups_per_partition; + let key_layer = { + let (d0, d1, d2, d3) = key_layer.dims4()?; + key_layer + .unsqueeze(D::Minus2)? + .expand((d0, d1, d2, ratio, d3))? + .reshape(( + d0, + d1, + self.num_attention_heads_per_partition, + self.hidden_size_per_attention_head, + ))? + }; + let value_layer = { + let (d0, d1, d2, d3) = value_layer.dims4()?; + value_layer + .unsqueeze(D::Minus2)? + .expand((d0, d1, d2, ratio, d3))? + .reshape(( + d0, + d1, + self.num_attention_heads_per_partition, + self.hidden_size_per_attention_head, + ))? + }; + + let context_layer = + self.core_attention + .forward(&query_layer, &key_layer, &value_layer, attention_mask)?; + let output = context_layer.apply(&self.dense)?; + Ok(output) + } +} + +#[allow(clippy::upper_case_acronyms)] +#[derive(Debug, Clone)] +struct MLP { + dense_h_to_4h: Linear, + dense_4h_to_h: Linear, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense_h_to_4h = linear( + cfg.hidden_size, + cfg.ffn_hidden_size * 2, + cfg.add_bias_linear, + vb.pp("dense_h_to_4h"), + )?; + let dense_4h_to_h = linear( + cfg.ffn_hidden_size, + cfg.hidden_size, + cfg.add_bias_linear, + vb.pp("dense_4h_to_h"), + )?; + Ok(Self { + dense_4h_to_h, + dense_h_to_4h, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.dense_h_to_4h)? + .apply(&candle_nn::Activation::Swiglu)? + .apply(&self.dense_4h_to_h) + } +} + +#[derive(Debug, Clone)] +struct Block { + input_layernorm: candle_nn::LayerNorm, + self_attention: SelfAttention, + post_attention_layernorm: candle_nn::LayerNorm, + mlp: MLP, + apply_residual_connection_post_layernorm: bool, +} + +impl Block { + fn new(layer_number: usize, cfg: &Config, vb: VarBuilder) -> Result { + let input_layernorm = if cfg.rmsnorm { + candle_nn::rms_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("input_layernorm"), + )? + .into_inner() + } else { + candle_nn::layer_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("input_layernorm"), + )? + }; + let post_attention_layernorm = if cfg.rmsnorm { + candle_nn::rms_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("post_attention_layernorm"), + )? + .into_inner() + } else { + candle_nn::layer_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("post_attention_layernorm"), + )? + }; + let self_attention = SelfAttention::new(layer_number, cfg, vb.pp("self_attention"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + Ok(Self { + input_layernorm, + self_attention, + post_attention_layernorm, + mlp, + apply_residual_connection_post_layernorm: cfg.apply_residual_connection_post_layernorm, + }) + } + + fn reset_kv_cache(&mut self) { + self.self_attention.reset_kv_cache() + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: &Option, + rotary_emb: &RotaryEmbedding, + ) -> Result { + let layernorm_output = xs.apply(&self.input_layernorm)?; + let attention_output = + self.self_attention + .forward(&layernorm_output, attention_mask, rotary_emb)?; + let residual = if self.apply_residual_connection_post_layernorm { + &layernorm_output + } else { + xs + }; + let layernorm_input = (residual + attention_output)?; + let layernorm_output = layernorm_input.apply(&self.post_attention_layernorm)?; + let mlp_output = layernorm_output.apply(&self.mlp)?; + let residual = if self.apply_residual_connection_post_layernorm { + &layernorm_output + } else { + &layernorm_input + }; + mlp_output + residual + } +} + +#[derive(Debug, Clone)] +struct Transformer { + layers: Vec, + final_layernorm: Option, + rotary_emb: RotaryEmbedding, +} + +impl Transformer { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_l = vb.pp("layers"); + let mut layers = Vec::with_capacity(cfg.num_layers); + for layer_index in 0..cfg.num_layers { + let block = Block::new(layer_index + 1, cfg, vb_l.pp(layer_index))?; + layers.push(block) + } + let final_layernorm = if cfg.post_layer_norm { + let ln = if cfg.rmsnorm { + candle_nn::rms_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("final_layernorm"), + )? + .into_inner() + } else { + candle_nn::layer_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("final_layernorm"), + )? + }; + Some(ln) + } else { + None + }; + let rotary_emb = RotaryEmbedding::new(cfg, vb.dtype(), vb.device())?; + Ok(Self { + layers, + final_layernorm, + rotary_emb, + }) + } + + fn reset_kv_cache(&mut self) { + for block in self.layers.iter_mut() { + block.reset_kv_cache() + } + } + + fn forward(&mut self, xs: &Tensor, attention_mask: &Option) -> Result { + let mut xs = xs.clone(); + for block in self.layers.iter_mut() { + xs = block.forward(&xs, attention_mask, &self.rotary_emb)? + } + match self.final_layernorm.as_ref() { + None => Ok(xs), + Some(ln) => xs.apply(ln), + } + } +} + +#[derive(Debug, Clone)] +struct Embedding { + word_embeddings: candle_nn::Embedding, + fp32_residual_connection: bool, +} + +impl Embedding { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let word_embeddings = candle_nn::embedding( + cfg.padded_vocab_size, + cfg.hidden_size, + vb.pp("word_embeddings"), + )?; + Ok(Self { + word_embeddings, + fp32_residual_connection: cfg.fp32_residual_connection, + }) + } +} + +impl Module for Embedding { + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.word_embeddings.forward(xs)?.transpose(0, 1)?; // b,s,h -> s,b,h + if self.fp32_residual_connection { + xs.to_dtype(candle::DType::F32) + } else { + xs.contiguous() + } + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embedding: Embedding, + encoder: Transformer, + output_layer: Linear, +} + +fn get_mask(size: usize, device: &Device) -> Result { + let mask: Vec<_> = (0..size) + .flat_map(|i| (0..size).map(move |j| u8::from(j > i))) + .collect(); + Tensor::from_slice(&mask, (size, size), device) +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb = vb.pp("transformer"); + let embedding = Embedding::new(cfg, vb.pp("embedding"))?; + let encoder = Transformer::new(cfg, vb.pp("encoder"))?; + let output_layer = linear( + cfg.hidden_size, + cfg.padded_vocab_size, + false, + vb.pp("output_layer"), + )?; + + Ok(Self { + embedding, + encoder, + output_layer, + }) + } + + pub fn reset_kv_cache(&mut self) { + self.encoder.reset_kv_cache() + } + + pub fn forward(&mut self, xs: &Tensor) -> Result { + let (_b_size, seq_len) = xs.dims2()?; + let input_embeds = xs.apply(&self.embedding)?; + let attention_mask = if seq_len <= 1 { + None + } else { + Some(get_mask(seq_len, xs.device())?) + }; + let xs = self.encoder.forward(&input_embeds, &attention_mask)?; + let lm_logits = xs.i(seq_len - 1)?.apply(&self.output_layer)?; + Ok(lm_logits) + } +} diff --git a/patches/candle-transformers/src/models/colpali.rs b/patches/candle-transformers/src/models/colpali.rs new file mode 100644 index 0000000000..16ca4eb304 --- /dev/null +++ b/patches/candle-transformers/src/models/colpali.rs @@ -0,0 +1,47 @@ +//! Colpali Model for text/image similarity scoring. +//! +//! Colpali combines a vision encoder with an efficient LM for retrieving content. +//! + +use candle::{Module, Result, Tensor}; +use candle_nn::VarBuilder; + +use super::paligemma; +use candle_nn::{linear, Linear}; + +pub struct Model { + pub model: paligemma::Model, + pub custom_text_projection: Linear, +} + +impl Model { + pub fn new(config: &paligemma::Config, vb: VarBuilder) -> Result { + let model = paligemma::Model::new(config, vb.pp("model"))?; + let custom_text_projection = linear( + config.text_config.hidden_size, + 128, + vb.pp("custom_text_proj"), + )?; + + Ok(Self { + model, + custom_text_projection, + }) + } + + pub fn forward_images(&mut self, pixel_values: &Tensor, input_ids: &Tensor) -> Result { + let outputs = self + .model + .setup_without_projection(pixel_values, input_ids)?; + let outputs = self.custom_text_projection.forward(&outputs)?; + let outputs = outputs.broadcast_div(&outputs.sqr()?.sum_keepdim(2)?.sqrt()?)?; + Ok(outputs) + } + + pub fn forward_text(&mut self, input_ids: &Tensor) -> Result { + let outputs = self.model.forward_without_projection(input_ids)?; + let outputs = self.custom_text_projection.forward(&outputs)?; + let outputs = outputs.broadcast_div(&outputs.sqr()?.sum_keepdim(2)?.sqrt()?)?; + Ok(outputs) + } +} diff --git a/patches/candle-transformers/src/models/convmixer.rs b/patches/candle-transformers/src/models/convmixer.rs new file mode 100644 index 0000000000..77570a6aaf --- /dev/null +++ b/patches/candle-transformers/src/models/convmixer.rs @@ -0,0 +1,89 @@ +//! ConvMixer implementation. +//! +//! See "Patches Are All You Need?" by Trockman et al. 2022 +//! +//! - 📝 [Arxiv](https://arxiv.org/abs/2201.09792) +//! - 💻 [GitHub](https://github.com/locuslab/convmixer) +//! +use candle::Result; +use candle_nn::{batch_norm, Conv2dConfig, Module, VarBuilder}; + +#[allow(clippy::many_single_char_names)] +fn conv2d_same( + i: usize, + o: usize, + k: usize, + c: Conv2dConfig, + vb: VarBuilder, +) -> Result { + let conv2d = candle_nn::conv2d(i, o, k, c, vb)?; + let s = c.stride; + let module = candle_nn::func(move |xs| { + let ih = xs.dim(2)?; + let iw = xs.dim(3)?; + let oh = ih.div_ceil(s); + let ow = iw.div_ceil(s); + let pad_h = usize::max((oh - 1) * s + k - ih, 0); + let pad_w = usize::max((ow - 1) * s + k - iw, 0); + if pad_h > 0 || pad_w > 0 { + xs.pad_with_zeros(3, pad_w / 2, pad_w - pad_w / 2)? + .pad_with_zeros(2, pad_h / 2, pad_h - pad_h / 2)? + .apply(&conv2d) + } else { + xs.apply(&conv2d) + } + }); + Ok(module) +} + +fn block(dim: usize, kernel_size: usize, vb: VarBuilder) -> Result { + let conv2d_cfg = Conv2dConfig { + groups: dim, + ..Default::default() + }; + let vb_fn = vb.pp(0).pp("fn"); + let conv1 = conv2d_same(dim, dim, kernel_size, conv2d_cfg, vb_fn.pp(0))?; + let bn1 = batch_norm(dim, 1e-5, vb_fn.pp(2))?; + let conv2 = candle_nn::conv2d(dim, dim, 1, Default::default(), vb.pp(1))?; + let bn2 = batch_norm(dim, 1e-5, vb.pp(3))?; + Ok(candle_nn::func(move |xs| { + let ys = xs.apply(&conv1)?.gelu_erf()?.apply_t(&bn1, false)?; + (xs + ys)?.apply(&conv2)?.gelu_erf()?.apply_t(&bn2, false) + })) +} + +fn convmixer( + nclasses: usize, + dim: usize, + depth: usize, + kernel_size: usize, + patch_size: usize, + vb: VarBuilder, +) -> Result> { + let conv2d_cfg = Conv2dConfig { + stride: patch_size, + ..Default::default() + }; + let conv1 = candle_nn::conv2d(3, dim, patch_size, conv2d_cfg, vb.pp(0))?; + let bn1 = batch_norm(dim, 1e-5, vb.pp(2))?; + let blocks: Vec<_> = (0..depth) + .map(|index| block(dim, kernel_size, vb.pp(3 + index))) + .collect::>>()?; + let fc = candle_nn::linear(dim, nclasses, vb.pp(25))?; + Ok(candle_nn::func(move |xs| { + let mut xs = xs.apply(&conv1)?.gelu_erf()?.apply_t(&bn1, false)?; + for block in blocks.iter() { + xs = xs.apply(block)? + } + // This performs the adaptive average pooling with a target size of (1, 1). + xs.mean(3)?.mean(2)?.apply(&fc) + })) +} + +pub fn c1536_20(nclasses: usize, vb: VarBuilder) -> Result> { + convmixer(nclasses, 1536, 20, 9, 7, vb) +} + +pub fn c1024_20(nclasses: usize, vb: VarBuilder) -> Result> { + convmixer(nclasses, 1024, 20, 9, 14, vb) +} diff --git a/patches/candle-transformers/src/models/convnext.rs b/patches/candle-transformers/src/models/convnext.rs new file mode 100644 index 0000000000..727e11381c --- /dev/null +++ b/patches/candle-transformers/src/models/convnext.rs @@ -0,0 +1,340 @@ +//! ConvNeXt implementation. +//! +//! This candle implementation uses a pre-trained ConvNeXt network for inference. The +//! classification head has been trained on the ImageNet dataset and returns the +//! probabilities for the top-5 classes. +//! +//! Original code: +//! - 💻 [ConvNeXt](https://github.com/facebookresearch/ConvNeXt/) +//! - 💻 [ConvNeXt-V2](https://github.com/facebookresearch/ConvNeXt-V2/) +//! - 💻 [timm](https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/convnext.py) +//! - 📝 [Paper](https://arxiv.org/abs/2201.03545) A ConvNet for the 2020s +//! - 📝 [Paper](https://arxiv.org/abs/2301.00808) ConvNeXt V2: Co-designing and Scaling ConvNets with Masked Autoencoders +//! + +use candle::shape::ShapeWithOneHole; +use candle::{Result, D}; +use candle_nn::{conv2d, layer_norm, linear, Conv2dConfig, Func, VarBuilder}; + +#[derive(Clone)] +pub struct Config { + blocks: [usize; 4], + channels: [usize; 4], + use_conv_mlp: bool, +} + +impl Config { + pub fn atto() -> Self { + Self { + blocks: [2, 2, 6, 2], + channels: [40, 80, 160, 320], + use_conv_mlp: true, + } + } + + pub fn femto() -> Self { + Self { + blocks: [2, 2, 6, 2], + channels: [48, 96, 192, 384], + use_conv_mlp: true, + } + } + + pub fn pico() -> Self { + Self { + blocks: [2, 2, 6, 2], + channels: [64, 128, 256, 512], + use_conv_mlp: true, + } + } + + pub fn nano() -> Self { + Self { + blocks: [2, 2, 8, 2], + channels: [80, 160, 320, 640], + use_conv_mlp: true, + } + } + + pub fn tiny() -> Self { + Self { + blocks: [3, 3, 9, 3], + channels: [96, 192, 384, 768], + use_conv_mlp: false, + } + } + + pub fn small() -> Self { + Self { + blocks: [3, 3, 27, 3], + channels: [96, 192, 384, 768], + use_conv_mlp: false, + } + } + + pub fn base() -> Self { + Self { + blocks: [3, 3, 27, 3], + channels: [128, 256, 512, 1024], + use_conv_mlp: false, + } + } + + pub fn large() -> Self { + Self { + blocks: [3, 3, 27, 3], + channels: [192, 384, 768, 1536], + use_conv_mlp: false, + } + } + + pub fn xlarge() -> Self { + Self { + blocks: [3, 3, 27, 3], + channels: [256, 512, 1024, 2048], + use_conv_mlp: false, + } + } + + pub fn huge() -> Self { + Self { + blocks: [3, 3, 27, 3], + channels: [352, 704, 1408, 2816], + use_conv_mlp: false, + } + } +} + +// Layer norm for data in channels-last format. +fn layer_norm_cl(dim: usize, vb: VarBuilder) -> Result> { + let norm = layer_norm(dim, 1e-6, vb)?; + + Ok(Func::new(move |xs| xs.apply(&norm))) +} + +// Layer norm for data in channels-first format. +fn layer_norm_cf(dim: usize, vb: VarBuilder) -> Result> { + let norm = layer_norm(dim, 1e-6, vb)?; + + Ok(Func::new(move |xs| { + let xs = xs + .permute((0, 2, 3, 1))? + .apply(&norm)? + .permute((0, 3, 1, 2))?; + Ok(xs) + })) +} + +// Global response normalization layer +// Based on https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/grn.py +fn convnext2_grn(dim: usize, channels_last: bool, vb: VarBuilder) -> Result> { + let (shape, spatial_dim, channel_dim) = if channels_last { + ((1, 1, 1, ()).into_shape(dim)?, [1, 2], 3) + } else { + ((1, (), 1, 1).into_shape(dim)?, [2, 3], 1) + }; + + let gamma = vb.get(dim, "weight")?.reshape(&shape)?; + let beta = vb.get(dim, "bias")?.reshape(&shape)?; + + Ok(Func::new(move |xs| { + let residual = xs; + let gx = xs + .sqr()? + .sum_keepdim(spatial_dim)? + .mean_keepdim(spatial_dim)? + .sqrt()?; + + let gxmean = gx.mean_keepdim(channel_dim)?; + let nx = gx.broadcast_div(&(gxmean + 1e-6)?)?; + let xs = xs + .broadcast_mul(&nx)? + .broadcast_mul(&gamma)? + .broadcast_add(&beta)?; + + xs + residual + })) +} + +// Initial downsampling via a patchify layer. +fn convnext_stem(out_channels: usize, vb: VarBuilder) -> Result> { + let conv2d_cfg = Conv2dConfig { + stride: 4, + ..Default::default() + }; + let patchify = conv2d(3, out_channels, 4, conv2d_cfg, vb.pp(0))?; + let norm = layer_norm_cf(out_channels, vb.pp(1))?; + + Ok(Func::new(move |xs| xs.apply(&patchify)?.apply(&norm))) +} + +// Downsampling applied after the stages. +fn convnext_downsample(dim: usize, vb: VarBuilder) -> Result> { + let conv2d_cfg = Conv2dConfig { + stride: 2, + ..Default::default() + }; + let norm = layer_norm_cf(dim / 2, vb.pp(0))?; + let conv = conv2d(dim / 2, dim, 2, conv2d_cfg, vb.pp(1))?; + + Ok(Func::new(move |xs| xs.apply(&norm)?.apply(&conv))) +} + +// MLP block from the original paper with optional GRN layer (v2 models). +fn convnext_mlp(dim: usize, vb: VarBuilder) -> Result> { + let fc1 = linear(dim, 4 * dim, vb.pp("fc1"))?; + let fc2 = linear(4 * dim, dim, vb.pp("fc2"))?; + let grn = convnext2_grn(4 * dim, true, vb.pp("grn")); + + Ok(Func::new(move |xs| { + let mut xs = xs.apply(&fc1)?.gelu_erf()?; + if let Ok(g) = &grn { + xs = xs.apply(g)?; + } + xs = xs.apply(&fc2)?; + Ok(xs) + })) +} + +// MLP block using pointwise convolutions, with optional GRN layer (v2 models). +fn convnext_conv_mlp(dim: usize, vb: VarBuilder) -> Result> { + let conv2d_cfg = Conv2dConfig { + ..Default::default() + }; + let fc1 = conv2d(dim, 4 * dim, 1, conv2d_cfg, vb.pp("fc1"))?; + let fc2 = conv2d(4 * dim, dim, 1, conv2d_cfg, vb.pp("fc2"))?; + + let grn = convnext2_grn(4 * dim, false, vb.pp("grn")); + Ok(Func::new(move |xs| { + let mut xs = xs.apply(&fc1)?.gelu_erf()?; + if let Ok(g) = &grn { + xs = xs.apply(g)?; + } + xs = xs.apply(&fc2)?; + Ok(xs) + })) +} + +// A block consisting of a depthwise convolution, a MLP and layer scaling (v1 models only). +fn convnext_block(dim: usize, use_conv_mlp: bool, vb: VarBuilder) -> Result> { + let conv2d_cfg = Conv2dConfig { + groups: dim, + padding: 3, + ..Default::default() + }; + + let conv_dw = conv2d(dim, dim, 7, conv2d_cfg, vb.pp("conv_dw"))?; + let gamma = vb.get(dim, "gamma"); + + let (mlp, norm) = if use_conv_mlp { + ( + convnext_conv_mlp(dim, vb.pp("mlp"))?, + layer_norm_cf(dim, vb.pp("norm"))?, + ) + } else { + ( + convnext_mlp(dim, vb.pp("mlp"))?, + layer_norm_cl(dim, vb.pp("norm"))?, + ) + }; + + Ok(Func::new(move |xs| { + let residual = xs; + let mut xs = xs.apply(&conv_dw)?; + + xs = if use_conv_mlp { + xs.apply(&norm)?.apply(&mlp)? + } else { + xs.permute((0, 2, 3, 1))? + .apply(&norm)? + .apply(&mlp)? + .permute((0, 3, 1, 2))? + }; + + if let Ok(g) = &gamma { + xs = xs.broadcast_mul(&g.reshape((1, (), 1, 1))?)?; + }; + + xs + residual + })) +} + +// Each stage contains blocks and a downsampling layer for the previous stage. +fn convnext_stage(cfg: &Config, stage_idx: usize, vb: VarBuilder) -> Result> { + let nblocks = cfg.blocks[stage_idx]; + let mut blocks = Vec::with_capacity(nblocks); + + let dim = cfg.channels[stage_idx]; + + if stage_idx > 0 { + blocks.push(convnext_downsample(dim, vb.pp("downsample"))?); + } + + for block_idx in 0..nblocks { + blocks.push(convnext_block( + dim, + cfg.use_conv_mlp, + vb.pp(format!("blocks.{block_idx}")), + )?); + } + + Ok(Func::new(move |xs| { + let mut xs = xs.clone(); + for block in blocks.iter() { + xs = xs.apply(block)? + } + Ok(xs) + })) +} + +// Classification head. +fn convnext_head(outputs: usize, nclasses: usize, vb: VarBuilder) -> Result> { + let norm = layer_norm_cl(outputs, vb.pp("norm"))?; + let linear = linear(outputs, nclasses, vb.pp("fc"))?; + Ok(Func::new(move |xs| xs.apply(&norm)?.apply(&linear))) +} + +// Build a convnext model for a given configuration. +fn convnext_model( + config: &Config, + nclasses: Option, + vb: VarBuilder, +) -> Result> { + let head = match nclasses { + None => None, + Some(nclasses) => { + let head = convnext_head(config.channels[3], nclasses, vb.pp("head"))?; + Some(head) + } + }; + + let stem = convnext_stem(config.channels[0], vb.pp("stem"))?; + let vb = vb.pp("stages"); + let stage1 = convnext_stage(config, 0, vb.pp(0))?; + let stage2 = convnext_stage(config, 1, vb.pp(1))?; + let stage3 = convnext_stage(config, 2, vb.pp(2))?; + let stage4 = convnext_stage(config, 3, vb.pp(3))?; + + Ok(Func::new(move |xs| { + let xs = xs + .apply(&stem)? + .apply(&stage1)? + .apply(&stage2)? + .apply(&stage3)? + .apply(&stage4)? + .mean(D::Minus2)? + .mean(D::Minus1)?; + match &head { + None => Ok(xs), + Some(head) => xs.apply(head), + } + })) +} + +pub fn convnext(cfg: &Config, nclasses: usize, vb: VarBuilder) -> Result> { + convnext_model(cfg, Some(nclasses), vb) +} + +pub fn convnext_no_final_layer(cfg: &Config, vb: VarBuilder) -> Result> { + convnext_model(cfg, None, vb) +} diff --git a/patches/candle-transformers/src/models/csm.rs b/patches/candle-transformers/src/models/csm.rs new file mode 100644 index 0000000000..28267ecc7a --- /dev/null +++ b/patches/candle-transformers/src/models/csm.rs @@ -0,0 +1,533 @@ +//! Implementation of the Conversational Speech Model (CSM) from Sesame +//! +//! See: [CSM](Conversational Speech Model) +//! +/// CSM (Conversational Speech Model) is a speech generation model from Sesame that generates RVQ +/// audio codes from text and audio inputs. The model architecture employs a Llama backbone and a +/// smaller audio decoder that produces Mimi audio codes. +/// +use crate::generation::LogitsProcessor; +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{embedding, linear_b, Embedding, Linear, RmsNorm, VarBuilder}; +use std::sync::Arc; + +#[derive(serde::Deserialize, Debug, Clone, Copy, PartialEq, Eq)] +pub enum Flavor { + #[serde(rename = "llama-1B")] + Llama1B, + #[serde(rename = "llama-100M")] + Llama100M, +} + +#[derive(serde::Deserialize, Debug, Clone)] +pub struct Config { + pub audio_num_codebooks: usize, + pub audio_vocab_size: usize, + pub backbone_flavor: Flavor, + pub decoder_flavor: Flavor, + pub text_vocab_size: usize, +} + +#[allow(unused)] +#[derive(Debug, Clone)] +pub struct LlamaConfig { + vocab_size: usize, + num_layers: usize, + num_heads: usize, + num_kv_heads: usize, + embed_dim: usize, + max_seq_len: usize, + intermediate_dim: usize, + norm_eps: f64, + rope_base: f32, + scale_factor: usize, +} + +impl LlamaConfig { + pub fn from_flavor(flavor: Flavor) -> Self { + match flavor { + Flavor::Llama1B => Self { + vocab_size: 128256, + num_layers: 16, + num_heads: 32, + num_kv_heads: 8, + embed_dim: 2048, + max_seq_len: 2048, + intermediate_dim: 8192, + norm_eps: 1e-5, + rope_base: 500_000., + scale_factor: 32, + }, + Flavor::Llama100M => Self { + vocab_size: 128256, + num_layers: 4, + num_heads: 8, + num_kv_heads: 2, + embed_dim: 1024, + max_seq_len: 2048, + intermediate_dim: 8192, + norm_eps: 1e-5, + rope_base: 500_000., + scale_factor: 32, + }, + } + } +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +fn calculate_default_inv_freq(cfg: &LlamaConfig) -> Vec { + let head_dim = cfg.embed_dim / cfg.num_heads; + (0..head_dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_base.powf(i as f32 / head_dim as f32)) + .collect() +} + +impl RotaryEmbedding { + fn new(dtype: DType, cfg: &LlamaConfig, dev: &Device) -> Result { + let low_freq_factor = 1.0; + let high_freq_factor = 4.0; + let original_max_position_embeddings = 8192; + let scale_factor = cfg.scale_factor as f32; + let theta = { + let low_freq_wavelen = original_max_position_embeddings as f32 / low_freq_factor; + let high_freq_wavelen = original_max_position_embeddings as f32 / high_freq_factor; + + calculate_default_inv_freq(cfg) + .into_iter() + .map(|freq| { + let wavelen = 2. * std::f32::consts::PI / freq; + if wavelen < high_freq_wavelen { + freq + } else if wavelen > low_freq_wavelen { + freq / scale_factor + } else { + let smooth = (original_max_position_embeddings as f32 / wavelen + - low_freq_factor) + / (high_freq_factor - low_freq_factor); + (1. - smooth) * freq / scale_factor + smooth * freq + } + }) + .collect::>() + }; + + let theta = Tensor::new(theta, dev)?; + let idx_theta = Tensor::arange(0, cfg.max_seq_len as u32, dev)? + .to_dtype(DType::F32)? + .reshape((cfg.max_seq_len, 1))? + .matmul(&theta.reshape((1, theta.elem_count()))?)?; + // This is different from the paper, see: + // https://github.com/huggingface/transformers/blob/6112b1c6442aaf7affd2b0676a1cd4eee30c45cf/src/transformers/models/llama/modeling_llama.py#L112 + let cos = idx_theta.cos()?.to_dtype(dtype)?; + let sin = idx_theta.sin()?.to_dtype(dtype)?; + Ok(Self { cos, sin }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope_i(q, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope_i(k, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} +fn rms_norm(hidden_size: usize, eps: f64, vb: VarBuilder) -> Result { + let weight = vb.get((hidden_size,), "scale")?; + Ok(RmsNorm::new(weight, eps)) +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, + num_heads: usize, + head_dim: usize, + num_kv_heads: usize, + num_kv_groups: usize, +} + +impl Attention { + fn new(cfg: &LlamaConfig, rotary_emb: Arc, vb: VarBuilder) -> Result { + let head_dim = cfg.embed_dim / cfg.num_heads; + let kv_dim = cfg.num_kv_heads * head_dim; + + let q_proj = linear_b(cfg.embed_dim, cfg.embed_dim, false, vb.pp("q_proj"))?; + let k_proj = linear_b(cfg.embed_dim, kv_dim, false, vb.pp("k_proj"))?; + let v_proj = linear_b(cfg.embed_dim, kv_dim, false, vb.pp("v_proj"))?; + let o_proj = linear_b(cfg.embed_dim, cfg.embed_dim, false, vb.pp("output_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + rotary_emb, + kv_cache: None, + num_heads: cfg.num_heads, + num_kv_heads: cfg.num_kv_heads, + num_kv_groups: cfg.num_heads / cfg.num_kv_heads, + head_dim, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + + let (query_states, key_states) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states, seqlen_offset)?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?; + let value_states = crate::utils::repeat_kv(value_states, self.num_kv_groups)?; + + let attn_output = { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + }; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.num_heads * self.head_dim))? + .apply(&self.o_proj) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +struct Mlp { + w1: Linear, + w2: Linear, + w3: Linear, +} + +impl Mlp { + fn new(cfg: &LlamaConfig, vb: VarBuilder) -> Result { + let w1 = linear_b(cfg.embed_dim, cfg.intermediate_dim, false, vb.pp("w1"))?; + let w2 = linear_b(cfg.intermediate_dim, cfg.embed_dim, false, vb.pp("w2"))?; + let w3 = linear_b(cfg.embed_dim, cfg.intermediate_dim, false, vb.pp("w3"))?; + Ok(Self { w1, w2, w3 }) + } +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let lhs = xs.apply(&self.w1)?.silu()?; + let rhs = xs.apply(&self.w3)?; + (lhs * rhs)?.apply(&self.w2) + } +} + +#[derive(Debug, Clone)] +struct Layer { + mlp_norm: RmsNorm, + sa_norm: RmsNorm, + attn: Attention, + mlp: Mlp, +} + +impl Layer { + fn new(cfg: &LlamaConfig, rotary_emb: Arc, vb: VarBuilder) -> Result { + let mlp_norm = rms_norm(cfg.embed_dim, cfg.norm_eps, vb.pp("mlp_norm"))?; + let sa_norm = rms_norm(cfg.embed_dim, cfg.norm_eps, vb.pp("sa_norm"))?; + let attn = Attention::new(cfg, rotary_emb, vb.pp("attn"))?; + let mlp = Mlp::new(cfg, vb.pp("mlp"))?; + Ok(Self { + mlp_norm, + sa_norm, + attn, + mlp, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.sa_norm.forward(xs)?; + let xs = self.attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.mlp_norm)?.apply(&self.mlp)?; + residual + xs + } + + fn clear_kv_cache(&mut self) { + self.attn.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct LlamaModel { + layers: Vec, + norm: RmsNorm, + device: Device, + dtype: DType, +} + +impl LlamaModel { + pub fn new(cfg: &LlamaConfig, vb: VarBuilder) -> Result { + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb.device())?); + let mut layers = Vec::with_capacity(cfg.num_layers); + let vb_l = vb.pp("layers"); + for layer_idx in 0..cfg.num_layers { + let layer = Layer::new(cfg, rotary_emb.clone(), vb_l.pp(layer_idx))?; + layers.push(layer); + } + let norm = rms_norm(cfg.embed_dim, cfg.norm_eps, vb.pp("norm"))?; + Ok(Self { + layers, + norm, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } + + fn prepare_decoder_attention_mask( + &self, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { f32::NEG_INFINITY } else { 0. })) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((1, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + pub fn forward(&mut self, xs: &Tensor, seqlen_offset: usize) -> Result { + let (_b_size, seq_len, _embed_dim) = xs.dims3()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(seq_len, seqlen_offset)?; + Some(mask) + }; + let mut xs = xs.clone(); + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)?; + } + let ys = xs.narrow(1, seq_len - 1, 1)?.apply(&self.norm)?; + Ok(ys) + } +} + +#[derive(Debug, Clone)] +pub struct Model { + backbone: LlamaModel, + decoder: LlamaModel, + codebook0_head: Linear, + audio_embeddings: Embedding, + text_embeddings: Embedding, + projection: Linear, + audio_head: Tensor, + config: Config, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let backbone_cfg = LlamaConfig::from_flavor(cfg.backbone_flavor); + let backbone = LlamaModel::new(&backbone_cfg, vb.pp("backbone"))?; + let decoder_cfg = LlamaConfig::from_flavor(cfg.decoder_flavor); + let decoder = LlamaModel::new(&decoder_cfg, vb.pp("decoder"))?; + let backbone_dim = backbone_cfg.embed_dim; + let decoder_dim = decoder_cfg.embed_dim; + let audio_embeddings = embedding( + cfg.audio_vocab_size * cfg.audio_num_codebooks, + backbone_dim, + vb.pp("audio_embeddings"), + )?; + let text_embeddings = + embedding(cfg.text_vocab_size, backbone_dim, vb.pp("text_embeddings"))?; + let projection = linear_b(backbone_dim, decoder_dim, false, vb.pp("projection"))?; + let codebook0_head = linear_b( + backbone_dim, + cfg.audio_vocab_size, + false, + vb.pp("codebook0_head"), + )?; + let audio_head = vb.get( + ( + cfg.audio_num_codebooks - 1, + decoder_dim, + cfg.audio_vocab_size, + ), + "audio_head", + )?; + Ok(Self { + backbone, + decoder, + codebook0_head, + audio_embeddings, + text_embeddings, + projection, + audio_head, + config: cfg.clone(), + }) + } + + pub fn clear_kv_cache(&mut self) { + self.backbone.clear_kv_cache(); + self.decoder.clear_kv_cache(); + } + + pub fn generate_frame( + &mut self, + tokens: &Tensor, + tokens_mask: &Tensor, + input_pos: usize, + lp: &mut LogitsProcessor, + ) -> Result> { + let (b_sz, seq_len, _cb_plus_one) = tokens.dims3()?; + let audio_tokens = tokens.narrow(2, 0, self.config.audio_num_codebooks)?; + let text_tokens = tokens.narrow(2, self.config.audio_num_codebooks, 1)?; + let text_embeds = self.text_embeddings.forward(&text_tokens)?; + let arange = (Tensor::arange( + 0u32, + self.config.audio_num_codebooks as u32, + &self.decoder.device, + )? * self.config.audio_vocab_size as f64)?; + let audio_tokens = audio_tokens.broadcast_add(&arange.reshape((1, 1, ()))?)?; + let audio_embeds = self.audio_embeddings.forward(&audio_tokens)?.reshape(( + b_sz, + seq_len, + self.config.audio_num_codebooks, + (), + ))?; + let embeds = Tensor::cat(&[&audio_embeds, &text_embeds], D::Minus2)?; + let embeds = embeds.broadcast_mul( + &tokens_mask + .to_dtype(self.backbone.dtype)? + .unsqueeze(D::Minus1)?, + )?; + let embeds = embeds.sum(2)?; + let h = self.backbone.forward(&embeds, input_pos)?; + let c0_logits = h.apply(&self.codebook0_head)?; + let c0_sample = lp.sample(&c0_logits.i((0, 0))?)?; + let mut all_samples = vec![c0_sample]; + let c0_sample = Tensor::from_slice(&[c0_sample], (1, 1), &self.decoder.device)?; + let c0_embed = self.audio_embeddings.forward(&c0_sample)?; + let mut curr_h = Tensor::cat(&[h, c0_embed], 1)?; + + self.decoder.clear_kv_cache(); + let mut decoder_pos = 0; + for i in 1..self.config.audio_num_codebooks { + let proj_h = curr_h.apply(&self.projection)?; + let decoder_h = self.decoder.forward(&proj_h, decoder_pos)?; + decoder_pos += curr_h.dim(1)?; + let ci_logits = decoder_h.broadcast_matmul(&self.audio_head.get(i - 1)?)?; + let ci_sample = lp.sample(&ci_logits.i((0, 0))?)?; + all_samples.push(ci_sample); + let ci_sample = Tensor::from_slice( + &[ci_sample + (i * self.config.audio_vocab_size) as u32], + (1, 1), + &self.decoder.device, + )?; + let ci_embed = self.audio_embeddings.forward(&ci_sample)?; + curr_h = ci_embed + } + Ok(all_samples) + } + + pub fn audio_tokens_and_mask(&self, mut frame: Vec) -> Result<(Tensor, Tensor)> { + let cb = self.config.audio_num_codebooks; + let device = &self.backbone.device; + let mut mask = vec![1u8; cb]; + mask.push(0); + let mask = Tensor::from_vec(mask, (1, 1, cb + 1), device)?; + + frame.push(0); + let tokens = Tensor::from_vec(frame, (1, 1, cb + 1), device)?; + Ok((tokens, mask)) + } + + pub fn text_tokens_and_mask(&self, ids: &[u32]) -> Result<(Tensor, Tensor)> { + let cb = self.config.audio_num_codebooks; + let device = &self.backbone.device; + let mut tokens = vec![]; + let mut mask = vec![]; + for &v in ids.iter() { + let mut token = vec![0; cb]; + token.push(v); + let token = Tensor::from_vec(token, (1, 1, cb + 1), device)?; + tokens.push(token); + let mut m = vec![0u8; cb]; + m.push(1); + let m = Tensor::from_vec(m, (1, 1, cb + 1), device)?; + mask.push(m); + } + let tokens = Tensor::cat(&tokens, 1)?; + let mask = Tensor::cat(&mask, 1)?; + Ok((tokens, mask)) + } +} diff --git a/patches/candle-transformers/src/models/dac.rs b/patches/candle-transformers/src/models/dac.rs new file mode 100644 index 0000000000..21cba02e87 --- /dev/null +++ b/patches/candle-transformers/src/models/dac.rs @@ -0,0 +1,381 @@ +//! Implementation of the Descript Audio Codec (DAC) model +//! +//! See: [Descript Audio Codec](https://github.com/descriptinc/descript-audio-codec) +//! +/// An efficient neural codec for compressing/decompressing audio +/// +use crate::models::encodec; +use candle::{IndexOp, Result, Tensor, D}; +use candle_nn::{Conv1d, Conv1dConfig, ConvTranspose1d, ConvTranspose1dConfig, VarBuilder}; + +#[derive(serde::Deserialize, Debug, Clone)] +pub struct Config { + pub num_codebooks: usize, + pub model_bitrate: u32, + pub codebook_size: usize, + pub latent_dim: usize, + pub frame_rate: u32, + pub sampling_rate: u32, +} + +#[derive(Debug, Clone)] +pub struct Snake1d { + alpha: Tensor, +} + +impl Snake1d { + pub fn new(channels: usize, vb: VarBuilder) -> Result { + let alpha = vb.get((1, channels, 1), "alpha")?; + Ok(Self { alpha }) + } +} + +impl candle::Module for Snake1d { + fn forward(&self, xs: &Tensor) -> Result { + let xs_shape = xs.shape(); + let xs = xs.flatten_from(2)?; + let sin = self.alpha.broadcast_mul(&xs)?.sin()?; + let sin = (&sin * &sin)?; + (xs + (&self.alpha + 1e-9)?.recip()?.broadcast_mul(&sin)?)?.reshape(xs_shape) + } +} + +#[derive(Debug, Clone)] +pub struct ResidualUnit { + snake1: Snake1d, + conv1: Conv1d, + snake2: Snake1d, + conv2: Conv1d, +} + +impl ResidualUnit { + pub fn new(dim: usize, dilation: usize, vb: VarBuilder) -> Result { + let pad = ((7 - 1) * dilation) / 2; + let vb = vb.pp("block"); + let snake1 = Snake1d::new(dim, vb.pp(0))?; + let cfg1 = Conv1dConfig { + dilation, + padding: pad, + ..Default::default() + }; + let conv1 = encodec::conv1d_weight_norm(dim, dim, 7, cfg1, vb.pp(1))?; + let snake2 = Snake1d::new(dim, vb.pp(2))?; + let conv2 = encodec::conv1d_weight_norm(dim, dim, 1, Default::default(), vb.pp(3))?; + Ok(Self { + snake1, + conv1, + snake2, + conv2, + }) + } +} + +impl candle::Module for ResidualUnit { + fn forward(&self, xs: &Tensor) -> Result { + let ys = xs + .apply(&self.snake1)? + .apply(&self.conv1)? + .apply(&self.snake2)? + .apply(&self.conv2)?; + let pad = (xs.dim(D::Minus1)? - ys.dim(D::Minus1)?) / 2; + if pad > 0 { + &ys + xs.narrow(D::Minus1, pad, ys.dim(D::Minus1)?) + } else { + ys + xs + } + } +} + +#[derive(Debug, Clone)] +pub struct EncoderBlock { + res1: ResidualUnit, + res2: ResidualUnit, + res3: ResidualUnit, + snake1: Snake1d, + conv1: Conv1d, +} + +impl EncoderBlock { + pub fn new(dim: usize, stride: usize, vb: VarBuilder) -> Result { + let vb = vb.pp("block"); + let res1 = ResidualUnit::new(dim / 2, 1, vb.pp(0))?; + let res2 = ResidualUnit::new(dim / 2, 3, vb.pp(1))?; + let res3 = ResidualUnit::new(dim / 2, 9, vb.pp(2))?; + let snake1 = Snake1d::new(dim / 2, vb.pp(3))?; + let cfg1 = Conv1dConfig { + stride, + padding: stride.div_ceil(2), + ..Default::default() + }; + let conv1 = encodec::conv1d_weight_norm(dim / 2, dim, 2 * stride, cfg1, vb.pp(4))?; + Ok(Self { + res1, + res2, + res3, + snake1, + conv1, + }) + } +} + +impl candle::Module for EncoderBlock { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.res1)? + .apply(&self.res2)? + .apply(&self.res3)? + .apply(&self.snake1)? + .apply(&self.conv1) + } +} + +#[derive(Debug, Clone)] +pub struct Encoder { + conv1: Conv1d, + blocks: Vec, + snake1: Snake1d, + conv2: Conv1d, +} + +impl candle::Module for Encoder { + fn forward(&self, xs: &Tensor) -> Result { + let mut xs = xs.apply(&self.conv1)?; + for block in self.blocks.iter() { + xs = xs.apply(block)? + } + xs.apply(&self.snake1)?.apply(&self.conv2) + } +} + +impl Encoder { + pub fn new( + mut d_model: usize, + strides: &[usize], + d_latent: usize, + vb: VarBuilder, + ) -> Result { + let vb = vb.pp("block"); + let cfg1 = Conv1dConfig { + padding: 3, + ..Default::default() + }; + let conv1 = encodec::conv1d_weight_norm(1, d_model, 7, cfg1, vb.pp(0))?; + let mut blocks = Vec::with_capacity(strides.len()); + for (block_idx, stride) in strides.iter().enumerate() { + d_model *= 2; + let block = EncoderBlock::new(d_model, *stride, vb.pp(block_idx + 1))?; + blocks.push(block) + } + let snake1 = Snake1d::new(d_model, vb.pp(strides.len() + 1))?; + let cfg2 = Conv1dConfig { + padding: 1, + ..Default::default() + }; + let conv2 = + encodec::conv1d_weight_norm(d_model, d_latent, 3, cfg2, vb.pp(strides.len() + 2))?; + Ok(Self { + conv1, + blocks, + snake1, + conv2, + }) + } +} + +#[derive(Debug, Clone)] +pub struct DecoderBlock { + snake1: Snake1d, + conv_tr1: ConvTranspose1d, + res1: ResidualUnit, + res2: ResidualUnit, + res3: ResidualUnit, +} + +impl DecoderBlock { + pub fn new(in_dim: usize, out_dim: usize, stride: usize, vb: VarBuilder) -> Result { + let vb = vb.pp("block"); + let snake1 = Snake1d::new(in_dim, vb.pp(0))?; + let cfg = ConvTranspose1dConfig { + stride, + padding: stride.div_ceil(2), + ..Default::default() + }; + let conv_tr1 = encodec::conv_transpose1d_weight_norm( + in_dim, + out_dim, + 2 * stride, + true, + cfg, + vb.pp(1), + )?; + let res1 = ResidualUnit::new(out_dim, 1, vb.pp(2))?; + let res2 = ResidualUnit::new(out_dim, 3, vb.pp(3))?; + let res3 = ResidualUnit::new(out_dim, 9, vb.pp(4))?; + Ok(Self { + snake1, + conv_tr1, + res1, + res2, + res3, + }) + } +} + +impl candle_nn::Module for DecoderBlock { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.snake1)? + .apply(&self.conv_tr1)? + .apply(&self.res1)? + .apply(&self.res2)? + .apply(&self.res3) + } +} + +#[derive(Debug, Clone)] +pub struct Decoder { + conv1: Conv1d, + blocks: Vec, + snake1: Snake1d, + conv2: Conv1d, +} + +impl Decoder { + pub fn new( + in_c: usize, + mut channels: usize, + rates: &[usize], + d_out: usize, + vb: VarBuilder, + ) -> Result { + let vb = vb.pp("model"); + let cfg1 = Conv1dConfig { + padding: 3, + ..Default::default() + }; + let conv1 = encodec::conv1d_weight_norm(in_c, channels, 7, cfg1, vb.pp(0))?; + let mut blocks = Vec::with_capacity(rates.len()); + for (idx, stride) in rates.iter().enumerate() { + let block = DecoderBlock::new(channels, channels / 2, *stride, vb.pp(idx + 1))?; + channels /= 2; + blocks.push(block) + } + let snake1 = Snake1d::new(channels, vb.pp(rates.len() + 1))?; + let conv2 = encodec::conv1d_weight_norm(channels, d_out, 7, cfg1, vb.pp(rates.len() + 2))?; + Ok(Self { + conv1, + blocks, + snake1, + conv2, + }) + } +} + +impl candle::Module for Decoder { + fn forward(&self, xs: &Tensor) -> Result { + let mut xs = xs.apply(&self.conv1)?; + for block in self.blocks.iter() { + xs = xs.apply(block)? + } + xs.apply(&self.snake1)?.apply(&self.conv2) + } +} + +#[allow(unused)] +#[derive(Clone, Debug)] +pub struct VectorQuantizer { + in_proj: Conv1d, + out_proj: Conv1d, + codebook: candle_nn::Embedding, +} + +impl VectorQuantizer { + pub fn new(in_dim: usize, cb_size: usize, cb_dim: usize, vb: VarBuilder) -> Result { + let in_proj = + encodec::conv1d_weight_norm(in_dim, cb_dim, 1, Default::default(), vb.pp("in_proj"))?; + let out_proj = + encodec::conv1d_weight_norm(cb_dim, in_dim, 1, Default::default(), vb.pp("out_proj"))?; + let codebook = candle_nn::embedding(cb_size, cb_dim, vb.pp("codebook"))?; + Ok(Self { + in_proj, + out_proj, + codebook, + }) + } + + pub fn embed_code(&self, embed_id: &Tensor) -> Result { + embed_id.apply(&self.codebook) + } + + pub fn decode_code(&self, embed_id: &Tensor) -> Result { + self.embed_code(embed_id)?.transpose(1, 2) + } +} + +#[derive(Clone, Debug)] +pub struct ResidualVectorQuantizer { + quantizers: Vec, +} + +impl ResidualVectorQuantizer { + pub fn new( + input_dim: usize, + n_codebooks: usize, + cb_size: usize, + cb_dim: usize, + vb: VarBuilder, + ) -> Result { + let vb = &vb.pp("quantizers"); + let quantizers = (0..n_codebooks) + .map(|i| VectorQuantizer::new(input_dim, cb_size, cb_dim, vb.pp(i))) + .collect::>>()?; + Ok(Self { quantizers }) + } + + #[allow(clippy::wrong_self_convention)] + pub fn from_codes(&self, codes: &Tensor) -> Result { + let mut sum = None; + for (idx, quantizer) in self.quantizers.iter().enumerate() { + let z_p_i = quantizer.decode_code(&codes.i((.., idx))?)?; + let z_q_i = z_p_i.apply(&quantizer.out_proj)?; + let s = match sum { + None => z_q_i, + Some(s) => (s + z_q_i)?, + }; + sum = Some(s) + } + match sum { + Some(s) => Ok(s), + None => candle::bail!("empty codebooks"), + } + } +} + +#[derive(Debug, Clone)] +pub struct Model { + pub encoder: Encoder, + pub quantizer: ResidualVectorQuantizer, + pub decoder: Decoder, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let encoder = Encoder::new(64, &[2, 4, 8, 8], cfg.latent_dim, vb.pp("encoder"))?; + let quantizer = ResidualVectorQuantizer::new( + cfg.latent_dim, + cfg.num_codebooks, + cfg.codebook_size, + 8, + vb.pp("quantizer"), + )?; + let decoder = Decoder::new(cfg.latent_dim, 1536, &[8, 8, 4, 2], 1, vb.pp("decoder"))?; + Ok(Self { + encoder, + decoder, + quantizer, + }) + } + + pub fn decode_codes(&self, audio_codes: &Tensor) -> Result { + let audio_values = self.quantizer.from_codes(audio_codes)?; + audio_values.apply(&self.decoder) + } +} diff --git a/patches/candle-transformers/src/models/debertav2.rs b/patches/candle-transformers/src/models/debertav2.rs new file mode 100644 index 0000000000..4f19d3b419 --- /dev/null +++ b/patches/candle-transformers/src/models/debertav2.rs @@ -0,0 +1,1444 @@ +use std::collections::HashMap; + +use candle::{bail, Context, DType, Device, Module, Result, Tensor, D}; +use candle_nn::{ + conv1d, embedding, layer_norm, Conv1d, Conv1dConfig, Embedding, LayerNorm, VarBuilder, +}; +use serde::{Deserialize, Deserializer}; + +pub const DTYPE: DType = DType::F32; + +// NOTE: HiddenAct and HiddenActLayer are both direct copies from bert.rs. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum HiddenAct { + Gelu, + GeluApproximate, + Relu, +} + +pub struct HiddenActLayer { + act: HiddenAct, + span: tracing::Span, +} + +impl HiddenActLayer { + fn new(act: HiddenAct) -> Self { + let span = tracing::span!(tracing::Level::TRACE, "hidden-act"); + Self { act, span } + } + + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + match self.act { + // https://github.com/huggingface/transformers/blob/cd4584e3c809bb9e1392ccd3fe38b40daba5519a/src/transformers/activations.py#L213 + HiddenAct::Gelu => xs.gelu_erf(), + HiddenAct::GeluApproximate => xs.gelu(), + HiddenAct::Relu => xs.relu(), + } + } +} + +pub type Id2Label = HashMap; +pub type Label2Id = HashMap; + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub intermediate_size: usize, + pub hidden_act: HiddenAct, + pub hidden_dropout_prob: f64, + pub attention_probs_dropout_prob: f64, + pub max_position_embeddings: usize, + pub type_vocab_size: usize, + pub initializer_range: f64, + pub layer_norm_eps: f64, + pub relative_attention: bool, + pub max_relative_positions: isize, + pub pad_token_id: Option, + pub position_biased_input: bool, + #[serde(deserialize_with = "deserialize_pos_att_type")] + pub pos_att_type: Vec, + pub position_buckets: Option, + pub share_att_key: Option, + pub attention_head_size: Option, + pub embedding_size: Option, + pub norm_rel_ebd: Option, + pub conv_kernel_size: Option, + pub conv_groups: Option, + pub conv_act: Option, + pub id2label: Option, + pub label2id: Option, + pub pooler_dropout: Option, + pub pooler_hidden_act: Option, + pub pooler_hidden_size: Option, + pub cls_dropout: Option, +} + +fn deserialize_pos_att_type<'de, D>(deserializer: D) -> std::result::Result, D::Error> +where + D: Deserializer<'de>, +{ + #[derive(Deserialize, Debug)] + #[serde(untagged)] + enum StringOrVec { + String(String), + Vec(Vec), + } + + match StringOrVec::deserialize(deserializer)? { + StringOrVec::String(s) => Ok(s.split('|').map(String::from).collect()), + StringOrVec::Vec(v) => Ok(v), + } +} + +// NOTE: Dropout is probably not needed for now since this will primarily be used +// in inferencing. However, for training/fine-tuning it will be necessary. +pub struct StableDropout { + _drop_prob: f64, + _count: usize, +} + +impl StableDropout { + pub fn new(drop_prob: f64) -> Self { + Self { + _drop_prob: drop_prob, + _count: 0, + } + } + + pub fn forward(&self, x: &Tensor) -> Result { + Ok(x.clone()) + } +} + +// https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/deberta_v2/modeling_deberta_v2.py#L823 +pub struct DebertaV2Embeddings { + device: Device, + word_embeddings: Embedding, + position_embeddings: Option, + token_type_embeddings: Option, + layer_norm: LayerNorm, + dropout: StableDropout, + position_ids: Tensor, + config: Config, + embedding_size: usize, + embed_proj: Option, +} + +impl DebertaV2Embeddings { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let device = vb.device().clone(); + let config = config.clone(); + + let embedding_size = config.embedding_size.unwrap_or(config.hidden_size); + + let word_embeddings = + embedding(config.vocab_size, embedding_size, vb.pp("word_embeddings"))?; + + let position_embeddings = if config.position_biased_input { + Some(embedding( + config.max_position_embeddings, + embedding_size, + vb.pp("position_embeddings"), + )?) + } else { + None + }; + + let token_type_embeddings: Option = if config.type_vocab_size > 0 { + Some(candle_nn::embedding( + config.type_vocab_size, + config.hidden_size, + vb.pp("token_type_embeddings"), + )?) + } else { + None + }; + + let embed_proj: Option = if embedding_size != config.hidden_size { + Some(candle_nn::linear_no_bias( + embedding_size, + config.hidden_size, + vb.pp("embed_proj"), + )?) + } else { + None + }; + + let layer_norm = layer_norm( + config.hidden_size, + config.layer_norm_eps, + vb.pp("LayerNorm"), + )?; + + let dropout = StableDropout::new(config.hidden_dropout_prob); + + let position_ids = + Tensor::arange(0, config.max_position_embeddings as u32, &device)?.unsqueeze(0)?; + + Ok(Self { + word_embeddings, + position_embeddings, + token_type_embeddings, + layer_norm, + dropout, + position_ids, + device, + config, + embedding_size, + embed_proj, + }) + } + + pub fn forward( + &self, + input_ids: Option<&Tensor>, + token_type_ids: Option<&Tensor>, + position_ids: Option<&Tensor>, + mask: Option<&Tensor>, + inputs_embeds: Option<&Tensor>, + ) -> Result { + let (input_shape, input_embeds) = match (input_ids, inputs_embeds) { + (Some(ids), None) => { + let embs = self.word_embeddings.forward(ids)?; + (ids.dims(), embs) + } + (None, Some(e)) => (e.dims(), e.clone()), + (None, None) => { + bail!("Must specify either input_ids or inputs_embeds") + } + (Some(_), Some(_)) => { + bail!("Can't specify both input_ids and inputs_embeds") + } + }; + + let seq_length = match input_shape.last() { + Some(v) => *v, + None => bail!("DebertaV2Embeddings invalid input shape"), + }; + + let position_ids = match position_ids { + Some(v) => v.clone(), + None => self.position_ids.narrow(1, 0, seq_length)?, + }; + + let token_type_ids = match token_type_ids { + Some(ids) => ids.clone(), + None => Tensor::zeros(input_shape, DType::U32, &self.device)?, + }; + + let position_embeddings = match &self.position_embeddings { + Some(emb) => emb.forward(&position_ids)?, + None => Tensor::zeros_like(&input_embeds)?, + }; + + let mut embeddings = input_embeds; + + if self.config.position_biased_input { + embeddings = embeddings.add(&position_embeddings)?; + } + + if self.config.type_vocab_size > 0 { + embeddings = self.token_type_embeddings.as_ref().map_or_else( + || bail!("token_type_embeddings must be set when type_vocab_size > 0"), + |token_type_embeddings| { + embeddings.add(&token_type_embeddings.forward(&token_type_ids)?) + }, + )?; + } + + if self.embedding_size != self.config.hidden_size { + embeddings = if let Some(embed_proj) = &self.embed_proj { + embed_proj.forward(&embeddings)? + } else { + bail!("embed_proj must exist if embedding_size != config.hidden_size"); + } + } + + embeddings = self.layer_norm.forward(&embeddings)?; + + if let Some(mask) = mask { + let mut mask = mask.clone(); + if mask.dims() != embeddings.dims() { + if mask.dims().len() == 4 { + mask = mask.squeeze(1)?.squeeze(1)?; + } + mask = mask.unsqueeze(2)?; + } + + mask = mask.to_dtype(embeddings.dtype())?; + embeddings = embeddings.broadcast_mul(&mask)?; + } + + self.dropout.forward(&embeddings) + } +} + +// https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/deberta_v2/modeling_deberta_v2.py#L72 +struct XSoftmax {} + +impl XSoftmax { + pub fn apply(input: &Tensor, mask: &Tensor, dim: D, device: &Device) -> Result { + // NOTE: At the time of this writing, candle does not have a logical-not operator. + let mut rmask = mask.broadcast_as(input.shape())?.to_dtype(DType::F32)?; + + rmask = rmask + .broadcast_lt(&Tensor::new(&[1.0_f32], device)?)? + .to_dtype(DType::U8)?; + + let min_value_tensor = Tensor::new(&[f32::MIN], device)?.broadcast_as(input.shape())?; + let mut output = rmask.where_cond(&min_value_tensor, input)?; + + output = candle_nn::ops::softmax(&output, dim)?; + + let t_zeroes = Tensor::new(&[0f32], device)?.broadcast_as(input.shape())?; + output = rmask.where_cond(&t_zeroes, &output)?; + + Ok(output) + } +} + +// https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/deberta_v2/modeling_deberta_v2.py#L605 +pub struct DebertaV2DisentangledSelfAttention { + config: Config, + num_attention_heads: usize, + query_proj: candle_nn::Linear, + key_proj: candle_nn::Linear, + value_proj: candle_nn::Linear, + dropout: StableDropout, + device: Device, + relative_attention: bool, + pos_dropout: Option, + position_buckets: isize, + max_relative_positions: isize, + pos_ebd_size: isize, + share_att_key: bool, + pos_key_proj: Option, + pos_query_proj: Option, +} + +impl DebertaV2DisentangledSelfAttention { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let config = config.clone(); + let vb = vb.clone(); + + if !config + .hidden_size + .is_multiple_of(config.num_attention_heads) + { + return Err(candle::Error::Msg(format!( + "The hidden size {} is not a multiple of the number of attention heads {}", + config.hidden_size, config.num_attention_heads + ))); + } + + let num_attention_heads = config.num_attention_heads; + + let attention_head_size = config + .attention_head_size + .unwrap_or(config.hidden_size / config.num_attention_heads); + + let all_head_size = num_attention_heads * attention_head_size; + + let query_proj = candle_nn::linear(config.hidden_size, all_head_size, vb.pp("query_proj"))?; + let key_proj = candle_nn::linear(config.hidden_size, all_head_size, vb.pp("key_proj"))?; + let value_proj = candle_nn::linear(config.hidden_size, all_head_size, vb.pp("value_proj"))?; + + let share_att_key = config.share_att_key.unwrap_or(false); + let relative_attention = config.relative_attention; + let mut max_relative_positions = config.max_relative_positions; + + let mut pos_ebd_size: isize = 0; + let position_buckets = config.position_buckets.unwrap_or(-1); + let mut pos_dropout: Option = None; + let mut pos_key_proj: Option = None; + let mut pos_query_proj: Option = None; + + if relative_attention { + if max_relative_positions < 1 { + max_relative_positions = config.max_position_embeddings as isize; + } + pos_ebd_size = max_relative_positions; + if position_buckets > 0 { + pos_ebd_size = position_buckets + } + + pos_dropout = Some(StableDropout::new(config.hidden_dropout_prob)); + + if !share_att_key { + if config.pos_att_type.iter().any(|s| s == "c2p") { + pos_key_proj = Some(candle_nn::linear( + config.hidden_size, + all_head_size, + vb.pp("pos_key_proj"), + )?); + } + if config.pos_att_type.iter().any(|s| s == "p2c") { + pos_query_proj = Some(candle_nn::linear( + config.hidden_size, + all_head_size, + vb.pp("pos_query_proj"), + )?); + } + } + } + + let dropout = StableDropout::new(config.attention_probs_dropout_prob); + let device = vb.device().clone(); + + Ok(Self { + config, + num_attention_heads, + query_proj, + key_proj, + value_proj, + dropout, + device, + relative_attention, + pos_dropout, + position_buckets, + max_relative_positions, + pos_ebd_size, + share_att_key, + pos_key_proj, + pos_query_proj, + }) + } + + pub fn forward( + &self, + hidden_states: &Tensor, + attention_mask: &Tensor, + query_states: Option<&Tensor>, + relative_pos: Option<&Tensor>, + rel_embeddings: Option<&Tensor>, + ) -> Result { + let query_states = match query_states { + Some(qs) => qs, + None => hidden_states, + }; + + let query_layer = self.transpose_for_scores(&self.query_proj.forward(query_states)?)?; + let key_layer = self.transpose_for_scores(&self.key_proj.forward(query_states)?)?; + let value_layer = self.transpose_for_scores(&self.value_proj.forward(query_states)?)?; + + let mut rel_att: Option = None; + + let mut scale_factor: usize = 1; + + if self.config.pos_att_type.iter().any(|s| s == "c2p") { + scale_factor += 1; + } + + if self.config.pos_att_type.iter().any(|s| s == "p2c") { + scale_factor += 1; + } + + let scale = { + let q_size = query_layer.dim(D::Minus1)?; + Tensor::new(&[(q_size * scale_factor) as f32], &self.device)?.sqrt()? + }; + + let mut attention_scores: Tensor = { + let key_layer_transposed = key_layer.t()?; + let div = key_layer_transposed + .broadcast_div(scale.to_dtype(query_layer.dtype())?.as_ref())?; + query_layer.matmul(&div)? + }; + + if self.relative_attention { + if let Some(rel_embeddings) = rel_embeddings { + let rel_embeddings = self + .pos_dropout + .as_ref() + .context("relative_attention requires pos_dropout")? + .forward(rel_embeddings)?; + rel_att = Some(self.disentangled_attention_bias( + query_layer, + key_layer, + relative_pos, + rel_embeddings, + scale_factor, + )?); + } + } + + if let Some(rel_att) = rel_att { + attention_scores = attention_scores.broadcast_add(&rel_att)?; + } + + attention_scores = attention_scores.reshape(( + (), + self.num_attention_heads, + attention_scores.dim(D::Minus2)?, + attention_scores.dim(D::Minus1)?, + ))?; + + let mut attention_probs = + XSoftmax::apply(&attention_scores, attention_mask, D::Minus1, &self.device)?; + + attention_probs = self.dropout.forward(&attention_probs)?; + + let mut context_layer = attention_probs + .reshape(( + (), + attention_probs.dim(D::Minus2)?, + attention_probs.dim(D::Minus1)?, + ))? + .matmul(&value_layer)?; + + context_layer = context_layer + .reshape(( + (), + self.num_attention_heads, + context_layer.dim(D::Minus2)?, + context_layer.dim(D::Minus1)?, + ))? + .permute((0, 2, 1, 3))? + .contiguous()?; + + let dims = context_layer.dims(); + + context_layer = match dims.len() { + 2 => context_layer.reshape(())?, + 3 => context_layer.reshape((dims[0], ()))?, + 4 => context_layer.reshape((dims[0], dims[1], ()))?, + 5 => context_layer.reshape((dims[0], dims[1], dims[2], ()))?, + _ => { + bail!( + "Invalid shape for DisentabgledSelfAttention context layer: {:?}", + dims + ) + } + }; + + Ok(context_layer) + } + + fn transpose_for_scores(&self, xs: &Tensor) -> Result { + let dims = xs.dims().to_vec(); + match dims.len() { + 3 => { + let reshaped = xs.reshape((dims[0], dims[1], self.num_attention_heads, ()))?; + + reshaped.transpose(1, 2)?.contiguous()?.reshape(( + (), + reshaped.dim(1)?, + reshaped.dim(D::Minus1)?, + )) + } + shape => { + bail!("Invalid shape for transpose_for_scores. Expected 3 dimensions, got {shape}") + } + } + } + + fn disentangled_attention_bias( + &self, + query_layer: Tensor, + key_layer: Tensor, + relative_pos: Option<&Tensor>, + rel_embeddings: Tensor, + scale_factor: usize, + ) -> Result { + let mut relative_pos = relative_pos.map_or( + build_relative_position( + query_layer.dim(D::Minus2)?, + key_layer.dim(D::Minus2)?, + &self.device, + Some(self.position_buckets), + Some(self.max_relative_positions), + )?, + |pos| pos.clone(), + ); + + relative_pos = match relative_pos.dims().len() { + 2 => relative_pos.unsqueeze(0)?.unsqueeze(0)?, + 3 => relative_pos.unsqueeze(1)?, + other => { + bail!("Relative position ids must be of dim 2 or 3 or 4. Got dim of size {other}") + } + }; + + let att_span = self.pos_ebd_size; + + let rel_embeddings = rel_embeddings + .narrow(0, 0, (att_span * 2) as usize)? + .unsqueeze(0)?; + + let mut pos_query_layer: Option = None; + let mut pos_key_layer: Option = None; + + let repeat_with = query_layer.dim(0)? / self.num_attention_heads; + if self.share_att_key { + pos_query_layer = Some( + self.transpose_for_scores(&self.query_proj.forward(&rel_embeddings)?)? + .repeat(repeat_with)?, + ); + + pos_key_layer = Some( + self.transpose_for_scores(&self.key_proj.forward(&rel_embeddings)?)? + .repeat(repeat_with)?, + ) + } else { + if self.config.pos_att_type.iter().any(|s| s == "c2p") { + pos_key_layer = Some( + self.transpose_for_scores( + &self + .pos_key_proj + .as_ref() + .context( + "Need pos_key_proj when share_att_key is false or not specified", + )? + .forward(&rel_embeddings)?, + )? + .repeat(repeat_with)?, + ) + } + if self.config.pos_att_type.iter().any(|s| s == "p2c") { + pos_query_layer = Some(self.transpose_for_scores(&self + .pos_query_proj + .as_ref() + .context("Need a pos_query_proj when share_att_key is false or not specified")? + .forward(&rel_embeddings)?)?.repeat(repeat_with)?) + } + } + + let mut score = Tensor::new(&[0 as f32], &self.device)?; + + if self.config.pos_att_type.iter().any(|s| s == "c2p") { + let pos_key_layer = pos_key_layer.context("c2p without pos_key_layer")?; + + let scale = Tensor::new( + &[(pos_key_layer.dim(D::Minus1)? * scale_factor) as f32], + &self.device, + )? + .sqrt()?; + + let mut c2p_att = query_layer.matmul(&pos_key_layer.t()?)?; + + let c2p_pos = relative_pos + .broadcast_add(&Tensor::new(&[att_span as i64], &self.device)?)? + .clamp(0 as f32, (att_span * 2 - 1) as f32)?; + + c2p_att = c2p_att.gather( + &c2p_pos + .squeeze(0)? + .expand(&[ + query_layer.dim(0)?, + query_layer.dim(1)?, + relative_pos.dim(D::Minus1)?, + ])? + .contiguous()?, + D::Minus1, + )?; + + score = score.broadcast_add( + &c2p_att.broadcast_div(scale.to_dtype(c2p_att.dtype())?.as_ref())?, + )?; + } + + if self.config.pos_att_type.iter().any(|s| s == "p2c") { + let pos_query_layer = pos_query_layer.context("p2c without pos_key_layer")?; + + let scale = Tensor::new( + &[(pos_query_layer.dim(D::Minus1)? * scale_factor) as f32], + &self.device, + )? + .sqrt()?; + + let r_pos = { + if key_layer.dim(D::Minus2)? != query_layer.dim(D::Minus2)? { + build_relative_position( + key_layer.dim(D::Minus2)?, + key_layer.dim(D::Minus2)?, + &self.device, + Some(self.position_buckets), + Some(self.max_relative_positions), + )? + .unsqueeze(0)? + } else { + relative_pos + } + }; + + let p2c_pos = r_pos + .to_dtype(DType::F32)? + .neg()? + .broadcast_add(&Tensor::new(&[att_span as f32], &self.device)?)? + .clamp(0f32, (att_span * 2 - 1) as f32)?; + + let p2c_att = key_layer + .matmul(&pos_query_layer.t()?)? + .gather( + &p2c_pos + .squeeze(0)? + .expand(&[ + query_layer.dim(0)?, + key_layer.dim(D::Minus2)?, + key_layer.dim(D::Minus2)?, + ])? + .contiguous()? + .to_dtype(DType::U32)?, + D::Minus1, + )? + .t()?; + + score = + score.broadcast_add(&p2c_att.broadcast_div(&scale.to_dtype(p2c_att.dtype())?)?)?; + } + + Ok(score) + } +} + +// https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/deberta_v2/modeling_deberta_v2.py#L270 +pub struct DebertaV2Attention { + dsa: DebertaV2DisentangledSelfAttention, + output: DebertaV2SelfOutput, +} + +impl DebertaV2Attention { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let dsa = DebertaV2DisentangledSelfAttention::load(vb.pp("attention.self"), config)?; + let output = DebertaV2SelfOutput::load(vb.pp("attention.output"), config)?; + Ok(Self { dsa, output }) + } + + fn forward( + &self, + hidden_states: &Tensor, + attention_mask: &Tensor, + query_states: Option<&Tensor>, + relative_pos: Option<&Tensor>, + rel_embeddings: Option<&Tensor>, + ) -> Result { + let self_output = self.dsa.forward( + hidden_states, + attention_mask, + query_states, + relative_pos, + rel_embeddings, + )?; + + self.output + .forward(&self_output, query_states.unwrap_or(hidden_states)) + } +} + +// https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/deberta_v2/modeling_deberta_v2.py#L255 +pub struct DebertaV2SelfOutput { + dense: candle_nn::Linear, + layer_norm: LayerNorm, + dropout: StableDropout, +} + +impl DebertaV2SelfOutput { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let dense = candle_nn::linear(config.hidden_size, config.hidden_size, vb.pp("dense"))?; + let layer_norm = candle_nn::layer_norm( + config.hidden_size, + config.layer_norm_eps, + vb.pp("LayerNorm"), + )?; + let dropout = StableDropout::new(config.hidden_dropout_prob); + Ok(Self { + dense, + layer_norm, + dropout, + }) + } + + pub fn forward(&self, hidden_states: &Tensor, input_tensor: &Tensor) -> Result { + let mut hidden_states = self.dense.forward(hidden_states)?; + hidden_states = self.dropout.forward(&hidden_states)?; + self.layer_norm + .forward(&hidden_states.broadcast_add(input_tensor)?) + } +} + +// https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/deberta_v2/modeling_deberta_v2.py#L307 +pub struct DebertaV2Intermediate { + dense: candle_nn::Linear, + intermediate_act: HiddenActLayer, +} + +impl DebertaV2Intermediate { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let dense = candle_nn::linear( + config.hidden_size, + config.intermediate_size, + vb.pp("intermediate.dense"), + )?; + let intermediate_act = HiddenActLayer::new(config.hidden_act); + Ok(Self { + dense, + intermediate_act, + }) + } + + pub fn forward(&self, hidden_states: &Tensor) -> Result { + self.intermediate_act + .forward(&self.dense.forward(hidden_states)?) + } +} + +// https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/deberta_v2/modeling_deberta_v2.py#L323 +pub struct DebertaV2Output { + dense: candle_nn::Linear, + layer_norm: LayerNorm, + dropout: StableDropout, +} + +impl DebertaV2Output { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let dense = candle_nn::linear( + config.intermediate_size, + config.hidden_size, + vb.pp("output.dense"), + )?; + let layer_norm = candle_nn::layer_norm( + config.hidden_size, + config.layer_norm_eps, + vb.pp("output.LayerNorm"), + )?; + let dropout = StableDropout::new(config.hidden_dropout_prob); + Ok(Self { + dense, + layer_norm, + dropout, + }) + } + + pub fn forward(&self, hidden_states: &Tensor, input_tensor: &Tensor) -> Result { + let mut hidden_states = self.dense.forward(hidden_states)?; + hidden_states = self.dropout.forward(&hidden_states)?; + hidden_states = { + let to_norm = hidden_states.broadcast_add(input_tensor)?; + self.layer_norm.forward(&to_norm)? + }; + Ok(hidden_states) + } +} + +// https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/deberta_v2/modeling_deberta_v2.py#L339 +pub struct DebertaV2Layer { + attention: DebertaV2Attention, + intermediate: DebertaV2Intermediate, + output: DebertaV2Output, +} + +impl DebertaV2Layer { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let attention = DebertaV2Attention::load(vb.clone(), config)?; + let intermediate = DebertaV2Intermediate::load(vb.clone(), config)?; + let output = DebertaV2Output::load(vb.clone(), config)?; + Ok(Self { + attention, + intermediate, + output, + }) + } + + fn forward( + &self, + hidden_states: &Tensor, + attention_mask: &Tensor, + query_states: Option<&Tensor>, + relative_pos: Option<&Tensor>, + rel_embeddings: Option<&Tensor>, + ) -> Result { + let attention_output = self.attention.forward( + hidden_states, + attention_mask, + query_states, + relative_pos, + rel_embeddings, + )?; + + let intermediate_output = self.intermediate.forward(&attention_output)?; + + let layer_output = self + .output + .forward(&intermediate_output, &attention_output)?; + + Ok(layer_output) + } +} + +// TODO: In order to fully test ConvLayer a model needs to be found has a configuration where `conv_kernel_size` exists and is > 0 +// https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/deberta_v2/modeling_deberta_v2.py#L373 +pub struct ConvLayer { + _conv_act: String, + _conv: Conv1d, + _layer_norm: LayerNorm, + _dropout: StableDropout, + _config: Config, +} + +impl ConvLayer { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let config = config.clone(); + let kernel_size = config.conv_kernel_size.unwrap_or(3); + let groups = config.conv_groups.unwrap_or(1); + let conv_act: String = config.conv_act.clone().unwrap_or("tanh".to_string()); + + let conv_conf = Conv1dConfig { + padding: (kernel_size - 1) / 2, + groups, + ..Default::default() + }; + + let conv = conv1d( + config.hidden_size, + config.hidden_size, + kernel_size, + conv_conf, + vb.pp("conv"), + )?; + + let layer_norm = layer_norm( + config.hidden_size, + config.layer_norm_eps, + vb.pp("LayerNorm"), + )?; + + let dropout = StableDropout::new(config.hidden_dropout_prob); + + Ok(Self { + _conv_act: conv_act, + _conv: conv, + _layer_norm: layer_norm, + _dropout: dropout, + _config: config, + }) + } + + pub fn forward( + &self, + _hidden_states: &Tensor, + _residual_states: &Tensor, + _input_mask: &Tensor, + ) -> Result { + todo!("Need a model that contains a conv layer to test against.") + } +} + +// https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/deberta_v2/modeling_deberta_v2.py#L409 +pub struct DebertaV2Encoder { + layer: Vec, + relative_attention: bool, + max_relative_positions: isize, + position_buckets: isize, + rel_embeddings: Option, + norm_rel_ebd: String, + layer_norm: Option, + conv: Option, + device: Device, +} + +impl DebertaV2Encoder { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let layer = (0..config.num_hidden_layers) + .map(|index| DebertaV2Layer::load(vb.pp(format!("layer.{index}")), config)) + .collect::>>()?; + + let relative_attention = config.relative_attention; + let mut max_relative_positions = config.max_relative_positions; + + let position_buckets = config.position_buckets.unwrap_or(-1); + + let mut rel_embeddings: Option = None; + + if relative_attention { + if max_relative_positions < 1 { + max_relative_positions = config.max_position_embeddings as isize; + } + + let mut pos_ebd_size = max_relative_positions * 2; + + if position_buckets > 0 { + pos_ebd_size = position_buckets * 2; + } + + rel_embeddings = Some(embedding( + pos_ebd_size as usize, + config.hidden_size, + vb.pp("rel_embeddings"), + )?); + } + + // NOTE: The Python code assumes that the config attribute "norm_rel_ebd" is an array of some kind, but most examples have it as a string. + // So it might need to be updated at some point. + let norm_rel_ebd = match config.norm_rel_ebd.as_ref() { + Some(nre) => nre.trim().to_string(), + None => "none".to_string(), + }; + + let layer_norm: Option = if norm_rel_ebd == "layer_norm" { + Some(layer_norm( + config.hidden_size, + config.layer_norm_eps, + vb.pp("LayerNorm"), + )?) + } else { + None + }; + + let conv: Option = if config.conv_kernel_size.unwrap_or(0) > 0 { + Some(ConvLayer::load(vb.pp("conv"), config)?) + } else { + None + }; + + Ok(Self { + layer, + relative_attention, + max_relative_positions, + position_buckets, + rel_embeddings, + norm_rel_ebd, + layer_norm, + conv, + device: vb.device().clone(), + }) + } + + pub fn forward( + &self, + hidden_states: &Tensor, + attention_mask: &Tensor, + query_states: Option<&Tensor>, + relative_pos: Option<&Tensor>, + ) -> Result { + let input_mask = if attention_mask.dims().len() <= 2 { + attention_mask.clone() + } else { + attention_mask + .sum_keepdim(attention_mask.rank() - 2)? + .gt(0.)? + }; + + let attention_mask = self.get_attention_mask(attention_mask.clone())?; + + let relative_pos = self.get_rel_pos(hidden_states, query_states, relative_pos)?; + + let mut next_kv: Tensor = hidden_states.clone(); + let rel_embeddings = self.get_rel_embedding()?; + let mut output_states = next_kv.to_owned(); + let mut query_states: Option = query_states.cloned(); + + for (i, layer_module) in self.layer.iter().enumerate() { + // NOTE: The original python code branches here if this model is being + // used for training vs. inferencing. For now, we will only handle the + // inferencing side of things + + output_states = layer_module.forward( + next_kv.as_ref(), + &attention_mask, + query_states.as_ref(), + relative_pos.as_ref(), + rel_embeddings.as_ref(), + )?; + + if i == 0 { + if let Some(conv) = &self.conv { + output_states = conv.forward(hidden_states, &output_states, &input_mask)?; + } + } + + if query_states.is_some() { + query_states = Some(output_states.clone()); + } else { + next_kv = output_states.clone(); + } + } + + Ok(output_states) + } + + fn get_attention_mask(&self, mut attention_mask: Tensor) -> Result { + match attention_mask.dims().len() { + 0..=2 => { + let extended_attention_mask = attention_mask.unsqueeze(1)?.unsqueeze(2)?; + attention_mask = extended_attention_mask.broadcast_mul( + &extended_attention_mask + .squeeze(D::Minus2)? + .unsqueeze(D::Minus1)?, + )?; + } + 3 => attention_mask = attention_mask.unsqueeze(1)?, + len => bail!("Unsupported attentiom mask size length: {len}"), + } + + Ok(attention_mask) + } + + fn get_rel_pos( + &self, + hidden_states: &Tensor, + query_states: Option<&Tensor>, + relative_pos: Option<&Tensor>, + ) -> Result> { + if self.relative_attention && relative_pos.is_none() { + let q = if let Some(query_states) = query_states { + query_states.dim(D::Minus2)? + } else { + hidden_states.dim(D::Minus2)? + }; + + return Ok(Some(build_relative_position( + q, + hidden_states.dim(D::Minus2)?, + &self.device, + Some(self.position_buckets), + Some(self.max_relative_positions), + )?)); + } + + if relative_pos.is_some() { + Ok(relative_pos.cloned()) + } else { + Ok(None) + } + } + fn get_rel_embedding(&self) -> Result> { + if !self.relative_attention { + return Ok(None); + } + + let rel_embeddings = self + .rel_embeddings + .as_ref() + .context("self.rel_embeddings not present when using relative_attention")? + .embeddings() + .clone(); + + if !self.norm_rel_ebd.contains("layer_norm") { + return Ok(Some(rel_embeddings)); + } + + let layer_normed_embeddings = self + .layer_norm + .as_ref() + .context("DebertaV2Encoder layer_norm is None when norm_rel_ebd contains layer_norm")? + .forward(&rel_embeddings)?; + + Ok(Some(layer_normed_embeddings)) + } +} + +// https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/deberta_v2/modeling_deberta_v2.py#L991 +pub struct DebertaV2Model { + embeddings: DebertaV2Embeddings, + encoder: DebertaV2Encoder, + z_steps: usize, + pub device: Device, +} + +impl DebertaV2Model { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let vb = vb.clone(); + let embeddings = DebertaV2Embeddings::load(vb.pp("embeddings"), config)?; + let encoder = DebertaV2Encoder::load(vb.pp("encoder"), config)?; + let z_steps: usize = 0; + + Ok(Self { + embeddings, + encoder, + z_steps, + device: vb.device().clone(), + }) + } + + pub fn forward( + &self, + input_ids: &Tensor, + token_type_ids: Option, + attention_mask: Option, + ) -> Result { + let input_ids_shape = input_ids.shape(); + + let attention_mask = match attention_mask { + Some(mask) => mask, + None => Tensor::ones(input_ids_shape, DType::I64, &self.device)?, + }; + + let token_type_ids = match token_type_ids { + Some(ids) => ids, + None => Tensor::zeros(input_ids_shape, DType::U32, &self.device)?, + }; + + let embedding_output = self.embeddings.forward( + Some(input_ids), + Some(&token_type_ids), + None, + Some(&attention_mask), + None, + )?; + + let encoder_output = + self.encoder + .forward(&embedding_output, &attention_mask, None, None)?; + + if self.z_steps > 1 { + todo!("Complete DebertaV2Model forward() when z_steps > 1 -- Needs a model to test this situation.") + } + + Ok(encoder_output) + } +} + +#[derive(Debug)] +pub struct NERItem { + pub entity: String, + pub word: String, + pub score: f32, + pub start: usize, + pub end: usize, + pub index: usize, +} + +#[derive(Debug)] +pub struct TextClassificationItem { + pub label: String, + pub score: f32, +} + +pub struct DebertaV2NERModel { + pub device: Device, + deberta: DebertaV2Model, + dropout: candle_nn::Dropout, + classifier: candle_nn::Linear, +} + +fn id2label_len(config: &Config, id2label: Option>) -> Result { + let id2label_len = match (&config.id2label, id2label) { + (None, None) => bail!("Id2Label is either not present in the model configuration or not passed into DebertaV2NERModel::load as a parameter"), + (None, Some(id2label_p)) => id2label_p.len(), + (Some(id2label_c), None) => id2label_c.len(), + (Some(id2label_c), Some(id2label_p)) => { + if *id2label_c == id2label_p { + id2label_c.len() + } else { + bail!("Id2Label is both present in the model configuration and provided as a parameter, and they are different.") + } + } + }; + Ok(id2label_len) +} + +impl DebertaV2NERModel { + pub fn load(vb: VarBuilder, config: &Config, id2label: Option) -> Result { + let id2label_len = id2label_len(config, id2label)?; + + let deberta = DebertaV2Model::load(vb.clone(), config)?; + let dropout = candle_nn::Dropout::new(config.hidden_dropout_prob as f32); + let classifier: candle_nn::Linear = candle_nn::linear_no_bias( + config.hidden_size, + id2label_len, + vb.root().pp("classifier"), + )?; + + Ok(Self { + device: vb.device().clone(), + deberta, + dropout, + classifier, + }) + } + + pub fn forward( + &self, + input_ids: &Tensor, + token_type_ids: Option, + attention_mask: Option, + ) -> Result { + let output = self + .deberta + .forward(input_ids, token_type_ids, attention_mask)?; + let output = self.dropout.forward(&output, false)?; + self.classifier.forward(&output) + } +} + +pub struct DebertaV2SeqClassificationModel { + pub device: Device, + deberta: DebertaV2Model, + dropout: StableDropout, + pooler: DebertaV2ContextPooler, + classifier: candle_nn::Linear, +} + +impl DebertaV2SeqClassificationModel { + pub fn load(vb: VarBuilder, config: &Config, id2label: Option) -> Result { + let id2label_len = id2label_len(config, id2label)?; + let deberta = DebertaV2Model::load(vb.clone(), config)?; + let pooler = DebertaV2ContextPooler::load(vb.clone(), config)?; + let output_dim = pooler.output_dim()?; + let classifier = candle_nn::linear(output_dim, id2label_len, vb.root().pp("classifier"))?; + let dropout = match config.cls_dropout { + Some(cls_dropout) => StableDropout::new(cls_dropout), + None => StableDropout::new(config.hidden_dropout_prob), + }; + + Ok(Self { + device: vb.device().clone(), + deberta, + dropout, + pooler, + classifier, + }) + } + + pub fn forward( + &self, + input_ids: &Tensor, + token_type_ids: Option, + attention_mask: Option, + ) -> Result { + let encoder_layer = self + .deberta + .forward(input_ids, token_type_ids, attention_mask)?; + let pooled_output = self.pooler.forward(&encoder_layer)?; + let pooled_output = self.dropout.forward(&pooled_output)?; + self.classifier.forward(&pooled_output) + } +} + +pub struct DebertaV2ContextPooler { + dense: candle_nn::Linear, + dropout: StableDropout, + config: Config, +} + +// https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/deberta_v2/modeling_deberta_v2.py#L49 +impl DebertaV2ContextPooler { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let pooler_hidden_size = config + .pooler_hidden_size + .context("config.pooler_hidden_size is required for DebertaV2ContextPooler")?; + + let pooler_dropout = config + .pooler_dropout + .context("config.pooler_dropout is required for DebertaV2ContextPooler")?; + + let dense = candle_nn::linear( + pooler_hidden_size, + pooler_hidden_size, + vb.root().pp("pooler.dense"), + )?; + + let dropout = StableDropout::new(pooler_dropout); + + Ok(Self { + dense, + dropout, + config: config.clone(), + }) + } + + pub fn forward(&self, hidden_states: &Tensor) -> Result { + let context_token = hidden_states.narrow(1, 0, 1)?.squeeze(1)?; + let context_token = self.dropout.forward(&context_token)?; + + let pooled_output = self.dense.forward(&context_token.contiguous()?)?; + let pooler_hidden_act = self + .config + .pooler_hidden_act + .context("Could not obtain pooler hidden act from config")?; + + HiddenActLayer::new(pooler_hidden_act).forward(&pooled_output) + } + + pub fn output_dim(&self) -> Result { + self.config.pooler_hidden_size.context("DebertaV2ContextPooler cannot return output_dim (pooler_hidden_size) since it is not specified in the model config") + } +} + +// https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/deberta_v2/modeling_deberta_v2.py#L557 +pub(crate) fn build_relative_position( + query_size: usize, + key_size: usize, + device: &Device, + bucket_size: Option, + max_position: Option, +) -> Result { + let q_ids = Tensor::arange(0, query_size as i64, device)?.unsqueeze(0)?; + let k_ids: Tensor = Tensor::arange(0, key_size as i64, device)?.unsqueeze(D::Minus1)?; + let mut rel_pos_ids = k_ids.broadcast_sub(&q_ids)?; + let bucket_size = bucket_size.unwrap_or(-1); + let max_position = max_position.unwrap_or(-1); + + if bucket_size > 0 && max_position > 0 { + rel_pos_ids = make_log_bucket_position(rel_pos_ids, bucket_size, max_position, device)?; + } + + rel_pos_ids = rel_pos_ids.to_dtype(DType::I64)?; + rel_pos_ids = rel_pos_ids.narrow(0, 0, query_size)?; + rel_pos_ids.unsqueeze(0) +} + +// https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/deberta_v2/modeling_deberta_v2.py#L542 +pub(crate) fn make_log_bucket_position( + relative_pos: Tensor, + bucket_size: isize, + max_position: isize, + device: &Device, +) -> Result { + let sign = relative_pos.to_dtype(DType::F32)?.sign()?; + + let mid = bucket_size / 2; + + let lt_mid = relative_pos.lt(mid as i64)?; + let gt_neg_mid = relative_pos.gt(-mid as i64)?; + + let condition = lt_mid + .to_dtype(candle::DType::F32)? + .mul(>_neg_mid.to_dtype(candle::DType::F32)?)? + .to_dtype(DType::U8)?; + + let on_true = Tensor::new(&[(mid - 1) as u32], device)? + .broadcast_as(relative_pos.shape())? + .to_dtype(relative_pos.dtype())?; + + let on_false = relative_pos + .to_dtype(DType::F32)? + .abs()? + .to_dtype(DType::I64)?; + + let abs_pos = condition.where_cond(&on_true, &on_false)?; + + let mid_as_tensor = Tensor::from_slice(&[mid as f32], (1,), device)?; + + let log_pos = { + let first_log = abs_pos + .to_dtype(DType::F32)? + .broadcast_div(&mid_as_tensor)? + .log()?; + + let second_log = + Tensor::from_slice(&[((max_position as f32 - 1.0) / mid as f32)], (1,), device)? + .log()?; + + let first_div_second = first_log.broadcast_div(&second_log)?; + + let to_ceil = first_div_second + .broadcast_mul(Tensor::from_slice(&[(mid - 1) as f32], (1,), device)?.as_ref())?; + + let ceil = to_ceil.ceil()?; + + ceil.broadcast_add(&mid_as_tensor)? + }; + + Ok({ + let abs_pos_lte_mid = abs_pos.to_dtype(DType::F32)?.broadcast_le(&mid_as_tensor)?; + let relative_pos = relative_pos.to_dtype(relative_pos.dtype())?; + let log_pos_mul_sign = log_pos.broadcast_mul(&sign.to_dtype(DType::F32)?)?; + abs_pos_lte_mid.where_cond(&relative_pos.to_dtype(DType::F32)?, &log_pos_mul_sign)? + }) +} diff --git a/patches/candle-transformers/src/models/deepseek2.rs b/patches/candle-transformers/src/models/deepseek2.rs new file mode 100644 index 0000000000..5260dfa62d --- /dev/null +++ b/patches/candle-transformers/src/models/deepseek2.rs @@ -0,0 +1,1073 @@ +#![allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] + +use std::{f32::consts::PI, sync::Arc}; + +use candle::{ + shape::Dim, CpuStorage, CustomOp1, DType, Device, Error, IndexOp, Layout, Result, Shape, + Tensor, WithDType, D, +}; +use candle_nn::{embedding, rms_norm, Activation, Embedding, Linear, Module, RmsNorm, VarBuilder}; +use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; +use serde::Deserialize; + +struct NonZero {} + +impl NonZero { + // Sequential version + fn nonzero(&self, vs: &[T], layout: &Layout) -> Vec { + let n = layout.dims().len(); + let mut result = Vec::new(); + let mut indices = vec![0u32; n]; + for (i, v) in vs.iter().enumerate() { + if !v.is_zero() { + let mut idx = i; + for (dim_index, dim) in layout.dims().iter().enumerate().rev() { + let d = idx % dim; + indices[dim_index] = u32::try_from(d).unwrap(); + idx /= dim; + } + result.extend_from_slice(&indices); + } + } + result + } +} + +impl CustomOp1 for NonZero { + fn name(&self) -> &'static str { + "nonzero" + } + + fn cpu_fwd(&self, storage: &CpuStorage, layout: &Layout) -> Result<(CpuStorage, Shape)> { + if !layout.is_contiguous() { + return Err(Error::RequiresContiguous { op: "nonzero" }); + } + let result = match storage { + candle::CpuStorage::U8(vs) => self.nonzero(vs, layout), + candle::CpuStorage::U32(vs) => self.nonzero(vs, layout), + candle::CpuStorage::I16(vs) => self.nonzero(vs, layout), + candle::CpuStorage::I32(vs) => self.nonzero(vs, layout), + candle::CpuStorage::I64(vs) => self.nonzero(vs, layout), + candle::CpuStorage::BF16(vs) => self.nonzero(vs, layout), + candle::CpuStorage::F16(vs) => self.nonzero(vs, layout), + candle::CpuStorage::F32(vs) => self.nonzero(vs, layout), + candle::CpuStorage::F64(vs) => self.nonzero(vs, layout), + candle::CpuStorage::F8E4M3(vs) => self.nonzero(vs, layout), + // Dummy types don't support nonzero operation + candle::CpuStorage::F6E2M3(_) => { + return Err( + candle::Error::UnsupportedDTypeForOp(candle::DType::F6E2M3, "nonzero").bt(), + ) + } + candle::CpuStorage::F6E3M2(_) => { + return Err( + candle::Error::UnsupportedDTypeForOp(candle::DType::F6E3M2, "nonzero").bt(), + ) + } + candle::CpuStorage::F4(_) => { + return Err(candle::Error::UnsupportedDTypeForOp(candle::DType::F4, "nonzero").bt()) + } + candle::CpuStorage::F8E8M0(_) => { + return Err( + candle::Error::UnsupportedDTypeForOp(candle::DType::F8E8M0, "nonzero").bt(), + ) + } + }; + let index_len = layout.dims().len(); + let result_len = result.len() / index_len; + let result = CpuStorage::U32(result); + let shape = Shape::from_dims(&[result_len, index_len]); + Ok((result, shape)) + } +} + +pub trait NonZeroOp { + fn nonzero(&self) -> Result; +} + +impl NonZeroOp for Tensor { + fn nonzero(&self) -> Result { + if !self.is_contiguous() { + return Err(candle::Error::RequiresContiguous { op: "nonzero" }); + } + let original_device = self.device(); + self.to_device(&candle::Device::Cpu)? + .apply_op1_no_bwd(&NonZero {})? + .to_device(original_device) + } +} + +pub struct TopKOutput { + pub values: Tensor, + pub indices: Tensor, +} + +pub trait TopKLastDimOp { + /// Topk in the last dim. `values` retains a gradient but `indices` has none w.r.t self. + /// This expects a contiguous tensor. + /// Note: this implements torch.topk with sorted=True. + fn topk(&self, topk: usize) -> Result; + + /// Topk in the last dim. `values` retains a gradient but `indices` has none w.r.t self. + /// This expects a contiguous tensor. + /// Note: this implements torch.topk with sorted=False. + fn topk_unsorted(&self, topk: usize) -> Result; +} + +impl TopKLastDimOp for Tensor { + fn topk(&self, topk: usize) -> Result { + // Sorted descending + let sorted_indices = self.arg_sort_last_dim(false)?; + let topk_indices = sorted_indices.narrow(D::Minus1, 0, topk)?.contiguous()?; + Ok(TopKOutput { + values: self.gather(&topk_indices, D::Minus1)?, + indices: topk_indices, + }) + } + + fn topk_unsorted(&self, topk: usize) -> Result { + // Sorted descending + let sorted_indices_all = self.arg_sort_last_dim(false)?; + let topk_indices_sorted = sorted_indices_all + .narrow(D::Minus1, 0, topk)? + .contiguous()?; + let topk_values_sorted = self.gather(&topk_indices_sorted, D::Minus1)?; + + // Reorder the indices ascending + let reorder_indices = topk_indices_sorted.arg_sort_last_dim(true)?; + let topk_indices_unsorted = topk_indices_sorted.gather(&reorder_indices, D::Minus1)?; + let topk_values_unsorted = topk_values_sorted.gather(&reorder_indices, D::Minus1)?; + Ok(TopKOutput { + values: topk_values_unsorted, + indices: topk_indices_unsorted, + }) + } +} + +pub trait SplitOp { + fn split(&self, splits: &[usize], dim: D) -> Result>; +} + +impl SplitOp for Tensor { + fn split(&self, splits: &[usize], dim: D) -> Result> { + let dim = dim.to_index(self.shape(), "split")?; + let mut split_res = Vec::new(); + let mut index = 0; + for split in splits { + split_res.push(self.narrow(dim, index, *split)?); + index += *split; + } + Ok(split_res) + } +} + +pub trait BincountOp { + fn bincount(&self, minlength: u32) -> Result>; +} + +fn bincount(values: &[u32], minlength: u32) -> Vec { + // Find the maximum value in `values` (or zero if empty) + let max_val = values.par_iter().max().copied().unwrap_or(0); + + // The final size of the bin counts must be at least `minlength` + // and large enough to include the largest value in `values`. + let result_len = (max_val + 1).max(minlength); + + // Each thread creates a local histogram (`fold`), + // and then they are merged together (`reduce`). + values + .par_iter() + .fold( + // Create a local histogram + || vec![0u32; result_len as usize], + // Update the local histogram + |mut local_counts, &val| { + local_counts[val as usize] += 1; + local_counts + }, + ) + // Merge histograms from all threads + .reduce( + // Identity (empty histogram) + || vec![0u32; result_len as usize], + // Combine two histograms + |mut global_counts, local_counts| { + for (g, l) in global_counts.iter_mut().zip(local_counts) { + *g += l; + } + global_counts + }, + ) +} + +impl BincountOp for Tensor { + fn bincount(&self, minlength: u32) -> Result> { + let values = self.to_vec1::()?; + + Ok(bincount(&values, minlength)) + } +} + +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) +} + +#[doc(hidden)] +#[macro_export] +macro_rules! serde_default_fn { + ($t:ty, $name:ident, $v:expr) => { + fn $name() -> $t { + $v + } + }; +} + +serde_default_fn!(f64, routed_scaling_factor, 1.0); +serde_default_fn!(TopkMethod, topk_method, TopkMethod::Greedy); +serde_default_fn!(usize, moe_layer_freq, 1); +serde_default_fn!(usize, first_k_dense_replace, 0); +serde_default_fn!(bool, norm_topk_prob, false); +serde_default_fn!(ScoringFunc, scoring_func, ScoringFunc::Softmax); +serde_default_fn!(Activation, hidden_act, Activation::Silu); +serde_default_fn!(bool, tie_word_embeddings, false); + +#[derive(Deserialize, Clone, Debug)] +enum TopkMethod { + #[serde(rename = "greedy")] + Greedy, + #[serde(rename = "group_limited_greedy")] + GroupLimitedGreedy, +} + +#[derive(Deserialize, Clone, Debug)] +enum ScoringFunc { + #[serde(rename = "softmax")] + Softmax, +} + +#[derive(Deserialize, Clone, Debug)] +pub struct DeepSeekV2Config { + pub(crate) vocab_size: usize, + pub(crate) hidden_size: usize, + pub(crate) intermediate_size: usize, + pub(crate) moe_intermediate_size: usize, + pub(crate) num_hidden_layers: usize, + pub(crate) num_attention_heads: usize, + pub(crate) n_shared_experts: Option, + pub(crate) n_routed_experts: Option, + #[serde(default = "routed_scaling_factor")] + pub(crate) routed_scaling_factor: f64, + #[serde(default = "topk_method")] + topk_method: TopkMethod, + pub(crate) num_experts_per_tok: Option, + #[serde(default = "moe_layer_freq")] + pub(crate) moe_layer_freq: usize, + #[serde(default = "first_k_dense_replace")] + pub(crate) first_k_dense_replace: usize, + // k dense layers + #[serde(default = "norm_topk_prob")] + pub(crate) norm_topk_prob: bool, + #[serde(default = "scoring_func")] + scoring_func: ScoringFunc, + #[serde(default = "hidden_act")] + pub(crate) hidden_act: Activation, + pub(crate) max_position_embeddings: usize, + pub(crate) rms_norm_eps: f64, + #[serde(default = "tie_word_embeddings")] + pub(crate) tie_word_embeddings: bool, + pub(crate) rope_theta: f32, + pub(crate) rope_scaling: Option, + pub(crate) attention_bias: bool, + pub(crate) q_lora_rank: Option, + pub(crate) qk_rope_head_dim: usize, + pub(crate) kv_lora_rank: usize, + pub(crate) v_head_dim: usize, + pub(crate) qk_nope_head_dim: usize, + pub(crate) n_group: usize, + pub(crate) topk_group: usize, +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum ScaledRopeType { + #[serde(alias = "su")] + #[serde(alias = "longrope")] + Su, + #[serde(alias = "yarn")] + Yarn, + #[serde(alias = "dynamic")] + Dynamic, + #[serde(alias = "linear")] + Linear, +} + +#[derive(Debug, Clone)] +pub struct DeepSeekV2RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(untagged)] +pub enum DeepSeekV2RopeScaling { + Yarn { + original_max_position_embeddings: usize, + beta_fast: f32, + beta_slow: f32, + mscale: f32, + mscale_all_dim: f32, + factor: f32, + #[serde(rename = "type")] + scaling_type: ScaledRopeType, + }, + LinearOrDynamic { + #[serde(rename = "type")] + scaling_type: ScaledRopeType, + factor: f64, + }, +} + +pub struct DeepSeekV2RopeConfig { + pub rope_scaling: Option, + pub max_position_embeddings: usize, + pub rope_theta: f32, + pub qk_rope_head_dim: usize, +} + +impl DeepSeekV2RotaryEmbedding { + fn new_unscaled(cfg: &DeepSeekV2RopeConfig, dtype: DType, dev: &Device) -> Result { + let max_seq_len = cfg.max_position_embeddings; + let dim = cfg.qk_rope_head_dim; + + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f32 / dim as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(DType::F32)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + + let sin = freqs.sin()?.to_dtype(dtype)?; + let cos = freqs.cos()?.to_dtype(dtype)?; + + Ok(Self { sin, cos }) + } + + fn yarn_find_correction_dim( + num_rot: f32, + dim: usize, + base: f32, + max_position_embeddings: usize, + ) -> f32 { + (dim as f32 * (max_position_embeddings as f32 / (num_rot * 2. * PI)).ln()) + / (2. * base.ln()) + } + + fn yarn_find_correction_range( + low_rot: f32, + high_rot: f32, + dim: usize, + base: f32, + max_position_embeddings: usize, + ) -> (f32, f32) { + let low = + Self::yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings).floor(); + let high = + Self::yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings).ceil(); + (low.max(0.), high.min(dim as f32 - 1.)) + } + + fn yarn_linear_ramp_mask(min: f32, mut max: f32, dim: usize, dev: &Device) -> Result { + if min == max { + // https://huggingface.co/deepseek-ai/DeepSeek-V2-Lite/blob/604d5664dddd88a0433dbae533b7fe9472482de0/modeling_deepseek.py#L255 + max += 0.001; + } + let linear_func = + ((Tensor::arange(0f32, dim as f32, dev)? - min as f64)? / (max as f64 - min as f64))?; + linear_func.clamp(0., 1.) + } + + pub(crate) fn yarn_get_mscale(scale: f32, mscale: f32) -> f32 { + if scale <= 1. { + return 1.; + } + 0.1 * mscale * scale.ln() + 1. + } + + #[allow(clippy::too_many_arguments)] + fn new_yarn( + cfg: &DeepSeekV2RopeConfig, + dtype: DType, + dev: &Device, + original_max_position_embeddings: usize, + beta_fast: f32, + beta_slow: f32, + factor: f32, + mscale: f32, + mscale_all_dim: f32, + ) -> Result { + let freq_extra: Vec<_> = (0..cfg.qk_rope_head_dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f32 / cfg.qk_rope_head_dim as f32)) + .collect(); + let freq_extra_len = freq_extra.len(); + let freq_extra = Tensor::from_vec(freq_extra, freq_extra_len, dev)?; + let freq_inter: Vec<_> = (0..cfg.qk_rope_head_dim) + .step_by(2) + .map(|i| 1f32 / (factor * cfg.rope_theta.powf(i as f32 / cfg.qk_rope_head_dim as f32))) + .collect(); + let freq_inter_len = freq_inter.len(); + let freq_inter = Tensor::from_vec(freq_inter, (1, freq_inter_len), dev)?; + + let (low, high) = Self::yarn_find_correction_range( + beta_fast, + beta_slow, + cfg.qk_rope_head_dim, + cfg.rope_theta, + original_max_position_embeddings, + ); + let inv_freq_mask = + (1. - Self::yarn_linear_ramp_mask(low, high, cfg.qk_rope_head_dim / 2, dev)?)?; + let inv_freq = freq_inter + .broadcast_mul(&(1. - &inv_freq_mask)?)? + .broadcast_add(&freq_extra.broadcast_mul(&inv_freq_mask)?)?; + + let t = Tensor::arange(0u32, cfg.max_position_embeddings as u32, dev)? + .to_dtype(DType::F32)? + .reshape((cfg.max_position_embeddings, 1))?; + let freqs = t.matmul(&inv_freq)?; + + let mscale = + Self::yarn_get_mscale(factor, mscale) / Self::yarn_get_mscale(factor, mscale_all_dim); + let sin = (freqs.sin()? * mscale as f64)?.to_dtype(dtype)?; + let cos = (freqs.cos()? * mscale as f64)?.to_dtype(dtype)?; + + Ok(Self { sin, cos }) + } + + pub fn new(cfg: &DeepSeekV2RopeConfig, dtype: DType, dev: &Device) -> Result { + match &cfg.rope_scaling { + Some(DeepSeekV2RopeScaling::LinearOrDynamic { + scaling_type: _, + factor: _, + }) => candle::bail!("linear and dynamic rope are not implemented yet!"), + Some(DeepSeekV2RopeScaling::Yarn { + original_max_position_embeddings, + beta_fast, + beta_slow, + factor, + mscale, + mscale_all_dim, + scaling_type: _, + }) => Self::new_yarn( + cfg, + dtype, + dev, + *original_max_position_embeddings, + *beta_fast, + *beta_slow, + *factor, + *mscale, + *mscale_all_dim, + ), + None => Self::new_unscaled(cfg, dtype, dev), + } + } + + pub fn forward( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + + let q_embed = candle_nn::rotary_emb::rope_i(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope_i(&k.contiguous()?, &cos, &sin)?; + + Ok((q_embed, k_embed)) + } +} + +impl DeepSeekV2Config { + pub(crate) fn q_head_dim(&self) -> usize { + self.qk_rope_head_dim + self.qk_nope_head_dim + } + + fn softmax_scale(&self) -> f32 { + let mut softmax_scale = 1.0 / (self.q_head_dim() as f32).sqrt(); + if let Some(DeepSeekV2RopeScaling::Yarn { + mscale_all_dim, + factor, + .. + }) = self.rope_scaling + { + let mscale = DeepSeekV2RotaryEmbedding::yarn_get_mscale(factor, mscale_all_dim); + softmax_scale = softmax_scale * mscale * mscale; + } + softmax_scale + } +} + +enum QProj { + Plain(Linear), + Lora { a: Linear, norm: RmsNorm, b: Linear }, +} + +impl QProj { + fn forward(&self, xs: &Tensor) -> Result { + match self { + Self::Lora { a, norm, b } => b.forward(&norm.forward(&a.forward(xs)?)?), + Self::Plain(lin) => lin.forward(xs), + } + } +} + +struct Attention { + q: QProj, + kv_a_proj_with_mqa: Linear, + kv_a_layernorm: RmsNorm, + kv_b_proj: Linear, + o_proj: Linear, + rotary_emb: Arc, + cfg: DeepSeekV2Config, + q_head_dim: usize, + softmax_scale: f64, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl Attention { + fn new( + rotary_emb: Arc, + cfg: &DeepSeekV2Config, + vb: VarBuilder, + ) -> Result { + let q_head_dim = cfg.q_head_dim(); + let q = match cfg.q_lora_rank { + Some(lora_rank) => { + let a = candle_nn::linear_b( + cfg.hidden_size, + lora_rank, + cfg.attention_bias, + vb.pp("q_a_proj"), + )?; + let norm = rms_norm(lora_rank, cfg.rms_norm_eps, vb.pp("q_a_layernorm"))?; + let b = candle_nn::linear_no_bias( + lora_rank, + cfg.num_attention_heads * q_head_dim, + vb.pp("q_b_proj"), + )?; + QProj::Lora { a, norm, b } + } + None => QProj::Plain(candle_nn::linear_no_bias( + cfg.hidden_size, + cfg.num_attention_heads * q_head_dim, + vb.pp("q_proj"), + )?), + }; + + let kv_a_proj_with_mqa = candle_nn::linear_b( + cfg.hidden_size, + cfg.kv_lora_rank + cfg.qk_rope_head_dim, + cfg.attention_bias, + vb.pp("kv_a_proj_with_mqa"), + )?; + let kv_a_layernorm = rms_norm(cfg.kv_lora_rank, cfg.rms_norm_eps, vb.pp("kv_a_layernorm"))?; + let kv_b_proj = candle_nn::linear_no_bias( + cfg.kv_lora_rank, + cfg.num_attention_heads * (q_head_dim - cfg.qk_rope_head_dim + cfg.v_head_dim), + vb.pp("kv_b_proj"), + )?; + + let o_proj = candle_nn::linear_b( + cfg.num_attention_heads * cfg.v_head_dim, + cfg.hidden_size, + cfg.attention_bias, + vb.pp("o_proj"), + )?; + + Ok(Self { + q, + kv_a_proj_with_mqa, + kv_a_layernorm, + kv_b_proj, + o_proj, + rotary_emb, + cfg: cfg.clone(), + q_head_dim, + softmax_scale: cfg.softmax_scale() as f64, + kv_cache: None, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (bs, seq_len, _) = xs.dims3()?; + + let q = { + let q = self.q.forward(xs)?; + q.reshape((bs, seq_len, self.cfg.num_attention_heads, self.q_head_dim))? + .transpose(1, 2)? + }; + let q_split = q.split( + &[self.cfg.qk_nope_head_dim, self.cfg.qk_rope_head_dim], + D::Minus1, + )?; + let q_nope = q_split[0].clone(); + let q_pe = q_split[1].clone(); + + let compressed_kv = self.kv_a_proj_with_mqa.forward(xs)?; + let ckv_split = compressed_kv.split( + &[self.cfg.kv_lora_rank, self.cfg.qk_rope_head_dim], + D::Minus1, + )?; + let compressed_kv = ckv_split[0].clone(); + let k_pe = { + let k_pe = ckv_split[1].clone(); + k_pe.reshape((bs, seq_len, 1, self.cfg.qk_rope_head_dim))? + .transpose(1, 2)? + }; + let kv = { + let kv = self + .kv_b_proj + .forward(&self.kv_a_layernorm.forward(&compressed_kv)?)?; + kv.reshape(( + bs, + seq_len, + self.cfg.num_attention_heads, + self.cfg.qk_nope_head_dim + self.cfg.v_head_dim, + ))? + .transpose(1, 2)? + }; + + let kv_split = kv.split(&[self.cfg.qk_nope_head_dim, self.cfg.v_head_dim], D::Minus1)?; + let k_nope = kv_split[0].clone(); + let v = kv_split[1].clone(); + + let (q_pe, k_pe) = self.rotary_emb.forward(&q_pe, &k_pe, seqlen_offset)?; + + let q = Tensor::cat(&[q_nope, q_pe], D::Minus1)?; + let k = Tensor::cat(&[k_nope, k_pe.repeat((1, q.dim(1)?, 1, 1))?], D::Minus1)?; + + let (k, v) = match &self.kv_cache { + None => (k, v), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &k], 2)?; + let value_states = Tensor::cat(&[prev_v, &v], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((k.clone(), v.clone())); + + let attn_out = { + let att = (q.contiguous()?.matmul(&k.t()?.contiguous()?)? * self.softmax_scale)?; + let att = match attention_mask { + Some(mask) => att.broadcast_add(mask)?, + None => att, + }; + + let att = candle_nn::ops::softmax_last_dim(&att)?; + // Convert to contiguous as matmul doesn't support strided vs for now. + att.matmul(&v.contiguous()?)? + }; + + let attn_out = if attention_mask.is_some() { + attn_out.transpose(1, 2)?.reshape((bs, seq_len, ()))? + } else { + attn_out.reshape((bs, seq_len, ()))? + }; + + self.o_proj.forward(&attn_out) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +struct Mlp { + gate: Linear, + up: Linear, + down: Linear, + act: Activation, +} + +impl Mlp { + fn new( + cfg: &DeepSeekV2Config, + vb: VarBuilder, + hidden_size: Option, + intermediate_size: Option, + ) -> Result { + let hidden_size = hidden_size.unwrap_or(cfg.hidden_size); + let intermediate_size = intermediate_size.unwrap_or(cfg.intermediate_size); + + Ok(Self { + gate: candle_nn::linear_no_bias(hidden_size, intermediate_size, vb.pp("gate_proj"))?, + up: candle_nn::linear_no_bias(hidden_size, intermediate_size, vb.pp("up_proj"))?, + down: candle_nn::linear_no_bias(intermediate_size, hidden_size, vb.pp("down_proj"))?, + act: cfg.hidden_act, + }) + } + + fn forward(&self, xs: &Tensor) -> Result { + let lhs = self.gate.forward(xs)?.apply(&self.act)?; + let rhs = self.up.forward(xs)?; + self.down.forward(&(&lhs * &rhs)?) + } +} + +struct MoeGate { + weight: Tensor, + cfg: DeepSeekV2Config, + top_k: usize, + n_routed_experts: usize, +} + +impl MoeGate { + fn new(cfg: &DeepSeekV2Config, vb: VarBuilder, n_routed_experts: usize) -> Result { + let weight = vb.get((n_routed_experts, cfg.hidden_size), "weight")?; + Ok(Self { + weight, + cfg: cfg.clone(), + top_k: cfg.num_experts_per_tok.unwrap(), + n_routed_experts, + }) + } + + /// (topk_idx, topk_weight) + fn forward(&self, xs: &Tensor) -> Result<(Tensor, Tensor)> { + let (bs, seq_len, h) = xs.dims3()?; + // Compute gating score + let xs = xs.reshape(((), h))?; + let logits = xs + .to_dtype(DType::F32)? + .broadcast_matmul(&self.weight.t()?.to_dtype(DType::F32)?)?; + let scores = match self.cfg.scoring_func { + ScoringFunc::Softmax => candle_nn::ops::softmax_last_dim(&logits)?, + }; + + // Select top-k experts + let (mut topk_weight, topk_idx) = match self.cfg.topk_method { + TopkMethod::Greedy => { + let TopKOutput { values, indices } = scores.topk_unsorted(self.top_k)?; + (values, indices) + } + TopkMethod::GroupLimitedGreedy => { + // (n, n_group) + let group_scores = scores + .reshape((bs * seq_len, self.cfg.n_group, ()))? + .max(D::Minus1)?; + // (n, topk_group) + let group_idx = scores.topk_unsorted(self.cfg.topk_group)?.indices; + // (n, n_group) + let group_mask = group_scores.zeros_like()?.scatter_add( + &group_idx, + &group_idx.ones_like()?.to_dtype(group_scores.dtype())?, + 1, + )?; + // (n, e) + let score_mask = group_mask + .unsqueeze(D::Minus1)? + .expand(( + bs * seq_len, + self.cfg.n_group, + self.n_routed_experts / self.cfg.n_group, + ))? + .reshape((bs, seq_len, ()))?; + // (n, e) + // Invert the mask + let tmp_scores = masked_fill(&score_mask, &(1. - &score_mask.ne(0.)?)?, 0.)?; + let TopKOutput { values, indices } = tmp_scores.topk_unsorted(self.top_k)?; + (values, indices) + } + }; + + if self.top_k > 1 && self.cfg.norm_topk_prob { + let denominator = (topk_weight.sum_keepdim(D::Minus1)? + 1e-20)?; + topk_weight = (topk_weight / denominator)?; + } else { + topk_weight = (topk_weight * self.cfg.routed_scaling_factor)?; + } + Ok((topk_idx, topk_weight)) + } +} + +struct Moe { + experts: Vec, + shared_experts: Option, + gate: MoeGate, +} + +impl Moe { + fn new( + cfg: &DeepSeekV2Config, + vb: VarBuilder, + + n_shared_experts: Option, + n_routed_experts: usize, + ) -> Result { + let mut experts = Vec::with_capacity(n_routed_experts); + for i in 0..n_routed_experts { + let vb_e = vb.pp("experts").pp(i); + experts.push(Mlp::new(cfg, vb_e, None, Some(cfg.moe_intermediate_size))?); + } + let shared_experts = if let Some(n_shared_experts) = n_shared_experts { + let intermediate_size = cfg.moe_intermediate_size * n_shared_experts; + Some(Mlp::new( + cfg, + vb.pp("shared_experts"), + None, + Some(intermediate_size), + )?) + } else { + None + }; + let gate = MoeGate::new(cfg, vb.pp("gate"), n_routed_experts)?; + Ok(Self { + experts, + shared_experts, + gate, + }) + } + + fn moe_infer(&self, xs: &Tensor, topk_ids: &Tensor, topk_weight: &Tensor) -> Result { + let mut y = xs.zeros_like()?; + let counts = topk_ids + .flatten_all()? + .bincount(self.experts.len() as u32)?; + for (i, expert) in self.experts.iter().enumerate() { + if counts[i] == 0 { + continue; + } + let idx_top = topk_ids.eq(i as f64)?.nonzero()?.t()?; + let idx = &idx_top.i(0)?.contiguous()?; + let top = &idx_top.i(1)?.contiguous()?; + + y = y.index_add( + idx, + &expert.forward(&xs.index_select(idx, 0)?)?.broadcast_mul( + &topk_weight + .index_select(idx, 0)? + .gather(&top.unsqueeze(1)?, 1)? + .squeeze(1)? + .unsqueeze(D::Minus1)? + .to_dtype(xs.dtype())?, + )?, + 0, + )?; + } + + Ok(y) + } + + fn forward(&self, xs: &Tensor) -> Result { + let identity = xs.clone(); + let orig_shape = xs.shape(); + let (topk_idx, topk_weight) = self.gate.forward(xs)?; + let xs = xs.reshape(((), xs.dim(D::Minus1)?))?; + + let mut y = self + .moe_infer(&xs, &topk_idx, &topk_weight)? + .reshape(orig_shape)?; + if let Some(ref shared_experts) = self.shared_experts { + y = (y + shared_experts.forward(&identity)?)?; + } + Ok(y) + } +} + +enum MoeOrMlp { + Moe(Box), + Mlp(Box), +} + +impl MoeOrMlp { + fn forward(&self, xs: &Tensor) -> Result { + match self { + Self::Mlp(mlp) => mlp.forward(xs), + Self::Moe(moe) => moe.forward(xs), + } + } +} + +struct DecoderLayer { + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, + attn: Attention, + moe_or_mlp: MoeOrMlp, +} + +impl DecoderLayer { + fn new( + rotary_emb: Arc, + cfg: &DeepSeekV2Config, + vb: VarBuilder, + layer_idx: usize, + ) -> Result { + let attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let input_layernorm = + rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let post_attention_layernorm = rms_norm( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + let moe_or_mlp = if let Some(n_routed_experts) = cfg.n_routed_experts { + if layer_idx >= cfg.first_k_dense_replace + && layer_idx.is_multiple_of(cfg.moe_layer_freq) + { + MoeOrMlp::Moe( + Moe::new(cfg, vb.pp("mlp"), cfg.n_shared_experts, n_routed_experts)?.into(), + ) + } else { + MoeOrMlp::Mlp(Mlp::new(cfg, vb.pp("mlp"), None, None)?.into()) + } + } else { + MoeOrMlp::Mlp(Mlp::new(cfg, vb.pp("mlp"), None, None)?.into()) + }; + + Ok(Self { + input_layernorm, + post_attention_layernorm, + attn, + moe_or_mlp, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = self + .moe_or_mlp + .forward(&xs.apply(&self.post_attention_layernorm)?)?; + residual + xs + } + + fn clear_kv_cache(&mut self) { + self.attn.clear_kv_cache(); + } +} + +pub struct DeepSeekV2 { + lm_head: Linear, + embed_tokens: Embedding, + norm: RmsNorm, + layers: Vec, + dtype: DType, + device: Device, +} + +impl DeepSeekV2 { + pub fn new(cfg: &DeepSeekV2Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + + let embed_tokens = embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let lm_head = if !cfg.tie_word_embeddings { + candle_nn::linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))? + } else { + candle_nn::Linear::new(embed_tokens.embeddings().clone(), None) + }; + let norm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb_m.pp("norm"))?; + + let rope_cfg = DeepSeekV2RopeConfig { + rope_scaling: cfg.rope_scaling.clone(), + max_position_embeddings: cfg.max_position_embeddings, + rope_theta: cfg.rope_theta, + qk_rope_head_dim: cfg.qk_rope_head_dim, + }; + let rotary_emb = Arc::new(DeepSeekV2RotaryEmbedding::new( + &rope_cfg, + vb.dtype(), + vb.device(), + )?); + + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(rotary_emb.clone(), cfg, vb_l.pp(layer_idx), layer_idx)?; + layers.push(layer) + } + + Ok(Self { + lm_head, + embed_tokens, + norm, + layers, + dtype: vb.dtype(), + device: vb.device().clone(), + }) + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { f32::NEG_INFINITY } else { 0. })) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let (bs, seq_len) = input_ids.dims2()?; + let mut xs = self.embed_tokens.forward(input_ids)?; + let attention_mask = if seq_len == 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(bs, seq_len, seqlen_offset)?; + Some(mask) + }; + for layer in &mut self.layers { + xs = layer.forward( + &xs, + attention_mask + .as_ref() + .map(|m| m.to_device(xs.device()).unwrap()) + .as_ref(), + seqlen_offset, + )?; + } + let xs = xs.apply(&self.norm)?; + let xs = xs.i((.., seq_len - 1, ..))?.contiguous()?; + let logits = self.lm_head.forward(&xs)?; + logits.to_dtype(DType::F32) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache(); + } + } +} diff --git a/patches/candle-transformers/src/models/depth_anything_v2.rs b/patches/candle-transformers/src/models/depth_anything_v2.rs new file mode 100644 index 0000000000..690d396bdc --- /dev/null +++ b/patches/candle-transformers/src/models/depth_anything_v2.rs @@ -0,0 +1,570 @@ +//! Implementation of the Depth Anything model from FAIR. +//! +//! See: +//! - ["Depth Anything: Unleashing the Power of Large-Scale Unlabeled Data"](https://github.com/LiheYoung/Depth-Anything) +//! + +use std::sync::Arc; + +use candle::D::Minus1; +use candle::{Module, Result, Tensor}; +use candle_nn::ops::Identity; +use candle_nn::{ + batch_norm, conv2d, conv2d_no_bias, conv_transpose2d, linear, seq, Activation, BatchNorm, + BatchNormConfig, Conv2d, Conv2dConfig, ConvTranspose2dConfig, Sequential, VarBuilder, +}; + +use crate::models::dinov2::DinoVisionTransformer; + +pub struct DepthAnythingV2Config { + out_channel_sizes: [usize; 4], + in_channel_size: usize, // embed_dim in the Dino model + num_features: usize, + use_batch_norm: bool, + use_class_token: bool, + layer_ids_vits: Vec, + input_image_size: usize, + target_patch_size: usize, +} + +impl DepthAnythingV2Config { + #[allow(clippy::too_many_arguments)] + pub fn new( + out_channel_sizes: [usize; 4], + in_channel_size: usize, + num_features: usize, + use_batch_norm: bool, + use_class_token: bool, + layer_ids_vits: Vec, + input_image_size: usize, + target_patch_size: usize, + ) -> Self { + Self { + out_channel_sizes, + in_channel_size, + num_features, + use_batch_norm, + use_class_token, + layer_ids_vits, + input_image_size, + target_patch_size, + } + } + + pub fn vit_small() -> Self { + Self { + out_channel_sizes: [48, 96, 192, 384], + in_channel_size: 384, + num_features: 64, + use_batch_norm: false, + use_class_token: false, + layer_ids_vits: vec![2, 5, 8, 11], + input_image_size: 518, + target_patch_size: 518 / 14, + } + } + + pub fn vit_base() -> Self { + Self { + out_channel_sizes: [96, 192, 384, 768], + in_channel_size: 768, + num_features: 128, + use_batch_norm: false, + use_class_token: false, + layer_ids_vits: vec![2, 5, 8, 11], + input_image_size: 518, + target_patch_size: 518 / 14, + } + } + + pub fn vit_large() -> Self { + Self { + out_channel_sizes: [256, 512, 1024, 1024], + in_channel_size: 1024, + num_features: 256, + use_batch_norm: false, + use_class_token: false, + layer_ids_vits: vec![4, 11, 17, 23], + input_image_size: 518, + target_patch_size: 518 / 14, + } + } + + pub fn vit_giant() -> Self { + Self { + out_channel_sizes: [1536, 1536, 1536, 1536], + in_channel_size: 1536, + num_features: 384, + use_batch_norm: false, + use_class_token: false, + layer_ids_vits: vec![9, 19, 29, 39], + input_image_size: 518, + target_patch_size: 518 / 14, + } + } +} + +pub struct ResidualConvUnit { + activation: Activation, + conv1: Conv2d, + conv2: Conv2d, + batch_norm1: Option, + batch_norm2: Option, +} + +impl ResidualConvUnit { + pub fn new( + conf: &DepthAnythingV2Config, + activation: Activation, + vb: VarBuilder, + ) -> Result { + const KERNEL_SIZE: usize = 3; + let conv_cfg = Conv2dConfig { + padding: 1, + stride: 1, + dilation: 1, + groups: 1, + cudnn_fwd_algo: None, + }; + let conv1 = conv2d( + conf.num_features, + conf.num_features, + KERNEL_SIZE, + conv_cfg, + vb.pp("conv1"), + )?; + let conv2 = conv2d( + conf.num_features, + conf.num_features, + KERNEL_SIZE, + conv_cfg, + vb.pp("conv2"), + )?; + + let (batch_norm1, batch_norm2) = match conf.use_batch_norm { + true => { + let batch_norm_cfg = BatchNormConfig { + eps: 1e-05, + remove_mean: false, + affine: true, + momentum: 0.1, + }; + ( + Some(batch_norm(conf.num_features, batch_norm_cfg, vb.pp("bn1"))?), + Some(batch_norm(conf.num_features, batch_norm_cfg, vb.pp("bn2"))?), + ) + } + false => (None, None), + }; + + Ok(Self { + activation, + conv1, + conv2, + batch_norm1, + batch_norm2, + }) + } +} + +impl Module for ResidualConvUnit { + fn forward(&self, xs: &Tensor) -> Result { + let out = self.activation.forward(xs)?; + let out = self.conv1.forward(&out)?; + let out = if let Some(batch_norm1) = &self.batch_norm1 { + batch_norm1.forward_train(&out)? + } else { + out + }; + + let out = self.activation.forward(&out)?; + let out = self.conv2.forward(&out)?; + let out = if let Some(batch_norm2) = &self.batch_norm2 { + batch_norm2.forward_train(&out)? + } else { + out + }; + + out + xs + } +} + +pub struct FeatureFusionBlock { + res_conv_unit1: ResidualConvUnit, + res_conv_unit2: ResidualConvUnit, + output_conv: Conv2d, + target_patch_size: usize, +} + +impl FeatureFusionBlock { + pub fn new( + conf: &DepthAnythingV2Config, + target_patch_size: usize, + activation: Activation, + vb: VarBuilder, + ) -> Result { + const KERNEL_SIZE: usize = 1; + let conv_cfg = Conv2dConfig { + padding: 0, + stride: 1, + dilation: 1, + groups: 1, + cudnn_fwd_algo: None, + }; + let output_conv = conv2d( + conf.num_features, + conf.num_features, + KERNEL_SIZE, + conv_cfg, + vb.pp("out_conv"), + )?; + let res_conv_unit1 = ResidualConvUnit::new(conf, activation, vb.pp("resConfUnit1"))?; + let res_conv_unit2 = ResidualConvUnit::new(conf, activation, vb.pp("resConfUnit2"))?; + + Ok(Self { + res_conv_unit1, + res_conv_unit2, + output_conv, + target_patch_size, + }) + } +} + +impl Module for FeatureFusionBlock { + fn forward(&self, xs: &Tensor) -> Result { + let out = self.res_conv_unit2.forward(xs)?; + let out = out.interpolate2d(self.target_patch_size, self.target_patch_size)?; + + self.output_conv.forward(&out) + } +} + +pub struct Scratch { + layer1_rn: Conv2d, + layer2_rn: Conv2d, + layer3_rn: Conv2d, + layer4_rn: Conv2d, + refine_net1: FeatureFusionBlock, + refine_net2: FeatureFusionBlock, + refine_net3: FeatureFusionBlock, + refine_net4: FeatureFusionBlock, + output_conv1: Conv2d, + output_conv2: Sequential, +} + +impl Scratch { + pub fn new(conf: &DepthAnythingV2Config, vb: VarBuilder) -> Result { + const KERNEL_SIZE: usize = 3; + let conv_cfg = Conv2dConfig { + padding: 1, + stride: 1, + dilation: 1, + groups: 1, + cudnn_fwd_algo: None, + }; + + let layer1_rn = conv2d_no_bias( + conf.out_channel_sizes[0], + conf.num_features, + KERNEL_SIZE, + conv_cfg, + vb.pp("layer1_rn"), + )?; + let layer2_rn = conv2d_no_bias( + conf.out_channel_sizes[1], + conf.num_features, + KERNEL_SIZE, + conv_cfg, + vb.pp("layer2_rn"), + )?; + let layer3_rn = conv2d_no_bias( + conf.out_channel_sizes[2], + conf.num_features, + KERNEL_SIZE, + conv_cfg, + vb.pp("layer3_rn"), + )?; + let layer4_rn = conv2d_no_bias( + conf.out_channel_sizes[3], + conf.num_features, + KERNEL_SIZE, + conv_cfg, + vb.pp("layer4_rn"), + )?; + + let refine_net1 = FeatureFusionBlock::new( + conf, + conf.target_patch_size * 8, + Activation::Relu, + vb.pp("refinenet1"), + )?; + let refine_net2 = FeatureFusionBlock::new( + conf, + conf.target_patch_size * 4, + Activation::Relu, + vb.pp("refinenet2"), + )?; + let refine_net3 = FeatureFusionBlock::new( + conf, + conf.target_patch_size * 2, + Activation::Relu, + vb.pp("refinenet3"), + )?; + let refine_net4 = FeatureFusionBlock::new( + conf, + conf.target_patch_size, + Activation::Relu, + vb.pp("refinenet4"), + )?; + + let conv_cfg = Conv2dConfig { + padding: 1, + stride: 1, + dilation: 1, + groups: 1, + cudnn_fwd_algo: None, + }; + let output_conv1 = conv2d( + conf.num_features, + conf.num_features / 2, + KERNEL_SIZE, + conv_cfg, + vb.pp("output_conv1"), + )?; + + let output_conv2 = seq(); + const HEAD_FEATURES_2: usize = 32; + const OUT_CHANNELS_2: usize = 1; + const KERNEL_SIZE_2: usize = 1; + let output_conv2 = output_conv2.add(conv2d( + conf.num_features / 2, + HEAD_FEATURES_2, + KERNEL_SIZE, + conv_cfg, + vb.pp("output_conv2").pp("0"), + )?); + let output_conv2 = output_conv2 + .add(Activation::Relu) + .add(conv2d( + HEAD_FEATURES_2, + OUT_CHANNELS_2, + KERNEL_SIZE_2, + conv_cfg, + vb.pp("output_conv2").pp("2"), + )?) + .add(Activation::Relu); + + Ok(Self { + layer1_rn, + layer2_rn, + layer3_rn, + layer4_rn, + refine_net1, + refine_net2, + refine_net3, + refine_net4, + output_conv1, + output_conv2, + }) + } +} + +const NUM_CHANNELS: usize = 4; + +pub struct DPTHead { + projections: Vec, + resize_layers: Vec>, + readout_projections: Vec, + scratch: Scratch, + use_class_token: bool, + input_image_size: usize, + target_patch_size: usize, +} + +impl DPTHead { + pub fn new(conf: &DepthAnythingV2Config, vb: VarBuilder) -> Result { + let mut projections: Vec = Vec::with_capacity(conf.out_channel_sizes.len()); + for (conv_index, out_channel_size) in conf.out_channel_sizes.iter().enumerate() { + projections.push(conv2d( + conf.in_channel_size, + *out_channel_size, + 1, + Default::default(), + vb.pp("projects").pp(conv_index.to_string()), + )?); + } + + let resize_layers: Vec> = vec![ + Box::new(conv_transpose2d( + conf.out_channel_sizes[0], + conf.out_channel_sizes[0], + 4, + ConvTranspose2dConfig { + padding: 0, + stride: 4, + dilation: 1, + output_padding: 0, + }, + vb.pp("resize_layers").pp("0"), + )?), + Box::new(conv_transpose2d( + conf.out_channel_sizes[1], + conf.out_channel_sizes[1], + 2, + ConvTranspose2dConfig { + padding: 0, + stride: 2, + dilation: 1, + output_padding: 0, + }, + vb.pp("resize_layers").pp("1"), + )?), + Box::new(Identity::new()), + Box::new(conv2d( + conf.out_channel_sizes[3], + conf.out_channel_sizes[3], + 3, + Conv2dConfig { + padding: 1, + stride: 2, + dilation: 1, + groups: 1, + cudnn_fwd_algo: None, + }, + vb.pp("resize_layers").pp("3"), + )?), + ]; + + let readout_projections = if conf.use_class_token { + let rop = Vec::with_capacity(NUM_CHANNELS); + for rop_index in 0..NUM_CHANNELS { + seq() + .add(linear( + 2 * conf.in_channel_size, + conf.in_channel_size, + vb.pp("readout_projects").pp(rop_index.to_string()), + )?) + .add(Activation::Gelu); + } + rop + } else { + vec![] + }; + + let scratch = Scratch::new(conf, vb.pp("scratch"))?; + + Ok(Self { + projections, + resize_layers, + readout_projections, + scratch, + use_class_token: conf.use_class_token, + input_image_size: conf.input_image_size, + target_patch_size: conf.target_patch_size, + }) + } +} + +impl Module for DPTHead { + fn forward(&self, xs: &Tensor) -> Result { + let mut out: Vec = Vec::with_capacity(NUM_CHANNELS); + for i in 0..NUM_CHANNELS { + let x = if self.use_class_token { + let x = xs.get(i)?.get(0)?; + let class_token = xs.get(i)?.get(1)?; + let readout = class_token.unsqueeze(1)?.expand(x.shape())?; + let to_cat = [x, readout]; + let cat = Tensor::cat(&to_cat, Minus1)?; + self.readout_projections[i].forward(&cat)? + } else { + xs.get(i)? + }; + let x_dims = x.dims(); + + let x = x.permute((0, 2, 1))?.reshape(( + x_dims[0], + x_dims[x_dims.len() - 1], + self.target_patch_size, + self.target_patch_size, + ))?; + let x = self.projections[i].forward(&x)?; + + let x = self.resize_layers[i].forward(&x)?; + out.push(x); + } + + let layer_1_rn = self.scratch.layer1_rn.forward(&out[0])?; + let layer_2_rn = self.scratch.layer2_rn.forward(&out[1])?; + let layer_3_rn = self.scratch.layer3_rn.forward(&out[2])?; + let layer_4_rn = self.scratch.layer4_rn.forward(&out[3])?; + + let path4 = self.scratch.refine_net4.forward(&layer_4_rn)?; + + let res3_out = self + .scratch + .refine_net3 + .res_conv_unit1 + .forward(&layer_3_rn)?; + let res3_out = path4.add(&res3_out)?; + let path3 = self.scratch.refine_net3.forward(&res3_out)?; + + let res2_out = self + .scratch + .refine_net2 + .res_conv_unit1 + .forward(&layer_2_rn)?; + let res2_out = path3.add(&res2_out)?; + let path2 = self.scratch.refine_net2.forward(&res2_out)?; + + let res1_out = self + .scratch + .refine_net1 + .res_conv_unit1 + .forward(&layer_1_rn)?; + let res1_out = path2.add(&res1_out)?; + let path1 = self.scratch.refine_net1.forward(&res1_out)?; + + let out = self.scratch.output_conv1.forward(&path1)?; + + let out = out.interpolate2d(self.input_image_size, self.input_image_size)?; + + self.scratch.output_conv2.forward(&out) + } +} + +pub struct DepthAnythingV2 { + pretrained: Arc, + depth_head: DPTHead, + conf: DepthAnythingV2Config, +} + +impl DepthAnythingV2 { + pub fn new( + pretrained: Arc, + conf: DepthAnythingV2Config, + vb: VarBuilder, + ) -> Result { + let depth_head = DPTHead::new(&conf, vb.pp("depth_head"))?; + + Ok(Self { + pretrained, + depth_head, + conf, + }) + } +} + +impl Module for DepthAnythingV2 { + fn forward(&self, xs: &Tensor) -> Result { + let features = self.pretrained.get_intermediate_layers( + xs, + &self.conf.layer_ids_vits, + false, + false, + true, + )?; + let depth = self.depth_head.forward(&features)?; + + depth.relu() + } +} diff --git a/patches/candle-transformers/src/models/dinov2.rs b/patches/candle-transformers/src/models/dinov2.rs new file mode 100644 index 0000000000..6dd0ab2dad --- /dev/null +++ b/patches/candle-transformers/src/models/dinov2.rs @@ -0,0 +1,396 @@ +//! Implementation of the DINOv2 models from Meta Research. +//! +//! This module implements the DINOv2 vision transformer model from Meta AI Research. +//! DINOv2 is a self-supervised learning model that can learn visual features +//! without using any labeled data. See: ["DINOv2: Learning Robust Visual Features without Supervision"](https://github.com/facebookresearch/dinov2) +//! +//! ## Running an example with color map and CUDA +//! +//! ```bash +//! cargo run \ +//! --features cuda,depth_anything_v2 \ +//! --package candle-examples \ +//! --example depth_anything_v2 \ +//! -- --color-map \ +//! --image candle-examples/examples/yolo-v8/assets/bike.jpg +//! ``` +//! +//! ## Running as an ImageNet classifier +//! +//! The model returns the probability for the image to belong to each of the 1000 ImageNet categories. +//! +//!
+//! +//!
+//! +//! ```bash +//! cargo run \ +//! --example dinov2 \ +//! --release \ +//! -- --image candle-examples/examples/yolo-v8/assets/bike.jpg +//! +//! > mountain bike, all-terrain bike, off-roader: 43.67% +//! > bicycle-built-for-two, tandem bicycle, tandem: 33.20% +//! > crash helmet : 13.23% +//! > unicycle, monocycle : 2.44% +//! > maillot : 2.42% +//! ``` +//! + +use candle::{IndexOp, Result, Tensor, D}; +use candle_nn::{layer_norm, LayerNorm, Linear, Module, VarBuilder}; + +const IMG_SIZE: usize = 518; +const PATCH_SIZE: usize = 14; +const NUM_CLASSES: usize = 1000; + +fn linear(vb: VarBuilder, in_dim: usize, out_dim: usize, bias: bool) -> Result { + if bias { + candle_nn::linear(in_dim, out_dim, vb) + } else { + candle_nn::linear_no_bias(in_dim, out_dim, vb) + } +} + +#[derive(Debug)] +struct Attention { + qkv: Linear, + proj: Linear, + num_heads: usize, + scale: f64, +} + +impl Attention { + fn new( + vb: VarBuilder, + dim: usize, + num_heads: usize, + qkv_bias: bool, + proj_bias: bool, + ) -> Result { + let qkv = linear(vb.pp("qkv"), dim, dim * 3, qkv_bias)?; + let proj = linear(vb.pp("proj"), dim, dim, proj_bias)?; + let scale = 1. / ((dim / num_heads) as f64).sqrt(); + Ok(Self { + qkv, + proj, + num_heads, + scale, + }) + } +} + +impl Module for Attention { + fn forward(&self, xs: &Tensor) -> Result { + let (b, n, c) = xs.dims3()?; + let qkv = self + .qkv + .forward(xs)? + .reshape((b, n, 3, self.num_heads, c / self.num_heads))? + .transpose(1, 2)? // 02134 + .transpose(0, 1)? // 20134 + .transpose(2, 3)?; // 20314 + let q = (qkv.i(0)? * self.scale)?; + let k = qkv.i(1)?.contiguous()?; + let v = qkv.i(2)?.contiguous()?; + let attn = candle_nn::ops::softmax(&q.matmul(&k.t()?)?, D::Minus1)?; + let attn = attn.matmul(&v)?.transpose(1, 2)?.reshape((b, n, c))?; + self.proj.forward(&attn) + } +} + +#[derive(Debug)] +struct LayerScale { + gamma: Tensor, +} + +impl LayerScale { + fn new(vb: VarBuilder, dim: usize) -> Result { + let gamma = vb.get(dim, "gamma")?; + Ok(Self { gamma }) + } +} + +impl Module for LayerScale { + fn forward(&self, xs: &Tensor) -> Result { + xs.broadcast_mul(&self.gamma) + } +} + +#[derive(Debug)] +struct Mlp { + fc1: Linear, + fc2: Linear, +} + +impl Mlp { + fn new(vb: VarBuilder, in_features: usize, hidden_features: usize, bias: bool) -> Result { + let out_features = in_features; + let fc1 = linear(vb.pp("fc1"), in_features, hidden_features, bias)?; + let fc2 = linear(vb.pp("fc2"), hidden_features, out_features, bias)?; + Ok(Self { fc1, fc2 }) + } +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.fc1.forward(xs)?.gelu()?; + self.fc2.forward(&xs) + } +} + +#[derive(Debug)] +struct Block { + norm1: LayerNorm, + attn: Attention, + ls1: LayerScale, + norm2: LayerNorm, + mlp: Mlp, + ls2: LayerScale, +} + +impl Block { + fn new(vb: VarBuilder, dim: usize, num_heads: usize) -> Result { + let norm1 = layer_norm(dim, 1e-5, vb.pp("norm1"))?; + let attn = Attention::new(vb.pp("attn"), dim, num_heads, true, true)?; + let ls1 = LayerScale::new(vb.pp("ls1"), dim)?; + let norm2 = layer_norm(dim, 1e-5, vb.pp("norm2"))?; + let mlp = Mlp::new(vb.pp("mlp"), dim, dim * 4, true)?; + let ls2 = LayerScale::new(vb.pp("ls2"), dim)?; + Ok(Self { + norm1, + attn, + ls1, + norm2, + mlp, + ls2, + }) + } +} + +impl Module for Block { + fn forward(&self, xs: &Tensor) -> Result { + let residual = xs; + let xs = self + .ls1 + .forward(&self.attn.forward(&self.norm1.forward(xs)?)?)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = self + .ls2 + .forward(&self.mlp.forward(&self.norm2.forward(&xs)?)?)?; + xs + residual + } +} + +#[derive(Debug)] +struct PatchEmbed { + proj: candle_nn::Conv2d, + patch_size: (usize, usize), + num_patches: usize, +} + +impl PatchEmbed { + fn new( + vb: VarBuilder, + img_size: usize, + patch_size: usize, + in_chans: usize, + embed_dim: usize, + ) -> Result { + let config = candle_nn::Conv2dConfig { + stride: patch_size, + ..Default::default() + }; + let proj = candle_nn::conv2d(in_chans, embed_dim, patch_size, config, vb.pp("proj"))?; + let num_patches = (img_size / patch_size) * (img_size / patch_size); + Ok(Self { + proj, + patch_size: (patch_size, patch_size), + num_patches, + }) + } +} + +impl Module for PatchEmbed { + fn forward(&self, xs: &Tensor) -> Result { + let (_b, _c, h, w) = xs.dims4()?; + let (patch_h, patch_w) = self.patch_size; + if (h % patch_h) != 0 { + candle::bail!("image height {h} is not a multiple of patch height {patch_h}") + } + if (w % patch_w) != 0 { + candle::bail!("image width {w} is not a multiple of patch width {patch_w}") + } + let xs = self.proj.forward(xs)?; + let (b, c, h, w) = xs.dims4()?; + // flatten embeddings. + xs.reshape((b, c, h * w))?.transpose(1, 2) + } +} + +#[derive(Debug)] +pub struct DinoVisionTransformer { + patch_embed: PatchEmbed, + cls_token: Tensor, + pos_embed: Tensor, + blocks: Vec, + norm: LayerNorm, + head: Linear, +} + +impl DinoVisionTransformer { + pub fn new(vb: VarBuilder, depth: usize, embed_dim: usize, num_heads: usize) -> Result { + let patch_embed = + PatchEmbed::new(vb.pp("patch_embed"), IMG_SIZE, PATCH_SIZE, 3, embed_dim)?; + let cls_token = vb.get((1, 1, embed_dim), "cls_token")?; + let num_tokens = 1; + let pos_embed = vb.get( + (1, patch_embed.num_patches + num_tokens, embed_dim), + "pos_embed", + )?; + let head = linear(vb.pp("head"), 2 * embed_dim, NUM_CLASSES, true)?; + let norm = layer_norm(embed_dim, 1e-5, vb.pp("norm"))?; + let vb_b = vb.pp("blocks"); + let blocks = (0..depth) + .map(|i| Block::new(vb_b.pp(i.to_string()), embed_dim, num_heads)) + .collect::>>()?; + Ok(Self { + patch_embed, + cls_token, + pos_embed, + blocks, + norm, + head, + }) + } + + fn interpolate_pos_encoding(&self, xs: &Tensor, w: usize, h: usize) -> Result { + let npatch = xs.dim(1)? - 1; + let n = self.pos_embed.dim(1)? - 1; + let sqrt_n = (n as f64).sqrt(); + if npatch == n && w == h { + return Ok(self.pos_embed.clone()); + } + let class_pos_embed = self.pos_embed.i((.., ..1))?; + let patch_pos_embed = self.pos_embed.i((.., 1..))?; + let dim = xs.dim(D::Minus1)?; + let (w0, h0) = ((w / PATCH_SIZE) as f64 + 0.1, (h / PATCH_SIZE) as f64 + 0.1); + let patch_pos_embed = patch_pos_embed + .reshape((1, sqrt_n as usize, sqrt_n as usize, dim))? + .transpose(2, 3)? + .transpose(1, 2)?; + // This uses bicubic interpolation in the original implementation. + let patch_pos_embed = patch_pos_embed.upsample_nearest2d(h0 as usize, w0 as usize)?; + let el_count = patch_pos_embed.shape().elem_count(); + let patch_pos_embed = + patch_pos_embed + .transpose(1, 2)? + .transpose(2, 3)? + .reshape((1, el_count / dim, dim))?; + Tensor::cat(&[&class_pos_embed, &patch_pos_embed], 1) + } + + fn prepare_tokens_with_mask(&self, xs: &Tensor) -> Result { + let (_b, _nc, w, h) = xs.dims4()?; + let xs = self.patch_embed.forward(xs)?; + let xs = Tensor::cat(&[&self.cls_token, &xs], 1)?; + &xs + &self.interpolate_pos_encoding(&xs, w, h)? + } + + fn get_intermediate_layers_not_chunked( + &self, + xs: &Tensor, + blocks_to_take: &[usize], + ) -> Result> { + let mut xs = self.prepare_tokens_with_mask(xs)?; + let mut output = Vec::new(); + for (i, blk) in self.blocks.iter().enumerate() { + xs = blk.forward(&xs)?; + if blocks_to_take.contains(&i) { + output.push(xs.clone()); + } + } + if output.len() != blocks_to_take.len() { + candle::bail!( + "only {} / {} blocks found", + output.len(), + blocks_to_take.len() + ); + } + Ok(output) + } + + pub fn get_intermediate_layers( + &self, + xs: &Tensor, + blocks_to_take: &[usize], + reshape: bool, + return_class_token: bool, + norm: bool, + ) -> Result { + let outputs = self.get_intermediate_layers_not_chunked(xs, blocks_to_take)?; + let outputs = if norm { + outputs + .iter() + .map(|out| self.norm.forward(out)) + .collect::>>()? + } else { + outputs + }; + let class_tokens = outputs + .iter() + .map(|out| out.i((.., 0))) + .collect::>>()?; + let outputs = outputs + .iter() + .map(|out| out.i((.., 1..))) + .collect::>>()?; + + let outputs = if reshape { + let (b, _c, w, h) = xs.dims4()?; + let patch_size = self.patch_embed.patch_size.0; + let num_channels = outputs[0].elem_count() / (b * (w / patch_size) * (h / patch_size)); + outputs + .iter() + .map(|out| { + out.reshape((b, w / patch_size, h / patch_size, num_channels))? + .transpose(2, 3)? + .transpose(1, 2) + }) + .collect::>>()? + } else { + outputs + }; + + let outputs = if return_class_token { + outputs + .iter() + .zip(class_tokens.iter()) + .map(|(out, class_token)| Tensor::cat(&[out, class_token], D::Minus1)) + .collect::>>()? + } else { + outputs + }; + + Tensor::stack(&outputs[..], 0) + } +} + +impl Module for DinoVisionTransformer { + fn forward(&self, xs: &Tensor) -> Result { + let mut xs = self.prepare_tokens_with_mask(xs)?; + for blk in self.blocks.iter() { + xs = blk.forward(&xs)? + } + let xs = self.norm.forward(&xs)?; + let xs_norm_clstoken = xs.i((.., 0))?; + let xs_norm_patchtokens = xs.i((.., 1..))?.mean(1)?; + let xs = Tensor::cat(&[xs_norm_clstoken, xs_norm_patchtokens], D::Minus1)?; + self.head.forward(&xs) + } +} + +pub fn vit_small(vb: VarBuilder) -> Result { + DinoVisionTransformer::new(vb, 12, 384, 6) +} diff --git a/patches/candle-transformers/src/models/dinov2reg4.rs b/patches/candle-transformers/src/models/dinov2reg4.rs new file mode 100644 index 0000000000..549f2c3ce5 --- /dev/null +++ b/patches/candle-transformers/src/models/dinov2reg4.rs @@ -0,0 +1,313 @@ +//! Implementation of the DINOv2 revision (4 regularization) +//! +//! The DINOv2-reg4 model is a variant of DINOv2 that adds 4 regularization tokens to the +//! original architecture. This implementation is specifically trained for plant species +//! classification on the PlantCLEF2024 dataset with 7,806 classes. +//! +//! - [Paper](https://arxiv.org/abs/2309.16588). DINOv2: Learning Robust Visual Features without Supervision +//! - [GH Repo](https://github.com/facebookresearch/dinov2) +//! +//! # Example +//! +//! ```bash +//! # Download classes names and a plant picture to identify +//! # see candle/examples/dinov2reg4 for full code. +//! +//! # Perform inference +//! cargo run \ +//! --example dinov2reg4 \ +//! --release -- \ +//! --image +//! +//! > Orchis simia Lam. : 45.55% +//! > Orchis × bergonii Nanteuil: 9.80% +//! > Orchis italica Poir. : 9.66% +//! > Orchis × angusticruris Franch.: 2.76% +//! > Orchis × bivonae Tod. : 2.54% +//! ``` +//! +//!
+//! +//!
+//! +use candle::{IndexOp, Result, Tensor, D}; +use candle_nn::{layer_norm, LayerNorm, Linear, Module, VarBuilder}; + +const IMG_SIZE: usize = 518; +const PATCH_SIZE: usize = 14; +const NUM_CLASSES: usize = 7806; // PlantCLEF2024 DINOv2 (https://zenodo.org/records/10848263) + +fn linear(vb: VarBuilder, in_dim: usize, out_dim: usize, bias: bool) -> Result { + if bias { + candle_nn::linear(in_dim, out_dim, vb) + } else { + candle_nn::linear_no_bias(in_dim, out_dim, vb) + } +} + +#[derive(Debug)] +struct Attention { + qkv: Linear, + proj: Linear, + num_heads: usize, + scale: f64, +} + +impl Attention { + fn new( + vb: VarBuilder, + dim: usize, + num_heads: usize, + qkv_bias: bool, + proj_bias: bool, + ) -> Result { + let qkv = linear(vb.pp("qkv"), dim, dim * 3, qkv_bias)?; + let proj = linear(vb.pp("proj"), dim, dim, proj_bias)?; + let scale = 1. / ((dim / num_heads) as f64).sqrt(); + Ok(Self { + qkv, + proj, + num_heads, + scale, + }) + } +} + +impl Module for Attention { + fn forward(&self, xs: &Tensor) -> Result { + let (b, n, c) = xs.dims3()?; + let qkv = self + .qkv + .forward(xs)? + .reshape((b, n, 3, self.num_heads, c / self.num_heads))? + .transpose(1, 2)? // 02134 + .transpose(0, 1)? // 20134 + .transpose(2, 3)?; // 20314 + let q = (qkv.i(0)? * self.scale)?; + let k = qkv.i(1)?.contiguous()?; + let v = qkv.i(2)?.contiguous()?; + let attn = candle_nn::ops::softmax(&q.matmul(&k.t()?)?, D::Minus1)?; + let attn = attn.matmul(&v)?.transpose(1, 2)?.reshape((b, n, c))?; + self.proj.forward(&attn) + } +} + +#[derive(Debug)] +struct LayerScale { + gamma: Tensor, +} + +impl LayerScale { + fn new(vb: VarBuilder, dim: usize) -> Result { + let gamma = vb.get(dim, "gamma")?; + Ok(Self { gamma }) + } +} + +impl Module for LayerScale { + fn forward(&self, xs: &Tensor) -> Result { + xs.broadcast_mul(&self.gamma) + } +} + +#[derive(Debug)] +struct Mlp { + fc1: Linear, + fc2: Linear, +} + +impl Mlp { + fn new(vb: VarBuilder, in_features: usize, hidden_features: usize, bias: bool) -> Result { + let out_features = in_features; + let fc1 = linear(vb.pp("fc1"), in_features, hidden_features, bias)?; + let fc2 = linear(vb.pp("fc2"), hidden_features, out_features, bias)?; + Ok(Self { fc1, fc2 }) + } +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.fc1.forward(xs)?.gelu()?; + self.fc2.forward(&xs) + } +} + +#[derive(Debug)] +struct Block { + norm1: LayerNorm, + attn: Attention, + ls1: LayerScale, + norm2: LayerNorm, + mlp: Mlp, + ls2: LayerScale, +} + +impl Block { + fn new(vb: VarBuilder, dim: usize, num_heads: usize) -> Result { + let norm1 = layer_norm(dim, 1e-6, vb.pp("norm1"))?; + let attn = Attention::new(vb.pp("attn"), dim, num_heads, true, true)?; + let ls1 = LayerScale::new(vb.pp("ls1"), dim)?; + let norm2 = layer_norm(dim, 1e-6, vb.pp("norm2"))?; + let mlp = Mlp::new(vb.pp("mlp"), dim, dim * 4, true)?; + let ls2 = LayerScale::new(vb.pp("ls2"), dim)?; + Ok(Self { + norm1, + attn, + ls1, + norm2, + mlp, + ls2, + }) + } +} + +impl Module for Block { + fn forward(&self, xs: &Tensor) -> Result { + let residual = xs; + let xs = self + .ls1 + .forward(&self.attn.forward(&self.norm1.forward(xs)?)?)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = self + .ls2 + .forward(&self.mlp.forward(&self.norm2.forward(&xs)?)?)?; + xs + residual + } +} + +#[derive(Debug)] +struct PatchEmbed { + proj: candle_nn::Conv2d, + patch_size: (usize, usize), + num_patches: usize, +} + +impl PatchEmbed { + fn new( + vb: VarBuilder, + img_size: usize, + patch_size: usize, + in_chans: usize, + embed_dim: usize, + ) -> Result { + let config = candle_nn::Conv2dConfig { + stride: patch_size, + ..Default::default() + }; + let proj = candle_nn::conv2d(in_chans, embed_dim, patch_size, config, vb.pp("proj"))?; + let num_patches = (img_size / patch_size) * (img_size / patch_size); + Ok(Self { + proj, + patch_size: (patch_size, patch_size), + num_patches, + }) + } +} + +impl Module for PatchEmbed { + fn forward(&self, xs: &Tensor) -> Result { + let (_b, _c, h, w) = xs.dims4()?; + let (patch_h, patch_w) = self.patch_size; + if (h % patch_h) != 0 { + candle::bail!("image height {h} is not a multiple of patch height {patch_h}") + } + if (w % patch_w) != 0 { + candle::bail!("image width {w} is not a multiple of patch width {patch_w}") + } + let xs = self.proj.forward(xs)?; + let (b, c, h, w) = xs.dims4()?; + // flatten embeddings. + xs.reshape((b, c, h * w))?.transpose(1, 2) + } +} + +#[derive(Debug)] +pub struct DinoVisionTransformer { + patch_embed: PatchEmbed, + cls_token: Tensor, + reg_token: Tensor, + pos_embed: Tensor, + blocks: Vec, + norm: LayerNorm, + head: Linear, +} + +impl DinoVisionTransformer { + pub fn new(vb: VarBuilder, depth: usize, embed_dim: usize, num_heads: usize) -> Result { + let patch_embed = + PatchEmbed::new(vb.pp("patch_embed"), IMG_SIZE, PATCH_SIZE, 3, embed_dim)?; + let cls_token = vb.get((1, 1, embed_dim), "cls_token")?; + let reg_token = vb.get((1, 4, embed_dim), "reg_token")?; + let pos_embed = vb.get((1, patch_embed.num_patches, embed_dim), "pos_embed")?; + let head = linear(vb.pp("head"), embed_dim, NUM_CLASSES, true)?; + let norm = layer_norm(embed_dim, 1e-6, vb.pp("norm"))?; + let vb_b = vb.pp("blocks"); + let blocks = (0..depth) + .map(|i| Block::new(vb_b.pp(i.to_string()), embed_dim, num_heads)) + .collect::>>()?; + Ok(Self { + patch_embed, + cls_token, + reg_token, + pos_embed, + blocks, + norm, + head, + }) + } + + fn interpolate_pos_encoding(&self, xs: &Tensor, w: usize, h: usize) -> Result { + let npatch = xs.dim(1)? - 1; + let n = self.pos_embed.dim(1)? - 1; + let sqrt_n = (n as f64).sqrt(); + if npatch == n && w == h { + return Ok(self.pos_embed.clone()); + } + let patch_pos_embed = &self.pos_embed; + let dim = xs.dim(D::Minus1)?; + let (w0, h0) = ((w / PATCH_SIZE) as f64 + 0.1, (h / PATCH_SIZE) as f64 + 0.1); + let patch_pos_embed = patch_pos_embed + .reshape((1, sqrt_n as usize, sqrt_n as usize, dim))? + .transpose(2, 3)? + .transpose(1, 2)?; + // This uses bicubic interpolation in the original implementation. + let patch_pos_embed = patch_pos_embed.upsample_nearest2d(h0 as usize, w0 as usize)?; + let el_count = patch_pos_embed.shape().elem_count(); + patch_pos_embed + .transpose(1, 2)? + .transpose(2, 3)? + .reshape((1, el_count / dim, dim)) + } + + fn prepare_tokens_with_mask(&self, xs: &Tensor) -> Result { + let (_b, _nc, w, h) = xs.dims4()?; + if (w != IMG_SIZE) || (h != IMG_SIZE) { + panic!("Error: The input tensor should have the shape: Bx3x518x518."); + } + let xs = self.patch_embed.forward(xs)?; + let xs = (&xs + &self.interpolate_pos_encoding(&xs, w, h)?)?; + let xs = Tensor::cat(&[&self.cls_token, &self.reg_token, &xs], 1)?; + Ok(xs) + } +} + +impl Module for DinoVisionTransformer { + fn forward(&self, xs: &Tensor) -> Result { + let mut xs = self.prepare_tokens_with_mask(xs)?; + for blk in self.blocks.iter() { + xs = blk.forward(&xs)? + } + let xs = self.norm.forward(&xs)?; + let xs_norm_clstoken = xs.i((.., 0))?; + self.head.forward(&xs_norm_clstoken) + } +} + +pub fn vit_small(vb: VarBuilder) -> Result { + DinoVisionTransformer::new(vb, 12, 384, 6) +} + +pub fn vit_base(vb: VarBuilder) -> Result { + DinoVisionTransformer::new(vb, 12, 768, 12) +} diff --git a/patches/candle-transformers/src/models/distilbert.rs b/patches/candle-transformers/src/models/distilbert.rs new file mode 100644 index 0000000000..abaffa81fb --- /dev/null +++ b/patches/candle-transformers/src/models/distilbert.rs @@ -0,0 +1,451 @@ +//! Implementation of DistilBert, a distilled version of BERT. +//! +//! See: +//! - ["DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter"](https://arxiv.org/abs/1910.01108) +//! +use super::with_tracing::{layer_norm, linear, LayerNorm, Linear}; +use candle::{DType, Device, Result, Tensor}; +use candle_nn::{Embedding, Module, VarBuilder}; +use serde::Deserialize; + +pub const DTYPE: DType = DType::F32; + +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) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum HiddenAct { + Gelu, + Relu, +} + +struct HiddenActLayer { + act: HiddenAct, + span: tracing::Span, +} + +impl HiddenActLayer { + fn new(act: HiddenAct) -> Self { + let span = tracing::span!(tracing::Level::TRACE, "hidden-act"); + Self { act, span } + } +} + +impl Module for HiddenActLayer { + fn forward(&self, xs: &Tensor) -> candle::Result { + let _enter = self.span.enter(); + match self.act { + // https://github.com/huggingface/transformers/blob/cd4584e3c809bb9e1392ccd3fe38b40daba5519a/src/transformers/activations.py#L213 + HiddenAct::Gelu => xs.gelu(), + HiddenAct::Relu => xs.relu(), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Default)] +#[serde(rename_all = "lowercase")] +pub enum PositionEmbeddingType { + #[default] + Absolute, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub dim: usize, + n_layers: usize, + n_heads: usize, + hidden_dim: usize, + activation: HiddenAct, + max_position_embeddings: usize, + initializer_range: f64, + pub pad_token_id: usize, + #[serde(default)] + position_embedding_type: PositionEmbeddingType, + #[serde(default)] + use_cache: bool, + model_type: Option, +} + +impl Default for Config { + fn default() -> Self { + Self { + vocab_size: 30522, + dim: 768, + n_layers: 12, + n_heads: 12, + hidden_dim: 3072, + activation: HiddenAct::Gelu, + max_position_embeddings: 512, + initializer_range: 0.02, + pad_token_id: 0, + position_embedding_type: PositionEmbeddingType::Absolute, + use_cache: true, + model_type: Some("distilbert".to_string()), + } + } +} + +struct Embeddings { + word_embeddings: Embedding, + position_embeddings: Embedding, + layer_norm: LayerNorm, + span: tracing::Span, +} + +impl Embeddings { + fn load(vb: VarBuilder, config: &Config) -> Result { + let word_embeddings = + candle_nn::embedding(config.vocab_size, config.dim, vb.pp("word_embeddings"))?; + let position_embeddings = candle_nn::embedding( + config.max_position_embeddings, + config.dim, + vb.pp("position_embeddings"), + )?; + let layer_norm = layer_norm(config.dim, 1e-12, vb.pp("LayerNorm"))?; + Ok(Self { + word_embeddings, + position_embeddings, + layer_norm, + span: tracing::span!(tracing::Level::TRACE, "embeddings"), + }) + } + + fn forward(&self, input_ids: &Tensor) -> Result { + let _enter = self.span.enter(); + let (_bsize, seq_len) = input_ids.dims2()?; + let input_embeddings = self.word_embeddings.forward(input_ids)?; + let position_ids = (0..seq_len as u32).collect::>(); + let position_ids = Tensor::new(&position_ids[..], input_ids.device())?; + let embeddings = + input_embeddings.broadcast_add(&self.position_embeddings.forward(&position_ids)?)?; + + let embeddings = self.layer_norm.forward(&embeddings)?; + Ok(embeddings) + } +} + +struct MultiHeadSelfAttention { + q_lin: Linear, + k_lin: Linear, + v_lin: Linear, + out_lin: Linear, + n_heads: usize, + attention_head_size: usize, + span: tracing::Span, +} + +impl MultiHeadSelfAttention { + fn load(vb: VarBuilder, config: &Config) -> Result { + let attention_head_size = config.dim / config.n_heads; + let all_head_size = config.n_heads * attention_head_size; + let dim = config.dim; + let q_lin = linear(dim, all_head_size, vb.pp("q_lin"))?; + let v_lin = linear(dim, all_head_size, vb.pp("v_lin"))?; + let k_lin = linear(dim, all_head_size, vb.pp("k_lin"))?; + let out_lin = linear(all_head_size, dim, vb.pp("out_lin"))?; + Ok(Self { + q_lin, + k_lin, + v_lin, + out_lin, + n_heads: config.n_heads, + attention_head_size, + span: tracing::span!(tracing::Level::TRACE, "attention"), + }) + } +} + +impl MultiHeadSelfAttention { + fn forward(&self, hidden_states: &Tensor, attention_mask: &Tensor) -> Result { + let _enter = self.span.enter(); + let (bs, q_length, _dim) = hidden_states.dims3()?; + + let dim_per_head = self.attention_head_size; + let q = self.q_lin.forward(hidden_states)?; + let k = self.k_lin.forward(hidden_states)?; + let v = self.v_lin.forward(hidden_states)?; + + let q = q + .reshape((bs, q_length, self.n_heads, dim_per_head))? + .transpose(1, 2)?; + let k = k + .reshape((bs, q_length, self.n_heads, dim_per_head))? + .transpose(1, 2)?; + let v = v + .reshape((bs, q_length, self.n_heads, dim_per_head))? + .transpose(1, 2)?; + + let q: Tensor = (q / (dim_per_head as f64).sqrt())?; + let scores = q.matmul(&k.transpose(2, 3)?.contiguous()?)?; + let mask = attention_mask.broadcast_as(scores.shape())?; + + let scores = masked_fill(&scores.to_dtype(DType::F32)?, &mask, f32::NEG_INFINITY)?; + let weights = candle_nn::ops::softmax(&scores, candle::D::Minus1)?; + + let context = weights.matmul(&v.contiguous()?)?; + let context = context + .transpose(1, 2)? + .reshape((bs, q_length, self.n_heads * dim_per_head))? + .contiguous()?; + let context = self.out_lin.forward(&context)?; + + Ok(context) + } +} + +#[allow(clippy::upper_case_acronyms)] +struct FFN { + lin1: Linear, + lin2: Linear, + activation: HiddenActLayer, + span: tracing::Span, +} + +impl FFN { + fn load(vb: VarBuilder, config: &Config) -> Result { + let lin1 = linear(config.dim, config.hidden_dim, vb.pp("lin1"))?; + let lin2 = linear(config.hidden_dim, config.dim, vb.pp("lin2"))?; + Ok(Self { + lin1, + lin2, + activation: HiddenActLayer::new(config.activation), + span: tracing::span!(tracing::Level::TRACE, "ffn"), + }) + } +} + +impl Module for FFN { + fn forward(&self, hidden_states: &Tensor) -> Result { + let _enter = self.span.enter(); + hidden_states + .apply(&self.lin1)? + .apply(&self.activation)? + .apply(&self.lin2) + } +} + +struct TransformerBlock { + attention: MultiHeadSelfAttention, + sa_layer_norm: LayerNorm, + ffn: FFN, + output_layer_norm: LayerNorm, + span: tracing::Span, +} + +impl TransformerBlock { + fn load(vb: VarBuilder, config: &Config) -> Result { + let attention = MultiHeadSelfAttention::load(vb.pp("attention"), config)?; + let sa_layer_norm = layer_norm(config.dim, 1e-12, vb.pp("sa_layer_norm"))?; + let ffn = FFN::load(vb.pp("ffn"), config)?; + let output_layer_norm = layer_norm(config.dim, 1e-12, vb.pp("output_layer_norm"))?; + Ok(Self { + attention, + sa_layer_norm, + ffn, + output_layer_norm, + span: tracing::span!(tracing::Level::TRACE, "layer"), + }) + } +} + +impl TransformerBlock { + fn forward(&self, hidden_states: &Tensor, attention_mask: &Tensor) -> Result { + let _enter = self.span.enter(); + let sa_output = self.attention.forward(hidden_states, attention_mask)?; + // TODO: Support cross-attention? + // https://github.com/huggingface/transformers/blob/6eedfa6dd15dc1e22a55ae036f681914e5a0d9a1/src/transformers/models/bert/modeling_bert.py#L523 + // TODO: Support something similar to `apply_chunking_to_forward`? + let sa_output = sa_output.broadcast_add(hidden_states)?; + let sa_output = self.sa_layer_norm.forward(&sa_output)?; + + let ffn_output = self.ffn.forward(&sa_output)?; + let ffn_output = (&ffn_output + sa_output)?; + let output = self.output_layer_norm.forward(&ffn_output)?; + Ok(output) + } +} + +// https://github.com/huggingface/transformers/blob/6eedfa6dd15dc1e22a55ae036f681914e5a0d9a1/src/transformers/models/bert/modeling_bert.py#L556 +struct Transformer { + layers: Vec, + span: tracing::Span, +} + +impl Transformer { + fn load(vb: VarBuilder, config: &Config) -> Result { + let layers = (0..config.n_layers) + .map(|index| TransformerBlock::load(vb.pp(format!("layer.{index}")), config)) + .collect::>>()?; + let span = tracing::span!(tracing::Level::TRACE, "encoder"); + Ok(Transformer { layers, span }) + } +} + +impl Transformer { + fn forward(&self, hidden_states: &Tensor, attention_mask: &Tensor) -> Result { + let _enter = self.span.enter(); + let mut hidden_states = hidden_states.clone(); + // Use a loop rather than a fold as it's easier to modify when adding debug/... + for layer in self.layers.iter() { + hidden_states = layer.forward(&hidden_states, attention_mask)?; + } + Ok(hidden_states) + } +} + +pub struct DistilBertModel { + embeddings: Embeddings, + transformer: Transformer, + pub device: Device, + span: tracing::Span, +} + +impl DistilBertModel { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let (embeddings, transformer) = match ( + Embeddings::load(vb.pp("embeddings"), config), + Transformer::load(vb.pp("transformer"), config), + ) { + (Ok(embeddings), Ok(encoder)) => (embeddings, encoder), + (Err(err), _) | (_, Err(err)) => { + if let Some(model_type) = &config.model_type { + if let (Ok(embeddings), Ok(encoder)) = ( + Embeddings::load(vb.pp(format!("{model_type}.embeddings")), config), + Transformer::load(vb.pp(format!("{model_type}.transformer")), config), + ) { + (embeddings, encoder) + } else { + return Err(err); + } + } else { + return Err(err); + } + } + }; + Ok(Self { + embeddings, + transformer, + device: vb.device().clone(), + span: tracing::span!(tracing::Level::TRACE, "model"), + }) + } + + pub fn forward(&self, input_ids: &Tensor, attention_mask: &Tensor) -> Result { + let _enter = self.span.enter(); + let embedding_output = self.embeddings.forward(input_ids)?; + let sequence_output = self + .transformer + .forward(&embedding_output, attention_mask)?; + Ok(sequence_output) + } +} + +struct DistilBertPredictionHeadTransform { + dense: Linear, + activation: HiddenActLayer, + layer_norm: LayerNorm, +} + +impl DistilBertPredictionHeadTransform { + fn load(vb: VarBuilder, config: &Config) -> Result { + let dense = linear(config.dim, config.dim, vb.pp("vocab_transform"))?; + let activation = HiddenActLayer::new(config.activation); + let layer_norm = layer_norm(config.dim, 1e-12, vb.pp("vocab_layer_norm"))?; + Ok(Self { + dense, + activation, + layer_norm, + }) + } +} + +impl Module for DistilBertPredictionHeadTransform { + fn forward(&self, hidden_states: &Tensor) -> Result { + let hidden_states = self + .activation + .forward(&self.dense.forward(hidden_states)?)?; + self.layer_norm.forward(&hidden_states) + } +} + +// https://github.com/huggingface/transformers/blob/1bd604d11c405dfb8b78bda4062d88fc75c17de0/src/transformers/models/bert/modeling_bert.py#L769C1-L790C1 +pub struct DistilBertLMPredictionHead { + transform: DistilBertPredictionHeadTransform, + decoder: Linear, +} + +impl DistilBertLMPredictionHead { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let transform = DistilBertPredictionHeadTransform::load(vb.clone(), config)?; + + // distil_bert_uncased uses the word embeddings for the vocab projector weight, but has a separate vocab_projector bias + let vocab_projector_weight_vb = vb.pp("distilbert.embeddings.word_embeddings"); + let init_ws = candle_nn::init::DEFAULT_KAIMING_NORMAL; + let ws = vocab_projector_weight_vb.get_with_hints( + (config.vocab_size, config.dim), + "weight", + init_ws, + )?; + let bound = 1. / (config.dim as f64).sqrt(); + let init_bs = candle_nn::Init::Uniform { + lo: -bound, + up: bound, + }; + + let vocab_projector_bias_vb = vb.pp("vocab_projector"); + let bs = vocab_projector_bias_vb.get_with_hints(config.vocab_size, "bias", init_bs)?; + + let decoder = Linear::from_weights(ws, Some(bs)); + + Ok(Self { transform, decoder }) + } +} + +impl Module for DistilBertLMPredictionHead { + fn forward(&self, hidden_states: &Tensor) -> Result { + self.decoder + .forward(&self.transform.forward(hidden_states)?) + } +} + +// https://github.com/huggingface/transformers/blob/1bd604d11c405dfb8b78bda4062d88fc75c17de0/src/transformers/models/bert/modeling_bert.py#L792 +pub struct DistilBertOnlyMLMHead { + predictions: DistilBertLMPredictionHead, +} + +impl DistilBertOnlyMLMHead { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let predictions = DistilBertLMPredictionHead::load(vb.clone(), config)?; + Ok(Self { predictions }) + } +} + +impl Module for DistilBertOnlyMLMHead { + fn forward(&self, sequence_output: &Tensor) -> Result { + self.predictions.forward(sequence_output) + } +} + +pub struct DistilBertForMaskedLM { + pub bert: DistilBertModel, + cls: DistilBertOnlyMLMHead, +} + +impl DistilBertForMaskedLM { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let bert = DistilBertModel::load(vb.pp("distilbert"), config)?; + let cls = DistilBertOnlyMLMHead::load(vb.clone(), config)?; + Ok(Self { bert, cls }) + } + + pub fn forward(&self, input_ids: &Tensor, attention_mask: &Tensor) -> Result { + let sequence_output = self.bert.forward(input_ids, attention_mask)?; + self.cls.forward(&sequence_output) + } +} diff --git a/patches/candle-transformers/src/models/efficientnet.rs b/patches/candle-transformers/src/models/efficientnet.rs new file mode 100644 index 0000000000..be69546057 --- /dev/null +++ b/patches/candle-transformers/src/models/efficientnet.rs @@ -0,0 +1,335 @@ +//! Implementation of EfficientBert, an efficient variant of BERT for computer vision tasks. +//! +//! See: +//! - ["EfficientBERT: Progressively Searching Multilayer Perceptron Architectures for BERT"](https://arxiv.org/abs/2201.00462) +//! +use candle::{Context, Result, Tensor, D}; +use candle_nn as nn; +use nn::{Module, VarBuilder}; + +// Based on the Python version from torchvision. +// https://github.com/pytorch/vision/blob/0d75d9e5516f446c9c0ef93bd4ed9fea13992d06/torchvision/models/efficientnet.py#L47 +#[derive(Debug, Clone, Copy)] +pub struct MBConvConfig { + expand_ratio: f64, + kernel: usize, + stride: usize, + input_channels: usize, + out_channels: usize, + num_layers: usize, +} + +fn make_divisible(v: f64, divisor: usize) -> usize { + let min_value = divisor; + let new_v = usize::max( + min_value, + (v + divisor as f64 * 0.5) as usize / divisor * divisor, + ); + if (new_v as f64) < 0.9 * v { + new_v + divisor + } else { + new_v + } +} + +fn bneck_confs(width_mult: f64, depth_mult: f64) -> Vec { + let bneck_conf = |e, k, s, i, o, n| { + let input_channels = make_divisible(i as f64 * width_mult, 8); + let out_channels = make_divisible(o as f64 * width_mult, 8); + let num_layers = (n as f64 * depth_mult).ceil() as usize; + MBConvConfig { + expand_ratio: e, + kernel: k, + stride: s, + input_channels, + out_channels, + num_layers, + } + }; + vec![ + bneck_conf(1., 3, 1, 32, 16, 1), + bneck_conf(6., 3, 2, 16, 24, 2), + bneck_conf(6., 5, 2, 24, 40, 2), + bneck_conf(6., 3, 2, 40, 80, 3), + bneck_conf(6., 5, 1, 80, 112, 3), + bneck_conf(6., 5, 2, 112, 192, 4), + bneck_conf(6., 3, 1, 192, 320, 1), + ] +} + +impl MBConvConfig { + pub fn b0() -> Vec { + bneck_confs(1.0, 1.0) + } + pub fn b1() -> Vec { + bneck_confs(1.0, 1.1) + } + pub fn b2() -> Vec { + bneck_confs(1.1, 1.2) + } + pub fn b3() -> Vec { + bneck_confs(1.2, 1.4) + } + pub fn b4() -> Vec { + bneck_confs(1.4, 1.8) + } + pub fn b5() -> Vec { + bneck_confs(1.6, 2.2) + } + pub fn b6() -> Vec { + bneck_confs(1.8, 2.6) + } + pub fn b7() -> Vec { + bneck_confs(2.0, 3.1) + } +} + +/// Conv2D with same padding. +#[derive(Debug)] +struct Conv2DSame { + conv2d: nn::Conv2d, + s: usize, + k: usize, +} + +impl Conv2DSame { + fn new( + vb: VarBuilder, + i: usize, + o: usize, + k: usize, + stride: usize, + groups: usize, + bias: bool, + ) -> Result { + let conv_config = nn::Conv2dConfig { + stride, + groups, + ..Default::default() + }; + let conv2d = if bias { + nn::conv2d(i, o, k, conv_config, vb)? + } else { + nn::conv2d_no_bias(i, o, k, conv_config, vb)? + }; + Ok(Self { + conv2d, + s: stride, + k, + }) + } +} + +impl Module for Conv2DSame { + fn forward(&self, xs: &Tensor) -> Result { + let s = self.s; + let k = self.k; + let (_, _, ih, iw) = xs.dims4()?; + let oh = ih.div_ceil(s); + let ow = iw.div_ceil(s); + let pad_h = usize::max((oh - 1) * s + k - ih, 0); + let pad_w = usize::max((ow - 1) * s + k - iw, 0); + if pad_h > 0 || pad_w > 0 { + let xs = xs.pad_with_zeros(2, pad_h / 2, pad_h - pad_h / 2)?; + let xs = xs.pad_with_zeros(3, pad_w / 2, pad_w - pad_w / 2)?; + self.conv2d.forward(&xs) + } else { + self.conv2d.forward(xs) + } + } +} + +#[derive(Debug)] +struct ConvNormActivation { + conv2d: Conv2DSame, + bn2d: nn::BatchNorm, + activation: bool, +} + +impl ConvNormActivation { + fn new( + vb: VarBuilder, + i: usize, + o: usize, + k: usize, + stride: usize, + groups: usize, + ) -> Result { + let conv2d = Conv2DSame::new(vb.pp("0"), i, o, k, stride, groups, false)?; + let bn2d = nn::batch_norm(o, 1e-3, vb.pp("1"))?; + Ok(Self { + conv2d, + bn2d, + activation: true, + }) + } + + fn no_activation(self) -> Self { + Self { + activation: false, + ..self + } + } +} + +impl Module for ConvNormActivation { + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.conv2d.forward(xs)?.apply_t(&self.bn2d, false)?; + if self.activation { + swish(&xs) + } else { + Ok(xs) + } + } +} + +#[derive(Debug)] +struct SqueezeExcitation { + fc1: Conv2DSame, + fc2: Conv2DSame, +} + +impl SqueezeExcitation { + fn new(vb: VarBuilder, in_channels: usize, squeeze_channels: usize) -> Result { + let fc1 = Conv2DSame::new(vb.pp("fc1"), in_channels, squeeze_channels, 1, 1, 1, true)?; + let fc2 = Conv2DSame::new(vb.pp("fc2"), squeeze_channels, in_channels, 1, 1, 1, true)?; + Ok(Self { fc1, fc2 }) + } +} + +impl Module for SqueezeExcitation { + fn forward(&self, xs: &Tensor) -> Result { + let residual = xs; + // equivalent to adaptive_avg_pool2d([1, 1]) + let xs = xs.mean_keepdim(D::Minus2)?.mean_keepdim(D::Minus1)?; + let xs = self.fc1.forward(&xs)?; + let xs = swish(&xs)?; + let xs = self.fc2.forward(&xs)?; + let xs = nn::ops::sigmoid(&xs)?; + residual.broadcast_mul(&xs) + } +} + +#[derive(Debug)] +struct MBConv { + expand_cna: Option, + depthwise_cna: ConvNormActivation, + squeeze_excitation: SqueezeExcitation, + project_cna: ConvNormActivation, + config: MBConvConfig, +} + +impl MBConv { + fn new(vb: VarBuilder, c: MBConvConfig) -> Result { + let vb = vb.pp("block"); + let exp = make_divisible(c.input_channels as f64 * c.expand_ratio, 8); + let expand_cna = if exp != c.input_channels { + Some(ConvNormActivation::new( + vb.pp("0"), + c.input_channels, + exp, + 1, + 1, + 1, + )?) + } else { + None + }; + let start_index = if expand_cna.is_some() { 1 } else { 0 }; + let depthwise_cna = + ConvNormActivation::new(vb.pp(start_index), exp, exp, c.kernel, c.stride, exp)?; + let squeeze_channels = usize::max(1, c.input_channels / 4); + let squeeze_excitation = + SqueezeExcitation::new(vb.pp(start_index + 1), exp, squeeze_channels)?; + let project_cna = + ConvNormActivation::new(vb.pp(start_index + 2), exp, c.out_channels, 1, 1, 1)? + .no_activation(); + Ok(Self { + expand_cna, + depthwise_cna, + squeeze_excitation, + project_cna, + config: c, + }) + } +} + +impl Module for MBConv { + fn forward(&self, xs: &Tensor) -> Result { + let use_res_connect = + self.config.stride == 1 && self.config.input_channels == self.config.out_channels; + let ys = match &self.expand_cna { + Some(expand_cna) => expand_cna.forward(xs)?, + None => xs.clone(), + }; + let ys = self.depthwise_cna.forward(&ys)?; + let ys = self.squeeze_excitation.forward(&ys)?; + let ys = self.project_cna.forward(&ys)?; + if use_res_connect { + ys + xs + } else { + Ok(ys) + } + } +} + +fn swish(s: &Tensor) -> Result { + s * nn::ops::sigmoid(s)? +} + +#[derive(Debug)] +pub struct EfficientNet { + init_cna: ConvNormActivation, + blocks: Vec, + final_cna: ConvNormActivation, + classifier: nn::Linear, +} + +impl EfficientNet { + pub fn new(p: VarBuilder, configs: Vec, nclasses: usize) -> Result { + let f_p = p.pp("features"); + let first_in_c = configs[0].input_channels; + let last_out_c = configs.last().context("no last")?.out_channels; + let final_out_c = 4 * last_out_c; + let init_cna = ConvNormActivation::new(f_p.pp(0), 3, first_in_c, 3, 2, 1)?; + let nconfigs = configs.len(); + let mut blocks = vec![]; + for (index, cnf) in configs.into_iter().enumerate() { + let f_p = f_p.pp(index + 1); + for r_index in 0..cnf.num_layers { + let cnf = if r_index == 0 { + cnf + } else { + MBConvConfig { + input_channels: cnf.out_channels, + stride: 1, + ..cnf + } + }; + blocks.push(MBConv::new(f_p.pp(r_index), cnf)?) + } + } + let final_cna = + ConvNormActivation::new(f_p.pp(nconfigs + 1), last_out_c, final_out_c, 1, 1, 1)?; + let classifier = nn::linear(final_out_c, nclasses, p.pp("classifier.1"))?; + Ok(Self { + init_cna, + blocks, + final_cna, + classifier, + }) + } +} + +impl Module for EfficientNet { + fn forward(&self, xs: &Tensor) -> Result { + let mut xs = self.init_cna.forward(xs)?; + for block in self.blocks.iter() { + xs = block.forward(&xs)? + } + let xs = self.final_cna.forward(&xs)?; + // Equivalent to adaptive_avg_pool2d([1, 1]) -> squeeze(-1) -> squeeze(-1) + let xs = xs.mean(D::Minus1)?.mean(D::Minus1)?; + self.classifier.forward(&xs) + } +} diff --git a/patches/candle-transformers/src/models/efficientvit.rs b/patches/candle-transformers/src/models/efficientvit.rs new file mode 100644 index 0000000000..4c231d7679 --- /dev/null +++ b/patches/candle-transformers/src/models/efficientvit.rs @@ -0,0 +1,490 @@ +//! EfficientViT (MSRA) inference implementation based on timm. +//! +//! This crate provides an implementation of the EfficientViT model from Microsoft Research Asia +//! for efficient image classification. The model uses cascaded group attention modules +//! to achieve strong performance while maintaining low memory usage. +//! +//! The model was originally described in the paper: +//! ["EfficientViT: Memory Efficient Vision Transformer with Cascaded Group Attention"](https://arxiv.org/abs/2305.07027) +//! +//! This implementation is based on the reference implementation from +//! [pytorch-image-models](https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/efficientvit_msra.py). +//! +//! # Example Usage +//! +//! This candle implementation uses a pre-trained EfficientViT (from Microsoft Research Asia) network for inference. +//! The classification head has been trained on the ImageNet dataset and returns the probabilities for the top-5 classes. +//! +//! +//! ```bash +//! cargo run +//! --example efficientvit \ +//! --release -- \ +//! --image candle-examples/examples/yolo-v8/assets/bike.jpg --which m1 +//! +//! > loaded image Tensor[dims 3, 224, 224; f32] +//! > model built +//! > mountain bike, all-terrain bike, off-roader: 69.80% +//! > unicycle, monocycle : 13.03% +//! > bicycle-built-for-two, tandem bicycle, tandem: 9.28% +//! > crash helmet : 2.25% +//! > alp : 0.46% +//! ``` +//! +//!
+//! +//!
+//! +use candle::{Result, Tensor, D}; +use candle_nn::{ + batch_norm, conv2d, conv2d_no_bias, linear, ops::sigmoid, ops::softmax, Conv2dConfig, Func, + VarBuilder, +}; + +#[derive(Clone)] +pub struct Config { + channels: [usize; 3], + blocks: [usize; 3], + heads: [usize; 3], + kernels: [usize; 4], +} + +impl Config { + pub fn m0() -> Self { + Self { + channels: [64, 128, 192], + blocks: [1, 2, 3], + heads: [4, 4, 4], + kernels: [5, 5, 5, 5], + } + } + pub fn m1() -> Self { + Self { + channels: [128, 144, 192], + blocks: [1, 2, 3], + heads: [2, 3, 3], + kernels: [7, 5, 3, 3], + } + } + pub fn m2() -> Self { + Self { + channels: [128, 192, 224], + blocks: [1, 2, 3], + heads: [4, 3, 2], + kernels: [7, 5, 3, 3], + } + } + pub fn m3() -> Self { + Self { + channels: [128, 240, 320], + blocks: [1, 2, 3], + heads: [4, 3, 4], + kernels: [5, 5, 5, 5], + } + } + pub fn m4() -> Self { + Self { + channels: [128, 256, 384], + blocks: [1, 2, 3], + heads: [4, 4, 4], + kernels: [7, 5, 3, 3], + } + } + + pub fn m5() -> Self { + Self { + channels: [192, 288, 384], + blocks: [1, 3, 4], + heads: [3, 3, 4], + kernels: [7, 5, 3, 3], + } + } +} + +fn efficientvit_stemblock( + in_channels: usize, + out_channels: usize, + vb: VarBuilder, +) -> Result> { + let conv2d_cfg = Conv2dConfig { + stride: 2, + padding: 1, + ..Default::default() + }; + + let bn = batch_norm(out_channels, 1e-5, vb.pp("bn"))?; + let conv = conv2d_no_bias(in_channels, out_channels, 3, conv2d_cfg, vb.pp("conv"))?; + + Ok(Func::new(move |xs| { + let xs = xs.apply(&conv)?.apply_t(&bn, false)?; + Ok(xs) + })) +} + +fn efficientvit_stem(dim: usize, vb: VarBuilder) -> Result> { + let conv1 = efficientvit_stemblock(3, dim / 8, vb.pp("conv1"))?; + let conv2 = efficientvit_stemblock(dim / 8, dim / 4, vb.pp("conv2"))?; + let conv3 = efficientvit_stemblock(dim / 4, dim / 2, vb.pp("conv3"))?; + let conv4 = efficientvit_stemblock(dim / 2, dim, vb.pp("conv4"))?; + + Ok(Func::new(move |xs| { + let xs = xs + .apply(&conv1)? + .relu()? + .apply(&conv2)? + .relu()? + .apply(&conv3)? + .relu()? + .apply(&conv4)?; + + Ok(xs) + })) +} + +fn depthwise_conv( + channels: usize, + kernel: usize, + stride: usize, + padding: usize, + vb: VarBuilder, +) -> Result> { + let conv2d_cfg = Conv2dConfig { + stride, + padding, + groups: channels, + ..Default::default() + }; + + let bn = batch_norm(channels, 1e-5, vb.pp("bn"))?; + let conv = conv2d_no_bias(channels, channels, kernel, conv2d_cfg, vb.pp("conv"))?; + + Ok(Func::new(move |xs| xs.apply(&conv)?.apply_t(&bn, false))) +} + +fn pointwise_conv( + in_channels: usize, + out_channels: usize, + vb: VarBuilder, +) -> Result> { + let conv2d_cfg = Conv2dConfig { + ..Default::default() + }; + + let bn = batch_norm(out_channels, 1e-5, vb.pp("bn"))?; + let conv = conv2d_no_bias(in_channels, out_channels, 1, conv2d_cfg, vb.pp("conv"))?; + + Ok(Func::new(move |xs| xs.apply(&conv)?.apply_t(&bn, false))) +} + +fn conv_mlp(in_channels: usize, out_channels: usize, vb: VarBuilder) -> Result> { + let pw1 = pointwise_conv(in_channels, out_channels, vb.pp("pw1"))?; + let pw2 = pointwise_conv(out_channels, in_channels, vb.pp("pw2"))?; + + Ok(Func::new(move |xs| { + let xs = xs.apply(&pw1)?.relu()?.apply(&pw2)?; + Ok(xs) + })) +} + +// Fixed per-stage resolutions +const RESOLUTIONS: [usize; 3] = [14, 7, 4]; + +// Attention block +fn efficientvit_attn( + cfg: &Config, + stage: usize, + in_channels: usize, + vb: VarBuilder, +) -> Result> { + let cga = cascaded_group_attn(cfg, stage, in_channels, vb)?; + + Ok(Func::new(move |xs| { + let mut xs = xs.clone(); + + let (b, c, h, w) = xs.dims4()?; + let win_res = 7; // Fixed window resolution + let pad_b = (win_res - h % win_res) % win_res; + let pad_r = (win_res - w % win_res) % win_res; + let ph = h + pad_b; + let pw = w + pad_r; + let nh = ph / win_res; + let nw = pw / win_res; + + if RESOLUTIONS[stage] > win_res { + xs = xs.permute((0, 2, 3, 1))?; + xs = xs.pad_with_zeros(D::Minus1, 0, pad_r)?; + xs = xs.pad_with_zeros(D::Minus2, 0, pad_b)?; + xs = xs + .reshape((b, nh, win_res, nw, win_res, c))? + .transpose(2, 3)?; + xs = xs + .reshape((b * nh * nw, win_res, win_res, c))? + .permute((0, 3, 1, 2))?; + } + + xs = xs.apply(&cga)?; + + if RESOLUTIONS[stage] > win_res { + xs = xs + .permute((0, 2, 3, 1))? + .reshape((b, nh, nw, win_res, win_res, c))?; + xs = xs.transpose(2, 3)?.reshape((b, ph, pw, c))?; + xs = xs.permute((0, 3, 1, 2))?; + } + + Ok(xs) + })) +} + +// Cascaded group attention +fn cascaded_group_attn( + cfg: &Config, + stage: usize, + in_channels: usize, + vb: VarBuilder, +) -> Result> { + let heads = cfg.heads[stage]; + let key_dim = 16; + + let val_dim = in_channels / heads; + + let scale = (key_dim as f64).powf(-0.5); + + let mut dws = Vec::with_capacity(heads); + let mut qkvs = Vec::with_capacity(heads); + for i in 0..heads { + dws.push(depthwise_conv( + key_dim, + cfg.kernels[i], + 1, + cfg.kernels[i] / 2, + vb.pp(format!("dws.{i}")), + )?); + + qkvs.push(pointwise_conv( + in_channels / heads, + in_channels / heads + 2 * key_dim, + vb.pp(format!("qkvs.{i}")), + )?); + } + let proj = pointwise_conv(in_channels, in_channels, vb.pp("proj.1"))?; + + Ok(Func::new(move |xs| { + let (b, _, h, w) = xs.dims4()?; + let feats_in = xs.chunk(heads, 1)?; + let mut feats_out = Vec::with_capacity(heads); + let mut feat = feats_in[0].clone(); + + for i in 0..heads { + if i > 0 { + feat = (&feat + &feats_in[i])?; + } + feat = feat.apply(&qkvs[i])?; + let res = feat.reshape((b, (), h, w))?; + let q = res.narrow(1, 0, key_dim)?; + let k = res.narrow(1, key_dim, key_dim)?; + let v = res.narrow(1, 2 * key_dim, val_dim)?; + + let q = q.apply(&dws[i])?; + + let q = q.flatten_from(2)?; + let k = k.flatten_from(2)?; + let v = v.flatten_from(2)?; + let q = (q * scale)?; + + let att = q.transpose(D::Minus2, D::Minus1)?.matmul(&k)?; + let att = softmax(&att, D::Minus1)?; + feat = v.matmul(&att.transpose(D::Minus2, D::Minus1)?)?; + feat = feat.reshape((b, val_dim, h, w))?; + feats_out.push(feat.clone()); + } + + let xs = Tensor::cat(&feats_out, 1)?; + let xs = xs.relu()?.apply(&proj)?; + + Ok(xs) + })) +} + +// Used by the downsampling layer +fn squeeze_and_excitation( + in_channels: usize, + squeeze_channels: usize, + vb: VarBuilder, +) -> Result> { + let conv2d_cfg = Conv2dConfig { + ..Default::default() + }; + let fc1 = conv2d(in_channels, squeeze_channels, 1, conv2d_cfg, vb.pp("fc1"))?; + let fc2 = conv2d(squeeze_channels, in_channels, 1, conv2d_cfg, vb.pp("fc2"))?; + + Ok(Func::new(move |xs| { + let residual = xs; + let xs = xs.mean_keepdim(D::Minus2)?.mean_keepdim(D::Minus1)?; + let xs = sigmoid(&xs.apply(&fc1)?.relu()?.apply(&fc2)?)?; + + residual.broadcast_mul(&xs) + })) +} + +// Used by the downsampling layer +fn patchmerge(in_channels: usize, out_channels: usize, vb: VarBuilder) -> Result> { + let dim = in_channels; + let hid_dim = in_channels * 4; + let conv1 = pointwise_conv(dim, hid_dim, vb.pp("conv1"))?; + let conv2 = depthwise_conv(hid_dim, 3, 2, 1, vb.pp("conv2"))?; + let conv3 = pointwise_conv(hid_dim, out_channels, vb.pp("conv3"))?; + let se = squeeze_and_excitation(hid_dim, hid_dim / 4, vb.pp("se"))?; + Ok(Func::new(move |xs| { + let xs = xs + .apply(&conv1)? + .relu()? + .apply(&conv2)? + .relu()? + .apply(&se)? + .apply(&conv3)?; + Ok(xs) + })) +} + +// Used by the downsampling layer +fn res(dim: usize, vb: VarBuilder) -> Result> { + let dw = depthwise_conv(dim, 3, 1, 1, vb.pp("0.m"))?; + let mlp = conv_mlp(dim, dim * 2, vb.pp("1.m"))?; + Ok(Func::new(move |xs| { + let mut xs = xs.clone(); + xs = (&xs + &xs.apply(&dw)?)?; + xs = (&xs + &xs.apply(&mlp)?)?; + Ok(xs) + })) +} + +// Downsampling +fn efficientvit_downsample( + in_channels: usize, + out_channels: usize, + vb: VarBuilder, +) -> Result> { + let res1 = res(in_channels, vb.pp("res1"))?; + let res2 = res(out_channels, vb.pp("res2"))?; + let patchmerge = patchmerge(in_channels, out_channels, vb.pp("patchmerge"))?; + Ok(Func::new(move |xs| { + let xs = xs.apply(&res1)?.apply(&patchmerge)?.apply(&res2)?; + Ok(xs) + })) +} + +fn efficientvit_block( + cfg: &Config, + stage: usize, + dim: usize, + vb: VarBuilder, +) -> Result> { + let dw0 = depthwise_conv(dim, 3, 1, 1, vb.pp("dw0.m"))?; + let dw1 = depthwise_conv(dim, 3, 1, 1, vb.pp("dw1.m"))?; + let ffn0 = conv_mlp(dim, dim * 2, vb.pp("ffn0.m"))?; + let ffn1 = conv_mlp(dim, dim * 2, vb.pp("ffn1.m"))?; + let attn = efficientvit_attn(cfg, stage, dim, vb.pp("mixer.m.attn"))?; + Ok(Func::new(move |xs| { + let mut xs = xs.clone(); + xs = (&xs + &xs.apply(&dw0)?)?; + xs = (&xs + &xs.apply(&ffn0)?)?; + xs = (&xs + &xs.apply(&attn)?)?; + xs = (&xs + &xs.apply(&dw1)?)?; + xs = (&xs + &xs.apply(&ffn1)?)?; + Ok(xs) + })) +} + +// Each stage is made of blocks. There is a downsampling layer between stages. +fn efficientvit_stage(cfg: &Config, stage: usize, vb: VarBuilder) -> Result> { + let nblocks = cfg.blocks[stage]; + let mut blocks = Vec::with_capacity(nblocks + 1); + + let in_channels = if stage > 0 { + cfg.channels[stage - 1] + } else { + cfg.channels[0] + }; + let out_channels = cfg.channels[stage]; + + if stage > 0 { + blocks.push(efficientvit_downsample( + in_channels, + out_channels, + vb.pp("downsample"), + )?); + } + + for i in 0..nblocks { + blocks.push(efficientvit_block( + cfg, + stage, + out_channels, + vb.pp(format!("blocks.{i}")), + )?); + } + + Ok(Func::new(move |xs| { + let mut xs = xs.clone(); + for block in blocks.iter() { + xs = xs.apply(block)? + } + Ok(xs) + })) +} + +// Classification head. +fn efficientvit_head(outputs: usize, nclasses: usize, vb: VarBuilder) -> Result> { + let norm = batch_norm(outputs, 1e-6, vb.pp("bn"))?; + let linear = linear(outputs, nclasses, vb.pp("linear"))?; + Ok(Func::new(move |xs| { + xs.apply_t(&norm, false)?.apply(&linear) + })) +} + +// Build a efficientvit model for a given configuration. +fn efficientvit_model( + config: &Config, + nclasses: Option, + vb: VarBuilder, +) -> Result> { + let cls = match nclasses { + None => None, + Some(nclasses) => { + let outputs = config.channels[2]; + let head = efficientvit_head(outputs, nclasses, vb.pp("head"))?; + Some(head) + } + }; + + let stem_dim = config.channels[0]; + let stem = efficientvit_stem(stem_dim, vb.pp("patch_embed"))?; + + let vb = vb.pp("stages"); + let stage1 = efficientvit_stage(config, 0, vb.pp(0))?; + let stage2 = efficientvit_stage(config, 1, vb.pp(1))?; + let stage3 = efficientvit_stage(config, 2, vb.pp(2))?; + + Ok(Func::new(move |xs| { + let xs = xs + .apply(&stem)? + .apply(&stage1)? + .apply(&stage2)? + .apply(&stage3)? + .mean(D::Minus2)? + .mean(D::Minus1)?; + match &cls { + None => Ok(xs), + Some(cls) => xs.apply(cls), + } + })) +} + +pub fn efficientvit(cfg: &Config, nclasses: usize, vb: VarBuilder) -> Result> { + efficientvit_model(cfg, Some(nclasses), vb) +} + +pub fn efficientvit_no_final_layer(cfg: &Config, vb: VarBuilder) -> Result> { + efficientvit_model(cfg, None, vb) +} diff --git a/patches/candle-transformers/src/models/encodec.rs b/patches/candle-transformers/src/models/encodec.rs new file mode 100644 index 0000000000..de280a570a --- /dev/null +++ b/patches/candle-transformers/src/models/encodec.rs @@ -0,0 +1,794 @@ +//! EnCodec neural audio codec based on the Encodec implementation. +//! +//! See ["High Fidelity Neural Audio Compression"](https://arxiv.org/abs/2210.13438) +//! +//! Based on implementation from [huggingface/transformers](https://github.com/huggingface/transformers/blob/main/src/transformers/models/encodec/modeling_encodec.py) + +use candle::{DType, IndexOp, Layout, Module, Result, Shape, Tensor, D}; +use candle_nn::{conv1d, Conv1d, ConvTranspose1d, VarBuilder}; + +// Encodec Model +// https://github.com/huggingface/transformers/blob/main/src/transformers/models/encodec/modeling_encodec.py + +#[derive(Debug, Copy, Clone, PartialEq, Eq, serde::Deserialize)] +pub enum NormType { + WeightNorm, + TimeGroupNorm, + None, +} + +#[derive(Debug, Copy, Clone, PartialEq, Eq, serde::Deserialize)] +pub enum PadMode { + Constant, + Reflect, + Replicate, +} + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Config { + pub target_bandwidths: Vec, + pub sampling_rate: usize, + pub audio_channels: usize, + pub normalize: bool, + pub chunk_length_s: Option, + pub overlap: Option, + pub hidden_size: usize, + pub num_filters: usize, + pub num_residual_layers: usize, + pub upsampling_ratios: Vec, + pub norm_type: NormType, + pub kernel_size: usize, + pub last_kernel_size: usize, + pub residual_kernel_size: usize, + pub dilation_growth_rate: usize, + pub use_causal_conv: bool, + pub pad_mode: PadMode, + pub compress: usize, + pub num_lstm_layers: usize, + pub trim_right_ratio: f64, + pub codebook_size: usize, + pub codebook_dim: Option, + pub use_conv_shortcut: bool, +} + +impl Default for Config { + fn default() -> Self { + Self { + target_bandwidths: vec![1.5, 3.0, 6.0, 12.0, 24.0], + sampling_rate: 24_000, + audio_channels: 1, + normalize: false, + chunk_length_s: None, + overlap: None, + hidden_size: 128, + num_filters: 32, + num_residual_layers: 1, + upsampling_ratios: vec![8, 5, 4, 2], + norm_type: NormType::WeightNorm, + kernel_size: 7, + last_kernel_size: 7, + residual_kernel_size: 3, + dilation_growth_rate: 2, + use_causal_conv: true, + // This should be PadMode::Reflect which is currently unsupported in candle. + pad_mode: PadMode::Replicate, + compress: 2, + num_lstm_layers: 2, + trim_right_ratio: 1.0, + codebook_size: 1024, + codebook_dim: None, + use_conv_shortcut: true, + } + } +} + +impl Config { + fn codebook_dim(&self) -> usize { + self.codebook_dim.unwrap_or(self.hidden_size) + } + + fn frame_rate(&self) -> usize { + let hop_length: usize = self.upsampling_ratios.iter().product(); + self.sampling_rate.div_ceil(hop_length) + } + + fn num_quantizers(&self) -> usize { + let num = 1000f64 + * self + .target_bandwidths + .last() + .expect("empty target_bandwidths"); + (num as usize) / (self.frame_rate() * 10) + } +} + +fn get_extra_padding_for_conv1d( + xs: &Tensor, + k_size: usize, + stride: usize, + padding_total: usize, +) -> Result { + let len = xs.dim(D::Minus1)?; + let n_frames = (len + padding_total).saturating_sub(k_size) as f64 / stride as f64 + 1.0; + let ideal_len = + ((n_frames.ceil() as usize - 1) * stride + k_size).saturating_sub(padding_total); + Ok(ideal_len.saturating_sub(len)) +} + +fn pad1d(xs: &Tensor, pad_l: usize, pad_r: usize, mode: PadMode) -> Result { + match mode { + PadMode::Constant => xs.pad_with_zeros(D::Minus1, pad_l, pad_r), + PadMode::Reflect => candle::bail!("pad-mode 'reflect' is not supported"), + PadMode::Replicate => xs.pad_with_same(D::Minus1, pad_l, pad_r), + } +} + +// Applies weight norm for inference by recomputing the weight tensor. This +// does not apply to training. +// https://pytorch.org/docs/stable/generated/torch.nn.utils.weight_norm.html +pub fn conv1d_weight_norm( + in_c: usize, + out_c: usize, + kernel_size: usize, + config: candle_nn::Conv1dConfig, + vb: VarBuilder, +) -> Result { + let weight_g = vb.get((out_c, 1, 1), "weight_g")?; + let weight_v = vb.get((out_c, in_c, kernel_size), "weight_v")?; + let norm_v = weight_v.sqr()?.sum_keepdim((1, 2))?.sqrt()?; + let weight = weight_v.broadcast_mul(&weight_g)?.broadcast_div(&norm_v)?; + let bias = vb.get(out_c, "bias")?; + Ok(Conv1d::new(weight, Some(bias), config)) +} + +pub fn conv1d_weight_norm_no_bias( + in_c: usize, + out_c: usize, + kernel_size: usize, + config: candle_nn::Conv1dConfig, + vb: VarBuilder, +) -> Result { + let weight_g = vb.get((out_c, 1, 1), "weight_g")?; + let weight_v = vb.get((out_c, in_c, kernel_size), "weight_v")?; + let norm_v = weight_v.sqr()?.sum_keepdim((1, 2))?.sqrt()?; + let weight = weight_v.broadcast_mul(&weight_g)?.broadcast_div(&norm_v)?; + Ok(Conv1d::new(weight, None, config)) +} + +pub fn conv_transpose1d_weight_norm( + in_c: usize, + out_c: usize, + kernel_size: usize, + bias: bool, + config: candle_nn::ConvTranspose1dConfig, + vb: VarBuilder, +) -> Result { + let weight_g = vb.get((in_c, 1, 1), "weight_g")?; + let weight_v = vb.get((in_c, out_c, kernel_size), "weight_v")?; + let norm_v = weight_v.sqr()?.sum_keepdim((1, 2))?.sqrt()?; + let weight = weight_v.broadcast_mul(&weight_g)?.broadcast_div(&norm_v)?; + let bias = if bias { + Some(vb.get(out_c, "bias")?) + } else { + None + }; + Ok(ConvTranspose1d::new(weight, bias, config)) +} + +struct CodebookEncode; + +impl candle::CustomOp2 for CodebookEncode { + fn name(&self) -> &'static str { + "cb" + } + + fn cpu_fwd( + &self, + lhs_storage: &candle::CpuStorage, + lhs_layout: &Layout, + rhs_storage: &candle::CpuStorage, + rhs_layout: &Layout, + ) -> Result<(candle::CpuStorage, Shape)> { + use rayon::prelude::*; + + let (lhs_dim1, lhs_dim2) = lhs_layout.shape().dims2()?; + let (rhs_dim1, rhs_dim2) = rhs_layout.shape().dims2()?; + if lhs_dim2 != rhs_dim2 { + candle::bail!("CodebookEncode, mismatch on last dim, {lhs_layout:?} {rhs_layout:?}"); + } + if lhs_dim2 == 0 { + candle::bail!("CodebookEncode, empty last dim {lhs_layout:?}") + } + let lhs = match lhs_layout.contiguous_offsets() { + None => candle::bail!("CodebookEncode, lhs has to be contiguous, got {lhs_layout:?}"), + Some((o1, o2)) => { + let slice = lhs_storage.as_slice::()?; + &slice[o1..o2] + } + }; + let rhs = match rhs_layout.contiguous_offsets() { + None => candle::bail!("CodebookEncode, rhs has to be contiguous, got {rhs_layout:?}"), + Some((o1, o2)) => { + let slice = rhs_storage.as_slice::()?; + &slice[o1..o2] + } + }; + let dst = (0..lhs_dim1) + .into_par_iter() + .map(|idx1| { + let mut where_min = 0; + let mut min_dist = f32::INFINITY; + let lhs = &lhs[idx1 * lhs_dim2..(idx1 + 1) * lhs_dim2]; + for idx2 in 0..rhs_dim1 { + let rhs = &rhs[idx2 * rhs_dim2..(idx2 + 1) * rhs_dim2]; + let mut dist = 0f32; + for (a, b) in lhs.iter().zip(rhs.iter()) { + dist += (a - b) * (a - b) + } + if dist < min_dist { + min_dist = dist; + where_min = idx2; + } + } + where_min as u32 + }) + .collect(); + let storage = candle::WithDType::to_cpu_storage_owned(dst); + Ok((storage, (lhs_dim1,).into())) + } +} + +// https://github.com/huggingface/transformers/blob/abaca9f9432a84cfaa95531de4c72334f38a42f2/src/transformers/models/encodec/modeling_encodec.py#L340 +#[allow(unused)] +#[derive(Clone, Debug)] +pub struct EuclideanCodebook { + inited: Tensor, + cluster_size: Tensor, + embed: candle_nn::Embedding, + embed_avg: Tensor, + c2: Tensor, +} + +impl EuclideanCodebook { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let inited = vb.get(1, "inited")?; + let cluster_size = vb.get(cfg.codebook_size, "cluster_size")?; + let e_shape = (cfg.codebook_size, cfg.codebook_dim()); + let embed = vb.get(e_shape, "embed")?; + let c2 = ((&embed * &embed)?.sum(D::Minus1)? / 2.0)?; + let embed_avg = vb.get(e_shape, "embed_avg")?; + Ok(Self { + inited, + cluster_size, + embed: candle_nn::Embedding::new(embed, cfg.codebook_dim()), + embed_avg, + c2, + }) + } + + pub fn encode_slow(&self, xs: &Tensor) -> Result { + let mut target_shape = xs.dims().to_vec(); + target_shape.pop(); + let xs = xs.flatten_to(D::Minus2)?; + let _ = xs.dims2()?; + let dot_prod = xs.matmul(&self.embed.embeddings().t()?)?; + let codes = self.c2.broadcast_sub(&dot_prod)?.argmin(D::Minus1)?; + codes.reshape(target_shape) + } + + pub fn encode(&self, xs: &Tensor) -> Result { + let mut target_shape = xs.dims().to_vec(); + target_shape.pop(); + let xs = xs.flatten_to(D::Minus2)?; + let _ = xs.dims2()?; + let codes = Tensor::apply_op2(&xs, self.embed.embeddings(), CodebookEncode)?; + codes.reshape(target_shape) + } + + pub fn decode(&self, embed_ind: &Tensor) -> Result { + let quantize = self.embed.forward(embed_ind)?; + Ok(quantize) + } +} + +#[derive(Clone, Debug)] +pub struct VectorQuantization { + codebook: EuclideanCodebook, +} + +impl VectorQuantization { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let codebook = EuclideanCodebook::new(cfg, vb.pp("codebook"))?; + Ok(Self { codebook }) + } + + pub fn encode(&self, xs: &Tensor) -> Result { + let xs = xs.transpose(1, 2)?; + self.codebook.encode_slow(&xs) + } + + pub fn decode(&self, embed_ind: &Tensor) -> Result { + let quantize = self.codebook.decode(embed_ind)?; + let quantize = quantize.transpose(1, 2)?; + Ok(quantize) + } +} + +#[derive(Clone, Debug)] +pub struct ResidualVectorQuantizer { + layers: Vec, + dtype: DType, +} + +impl ResidualVectorQuantizer { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb = &vb.pp("layers"); + let layers = (0..cfg.num_quantizers()) + .map(|i| VectorQuantization::new(cfg, vb.pp(i))) + .collect::>>()?; + Ok(Self { + layers, + dtype: vb.dtype(), + }) + } + + pub fn encode(&self, xs: &Tensor) -> Result { + let mut codes = Vec::with_capacity(self.layers.len()); + let mut residual = xs.clone(); + for layer in self.layers.iter() { + let indices = layer.encode(&residual)?; + let quantized = layer.decode(&indices)?; + residual = (residual - quantized)?; + codes.push(indices) + } + Tensor::stack(&codes, 0) + } + + pub fn decode(&self, codes: &Tensor) -> Result { + let mut quantized_out = Tensor::zeros((), self.dtype, codes.device())?; + let ncodes = codes.dim(0)?; + if ncodes > self.layers.len() { + candle::bail!( + "codes shape {:?} does not match the number of quantization layers {}", + codes.shape(), + self.layers.len() + ) + } + for (i, layer) in self.layers.iter().take(ncodes).enumerate() { + let quantized = layer.decode(&codes.i(i)?)?; + quantized_out = quantized.broadcast_add(&quantized_out)?; + } + Ok(quantized_out) + } +} + +// https://github.com/huggingface/transformers/blob/abaca9f9432a84cfaa95531de4c72334f38a42f2/src/transformers/models/encodec/modeling_encodec.py#L226 +#[derive(Clone, Debug)] +pub struct EncodecLSTM { + layers: Vec, +} + +impl EncodecLSTM { + pub fn new(dim: usize, cfg: &Config, vb: VarBuilder) -> Result { + let vb = &vb.pp("lstm"); + let mut layers = vec![]; + for layer_idx in 0..cfg.num_lstm_layers { + let config = candle_nn::LSTMConfig { + layer_idx, + ..Default::default() + }; + let lstm = candle_nn::lstm(dim, dim, config, vb.clone())?; + layers.push(lstm) + } + Ok(Self { layers }) + } +} + +impl Module for EncodecLSTM { + fn forward(&self, xs: &Tensor) -> Result { + use candle_nn::RNN; + // This is different from the Python transformers version as candle LSTM is batch first. + let xs = xs.t()?; + let residual = &xs; + let mut xs = xs.clone(); + for layer in self.layers.iter() { + let states = layer.seq(&xs)?; + xs = layer.states_to_tensor(&states)?; + } + let xs = (xs + residual)?.t()?; + Ok(xs) + } +} + +#[derive(Clone, Debug)] +pub struct EncodecConvTranspose1d { + conv: ConvTranspose1d, +} + +impl EncodecConvTranspose1d { + fn new( + in_c: usize, + out_c: usize, + k: usize, + stride: usize, + _cfg: &Config, + vb: VarBuilder, + ) -> Result { + let cfg = candle_nn::ConvTranspose1dConfig { + stride, + ..Default::default() + }; + let conv = conv_transpose1d_weight_norm(in_c, out_c, k, true, cfg, vb.pp("conv"))?; + Ok(Self { conv }) + } +} + +impl Module for EncodecConvTranspose1d { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.conv) + } +} + +#[derive(Clone, Debug)] +pub struct EncodecConv1d { + causal: bool, + conv: Conv1d, + norm: Option, + pad_mode: PadMode, +} + +impl EncodecConv1d { + pub fn new( + in_c: usize, + out_c: usize, + kernel_size: usize, + stride: usize, + dilation: usize, + cfg: &Config, + vb: VarBuilder, + ) -> Result { + let conv = match cfg.norm_type { + NormType::WeightNorm => conv1d_weight_norm( + in_c, + out_c, + kernel_size, + candle_nn::Conv1dConfig { + stride, + dilation, + ..Default::default() + }, + vb.pp("conv"), + )?, + NormType::None | NormType::TimeGroupNorm => conv1d( + in_c, + out_c, + kernel_size, + candle_nn::Conv1dConfig { + padding: 0, + stride, + groups: 1, + dilation: 1, + cudnn_fwd_algo: None, + }, + vb.pp("conv"), + )?, + }; + let norm = match cfg.norm_type { + NormType::None | NormType::WeightNorm => None, + NormType::TimeGroupNorm => { + let gn = candle_nn::group_norm(1, out_c, 1e-5, vb.pp("norm"))?; + Some(gn) + } + }; + Ok(Self { + causal: cfg.use_causal_conv, + conv, + norm, + pad_mode: cfg.pad_mode, + }) + } +} + +impl Module for EncodecConv1d { + fn forward(&self, xs: &Tensor) -> Result { + let (_b, _t, _c) = xs.dims3()?; + let k_size = self.conv.weight().dim(D::Minus1)?; + let conv_cfg = self.conv.config(); + // Effective kernel size with dilations. + let k_size = (k_size - 1) * conv_cfg.dilation + 1; + let padding_total = k_size - conv_cfg.stride; + let extra_padding = + get_extra_padding_for_conv1d(xs, k_size, conv_cfg.stride, padding_total)?; + let xs = if self.causal { + pad1d(xs, padding_total, extra_padding, self.pad_mode)? + } else { + let padding_right = padding_total / 2; + let padding_left = padding_total - padding_right; + pad1d( + xs, + padding_left, + padding_right + extra_padding, + self.pad_mode, + )? + }; + let xs = self.conv.forward(&xs)?; + match &self.norm { + None => Ok(xs), + Some(norm) => xs.apply(norm), + } + } +} + +#[derive(Clone, Debug)] +pub struct EncodecResnetBlock { + block_conv1: EncodecConv1d, + block_conv2: EncodecConv1d, + shortcut: Option, +} + +impl EncodecResnetBlock { + pub fn new( + dim: usize, + (dilation1, dilation2): (usize, usize), + cfg: &Config, + vb: VarBuilder, + ) -> Result { + let h = dim / cfg.compress; + let mut layer = Layer::new(vb.pp("block")); + // TODO: Apply dilations! + layer.inc(); + let block_conv1 = EncodecConv1d::new( + dim, + h, + cfg.residual_kernel_size, + 1, + dilation1, + cfg, + layer.next(), + )?; + layer.inc(); + let block_conv2 = EncodecConv1d::new(h, dim, 1, 1, dilation2, cfg, layer.next())?; + let shortcut = if cfg.use_conv_shortcut { + let conv = EncodecConv1d::new(dim, dim, 1, 1, 1, cfg, vb.pp("shortcut"))?; + Some(conv) + } else { + None + }; + Ok(Self { + block_conv1, + block_conv2, + shortcut, + }) + } +} + +impl Module for EncodecResnetBlock { + fn forward(&self, xs: &Tensor) -> Result { + let residual = xs.clone(); + let xs = xs.elu(1.)?; + let xs = self.block_conv1.forward(&xs)?; + let xs = xs.elu(1.)?; + let xs = self.block_conv2.forward(&xs)?; + let xs = match &self.shortcut { + None => (xs + residual)?, + Some(shortcut) => xs.add(&shortcut.forward(&residual)?)?, + }; + Ok(xs) + } +} + +struct Layer<'a> { + vb: VarBuilder<'a>, + cnt: usize, +} + +impl<'a> Layer<'a> { + fn new(vb: VarBuilder<'a>) -> Self { + Self { vb, cnt: 0 } + } + + fn inc(&mut self) { + self.cnt += 1; + } + + fn next(&mut self) -> VarBuilder<'_> { + let vb = self.vb.pp(self.cnt.to_string()); + self.cnt += 1; + vb + } +} + +#[derive(Clone, Debug)] +pub struct Encoder { + init_conv: EncodecConv1d, + sampling_layers: Vec<(Vec, EncodecConv1d)>, + final_lstm: EncodecLSTM, + final_conv: EncodecConv1d, +} + +impl Encoder { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let mut layer = Layer::new(vb.pp("layers")); + let init_conv = EncodecConv1d::new( + cfg.audio_channels, + cfg.num_filters, + cfg.kernel_size, + 1, + 1, + cfg, + layer.next(), + )?; + let mut sampling_layers = vec![]; + let mut scaling = 1; + for &ratio in cfg.upsampling_ratios.iter().rev() { + let current_scale = scaling * cfg.num_filters; + let mut resnets = vec![]; + for j in 0..(cfg.num_residual_layers as u32) { + let resnet = EncodecResnetBlock::new( + current_scale, + (cfg.dilation_growth_rate.pow(j), 1), + cfg, + layer.next(), + )?; + resnets.push(resnet) + } + layer.inc(); // ELU + let conv1d = EncodecConv1d::new( + current_scale, + current_scale * 2, + ratio * 2, + ratio, + 1, + cfg, + layer.next(), + )?; + sampling_layers.push((resnets, conv1d)); + scaling *= 2; + } + let final_lstm = EncodecLSTM::new(cfg.num_filters * scaling, cfg, layer.next())?; + layer.inc(); // ELU + let final_conv = EncodecConv1d::new( + cfg.num_filters * scaling, + cfg.hidden_size, + cfg.last_kernel_size, + 1, + 1, + cfg, + layer.next(), + )?; + Ok(Self { + init_conv, + sampling_layers, + final_conv, + final_lstm, + }) + } +} + +impl Module for Encoder { + fn forward(&self, xs: &Tensor) -> Result { + let mut xs = xs.apply(&self.init_conv)?; + for (resnets, conv) in self.sampling_layers.iter() { + for resnet in resnets.iter() { + xs = xs.apply(resnet)?; + } + xs = xs.elu(1.0)?.apply(conv)?; + } + xs.apply(&self.final_lstm)? + .elu(1.0)? + .apply(&self.final_conv) + } +} + +#[derive(Clone, Debug)] +pub struct Decoder { + init_conv: EncodecConv1d, + init_lstm: EncodecLSTM, + sampling_layers: Vec<(EncodecConvTranspose1d, Vec)>, + final_conv: EncodecConv1d, +} + +impl Decoder { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let mut layer = Layer::new(vb.pp("layers")); + let mut scaling = usize::pow(2, cfg.upsampling_ratios.len() as u32); + let init_conv = EncodecConv1d::new( + cfg.hidden_size, + cfg.num_filters * scaling, + cfg.last_kernel_size, + 1, + 1, + cfg, + layer.next(), + )?; + let init_lstm = EncodecLSTM::new(cfg.num_filters * scaling, cfg, layer.next())?; + let mut sampling_layers = vec![]; + for &ratio in cfg.upsampling_ratios.iter() { + let current_scale = scaling * cfg.num_filters; + layer.inc(); // ELU + let conv1d = EncodecConvTranspose1d::new( + current_scale, + current_scale / 2, + ratio * 2, + ratio, + cfg, + layer.next(), + )?; + let mut resnets = vec![]; + for j in 0..(cfg.num_residual_layers as u32) { + let resnet = EncodecResnetBlock::new( + current_scale / 2, + (cfg.dilation_growth_rate.pow(j), 1), + cfg, + layer.next(), + )?; + resnets.push(resnet) + } + sampling_layers.push((conv1d, resnets)); + scaling /= 2; + } + layer.inc(); // ELU + let final_conv = EncodecConv1d::new( + cfg.num_filters, + cfg.audio_channels, + cfg.last_kernel_size, + 1, + 1, + cfg, + layer.next(), + )?; + Ok(Self { + init_conv, + init_lstm, + sampling_layers, + final_conv, + }) + } +} + +impl Module for Decoder { + fn forward(&self, xs: &Tensor) -> Result { + let mut xs = xs.apply(&self.init_conv)?.apply(&self.init_lstm)?; + for (conv, resnets) in self.sampling_layers.iter() { + xs = xs.elu(1.)?.apply(conv)?; + for resnet in resnets.iter() { + xs = xs.apply(resnet)? + } + } + xs.elu(1.)?.apply(&self.final_conv) + } +} + +#[derive(Debug)] +pub struct Model { + encoder: Encoder, + decoder: Decoder, + quantizer: ResidualVectorQuantizer, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let encoder = Encoder::new(cfg, vb.pp("encoder"))?; + let decoder = Decoder::new(cfg, vb.pp("decoder"))?; + let quantizer = ResidualVectorQuantizer::new(cfg, vb.pp("quantizer"))?; + Ok(Self { + encoder, + decoder, + quantizer, + }) + } + + pub fn encode(&self, xs: &Tensor) -> Result { + let xs = self.encoder.forward(xs)?; + let codes = self.quantizer.encode(&xs)?; + codes.transpose(0, 1) + } + + pub fn decode(&self, codes: &Tensor) -> Result { + let (_b_sz, _codebooks, _seqlen) = codes.dims3()?; + let codes = codes.transpose(0, 1)?; + let embeddings = self.quantizer.decode(&codes)?; + let outputs = self.decoder.forward(&embeddings)?; + Ok(outputs) + } +} diff --git a/patches/candle-transformers/src/models/eva2.rs b/patches/candle-transformers/src/models/eva2.rs new file mode 100644 index 0000000000..9e31f58c73 --- /dev/null +++ b/patches/candle-transformers/src/models/eva2.rs @@ -0,0 +1,439 @@ +//! EVA-2 inference implementation. +//! +//! EVA-02 is a computer vision model that can be used as an ImageNet classifier. +//! The model returns the probability for an image to belong to each of the 1000 +//! ImageNet categories. +//! +//! - [Paper](https://arxiv.org/abs/2303.11331). EVA-02: A Visual Representation for Neon Genesis +//! - [Code](https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/eva2.py) +//! +//! # Example +//! +//! ```bash +//! cargo run \ +//! --example eva2 \ +//! --release -- \ +//! --image candle-examples/examples/yolo-v8/assets/bike.jpg +//! +//! > mountain bike, all-terrain bike, off-roader: 37.09% +//! > maillot : 8.30% +//! > alp : 2.13% +//! > bicycle-built-for-two, tandem bicycle, tandem: 0.84% +//! > crash helmet : 0.73% +//! ``` +//! +//!
+//! +//!
+//! +use candle::{IndexOp, Result, Tensor, D}; +use candle_nn::{layer_norm, LayerNorm, Linear, Module, VarBuilder}; + +const IMG_SIZE: usize = 448; +const PATCH_SIZE: usize = 14; +const NUM_CLASSES: usize = 1000; + +fn linear(vb: VarBuilder, in_dim: usize, out_dim: usize, bias: bool) -> Result { + if bias { + candle_nn::linear(in_dim, out_dim, vb) + } else { + candle_nn::linear_no_bias(in_dim, out_dim, vb) + } +} + +#[derive(Debug)] +struct Attention { + q: Linear, + k: Linear, + v: Linear, + proj: Linear, + rot_pos_embed: Tensor, + num_heads: usize, + scale: f64, +} + +impl Attention { + fn new( + vb: VarBuilder, + dim: usize, + num_heads: usize, + qkv_bias: bool, + proj_bias: bool, + rot_pos_embed: &Tensor, + ) -> Result { + let q = linear(vb.pp("q_proj"), dim, dim, qkv_bias)?; + let k = linear(vb.pp("k_proj"), dim, dim, false)?; // no bias for Key + let v = linear(vb.pp("v_proj"), dim, dim, qkv_bias)?; + let proj = linear(vb.pp("proj"), dim, dim, proj_bias)?; + let rot_pos_embed = rot_pos_embed.clone(); + let scale = 1. / ((dim / num_heads) as f64).sqrt(); + Ok(Self { + q, + k, + v, + proj, + rot_pos_embed, + num_heads, + scale, + }) + } +} + +impl Attention { + // See: https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/pos_embed_sincos.py#L210 + fn apply_rot_embed_cat(x: &Tensor, emb: &Tensor) -> Result { + let cos_emb = emb.i((0.., 64..128))?; //.transpose(0, 1)?; + let sin_emb = emb.i((0.., 0..64))?; //.transpose(0, 1)?; + let index_even: [u32; 32] = (0u32..=63) + .step_by(2) + .collect::>() + .try_into() + .expect("wrong size iterator"); + let index_odd: [u32; 32] = (1u32..=63) + .step_by(2) + .collect::>() + .try_into() + .expect("wrong size iterator"); + let t_index_even = Tensor::new(&index_even, x.device())?; + let t_index_odd = Tensor::new(&index_odd, x.device())?; + let x_c = x.contiguous()?; + let rot_x_even = x_c.index_select(&t_index_even, D::Minus1)?; + let rot_x_odd_minus = (-1.0 * x_c.index_select(&t_index_odd, D::Minus1)?)?; + let rot_x = + Tensor::stack(&[&rot_x_odd_minus, &rot_x_even], D::Minus1)?.reshape(x.shape())?; + x.broadcast_mul(&cos_emb)? + rot_x.broadcast_mul(&sin_emb)? + } +} + +impl Module for Attention { + fn forward(&self, xs: &Tensor) -> Result { + let (b, n, c) = xs.dims3()?; + let qkv = Tensor::cat( + &[ + &self.q.forward(xs)?, + &self.k.forward(xs)?, + &self.v.forward(xs)?, + ], + 2, + )? + .reshape((b, n, 3, self.num_heads, c / self.num_heads))? + .transpose(1, 2)? // 02134 + .transpose(0, 1)? // 20134 + .transpose(2, 3)?; // 20314 + let q = qkv.i(0)?; + let k = qkv.i(1)?.contiguous()?; + let v = qkv.i(2)?.contiguous()?; + + let npt = 1; // num_prefix_tokens = 1 for CLS token + let q = Tensor::cat( + &[ + &q.i((0.., 0.., ..npt, 0..))?, + &Self::apply_rot_embed_cat(&q.i((0.., 0.., npt.., 0..))?, &self.rot_pos_embed)?, + ], + 2, + )?; + let k = Tensor::cat( + &[ + &k.i((0.., 0.., ..npt, 0..))?, + &Self::apply_rot_embed_cat(&k.i((0.., 0.., npt.., 0..))?, &self.rot_pos_embed)?, + ], + 2, + )?; + + let q = (q * self.scale)?; + let attn = &q.matmul(&k.t()?)?; + let attn = candle_nn::ops::softmax(attn, D::Minus1)?; + let attn = attn.matmul(&v)?.transpose(1, 2)?.reshape((b, n, c))?; + self.proj.forward(&attn) + } +} + +#[derive(Debug)] +struct Mlp { + fc1_g: Linear, + fc1_x: Linear, + norm: LayerNorm, + fc2: Linear, +} + +impl Mlp { + fn new(vb: VarBuilder, in_features: usize, hidden_features: usize, bias: bool) -> Result { + let out_features = in_features; + let fc1_g = linear(vb.pp("fc1_g"), in_features, hidden_features, bias)?; + let fc1_x = linear(vb.pp("fc1_x"), in_features, hidden_features, bias)?; + let norm = layer_norm(hidden_features, 1e-6, vb.pp("norm"))?; + let fc2 = linear(vb.pp("fc2"), hidden_features, out_features, bias)?; + Ok(Self { + fc1_g, + fc1_x, + norm, + fc2, + }) + } +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let xs_g = self.fc1_g.forward(xs)?.silu()?; + let xs = self.fc1_x.forward(xs)?; + let xs = self.norm.forward(&(xs_g.mul(&xs)?))?; + self.fc2.forward(&xs) + } +} + +#[derive(Debug)] +struct Block { + norm1: LayerNorm, + attn: Attention, + norm2: LayerNorm, + mlp: Mlp, +} + +impl Block { + fn new(vb: VarBuilder, dim: usize, num_heads: usize, rot_pos_embed: &Tensor) -> Result { + let norm1 = layer_norm(dim, 1e-6, vb.pp("norm1"))?; + let attn = Attention::new(vb.pp("attn"), dim, num_heads, true, true, rot_pos_embed)?; + let norm2 = layer_norm(dim, 1e-6, vb.pp("norm2"))?; + let hidden_dim = dim * 4 * 2 / 3; // 768 * 4 * 2 / 3 = 3072 * 2 / 3 = 2048 + let mlp = Mlp::new(vb.pp("mlp"), dim, hidden_dim, true)?; + Ok(Self { + norm1, + attn, + norm2, + mlp, + }) + } +} + +impl Module for Block { + fn forward(&self, xs: &Tensor) -> Result { + let residual = xs; + let xs = &self.attn.forward(&self.norm1.forward(xs)?)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = &self.mlp.forward(&self.norm2.forward(&xs)?)?; + xs + residual + } +} + +#[derive(Debug)] +struct PatchEmbed { + proj: candle_nn::Conv2d, + patch_size: (usize, usize), + num_patches: usize, +} + +impl PatchEmbed { + fn new( + vb: VarBuilder, + img_size: usize, + patch_size: usize, + in_chans: usize, + embed_dim: usize, + ) -> Result { + let config = candle_nn::Conv2dConfig { + stride: patch_size, + ..Default::default() + }; + let proj = candle_nn::conv2d(in_chans, embed_dim, patch_size, config, vb.pp("proj"))?; + let num_patches = (img_size / patch_size) * (img_size / patch_size); + Ok(Self { + proj, + patch_size: (patch_size, patch_size), + num_patches, + }) + } +} + +impl Module for PatchEmbed { + fn forward(&self, xs: &Tensor) -> Result { + let (_b, _c, h, w) = xs.dims4()?; + let (patch_h, patch_w) = self.patch_size; + if (h % patch_h) != 0 { + candle::bail!("image height {h} is not a multiple of patch height {patch_h}") + } + if (w % patch_w) != 0 { + candle::bail!("image width {w} is not a multiple of patch width {patch_w}") + } + let xs = self.proj.forward(xs)?; + let (b, c, h, w) = xs.dims4()?; + // flatten embeddings. + xs.reshape((b, c, h * w))?.transpose(1, 2) + } +} + +#[derive(Debug)] +pub struct EVA2VisionTransformer { + patch_embed: PatchEmbed, + cls_token: Tensor, + pos_embed: Tensor, + blocks: Vec, + norm: LayerNorm, + head: Linear, +} + +impl EVA2VisionTransformer { + pub fn new(vb: VarBuilder, depth: usize, embed_dim: usize, num_heads: usize) -> Result { + let patch_embed = + PatchEmbed::new(vb.pp("patch_embed"), IMG_SIZE, PATCH_SIZE, 3, embed_dim)?; + let cls_token = vb.get((1, 1, embed_dim), "cls_token")?; + let pos_embed = vb.get((1, patch_embed.num_patches + 1, embed_dim), "pos_embed")?; + let rot_pos_embed = vb.get((patch_embed.num_patches, 128), "rot_pos_embed")?; + let head = linear(vb.pp("head"), embed_dim, NUM_CLASSES, true)?; + let norm = layer_norm(embed_dim, 1e-6, vb.pp("norm"))?; + let vb_b = vb.pp("blocks"); + let blocks = (0..depth) + .map(|i| Block::new(vb_b.pp(i.to_string()), embed_dim, num_heads, &rot_pos_embed)) + .collect::>>()?; + Ok(Self { + patch_embed, + cls_token, + pos_embed, + blocks, + norm, + head, + }) + } + + fn interpolate_pos_encoding( + &self, + xs: &Tensor, + w: usize, + h: usize, + num_prefix_tokens: usize, + ) -> Result { + let npatch = xs.dim(1)? - 1; + let n = self.pos_embed.dim(1)? - 1; + let sqrt_n = (n as f64).sqrt(); + if npatch == n && w == h { + return Ok(self.pos_embed.clone()); + } + // Interpolate only local tokens, i.e. those after the CLS token + let prefix_tokens_pos_embed = self.pos_embed.i((0.., ..num_prefix_tokens, 0..))?.clone(); + let patch_pos_embed = &self.pos_embed.i((0.., num_prefix_tokens.., 0..))?; + let dim = xs.dim(D::Minus1)?; + let (w0, h0) = ((w / PATCH_SIZE) as f64 + 0.1, (h / PATCH_SIZE) as f64 + 0.1); + let patch_pos_embed = patch_pos_embed + .reshape((1, sqrt_n as usize, sqrt_n as usize, dim))? + .transpose(2, 3)? + .transpose(1, 2)?; + // This uses bicubic interpolation in the original implementation. + let patch_pos_embed = patch_pos_embed.upsample_nearest2d(h0 as usize, w0 as usize)?; + let el_count = patch_pos_embed.shape().elem_count(); + let patch_pos_embed = + patch_pos_embed + .transpose(1, 2)? + .transpose(2, 3)? + .reshape((1, el_count / dim, dim))?; + Tensor::cat(&[&prefix_tokens_pos_embed, &patch_pos_embed], 1) + } + + fn prepare_tokens_with_mask(&self, xs: &Tensor) -> Result { + let (_b, _nc, w, h) = xs.dims4()?; + if (w != IMG_SIZE) || (h != IMG_SIZE) { + panic!("Error: The input tensor should have the shape: Bx3x518x518."); + } + let xs = self.patch_embed.forward(xs)?; + let xs = Tensor::cat(&[&self.cls_token, &xs], 1)?; + let xs = (&xs + &self.interpolate_pos_encoding(&xs, w, h, 1)?)?; + Ok(xs) + } + + fn get_intermediate_layers_not_chunked( + &self, + xs: &Tensor, + blocks_to_take: &[usize], + ) -> Result> { + let mut xs = self.prepare_tokens_with_mask(xs)?; + let mut output = Vec::new(); + for (i, blk) in self.blocks.iter().enumerate() { + xs = blk.forward(&xs)?; + if blocks_to_take.contains(&i) { + output.push(xs.clone()); + } + } + if output.len() != blocks_to_take.len() { + candle::bail!( + "only {} / {} blocks found", + output.len(), + blocks_to_take.len() + ); + } + Ok(output) + } + + pub fn get_intermediate_layers( + &self, + xs: &Tensor, + blocks_to_take: &[usize], + reshape: bool, + return_class_token: bool, + norm: bool, + ) -> Result { + let outputs = self.get_intermediate_layers_not_chunked(xs, blocks_to_take)?; + let outputs = if norm { + outputs + .iter() + .map(|out| self.norm.forward(out)) + .collect::>>()? + } else { + outputs + }; + let class_tokens = outputs + .iter() + .map(|out| out.i((.., 0))) + .collect::>>()?; + let outputs = outputs + .iter() + .map(|out| out.i((.., 1..))) + .collect::>>()?; + + let outputs = if reshape { + let (b, _c, w, h) = xs.dims4()?; + let patch_size = self.patch_embed.patch_size.0; + let num_channels = outputs[0].elem_count() / (b * (w / patch_size) * (h / patch_size)); + outputs + .iter() + .map(|out| { + out.reshape((b, w / patch_size, h / patch_size, num_channels))? + .transpose(2, 3)? + .transpose(1, 2) + }) + .collect::>>()? + } else { + outputs + }; + + let outputs = if return_class_token { + outputs + .iter() + .zip(class_tokens.iter()) + .map(|(out, class_token)| Tensor::cat(&[out, class_token], D::Minus1)) + .collect::>>()? + } else { + outputs + }; + + Tensor::stack(&outputs[..], 0) + } +} + +impl Module for EVA2VisionTransformer { + fn forward(&self, xs: &Tensor) -> Result { + let mut xs = self.prepare_tokens_with_mask(xs)?; + for blk in self.blocks.iter() { + xs = blk.forward(&xs)? + } + let xs_moy_local_tokens = xs.i((.., 1..))?.mean(1)?; + let xs_norm = self.norm.forward(&xs_moy_local_tokens)?; + self.head.forward(&xs_norm) + } +} + +pub fn vit_base(vb: VarBuilder) -> Result { + EVA2VisionTransformer::new(vb, 12, 768, 12) +} + +pub fn vit_large(vb: VarBuilder) -> Result { + EVA2VisionTransformer::new(vb, 24, 1024, 16) +} diff --git a/patches/candle-transformers/src/models/falcon.rs b/patches/candle-transformers/src/models/falcon.rs new file mode 100644 index 0000000000..c75b4d70d3 --- /dev/null +++ b/patches/candle-transformers/src/models/falcon.rs @@ -0,0 +1,500 @@ +//! Falcon language model inference implementation +//! +//! See ["Falcon: a new approach to large language models"](https://huggingface.co/blog/falcon) +//! +//! Based on implementation from [Huggingface Transformers](https://github.com/huggingface/transformers/blob/main/src/transformers/models/falcon) + +use candle::{DType, Device, Result, Tensor, D}; +use candle_nn::{embedding, linear_b as linear, Embedding, LayerNorm, Linear, Module, VarBuilder}; +use serde::Deserialize; + +const MAX_SEQ_LEN: usize = 5000; + +fn layer_norm(size: usize, eps: f64, vb: VarBuilder) -> Result { + let (weight, bias) = match (vb.get(size, "weight"), vb.get(size, "bias")) { + (Ok(weight), Ok(bias)) => (weight, bias), + (Err(err), _) | (_, Err(err)) => { + if let (Ok(weight), Ok(bias)) = (vb.get(size, "gamma"), vb.get(size, "beta")) { + (weight, bias) + } else { + return Err(err); + } + } + }; + Ok(LayerNorm::new(weight, bias, eps)) +} + +// https://raw.githubusercontent.com/huggingface/transformers/030c863aaa0165e98352b61697430bf69bf33755/src/transformers/models/falcon/configuration_falcon.py +#[derive(Clone, Debug, Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub layer_norm_epsilon: f64, + pub initializer_range: f64, + pub use_cache: bool, + pub bos_token_id: u32, + pub eos_token_id: u32, + pub hidden_dropout: f64, + pub attention_dropout: f64, + pub n_head_kv: Option, + pub alibi: bool, + pub new_decoder_architecture: bool, + pub multi_query: bool, + pub parallel_attn: bool, + pub bias: bool, +} + +impl Default for Config { + fn default() -> Self { + Self { + vocab_size: 65024, + hidden_size: 4544, + num_hidden_layers: 32, + num_attention_heads: 71, + layer_norm_epsilon: 1e-5, + initializer_range: 0.02, + use_cache: true, + bos_token_id: 11, + eos_token_id: 11, + hidden_dropout: 0.0, + attention_dropout: 0.0, + n_head_kv: None, + alibi: false, + new_decoder_architecture: false, + multi_query: true, + parallel_attn: true, + bias: false, + } + } +} + +impl Config { + pub fn validate(&self) -> Result<()> { + if self.alibi { + candle::bail!("alibi is not supported"); + } + if self.new_decoder_architecture { + candle::bail!("new_decoder_architecture is not supported"); + } + if self.n_head_kv.is_some() { + candle::bail!("n_head_kv is not supported"); + } + Ok(()) + } + + // https://huggingface.co/tiiuae/falcon-7b/blob/main/config.json + pub fn falcon7b() -> Self { + // This is currently on par with the defaults, the defaults come from the Python default + // arguments for the config initialization whereas the following come from the json config. + Self { + vocab_size: 65024, + hidden_size: 4544, + num_hidden_layers: 32, + num_attention_heads: 71, + layer_norm_epsilon: 1e-5, + initializer_range: 0.02, + use_cache: true, + bos_token_id: 11, + eos_token_id: 11, + hidden_dropout: 0., + attention_dropout: 0., + n_head_kv: None, + alibi: false, + new_decoder_architecture: false, + multi_query: true, + parallel_attn: true, + bias: false, + } + } + + fn head_dim(&self) -> usize { + self.hidden_size / self.num_attention_heads + } + + fn rotary(&self) -> bool { + !self.alibi + } +} + +fn rotate_half(x: &Tensor) -> Result { + let l = x.dim(D::Minus1)?; + let x1 = x.narrow(D::Minus1, 0, l / 2)?; + let x2 = x.narrow(D::Minus1, l / 2, l - l / 2)?; + let x21 = Tensor::cat(&[&x2.neg()?, &x1], D::Minus1)?; + Ok(x21) +} + +#[derive(Debug, Clone)] +struct FalconRotaryEmbedding { + inv_freq: Tensor, + cache: Option<(usize, Tensor, Tensor)>, +} + +impl FalconRotaryEmbedding { + fn load(device: &Device, cfg: &Config) -> Result { + let head_dim = cfg.head_dim(); + let inv_freq: Vec<_> = (0..head_dim) + .step_by(2) + .map(|i| 1f32 / 10000f32.powf(i as f32 / head_dim as f32)) + .collect(); + Ok(Self { + inv_freq: Tensor::new(inv_freq.as_slice(), device)?, + cache: None, + }) + } + + fn cos_sin( + &mut self, + seq_len: usize, + device: &Device, + dtype: DType, + ) -> Result<(Tensor, Tensor)> { + match &self.cache { + Some((s, cos, sin)) if *s == seq_len => { + return Ok((cos.clone(), sin.clone())); + } + _ => {} + } + let t = Tensor::arange(0, seq_len as u32, device)?.to_dtype(dtype)?; + let inv_freq = self.inv_freq.to_dtype(dtype)?; + let freqs = t.unsqueeze(1)?.matmul(&inv_freq.unsqueeze(0)?)?; + let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?; + let cos = emb.cos()?; + let sin = emb.sin()?; + self.cache = Some((seq_len, cos.clone(), sin.clone())); + Ok((cos, sin)) + } + + fn forward( + &mut self, + query: &Tensor, + key: &Tensor, + past_kv_len: usize, + ) -> Result<(Tensor, Tensor)> { + let (_batch, seq_len, _head_dim) = query.dims3()?; + let (cos, sin) = self.cos_sin(MAX_SEQ_LEN, query.device(), query.dtype())?; + let cos = cos.narrow(0, past_kv_len, seq_len)?; + let sin = sin.narrow(0, past_kv_len, seq_len)?; + let qs = (query.broadcast_mul(&cos)? + &rotate_half(query)?.broadcast_mul(&sin)?)?; + let ks = (key.broadcast_mul(&cos)? + &rotate_half(key)?.broadcast_mul(&sin)?)?; + Ok((qs, ks)) + } +} + +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())? + .to_dtype(on_false.dtype())? + .broadcast_as(shape.dims())?; + let m = mask.where_cond(&on_true, on_false)?; + Ok(m) +} + +#[derive(Debug, Clone)] +struct FalconAttention { + query_key_value: Linear, + dense: Linear, + maybe_rotary: Option, + kv_cache: Option<(Tensor, Tensor)>, + inv_norm_factor: f64, + multi_query: bool, + use_cache: bool, + num_heads: usize, + head_dim: usize, + n_head_kv: usize, +} + +impl FalconAttention { + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let maybe_rotary = if cfg.rotary() { + let rotary = FalconRotaryEmbedding::load(vb.device(), cfg)?; + Some(rotary) + } else { + None + }; + let head_dim = cfg.head_dim(); + let hidden_size = cfg.hidden_size; + let qkv_out_dim = if cfg.multi_query { + hidden_size + 2 * head_dim + } else { + 3 * hidden_size + }; + let query_key_value = linear(hidden_size, qkv_out_dim, cfg.bias, vb.pp("query_key_value"))?; + let dense = linear(hidden_size, hidden_size, cfg.bias, vb.pp("dense"))?; + Ok(Self { + query_key_value, + dense, + maybe_rotary, + kv_cache: None, + inv_norm_factor: 1. / (head_dim as f64).sqrt(), + multi_query: cfg.multi_query, + use_cache: cfg.use_cache, + num_heads: cfg.num_attention_heads, + n_head_kv: cfg.n_head_kv.unwrap_or(1), + head_dim, + }) + } + + fn split_heads(&self, fused_qkv: &Tensor) -> Result<(Tensor, Tensor, Tensor)> { + let (b_sz, seq_len, _) = fused_qkv.dims3()?; + if !self.multi_query { + let fused_qkv = fused_qkv.reshape((b_sz, seq_len, self.num_heads, 3, self.head_dim))?; + let q = fused_qkv.narrow(D::Minus2, 0, 1)?.squeeze(D::Minus2)?; + let k = fused_qkv.narrow(D::Minus2, 1, 1)?.squeeze(D::Minus2)?; + let v = fused_qkv.narrow(D::Minus2, 2, 1)?.squeeze(D::Minus2)?; + Ok((q, k, v)) + } else { + let fused_qkv = + fused_qkv.reshape((b_sz, seq_len, self.num_heads + 2, self.head_dim))?; + let d = fused_qkv.dim(D::Minus2)?; + let q = fused_qkv.narrow(D::Minus2, 0, d - 2)?; + let k = fused_qkv.narrow(D::Minus2, d - 2, 1)?; + let v = fused_qkv.narrow(D::Minus2, d - 1, 1)?; + Ok((q, k, v)) + } + } + + fn forward(&mut self, x: &Tensor, mask: Option<&Tensor>, past_kv_len: usize) -> Result { + let fused_qkv = self.query_key_value.forward(x)?; + let head_dim = self.head_dim; + let (query, key, value) = self.split_heads(&fused_qkv)?; + let (b_sz, seq_len, _, _) = query.dims4()?; + let query = query + .transpose(1, 2)? + .reshape((b_sz * self.num_heads, seq_len, head_dim))?; + let key = key + .transpose(1, 2)? + .reshape((b_sz * self.n_head_kv, seq_len, head_dim))?; + let value = value + .transpose(1, 2)? + .reshape((b_sz * self.n_head_kv, seq_len, head_dim))?; + let (query, key) = if let Some(r) = &mut self.maybe_rotary { + r.forward(&query, &key, past_kv_len)? + } else { + (query, key) + }; + let (mut key, mut value) = (key, value); + if self.use_cache { + if let Some((cache_k, cache_v)) = &self.kv_cache { + // TODO: we could trim the tensors to MAX_SEQ_LEN so that this would work for + // arbitrarily large sizes. + key = Tensor::cat(&[cache_k, &key], 1)?.contiguous()?; + value = Tensor::cat(&[cache_v, &value], 1)?.contiguous()?; + } + self.kv_cache = Some((key.clone(), value.clone())) + } + let query = query.reshape((b_sz * self.num_heads, seq_len, head_dim))?; + let all_len = past_kv_len + seq_len; + let key = key.reshape((b_sz * self.n_head_kv, all_len, head_dim))?; + let value = value.reshape((b_sz * self.n_head_kv, all_len, head_dim))?; + + let (key, value) = if self.n_head_kv == 1 { + ( + key.broadcast_as((b_sz * self.num_heads, all_len, head_dim))?, + value.broadcast_as((b_sz * self.num_heads, all_len, head_dim))?, + ) + } else { + (key, value) + }; + + // Only handle the case where alibi is None here, and non-flash attention. + let attention_scores = (query.matmul(&key.t()?)? * self.inv_norm_factor)?; + let attention_scores = match mask { + None => attention_scores, + Some(mask) => { + let mask = masked_fill(&mask.to_dtype(DType::F32)?, mask, -1e9)? + .to_dtype(query.dtype())?; + attention_scores.broadcast_add(&mask.squeeze(1)?)? + } + }; + + let attention_scores = + candle_nn::ops::softmax(&attention_scores.to_dtype(DType::F32)?, D::Minus1)? + .to_dtype(x.dtype())?; + let attn_output = attention_scores + .matmul(&value)? + .reshape((b_sz, self.num_heads, seq_len, head_dim))? + .transpose(1, 2)? + .reshape((b_sz, seq_len, self.num_heads * head_dim))?; + let attn_output = self.dense.forward(&attn_output)?; + Ok(attn_output) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +struct FalconMlp { + dense_h_to_4h: Linear, + dense_4h_to_h: Linear, +} + +impl FalconMlp { + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let h = cfg.hidden_size; + let b = cfg.bias; + let dense_h_to_4h = linear(h, 4 * h, b, vb.pp("dense_h_to_4h"))?; + let dense_4h_to_h = linear(4 * h, h, b, vb.pp("dense_4h_to_h"))?; + Ok(Self { + dense_h_to_4h, + dense_4h_to_h, + }) + } + + fn forward(&self, x: &Tensor) -> Result { + let x = self.dense_h_to_4h.forward(x)?.gelu()?; + let x = self.dense_4h_to_h.forward(&x)?; + Ok(x) + } +} + +#[derive(Debug, Clone)] +struct FalconDecoderLayer { + inp_layernorm: LayerNorm, + self_attention: FalconAttention, + post_attention_layernorm: Option, + mlp: FalconMlp, + parallel_attn: bool, +} + +impl FalconDecoderLayer { + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let mlp = FalconMlp::load(vb.pp("mlp"), cfg)?; + let inp_layernorm = layer_norm( + cfg.hidden_size, + cfg.layer_norm_epsilon, + vb.pp("input_layernorm"), + )?; + let self_attention = FalconAttention::load(vb.pp("self_attention"), cfg)?; + let post_attention_layernorm = if cfg.parallel_attn { + None + } else { + let ln = layer_norm( + cfg.hidden_size, + cfg.layer_norm_epsilon, + vb.pp("post_attention_layernorm"), + )?; + Some(ln) + }; + Ok(Self { + inp_layernorm, + self_attention, + post_attention_layernorm, + mlp, + parallel_attn: cfg.parallel_attn, + }) + } + + fn forward(&mut self, x: &Tensor, mask: Option<&Tensor>, past_kv_len: usize) -> Result { + let residual = x.clone(); + let ln_attn = self.inp_layernorm.forward(x)?; + let attn_output = self.self_attention.forward(&ln_attn, mask, past_kv_len)?; + let (residual, ln_mlp) = match &self.post_attention_layernorm { + None => (residual, ln_attn), + Some(pal) => { + // This should include some dropout. + let residual = (&attn_output + &residual)?; + let ln_mlp = pal.forward(&residual)?; + (residual, ln_mlp) + } + }; + let mlp_output = self.mlp.forward(&ln_mlp)?; + + let mlp_output = if self.parallel_attn { + (mlp_output + attn_output)? + } else { + mlp_output + }; + let output = (mlp_output + residual)?; + Ok(output) + } + + pub fn clear_kv_cache(&mut self) { + self.self_attention.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct Falcon { + word_embeddings: Embedding, + blocks: Vec, + ln_f: LayerNorm, + lm_head: Linear, + config: Config, +} + +fn make_causal_mask(t: usize) -> Result { + 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)?; + Ok(mask) +} + +fn prepare_attn_mask(b_sz: usize, seq_len: usize) -> Result { + // let mask = Tensor::ones((b_sz, seq_len), DType::U32, &Device::Cpu)?; + let mask = make_causal_mask(seq_len)?; + let mask = mask.broadcast_as((b_sz, 1, seq_len, seq_len))?; + Ok(mask) +} + +impl Falcon { + pub fn config(&self) -> &Config { + &self.config + } + + pub fn load(vb: VarBuilder, cfg: Config) -> Result { + let word_embeddings = embedding( + cfg.vocab_size, + cfg.hidden_size, + vb.pp("transformer.word_embeddings"), + )?; + let blocks = (0..cfg.num_hidden_layers) + .map(|i| FalconDecoderLayer::load(vb.pp(format!("transformer.h.{i}")), &cfg)) + .collect::>>()?; + let ln_f = layer_norm( + cfg.hidden_size, + cfg.layer_norm_epsilon, + vb.pp("transformer.ln_f"), + )?; + let lm_head = linear(cfg.hidden_size, cfg.vocab_size, false, vb.pp("lm_head"))?; + Ok(Self { + word_embeddings, + blocks, + ln_f, + lm_head, + config: cfg, + }) + } + + pub fn forward(&mut self, input_ids: &Tensor) -> Result { + let (b_sz, seq_len) = input_ids.dims2()?; + let mut hidden_state = self.word_embeddings.forward(input_ids)?; + let past_kv_len = match &self.blocks[0].self_attention.kv_cache { + Some((k, _)) => k.dim(1)?, + None => 0, + }; + let causal_mask = if seq_len <= 1 { + None + } else { + Some(prepare_attn_mask(b_sz, seq_len)?.to_device(input_ids.device())?) + }; + for block in self.blocks.iter_mut() { + hidden_state = block.forward(&hidden_state, causal_mask.as_ref(), past_kv_len)?; + } + let hidden_state = self.ln_f.forward(&hidden_state)?; + let hidden_state = hidden_state.narrow(1, seq_len - 1, 1)?; + let logits = self.lm_head.forward(&hidden_state)?.squeeze(1)?; + Ok(logits) + } + + pub fn clear_kv_cache(&mut self) { + for block in self.blocks.iter_mut() { + block.clear_kv_cache() + } + } +} diff --git a/patches/candle-transformers/src/models/fastvit.rs b/patches/candle-transformers/src/models/fastvit.rs new file mode 100644 index 0000000000..3f8664d9ba --- /dev/null +++ b/patches/candle-transformers/src/models/fastvit.rs @@ -0,0 +1,511 @@ +//! # FastViT inference implementation based on timm +//! +//! ## Description +//! See ["FastViT: A Fast Hybrid Vision Transformer using Structural Reparameterization"](https://arxiv.org/pdf/2303.14189) +//! +//! Implementation based on [timm model](https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/fastvit.py) + +use candle::{Context, DType, Result, Tensor, D}; +use candle_nn::{ + batch_norm, conv2d, conv2d_no_bias, linear, linear_no_bias, ops::sigmoid, ops::softmax, + BatchNorm, Conv2d, Conv2dConfig, Func, VarBuilder, +}; + +#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)] +pub struct Config { + pub exp_ratio: usize, + pub in_channels: usize, + pub blocks: [usize; 4], + pub attn: bool, + pub lkc_use_act: bool, +} + +impl Config { + pub fn t8() -> Self { + Self { + exp_ratio: 3, + in_channels: 48, + blocks: [2, 2, 4, 2], + attn: false, + lkc_use_act: false, + } + } + + pub fn t12() -> Self { + Self { + exp_ratio: 3, + in_channels: 64, + blocks: [2, 2, 6, 2], + attn: false, + lkc_use_act: false, + } + } + pub fn s12() -> Self { + Self { + exp_ratio: 4, + in_channels: 64, + blocks: [2, 2, 6, 2], + attn: false, + lkc_use_act: false, + } + } + pub fn sa12() -> Self { + Self { + exp_ratio: 4, + in_channels: 64, + blocks: [2, 2, 6, 2], + attn: true, + lkc_use_act: false, + } + } + pub fn sa24() -> Self { + Self { + exp_ratio: 4, + in_channels: 64, + blocks: [4, 4, 12, 4], + attn: true, + lkc_use_act: false, + } + } + pub fn sa36() -> Self { + Self { + exp_ratio: 4, + in_channels: 64, + blocks: [6, 6, 18, 6], + attn: true, + lkc_use_act: false, + } + } + pub fn ma36() -> Self { + Self { + exp_ratio: 4, + in_channels: 76, + blocks: [6, 6, 18, 6], + attn: true, + lkc_use_act: false, + } + } + + // configs used by MobileCLIP's image encoder + pub fn mci0() -> Self { + Self { + exp_ratio: 3, + in_channels: 64, + blocks: [2, 6, 10, 2], + attn: true, + lkc_use_act: true, + } + } + pub fn mci1() -> Self { + Self { + exp_ratio: 3, + in_channels: 64, + blocks: [4, 12, 20, 4], + attn: true, + lkc_use_act: true, + } + } + pub fn mci2() -> Self { + Self { + exp_ratio: 3, + in_channels: 80, + blocks: [4, 12, 24, 4], + attn: true, + lkc_use_act: true, + } + } +} + +fn conv_norm( + in_channels: usize, + out_channels: usize, + kernel: usize, + stride: usize, + vb: VarBuilder, +) -> Result> { + let conv2d_cfg = Conv2dConfig { + stride, + padding: kernel / 2, + groups: in_channels, + ..Default::default() + }; + + let bn = batch_norm(out_channels, 1e-5, vb.pp("bn"))?; + let conv = conv2d_no_bias(in_channels, out_channels, kernel, conv2d_cfg, vb.pp("conv"))?; + let conv = conv.absorb_bn(&bn)?; + Ok(Func::new(move |xs| { + let xs = xs.apply(&conv)?; + Ok(xs) + })) +} + +fn conv_mlp(dim: usize, exp_ratio: usize, vb: VarBuilder) -> Result> { + let conv2d_cfg = Conv2dConfig { + ..Default::default() + }; + + let conv = conv_norm(dim, dim, 7, 1, vb.pp("conv"))?; + let fc1 = conv2d(dim, dim * exp_ratio, 1, conv2d_cfg, vb.pp("fc1"))?; + let fc2 = conv2d(dim * exp_ratio, dim, 1, conv2d_cfg, vb.pp("fc2"))?; + + Ok(Func::new(move |xs| { + let xs = xs.apply(&conv)?.apply(&fc1)?.gelu_erf()?.apply(&fc2)?; + Ok(xs) + })) +} + +fn squeeze_and_excitation( + in_channels: usize, + squeeze_channels: usize, + vb: VarBuilder, +) -> Result> { + let conv2d_cfg = Conv2dConfig { + ..Default::default() + }; + let fc1 = conv2d(in_channels, squeeze_channels, 1, conv2d_cfg, vb.pp("fc1"))?; + let fc2 = conv2d(squeeze_channels, in_channels, 1, conv2d_cfg, vb.pp("fc2"))?; + + Ok(Func::new(move |xs| { + let residual = xs; + let xs = xs.mean_keepdim(D::Minus2)?.mean_keepdim(D::Minus1)?; + let xs = sigmoid(&xs.apply(&fc1)?.relu()?.apply(&fc2)?)?; + + residual.broadcast_mul(&xs) + })) +} + +// fuses a convolutional kernel and a batchnorm layer into a convolutional layer +// based on the _fuse_bn_tensor method in timm +// see https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/byobnet.py#L602 +fn fuse_conv_bn(weights: &Tensor, bn: BatchNorm) -> Result<(Tensor, Tensor)> { + let (gamma, beta) = bn.weight_and_bias().context("no weight-bias")?; + let mu = bn.running_mean(); + let sigma = (bn.running_var() + bn.eps())?.sqrt(); + let gps = (gamma / sigma)?; + let bias = (beta - mu * &gps)?; + let weights = weights.broadcast_mul(&gps.reshape(((), 1, 1, 1))?)?; + + Ok((weights, bias)) +} + +fn mobileone_block( + in_channels: usize, + out_channels: usize, + kernel: usize, + stride: usize, + group_size: usize, + use_act: bool, + vb: VarBuilder, +) -> Result> { + let groups = if group_size == 0 { + 1 + } else { + in_channels / group_size + }; + + let padding = kernel / 2; + let conv2d_cfg = Conv2dConfig { + stride, + groups, + padding, + ..Default::default() + }; + + let mut w = Tensor::zeros( + (out_channels, in_channels / groups, kernel, kernel), + DType::F32, + vb.device(), + )?; + let dim = out_channels; + + let mut b = Tensor::zeros(dim, DType::F32, vb.device())?; + + let conv_kxk_bn = batch_norm(dim, 1e-5, vb.pp("conv_kxk.0.bn")); + let conv_kxk = conv2d_no_bias( + in_channels, + out_channels, + kernel, + conv2d_cfg, + vb.pp("conv_kxk.0.conv"), + ); + + if let (Ok(conv), Ok(bn)) = (conv_kxk, conv_kxk_bn) { + let (wk, bk) = fuse_conv_bn(conv.weight(), bn)?; + w = (w + wk)?; + b = (b + bk)?; + }; + + let conv_scale_bn = batch_norm(dim, 1e-5, vb.pp("conv_scale.bn")); + let conv_scale = conv2d_no_bias( + in_channels, + out_channels, + 1, + conv2d_cfg, + vb.pp("conv_scale.conv"), + ); + + if let (Ok(conv), Ok(bn)) = (conv_scale, conv_scale_bn) { + let (ws, bs) = fuse_conv_bn(conv.weight(), bn)?; + // pad to 3x3 + let ws = ws + .pad_with_zeros(D::Minus1, 1, 1)? + .pad_with_zeros(D::Minus2, 1, 1)?; + + w = (w + ws)?; + b = (b + bs)?; + }; + + let se = squeeze_and_excitation(out_channels, out_channels / 16, vb.pp("se")); + + // read and reparameterize the identity bn into wi and bi + let identity_bn = batch_norm(dim, 1e-5, vb.pp("identity")); + + if let Ok(id_bn) = identity_bn { + let mut weights: Vec = vec![0.0; w.elem_count()]; + let id = in_channels / groups; + // See https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/byobnet.py#L809 + for i in 0..in_channels { + if kernel > 1 { + weights[i * kernel * kernel + 4] = 1.0; + } else { + weights[i * (id + 1)] = 1.0; + } + } + + let weights = &Tensor::from_vec(weights, w.shape(), w.device())?; + let (wi, bi) = fuse_conv_bn(weights, id_bn)?; + + w = (w + wi)?; + b = (b + bi)?; + }; + let reparam_conv = Conv2d::new(w, Some(b), conv2d_cfg); + + Ok(Func::new(move |xs| { + let mut xs = xs.apply(&reparam_conv)?; + if let Ok(f) = &se { + xs = xs.apply(f)?; + } + if use_act { + xs = xs.gelu_erf()?; + }; + Ok(xs) + })) +} + +fn repmixer(dim: usize, kernel: usize, vb: VarBuilder) -> Result> { + let gamma = vb.get((dim, 1, 1), "layer_scale.gamma")?; + let norm = mobileone_block(dim, dim, kernel, 1, 1, false, vb.pp("norm"))?; + let mixer = mobileone_block(dim, dim, kernel, 1, 1, false, vb.pp("mixer"))?; + + Ok(Func::new(move |xs| { + let residual = xs.clone(); + let xs = (xs.apply(&mixer)? - xs.apply(&norm)?)?; + let xs = xs.broadcast_mul(&gamma.reshape((1, (), 1, 1))?)?; + let xs = (xs + residual)?; + Ok(xs) + })) +} + +fn repmixer_block(dim: usize, exp_ratio: usize, vb: VarBuilder) -> Result> { + let gamma = vb.get((dim, 1, 1), "layer_scale.gamma")?; + let token_mixer = repmixer(dim, 3, vb.pp("token_mixer"))?; + let mlp = conv_mlp(dim, exp_ratio, vb.pp("mlp"))?; + + Ok(Func::new(move |xs| { + let residual = xs.apply(&token_mixer)?; + let mut xs = residual.apply(&mlp)?; + xs = xs.broadcast_mul(&gamma.reshape((1, (), 1, 1))?)?; + let xs = (xs + residual)?; + Ok(xs) + })) +} + +fn positional_encoding(dim: usize, vb: VarBuilder) -> Result> { + let conv2d_cfg = Conv2dConfig { + stride: 1, + padding: 3, + groups: dim, + ..Default::default() + }; + + let conv = conv2d(dim, dim, 7, conv2d_cfg, vb.pp("pos_enc"))?; + + Ok(Func::new(move |xs| { + let xs = (xs + xs.apply(&conv)?)?; + Ok(xs) + })) +} + +fn attention(dim: usize, vb: VarBuilder) -> Result> { + let qkv = linear_no_bias(dim, dim * 3, vb.pp("qkv"))?; + let proj = linear(dim, dim, vb.pp("proj"))?; + let head_dim = 32; + let num_heads = dim / head_dim; + let scale = (head_dim as f64).powf(-0.5); + + Ok(Func::new(move |xs| { + let xs = xs.clone(); + let (b, c, h, w) = xs.dims4()?; + let n = h * w; + let xs = xs.flatten_from(2)?.transpose(D::Minus1, D::Minus2)?; + let qkv = xs + .apply(&qkv)? + .reshape((b, n, 3, num_heads, head_dim))? + .permute((2, 0, 3, 1, 4))?; + + let q = qkv.get(0)?; + let k = qkv.get(1)?; + let v = qkv.get(2)?; + + let q = (q * scale)?; + + let att = q.matmul(&k.transpose(D::Minus2, D::Minus1)?)?; + let att = softmax(&att, D::Minus1)?; + let xs = att.matmul(&v)?; + + let xs = xs.transpose(1, 2)?.reshape((b, n, c))?; + let xs = xs.apply(&proj)?; + let xs = xs.transpose(D::Minus1, D::Minus2)?.reshape((b, c, h, w))?; + + Ok(xs) + })) +} + +fn attention_block(dim: usize, exp_ratio: usize, vb: VarBuilder) -> Result> { + let gamma1 = vb.get((dim, 1, 1), "layer_scale_1.gamma")?; + let gamma2 = vb.get((dim, 1, 1), "layer_scale_2.gamma")?; + let norm = batch_norm(dim, 1e-5, vb.pp("norm"))?; + let token_mixer = attention(dim, vb.pp("token_mixer"))?; + let mlp = conv_mlp(dim, exp_ratio, vb.pp("mlp"))?; + + Ok(Func::new(move |xs| { + let xs = xs.clone(); + let xs = (&xs + + &xs + .apply_t(&norm, false)? + .apply(&token_mixer)? + .broadcast_mul(&gamma1.reshape((1, (), 1, 1))?)?)?; + + let xs = (&xs + + &xs + .apply(&mlp)? + .broadcast_mul(&gamma2.reshape((1, (), 1, 1))?)?)?; + + Ok(xs) + })) +} + +fn fastvit_stage(cfg: &Config, idx: usize, vb: VarBuilder) -> Result> { + let nblocks = cfg.blocks[idx]; + let mut blocks = Vec::with_capacity(nblocks); + + let dim = cfg.in_channels << idx; + let downsample = fastvit_patch_embed(dim / 2, dim, cfg.lkc_use_act, vb.pp("downsample")); + for block_idx in 0..nblocks { + let block = if cfg.attn && idx == 3 { + attention_block(dim, cfg.exp_ratio, vb.pp(format!("blocks.{block_idx}")))? + } else { + repmixer_block(dim, cfg.exp_ratio, vb.pp(format!("blocks.{block_idx}")))? + }; + blocks.push(block); + } + let pos_emb = positional_encoding(dim, vb.pp("pos_emb")); + + Ok(Func::new(move |xs| { + let mut xs = xs.clone(); + if let Ok(ds) = &downsample { + xs = xs.apply(ds)?; + } + if let Ok(pos) = &pos_emb { + xs = xs.apply(pos)?; + } + for block in blocks.iter() { + xs = xs.apply(block)?; + } + Ok(xs) + })) +} + +fn fastvit_patch_embed( + in_channels: usize, + out_channels: usize, + use_act: bool, + vb: VarBuilder, +) -> Result> { + let lk = conv_norm(in_channels, out_channels, 7, 2, vb.pp("proj.0.large_conv"))?; + let sk = conv_norm(in_channels, out_channels, 3, 2, vb.pp("proj.0.small_conv"))?; + let se = squeeze_and_excitation(out_channels, out_channels / 4, vb.pp("proj.0.se")); + let mb = mobileone_block(out_channels, out_channels, 1, 1, 0, true, vb.pp("proj.1"))?; + + Ok(Func::new(move |xs| { + let mut xs = (xs.apply(&lk)? + xs.apply(&sk)?)?; + if let Ok(f) = &se { + xs = xs.apply(f)?; + } + if use_act { + xs = xs.gelu_erf()?; + }; + let xs = xs.apply(&mb)?; + Ok(xs) + })) +} + +fn fastvit_stem(in_channels: usize, out_channels: usize, vb: VarBuilder) -> Result> { + let mb0 = mobileone_block(in_channels, out_channels, 3, 2, 0, true, vb.pp(0))?; + let mb1 = mobileone_block(out_channels, out_channels, 3, 2, 1, true, vb.pp(1))?; + let mb2 = mobileone_block(out_channels, out_channels, 1, 1, 0, true, vb.pp(2))?; + Ok(Func::new(move |xs| { + let xs = xs.apply(&mb0)?.apply(&mb1)?.apply(&mb2)?; + Ok(xs) + })) +} + +// Build a fastvit model for a given configuration. +fn fastvit_model(cfg: &Config, nclasses: Option, vb: VarBuilder) -> Result> { + let cls = match nclasses { + None => None, + Some(nclasses) => { + let linear = linear(cfg.in_channels * 16, nclasses, vb.pp("head.fc"))?; + Some(linear) + } + }; + + let stem = fastvit_stem(3, cfg.in_channels, vb.pp("stem"))?; + let final_conv = mobileone_block( + cfg.in_channels * 8, + cfg.in_channels * 16, + 3, + 1, + 1, + true, + vb.pp("final_conv"), + )?; + + let vb = vb.pp("stages"); + let stage1 = fastvit_stage(cfg, 0, vb.pp(0))?; + let stage2 = fastvit_stage(cfg, 1, vb.pp(1))?; + let stage3 = fastvit_stage(cfg, 2, vb.pp(2))?; + let stage4 = fastvit_stage(cfg, 3, vb.pp(3))?; + + Ok(Func::new(move |xs| { + let xs = xs + .apply(&stem)? + .apply(&stage1)? + .apply(&stage2)? + .apply(&stage3)? + .apply(&stage4)? + .apply(&final_conv)?; + match &cls { + None => Ok(xs), + Some(cls) => xs.mean(D::Minus2)?.mean(D::Minus1)?.apply(cls), + } + })) +} + +pub fn fastvit(cfg: &Config, nclasses: usize, vb: VarBuilder) -> Result> { + fastvit_model(cfg, Some(nclasses), vb) +} + +pub fn fastvit_no_final_layer(cfg: &Config, vb: VarBuilder) -> Result> { + fastvit_model(cfg, None, vb) +} diff --git a/patches/candle-transformers/src/models/flux/autoencoder.rs b/patches/candle-transformers/src/models/flux/autoencoder.rs new file mode 100644 index 0000000000..8c2aebbdc4 --- /dev/null +++ b/patches/candle-transformers/src/models/flux/autoencoder.rs @@ -0,0 +1,440 @@ +use candle::{Result, Tensor, D}; +use candle_nn::{conv2d, group_norm, Conv2d, GroupNorm, VarBuilder}; + +// https://github.com/black-forest-labs/flux/blob/727e3a71faf37390f318cf9434f0939653302b60/src/flux/modules/autoencoder.py#L9 +#[derive(Debug, Clone)] +pub struct Config { + pub resolution: usize, + pub in_channels: usize, + pub ch: usize, + pub out_ch: usize, + pub ch_mult: Vec, + pub num_res_blocks: usize, + pub z_channels: usize, + pub scale_factor: f64, + pub shift_factor: f64, +} + +impl Config { + // https://github.com/black-forest-labs/flux/blob/727e3a71faf37390f318cf9434f0939653302b60/src/flux/util.py#L47 + pub fn dev() -> Self { + Self { + resolution: 256, + in_channels: 3, + ch: 128, + out_ch: 3, + ch_mult: vec![1, 2, 4, 4], + num_res_blocks: 2, + z_channels: 16, + scale_factor: 0.3611, + shift_factor: 0.1159, + } + } + + // https://github.com/black-forest-labs/flux/blob/727e3a71faf37390f318cf9434f0939653302b60/src/flux/util.py#L79 + pub fn schnell() -> Self { + Self { + resolution: 256, + in_channels: 3, + ch: 128, + out_ch: 3, + ch_mult: vec![1, 2, 4, 4], + num_res_blocks: 2, + z_channels: 16, + scale_factor: 0.3611, + shift_factor: 0.1159, + } + } +} + +fn scaled_dot_product_attention(q: &Tensor, k: &Tensor, v: &Tensor) -> Result { + let dim = q.dim(D::Minus1)?; + let scale_factor = 1.0 / (dim as f64).sqrt(); + let attn_weights = (q.matmul(&k.t()?)? * scale_factor)?; + candle_nn::ops::softmax_last_dim(&attn_weights)?.matmul(v) +} + +#[derive(Debug, Clone)] +struct AttnBlock { + q: Conv2d, + k: Conv2d, + v: Conv2d, + proj_out: Conv2d, + norm: GroupNorm, +} + +impl AttnBlock { + fn new(in_c: usize, vb: VarBuilder) -> Result { + let q = conv2d(in_c, in_c, 1, Default::default(), vb.pp("q"))?; + let k = conv2d(in_c, in_c, 1, Default::default(), vb.pp("k"))?; + let v = conv2d(in_c, in_c, 1, Default::default(), vb.pp("v"))?; + let proj_out = conv2d(in_c, in_c, 1, Default::default(), vb.pp("proj_out"))?; + let norm = group_norm(32, in_c, 1e-6, vb.pp("norm"))?; + Ok(Self { + q, + k, + v, + proj_out, + norm, + }) + } +} + +impl candle::Module for AttnBlock { + fn forward(&self, xs: &Tensor) -> Result { + let init_xs = xs; + let xs = xs.apply(&self.norm)?; + let q = xs.apply(&self.q)?; + let k = xs.apply(&self.k)?; + let v = xs.apply(&self.v)?; + let (b, c, h, w) = q.dims4()?; + let q = q.flatten_from(2)?.t()?.unsqueeze(1)?; + let k = k.flatten_from(2)?.t()?.unsqueeze(1)?; + let v = v.flatten_from(2)?.t()?.unsqueeze(1)?; + let xs = scaled_dot_product_attention(&q, &k, &v)?; + let xs = xs.squeeze(1)?.t()?.reshape((b, c, h, w))?; + xs.apply(&self.proj_out)? + init_xs + } +} + +#[derive(Debug, Clone)] +struct ResnetBlock { + norm1: GroupNorm, + conv1: Conv2d, + norm2: GroupNorm, + conv2: Conv2d, + nin_shortcut: Option, +} + +impl ResnetBlock { + fn new(in_c: usize, out_c: usize, vb: VarBuilder) -> Result { + let conv_cfg = candle_nn::Conv2dConfig { + padding: 1, + ..Default::default() + }; + let norm1 = group_norm(32, in_c, 1e-6, vb.pp("norm1"))?; + let conv1 = conv2d(in_c, out_c, 3, conv_cfg, vb.pp("conv1"))?; + let norm2 = group_norm(32, out_c, 1e-6, vb.pp("norm2"))?; + let conv2 = conv2d(out_c, out_c, 3, conv_cfg, vb.pp("conv2"))?; + let nin_shortcut = if in_c == out_c { + None + } else { + Some(conv2d( + in_c, + out_c, + 1, + Default::default(), + vb.pp("nin_shortcut"), + )?) + }; + Ok(Self { + norm1, + conv1, + norm2, + conv2, + nin_shortcut, + }) + } +} + +impl candle::Module for ResnetBlock { + fn forward(&self, xs: &Tensor) -> Result { + let h = xs + .apply(&self.norm1)? + .apply(&candle_nn::Activation::Swish)? + .apply(&self.conv1)? + .apply(&self.norm2)? + .apply(&candle_nn::Activation::Swish)? + .apply(&self.conv2)?; + match self.nin_shortcut.as_ref() { + None => xs + h, + Some(c) => xs.apply(c)? + h, + } + } +} + +#[derive(Debug, Clone)] +struct Downsample { + conv: Conv2d, +} + +impl Downsample { + fn new(in_c: usize, vb: VarBuilder) -> Result { + let conv_cfg = candle_nn::Conv2dConfig { + stride: 2, + ..Default::default() + }; + let conv = conv2d(in_c, in_c, 3, conv_cfg, vb.pp("conv"))?; + Ok(Self { conv }) + } +} + +impl candle::Module for Downsample { + fn forward(&self, xs: &Tensor) -> Result { + let xs = xs.pad_with_zeros(D::Minus1, 0, 1)?; + let xs = xs.pad_with_zeros(D::Minus2, 0, 1)?; + xs.apply(&self.conv) + } +} + +#[derive(Debug, Clone)] +struct Upsample { + conv: Conv2d, +} + +impl Upsample { + fn new(in_c: usize, vb: VarBuilder) -> Result { + let conv_cfg = candle_nn::Conv2dConfig { + padding: 1, + ..Default::default() + }; + let conv = conv2d(in_c, in_c, 3, conv_cfg, vb.pp("conv"))?; + Ok(Self { conv }) + } +} + +impl candle::Module for Upsample { + fn forward(&self, xs: &Tensor) -> Result { + let (_, _, h, w) = xs.dims4()?; + xs.upsample_nearest2d(h * 2, w * 2)?.apply(&self.conv) + } +} + +#[derive(Debug, Clone)] +struct DownBlock { + block: Vec, + downsample: Option, +} + +#[derive(Debug, Clone)] +pub struct Encoder { + conv_in: Conv2d, + mid_block_1: ResnetBlock, + mid_attn_1: AttnBlock, + mid_block_2: ResnetBlock, + norm_out: GroupNorm, + conv_out: Conv2d, + down: Vec, +} + +impl Encoder { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let conv_cfg = candle_nn::Conv2dConfig { + padding: 1, + ..Default::default() + }; + let mut block_in = cfg.ch; + let conv_in = conv2d(cfg.in_channels, block_in, 3, conv_cfg, vb.pp("conv_in"))?; + + let mut down = Vec::with_capacity(cfg.ch_mult.len()); + let vb_d = vb.pp("down"); + for (i_level, ch_mult) in cfg.ch_mult.iter().enumerate() { + let mut block = Vec::with_capacity(cfg.num_res_blocks); + let vb_d = vb_d.pp(i_level); + let vb_b = vb_d.pp("block"); + let in_ch_mult = if i_level == 0 { + 1 + } else { + cfg.ch_mult[i_level - 1] + }; + block_in = cfg.ch * in_ch_mult; + let block_out = cfg.ch * ch_mult; + for i_block in 0..cfg.num_res_blocks { + let b = ResnetBlock::new(block_in, block_out, vb_b.pp(i_block))?; + block.push(b); + block_in = block_out; + } + let downsample = if i_level != cfg.ch_mult.len() - 1 { + Some(Downsample::new(block_in, vb_d.pp("downsample"))?) + } else { + None + }; + let block = DownBlock { block, downsample }; + down.push(block) + } + + let mid_block_1 = ResnetBlock::new(block_in, block_in, vb.pp("mid.block_1"))?; + let mid_attn_1 = AttnBlock::new(block_in, vb.pp("mid.attn_1"))?; + let mid_block_2 = ResnetBlock::new(block_in, block_in, vb.pp("mid.block_2"))?; + let conv_out = conv2d(block_in, 2 * cfg.z_channels, 3, conv_cfg, vb.pp("conv_out"))?; + let norm_out = group_norm(32, block_in, 1e-6, vb.pp("norm_out"))?; + Ok(Self { + conv_in, + mid_block_1, + mid_attn_1, + mid_block_2, + norm_out, + conv_out, + down, + }) + } +} + +impl candle_nn::Module for Encoder { + fn forward(&self, xs: &Tensor) -> Result { + let mut h = xs.apply(&self.conv_in)?; + for block in self.down.iter() { + for b in block.block.iter() { + h = h.apply(b)? + } + if let Some(ds) = block.downsample.as_ref() { + h = h.apply(ds)? + } + } + h.apply(&self.mid_block_1)? + .apply(&self.mid_attn_1)? + .apply(&self.mid_block_2)? + .apply(&self.norm_out)? + .apply(&candle_nn::Activation::Swish)? + .apply(&self.conv_out) + } +} + +#[derive(Debug, Clone)] +struct UpBlock { + block: Vec, + upsample: Option, +} + +#[derive(Debug, Clone)] +pub struct Decoder { + conv_in: Conv2d, + mid_block_1: ResnetBlock, + mid_attn_1: AttnBlock, + mid_block_2: ResnetBlock, + norm_out: GroupNorm, + conv_out: Conv2d, + up: Vec, +} + +impl Decoder { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let conv_cfg = candle_nn::Conv2dConfig { + padding: 1, + ..Default::default() + }; + let mut block_in = cfg.ch * cfg.ch_mult.last().unwrap_or(&1); + let conv_in = conv2d(cfg.z_channels, block_in, 3, conv_cfg, vb.pp("conv_in"))?; + let mid_block_1 = ResnetBlock::new(block_in, block_in, vb.pp("mid.block_1"))?; + let mid_attn_1 = AttnBlock::new(block_in, vb.pp("mid.attn_1"))?; + let mid_block_2 = ResnetBlock::new(block_in, block_in, vb.pp("mid.block_2"))?; + + let mut up = Vec::with_capacity(cfg.ch_mult.len()); + let vb_u = vb.pp("up"); + for (i_level, ch_mult) in cfg.ch_mult.iter().enumerate().rev() { + let block_out = cfg.ch * ch_mult; + let vb_u = vb_u.pp(i_level); + let vb_b = vb_u.pp("block"); + let mut block = Vec::with_capacity(cfg.num_res_blocks + 1); + for i_block in 0..=cfg.num_res_blocks { + let b = ResnetBlock::new(block_in, block_out, vb_b.pp(i_block))?; + block.push(b); + block_in = block_out; + } + let upsample = if i_level != 0 { + Some(Upsample::new(block_in, vb_u.pp("upsample"))?) + } else { + None + }; + let block = UpBlock { block, upsample }; + up.push(block) + } + up.reverse(); + + let norm_out = group_norm(32, block_in, 1e-6, vb.pp("norm_out"))?; + let conv_out = conv2d(block_in, cfg.out_ch, 3, conv_cfg, vb.pp("conv_out"))?; + Ok(Self { + conv_in, + mid_block_1, + mid_attn_1, + mid_block_2, + norm_out, + conv_out, + up, + }) + } +} + +impl candle_nn::Module for Decoder { + fn forward(&self, xs: &Tensor) -> Result { + let h = xs.apply(&self.conv_in)?; + let mut h = h + .apply(&self.mid_block_1)? + .apply(&self.mid_attn_1)? + .apply(&self.mid_block_2)?; + for block in self.up.iter().rev() { + for b in block.block.iter() { + h = h.apply(b)? + } + if let Some(us) = block.upsample.as_ref() { + h = h.apply(us)? + } + } + h.apply(&self.norm_out)? + .apply(&candle_nn::Activation::Swish)? + .apply(&self.conv_out) + } +} + +#[derive(Debug, Clone)] +pub struct DiagonalGaussian { + sample: bool, + chunk_dim: usize, +} + +impl DiagonalGaussian { + pub fn new(sample: bool, chunk_dim: usize) -> Result { + Ok(Self { sample, chunk_dim }) + } +} + +impl candle_nn::Module for DiagonalGaussian { + fn forward(&self, xs: &Tensor) -> Result { + let chunks = xs.chunk(2, self.chunk_dim)?; + if self.sample { + let std = (&chunks[1] * 0.5)?.exp()?; + &chunks[0] + (std * chunks[0].randn_like(0., 1.))? + } else { + Ok(chunks[0].clone()) + } + } +} + +#[derive(Debug, Clone)] +pub struct AutoEncoder { + encoder: Encoder, + decoder: Decoder, + reg: DiagonalGaussian, + shift_factor: f64, + scale_factor: f64, +} + +impl AutoEncoder { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let encoder = Encoder::new(cfg, vb.pp("encoder"))?; + let decoder = Decoder::new(cfg, vb.pp("decoder"))?; + let reg = DiagonalGaussian::new(true, 1)?; + Ok(Self { + encoder, + decoder, + reg, + scale_factor: cfg.scale_factor, + shift_factor: cfg.shift_factor, + }) + } + + pub fn encode(&self, xs: &Tensor) -> Result { + let z = xs.apply(&self.encoder)?.apply(&self.reg)?; + (z - self.shift_factor)? * self.scale_factor + } + pub fn decode(&self, xs: &Tensor) -> Result { + let xs = ((xs / self.scale_factor)? + self.shift_factor)?; + xs.apply(&self.decoder) + } +} + +impl candle::Module for AutoEncoder { + fn forward(&self, xs: &Tensor) -> Result { + self.decode(&self.encode(xs)?) + } +} diff --git a/patches/candle-transformers/src/models/flux/mod.rs b/patches/candle-transformers/src/models/flux/mod.rs new file mode 100644 index 0000000000..1d2fa4ef33 --- /dev/null +++ b/patches/candle-transformers/src/models/flux/mod.rs @@ -0,0 +1,43 @@ +//! Flux Model +//! +//! Flux is a 12B rectified flow transformer capable of generating images from text descriptions. +//! +//! - 🤗 [Hugging Face Model](https://huggingface.co/black-forest-labs/FLUX.1-schnell) +//! - 💻 [GitHub Repository](https://github.com/black-forest-labs/flux) +//! - 📝 [Blog Post](https://blackforestlabs.ai/announcing-black-forest-labs/) +//! +//! # Usage +//! +//! ```bash +//! cargo run --features cuda \ +//! --example flux -r -- \ +//! --height 1024 --width 1024 \ +//! --prompt "a rusty robot walking on a beach holding a small torch, \ +//! the robot has the word \"rust\" written on it, high quality, 4k" +//! ``` +//! +//!
+//! +//!
+//! + +use candle::{Result, Tensor}; + +pub trait WithForward { + #[allow(clippy::too_many_arguments)] + fn forward( + &self, + img: &Tensor, + img_ids: &Tensor, + txt: &Tensor, + txt_ids: &Tensor, + timesteps: &Tensor, + y: &Tensor, + guidance: Option<&Tensor>, + ) -> Result; +} + +pub mod autoencoder; +pub mod model; +pub mod quantized_model; +pub mod sampling; diff --git a/patches/candle-transformers/src/models/flux/model.rs b/patches/candle-transformers/src/models/flux/model.rs new file mode 100644 index 0000000000..17b4eb2532 --- /dev/null +++ b/patches/candle-transformers/src/models/flux/model.rs @@ -0,0 +1,626 @@ +use candle::{DType, IndexOp, Result, Tensor, D}; +use candle_nn::{LayerNorm, Linear, RmsNorm, VarBuilder}; + +// https://github.com/black-forest-labs/flux/blob/727e3a71faf37390f318cf9434f0939653302b60/src/flux/model.py#L12 +#[derive(Debug, Clone)] +pub struct Config { + pub in_channels: usize, + pub vec_in_dim: usize, + pub context_in_dim: usize, + pub hidden_size: usize, + pub mlp_ratio: f64, + pub num_heads: usize, + pub depth: usize, + pub depth_single_blocks: usize, + pub axes_dim: Vec, + pub theta: usize, + pub qkv_bias: bool, + pub guidance_embed: bool, +} + +impl Config { + // https://github.com/black-forest-labs/flux/blob/727e3a71faf37390f318cf9434f0939653302b60/src/flux/util.py#L32 + pub fn dev() -> Self { + Self { + in_channels: 64, + vec_in_dim: 768, + context_in_dim: 4096, + hidden_size: 3072, + mlp_ratio: 4.0, + num_heads: 24, + depth: 19, + depth_single_blocks: 38, + axes_dim: vec![16, 56, 56], + theta: 10_000, + qkv_bias: true, + guidance_embed: true, + } + } + + // https://github.com/black-forest-labs/flux/blob/727e3a71faf37390f318cf9434f0939653302b60/src/flux/util.py#L64 + pub fn schnell() -> Self { + Self { + in_channels: 64, + vec_in_dim: 768, + context_in_dim: 4096, + hidden_size: 3072, + mlp_ratio: 4.0, + num_heads: 24, + depth: 19, + depth_single_blocks: 38, + axes_dim: vec![16, 56, 56], + theta: 10_000, + qkv_bias: true, + guidance_embed: false, + } + } +} + +fn layer_norm(dim: usize, vb: VarBuilder) -> Result { + let ws = Tensor::ones(dim, vb.dtype(), vb.device())?; + Ok(LayerNorm::new_no_bias(ws, 1e-6)) +} + +fn scaled_dot_product_attention(q: &Tensor, k: &Tensor, v: &Tensor) -> Result { + let dim = q.dim(D::Minus1)?; + let scale_factor = 1.0 / (dim as f64).sqrt(); + let mut batch_dims = q.dims().to_vec(); + batch_dims.pop(); + batch_dims.pop(); + let q = q.flatten_to(batch_dims.len() - 1)?; + let k = k.flatten_to(batch_dims.len() - 1)?; + let v = v.flatten_to(batch_dims.len() - 1)?; + let attn_weights = (q.matmul(&k.t()?)? * scale_factor)?; + let attn_scores = candle_nn::ops::softmax_last_dim(&attn_weights)?.matmul(&v)?; + batch_dims.push(attn_scores.dim(D::Minus2)?); + batch_dims.push(attn_scores.dim(D::Minus1)?); + attn_scores.reshape(batch_dims) +} + +fn rope(pos: &Tensor, dim: usize, theta: usize) -> Result { + if dim % 2 == 1 { + candle::bail!("dim {dim} is odd") + } + let dev = pos.device(); + let theta = theta as f64; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, 1, inv_freq_len), dev)?; + let inv_freq = inv_freq.to_dtype(pos.dtype())?; + let freqs = pos.unsqueeze(2)?.broadcast_mul(&inv_freq)?; + let cos = freqs.cos()?; + let sin = freqs.sin()?; + let out = Tensor::stack(&[&cos, &sin.neg()?, &sin, &cos], 3)?; + let (b, n, d, _ij) = out.dims4()?; + out.reshape((b, n, d, 2, 2)) +} + +fn apply_rope(x: &Tensor, freq_cis: &Tensor) -> Result { + let dims = x.dims(); + let (b_sz, n_head, seq_len, n_embd) = x.dims4()?; + 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 fr0 = freq_cis.get_on_dim(D::Minus1, 0)?; + let fr1 = freq_cis.get_on_dim(D::Minus1, 1)?; + (fr0.broadcast_mul(&x0)? + fr1.broadcast_mul(&x1)?)?.reshape(dims.to_vec()) +} + +pub(crate) fn attention(q: &Tensor, k: &Tensor, v: &Tensor, pe: &Tensor) -> Result { + let q = apply_rope(q, pe)?.contiguous()?; + let k = apply_rope(k, pe)?.contiguous()?; + let x = scaled_dot_product_attention(&q, &k, v)?; + x.transpose(1, 2)?.flatten_from(2) +} + +pub(crate) fn timestep_embedding(t: &Tensor, dim: usize, dtype: DType) -> Result { + const TIME_FACTOR: f64 = 1000.; + const MAX_PERIOD: f64 = 10000.; + if dim % 2 == 1 { + candle::bail!("{dim} is odd") + } + let dev = t.device(); + let half = dim / 2; + let t = (t * TIME_FACTOR)?; + let arange = Tensor::arange(0, half as u32, dev)?.to_dtype(candle::DType::F32)?; + let freqs = (arange * (-MAX_PERIOD.ln() / half as f64))?.exp()?; + let args = t + .unsqueeze(1)? + .to_dtype(candle::DType::F32)? + .broadcast_mul(&freqs.unsqueeze(0)?)?; + let emb = Tensor::cat(&[args.cos()?, args.sin()?], D::Minus1)?.to_dtype(dtype)?; + Ok(emb) +} + +#[derive(Debug, Clone)] +pub struct EmbedNd { + #[allow(unused)] + dim: usize, + theta: usize, + axes_dim: Vec, +} + +impl EmbedNd { + pub fn new(dim: usize, theta: usize, axes_dim: Vec) -> Self { + Self { + dim, + theta, + axes_dim, + } + } +} + +impl candle::Module for EmbedNd { + fn forward(&self, ids: &Tensor) -> Result { + let n_axes = ids.dim(D::Minus1)?; + let mut emb = Vec::with_capacity(n_axes); + for idx in 0..n_axes { + let r = rope( + &ids.get_on_dim(D::Minus1, idx)?, + self.axes_dim[idx], + self.theta, + )?; + emb.push(r) + } + let emb = Tensor::cat(&emb, 2)?; + emb.unsqueeze(1) + } +} + +#[derive(Debug, Clone)] +pub struct MlpEmbedder { + in_layer: Linear, + out_layer: Linear, +} + +impl MlpEmbedder { + fn new(in_sz: usize, h_sz: usize, vb: VarBuilder) -> Result { + let in_layer = candle_nn::linear(in_sz, h_sz, vb.pp("in_layer"))?; + let out_layer = candle_nn::linear(h_sz, h_sz, vb.pp("out_layer"))?; + Ok(Self { + in_layer, + out_layer, + }) + } +} + +impl candle::Module for MlpEmbedder { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.in_layer)?.silu()?.apply(&self.out_layer) + } +} + +#[derive(Debug, Clone)] +pub struct QkNorm { + query_norm: RmsNorm, + key_norm: RmsNorm, +} + +impl QkNorm { + fn new(dim: usize, vb: VarBuilder) -> Result { + let query_norm = vb.get(dim, "query_norm.scale")?; + let query_norm = RmsNorm::new(query_norm, 1e-6); + let key_norm = vb.get(dim, "key_norm.scale")?; + let key_norm = RmsNorm::new(key_norm, 1e-6); + Ok(Self { + query_norm, + key_norm, + }) + } +} + +struct ModulationOut { + shift: Tensor, + scale: Tensor, + gate: Tensor, +} + +impl ModulationOut { + fn scale_shift(&self, xs: &Tensor) -> Result { + xs.broadcast_mul(&(&self.scale + 1.)?)? + .broadcast_add(&self.shift) + } + + fn gate(&self, xs: &Tensor) -> Result { + self.gate.broadcast_mul(xs) + } +} + +#[derive(Debug, Clone)] +struct Modulation1 { + lin: Linear, +} + +impl Modulation1 { + fn new(dim: usize, vb: VarBuilder) -> Result { + let lin = candle_nn::linear(dim, 3 * dim, vb.pp("lin"))?; + Ok(Self { lin }) + } + + fn forward(&self, vec_: &Tensor) -> Result { + let ys = vec_ + .silu()? + .apply(&self.lin)? + .unsqueeze(1)? + .chunk(3, D::Minus1)?; + if ys.len() != 3 { + candle::bail!("unexpected len from chunk {ys:?}") + } + Ok(ModulationOut { + shift: ys[0].clone(), + scale: ys[1].clone(), + gate: ys[2].clone(), + }) + } +} + +#[derive(Debug, Clone)] +struct Modulation2 { + lin: Linear, +} + +impl Modulation2 { + fn new(dim: usize, vb: VarBuilder) -> Result { + let lin = candle_nn::linear(dim, 6 * dim, vb.pp("lin"))?; + Ok(Self { lin }) + } + + fn forward(&self, vec_: &Tensor) -> Result<(ModulationOut, ModulationOut)> { + let ys = vec_ + .silu()? + .apply(&self.lin)? + .unsqueeze(1)? + .chunk(6, D::Minus1)?; + if ys.len() != 6 { + candle::bail!("unexpected len from chunk {ys:?}") + } + let mod1 = ModulationOut { + shift: ys[0].clone(), + scale: ys[1].clone(), + gate: ys[2].clone(), + }; + let mod2 = ModulationOut { + shift: ys[3].clone(), + scale: ys[4].clone(), + gate: ys[5].clone(), + }; + Ok((mod1, mod2)) + } +} + +#[derive(Debug, Clone)] +pub struct SelfAttention { + qkv: Linear, + norm: QkNorm, + proj: Linear, + num_heads: usize, +} + +impl SelfAttention { + fn new(dim: usize, num_heads: usize, qkv_bias: bool, vb: VarBuilder) -> Result { + let head_dim = dim / num_heads; + let qkv = candle_nn::linear_b(dim, dim * 3, qkv_bias, vb.pp("qkv"))?; + let norm = QkNorm::new(head_dim, vb.pp("norm"))?; + let proj = candle_nn::linear(dim, dim, vb.pp("proj"))?; + Ok(Self { + qkv, + norm, + proj, + num_heads, + }) + } + + fn qkv(&self, xs: &Tensor) -> Result<(Tensor, Tensor, Tensor)> { + let qkv = xs.apply(&self.qkv)?; + let (b, l, _khd) = qkv.dims3()?; + let qkv = qkv.reshape((b, l, 3, self.num_heads, ()))?; + let q = qkv.i((.., .., 0))?.transpose(1, 2)?; + let k = qkv.i((.., .., 1))?.transpose(1, 2)?; + let v = qkv.i((.., .., 2))?.transpose(1, 2)?; + let q = q.apply(&self.norm.query_norm)?; + let k = k.apply(&self.norm.key_norm)?; + Ok((q, k, v)) + } + + #[allow(unused)] + fn forward(&self, xs: &Tensor, pe: &Tensor) -> Result { + let (q, k, v) = self.qkv(xs)?; + attention(&q, &k, &v, pe)?.apply(&self.proj) + } +} + +#[derive(Debug, Clone)] +struct Mlp { + lin1: Linear, + lin2: Linear, +} + +impl Mlp { + fn new(in_sz: usize, mlp_sz: usize, vb: VarBuilder) -> Result { + let lin1 = candle_nn::linear(in_sz, mlp_sz, vb.pp("0"))?; + let lin2 = candle_nn::linear(mlp_sz, in_sz, vb.pp("2"))?; + Ok(Self { lin1, lin2 }) + } +} + +impl candle::Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.lin1)?.gelu()?.apply(&self.lin2) + } +} + +#[derive(Debug, Clone)] +pub struct DoubleStreamBlock { + img_mod: Modulation2, + img_norm1: LayerNorm, + img_attn: SelfAttention, + img_norm2: LayerNorm, + img_mlp: Mlp, + txt_mod: Modulation2, + txt_norm1: LayerNorm, + txt_attn: SelfAttention, + txt_norm2: LayerNorm, + txt_mlp: Mlp, +} + +impl DoubleStreamBlock { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let h_sz = cfg.hidden_size; + let mlp_sz = (h_sz as f64 * cfg.mlp_ratio) as usize; + let img_mod = Modulation2::new(h_sz, vb.pp("img_mod"))?; + let img_norm1 = layer_norm(h_sz, vb.pp("img_norm1"))?; + let img_attn = SelfAttention::new(h_sz, cfg.num_heads, cfg.qkv_bias, vb.pp("img_attn"))?; + let img_norm2 = layer_norm(h_sz, vb.pp("img_norm2"))?; + let img_mlp = Mlp::new(h_sz, mlp_sz, vb.pp("img_mlp"))?; + let txt_mod = Modulation2::new(h_sz, vb.pp("txt_mod"))?; + let txt_norm1 = layer_norm(h_sz, vb.pp("txt_norm1"))?; + let txt_attn = SelfAttention::new(h_sz, cfg.num_heads, cfg.qkv_bias, vb.pp("txt_attn"))?; + let txt_norm2 = layer_norm(h_sz, vb.pp("txt_norm2"))?; + let txt_mlp = Mlp::new(h_sz, mlp_sz, vb.pp("txt_mlp"))?; + Ok(Self { + img_mod, + img_norm1, + img_attn, + img_norm2, + img_mlp, + txt_mod, + txt_norm1, + txt_attn, + txt_norm2, + txt_mlp, + }) + } + + fn forward( + &self, + img: &Tensor, + txt: &Tensor, + vec_: &Tensor, + pe: &Tensor, + ) -> Result<(Tensor, Tensor)> { + let (img_mod1, img_mod2) = self.img_mod.forward(vec_)?; // shift, scale, gate + let (txt_mod1, txt_mod2) = self.txt_mod.forward(vec_)?; // shift, scale, gate + let img_modulated = img.apply(&self.img_norm1)?; + let img_modulated = img_mod1.scale_shift(&img_modulated)?; + let (img_q, img_k, img_v) = self.img_attn.qkv(&img_modulated)?; + + let txt_modulated = txt.apply(&self.txt_norm1)?; + let txt_modulated = txt_mod1.scale_shift(&txt_modulated)?; + let (txt_q, txt_k, txt_v) = self.txt_attn.qkv(&txt_modulated)?; + + let q = Tensor::cat(&[txt_q, img_q], 2)?; + let k = Tensor::cat(&[txt_k, img_k], 2)?; + let v = Tensor::cat(&[txt_v, img_v], 2)?; + + let attn = attention(&q, &k, &v, pe)?; + let txt_attn = attn.narrow(1, 0, txt.dim(1)?)?; + let img_attn = attn.narrow(1, txt.dim(1)?, attn.dim(1)? - txt.dim(1)?)?; + + let img = (img + img_mod1.gate(&img_attn.apply(&self.img_attn.proj)?))?; + let img = (&img + + img_mod2.gate( + &img_mod2 + .scale_shift(&img.apply(&self.img_norm2)?)? + .apply(&self.img_mlp)?, + )?)?; + + let txt = (txt + txt_mod1.gate(&txt_attn.apply(&self.txt_attn.proj)?))?; + let txt = (&txt + + txt_mod2.gate( + &txt_mod2 + .scale_shift(&txt.apply(&self.txt_norm2)?)? + .apply(&self.txt_mlp)?, + )?)?; + + Ok((img, txt)) + } +} + +#[derive(Debug, Clone)] +pub struct SingleStreamBlock { + linear1: Linear, + linear2: Linear, + norm: QkNorm, + pre_norm: LayerNorm, + modulation: Modulation1, + h_sz: usize, + mlp_sz: usize, + num_heads: usize, +} + +impl SingleStreamBlock { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let h_sz = cfg.hidden_size; + let mlp_sz = (h_sz as f64 * cfg.mlp_ratio) as usize; + let head_dim = h_sz / cfg.num_heads; + let linear1 = candle_nn::linear(h_sz, h_sz * 3 + mlp_sz, vb.pp("linear1"))?; + let linear2 = candle_nn::linear(h_sz + mlp_sz, h_sz, vb.pp("linear2"))?; + let norm = QkNorm::new(head_dim, vb.pp("norm"))?; + let pre_norm = layer_norm(h_sz, vb.pp("pre_norm"))?; + let modulation = Modulation1::new(h_sz, vb.pp("modulation"))?; + Ok(Self { + linear1, + linear2, + norm, + pre_norm, + modulation, + h_sz, + mlp_sz, + num_heads: cfg.num_heads, + }) + } + + fn forward(&self, xs: &Tensor, vec_: &Tensor, pe: &Tensor) -> Result { + let mod_ = self.modulation.forward(vec_)?; + let x_mod = mod_.scale_shift(&xs.apply(&self.pre_norm)?)?; + let x_mod = x_mod.apply(&self.linear1)?; + let qkv = x_mod.narrow(D::Minus1, 0, 3 * self.h_sz)?; + let (b, l, _khd) = qkv.dims3()?; + let qkv = qkv.reshape((b, l, 3, self.num_heads, ()))?; + let q = qkv.i((.., .., 0))?.transpose(1, 2)?; + let k = qkv.i((.., .., 1))?.transpose(1, 2)?; + let v = qkv.i((.., .., 2))?.transpose(1, 2)?; + let mlp = x_mod.narrow(D::Minus1, 3 * self.h_sz, self.mlp_sz)?; + let q = q.apply(&self.norm.query_norm)?; + let k = k.apply(&self.norm.key_norm)?; + let attn = attention(&q, &k, &v, pe)?; + let output = Tensor::cat(&[attn, mlp.gelu()?], 2)?.apply(&self.linear2)?; + xs + mod_.gate(&output) + } +} + +#[derive(Debug, Clone)] +pub struct LastLayer { + norm_final: LayerNorm, + linear: Linear, + ada_ln_modulation: Linear, +} + +impl LastLayer { + fn new(h_sz: usize, p_sz: usize, out_c: usize, vb: VarBuilder) -> Result { + let norm_final = layer_norm(h_sz, vb.pp("norm_final"))?; + let linear = candle_nn::linear(h_sz, p_sz * p_sz * out_c, vb.pp("linear"))?; + let ada_ln_modulation = candle_nn::linear(h_sz, 2 * h_sz, vb.pp("adaLN_modulation.1"))?; + Ok(Self { + norm_final, + linear, + ada_ln_modulation, + }) + } + + fn forward(&self, xs: &Tensor, vec: &Tensor) -> Result { + let chunks = vec.silu()?.apply(&self.ada_ln_modulation)?.chunk(2, 1)?; + let (shift, scale) = (&chunks[0], &chunks[1]); + let xs = xs + .apply(&self.norm_final)? + .broadcast_mul(&(scale.unsqueeze(1)? + 1.0)?)? + .broadcast_add(&shift.unsqueeze(1)?)?; + xs.apply(&self.linear) + } +} + +#[derive(Debug, Clone)] +pub struct Flux { + img_in: Linear, + txt_in: Linear, + time_in: MlpEmbedder, + vector_in: MlpEmbedder, + guidance_in: Option, + pe_embedder: EmbedNd, + double_blocks: Vec, + single_blocks: Vec, + final_layer: LastLayer, +} + +impl Flux { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let img_in = candle_nn::linear(cfg.in_channels, cfg.hidden_size, vb.pp("img_in"))?; + let txt_in = candle_nn::linear(cfg.context_in_dim, cfg.hidden_size, vb.pp("txt_in"))?; + let mut double_blocks = Vec::with_capacity(cfg.depth); + let vb_d = vb.pp("double_blocks"); + for idx in 0..cfg.depth { + let db = DoubleStreamBlock::new(cfg, vb_d.pp(idx))?; + double_blocks.push(db) + } + let mut single_blocks = Vec::with_capacity(cfg.depth_single_blocks); + let vb_s = vb.pp("single_blocks"); + for idx in 0..cfg.depth_single_blocks { + let sb = SingleStreamBlock::new(cfg, vb_s.pp(idx))?; + single_blocks.push(sb) + } + let time_in = MlpEmbedder::new(256, cfg.hidden_size, vb.pp("time_in"))?; + let vector_in = MlpEmbedder::new(cfg.vec_in_dim, cfg.hidden_size, vb.pp("vector_in"))?; + let guidance_in = if cfg.guidance_embed { + let mlp = MlpEmbedder::new(256, cfg.hidden_size, vb.pp("guidance_in"))?; + Some(mlp) + } else { + None + }; + let final_layer = + LastLayer::new(cfg.hidden_size, 1, cfg.in_channels, vb.pp("final_layer"))?; + let pe_dim = cfg.hidden_size / cfg.num_heads; + let pe_embedder = EmbedNd::new(pe_dim, cfg.theta, cfg.axes_dim.to_vec()); + Ok(Self { + img_in, + txt_in, + time_in, + vector_in, + guidance_in, + pe_embedder, + double_blocks, + single_blocks, + final_layer, + }) + } +} + +impl super::WithForward for Flux { + #[allow(clippy::too_many_arguments)] + fn forward( + &self, + img: &Tensor, + img_ids: &Tensor, + txt: &Tensor, + txt_ids: &Tensor, + timesteps: &Tensor, + y: &Tensor, + guidance: Option<&Tensor>, + ) -> Result { + if txt.rank() != 3 { + candle::bail!("unexpected shape for txt {:?}", txt.shape()) + } + if img.rank() != 3 { + candle::bail!("unexpected shape for img {:?}", img.shape()) + } + let dtype = img.dtype(); + let pe = { + let ids = Tensor::cat(&[txt_ids, img_ids], 1)?; + ids.apply(&self.pe_embedder)? + }; + let mut txt = txt.apply(&self.txt_in)?; + let mut img = img.apply(&self.img_in)?; + let vec_ = timestep_embedding(timesteps, 256, dtype)?.apply(&self.time_in)?; + let vec_ = match (self.guidance_in.as_ref(), guidance) { + (Some(g_in), Some(guidance)) => { + (vec_ + timestep_embedding(guidance, 256, dtype)?.apply(g_in))? + } + _ => vec_, + }; + let vec_ = (vec_ + y.apply(&self.vector_in))?; + + // Double blocks + for block in self.double_blocks.iter() { + (img, txt) = block.forward(&img, &txt, &vec_, &pe)? + } + // Single blocks + let mut img = Tensor::cat(&[&txt, &img], 1)?; + for block in self.single_blocks.iter() { + img = block.forward(&img, &vec_, &pe)?; + } + let img = img.i((.., txt.dim(1)?..))?; + self.final_layer.forward(&img, &vec_) + } +} diff --git a/patches/candle-transformers/src/models/flux/quantized_model.rs b/patches/candle-transformers/src/models/flux/quantized_model.rs new file mode 100644 index 0000000000..0efeeab573 --- /dev/null +++ b/patches/candle-transformers/src/models/flux/quantized_model.rs @@ -0,0 +1,465 @@ +use super::model::{attention, timestep_embedding, Config, EmbedNd}; +use crate::quantized_nn::{linear, linear_b, Linear}; +use crate::quantized_var_builder::VarBuilder; +use candle::{DType, IndexOp, Result, Tensor, D}; +use candle_nn::{LayerNorm, RmsNorm}; + +fn layer_norm(dim: usize, vb: VarBuilder) -> Result { + let ws = Tensor::ones(dim, DType::F32, vb.device())?; + Ok(LayerNorm::new_no_bias(ws, 1e-6)) +} + +#[derive(Debug, Clone)] +pub struct MlpEmbedder { + in_layer: Linear, + out_layer: Linear, +} + +impl MlpEmbedder { + fn new(in_sz: usize, h_sz: usize, vb: VarBuilder) -> Result { + let in_layer = linear(in_sz, h_sz, vb.pp("in_layer"))?; + let out_layer = linear(h_sz, h_sz, vb.pp("out_layer"))?; + Ok(Self { + in_layer, + out_layer, + }) + } +} + +impl candle::Module for MlpEmbedder { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.in_layer)?.silu()?.apply(&self.out_layer) + } +} + +#[derive(Debug, Clone)] +pub struct QkNorm { + query_norm: RmsNorm, + key_norm: RmsNorm, +} + +impl QkNorm { + fn new(dim: usize, vb: VarBuilder) -> Result { + let query_norm = vb.get(dim, "query_norm.scale")?.dequantize(vb.device())?; + let query_norm = RmsNorm::new(query_norm, 1e-6); + let key_norm = vb.get(dim, "key_norm.scale")?.dequantize(vb.device())?; + let key_norm = RmsNorm::new(key_norm, 1e-6); + Ok(Self { + query_norm, + key_norm, + }) + } +} + +struct ModulationOut { + shift: Tensor, + scale: Tensor, + gate: Tensor, +} + +impl ModulationOut { + fn scale_shift(&self, xs: &Tensor) -> Result { + xs.broadcast_mul(&(&self.scale + 1.)?)? + .broadcast_add(&self.shift) + } + + fn gate(&self, xs: &Tensor) -> Result { + self.gate.broadcast_mul(xs) + } +} + +#[derive(Debug, Clone)] +struct Modulation1 { + lin: Linear, +} + +impl Modulation1 { + fn new(dim: usize, vb: VarBuilder) -> Result { + let lin = linear(dim, 3 * dim, vb.pp("lin"))?; + Ok(Self { lin }) + } + + fn forward(&self, vec_: &Tensor) -> Result { + let ys = vec_ + .silu()? + .apply(&self.lin)? + .unsqueeze(1)? + .chunk(3, D::Minus1)?; + if ys.len() != 3 { + candle::bail!("unexpected len from chunk {ys:?}") + } + Ok(ModulationOut { + shift: ys[0].clone(), + scale: ys[1].clone(), + gate: ys[2].clone(), + }) + } +} + +#[derive(Debug, Clone)] +struct Modulation2 { + lin: Linear, +} + +impl Modulation2 { + fn new(dim: usize, vb: VarBuilder) -> Result { + let lin = linear(dim, 6 * dim, vb.pp("lin"))?; + Ok(Self { lin }) + } + + fn forward(&self, vec_: &Tensor) -> Result<(ModulationOut, ModulationOut)> { + let ys = vec_ + .silu()? + .apply(&self.lin)? + .unsqueeze(1)? + .chunk(6, D::Minus1)?; + if ys.len() != 6 { + candle::bail!("unexpected len from chunk {ys:?}") + } + let mod1 = ModulationOut { + shift: ys[0].clone(), + scale: ys[1].clone(), + gate: ys[2].clone(), + }; + let mod2 = ModulationOut { + shift: ys[3].clone(), + scale: ys[4].clone(), + gate: ys[5].clone(), + }; + Ok((mod1, mod2)) + } +} + +#[derive(Debug, Clone)] +pub struct SelfAttention { + qkv: Linear, + norm: QkNorm, + proj: Linear, + num_heads: usize, +} + +impl SelfAttention { + fn new(dim: usize, num_heads: usize, qkv_bias: bool, vb: VarBuilder) -> Result { + let head_dim = dim / num_heads; + let qkv = linear_b(dim, dim * 3, qkv_bias, vb.pp("qkv"))?; + let norm = QkNorm::new(head_dim, vb.pp("norm"))?; + let proj = linear(dim, dim, vb.pp("proj"))?; + Ok(Self { + qkv, + norm, + proj, + num_heads, + }) + } + + fn qkv(&self, xs: &Tensor) -> Result<(Tensor, Tensor, Tensor)> { + let qkv = xs.apply(&self.qkv)?; + let (b, l, _khd) = qkv.dims3()?; + let qkv = qkv.reshape((b, l, 3, self.num_heads, ()))?; + let q = qkv.i((.., .., 0))?.transpose(1, 2)?; + let k = qkv.i((.., .., 1))?.transpose(1, 2)?; + let v = qkv.i((.., .., 2))?.transpose(1, 2)?; + let q = q.apply(&self.norm.query_norm)?; + let k = k.apply(&self.norm.key_norm)?; + Ok((q, k, v)) + } + + #[allow(unused)] + fn forward(&self, xs: &Tensor, pe: &Tensor) -> Result { + let (q, k, v) = self.qkv(xs)?; + attention(&q, &k, &v, pe)?.apply(&self.proj) + } +} + +#[derive(Debug, Clone)] +struct Mlp { + lin1: Linear, + lin2: Linear, +} + +impl Mlp { + fn new(in_sz: usize, mlp_sz: usize, vb: VarBuilder) -> Result { + let lin1 = linear(in_sz, mlp_sz, vb.pp("0"))?; + let lin2 = linear(mlp_sz, in_sz, vb.pp("2"))?; + Ok(Self { lin1, lin2 }) + } +} + +impl candle::Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.lin1)?.gelu()?.apply(&self.lin2) + } +} + +#[derive(Debug, Clone)] +pub struct DoubleStreamBlock { + img_mod: Modulation2, + img_norm1: LayerNorm, + img_attn: SelfAttention, + img_norm2: LayerNorm, + img_mlp: Mlp, + txt_mod: Modulation2, + txt_norm1: LayerNorm, + txt_attn: SelfAttention, + txt_norm2: LayerNorm, + txt_mlp: Mlp, +} + +impl DoubleStreamBlock { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let h_sz = cfg.hidden_size; + let mlp_sz = (h_sz as f64 * cfg.mlp_ratio) as usize; + let img_mod = Modulation2::new(h_sz, vb.pp("img_mod"))?; + let img_norm1 = layer_norm(h_sz, vb.pp("img_norm1"))?; + let img_attn = SelfAttention::new(h_sz, cfg.num_heads, cfg.qkv_bias, vb.pp("img_attn"))?; + let img_norm2 = layer_norm(h_sz, vb.pp("img_norm2"))?; + let img_mlp = Mlp::new(h_sz, mlp_sz, vb.pp("img_mlp"))?; + let txt_mod = Modulation2::new(h_sz, vb.pp("txt_mod"))?; + let txt_norm1 = layer_norm(h_sz, vb.pp("txt_norm1"))?; + let txt_attn = SelfAttention::new(h_sz, cfg.num_heads, cfg.qkv_bias, vb.pp("txt_attn"))?; + let txt_norm2 = layer_norm(h_sz, vb.pp("txt_norm2"))?; + let txt_mlp = Mlp::new(h_sz, mlp_sz, vb.pp("txt_mlp"))?; + Ok(Self { + img_mod, + img_norm1, + img_attn, + img_norm2, + img_mlp, + txt_mod, + txt_norm1, + txt_attn, + txt_norm2, + txt_mlp, + }) + } + + fn forward( + &self, + img: &Tensor, + txt: &Tensor, + vec_: &Tensor, + pe: &Tensor, + ) -> Result<(Tensor, Tensor)> { + let (img_mod1, img_mod2) = self.img_mod.forward(vec_)?; // shift, scale, gate + let (txt_mod1, txt_mod2) = self.txt_mod.forward(vec_)?; // shift, scale, gate + let img_modulated = img.apply(&self.img_norm1)?; + let img_modulated = img_mod1.scale_shift(&img_modulated)?; + let (img_q, img_k, img_v) = self.img_attn.qkv(&img_modulated)?; + + let txt_modulated = txt.apply(&self.txt_norm1)?; + let txt_modulated = txt_mod1.scale_shift(&txt_modulated)?; + let (txt_q, txt_k, txt_v) = self.txt_attn.qkv(&txt_modulated)?; + + let q = Tensor::cat(&[txt_q, img_q], 2)?; + let k = Tensor::cat(&[txt_k, img_k], 2)?; + let v = Tensor::cat(&[txt_v, img_v], 2)?; + + let attn = attention(&q, &k, &v, pe)?; + let txt_attn = attn.narrow(1, 0, txt.dim(1)?)?; + let img_attn = attn.narrow(1, txt.dim(1)?, attn.dim(1)? - txt.dim(1)?)?; + + let img = (img + img_mod1.gate(&img_attn.apply(&self.img_attn.proj)?))?; + let img = (&img + + img_mod2.gate( + &img_mod2 + .scale_shift(&img.apply(&self.img_norm2)?)? + .apply(&self.img_mlp)?, + )?)?; + + let txt = (txt + txt_mod1.gate(&txt_attn.apply(&self.txt_attn.proj)?))?; + let txt = (&txt + + txt_mod2.gate( + &txt_mod2 + .scale_shift(&txt.apply(&self.txt_norm2)?)? + .apply(&self.txt_mlp)?, + )?)?; + + Ok((img, txt)) + } +} + +#[derive(Debug, Clone)] +pub struct SingleStreamBlock { + linear1: Linear, + linear2: Linear, + norm: QkNorm, + pre_norm: LayerNorm, + modulation: Modulation1, + h_sz: usize, + mlp_sz: usize, + num_heads: usize, +} + +impl SingleStreamBlock { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let h_sz = cfg.hidden_size; + let mlp_sz = (h_sz as f64 * cfg.mlp_ratio) as usize; + let head_dim = h_sz / cfg.num_heads; + let linear1 = linear(h_sz, h_sz * 3 + mlp_sz, vb.pp("linear1"))?; + let linear2 = linear(h_sz + mlp_sz, h_sz, vb.pp("linear2"))?; + let norm = QkNorm::new(head_dim, vb.pp("norm"))?; + let pre_norm = layer_norm(h_sz, vb.pp("pre_norm"))?; + let modulation = Modulation1::new(h_sz, vb.pp("modulation"))?; + Ok(Self { + linear1, + linear2, + norm, + pre_norm, + modulation, + h_sz, + mlp_sz, + num_heads: cfg.num_heads, + }) + } + + fn forward(&self, xs: &Tensor, vec_: &Tensor, pe: &Tensor) -> Result { + let mod_ = self.modulation.forward(vec_)?; + let x_mod = mod_.scale_shift(&xs.apply(&self.pre_norm)?)?; + let x_mod = x_mod.apply(&self.linear1)?; + let qkv = x_mod.narrow(D::Minus1, 0, 3 * self.h_sz)?; + let (b, l, _khd) = qkv.dims3()?; + let qkv = qkv.reshape((b, l, 3, self.num_heads, ()))?; + let q = qkv.i((.., .., 0))?.transpose(1, 2)?; + let k = qkv.i((.., .., 1))?.transpose(1, 2)?; + let v = qkv.i((.., .., 2))?.transpose(1, 2)?; + let mlp = x_mod.narrow(D::Minus1, 3 * self.h_sz, self.mlp_sz)?; + let q = q.apply(&self.norm.query_norm)?; + let k = k.apply(&self.norm.key_norm)?; + let attn = attention(&q, &k, &v, pe)?; + let output = Tensor::cat(&[attn, mlp.gelu()?], 2)?.apply(&self.linear2)?; + xs + mod_.gate(&output) + } +} + +#[derive(Debug, Clone)] +pub struct LastLayer { + norm_final: LayerNorm, + linear: Linear, + ada_ln_modulation: Linear, +} + +impl LastLayer { + fn new(h_sz: usize, p_sz: usize, out_c: usize, vb: VarBuilder) -> Result { + let norm_final = layer_norm(h_sz, vb.pp("norm_final"))?; + let linear_ = linear(h_sz, p_sz * p_sz * out_c, vb.pp("linear"))?; + let ada_ln_modulation = linear(h_sz, 2 * h_sz, vb.pp("adaLN_modulation.1"))?; + Ok(Self { + norm_final, + linear: linear_, + ada_ln_modulation, + }) + } + + fn forward(&self, xs: &Tensor, vec: &Tensor) -> Result { + let chunks = vec.silu()?.apply(&self.ada_ln_modulation)?.chunk(2, 1)?; + let (shift, scale) = (&chunks[0], &chunks[1]); + let xs = xs + .apply(&self.norm_final)? + .broadcast_mul(&(scale.unsqueeze(1)? + 1.0)?)? + .broadcast_add(&shift.unsqueeze(1)?)?; + xs.apply(&self.linear) + } +} + +#[derive(Debug, Clone)] +pub struct Flux { + img_in: Linear, + txt_in: Linear, + time_in: MlpEmbedder, + vector_in: MlpEmbedder, + guidance_in: Option, + pe_embedder: EmbedNd, + double_blocks: Vec, + single_blocks: Vec, + final_layer: LastLayer, +} + +impl Flux { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let img_in = linear(cfg.in_channels, cfg.hidden_size, vb.pp("img_in"))?; + let txt_in = linear(cfg.context_in_dim, cfg.hidden_size, vb.pp("txt_in"))?; + let mut double_blocks = Vec::with_capacity(cfg.depth); + let vb_d = vb.pp("double_blocks"); + for idx in 0..cfg.depth { + let db = DoubleStreamBlock::new(cfg, vb_d.pp(idx))?; + double_blocks.push(db) + } + let mut single_blocks = Vec::with_capacity(cfg.depth_single_blocks); + let vb_s = vb.pp("single_blocks"); + for idx in 0..cfg.depth_single_blocks { + let sb = SingleStreamBlock::new(cfg, vb_s.pp(idx))?; + single_blocks.push(sb) + } + let time_in = MlpEmbedder::new(256, cfg.hidden_size, vb.pp("time_in"))?; + let vector_in = MlpEmbedder::new(cfg.vec_in_dim, cfg.hidden_size, vb.pp("vector_in"))?; + let guidance_in = if cfg.guidance_embed { + let mlp = MlpEmbedder::new(256, cfg.hidden_size, vb.pp("guidance_in"))?; + Some(mlp) + } else { + None + }; + let final_layer = + LastLayer::new(cfg.hidden_size, 1, cfg.in_channels, vb.pp("final_layer"))?; + let pe_dim = cfg.hidden_size / cfg.num_heads; + let pe_embedder = EmbedNd::new(pe_dim, cfg.theta, cfg.axes_dim.to_vec()); + Ok(Self { + img_in, + txt_in, + time_in, + vector_in, + guidance_in, + pe_embedder, + double_blocks, + single_blocks, + final_layer, + }) + } +} + +impl super::WithForward for Flux { + #[allow(clippy::too_many_arguments)] + fn forward( + &self, + img: &Tensor, + img_ids: &Tensor, + txt: &Tensor, + txt_ids: &Tensor, + timesteps: &Tensor, + y: &Tensor, + guidance: Option<&Tensor>, + ) -> Result { + if txt.rank() != 3 { + candle::bail!("unexpected shape for txt {:?}", txt.shape()) + } + if img.rank() != 3 { + candle::bail!("unexpected shape for img {:?}", img.shape()) + } + let dtype = img.dtype(); + let pe = { + let ids = Tensor::cat(&[txt_ids, img_ids], 1)?; + ids.apply(&self.pe_embedder)? + }; + let mut txt = txt.apply(&self.txt_in)?; + let mut img = img.apply(&self.img_in)?; + let vec_ = timestep_embedding(timesteps, 256, dtype)?.apply(&self.time_in)?; + let vec_ = match (self.guidance_in.as_ref(), guidance) { + (Some(g_in), Some(guidance)) => { + (vec_ + timestep_embedding(guidance, 256, dtype)?.apply(g_in))? + } + _ => vec_, + }; + let vec_ = (vec_ + y.apply(&self.vector_in))?; + + // Double blocks + for block in self.double_blocks.iter() { + (img, txt) = block.forward(&img, &txt, &vec_, &pe)? + } + // Single blocks + let mut img = Tensor::cat(&[&txt, &img], 1)?; + for block in self.single_blocks.iter() { + img = block.forward(&img, &vec_, &pe)?; + } + let img = img.i((.., txt.dim(1)?..))?; + self.final_layer.forward(&img, &vec_) + } +} diff --git a/patches/candle-transformers/src/models/flux/sampling.rs b/patches/candle-transformers/src/models/flux/sampling.rs new file mode 100644 index 0000000000..cdfef043ed --- /dev/null +++ b/patches/candle-transformers/src/models/flux/sampling.rs @@ -0,0 +1,119 @@ +use candle::{Device, Result, Tensor}; + +pub fn get_noise( + num_samples: usize, + height: usize, + width: usize, + device: &Device, +) -> Result { + let height = height.div_ceil(16) * 2; + let width = width.div_ceil(16) * 2; + Tensor::randn(0f32, 1., (num_samples, 16, height, width), device) +} + +#[derive(Debug, Clone)] +pub struct State { + pub img: Tensor, + pub img_ids: Tensor, + pub txt: Tensor, + pub txt_ids: Tensor, + pub vec: Tensor, +} + +impl State { + pub fn new(t5_emb: &Tensor, clip_emb: &Tensor, img: &Tensor) -> Result { + let dtype = img.dtype(); + let (bs, c, h, w) = img.dims4()?; + let dev = img.device(); + let img = img.reshape((bs, c, h / 2, 2, w / 2, 2))?; // (b, c, h, ph, w, pw) + let img = img.permute((0, 2, 4, 1, 3, 5))?; // (b, h, w, c, ph, pw) + let img = img.reshape((bs, h / 2 * w / 2, c * 4))?; + let img_ids = Tensor::stack( + &[ + Tensor::full(0u32, (h / 2, w / 2), dev)?, + Tensor::arange(0u32, h as u32 / 2, dev)? + .reshape(((), 1))? + .broadcast_as((h / 2, w / 2))?, + Tensor::arange(0u32, w as u32 / 2, dev)? + .reshape((1, ()))? + .broadcast_as((h / 2, w / 2))?, + ], + 2, + )? + .to_dtype(dtype)?; + let img_ids = img_ids.reshape((1, h / 2 * w / 2, 3))?; + let img_ids = img_ids.repeat((bs, 1, 1))?; + let txt = t5_emb.repeat(bs)?; + let txt_ids = Tensor::zeros((bs, txt.dim(1)?, 3), dtype, dev)?; + let vec = clip_emb.repeat(bs)?; + Ok(Self { + img, + img_ids, + txt, + txt_ids, + vec, + }) + } +} + +fn time_shift(mu: f64, sigma: f64, t: f64) -> f64 { + let e = mu.exp(); + e / (e + (1. / t - 1.).powf(sigma)) +} + +/// `shift` is a triple `(image_seq_len, base_shift, max_shift)`. +pub fn get_schedule(num_steps: usize, shift: Option<(usize, f64, f64)>) -> Vec { + let timesteps: Vec = (0..=num_steps) + .map(|v| v as f64 / num_steps as f64) + .rev() + .collect(); + match shift { + None => timesteps, + Some((image_seq_len, y1, y2)) => { + let (x1, x2) = (256., 4096.); + let m = (y2 - y1) / (x2 - x1); + let b = y1 - m * x1; + let mu = m * image_seq_len as f64 + b; + timesteps + .into_iter() + .map(|v| time_shift(mu, 1., v)) + .collect() + } + } +} + +pub fn unpack(xs: &Tensor, height: usize, width: usize) -> Result { + let (b, _h_w, c_ph_pw) = xs.dims3()?; + let height = height.div_ceil(16); + let width = width.div_ceil(16); + xs.reshape((b, height, width, c_ph_pw / 4, 2, 2))? // (b, h, w, c, ph, pw) + .permute((0, 3, 1, 4, 2, 5))? // (b, c, h, ph, w, pw) + .reshape((b, c_ph_pw / 4, height * 2, width * 2)) +} + +#[allow(clippy::too_many_arguments)] +pub fn denoise( + model: &M, + img: &Tensor, + img_ids: &Tensor, + txt: &Tensor, + txt_ids: &Tensor, + vec_: &Tensor, + timesteps: &[f64], + guidance: f64, +) -> Result { + let b_sz = img.dim(0)?; + let dev = img.device(); + let guidance = Tensor::full(guidance as f32, b_sz, dev)?; + let mut img = img.clone(); + for window in timesteps.windows(2) { + let (t_curr, t_prev) = match window { + [a, b] => (a, b), + _ => continue, + }; + let t_vec = Tensor::full(*t_curr as f32, b_sz, dev)?; + let pred = model.forward(&img, img_ids, txt, txt_ids, &t_vec, vec_, Some(&guidance))?; + img = (img + pred * (t_prev - t_curr))? + } + Ok(img) +} diff --git a/patches/candle-transformers/src/models/gemma.rs b/patches/candle-transformers/src/models/gemma.rs new file mode 100644 index 0000000000..4b656d6a7f --- /dev/null +++ b/patches/candle-transformers/src/models/gemma.rs @@ -0,0 +1,448 @@ +//! Gemma inference implementation. +//! +//! See ["Gemma: Open Models Based on Gemini Technology"](https://blog.google/technology/developers/gemma-open-ai-model/) +//! +//! Based on implementation from Google and PyTorch + +use std::sync::Arc; + +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::{linear_b as linear, Activation, Linear, VarBuilder}; + +fn default_max_position_embeddings() -> usize { + 4096 +} + +#[derive(serde::Deserialize, Debug, Clone)] +pub struct Config { + pub attention_bias: bool, + pub head_dim: usize, + // The code gemma configs include both hidden_act and hidden_activation. + pub hidden_act: Option, + pub hidden_activation: Option, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_attention_heads: usize, + pub num_hidden_layers: usize, + pub num_key_value_heads: usize, + pub rms_norm_eps: f64, + pub rope_theta: f64, + pub vocab_size: usize, + + #[serde(default = "default_max_position_embeddings")] + pub max_position_embeddings: usize, +} + +impl Config { + fn hidden_act(&self) -> Result { + match (self.hidden_act, self.hidden_activation) { + (None, Some(act)) | (Some(act), None) => Ok(act), + (Some(_), Some(_)) => candle::bail!("both hidden_act and hidden_activation are set"), + (None, None) => candle::bail!("none of hidden_act and hidden_activation are set"), + } + } +} + +#[derive(Debug, Clone)] +struct RmsNorm { + weight: Tensor, + eps: f64, +} + +impl RmsNorm { + fn new(dim: usize, eps: f64, vb: VarBuilder) -> Result { + let weight = vb.get(dim, "weight")?; + Ok(Self { weight, eps }) + } +} + +impl Module for RmsNorm { + fn forward(&self, x: &Tensor) -> Result { + let x_dtype = x.dtype(); + let internal_dtype = match x_dtype { + DType::F16 | DType::BF16 => DType::F32, + d => d, + }; + let hidden_size = x.dim(D::Minus1)?; + let x = x.to_dtype(internal_dtype)?; + let norm_x = (x.sqr()?.sum_keepdim(D::Minus1)? / hidden_size as f64)?; + let x_normed = x.broadcast_div(&(norm_x + self.eps)?.sqrt()?)?; + x_normed + .to_dtype(x_dtype)? + .broadcast_mul(&(&self.weight + 1.0)?) + } +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let dim = cfg.head_dim; + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: candle_nn::Activation, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + let gate_proj = linear(hidden_sz, intermediate_sz, false, vb.pp("gate_proj"))?; + let up_proj = linear(hidden_sz, intermediate_sz, false, vb.pp("up_proj"))?; + let down_proj = linear(intermediate_sz, hidden_sz, false, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_act()?, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = xs.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, + use_flash_attn: bool, +} + +impl Attention { + fn new( + rotary_emb: Arc, + use_flash_attn: bool, + cfg: &Config, + vb: VarBuilder, + ) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + let head_dim = cfg.head_dim; + let bias = cfg.attention_bias; + let q_proj = linear(hidden_sz, num_heads * head_dim, bias, vb.pp("q_proj"))?; + let k_proj = linear(hidden_sz, num_kv_heads * head_dim, bias, vb.pp("k_proj"))?; + let v_proj = linear(hidden_sz, num_kv_heads * head_dim, bias, vb.pp("v_proj"))?; + let o_proj = linear(num_heads * head_dim, hidden_sz, bias, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + rotary_emb, + kv_cache: None, + use_flash_attn, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let (query_states, key_states) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states, seqlen_offset)?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?.contiguous()?; + let value_states = + crate::utils::repeat_kv(value_states, self.num_kv_groups)?.contiguous()?; + + let attn_output = if self.use_flash_attn { + // flash-attn expects (b_sz, seq_len, nheads, head_dim) + let q = query_states.transpose(1, 2)?; + let k = key_states.transpose(1, 2)?; + let v = value_states.transpose(1, 2)?; + let scale = 1f32 / (self.head_dim as f32).sqrt(); + flash_attn(&q, &k, &v, scale, attention_mask.is_some())?.transpose(1, 2)? + } else { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + }; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, ()))? + .apply(&self.o_proj) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: MLP, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl DecoderLayer { + fn new( + rotary_emb: Arc, + use_flash_attn: bool, + cfg: &Config, + vb: VarBuilder, + ) -> Result { + let self_attn = Attention::new(rotary_emb, use_flash_attn, cfg, vb.pp("self_attn"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let input_layernorm = + RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let post_attention_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; + residual + xs + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: Linear, + device: Device, + dtype: DType, + hidden_size: usize, +} + +impl Model { + pub fn new(use_flash_attn: bool, cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb_m.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = + DecoderLayer::new(rotary_emb.clone(), use_flash_attn, cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb_m.pp("norm"))?; + let lm_head = Linear::new(embed_tokens.embeddings().clone(), None); + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + device: vb.device().clone(), + dtype: vb.dtype(), + hidden_size: cfg.hidden_size, + }) + } + + pub fn embed_tokens(&self) -> &candle_nn::Embedding { + &self.embed_tokens + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { f32::NEG_INFINITY } else { 0. })) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let (b_size, seq_len) = input_ids.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(b_size, seq_len, seqlen_offset)?; + Some(mask) + }; + let xs = self.embed_tokens.forward(input_ids)?; + let mut xs = (xs * (self.hidden_size as f64).sqrt())?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)? + } + xs.narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head) + } + pub fn forward_embeds( + &mut self, + xs: &Tensor, + attn_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (_, seq_len, _) = xs.dims3()?; + let mut xs = (xs * (self.hidden_size as f64).sqrt())?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attn_mask, seqlen_offset)? + } + xs.narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head) + } + + // Forward the model and return the hidden states without the lm_head + pub fn forward_embeds_without_projection( + &mut self, + xs: &Tensor, + attn_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (_, _, _) = xs.dims3()?; + let mut xs = (xs * (self.hidden_size as f64).sqrt())?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attn_mask, seqlen_offset)? + } + Ok(xs) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +} diff --git a/patches/candle-transformers/src/models/gemma2.rs b/patches/candle-transformers/src/models/gemma2.rs new file mode 100644 index 0000000000..ec23efc529 --- /dev/null +++ b/patches/candle-transformers/src/models/gemma2.rs @@ -0,0 +1,455 @@ +//! Gemma LLM architecture (Google) inference implementation. +//! +//! See ["Gemma: Open Models Based on Gemini Technology"](https://blog.google/technology/developers/gemma-open-models/) +//! +//! Based on implementations from Google and OpenLLM + +use std::sync::Arc; + +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::{linear_b as linear, Activation, Linear, VarBuilder}; + +fn default_max_position_embeddings() -> usize { + 4096 +} + +#[derive(serde::Deserialize, Debug, Clone)] +pub struct Config { + pub attention_bias: bool, + pub head_dim: usize, + pub hidden_activation: Activation, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_attention_heads: usize, + pub num_hidden_layers: usize, + pub num_key_value_heads: usize, + pub rms_norm_eps: f64, + pub rope_theta: f64, + pub vocab_size: usize, + pub final_logit_softcapping: Option, + pub attn_logit_softcapping: Option, + pub query_pre_attn_scalar: usize, + // TODO: Handle the sliding window in the attention mask. + pub sliding_window: Option, + + #[serde(default = "default_max_position_embeddings")] + pub max_position_embeddings: usize, +} + +#[derive(Debug, Clone)] +struct RmsNorm { + weight: Tensor, + eps: f64, +} + +impl RmsNorm { + fn new(dim: usize, eps: f64, vb: VarBuilder) -> Result { + let weight = vb.get(dim, "weight")?; + Ok(Self { weight, eps }) + } +} + +impl Module for RmsNorm { + fn forward(&self, x: &Tensor) -> Result { + let x_dtype = x.dtype(); + let internal_dtype = match x_dtype { + DType::F16 | DType::BF16 => DType::F32, + d => d, + }; + let hidden_size = x.dim(D::Minus1)?; + let x = x.to_dtype(internal_dtype)?; + let norm_x = (x.sqr()?.sum_keepdim(D::Minus1)? / hidden_size as f64)?; + let x_normed = x.broadcast_div(&(norm_x + self.eps)?.sqrt()?)?; + x_normed + .to_dtype(x_dtype)? + .broadcast_mul(&(&self.weight + 1.0)?) + } +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let dim = cfg.head_dim; + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: candle_nn::Activation, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + let gate_proj = linear(hidden_sz, intermediate_sz, false, vb.pp("gate_proj"))?; + let up_proj = linear(hidden_sz, intermediate_sz, false, vb.pp("up_proj"))?; + let down_proj = linear(intermediate_sz, hidden_sz, false, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_activation, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = xs.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + attn_logit_softcapping: Option, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, + use_flash_attn: bool, +} + +impl Attention { + fn new( + rotary_emb: Arc, + use_flash_attn: bool, + cfg: &Config, + vb: VarBuilder, + ) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + let head_dim = cfg.head_dim; + let bias = cfg.attention_bias; + let q_proj = linear(hidden_sz, num_heads * head_dim, bias, vb.pp("q_proj"))?; + let k_proj = linear(hidden_sz, num_kv_heads * head_dim, bias, vb.pp("k_proj"))?; + let v_proj = linear(hidden_sz, num_kv_heads * head_dim, bias, vb.pp("v_proj"))?; + let o_proj = linear(num_heads * head_dim, hidden_sz, bias, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + attn_logit_softcapping: cfg.attn_logit_softcapping, + rotary_emb, + kv_cache: None, + use_flash_attn, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let (query_states, key_states) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states, seqlen_offset)?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?.contiguous()?; + let value_states = + crate::utils::repeat_kv(value_states, self.num_kv_groups)?.contiguous()?; + + let attn_output = if self.use_flash_attn { + // flash-attn expects (b_sz, seq_len, nheads, head_dim) + let q = query_states.transpose(1, 2)?; + let k = key_states.transpose(1, 2)?; + let v = value_states.transpose(1, 2)?; + let scale = 1f32 / (self.head_dim as f32).sqrt(); + flash_attn(&q, &k, &v, scale, attention_mask.is_some())?.transpose(1, 2)? + } else { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match self.attn_logit_softcapping { + None => attn_weights, + Some(sc) => ((attn_weights / sc)?.tanh()? * sc)?, + }; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + }; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, ()))? + .apply(&self.o_proj) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: MLP, + input_layernorm: RmsNorm, + pre_feedforward_layernorm: RmsNorm, + post_feedforward_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl DecoderLayer { + fn new( + rotary_emb: Arc, + use_flash_attn: bool, + cfg: &Config, + vb: VarBuilder, + ) -> Result { + let self_attn = Attention::new(rotary_emb, use_flash_attn, cfg, vb.pp("self_attn"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let input_layernorm = + RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let pre_feedforward_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("pre_feedforward_layernorm"), + )?; + let post_feedforward_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_feedforward_layernorm"), + )?; + let post_attention_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + pre_feedforward_layernorm, + post_feedforward_layernorm, + post_attention_layernorm, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = xs.apply(&self.post_attention_layernorm)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.pre_feedforward_layernorm)?; + let xs = xs.apply(&self.mlp)?; + let xs = xs.apply(&self.post_feedforward_layernorm)?; + residual + xs + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: Linear, + final_logit_softcapping: Option, + device: Device, + dtype: DType, + hidden_size: usize, + sliding_window: Option, +} + +impl Model { + pub fn new(use_flash_attn: bool, cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb_m.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = + DecoderLayer::new(rotary_emb.clone(), use_flash_attn, cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb_m.pp("norm"))?; + let lm_head = Linear::new(embed_tokens.embeddings().clone(), None); + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + final_logit_softcapping: cfg.final_logit_softcapping, + device: vb.device().clone(), + dtype: vb.dtype(), + hidden_size: cfg.hidden_size, + sliding_window: cfg.sliding_window, + }) + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + let mask: Vec<_> = match self.sliding_window { + None => (0..tgt_len) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { f32::NEG_INFINITY } else { 0. })) + .collect(), + Some(sliding_window) => (0..tgt_len) + .flat_map(|i| { + (0..tgt_len).map(move |j| { + if i < j || j + sliding_window < i { + f32::NEG_INFINITY + } else { + 0. + } + }) + }) + .collect(), + }; + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let (b_size, seq_len) = input_ids.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(b_size, seq_len, seqlen_offset)?; + Some(mask) + }; + let xs = self.embed_tokens.forward(input_ids)?; + let mut xs = (xs * (self.hidden_size as f64).sqrt())?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)? + } + let logits = xs + .narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head)?; + let logits = match self.final_logit_softcapping { + None => logits, + Some(sc) => ((logits / sc)?.tanh()? * sc)?, + }; + + Ok(logits) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +} diff --git a/patches/candle-transformers/src/models/gemma3.rs b/patches/candle-transformers/src/models/gemma3.rs new file mode 100644 index 0000000000..08b4e5ad6e --- /dev/null +++ b/patches/candle-transformers/src/models/gemma3.rs @@ -0,0 +1,536 @@ +//! Gemma LLM architecture (Google) inference implementation. +//! +//! See ["Introducing Gemma 3: The most capable model you can run on a single GPU or TPU"](https://blog.google/technology/developers/gemma-3/) +//! +//! Based on implementations from HuggingFace transformers. + +use std::sync::Arc; + +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::{linear_b as linear, Activation, Linear, VarBuilder}; + +#[derive(serde::Deserialize, Debug, Clone)] +pub struct Config { + pub attention_bias: bool, + pub head_dim: usize, + pub hidden_activation: Activation, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_attention_heads: usize, + pub num_hidden_layers: usize, + pub num_key_value_heads: usize, + pub rms_norm_eps: f64, + pub rope_theta: f64, + pub rope_local_base_freq: f64, + pub vocab_size: usize, + pub final_logit_softcapping: Option, + pub attn_logit_softcapping: Option, + pub query_pre_attn_scalar: usize, + pub sliding_window: usize, + pub sliding_window_pattern: usize, + pub max_position_embeddings: usize, +} + +#[derive(Debug, Clone)] +struct RmsNorm { + weight: Tensor, + eps: f64, +} + +impl RmsNorm { + fn new(dim: usize, eps: f64, vb: VarBuilder) -> Result { + let weight = vb.get(dim, "weight")?; + Ok(Self { weight, eps }) + } +} + +impl Module for RmsNorm { + fn forward(&self, x: &Tensor) -> Result { + let x_dtype = x.dtype(); + let internal_dtype = match x_dtype { + DType::F16 | DType::BF16 => DType::F32, + d => d, + }; + let hidden_size = x.dim(D::Minus1)?; + let x = x.to_dtype(internal_dtype)?; + let norm_x = (x.sqr()?.sum_keepdim(D::Minus1)? / hidden_size as f64)?; + let x_normed = x.broadcast_div(&(norm_x + self.eps)?.sqrt()?)?; + x_normed + .to_dtype(x_dtype)? + .broadcast_mul(&(&self.weight + 1.0)?) + } +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new( + dtype: DType, + cfg: &Config, + dev: &Device, + sliding_window: Option, + ) -> Result { + let dim = cfg.head_dim; + let max_seq_len = cfg.max_position_embeddings; + let rope_freq = if sliding_window.is_some() { + cfg.rope_local_base_freq + } else { + cfg.rope_theta + }; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / rope_freq.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: candle_nn::Activation, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + let gate_proj = linear(hidden_sz, intermediate_sz, false, vb.pp("gate_proj"))?; + let up_proj = linear(hidden_sz, intermediate_sz, false, vb.pp("up_proj"))?; + let down_proj = linear(intermediate_sz, hidden_sz, false, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_activation, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = xs.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +enum KvCache { + Normal(candle_nn::kv_cache::KvCache), + Rotating(candle_nn::kv_cache::RotatingKvCache), +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + q_norm: RmsNorm, + k_norm: RmsNorm, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + attn_logit_softcapping: Option, + rotary_emb: Arc, + kv_cache: KvCache, + use_flash_attn: bool, +} + +impl Attention { + fn new( + rotary_emb: Arc, + use_flash_attn: bool, + cfg: &Config, + sliding_window: Option, + vb: VarBuilder, + ) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + let head_dim = cfg.head_dim; + let bias = cfg.attention_bias; + let q_proj = linear(hidden_sz, num_heads * head_dim, bias, vb.pp("q_proj"))?; + let k_proj = linear(hidden_sz, num_kv_heads * head_dim, bias, vb.pp("k_proj"))?; + let v_proj = linear(hidden_sz, num_kv_heads * head_dim, bias, vb.pp("v_proj"))?; + let o_proj = linear(num_heads * head_dim, hidden_sz, bias, vb.pp("o_proj"))?; + let q_norm = RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("q_norm"))?; + let k_norm = RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("k_norm"))?; + let kv_cache = if let Some(sliding_window) = sliding_window { + KvCache::Rotating(candle_nn::kv_cache::RotatingKvCache::new(2, sliding_window)) + } else { + KvCache::Normal(candle_nn::kv_cache::KvCache::new( + 2, + cfg.max_position_embeddings, + )) + }; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + q_norm, + k_norm, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + attn_logit_softcapping: cfg.attn_logit_softcapping, + rotary_emb, + kv_cache, + use_flash_attn, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let query_states = self.q_norm.forward(&query_states)?; + let key_states = self.k_norm.forward(&key_states)?; + + let (query_states, key_states) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states, seqlen_offset)?; + + let (key_states, value_states) = match &mut self.kv_cache { + KvCache::Normal(cache) => cache.append(&key_states, &value_states)?, + KvCache::Rotating(cache) => cache.append(&key_states, &value_states)?, + }; + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?.contiguous()?; + let value_states = + crate::utils::repeat_kv(value_states, self.num_kv_groups)?.contiguous()?; + + let attn_output = if self.use_flash_attn { + // flash-attn expects (b_sz, seq_len, nheads, head_dim) + let q = query_states.transpose(1, 2)?; + let k = key_states.transpose(1, 2)?; + let v = value_states.transpose(1, 2)?; + let scale = 1f32 / (self.head_dim as f32).sqrt(); + flash_attn(&q, &k, &v, scale, attention_mask.is_some())?.transpose(1, 2)? + } else { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match self.attn_logit_softcapping { + None => attn_weights, + Some(sc) => ((attn_weights / sc)?.tanh()? * sc)?, + }; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + }; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, ()))? + .apply(&self.o_proj) + } + + fn clear_kv_cache(&mut self) { + match &mut self.kv_cache { + KvCache::Normal(c) => c.reset(), + KvCache::Rotating(c) => c.reset(), + } + } +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: MLP, + input_layernorm: RmsNorm, + pre_feedforward_layernorm: RmsNorm, + post_feedforward_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, + sliding_window: Option, +} + +impl DecoderLayer { + fn new( + use_flash_attn: bool, + cfg: &Config, + vb: VarBuilder, + sliding_window: Option, + ) -> Result { + let rotary_emb = Arc::new(RotaryEmbedding::new( + vb.dtype(), + cfg, + vb.device(), + sliding_window, + )?); + let self_attn = Attention::new( + rotary_emb, + use_flash_attn, + cfg, + sliding_window, + vb.pp("self_attn"), + )?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let input_layernorm = + RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let pre_feedforward_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("pre_feedforward_layernorm"), + )?; + let post_feedforward_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_feedforward_layernorm"), + )?; + let post_attention_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + pre_feedforward_layernorm, + post_feedforward_layernorm, + post_attention_layernorm, + sliding_window, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = xs.apply(&self.post_attention_layernorm)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.pre_feedforward_layernorm)?; + let xs = xs.apply(&self.mlp)?; + let xs = xs.apply(&self.post_feedforward_layernorm)?; + residual + xs + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } +} + +fn prepare_decoder_attention_mask( + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + sliding_window: Option, + dtype: DType, + device: &Device, +) -> Result { + let mask: Vec<_> = if let Some(sliding_window) = sliding_window { + (0..tgt_len) + .flat_map(|i| { + (0..tgt_len).map(move |j| { + if i < j || j + sliding_window < i { + f32::NEG_INFINITY + } else { + 0. + } + }) + }) + .collect() + } else { + (0..tgt_len) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { f32::NEG_INFINITY } else { 0f32 })) + .collect() + }; + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(dtype) +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: Linear, + final_logit_softcapping: Option, + device: Device, + dtype: DType, + hidden_size: usize, + sliding_window: usize, +} + +impl Model { + pub fn new(use_flash_attn: bool, cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let sliding_window = (layer_idx + 1) % cfg.sliding_window_pattern > 0; + let layer = DecoderLayer::new( + use_flash_attn, + cfg, + vb_l.pp(layer_idx), + sliding_window.then_some(cfg.sliding_window), + )?; + layers.push(layer) + } + let norm = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb_m.pp("norm"))?; + let lm_head = Linear::new(embed_tokens.embeddings().clone(), None); + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + final_logit_softcapping: cfg.final_logit_softcapping, + device: vb.device().clone(), + dtype: vb.dtype(), + hidden_size: cfg.hidden_size, + sliding_window: cfg.sliding_window, + }) + } + + fn create_attention_masks( + &self, + batch_size: usize, + seq_len: usize, + seqlen_offset: usize, + ) -> Result<(Option, Option)> { + if seq_len <= 1 { + return Ok((None, None)); + } + + let mask = prepare_decoder_attention_mask( + batch_size, + seq_len, + seqlen_offset, + None, + self.dtype, + &self.device, + )?; + + let sliding_mask = prepare_decoder_attention_mask( + batch_size, + seq_len, + seqlen_offset, + Some(self.sliding_window), + self.dtype, + &self.device, + )?; + + Ok((Some(mask), Some(sliding_mask))) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let (b_size, seq_len) = input_ids.dims2()?; + let xs = self.embed_tokens.forward(input_ids)?; + let mut xs = (xs * (self.hidden_size as f64).sqrt())?; + + let (attention_mask, sliding_attention_mask) = + self.create_attention_masks(b_size, seq_len, seqlen_offset)?; + + for layer in self.layers.iter_mut() { + let mask = if layer.sliding_window.is_some() { + &sliding_attention_mask + } else { + &attention_mask + }; + xs = layer.forward(&xs, mask.as_ref(), seqlen_offset)? + } + let logits = xs + .narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head)?; + let logits = match self.final_logit_softcapping { + None => logits, + Some(sc) => ((logits / sc)?.tanh()? * sc)?, + }; + + Ok(logits) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +} diff --git a/patches/candle-transformers/src/models/glm4.rs b/patches/candle-transformers/src/models/glm4.rs new file mode 100644 index 0000000000..969325f2c9 --- /dev/null +++ b/patches/candle-transformers/src/models/glm4.rs @@ -0,0 +1,633 @@ +//! GLM-4 inference implementation. +//! +//! An open bilingual language model with 130B parameters. +//! +//! Based on implementation from [ChatGLM-6B](https://github.com/THUDM/ChatGLM-6B) + +use crate::models::with_tracing::{linear_b as linear, Linear}; +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::VarBuilder; +use serde::de::{self, Deserializer, Visitor}; +use serde::Deserialize; +use std::fmt; + +#[derive(Debug, Clone)] +pub enum EosTokenId { + Single(u32), + Multiple(Vec), +} + +impl<'de> Deserialize<'de> for EosTokenId { + fn deserialize(deserializer: D) -> std::result::Result + where + D: Deserializer<'de>, + { + struct EosTokenIdVisitor; + + impl<'de> Visitor<'de> for EosTokenIdVisitor { + type Value = EosTokenId; + + fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter.write_str("an integer or a list of integers") + } + + fn visit_u64(self, value: u64) -> std::result::Result + where + E: de::Error, + { + if value <= u32::MAX as u64 { + Ok(EosTokenId::Single(value as u32)) + } else { + Err(de::Error::custom("value too large for u32")) + } + } + + fn visit_seq(self, mut seq: A) -> std::result::Result + where + A: serde::de::SeqAccess<'de>, + { + let mut values = Vec::new(); + while let Some(value) = seq.next_element::()? { + values.push(value); + } + Ok(EosTokenId::Multiple(values)) + } + } + + deserializer.deserialize_any(EosTokenIdVisitor) + } +} + +fn default_one() -> usize { + 1 +} + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + pub num_layers: usize, + pub padded_vocab_size: usize, + pub hidden_size: usize, + pub ffn_hidden_size: usize, + pub kv_channels: usize, + pub num_attention_heads: usize, + pub seq_length: usize, + pub layernorm_epsilon: f64, + pub rmsnorm: bool, + pub apply_residual_connection_post_layernorm: bool, + pub post_layer_norm: bool, + pub add_bias_linear: bool, + pub add_qkv_bias: bool, + pub bias_dropout_fusion: bool, + pub multi_query_attention: bool, + pub multi_query_group_num: usize, + pub apply_query_key_layer_scaling: bool, + pub attention_softmax_in_fp32: bool, + pub fp32_residual_connection: bool, + #[serde(default = "default_one")] + pub rope_ratio: usize, + pub eos_token_id: Option, +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + cache: Tensor, +} + +impl RotaryEmbedding { + fn new(cfg: &Config, dtype: DType, dev: &Device) -> Result { + let rotary_dim = cfg.kv_channels; + let n_elem = rotary_dim / 2; + let base = 10_000f64 * cfg.rope_ratio as f64; + let inv_freq: Vec<_> = (0..n_elem) + .step_by(2) + .map(|i| 1f32 / base.powf(i as f64 / n_elem as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, cfg.seq_length as u32, dev)? + .to_dtype(dtype)? + .reshape((cfg.seq_length, 1))?; + let freqs = t.matmul(&inv_freq)?; + let cache = Tensor::stack(&[&freqs.cos()?, &freqs.sin()?], D::Minus1)?; + Ok(Self { cache }) + } + + fn apply(&self, xs: &Tensor, seqlen_offset: usize) -> Result { + let (seqlen, _b, np, _hn) = xs.dims4()?; + let cache = self.cache.narrow(0, seqlen_offset, seqlen)?; + let rot_dim = cache.dim(D::Minus2)? * 2; + let (xs, xs_pass) = ( + xs.narrow(D::Minus1, 0, rot_dim)?, + xs.narrow(D::Minus1, rot_dim, rot_dim)?, + ); + let xshaped = xs.reshape((seqlen, (), np, rot_dim / 2, 2))?; + let cache = cache.reshape((seqlen, (), 1, rot_dim / 2, 2))?; + let (xshaped0, xshaped1) = ( + xshaped.i((.., .., .., .., 0))?, + xshaped.i((.., .., .., .., 1))?, + ); + let (cache0, cache1) = (cache.i((.., .., .., .., 0))?, cache.i((.., .., .., .., 1))?); + let xs_out = Tensor::stack( + &[ + (xshaped0.broadcast_mul(&cache0)? - xshaped1.broadcast_mul(&cache1)?)?, + (xshaped1.broadcast_mul(&cache0)? + xshaped0.broadcast_mul(&cache1)?)?, + ], + D::Minus1, + )?; + let xs_out = xs_out.flatten_from(3)?; + Tensor::cat(&[xs_out, xs_pass], D::Minus1) + } +} + +#[derive(Debug, Clone)] +struct CoreAttention { + coeff: Option, + norm_factor: f64, + dtype: DType, +} + +fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: f32, dtype: DType) -> 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.to_dtype(dtype)?, on_false)?; + Ok(m) +} + +impl CoreAttention { + fn new(layer_number: usize, cfg: &Config, dtype: DType) -> Result { + let norm_factor = (cfg.kv_channels as f64).sqrt(); + let (norm_factor, coeff) = if cfg.apply_query_key_layer_scaling { + let coeff = f64::max(1.0, layer_number as f64); + (norm_factor * coeff, Some(coeff)) + } else { + (norm_factor, None) + }; + Ok(Self { + coeff, + norm_factor, + dtype, + }) + } + + fn forward( + &self, + query_layer: &Tensor, + key_layer: &Tensor, + value_layer: &Tensor, + attention_mask: &Option, + ) -> Result { + let output_size = ( + query_layer.dim(1)?, // b + query_layer.dim(2)?, // np + query_layer.dim(0)?, // sq + key_layer.dim(0)?, // sk + ); + let query_layer = + query_layer.reshape((output_size.2, output_size.0 * output_size.1, ()))?; + let key_layer = key_layer.reshape((output_size.3, output_size.0 * output_size.1, ()))?; + let matmul_result = Tensor::matmul( + &query_layer.transpose(0, 1)?.contiguous()?, + &key_layer.transpose(0, 1)?.transpose(1, 2)?.contiguous()?, + )?; + let matmul_result = (matmul_result / self.norm_factor)?.reshape(output_size)?; + let matmul_result = match self.coeff { + None => matmul_result, + Some(coeff) => (matmul_result * coeff)?, + }; + let attention_scores = match attention_mask { + Some(mask) => masked_fill( + &matmul_result, + &mask.broadcast_left((matmul_result.dim(0)?, matmul_result.dim(1)?))?, + f32::NEG_INFINITY, + self.dtype, + )?, + None => matmul_result, + }; + let attention_probs = candle_nn::ops::softmax_last_dim(&attention_scores)?; + + let output_size = ( + value_layer.dim(1)?, + value_layer.dim(2)?, + query_layer.dim(0)?, + value_layer.dim(3)?, + ); + let value_layer = + value_layer.reshape((value_layer.dim(0)?, output_size.0 * output_size.1, ()))?; + let attention_probs = + attention_probs.reshape((output_size.0 * output_size.1, output_size.2, ()))?; + let context_layer = Tensor::matmul( + &attention_probs.contiguous()?, + &value_layer.transpose(0, 1)?.contiguous()?, + )?; + let context_layer = context_layer.reshape(output_size)?; + let context_layer = context_layer.permute((2, 0, 1, 3))?.contiguous()?; + context_layer.flatten_from(D::Minus2) + } +} + +#[derive(Debug, Clone)] +struct SelfAttention { + query_key_value: Linear, + core_attention: CoreAttention, + dense: Linear, + multi_query_attention: bool, + num_attention_heads_per_partition: usize, + num_multi_query_groups_per_partition: usize, + hidden_size_per_attention_head: usize, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl SelfAttention { + fn new(layer_number: usize, cfg: &Config, vb: VarBuilder) -> Result { + let projection_size = cfg.kv_channels * cfg.num_attention_heads; + let hidden_size_per_attention_head = projection_size / cfg.num_attention_heads; + let qkv_hidden_size = if cfg.multi_query_attention { + projection_size + 2 * hidden_size_per_attention_head * cfg.multi_query_group_num + } else { + 3 * projection_size + }; + let query_key_value = linear( + cfg.hidden_size, + qkv_hidden_size, + cfg.add_bias_linear || cfg.add_qkv_bias, + vb.pp("query_key_value"), + )?; + let core_attention = CoreAttention::new(layer_number, cfg, vb.dtype())?; + let dense = linear( + cfg.hidden_size, + cfg.hidden_size, + cfg.add_bias_linear, + vb.pp("dense"), + )?; + Ok(Self { + query_key_value, + core_attention, + dense, + multi_query_attention: cfg.multi_query_attention, + num_attention_heads_per_partition: cfg.num_attention_heads, + num_multi_query_groups_per_partition: cfg.multi_query_group_num, + hidden_size_per_attention_head: cfg.kv_channels, + kv_cache: None, + }) + } + + fn reset_kv_cache(&mut self) { + self.kv_cache = None + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: &Option, + rotary_emb: &RotaryEmbedding, + ) -> Result { + let mixed_x_layer = xs.apply(&self.query_key_value)?; + if !self.multi_query_attention { + candle::bail!("only multi_query_attention=true is supported") + } + let hpa = self.hidden_size_per_attention_head; + let query_layer = + mixed_x_layer.narrow(D::Minus1, 0, self.num_attention_heads_per_partition * hpa)?; + let key_layer = mixed_x_layer.narrow( + D::Minus1, + self.num_attention_heads_per_partition * hpa, + self.num_multi_query_groups_per_partition * hpa, + )?; + let value_layer = mixed_x_layer.narrow( + D::Minus1, + self.num_attention_heads_per_partition * hpa + + self.num_multi_query_groups_per_partition * hpa, + self.num_multi_query_groups_per_partition * hpa, + )?; + let query_layer = query_layer.reshape(( + query_layer.dim(0)?, + query_layer.dim(1)?, + self.num_attention_heads_per_partition, + hpa, + ))?; + let key_layer = key_layer.reshape(( + key_layer.dim(0)?, + key_layer.dim(1)?, + self.num_multi_query_groups_per_partition, + hpa, + ))?; + let value_layer = value_layer.reshape(( + value_layer.dim(0)?, + value_layer.dim(1)?, + self.num_multi_query_groups_per_partition, + hpa, + ))?; + + // Rotary embeddings. + let seqlen_offset = match &self.kv_cache { + None => 0, + Some((prev_k, _)) => prev_k.dim(0)?, + }; + let query_layer = rotary_emb.apply(&query_layer, seqlen_offset)?; + let key_layer = rotary_emb.apply(&key_layer, seqlen_offset)?; + + // KV cache. + let (key_layer, value_layer) = match &self.kv_cache { + None => (key_layer, value_layer), + Some((prev_k, prev_v)) => { + let k = Tensor::cat(&[prev_k, &key_layer], 0)?; + let v = Tensor::cat(&[prev_v, &value_layer], 0)?; + (k, v) + } + }; + self.kv_cache = Some((key_layer.clone(), value_layer.clone())); + + // Repeat KV. + let ratio = + self.num_attention_heads_per_partition / self.num_multi_query_groups_per_partition; + let key_layer = { + let (d0, d1, d2, d3) = key_layer.dims4()?; + key_layer + .unsqueeze(D::Minus2)? + .expand((d0, d1, d2, ratio, d3))? + .reshape(( + d0, + d1, + self.num_attention_heads_per_partition, + self.hidden_size_per_attention_head, + ))? + }; + let value_layer = { + let (d0, d1, d2, d3) = value_layer.dims4()?; + value_layer + .unsqueeze(D::Minus2)? + .expand((d0, d1, d2, ratio, d3))? + .reshape(( + d0, + d1, + self.num_attention_heads_per_partition, + self.hidden_size_per_attention_head, + ))? + }; + + let context_layer = + self.core_attention + .forward(&query_layer, &key_layer, &value_layer, attention_mask)?; + let output = context_layer.apply(&self.dense)?; + Ok(output) + } +} + +#[allow(clippy::upper_case_acronyms)] +#[derive(Debug, Clone)] +struct MLP { + dense_h_to_4h: Linear, + dense_4h_to_h: Linear, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense_h_to_4h = linear( + cfg.hidden_size, + cfg.ffn_hidden_size * 2, + cfg.add_bias_linear, + vb.pp("dense_h_to_4h"), + )?; + let dense_4h_to_h = linear( + cfg.ffn_hidden_size, + cfg.hidden_size, + cfg.add_bias_linear, + vb.pp("dense_4h_to_h"), + )?; + Ok(Self { + dense_4h_to_h, + dense_h_to_4h, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.dense_h_to_4h)? + .apply(&candle_nn::Activation::Swiglu)? + .apply(&self.dense_4h_to_h) + } +} + +#[derive(Debug, Clone)] +struct Block { + input_layernorm: candle_nn::LayerNorm, + self_attention: SelfAttention, + post_attention_layernorm: candle_nn::LayerNorm, + mlp: MLP, + apply_residual_connection_post_layernorm: bool, +} + +impl Block { + fn new(layer_number: usize, cfg: &Config, vb: VarBuilder) -> Result { + let input_layernorm = if cfg.rmsnorm { + candle_nn::rms_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("input_layernorm"), + )? + .into_inner() + } else { + candle_nn::layer_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("input_layernorm"), + )? + }; + let post_attention_layernorm = if cfg.rmsnorm { + candle_nn::rms_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("post_attention_layernorm"), + )? + .into_inner() + } else { + candle_nn::layer_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("post_attention_layernorm"), + )? + }; + let self_attention = SelfAttention::new(layer_number, cfg, vb.pp("self_attention"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + Ok(Self { + input_layernorm, + self_attention, + post_attention_layernorm, + mlp, + apply_residual_connection_post_layernorm: cfg.apply_residual_connection_post_layernorm, + }) + } + + fn reset_kv_cache(&mut self) { + self.self_attention.reset_kv_cache() + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: &Option, + rotary_emb: &RotaryEmbedding, + ) -> Result { + let layernorm_output = xs.apply(&self.input_layernorm)?; + let attention_output = + self.self_attention + .forward(&layernorm_output, attention_mask, rotary_emb)?; + let residual = if self.apply_residual_connection_post_layernorm { + &layernorm_output + } else { + xs + }; + let layernorm_input = (residual + attention_output)?; + let layernorm_output = layernorm_input.apply(&self.post_attention_layernorm)?; + let mlp_output = layernorm_output.apply(&self.mlp)?; + let residual = if self.apply_residual_connection_post_layernorm { + &layernorm_output + } else { + &layernorm_input + }; + mlp_output + residual + } +} + +#[derive(Debug, Clone)] +struct Transformer { + layers: Vec, + final_layernorm: Option, + rotary_emb: RotaryEmbedding, +} + +impl Transformer { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_l = vb.pp("layers"); + let mut layers = Vec::with_capacity(cfg.num_layers); + for layer_index in 0..cfg.num_layers { + let block = Block::new(layer_index + 1, cfg, vb_l.pp(layer_index))?; + layers.push(block) + } + let final_layernorm = if cfg.post_layer_norm { + let ln = if cfg.rmsnorm { + candle_nn::rms_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("final_layernorm"), + )? + .into_inner() + } else { + candle_nn::layer_norm( + cfg.hidden_size, + cfg.layernorm_epsilon, + vb.pp("final_layernorm"), + )? + }; + Some(ln) + } else { + None + }; + let rotary_emb = RotaryEmbedding::new(cfg, vb.dtype(), vb.device())?; + Ok(Self { + layers, + final_layernorm, + rotary_emb, + }) + } + + fn reset_kv_cache(&mut self) { + for block in self.layers.iter_mut() { + block.reset_kv_cache() + } + } + + fn forward(&mut self, xs: &Tensor, attention_mask: &Option) -> Result { + let mut xs = xs.clone(); + for block in self.layers.iter_mut() { + xs = block.forward(&xs, attention_mask, &self.rotary_emb)? + } + match self.final_layernorm.as_ref() { + None => Ok(xs), + Some(ln) => xs.apply(ln), + } + } +} + +#[derive(Debug, Clone)] +struct Embedding { + word_embeddings: candle_nn::Embedding, + fp32_residual_connection: bool, +} + +impl Embedding { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let word_embeddings = candle_nn::embedding( + cfg.padded_vocab_size, + cfg.hidden_size, + vb.pp("word_embeddings"), + )?; + Ok(Self { + word_embeddings, + fp32_residual_connection: cfg.fp32_residual_connection, + }) + } +} + +impl Module for Embedding { + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.word_embeddings.forward(xs)?.transpose(0, 1)?; // b,s,h -> s,b,h + if self.fp32_residual_connection { + xs.to_dtype(candle::DType::F32) + } else { + xs.contiguous() + } + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embedding: Embedding, + encoder: Transformer, + output_layer: Linear, +} + +fn get_mask(size: usize, device: &Device) -> Result { + let mask: Vec<_> = (0..size) + .flat_map(|i| (0..size).map(move |j| u8::from(j > i))) + .collect(); + Tensor::from_slice(&mask, (size, size), device) +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb = vb.pp("transformer"); + let embedding = Embedding::new(cfg, vb.pp("embedding"))?; + let encoder = Transformer::new(cfg, vb.pp("encoder"))?; + let output_layer = linear( + cfg.hidden_size, + cfg.padded_vocab_size, + false, + vb.pp("output_layer"), + )?; + + Ok(Self { + embedding, + encoder, + output_layer, + }) + } + + pub fn reset_kv_cache(&mut self) { + self.encoder.reset_kv_cache() + } + + pub fn forward(&mut self, xs: &Tensor) -> Result { + let (_b_size, seq_len) = xs.dims2()?; + let input_embeds = xs.apply(&self.embedding)?; + let attention_mask = if seq_len <= 1 { + None + } else { + Some(get_mask(seq_len, xs.device())?) + }; + let xs = self.encoder.forward(&input_embeds, &attention_mask)?; + let lm_logits = xs.i(seq_len - 1)?.apply(&self.output_layer)?; + Ok(lm_logits) + } +} diff --git a/patches/candle-transformers/src/models/glm4_new.rs b/patches/candle-transformers/src/models/glm4_new.rs new file mode 100644 index 0000000000..bee255327c --- /dev/null +++ b/patches/candle-transformers/src/models/glm4_new.rs @@ -0,0 +1,404 @@ +use crate::models::glm4::EosTokenId; +use crate::{ + models::with_tracing::{linear_b, linear_no_bias, Linear, RmsNorm}, + utils::repeat_kv, +}; +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{kv_cache::KvCache, Activation, VarBuilder}; +use std::sync::Arc; + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub head_dim: Option, + pub partial_rotary_factor: Option, + pub attention_bias: Option, + pub num_key_value_heads: usize, + pub max_position_embeddings: usize, + pub sliding_window: Option, + pub tie_word_embeddings: bool, + pub rope_theta: f64, + pub rms_norm_eps: f64, + pub hidden_act: Activation, + pub eos_token_id: Option, +} + +#[derive(Debug, Clone)] +pub(crate) struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, + rotary_dim: usize, +} + +impl RotaryEmbedding { + pub(crate) fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let dim = cfg + .head_dim + .unwrap_or(cfg.hidden_size / cfg.num_attention_heads); + let rotary_dim = if let Some(factor) = cfg.partial_rotary_factor { + (factor * dim as f32) as usize + } else { + dim + }; + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..rotary_dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / rotary_dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + rotary_dim, + }) + } + + pub(crate) fn apply(&self, xs: &Tensor, offset: usize) -> Result { + let (_, _, seq_len, _) = xs.dims4()?; + let (s, e) = (offset, offset + seq_len); + let cos = self.cos.i((s..e, ..))?.contiguous()?; + let sin = self.sin.i((s..e, ..))?.contiguous()?; + let xs_rot = xs + .i((0, .., .., ..self.rotary_dim))? + .unsqueeze(0)? + .contiguous()?; + let xs_pass = xs.i((0, .., .., self.rotary_dim..))?.unsqueeze(0)?; + let xs_rot = candle_nn::rotary_emb::rope_i(&xs_rot, &cos, &sin).unwrap(); + Tensor::cat(&[&xs_rot, &xs_pass], D::Minus1)?.contiguous() + } +} + +#[derive(Debug, Clone)] +pub(crate) struct Mlp { + gate_up_proj: Linear, + down_proj: Linear, + act_fn: Activation, +} + +impl Mlp { + pub(crate) fn new(cfg: &Config, vb: VarBuilder) -> Result { + Ok(Self { + gate_up_proj: linear_no_bias( + cfg.hidden_size, + cfg.intermediate_size * 2, + vb.pp("gate_up_proj"), + )?, + down_proj: linear_no_bias(cfg.intermediate_size, cfg.hidden_size, vb.pp("down_proj"))?, + act_fn: cfg.hidden_act, + }) + } +} + +impl Module for Mlp { + fn forward(&self, x: &Tensor) -> Result { + let w = self.gate_up_proj.forward(x)?; + let dim = w.dims().len() - 1; + let gate = w.narrow(dim, 0, w.dim(dim)? / 2)?.contiguous()?; + let gate = gate.apply(&self.act_fn)?; + let up_states = w + .narrow(dim, w.dim(dim)? / 2, w.dim(dim)? / 2)? + .contiguous()?; + self.down_proj.forward(&(gate * up_states)?) + } +} + +#[derive(Debug, Clone)] +pub(crate) struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + hidden_size: usize, + rotary_emb: Arc, + kv_cache: KvCache, +} + +impl Attention { + pub(crate) fn new( + cfg: &Config, + rotary_emb: Arc, + vb: VarBuilder, + ) -> Result { + let head_dim = cfg + .head_dim + .unwrap_or(cfg.hidden_size / cfg.num_attention_heads); + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + + let q_proj = linear_b( + cfg.hidden_size, + num_heads * head_dim, + cfg.attention_bias.unwrap_or(false), + vb.pp("q_proj"), + )?; + let k_proj = linear_b( + cfg.hidden_size, + num_kv_heads * head_dim, + cfg.attention_bias.unwrap_or(false), + vb.pp("k_proj"), + )?; + let v_proj = linear_b( + cfg.hidden_size, + num_kv_heads * head_dim, + cfg.attention_bias.unwrap_or(false), + vb.pp("v_proj"), + )?; + let o_proj = linear_b( + num_heads * head_dim, + cfg.hidden_size, + false, + vb.pp("o_proj"), + )?; + + // Necessary because the hidden_size in the config isn't always accurate + let hidden_size = head_dim * cfg.num_attention_heads; + + // Initialize KV cache with 512 tokens capacity to reduce initial memory allocation. + // The cache will grow in chunks of 512 tokens when needed. + let kv_cache = KvCache::new(2, 512); + + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + hidden_size, + rotary_emb, + kv_cache, + }) + } + + pub(crate) fn forward( + &mut self, + x: &Tensor, + attn_mask: Option<&Tensor>, + offset: usize, + ) -> Result { + let (b, l, _) = x.dims3()?; + + let q = self.q_proj.forward(x)?; + let k = self.k_proj.forward(x)?; + let v = self.v_proj.forward(x)?; + + let q = q + .reshape((b, l, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let k = k + .reshape((b, l, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let v = v + .reshape((b, l, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let q = self.rotary_emb.apply(&q, offset)?; + let k = self.rotary_emb.apply(&k, offset)?; + + let (k, v) = self.kv_cache.append(&k.contiguous()?, &v.contiguous()?)?; + + let k = repeat_kv(k, self.num_kv_groups)?; + let v = repeat_kv(v, self.num_kv_groups)?; + + let scale = 1.0 / (self.head_dim as f64).sqrt(); + let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?; + if let Some(m) = attn_mask { + scores = scores.broadcast_add(m)?; + } + let probs = candle_nn::ops::softmax_last_dim(&scores)?; + let ctx = probs.matmul(&v)?; + + ctx.transpose(1, 2)? + .reshape((b, l, self.hidden_size))? + .apply(&self.o_proj) + } + + pub(crate) fn clear_kv_cache(&mut self) { + self.kv_cache.reset(); + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: Mlp, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, + post_mlp_layernorm: RmsNorm, + post_self_attn_layernorm: RmsNorm, +} + +impl DecoderLayer { + fn new(cfg: &Config, rotary: Arc, vb: VarBuilder) -> Result { + let self_attn = Attention::new(cfg, rotary, vb.pp("self_attn"))?; + let mlp = Mlp::new(cfg, vb.pp("mlp"))?; + + let input_layernorm = + RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let post_attention_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + let post_self_attn_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_self_attn_layernorm"), + )?; + let post_mlp_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_mlp_layernorm"), + )?; + + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + post_self_attn_layernorm, + post_mlp_layernorm, + }) + } + + fn forward(&mut self, xs: &Tensor, mask: Option<&Tensor>, offset: usize) -> Result { + let residual = xs; + let hidden_states = self.input_layernorm.forward(xs)?; + let hidden_states = self.self_attn.forward(&hidden_states, mask, offset)?; + let hidden_states = self.post_self_attn_layernorm.forward(&hidden_states)?; + let hidden_states = (residual + hidden_states)?; + let residual = &hidden_states; + let hidden_states = self.post_attention_layernorm.forward(&hidden_states)?; + let hidden_states = self.mlp.forward(&hidden_states)?; + let hidden_states = self.post_mlp_layernorm.forward(&hidden_states)?; + residual + hidden_states + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache(); + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + device: Device, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("model.embed_tokens"))?; + let rotary = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb.pp("model.layers"); + for i in 0..cfg.num_hidden_layers { + layers.push(DecoderLayer::new(cfg, rotary.clone(), vb_l.pp(i))?); + } + Ok(Self { + embed_tokens, + layers, + norm: RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("model.norm"))?, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + fn clear_kv_cache(&mut self) { + for l in &mut self.layers { + l.clear_kv_cache(); + } + } + + fn causal_mask( + &self, + b: usize, + tgt: usize, + offset: usize, + sw: Option, + ) -> Result { + let minf = f32::NEG_INFINITY; + let mask: Vec<_> = (0..tgt) + .flat_map(|i| { + (0..(tgt + offset)).map(move |j| { + let past_ok = j <= i + offset; + let sw_ok = match sw { + Some(w) => (i + offset) as i64 - j as i64 <= w as i64, + None => true, + }; + if past_ok && sw_ok { + 0. + } else { + minf + } + }) + }) + .collect(); + Tensor::from_slice(&mask, (b, 1, tgt, tgt + offset), &self.device)?.to_dtype(self.dtype) + } + + pub fn forward(&mut self, input: &Tensor, offset: usize) -> Result { + let (b, l) = input.dims2()?; + let mut h = self.embed_tokens.forward(input)?; + + let causal = if l == 1 { + None + } else { + Some(self.causal_mask(b, l, offset, None)?) + }; + + for layer in &mut self.layers { + h = layer.forward(&h, causal.as_ref(), offset)?; + } + self.norm.forward(&h) + } +} + +#[derive(Debug, Clone)] +pub struct ModelForCausalLM { + base: Model, + lm_head: Linear, +} + +impl ModelForCausalLM { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let base = Model::new(cfg, vb.clone())?; + let lm_head = if cfg.tie_word_embeddings { + Linear::from_weights(base.embed_tokens.embeddings().clone(), None) + } else { + linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))? + }; + Ok(Self { base, lm_head }) + } + + pub fn forward(&mut self, input: &Tensor, offset: usize) -> Result { + let (_, l) = input.dims2()?; + self.base + .forward(input, offset)? + .narrow(1, l - 1, 1)? + .apply(&self.lm_head) + } + + pub fn clear_kv_cache(&mut self) { + self.base.clear_kv_cache(); + } +} diff --git a/patches/candle-transformers/src/models/granite.rs b/patches/candle-transformers/src/models/granite.rs new file mode 100644 index 0000000000..95b188e08d --- /dev/null +++ b/patches/candle-transformers/src/models/granite.rs @@ -0,0 +1,463 @@ +//! Granite is a Long Context Transformer Language Model. +//! +//! A high performance transformer model optimized for efficient processing +//! of very long context sequences + +use super::with_tracing::{linear_no_bias as linear, Linear, RmsNorm}; +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::{embedding, Embedding, Module, VarBuilder}; +use std::{collections::HashMap, f32::consts::PI}; + +pub const DEFAULT_MAX_SEQ_LEN: usize = 4096; + +#[derive(Debug, Clone, serde::Deserialize, Default)] +pub enum GraniteRopeType { + #[serde(rename = "granite")] + Granite, + #[default] + #[serde(rename = "default")] + Default, +} + +#[derive(Debug, Clone, serde::Deserialize, Default)] +pub struct GraniteRopeConfig { + pub factor: f32, + pub low_freq_factor: f32, + pub high_freq_factor: f32, + pub original_max_position_embeddings: usize, + pub rope_type: GraniteRopeType, +} +#[derive(Debug, Clone, serde::Deserialize)] +#[serde(untagged)] +pub enum GraniteEosToks { + Single(u32), + Multiple(Vec), +} + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct GraniteConfig { + pub hidden_size: usize, + pub intermediate_size: usize, + pub vocab_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: Option, + pub rms_norm_eps: f64, + #[serde(default = "default_rope")] + pub rope_theta: f32, + pub bos_token_id: Option, + pub eos_token_id: Option, + pub rope_scaling: Option, + pub max_position_embeddings: usize, +} + +impl GraniteConfig { + pub fn num_key_value_heads(&self) -> usize { + self.num_key_value_heads.unwrap_or(self.num_attention_heads) + } +} + +fn default_rope() -> f32 { + 10_000.0 +} + +impl GraniteConfig { + pub fn into_config(self, use_flash_attn: bool) -> Config { + Config { + hidden_size: self.hidden_size, + intermediate_size: self.intermediate_size, + vocab_size: self.vocab_size, + num_hidden_layers: self.num_hidden_layers, + num_attention_heads: self.num_attention_heads, + num_key_value_heads: self.num_key_value_heads(), + rms_norm_eps: self.rms_norm_eps, + rope_theta: self.rope_theta, + use_flash_attn, + bos_token_id: self.bos_token_id, + eos_token_id: self.eos_token_id, + rope_scaling: self.rope_scaling, + max_position_embeddings: self.max_position_embeddings, + } + } +} + +#[derive(Debug, Clone)] +pub struct Config { + pub hidden_size: usize, + pub intermediate_size: usize, + pub vocab_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub use_flash_attn: bool, + pub rms_norm_eps: f64, + pub rope_theta: f32, + pub bos_token_id: Option, + pub eos_token_id: Option, + pub rope_scaling: Option, + pub max_position_embeddings: usize, +} + +#[derive(Debug, Clone)] +pub struct Cache { + masks: HashMap, + pub use_kv_cache: bool, + kvs: Vec>, + cos: Tensor, + sin: Tensor, + device: Device, +} + +fn calculate_default_inv_freq(cfg: &Config) -> Vec { + let head_dim = cfg.hidden_size / cfg.num_attention_heads; + (0..head_dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f32 / head_dim as f32)) + .collect() +} + +impl Cache { + pub fn new(use_kv_cache: bool, dtype: DType, config: &Config, device: &Device) -> Result { + // precompute freqs_cis + let theta = match &config.rope_scaling { + None + | Some(GraniteRopeConfig { + rope_type: GraniteRopeType::Default, + .. + }) => calculate_default_inv_freq(config), + Some(rope_scaling) => { + let low_freq_wavelen = rope_scaling.original_max_position_embeddings as f32 + / rope_scaling.low_freq_factor; + let high_freq_wavelen = rope_scaling.original_max_position_embeddings as f32 + / rope_scaling.high_freq_factor; + + calculate_default_inv_freq(config) + .into_iter() + .map(|freq| { + let wavelen = 2. * PI / freq; + if wavelen < high_freq_wavelen { + freq + } else if wavelen > low_freq_wavelen { + freq / rope_scaling.factor + } else { + let smooth = (rope_scaling.original_max_position_embeddings as f32 + / wavelen + - rope_scaling.low_freq_factor) + / (rope_scaling.high_freq_factor - rope_scaling.low_freq_factor); + (1. - smooth) * freq / rope_scaling.factor + smooth * freq + } + }) + .collect::>() + } + }; + + let theta = Tensor::new(theta, device)?; + + let idx_theta = Tensor::arange(0, config.max_position_embeddings as u32, device)? + .to_dtype(DType::F32)? + .reshape((config.max_position_embeddings, 1))? + .matmul(&theta.reshape((1, theta.elem_count()))?)?; + let cos = idx_theta.cos()?.to_dtype(dtype)?; + let sin = idx_theta.sin()?.to_dtype(dtype)?; + Ok(Self { + masks: HashMap::new(), + use_kv_cache, + kvs: vec![None; config.num_hidden_layers], + device: device.clone(), + cos, + sin, + }) + } + + 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), &self.device)?; + self.masks.insert(t, mask.clone()); + Ok(mask) + } + } +} + +#[derive(Debug, Clone)] +struct CausalSelfAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_attention_heads: usize, + num_key_value_heads: usize, + head_dim: usize, + use_flash_attn: bool, + span: tracing::Span, + span_rot: tracing::Span, + max_position_embeddings: usize, +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +impl CausalSelfAttention { + fn apply_rotary_emb(&self, x: &Tensor, index_pos: usize, cache: &Cache) -> Result { + let _enter = self.span_rot.enter(); + let (_b_sz, _, seq_len, _hidden_size) = x.dims4()?; + let cos = cache.cos.narrow(0, index_pos, seq_len)?; + let sin = cache.sin.narrow(0, index_pos, seq_len)?; + candle_nn::rotary_emb::rope(x, &cos, &sin) + } + + fn forward( + &self, + x: &Tensor, + index_pos: usize, + block_idx: usize, + cache: &mut Cache, + ) -> Result { + let _enter = self.span.enter(); + let (b_sz, seq_len, hidden_size) = x.dims3()?; + let q = self.q_proj.forward(x)?; + let k = self.k_proj.forward(x)?; + let v = self.v_proj.forward(x)?; + + let q = q + .reshape((b_sz, seq_len, self.num_attention_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let k = k + .reshape((b_sz, seq_len, self.num_key_value_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let mut v = v + .reshape((b_sz, seq_len, self.num_key_value_heads, self.head_dim))? + .transpose(1, 2)?; + + let q = self.apply_rotary_emb(&q, index_pos, cache)?; + let mut k = self.apply_rotary_emb(&k, index_pos, cache)?; + + if cache.use_kv_cache { + if let Some((cache_k, cache_v)) = &cache.kvs[block_idx] { + k = Tensor::cat(&[cache_k, &k], 2)?.contiguous()?; + v = Tensor::cat(&[cache_v, &v], 2)?.contiguous()?; + let k_seq_len = k.dims()[1]; + if k_seq_len > self.max_position_embeddings { + k = k + .narrow( + D::Minus1, + k_seq_len - self.max_position_embeddings, + self.max_position_embeddings, + )? + .contiguous()? + } + let v_seq_len = v.dims()[1]; + if v_seq_len > 2 * self.max_position_embeddings { + v = v + .narrow( + D::Minus1, + v_seq_len - self.max_position_embeddings, + self.max_position_embeddings, + )? + .contiguous()? + } + } + cache.kvs[block_idx] = Some((k.clone(), v.clone())) + } + + let k = self.repeat_kv(k)?; + let v = self.repeat_kv(v)?; + + let y = if self.use_flash_attn { + // flash-attn expects (b_sz, seq_len, nheads, head_dim) + let q = q.transpose(1, 2)?; + let k = k.transpose(1, 2)?; + let v = v.transpose(1, 2)?; + let softmax_scale = 1f32 / (self.head_dim as f32).sqrt(); + flash_attn(&q, &k, &v, softmax_scale, seq_len > 1)?.transpose(1, 2)? + } else { + let in_dtype = q.dtype(); + let q = q.to_dtype(DType::F32)?; + let k = k.to_dtype(DType::F32)?; + let v = v.to_dtype(DType::F32)?; + let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?; + let att = if seq_len == 1 { + att + } else { + let mask = cache.mask(seq_len)?.broadcast_as(att.shape())?; + masked_fill(&att, &mask, f32::NEG_INFINITY)? + }; + let att = candle_nn::ops::softmax(&att, D::Minus1)?; + // Convert to contiguous as matmul doesn't support strided vs for now. + att.matmul(&v.contiguous()?)?.to_dtype(in_dtype)? + }; + let y = y.transpose(1, 2)?.reshape(&[b_sz, seq_len, hidden_size])?; + let y = self.o_proj.forward(&y)?; + Ok(y) + } + + fn repeat_kv(&self, x: Tensor) -> Result { + crate::utils::repeat_kv(x, self.num_attention_heads / self.num_key_value_heads) + } + + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "attn"); + let span_rot = tracing::span!(tracing::Level::TRACE, "attn-rot"); + let size_in = cfg.hidden_size; + let size_q = (cfg.hidden_size / cfg.num_attention_heads) * cfg.num_attention_heads; + let size_kv = (cfg.hidden_size / cfg.num_attention_heads) * cfg.num_key_value_heads; + let q_proj = linear(size_in, size_q, vb.pp("q_proj"))?; + let k_proj = linear(size_in, size_kv, vb.pp("k_proj"))?; + let v_proj = linear(size_in, size_kv, vb.pp("v_proj"))?; + let o_proj = linear(size_q, size_in, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_attention_heads: cfg.num_attention_heads, + num_key_value_heads: cfg.num_key_value_heads, + head_dim: cfg.hidden_size / cfg.num_attention_heads, + use_flash_attn: cfg.use_flash_attn, + span, + span_rot, + max_position_embeddings: cfg.max_position_embeddings, + }) + } +} + +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) +} + +#[derive(Debug, Clone)] +struct Mlp { + c_fc1: Linear, + c_fc2: Linear, + c_proj: Linear, + span: tracing::Span, +} + +impl Mlp { + fn forward(&self, x: &Tensor) -> Result { + let _enter = self.span.enter(); + let x = (candle_nn::ops::silu(&self.c_fc1.forward(x)?)? * self.c_fc2.forward(x)?)?; + self.c_proj.forward(&x) + } + + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "mlp"); + let h_size = cfg.hidden_size; + let i_size = cfg.intermediate_size; + let c_fc1 = linear(h_size, i_size, vb.pp("gate_proj"))?; + let c_fc2 = linear(h_size, i_size, vb.pp("up_proj"))?; + let c_proj = linear(i_size, h_size, vb.pp("down_proj"))?; + Ok(Self { + c_fc1, + c_fc2, + c_proj, + span, + }) + } +} + +#[derive(Debug, Clone)] +struct Block { + rms_1: RmsNorm, + attn: CausalSelfAttention, + rms_2: RmsNorm, + mlp: Mlp, + span: tracing::Span, +} + +impl Block { + fn forward( + &self, + x: &Tensor, + index_pos: usize, + block_idx: usize, + cache: &mut Cache, + ) -> Result { + let _enter = self.span.enter(); + let residual = x; + let x = self.rms_1.forward(x)?; + let x = (self.attn.forward(&x, index_pos, block_idx, cache)? + residual)?; + let residual = &x; + let x = (self.mlp.forward(&self.rms_2.forward(&x)?)? + residual)?; + Ok(x) + } + + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "block"); + let attn = CausalSelfAttention::load(vb.pp("self_attn"), cfg)?; + let mlp = Mlp::load(vb.pp("mlp"), cfg)?; + let rms_1 = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let rms_2 = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + rms_1, + attn, + rms_2, + mlp, + span, + }) + } +} + +#[derive(Debug, Clone)] +pub struct Granite { + wte: Embedding, + blocks: Vec, + ln_f: RmsNorm, + lm_head: Linear, +} + +impl Granite { + pub fn forward(&self, x: &Tensor, index_pos: usize, cache: &mut Cache) -> Result { + let (_b_sz, seq_len) = x.dims2()?; + let mut x = self.wte.forward(x)?; + for (block_idx, block) in self.blocks.iter().enumerate() { + x = block.forward(&x, index_pos, block_idx, cache)?; + } + let x = self.ln_f.forward(&x)?; + let x = x.i((.., seq_len - 1, ..))?.contiguous()?; + let logits = self.lm_head.forward(&x)?; + logits.to_dtype(DType::F32) + } + + pub fn load(vb: VarBuilder, cfg: &Config) -> Result { + let wte = embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("model.embed_tokens"))?; + let lm_head = linear(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))?; + let ln_f = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("model.norm"))?; + let blocks: Vec<_> = (0..cfg.num_hidden_layers) + .map(|i| Block::load(vb.pp(format!("model.layers.{i}")), cfg).unwrap()) + .collect(); + + Ok(Self { + wte, + blocks, + ln_f, + lm_head, + }) + } +} diff --git a/patches/candle-transformers/src/models/granitemoehybrid.rs b/patches/candle-transformers/src/models/granitemoehybrid.rs new file mode 100644 index 0000000000..30ddeff2c1 --- /dev/null +++ b/patches/candle-transformers/src/models/granitemoehybrid.rs @@ -0,0 +1,586 @@ +//! GraniteMoeHybrid is a Long Context Transformer Language Model. +//! +//! A high performance transformer model optimized for efficient processing +//! of very long context sequences + +use super::with_tracing::{linear_no_bias as linear, Linear, RmsNorm}; +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::{embedding, Embedding, Module, VarBuilder}; +use std::iter::repeat_n; +use std::{collections::HashMap, f32::consts::PI}; + +pub const DEFAULT_MAX_SEQ_LEN: usize = 4096; + +#[derive(Debug, Clone, serde::Deserialize, Default)] +pub enum GraniteMoeHybridRopeType { + #[serde(rename = "granite")] + Granite, + #[default] + #[serde(rename = "default")] + Default, +} + +#[derive(Debug, Clone, serde::Deserialize, Default)] +pub struct GraniteMoeHybridRopeConfig { + pub factor: f32, + pub low_freq_factor: f32, + pub high_freq_factor: f32, + pub original_max_position_embeddings: usize, + pub rope_type: GraniteMoeHybridRopeType, +} + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct GraniteMoeHybridConfig { + pub hidden_size: usize, + pub intermediate_size: usize, + pub vocab_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: Option, + pub rms_norm_eps: f64, + #[serde(default = "default_rope")] + pub rope_theta: f32, + pub bos_token_id: Option, + pub eos_token_id: Option, + pub rope_scaling: Option, + pub max_position_embeddings: usize, + #[serde(default)] + pub layer_types: Vec, + #[serde(default = "default_one")] + pub attention_multiplier: f32, + #[serde(default = "default_one")] + pub embedding_multiplier: f32, + #[serde(default = "default_one")] + pub residual_multiplier: f32, + #[serde(default = "default_one")] + pub logits_scaling: f32, + #[serde(default)] + pub shared_intermediate_size: Option, +} + +impl GraniteMoeHybridConfig { + pub fn num_key_value_heads(&self) -> usize { + self.num_key_value_heads.unwrap_or(self.num_attention_heads) + } +} + +fn default_rope() -> f32 { + 10_000.0 +} + +fn default_one() -> f32 { + 1.0 +} + +#[derive(Debug, Clone, serde::Deserialize, Default)] +#[serde(rename_all = "lowercase")] +pub enum GraniteMoeHybridLayerType { + #[default] + Attention, + Mamba, +} + +impl GraniteMoeHybridConfig { + pub fn into_config(self, use_flash_attn: bool) -> GraniteMoeHybridInternalConfig { + let layer_types = if self.layer_types.is_empty() { + vec![GraniteMoeHybridLayerType::Attention; self.num_hidden_layers] + } else { + self.layer_types.clone() + }; + let shared_intermediate_size = self + .shared_intermediate_size + .unwrap_or(self.intermediate_size); + GraniteMoeHybridInternalConfig { + hidden_size: self.hidden_size, + intermediate_size: self.intermediate_size, + shared_intermediate_size, + vocab_size: self.vocab_size, + num_hidden_layers: self.num_hidden_layers, + num_attention_heads: self.num_attention_heads, + num_key_value_heads: self.num_key_value_heads(), + use_flash_attn, + rms_norm_eps: self.rms_norm_eps, + rope_theta: self.rope_theta, + bos_token_id: self.bos_token_id, + eos_token_id: self.eos_token_id, + rope_scaling: self.rope_scaling, + max_position_embeddings: self.max_position_embeddings, + layer_types, + attention_multiplier: self.attention_multiplier, + embedding_multiplier: self.embedding_multiplier, + residual_multiplier: self.residual_multiplier, + logits_scaling: self.logits_scaling, + } + } +} + +#[derive(Debug, Clone)] +pub struct GraniteMoeHybridInternalConfig { + pub hidden_size: usize, + pub intermediate_size: usize, + pub shared_intermediate_size: usize, + pub vocab_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub use_flash_attn: bool, + pub rms_norm_eps: f64, + pub rope_theta: f32, + pub bos_token_id: Option, + pub eos_token_id: Option, + pub rope_scaling: Option, + pub max_position_embeddings: usize, + pub layer_types: Vec, + pub attention_multiplier: f32, + pub embedding_multiplier: f32, + pub residual_multiplier: f32, + pub logits_scaling: f32, +} + +#[derive(Debug, Clone)] +pub struct GraniteMoeHybridCache { + masks: HashMap, + pub use_kv_cache: bool, + kvs: Vec>, + cos: Tensor, + sin: Tensor, + device: Device, +} + +fn calculate_default_inv_freq(cfg: &GraniteMoeHybridInternalConfig) -> Vec { + let head_dim = cfg.hidden_size / cfg.num_attention_heads; + (0..head_dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f32 / head_dim as f32)) + .collect() +} + +impl GraniteMoeHybridCache { + pub fn new( + use_kv_cache: bool, + dtype: DType, + config: &GraniteMoeHybridInternalConfig, + device: &Device, + ) -> Result { + // precompute freqs_cis + let theta = match &config.rope_scaling { + None + | Some(GraniteMoeHybridRopeConfig { + rope_type: GraniteMoeHybridRopeType::Default, + .. + }) => calculate_default_inv_freq(config), + Some(rope_scaling) => { + let low_freq_wavelen = rope_scaling.original_max_position_embeddings as f32 + / rope_scaling.low_freq_factor; + let high_freq_wavelen = rope_scaling.original_max_position_embeddings as f32 + / rope_scaling.high_freq_factor; + + calculate_default_inv_freq(config) + .into_iter() + .map(|freq| { + let wavelen = 2. * PI / freq; + if wavelen < high_freq_wavelen { + freq + } else if wavelen > low_freq_wavelen { + freq / rope_scaling.factor + } else { + let smooth = (rope_scaling.original_max_position_embeddings as f32 + / wavelen + - rope_scaling.low_freq_factor) + / (rope_scaling.high_freq_factor - rope_scaling.low_freq_factor); + (1. - smooth) * freq / rope_scaling.factor + smooth * freq + } + }) + .collect::>() + } + }; + + let theta = Tensor::new(theta, device)?; + + let idx_theta = Tensor::arange(0, config.max_position_embeddings as u32, device)? + .to_dtype(DType::F32)? + .reshape((config.max_position_embeddings, 1))? + .matmul(&theta.reshape((1, theta.elem_count()))?)?; + let cos = idx_theta.cos()?.to_dtype(dtype)?; + let sin = idx_theta.sin()?.to_dtype(dtype)?; + Ok(Self { + masks: HashMap::new(), + use_kv_cache, + kvs: vec![None; config.num_hidden_layers], + device: device.clone(), + cos, + sin, + }) + } + + fn mask(&mut self, t: usize) -> Result { + if let Some(mask) = self.masks.get(&t) { + Ok(mask.clone()) + } else { + let mut mask: Vec = Vec::with_capacity(t * t); + (0..t).for_each(|i| { + mask.extend(repeat_n(0, i + 1)); + mask.extend(repeat_n(1, t - i - 1)); + }); + let mask = Tensor::from_slice(&mask, (t, t), &self.device)?; + self.masks.insert(t, mask.clone()); + Ok(mask) + } + } +} + +#[derive(Debug, Clone)] +struct CausalSelfAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_attention_heads: usize, + num_key_value_heads: usize, + head_dim: usize, + use_flash_attn: bool, + span: tracing::Span, + span_rot: tracing::Span, + max_position_embeddings: usize, + attention_multiplier: f32, +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +impl CausalSelfAttention { + fn apply_rotary_emb( + &self, + x: &Tensor, + index_pos: usize, + cache: &GraniteMoeHybridCache, + ) -> Result { + let _enter = self.span_rot.enter(); + let (_b_sz, _, seq_len, _hidden_size) = x.dims4()?; + let cos = cache.cos.narrow(0, index_pos, seq_len)?; + let sin = cache.sin.narrow(0, index_pos, seq_len)?; + candle_nn::rotary_emb::rope(x, &cos, &sin) + } + + fn forward( + &self, + x: &Tensor, + index_pos: usize, + block_idx: usize, + cache: &mut GraniteMoeHybridCache, + ) -> Result { + let _enter = self.span.enter(); + let (b_sz, seq_len, hidden_size) = x.dims3()?; + let q = self.q_proj.forward(x)?; + let k = self.k_proj.forward(x)?; + let v = self.v_proj.forward(x)?; + + let q = q + .reshape((b_sz, seq_len, self.num_attention_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let k = k + .reshape((b_sz, seq_len, self.num_key_value_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let mut v = v + .reshape((b_sz, seq_len, self.num_key_value_heads, self.head_dim))? + .transpose(1, 2)?; + + let q = self.apply_rotary_emb(&q, index_pos, cache)?; + let mut k = self.apply_rotary_emb(&k, index_pos, cache)?; + + if cache.use_kv_cache { + if let Some((cache_k, cache_v)) = &cache.kvs[block_idx] { + k = Tensor::cat(&[cache_k, &k], 2)?.contiguous()?; + v = Tensor::cat(&[cache_v, &v], 2)?.contiguous()?; + let k_seq_len = k.dims()[1]; + if k_seq_len > self.max_position_embeddings { + k = k + .narrow( + D::Minus1, + k_seq_len - self.max_position_embeddings, + self.max_position_embeddings, + )? + .contiguous()? + } + let v_seq_len = v.dims()[1]; + if v_seq_len > 2 * self.max_position_embeddings { + v = v + .narrow( + D::Minus1, + v_seq_len - self.max_position_embeddings, + self.max_position_embeddings, + )? + .contiguous()? + } + } + cache.kvs[block_idx] = Some((k.clone(), v.clone())) + } + + let k = self.repeat_kv(k)?; + let v = self.repeat_kv(v)?; + + let y = if self.use_flash_attn { + // flash-attn expects (b_sz, seq_len, nheads, head_dim) + let q = q.transpose(1, 2)?; + let k = k.transpose(1, 2)?; + let v = v.transpose(1, 2)?; + flash_attn(&q, &k, &v, self.attention_multiplier, seq_len > 1)?.transpose(1, 2)? + } else { + let in_dtype = q.dtype(); + let q = q.to_dtype(DType::F32)?; + let k = k.to_dtype(DType::F32)?; + let v = v.to_dtype(DType::F32)?; + let att = q + .matmul(&k.t()?)? + .affine(self.attention_multiplier as f64, 0.)?; + let att = if seq_len == 1 { + att + } else { + let mask = cache.mask(seq_len)?.broadcast_as(att.shape())?; + masked_fill(&att, &mask, f32::NEG_INFINITY)? + }; + let att = candle_nn::ops::softmax(&att, D::Minus1)?; + // Convert to contiguous as matmul doesn't support strided vs for now. + att.matmul(&v.contiguous()?)?.to_dtype(in_dtype)? + }; + let y = y.transpose(1, 2)?.reshape(&[b_sz, seq_len, hidden_size])?; + let y = self.o_proj.forward(&y)?; + Ok(y) + } + + fn repeat_kv(&self, x: Tensor) -> Result { + crate::utils::repeat_kv(x, self.num_attention_heads / self.num_key_value_heads) + } + + fn load(vb: VarBuilder, cfg: &GraniteMoeHybridInternalConfig) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "attn"); + let span_rot = tracing::span!(tracing::Level::TRACE, "attn-rot"); + let size_in = cfg.hidden_size; + let size_q = (cfg.hidden_size / cfg.num_attention_heads) * cfg.num_attention_heads; + let size_kv = (cfg.hidden_size / cfg.num_attention_heads) * cfg.num_key_value_heads; + let q_proj = linear(size_in, size_q, vb.pp("q_proj"))?; + let k_proj = linear(size_in, size_kv, vb.pp("k_proj"))?; + let v_proj = linear(size_in, size_kv, vb.pp("v_proj"))?; + let o_proj = linear(size_q, size_in, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_attention_heads: cfg.num_attention_heads, + num_key_value_heads: cfg.num_key_value_heads, + head_dim: cfg.hidden_size / cfg.num_attention_heads, + use_flash_attn: cfg.use_flash_attn, + span, + span_rot, + max_position_embeddings: cfg.max_position_embeddings, + attention_multiplier: cfg.attention_multiplier, + }) + } +} + +/// Utility function to fill elements of a tensor based on a boolean mask. +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) +} + +// A simple feed forward network with a gated activation +// (GeLU, SiLU, etc.). The goal is to add non-linearity and +// increase the model's capacity to learn complex patterns. +#[derive(Debug, Clone)] +struct MultiLayerPercepton { + input_linear: Linear, + output_linear: Linear, + span: tracing::Span, +} + +impl MultiLayerPercepton { + fn forward(&self, x: &Tensor) -> Result { + let _enter = self.span.enter(); + let projected = self.input_linear.forward(x)?; + let chunks = projected.chunk(2, D::Minus1)?; + let (left, right) = (&chunks[0], &chunks[1]); + let gated = (candle_nn::ops::silu(left)? * right)?; + self.output_linear.forward(&gated) + } + + fn load(vb: VarBuilder, cfg: &GraniteMoeHybridInternalConfig) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "mlp"); + let h_size = cfg.hidden_size; + let inter_size = cfg.shared_intermediate_size; + let input_linear = linear(h_size, inter_size * 2, vb.pp("shared_mlp.input_linear"))?; + let output_linear = linear(inter_size, h_size, vb.pp("shared_mlp.output_linear"))?; + Ok(Self { + input_linear, + output_linear, + span, + }) + } +} + +// A Block is a actually a Transformer layer, consisting of +// a self-attention mechanism followed by a feed-forward neural network (MLP). +#[derive(Debug, Clone)] +struct Block { + rms_1: RmsNorm, + attn: CausalSelfAttention, + rms_2: RmsNorm, + multi_layer_percepton: MultiLayerPercepton, + span: tracing::Span, + residual_scale: f32, +} + +impl Block { + fn forward( + &self, + x: &Tensor, + index_pos: usize, + block_idx: usize, + cache: &mut GraniteMoeHybridCache, + ) -> Result { + let _enter = self.span.enter(); + let residual = x; + let x = self.rms_1.forward(x)?; + let attn = self.attn.forward(&x, index_pos, block_idx, cache)?; + let attn = scale_tensor(attn, self.residual_scale)?; + let x = (attn + residual)?; + let residual = &x; + let multi_layer_percepton_out = self + .multi_layer_percepton + .forward(&self.rms_2.forward(&x)?)?; + let multi_layer_percepton_out = + scale_tensor(multi_layer_percepton_out, self.residual_scale)?; + let x = (multi_layer_percepton_out + residual)?; + Ok(x) + } + + fn load(vb: VarBuilder, cfg: &GraniteMoeHybridInternalConfig) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "block"); + let attn = CausalSelfAttention::load(vb.pp("self_attn"), cfg)?; + let multi_layer_percepton = MultiLayerPercepton::load(vb.clone(), cfg)?; + let rms_1 = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let rms_2 = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + rms_1, + attn, + rms_2, + multi_layer_percepton, + span, + residual_scale: cfg.residual_multiplier, + }) + } +} + +#[derive(Debug, Clone)] +pub struct GraniteMoeHybrid { + word_token_embedding: Embedding, + blocks: Vec, + ln_f: RmsNorm, + logits_scale: f32, + embedding_scale: f32, +} + +impl GraniteMoeHybrid { + pub fn forward( + &self, + x: &Tensor, + index_pos: usize, + cache: &mut GraniteMoeHybridCache, + ) -> Result { + let (_b_sz, seq_len) = x.dims2()?; + let x = self.word_token_embedding.forward(x)?; + let x = scale_tensor(x, self.embedding_scale)?; + let x = self + .blocks + .iter() + .enumerate() + .try_fold(x, |x, (block_idx, block)| { + block.forward(&x, index_pos, block_idx, cache) + })?; + // Final normalization + let x = self.ln_f.forward(&x)?; + let x = x.i((.., seq_len - 1, ..))?.contiguous()?; + // Project to vocabulary size + let logits = x.matmul(&self.word_token_embedding.embeddings().t()?)?; + let logits = logits.to_dtype(DType::F32)?; + // Scale the logits if needed (that's also different from Granite 1) + let scaled_logits = if (self.logits_scale - 1.0).abs() < f32::EPSILON { + logits + } else { + logits.affine(self.logits_scale as f64, 0.)? + }; + + Ok(scaled_logits) + } + + pub fn load(vb: VarBuilder, cfg: &GraniteMoeHybridInternalConfig) -> Result { + let wte = embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("model.embed_tokens"))?; + let ln_f = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("model.norm"))?; + if cfg.layer_types.len() != cfg.num_hidden_layers { + candle::bail!( + "layer_types length {} does not match num_hidden_layers {}", + cfg.layer_types.len(), + cfg.num_hidden_layers + ); + } + let blocks = cfg + .layer_types + .iter() + .enumerate() + .map(|(idx, layer_ty)| match layer_ty { + GraniteMoeHybridLayerType::Attention => { + Block::load(vb.pp(format!("model.layers.{idx}")), cfg) + } + GraniteMoeHybridLayerType::Mamba => { + // TODO: Not supprting Mamba layers (blocks) for now, + // so we only iterate over attention layers. + candle::bail!( + "mamba layers are not yet supported in GraniteMoeHybrid inference" + ) + } + }) + .collect::>>()?; + + Ok(Self { + word_token_embedding: wte, + blocks, + ln_f, + logits_scale: if cfg.logits_scaling == 0.0 { + 1.0 + } else { + 1.0 / cfg.logits_scaling + }, + embedding_scale: cfg.embedding_multiplier, + }) + } +} + +fn scale_tensor(tensor: Tensor, scale: f32) -> Result { + if (scale - 1.0).abs() < f32::EPSILON { + Ok(tensor) + } else { + tensor.affine(scale as f64, 0.) + } +} diff --git a/patches/candle-transformers/src/models/helium.rs b/patches/candle-transformers/src/models/helium.rs new file mode 100644 index 0000000000..40cff396e7 --- /dev/null +++ b/patches/candle-transformers/src/models/helium.rs @@ -0,0 +1,395 @@ +//! Helium inference implementation. +//! +//! See the model card on Hugging Face's [hub](https://huggingface.co/kmhf/helium-2b). + +use super::with_tracing::{linear_b as linear, Linear, RmsNorm}; +use candle::{DType, Device, Result, Tensor, D}; +use candle_nn::{Module, VarBuilder}; +use std::sync::Arc; + +fn default_use_flash_attn() -> bool { + false +} + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + pub attention_bias: bool, + pub bos_token_id: u32, + pub eos_token_id: u32, + pub head_dim: usize, + pub hidden_act: candle_nn::Activation, + pub hidden_size: usize, + pub intermediate_size: usize, + pub max_position_embeddings: usize, + pub mlp_bias: bool, + pub num_attention_heads: usize, + pub num_hidden_layers: usize, + pub num_key_value_heads: usize, + pub rms_norm_eps: f64, + pub rope_theta: f64, + pub tie_word_embeddings: bool, + pub vocab_size: usize, + #[serde(default = "default_use_flash_attn")] + pub use_flash_attn: bool, +} + +impl Config { + pub fn config_2b(use_flash_attn: bool) -> Self { + Self { + attention_bias: false, + bos_token_id: 1, + eos_token_id: 2, + head_dim: 128, + hidden_act: candle_nn::Activation::Silu, + hidden_size: 2560, + intermediate_size: 7040, + max_position_embeddings: 4096, + mlp_bias: false, + num_attention_heads: 20, + num_hidden_layers: 24, + num_key_value_heads: 20, + rms_norm_eps: 1e-08, + rope_theta: 100000.0, + tie_word_embeddings: false, + vocab_size: 48000, + use_flash_attn, + } + } +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let rope_theta = cfg.rope_theta as f32; + let dim = cfg.head_dim; + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / rope_theta.powf(i as f32 / dim as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(DType::F32)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(DType::F32)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?.to_dtype(dtype)?, + cos: freqs.cos()?.to_dtype(dtype)?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope_i(q, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope_i(k, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: candle_nn::Activation, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + let bias = cfg.mlp_bias; + let gate_proj = linear(hidden_sz, intermediate_sz, bias, vb.pp("gate_proj"))?; + let up_proj = linear(hidden_sz, intermediate_sz, bias, vb.pp("up_proj"))?; + let down_proj = linear(intermediate_sz, hidden_sz, bias, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_act, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = xs.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, + use_flash_attn: bool, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + let head_dim = cfg.head_dim; + let bias = cfg.attention_bias; + let q_proj = linear(hidden_sz, num_heads * head_dim, bias, vb.pp("q_proj"))?; + let k_proj = linear(hidden_sz, num_kv_heads * head_dim, bias, vb.pp("k_proj"))?; + let v_proj = linear(hidden_sz, num_kv_heads * head_dim, bias, vb.pp("v_proj"))?; + let o_proj = linear(num_heads * head_dim, hidden_sz, bias, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + rotary_emb, + kv_cache: None, + use_flash_attn: cfg.use_flash_attn, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + + let (query_states, key_states) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states, seqlen_offset)?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?; + let value_states = crate::utils::repeat_kv(value_states, self.num_kv_groups)?; + + let attn_output = if self.use_flash_attn { + // flash-attn expects (b_sz, seq_len, nheads, head_dim) + let q = query_states.transpose(1, 2)?; + let k = key_states.transpose(1, 2)?; + let v = value_states.transpose(1, 2)?; + let softmax_scale = 1f32 / (self.head_dim as f32).sqrt(); + flash_attn(&q, &k, &v, softmax_scale, q_len > 1)?.transpose(1, 2)? + } else { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + }; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.num_heads * self.head_dim))? + .apply(&self.o_proj) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: MLP, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl DecoderLayer { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let self_attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let input_layernorm = + RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let post_attention_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; + residual + xs + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: Linear, + device: Device, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb_m.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(rotary_emb.clone(), cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb_m.pp("norm"))?; + let lm_head = if cfg.tie_word_embeddings { + Linear::from_weights(embed_tokens.embeddings().clone(), None) + } else { + linear(cfg.hidden_size, cfg.vocab_size, false, vb.pp("lm_head"))? + }; + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + fn prepare_decoder_attention_mask( + &self, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { f32::NEG_INFINITY } else { 0. })) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((1, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + pub fn embed_tokens(&self) -> &candle_nn::Embedding { + &self.embed_tokens + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let (_b_size, seq_len) = input_ids.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(seq_len, seqlen_offset)?; + Some(mask) + }; + let mut xs = self.embed_tokens.forward(input_ids)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)? + } + xs.narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +} diff --git a/patches/candle-transformers/src/models/hiera.rs b/patches/candle-transformers/src/models/hiera.rs new file mode 100644 index 0000000000..98ad825737 --- /dev/null +++ b/patches/candle-transformers/src/models/hiera.rs @@ -0,0 +1,301 @@ +//! Hiera inference implementation based on timm. +//! +//! +//! - 💻 [Hiera](https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/hiera.py) +//! - 📝 [Paper](https://arxiv.org/abs/2306.00989). Hiera: A Hierarchical Vision Transformer without the Bells-and-Whistles + +use candle::{Result, D}; +use candle_nn::{conv2d, layer_norm, linear, ops::softmax, Conv2dConfig, Func, VarBuilder}; + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + channels: usize, + heads: usize, + stages: [usize; 4], +} + +impl Config { + pub fn tiny() -> Self { + Self { + channels: 96, + heads: 1, + stages: [1, 2, 7, 2], + } + } + pub fn small() -> Self { + Self { + channels: 96, + heads: 1, + stages: [1, 2, 11, 2], + } + } + pub fn base() -> Self { + Self { + channels: 96, + heads: 1, + stages: [2, 3, 16, 3], + } + } + pub fn base_plus() -> Self { + Self { + channels: 112, + heads: 2, + stages: [2, 3, 16, 3], + } + } + pub fn large() -> Self { + Self { + channels: 144, + heads: 2, + stages: [2, 6, 36, 4], + } + } + pub fn huge() -> Self { + Self { + channels: 256, + heads: 4, + stages: [2, 6, 36, 4], + } + } +} + +const NUM_TOKENS: usize = 56 * 56; + +fn hiera_embeddings(channels: usize, vb: VarBuilder) -> Result> { + let conv_cfg = Conv2dConfig { + stride: 4, + padding: 3, + ..Default::default() + }; + let proj = conv2d(3, channels, 7, conv_cfg, vb.pp("patch_embed.proj"))?; + + let pos_embed = vb.get((1, NUM_TOKENS, channels), "pos_embed")?; + + Ok(Func::new(move |xs| { + let xs = xs.apply(&proj)?; + let (b, c, _, _) = xs.dims4()?; + let xs = xs.reshape((b, c, ()))?.transpose(1, 2)?; + let xs = xs.broadcast_add(&pos_embed)?; + Ok(xs) + })) +} + +fn hiera_unroll() -> Result> { + Ok(Func::new(move |xs| { + let mut xs = xs.clone(); + let (mut b, _, c) = xs.dims3()?; + let mut size = 56; + + xs = xs.reshape((b, size, size, c))?; + for _ in 0..3 { + size /= 2; + let new_shape = &[b, size, 2, size, 2, c]; + xs = xs.reshape(new_shape)?; + xs = xs.permute((0, 2, 4, 1, 3, 5))?; + xs = xs.flatten(0, 2)?; + b *= 4; + } + xs = xs.reshape(((), NUM_TOKENS, c))?; + + Ok(xs) + })) +} + +fn hiera_mlp(in_channels: usize, out_channels: usize, vb: VarBuilder) -> Result> { + let fc1 = linear(in_channels, out_channels, vb.pp("fc1"))?; + let fc2 = linear(out_channels, in_channels, vb.pp("fc2"))?; + + Ok(Func::new(move |xs| { + let xs = xs.apply(&fc1)?.gelu()?.apply(&fc2)?; + Ok(xs) + })) +} + +fn hiera_attention( + in_channels: usize, + out_channels: usize, + heads: usize, + q_stride: usize, + window_size: usize, + use_mask_attention: bool, + vb: VarBuilder, +) -> Result> { + let head_dim = out_channels / heads; + + let scale = (head_dim as f64).powf(-0.5); + + let proj = linear(out_channels, out_channels, vb.pp("proj"))?; + let qkv = linear(in_channels, out_channels * 3, vb.pp("qkv"))?; + + Ok(Func::new(move |xs| { + let (b, n, _) = xs.dims3()?; + + let num_windows = if use_mask_attention { + n / (q_stride * window_size) + } else { + 1 + }; + let qkv = xs.apply(&qkv)?; + + let ec = qkv.elem_count(); + let s = ec / (b * num_windows * 3 * heads * head_dim); + let qkv = qkv + .reshape((b, s, num_windows, 3, heads, head_dim))? + .permute((3, 0, 4, 2, 1, 5))?; + + let mut q = qkv.get(0)?; + let k = qkv.get(1)?; + let v = qkv.get(2)?; + + if q_stride > 1 { + let ec = q.elem_count(); + let s = ec / (b * num_windows * q_stride * heads * head_dim); + q = q + .reshape((b, heads, num_windows, q_stride, s, head_dim))? + .max(3)?; + } + + let q = (q * scale)?; + + // Q, K and V are 6 dimensional with the first dimension being 1. + // Squeeze them for the attention calculation since 6 dimensional matmuls are not supported. + let att = q + .squeeze(0)? + .matmul(&k.squeeze(0)?.transpose(D::Minus2, D::Minus1)?)?; + let att = softmax(&att, D::Minus1)?; + let xs = att.matmul(&v.squeeze(0)?)?.unsqueeze(0)?; + + let xs = xs.transpose(1, 3)?.reshape((b, (), out_channels))?; + let xs = xs.apply(&proj)?; + + Ok(xs) + })) +} + +fn hiera_block( + heads: usize, + in_channels: usize, + out_channels: usize, + q_stride: usize, + window_size: usize, + use_mask_attention: bool, + vb: VarBuilder, +) -> Result> { + let norm1 = layer_norm(in_channels, 1e-6, vb.pp("norm1"))?; + let norm2 = layer_norm(out_channels, 1e-6, vb.pp("norm2"))?; + let proj = linear(in_channels, out_channels, vb.pp("proj")); + let stride = 4; + let mlp = hiera_mlp(out_channels, out_channels * 4, vb.pp("mlp"))?; + let attn = hiera_attention( + in_channels, + out_channels, + heads, + q_stride, + window_size, + use_mask_attention, + vb.pp("attn"), + )?; + + Ok(Func::new(move |xs| { + let mut xs = xs.clone(); + let xs_norm = xs.apply_t(&norm1, false)?; + if let Ok(p) = &proj { + xs = xs_norm.apply(p)?; + let (a, _, d) = xs.dims3()?; + xs = xs.reshape((a, stride, (), d))?.max(1)?; + } + let xs = (xs + &xs_norm.apply(&attn)?)?; + + let xs = (&xs + &xs.apply_t(&norm2, false)?.apply(&mlp)?)?; + + Ok(xs) + })) +} + +fn hiera_blocks(cfg: &Config, vb: VarBuilder) -> Result> { + let nblocks = cfg.stages.iter().sum(); + let mut blocks = Vec::with_capacity(nblocks); + + let mut out_channels = cfg.channels; + let mut in_channels = out_channels; + let mut heads = cfg.heads; + let mut b = 0; + + let mut q_stride = 1; + let mut window_size = 64; + + for s in 0..4 { + let use_mask_attention = s < 2; + + for _ in 0..cfg.stages[s] { + blocks.push(hiera_block( + heads, + in_channels, + out_channels, + q_stride, + window_size, + use_mask_attention, + vb.pp(b), + )?); + b += 1; + in_channels = out_channels; + q_stride = 1; + } + q_stride = 4; + out_channels *= 2; + heads *= 2; + window_size /= 4; + } + + Ok(Func::new(move |xs| { + let mut xs = xs.clone(); + for block in blocks.iter() { + xs = xs.apply(block)? + } + Ok(xs) + })) +} + +fn hiera_head(outputs: usize, nclasses: usize, vb: VarBuilder) -> Result> { + let norm = layer_norm(outputs, 1e-6, vb.pp("norm"))?; + let linear = linear(outputs, nclasses, vb.pp("fc"))?; + Ok(Func::new(move |xs| { + xs.apply_t(&norm, false)?.apply(&linear) + })) +} + +// Build a hiera model for a given configuration. +fn hiera_model(cfg: &Config, nclasses: Option, vb: VarBuilder) -> Result> { + let cls = match nclasses { + None => None, + Some(nclasses) => { + let outputs = cfg.channels * 8; + let head = hiera_head(outputs, nclasses, vb.pp("head"))?; + Some(head) + } + }; + + let embeddings = hiera_embeddings(cfg.channels, vb.clone())?; + let unroll = hiera_unroll()?; + let blocks = hiera_blocks(cfg, vb.pp("blocks"))?; + + Ok(Func::new(move |xs| { + let xs = xs + .apply(&embeddings)? + .apply(&unroll)? + .apply(&blocks)? + .mean(1)?; + match &cls { + None => Ok(xs), + Some(cls) => xs.apply(cls), + } + })) +} + +pub fn hiera(cfg: &Config, nclasses: usize, vb: VarBuilder) -> Result> { + hiera_model(cfg, Some(nclasses), vb) +} + +pub fn hiera_no_final_layer(cfg: &Config, vb: VarBuilder) -> Result> { + hiera_model(cfg, None, vb) +} diff --git a/patches/candle-transformers/src/models/jina_bert.rs b/patches/candle-transformers/src/models/jina_bert.rs new file mode 100644 index 0000000000..40535a8bb9 --- /dev/null +++ b/patches/candle-transformers/src/models/jina_bert.rs @@ -0,0 +1,406 @@ +//! # JinaBERT inference implementation +//! +//! Based on implementation from huggingface for Jina BERT and its variants +//! +//! See: [Jina Embeddings on HuggingFace](https://huggingface.co/jinaai/jina-embeddings-v2-base-en) + +use super::with_tracing::{linear, linear_no_bias, Embedding, Linear}; +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::{layer_norm, LayerNorm, Module, VarBuilder}; +use serde::Deserialize; + +pub const DTYPE: DType = DType::F32; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum PositionEmbeddingType { + Absolute, + Alibi, +} + +// https://huggingface.co/jinaai/jina-bert-implementation/blob/main/configuration_bert.py +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub intermediate_size: usize, + pub hidden_act: candle_nn::Activation, + pub max_position_embeddings: usize, + pub type_vocab_size: usize, + pub initializer_range: f64, + pub layer_norm_eps: f64, + pub pad_token_id: usize, + pub position_embedding_type: PositionEmbeddingType, +} + +impl Config { + pub fn v2_base() -> Self { + // https://huggingface.co/jinaai/jina-embeddings-v2-base-en/blob/main/config.json + Self { + vocab_size: 30528, + hidden_size: 768, + num_hidden_layers: 12, + num_attention_heads: 12, + intermediate_size: 3072, + hidden_act: candle_nn::Activation::Gelu, + max_position_embeddings: 8192, + type_vocab_size: 2, + initializer_range: 0.02, + layer_norm_eps: 1e-12, + pad_token_id: 0, + position_embedding_type: PositionEmbeddingType::Alibi, + } + } + + #[allow(clippy::too_many_arguments)] + pub fn new( + vocab_size: usize, + hidden_size: usize, + num_hidden_layers: usize, + num_attention_heads: usize, + intermediate_size: usize, + hidden_act: candle_nn::Activation, + max_position_embeddings: usize, + type_vocab_size: usize, + initializer_range: f64, + layer_norm_eps: f64, + pad_token_id: usize, + position_embedding_type: PositionEmbeddingType, + ) -> Self { + Config { + vocab_size, + hidden_size, + num_hidden_layers, + num_attention_heads, + intermediate_size, + hidden_act, + max_position_embeddings, + type_vocab_size, + initializer_range, + layer_norm_eps, + pad_token_id, + position_embedding_type, + } + } +} + +#[derive(Clone, Debug)] +struct BertEmbeddings { + word_embeddings: Embedding, + // no position_embeddings as we only support alibi. + token_type_embeddings: Embedding, + layer_norm: LayerNorm, + span: tracing::Span, +} + +impl BertEmbeddings { + fn new(vb: VarBuilder, cfg: &Config) -> Result { + let word_embeddings = + Embedding::new(cfg.vocab_size, cfg.hidden_size, vb.pp("word_embeddings"))?; + let token_type_embeddings = Embedding::new( + cfg.type_vocab_size, + cfg.hidden_size, + vb.pp("token_type_embeddings"), + )?; + let layer_norm = layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("LayerNorm"))?; + Ok(Self { + word_embeddings, + token_type_embeddings, + layer_norm, + span: tracing::span!(tracing::Level::TRACE, "embeddings"), + }) + } +} + +impl Module for BertEmbeddings { + fn forward(&self, input_ids: &Tensor) -> Result { + let _enter = self.span.enter(); + let (b_size, seq_len) = input_ids.dims2()?; + let input_embeddings = self.word_embeddings.forward(input_ids)?; + let token_type_embeddings = Tensor::zeros(seq_len, DType::U32, input_ids.device())? + .broadcast_left(b_size)? + .apply(&self.token_type_embeddings)?; + let embeddings = (&input_embeddings + token_type_embeddings)?; + let embeddings = self.layer_norm.forward(&embeddings)?; + Ok(embeddings) + } +} + +#[derive(Clone, Debug)] +struct BertSelfAttention { + query: Linear, + key: Linear, + value: Linear, + num_attention_heads: usize, + attention_head_size: usize, + span: tracing::Span, + span_softmax: tracing::Span, +} + +impl BertSelfAttention { + fn new(vb: VarBuilder, cfg: &Config) -> Result { + let attention_head_size = cfg.hidden_size / cfg.num_attention_heads; + let all_head_size = cfg.num_attention_heads * attention_head_size; + let hidden_size = cfg.hidden_size; + let query = linear(hidden_size, all_head_size, vb.pp("query"))?; + let value = linear(hidden_size, all_head_size, vb.pp("value"))?; + let key = linear(hidden_size, all_head_size, vb.pp("key"))?; + Ok(Self { + query, + key, + value, + num_attention_heads: cfg.num_attention_heads, + attention_head_size, + span: tracing::span!(tracing::Level::TRACE, "self-attn"), + span_softmax: tracing::span!(tracing::Level::TRACE, "softmax"), + }) + } + + fn transpose_for_scores(&self, xs: &Tensor) -> Result { + let mut x_shape = xs.dims().to_vec(); + x_shape.pop(); + x_shape.push(self.num_attention_heads); + x_shape.push(self.attention_head_size); + xs.reshape(x_shape)?.transpose(1, 2)?.contiguous() + } + + fn forward(&self, xs: &Tensor, bias: &Tensor) -> Result { + let _enter = self.span.enter(); + let query_layer = self.query.forward(xs)?; + let key_layer = self.key.forward(xs)?; + let value_layer = self.value.forward(xs)?; + + let query_layer = self.transpose_for_scores(&query_layer)?; + let key_layer = self.transpose_for_scores(&key_layer)?; + let value_layer = self.transpose_for_scores(&value_layer)?; + + let attention_scores = query_layer.matmul(&key_layer.t()?)?; + let attention_scores = (attention_scores / (self.attention_head_size as f64).sqrt())?; + let attention_scores = attention_scores.broadcast_add(bias)?; + let attention_probs = { + let _enter_sm = self.span_softmax.enter(); + candle_nn::ops::softmax_last_dim(&attention_scores)? + }; + let context_layer = attention_probs.matmul(&value_layer)?; + let context_layer = context_layer.transpose(1, 2)?.contiguous()?; + let context_layer = context_layer.flatten_from(D::Minus2)?; + Ok(context_layer) + } +} + +#[derive(Clone, Debug)] +struct BertSelfOutput { + dense: Linear, + layer_norm: LayerNorm, + span: tracing::Span, +} + +impl BertSelfOutput { + fn new(vb: VarBuilder, cfg: &Config) -> Result { + let dense = linear(cfg.hidden_size, cfg.hidden_size, vb.pp("dense"))?; + let layer_norm = layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("LayerNorm"))?; + Ok(Self { + dense, + layer_norm, + span: tracing::span!(tracing::Level::TRACE, "self-out"), + }) + } + + fn forward(&self, xs: &Tensor, input_tensor: &Tensor) -> Result { + let _enter = self.span.enter(); + let xs = self.dense.forward(xs)?; + self.layer_norm.forward(&(xs + input_tensor)?) + } +} + +#[derive(Clone, Debug)] +struct BertAttention { + self_attention: BertSelfAttention, + self_output: BertSelfOutput, + span: tracing::Span, +} + +impl BertAttention { + fn new(vb: VarBuilder, cfg: &Config) -> Result { + let self_attention = BertSelfAttention::new(vb.pp("self"), cfg)?; + let self_output = BertSelfOutput::new(vb.pp("output"), cfg)?; + Ok(Self { + self_attention, + self_output, + span: tracing::span!(tracing::Level::TRACE, "attn"), + }) + } + + fn forward(&self, xs: &Tensor, bias: &Tensor) -> Result { + let _enter = self.span.enter(); + let self_outputs = self.self_attention.forward(xs, bias)?; + let attention_output = self.self_output.forward(&self_outputs, xs)?; + Ok(attention_output) + } +} + +#[derive(Clone, Debug)] +struct BertGLUMLP { + gated_layers: Linear, + act: candle_nn::Activation, + wo: Linear, + layernorm: LayerNorm, + intermediate_size: usize, +} + +impl BertGLUMLP { + fn new(vb: VarBuilder, cfg: &Config) -> Result { + let gated_layers = linear_no_bias( + cfg.hidden_size, + cfg.intermediate_size * 2, + vb.pp("gated_layers"), + )?; + let act = candle_nn::Activation::Gelu; // geglu + let wo = linear(cfg.intermediate_size, cfg.hidden_size, vb.pp("wo"))?; + let layernorm = layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("layernorm"))?; + Ok(Self { + gated_layers, + act, + wo, + layernorm, + intermediate_size: cfg.intermediate_size, + }) + } +} + +impl Module for BertGLUMLP { + fn forward(&self, xs: &Tensor) -> Result { + let residual = xs; + let xs = xs.apply(&self.gated_layers)?; + let gated = xs.narrow(D::Minus1, 0, self.intermediate_size)?; + let non_gated = xs.narrow(D::Minus1, self.intermediate_size, self.intermediate_size)?; + let xs = (gated.apply(&self.act) * non_gated)?.apply(&self.wo); + (xs + residual)?.apply(&self.layernorm) + } +} + +#[derive(Clone, Debug)] +struct BertLayer { + attention: BertAttention, + mlp: BertGLUMLP, + span: tracing::Span, +} + +impl BertLayer { + fn new(vb: VarBuilder, cfg: &Config) -> Result { + let attention = BertAttention::new(vb.pp("attention"), cfg)?; + let mlp = BertGLUMLP::new(vb.pp("mlp"), cfg)?; + Ok(Self { + attention, + mlp, + span: tracing::span!(tracing::Level::TRACE, "layer"), + }) + } + + fn forward(&self, xs: &Tensor, bias: &Tensor) -> Result { + let _enter = self.span.enter(); + self.attention.forward(xs, bias)?.apply(&self.mlp) + } +} + +fn build_alibi_bias(cfg: &Config) -> Result { + let n_heads = cfg.num_attention_heads; + let seq_len = cfg.max_position_embeddings; + let alibi_bias = Tensor::arange(0, seq_len as i64, &Device::Cpu)?.to_dtype(DType::F32)?; + let alibi_bias = { + let a1 = alibi_bias.reshape((1, seq_len))?; + let a2 = alibi_bias.reshape((seq_len, 1))?; + a1.broadcast_sub(&a2)?.abs()?.broadcast_left(n_heads)? + }; + let mut n_heads2 = 1; + while n_heads2 < n_heads { + n_heads2 *= 2 + } + let slopes = (1..=n_heads2) + .map(|v| -1f32 / 2f32.powf((v * 8) as f32 / n_heads2 as f32)) + .collect::>(); + let slopes = if n_heads2 == n_heads { + slopes + } else { + slopes + .iter() + .skip(1) + .step_by(2) + .chain(slopes.iter().step_by(2)) + .take(n_heads) + .cloned() + .collect::>() + }; + let slopes = Tensor::new(slopes, &Device::Cpu)?.reshape((1, (), 1, 1))?; + alibi_bias.to_dtype(DType::F32)?.broadcast_mul(&slopes) +} + +#[derive(Clone, Debug)] +struct BertEncoder { + alibi: Tensor, + layers: Vec, + span: tracing::Span, +} + +impl BertEncoder { + fn new(vb: VarBuilder, cfg: &Config) -> Result { + if cfg.position_embedding_type != PositionEmbeddingType::Alibi { + candle::bail!("only alibi is supported as a position-embedding-type") + } + let layers = (0..cfg.num_hidden_layers) + .map(|index| BertLayer::new(vb.pp(format!("layer.{index}")), cfg)) + .collect::>>()?; + let span = tracing::span!(tracing::Level::TRACE, "encoder"); + let alibi = build_alibi_bias(cfg)?.to_device(vb.device())?; + Ok(Self { + alibi, + layers, + span, + }) + } +} + +impl Module for BertEncoder { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let seq_len = xs.dim(1)?; + let alibi_bias = self.alibi.i((.., .., ..seq_len, ..seq_len))?; + let mut xs = xs.clone(); + for layer in self.layers.iter() { + xs = layer.forward(&xs, &alibi_bias)? + } + Ok(xs) + } +} + +#[derive(Clone, Debug)] +pub struct BertModel { + embeddings: BertEmbeddings, + encoder: BertEncoder, + pub device: Device, + span: tracing::Span, +} + +impl BertModel { + pub fn new(vb: VarBuilder, cfg: &Config) -> Result { + let embeddings = BertEmbeddings::new(vb.pp("embeddings"), cfg)?; + let encoder = BertEncoder::new(vb.pp("encoder"), cfg)?; + Ok(Self { + embeddings, + encoder, + device: vb.device().clone(), + span: tracing::span!(tracing::Level::TRACE, "model"), + }) + } +} + +impl Module for BertModel { + fn forward(&self, input_ids: &Tensor) -> Result { + let _enter = self.span.enter(); + let embedding_output = self.embeddings.forward(input_ids)?; + let sequence_output = self.encoder.forward(&embedding_output)?; + Ok(sequence_output) + } +} diff --git a/patches/candle-transformers/src/models/llama.rs b/patches/candle-transformers/src/models/llama.rs new file mode 100644 index 0000000000..4396063ff7 --- /dev/null +++ b/patches/candle-transformers/src/models/llama.rs @@ -0,0 +1,536 @@ +//! Llama inference implementation. +//! +//! See ["LLaMA: Open and Efficient Foundation Language Models"](https://arxiv.org/abs/2302.13971) +//! +//! Implementation based on Hugging Face's [transformers](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py) + +use super::with_tracing::{linear_no_bias as linear, Linear, RmsNorm}; +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::{embedding, Embedding, Module, VarBuilder}; +use std::{collections::HashMap, f32::consts::PI}; + +pub const DEFAULT_MAX_SEQ_LEN: usize = 4096; + +#[derive(Debug, Clone, serde::Deserialize, Default)] +pub enum Llama3RopeType { + #[serde(rename = "llama3")] + Llama3, + #[default] + #[serde(rename = "default")] + Default, +} + +#[derive(Debug, Clone, serde::Deserialize, Default)] +pub struct Llama3RopeConfig { + pub factor: f32, + pub low_freq_factor: f32, + pub high_freq_factor: f32, + pub original_max_position_embeddings: usize, + pub rope_type: Llama3RopeType, +} +#[derive(Debug, Clone, serde::Deserialize)] +#[serde(untagged)] +pub enum LlamaEosToks { + Single(u32), + Multiple(Vec), +} + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct LlamaConfig { + pub hidden_size: usize, + pub intermediate_size: usize, + pub vocab_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: Option, + pub rms_norm_eps: f64, + #[serde(default = "default_rope")] + pub rope_theta: f32, + pub bos_token_id: Option, + pub eos_token_id: Option, + pub rope_scaling: Option, + pub max_position_embeddings: usize, + pub tie_word_embeddings: Option, +} + +impl LlamaConfig { + pub fn num_key_value_heads(&self) -> usize { + self.num_key_value_heads.unwrap_or(self.num_attention_heads) + } +} + +fn default_rope() -> f32 { + 10_000.0 +} + +impl LlamaConfig { + pub fn into_config(self, use_flash_attn: bool) -> Config { + Config { + hidden_size: self.hidden_size, + intermediate_size: self.intermediate_size, + vocab_size: self.vocab_size, + num_hidden_layers: self.num_hidden_layers, + num_attention_heads: self.num_attention_heads, + num_key_value_heads: self.num_key_value_heads(), + rms_norm_eps: self.rms_norm_eps, + rope_theta: self.rope_theta, + use_flash_attn, + bos_token_id: self.bos_token_id, + eos_token_id: self.eos_token_id, + rope_scaling: self.rope_scaling, + max_position_embeddings: self.max_position_embeddings, + tie_word_embeddings: self.tie_word_embeddings.unwrap_or(false), + } + } +} + +#[derive(Debug, Clone)] +pub struct Config { + pub hidden_size: usize, + pub intermediate_size: usize, + pub vocab_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub use_flash_attn: bool, + pub rms_norm_eps: f64, + pub rope_theta: f32, + pub bos_token_id: Option, + pub eos_token_id: Option, + pub rope_scaling: Option, + pub max_position_embeddings: usize, + pub tie_word_embeddings: bool, +} + +impl Config { + pub fn config_7b_v1(use_flash_attn: bool) -> Self { + Self { + hidden_size: 4096, + intermediate_size: 11008, + vocab_size: 32000, + num_hidden_layers: 32, + num_attention_heads: 32, + num_key_value_heads: 32, + use_flash_attn, + rms_norm_eps: 1e-6, + rope_theta: 10_000.0, + bos_token_id: None, + eos_token_id: None, + rope_scaling: None, + max_position_embeddings: DEFAULT_MAX_SEQ_LEN, + tie_word_embeddings: false, + } + } + + pub fn config_7b_v2(use_flash_attn: bool) -> Self { + Self { + hidden_size: 4096, + intermediate_size: 11008, + vocab_size: 32000, + num_hidden_layers: 32, + num_attention_heads: 32, + num_key_value_heads: 32, + use_flash_attn, + rms_norm_eps: 1e-5, + rope_theta: 10_000.0, + bos_token_id: None, + eos_token_id: None, + rope_scaling: None, + max_position_embeddings: DEFAULT_MAX_SEQ_LEN, + tie_word_embeddings: false, + } + } +} + +#[derive(Debug, Clone)] +pub struct Cache { + masks: HashMap, + pub use_kv_cache: bool, + kvs: Vec>, + cos: Tensor, + sin: Tensor, + device: Device, +} + +fn calculate_default_inv_freq(cfg: &Config) -> Vec { + let head_dim = cfg.hidden_size / cfg.num_attention_heads; + (0..head_dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f32 / head_dim as f32)) + .collect() +} + +impl Cache { + pub fn new(use_kv_cache: bool, dtype: DType, config: &Config, device: &Device) -> Result { + // precompute freqs_cis + let theta = match &config.rope_scaling { + None + | Some(Llama3RopeConfig { + rope_type: Llama3RopeType::Default, + .. + }) => calculate_default_inv_freq(config), + Some(rope_scaling) => { + let low_freq_wavelen = rope_scaling.original_max_position_embeddings as f32 + / rope_scaling.low_freq_factor; + let high_freq_wavelen = rope_scaling.original_max_position_embeddings as f32 + / rope_scaling.high_freq_factor; + + calculate_default_inv_freq(config) + .into_iter() + .map(|freq| { + let wavelen = 2. * PI / freq; + if wavelen < high_freq_wavelen { + freq + } else if wavelen > low_freq_wavelen { + freq / rope_scaling.factor + } else { + let smooth = (rope_scaling.original_max_position_embeddings as f32 + / wavelen + - rope_scaling.low_freq_factor) + / (rope_scaling.high_freq_factor - rope_scaling.low_freq_factor); + (1. - smooth) * freq / rope_scaling.factor + smooth * freq + } + }) + .collect::>() + } + }; + + let theta = Tensor::new(theta, device)?; + + let idx_theta = Tensor::arange(0, config.max_position_embeddings as u32, device)? + .to_dtype(DType::F32)? + .reshape((config.max_position_embeddings, 1))? + .matmul(&theta.reshape((1, theta.elem_count()))?)?; + // This is different from the paper, see: + // https://github.com/huggingface/transformers/blob/6112b1c6442aaf7affd2b0676a1cd4eee30c45cf/src/transformers/models/llama/modeling_llama.py#L112 + let cos = idx_theta.cos()?.to_dtype(dtype)?; + let sin = idx_theta.sin()?.to_dtype(dtype)?; + Ok(Self { + masks: HashMap::new(), + use_kv_cache, + kvs: vec![None; config.num_hidden_layers], + device: device.clone(), + cos, + sin, + }) + } + + 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), &self.device)?; + self.masks.insert(t, mask.clone()); + Ok(mask) + } + } +} + +#[derive(Debug, Clone)] +struct CausalSelfAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_attention_heads: usize, + num_key_value_heads: usize, + head_dim: usize, + use_flash_attn: bool, + span: tracing::Span, + span_rot: tracing::Span, + max_position_embeddings: usize, +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +impl CausalSelfAttention { + fn apply_rotary_emb(&self, x: &Tensor, index_pos: usize, cache: &Cache) -> Result { + let _enter = self.span_rot.enter(); + let (_b_sz, _, seq_len, _hidden_size) = x.dims4()?; + let cos = cache.cos.narrow(0, index_pos, seq_len)?; + let sin = cache.sin.narrow(0, index_pos, seq_len)?; + candle_nn::rotary_emb::rope(x, &cos, &sin) + } + + fn forward( + &self, + x: &Tensor, + index_pos: usize, + block_idx: usize, + cache: &mut Cache, + ) -> Result { + let _enter = self.span.enter(); + let (b_sz, seq_len, hidden_size) = x.dims3()?; + let q = self.q_proj.forward(x)?; + let k = self.k_proj.forward(x)?; + let v = self.v_proj.forward(x)?; + + let q = q + .reshape((b_sz, seq_len, self.num_attention_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let k = k + .reshape((b_sz, seq_len, self.num_key_value_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let mut v = v + .reshape((b_sz, seq_len, self.num_key_value_heads, self.head_dim))? + .transpose(1, 2)?; + + let q = self.apply_rotary_emb(&q, index_pos, cache)?; + let mut k = self.apply_rotary_emb(&k, index_pos, cache)?; + + if cache.use_kv_cache { + if let Some((cache_k, cache_v)) = &cache.kvs[block_idx] { + k = Tensor::cat(&[cache_k, &k], 2)?.contiguous()?; + v = Tensor::cat(&[cache_v, &v], 2)?.contiguous()?; + let k_seq_len = k.dims()[1]; + if k_seq_len > self.max_position_embeddings { + k = k + .narrow( + D::Minus1, + k_seq_len - self.max_position_embeddings, + self.max_position_embeddings, + )? + .contiguous()? + } + let v_seq_len = v.dims()[1]; + if v_seq_len > 2 * self.max_position_embeddings { + v = v + .narrow( + D::Minus1, + v_seq_len - self.max_position_embeddings, + self.max_position_embeddings, + )? + .contiguous()? + } + } + cache.kvs[block_idx] = Some((k.clone(), v.clone())) + } + + let k = self.repeat_kv(k)?; + let v = self.repeat_kv(v)?; + + let y = if self.use_flash_attn { + // flash-attn expects (b_sz, seq_len, nheads, head_dim) + let q = q.transpose(1, 2)?; + let k = k.transpose(1, 2)?; + let v = v.transpose(1, 2)?; + let softmax_scale = 1f32 / (self.head_dim as f32).sqrt(); + flash_attn(&q, &k, &v, softmax_scale, seq_len > 1)?.transpose(1, 2)? + } else { + let in_dtype = q.dtype(); + let q = q.to_dtype(DType::F32)?; + let k = k.to_dtype(DType::F32)?; + let v = v.to_dtype(DType::F32)?; + let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?; + let att = if seq_len == 1 { + att + } else { + let mask = cache.mask(seq_len)?.broadcast_as(att.shape())?; + 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. + att.matmul(&v.contiguous()?)?.to_dtype(in_dtype)? + }; + let y = y.transpose(1, 2)?.reshape(&[b_sz, seq_len, hidden_size])?; + let y = self.o_proj.forward(&y)?; + Ok(y) + } + + fn repeat_kv(&self, x: Tensor) -> Result { + crate::utils::repeat_kv(x, self.num_attention_heads / self.num_key_value_heads) + } + + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "attn"); + let span_rot = tracing::span!(tracing::Level::TRACE, "attn-rot"); + let size_in = cfg.hidden_size; + let size_q = (cfg.hidden_size / cfg.num_attention_heads) * cfg.num_attention_heads; + let size_kv = (cfg.hidden_size / cfg.num_attention_heads) * cfg.num_key_value_heads; + let q_proj = linear(size_in, size_q, vb.pp("q_proj"))?; + let k_proj = linear(size_in, size_kv, vb.pp("k_proj"))?; + let v_proj = linear(size_in, size_kv, vb.pp("v_proj"))?; + let o_proj = linear(size_q, size_in, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_attention_heads: cfg.num_attention_heads, + num_key_value_heads: cfg.num_key_value_heads, + head_dim: cfg.hidden_size / cfg.num_attention_heads, + use_flash_attn: cfg.use_flash_attn, + span, + span_rot, + max_position_embeddings: cfg.max_position_embeddings, + }) + } +} + +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) +} + +#[derive(Debug, Clone)] +struct Mlp { + c_fc1: Linear, + c_fc2: Linear, + c_proj: Linear, + span: tracing::Span, +} + +impl Mlp { + fn forward(&self, x: &Tensor) -> Result { + let _enter = self.span.enter(); + let x = (candle_nn::ops::silu(&self.c_fc1.forward(x)?)? * self.c_fc2.forward(x)?)?; + self.c_proj.forward(&x) + } + + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "mlp"); + let h_size = cfg.hidden_size; + let i_size = cfg.intermediate_size; + let c_fc1 = linear(h_size, i_size, vb.pp("gate_proj"))?; + let c_fc2 = linear(h_size, i_size, vb.pp("up_proj"))?; + let c_proj = linear(i_size, h_size, vb.pp("down_proj"))?; + Ok(Self { + c_fc1, + c_fc2, + c_proj, + span, + }) + } +} + +#[derive(Debug, Clone)] +struct Block { + rms_1: RmsNorm, + attn: CausalSelfAttention, + rms_2: RmsNorm, + mlp: Mlp, + span: tracing::Span, +} + +impl Block { + fn forward( + &self, + x: &Tensor, + index_pos: usize, + block_idx: usize, + cache: &mut Cache, + ) -> Result { + let _enter = self.span.enter(); + let residual = x; + let x = self.rms_1.forward(x)?; + let x = (self.attn.forward(&x, index_pos, block_idx, cache)? + residual)?; + let residual = &x; + let x = (self.mlp.forward(&self.rms_2.forward(&x)?)? + residual)?; + Ok(x) + } + + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "block"); + let attn = CausalSelfAttention::load(vb.pp("self_attn"), cfg)?; + let mlp = Mlp::load(vb.pp("mlp"), cfg)?; + let rms_1 = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let rms_2 = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + rms_1, + attn, + rms_2, + mlp, + span, + }) + } +} + +#[derive(Debug, Clone)] +pub struct Llama { + wte: Embedding, + blocks: Vec, + ln_f: RmsNorm, + lm_head: Linear, +} + +impl Llama { + // required by LLaVA + pub fn embed(&self, x: &Tensor) -> Result { + self.wte.forward(x) + } + // required by LLaVA + pub fn forward_input_embed( + &self, + input_embed: &Tensor, + index_pos: usize, + cache: &mut Cache, + ) -> Result { + let (_, seq_len, _) = input_embed.dims3()?; + let mut x = input_embed.clone(); + for (block_idx, block) in self.blocks.iter().enumerate() { + x = block.forward(&x, index_pos, block_idx, cache)?; + } + let x = self.ln_f.forward(&x)?; + let x = x.i((.., seq_len - 1, ..))?.contiguous()?; + let logits = self.lm_head.forward(&x)?; + logits.to_dtype(DType::F32) + } + + pub fn forward(&self, x: &Tensor, index_pos: usize, cache: &mut Cache) -> Result { + let (_b_sz, seq_len) = x.dims2()?; + let mut x = self.wte.forward(x)?; + for (block_idx, block) in self.blocks.iter().enumerate() { + x = block.forward(&x, index_pos, block_idx, cache)?; + } + let x = self.ln_f.forward(&x)?; + let x = x.i((.., seq_len - 1, ..))?.contiguous()?; + let logits = self.lm_head.forward(&x)?; + logits.to_dtype(DType::F32) + } + + pub fn load(vb: VarBuilder, cfg: &Config) -> Result { + let wte = embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("model.embed_tokens"))?; + let lm_head = if cfg.tie_word_embeddings { + Linear::from_weights(wte.embeddings().clone(), None) + } else { + linear(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))? + }; + let ln_f = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("model.norm"))?; + let blocks: Vec<_> = (0..cfg.num_hidden_layers) + .map(|i| Block::load(vb.pp(format!("model.layers.{i}")), cfg).unwrap()) + .collect(); + + Ok(Self { + wte, + blocks, + ln_f, + lm_head, + }) + } +} diff --git a/patches/candle-transformers/src/models/llama2_c.rs b/patches/candle-transformers/src/models/llama2_c.rs new file mode 100644 index 0000000000..930c8b8aa6 --- /dev/null +++ b/patches/candle-transformers/src/models/llama2_c.rs @@ -0,0 +1,375 @@ +//! Llama2 inference implementation. +//! +//! See ["LLaMA 2: Open Foundation and Fine-Tuned Chat Models"](https://arxiv.org/abs/2307.09288) +//! +//! - ⚡ [Interactive Wasm Example](https://huggingface.co/spaces/lmz/candle-llama2) +//! - 💻 llama2.c [GH Link](https://github.com/karpathy/llama2.c) +//! + +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::linear_no_bias as linear; +use candle_nn::{embedding, rms_norm, Embedding, Linear, Module, RmsNorm, VarBuilder}; +use std::collections::HashMap; + +#[derive(Debug, Clone)] +pub struct Config { + pub dim: usize, // transformer dimension + pub hidden_dim: usize, // for ffn layers + pub n_layers: usize, // number of layers + pub n_heads: usize, // number of query heads + pub n_kv_heads: usize, // number of key/value heads (can be < query heads because of multiquery) + pub vocab_size: usize, // vocabulary size, usually 256 (byte-level) + pub seq_len: usize, // max sequence length + pub norm_eps: f64, +} + +impl Config { + pub fn tiny_260k() -> Self { + Self { + dim: 64, + hidden_dim: 768, + n_layers: 5, + n_heads: 8, + n_kv_heads: 4, + vocab_size: 32000, + seq_len: 512, + norm_eps: 1e-5, + } + } + + pub fn tiny_15m() -> Self { + Self { + dim: 288, + hidden_dim: 768, + n_layers: 6, + n_heads: 6, + n_kv_heads: 6, + vocab_size: 32000, + seq_len: 256, + norm_eps: 1e-5, + } + } + + pub fn tiny_42m() -> Self { + Self { + dim: 512, + hidden_dim: 768, + n_layers: 8, + n_heads: 8, + n_kv_heads: 8, + vocab_size: 32000, + seq_len: 1024, + norm_eps: 1e-5, + } + } + + pub fn tiny_110m() -> Self { + Self { + dim: 768, + hidden_dim: 768, + n_layers: 12, + n_heads: 12, + n_kv_heads: 12, + vocab_size: 32000, + seq_len: 1024, + norm_eps: 1e-5, + } + } +} + +#[derive(Debug, Clone)] +pub struct Cache { + masks: HashMap, + pub use_kv_cache: bool, + pub kvs: Vec>, + pub cos: Tensor, + pub sin: Tensor, + device: Device, +} + +impl Cache { + pub fn new(use_kv_cache: bool, cfg: &Config, vb: VarBuilder) -> Result { + let n_elem = cfg.dim / cfg.n_heads; + let theta: Vec<_> = (0..n_elem) + .step_by(2) + .map(|i| 1f32 / 10000f32.powf(i as f32 / n_elem as f32)) + .collect(); + let theta = Tensor::new(theta.as_slice(), vb.device())?; + let idx_theta = Tensor::arange(0, cfg.seq_len as u32, vb.device())? + .to_dtype(DType::F32)? + .reshape((cfg.seq_len, 1))? + .matmul(&theta.reshape((1, theta.elem_count()))?)?; + let precomputed_cos = idx_theta.cos()?; + let precomputed_sin = idx_theta.sin()?; + + let freq_cis_real = vb + .get((cfg.seq_len, cfg.head_size() / 2), "freq_cis_real") + .unwrap_or(precomputed_cos); + let freq_cis_imag = vb + .get((cfg.seq_len, cfg.head_size() / 2), "freq_cis_imag") + .unwrap_or(precomputed_sin); + let cos = freq_cis_real.reshape((cfg.seq_len, cfg.head_size() / 2, 1))?; + let sin = freq_cis_imag.reshape((cfg.seq_len, cfg.head_size() / 2, 1))?; + Ok(Self { + masks: HashMap::new(), + use_kv_cache, + kvs: vec![None; cfg.n_layers], + cos, + sin, + device: vb.device().clone(), + }) + } + + pub 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), &self.device)?; + self.masks.insert(t, mask.clone()); + Ok(mask) + } + } +} + +fn silu(xs: &Tensor) -> Result { + xs / (xs.neg()?.exp()? + 1.0)? +} + +#[derive(Debug, Clone)] +struct CausalSelfAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + n_head: usize, + n_key_value_head: usize, + head_dim: usize, +} + +impl CausalSelfAttention { + fn apply_rotary_emb(&self, x: &Tensor, index_pos: usize, cache: &Cache) -> Result { + let (b_sz, seq_len, h, n_embd) = x.dims4()?; + let cos = cache.cos.i(index_pos..index_pos + seq_len)?; + let sin = cache.sin.i(index_pos..index_pos + seq_len)?; + let cos = cos.unsqueeze(1)?; + let sin = sin.unsqueeze(1)?; + let cos = cos.broadcast_as((b_sz, seq_len, 1, n_embd / 2, 1))?; + let sin = sin.broadcast_as((b_sz, seq_len, 1, n_embd / 2, 1))?; + let x = x.reshape((b_sz, seq_len, h, n_embd / 2, 2))?; + let x0 = x.narrow(D::Minus1, 0, 1)?; + let x1 = x.narrow(D::Minus1, 1, 1)?; + let dst0 = (x0.broadcast_mul(&cos)? - x1.broadcast_mul(&sin)?)?; + let dst1 = (x0.broadcast_mul(&sin)? + x1.broadcast_mul(&cos)?)?; + let rope = Tensor::cat(&[&dst0, &dst1], D::Minus1)?.reshape((b_sz, seq_len, h, n_embd))?; + Ok(rope) + } + + fn forward( + &self, + x: &Tensor, + index_pos: usize, + block_idx: usize, + cache: &mut Cache, + ) -> Result { + let (b_sz, seq_len, n_embd) = x.dims3()?; + let q = self.q_proj.forward(x)?; + let k = self.k_proj.forward(x)?; + let v = self.v_proj.forward(x)?; + + let q = q.reshape((b_sz, seq_len, self.n_head, self.head_dim))?; + let k = k.reshape((b_sz, seq_len, self.n_key_value_head, self.head_dim))?; + let mut v = v.reshape((b_sz, seq_len, self.n_key_value_head, self.head_dim))?; + + let q = self.apply_rotary_emb(&q, index_pos, cache)?; + let mut k = self.apply_rotary_emb(&k, index_pos, cache)?; + + if cache.use_kv_cache { + if let Some((cache_k, cache_v)) = &cache.kvs[block_idx] { + k = Tensor::cat(&[cache_k, &k], 1)?.contiguous()?; + v = Tensor::cat(&[cache_v, &v], 1)?.contiguous()?; + } + cache.kvs[block_idx] = Some((k.clone(), v.clone())) + } + + let k = self.repeat_kv(k)?; + let v = self.repeat_kv(v)?; + + let q = q.transpose(1, 2)?.contiguous()?; + let k = k.transpose(1, 2)?.contiguous()?; + let v = v.transpose(1, 2)?.contiguous()?; + + let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?; + let att = if seq_len <= 1 { + att + } else { + let mask = cache.mask(seq_len)?.broadcast_as(att.shape())?; + masked_fill(&att, &mask, f32::NEG_INFINITY)? + }; + let att = candle_nn::ops::softmax(&att, D::Minus1)?; + // 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.o_proj.forward(&y)?; + Ok(y) + } + + fn repeat_kv(&self, x: Tensor) -> Result { + let n_rep = self.n_head / self.n_key_value_head; + if n_rep == 1 { + Ok(x) + } else { + let (b_sz, seq_len, n_kv_head, head_dim) = x.dims4()?; + let x = x + .unsqueeze(3)? + .expand((b_sz, seq_len, n_kv_head, n_rep, head_dim))? + .reshape((b_sz, seq_len, n_kv_head * n_rep, head_dim))?; + Ok(x) + } + } + + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let size_in = cfg.dim; + let size_q = (cfg.dim / cfg.n_heads) * cfg.n_heads; + let size_kv = (cfg.dim / cfg.n_heads) * cfg.n_kv_heads; + let q_proj = linear(size_in, size_q, vb.pp("q_proj"))?; + let k_proj = linear(size_in, size_kv, vb.pp("k_proj"))?; + let v_proj = linear(size_in, size_kv, vb.pp("v_proj"))?; + let o_proj = linear(size_q, size_in, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + n_head: cfg.n_heads, + n_key_value_head: cfg.n_kv_heads, + head_dim: cfg.dim / cfg.n_heads, + }) + } +} + +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) +} + +#[derive(Debug, Clone)] +struct Mlp { + c_fc1: Linear, + c_fc2: Linear, + c_proj: Linear, +} + +impl Mlp { + fn new(c_fc1: Linear, c_fc2: Linear, c_proj: Linear) -> Self { + Self { + c_fc1, + c_fc2, + c_proj, + } + } + + fn forward(&self, x: &Tensor) -> Result { + let x = (silu(&self.c_fc1.forward(x)?)? * self.c_fc2.forward(x)?)?; + self.c_proj.forward(&x) + } + + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let h_size = cfg.dim; + let i_size = cfg.hidden_dim; + let c_fc1 = linear(h_size, i_size, vb.pp("gate_proj"))?; + let c_fc2 = linear(h_size, i_size, vb.pp("up_proj"))?; + let c_proj = linear(i_size, h_size, vb.pp("down_proj"))?; + Ok(Self::new(c_fc1, c_fc2, c_proj)) + } +} + +#[derive(Debug, Clone)] +struct Block { + rms_1: RmsNorm, + attn: CausalSelfAttention, + rms_2: RmsNorm, + mlp: Mlp, +} + +impl Block { + fn new(rms_1: RmsNorm, attn: CausalSelfAttention, rms_2: RmsNorm, mlp: Mlp) -> Self { + Self { + rms_1, + attn, + rms_2, + mlp, + } + } + + fn forward( + &self, + x: &Tensor, + index_pos: usize, + block_idx: usize, + cache: &mut Cache, + ) -> Result { + let residual = x; + let x = self.rms_1.forward(x)?; + let x = (self.attn.forward(&x, index_pos, block_idx, cache)? + residual)?; + let residual = &x; + let x = (self.mlp.forward(&self.rms_2.forward(&x)?)? + residual)?; + Ok(x) + } + + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let attn = CausalSelfAttention::load(vb.pp("self_attn"), cfg)?; + let mlp = Mlp::load(vb.pp("mlp"), cfg)?; + let input_layernorm = rms_norm(cfg.dim, cfg.norm_eps, vb.pp("input_layernorm"))?; + let post_attention_layernorm = + rms_norm(cfg.dim, cfg.norm_eps, vb.pp("post_attention_layernorm"))?; + Ok(Self::new( + input_layernorm, + attn, + post_attention_layernorm, + mlp, + )) + } +} + +#[derive(Debug, Clone)] +pub struct Llama { + wte: Embedding, + blocks: Vec, + ln_f: RmsNorm, + lm_head: Linear, + pub config: Config, +} + +impl Llama { + pub fn forward(&self, x: &Tensor, index_pos: usize, cache: &mut Cache) -> Result { + let (_b_sz, _seq_len) = x.dims2()?; + let mut x = self.wte.forward(x)?; + for (block_idx, block) in self.blocks.iter().enumerate() { + x = block.forward(&x, index_pos, block_idx, cache)?; + } + let x = self.ln_f.forward(&x)?; + let logits = self.lm_head.forward(&x)?; + logits.to_dtype(DType::F32) + } + + pub fn load(vb: VarBuilder, cfg: Config) -> Result { + let wte = embedding(cfg.vocab_size, cfg.dim, vb.pp("model.embed_tokens"))?; + let lm_head = linear(cfg.dim, cfg.vocab_size, vb.pp("lm_head"))?; + let ln_f = rms_norm(cfg.dim, cfg.norm_eps, vb.pp("model.norm"))?; + let blocks: Vec<_> = (0..cfg.n_layers) + .map(|i| Block::load(vb.pp(format!("model.layers.{i}")), &cfg).unwrap()) + .collect(); + Ok(Self { + wte, + blocks, + ln_f, + lm_head, + config: cfg, + }) + } +} diff --git a/patches/candle-transformers/src/models/llama2_c_weights.rs b/patches/candle-transformers/src/models/llama2_c_weights.rs new file mode 100644 index 0000000000..8149c214c9 --- /dev/null +++ b/patches/candle-transformers/src/models/llama2_c_weights.rs @@ -0,0 +1,173 @@ +//! Llama2 inference implementation. +//! +//! See ["LLaMA 2: Open Foundation and Fine-Tuned Chat Models"](https://arxiv.org/abs/2307.09288) +//! +//! Based on the [llama2.c](https://github.com/karpathy/llama2.c) implementation + +use byteorder::{LittleEndian, ReadBytesExt}; +use candle::{DType, Device, IndexOp, Result, Shape, Tensor}; +use candle_nn::VarBuilder; + +use super::llama2_c::Config; + +pub struct TransformerWeights { + // token embedding table + token_embedding_table: Tensor, // (vocab_size, dim) + // weights for rmsnorms + rms_att_weight: Tensor, // (layer, dim) rmsnorm weights + rms_ffn_weight: Tensor, // (layer, dim) + // weights for matmuls + wq: Tensor, // (layer, dim, dim) + wk: Tensor, // (layer, dim, dim) + wv: Tensor, // (layer, dim, dim) + wo: Tensor, // (layer, dim, dim) + // weights for ffn + w1: Tensor, // (layer, hidden_dim, dim) + w2: Tensor, // (layer, dim, hidden_dim) + w3: Tensor, // (layer, hidden_dim, dim) + // final rmsnorm + rms_final_weight: Tensor, // (dim,) + // freq_cis for RoPE relatively positional embeddings + freq_cis_real: Tensor, // (seq_len, head_size/2) + freq_cis_imag: Tensor, // (seq_len, head_size/2) +} + +fn read_i32(r: &mut R) -> Result { + let mut buf = [0u8; 4]; + r.read_exact(&mut buf)?; + Ok(i32::from_le_bytes(buf)) +} + +fn read_tensor>( + r: &mut R, + shape: S, + dev: &Device, +) -> Result { + let shape = shape.into(); + let mut data_t = vec![0f32; shape.elem_count()]; + r.read_f32_into::(&mut data_t)?; + let tensor = Tensor::from_vec(data_t, shape, dev)?; + Ok(tensor) +} + +impl Config { + pub fn from_reader(r: &mut R) -> Result { + let dim = read_i32(r)? as usize; + let hidden_dim = read_i32(r)? as usize; + let n_layers = read_i32(r)? as usize; + let n_heads = read_i32(r)? as usize; + let n_kv_heads = read_i32(r)? as usize; + let vocab_size = read_i32(r)? as usize; + let seq_len = read_i32(r)? as usize; + Ok(Self { + dim, + hidden_dim, + n_layers, + n_heads, + n_kv_heads, + vocab_size, + seq_len, + norm_eps: 1e-5, + }) + } + + pub fn head_size(&self) -> usize { + self.dim / self.n_heads + } +} + +impl TransformerWeights { + pub fn from_reader(r: &mut R, c: &Config, dev: &Device) -> Result { + let token_embedding_table = read_tensor(r, (c.vocab_size, c.dim), dev)?; + let rms_att_weight = read_tensor(r, (c.n_layers, c.dim), dev)?; + let wq = read_tensor(r, (c.n_layers, c.dim, c.dim), dev)?; + let wk = read_tensor(r, (c.n_layers, c.dim, c.dim), dev)?; + let wv = read_tensor(r, (c.n_layers, c.dim, c.dim), dev)?; + let wo = read_tensor(r, (c.n_layers, c.dim, c.dim), dev)?; + let rms_ffn_weight = read_tensor(r, (c.n_layers, c.dim), dev)?; + let w1 = read_tensor(r, (c.n_layers, c.hidden_dim, c.dim), dev)?; + let w2 = read_tensor(r, (c.n_layers, c.dim, c.hidden_dim), dev)?; + let w3 = read_tensor(r, (c.n_layers, c.hidden_dim, c.dim), dev)?; + let rms_final_weight = read_tensor(r, c.dim, dev)?; + let head_size = c.head_size(); + let freq_cis_real = read_tensor(r, (c.seq_len, head_size / 2), dev)?; + let freq_cis_imag = read_tensor(r, (c.seq_len, head_size / 2), dev)?; + Ok(Self { + token_embedding_table, + rms_att_weight, + wq, + wk, + wv, + wo, + rms_ffn_weight, + w1, + w2, + w3, + rms_final_weight, + freq_cis_real, + freq_cis_imag, + }) + } + + pub fn var_builder(&self, cfg: &Config, device: &Device) -> Result> { + // TODO: As of 2023-08-04, gemm is slower than expected when multiplying a matrix of + // size (1, k) with the transpose of a matrix of size (k, n) as it ends up transposing the + // second matrix back. We detect this case here and as a temporary hack make the weight + // matrix column major rather than row major. This ends up speeding up text generation from + // 120 token/s to 220 token/s on a Ryzen 2600X. + let tr = device.is_cpu() && !candle::utils::has_mkl(); + let tr = |x: Tensor| if tr { x.t()?.contiguous()?.t() } else { Ok(x) }; + let mut ws = std::collections::HashMap::new(); + let mut insert = |name: &str, t: Tensor| { + ws.insert(name.to_string(), t); + }; + insert("rot.freq_cis_real", self.freq_cis_real.clone()); + insert("rot.freq_cis_imag", self.freq_cis_imag.clone()); + insert( + "model.embed_tokens.weight", + self.token_embedding_table.clone(), + ); + insert("lm_head.weight", tr(self.token_embedding_table.clone())?); + insert("model.norm.weight", self.rms_final_weight.clone()); + for layer in 0..cfg.n_layers { + ws.insert( + format!("model.layers.{layer}.self_attn.q_proj.weight"), + tr(self.wq.i(layer)?)?, + ); + ws.insert( + format!("model.layers.{layer}.self_attn.k_proj.weight"), + tr(self.wk.i(layer)?)?, + ); + ws.insert( + format!("model.layers.{layer}.self_attn.v_proj.weight"), + tr(self.wv.i(layer)?)?, + ); + ws.insert( + format!("model.layers.{layer}.self_attn.o_proj.weight"), + tr(self.wo.i(layer)?)?, + ); + ws.insert( + format!("model.layers.{layer}.mlp.gate_proj.weight"), + tr(self.w1.i(layer)?)?, + ); + ws.insert( + format!("model.layers.{layer}.mlp.down_proj.weight"), + tr(self.w2.i(layer)?)?, + ); + ws.insert( + format!("model.layers.{layer}.mlp.up_proj.weight"), + tr(self.w3.i(layer)?)?, + ); + ws.insert( + format!("model.layers.{layer}.input_layernorm.weight"), + self.rms_att_weight.i(layer)?, + ); + ws.insert( + format!("model.layers.{layer}.post_attention_layernorm.weight"), + self.rms_ffn_weight.i(layer)?, + ); + } + let vb = VarBuilder::from_tensors(ws, DType::F32, device); + Ok(vb) + } +} diff --git a/patches/candle-transformers/src/models/llava/config.rs b/patches/candle-transformers/src/models/llava/config.rs new file mode 100644 index 0000000000..405eedb934 --- /dev/null +++ b/patches/candle-transformers/src/models/llava/config.rs @@ -0,0 +1,272 @@ +use std::collections::HashMap; + +use crate::models::{ + clip::{text_model::Activation, vision_model::ClipVisionConfig}, + llama::{Config, LlamaEosToks}, +}; +use serde::{Deserialize, Serialize}; + +// original config from liuhaotian/llava +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct LLaVAConfig { + pub architectures: Vec, + pub bos_token_id: usize, + pub eos_token_id: usize, + pub hidden_size: usize, + #[serde(default = "default_image_aspect_ratio")] + pub image_aspect_ratio: String, + pub image_crop_resolution: usize, + pub image_grid_pinpoints: Vec<(u32, u32)>, + pub image_split_resolution: usize, + pub intermediate_size: usize, + pub max_position_embeddings: usize, + pub mm_hidden_size: usize, + #[serde(default = "default_mm_patch_merge_type")] + pub mm_patch_merge_type: String, + pub mm_projector_type: String, + pub mm_use_im_start_end: bool, + pub mm_vision_select_feature: String, + pub mm_vision_select_layer: isize, + pub mm_vision_tower: Option, + pub model_type: String, + pub num_attention_heads: usize, + pub num_hidden_layers: usize, + pub num_key_value_heads: usize, + pub pad_token_id: usize, + pub rms_norm_eps: f32, + pub rope_theta: f32, + pub tokenizer_model_max_length: Option, + pub torch_dtype: String, + pub use_cache: bool, + pub vocab_size: usize, + #[serde(default = "default_image_token_index")] + pub image_token_index: isize, + #[serde(default = "default_hf")] + pub hf: bool, + pub tie_word_embeddings: Option, +} + +fn default_hf() -> bool { + false +} + +fn default_image_token_index() -> isize { + -200 +} + +fn default_mm_patch_merge_type() -> String { + "flat".to_string() +} + +fn default_image_aspect_ratio() -> String { + "square".to_string() +} + +impl LLaVAConfig { + pub fn to_llama_config(&self) -> Config { + Config { + hidden_size: self.hidden_size, + intermediate_size: self.intermediate_size, + vocab_size: self.vocab_size, + num_hidden_layers: self.num_hidden_layers, + num_attention_heads: self.num_attention_heads, + num_key_value_heads: self.num_key_value_heads, + rms_norm_eps: self.rms_norm_eps as f64, + rope_theta: self.rope_theta, + bos_token_id: Some(self.bos_token_id as u32), + eos_token_id: Some(LlamaEosToks::Single(self.eos_token_id as u32)), + use_flash_attn: false, + rope_scaling: None, // Assume we don't have LLaVA for Llama 3.1 + max_position_embeddings: self.max_position_embeddings, + tie_word_embeddings: self.tie_word_embeddings.unwrap_or(false), + } + } +} + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct HFLLaVATextConfig { + pub architectures: Vec, + #[serde(default = "default_hidden_size")] + pub hidden_size: usize, + #[serde(default = "default_intermediate_size")] + pub intermediate_size: usize, + #[serde(default = "default_max_length")] + pub max_length: usize, + pub max_position_embeddings: usize, + pub model_type: String, + #[serde(default = "default_num_attention_heads")] + pub num_attention_heads: usize, + #[serde(default = "default_num_hidden_layers")] + pub num_hidden_layers: usize, + #[serde(default = "default_num_key_value_heads")] + pub num_key_value_heads: usize, + pub pad_token_id: usize, + pub rms_norm_eps: f32, + #[serde(default = "default_rope_theta")] + pub rope_theta: f32, + pub torch_dtype: String, + #[serde(default = "default_use_cache")] + pub use_cache: bool, + pub vocab_size: usize, +} + +fn default_num_hidden_layers() -> usize { + 32 +} + +fn default_use_cache() -> bool { + true +} + +fn default_hidden_size() -> usize { + 4096 +} + +fn default_intermediate_size() -> usize { + 11008 +} + +fn default_max_length() -> usize { + 4096 +} + +fn default_num_attention_heads() -> usize { + 32 +} + +fn default_num_key_value_heads() -> usize { + 32 +} + +fn default_rope_theta() -> f32 { + 10000.0 +} + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct HFLLaVAVisionConfig { + pub hidden_size: usize, + pub image_size: usize, + pub intermediate_size: usize, + pub model_type: String, + pub num_attention_heads: usize, + pub num_hidden_layers: usize, + pub patch_size: usize, + pub projection_dim: usize, + pub vocab_size: usize, +} + +// config from llava-v1.6-vicuna-7b-hf +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct HFLLaVAConfig { + pub architectures: Vec, + pub ignore_index: isize, + pub image_grid_pinpoints: Vec<(u32, u32)>, + pub image_token_index: isize, + pub model_type: String, + pub projector_hidden_act: String, + pub text_config: HFLLaVATextConfig, + pub torch_dtype: String, + pub use_image_newline_parameter: bool, + pub vision_config: HFLLaVAVisionConfig, + pub vision_feature_layer: isize, + pub vision_feature_select_strategy: String, + pub vocab_size: usize, +} + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct HFGenerationConfig { + pub bos_token_id: usize, + pub eos_token_id: usize, + #[serde(default = "default_max_length")] + pub max_length: usize, + pub pad_token_id: usize, +} + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct HFPreProcessorConfig { + pub aspect_ratio_setting: String, + pub crop_size: HashMap, + pub do_center_crop: bool, + pub do_convert_rgb: bool, + pub do_normalize: bool, + pub do_rescale: bool, + pub do_resize: bool, + pub image_mean: Vec, + pub image_std: Vec, + pub resample: u32, + pub rescale_factor: f32, + pub size: HashMap, +} + +impl HFLLaVAConfig { + pub fn to_clip_vision_config(&self) -> ClipVisionConfig { + ClipVisionConfig { + embed_dim: self.vision_config.hidden_size, + activation: Activation::QuickGelu, + intermediate_size: self.vision_config.intermediate_size, + num_hidden_layers: self.vision_config.num_hidden_layers, + num_attention_heads: self.vision_config.num_attention_heads, + projection_dim: self.vision_config.projection_dim, + num_channels: 3, + image_size: self.vision_config.image_size, + patch_size: self.vision_config.patch_size, + } + } + fn map_projector_type(s: &str) -> String { + if s == "gelu" { + "mlp2x_gelu".to_string() + } else { + s.to_string() + } + } + + fn map_select_feature(s: &str) -> String { + if s == "default" { + "patch".to_string() + } else { + "cls_patch".to_string() + } + } + + pub fn to_llava_config( + &self, + generation_config: &HFGenerationConfig, + preprocessor_config: &HFPreProcessorConfig, + ) -> LLaVAConfig { + LLaVAConfig { + hf: true, + architectures: self.architectures.clone(), + bos_token_id: generation_config.bos_token_id, + eos_token_id: generation_config.eos_token_id, + hidden_size: self.text_config.hidden_size, + image_aspect_ratio: preprocessor_config.aspect_ratio_setting.clone(), + image_crop_resolution: 224, + image_grid_pinpoints: self.image_grid_pinpoints.clone(), + image_split_resolution: 224, + intermediate_size: self.text_config.intermediate_size, + max_position_embeddings: self.text_config.max_position_embeddings, + mm_hidden_size: 1024, + mm_patch_merge_type: "spatial_unpad".to_string(), + mm_projector_type: Self::map_projector_type(&self.projector_hidden_act), + mm_use_im_start_end: false, + mm_vision_select_feature: Self::map_select_feature( + &self.vision_feature_select_strategy, + ), + mm_vision_select_layer: self.vision_feature_layer, + mm_vision_tower: None, + model_type: self.model_type.clone(), + num_attention_heads: self.text_config.num_attention_heads, + num_hidden_layers: self.text_config.num_hidden_layers, + num_key_value_heads: self.text_config.num_key_value_heads, + pad_token_id: self.text_config.pad_token_id, + rms_norm_eps: self.text_config.rms_norm_eps, + rope_theta: self.text_config.rope_theta, + tokenizer_model_max_length: Some(4096), + torch_dtype: self.torch_dtype.clone(), + use_cache: self.text_config.use_cache, + vocab_size: self.vocab_size, + image_token_index: self.image_token_index, + tie_word_embeddings: None, + } + } +} diff --git a/patches/candle-transformers/src/models/llava/mod.rs b/patches/candle-transformers/src/models/llava/mod.rs new file mode 100644 index 0000000000..cc40ed357f --- /dev/null +++ b/patches/candle-transformers/src/models/llava/mod.rs @@ -0,0 +1,412 @@ +//! The LLaVA (Large Language and Vision Assistant) model. +//! +//! This provides the main model implementation combining a vision tower (CLIP) with +//! language model (Llama) for multimodal capabilities. The architecture implements the training-free projection technique. +//! +//! - 💻[GH Link](https://github.com/haotian-liu/LLaVA/tree/main) +//! - 📝 [Paper](https://arxiv.org/abs/2304.08485)/ Visual Instruction Tuning +//! + +pub mod config; +pub mod utils; + +use crate::models::clip::vision_model::{ClipVisionConfig, ClipVisionTransformer}; +use crate::models::llama::{Cache, Llama}; +use crate::models::with_tracing::linear; + +use candle::{bail, Context, Device, IndexOp, Result, Tensor}; +use candle_nn::{seq, Activation, Module, Sequential, VarBuilder}; +use fancy_regex::Regex; +use utils::get_anyres_image_grid_shape; + +use config::LLaVAConfig; + +fn mlp_gelu_match(mm_projector_type: &str) -> Option { + let mlp_gelu_regex = Regex::new(r"^mlp(\d+)x_gelu$").unwrap(); + + if let Ok(Some(captures)) = mlp_gelu_regex.captures(mm_projector_type) { + if let Some(match_str) = captures.get(1) { + let match_str = match_str.as_str(); + match_str.parse::().ok() + } else { + None + } + } else { + None + } +} + +fn unpad_image(tensor: &Tensor, original_size: &(u32, u32)) -> Result { + assert_eq!(tensor.dims().len(), 3); + let (original_width, original_height) = *original_size; + let tensor_dims = tensor.dims(); + let current_height = tensor_dims[1]; + let current_width = tensor_dims[2]; + let original_aspect_ratio = (original_width as f32) / (original_height as f32); + let current_aspect_ratio = (current_width as f32) / (current_height as f32); + if original_aspect_ratio > current_aspect_ratio { + let scale_factor = (current_width as f32) / (original_width as f32); + let new_height = (original_height as f32 * scale_factor).floor() as usize; + let padding = (current_height - new_height) / 2; + tensor.i((.., padding..current_width - padding, ..)) + } else { + let scale_factor = (current_height as f32) / (original_height as f32); + let new_width = (original_width as f32 * scale_factor).floor() as usize; + let padding = (current_width - new_width) / 2; + tensor.i((.., .., padding..current_width - padding)) + } +} + +pub struct IdentityMap {} + +impl Module for IdentityMap { + fn forward(&self, x: &Tensor) -> Result { + Ok(x.clone()) + } +} + +pub struct MMProjector { + pub modules: Sequential, +} + +impl MMProjector { + pub fn load(vb: &VarBuilder, config: &LLaVAConfig) -> Result { + if config.mm_projector_type == "linear" { + let vb_prefix = if config.hf { + "multi_modal_projector.linear_1" + } else { + "model.mm_projector.0" + }; + let linear = linear(config.mm_hidden_size, config.hidden_size, vb.pp(vb_prefix))?; + let modules = seq().add(linear); + Ok(Self { modules }) + } else if let Some(mlp_depth) = mlp_gelu_match(&config.mm_projector_type) { + let modules = if config.hf { + let mut modules = seq().add(linear( + config.mm_hidden_size, + config.hidden_size, + vb.pp("multi_modal_projector.linear_1"), + )?); + for i in 1..mlp_depth { + modules = modules.add(Activation::Gelu).add(linear( + config.hidden_size, + config.hidden_size, + vb.pp(format!("multi_modal_projector.linear_{}", i + 1)), + )?); + } + modules + } else { + let mut modules = seq().add(linear( + config.mm_hidden_size, + config.hidden_size, + vb.pp("model.mm_projector.0"), + )?); + for i in 1..mlp_depth { + modules = modules.add(Activation::Gelu).add(linear( + config.hidden_size, + config.hidden_size, + vb.pp(format!("model.mm_projector.{}", i * 2)), + )?); + } + modules + }; + Ok(Self { modules }) + } else if config.mm_projector_type == "identity" { + Ok(Self { + modules: seq().add(IdentityMap {}), + }) + } else { + bail!( + "Unsupported MM projector type: {}", + config.mm_projector_type + ) + } + } + + pub fn forward(&self, x: &Tensor) -> Result { + self.modules.forward(x) + } +} + +pub struct ClipVisionTower { + model: ClipVisionTransformer, + select_layer: isize, + select_feature_method: String, + pub config: ClipVisionConfig, +} + +impl ClipVisionTower { + pub fn new( + vb: VarBuilder, + select_layer: isize, + select_feature_method: &str, + config: &Option, + ) -> Result { + let config = if config.is_none() { + ClipVisionConfig::clip_vit_large_patch14_336() + } else { + config.clone().context("no config")? + }; + let select_layer = match select_layer { + -1 | -2 => select_layer, + _ => bail!("Unsupported select layer: {}", select_layer), + }; + let model = ClipVisionTransformer::new(vb, &config)?; + Ok(Self { + model, + select_layer, + select_feature_method: select_feature_method.to_string(), + config, + }) + } + + pub fn forward(&self, x: &Tensor) -> Result { + let result = self.model.output_hidden_states(x)?; + let index = result.len() as isize + self.select_layer; + let result = result[index as usize].clone(); + if self.select_feature_method == "cls_patch" { + Ok(result) + } else { + result.i((.., 1..)) + } + } + + pub fn num_patches_per_side(&self) -> usize { + self.config.image_size / self.config.patch_size + } +} + +pub struct LLaVA { + pub clip_vision_tower: ClipVisionTower, + pub image_newline: Tensor, + pub mm_projector: MMProjector, + pub llama: Llama, + config: LLaVAConfig, + device: Device, +} + +impl LLaVA { + pub fn load( + vb: VarBuilder, + config: &LLaVAConfig, + clip_vision_config: Option, + ) -> Result { + let device = vb.device().clone(); + let llama_config = config.to_llama_config(); + let mm_projector = MMProjector::load(&vb, config)?; + let (clip_vision_tower, image_newline, llama) = if config.hf { + ( + ClipVisionTower::new( + vb.pp("vision_tower.vision_model"), + config.mm_vision_select_layer, + &config.mm_vision_select_feature, + &clip_vision_config, + )?, + vb.get(&[config.hidden_size], "image_newline")? + .to_device(&device)?, + Llama::load(vb.pp("language_model"), &llama_config)?, + ) + } else { + ( + ClipVisionTower::new( + vb.pp("model.vision_tower.vision_tower.vision_model"), + config.mm_vision_select_layer, + &config.mm_vision_select_feature, + &clip_vision_config, + )?, + vb.get(&[config.hidden_size], "model.image_newline")? + .to_device(&device)?, + Llama::load(vb, &llama_config)?, + ) + }; + Ok(Self { + clip_vision_tower, + image_newline, + mm_projector, + llama, + config: (*config).clone(), + device, + }) + } + + pub fn encode_images(&self, x: &Tensor) -> Result { + let image_features = self.clip_vision_tower.forward(x)?; + let image_features = self.mm_projector.forward(&image_features)?; + Ok(image_features) + } + // currently only for single image, 4 dim tensor + pub fn prepare_inputs_labels_for_multimodal( + &self, + input_ids: &Tensor, + images: &[Tensor], + image_sizes: &[(u32, u32)], + ) -> Result { + //TODO: process of multiple images/ new line + // 576: 336(input size)/14(patch size)=24 24*24+1(class)=577 577-1=576 + let concat_images = Tensor::cat(images, 0)?; + let image_features_together = self.encode_images(&concat_images)?; + let split_sizes = images + .iter() + .map(|x| x.shape().dims()[0]) + .collect::>(); + // can be replaced by split + let mut index_pos = 0; + let mut image_features = Vec::new(); + for split_size in split_sizes.iter() { + image_features.push(image_features_together.i(index_pos..index_pos + (*split_size))?); + index_pos += *split_size; + } + let mm_patch_merge_type = &self.config.mm_patch_merge_type; + let image_aspect_ratio = &self.config.image_aspect_ratio; + + let image_features = if mm_patch_merge_type == "flat" { + image_features + .iter() + .map(|x| x.flatten(0, 1)) + .collect::>>()? + } else if mm_patch_merge_type.starts_with("spatial") { + let mut new_image_features = Vec::new(); + for (image_idx, image_feature) in image_features.iter().enumerate() { + let new_image_feature = if image_feature.dims()[0] > 1 { + let base_image_feature = image_feature.get(0)?; + let patch_image_feature = image_feature.i(1..)?; + let height = self.clip_vision_tower.num_patches_per_side(); + let width = height; + assert_eq!(height * width, base_image_feature.dims()[0]); + let image_size = image_sizes[image_idx]; + let new_image_feature = if image_aspect_ratio == "anyres" { + let (num_patch_width, num_patch_height) = get_anyres_image_grid_shape( + image_size, + &self.config.image_grid_pinpoints, + self.clip_vision_tower.config.image_size as u32, + ); + patch_image_feature.reshape(( + num_patch_height as usize, + num_patch_width as usize, + height, + width, + (), + ))? + } else { + bail!("not implemented in original python LLaVA yet") + }; + let new_image_feature = if mm_patch_merge_type.contains("unpad") { + let new_image_feature = new_image_feature + .permute((4, 0, 2, 1, 3))? + .flatten(1, 2)? + .flatten(2, 3)?; + let new_image_feature = unpad_image(&new_image_feature, &image_size)?; + let new_image_feature_dims = new_image_feature.dims(); + let image_new_line = self + .image_newline + .reshape((self.config.hidden_size, 1, 1))? + .broadcast_as(( + new_image_feature_dims[0], + new_image_feature_dims[1], + 1, + ))?; + let new_image_feature = + Tensor::cat(&[new_image_feature, image_new_line], 2)?; + new_image_feature.flatten(1, 2)?.transpose(0, 1)? + } else { + new_image_feature.permute((0, 2, 1, 3, 4))?.flatten(0, 3)? + }; + Tensor::cat(&[base_image_feature, new_image_feature], 0)? + } else { + let new_image_feature = image_feature.get(0)?; + if mm_patch_merge_type.contains("unpad") { + Tensor::cat( + &[new_image_feature, self.image_newline.clone().unsqueeze(0)?], + 0, + )? + } else { + new_image_feature + } + }; + new_image_features.push(new_image_feature); + } + new_image_features + } else { + bail!("Unexpected mm_patch_merge_type: {mm_patch_merge_type}") + }; + // can easily be replaced by nonzero if it is implemented in candle + let input_ids_vec = input_ids.squeeze(0)?.to_vec1::()?; + let mut image_indices = { + let mut image_indices = vec![0_i64]; + image_indices.extend( + input_ids_vec + .iter() + .enumerate() + .filter_map(|(i, x)| { + if *x == self.config.image_token_index as i64 { + Some(i as i64) + } else { + None + } + }) + .collect::>(), + ); + image_indices + }; + if image_indices.len() == 1 { + //no image, only [0], + return self.llama.embed(input_ids); + } + + let input_ids_noim = input_ids_vec + .iter() + .filter_map(|x| { + if *x != self.config.image_token_index as i64 { + Some(*x) + } else { + None + } + }) + .collect::>(); + let input_ids_noim_len = input_ids_noim.len(); + image_indices.push((input_ids_noim_len) as i64); + let input_ids_noim = Tensor::from_vec(input_ids_noim, input_ids_noim_len, &self.device)?; + let cur_input_embeds = self.llama.embed(&input_ids_noim)?; + // can be replace by split if it is implemented in candle + let input_embed_no_ims = { + let mut input_embeds = Vec::new(); + for i in 0..image_indices.len() - 1 { + let start = (image_indices[i]) as usize; + let end = image_indices[i + 1] as usize; + input_embeds.push(cur_input_embeds.i((start..end, ..))?) + } + input_embeds + }; + + let mut cur_new_input_embeds = Vec::new(); + for (i, image_feature) in image_features.iter().enumerate() { + cur_new_input_embeds.push(input_embed_no_ims[i].clone()); + cur_new_input_embeds.push(image_feature.clone()); + } + cur_new_input_embeds.push(input_embed_no_ims[image_features.len()].clone()); + let new_input_embeds = Tensor::cat(&cur_new_input_embeds, 0)?; + //truncate + let new_input_embeds = + if let Some(tokenizer_model_max_length) = self.config.tokenizer_model_max_length { + let (new_input_embeds_length, _) = new_input_embeds.shape().dims2()?; + if new_input_embeds_length > tokenizer_model_max_length { + new_input_embeds.i((..tokenizer_model_max_length, ..))? + } else { + new_input_embeds + } + } else { + new_input_embeds + }; + new_input_embeds.unsqueeze(0) + } + + pub fn forward( + &self, + input_embeds: &Tensor, + position_id: usize, + cache: &mut Cache, + ) -> Result { + self.llama + .forward_input_embed(input_embeds, position_id, cache) + } +} diff --git a/patches/candle-transformers/src/models/llava/utils.rs b/patches/candle-transformers/src/models/llava/utils.rs new file mode 100644 index 0000000000..3b4c18bb60 --- /dev/null +++ b/patches/candle-transformers/src/models/llava/utils.rs @@ -0,0 +1,41 @@ +pub fn get_anyres_image_grid_shape( + image_size: (u32, u32), + grid_pinpoints: &[(u32, u32)], + patch_size: u32, +) -> (u32, u32) { + let (width, height) = select_best_resolution(image_size, grid_pinpoints); + (width / patch_size, height / patch_size) +} + +pub fn select_best_resolution( + original_size: (u32, u32), + possible_resolutions: &[(u32, u32)], +) -> (u32, u32) { + let (original_width, original_height) = original_size; + let mut best_fit = (0, 0); + let original_width_f = original_width as f32; + let original_height_f = original_height as f32; + let mut max_effective_resolution = 0_u32; + let mut min_wasted_resolution = u32::MAX; + for (width, height) in possible_resolutions { + let width_f = *width as f32; + let height_f = *height as f32; + let scale = (width_f / original_width_f).min(height_f / original_height_f); + let (downscaled_width, downscaled_height) = ( + (original_width_f * scale) as u32, + (original_height_f * scale) as u32, + ); + let effective_resolution = + std::cmp::min((*width) * (*height), downscaled_width * downscaled_height); + let wasted_resolution = (*width) * (*height) - effective_resolution; + if effective_resolution > max_effective_resolution + || (effective_resolution == max_effective_resolution + && wasted_resolution < min_wasted_resolution) + { + best_fit = (*width, *height); + max_effective_resolution = effective_resolution; + min_wasted_resolution = wasted_resolution; + } + } + best_fit +} diff --git a/patches/candle-transformers/src/models/mamba.rs b/patches/candle-transformers/src/models/mamba.rs new file mode 100644 index 0000000000..dfae0af398 --- /dev/null +++ b/patches/candle-transformers/src/models/mamba.rs @@ -0,0 +1,221 @@ +//! Mamba inference implementation. +//! +//! See ["Mamba: Linear-Time Sequence Modeling with Selective State Spaces"](https://arxiv.org/abs/2312.00752) +//! +//! Based on reference implementation from the AlbertMamba project +//! A fast implementation of mamba for inference only. +//! Based on Laurent Mazare's rust implementation: [mamba.rs](https://github.com/LaurentMazare/mamba.rs) +use crate::models::with_tracing::{linear, linear_no_bias, Linear}; +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{RmsNorm, VarBuilder}; + +const D_CONV: usize = 4; +const D_STATE: usize = 16; + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + pub d_model: usize, + pub n_layer: usize, + pub vocab_size: usize, + pub pad_vocab_size_multiple: usize, +} + +impl Config { + fn vocab_size(&self) -> usize { + let pad = self.pad_vocab_size_multiple; + self.vocab_size.div_ceil(pad) * pad + } + + fn dt_rank(&self) -> usize { + self.d_model.div_ceil(16) + } + + fn d_inner(&self) -> usize { + self.d_model * 2 + } +} + +pub struct State { + pub hs: Vec, + pub prev_xs: Vec<[Tensor; D_CONV]>, + pub pos: usize, +} + +impl State { + pub fn new(batch_size: usize, cfg: &Config, dtype: DType, device: &Device) -> Result { + let mut hs = Vec::with_capacity(cfg.n_layer); + let mut prev_xs = Vec::with_capacity(cfg.n_layer); + for _i in 0..cfg.n_layer { + let h = Tensor::zeros((batch_size, cfg.d_inner(), D_STATE), dtype, device)?; + let x = Tensor::zeros((batch_size, cfg.d_inner()), dtype, device)?; + hs.push(h); + prev_xs.push([x.clone(), x.clone(), x.clone(), x.clone()]); + } + Ok(Self { + hs, + prev_xs, + pos: 0, + }) + } +} + +#[derive(Clone, Debug)] +pub struct MambaBlock { + in_proj: Linear, + conv1d_bias: Tensor, + conv1d_weights: [Tensor; D_CONV], + x_proj: Linear, + dt_proj: Linear, + a_log: Tensor, + d: Tensor, + out_proj: Linear, + dt_rank: usize, + layer_index: usize, + d_inner: usize, +} + +impl MambaBlock { + pub fn new(layer_index: usize, cfg: &Config, vb: VarBuilder) -> Result { + let d_inner = cfg.d_inner(); + let dt_rank = cfg.dt_rank(); + let in_proj = linear_no_bias(cfg.d_model, d_inner * 2, vb.pp("in_proj"))?; + let x_proj = linear_no_bias(d_inner, dt_rank + D_STATE * 2, vb.pp("x_proj"))?; + let dt_proj = linear(dt_rank, d_inner, vb.pp("dt_proj"))?; + let a_log = vb.get((d_inner, D_STATE), "A_log")?; + let d = vb.get(d_inner, "D")?; + let out_proj = linear_no_bias(d_inner, cfg.d_model, vb.pp("out_proj"))?; + let conv1d_bias = vb.get(d_inner, "conv1d.bias")?; + let conv1d_weight = vb.get((d_inner, 1, D_CONV), "conv1d.weight")?; + let conv1d_weights = [ + conv1d_weight.i((.., 0, 0))?, + conv1d_weight.i((.., 0, 1))?, + conv1d_weight.i((.., 0, 2))?, + conv1d_weight.i((.., 0, 3))?, + ]; + Ok(Self { + in_proj, + conv1d_bias, + conv1d_weights, + x_proj, + dt_proj, + a_log, + d, + out_proj, + dt_rank, + layer_index, + d_inner, + }) + } + + pub fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let (b_sz, _dim) = xs.dims2()?; + let li = self.layer_index; + let mut xs = xs.apply(&self.in_proj)?.chunk(2, D::Minus1)?; + let proj_for_silu = xs.remove(1); + state.prev_xs[li][state.pos % D_CONV] = xs.remove(0); + let mut proj_for_conv = self.conv1d_bias.broadcast_as((b_sz, self.d_inner))?; + for d_c in 0..D_CONV { + proj_for_conv = (proj_for_conv + + self.conv1d_weights[d_c] + .broadcast_mul(&state.prev_xs[li][(d_c + 1 + state.pos) % D_CONV])?)?; + } + let proj_for_conv = candle_nn::ops::silu(&proj_for_conv)?; + // SSM + Selection, we're doing inference here so only need the last step of + // the sequence. + // Algorithm 3.2 on page 6, https://arxiv.org/pdf/2312.00752.pdf + + let x_proj = self.x_proj.forward(&proj_for_conv)?; + let delta = x_proj.narrow(D::Minus1, 0, self.dt_rank)?.contiguous()?; + let b = x_proj.narrow(D::Minus1, self.dt_rank, D_STATE)?; + let c = x_proj.narrow(D::Minus1, self.dt_rank + D_STATE, D_STATE)?; + + let delta = delta.apply(&self.dt_proj)?; + // softplus + let delta = (delta.exp()? + 1.)?.log()?; + let a = self.a_log.to_dtype(delta.dtype())?.exp()?.neg()?; + let d = self.d.to_dtype(delta.dtype())?; + + // Selective scan part + // Eqn (2a), page 3, h_t = Ab h_{t-1} + Bb x_t + let delta = delta + .unsqueeze(D::Minus1)? + .broadcast_as((b_sz, self.d_inner, D_STATE))?; + let a = a.broadcast_as((b_sz, self.d_inner, D_STATE))?; + let b = b.broadcast_as((b_sz, self.d_inner, D_STATE))?; + let proj_for_conv_b = + proj_for_conv + .unsqueeze(D::Minus1)? + .broadcast_as((b_sz, self.d_inner, D_STATE))?; + state.hs[li] = ((&state.hs[li] * (&delta * &a)?.exp()?)? + &delta * &b * &proj_for_conv_b)?; + let ss = (state.hs[li] + .matmul(&c.unsqueeze(D::Minus1)?)? + .squeeze(D::Minus1)? + + proj_for_conv.broadcast_mul(&d)?)?; + + let ys = (ss * candle_nn::ops::silu(&proj_for_silu))?; + ys.apply(&self.out_proj) + } +} + +#[derive(Clone, Debug)] +pub struct ResidualBlock { + mixer: MambaBlock, + norm: RmsNorm, +} + +impl ResidualBlock { + pub fn new(layer_index: usize, cfg: &Config, vb: VarBuilder) -> Result { + let norm = candle_nn::rms_norm(cfg.d_model, 1e-5, vb.pp("norm"))?; + let mixer = MambaBlock::new(layer_index, cfg, vb.pp("mixer"))?; + Ok(Self { mixer, norm }) + } + + fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + self.mixer.forward(&xs.apply(&self.norm)?, state)? + xs + } +} + +// https://github.com/johnma2006/mamba-minimal/blob/61f01953ca153f8c4a850d7111beecbf4be9cee1/model.py#L56 +#[derive(Clone, Debug)] +pub struct Model { + embedding: candle_nn::Embedding, + layers: Vec, + norm_f: RmsNorm, + lm_head: Linear, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let embedding = candle_nn::embedding(cfg.vocab_size(), cfg.d_model, vb.pp("embedding"))?; + let mut layers = Vec::with_capacity(cfg.n_layer); + let vb_l = vb.pp("layers"); + for layer_idx in 0..cfg.n_layer { + let layer = ResidualBlock::new(layer_idx, cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm_f = candle_nn::rms_norm(cfg.d_model, 1e-5, vb.pp("norm_f"))?; + let lm_head = Linear::from_weights(embedding.embeddings().clone(), None); + Ok(Self { + embedding, + layers, + norm_f, + lm_head, + dtype: vb.dtype(), + }) + } + + pub fn forward(&self, input_ids: &Tensor, state: &mut State) -> Result { + let _b_size = input_ids.dims1()?; + let mut xs = self.embedding.forward(input_ids)?; + for layer in self.layers.iter() { + xs = layer.forward(&xs, state)? + } + state.pos += 1; + xs.apply(&self.norm_f)?.apply(&self.lm_head) + } + + pub fn dtype(&self) -> DType { + self.dtype + } +} diff --git a/patches/candle-transformers/src/models/mamba2.rs b/patches/candle-transformers/src/models/mamba2.rs new file mode 100644 index 0000000000..834510769b --- /dev/null +++ b/patches/candle-transformers/src/models/mamba2.rs @@ -0,0 +1,647 @@ +//! Mamba2 inference implementation. +//! +//! See ["Transformers are SSMs: Generalized Models and Efficient Algorithms +//! Through Structured State Space Duality"](https://arxiv.org/abs/2405.21060) + +use crate::models::with_tracing::{linear_no_bias, Linear}; +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{RmsNorm, VarBuilder}; + +const D_CONV: usize = 4; + +/// Segment sum for SSD: computes cumsum[i] - cumsum[j] with lower triangular mask. +/// See Algorithm 1 in the Mamba2 paper. +fn segsum(x: &Tensor) -> Result { + let device = x.device(); + let dtype = x.dtype(); + let t = x.dim(D::Minus1)?; + + let x_cumsum = x.cumsum(D::Minus1)?; + + let target_shape: Vec = { + let mut shape = x.dims().to_vec(); + shape.push(t); + shape + }; + + let x_cumsum_row = x_cumsum + .unsqueeze(D::Minus1)? + .broadcast_as(target_shape.as_slice())?; + let x_cumsum_col = x_cumsum + .unsqueeze(x.rank() - 1)? + .broadcast_as(target_shape.as_slice())?; + let x_segsum = (&x_cumsum_row - &x_cumsum_col)?; + + let mask_lower = Tensor::tril2(t, DType::U8, device)?; + let neg_inf = Tensor::new(f32::NEG_INFINITY, device)? + .to_dtype(dtype)? + .broadcast_as(x_segsum.shape())?; + + mask_lower + .broadcast_as(x_segsum.shape())? + .where_cond(&x_segsum, &neg_inf) +} + +fn pad_to_chunk_size(x: &Tensor, chunk_size: usize) -> Result<(Tensor, usize)> { + let seq_len = x.dim(1)?; + let pad_len = (chunk_size - (seq_len % chunk_size)) % chunk_size; + if pad_len == 0 { + return Ok((x.clone(), 0)); + } + + let mut pad_shape = x.dims().to_vec(); + pad_shape[1] = pad_len; + let padding = Tensor::zeros(pad_shape, x.dtype(), x.device())?; + Ok((Tensor::cat(&[x, &padding], 1)?, pad_len)) +} + +fn reshape_into_chunks(x: &Tensor, chunk_size: usize) -> Result { + let dims = x.dims(); + let b = dims[0]; + let l = dims[1]; + let n_chunks = l / chunk_size; + + let mut new_shape = vec![b, n_chunks, chunk_size]; + new_shape.extend_from_slice(&dims[2..]); + x.reshape(new_shape) +} + +fn reshape_from_chunks(x: &Tensor) -> Result { + let dims = x.dims(); + let b = dims[0]; + let n_chunks = dims[1]; + let chunk_size = dims[2]; + + let mut new_shape = vec![b, n_chunks * chunk_size]; + new_shape.extend_from_slice(&dims[3..]); + x.reshape(new_shape) +} + +fn default_d_state() -> usize { + 64 +} +fn default_expand() -> usize { + 2 +} +fn default_headdim() -> usize { + 64 +} +fn default_ngroups() -> usize { + 1 +} +fn default_pad_vocab_size_multiple() -> usize { + 16 +} + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + #[serde(alias = "hidden_size")] + pub d_model: usize, + #[serde(alias = "num_hidden_layers")] + pub n_layer: usize, + pub vocab_size: usize, + #[serde(alias = "state_size", default = "default_d_state")] + pub d_state: usize, + #[serde(default = "default_expand")] + pub expand: usize, + #[serde(alias = "head_dim", default = "default_headdim")] + pub headdim: usize, + #[serde(alias = "n_groups", default = "default_ngroups")] + pub ngroups: usize, + #[serde(default = "default_pad_vocab_size_multiple")] + pub pad_vocab_size_multiple: usize, +} + +impl Config { + fn vocab_size(&self) -> usize { + let pad = self.pad_vocab_size_multiple; + self.vocab_size.div_ceil(pad) * pad + } + + fn d_inner(&self) -> usize { + self.d_model * self.expand + } + + fn d_xbc(&self) -> usize { + self.d_inner() + 2 * self.ngroups * self.d_state + } + + fn nheads(&self) -> usize { + self.d_inner() / self.headdim + } +} + +pub struct State { + pub hs: Vec, + pub conv_states: Vec, + pub pos: usize, +} + +impl State { + pub fn new(batch_size: usize, cfg: &Config, dtype: DType, device: &Device) -> Result { + let d_xbc = cfg.d_xbc(); + let nheads = cfg.nheads(); + let mut hs = Vec::with_capacity(cfg.n_layer); + let mut conv_states = Vec::with_capacity(cfg.n_layer); + for _ in 0..cfg.n_layer { + let h = Tensor::zeros( + (batch_size, nheads, cfg.headdim, cfg.d_state), + dtype, + device, + )?; + let conv = Tensor::zeros((batch_size, d_xbc, D_CONV), dtype, device)?; + hs.push(h); + conv_states.push(conv); + } + Ok(Self { + hs, + conv_states, + pos: 0, + }) + } +} + +#[derive(Clone, Debug)] +pub struct Mamba2Block { + in_proj: Linear, + conv1d_weight: Tensor, + conv1d_bias: Tensor, + a_log: Tensor, + d: Tensor, + dt_bias: Tensor, + out_proj: Linear, + norm: RmsNorm, + d_inner: usize, + d_state: usize, + d_xbc: usize, + headdim: usize, + nheads: usize, + ngroups: usize, + layer_idx: usize, +} + +impl Mamba2Block { + pub fn new(layer_idx: usize, cfg: &Config, vb: VarBuilder) -> Result { + let d_inner = cfg.d_inner(); + let nheads = cfg.nheads(); + let ngroups = cfg.ngroups; + let d_state = cfg.d_state; + let d_xbc = cfg.d_xbc(); + + let proj_size = d_inner + d_xbc + nheads; + let in_proj = linear_no_bias(cfg.d_model, proj_size, vb.pp("in_proj"))?; + + let conv1d_weight = vb.get((d_xbc, 1, D_CONV), "conv1d.weight")?; + let conv1d_bias = vb.get(d_xbc, "conv1d.bias")?; + + let a_log = vb.get(nheads, "A_log")?; + let d = vb.get(nheads, "D")?; + let dt_bias = vb.get(nheads, "dt_bias")?; + + let out_proj = linear_no_bias(d_inner, cfg.d_model, vb.pp("out_proj"))?; + let norm = candle_nn::rms_norm(d_inner, 1e-5, vb.pp("norm"))?; + + Ok(Self { + in_proj, + conv1d_weight, + conv1d_bias, + a_log, + d, + dt_bias, + out_proj, + norm, + d_inner, + d_state, + d_xbc, + headdim: cfg.headdim, + nheads, + ngroups, + layer_idx, + }) + } + + pub fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let (b_sz, _dim) = xs.dims2()?; + + let proj = self.in_proj.forward(xs)?; + + let z = proj.narrow(D::Minus1, 0, self.d_inner)?; + let xbc = proj.narrow(D::Minus1, self.d_inner, self.d_xbc)?; + let dt = proj.narrow(D::Minus1, self.d_inner + self.d_xbc, self.nheads)?; + + let xbc_conv = self.apply_conv1d(&xbc, &mut state.conv_states[self.layer_idx])?; + let xbc_conv = candle_nn::ops::silu(&xbc_conv)?; + + let x_conv = xbc_conv.narrow(D::Minus1, 0, self.d_inner)?; + let b = xbc_conv.narrow(D::Minus1, self.d_inner, self.ngroups * self.d_state)?; + let c = xbc_conv.narrow( + D::Minus1, + self.d_inner + self.ngroups * self.d_state, + self.ngroups * self.d_state, + )?; + + let dt_bias = self.dt_bias.broadcast_as(dt.shape())?; + let dt = ((&dt + &dt_bias)?.exp()? + 1.)?.log()?; // softplus + + let a = self.a_log.exp()?.neg()?; + + let y = self.ssm_step(&x_conv, &a, &b, &c, &dt, state)?; + + let d = self.d.broadcast_as((b_sz, self.nheads))?; + let x_skip = x_conv.reshape((b_sz, self.nheads, self.headdim))?; + let y = (&y + x_skip.broadcast_mul(&d.unsqueeze(D::Minus1)?)?)?; + let y = y.reshape((b_sz, self.d_inner))?; + + // Mamba2 applies gate before norm (MambaRMSNormGated) + let y = (y * candle_nn::ops::silu(&z)?)?; + let y = self.norm.forward(&y)?; + + self.out_proj.forward(&y) + } + + fn apply_conv1d(&self, xbc: &Tensor, conv_state: &mut Tensor) -> Result { + let (b_sz, d_xbc) = xbc.dims2()?; + + let shifted = conv_state.narrow(D::Minus1, 1, D_CONV - 1)?; + let xbc_expanded = xbc.unsqueeze(D::Minus1)?; + *conv_state = Tensor::cat(&[shifted, xbc_expanded], D::Minus1)?; + + let mut result = self.conv1d_bias.broadcast_as((b_sz, d_xbc))?; + for i in 0..D_CONV { + let w = self.conv1d_weight.i((.., 0, i))?; + let xbc_i = conv_state.i((.., .., i))?; + result = (result + w.broadcast_mul(&xbc_i)?)?; + } + Ok(result) + } + + fn ssm_step( + &self, + x: &Tensor, + a: &Tensor, + b: &Tensor, + c: &Tensor, + dt: &Tensor, + state: &mut State, + ) -> Result { + let (b_sz, _) = x.dims2()?; + let h = &mut state.hs[self.layer_idx]; + + let x = x.reshape((b_sz, self.nheads, self.headdim))?; + + let b = b.reshape((b_sz, self.ngroups, self.d_state))?; + let c = c.reshape((b_sz, self.ngroups, self.d_state))?; + let heads_per_group = self.nheads / self.ngroups; + let b = + b.unsqueeze(2)? + .broadcast_as((b_sz, self.ngroups, heads_per_group, self.d_state))?; + let b = b.reshape((b_sz, self.nheads, self.d_state))?; + let c = + c.unsqueeze(2)? + .broadcast_as((b_sz, self.ngroups, heads_per_group, self.d_state))?; + let c = c.reshape((b_sz, self.nheads, self.d_state))?; + + let dt_a = dt.broadcast_mul(a)?; + let decay = dt_a.exp()?; + let decay = decay.unsqueeze(D::Minus1)?.unsqueeze(D::Minus1)?; + let decay = decay.broadcast_as((b_sz, self.nheads, self.headdim, self.d_state))?; + + let x_unsq = x.unsqueeze(D::Minus1)?; + let b_unsq = b.unsqueeze(2)?; + let x_b = x_unsq.broadcast_mul(&b_unsq)?; + + let dt_expanded = dt.unsqueeze(D::Minus1)?.unsqueeze(D::Minus1)?; + let dt_expanded = + dt_expanded.broadcast_as((b_sz, self.nheads, self.headdim, self.d_state))?; + + // SSM recurrence: h = exp(A*dt) * h + dt * (x ⊗ B) + *h = ((&*h * &decay)? + (&dt_expanded * &x_b)?)?; + + let c_unsq = c.unsqueeze(2)?; + let c_broadcast = c_unsq.broadcast_as(h.shape())?; + let y = (&*h * &c_broadcast)?.sum(D::Minus1)?; + + Ok(y) + } + + /// Chunked SSD algorithm for parallel prefill (Algorithm 1 in Mamba2 paper). + fn ssd_chunked( + &self, + x: &Tensor, + a: &Tensor, + b: &Tensor, + c: &Tensor, + chunk_size: usize, + initial_state: Option<&Tensor>, + ) -> Result<(Tensor, Tensor)> { + let device = x.device(); + let dtype = x.dtype(); + let (batch, seq_len, nheads, headdim) = x.dims4()?; + let d_state = self.d_state; + let n_chunks = seq_len / chunk_size; + + let x = reshape_into_chunks(x, chunk_size)?; + let a = reshape_into_chunks(a, chunk_size)?; + let b = reshape_into_chunks(b, chunk_size)?; + let c = reshape_into_chunks(c, chunk_size)?; + + // contiguous() required for Metal: cumsum uses matmul internally + let a = a.permute((0, 3, 1, 2))?.contiguous()?; + let a_cumsum = a.cumsum(D::Minus1)?; + + // Intra-chunk (diagonal blocks) + let l = segsum(&a)?.exp()?; + + let c_expanded = c.unsqueeze(3)?; + let b_expanded = b.unsqueeze(2)?; + let cb_shape = (batch, n_chunks, chunk_size, chunk_size, nheads, d_state); + let cb = (c_expanded.broadcast_as(cb_shape)? * b_expanded.broadcast_as(cb_shape)?)? + .sum(D::Minus1)?; + let cb = cb.permute((0, 1, 4, 2, 3))?; + + let l_t = l.permute((0, 2, 1, 3, 4))?; + let cb_l = (&cb * &l_t)?; + + let x_t = x.permute((0, 1, 3, 2, 4))?; + let y_diag_shape = (batch, n_chunks, nheads, chunk_size, chunk_size, headdim); + let y_diag = (cb_l.unsqueeze(D::Minus1)?.broadcast_as(y_diag_shape)? + * x_t.unsqueeze(3)?.broadcast_as(y_diag_shape)?)? + .sum(4)? + .permute((0, 1, 3, 2, 4))?; + + // Intra-chunk states + let a_last = a_cumsum.narrow(D::Minus1, chunk_size - 1, 1)?; + let decay_states = (a_last.broadcast_as(a_cumsum.shape())? - &a_cumsum)?.exp()?; + + let decay_s = decay_states.permute((0, 2, 1, 3))?.unsqueeze(D::Minus1)?; + let b_t = b.permute((0, 1, 3, 2, 4))?; + let b_weighted = b_t.broadcast_mul(&decay_s)?; + + let x_t2 = x.permute((0, 1, 3, 2, 4))?; + let states_shape = (batch, n_chunks, nheads, chunk_size, headdim, d_state); + let states = (x_t2.unsqueeze(D::Minus1)?.broadcast_as(states_shape)? + * b_weighted.unsqueeze(4)?.broadcast_as(states_shape)?)? + .sum(3)?; + + // Inter-chunk recurrence + let init_state = match initial_state { + Some(s) => s.unsqueeze(1)?, + None => Tensor::zeros((batch, 1, nheads, headdim, d_state), dtype, device)?, + }; + let states_with_init = Tensor::cat(&[&init_state, &states], 1)?; + + let a_chunk = a_cumsum + .narrow(D::Minus1, chunk_size - 1, 1)? + .squeeze(D::Minus1)?; + let zeros = Tensor::zeros((batch, nheads, 1), dtype, device)?; + let a_chunk_padded = Tensor::cat(&[&zeros, &a_chunk], D::Minus1)?; + let decay_chunk = segsum(&a_chunk_padded)?.exp()?; + + let states_p = states_with_init.permute((0, 2, 1, 3, 4))?; + let inter_shape = (batch, nheads, n_chunks + 1, n_chunks + 1, headdim, d_state); + let new_states = (decay_chunk + .unsqueeze(D::Minus1)? + .unsqueeze(D::Minus1)? + .broadcast_as(inter_shape)? + * states_p.unsqueeze(2)?.broadcast_as(inter_shape)?)? + .sum(3)? + .permute((0, 2, 1, 3, 4))?; + + let states_out = new_states.narrow(1, 0, n_chunks)?; + let final_state = new_states.narrow(1, n_chunks, 1)?.squeeze(1)?; + + // State-to-output (off-diagonal blocks) + let state_decay_out = a_cumsum.exp()?; + + let c_t2 = c.permute((0, 1, 3, 2, 4))?; + let off_shape = (batch, n_chunks, nheads, chunk_size, headdim, d_state); + let c_states = (c_t2.unsqueeze(4)?.broadcast_as(off_shape)? + * states_out.unsqueeze(3)?.broadcast_as(off_shape)?)? + .sum(D::Minus1)?; + + let decay_out = state_decay_out + .permute((0, 2, 1, 3))? + .unsqueeze(D::Minus1)?; + let y_off = c_states + .broadcast_mul(&decay_out)? + .permute((0, 1, 3, 2, 4))?; + + let y = (&y_diag + &y_off)?; + let y = reshape_from_chunks(&y)?; + + Ok((y, final_state)) + } + + pub fn forward_prefill( + &self, + xs: &Tensor, + state: &mut State, + chunk_size: usize, + ) -> Result { + let (b_sz, seq_len, _) = xs.dims3()?; + + let (xs, pad_len) = pad_to_chunk_size(xs, chunk_size)?; + let padded_len = xs.dim(1)?; + + let proj = xs.apply(&self.in_proj)?; + + let z = proj.narrow(D::Minus1, 0, self.d_inner)?; + let xbc = proj.narrow(D::Minus1, self.d_inner, self.d_xbc)?; + let dt = proj.narrow(D::Minus1, self.d_inner + self.d_xbc, self.nheads)?; + + let xbc_t = xbc.transpose(1, 2)?; + let pad = Tensor::zeros((b_sz, self.d_xbc, D_CONV - 1), xbc.dtype(), xbc.device())?; + let xbc_padded = Tensor::cat(&[&pad, &xbc_t], D::Minus1)?; + let xbc_conv = xbc_padded.conv1d(&self.conv1d_weight, 0, 1, 1, self.d_xbc)?; + let xbc_conv = xbc_conv + .broadcast_add(&self.conv1d_bias.reshape((1, self.d_xbc, 1))?)? + .transpose(1, 2)?; + let xbc_conv = candle_nn::ops::silu(&xbc_conv)?; + + // Update conv_state from real sequence tokens (not padding) for correct autoregressive behavior + let start = seq_len.saturating_sub(D_CONV); + let count = D_CONV.min(seq_len); + let last_tokens = xbc.narrow(1, start, count)?; + let last_tokens = last_tokens.transpose(1, 2)?; + if count >= D_CONV { + state.conv_states[self.layer_idx] = last_tokens.contiguous()?; + } else { + let existing = + state.conv_states[self.layer_idx].narrow(D::Minus1, count, D_CONV - count)?; + state.conv_states[self.layer_idx] = Tensor::cat(&[&existing, &last_tokens], D::Minus1)?; + } + + let x_conv = xbc_conv.narrow(D::Minus1, 0, self.d_inner)?; + let b = xbc_conv.narrow(D::Minus1, self.d_inner, self.ngroups * self.d_state)?; + let c = xbc_conv.narrow( + D::Minus1, + self.d_inner + self.ngroups * self.d_state, + self.ngroups * self.d_state, + )?; + + let dt_bias = self.dt_bias.broadcast_as(dt.shape())?; + let dt = ((&dt + &dt_bias)?.exp()? + 1.)?.log()?; + + let a = self.a_log.exp()?.neg()?; + let mut a_dt = dt.broadcast_mul(&a)?; + + let mut x_ssd = x_conv.reshape((b_sz, padded_len, self.nheads, self.headdim))?; + + // Zero out padding to prevent it from affecting chunk state computation + if pad_len > 0 { + let mask_ones = Tensor::ones( + (b_sz, seq_len, self.nheads, self.headdim), + x_ssd.dtype(), + x_ssd.device(), + )?; + let mask_zeros = Tensor::zeros( + (b_sz, pad_len, self.nheads, self.headdim), + x_ssd.dtype(), + x_ssd.device(), + )?; + let mask = Tensor::cat(&[&mask_ones, &mask_zeros], 1)?; + x_ssd = x_ssd.broadcast_mul(&mask)?; + + let mask_ones_a = + Tensor::ones((b_sz, seq_len, self.nheads), a_dt.dtype(), a_dt.device())?; + let mask_zeros_a = + Tensor::zeros((b_sz, pad_len, self.nheads), a_dt.dtype(), a_dt.device())?; + let mask_a = Tensor::cat(&[&mask_ones_a, &mask_zeros_a], 1)?; + a_dt = a_dt.broadcast_mul(&mask_a)?; + } + + let heads_per_group = self.nheads / self.ngroups; + let b = b.reshape((b_sz, padded_len, self.ngroups, self.d_state))?; + let b = b + .unsqueeze(3)? + .broadcast_as(( + b_sz, + padded_len, + self.ngroups, + heads_per_group, + self.d_state, + ))? + .reshape((b_sz, padded_len, self.nheads, self.d_state))?; + // Discretize B: B_bar = dt * B (ZOH discretization absorbed into ssd_chunked) + let b = b.broadcast_mul(&dt.unsqueeze(D::Minus1)?)?; + let c = c.reshape((b_sz, padded_len, self.ngroups, self.d_state))?; + let c = c + .unsqueeze(3)? + .broadcast_as(( + b_sz, + padded_len, + self.ngroups, + heads_per_group, + self.d_state, + ))? + .reshape((b_sz, padded_len, self.nheads, self.d_state))?; + + let initial_state = Some(&state.hs[self.layer_idx]); + let (y, final_state) = + self.ssd_chunked(&x_ssd, &a_dt, &b, &c, chunk_size, initial_state)?; + state.hs[self.layer_idx] = final_state; + + let y = y.reshape((b_sz, padded_len, self.d_inner))?; + + let d = self.d.unsqueeze(0)?.unsqueeze(0)?; + let x_skip = x_conv.reshape((b_sz, padded_len, self.nheads, self.headdim))?; + let y = (&y.reshape((b_sz, padded_len, self.nheads, self.headdim))? + + x_skip.broadcast_mul(&d.unsqueeze(D::Minus1)?)?)?; + let y = y.reshape((b_sz, padded_len, self.d_inner))?; + + let y = (y * candle_nn::ops::silu(&z)?)?; + let y = y.reshape((b_sz * padded_len, self.d_inner))?; + let y = self.norm.forward(&y)?; + let y = y.reshape((b_sz, padded_len, self.d_inner))?; + + let y = y.apply(&self.out_proj)?; + + if pad_len > 0 { + y.narrow(1, 0, seq_len) + } else { + Ok(y) + } + } +} + +#[derive(Clone, Debug)] +pub struct ResidualBlock { + mixer: Mamba2Block, + norm: RmsNorm, +} + +impl ResidualBlock { + pub fn new(layer_idx: usize, cfg: &Config, vb: VarBuilder) -> Result { + let norm = candle_nn::rms_norm(cfg.d_model, 1e-5, vb.pp("norm"))?; + let mixer = Mamba2Block::new(layer_idx, cfg, vb.pp("mixer"))?; + Ok(Self { mixer, norm }) + } + + fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + self.mixer.forward(&xs.apply(&self.norm)?, state)? + xs + } + + fn forward_prefill(&self, xs: &Tensor, state: &mut State, chunk_size: usize) -> Result { + let normed = xs.apply(&self.norm)?; + self.mixer.forward_prefill(&normed, state, chunk_size)? + xs + } +} + +#[derive(Clone, Debug)] +pub struct Model { + embedding: candle_nn::Embedding, + layers: Vec, + norm_f: RmsNorm, + lm_head: Linear, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let embedding = candle_nn::embedding(cfg.vocab_size(), cfg.d_model, vb.pp("embeddings"))?; + let mut layers = Vec::with_capacity(cfg.n_layer); + let vb_l = vb.pp("layers"); + for layer_idx in 0..cfg.n_layer { + layers.push(ResidualBlock::new(layer_idx, cfg, vb_l.pp(layer_idx))?); + } + let norm_f = candle_nn::rms_norm(cfg.d_model, 1e-5, vb.pp("norm_f"))?; + let lm_head = Linear::from_weights(embedding.embeddings().clone(), None); + Ok(Self { + embedding, + layers, + norm_f, + lm_head, + dtype: vb.dtype(), + }) + } + + pub fn forward(&self, input_ids: &Tensor, state: &mut State) -> Result { + let mut xs = self.embedding.forward(input_ids)?; + for layer in self.layers.iter() { + xs = layer.forward(&xs, state)?; + } + state.pos += 1; + xs.apply(&self.norm_f)?.apply(&self.lm_head) + } + + pub fn forward_prefill( + &self, + input_ids: &Tensor, + state: &mut State, + chunk_size: usize, + ) -> Result { + let (b_sz, seq_len) = input_ids.dims2()?; + let mut xs = self.embedding.forward(input_ids)?; + for layer in self.layers.iter() { + xs = layer.forward_prefill(&xs, state, chunk_size)?; + } + state.pos += seq_len; + let xs = xs.reshape((b_sz * seq_len, xs.dim(D::Minus1)?))?; + let logits = xs.apply(&self.norm_f)?.apply(&self.lm_head)?; + logits.reshape((b_sz, seq_len, logits.dim(D::Minus1)?)) + } + + pub fn dtype(&self) -> DType { + self.dtype + } +} diff --git a/patches/candle-transformers/src/models/marian.rs b/patches/candle-transformers/src/models/marian.rs new file mode 100644 index 0000000000..ad57b876e1 --- /dev/null +++ b/patches/candle-transformers/src/models/marian.rs @@ -0,0 +1,646 @@ +//! Marian Neural Machine Translation +//! +//! See "Marian: Fast Neural Machine Translation in C++" Junczys-Dowmunt et al. 2018 +//! - [ACL Anthology](https://aclanthology.org/P18-4020/) +//! - [GitHub](https://github.com/marian-nmt/marian) +//! +use super::with_tracing::{linear, Embedding, Linear}; +use candle::{Result, Tensor}; +use candle_nn::{layer_norm, LayerNorm, VarBuilder}; + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub decoder_vocab_size: Option, + pub max_position_embeddings: usize, + pub encoder_layers: usize, + pub encoder_ffn_dim: usize, + pub encoder_attention_heads: usize, + pub decoder_layers: usize, + pub decoder_ffn_dim: usize, + pub decoder_attention_heads: usize, + pub use_cache: bool, + pub is_encoder_decoder: bool, + pub activation_function: candle_nn::Activation, + pub d_model: usize, + pub decoder_start_token_id: u32, + pub scale_embedding: bool, + pub pad_token_id: u32, + pub eos_token_id: u32, + pub forced_eos_token_id: u32, + pub share_encoder_decoder_embeddings: bool, +} + +impl Config { + // https://huggingface.co/Helsinki-NLP/opus-mt-tc-big-fr-en/blob/main/config.json + pub fn opus_mt_tc_big_fr_en() -> Self { + Self { + activation_function: candle_nn::Activation::Relu, + d_model: 1024, + decoder_attention_heads: 16, + decoder_ffn_dim: 4096, + decoder_layers: 6, + decoder_start_token_id: 53016, + decoder_vocab_size: Some(53017), + encoder_attention_heads: 16, + encoder_ffn_dim: 4096, + encoder_layers: 6, + eos_token_id: 43311, + forced_eos_token_id: 43311, + is_encoder_decoder: true, + max_position_embeddings: 1024, + pad_token_id: 53016, + scale_embedding: true, + share_encoder_decoder_embeddings: true, + use_cache: true, + vocab_size: 53017, + } + } + + // https://huggingface.co/Helsinki-NLP/opus-mt-fr-en/blob/main/config.json + pub fn opus_mt_fr_en() -> Self { + Self { + activation_function: candle_nn::Activation::Swish, + d_model: 512, + decoder_attention_heads: 8, + decoder_ffn_dim: 2048, + decoder_layers: 6, + decoder_start_token_id: 59513, + decoder_vocab_size: Some(59514), + encoder_attention_heads: 8, + encoder_ffn_dim: 2048, + encoder_layers: 6, + eos_token_id: 0, + forced_eos_token_id: 0, + is_encoder_decoder: true, + max_position_embeddings: 512, + pad_token_id: 59513, + scale_embedding: true, + share_encoder_decoder_embeddings: true, + use_cache: true, + vocab_size: 59514, + } + } + + pub fn opus_mt_en_zh() -> Self { + Self { + activation_function: candle_nn::Activation::Swish, + d_model: 512, + decoder_attention_heads: 8, + decoder_ffn_dim: 2048, + decoder_layers: 6, + decoder_start_token_id: 65000, + decoder_vocab_size: Some(65001), + encoder_attention_heads: 8, + encoder_ffn_dim: 2048, + encoder_layers: 6, + eos_token_id: 0, + forced_eos_token_id: 0, + is_encoder_decoder: true, + max_position_embeddings: 512, + pad_token_id: 65000, + scale_embedding: true, + share_encoder_decoder_embeddings: true, + use_cache: true, + vocab_size: 65001, + } + } + + pub fn opus_mt_en_hi() -> Self { + Self { + activation_function: candle_nn::Activation::Swish, + d_model: 512, + decoder_attention_heads: 8, + decoder_ffn_dim: 2048, + decoder_layers: 6, + decoder_start_token_id: 61949, + decoder_vocab_size: Some(61950), + encoder_attention_heads: 8, + encoder_ffn_dim: 2048, + encoder_layers: 6, + eos_token_id: 0, + forced_eos_token_id: 0, + is_encoder_decoder: true, + max_position_embeddings: 512, + pad_token_id: 61949, + scale_embedding: true, + share_encoder_decoder_embeddings: true, + use_cache: true, + vocab_size: 61950, + } + } + + pub fn opus_mt_en_es() -> Self { + Self { + activation_function: candle_nn::Activation::Swish, + d_model: 512, + decoder_attention_heads: 8, + decoder_ffn_dim: 2048, + decoder_layers: 6, + decoder_start_token_id: 65000, + decoder_vocab_size: Some(65001), + encoder_attention_heads: 8, + encoder_ffn_dim: 2048, + encoder_layers: 6, + eos_token_id: 0, + forced_eos_token_id: 0, + is_encoder_decoder: true, + max_position_embeddings: 512, + pad_token_id: 65000, + scale_embedding: true, + share_encoder_decoder_embeddings: true, + use_cache: true, + vocab_size: 65001, + } + } + + pub fn opus_mt_en_fr() -> Self { + Self { + activation_function: candle_nn::Activation::Swish, + d_model: 512, + decoder_attention_heads: 8, + decoder_ffn_dim: 2048, + decoder_layers: 6, + decoder_start_token_id: 59513, + decoder_vocab_size: Some(59514), + encoder_attention_heads: 8, + encoder_ffn_dim: 2048, + encoder_layers: 6, + eos_token_id: 0, + forced_eos_token_id: 0, + is_encoder_decoder: true, + max_position_embeddings: 512, + pad_token_id: 59513, + scale_embedding: true, + share_encoder_decoder_embeddings: true, + use_cache: true, + vocab_size: 59514, + } + } + + pub fn opus_mt_en_ru() -> Self { + Self { + activation_function: candle_nn::Activation::Swish, + d_model: 512, + decoder_attention_heads: 8, + decoder_ffn_dim: 2048, + decoder_layers: 6, + decoder_start_token_id: 62517, + decoder_vocab_size: Some(62518), + encoder_attention_heads: 8, + encoder_ffn_dim: 2048, + encoder_layers: 6, + eos_token_id: 0, + forced_eos_token_id: 0, + is_encoder_decoder: true, + max_position_embeddings: 512, + pad_token_id: 62517, + scale_embedding: true, + share_encoder_decoder_embeddings: true, + use_cache: true, + vocab_size: 62518, + } + } +} + +#[derive(Debug, Clone)] +struct SinusoidalPositionalEmbedding { + emb: Embedding, +} + +impl SinusoidalPositionalEmbedding { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dev = vb.device(); + let dtype = vb.dtype(); + let num_positions = cfg.max_position_embeddings; + let dim = cfg.d_model; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / 10000f32.powf(i as f32 / dim as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, num_positions as u32, dev)? + .to_dtype(dtype)? + .reshape((num_positions, 1))?; + let freqs = t.matmul(&inv_freq)?; + let sin = freqs.sin()?; + let cos = freqs.cos()?; + let weights = Tensor::cat(&[&sin, &cos], 1)?.contiguous()?; + let emb = Embedding::from_weights(weights)?; + Ok(Self { emb }) + } + + fn forward(&self, input_ids: &Tensor, past_kv_len: usize) -> Result { + let seq_len = input_ids.dim(1)?; + Tensor::arange( + past_kv_len as u32, + (past_kv_len + seq_len) as u32, + input_ids.device(), + )? + .apply(&self.emb) + } +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + out_proj: Linear, + scaling: f64, + num_heads: usize, + head_dim: usize, + kv_cache: Option<(Tensor, Tensor)>, + is_decoder: bool, +} + +impl Attention { + fn new(cfg: &Config, is_decoder: bool, vb: VarBuilder) -> Result { + let num_heads = if is_decoder { + cfg.decoder_attention_heads + } else { + cfg.encoder_attention_heads + }; + let embed_dim = cfg.d_model; + let head_dim = embed_dim / num_heads; + let scaling = (head_dim as f64).powf(-0.5); + let q_proj = linear(embed_dim, embed_dim, vb.pp("q_proj"))?; + let k_proj = linear(embed_dim, embed_dim, vb.pp("k_proj"))?; + let v_proj = linear(embed_dim, embed_dim, vb.pp("v_proj"))?; + let out_proj = linear(embed_dim, embed_dim, vb.pp("out_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + out_proj, + scaling, + num_heads, + head_dim, + kv_cache: None, + is_decoder, + }) + } + + fn _shape(&self, tensor: &Tensor, bsz: usize) -> Result { + tensor + .reshape((bsz, (), self.num_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous() + } + + fn forward( + &mut self, + xs: &Tensor, + kv_states: Option<&Tensor>, + attn_mask: Option<&Tensor>, + ) -> Result { + let (b_sz, tgt_len, _) = xs.dims3()?; + let query_states = (xs.apply(&self.q_proj)? * self.scaling)?; + let (key_states, value_states) = match kv_states { + None => { + let key_states = self._shape(&xs.apply(&self.k_proj)?, b_sz)?; + let value_states = self._shape(&xs.apply(&self.v_proj)?, b_sz)?; + if self.is_decoder { + let kv_states = match &self.kv_cache { + None => (key_states, value_states), + Some((p_key_states, p_value_states)) => { + let key_states = Tensor::cat(&[p_key_states, &key_states], 2)?; + let value_states = Tensor::cat(&[p_value_states, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some(kv_states.clone()); + kv_states + } else { + (key_states, value_states) + } + } + Some(kv_states) => { + let key_states = self._shape(&kv_states.apply(&self.k_proj)?, b_sz)?; + let value_states = self._shape(&kv_states.apply(&self.v_proj)?, b_sz)?; + (key_states, value_states) + } + }; + let proj_shape = (b_sz * self.num_heads, (), self.head_dim); + let query_states = self._shape(&query_states, b_sz)?.reshape(proj_shape)?; + let key_states = key_states.reshape(proj_shape)?; + let value_states = value_states.reshape(proj_shape)?; + let attn_weights = query_states.matmul(&key_states.transpose(1, 2)?)?; + let attn_weights = match attn_mask { + None => attn_weights, + Some(attn_mask) => attn_weights.broadcast_add(attn_mask)?, + }; + let attn_probs = candle_nn::ops::softmax_last_dim(&attn_weights)?; + let attn_output = attn_probs.matmul(&value_states)?; + attn_output + .reshape((b_sz, self.num_heads, tgt_len, self.head_dim))? + .transpose(1, 2)? + .reshape((b_sz, tgt_len, self.head_dim * self.num_heads))? + .apply(&self.out_proj) + } + + fn reset_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +struct EncoderLayer { + self_attn: Attention, + self_attn_layer_norm: LayerNorm, + activation_fn: candle_nn::Activation, + fc1: Linear, + fc2: Linear, + final_layer_norm: LayerNorm, +} + +impl EncoderLayer { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let self_attn = Attention::new(cfg, true, vb.pp("self_attn"))?; + let self_attn_layer_norm = layer_norm(cfg.d_model, 1e-5, vb.pp("self_attn_layer_norm"))?; + let fc1 = linear(cfg.d_model, cfg.encoder_ffn_dim, vb.pp("fc1"))?; + let fc2 = linear(cfg.encoder_ffn_dim, cfg.d_model, vb.pp("fc2"))?; + let final_layer_norm = layer_norm(cfg.d_model, 1e-5, vb.pp("final_layer_norm"))?; + Ok(Self { + self_attn, + self_attn_layer_norm, + activation_fn: cfg.activation_function, + fc1, + fc2, + final_layer_norm, + }) + } + + fn forward(&mut self, xs: &Tensor) -> Result { + let residual = xs; + let xs = (self.self_attn.forward(xs, None, None)? + residual)? + .apply(&self.self_attn_layer_norm)?; + let residual = &xs; + let xs = xs + .apply(&self.fc1)? + .apply(&self.activation_fn)? + .apply(&self.fc2)?; + (xs + residual)?.apply(&self.final_layer_norm) + } + + fn reset_kv_cache(&mut self) { + self.self_attn.reset_kv_cache() + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + self_attn_layer_norm: LayerNorm, + activation_fn: candle_nn::Activation, + encoder_attn: Attention, + encoder_attn_layer_norm: LayerNorm, + fc1: Linear, + fc2: Linear, + final_layer_norm: LayerNorm, +} + +impl DecoderLayer { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let self_attn = Attention::new(cfg, true, vb.pp("self_attn"))?; + let self_attn_layer_norm = layer_norm(cfg.d_model, 1e-5, vb.pp("self_attn_layer_norm"))?; + let encoder_attn = Attention::new(cfg, true, vb.pp("encoder_attn"))?; + let encoder_attn_layer_norm = + layer_norm(cfg.d_model, 1e-5, vb.pp("encoder_attn_layer_norm"))?; + let fc1 = linear(cfg.d_model, cfg.decoder_ffn_dim, vb.pp("fc1"))?; + let fc2 = linear(cfg.decoder_ffn_dim, cfg.d_model, vb.pp("fc2"))?; + let final_layer_norm = layer_norm(cfg.d_model, 1e-5, vb.pp("final_layer_norm"))?; + Ok(Self { + self_attn, + self_attn_layer_norm, + activation_fn: cfg.activation_function, + encoder_attn, + encoder_attn_layer_norm, + fc1, + fc2, + final_layer_norm, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + encoder_xs: Option<&Tensor>, + attn_mask: &Tensor, + ) -> Result { + let residual = xs; + let xs = (self.self_attn.forward(xs, None, Some(attn_mask))? + residual)? + .apply(&self.self_attn_layer_norm)?; + let xs = match encoder_xs { + None => xs, + Some(encoder_xs) => { + let residual = &xs; + let xs = self.encoder_attn.forward(&xs, Some(encoder_xs), None)?; + (residual + xs)?.apply(&self.encoder_attn_layer_norm)? + } + }; + let residual = &xs; + let xs = xs + .apply(&self.fc1)? + .apply(&self.activation_fn)? + .apply(&self.fc2)?; + let xs = (xs + residual)?.apply(&self.final_layer_norm)?; + Ok(xs) + } + + fn reset_kv_cache(&mut self) { + self.self_attn.reset_kv_cache(); + self.encoder_attn.reset_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct Encoder { + embed_tokens: Embedding, + embed_positions: SinusoidalPositionalEmbedding, + layers: Vec, + embed_scale: Option, +} + +impl Encoder { + fn new(cfg: &Config, embed_tokens: &Embedding, vb: VarBuilder) -> Result { + let embed_positions = SinusoidalPositionalEmbedding::new(cfg, vb.pp("embed_positions"))?; + let mut layers = Vec::with_capacity(cfg.encoder_layers); + let vb_l = vb.pp("layers"); + for idx in 0..cfg.encoder_layers { + let layer = EncoderLayer::new(cfg, vb_l.pp(idx))?; + layers.push(layer) + } + let embed_scale = if cfg.scale_embedding { + Some((cfg.d_model as f64).sqrt()) + } else { + None + }; + Ok(Self { + embed_tokens: embed_tokens.clone(), + embed_positions, + layers, + embed_scale, + }) + } + + pub fn forward(&mut self, xs: &Tensor, past_kv_len: usize) -> Result { + let xs = xs.apply(&self.embed_tokens)?; + let xs = match self.embed_scale { + None => xs, + Some(scale) => (xs * scale)?, + }; + let embed_pos = self + .embed_positions + .forward(&xs, past_kv_len)? + .unsqueeze(0)?; + let mut xs = xs.broadcast_add(&embed_pos)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs)? + } + Ok(xs) + } + + pub fn reset_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.reset_kv_cache() + } + } +} + +#[derive(Debug, Clone)] +pub struct Decoder { + embed_tokens: Embedding, + embed_positions: SinusoidalPositionalEmbedding, + layers: Vec, + embed_scale: Option, +} + +impl Decoder { + fn new(cfg: &Config, embed_tokens: &Embedding, vb: VarBuilder) -> Result { + let embed_positions = SinusoidalPositionalEmbedding::new(cfg, vb.pp("embed_positions"))?; + let mut layers = Vec::with_capacity(cfg.decoder_layers); + let vb_l = vb.pp("layers"); + for idx in 0..cfg.decoder_layers { + let layer = DecoderLayer::new(cfg, vb_l.pp(idx))?; + layers.push(layer) + } + let embed_scale = if cfg.scale_embedding { + Some((cfg.d_model as f64).sqrt()) + } else { + None + }; + Ok(Self { + embed_tokens: embed_tokens.clone(), + embed_positions, + layers, + embed_scale, + }) + } + + pub fn forward( + &mut self, + xs: &Tensor, + encoder_xs: Option<&Tensor>, + past_kv_len: usize, + attn_mask: &Tensor, + ) -> Result { + let xs = xs.apply(&self.embed_tokens)?; + let xs = match self.embed_scale { + None => xs, + Some(scale) => (xs * scale)?, + }; + let embed_pos = self + .embed_positions + .forward(&xs, past_kv_len)? + .unsqueeze(0)?; + let mut xs = xs.broadcast_add(&embed_pos)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, encoder_xs, attn_mask)?; + } + Ok(xs) + } + + pub fn reset_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.reset_kv_cache() + } + } +} + +#[derive(Debug, Clone)] +struct Model { + shared: Embedding, + encoder: Encoder, + decoder: Decoder, +} + +impl Model { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let shared = Embedding::new(cfg.vocab_size, cfg.d_model, vb.pp("shared"))?; + let encoder = Encoder::new(cfg, &shared, vb.pp("encoder"))?; + let decoder = Decoder::new(cfg, &shared, vb.pp("decoder"))?; + Ok(Self { + shared, + encoder, + decoder, + }) + } + + fn reset_kv_cache(&mut self) { + self.encoder.reset_kv_cache(); + self.decoder.reset_kv_cache(); + } +} + +#[derive(Debug, Clone)] +pub struct MTModel { + model: Model, + lm_head: Linear, + final_logits_bias: Tensor, +} + +impl MTModel { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let target_vocab_size = cfg.decoder_vocab_size.unwrap_or(cfg.vocab_size); + let final_logits_bias = vb.get((1, target_vocab_size), "final_logits_bias")?; + let model = Model::new(cfg, vb.pp("model"))?; + let lm_head = Linear::from_weights(model.shared.embeddings().clone(), None); + Ok(Self { + model, + lm_head, + final_logits_bias, + }) + } + + pub fn encoder(&mut self) -> &mut Encoder { + &mut self.model.encoder + } + + pub fn decoder(&mut self) -> &mut Decoder { + &mut self.model.decoder + } + + pub fn decode( + &mut self, + xs: &Tensor, + encoder_xs: &Tensor, + past_kv_len: usize, + ) -> Result { + let seq_len = xs.dim(1)?; + let mask: Vec<_> = (0..seq_len) + .flat_map(|i| (0..seq_len).map(move |j| if j > i { f32::NEG_INFINITY } else { 0f32 })) + .collect(); + let mask = Tensor::from_vec(mask, (seq_len, seq_len), xs.device())?; + self.model + .decoder + .forward(xs, Some(encoder_xs), past_kv_len, &mask)? + .apply(&self.lm_head)? + .broadcast_add(&self.final_logits_bias) + } + + pub fn reset_kv_cache(&mut self) { + self.model.reset_kv_cache(); + } +} diff --git a/patches/candle-transformers/src/models/metavoice.rs b/patches/candle-transformers/src/models/metavoice.rs new file mode 100644 index 0000000000..722aa9e671 --- /dev/null +++ b/patches/candle-transformers/src/models/metavoice.rs @@ -0,0 +1,1032 @@ +//! MetaVoice Studio ML Models +//! +//! See MetaVoice's TTS and voice cloning models: +//! - [GitHub](https://github.com/metavoiceio/metavoice-src) +//! - [Website](https://studio.metavoice.ai/) + +use candle::{DType, Device, Error as E, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{embedding, linear_b, rms_norm, Embedding, Linear, RmsNorm, VarBuilder}; + +// Equivalent to torch.repeat_interleave +pub(crate) fn repeat_interleave(img: &Tensor, repeats: usize, dim: usize) -> Result { + let img = img.unsqueeze(dim + 1)?; + let mut dims = img.dims().to_vec(); + dims[dim + 1] = repeats; + img.broadcast_as(dims)?.flatten(dim, dim + 1) +} +pub mod speaker_encoder { + use super::*; + + #[derive(Debug, Clone, serde::Deserialize)] + pub struct Config { + pub sampling_rate: usize, + pub partial_n_frames: usize, + pub model_hidden_size: usize, + pub model_embedding_size: usize, + pub model_num_layers: usize, + pub mel_window_length: usize, + pub mel_window_step: usize, + pub mel_n_channels: usize, + } + + impl Config { + pub fn cfg() -> Self { + Self { + sampling_rate: 16_000, + partial_n_frames: 160, + model_hidden_size: 256, + model_embedding_size: 256, + model_num_layers: 3, + mel_window_length: 25, + mel_window_step: 10, + mel_n_channels: 40, + } + } + } + + pub struct Model { + lstms: Vec, + linear: Linear, + cfg: Config, + } + + type Slice = (usize, usize); + + impl Model { + pub fn new(cfg: Config, vb: VarBuilder) -> Result { + let mut lstms = Vec::with_capacity(cfg.model_num_layers); + let vb_l = vb.pp("lstm"); + for layer_idx in 0..cfg.model_num_layers { + let c = candle_nn::LSTMConfig { + layer_idx, + ..Default::default() + }; + let lstm = candle_nn::lstm( + cfg.mel_n_channels, + cfg.model_hidden_size, + c, + vb_l.pp(layer_idx), + )?; + lstms.push(lstm) + } + let linear = linear_b( + cfg.model_hidden_size, + cfg.model_embedding_size, + true, + vb.pp("linear"), + )?; + Ok(Self { lstms, linear, cfg }) + } + + fn compute_partial_slices( + &self, + n_samples: usize, + rate: f64, + min_coverage: f64, + ) -> (Vec, Vec) { + let c = &self.cfg; + // Compute how many frames separate two partial utterances + let samples_per_frame = c.sampling_rate * c.mel_window_step / 1000; + let n_frames = n_samples / samples_per_frame + 1; + let frame_step = + (c.sampling_rate as f64 / rate / samples_per_frame as f64).round() as usize; + let steps = (n_frames + frame_step).saturating_sub(c.partial_n_frames) + 1; + // Compute the slices. + let mut wav_slices = vec![]; + let mut mel_slices = vec![]; + for i in (0..steps).step_by(frame_step) { + let mel_range = (i, i + c.partial_n_frames); + let wav_range = ( + i * samples_per_frame, + (i + c.partial_n_frames) * samples_per_frame, + ); + mel_slices.push(mel_range); + wav_slices.push(wav_range); + } + // Evaluate whether extra padding is warranted or not. + let last_wav_range = match wav_slices.last() { + None => return (wav_slices, mel_slices), + Some(l) => *l, + }; + let coverage = (n_samples - last_wav_range.0) as f64 + / (last_wav_range.1 - last_wav_range.0) as f64; + if coverage > min_coverage && mel_slices.len() > 1 { + mel_slices.pop(); + wav_slices.pop(); + } + (wav_slices, mel_slices) + } + + pub fn embed_utterance( + &self, + wav: &[f32], + mel_filters: &[f32], + rate: f64, + min_c: f64, + device: &Device, + ) -> Result { + let (wav_slices, mel_slices) = self.compute_partial_slices(wav.len(), rate, min_c); + let max_wave_length = match wav_slices.last() { + Some(v) => v.1, + None => candle::bail!("empty wav slices"), + }; + let wav = if max_wave_length > wav.len() { + let mut wav = wav.to_vec(); + wav.resize(max_wave_length - wav.len(), 0.0); + std::borrow::Cow::Owned(wav) + } else { + std::borrow::Cow::Borrowed(wav) + }; + let mel = crate::models::whisper::audio::log_mel_spectrogram_( + wav.as_ref(), + mel_filters, + /* fft_size */ self.cfg.mel_window_length, + /* fft_step */ self.cfg.mel_window_step, + self.cfg.mel_n_channels, + false, + ); + let mels = mel_slices + .iter() + .flat_map(|s| [mel[s.0], mel[s.1]]) + .collect::>(); + let mels = Tensor::from_vec(mels, (mel_slices.len(), 2), device)?; + let partial_embeds = self.forward(&mels)?; + let raw_embed = partial_embeds.mean(0)?; + let norm = raw_embed.sqr()?.sum_all()?.sqrt()?; + raw_embed.broadcast_div(&norm) + } + } + + impl Module for Model { + fn forward(&self, xs: &Tensor) -> Result { + use candle_nn::RNN; + + // This is different from the Python transformers version as candle LSTM is batch first. + let xs = xs.t()?; + let mut xs = xs.clone(); + for layer in self.lstms.iter() { + let states = layer.seq(&xs)?; + xs = layer.states_to_tensor(&states)?; + } + let xs = xs.t()?; + let embeds_raw = xs.apply(&self.linear)?.relu()?; + let norm = embeds_raw.sqr()?.sum_keepdim(1)?.sqrt()?; + embeds_raw.broadcast_div(&norm) + } + } +} + +type Rank = u32; + +pub mod tokenizers { + use super::*; + use std::collections::HashMap; + + pub struct BPE { + pub re: fancy_regex::Regex, + pub end_of_text: usize, + pub offset: usize, + pub ranks: HashMap, Rank>, + span: tracing::Span, + } + + impl BPE { + pub fn from_json(json: &serde_json::Value, end_of_text: usize) -> Result { + let json = match json.as_object() { + None => candle::bail!("json value is not an object"), + Some(json) => json, + }; + let re = match json.get("pat_str") { + None => candle::bail!("json object has no pat_str field"), + Some(pat_str) => match pat_str.as_str() { + None => candle::bail!("pat_str field is not a string"), + Some(pat_str) => fancy_regex::Regex::new(pat_str).map_err(E::wrap)?, + }, + }; + let offset = match json.get("offset") { + None => candle::bail!("json object has no offset field"), + Some(offset) => match offset.as_u64() { + None => candle::bail!("offset field is not a positive int"), + Some(offset) => offset as usize, + }, + }; + let mut ranks = HashMap::new(); + for id in 0u8..=255 { + ranks.insert(vec![id], id as u32); + } + let mergeable_ranks = match json.get("mergeable_ranks") { + None => candle::bail!("json object has no mergeable_ranks field"), + Some(mr) => match mr.as_object() { + None => candle::bail!("mergeable_ranks is not an object"), + Some(mr) => mr, + }, + }; + for (key, value) in mergeable_ranks.iter() { + let value = match value.as_u64() { + None => candle::bail!("mergeable_ranks '{key}' is not a u64"), + Some(value) => value as u32, + }; + if value < 256 { + continue; + } + // No escaping for other keys. + let key = key.as_bytes().to_vec(); + ranks.insert(key, value); + } + Ok(Self { + re, + end_of_text, + offset, + ranks, + span: tracing::span!(tracing::Level::TRACE, "bpe"), + }) + } + + // Taken from: + // https://github.com/openai/tiktoken/blob/1b9faf2779855124f05174adf1383e53689ed94b/src/lib.rs#L16C1-L82C2 + fn _byte_pair_merge(&self, piece: &[u8]) -> Vec<(usize, Rank)> { + // This is a vector of (start, rank). + // The rank is of the pair starting at position start. + let mut parts = Vec::with_capacity(piece.len() + 1); + + // Note that we hash bytes when indexing into `ranks`, not token pairs. As long as we train BPE + // the way we currently do, this is equivalent. An easy way to break this would be to decouple + // merge priority from token index or to prevent specific token merges. + let mut min_rank: (Rank, usize) = (Rank::MAX, usize::MAX); + for i in 0..piece.len() - 1 { + let rank = *self.ranks.get(&piece[i..i + 2]).unwrap_or(&Rank::MAX); + if rank < min_rank.0 { + min_rank = (rank, i); + } + parts.push((i, rank)); + } + parts.push((piece.len() - 1, Rank::MAX)); + parts.push((piece.len(), Rank::MAX)); + + let get_rank = { + #[inline(always)] + |parts: &Vec<(usize, Rank)>, i: usize| { + if (i + 3) < parts.len() { + // Similar to `piece[i..i + 2]` above. The +3 is because we haven't yet deleted + // parts[i + 1], see comment in the main loop. + *self + .ranks + .get(&piece[parts[i].0..parts[i + 3].0]) + .unwrap_or(&Rank::MAX) + } else { + Rank::MAX + } + } + }; + + // If you have n parts and m merges, this does O(mn) work. + // We could do something with a heap and do O(m log n) work. + // n is often very small so considerations like cache-locality outweigh the algorithmic + // complexity downsides of the `parts` vector. + while min_rank.0 != Rank::MAX { + let i = min_rank.1; + // Update parts[i] and parts[i - 1] before removing parts[i + 1], since + // `parts.remove(i + 1)` will thrash the cache. + if i > 0 { + parts[i - 1].1 = get_rank(&parts, i - 1); + } + parts[i].1 = get_rank(&parts, i); + parts.remove(i + 1); + + min_rank = (Rank::MAX, usize::MAX); + for (i, &(_, rank)) in parts[..parts.len() - 1].iter().enumerate() { + if rank < min_rank.0 { + min_rank = (rank, i); + } + } + } + parts + } + + pub fn byte_pair_encode(&self, piece: &[u8]) -> Vec { + if piece.is_empty() { + return Vec::new(); + } + if piece.len() == 1 { + return vec![self.ranks[piece]]; + } + assert!(piece.len() > 1); + self._byte_pair_merge(piece) + .windows(2) + .map(|part| self.ranks[&piece[part[0].0..part[1].0]]) + .collect() + } + + pub fn encode(&self, text: &str) -> Result> { + let _enter = self.span.enter(); + let mut bpe_tokens: Vec = Vec::new(); + for word in self.re.find_iter(text) { + let word = word.map_err(E::wrap)?; + let word_tokens = self.byte_pair_encode(word.as_str().as_bytes()); + for &token in word_tokens.iter() { + bpe_tokens.push(token + self.offset as u32) + } + } + bpe_tokens.push((self.end_of_text + self.offset) as u32); + Ok(bpe_tokens) + } + } +} + +pub mod gpt { + use super::*; + + #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)] + pub enum NormType { + LayerNorm, + RMSNorm, + } + + #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)] + pub enum AttnKernelType { + Fa2, + TorchAttn, + Hand, + } + + #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)] + pub enum NonLinearityType { + Gelu, + Swiglu, + } + + enum Norm { + RMSNorm(candle_nn::RmsNorm), + LayerNorm(candle_nn::LayerNorm), + } + + // https://github.com/metavoiceio/metavoice-src/blob/11550bb4e8a1ad032cc1556cc924f7a4e767cbfa/fam/llm/model.py#L27 + #[derive(Debug, Clone)] + pub struct Config { + pub block_size: usize, + pub vocab_sizes: Vec, + pub target_vocab_sizes: Vec, + pub n_layer: usize, + pub n_head: usize, + pub n_embd: usize, + pub bias: bool, + pub causal: bool, + pub spk_emb_on_text: bool, + pub norm_type: NormType, + pub rmsnorm_eps: f64, + pub nonlinearity_type: NonLinearityType, + pub swiglu_multiple_of: Option, + pub attn_kernel_type: AttnKernelType, + pub kv_cache_enabled: bool, + } + + impl Config { + pub fn cfg1b_v0_1() -> Self { + Self { + n_layer: 6, + n_head: 6, + n_embd: 384, + block_size: 1024, + bias: false, + vocab_sizes: vec![1538, 1025], + causal: false, + target_vocab_sizes: vec![1025, 1025, 1025, 1025, 1025, 1025], + swiglu_multiple_of: Some(256), + norm_type: NormType::LayerNorm, + kv_cache_enabled: false, + attn_kernel_type: AttnKernelType::TorchAttn, + spk_emb_on_text: true, + nonlinearity_type: NonLinearityType::Gelu, + rmsnorm_eps: 1e-5, + } + } + } + + impl Norm { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + match cfg.norm_type { + NormType::RMSNorm => { + let rms_norm = candle_nn::rms_norm(cfg.n_embd, cfg.rmsnorm_eps, vb)?; + Ok(Self::RMSNorm(rms_norm)) + } + NormType::LayerNorm => { + let ln_cfg = candle_nn::LayerNormConfig { + affine: cfg.bias, + ..Default::default() + }; + let layer_norm = candle_nn::layer_norm(cfg.n_embd, ln_cfg, vb)?; + Ok(Self::LayerNorm(layer_norm)) + } + } + } + } + + impl Module for Norm { + fn forward(&self, xs: &Tensor) -> Result { + match self { + Self::RMSNorm(m) => m.forward(xs), + Self::LayerNorm(m) => m.forward(xs), + } + } + } + + // https://github.com/metavoiceio/metavoice-src/blob/11550bb4e8a1ad032cc1556cc924f7a4e767cbfa/fam/llm/layers/attn.py#L18 + struct SelfAttention { + c_attn: Linear, + c_proj: Linear, + n_head: usize, + span: tracing::Span, + } + + impl SelfAttention { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + // The different attention variants are likely to be identical but still we only accept + // TorchAttn for now. + if cfg.attn_kernel_type != AttnKernelType::TorchAttn { + candle::bail!("only TorchAttn is supported") + } + if cfg.kv_cache_enabled { + candle::bail!("kv_cache_enabled=true is not supported") + } + let c_attn = linear_b(cfg.n_embd, cfg.n_embd * 3, cfg.bias, vb.pp("c_attn"))?; + let c_proj = linear_b(cfg.n_embd, cfg.n_embd, cfg.bias, vb.pp("c_proj"))?; + Ok(Self { + c_attn, + c_proj, + n_head: cfg.n_head, + span: tracing::span!(tracing::Level::TRACE, "self-attn"), + }) + } + } + + impl Module for SelfAttention { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let (b, t, c) = xs.dims3()?; + let c_x = xs + .apply(&self.c_attn)? + .reshape((b, t, 3, self.n_head, c / self.n_head))?; + let q = c_x.i((.., .., 0))?; + let k = c_x.i((.., .., 1))?; + let v = c_x.i((.., .., 2))?; + let q = q.transpose(1, 2)?.contiguous()?; + let k = k.transpose(1, 2)?.contiguous()?; + let v = v.transpose(1, 2)?.contiguous()?; + let att = (q.matmul(&k.t()?)? / (k.dim(D::Minus1)? as f64).sqrt())?; + // TODO: causal mask + let att = candle_nn::ops::softmax_last_dim(&att)?; + let att = att.matmul(&v)?.transpose(1, 2)?; + att.reshape((b, t, c))?.apply(&self.c_proj) + } + } + + // https://github.com/metavoiceio/metavoice-src/blob/11550bb4e8a1ad032cc1556cc924f7a4e767cbfa/fam/llm/layers/layers.py#L43 + #[allow(clippy::upper_case_acronyms)] + enum MLP { + Gelu { + c_fc: Linear, + c_proj: Linear, + span: tracing::Span, + }, + Swiglu { + w1: Linear, + w3: Linear, + c_proj: Linear, + span: tracing::Span, + }, + } + + impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_dim = 4 * cfg.n_embd; + let slf = match cfg.nonlinearity_type { + NonLinearityType::Gelu => { + let c_fc = linear_b(cfg.n_embd, hidden_dim, cfg.bias, vb.pp("c_fc"))?; + let c_proj = linear_b(hidden_dim, cfg.n_embd, cfg.bias, vb.pp("c_proj"))?; + Self::Gelu { + c_fc, + c_proj, + span: tracing::span!(tracing::Level::TRACE, "mlp-gelu"), + } + } + NonLinearityType::Swiglu => { + let hidden_dim = (2 * hidden_dim) / 3; + let swiglu_multiple_of = match cfg.swiglu_multiple_of { + None => candle::bail!("swiglu-multiple-of has to be set"), + Some(smo) => smo, + }; + let hidden_dim = swiglu_multiple_of * (hidden_dim + swiglu_multiple_of - 1) + / swiglu_multiple_of; + let w1 = linear_b(cfg.n_embd, hidden_dim, cfg.bias, vb.pp("w1"))?; + let w3 = linear_b(cfg.n_embd, hidden_dim, cfg.bias, vb.pp("w3"))?; + let c_proj = linear_b(hidden_dim, cfg.n_embd, cfg.bias, vb.pp("c_proj"))?; + Self::Swiglu { + w1, + w3, + c_proj, + span: tracing::span!(tracing::Level::TRACE, "mlp-swiglu"), + } + } + }; + Ok(slf) + } + } + + impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + match self { + Self::Gelu { c_fc, c_proj, span } => { + let _enter = span.enter(); + xs.apply(c_fc)?.gelu()?.apply(c_proj) + } + Self::Swiglu { + w1, + w3, + c_proj, + span, + } => { + let _enter = span.enter(); + let w1 = xs.apply(w1)?; + let w3 = xs.apply(w3)?; + (w1.silu()? * w3)?.apply(c_proj) + } + } + } + } + + // https://github.com/metavoiceio/metavoice-src/blob/11550bb4e8a1ad032cc1556cc924f7a4e767cbfa/fam/llm/layers/combined.py#L7 + struct Block { + ln_1: Norm, + ln_2: Norm, + attn: SelfAttention, + mlp: MLP, + span: tracing::Span, + } + + impl Block { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let ln_1 = Norm::new(cfg, vb.pp("ln_1"))?; + let ln_2 = Norm::new(cfg, vb.pp("ln_2"))?; + let attn = SelfAttention::new(cfg, vb.pp("attn"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + Ok(Block { + ln_1, + ln_2, + attn, + mlp, + span: tracing::span!(tracing::Level::TRACE, "gpt-block"), + }) + } + } + + impl Module for Block { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let xs = (xs + xs.apply(&self.ln_1)?.apply(&self.attn))?; + let xs = (&xs + xs.apply(&self.ln_2)?.apply(&self.mlp))?; + Ok(xs) + } + } + + // https://github.com/metavoiceio/metavoice-src/blob/11550bb4e8a1ad032cc1556cc924f7a4e767cbfa/fam/llm/model.py#L79 + #[allow(clippy::upper_case_acronyms)] + pub struct Model { + wtes: Vec, + wpe: candle_nn::Embedding, + h: Vec, + ln_f: Norm, + lm_heads: Vec, + cfg: Config, + dtype: DType, + span: tracing::Span, + } + + impl Model { + pub fn new(cfg: Config, vb: VarBuilder) -> Result { + let vb_t = vb.pp("transformer"); + let ln_f = Norm::new(&cfg, vb_t.pp("ln_f"))?; + let mut wtes = Vec::with_capacity(cfg.vocab_sizes.len()); + let vb_w = vb_t.pp("wtes"); + for (idx, vocab_size) in cfg.vocab_sizes.iter().enumerate() { + let wte = candle_nn::embedding(*vocab_size, cfg.n_embd, vb_w.pp(idx))?; + wtes.push(wte) + } + let wpe = candle_nn::embedding(cfg.block_size, cfg.n_embd, vb_t.pp("wpe"))?; + + let mut h = Vec::with_capacity(cfg.n_layer); + let vb_h = vb_t.pp("h"); + for idx in 0..cfg.n_layer { + let block = Block::new(&cfg, vb_h.pp(idx))?; + h.push(block) + } + + let mut lm_heads = Vec::with_capacity(cfg.target_vocab_sizes.len()); + let vb_l = vb.pp("lm_heads"); + for (idx, vocab_size) in cfg.target_vocab_sizes.iter().enumerate() { + let head = linear_b(cfg.n_embd, *vocab_size, false, vb_l.pp(idx))?; + lm_heads.push(head) + } + Ok(Self { + wtes, + wpe, + h, + ln_f, + lm_heads, + cfg, + dtype: vb.dtype(), + span: tracing::span!(tracing::Level::TRACE, "gpt"), + }) + } + + pub fn config(&self) -> &Config { + &self.cfg + } + + pub fn forward(&self, idx: &Tensor) -> Result> { + let _enter = self.span.enter(); + let device = idx.device(); + let (b, _num_hierarchies, t) = idx.dims3()?; + let pos = Tensor::arange(0u32, t as u32, device)?; + let pos_emb = pos.apply(&self.wpe)?; + let mut tok_emb = Tensor::zeros((b, t, self.cfg.n_embd), self.dtype, device)?; + for (wte_idx, wte) in self.wtes.iter().enumerate() { + let emb = idx.i((.., wte_idx, ..))?.apply(wte)?; + tok_emb = (tok_emb + emb)?; + } + // TODO: speaker embs. + let spk_emb = 0f64; + let mut xs = (pos_emb.broadcast_add(&tok_emb)? + spk_emb)?; + for block in self.h.iter() { + xs = xs.apply(block)? + } + let xs = xs.apply(&self.ln_f)?; + let mut logits = Vec::with_capacity(self.lm_heads.len()); + for lm_head in self.lm_heads.iter() { + // non-causal mode only. + let ys = xs.apply(lm_head)?; + logits.push(ys) + } + Ok(logits) + } + } +} + +pub mod transformer { + use super::*; + + #[derive(Debug, Clone, serde::Deserialize)] + pub struct Config { + pub block_size: usize, + pub vocab_size: usize, + pub n_layer: usize, + pub n_head: usize, + pub dim: usize, + pub speaker_emb_dim: usize, + pub intermediate_size: Option, + pub n_local_heads: Option, + pub norm_eps: f64, + } + + impl Config { + pub fn cfg1b_v0_1() -> Self { + Self { + n_layer: 24, + n_head: 16, + dim: 2048, + vocab_size: 2562, + speaker_emb_dim: 256, + block_size: 2048, + intermediate_size: None, + n_local_heads: None, + norm_eps: 1e-5, + } + } + + pub(crate) fn n_local_heads(&self) -> usize { + self.n_local_heads.unwrap_or(self.n_head) + } + + pub(crate) fn head_dim(&self) -> usize { + self.dim / self.n_head + } + + pub(crate) fn intermediate_size(&self) -> usize { + match self.intermediate_size { + Some(intermediate_size) => intermediate_size, + None => { + let hidden_dim = self.dim * 4; + let n_hidden = ((2 * hidden_dim) as f64 / 3.) as usize; + n_hidden.div_ceil(256) * 256 + } + } + } + } + + #[derive(Debug, Clone)] + struct FeedForward { + w1: Linear, + w2: Linear, + w3: Linear, + span: tracing::Span, + } + + impl FeedForward { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let i_size = cfg.intermediate_size(); + let w1 = linear_b(cfg.dim, i_size, false, vb.pp("swiglu.w1"))?; + let w2 = linear_b(i_size, cfg.dim, false, vb.pp("w2"))?; + let w3 = linear_b(cfg.dim, i_size, false, vb.pp("swiglu.w3"))?; + Ok(Self { + w1, + w2, + w3, + span: tracing::span!(tracing::Level::TRACE, "feed-forward"), + }) + } + } + + impl Module for FeedForward { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let swiglu = (candle_nn::ops::silu(&xs.apply(&self.w1)?)? * xs.apply(&self.w3))?; + swiglu.apply(&self.w2) + } + } + + #[derive(Debug, Clone)] + struct Attention { + wqkv: Linear, + wo: Linear, + dim: usize, + kv_size: usize, + n_local_heads: usize, + head_dim: usize, + n_head: usize, + kv_cache: Option<(Tensor, Tensor)>, + span: tracing::Span, + } + + impl Attention { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let n_local_heads = cfg.n_local_heads(); + let head_dim = cfg.head_dim(); + let total_head_dim = (cfg.n_head + 2 * n_local_heads) * head_dim; + let wqkv = linear_b(cfg.dim, total_head_dim, false, vb.pp("wqkv"))?; + let wo = linear_b(cfg.dim, cfg.dim, false, vb.pp("wo"))?; + Ok(Self { + wqkv, + wo, + dim: cfg.dim, + kv_size: n_local_heads * head_dim, + n_local_heads, + head_dim, + n_head: cfg.n_head, + kv_cache: None, + span: tracing::span!(tracing::Level::TRACE, "feed-forward"), + }) + } + + fn forward(&mut self, xs: &Tensor, _pos: usize, mask: &Tensor) -> Result { + let _enter = self.span.enter(); + let (b_sz, seqlen, _) = xs.dims3()?; + + let qkv = xs.apply(&self.wqkv)?; + let q = qkv.narrow(D::Minus1, 0, self.dim)?; + let k = qkv.narrow(D::Minus1, self.dim, self.kv_size)?; + let v = qkv.narrow(D::Minus1, self.dim + self.kv_size, self.kv_size)?; + let q = q + .reshape((b_sz, seqlen, self.n_head, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let k = k + .reshape((b_sz, seqlen, self.n_local_heads, self.head_dim))? + .transpose(1, 2)?; + let v = v + .reshape((b_sz, seqlen, self.n_local_heads, self.head_dim))? + .transpose(1, 2)?; + + let (k, v) = match &self.kv_cache { + None => (k, v), + Some((prev_k, prev_v)) => { + let k = Tensor::cat(&[prev_k, &k], 2)?; + let v = Tensor::cat(&[prev_v, &v], 2)?; + (k, v) + } + }; + self.kv_cache = Some((k.clone(), v.clone())); + + let k = repeat_interleave(&k, self.n_head / self.n_local_heads, 1)?; + let v = repeat_interleave(&v, self.n_head / self.n_local_heads, 1)?; + + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (q.matmul(&k.transpose(2, 3)?)? * scale)?; + + let attn_weights = attn_weights.broadcast_add(mask)?; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + let attn_output = attn_weights.matmul(&v)?; + attn_output + .transpose(1, 2)? + .reshape((b_sz, seqlen, self.dim))? + .apply(&self.wo) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } + } + + #[derive(Debug, Clone)] + struct Block { + attention: Attention, + feed_forward: FeedForward, + ffn_norm: RmsNorm, + attention_norm: RmsNorm, + span: tracing::Span, + } + + impl Block { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let attention = Attention::new(cfg, vb.pp("attention"))?; + let feed_forward = FeedForward::new(cfg, vb.pp("feed_forward"))?; + let ffn_norm = rms_norm(cfg.dim, cfg.norm_eps, vb.pp("ffn_norm"))?; + let attention_norm = rms_norm(cfg.dim, cfg.norm_eps, vb.pp("attention_norm"))?; + Ok(Self { + attention, + feed_forward, + ffn_norm, + attention_norm, + span: tracing::span!(tracing::Level::TRACE, "block"), + }) + } + + fn forward(&mut self, xs: &Tensor, pos: usize, mask: &Tensor) -> Result { + let _enter = self.span.enter(); + let hs = xs.apply(&self.attention_norm)?; + let hs = (xs + self.attention.forward(&hs, pos, mask))?; + &hs + hs.apply(&self.ffn_norm)?.apply(&self.feed_forward) + } + + fn clear_kv_cache(&mut self) { + self.attention.clear_kv_cache() + } + } + + #[derive(Debug, Clone)] + pub struct Model { + tok_embeddings: Embedding, + pos_embeddings: Embedding, + speaker_cond_pos: Linear, + layers: Vec, + norm: RmsNorm, + output: Linear, + spk_cond_mask: Tensor, + span: tracing::Span, + } + + impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let tok_embeddings = embedding(cfg.vocab_size, cfg.dim, vb.pp("tok_embeddings"))?; + let pos_embeddings = embedding(cfg.block_size, cfg.dim, vb.pp("pos_embeddings"))?; + let speaker_cond_pos = linear_b( + cfg.speaker_emb_dim, + cfg.dim, + false, + vb.pp("speaker_cond_pos"), + )?; + let mut layers = Vec::with_capacity(cfg.n_layer); + let vb_l = vb.pp("layers"); + for layer_idx in 0..cfg.n_layer { + let layer = Block::new(cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = rms_norm(cfg.dim, cfg.norm_eps, vb.pp("norm"))?; + let output = linear_b(cfg.dim, cfg.vocab_size, false, vb.pp("output"))?; + let dtype = vb.dtype(); + let spk_cond_mask = Tensor::cat( + &[ + Tensor::ones((1, 1, cfg.dim), dtype, vb.device())?, + Tensor::zeros((1, 1, cfg.dim), dtype, vb.device())?, + ], + 0, + )?; + Ok(Self { + tok_embeddings, + pos_embeddings, + speaker_cond_pos, + layers, + norm, + output, + spk_cond_mask, + span: tracing::span!(tracing::Level::TRACE, "transformer"), + }) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } + + pub fn forward(&mut self, xs: &Tensor, spk_emb: &Tensor, pos: usize) -> Result { + let _enter = self.span.enter(); + let (_b_sz, seqlen) = xs.dims2()?; + let mask: Vec<_> = (0..seqlen) + .flat_map(|i| (0..seqlen).map(move |j| if i < j { f32::NEG_INFINITY } else { 0. })) + .collect(); + let mask = Tensor::from_slice(&mask, (1, 1, seqlen, seqlen), xs.device())?; + let input_pos = Tensor::arange(pos as u32, (pos + seqlen) as u32, xs.device())?; + let tok_embeddings = xs.apply(&self.tok_embeddings)?; + let pos_embeddings = input_pos.apply(&self.pos_embeddings)?; + let mut xs = tok_embeddings + .broadcast_add(&pos_embeddings)? + .broadcast_add( + &spk_emb + .apply(&self.speaker_cond_pos)? + .broadcast_mul(&self.spk_cond_mask)?, + )?; + let mask = mask.to_dtype(xs.dtype())?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, pos, &mask)? + } + xs.narrow(1, seqlen - 1, 1)? + .apply(&self.norm)? + .apply(&self.output) + } + } +} + +pub mod adapters { + // https://github.com/metavoiceio/metavoice-src/blob/9078234c496d76adbec06df789b6b04b1875f129/fam/llm/adapters/tilted_encodec.py + pub struct TiltedEncodec { + end_of_audio_token: u32, + span: tracing::Span, + } + + impl TiltedEncodec { + pub fn new(end_of_audio_token: u32) -> Self { + Self { + end_of_audio_token, + span: tracing::span!(tracing::Level::TRACE, "tilted-encodec"), + } + } + + pub fn decode(&self, tokens: &[Vec]) -> (Vec, Vec>) { + let _enter = self.span.enter(); + let mut text_ids = vec![]; + let mut extracted_audio_ids = vec![]; + let mut min_audio_ids_len = usize::MAX; + for (book_id, tokens) in tokens.iter().enumerate() { + let mut audio_ids = vec![]; + for &t in tokens.iter() { + #[allow(clippy::comparison_chain)] + if t > self.end_of_audio_token { + if book_id == 0 { + text_ids.push(t) + } + } else if t < self.end_of_audio_token { + audio_ids.push(t) + } + } + min_audio_ids_len = usize::min(min_audio_ids_len, audio_ids.len()); + extracted_audio_ids.push(audio_ids) + } + for audio_ids in extracted_audio_ids.iter_mut() { + audio_ids.truncate(min_audio_ids_len) + } + (text_ids, extracted_audio_ids) + } + } + + // https://github.com/metavoiceio/metavoice-src/blob/9078234c496d76adbec06df789b6b04b1875f129/fam/llm/adapters/flattened_encodec.py#L4 + pub struct FlattenedInterleavedEncodec2Codebook { + end_of_audio_token: u32, + span: tracing::Span, + } + + impl FlattenedInterleavedEncodec2Codebook { + pub fn new(end_of_audio_token: u32) -> Self { + Self { + end_of_audio_token, + span: tracing::span!(tracing::Level::TRACE, "encodec2codebook"), + } + } + + pub fn decode(&self, tokens: &[u32]) -> (Vec, Vec, Vec) { + let _enter = self.span.enter(); + let mut text_ids = vec![]; + let mut audio_ids1 = vec![]; + let mut audio_ids2 = vec![]; + for &t in tokens.iter() { + #[allow(clippy::comparison_chain)] + if t < self.end_of_audio_token { + audio_ids1.push(t) + } else if t < 2 * self.end_of_audio_token { + audio_ids2.push(t - self.end_of_audio_token) + } else { + text_ids.push(t) + } + } + (text_ids, audio_ids1, audio_ids2) + } + } +} diff --git a/patches/candle-transformers/src/models/mimi/conv.rs b/patches/candle-transformers/src/models/mimi/conv.rs new file mode 100644 index 0000000000..695c0de66f --- /dev/null +++ b/patches/candle-transformers/src/models/mimi/conv.rs @@ -0,0 +1,671 @@ +// Copyright (c) Kyutai, all rights reserved. +// This source code is licensed under the license found in the +// LICENSE file in the root directory of this source tree. + +use candle::{Module, Result, StreamTensor, StreamingModule, Tensor, D}; +use candle_nn::{Conv1d, VarBuilder}; + +#[allow(clippy::enum_variant_names)] +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum Norm { + WeightNorm, + SpectralNorm, + TimeGroupNorm, +} + +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum PadMode { + Constant, + Reflect, + Replicate, +} + +// Applies weight norm for inference by recomputing the weight tensor. This +// does not apply to training. +// https://pytorch.org/docs/stable/generated/torch.nn.utils.weight_norm.html +fn conv1d_weight_norm( + in_c: usize, + out_c: usize, + kernel_size: usize, + bias: bool, + config: candle_nn::Conv1dConfig, + vb: VarBuilder, +) -> Result { + let weight = if vb.contains_tensor("weight") { + vb.get((out_c, in_c, kernel_size), "weight")? + } else { + let weight_g = vb.get((out_c, 1, 1), "weight_g")?; + let weight_v = vb.get((out_c, in_c, kernel_size), "weight_v")?; + let norm_v = weight_v.sqr()?.sum_keepdim((1, 2))?.sqrt()?; + weight_v.broadcast_mul(&weight_g)?.broadcast_div(&norm_v)? + }; + let bias = if bias { + Some(vb.get(out_c, "bias")?) + } else { + None + }; + Ok(Conv1d::new(weight, bias, config)) +} + +#[derive(Debug, Clone)] +pub struct NormConv1d { + conv: Conv1d, + norm: Option, + span: tracing::Span, +} + +impl NormConv1d { + #[allow(clippy::too_many_arguments)] + pub fn new( + in_c: usize, + out_c: usize, + k_size: usize, + causal: bool, + norm: Option, + bias: bool, + cfg: candle_nn::Conv1dConfig, + vb: VarBuilder, + ) -> Result { + let conv = match norm { + None | Some(Norm::TimeGroupNorm) => { + if bias { + candle_nn::conv1d(in_c, out_c, k_size, cfg, vb.pp("conv"))? + } else { + candle_nn::conv1d_no_bias(in_c, out_c, k_size, cfg, vb.pp("conv"))? + } + } + Some(Norm::WeightNorm) => { + conv1d_weight_norm(in_c, out_c, k_size, bias, cfg, vb.pp("conv"))? + } + Some(Norm::SpectralNorm) => candle::bail!("SpectralNorm is not supported yet."), + }; + let norm = match norm { + None | Some(Norm::WeightNorm) | Some(Norm::SpectralNorm) => None, + Some(Norm::TimeGroupNorm) => { + if causal { + candle::bail!("GroupNorm doesn't support causal evaluation.") + } + let norm = candle_nn::group_norm(1, out_c, 1e-5, vb.pp("norm"))?; + Some(norm) + } + }; + Ok(Self { + conv, + norm, + span: tracing::span!(tracing::Level::TRACE, "norm-conv1d"), + }) + } +} + +impl Module for NormConv1d { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let xs = xs.apply(&self.conv)?; + match self.norm.as_ref() { + None => Ok(xs), + Some(norm) => xs.apply(norm), + } + } +} + +#[derive(Debug, Clone)] +pub struct NormConvTranspose1d { + ws: Tensor, + bs: Option, + k_size: usize, + stride: usize, + groups: usize, + norm: Option, + span: tracing::Span, +} + +impl NormConvTranspose1d { + #[allow(clippy::too_many_arguments)] + pub fn new( + in_c: usize, + out_c: usize, + k_size: usize, + causal: bool, + norm: Option, + bias: bool, + stride: usize, + groups: usize, + vb: VarBuilder, + ) -> Result { + let vb = vb.pp("conv"); + let bs = if bias { + Some(vb.get(out_c, "bias")?) + } else { + None + }; + let ws = match norm { + None | Some(Norm::TimeGroupNorm) => vb.get((in_c, out_c / groups, k_size), "weight")?, + Some(Norm::WeightNorm) => { + if vb.contains_tensor("weight") { + vb.get((in_c, out_c, k_size), "weight")? + } else { + let weight_g = vb.get((in_c, 1, 1), "weight_g")?; + let weight_v = vb.get((in_c, out_c, k_size), "weight_v")?; + let norm_v = weight_v.sqr()?.sum_keepdim((1, 2))?.sqrt()?; + weight_v.broadcast_mul(&weight_g)?.broadcast_div(&norm_v)? + } + } + Some(Norm::SpectralNorm) => candle::bail!("SpectralNorm is not supported yet."), + }; + let (ws, groups) = if groups == out_c && in_c == out_c { + let eye = Tensor::eye(out_c, ws.dtype(), ws.device())?; + let ws = ws + .repeat((1, out_c, 1))? + .mul(&eye.unsqueeze(2)?.repeat((1, 1, k_size))?)?; + (ws, 1) + } else { + (ws, groups) + }; + let norm = match norm { + None | Some(Norm::WeightNorm) | Some(Norm::SpectralNorm) => None, + Some(Norm::TimeGroupNorm) => { + if causal { + candle::bail!("GroupNorm doesn't support causal evaluation.") + } + let norm = candle_nn::group_norm(1, out_c, 1e-5, vb.pp("norm"))?; + Some(norm) + } + }; + Ok(Self { + ws, + bs, + k_size, + stride, + groups, + norm, + span: tracing::span!(tracing::Level::TRACE, "norm-conv-tr1d"), + }) + } +} + +impl Module for NormConvTranspose1d { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + // conv-transpose1d seems to be broken on metal after enough iterations. Causing + // the following error: + // _status < MTLCommandBufferStatusCommitted > + // -[IOGPUMetalCommandBuffer setCurrentCommandEncoder:] + // This is now fixed in candle. + let xs = Tensor::conv_transpose1d(xs, &self.ws, 0, 0, self.stride, 1, self.groups)?; + let xs = match &self.bs { + None => xs, + Some(bias) => { + let b = bias.dims1()?; + let bias = bias.reshape((1, b, 1))?; + xs.broadcast_add(&bias)? + } + }; + match self.norm.as_ref() { + None => Ok(xs), + Some(norm) => xs.apply(norm), + } + } +} + +fn get_extra_padding_for_conv1d( + xs: &Tensor, + k_size: usize, + stride: usize, + padding_total: usize, +) -> Result { + let len = xs.dim(D::Minus1)?; + let n_frames = (len + padding_total).saturating_sub(k_size) as f64 / stride as f64 + 1.0; + let ideal_len = + ((n_frames.ceil() as usize - 1) * stride + k_size).saturating_sub(padding_total); + Ok(ideal_len.saturating_sub(len)) +} + +fn pad1d(xs: &Tensor, pad_l: usize, pad_r: usize, mode: PadMode) -> Result { + match mode { + PadMode::Constant => xs.pad_with_zeros(D::Minus1, pad_l, pad_r), + PadMode::Reflect => candle::bail!("pad-mode 'reflect' is not supported"), + PadMode::Replicate => xs.pad_with_same(D::Minus1, pad_l, pad_r), + } +} + +fn unpad1d(xs: &Tensor, unpad_l: usize, unpad_r: usize) -> Result { + let len = xs.dim(D::Minus1)?; + if len < unpad_l + unpad_r { + candle::bail!("unpad1d: tensor len {len} is too low, {unpad_l} + {unpad_r}") + } + xs.narrow(D::Minus1, unpad_l, len - (unpad_l + unpad_r)) +} + +#[derive(Debug, Clone)] +pub struct StreamableConv1d { + conv: NormConv1d, + causal: bool, + pad_mode: PadMode, + state_prev_xs: StreamTensor, + left_pad_applied: bool, + kernel_size: usize, + span: tracing::Span, +} + +impl StreamableConv1d { + #[allow(clippy::too_many_arguments)] + pub fn new( + in_c: usize, + out_c: usize, + k_size: usize, + stride: usize, + dilation: usize, + groups: usize, + bias: bool, + causal: bool, + norm: Option, + pad_mode: PadMode, + vb: VarBuilder, + ) -> Result { + let cfg = candle_nn::Conv1dConfig { + padding: 0, + stride, + dilation, + groups, + cudnn_fwd_algo: None, + }; + let conv = NormConv1d::new(in_c, out_c, k_size, causal, norm, bias, cfg, vb)?; + if k_size < stride { + candle::bail!("kernel-size {k_size} is smaller than stride {stride}") + } + Ok(Self { + conv, + causal, + pad_mode, + state_prev_xs: StreamTensor::empty(), + left_pad_applied: false, + kernel_size: k_size, + span: tracing::span!(tracing::Level::TRACE, "streamable-conv1d"), + }) + } +} + +impl Module for StreamableConv1d { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let (_b, _t, _c) = xs.dims3()?; + let k_size = self.conv.conv.weight().dim(D::Minus1)?; + let conv_cfg = self.conv.conv.config(); + // Effective kernel size with dilations. + let k_size = (k_size - 1) * conv_cfg.dilation + 1; + let padding_total = k_size - conv_cfg.stride; + let extra_padding = + get_extra_padding_for_conv1d(xs, k_size, conv_cfg.stride, padding_total)?; + let xs = if self.causal { + pad1d(xs, padding_total, extra_padding, self.pad_mode)? + } else { + let padding_right = padding_total / 2; + let padding_left = padding_total - padding_right; + pad1d( + xs, + padding_left, + padding_right + extra_padding, + self.pad_mode, + )? + }; + xs.apply(&self.conv) + } +} + +impl StreamingModule for StreamableConv1d { + fn reset_state(&mut self) { + self.state_prev_xs.reset(); + self.left_pad_applied = false; + } + + fn step(&mut self, xs: &StreamTensor) -> Result { + let _enter = self.span.enter(); + let xs = match xs.as_option() { + None => return Ok(().into()), + Some(xs) => xs.clone(), + }; + let xs = if self.left_pad_applied { + xs + } else { + self.left_pad_applied = true; + let k_size = self.conv.conv.weight().dim(D::Minus1)?; + let conv_cfg = self.conv.conv.config(); + let k_size = (k_size - 1) * conv_cfg.dilation + 1; + let padding_total = k_size - conv_cfg.stride; + pad1d(&xs, padding_total, 0, self.pad_mode)? + }; + let cfg = self.conv.conv.config(); + let stride = cfg.stride; + let dilation = cfg.dilation; + let kernel = (self.kernel_size - 1) * dilation + 1; + let xs = StreamTensor::cat2(&self.state_prev_xs, &xs.into(), D::Minus1)?; + let seq_len = xs.seq_len(D::Minus1)?; + let num_frames = (seq_len + stride).saturating_sub(kernel) / stride; + if num_frames > 0 { + let offset = num_frames * stride; + self.state_prev_xs = xs.narrow(D::Minus1, offset, seq_len - offset)?; + let in_l = (num_frames - 1) * stride + kernel; + let xs = xs.narrow(D::Minus1, 0, in_l)?; + // We apply the underlying convtr directly rather than through forward so as + // not to apply any padding here. + xs.apply(&self.conv.conv) + } else { + self.state_prev_xs = xs; + Ok(StreamTensor::empty()) + } + } +} + +#[derive(Debug, Clone)] +pub struct StreamableConvTranspose1d { + convtr: NormConvTranspose1d, + causal: bool, + state_prev_ys: StreamTensor, + kernel_size: usize, + span: tracing::Span, +} + +impl StreamableConvTranspose1d { + #[allow(clippy::too_many_arguments)] + pub fn new( + in_c: usize, + out_c: usize, + k_size: usize, + stride: usize, + groups: usize, + bias: bool, + causal: bool, + norm: Option, + vb: VarBuilder, + ) -> Result { + let convtr = + NormConvTranspose1d::new(in_c, out_c, k_size, causal, norm, bias, stride, groups, vb)?; + Ok(Self { + convtr, + causal, + kernel_size: k_size, + state_prev_ys: StreamTensor::empty(), + span: tracing::span!(tracing::Level::TRACE, "streamable-conv-tr1d"), + }) + } +} + +impl Module for StreamableConvTranspose1d { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let k_size = self.convtr.k_size; + let stride = self.convtr.stride; + let padding_total = k_size.saturating_sub(stride); + let xs = xs.apply(&self.convtr)?; + if self.causal { + // This corresponds to trim_right_ratio = 1. + unpad1d(&xs, 0, padding_total) + } else { + let padding_right = padding_total / 2; + let padding_left = padding_total - padding_right; + unpad1d(&xs, padding_left, padding_right) + } + } +} + +impl StreamingModule for StreamableConvTranspose1d { + fn reset_state(&mut self) { + self.state_prev_ys.reset() + } + + fn step(&mut self, xs: &StreamTensor) -> Result { + let _enter = self.span.enter(); + let xs = match xs.as_option() { + Some(xs) => xs, + None => return Ok(StreamTensor::empty()), + }; + let stride = self.convtr.stride; + // We apply the underlying convtr directly rather than through forward so as + // not to apply any padding here. + let ys = self.convtr.forward(xs)?; + let ot = ys.dim(D::Minus1)?; + let ys = match self.state_prev_ys.as_option() { + None => ys, + Some(prev_ys) => { + let pt = prev_ys.dim(D::Minus1)?; + // Remove the bias as it will be applied multiple times. + let prev_ys = match &self.convtr.bs { + None => prev_ys.clone(), + Some(bias) => { + let bias = bias.reshape((1, (), 1))?; + prev_ys.broadcast_sub(&bias)? + } + }; + let ys1 = (ys.narrow(D::Minus1, 0, pt)? + prev_ys)?; + let ys2 = ys.narrow(D::Minus1, pt, ot - pt)?; + Tensor::cat(&[ys1, ys2], D::Minus1)? + } + }; + let invalid_steps = self.kernel_size - stride; + let (ys, prev_ys) = StreamTensor::from(ys).split(D::Minus1, ot - invalid_steps)?; + self.state_prev_ys = prev_ys; + Ok(ys) + } +} + +#[derive(Debug, Clone)] +pub struct ConvDownsample1d { + conv: StreamableConv1d, +} + +impl ConvDownsample1d { + pub fn new( + stride: usize, + dim: usize, + causal: bool, + learnt: bool, + vb: VarBuilder, + ) -> Result { + if !learnt { + candle::bail!("only learnt=true is supported") + } + let conv = StreamableConv1d::new( + /* in_c */ dim, + /* out_c */ dim, + /* k_size_c */ 2 * stride, + /* stride */ stride, + /* dilation */ 1, + /* groups */ 1, // channel_wise = false + /* bias */ false, + /* causal */ causal, + /* norm */ None, + /* pad_mode */ PadMode::Replicate, + vb, + )?; + Ok(Self { conv }) + } +} + +impl Module for ConvDownsample1d { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.conv) + } +} + +impl StreamingModule for ConvDownsample1d { + fn reset_state(&mut self) { + self.conv.reset_state() + } + + fn step(&mut self, xs: &StreamTensor) -> Result { + self.conv.step(xs) + } +} + +#[derive(Debug, Clone)] +pub struct ConvTrUpsample1d { + convtr: StreamableConvTranspose1d, +} + +impl ConvTrUpsample1d { + pub fn new( + stride: usize, + dim: usize, + causal: bool, + learnt: bool, + vb: VarBuilder, + ) -> Result { + if !learnt { + candle::bail!("only learnt=true is supported") + } + let convtr = StreamableConvTranspose1d::new( + dim, + dim, + /* k_size */ 2 * stride, + /* stride */ stride, + /* groups */ dim, + /* bias */ false, + /* causal */ causal, + /* norm */ None, + vb, + )?; + Ok(Self { convtr }) + } +} + +impl Module for ConvTrUpsample1d { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.convtr) + } +} + +impl StreamingModule for ConvTrUpsample1d { + fn reset_state(&mut self) { + self.convtr.reset_state() + } + + fn step(&mut self, xs: &StreamTensor) -> Result { + self.convtr.step(xs) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use candle::IndexOp; + + fn run_conv1d( + k_size: usize, + stride: usize, + dilation: usize, + step_size: usize, + len: usize, + bias: bool, + ) -> Result<()> { + // TODO: We should ensure for the seed to be constant when running these tests. + let dev = &candle::Device::Cpu; + let vm = candle_nn::VarMap::new(); + let vb = VarBuilder::from_varmap(&vm, candle::DType::F32, dev); + let conv1d = StreamableConv1d::new( + /* in_c */ 2, + /* out_c */ 3, + /* k_size */ k_size, + /* stride */ stride, + /* dilation */ dilation, + /* groups */ 1, + /* bias */ bias, + /* causal */ true, + /* norm */ None, + /* pad_mode */ PadMode::Constant, + vb, + )?; + let xs = Tensor::randn(0f32, 1., (1, 2, step_size * len), dev)?; + let ys = conv1d.forward(&xs)?; + let mut conv1d = conv1d; + let mut ys_steps = vec![]; + for idx in 0..len { + let xs = xs.i((.., .., step_size * idx..step_size * (idx + 1)))?; + let ys = conv1d.step(&xs.into())?; + if let Some(ys) = ys.as_option() { + ys_steps.push(ys.clone()) + } + } + let ys_steps = Tensor::cat(&ys_steps, D::Minus1)?; + let diff = (&ys - &ys_steps)? + .abs()? + .flatten_all()? + .max(0)? + .to_vec0::()?; + if diff > 1e-5 { + println!("{xs}"); + println!("{ys}"); + println!("{ys_steps}"); + candle::bail!("larger diff than expected {diff}") + } + Ok(()) + } + + fn run_conv_tr1d( + k_size: usize, + stride: usize, + step_size: usize, + len: usize, + bias: bool, + ) -> Result<()> { + // TODO: We should ensure for the seed to be constant when running these tests. + let dev = &candle::Device::Cpu; + let vm = candle_nn::VarMap::new(); + let vb = VarBuilder::from_varmap(&vm, candle::DType::F32, dev); + let conv1d = StreamableConvTranspose1d::new( + /* in_c */ 2, /* out_c */ 3, /* k_size */ k_size, + /* stride */ stride, /* groups */ 1, /* bias */ bias, + /* causal */ true, /* norm */ None, vb, + )?; + let xs = Tensor::randn(0f32, 1., (1, 2, step_size * len), dev)?; + let ys = conv1d.forward(&xs)?; + let mut conv1d = conv1d; + let mut ys_steps = vec![]; + for idx in 0..len { + let xs = xs.i((.., .., step_size * idx..step_size * (idx + 1)))?; + let ys = conv1d.step(&xs.into())?; + if let Some(ys) = ys.as_option() { + ys_steps.push(ys.clone()) + } + } + let ys_steps = Tensor::cat(&ys_steps, D::Minus1)?; + let diff = (&ys - &ys_steps)? + .abs()? + .flatten_all()? + .max(0)? + .to_vec0::()?; + if diff > 1e-5 { + println!("{xs}"); + println!("{ys}"); + println!("{ys_steps}"); + candle::bail!("larger diff than expected {diff}") + } + Ok(()) + } + + #[test] + fn conv1d() -> Result<()> { + for step_size in [1, 2, 3] { + for bias in [false, true] { + run_conv1d(1, 1, 1, step_size, 5, bias)?; + run_conv1d(2, 1, 1, step_size, 5, bias)?; + run_conv1d(2, 2, 1, step_size, 6, bias)?; + run_conv1d(3, 2, 1, step_size, 8, bias)?; + run_conv1d(3, 2, 2, step_size, 8, bias)?; + } + } + Ok(()) + } + + #[test] + fn conv_tr1d() -> Result<()> { + for step_size in [1, 2, 3] { + for bias in [false, true] { + run_conv_tr1d(1, 1, step_size, 5, bias)?; + run_conv_tr1d(2, 1, step_size, 5, bias)?; + run_conv_tr1d(3, 1, step_size, 5, bias)?; + run_conv_tr1d(3, 2, step_size, 5, bias)?; + } + } + Ok(()) + } +} diff --git a/patches/candle-transformers/src/models/mimi/encodec.rs b/patches/candle-transformers/src/models/mimi/encodec.rs new file mode 100644 index 0000000000..f659da3a1f --- /dev/null +++ b/patches/candle-transformers/src/models/mimi/encodec.rs @@ -0,0 +1,229 @@ +// Copyright (c) Kyutai, all rights reserved. +// This source code is licensed under the license found in the +// LICENSE file in the root directory of this source tree. + +use super::{conv, quantization, seanet, transformer}; +use candle::{DType, Device, Module, Result, StreamTensor, StreamingModule, Tensor}; +use candle_nn::VarBuilder; + +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum ResampleMethod { + Conv, + Interpolate, +} + +#[derive(Debug, Clone)] +pub struct Config { + pub channels: usize, + pub sample_rate: f64, + pub frame_rate: f64, + pub renormalize: bool, + pub resample_method: ResampleMethod, + pub seanet: seanet::Config, + pub transformer: transformer::Config, + pub quantizer_n_q: usize, + pub quantizer_bins: usize, + pub quantizer_dim: usize, +} + +impl Config { + // /lustre/scwpod02/client/kyutai/alex/mimi_exp/xps/b7d2bd5a/.hydra/config.yaml + pub fn v0_1(num_codebooks: Option) -> Self { + let seanet_cfg = seanet::Config { + dimension: 512, + channels: 1, + causal: true, + n_filters: 64, + n_residual_layers: 1, + activation: candle_nn::Activation::Elu(1.), + compress: 2, + dilation_base: 2, + disable_norm_outer_blocks: 0, + final_activation: None, + kernel_size: 7, + residual_kernel_size: 3, + last_kernel_size: 3, + lstm: 0, + norm: conv::Norm::WeightNorm, + pad_mode: conv::PadMode::Constant, + ratios: vec![8, 6, 5, 4], + true_skip: true, + }; + let transformer_cfg = transformer::Config { + d_model: seanet_cfg.dimension, + num_heads: 8, + num_layers: 8, + causal: true, + norm_first: true, + bias_ff: false, + bias_attn: false, + layer_scale: Some(0.01), + context: 250, + conv_kernel_size: 5, + use_conv_bias: true, + use_conv_block: false, + cross_attention: false, + max_period: 10000, + gating: None, + norm: super::NormType::LayerNorm, + positional_embedding: transformer::PositionalEmbedding::Rope, + + dim_feedforward: 2048, + kv_repeat: 1, + conv_layout: true, // see builders.py + max_seq_len: 8192, // the transformer works at 25hz so this is ~5 mins. + }; + Config { + channels: 1, + sample_rate: 24_000., + frame_rate: 12.5, + renormalize: true, + resample_method: ResampleMethod::Conv, + seanet: seanet_cfg, + transformer: transformer_cfg, + quantizer_n_q: num_codebooks.unwrap_or(16), + quantizer_bins: 2048, + quantizer_dim: 256, + } + } +} + +#[derive(Debug, Clone)] +pub struct Encodec { + encoder: seanet::SeaNetEncoder, + decoder: seanet::SeaNetDecoder, + encoder_transformer: transformer::ProjectedTransformer, + decoder_transformer: transformer::ProjectedTransformer, + downsample: conv::ConvDownsample1d, + upsample: conv::ConvTrUpsample1d, + quantizer: quantization::SplitResidualVectorQuantizer, + config: Config, +} + +impl Encodec { + pub fn new(cfg: Config, vb: VarBuilder) -> Result { + let dim = cfg.seanet.dimension; + let encoder = seanet::SeaNetEncoder::new(&cfg.seanet, vb.pp("encoder"))?; + let decoder = seanet::SeaNetDecoder::new(&cfg.seanet, vb.pp("decoder"))?; + let encoder_transformer = transformer::ProjectedTransformer::new( + dim, + &[dim], + &cfg.transformer, + vb.pp("encoder_transformer"), + )?; + let decoder_transformer = transformer::ProjectedTransformer::new( + dim, + &[dim], + &cfg.transformer, + vb.pp("decoder_transformer"), + )?; + let quantizer = quantization::SplitResidualVectorQuantizer::new( + /* dim */ cfg.quantizer_dim, + /* input_dim */ Some(dim), + /* output_dim */ Some(dim), + /* n_q */ cfg.quantizer_n_q, + /* bins */ cfg.quantizer_bins, + vb.pp("quantizer"), + )?; + let encoder_frame_rate = + cfg.sample_rate / cfg.seanet.ratios.iter().product::() as f64; + + let downsample_stride = (encoder_frame_rate / cfg.frame_rate) as usize; + // `upsample` and `downsample` only apply if frame_rate is different from encoder_frame_rate. + let downsample = conv::ConvDownsample1d::new( + /* stride */ downsample_stride, + /* dim */ dim, + /* causal */ true, + /* learnt */ true, + vb.pp("downsample"), + )?; + let upsample = conv::ConvTrUpsample1d::new( + /* stride */ downsample_stride, + /* dim */ dim, + /* causal */ true, + /* learnt */ true, + vb.pp("upsample"), + )?; + + Ok(Self { + encoder, + decoder, + encoder_transformer, + decoder_transformer, + quantizer, + downsample, + upsample, + config: cfg, + }) + } + + pub fn config(&self) -> &Config { + &self.config + } + + pub fn encode_pre_quantize(&mut self, xs: &Tensor) -> Result { + let xs = self.encoder.forward(xs)?; + self.encoder_transformer.reset_state(); + let xs = self.encoder_transformer.forward(&xs)?; + let xs = &xs[0]; + xs.apply(&self.downsample) + } + + pub fn encode(&mut self, xs: &Tensor) -> Result { + let xs = self.encoder.forward(xs)?; + self.encoder_transformer.reset_state(); + let xs = self.encoder_transformer.forward(&xs)?; + let xs = &xs[0]; + let xs = xs.apply(&self.downsample)?; + let codes = self.quantizer.encode(&xs)?; + Ok(codes) + } + + pub fn encode_step(&mut self, xs: &StreamTensor) -> Result { + let xs = self.encoder.step(xs)?; + let xs = self.encoder_transformer.step(&xs)?; + let xs = self.downsample.step(&xs)?; + match xs.as_option() { + None => Ok(().into()), + Some(xs) => { + let codes = self.quantizer.encode(xs)?; + Ok(codes.into()) + } + } + } + + pub fn decode(&mut self, codes: &Tensor) -> Result { + let emb = self.quantizer.decode(codes)?; + let emb = emb.apply(&self.upsample)?; + self.decoder_transformer.reset_state(); + let outs = self.decoder_transformer.forward(&emb)?; + let out = &outs[0]; + self.decoder.forward(out) + } + + pub fn decode_step(&mut self, codes: &StreamTensor) -> Result { + let emb = match codes.as_option() { + Some(codes) => StreamTensor::from_tensor(self.quantizer.decode(codes)?), + None => StreamTensor::empty(), + }; + let emb = self.upsample.step(&emb)?; + let out = self.decoder_transformer.step(&emb)?; + self.decoder.step(&out) + } + + pub fn reset_state(&mut self) { + self.encoder.reset_state(); + self.encoder_transformer.reset_state(); + self.decoder.reset_state(); + self.decoder_transformer.reset_state(); + self.upsample.reset_state(); + } +} + +pub fn load(model_file: &str, num_codebooks: Option, dev: &Device) -> Result { + let vb = + unsafe { candle_nn::VarBuilder::from_mmaped_safetensors(&[model_file], DType::F32, dev)? }; + let cfg = Config::v0_1(num_codebooks); + let encodec = Encodec::new(cfg, vb)?; + Ok(encodec) +} diff --git a/patches/candle-transformers/src/models/mimi/mod.rs b/patches/candle-transformers/src/models/mimi/mod.rs new file mode 100644 index 0000000000..8945abfb03 --- /dev/null +++ b/patches/candle-transformers/src/models/mimi/mod.rs @@ -0,0 +1,45 @@ +//! mimi model +//! +//! [Mimi](https://huggingface.co/kyutai/mimi) is a state of the art audio +//! compression model using an encoder/decoder architecture with residual vector +//! quantization. The candle implementation supports streaming meaning that it's +//! possible to encode or decode a stream of audio tokens on the flight to provide +//! low latency interaction with an audio model. +//! +//! - 🤗 [HuggingFace Model Card](https://huggingface.co/kyutai/mimi) +//! - 💻 [GitHub](https://github.com/kyutai-labs/moshi) +//! +//! +//! # Example +//! ```bash +//! # Generating some audio tokens from an audio files. +//! wget https://github.com/metavoiceio/metavoice-src/raw/main/assets/bria.mp3 +//! cargo run --example mimi \ +//! --features mimi --release -- \ +//! audio-to-code bria.mp3 bria.safetensors +//! +//! # And decoding the audio tokens back into a sound file. +//! cargo run --example mimi +//! --features mimi --release -- \ +//! code-to-audio bria.safetensors bria.wav +//! + +// Copyright (c) Kyutai, all rights reserved. +// This source code is licensed under the license found in the +// LICENSE file in the root directory of this source tree. +pub use candle; +pub use candle_nn; + +pub mod conv; +pub mod encodec; +pub mod quantization; +pub mod seanet; +pub mod transformer; + +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum NormType { + RmsNorm, + LayerNorm, +} + +pub use encodec::{load, Config, Encodec as Model}; diff --git a/patches/candle-transformers/src/models/mimi/quantization.rs b/patches/candle-transformers/src/models/mimi/quantization.rs new file mode 100644 index 0000000000..3fde16472b --- /dev/null +++ b/patches/candle-transformers/src/models/mimi/quantization.rs @@ -0,0 +1,404 @@ +// Copyright (c) Kyutai, all rights reserved. +// This source code is licensed under the license found in the +// LICENSE file in the root directory of this source tree. + +use candle::{IndexOp, Layout, Result, Shape, Tensor, D}; +use candle_nn::{linear, Linear, VarBuilder}; + +struct CodebookEncode; + +impl candle::CustomOp2 for CodebookEncode { + fn name(&self) -> &'static str { + "cb" + } + + fn cpu_fwd( + &self, + lhs_storage: &candle::CpuStorage, + lhs_layout: &Layout, + rhs_storage: &candle::CpuStorage, + rhs_layout: &Layout, + ) -> Result<(candle::CpuStorage, Shape)> { + use rayon::prelude::*; + + let (lhs_dim1, lhs_dim2) = lhs_layout.shape().dims2()?; + let (rhs_dim1, rhs_dim2) = rhs_layout.shape().dims2()?; + if lhs_dim2 != rhs_dim2 { + candle::bail!("CodebookEncode, mismatch on last dim, {lhs_layout:?} {rhs_layout:?}"); + } + if lhs_dim2 == 0 { + candle::bail!("CodebookEncode, empty last dim {lhs_layout:?}") + } + let lhs = match lhs_layout.contiguous_offsets() { + None => candle::bail!("CodebookEncode, lhs has to be contiguous, got {lhs_layout:?}"), + Some((o1, o2)) => { + let slice = lhs_storage.as_slice::()?; + &slice[o1..o2] + } + }; + let rhs = match rhs_layout.contiguous_offsets() { + None => candle::bail!("CodebookEncode, rhs has to be contiguous, got {rhs_layout:?}"), + Some((o1, o2)) => { + let slice = rhs_storage.as_slice::()?; + &slice[o1..o2] + } + }; + let dst = (0..lhs_dim1) + .into_par_iter() + .map(|idx1| { + let mut where_min = 0; + let mut min_dist = f32::INFINITY; + let lhs = &lhs[idx1 * lhs_dim2..(idx1 + 1) * lhs_dim2]; + for idx2 in 0..rhs_dim1 { + let rhs = &rhs[idx2 * rhs_dim2..(idx2 + 1) * rhs_dim2]; + let mut dist = 0f32; + for (a, b) in lhs.iter().zip(rhs.iter()) { + dist += (a - b) * (a - b) + } + if dist < min_dist { + min_dist = dist; + where_min = idx2; + } + } + where_min as u32 + }) + .collect(); + let storage = candle::WithDType::to_cpu_storage_owned(dst); + Ok((storage, (lhs_dim1,).into())) + } +} + +#[allow(unused)] +#[derive(Debug, Clone)] +pub struct EuclideanCodebook { + initialized: Tensor, + cluster_usage: Tensor, + embedding_sum: Tensor, + embedding: Tensor, + c2: Tensor, + epsilon: f64, + dim: usize, + span_encode: tracing::Span, + span_decode: tracing::Span, +} + +impl EuclideanCodebook { + pub fn new(dim: usize, codebook_size: usize, vb: VarBuilder) -> Result { + let epsilon = 1e-5; + let initialized = vb.get(1, "initialized")?; + let cluster_usage = vb.get(codebook_size, "cluster_usage")?; + let embedding_sum = vb.get((codebook_size, dim), "embed_sum")?; + let embedding = { + let cluster_usage = cluster_usage.maximum(epsilon)?.unsqueeze(1)?; + embedding_sum.broadcast_div(&cluster_usage)? + }; + let c2 = ((&embedding * &embedding)?.sum(D::Minus1)? / 2.0)?; + Ok(Self { + initialized, + cluster_usage, + embedding_sum, + embedding, + c2, + epsilon, + dim, + span_encode: tracing::span!(tracing::Level::TRACE, "euclidean-encode"), + span_decode: tracing::span!(tracing::Level::TRACE, "euclidean-encode"), + }) + } + + pub fn encode_very_slow(&self, xs: &Tensor) -> Result { + let _enter = self.span_encode.enter(); + let mut target_shape = xs.dims().to_vec(); + target_shape.pop(); + let xs = xs.flatten_to(D::Minus2)?; + let _ = xs.dims2()?; + // TODO: avoid repeating this. + let cluster_usage = self.cluster_usage.maximum(self.epsilon)?.unsqueeze(1)?; + let embedding = self.embedding_sum.broadcast_div(&cluster_usage)?; + // Manual cdist implementation. + let diff = xs.unsqueeze(1)?.broadcast_sub(&embedding.unsqueeze(0)?)?; + let dists = diff.sqr()?.sum(D::Minus1)?; + let codes = dists.argmin(D::Minus1)?; + codes.reshape(target_shape) + } + + pub fn encode_slow(&self, xs: &Tensor) -> Result { + let _enter = self.span_encode.enter(); + let mut target_shape = xs.dims().to_vec(); + target_shape.pop(); + let xs = xs.flatten_to(D::Minus2)?; + let _ = xs.dims2()?; + let dot_prod = xs.matmul(&self.embedding.t()?)?; + let codes = self.c2.broadcast_sub(&dot_prod)?.argmin(D::Minus1)?; + codes.reshape(target_shape) + } + + pub fn encode(&self, xs: &Tensor) -> Result { + let _enter = self.span_encode.enter(); + let mut target_shape = xs.dims().to_vec(); + target_shape.pop(); + let xs = xs.flatten_to(D::Minus2)?; + let _ = xs.dims2()?; + let codes = Tensor::apply_op2(&xs, &self.embedding, CodebookEncode)?; + codes.reshape(target_shape) + } + + pub fn decode(&self, indexes: &Tensor) -> Result { + let _enter = self.span_decode.enter(); + // let ys = candle_nn::Embedding::new(self.embedding.clone(), self.dim).forward(xs)?; + let mut final_dims = indexes.dims().to_vec(); + final_dims.push(self.dim); + let indexes = indexes.flatten_all()?; + let values = self.embedding.index_select(&indexes, 0)?; + let values = values.reshape(final_dims)?; + Ok(values) + } +} + +#[allow(unused)] +#[derive(Debug, Clone)] +pub struct VectorQuantization { + project_in: Option, + project_out: Option, + codebook: EuclideanCodebook, +} + +impl VectorQuantization { + pub fn new( + dim: usize, + codebook_size: usize, + codebook_dim: Option, + vb: VarBuilder, + ) -> Result { + let codebook_dim = codebook_dim.unwrap_or(dim); + let (project_in, project_out) = if codebook_dim == dim { + (None, None) + } else { + let p_in = linear(dim, codebook_dim, vb.pp("project_in"))?; + let p_out = linear(codebook_dim, dim, vb.pp("project_out"))?; + (Some(p_in), Some(p_out)) + }; + let codebook = EuclideanCodebook::new(codebook_dim, codebook_size, vb.pp("codebook"))?; + Ok(Self { + project_in, + project_out, + codebook, + }) + } + + pub fn encode(&self, xs: &Tensor) -> Result { + let xs = xs.t()?.apply(&self.project_in.as_ref())?; + self.codebook.encode_slow(&xs) + } + + pub fn decode(&self, codes: &Tensor) -> Result { + let quantized = self.codebook.decode(codes)?; + let quantized = match &self.project_out { + None => quantized, + Some(p) => quantized.apply(p)?, + }; + quantized.t() + } +} + +#[derive(Debug, Clone)] +pub struct ResidualVectorQuantization { + layers: Vec, +} + +impl ResidualVectorQuantization { + pub fn new( + n_q: usize, + dim: usize, + codebook_size: usize, + codebook_dim: Option, + vb: VarBuilder, + ) -> Result { + let vb = vb.pp("layers"); + let mut layers = Vec::with_capacity(n_q); + for i in 0..n_q { + let layer = VectorQuantization::new(dim, codebook_size, codebook_dim, vb.pp(i))?; + layers.push(layer) + } + Ok(Self { layers }) + } + + pub fn encode(&self, xs: &Tensor) -> Result { + let mut codes = Vec::with_capacity(self.layers.len()); + let mut residual = xs.clone(); + for layer in self.layers.iter() { + let indices = layer.encode(&residual)?; + let quantized = layer.decode(&indices)?; + residual = (residual - quantized)?; + codes.push(indices) + } + Tensor::stack(&codes, 0) + } + + pub fn decode(&self, xs: &Tensor) -> Result { + if self.layers.is_empty() { + candle::bail!("empty layers in ResidualVectorQuantization") + } + if self.layers.len() != xs.dim(0)? { + candle::bail!( + "mismatch between the number of layers {} and the code shape {:?}", + self.layers.len(), + xs.shape() + ) + } + let mut quantized = self.layers[0].decode(&xs.i(0)?)?; + for (i, layer) in self.layers.iter().enumerate().skip(1) { + let xs = xs.i(i)?; + quantized = (quantized + layer.decode(&xs))? + } + Ok(quantized) + } +} + +#[allow(unused)] +#[derive(Debug, Clone)] +pub struct ResidualVectorQuantizer { + vq: ResidualVectorQuantization, + input_proj: Option, + output_proj: Option, +} + +impl ResidualVectorQuantizer { + pub fn new( + dim: usize, + input_dim: Option, + output_dim: Option, + n_q: usize, + bins: usize, + force_projection: bool, + vb: VarBuilder, + ) -> Result { + let input_dim = input_dim.unwrap_or(dim); + let output_dim = output_dim.unwrap_or(dim); + + let input_proj = if input_dim == dim && !force_projection { + None + } else { + let c = candle_nn::conv1d_no_bias( + input_dim, + dim, + 1, + Default::default(), + vb.pp("input_proj"), + )?; + Some(c) + }; + let output_proj = if output_dim == dim && !force_projection { + None + } else { + let c = candle_nn::conv1d_no_bias( + dim, + output_dim, + 1, + Default::default(), + vb.pp("output_proj"), + )?; + Some(c) + }; + + let vq = ResidualVectorQuantization::new( + n_q, dim, /* codebook_size */ bins, /* codebook_dim */ None, vb, + )?; + Ok(Self { + vq, + input_proj, + output_proj, + }) + } + + pub fn encode(&self, xs: &Tensor) -> Result { + let codes = self.vq.encode(&xs.apply(&self.input_proj.as_ref())?)?; + codes.transpose(0, 1) + } + + pub fn decode(&self, codes: &Tensor) -> Result { + // codes is [B, K, T], with T frames, K nb of codebooks, vq.decode expects [K, B, T]. + let codes = codes.transpose(0, 1)?; + let quantized = self.vq.decode(&codes)?; + match &self.output_proj { + None => Ok(quantized), + Some(p) => quantized.apply(p), + } + } +} + +// we do not use any codebook_offset at the moment. When reconstructing the codes, we could just +// concatenate the indexes. +#[derive(Debug, Clone)] +pub struct SplitResidualVectorQuantizer { + rvq_first: ResidualVectorQuantizer, + rvq_rest: ResidualVectorQuantizer, + n_q: usize, + span_encode: tracing::Span, + span_decode: tracing::Span, +} + +impl SplitResidualVectorQuantizer { + pub fn new( + dim: usize, + input_dim: Option, + output_dim: Option, + n_q: usize, + bins: usize, + vb: VarBuilder, + ) -> Result { + let rvq_first = ResidualVectorQuantizer::new( + dim, + input_dim, + output_dim, + 1, + bins, + true, + vb.pp("semantic_residual_vector_quantizer"), + )?; + let rvq_rest = ResidualVectorQuantizer::new( + dim, + input_dim, + output_dim, + n_q - 1, + bins, + true, + vb.pp("acoustic_residual_vector_quantizer"), + )?; + let span_encode = tracing::span!(tracing::Level::TRACE, "split-rvq-encode"); + let span_decode = tracing::span!(tracing::Level::TRACE, "split-rvq-decode"); + Ok(Self { + rvq_first, + rvq_rest, + n_q, + span_encode, + span_decode, + }) + } + + pub fn encode(&self, xs: &Tensor) -> Result { + let _enter = self.span_encode.enter(); + let codes = self.rvq_first.encode(xs)?; + if self.n_q > 1 { + // We encode xs again here rather than the residual. The decomposition is not + // hierarchical but rather having semantic tokens for rvq_first and the acoustic tokens + // for rvq_rest. + let rest_codes = self.rvq_rest.encode(xs)?; + Tensor::cat(&[codes, rest_codes], 1) + } else { + Ok(codes) + } + } + + pub fn decode(&self, codes: &Tensor) -> Result { + // codes is [B, K, T], with T frames, K nb of codebooks. + let _enter = self.span_decode.enter(); + let quantized = self.rvq_first.decode(&codes.i((.., ..1))?)?; + let quantized = if self.n_q > 1 { + (quantized + self.rvq_rest.decode(&codes.i((.., 1..))?))? + } else { + quantized + }; + Ok(quantized) + } +} diff --git a/patches/candle-transformers/src/models/mimi/seanet.rs b/patches/candle-transformers/src/models/mimi/seanet.rs new file mode 100644 index 0000000000..aa5c7d2139 --- /dev/null +++ b/patches/candle-transformers/src/models/mimi/seanet.rs @@ -0,0 +1,465 @@ +// Copyright (c) Kyutai, all rights reserved. +// This source code is licensed under the license found in the +// LICENSE file in the root directory of this source tree. + +use candle::{streaming, Module, Result, StreamTensor, StreamingModule, Tensor}; +use candle_nn::VarBuilder; + +use super::conv::{StreamableConv1d, StreamableConvTranspose1d}; + +#[derive(Debug, Clone)] +pub struct Config { + pub dimension: usize, + pub channels: usize, + pub causal: bool, + pub n_filters: usize, + pub n_residual_layers: usize, + pub ratios: Vec, + pub activation: candle_nn::Activation, + pub norm: super::conv::Norm, + pub kernel_size: usize, + pub residual_kernel_size: usize, + pub last_kernel_size: usize, + pub dilation_base: usize, + pub pad_mode: super::conv::PadMode, + pub true_skip: bool, + pub compress: usize, + pub lstm: usize, + pub disable_norm_outer_blocks: usize, + pub final_activation: Option, +} + +#[derive(Debug, Clone)] +pub struct SeaNetResnetBlock { + block: Vec, + shortcut: Option, + activation: candle_nn::Activation, + skip_op: candle::StreamingBinOp, + span: tracing::Span, +} + +impl SeaNetResnetBlock { + #[allow(clippy::too_many_arguments)] + pub fn new( + dim: usize, + k_sizes_and_dilations: &[(usize, usize)], + activation: candle_nn::Activation, + norm: Option, + causal: bool, + pad_mode: super::conv::PadMode, + compress: usize, + true_skip: bool, + vb: VarBuilder, + ) -> Result { + let mut block = Vec::with_capacity(k_sizes_and_dilations.len()); + let hidden = dim / compress; + let vb_b = vb.pp("block"); + for (i, (k_size, dilation)) in k_sizes_and_dilations.iter().enumerate() { + let in_c = if i == 0 { dim } else { hidden }; + let out_c = if i == k_sizes_and_dilations.len() - 1 { + dim + } else { + hidden + }; + let c = StreamableConv1d::new( + in_c, + out_c, + /* k_size */ *k_size, + /* stride */ 1, + /* dilation */ *dilation, + /* groups */ 1, + /* bias */ true, + /* causal */ causal, + /* norm */ norm, + /* pad_mode */ pad_mode, + vb_b.pp(2 * i + 1), + )?; + block.push(c) + } + let shortcut = if true_skip { + None + } else { + let c = StreamableConv1d::new( + dim, + dim, + /* k_size */ 1, + /* stride */ 1, + /* dilation */ 1, + /* groups */ 1, + /* bias */ true, + /* causal */ causal, + /* norm */ norm, + /* pad_mode */ pad_mode, + vb.pp("shortcut"), + )?; + Some(c) + }; + Ok(Self { + block, + shortcut, + activation, + skip_op: streaming::StreamingBinOp::new(streaming::BinOp::Add, candle::D::Minus1), + span: tracing::span!(tracing::Level::TRACE, "sea-resnet"), + }) + } +} + +impl Module for SeaNetResnetBlock { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let mut ys = xs.clone(); + for block in self.block.iter() { + ys = ys.apply(&self.activation)?.apply(block)?; + } + match self.shortcut.as_ref() { + None => ys + xs, + Some(shortcut) => ys + xs.apply(shortcut), + } + } +} + +impl StreamingModule for SeaNetResnetBlock { + fn reset_state(&mut self) { + for block in self.block.iter_mut() { + block.reset_state() + } + if let Some(shortcut) = self.shortcut.as_mut() { + shortcut.reset_state() + } + } + + fn step(&mut self, xs: &StreamTensor) -> Result { + let _enter = self.span.enter(); + let mut ys = xs.clone(); + for block in self.block.iter_mut() { + ys = block.step(&ys.apply(&self.activation)?)?; + } + match self.shortcut.as_ref() { + None => self.skip_op.step(&ys, xs), + Some(shortcut) => self.skip_op.step(&ys, &xs.apply(shortcut)?), + } + } +} + +#[derive(Debug, Clone)] +struct EncoderLayer { + residuals: Vec, + downsample: StreamableConv1d, +} + +#[derive(Debug, Clone)] +pub struct SeaNetEncoder { + init_conv1d: StreamableConv1d, + activation: candle_nn::Activation, + layers: Vec, + final_conv1d: StreamableConv1d, + span: tracing::Span, +} + +impl SeaNetEncoder { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + if cfg.lstm > 0 { + candle::bail!("seanet lstm is not supported") + } + let n_blocks = 2 + cfg.ratios.len(); + let mut mult = 1usize; + let init_norm = if cfg.disable_norm_outer_blocks >= 1 { + None + } else { + Some(cfg.norm) + }; + let mut layer_idx = 0; + let vb = vb.pp("layers"); + let init_conv1d = StreamableConv1d::new( + cfg.channels, + mult * cfg.n_filters, + cfg.kernel_size, + /* stride */ 1, + /* dilation */ 1, + /* groups */ 1, + /* bias */ true, + /* causal */ cfg.causal, + /* norm */ init_norm, + /* pad_mode */ cfg.pad_mode, + vb.pp(layer_idx), + )?; + layer_idx += 1; + let mut layers = Vec::with_capacity(cfg.ratios.len()); + + for (i, &ratio) in cfg.ratios.iter().rev().enumerate() { + let norm = if cfg.disable_norm_outer_blocks >= i + 2 { + None + } else { + Some(cfg.norm) + }; + let mut residuals = Vec::with_capacity(cfg.n_residual_layers); + for j in 0..cfg.n_residual_layers { + let resnet_block = SeaNetResnetBlock::new( + mult * cfg.n_filters, + &[ + (cfg.residual_kernel_size, cfg.dilation_base.pow(j as u32)), + (1, 1), + ], + cfg.activation, + norm, + cfg.causal, + cfg.pad_mode, + cfg.compress, + cfg.true_skip, + vb.pp(layer_idx), + )?; + residuals.push(resnet_block); + layer_idx += 1; + } + let downsample = StreamableConv1d::new( + mult * cfg.n_filters, + mult * cfg.n_filters * 2, + /* k_size */ ratio * 2, + /* stride */ ratio, + /* dilation */ 1, + /* groups */ 1, + /* bias */ true, + /* causal */ true, + /* norm */ norm, + /* pad_mode */ cfg.pad_mode, + vb.pp(layer_idx + 1), + )?; + layer_idx += 2; + let layer = EncoderLayer { + downsample, + residuals, + }; + layers.push(layer); + mult *= 2 + } + + let final_norm = if cfg.disable_norm_outer_blocks >= n_blocks { + None + } else { + Some(cfg.norm) + }; + let final_conv1d = StreamableConv1d::new( + mult * cfg.n_filters, + cfg.dimension, + cfg.last_kernel_size, + /* stride */ 1, + /* dilation */ 1, + /* groups */ 1, + /* bias */ true, + /* causal */ cfg.causal, + /* norm */ final_norm, + /* pad_mode */ cfg.pad_mode, + vb.pp(layer_idx + 1), + )?; + Ok(Self { + init_conv1d, + activation: cfg.activation, + layers, + final_conv1d, + span: tracing::span!(tracing::Level::TRACE, "sea-encoder"), + }) + } +} + +impl Module for SeaNetEncoder { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let mut xs = xs.apply(&self.init_conv1d)?; + for layer in self.layers.iter() { + for residual in layer.residuals.iter() { + xs = xs.apply(residual)? + } + xs = xs.apply(&self.activation)?.apply(&layer.downsample)?; + } + xs.apply(&self.activation)?.apply(&self.final_conv1d) + } +} + +impl StreamingModule for SeaNetEncoder { + fn reset_state(&mut self) { + self.init_conv1d.reset_state(); + self.layers.iter_mut().for_each(|v| { + v.residuals.iter_mut().for_each(|v| v.reset_state()); + v.downsample.reset_state() + }); + self.final_conv1d.reset_state(); + } + + fn step(&mut self, xs: &StreamTensor) -> Result { + let _enter = self.span.enter(); + let mut xs = self.init_conv1d.step(xs)?; + for layer in self.layers.iter_mut() { + for residual in layer.residuals.iter_mut() { + xs = residual.step(&xs)?; + } + xs = layer.downsample.step(&xs.apply(&self.activation)?)?; + } + self.final_conv1d.step(&xs.apply(&self.activation)?) + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + upsample: StreamableConvTranspose1d, + residuals: Vec, +} + +#[derive(Debug, Clone)] +pub struct SeaNetDecoder { + init_conv1d: StreamableConv1d, + activation: candle_nn::Activation, + layers: Vec, + final_conv1d: StreamableConv1d, + final_activation: Option, + span: tracing::Span, +} + +impl SeaNetDecoder { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + if cfg.lstm > 0 { + candle::bail!("seanet lstm is not supported") + } + let n_blocks = 2 + cfg.ratios.len(); + let mut mult = 1 << cfg.ratios.len(); + let init_norm = if cfg.disable_norm_outer_blocks == n_blocks { + None + } else { + Some(cfg.norm) + }; + let mut layer_idx = 0; + let vb = vb.pp("layers"); + let init_conv1d = StreamableConv1d::new( + cfg.dimension, + mult * cfg.n_filters, + cfg.kernel_size, + /* stride */ 1, + /* dilation */ 1, + /* groups */ 1, + /* bias */ true, + /* causal */ cfg.causal, + /* norm */ init_norm, + /* pad_mode */ cfg.pad_mode, + vb.pp(layer_idx), + )?; + layer_idx += 1; + let mut layers = Vec::with_capacity(cfg.ratios.len()); + for (i, &ratio) in cfg.ratios.iter().enumerate() { + let norm = if cfg.disable_norm_outer_blocks + i + 1 >= n_blocks { + None + } else { + Some(cfg.norm) + }; + let upsample = StreamableConvTranspose1d::new( + mult * cfg.n_filters, + mult * cfg.n_filters / 2, + /* k_size */ ratio * 2, + /* stride */ ratio, + /* groups */ 1, + /* bias */ true, + /* causal */ true, + /* norm */ norm, + vb.pp(layer_idx + 1), + )?; + layer_idx += 2; + + let mut residuals = Vec::with_capacity(cfg.n_residual_layers); + for j in 0..cfg.n_residual_layers { + let resnet_block = SeaNetResnetBlock::new( + mult * cfg.n_filters / 2, + &[ + (cfg.residual_kernel_size, cfg.dilation_base.pow(j as u32)), + (1, 1), + ], + cfg.activation, + norm, + cfg.causal, + cfg.pad_mode, + cfg.compress, + cfg.true_skip, + vb.pp(layer_idx), + )?; + residuals.push(resnet_block); + layer_idx += 1; + } + let layer = DecoderLayer { + upsample, + residuals, + }; + layers.push(layer); + mult /= 2 + } + let final_norm = if cfg.disable_norm_outer_blocks >= 1 { + None + } else { + Some(cfg.norm) + }; + let final_conv1d = StreamableConv1d::new( + cfg.n_filters, + cfg.channels, + cfg.last_kernel_size, + /* stride */ 1, + /* dilation */ 1, + /* groups */ 1, + /* bias */ true, + /* causal */ cfg.causal, + /* norm */ final_norm, + /* pad_mode */ cfg.pad_mode, + vb.pp(layer_idx + 1), + )?; + Ok(Self { + init_conv1d, + activation: cfg.activation, + layers, + final_conv1d, + final_activation: cfg.final_activation, + span: tracing::span!(tracing::Level::TRACE, "sea-decoder"), + }) + } +} + +impl Module for SeaNetDecoder { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let mut xs = xs.apply(&self.init_conv1d)?; + for layer in self.layers.iter() { + xs = xs.apply(&self.activation)?.apply(&layer.upsample)?; + for residual in layer.residuals.iter() { + xs = xs.apply(residual)? + } + } + let xs = xs.apply(&self.activation)?.apply(&self.final_conv1d)?; + let xs = match self.final_activation.as_ref() { + None => xs, + Some(act) => xs.apply(act)?, + }; + Ok(xs) + } +} + +impl StreamingModule for SeaNetDecoder { + fn reset_state(&mut self) { + self.init_conv1d.reset_state(); + self.layers.iter_mut().for_each(|v| { + v.residuals.iter_mut().for_each(|v| v.reset_state()); + v.upsample.reset_state() + }); + self.final_conv1d.reset_state(); + } + + fn step(&mut self, xs: &StreamTensor) -> Result { + let _enter = self.span.enter(); + let mut xs = self.init_conv1d.step(xs)?; + for layer in self.layers.iter_mut() { + xs = layer.upsample.step(&xs.apply(&self.activation)?)?; + for residual in layer.residuals.iter_mut() { + xs = residual.step(&xs)?; + } + } + let xs = self.final_conv1d.step(&xs.apply(&self.activation)?)?; + let xs = match self.final_activation.as_ref() { + None => xs, + Some(act) => xs.apply(act)?, + }; + Ok(xs) + } +} diff --git a/patches/candle-transformers/src/models/mimi/transformer.rs b/patches/candle-transformers/src/models/mimi/transformer.rs new file mode 100644 index 0000000000..6ccbc8d12a --- /dev/null +++ b/patches/candle-transformers/src/models/mimi/transformer.rs @@ -0,0 +1,777 @@ +// Copyright (c) Kyutai, all rights reserved. +// This source code is licensed under the license found in the +// LICENSE file in the root directory of this source tree. + +use candle::{DType, Device, IndexOp, Module, Result, StreamTensor, StreamingModule, Tensor, D}; +use candle_nn::{linear_no_bias, Linear, VarBuilder}; +use std::sync::Arc; + +fn linear(in_d: usize, out_d: usize, bias: bool, vb: VarBuilder) -> Result { + if bias { + candle_nn::linear(in_d, out_d, vb) + } else { + linear_no_bias(in_d, out_d, vb) + } +} + +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum PositionalEmbedding { + Rope, + Sin, + None, +} + +#[derive(Debug, Clone)] +pub struct Config { + pub d_model: usize, + pub num_heads: usize, + pub num_layers: usize, + pub causal: bool, + pub norm_first: bool, + pub bias_ff: bool, + pub bias_attn: bool, + pub layer_scale: Option, + pub positional_embedding: PositionalEmbedding, + pub use_conv_block: bool, + pub cross_attention: bool, + pub conv_kernel_size: usize, + pub use_conv_bias: bool, + pub gating: Option, + pub norm: super::NormType, + pub context: usize, + pub max_period: usize, + pub max_seq_len: usize, + + pub kv_repeat: usize, + pub dim_feedforward: usize, + pub conv_layout: bool, +} + +#[derive(Debug, Clone)] +pub struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, + span: tracing::Span, +} + +impl RotaryEmbedding { + pub fn new(dim: usize, max_seq_len: usize, theta: f32, dev: &Device) -> Result { + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / theta.powf(i as f32 / dim as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(DType::F32)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + span: tracing::span!(tracing::Level::TRACE, "rot"), + }) + } + + pub fn apply_rotary_emb(&self, qk: &Tensor, seqlen_offset: usize) -> Result { + let _enter = self.span.enter(); + let (_b_size, _nheads, seqlen, _headdim) = qk.dims4()?; + let qk_dtype = qk.dtype(); + let c = self.cos.narrow(0, seqlen_offset, seqlen)?; + let s = self.sin.narrow(0, seqlen_offset, seqlen)?; + candle_nn::rotary_emb::rope_i(&qk.to_dtype(DType::F32)?, &c, &s)?.to_dtype(qk_dtype) + } +} + +#[derive(Debug, Clone)] +pub struct LayerScale { + scale: Tensor, +} + +impl LayerScale { + pub fn new(d_model: usize, _init: f64, vb: VarBuilder) -> Result { + let scale = vb.get(d_model, "scale")?; + Ok(Self { scale }) + } +} + +impl Module for LayerScale { + fn forward(&self, xs: &Tensor) -> Result { + xs.broadcast_mul(&self.scale) + } +} + +#[derive(Debug, Clone)] +pub struct StreamingMultiheadAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + out_proj: Linear, + kv_repeat: usize, + num_heads: usize, + context: usize, + neg_inf: Tensor, + rope: Option>, + kv_cache: candle_nn::kv_cache::RotatingKvCache, + pos: usize, + use_flash_attn: bool, + span: tracing::Span, +} + +impl StreamingMultiheadAttention { + pub fn new(rope: &Option>, cfg: &Config, vb: VarBuilder) -> Result { + let embed_dim = cfg.d_model; + let num_kv = cfg.num_heads / cfg.kv_repeat; + let kv_dim = num_kv * (embed_dim / cfg.num_heads); + let q_proj = linear(embed_dim, embed_dim, cfg.bias_attn, vb.pp("q_proj"))?; + let k_proj = linear(embed_dim, kv_dim, cfg.bias_attn, vb.pp("k_proj"))?; + let v_proj = linear(embed_dim, kv_dim, cfg.bias_attn, vb.pp("v_proj"))?; + let out_proj = linear(embed_dim, embed_dim, cfg.bias_attn, vb.pp("o_proj"))?; + let neg_inf = Tensor::new(f32::NEG_INFINITY, vb.device())?.to_dtype(vb.dtype())?; + Ok(Self { + q_proj, + k_proj, + v_proj, + out_proj, + rope: rope.clone(), + kv_repeat: cfg.kv_repeat, + num_heads: cfg.num_heads, + context: cfg.context, + neg_inf, + kv_cache: candle_nn::kv_cache::RotatingKvCache::new(2, cfg.context), + pos: 0, + use_flash_attn: false, + span: tracing::span!(tracing::Level::TRACE, "mha"), + }) + } + + pub fn forward(&mut self, xs: &Tensor, mask: Option<&Tensor>) -> Result { + let _enter = self.span.enter(); + if self.kv_repeat != 1 { + candle::bail!("only kv-repeat = 1 is supported") + } + let (b, t, hd) = xs.dims3()?; + let head_dim = hd / self.num_heads; + let q = xs + .apply(&self.q_proj)? + .reshape((b, t, self.num_heads, head_dim))?; + let k = xs + .apply(&self.k_proj)? + .reshape((b, t, self.num_heads, head_dim))?; + let v = xs + .apply(&self.v_proj)? + .reshape((b, t, self.num_heads, head_dim))?; + // qk_layer_norm = None + // kv_repeat = 1, otherwise we would need repeat_kv + let mut q = q.transpose(1, 2)?.contiguous()?; // b,h,t,d + let mut k = k.transpose(1, 2)?.contiguous()?; // b,h,k,d + let v = v.transpose(1, 2)?.contiguous()?; // b,h,k,d + if let Some(rope) = &self.rope { + q = rope.apply_rotary_emb(&q, self.pos)?; + k = rope.apply_rotary_emb(&k, self.pos)?; + } + + let (k, v) = { + self.pos += k.dim(2)?; + self.kv_cache.append(&k.contiguous()?, &v.contiguous()?)? + }; + // The KV cache keeps all the data at the moment, we want to trim + // down the part that comes from the cache to at most context to + // be coherent with the mask shape we provide. + let k_len = k.dim(2)?; + let k_target_len = t + usize::min(self.context, k_len - t); + let (k, v) = if k_target_len < k_len { + let k = k.narrow(2, k_len - k_target_len, k_target_len)?; + let v = v.narrow(2, k_len - k_target_len, k_target_len)?; + (k, v) + } else { + (k.clone(), v.clone()) + }; + + let xs = if q.dtype() == DType::BF16 && self.use_flash_attn { + let q = q.transpose(1, 2)?; + let k = k.transpose(1, 2)?; + let v = v.transpose(1, 2)?; + let softmax_scale = 1f32 / (head_dim as f32).sqrt(); + flash_attn(&q, &k, &v, softmax_scale, t > 1)?.transpose(1, 2)? + } else { + let pre_ws = q.matmul(&k.t()?)?; // b,h,t,k + let pre_ws = (pre_ws * (head_dim as f64).powf(-0.5))?; + + let pre_ws = match mask { + None => pre_ws, + Some(mask) => { + let mask = mask.broadcast_left((b, self.num_heads))?; + let neg_inf = self.neg_inf.broadcast_as(pre_ws.shape())?; + mask.where_cond(&neg_inf, &pre_ws)? + } + }; + + let ws = candle_nn::ops::softmax_last_dim(&pre_ws)?; // b,h,t,k + ws.matmul(&v)? // b,h,t,d + }; + let xs = xs + .transpose(1, 2)? // b,t,h,d + .reshape((b, t, hd))? + .apply(&self.out_proj)?; + Ok(xs) + } + + pub fn reset_kv_cache(&mut self) { + self.kv_cache.reset() + } + + pub fn set_kv_cache(&mut self, kv_cache: candle_nn::kv_cache::RotatingKvCache) { + self.kv_cache = kv_cache + } +} + +#[derive(Debug, Clone)] +pub struct StreamingMultiheadCrossAttention { + in_proj_q: Linear, + in_proj_k: Linear, + in_proj_v: Linear, + out_proj: Linear, + kv_repeat: usize, + num_heads: usize, + neg_inf: Tensor, + span: tracing::Span, +} + +impl StreamingMultiheadCrossAttention { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let embed_dim = cfg.d_model; + let num_kv = cfg.num_heads / cfg.kv_repeat; + let kv_dim = num_kv * (embed_dim / cfg.num_heads); + let out_dim = embed_dim + 2 * kv_dim; + let in_proj_weight = vb.get((out_dim, embed_dim), "in_proj_weight")?; + let in_proj_weight_q = in_proj_weight.narrow(0, 0, embed_dim)?; + let in_proj_weight_k = in_proj_weight.narrow(0, embed_dim, kv_dim)?; + let in_proj_weight_v = in_proj_weight.narrow(0, embed_dim + kv_dim, kv_dim)?; + let (in_proj_bias_q, in_proj_bias_k, in_proj_bias_v) = if cfg.bias_attn { + let b = vb.get(out_dim, "in_proj_bias")?; + let q = b.narrow(0, 0, embed_dim)?; + let k = b.narrow(0, embed_dim, kv_dim)?; + let v = b.narrow(0, embed_dim + kv_dim, kv_dim)?; + (Some(q), Some(k), Some(v)) + } else { + (None, None, None) + }; + let in_proj_q = Linear::new(in_proj_weight_q, in_proj_bias_q); + let in_proj_k = Linear::new(in_proj_weight_k, in_proj_bias_k); + let in_proj_v = Linear::new(in_proj_weight_v, in_proj_bias_v); + let out_proj = linear(embed_dim, embed_dim, cfg.bias_attn, vb.pp("out_proj"))?; + let neg_inf = Tensor::new(f32::NEG_INFINITY, vb.device())?.to_dtype(vb.dtype())?; + Ok(Self { + in_proj_q, + in_proj_k, + in_proj_v, + out_proj, + kv_repeat: cfg.kv_repeat, + num_heads: cfg.num_heads, + neg_inf, + span: tracing::span!(tracing::Level::TRACE, "mhca"), + }) + } + + pub fn forward(&self, xs: &Tensor, ca_src: &Tensor, mask: Option<&Tensor>) -> Result { + let _enter = self.span.enter(); + if self.kv_repeat != 1 { + candle::bail!("only kv-repeat = 1 is supported") + } + let (b, t, hd) = xs.dims3()?; + let head_dim = hd / self.num_heads; + // time_dim = 1, layout: b,t,h,d + let q = xs.apply(&self.in_proj_q)?; + let k = ca_src.apply(&self.in_proj_k)?; + let v = ca_src.apply(&self.in_proj_v)?; + let (ca_b, ca_t, ca_dim) = k.dims3()?; + let q = q.reshape((b, t, self.num_heads, head_dim))?; + let k = k.reshape((ca_b, ca_t, ca_dim / head_dim, head_dim))?; + let v = v.reshape((ca_b, ca_t, ca_dim / head_dim, head_dim))?; + // qk_layer_norm = None + // kv_repeat = 1, otherwise we would need repeat_kv + let q = q.transpose(1, 2)?.contiguous()?; // b,h,t,d + let k = k.transpose(1, 2)?.contiguous()?; // b,h,k,d + let v = v.transpose(1, 2)?.contiguous()?; // b,h,k,d + + let pre_ws = q.matmul(&k.t()?)?; // b,h,t,k + let pre_ws = (pre_ws * (head_dim as f64).powf(-0.5))?; + + let pre_ws = match mask { + None => pre_ws, + Some(mask) => { + let mask = mask.broadcast_left((b, self.num_heads))?; + let neg_inf = self.neg_inf.broadcast_as(pre_ws.shape())?; + mask.where_cond(&neg_inf, &pre_ws)? + } + }; + + let ws = candle_nn::ops::softmax_last_dim(&pre_ws)?; // b,h,t,k + let xs = ws.matmul(&v)?; // b,h,t,d + let xs = xs + .transpose(1, 2)? // b,t,h,d + .reshape((b, t, hd))? + .apply(&self.out_proj)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +pub enum Mlp { + NoGating { + span1: tracing::Span, + linear1: Linear, + span2: tracing::Span, + linear2: Linear, + span: tracing::Span, + }, + Gating { + linear_in: Linear, + linear_out: Linear, + activation: candle_nn::Activation, + span: tracing::Span, + }, +} + +impl Mlp { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let d_model = cfg.d_model; + let span = tracing::span!(tracing::Level::TRACE, "mlp"); + + match cfg.gating { + None => { + let span1 = tracing::span!(tracing::Level::TRACE, "lin1"); + let span2 = tracing::span!(tracing::Level::TRACE, "lin2"); + let linear1 = linear(d_model, cfg.dim_feedforward, cfg.bias_ff, vb.pp("mlp.fc1"))?; + let linear2 = linear(cfg.dim_feedforward, d_model, cfg.bias_ff, vb.pp("mlp.fc2"))?; + Ok(Self::NoGating { + linear1, + linear2, + span, + span1, + span2, + }) + } + Some(activation) => { + let vb = vb.pp("gating"); + let hidden = if cfg.dim_feedforward == 4 * d_model { + 11 * d_model / 4 + } else { + 2 * cfg.dim_feedforward / 3 + }; + // TODO: Maybe use bias_ff here? + let linear_in = linear(d_model, 2 * hidden, false, vb.pp("linear_in"))?; + let linear_out = linear(hidden, d_model, false, vb.pp("linear_out"))?; + Ok(Self::Gating { + linear_in, + linear_out, + activation, + span, + }) + } + } + } +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + match self { + Self::NoGating { + linear1, + linear2, + span, + span1, + span2, + } => { + let _enter = span.enter(); + let xs = { + let _enter = span1.enter(); + xs.apply(linear1)? + }; + let xs = xs.gelu_erf()?; + { + let _enter = span2.enter(); + xs.apply(linear2) + } + } + Self::Gating { + linear_in, + linear_out, + activation, + span, + } => { + let _enter = span.enter(); + let xs = xs.apply(linear_in)?; + let (b, t, _) = xs.dims3()?; + let xs = xs.reshape((b, t, 2, ()))?; + let xs = (xs.i((.., .., 0))?.apply(activation)? * xs.i((.., .., 1))?)?; + xs.apply(linear_out) + } + } + } +} + +#[derive(Debug, Clone)] +pub struct RmsNorm { + pub(crate) alpha: Tensor, + pub(crate) eps: f32, +} + +impl RmsNorm { + pub fn new(d_model: usize, eps: f32, vb: VarBuilder) -> Result { + let alpha = vb.get((1, 1, d_model), "alpha")?.reshape(d_model)?; + Ok(Self { alpha, eps }) + } +} + +impl Module for RmsNorm { + fn forward(&self, xs: &Tensor) -> Result { + candle_nn::ops::rms_norm(xs, &self.alpha, self.eps) + } +} + +#[derive(Debug, Clone)] +pub enum Norm { + LayerNorm(candle_nn::LayerNorm), + RmsNorm(RmsNorm), +} + +impl Norm { + pub fn new(d_model: usize, cfg: &Config, vb: VarBuilder) -> Result { + let norm = match cfg.norm { + super::NormType::LayerNorm => { + let norm = candle_nn::layer_norm(d_model, 1e-5, vb)?; + Self::LayerNorm(norm) + } + super::NormType::RmsNorm => { + let norm = RmsNorm::new(d_model, 1e-8, vb)?; + Self::RmsNorm(norm) + } + }; + Ok(norm) + } +} + +impl Module for Norm { + fn forward(&self, xs: &Tensor) -> Result { + match self { + Self::LayerNorm(m) => m.forward(xs), + Self::RmsNorm(m) => m.forward(xs), + } + } +} + +#[derive(Debug, Clone)] +pub struct StreamingTransformerLayer { + self_attn: StreamingMultiheadAttention, + mlp: Mlp, + norm1: Norm, + norm2: Norm, + layer_scale_1: Option, + layer_scale_2: Option, + cross_attn: Option<(candle_nn::LayerNorm, StreamingMultiheadCrossAttention)>, + norm_first: bool, + span: tracing::Span, +} + +impl StreamingTransformerLayer { + pub fn new(rope: &Option>, cfg: &Config, vb: VarBuilder) -> Result { + if cfg.use_conv_block { + candle::bail!("conv-block is not supported") + } + let d_model = cfg.d_model; + let mlp = Mlp::new(cfg, vb.clone())?; + let (norm1, norm2) = match cfg.norm { + super::NormType::LayerNorm => { + let norm1 = candle_nn::layer_norm(d_model, 1e-5, vb.pp("input_layernorm"))?; + let norm2 = + candle_nn::layer_norm(d_model, 1e-5, vb.pp("post_attention_layernorm"))?; + (Norm::LayerNorm(norm1), Norm::LayerNorm(norm2)) + } + super::NormType::RmsNorm => { + let norm1 = RmsNorm::new(d_model, 1e-8, vb.pp("input_rmsnorm"))?; + let norm2 = RmsNorm::new(d_model, 1e-8, vb.pp("post_attention_rmsnorm"))?; + (Norm::RmsNorm(norm1), Norm::RmsNorm(norm2)) + } + }; + let layer_scale_1 = match cfg.layer_scale { + None => None, + Some(ls) => { + let ls = LayerScale::new(d_model, ls, vb.pp("self_attn_layer_scale"))?; + Some(ls) + } + }; + let layer_scale_2 = match cfg.layer_scale { + None => None, + Some(ls) => { + let ls = LayerScale::new(d_model, ls, vb.pp("mlp_layer_scale"))?; + Some(ls) + } + }; + let self_attn = StreamingMultiheadAttention::new(rope, cfg, vb.pp("self_attn"))?; + let cross_attn = if cfg.cross_attention { + let norm_cross = candle_nn::layer_norm(cfg.d_model, 1e-5, vb.pp("norm_cross"))?; + let cross_attn = StreamingMultiheadCrossAttention::new(cfg, vb.pp("cross_attention"))?; + Some((norm_cross, cross_attn)) + } else { + None + }; + Ok(Self { + self_attn, + mlp, + norm1, + norm2, + layer_scale_1, + layer_scale_2, + cross_attn, + norm_first: cfg.norm_first, + span: tracing::span!(tracing::Level::TRACE, "transformer-layer"), + }) + } + + pub fn forward( + &mut self, + xs: &Tensor, + ca_src: Option<&Tensor>, + mask: Option<&Tensor>, + ) -> Result { + let _enter = self.span.enter(); + if !self.norm_first { + candle::bail!("only norm_first = true is supported") + } + let norm1 = xs.apply(&self.norm1)?; + let xs = (xs + + self + .self_attn + .forward(&norm1, mask)? + .apply(&self.layer_scale_1.as_ref())?)?; + + let xs = match (&self.cross_attn, ca_src) { + (Some((norm_cross, cross_attn)), Some(ca_src)) => { + let residual = &xs; + let xs = xs.apply(norm_cross)?; + (residual + cross_attn.forward(&xs, ca_src, None)?)? + } + _ => xs, + }; + + let xs = (&xs + + xs.apply(&self.norm2)? + .apply(&self.mlp)? + .apply(&self.layer_scale_2.as_ref()))?; + Ok(xs) + } + + pub fn reset_kv_cache(&mut self) { + self.self_attn.reset_kv_cache() + } + + pub fn set_kv_cache(&mut self, kv_cache: candle_nn::kv_cache::RotatingKvCache) { + self.self_attn.set_kv_cache(kv_cache) + } +} + +#[derive(Debug, Clone)] +pub struct StreamingTransformer { + layers: Vec, + positional_embedding: PositionalEmbedding, + max_period: usize, +} + +impl StreamingTransformer { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_l = vb.pp("layers"); + let rope = match cfg.positional_embedding { + PositionalEmbedding::Rope => { + let rope = RotaryEmbedding::new( + cfg.d_model / cfg.num_heads, + cfg.max_seq_len, + cfg.max_period as f32, + vb.device(), + )?; + Some(Arc::new(rope)) + } + PositionalEmbedding::Sin | PositionalEmbedding::None => None, + }; + let mut layers = Vec::with_capacity(cfg.num_layers); + for layer_idx in 0..cfg.num_layers { + let layer = StreamingTransformerLayer::new(&rope, cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + Ok(Self { + layers, + positional_embedding: cfg.positional_embedding, + max_period: cfg.max_period, + }) + } + + pub fn forward(&mut self, xs: &Tensor) -> Result { + self.forward_ca(xs, None) + } + + pub fn forward_ca(&mut self, xs: &Tensor, ca_src: Option<&Tensor>) -> Result { + let (_b, t, c) = xs.dims3()?; + let pos = self.layers[0].self_attn.kv_cache.current_seq_len(); + let mask = self.layers[0] + .self_attn + .kv_cache + .attn_mask(t, xs.device())?; + let mut xs = match self.positional_embedding { + PositionalEmbedding::Rope | PositionalEmbedding::None => xs.clone(), + PositionalEmbedding::Sin => { + let dev = xs.device(); + let theta = self.max_period as f32; + let half_dim = c / 2; + let positions = Tensor::arange(pos as u32, (pos + t) as u32, dev)? + .unsqueeze(1)? + .to_dtype(DType::F32)?; + let inv_freq: Vec<_> = (0..half_dim) + .map(|i| 1f32 / theta.powf(i as f32 / (half_dim - 1) as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?; + let freqs = positions.broadcast_mul(&inv_freq)?; + let pos_emb = + Tensor::cat(&[freqs.cos()?, freqs.sin()?], D::Minus1)?.to_dtype(xs.dtype())?; + xs.broadcast_add(&pos_emb)? + } + }; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, ca_src, mask.as_ref())?; + } + Ok(xs) + } + + pub fn copy_state(&mut self, from: &Self) -> Result<()> { + if self.layers.len() != from.layers.len() { + candle::bail!("cannot copy kv-caches as the transformers have different depths") + } + self.layers + .iter_mut() + .zip(from.layers.iter()) + .for_each(|(v, w)| v.set_kv_cache(w.self_attn.kv_cache.clone())); + Ok(()) + } +} + +impl StreamingModule for StreamingTransformer { + fn reset_state(&mut self) { + self.layers.iter_mut().for_each(|v| v.reset_kv_cache()) + } + + fn step(&mut self, xs: &StreamTensor) -> Result { + match xs.as_option() { + None => Ok(StreamTensor::empty()), + Some(xs) => Ok(StreamTensor::from_tensor(self.forward(xs)?)), + } + } +} + +#[derive(Debug, Clone)] +pub struct ProjectedTransformer { + transformer: StreamingTransformer, + input_proj: Option, + output_projs: Vec>, + conv_layout: bool, + span: tracing::Span, +} + +impl ProjectedTransformer { + pub fn new( + input_dim: usize, + output_dims: &[usize], + cfg: &Config, + vb: VarBuilder, + ) -> Result { + let transformer = StreamingTransformer::new(cfg, vb.clone())?; + let input_proj = if input_dim == cfg.d_model { + None + } else { + let l = linear_no_bias(input_dim, cfg.d_model, vb.pp("input_proj"))?; + Some(l) + }; + let mut output_projs = Vec::with_capacity(output_dims.len()); + let vb_o = vb.pp("output_projs"); + for (i, &output_dim) in output_dims.iter().enumerate() { + let output_proj = if output_dim == cfg.d_model { + None + } else { + let l = linear_no_bias(cfg.d_model, output_dim, vb_o.pp(i))?; + Some(l) + }; + output_projs.push(output_proj) + } + Ok(Self { + transformer, + input_proj, + output_projs, + conv_layout: cfg.conv_layout, + span: tracing::span!(tracing::Level::TRACE, "proj-transformer"), + }) + } + + pub fn forward(&mut self, xs: &Tensor) -> Result> { + let _enter = self.span.enter(); + let xs = if self.conv_layout { + xs.transpose(1, 2)? + } else { + xs.clone() + }; + let xs = xs.apply(&self.input_proj.as_ref())?; + let xs = self.transformer.forward(&xs)?; + let mut ys = Vec::with_capacity(self.output_projs.len()); + for output_proj in self.output_projs.iter() { + let ys_ = xs.apply(&output_proj.as_ref())?; + let ys_ = if self.conv_layout { + ys_.transpose(1, 2)? + } else { + ys_ + }; + ys.push(ys_) + } + Ok(ys) + } +} + +impl StreamingModule for ProjectedTransformer { + fn reset_state(&mut self) { + self.transformer.reset_state() + } + + fn step(&mut self, xs: &StreamTensor) -> Result { + let xs = xs.apply(&|x: &Tensor| { + if self.conv_layout { + x.transpose(1, 2) + } else { + Ok(x.clone()) + } + })?; + let xs = xs.apply(&self.input_proj.as_ref())?; + let xs = self.transformer.step(&xs)?; + let ys = xs.apply(&self.output_projs[0].as_ref())?; + ys.apply(&|y: &Tensor| { + if self.conv_layout { + y.transpose(1, 2) + } else { + Ok(y.clone()) + } + }) + } +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} diff --git a/patches/candle-transformers/src/models/mistral.rs b/patches/candle-transformers/src/models/mistral.rs new file mode 100644 index 0000000000..23f982e990 --- /dev/null +++ b/patches/candle-transformers/src/models/mistral.rs @@ -0,0 +1,467 @@ +//! Mixtral Model, based on the Mistral architecture +//! +//! See Mistral and Mixtral at: +//! - [Hugging Face](https://huggingface.co/docs/transformers/model_doc/mixtral) +//! - [GitHub](https://github.com/mistralai/mistral-src) +//! + +use crate::models::with_tracing::{linear_no_bias, Linear, RmsNorm}; +/// Mistral LLM, https://github.com/mistralai/mistral-src +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::{Activation, VarBuilder}; +use std::sync::Arc; + +fn default_num_attention_heads() -> usize { + 32 +} + +fn default_use_flash_attn() -> bool { + false +} + +fn default_hidden_act() -> candle_nn::Activation { + candle_nn::Activation::Silu +} + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + #[serde(default = "default_num_attention_heads")] + pub num_attention_heads: usize, + pub head_dim: Option, + pub num_key_value_heads: usize, + #[serde(default = "default_hidden_act")] + pub hidden_act: Activation, + pub max_position_embeddings: usize, + pub rms_norm_eps: f64, + pub rope_theta: f64, + pub sliding_window: Option, + #[serde(default = "default_use_flash_attn")] + pub use_flash_attn: bool, +} + +impl Config { + // https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json + pub fn config_7b_v0_1(use_flash_attn: bool) -> Self { + Self { + vocab_size: 32000, + hidden_size: 4096, + intermediate_size: 14336, + num_hidden_layers: 32, + num_attention_heads: 32, + head_dim: None, + num_key_value_heads: 8, + hidden_act: Activation::Silu, + max_position_embeddings: 32768, + rms_norm_eps: 1e-5, + rope_theta: 10_000., + sliding_window: Some(4096), + use_flash_attn, + } + } + + // https://huggingface.co/Open-Orca/Mistral-7B-OpenOrca/blob/main/config.json + // https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/blob/main/config.json + pub fn config_chat_ml(use_flash_attn: bool) -> Self { + Self { + vocab_size: 32002, + hidden_size: 4096, + intermediate_size: 14336, + num_hidden_layers: 32, + num_attention_heads: 32, + head_dim: None, + num_key_value_heads: 8, + hidden_act: Activation::Silu, + max_position_embeddings: 32768, + rms_norm_eps: 1e-5, + rope_theta: 10_000., + sliding_window: Some(4096), + use_flash_attn, + } + } + + // https://huggingface.co/amazon/MistralLite/blob/main/config.json + pub fn config_amazon_mistral_lite(use_flash_attn: bool) -> Self { + Self { + vocab_size: 32003, + hidden_size: 4096, + intermediate_size: 14336, + num_hidden_layers: 32, + num_attention_heads: 32, + head_dim: None, + num_key_value_heads: 8, + hidden_act: Activation::Silu, + max_position_embeddings: 32768, + rms_norm_eps: 1e-5, + rope_theta: 10_000., + sliding_window: Some(4096), + use_flash_attn, + } + } + + fn head_dim(&self) -> usize { + self.head_dim + .unwrap_or(self.hidden_size / self.num_attention_heads) + } +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let rope_theta = cfg.rope_theta as f32; + let dim = cfg.head_dim(); + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / rope_theta.powf(i as f32 / dim as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(DType::F32)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(DType::F32)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?.to_dtype(dtype)?, + cos: freqs.cos()?.to_dtype(dtype)?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope(q, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(k, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: Activation, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + let gate_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("gate_proj"))?; + let up_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("up_proj"))?; + let down_proj = linear_no_bias(intermediate_sz, hidden_sz, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_act, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = xs.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, + use_flash_attn: bool, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + let head_dim = cfg.head_dim(); + let q_proj = linear_no_bias(hidden_sz, num_heads * head_dim, vb.pp("q_proj"))?; + let k_proj = linear_no_bias(hidden_sz, num_kv_heads * head_dim, vb.pp("k_proj"))?; + let v_proj = linear_no_bias(hidden_sz, num_kv_heads * head_dim, vb.pp("v_proj"))?; + let o_proj = linear_no_bias(num_heads * head_dim, hidden_sz, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + rotary_emb, + kv_cache: None, + use_flash_attn: cfg.use_flash_attn, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + + let (query_states, key_states) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states, seqlen_offset)?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?; + let value_states = crate::utils::repeat_kv(value_states, self.num_kv_groups)?; + + let attn_output = if self.use_flash_attn { + // flash-attn expects (b_sz, seq_len, nheads, head_dim) + let q = query_states.transpose(1, 2)?; + let k = key_states.transpose(1, 2)?; + let v = value_states.transpose(1, 2)?; + let softmax_scale = 1f32 / (self.head_dim as f32).sqrt(); + flash_attn(&q, &k, &v, softmax_scale, q_len > 1)?.transpose(1, 2)? + } else { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + }; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.num_heads * self.head_dim))? + .apply(&self.o_proj) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: MLP, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl DecoderLayer { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let self_attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let input_layernorm = + RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let post_attention_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; + residual + xs + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: Linear, + sliding_window: Option, + device: Device, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb_m.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(rotary_emb.clone(), cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb_m.pp("norm"))?; + let lm_head = linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))?; + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + sliding_window: cfg.sliding_window, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + fn prepare_decoder_attention_mask( + &self, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + let sliding_window = self.sliding_window.unwrap_or(tgt_len + 1); + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| { + (0..tgt_len).map(move |j| { + if i < j || j + sliding_window < i { + f32::NEG_INFINITY + } else { + 0. + } + }) + }) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((1, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + pub fn embed_tokens(&self) -> &candle_nn::Embedding { + &self.embed_tokens + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let (_b_size, seq_len) = input_ids.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(seq_len, seqlen_offset)?; + Some(mask) + }; + let mut xs = self.embed_tokens.forward(input_ids)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)? + } + xs.narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head) + } + + pub fn forward_embeds( + &mut self, + xs: &Tensor, + attn_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (_b_size, seq_len, _) = xs.dims3()?; + let mut xs = xs.clone(); + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attn_mask, seqlen_offset)? + } + xs.narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +} diff --git a/patches/candle-transformers/src/models/mixformer.rs b/patches/candle-transformers/src/models/mixformer.rs new file mode 100644 index 0000000000..797d75827e --- /dev/null +++ b/patches/candle-transformers/src/models/mixformer.rs @@ -0,0 +1,465 @@ +//! MixFormer (Microsoft's Phi Architecture) +//! +//! See "Textbooks Are All You Need II: phi-1.5 technical report", Lin et al. 2023 +//! - [Arxiv](https://arxiv.org/abs/2309.05463) +//! - [GitHub](https://huggingface.co/microsoft/phi-1_5) +//! + +use crate::models::with_tracing::{linear, Embedding as E, Linear}; +/// MixFormer model. +/// https://huggingface.co/microsoft/phi-1_5 +/// https://arxiv.org/abs/2309.05463 +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{Activation, VarBuilder}; +use serde::Deserialize; + +const MAX_SEQ_LEN: usize = 4096; + +// https://huggingface.co/microsoft/phi-1_5/blob/d38e6f954ec29b96fe2cf033937dad64e279b5d9/configuration_mixformer_sequential.py +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Config { + pub(crate) vocab_size: usize, + pub(crate) n_positions: usize, + pub(crate) n_embd: usize, + pub(crate) n_layer: usize, + pub(crate) n_inner: Option, + pub(crate) n_head: usize, + pub(crate) rotary_dim: usize, + pub(crate) activation_function: Activation, + pub(crate) layer_norm_epsilon: f64, + pub(crate) tie_word_embeddings: bool, + pub(crate) pad_vocab_size_multiple: usize, +} + +impl Config { + pub fn v1() -> Self { + Self { + vocab_size: 50304, + n_positions: 2048, + n_embd: 1024, + n_layer: 20, + n_inner: None, + n_head: 16, + rotary_dim: usize::min(32, 1024 / 16), + activation_function: Activation::Gelu, + layer_norm_epsilon: 1e-5, + tie_word_embeddings: false, + pad_vocab_size_multiple: 64, + } + } + + pub fn v1_5() -> Self { + Self { + vocab_size: 51200, + n_positions: 2048, + n_embd: 2048, + n_layer: 24, + n_inner: None, + n_head: 32, + rotary_dim: usize::min(32, 2048 / 32), + activation_function: Activation::Gelu, + layer_norm_epsilon: 1e-5, + tie_word_embeddings: false, + pad_vocab_size_multiple: 64, + } + } + + pub fn v2() -> Self { + Self { + vocab_size: 51200, + n_positions: 2048, + n_embd: 2560, + n_layer: 32, + n_inner: None, + n_head: 32, + rotary_dim: usize::min(32, 2560 / 32), + activation_function: Activation::Gelu, + layer_norm_epsilon: 1e-5, + tie_word_embeddings: false, + pad_vocab_size_multiple: 64, + } + } + + // https://huggingface.co/teknium/Puffin-Phi-v2/blob/main/config.json + pub fn puffin_phi_v2() -> Self { + Self { + vocab_size: 50304, + n_positions: 2048, + n_embd: 2048, + n_layer: 24, + n_inner: None, + n_head: 32, + rotary_dim: usize::min(32, 2048 / 32), + activation_function: Activation::Gelu, + layer_norm_epsilon: 1e-5, + tie_word_embeddings: false, + pad_vocab_size_multiple: 64, + } + } + + // https://huggingface.co/teknium/Phi-Hermes-1.3B/blob/main/config.json + pub fn phi_hermes_1_3b() -> Self { + Self { + vocab_size: 50304, + n_positions: 2048, + n_embd: 2048, + n_layer: 24, + n_inner: None, + n_head: 32, + rotary_dim: usize::min(32, 2048 / 32), + activation_function: Activation::NewGelu, + layer_norm_epsilon: 1e-5, + tie_word_embeddings: false, + pad_vocab_size_multiple: 64, + } + } +} + +#[derive(Debug, Clone)] +struct Embedding { + wte: E, +} + +impl Embedding { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let wte = E::new(cfg.vocab_size, cfg.n_embd, vb.pp("wte"))?; + Ok(Self { wte }) + } +} + +impl Module for Embedding { + fn forward(&self, xs: &Tensor) -> Result { + self.wte.forward(xs) + } +} + +fn get_mask(size: usize, dtype: DType, device: &Device) -> Result { + let mask: Vec<_> = (0..size) + .flat_map(|i| (0..size).map(move |j| if j > i { f32::NEG_INFINITY } else { 0. })) + .collect(); + Tensor::from_slice(&mask, (size, size), device)?.to_dtype(dtype) +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(dim: usize, max_seq_len: usize, dtype: DType, dev: &Device) -> Result { + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / 10000f32.powf(i as f32 / dim as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(DType::F32)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?.to_dtype(dtype)?, + cos: freqs.cos()?.to_dtype(dtype)?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + qkv: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor, Tensor)> { + let (_b_size, seqlen, three, _, _headdim) = qkv.dims5()?; + if three != 3 { + candle::bail!("unexpected shape for qkv {:?}", qkv.shape()) + } + let (_rotary_seqlen, rotary_dim) = self.cos.dims2()?; + let rotary_dim = rotary_dim * 2; + let q_rot = qkv.i((.., .., 0, .., ..rotary_dim))?.contiguous()?; + let q_pass = qkv.i((.., .., 0, .., rotary_dim..))?; + let k_rot = qkv.i((.., .., 1, .., ..rotary_dim))?.contiguous()?; + let k_pass = qkv.i((.., .., 1, .., rotary_dim..))?; + let c = self.cos.narrow(0, seqlen_offset, seqlen)?; + let s = self.sin.narrow(0, seqlen_offset, seqlen)?; + let q_rot = candle_nn::rotary_emb::rope_thd(&q_rot, &c, &s)?; + let k_rot = candle_nn::rotary_emb::rope_thd(&k_rot, &c, &s)?; + let q = Tensor::cat(&[&q_rot, &q_pass], D::Minus1)?; + let k = Tensor::cat(&[&k_rot, &k_pass], D::Minus1)?; + let v = qkv.i((.., .., 2))?; + Ok((q, k, v)) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + fc1: Linear, + fc2: Linear, + act: Activation, + span: tracing::Span, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let n_inner = cfg.n_inner.unwrap_or(4 * cfg.n_embd); + let fc1 = linear(cfg.n_embd, n_inner, vb.pp("fc1"))?; + let fc2 = linear(n_inner, cfg.n_embd, vb.pp("fc2"))?; + Ok(Self { + fc1, + fc2, + act: cfg.activation_function, + span: tracing::span!(tracing::Level::TRACE, "mlp"), + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + xs.apply(&self.fc1)?.apply(&self.act)?.apply(&self.fc2) + } +} + +#[derive(Debug, Clone)] +struct CausalLMHead { + ln: candle_nn::LayerNorm, + linear: Linear, +} + +impl CausalLMHead { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let ln = candle_nn::layer_norm(cfg.n_embd, cfg.layer_norm_epsilon, vb.pp("ln"))?; + let linear = linear(cfg.n_embd, cfg.vocab_size, vb.pp("linear"))?; + Ok(Self { ln, linear }) + } +} + +impl Module for CausalLMHead { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.ln)? + .apply(&self.linear)? + .to_dtype(DType::F32) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MHA { + wqkv: Linear, + out_proj: Linear, + rotary_emb: RotaryEmbedding, + kv_cache: Option<(Tensor, Tensor)>, + head_dim: usize, + softmax_scale: f64, + span: tracing::Span, + span_rope: tracing::Span, + span_mask: tracing::Span, + span_softmax: tracing::Span, +} + +impl MHA { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let head_dim = cfg.n_embd / cfg.n_head; + let op_size = cfg.n_embd; + let wqkv = linear(cfg.n_embd, 3 * op_size, vb.pp("Wqkv"))?; + let out_proj = linear(op_size, cfg.n_embd, vb.pp("out_proj"))?; + let rotary_emb = + RotaryEmbedding::new(cfg.rotary_dim, MAX_SEQ_LEN, vb.dtype(), vb.device())?; + let softmax_scale = 1f64 / (head_dim as f64).sqrt(); + Ok(Self { + wqkv, + out_proj, + head_dim, + kv_cache: None, + rotary_emb, + softmax_scale, + span: tracing::span!(tracing::Level::TRACE, "mha"), + span_rope: tracing::span!(tracing::Level::TRACE, "rope"), + span_mask: tracing::span!(tracing::Level::TRACE, "mask"), + span_softmax: tracing::span!(tracing::Level::TRACE, "softmax"), + }) + } + + fn forward(&mut self, xs: &Tensor, mask: Option<&Tensor>) -> Result { + let _enter = self.span.enter(); + let (b_size, seq_len, _n_embd) = xs.dims3()?; + let qkv = self + .wqkv + .forward(xs)? + .reshape((b_size, seq_len, 3, (), self.head_dim))?; + let seqlen_offset = match &self.kv_cache { + None => 0, + Some((prev_k, _)) => prev_k.dim(1)?, + }; + // In the python implementation, a single tensor is returned with the third axis of size 3. + let (q, k, v) = { + let _enter = self.span_rope.enter(); + self.rotary_emb.apply_rotary_emb_qkv(&qkv, seqlen_offset)? + }; + let (k, v) = match &self.kv_cache { + None => (k, v), + Some((prev_k, prev_v)) => { + let k = Tensor::cat(&[prev_k, &k], 1)?; + let v = Tensor::cat(&[prev_v, &v], 1)?; + (k, v) + } + }; + self.kv_cache = Some((k.clone(), v.clone())); + // scores = torch.einsum('bthd,bshd->bhts', q, k * softmax_scale) + let q = q.transpose(1, 2)?.flatten_to(1)?; // b*h, t, d + let k = k.transpose(1, 2)?.flatten_to(1)?; // b*h, s, d + let v = v.transpose(1, 2)?.flatten_to(1)?; // b*h, s, d + let attn_weights = (q.matmul(&k.t()?)? * self.softmax_scale)?; // b*h, t, s + + // causal_mask = torch.triu(torch.full((seqlen_q, seqlen_k), -10000.0, device=scores.device), 1) + // scores = scores + causal_mask.to(dtype=scores.dtype) + let attn_weights = match mask { + None => attn_weights, + Some(mask) => { + let _enter = self.span_mask.enter(); + attn_weights.broadcast_add(mask)? + } + }; + let attn_weights = { + let _enter = self.span_softmax.enter(); + candle_nn::ops::softmax_last_dim(&attn_weights)? + }; + + // output = torch.einsum('bhts,bshd->bthd', attention_drop, v) + // attn_weights: b*h,t,s, v: b*h,s,d + let attn_output = attn_weights.matmul(&v)?; + // b*h,t,d + let attn_output = attn_output + .reshape((b_size, (), seq_len, self.head_dim))? + .transpose(1, 2)? + .flatten_from(D::Minus2)?; + attn_output.apply(&self.out_proj) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +struct ParallelBlock { + ln: candle_nn::LayerNorm, + mixer: MHA, + mlp: MLP, + span: tracing::Span, +} + +impl ParallelBlock { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let ln = candle_nn::layer_norm(cfg.n_embd, cfg.layer_norm_epsilon, vb.pp("ln"))?; + let mixer = MHA::new(cfg, vb.pp("mixer"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + Ok(Self { + ln, + mixer, + mlp, + span: tracing::span!(tracing::Level::TRACE, "block"), + }) + } + + fn forward(&mut self, xs: &Tensor, mask: Option<&Tensor>) -> Result { + let _enter = self.span.enter(); + let residual = xs; + let xs = xs.apply(&self.ln)?; + let attn_outputs = self.mixer.forward(&xs, mask)?; + let feed_forward_hidden_states = self.mlp.forward(&xs)?; + attn_outputs + feed_forward_hidden_states + residual + } + + fn clear_kv_cache(&mut self) { + self.mixer.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct MixFormerSequentialForCausalLM { + embedding: Embedding, + blocks: Vec, + head: CausalLMHead, + span: tracing::Span, +} + +impl MixFormerSequentialForCausalLM { + pub fn new_v2(cfg: &Config, vb: VarBuilder) -> Result { + let vb_head = vb.pp("lm_head"); + let vb = vb.pp("transformer"); + let embedding = Embedding::new(cfg, vb.pp("embd"))?; + let mut blocks = Vec::new(); + for i in 0..cfg.n_layer { + let block = ParallelBlock::new(cfg, vb.pp("h").pp(i))?; + blocks.push(block) + } + let head = CausalLMHead::new(cfg, vb_head)?; + Ok(Self { + embedding, + blocks, + head, + span: tracing::span!(tracing::Level::TRACE, "mixformer"), + }) + } + + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb = vb.pp("layers"); + let embedding = Embedding::new(cfg, vb.pp(0))?; + let mut blocks = Vec::new(); + for i in 0..cfg.n_layer { + let block = ParallelBlock::new(cfg, vb.pp(i + 1))?; + blocks.push(block) + } + let head = CausalLMHead::new(cfg, vb.pp(cfg.n_layer + 1))?; + Ok(Self { + embedding, + blocks, + head, + span: tracing::span!(tracing::Level::TRACE, "mixformer"), + }) + } + + pub fn forward(&mut self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let (_b_size, seq_len) = xs.dims2()?; + let mut xs = xs.apply(&self.embedding)?; + let mask = if seq_len <= 1 { + None + } else { + Some(get_mask(seq_len, xs.dtype(), xs.device())?) + }; + for block in self.blocks.iter_mut() { + xs = block.forward(&xs, mask.as_ref())? + } + xs.narrow(1, seq_len - 1, 1)?.apply(&self.head)?.squeeze(1) + } + + pub fn forward_with_img( + &mut self, + bos_token: &Tensor, + xs: &Tensor, + img_embeds: &Tensor, + ) -> Result { + let _enter = self.span.enter(); + let xs = xs.apply(&self.embedding)?; + let bos_token = bos_token.apply(&self.embedding)?; + // Python implementation sequence order is + // https://github.com/vikhyat/moondream/blob/a9d788a20d1543fb1479edc54106e88cff7759d3/moondream/moondream.py#L43-L56 + let mut xs = Tensor::cat(&[bos_token, img_embeds.clone(), xs], 1)?; + let (_b_size, seq_len, _embds) = xs.dims3()?; + let mask = Some(get_mask(seq_len, xs.dtype(), xs.device())?); + for block in self.blocks.iter_mut() { + xs = block.forward(&xs, mask.as_ref())? + } + let xs = xs + .narrow(1, seq_len - 1, 1)? + .apply(&self.head)? + .squeeze(1)?; + Ok(xs) + } + + pub fn clear_kv_cache(&mut self) { + self.blocks.iter_mut().for_each(|b| b.clear_kv_cache()) + } +} diff --git a/patches/candle-transformers/src/models/mixtral.rs b/patches/candle-transformers/src/models/mixtral.rs new file mode 100644 index 0000000000..70115e10a3 --- /dev/null +++ b/patches/candle-transformers/src/models/mixtral.rs @@ -0,0 +1,483 @@ +//! Mixtral Model, a sparse mixture of expert model based on the Mistral architecture +//! +//! See Mixtral model details at: +//! - [Hugging Face](https://huggingface.co/docs/transformers/model_doc/mixtral) +//! - [Mixtral-8x7B Blog Post](https://mistral.ai/news/mixtral-of-experts/) +//! +//! The model uses a mixture of experts architecture with: +//! - 8 experts per layer +//! - Top 2 expert routing +//! - Sliding window attention +//! - RoPE embeddings +//! +//! References: +//! - [Hugging Face Implementation](https://github.com/huggingface/transformers/blob/main/src/transformers/models/mixtral/modeling_mixtral.py) +//! - [Mixtral Blog Post](https://mistral.ai/news/mixtral-of-experts/) +//! + +use crate::models::with_tracing::{linear_no_bias, Linear, RmsNorm}; +/// Mixtral Model +/// https://github.com/huggingface/transformers/blob/main/src/transformers/models/mixtral/modeling_mixtral.py +/// https://mistral.ai/news/mixtral-of-experts/ +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::{Activation, VarBuilder}; +use serde::Deserialize; +use std::sync::Arc; + +/// https://github.com/huggingface/transformers/blob/1a585c1222a56bcaecc070966d558d4a9d862e83/src/transformers/models/mixtral/configuration_mixtral.py#L113 +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Config { + pub(crate) vocab_size: usize, + pub(crate) hidden_size: usize, + pub(crate) intermediate_size: usize, + pub(crate) num_hidden_layers: usize, + pub(crate) num_attention_heads: usize, + pub(crate) num_key_value_heads: usize, + pub(crate) hidden_act: Activation, + pub(crate) max_position_embeddings: usize, + pub(crate) rms_norm_eps: f64, + pub(crate) rope_theta: f64, + pub(crate) sliding_window: usize, + pub(crate) num_experts_per_tok: usize, + pub(crate) num_local_experts: usize, + pub(crate) use_flash_attn: bool, +} + +impl Config { + /// https://huggingface.co/mistralai/Mixtral-8x7B-v0.1/blob/main/config.json + pub fn v0_1_8x7b(use_flash_attn: bool) -> Self { + Self { + vocab_size: 32000, + hidden_size: 4096, + intermediate_size: 14336, + num_hidden_layers: 32, + num_attention_heads: 32, + num_key_value_heads: 8, + hidden_act: Activation::Silu, + max_position_embeddings: 32768, + rms_norm_eps: 1e-5, + rope_theta: 1e6, + sliding_window: 4096, + num_experts_per_tok: 2, + num_local_experts: 8, + use_flash_attn, + } + } +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +fn rotate_half(xs: &Tensor) -> Result { + let last_dim = xs.dim(D::Minus1)?; + let xs1 = xs.narrow(D::Minus1, 0, last_dim / 2)?; + let xs2 = xs.narrow(D::Minus1, last_dim / 2, last_dim - last_dim / 2)?; + Tensor::cat(&[&xs2.neg()?, &xs1], D::Minus1) +} + +impl RotaryEmbedding { + fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let dim = cfg.hidden_size / cfg.num_attention_heads; + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / (cfg.rope_theta as f32).powf(i as f32 / dim as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + let freqs = Tensor::cat(&[&freqs, &freqs], D::Minus1)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let cos = cos.unsqueeze(0)?.unsqueeze(0)?; // (1, 1, seq_len, dim) + let sin = sin.unsqueeze(0)?.unsqueeze(0)?; // (1, 1, seq_len, dim) + let q_embed = (q.broadcast_mul(&cos)? + rotate_half(q)?.broadcast_mul(&sin))?; + let k_embed = (k.broadcast_mul(&cos)? + rotate_half(k)?.broadcast_mul(&sin))?; + Ok((q_embed, k_embed)) + } +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + hidden_size: usize, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, + use_flash_attn: bool, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + let head_dim = hidden_sz / num_heads; + let q_proj = linear_no_bias(hidden_sz, num_heads * head_dim, vb.pp("q_proj"))?; + let k_proj = linear_no_bias(hidden_sz, num_kv_heads * head_dim, vb.pp("k_proj"))?; + let v_proj = linear_no_bias(hidden_sz, num_kv_heads * head_dim, vb.pp("v_proj"))?; + let o_proj = linear_no_bias(num_heads * head_dim, hidden_sz, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + hidden_size: hidden_sz, + rotary_emb, + kv_cache: None, + use_flash_attn: cfg.use_flash_attn, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let (query_states, key_states) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states, seqlen_offset)?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?; + let value_states = crate::utils::repeat_kv(value_states, self.num_kv_groups)?; + + let attn_output = if self.use_flash_attn { + // flash-attn expects (b_sz, seq_len, nheads, head_dim) + let q = query_states.transpose(1, 2)?; + let k = key_states.transpose(1, 2)?; + let v = value_states.transpose(1, 2)?; + let softmax_scale = 1f32 / (self.head_dim as f32).sqrt(); + flash_attn(&q, &k, &v, softmax_scale, q_len > 1)?.transpose(1, 2)? + } else { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + }; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.hidden_size))? + .apply(&self.o_proj) + } +} + +#[derive(Debug, Clone)] +struct BlockSparseTop2MLP { + w1: Linear, + w2: Linear, + w3: Linear, + act_fn: Activation, +} + +impl BlockSparseTop2MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + let w1 = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("w1"))?; + let w2 = linear_no_bias(intermediate_sz, hidden_sz, vb.pp("w2"))?; + let w3 = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("w3"))?; + Ok(Self { + w1, + w2, + w3, + act_fn: cfg.hidden_act, + }) + } +} + +impl Module for BlockSparseTop2MLP { + fn forward(&self, xs: &Tensor) -> Result { + let lhs = xs.apply(&self.w1)?.apply(&self.act_fn)?; + let rhs = xs.apply(&self.w3)?; + (lhs * rhs)?.apply(&self.w2) + } +} + +#[derive(Debug, Clone)] +struct SparseMoeBlock { + gate: Linear, + experts: Vec, + num_experts_per_tok: usize, +} + +impl SparseMoeBlock { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let gate = linear_no_bias(cfg.hidden_size, cfg.num_local_experts, vb.pp("gate"))?; + let mut experts = Vec::with_capacity(cfg.num_local_experts); + let vb = vb.pp("experts"); + for idx in 0..cfg.num_local_experts { + let expert = BlockSparseTop2MLP::new(cfg, vb.pp(idx))?; + experts.push(expert) + } + Ok(SparseMoeBlock { + gate, + experts, + num_experts_per_tok: cfg.num_experts_per_tok, + }) + } +} + +impl Module for SparseMoeBlock { + fn forward(&self, xs: &Tensor) -> Result { + let (b_size, seq_len, hidden_dim) = xs.dims3()?; + let xs = xs.reshape(((), hidden_dim))?; + let router_logits = xs.apply(&self.gate)?; + 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![]; self.experts.len()]; + let mut selected_rws = vec![vec![]; self.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(self.num_experts_per_tok) { + 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(self.num_experts_per_tok) { + 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 self.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) + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + block_sparse_moe: SparseMoeBlock, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl DecoderLayer { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let self_attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let block_sparse_moe = SparseMoeBlock::new(cfg, vb.pp("block_sparse_moe"))?; + let input_layernorm = + RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let post_attention_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + block_sparse_moe, + input_layernorm, + post_attention_layernorm, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs + .apply(&self.post_attention_layernorm)? + .apply(&self.block_sparse_moe)?; + residual + xs + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: Linear, + sliding_window: usize, + device: Device, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb_m.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(rotary_emb.clone(), cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb_m.pp("norm"))?; + let lm_head = linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))?; + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + sliding_window: cfg.sliding_window, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + // Sliding window mask? + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| { + (0..tgt_len).map(move |j| { + if i < j || j + self.sliding_window < i { + f32::NEG_INFINITY + } else { + 0. + } + }) + }) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let (b_size, seq_len) = input_ids.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(b_size, seq_len, seqlen_offset)?; + Some(mask) + }; + let mut xs = self.embed_tokens.forward(input_ids)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)? + } + xs.narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head) + } +} diff --git a/patches/candle-transformers/src/models/mmdit/blocks.rs b/patches/candle-transformers/src/models/mmdit/blocks.rs new file mode 100644 index 0000000000..912e249835 --- /dev/null +++ b/patches/candle-transformers/src/models/mmdit/blocks.rs @@ -0,0 +1,497 @@ +use candle::{Module, Result, Tensor, D}; +use candle_nn as nn; + +use super::projections::{AttnProjections, Mlp, Qkv, QkvOnlyAttnProjections}; + +pub struct ModulateIntermediates { + gate_msa: Tensor, + shift_mlp: Tensor, + scale_mlp: Tensor, + gate_mlp: Tensor, +} + +pub struct DiTBlock { + norm1: LayerNormNoAffine, + attn: AttnProjections, + norm2: LayerNormNoAffine, + mlp: Mlp, + ada_ln_modulation: nn::Sequential, +} + +pub struct LayerNormNoAffine { + eps: f64, +} + +impl LayerNormNoAffine { + pub fn new(eps: f64) -> Self { + Self { eps } + } +} + +impl Module for LayerNormNoAffine { + fn forward(&self, x: &Tensor) -> Result { + nn::LayerNorm::new_no_bias(Tensor::ones_like(x)?, self.eps).forward(x) + } +} + +impl DiTBlock { + pub fn new(hidden_size: usize, num_heads: usize, vb: nn::VarBuilder) -> Result { + let norm1 = LayerNormNoAffine::new(1e-6); + let attn = AttnProjections::new(hidden_size, num_heads, vb.pp("attn"))?; + let norm2 = LayerNormNoAffine::new(1e-6); + let mlp_ratio = 4; + let mlp = Mlp::new(hidden_size, hidden_size * mlp_ratio, vb.pp("mlp"))?; + let n_mods = 6; + let ada_ln_modulation = nn::seq().add(nn::Activation::Silu).add(nn::linear( + hidden_size, + n_mods * hidden_size, + vb.pp("adaLN_modulation.1"), + )?); + + Ok(Self { + norm1, + attn, + norm2, + mlp, + ada_ln_modulation, + }) + } + + pub fn pre_attention(&self, x: &Tensor, c: &Tensor) -> Result<(Qkv, ModulateIntermediates)> { + let modulation = self.ada_ln_modulation.forward(c)?; + let chunks = modulation.chunk(6, D::Minus1)?; + let (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp) = ( + chunks[0].clone(), + chunks[1].clone(), + chunks[2].clone(), + chunks[3].clone(), + chunks[4].clone(), + chunks[5].clone(), + ); + + let norm_x = self.norm1.forward(x)?; + let modulated_x = modulate(&norm_x, &shift_msa, &scale_msa)?; + let qkv = self.attn.pre_attention(&modulated_x)?; + + Ok(( + qkv, + ModulateIntermediates { + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + }, + )) + } + + pub fn post_attention( + &self, + attn: &Tensor, + x: &Tensor, + mod_interm: &ModulateIntermediates, + ) -> Result { + let attn_out = self.attn.post_attention(attn)?; + let x = x.add(&attn_out.broadcast_mul(&mod_interm.gate_msa.unsqueeze(1)?)?)?; + + let norm_x = self.norm2.forward(&x)?; + let modulated_x = modulate(&norm_x, &mod_interm.shift_mlp, &mod_interm.scale_mlp)?; + let mlp_out = self.mlp.forward(&modulated_x)?; + let x = x.add(&mlp_out.broadcast_mul(&mod_interm.gate_mlp.unsqueeze(1)?)?)?; + + Ok(x) + } +} + +pub struct SelfAttnModulateIntermediates { + gate_msa: Tensor, + shift_mlp: Tensor, + scale_mlp: Tensor, + gate_mlp: Tensor, + gate_msa2: Tensor, +} + +pub struct SelfAttnDiTBlock { + norm1: LayerNormNoAffine, + attn: AttnProjections, + attn2: AttnProjections, + norm2: LayerNormNoAffine, + mlp: Mlp, + ada_ln_modulation: nn::Sequential, +} + +impl SelfAttnDiTBlock { + pub fn new(hidden_size: usize, num_heads: usize, vb: nn::VarBuilder) -> Result { + let norm1 = LayerNormNoAffine::new(1e-6); + let attn = AttnProjections::new(hidden_size, num_heads, vb.pp("attn"))?; + let attn2 = AttnProjections::new(hidden_size, num_heads, vb.pp("attn2"))?; + let norm2 = LayerNormNoAffine::new(1e-6); + let mlp_ratio = 4; + let mlp = Mlp::new(hidden_size, hidden_size * mlp_ratio, vb.pp("mlp"))?; + let n_mods = 9; + let ada_ln_modulation = nn::seq().add(nn::Activation::Silu).add(nn::linear( + hidden_size, + n_mods * hidden_size, + vb.pp("adaLN_modulation.1"), + )?); + + Ok(Self { + norm1, + attn, + attn2, + norm2, + mlp, + ada_ln_modulation, + }) + } + + pub fn pre_attention( + &self, + x: &Tensor, + c: &Tensor, + ) -> Result<(Qkv, Qkv, SelfAttnModulateIntermediates)> { + let modulation = self.ada_ln_modulation.forward(c)?; + let chunks = modulation.chunk(9, D::Minus1)?; + let ( + shift_msa, + scale_msa, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + shift_msa2, + scale_msa2, + gate_msa2, + ) = ( + chunks[0].clone(), + chunks[1].clone(), + chunks[2].clone(), + chunks[3].clone(), + chunks[4].clone(), + chunks[5].clone(), + chunks[6].clone(), + chunks[7].clone(), + chunks[8].clone(), + ); + + let norm_x = self.norm1.forward(x)?; + let modulated_x = modulate(&norm_x, &shift_msa, &scale_msa)?; + let qkv = self.attn.pre_attention(&modulated_x)?; + + let modulated_x2 = modulate(&norm_x, &shift_msa2, &scale_msa2)?; + let qkv2 = self.attn2.pre_attention(&modulated_x2)?; + + Ok(( + qkv, + qkv2, + SelfAttnModulateIntermediates { + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + gate_msa2, + }, + )) + } + + pub fn post_attention( + &self, + attn: &Tensor, + attn2: &Tensor, + x: &Tensor, + mod_interm: &SelfAttnModulateIntermediates, + ) -> Result { + let attn_out = self.attn.post_attention(attn)?; + let x = x.add(&attn_out.broadcast_mul(&mod_interm.gate_msa.unsqueeze(1)?)?)?; + let attn_out2 = self.attn2.post_attention(attn2)?; + let x = x.add(&attn_out2.broadcast_mul(&mod_interm.gate_msa2.unsqueeze(1)?)?)?; + + let norm_x = self.norm2.forward(&x)?; + let modulated_x = modulate(&norm_x, &mod_interm.shift_mlp, &mod_interm.scale_mlp)?; + let mlp_out = self.mlp.forward(&modulated_x)?; + let x = x.add(&mlp_out.broadcast_mul(&mod_interm.gate_mlp.unsqueeze(1)?)?)?; + Ok(x) + } +} + +pub struct QkvOnlyDiTBlock { + norm1: LayerNormNoAffine, + attn: QkvOnlyAttnProjections, + ada_ln_modulation: nn::Sequential, +} + +impl QkvOnlyDiTBlock { + pub fn new(hidden_size: usize, num_heads: usize, vb: nn::VarBuilder) -> Result { + let norm1 = LayerNormNoAffine::new(1e-6); + let attn = QkvOnlyAttnProjections::new(hidden_size, num_heads, vb.pp("attn"))?; + let n_mods = 2; + let ada_ln_modulation = nn::seq().add(nn::Activation::Silu).add(nn::linear( + hidden_size, + n_mods * hidden_size, + vb.pp("adaLN_modulation.1"), + )?); + + Ok(Self { + norm1, + attn, + ada_ln_modulation, + }) + } + + pub fn pre_attention(&self, x: &Tensor, c: &Tensor) -> Result { + let modulation = self.ada_ln_modulation.forward(c)?; + let chunks = modulation.chunk(2, D::Minus1)?; + let (shift_msa, scale_msa) = (chunks[0].clone(), chunks[1].clone()); + + let norm_x = self.norm1.forward(x)?; + let modulated_x = modulate(&norm_x, &shift_msa, &scale_msa)?; + self.attn.pre_attention(&modulated_x) + } +} + +pub struct FinalLayer { + norm_final: LayerNormNoAffine, + linear: nn::Linear, + ada_ln_modulation: nn::Sequential, +} + +impl FinalLayer { + pub fn new( + hidden_size: usize, + patch_size: usize, + out_channels: usize, + vb: nn::VarBuilder, + ) -> Result { + let norm_final = LayerNormNoAffine::new(1e-6); + let linear = nn::linear( + hidden_size, + patch_size * patch_size * out_channels, + vb.pp("linear"), + )?; + let ada_ln_modulation = nn::seq().add(nn::Activation::Silu).add(nn::linear( + hidden_size, + 2 * hidden_size, + vb.pp("adaLN_modulation.1"), + )?); + + Ok(Self { + norm_final, + linear, + ada_ln_modulation, + }) + } + + pub fn forward(&self, x: &Tensor, c: &Tensor) -> Result { + let modulation = self.ada_ln_modulation.forward(c)?; + let chunks = modulation.chunk(2, D::Minus1)?; + let (shift, scale) = (chunks[0].clone(), chunks[1].clone()); + + let norm_x = self.norm_final.forward(x)?; + let modulated_x = modulate(&norm_x, &shift, &scale)?; + let output = self.linear.forward(&modulated_x)?; + + Ok(output) + } +} + +fn modulate(x: &Tensor, shift: &Tensor, scale: &Tensor) -> Result { + let shift = shift.unsqueeze(1)?; + let scale = scale.unsqueeze(1)?; + let scale_plus_one = scale.add(&Tensor::ones_like(&scale)?)?; + shift.broadcast_add(&x.broadcast_mul(&scale_plus_one)?) +} + +pub trait JointBlock { + fn forward(&self, context: &Tensor, x: &Tensor, c: &Tensor) -> Result<(Tensor, Tensor)>; +} + +pub struct MMDiTJointBlock { + x_block: DiTBlock, + context_block: DiTBlock, + num_heads: usize, + use_flash_attn: bool, +} + +impl MMDiTJointBlock { + pub fn new( + hidden_size: usize, + num_heads: usize, + use_flash_attn: bool, + vb: nn::VarBuilder, + ) -> Result { + let x_block = DiTBlock::new(hidden_size, num_heads, vb.pp("x_block"))?; + let context_block = DiTBlock::new(hidden_size, num_heads, vb.pp("context_block"))?; + + Ok(Self { + x_block, + context_block, + num_heads, + use_flash_attn, + }) + } +} + +impl JointBlock for MMDiTJointBlock { + fn forward(&self, context: &Tensor, x: &Tensor, c: &Tensor) -> Result<(Tensor, Tensor)> { + let (context_qkv, context_interm) = self.context_block.pre_attention(context, c)?; + let (x_qkv, x_interm) = self.x_block.pre_attention(x, c)?; + let (context_attn, x_attn) = + joint_attn(&context_qkv, &x_qkv, self.num_heads, self.use_flash_attn)?; + let context_out = + self.context_block + .post_attention(&context_attn, context, &context_interm)?; + let x_out = self.x_block.post_attention(&x_attn, x, &x_interm)?; + Ok((context_out, x_out)) + } +} + +pub struct MMDiTXJointBlock { + x_block: SelfAttnDiTBlock, + context_block: DiTBlock, + num_heads: usize, + use_flash_attn: bool, +} + +impl MMDiTXJointBlock { + pub fn new( + hidden_size: usize, + num_heads: usize, + use_flash_attn: bool, + vb: nn::VarBuilder, + ) -> Result { + let x_block = SelfAttnDiTBlock::new(hidden_size, num_heads, vb.pp("x_block"))?; + let context_block = DiTBlock::new(hidden_size, num_heads, vb.pp("context_block"))?; + + Ok(Self { + x_block, + context_block, + num_heads, + use_flash_attn, + }) + } +} + +impl JointBlock for MMDiTXJointBlock { + fn forward(&self, context: &Tensor, x: &Tensor, c: &Tensor) -> Result<(Tensor, Tensor)> { + let (context_qkv, context_interm) = self.context_block.pre_attention(context, c)?; + let (x_qkv, x_qkv2, x_interm) = self.x_block.pre_attention(x, c)?; + let (context_attn, x_attn) = + joint_attn(&context_qkv, &x_qkv, self.num_heads, self.use_flash_attn)?; + let x_attn2 = attn(&x_qkv2, self.num_heads, self.use_flash_attn)?; + let context_out = + self.context_block + .post_attention(&context_attn, context, &context_interm)?; + let x_out = self + .x_block + .post_attention(&x_attn, &x_attn2, x, &x_interm)?; + Ok((context_out, x_out)) + } +} + +pub struct ContextQkvOnlyJointBlock { + x_block: DiTBlock, + context_block: QkvOnlyDiTBlock, + num_heads: usize, + use_flash_attn: bool, +} + +impl ContextQkvOnlyJointBlock { + pub fn new( + hidden_size: usize, + num_heads: usize, + use_flash_attn: bool, + vb: nn::VarBuilder, + ) -> Result { + let x_block = DiTBlock::new(hidden_size, num_heads, vb.pp("x_block"))?; + let context_block = QkvOnlyDiTBlock::new(hidden_size, num_heads, vb.pp("context_block"))?; + Ok(Self { + x_block, + context_block, + num_heads, + use_flash_attn, + }) + } + + pub fn forward(&self, context: &Tensor, x: &Tensor, c: &Tensor) -> Result { + let context_qkv = self.context_block.pre_attention(context, c)?; + let (x_qkv, x_interm) = self.x_block.pre_attention(x, c)?; + + let (_, x_attn) = joint_attn(&context_qkv, &x_qkv, self.num_heads, self.use_flash_attn)?; + + let x_out = self.x_block.post_attention(&x_attn, x, &x_interm)?; + Ok(x_out) + } +} + +// A QKV-attention that is compatible with the interface of candle_flash_attn::flash_attn +// Flash attention regards q, k, v dimensions as (batch_size, seqlen, nheads, headdim) +fn flash_compatible_attention( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, +) -> Result { + let q_dims_for_matmul = q.transpose(1, 2)?.dims().to_vec(); + let rank = q_dims_for_matmul.len(); + let q = q.transpose(1, 2)?.flatten_to(rank - 3)?; + let k = k.transpose(1, 2)?.flatten_to(rank - 3)?; + let v = v.transpose(1, 2)?.flatten_to(rank - 3)?; + let attn_weights = (q.matmul(&k.t()?)? * softmax_scale as f64)?; + let attn_scores = candle_nn::ops::softmax_last_dim(&attn_weights)?.matmul(&v)?; + attn_scores.reshape(q_dims_for_matmul)?.transpose(1, 2) +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +fn joint_attn( + context_qkv: &Qkv, + x_qkv: &Qkv, + num_heads: usize, + use_flash_attn: bool, +) -> Result<(Tensor, Tensor)> { + let qkv = Qkv { + q: Tensor::cat(&[&context_qkv.q, &x_qkv.q], 1)?, + k: Tensor::cat(&[&context_qkv.k, &x_qkv.k], 1)?, + v: Tensor::cat(&[&context_qkv.v, &x_qkv.v], 1)?, + }; + + let seqlen = qkv.q.dim(1)?; + let attn = attn(&qkv, num_heads, use_flash_attn)?; + let context_qkv_seqlen = context_qkv.q.dim(1)?; + let context_attn = attn.narrow(1, 0, context_qkv_seqlen)?; + let x_attn = attn.narrow(1, context_qkv_seqlen, seqlen - context_qkv_seqlen)?; + + Ok((context_attn, x_attn)) +} + +fn attn(qkv: &Qkv, num_heads: usize, use_flash_attn: bool) -> Result { + let batch_size = qkv.q.dim(0)?; + let seqlen = qkv.q.dim(1)?; + let qkv = Qkv { + q: qkv.q.reshape((batch_size, seqlen, num_heads, ()))?, + k: qkv.k.reshape((batch_size, seqlen, num_heads, ()))?, + v: qkv.v.clone(), + }; + + let headdim = qkv.q.dim(D::Minus1)?; + let softmax_scale = 1.0 / (headdim as f64).sqrt(); + let attn = if use_flash_attn { + flash_attn(&qkv.q, &qkv.k, &qkv.v, softmax_scale as f32, false)? + } else { + flash_compatible_attention(&qkv.q, &qkv.k, &qkv.v, softmax_scale as f32)? + }; + attn.reshape((batch_size, seqlen, ())) +} diff --git a/patches/candle-transformers/src/models/mmdit/embedding.rs b/patches/candle-transformers/src/models/mmdit/embedding.rs new file mode 100644 index 0000000000..eb88f8c3d7 --- /dev/null +++ b/patches/candle-transformers/src/models/mmdit/embedding.rs @@ -0,0 +1,197 @@ +use candle::{bail, DType, Module, Result, Tensor}; +use candle_nn as nn; + +pub struct PatchEmbedder { + proj: nn::Conv2d, +} + +impl PatchEmbedder { + pub fn new( + patch_size: usize, + in_channels: usize, + embed_dim: usize, + vb: nn::VarBuilder, + ) -> Result { + let proj = nn::conv2d( + in_channels, + embed_dim, + patch_size, + nn::Conv2dConfig { + stride: patch_size, + ..Default::default() + }, + vb.pp("proj"), + )?; + + Ok(Self { proj }) + } +} + +impl Module for PatchEmbedder { + fn forward(&self, x: &Tensor) -> Result { + let x = self.proj.forward(x)?; + + // flatten spatial dim and transpose to channels last + let (b, c, h, w) = x.dims4()?; + x.reshape((b, c, h * w))?.transpose(1, 2) + } +} + +pub struct Unpatchifier { + patch_size: usize, + out_channels: usize, +} + +impl Unpatchifier { + pub fn new(patch_size: usize, out_channels: usize) -> Result { + Ok(Self { + patch_size, + out_channels, + }) + } + + pub fn unpatchify(&self, x: &Tensor, h: usize, w: usize) -> Result { + let h = (h + 1) / self.patch_size; + let w = (w + 1) / self.patch_size; + + let x = x.reshape(( + x.dim(0)?, + h, + w, + self.patch_size, + self.patch_size, + self.out_channels, + ))?; + let x = x.permute((0, 5, 1, 3, 2, 4))?; // "nhwpqc->nchpwq" + x.reshape(( + x.dim(0)?, + self.out_channels, + self.patch_size * h, + self.patch_size * w, + )) + } +} + +pub struct PositionEmbedder { + pos_embed: Tensor, + patch_size: usize, + pos_embed_max_size: usize, +} + +impl PositionEmbedder { + pub fn new( + hidden_size: usize, + patch_size: usize, + pos_embed_max_size: usize, + vb: nn::VarBuilder, + ) -> Result { + let pos_embed = vb.get( + (1, pos_embed_max_size * pos_embed_max_size, hidden_size), + "pos_embed", + )?; + Ok(Self { + pos_embed, + patch_size, + pos_embed_max_size, + }) + } + pub fn get_cropped_pos_embed(&self, h: usize, w: usize) -> Result { + let h = (h + 1) / self.patch_size; + let w = (w + 1) / self.patch_size; + + if h > self.pos_embed_max_size || w > self.pos_embed_max_size { + bail!("Input size is too large for the position embedding") + } + + let top = (self.pos_embed_max_size - h) / 2; + let left = (self.pos_embed_max_size - w) / 2; + + let pos_embed = + self.pos_embed + .reshape((1, self.pos_embed_max_size, self.pos_embed_max_size, ()))?; + let pos_embed = pos_embed.narrow(1, top, h)?.narrow(2, left, w)?; + pos_embed.reshape((1, h * w, ())) + } +} + +pub struct TimestepEmbedder { + mlp: nn::Sequential, + frequency_embedding_size: usize, +} + +impl TimestepEmbedder { + pub fn new( + hidden_size: usize, + frequency_embedding_size: usize, + vb: nn::VarBuilder, + ) -> Result { + let mlp = nn::seq() + .add(nn::linear( + frequency_embedding_size, + hidden_size, + vb.pp("mlp.0"), + )?) + .add(nn::Activation::Silu) + .add(nn::linear(hidden_size, hidden_size, vb.pp("mlp.2"))?); + + Ok(Self { + mlp, + frequency_embedding_size, + }) + } + + fn timestep_embedding(t: &Tensor, dim: usize, max_period: f64) -> Result { + if !dim.is_multiple_of(2) { + bail!("Embedding dimension must be even") + } + + if t.dtype() != DType::F32 && t.dtype() != DType::F64 { + bail!("Input tensor must be floating point") + } + + let half = dim / 2; + let freqs = Tensor::arange(0f32, half as f32, t.device())? + .to_dtype(candle::DType::F32)? + .mul(&Tensor::full( + (-f64::ln(max_period) / half as f64) as f32, + half, + t.device(), + )?)? + .exp()?; + + let args = t + .unsqueeze(1)? + .to_dtype(candle::DType::F32)? + .matmul(&freqs.unsqueeze(0)?)?; + let embedding = Tensor::cat(&[args.cos()?, args.sin()?], 1)?; + embedding.to_dtype(candle::DType::F16) + } +} + +impl Module for TimestepEmbedder { + fn forward(&self, t: &Tensor) -> Result { + let t_freq = Self::timestep_embedding(t, self.frequency_embedding_size, 10000.0)?; + self.mlp.forward(&t_freq) + } +} + +pub struct VectorEmbedder { + mlp: nn::Sequential, +} + +impl VectorEmbedder { + pub fn new(input_dim: usize, hidden_size: usize, vb: nn::VarBuilder) -> Result { + let mlp = nn::seq() + .add(nn::linear(input_dim, hidden_size, vb.pp("mlp.0"))?) + .add(nn::Activation::Silu) + .add(nn::linear(hidden_size, hidden_size, vb.pp("mlp.2"))?); + + Ok(Self { mlp }) + } +} + +impl Module for VectorEmbedder { + fn forward(&self, x: &Tensor) -> Result { + self.mlp.forward(x) + } +} diff --git a/patches/candle-transformers/src/models/mmdit/mod.rs b/patches/candle-transformers/src/models/mmdit/mod.rs new file mode 100644 index 0000000000..88e73e1e3d --- /dev/null +++ b/patches/candle-transformers/src/models/mmdit/mod.rs @@ -0,0 +1,19 @@ +//! Mix of Multi-scale Dilated and Traditional Convolutions +//! +//! Mix of Multi-scale Dilated and Traditional Convolutions (MMDiT) is an architecture +//! introduced for Stable Diffusion 3, with the MMDiT-X variant used in Stable Diffusion 3.5. +//! +//! - 📝 [Research Paper](https://arxiv.org/abs/2403.03206) +//! - 💻 ComfyUI [reference implementation](https://github.com/comfyanonymous/ComfyUI/blob/78e133d0415784924cd2674e2ee48f3eeca8a2aa/comfy/ldm/modules/diffusionmodules/mmdit.py) +//! - 💻 Stability-AI [MMDiT-X implementation](https://github.com/Stability-AI/sd3.5/blob/4e484e05308d83fb77ae6f680028e6c313f9da54/mmditx.py) + +//! - ⚡ [Interactive Wasm Example](https://huggingface.co/spaces/radames/Candle-BLIP-Image-Captioning) +//! - 💻 [GH Link](https://github.com/salesforce/BLIP) +//! - 🤗 [HF Link](https://huggingface.co/Salesforce/blip-image-captioning-base) +//! - 📝 [Paper](https://arxiv.org/abs/2201.12086) +//! + +pub mod blocks; +pub mod embedding; +pub mod model; +pub mod projections; diff --git a/patches/candle-transformers/src/models/mmdit/model.rs b/patches/candle-transformers/src/models/mmdit/model.rs new file mode 100644 index 0000000000..2cf0dc9232 --- /dev/null +++ b/patches/candle-transformers/src/models/mmdit/model.rs @@ -0,0 +1,240 @@ +// Implement the MMDiT model originally introduced for Stable Diffusion 3 (https://arxiv.org/abs/2403.03206), +// as well as the MMDiT-X variant introduced for Stable Diffusion 3.5-medium (https://huggingface.co/stabilityai/stable-diffusion-3.5-medium) +// This follows the implementation of the MMDiT model in the ComfyUI repository. +// https://github.com/comfyanonymous/ComfyUI/blob/78e133d0415784924cd2674e2ee48f3eeca8a2aa/comfy/ldm/modules/diffusionmodules/mmdit.py#L1 +// with MMDiT-X support following the Stability-AI/sd3.5 repository. +// https://github.com/Stability-AI/sd3.5/blob/4e484e05308d83fb77ae6f680028e6c313f9da54/mmditx.py#L1 +use candle::{Module, Result, Tensor, D}; +use candle_nn as nn; + +use super::blocks::{ + ContextQkvOnlyJointBlock, FinalLayer, JointBlock, MMDiTJointBlock, MMDiTXJointBlock, +}; +use super::embedding::{ + PatchEmbedder, PositionEmbedder, TimestepEmbedder, Unpatchifier, VectorEmbedder, +}; + +#[derive(Debug, Clone)] +pub struct Config { + pub patch_size: usize, + pub in_channels: usize, + pub out_channels: usize, + pub depth: usize, + pub head_size: usize, + pub adm_in_channels: usize, + pub pos_embed_max_size: usize, + pub context_embed_size: usize, + pub frequency_embedding_size: usize, +} + +impl Config { + pub fn sd3_medium() -> Self { + Self { + patch_size: 2, + in_channels: 16, + out_channels: 16, + depth: 24, + head_size: 64, + adm_in_channels: 2048, + pos_embed_max_size: 192, + context_embed_size: 4096, + frequency_embedding_size: 256, + } + } + + pub fn sd3_5_medium() -> Self { + Self { + patch_size: 2, + in_channels: 16, + out_channels: 16, + depth: 24, + head_size: 64, + adm_in_channels: 2048, + pos_embed_max_size: 384, + context_embed_size: 4096, + frequency_embedding_size: 256, + } + } + + pub fn sd3_5_large() -> Self { + Self { + patch_size: 2, + in_channels: 16, + out_channels: 16, + depth: 38, + head_size: 64, + adm_in_channels: 2048, + pos_embed_max_size: 192, + context_embed_size: 4096, + frequency_embedding_size: 256, + } + } +} + +pub struct MMDiT { + core: MMDiTCore, + patch_embedder: PatchEmbedder, + pos_embedder: PositionEmbedder, + timestep_embedder: TimestepEmbedder, + vector_embedder: VectorEmbedder, + context_embedder: nn::Linear, + unpatchifier: Unpatchifier, +} + +impl MMDiT { + pub fn new(cfg: &Config, use_flash_attn: bool, vb: nn::VarBuilder) -> Result { + let hidden_size = cfg.head_size * cfg.depth; + let core = MMDiTCore::new( + cfg.depth, + hidden_size, + cfg.depth, + cfg.patch_size, + cfg.out_channels, + use_flash_attn, + vb.clone(), + )?; + let patch_embedder = PatchEmbedder::new( + cfg.patch_size, + cfg.in_channels, + hidden_size, + vb.pp("x_embedder"), + )?; + let pos_embedder = PositionEmbedder::new( + hidden_size, + cfg.patch_size, + cfg.pos_embed_max_size, + vb.clone(), + )?; + let timestep_embedder = TimestepEmbedder::new( + hidden_size, + cfg.frequency_embedding_size, + vb.pp("t_embedder"), + )?; + let vector_embedder = + VectorEmbedder::new(cfg.adm_in_channels, hidden_size, vb.pp("y_embedder"))?; + let context_embedder = nn::linear( + cfg.context_embed_size, + hidden_size, + vb.pp("context_embedder"), + )?; + let unpatchifier = Unpatchifier::new(cfg.patch_size, cfg.out_channels)?; + + Ok(Self { + core, + patch_embedder, + pos_embedder, + timestep_embedder, + vector_embedder, + context_embedder, + unpatchifier, + }) + } + + pub fn forward( + &self, + x: &Tensor, + t: &Tensor, + y: &Tensor, + context: &Tensor, + skip_layers: Option<&[usize]>, + ) -> Result { + // Following the convention of the ComfyUI implementation. + // https://github.com/comfyanonymous/ComfyUI/blob/78e133d0415784924cd2674e2ee48f3eeca8a2aa/comfy/ldm/modules/diffusionmodules/mmdit.py#L919 + // + // Forward pass of DiT. + // x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) + // t: (N,) tensor of diffusion timesteps + // y: (N,) tensor of class labels + let h = x.dim(D::Minus2)?; + let w = x.dim(D::Minus1)?; + let cropped_pos_embed = self.pos_embedder.get_cropped_pos_embed(h, w)?; + let x = self + .patch_embedder + .forward(x)? + .broadcast_add(&cropped_pos_embed)?; + let c = self.timestep_embedder.forward(t)?; + let y = self.vector_embedder.forward(y)?; + let c = (c + y)?; + let context = self.context_embedder.forward(context)?; + + let x = self.core.forward(&context, &x, &c, skip_layers)?; + let x = self.unpatchifier.unpatchify(&x, h, w)?; + x.narrow(2, 0, h)?.narrow(3, 0, w) + } +} + +pub struct MMDiTCore { + joint_blocks: Vec>, + context_qkv_only_joint_block: ContextQkvOnlyJointBlock, + final_layer: FinalLayer, +} + +impl MMDiTCore { + pub fn new( + depth: usize, + hidden_size: usize, + num_heads: usize, + patch_size: usize, + out_channels: usize, + use_flash_attn: bool, + vb: nn::VarBuilder, + ) -> Result { + let mut joint_blocks = Vec::with_capacity(depth - 1); + for i in 0..depth - 1 { + let joint_block_vb_pp = format!("joint_blocks.{i}"); + let joint_block: Box = + if vb.contains_tensor(&format!("{joint_block_vb_pp}.x_block.attn2.qkv.weight")) { + Box::new(MMDiTXJointBlock::new( + hidden_size, + num_heads, + use_flash_attn, + vb.pp(&joint_block_vb_pp), + )?) + } else { + Box::new(MMDiTJointBlock::new( + hidden_size, + num_heads, + use_flash_attn, + vb.pp(&joint_block_vb_pp), + )?) + }; + joint_blocks.push(joint_block); + } + + Ok(Self { + joint_blocks, + context_qkv_only_joint_block: ContextQkvOnlyJointBlock::new( + hidden_size, + num_heads, + use_flash_attn, + vb.pp(format!("joint_blocks.{}", depth - 1)), + )?, + final_layer: FinalLayer::new( + hidden_size, + patch_size, + out_channels, + vb.pp("final_layer"), + )?, + }) + } + + pub fn forward( + &self, + context: &Tensor, + x: &Tensor, + c: &Tensor, + skip_layers: Option<&[usize]>, + ) -> Result { + let (mut context, mut x) = (context.clone(), x.clone()); + for (i, joint_block) in self.joint_blocks.iter().enumerate() { + if let Some(skip_layers) = &skip_layers { + if skip_layers.contains(&i) { + continue; + } + } + (context, x) = joint_block.forward(&context, &x, c)?; + } + let x = self.context_qkv_only_joint_block.forward(&context, &x, c)?; + self.final_layer.forward(&x, c) + } +} diff --git a/patches/candle-transformers/src/models/mmdit/projections.rs b/patches/candle-transformers/src/models/mmdit/projections.rs new file mode 100644 index 0000000000..2775328596 --- /dev/null +++ b/patches/candle-transformers/src/models/mmdit/projections.rs @@ -0,0 +1,121 @@ +use candle::{Module, Result, Tensor}; +use candle_nn as nn; + +pub struct Qkv { + pub q: Tensor, + pub k: Tensor, + pub v: Tensor, +} + +pub struct Mlp { + fc1: nn::Linear, + act: nn::Activation, + fc2: nn::Linear, +} + +impl Mlp { + pub fn new( + in_features: usize, + hidden_features: usize, + vb: candle_nn::VarBuilder, + ) -> Result { + let fc1 = nn::linear(in_features, hidden_features, vb.pp("fc1"))?; + let act = nn::Activation::GeluPytorchTanh; + let fc2 = nn::linear(hidden_features, in_features, vb.pp("fc2"))?; + + Ok(Self { fc1, act, fc2 }) + } +} + +impl Module for Mlp { + fn forward(&self, x: &Tensor) -> Result { + let x = self.fc1.forward(x)?; + let x = self.act.forward(&x)?; + self.fc2.forward(&x) + } +} + +pub struct QkvOnlyAttnProjections { + qkv: nn::Linear, + head_dim: usize, +} + +impl QkvOnlyAttnProjections { + pub fn new(dim: usize, num_heads: usize, vb: nn::VarBuilder) -> Result { + let head_dim = dim / num_heads; + let qkv = nn::linear(dim, dim * 3, vb.pp("qkv"))?; + Ok(Self { qkv, head_dim }) + } + + pub fn pre_attention(&self, x: &Tensor) -> Result { + let qkv = self.qkv.forward(x)?; + split_qkv(&qkv, self.head_dim) + } +} + +pub struct AttnProjections { + head_dim: usize, + qkv: nn::Linear, + ln_k: Option, + ln_q: Option, + proj: nn::Linear, +} + +impl AttnProjections { + pub fn new(dim: usize, num_heads: usize, vb: nn::VarBuilder) -> Result { + let head_dim = dim / num_heads; + let qkv = nn::linear(dim, dim * 3, vb.pp("qkv"))?; + let proj = nn::linear(dim, dim, vb.pp("proj"))?; + let (ln_k, ln_q) = if vb.contains_tensor("ln_k.weight") { + let ln_k = candle_nn::rms_norm(head_dim, 1e-6, vb.pp("ln_k"))?; + let ln_q = candle_nn::rms_norm(head_dim, 1e-6, vb.pp("ln_q"))?; + (Some(ln_k), Some(ln_q)) + } else { + (None, None) + }; + Ok(Self { + head_dim, + qkv, + proj, + ln_k, + ln_q, + }) + } + + pub fn pre_attention(&self, x: &Tensor) -> Result { + let qkv = self.qkv.forward(x)?; + let Qkv { q, k, v } = split_qkv(&qkv, self.head_dim)?; + let q = match self.ln_q.as_ref() { + None => q, + Some(l) => { + let (b, t, h) = q.dims3()?; + l.forward(&q.reshape((b, t, (), self.head_dim))?)? + .reshape((b, t, h))? + } + }; + let k = match self.ln_k.as_ref() { + None => k, + Some(l) => { + let (b, t, h) = k.dims3()?; + l.forward(&k.reshape((b, t, (), self.head_dim))?)? + .reshape((b, t, h))? + } + }; + Ok(Qkv { q, k, v }) + } + + pub fn post_attention(&self, x: &Tensor) -> Result { + self.proj.forward(x) + } +} + +fn split_qkv(qkv: &Tensor, head_dim: usize) -> Result { + let (batch_size, seq_len, _) = qkv.dims3()?; + let qkv = qkv.reshape((batch_size, seq_len, 3, (), head_dim))?; + let q = qkv.get_on_dim(2, 0)?; + let q = q.reshape((batch_size, seq_len, ()))?; + let k = qkv.get_on_dim(2, 1)?; + let k = k.reshape((batch_size, seq_len, ()))?; + let v = qkv.get_on_dim(2, 2)?; + Ok(Qkv { q, k, v }) +} diff --git a/patches/candle-transformers/src/models/mobileclip.rs b/patches/candle-transformers/src/models/mobileclip.rs new file mode 100644 index 0000000000..f0baf9e10c --- /dev/null +++ b/patches/candle-transformers/src/models/mobileclip.rs @@ -0,0 +1,101 @@ +//! Mobile CLIP model, combining a lightweight vision encoder with a text encoder +//! +//! A mobile-optimized CLIP implementation that uses: +//! - FastViT as the vision encoder +//! - OpenCLIP text encoder +//! - Projection layers to align the feature spaces +//! +//! See model details at: +//! - [FastViT](https://arxiv.org/abs/2303.14189) +//! - [OpenCLIP](https://github.com/mlfoundations/open_clip) +//! +//! References: +//! - [MobileVLM](https://huggingface.co/mobileVLM) +//! - [MetaCLIP](https://arxiv.org/abs/2309.16671) +//! + +use super::fastvit; +use super::openclip::text_model; +use candle::{Result, Tensor, D}; +use candle_nn::{Func, VarBuilder}; + +#[derive(Clone, Debug)] +pub struct MobileClipModel { + text_model: text_model::OpenClipTextTransformer, + vision_model: Func<'static>, + text_projection: Tensor, + logit_scale: Tensor, +} + +#[derive(Clone, Debug)] +pub struct MobileClipConfig { + pub text_config: text_model::Config, + pub vision_config: fastvit::Config, + pub image_size: usize, +} + +impl MobileClipConfig { + pub fn s1() -> Self { + let text_config = text_model::Config::vit_base_patch32(); + let vision_config = fastvit::Config::mci1(); + Self { + text_config, + vision_config, + image_size: 256, + } + } + pub fn s2() -> Self { + let text_config = text_model::Config::vit_base_patch32(); + let vision_config = fastvit::Config::mci2(); + Self { + text_config, + vision_config, + image_size: 256, + } + } +} + +impl MobileClipModel { + pub fn new(vs: VarBuilder, c: &MobileClipConfig) -> Result { + let vision_model = fastvit::fastvit(&c.vision_config, 512, vs.pp("visual.trunk"))?; + let text_model = text_model::OpenClipTextTransformer::new(vs.pp("text"), &c.text_config)?; + let text_projection = vs.get( + (c.text_config.embed_dim, c.text_config.projection_dim), + "text.text_projection", + )?; + let logit_scale = vs.get(&[], "logit_scale")?; + Ok(Self { + text_model, + vision_model, + text_projection, + logit_scale, + }) + } + + pub fn get_text_features(&self, input_ids: &Tensor) -> Result { + input_ids + .apply(&self.text_model)? + .matmul(&self.text_projection) + } + + pub fn get_image_features(&self, pixel_values: &Tensor) -> Result { + pixel_values.apply(&self.vision_model) + } + + pub fn forward(&self, pixel_values: &Tensor, input_ids: &Tensor) -> Result<(Tensor, Tensor)> { + let image_features = self.get_image_features(pixel_values)?; + let text_features = self.get_text_features(input_ids)?; + let image_features_normalized = div_l2_norm(&image_features)?; + let text_features_normalized = div_l2_norm(&text_features)?; + let logits_per_text = text_features_normalized.matmul(&image_features_normalized.t()?)?; + let logit_scale = self.logit_scale.exp()?; + let logits_per_text = logits_per_text.broadcast_mul(&logit_scale)?; + let logits_per_image = logits_per_text.t()?; + Ok((logits_per_text, logits_per_image)) + } +} + +pub fn div_l2_norm(v: &Tensor) -> Result { + let l2_norm = v.sqr()?.sum_keepdim(D::Minus1)?.sqrt()?; + v.broadcast_div(&l2_norm) +} diff --git a/patches/candle-transformers/src/models/mobilenetv4.rs b/patches/candle-transformers/src/models/mobilenetv4.rs new file mode 100644 index 0000000000..ab1e70803f --- /dev/null +++ b/patches/candle-transformers/src/models/mobilenetv4.rs @@ -0,0 +1,805 @@ +//! # MobileNet-v4 +//! +//! MobileNet-v4 inference implementation based on timm. +//! +//! ## Paper +//! +//! ["MobileNetV4 - Universal Models for the Mobile Ecosystem"](https://arxiv.org/abs/2404.10518) +//! +//! ## References +//! +//! - [PyTorch Implementation](https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/mobilenetv3.py) + +use candle::{Result, Tensor, D}; +use candle_nn::{ + batch_norm, conv2d_no_bias, linear, ops::softmax, Activation, Conv2dConfig, Func, VarBuilder, +}; + +#[derive(Clone, Debug)] +enum BlockType { + Convolutional { + out_channels: usize, + kernel: usize, + stride: usize, + }, + UniversalBottleneck { + out_channels: usize, + start_kernel: usize, + mid_kernel: usize, + stride: usize, + expand: usize, + }, + EdgeResidual { + out_channels: usize, + kernel: usize, + stride: usize, + expand: usize, + }, + Attention { + out_channels: usize, + heads: usize, + kernel: usize, + stride: usize, + kv_dim: usize, + kv_stride: usize, + }, +} + +#[derive(Clone, Debug)] +pub struct Config { + stem_dim: usize, + activation: Activation, + stages: [Vec; 5], +} + +#[rustfmt::skip] +impl Config { + pub fn small() -> Self { + Self { + stem_dim: 32, + activation: Activation::Relu, + stages: [ + vec![ + BlockType::Convolutional { out_channels: 32, kernel: 3, stride: 2}, + BlockType::Convolutional { out_channels: 32, kernel: 1, stride: 1}, + ], + vec![ + BlockType::Convolutional { out_channels: 96, kernel: 3, stride: 2}, + BlockType::Convolutional { out_channels: 64, kernel: 1, stride: 1}, + ], + vec![ + BlockType::UniversalBottleneck { out_channels: 96, start_kernel: 5, mid_kernel: 5, stride: 2, expand: 3}, + BlockType::UniversalBottleneck { out_channels: 96, start_kernel: 0, mid_kernel: 3, stride: 1, expand: 2}, + BlockType::UniversalBottleneck { out_channels: 96, start_kernel: 0, mid_kernel: 3, stride: 1, expand: 2}, + BlockType::UniversalBottleneck { out_channels: 96, start_kernel: 0, mid_kernel: 3, stride: 1, expand: 2}, + BlockType::UniversalBottleneck { out_channels: 96, start_kernel: 0, mid_kernel: 3, stride: 1, expand: 2}, + BlockType::UniversalBottleneck { out_channels: 96, start_kernel: 3, mid_kernel: 0, stride: 1, expand: 4}, + ], + vec![ + BlockType::UniversalBottleneck { out_channels: 128, start_kernel: 3, mid_kernel: 3, stride: 2, expand: 6}, + BlockType::UniversalBottleneck { out_channels: 128, start_kernel: 5, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 128, start_kernel: 0, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 128, start_kernel: 0, mid_kernel: 5, stride: 1, expand: 3}, + BlockType::UniversalBottleneck { out_channels: 128, start_kernel: 0, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 128, start_kernel: 0, mid_kernel: 3, stride: 1, expand: 4}, + ], + vec![ + BlockType::Convolutional { out_channels: 960, kernel: 1, stride: 1}, + ], + ], + } + } + + pub fn medium() -> Self { + Self { + stem_dim: 32, + activation: Activation::Relu, + stages: [ + vec![ + BlockType::EdgeResidual { out_channels: 48, kernel: 3, stride: 2, expand: 4}, + ], + vec![ + BlockType::UniversalBottleneck { out_channels: 80, start_kernel: 3, mid_kernel: 5, stride: 2, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 80, start_kernel: 3, mid_kernel: 3, stride: 1, expand: 2}, + ], + vec![ + BlockType::UniversalBottleneck { out_channels: 160, start_kernel: 3, mid_kernel: 5, stride: 2, expand: 6}, + BlockType::UniversalBottleneck { out_channels: 160, start_kernel: 3, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 160, start_kernel: 3, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 160, start_kernel: 3, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 160, start_kernel: 3, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 160, start_kernel: 3, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 160, start_kernel: 0, mid_kernel: 0, stride: 1, expand: 2}, + BlockType::UniversalBottleneck { out_channels: 160, start_kernel: 3, mid_kernel: 0, stride: 1, expand: 4}, + ], + vec![ + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 5, mid_kernel: 5, stride: 2, expand: 6}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 5, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 3, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 3, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 0, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 3, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 3, mid_kernel: 5, stride: 1, expand: 2}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 5, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 0, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 0, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 5, mid_kernel: 0, stride: 1, expand: 2}, + + ], + vec![ + BlockType::Convolutional { out_channels: 960, kernel: 1, stride: 1}, + ], + ], + } + } + + pub fn hybrid_medium() -> Self { + Self { + stem_dim: 32, + activation: Activation::Relu, + stages: [ + vec![ + BlockType::EdgeResidual { out_channels: 48, kernel: 3, stride: 2, expand: 4}, + ], + vec![ + BlockType::UniversalBottleneck { out_channels: 80, start_kernel: 3, mid_kernel: 5, stride: 2, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 80, start_kernel: 3, mid_kernel: 3, stride: 1, expand: 2}, + ], + vec![ + BlockType::UniversalBottleneck { out_channels: 160, start_kernel: 3, mid_kernel: 5, stride: 2, expand: 6}, + BlockType::UniversalBottleneck { out_channels: 160, start_kernel: 0, mid_kernel: 0, stride: 1, expand: 2}, + BlockType::UniversalBottleneck { out_channels: 160, start_kernel: 3, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 160, start_kernel: 3, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::Attention { out_channels: 160, heads: 4, kernel: 3, stride: 1, kv_stride:2, kv_dim: 64}, + BlockType::UniversalBottleneck { out_channels: 160, start_kernel: 3, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::Attention { out_channels: 160, heads: 4, kernel: 3, stride: 1, kv_stride:2, kv_dim: 64}, + BlockType::UniversalBottleneck { out_channels: 160, start_kernel: 3, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::Attention { out_channels: 160, heads: 4, kernel: 3, stride: 1, kv_stride:2, kv_dim: 64}, + BlockType::UniversalBottleneck { out_channels: 160, start_kernel: 3, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::Attention { out_channels: 160, heads: 4, kernel: 3, stride: 1, kv_stride:2, kv_dim: 64}, + BlockType::UniversalBottleneck { out_channels: 160, start_kernel: 3, mid_kernel: 0, stride: 1, expand: 4}, + ], + + vec![ + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 5, mid_kernel: 5, stride: 2, expand: 6}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 5, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 3, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 3, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 0, mid_kernel: 0, stride: 1, expand: 2}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 3, mid_kernel: 5, stride: 1, expand: 2}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 0, mid_kernel: 0, stride: 1, expand: 2}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 0, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::Attention { out_channels: 256, heads: 4, kernel: 3, stride: 1, kv_stride:1, kv_dim: 64}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 3, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::Attention { out_channels: 256, heads: 4, kernel: 3, stride: 1, kv_stride:1, kv_dim: 64}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 5, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::Attention { out_channels: 256, heads: 4, kernel: 3, stride: 1, kv_stride:1, kv_dim: 64}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 5, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::Attention { out_channels: 256, heads: 4, kernel: 3, stride: 1, kv_stride:1, kv_dim: 64}, + BlockType::UniversalBottleneck { out_channels: 256, start_kernel: 5, mid_kernel: 0, stride: 1, expand: 4}, + ], + vec![ + BlockType::Convolutional { out_channels: 960, kernel: 1, stride: 1}, + ], + ], + } + } + + pub fn large() -> Self { + Self { + stem_dim: 24, + activation: Activation::Relu, + stages: [ + vec![ + BlockType::EdgeResidual { out_channels: 48, kernel: 3, stride: 2, expand: 4}, + ], + vec![ + BlockType::UniversalBottleneck { out_channels: 96, start_kernel: 3, mid_kernel: 5, stride: 2, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 96, start_kernel: 3, mid_kernel: 3, stride: 1, expand: 4}, + ], + vec![ + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 3, mid_kernel: 5, stride: 2, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 3, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 3, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 3, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 3, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 5, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 5, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 5, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 5, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 5, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 3, mid_kernel: 0, stride: 1, expand: 4}, + ], + vec![ + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 5, stride: 2, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 0, stride: 1, expand: 4}, + ], + vec![ + BlockType::Convolutional { out_channels: 960, kernel: 1, stride: 1}, + ], + ], + } + } + + pub fn hybrid_large() -> Self { + Self { + stem_dim: 24, + activation: Activation::Gelu, + stages: [ + vec![ + BlockType::EdgeResidual { out_channels: 48, kernel: 3, stride: 2, expand: 4}, + ], + vec![ + BlockType::UniversalBottleneck { out_channels: 96, start_kernel: 3, mid_kernel: 5, stride: 2, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 96, start_kernel: 3, mid_kernel: 3, stride: 1, expand: 4}, + ], + vec![ + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 3, mid_kernel: 5, stride: 2, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 3, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 3, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 3, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 3, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 5, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 5, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::Attention { out_channels: 192, heads: 8, kernel: 3, stride: 1, kv_stride:2, kv_dim: 48}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 5, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::Attention { out_channels: 192, heads: 8, kernel: 3, stride: 1, kv_stride:2, kv_dim: 48}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 5, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::Attention { out_channels: 192, heads: 8, kernel: 3, stride: 1, kv_stride:2, kv_dim: 48}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 5, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::Attention { out_channels: 192, heads: 8, kernel: 3, stride: 1, kv_stride:2, kv_dim: 48}, + BlockType::UniversalBottleneck { out_channels: 192, start_kernel: 3, mid_kernel: 0, stride: 1, expand: 4}, + ], + + vec![ + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 5, stride: 2, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 3, stride: 1, expand: 4}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 5, stride: 1, expand: 4}, + BlockType::Attention { out_channels: 512, heads: 8, kernel: 3, stride: 1, kv_stride:1, kv_dim: 64}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::Attention { out_channels: 512, heads: 8, kernel: 3, stride: 1, kv_stride:1, kv_dim: 64}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::Attention { out_channels: 512, heads: 8, kernel: 3, stride: 1, kv_stride:1, kv_dim: 64}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 0, stride: 1, expand: 4}, + BlockType::Attention { out_channels: 512, heads: 8, kernel: 3, stride: 1, kv_stride:1, kv_dim: 64}, + BlockType::UniversalBottleneck { out_channels: 512, start_kernel: 5, mid_kernel: 0, stride: 1, expand: 4}, + ], + vec![ + BlockType::Convolutional { out_channels: 960, kernel: 1, stride: 1}, + ], + ], + } + } +} + +fn depthwise_conv( + channels: usize, + kernel: usize, + stride: usize, + padding: usize, + vb: VarBuilder, +) -> Result> { + let conv2d_cfg = Conv2dConfig { + stride, + padding, + groups: channels, + ..Default::default() + }; + + let bn = batch_norm(channels, 1e-5, vb.pp("bn"))?; + let conv = conv2d_no_bias(channels, channels, kernel, conv2d_cfg, vb.pp("conv"))?; + + Ok(Func::new(move |xs| xs.apply(&conv)?.apply_t(&bn, false))) +} + +fn pointwise_conv( + in_channels: usize, + out_channels: usize, + vb: VarBuilder, +) -> Result> { + let conv2d_cfg = Conv2dConfig { + ..Default::default() + }; + + let bn = batch_norm(out_channels, 1e-5, vb.pp("bn"))?; + let conv = conv2d_no_bias(in_channels, out_channels, 1, conv2d_cfg, vb.pp("conv"))?; + + Ok(Func::new(move |xs| xs.apply(&conv)?.apply_t(&bn, false))) +} + +//Universal block that uses two pointwise convolutions and all combinations of two depthwise convolutions. +#[allow(clippy::too_many_arguments)] +fn universal_inverted_bottleneck_block( + cfg: &Config, + in_channels: usize, + out_channels: usize, + expand: usize, + start_kernel: usize, + mid_kernel: usize, + stride: usize, + vb: VarBuilder, +) -> Result> { + let act = cfg.activation; + let skip_connection = (in_channels == out_channels) && (stride == 1); + + let dw_start_stride = if mid_kernel > 0 { 1 } else { stride }; + let dw_start = depthwise_conv( + in_channels, + start_kernel, + dw_start_stride, + start_kernel / 2, + vb.pp("dw_start"), + ); + let pw_exp = pointwise_conv(in_channels, in_channels * expand, vb.pp("pw_exp"))?; + let dw_mid = depthwise_conv( + in_channels * expand, + mid_kernel, + stride, + mid_kernel / 2, + vb.pp("dw_mid"), + ); + let pw_proj = pointwise_conv(in_channels * expand, out_channels, vb.pp("pw_proj"))?; + + let gamma = vb.get(out_channels, "layer_scale.gamma"); + + Ok(Func::new(move |xs| { + let residual = xs.clone(); + + let mut xs = xs.clone(); + + if let Ok(f) = &dw_start { + xs = xs.apply(f)?; + } + + xs = xs.apply(&pw_exp)?.apply(&act)?; + + if let Ok(f) = &dw_mid { + xs = xs.apply(f)?.apply(&act)?; + } + + xs = xs.apply(&pw_proj)?; + + if let Ok(g) = &gamma { + xs = xs.broadcast_mul(&g.reshape((1, (), 1, 1))?)?; + }; + + if skip_connection { + xs = (xs + residual)?; + } + + Ok(xs) + })) +} + +// Convolutional block including norm and activation. +fn conv_block( + cfg: &Config, + in_channels: usize, + out_channels: usize, + kernel: usize, + stride: usize, + vb: VarBuilder, +) -> Result> { + let conv2d_cfg = Conv2dConfig { + stride, + padding: kernel / 2, + ..Default::default() + }; + + let act = cfg.activation; + let bn = batch_norm(out_channels, 1e-5, vb.pp("bn1"))?; + let conv = conv2d_no_bias(in_channels, out_channels, kernel, conv2d_cfg, vb.pp("conv"))?; + + Ok(Func::new(move |xs| { + xs.apply(&conv)?.apply_t(&bn, false)?.apply(&act) + })) +} + +fn edge_residual_block( + cfg: &Config, + in_channels: usize, + out_channels: usize, + kernel: usize, + stride: usize, + expand: usize, + vb: VarBuilder, +) -> Result> { + let conv_exp_cfg = Conv2dConfig { + stride, + padding: kernel / 2, + ..Default::default() + }; + + let conv_pwl_cfg = Conv2dConfig { + ..Default::default() + }; + + let act = cfg.activation; + let mid_channels = in_channels * expand; + let conv_exp = conv2d_no_bias( + in_channels, + mid_channels, + kernel, + conv_exp_cfg, + vb.pp("conv_exp"), + )?; + let bn1 = batch_norm(mid_channels, 1e-5, vb.pp("bn1"))?; + + let conv_pwl = conv2d_no_bias( + mid_channels, + out_channels, + 1, + conv_pwl_cfg, + vb.pp("conv_pwl"), + )?; + let bn2 = batch_norm(out_channels, 1e-5, vb.pp("bn2"))?; + + Ok(Func::new(move |xs| { + let xs = xs + .apply(&conv_exp)? + .apply_t(&bn1, false)? + .apply(&act)? + .apply(&conv_pwl)? + .apply_t(&bn2, false)?; + + Ok(xs) + })) +} + +fn reshape_kv(t: &Tensor) -> Result { + let d = t.dims4()?; + let t = t + .reshape((d.0, d.1, ()))? + .transpose(1, 2)? + .unsqueeze(1)? + .contiguous()?; + Ok(t) +} + +fn reshape_query(t: &Tensor, heads: usize, kv_dim: usize) -> Result { + let d = t.dims4()?; + + let t = t + .reshape((d.0, heads, kv_dim, ()))? + .transpose(D::Minus1, D::Minus2)? + .contiguous()?; + Ok(t) +} + +fn reshape_output(t: &Tensor, heads: usize, h: usize, w: usize) -> Result { + let d = t.dims4()?; + let t = t.transpose(1, 2)?; + let t = t + .reshape((d.0, h, w, d.3 * heads))? + .permute((0, 3, 1, 2))? + .contiguous()?; + Ok(t) +} + +// Mobile multi-query attention +#[allow(clippy::too_many_arguments)] +fn mqa_block( + in_channels: usize, + out_channels: usize, + heads: usize, + kernel: usize, + stride: usize, + kv_dim: usize, + kv_stride: usize, + vb: VarBuilder, +) -> Result> { + let down_conv2d_cfg = Conv2dConfig { + stride: kv_stride, + padding: kernel / 2, + groups: in_channels, + ..Default::default() + }; + + let proj_conv2d_cfg = Conv2dConfig { + stride, + ..Default::default() + }; + + let skip_connection = (in_channels == out_channels) && (stride == 1); + let gamma = vb.get(out_channels, "layer_scale.gamma"); + let norm = batch_norm(out_channels, 1e-5, vb.pp("norm"))?; + let scale = (kv_dim as f64).powf(-0.5); + + let vb = vb.pp("attn"); + + let query_proj = conv2d_no_bias( + out_channels, + kv_dim * heads, + 1, + proj_conv2d_cfg, + vb.pp("query.proj"), + )?; + + let key_down_conv = conv2d_no_bias( + in_channels, + out_channels, + kernel, + down_conv2d_cfg, + vb.pp("key.down_conv"), + ); + let key_norm = batch_norm(out_channels, 1e-5, vb.pp("key.norm")); + + let key_proj = conv2d_no_bias(out_channels, kv_dim, 1, proj_conv2d_cfg, vb.pp("key.proj"))?; + + let value_down_conv = conv2d_no_bias( + in_channels, + out_channels, + kernel, + down_conv2d_cfg, + vb.pp("value.down_conv"), + ); + + let value_norm = batch_norm(out_channels, 1e-5, vb.pp("value.norm")); + let value_proj = conv2d_no_bias( + out_channels, + kv_dim, + 1, + proj_conv2d_cfg, + vb.pp("value.proj"), + )?; + + let output_proj = conv2d_no_bias( + kv_dim * heads, + out_channels, + 1, + proj_conv2d_cfg, + vb.pp("output.proj"), + )?; + + Ok(Func::new(move |xs| { + let (_, _, h, w) = xs.dims4()?; + + let residual = xs.clone(); + + let xs = xs.apply_t(&norm, false)?; + + // Query + let q = xs.apply(&query_proj)?; + + let q = reshape_query(&q, heads, kv_dim)?; + let q = (q * scale)?; + + // Keys + let mut k = xs.clone(); + + if let (Ok(kd), Ok(n)) = (&key_down_conv, &key_norm) { + k = k.apply(kd)?.apply_t(n, false)?; + } + + let k = k.apply(&key_proj)?; + + let k = reshape_kv(&k)?; + + // Value + let mut v = xs.clone(); + + if let (Ok(vd), Ok(n)) = (&value_down_conv, &value_norm) { + v = v.apply(vd)?; + v = v.apply_t(n, false)?; + } + + let v = v.apply(&value_proj)?; + let v = reshape_kv(&v)?; + + let attn = q.broadcast_matmul(&(k.transpose(D::Minus2, D::Minus1)?))?; + let attn = softmax(&attn, D::Minus1)?; + let o = attn.broadcast_matmul(&v)?; + + let o = reshape_output(&o, heads, h, w)?; + + let mut xs = o.apply(&output_proj)?; + + // Layer scale + + if let Ok(g) = &gamma { + xs = xs.broadcast_mul(&g.reshape((1, (), 1, 1))?)?; + }; + + if skip_connection { + xs = (xs + residual)?; + } + Ok(xs) + })) +} + +// Stem. +fn mobilenetv4_stem(cfg: &Config, vb: VarBuilder) -> Result> { + let conv2d_cfg = Conv2dConfig { + stride: 2, + padding: 1, + ..Default::default() + }; + + let act = cfg.activation; + let out_channels = cfg.stem_dim; + let bn = batch_norm(out_channels, 1e-5, vb.pp("bn1"))?; + let conv = conv2d_no_bias(3, out_channels, 3, conv2d_cfg, vb.pp("conv_stem"))?; + + Ok(Func::new(move |xs| { + let xs = xs.apply(&conv)?.apply_t(&bn, false)?.apply(&act)?; + Ok(xs) + })) +} + +// The blocks in all the 5 stages of the model. +fn mobilenetv4_blocks(cfg: &Config, vb: VarBuilder) -> Result> { + let mut in_channels = cfg.stem_dim; + let mut blocks = Vec::new(); + + for stage in 0..5 { + let nblocks = cfg.stages[stage].len(); + + for block in 0..nblocks { + match cfg.stages[stage][block] { + BlockType::Convolutional { + out_channels, + kernel, + stride, + } => { + blocks.push(conv_block( + cfg, + in_channels, + out_channels, + kernel, + stride, + vb.pp(format!("{stage}.{block}")), + )?); + in_channels = out_channels; + } + + BlockType::EdgeResidual { + out_channels, + kernel, + stride, + expand, + } => { + blocks.push(edge_residual_block( + cfg, + in_channels, + out_channels, + kernel, + stride, + expand, + vb.pp(format!("{stage}.{block}")), + )?); + in_channels = out_channels; + } + + BlockType::UniversalBottleneck { + out_channels, + start_kernel, + mid_kernel, + stride, + expand, + } => { + blocks.push(universal_inverted_bottleneck_block( + cfg, + in_channels, + out_channels, + expand, + start_kernel, + mid_kernel, + stride, + vb.pp(format!("{stage}.{block}")), + )?); + in_channels = out_channels; + } + + BlockType::Attention { + out_channels, + heads, + kernel, + stride, + kv_dim, + kv_stride, + } => { + blocks.push(mqa_block( + in_channels, + out_channels, + heads, + kernel, + stride, + kv_dim, + kv_stride, + vb.pp(format!("{stage}.{block}")), + )?); + in_channels = out_channels; + } + } + } + } + + Ok(Func::new(move |xs| { + let mut xs = xs.clone(); + for block in blocks.iter() { + xs = xs.apply(block)? + } + Ok(xs) + })) +} + +// Classification head. +fn mobilenetv4_head( + cfg: &Config, + outputs: usize, + nclasses: usize, + vb: VarBuilder, +) -> Result> { + let conv2d_cfg = Conv2dConfig { + ..Default::default() + }; + + let act = cfg.activation; + let conv = conv2d_no_bias(960, outputs, 1, conv2d_cfg, vb.pp("conv_head"))?; + let norm = batch_norm(outputs, 1e-5, vb.pp("norm_head"))?; + let cls = linear(outputs, nclasses, vb.pp("classifier"))?; + + Ok(Func::new(move |xs| { + let mut xs = xs.clone(); + xs = xs.apply(&conv)?; + xs = xs.apply_t(&norm, false)?.apply(&act)?; + xs = xs.flatten_from(1)?; + xs = xs.apply(&cls)?; + Ok(xs) + })) +} + +// Build a mobilenetv4 model for a given configuration. +fn mobilenetv4_model( + cfg: &Config, + nclasses: Option, + vb: VarBuilder, +) -> Result> { + let cls = match nclasses { + None => None, + Some(nclasses) => { + let outputs = 1280; + let head = mobilenetv4_head(cfg, outputs, nclasses, vb.clone())?; + Some(head) + } + }; + + let stem = mobilenetv4_stem(cfg, vb.clone())?; + + let blocks = mobilenetv4_blocks(cfg, vb.pp("blocks"))?; + + Ok(Func::new(move |xs| { + let xs = xs.apply(&stem)?.apply(&blocks)?; + let xs = xs.mean_keepdim(D::Minus1)?.mean_keepdim(D::Minus2)?; + match &cls { + None => Ok(xs), + Some(cls) => xs.apply(cls), + } + })) +} + +pub fn mobilenetv4(cfg: &Config, nclasses: usize, vb: VarBuilder) -> Result> { + mobilenetv4_model(cfg, Some(nclasses), vb) +} + +pub fn mobilenetv4_no_final_layer(cfg: &Config, vb: VarBuilder) -> Result> { + mobilenetv4_model(cfg, None, vb) +} diff --git a/patches/candle-transformers/src/models/mobileone.rs b/patches/candle-transformers/src/models/mobileone.rs new file mode 100644 index 0000000000..e8836745b9 --- /dev/null +++ b/patches/candle-transformers/src/models/mobileone.rs @@ -0,0 +1,334 @@ +//! # MobileOne +//! +//! MobileOne inference implementation based on timm and candle-repvgg +//! +//! See ["MobileOne: An Improved One millisecond Mobile Backbone"](https://arxiv.org/abs/2206.04040) + +use candle::{DType, Result, Tensor, D}; +use candle_nn::{ + batch_norm, conv2d, conv2d_no_bias, linear, ops::sigmoid, BatchNorm, Conv2d, Conv2dConfig, + Func, VarBuilder, +}; + +struct StageConfig { + blocks: usize, + channels: usize, +} + +// The architecture in the paper has 6 stages. The timm implementation uses an equivalent form +// by concatenating the 5th stage (starts with stride 1) to the previous one. +const STAGES: [StageConfig; 5] = [ + StageConfig { + blocks: 1, + channels: 64, + }, + StageConfig { + blocks: 2, + channels: 64, + }, + StageConfig { + blocks: 8, + channels: 128, + }, + StageConfig { + blocks: 10, + channels: 256, + }, + StageConfig { + blocks: 1, + channels: 512, + }, +]; + +#[derive(Clone)] +pub struct Config { + /// overparameterization factor + k: usize, + /// per-stage channel number multipliers + alphas: [f32; 5], +} + +impl Config { + pub fn s0() -> Self { + Self { + k: 4, + alphas: [0.75, 0.75, 1.0, 1.0, 2.0], + } + } + pub fn s1() -> Self { + Self { + k: 1, + alphas: [1.5, 1.5, 1.5, 2.0, 2.5], + } + } + pub fn s2() -> Self { + Self { + k: 1, + alphas: [1.5, 1.5, 2.0, 2.5, 4.0], + } + } + pub fn s3() -> Self { + Self { + k: 1, + alphas: [2.0, 2.0, 2.5, 3.0, 4.0], + } + } + pub fn s4() -> Self { + Self { + k: 1, + alphas: [3.0, 3.0, 3.5, 3.5, 4.0], + } + } +} + +// SE blocks are used in the last stages of the s4 variant. +fn squeeze_and_excitation( + in_channels: usize, + squeeze_channels: usize, + vb: VarBuilder, +) -> Result> { + let conv2d_cfg = Conv2dConfig { + ..Default::default() + }; + let fc1 = conv2d(in_channels, squeeze_channels, 1, conv2d_cfg, vb.pp("fc1"))?; + let fc2 = conv2d(squeeze_channels, in_channels, 1, conv2d_cfg, vb.pp("fc2"))?; + + Ok(Func::new(move |xs| { + let residual = xs; + let xs = xs.mean_keepdim(D::Minus2)?.mean_keepdim(D::Minus1)?; + let xs = sigmoid(&xs.apply(&fc1)?.relu()?.apply(&fc2)?)?; + + residual.broadcast_mul(&xs) + })) +} + +// fuses a convolutional kernel and a batchnorm layer into a convolutional layer +// based on the _fuse_bn_tensor method in timm +// see https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/byobnet.py#L602 +fn fuse_conv_bn(weights: &Tensor, bn: BatchNorm) -> Result<(Tensor, Tensor)> { + let (gamma, beta) = bn.weight_and_bias().unwrap(); + let mu = bn.running_mean(); + let sigma = (bn.running_var() + bn.eps())?.sqrt(); + let gps = (gamma / sigma)?; + let bias = (beta - mu * &gps)?; + let weights = weights.broadcast_mul(&gps.reshape(((), 1, 1, 1))?)?; + + Ok((weights, bias)) +} + +// A mobileone block has a different training time and inference time architecture. +// The latter is a simple and efficient equivalent transformation of the former +// realized by a structural reparameterization technique, where convolutions +// along with identity branches and batchnorm layers are fused into a single convolution. +#[allow(clippy::too_many_arguments)] +fn mobileone_block( + has_identity: bool, + k: usize, + dim: usize, + stride: usize, + padding: usize, + groups: usize, + kernel: usize, + in_channels: usize, + out_channels: usize, + vb: VarBuilder, +) -> Result> { + let conv2d_cfg = Conv2dConfig { + stride, + padding, + groups, + ..Default::default() + }; + + let mut w = Tensor::zeros( + (out_channels, in_channels / groups, kernel, kernel), + DType::F32, + vb.device(), + )?; + let mut b = Tensor::zeros(dim, DType::F32, vb.device())?; + + // k is the training-time overparameterization factor, larger than 1 only in the s0 variant + for i in 0..k { + let conv_kxk_bn = batch_norm(dim, 1e-5, vb.pp(format!("conv_kxk.{i}.bn")))?; + let conv_kxk = conv2d_no_bias( + in_channels, + out_channels, + kernel, + conv2d_cfg, + vb.pp(format!("conv_kxk.{i}.conv")), + )?; + let (wk, bk) = fuse_conv_bn(conv_kxk.weight(), conv_kxk_bn)?; + w = (w + wk)?; + b = (b + bk)?; + } + + if kernel > 1 { + let conv_scale_bn = batch_norm(dim, 1e-5, vb.pp("conv_scale.bn"))?; + let conv_scale = conv2d_no_bias( + in_channels, + out_channels, + 1, + conv2d_cfg, + vb.pp("conv_scale.conv"), + )?; + + let (mut ws, bs) = fuse_conv_bn(conv_scale.weight(), conv_scale_bn)?; + // resize to 3x3 + ws = ws.pad_with_zeros(D::Minus1, 1, 1)?; + ws = ws.pad_with_zeros(D::Minus2, 1, 1)?; + + w = (w + ws)?; + b = (b + bs)?; + } + + // Use SE blocks if present (last layers of the s4 variant) + let se = squeeze_and_excitation(out_channels, out_channels / 16, vb.pp("attn")); + + // read and reparameterize the identity bn into wi and bi + if has_identity { + let identity_bn = batch_norm(dim, 1e-5, vb.pp("identity"))?; + + let mut weights: Vec = vec![0.0; w.elem_count()]; + + let id = in_channels / groups; + // See https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/byobnet.py#L809 + for i in 0..in_channels { + if kernel > 1 { + weights[i * kernel * kernel + 4] = 1.0; + } else { + weights[i * (id + 1)] = 1.0; + } + } + + let weights = &Tensor::from_vec(weights, w.shape(), w.device())?; + let (wi, bi) = fuse_conv_bn(weights, identity_bn)?; + + w = (w + wi)?; + b = (b + bi)?; + } + + let reparam_conv = Conv2d::new(w, Some(b), conv2d_cfg); + + Ok(Func::new(move |xs| { + let mut xs = xs.apply(&reparam_conv)?; + if let Ok(f) = &se { + xs = xs.apply(f)?; + } + xs = xs.relu()?; + Ok(xs) + })) +} + +// Get the number of output channels per stage taking into account the multipliers +fn output_channels_per_stage(cfg: &Config, stage: usize) -> usize { + let channels = STAGES[stage].channels as f32; + let alpha = cfg.alphas[stage]; + + match stage { + 0 => std::cmp::min(64, (channels * alpha) as usize), + _ => (channels * alpha) as usize, + } +} + +// Each stage is made of blocks. The first layer always downsamples with stride 2. +// All but the first block have a residual connection. +fn mobileone_stage(cfg: &Config, idx: usize, vb: VarBuilder) -> Result> { + let nblocks = STAGES[idx].blocks; + let mut blocks = Vec::with_capacity(nblocks); + + let mut in_channels = output_channels_per_stage(cfg, idx - 1); + + for block_idx in 0..nblocks { + let out_channels = output_channels_per_stage(cfg, idx); + let (has_identity, stride) = if block_idx == 0 { + (false, 2) + } else { + (true, 1) + }; + + // depthwise convolution layer + blocks.push(mobileone_block( + has_identity, + cfg.k, + in_channels, + stride, + 1, + in_channels, + 3, + in_channels, + in_channels, + vb.pp(block_idx * 2), + )?); + + // pointwise convolution layer + blocks.push(mobileone_block( + has_identity, + cfg.k, + out_channels, + 1, // stride + 0, // padding + 1, // groups + 1, // kernel + in_channels, + out_channels, + vb.pp(block_idx * 2 + 1), + )?); + + in_channels = out_channels; + } + + Ok(Func::new(move |xs| { + let mut xs = xs.clone(); + for block in blocks.iter() { + xs = xs.apply(block)? + } + Ok(xs) + })) +} + +// Build a mobileone model for a given configuration. +fn mobileone_model( + config: &Config, + nclasses: Option, + vb: VarBuilder, +) -> Result> { + let cls = match nclasses { + None => None, + Some(nclasses) => { + let outputs = output_channels_per_stage(config, 4); + let linear = linear(outputs, nclasses, vb.pp("head.fc"))?; + Some(linear) + } + }; + + let stem_dim = output_channels_per_stage(config, 0); + let stem = mobileone_block(false, 1, stem_dim, 2, 1, 1, 3, 3, stem_dim, vb.pp("stem"))?; + let vb = vb.pp("stages"); + let stage1 = mobileone_stage(config, 1, vb.pp(0))?; + let stage2 = mobileone_stage(config, 2, vb.pp(1))?; + let stage3 = mobileone_stage(config, 3, vb.pp(2))?; + let stage4 = mobileone_stage(config, 4, vb.pp(3))?; + + Ok(Func::new(move |xs| { + let xs = xs + .apply(&stem)? + .apply(&stage1)? + .apply(&stage2)? + .apply(&stage3)? + .apply(&stage4)? + .mean(D::Minus2)? + .mean(D::Minus1)?; + match &cls { + None => Ok(xs), + Some(cls) => xs.apply(cls), + } + })) +} + +pub fn mobileone(cfg: &Config, nclasses: usize, vb: VarBuilder) -> Result> { + mobileone_model(cfg, Some(nclasses), vb) +} + +pub fn mobileone_no_final_layer(cfg: &Config, vb: VarBuilder) -> Result> { + mobileone_model(cfg, None, vb) +} diff --git a/patches/candle-transformers/src/models/mod.rs b/patches/candle-transformers/src/models/mod.rs new file mode 100644 index 0000000000..309f67a230 --- /dev/null +++ b/patches/candle-transformers/src/models/mod.rs @@ -0,0 +1,136 @@ +//! Candle implementations for various deep learning models +//! +//! This crate provides implementations of popular machine learning models and architectures for different modalities. +//! +//! - Large language models: [`llama`], [`phi3`], [`mamba`], [`mixtral`], [`bert`], ... +//! - Text to text models: [`t5`], ... +//! - Image to text models: [`blip`], ... +//! - Text to image models: [`stable_diffusion`] and [`wuerstchen`], ... +//! - Audio models: [`whisper`], [`encodec`], [`metavoice`], [`parler_tts`], ... +//! - Computer vision models: [`dinov2`], [`convmixer`], [`efficientnet`], ... +//! +//! Some of the models also have quantized variants, e.g. [`quantized_blip`], [`quantized_llama`] and [`quantized_qwen2`]. +//! +//! The implementations aim to be readable while maintaining good performance. For more information +//! on each model see the model's module docs in the links below. + +pub mod based; +pub mod beit; +pub mod bert; +pub mod bigcode; +pub mod blip; +pub mod blip_text; +pub mod chatglm; +pub mod chinese_clip; +pub mod clip; +pub mod codegeex4_9b; +pub mod colpali; +pub mod convmixer; +pub mod convnext; +pub mod csm; +pub mod dac; +pub mod debertav2; +pub mod deepseek2; +pub mod depth_anything_v2; +pub mod dinov2; +pub mod dinov2reg4; +pub mod distilbert; +pub mod efficientnet; +pub mod efficientvit; +pub mod encodec; +pub mod eva2; +pub mod falcon; +pub mod fastvit; +pub mod flux; +pub mod gemma; +pub mod gemma2; +pub mod gemma3; +pub mod glm4; +pub mod glm4_new; +pub mod granite; +pub mod granitemoehybrid; +pub mod helium; +pub mod hiera; +pub mod jina_bert; +pub mod llama; +pub mod llama2_c; +pub mod llama2_c_weights; +pub mod llava; +pub mod mamba; +pub mod mamba2; +pub mod marian; +pub mod metavoice; +pub mod mimi; +pub mod mistral; +pub mod mixformer; +pub mod mixtral; +pub mod mmdit; +pub mod mobileclip; +pub mod mobilenetv4; +pub mod mobileone; +pub mod modernbert; +pub mod moondream; +pub mod mpt; +pub mod nvembed_v2; +pub mod olmo; +pub mod olmo2; +pub mod openclip; +pub mod paddleocr_vl; +pub mod paligemma; +pub mod parler_tts; +pub mod persimmon; +pub mod phi; +pub mod phi3; +pub mod pixtral; +pub mod quantized_blip; +pub mod quantized_blip_text; +pub mod quantized_gemma3; +pub mod quantized_glm4; +pub mod quantized_lfm2; +pub mod quantized_llama; +pub mod quantized_llama2_c; +pub mod quantized_metavoice; +pub mod quantized_mistral; +pub mod quantized_mixformer; +pub mod quantized_moondream; +pub mod quantized_mpt; +pub mod quantized_phi; +pub mod quantized_phi3; +pub mod quantized_qwen2; +pub mod quantized_qwen3; +pub mod quantized_qwen3_moe; +pub mod quantized_recurrent_gemma; +pub mod quantized_rwkv_v5; +pub mod quantized_rwkv_v6; +pub mod quantized_stable_lm; +pub mod quantized_t5; +pub mod qwen2; +pub mod qwen2_moe; +pub mod qwen3; +pub mod qwen3_moe; +pub mod qwen3_vl; +pub mod recurrent_gemma; +pub mod repvgg; +pub mod resnet; +pub mod rwkv_v5; +pub mod rwkv_v6; +pub mod segformer; +pub mod segment_anything; +pub mod siglip; +pub mod smol; +pub mod snac; +pub mod stable_diffusion; +pub mod stable_lm; +pub mod starcoder2; +pub mod stella_en_v5; +pub mod t5; +pub mod trocr; +pub mod vgg; +pub mod vit; +pub mod voxtral; +pub mod whisper; +pub mod with_tracing; +pub mod wuerstchen; +pub mod xlm_roberta; +pub mod yi; +pub mod z_image; diff --git a/patches/candle-transformers/src/models/modernbert.rs b/patches/candle-transformers/src/models/modernbert.rs new file mode 100644 index 0000000000..a1f0389aaf --- /dev/null +++ b/patches/candle-transformers/src/models/modernbert.rs @@ -0,0 +1,504 @@ +//! ModernBERT +//! +//! ModernBERT is a modernized bidirectional encoder-only Transformer model. +//! - [Arxiv](https://arxiv.org/abs/2412.13663) "Smarter, Better, Faster, Longer: A Modern Bidirectional Encoder for Fast, Memory Efficient, and Long Context Finetuning and Inference" +//! - Upstream [GitHub repo](https://github.com/AnswerDotAI/ModernBERT). +//! - See modernbert in [candle-examples](https://github.com/huggingface/candle/tree/main/candle-examples/) for runnable code +//! + +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::{ + embedding, layer_norm_no_bias, linear, linear_no_bias, ops::softmax, Embedding, LayerNorm, + Linear, Module, VarBuilder, +}; +use serde::Deserialize; + +use core::f32; +use std::collections::HashMap; +use std::sync::Arc; + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub intermediate_size: usize, + pub max_position_embeddings: usize, + pub layer_norm_eps: f64, + pub pad_token_id: u32, + pub global_attn_every_n_layers: usize, + pub global_rope_theta: f64, + pub local_attention: usize, + pub local_rope_theta: f64, + #[serde(default)] + #[serde(flatten)] + pub classifier_config: Option, +} + +#[derive(Debug, Clone, Deserialize, PartialEq, Copy, Default)] +#[serde(rename_all = "lowercase")] +pub enum ClassifierPooling { + #[default] + CLS, + MEAN, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct ClassifierConfig { + pub id2label: HashMap, + pub label2id: HashMap, + pub classifier_pooling: ClassifierPooling, +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(dtype: DType, config: &Config, rope_theta: f64, dev: &Device) -> Result { + let dim = config.hidden_size / config.num_attention_heads; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let max_seq_len = config.max_position_embeddings; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb_qkv(&self, q: &Tensor, k: &Tensor) -> Result<(Tensor, Tensor)> { + let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &self.cos, &self.sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &self.cos, &self.sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Clone)] +struct ModernBertAttention { + qkv: Linear, + proj: Linear, + num_attention_heads: usize, + attention_head_size: usize, + rotary_emb: Arc, +} + +impl ModernBertAttention { + fn load(vb: VarBuilder, config: &Config, rotary_emb: Arc) -> Result { + let num_attention_heads = config.num_attention_heads; + let attention_head_size = config.hidden_size / config.num_attention_heads; + + let qkv = linear_no_bias(config.hidden_size, config.hidden_size * 3, vb.pp("Wqkv"))?; + let proj = linear_no_bias(config.hidden_size, config.hidden_size, vb.pp("Wo"))?; + + Ok(Self { + qkv, + proj, + num_attention_heads, + attention_head_size, + rotary_emb, + }) + } + + fn forward(&self, hidden_states: &Tensor, attention_mask: &Tensor) -> Result { + let xs = hidden_states.clone(); + let (b, seq_len, d) = xs.dims3()?; + let qkv = xs + .apply(&self.qkv)? + .reshape(( + b, + seq_len, + 3, + self.num_attention_heads, + self.attention_head_size, + ))? + .permute((2, 0, 3, 1, 4))?; + + let q = qkv.get(0)?; + let k = qkv.get(1)?; + let v = qkv.get(2)?; + + let (q, k) = self.rotary_emb.apply_rotary_emb_qkv(&q, &k)?; + + let scale = (self.attention_head_size as f64).powf(-0.5); + let q = (q * scale)?; + + let att = q.matmul(&k.transpose(D::Minus2, D::Minus1)?)?; + + let att = att.broadcast_add(attention_mask)?; + let att = softmax(&att, D::Minus1)?; + + let xs = att.matmul(&v)?; + + let xs = xs.transpose(1, 2)?.reshape((b, seq_len, d))?; + let xs = xs.apply(&self.proj)?; + let xs = xs.reshape((b, seq_len, d))?; + + Ok(xs) + } +} + +#[derive(Clone)] +pub struct ModernBertMLP { + wi: Linear, + wo: Linear, +} + +impl ModernBertMLP { + fn load(vb: VarBuilder, config: &Config) -> Result { + let wi = linear_no_bias( + config.hidden_size, + config.intermediate_size * 2, + vb.pp("Wi"), + )?; + let wo = linear_no_bias(config.intermediate_size, config.hidden_size, vb.pp("Wo"))?; + Ok(Self { wi, wo }) + } +} + +impl Module for ModernBertMLP { + fn forward(&self, xs: &Tensor) -> Result { + let xs = xs.apply(&self.wi)?; + let xs = xs.chunk(2, D::Minus1)?; + let xs = (&xs[0].gelu_erf()? * &xs[1])?.apply(&self.wo)?; // GeGLU + Ok(xs) + } +} + +#[derive(Clone)] +pub struct ModernBertLayer { + attn: ModernBertAttention, + mlp: ModernBertMLP, + attn_norm: Option, + mlp_norm: LayerNorm, + uses_local_attention: bool, +} + +impl ModernBertLayer { + fn load( + vb: VarBuilder, + config: &Config, + rotary_emb: Arc, + uses_local_attention: bool, + ) -> Result { + let attn = ModernBertAttention::load(vb.pp("attn"), config, rotary_emb)?; + let mlp = ModernBertMLP::load(vb.pp("mlp"), config)?; + let attn_norm = layer_norm_no_bias( + config.hidden_size, + config.layer_norm_eps, + vb.pp("attn_norm"), + ) + .ok(); + let mlp_norm = + layer_norm_no_bias(config.hidden_size, config.layer_norm_eps, vb.pp("mlp_norm"))?; + Ok(Self { + attn, + mlp, + attn_norm, + mlp_norm, + uses_local_attention, + }) + } + + fn forward( + &self, + xs: &Tensor, + global_attention_mask: &Tensor, + local_attention_mask: &Tensor, + ) -> Result { + let residual = xs.clone(); + let mut xs = xs.clone(); + if let Some(norm) = &self.attn_norm { + xs = xs.apply(norm)?; + } + + let attention_mask = if self.uses_local_attention { + &global_attention_mask.broadcast_add(local_attention_mask)? + } else { + global_attention_mask + }; + let xs = self.attn.forward(&xs, attention_mask)?; + let xs = (xs + residual)?; + let mlp_out = xs.apply(&self.mlp_norm)?.apply(&self.mlp)?; + let xs = (xs + mlp_out)?; + Ok(xs) + } +} + +#[derive(Clone)] +pub struct ModernBertHead { + dense: Linear, + norm: LayerNorm, +} + +impl ModernBertHead { + fn load(vb: VarBuilder, config: &Config) -> Result { + let dense = linear_no_bias(config.hidden_size, config.hidden_size, vb.pp("dense"))?; + let norm = layer_norm_no_bias(config.hidden_size, config.layer_norm_eps, vb.pp("norm"))?; + Ok(Self { dense, norm }) + } +} + +impl Module for ModernBertHead { + fn forward(&self, xs: &Tensor) -> Result { + let xs = xs.apply(&self.dense)?.gelu_erf()?.apply(&self.norm)?; + Ok(xs) + } +} + +#[derive(Clone)] +pub struct ModernBertDecoder { + decoder: Linear, +} + +impl ModernBertDecoder { + fn load(vb: VarBuilder, config: &Config) -> Result { + // The decoder weights are tied with the embeddings layer weights + let decoder_weights = vb.get( + (config.vocab_size, config.hidden_size), + "model.embeddings.tok_embeddings.weight", + )?; + let decoder_bias = vb.get(config.vocab_size, "decoder.bias")?; + let decoder = Linear::new(decoder_weights, Some(decoder_bias)); + Ok(Self { decoder }) + } +} + +impl Module for ModernBertDecoder { + fn forward(&self, xs: &Tensor) -> Result { + let xs = xs.apply(&self.decoder)?; + Ok(xs) + } +} + +// Global attention mask calculated from padded token inputs +fn prepare_4d_attention_mask( + mask: &Tensor, + dtype: DType, + tgt_len: Option, +) -> Result { + let bsz = mask.dim(0)?; + let src_len = mask.dim(1)?; + let tgt_len = tgt_len.unwrap_or(src_len); + + let expanded_mask = mask + .unsqueeze(1)? + .unsqueeze(2)? + .expand((bsz, 1, tgt_len, src_len))? + .to_dtype(dtype)?; + + let inverted_mask = (1.0 - expanded_mask)?; + + (inverted_mask * f32::MIN as f64)?.to_dtype(dtype) +} + +// Attention mask caused by the sliding window +fn get_local_attention_mask( + seq_len: usize, + max_distance: usize, + device: &Device, +) -> Result { + let mask: Vec<_> = (0..seq_len) + .flat_map(|i| { + (0..seq_len).map(move |j| { + if (j as i32 - i as i32).abs() > max_distance as i32 { + f32::NEG_INFINITY + } else { + 0. + } + }) + }) + .collect(); + Tensor::from_slice(&mask, (seq_len, seq_len), device) +} + +// ModernBERT backbone +#[derive(Clone)] +pub struct ModernBert { + word_embeddings: Embedding, + norm: LayerNorm, + layers: Vec, + final_norm: LayerNorm, + local_attention_size: usize, +} + +impl ModernBert { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let word_embeddings = embedding( + config.vocab_size, + config.hidden_size, + vb.pp("model.embeddings.tok_embeddings"), + )?; + let norm = layer_norm_no_bias( + config.hidden_size, + config.layer_norm_eps, + vb.pp("model.embeddings.norm"), + )?; + let global_rotary_emb = Arc::new(RotaryEmbedding::new( + vb.dtype(), + config, + config.global_rope_theta, + vb.device(), + )?); + let local_rotary_emb = Arc::new(RotaryEmbedding::new( + vb.dtype(), + config, + config.local_rope_theta, + vb.device(), + )?); + + let mut layers = Vec::with_capacity(config.num_hidden_layers); + for layer_id in 0..config.num_hidden_layers { + let layer_uses_local_attention = layer_id % config.global_attn_every_n_layers != 0; + layers.push(ModernBertLayer::load( + vb.pp(format!("model.layers.{layer_id}")), + config, + if layer_uses_local_attention { + local_rotary_emb.clone() + } else { + global_rotary_emb.clone() + }, + layer_uses_local_attention, + )?); + } + + let final_norm = layer_norm_no_bias( + config.hidden_size, + config.layer_norm_eps, + vb.pp("model.final_norm"), + )?; + + Ok(Self { + word_embeddings, + norm, + layers, + final_norm, + local_attention_size: config.local_attention, + }) + } + + pub fn forward(&self, xs: &Tensor, mask: &Tensor) -> Result { + let seq_len = xs.shape().dims()[1]; + let global_attention_mask = + prepare_4d_attention_mask(mask, DType::F32, None)?.to_device(xs.device())?; + let local_attention_mask = + get_local_attention_mask(seq_len, self.local_attention_size / 2, xs.device())?; + let mut xs = xs.apply(&self.word_embeddings)?.apply(&self.norm)?; + for layer in self.layers.iter() { + xs = layer.forward(&xs, &global_attention_mask, &local_attention_mask)?; + } + let xs = xs.apply(&self.final_norm)?; + Ok(xs) + } +} + +// ModernBERT for the fill-mask task +#[derive(Clone)] +pub struct ModernBertForMaskedLM { + model: ModernBert, + decoder: ModernBertDecoder, + head: ModernBertHead, +} + +impl ModernBertForMaskedLM { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let model = ModernBert::load(vb.clone(), config)?; + let decoder = ModernBertDecoder::load(vb.clone(), config)?; + let head = ModernBertHead::load(vb.pp("head"), config)?; + Ok(Self { + model, + decoder, + head, + }) + } + + pub fn forward(&self, xs: &Tensor, mask: &Tensor) -> Result { + let xs = self + .model + .forward(xs, mask)? + .apply(&self.head)? + .apply(&self.decoder)?; + Ok(xs) + } +} + +#[derive(Clone)] +pub struct ModernBertClassifier { + classifier: Linear, +} + +impl ModernBertClassifier { + fn load(vb: VarBuilder, config: &Config) -> Result { + // The decoder weights are tied with the embeddings layer weights + let classifier = linear( + config.hidden_size, + config + .classifier_config + .as_ref() + .map(|cc| cc.id2label.len()) + .unwrap_or_default(), + vb.pp("classifier"), + )?; + Ok(Self { classifier }) + } +} + +impl Module for ModernBertClassifier { + fn forward(&self, xs: &Tensor) -> Result { + let xs = xs.apply(&self.classifier)?; + softmax(&xs, D::Minus1) + } +} + +#[derive(Clone)] +pub struct ModernBertForSequenceClassification { + model: ModernBert, + head: ModernBertHead, + classifier: ModernBertClassifier, + classifier_pooling: ClassifierPooling, +} + +impl ModernBertForSequenceClassification { + pub fn load(vb: VarBuilder, config: &Config) -> Result { + let model = ModernBert::load(vb.clone(), config)?; + let classifier = ModernBertClassifier::load(vb.clone(), config)?; + let head = ModernBertHead::load(vb.pp("head"), config)?; + Ok(Self { + model, + head, + classifier, + classifier_pooling: config + .classifier_config + .as_ref() + .map(|cc| cc.classifier_pooling) + .unwrap_or_default(), + }) + } + + pub fn forward(&self, xs: &Tensor, mask: &Tensor) -> Result { + let output = self.model.forward(xs, mask)?; + let last_hidden_state = match self.classifier_pooling { + ClassifierPooling::CLS => output.i((.., 0, ..))?.contiguous()?, + ClassifierPooling::MEAN => { + let unsqueezed_mask = &mask.unsqueeze(D::Minus1)?.to_dtype(DType::F32)?; + let sum_output = output.broadcast_mul(unsqueezed_mask)?.sum(1)?; + sum_output.broadcast_div(&mask.sum_keepdim(1)?.to_dtype(DType::F32)?)? + } + }; + let xs = self + .head + .forward(&last_hidden_state)? + .apply(&self.classifier)?; + Ok(xs) + } +} diff --git a/patches/candle-transformers/src/models/moondream.rs b/patches/candle-transformers/src/models/moondream.rs new file mode 100644 index 0000000000..4c0b30503e --- /dev/null +++ b/patches/candle-transformers/src/models/moondream.rs @@ -0,0 +1,366 @@ +//! MoonDream Model vision-to-text +//! +//! +//! Moondream is a computer-vision model that can answer real-world questions about images. +//! It's lightweight with only 1.6B parameters, enabling it to run on mobile phones and edge devices. +//! [MoonDream Original Implementation](https://github.com/vikhyat/moondream) +//! +//! The model consists of: +//! - Vision encoder using a ViT-style architecture +//! - Text decoder based on Microsoft's Phi model +//! - Vision projection module to align vision and text embeddings +//! +//! # Examples +//! +//! +//! +//! ```bash +//! # download an example image +//! wget https://raw.githubusercontent.com/vikhyat/moondream/main/assets/demo-1.jpg +//! +//! # Now you can run Moondream from the `candle-examples` crate: +//! cargo run --example moondream \ +//! --release -- \ +//! --prompt "What is the girl eating?" +//! --image "./demo-1.jpg" +//! +//! > avavx: false, neon: true, simd128: false, f16c: false +//! > temp: 0.00 repeat-penalty: 1.00 repeat-last-n: 64 +//! > retrieved the files in 3.395583ms +//! > Running on CPU, to run on GPU(metal), build this example with `--features metal` +//! > loaded the model in 5.485493792s +//! > loaded and encoded the image Tensor[dims 3, 378, 378; f32] in 4.801396417s +//! > starting the inference loop +//! > The girl is eating a hamburger.< +//! > 9 tokens generated (0.68 token/s) +//! ``` + +use crate::models::mixformer::{Config as PhiConfig, MixFormerSequentialForCausalLM as PhiModel}; +use crate::models::with_tracing::{layer_norm, linear_b, LayerNorm, Linear}; +use candle::{IndexOp, Module, Result, Tensor, D}; +use candle_nn::VarBuilder; + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + pub phi_config: PhiConfig, + pub vision_config: VisionConfig, +} + +impl Config { + pub fn v2() -> Self { + Self { + phi_config: PhiConfig::v1_5(), + vision_config: VisionConfig::v2(), + } + } +} + +fn scaled_dot_product_attention(q: &Tensor, k: &Tensor, v: &Tensor) -> Result { + let dim = q.dim(D::Minus1)?; + let scale_factor = 1.0 / (dim as f64).sqrt(); + let attn_weights = (q.matmul(&k.t()?)? * scale_factor)?; + candle_nn::ops::softmax_last_dim(&attn_weights)?.matmul(v) +} + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct VisionConfig { + pub(crate) image_embedding_dim: usize, + pub(crate) model_dim: usize, + pub(crate) hidden_dim: usize, + pub(crate) hidden_features: usize, + pub(crate) embed_len: usize, + pub(crate) embed_dim: usize, + pub(crate) num_blocks: usize, + pub(crate) num_heads: usize, + pub(crate) act: candle_nn::Activation, +} + +impl VisionConfig { + pub fn v2() -> Self { + Self { + image_embedding_dim: 1152, + model_dim: 2048, + hidden_dim: 2048 * 4, + hidden_features: 4304, + embed_len: 729, + embed_dim: 1152, + num_blocks: 27, + num_heads: 16, + act: candle_nn::Activation::GeluPytorchTanh, + } + } +} + +#[derive(Debug, Clone)] +struct LinearPatchEmbedding { + linear: Linear, +} + +impl LinearPatchEmbedding { + fn new(vb: VarBuilder) -> Result { + let linear = linear_b(588, 1152, true, vb.pp("linear"))?; + Ok(Self { linear }) + } +} + +impl Module for LinearPatchEmbedding { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.linear) + } +} + +#[derive(Debug, Clone)] +struct Attention { + num_heads: usize, + head_dim: usize, + qkv: Linear, + proj: Linear, + span: tracing::Span, +} + +impl Attention { + pub fn new(vb: VarBuilder, dim: usize, num_heads: usize) -> Result { + let qkv = linear_b(dim, dim * 3, true, vb.pp("qkv"))?; + let proj = linear_b(dim, dim, true, vb.pp("proj"))?; + Ok(Self { + num_heads, + head_dim: dim / num_heads, + qkv, + proj, + span: tracing::span!(tracing::Level::TRACE, "vit-attn"), + }) + } +} + +impl Module for Attention { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let (b, n, c) = xs.dims3()?; + let qkv = xs + .apply(&self.qkv)? + .reshape((b, n, 3, self.num_heads, self.head_dim))? + .permute((2, 0, 3, 1, 4))?; + let (q, k, v) = ( + qkv.i(0)?.contiguous()?, + qkv.i(1)?.contiguous()?, + qkv.i(2)?.contiguous()?, + ); + scaled_dot_product_attention(&q, &k, &v)? + .transpose(1, 2)? + .reshape((b, n, c))? + .apply(&self.proj) + } +} + +#[derive(Debug, Clone)] +struct VitBlock { + attn: Attention, + mlp: Mlp, + norm1: LayerNorm, + norm2: LayerNorm, + span: tracing::Span, +} + +impl VitBlock { + fn new(vb: VarBuilder, dim: usize, num_heads: usize, cfg: &VisionConfig) -> Result { + let attn = Attention::new(vb.pp("attn"), dim, num_heads)?; + let mlp = Mlp::new(vb.pp("mlp"), dim, cfg.hidden_features, dim, cfg.act)?; + let norm1 = layer_norm(dim, 1e-5, vb.pp("norm1"))?; + let norm2 = layer_norm(dim, 1e-5, vb.pp("norm2"))?; + Ok(Self { + attn, + mlp, + norm1, + norm2, + span: tracing::span!(tracing::Level::TRACE, "vit-block"), + }) + } +} + +impl Module for VitBlock { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let ys = xs.apply(&self.norm1)?.apply(&self.attn)?; + let xs = (xs + &ys)?; + let ys = xs.apply(&self.norm2)?.apply(&self.mlp)?; + let xs = (&xs + &ys)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct VisionTransformer { + patch_embed: LinearPatchEmbedding, + pos_embed: Tensor, + blocks: Vec, + norm: LayerNorm, + span: tracing::Span, +} + +impl VisionTransformer { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let patch_embed = LinearPatchEmbedding::new(vb.pp("patch_embed"))?; + let pos_embed = vb.get((1, cfg.embed_len, cfg.embed_dim), "pos_embed")?; + let blocks = (0..cfg.num_blocks) + .map(|i| { + VitBlock::new( + vb.pp(format!("blocks.{i}")), + cfg.embed_dim, + cfg.num_heads, + cfg, + ) + }) + .collect::>()?; + let norm = layer_norm(cfg.embed_dim, 1e-5, vb.pp("norm"))?; + Ok(Self { + patch_embed, + pos_embed, + blocks, + norm, + span: tracing::span!(tracing::Level::TRACE, "vit"), + }) + } +} + +impl Module for VisionTransformer { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let mut xs = (&xs.apply(&self.patch_embed)? + &self.pos_embed)?; + for block in self.blocks.iter() { + xs = xs.apply(block)?; + } + xs.apply(&self.norm) + } +} + +#[derive(Debug, Clone)] +pub struct Encoder { + model: VisionTransformer, +} + +impl Encoder { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let model = VisionTransformer::new(cfg, vb.pp("model.visual"))?; + Ok(Self { model }) + } +} + +impl Module for Encoder { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.model) + } +} + +#[derive(Debug, Clone)] +struct Mlp { + fc1: Linear, + act: candle_nn::Activation, + fc2: Linear, + span: tracing::Span, +} + +impl Mlp { + fn new( + vb: VarBuilder, + in_features: usize, + hidden_features: usize, + out_features: usize, + act: candle_nn::Activation, + ) -> Result { + let fc1 = linear_b(in_features, hidden_features, true, vb.pp("fc1"))?; + let fc2 = linear_b(hidden_features, out_features, true, vb.pp("fc2"))?; + Ok(Self { + fc1, + act, + fc2, + span: tracing::span!(tracing::Level::TRACE, "mlp"), + }) + } +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + xs.apply(&self.fc1)?.apply(&self.act)?.apply(&self.fc2) + } +} + +#[derive(Debug, Clone)] +struct VisionProjection { + mlp: Mlp, +} + +impl VisionProjection { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let mlp = Mlp::new( + vb.pp("mlp"), + cfg.image_embedding_dim, + cfg.hidden_dim, + cfg.model_dim, + cfg.act, + )?; + Ok(Self { mlp }) + } +} + +impl Module for VisionProjection { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.mlp) + } +} + +#[derive(Debug, Clone)] +pub struct VisionEncoder { + encoder: Encoder, + projection: VisionProjection, +} + +impl VisionEncoder { + pub fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let encoder = Encoder::new(cfg, vb.pp("encoder"))?; + let projection = VisionProjection::new(cfg, vb.pp("projection"))?; + Ok(Self { + encoder, + projection, + }) + } +} + +impl Module for VisionEncoder { + fn forward(&self, xs: &Tensor) -> Result { + let (b, c, hp1, wp2) = xs.dims4()?; + let (p1, p2) = (14, 14); + let h = hp1 / p1; + let w = wp2 / p2; + xs.reshape((b, c, h, p1, h, p2))? + .permute((0, 2, 4, 1, 3, 5))? + .reshape((b, h * w, c * p1 * p2))? + .apply(&self.encoder)? + .apply(&self.projection) + } +} + +#[derive(Debug, Clone)] +pub struct Model { + pub text_model: PhiModel, + pub vision_encoder: VisionEncoder, +} + +impl Model { + pub fn new(config: &Config, vb: VarBuilder) -> Result { + let text_model = PhiModel::new_v2(&config.phi_config, vb.pp("text_model"))?; + let vision_encoder = VisionEncoder::new(&config.vision_config, vb.pp("vision_encoder"))?; + Ok(Self { + text_model, + vision_encoder, + }) + } + + pub fn vision_encoder(&self) -> &VisionEncoder { + &self.vision_encoder + } + + pub fn text_model(&mut self) -> &mut PhiModel { + &mut self.text_model + } +} diff --git a/patches/candle-transformers/src/models/mpt.rs b/patches/candle-transformers/src/models/mpt.rs new file mode 100644 index 0000000000..d4170d6bff --- /dev/null +++ b/patches/candle-transformers/src/models/mpt.rs @@ -0,0 +1,298 @@ +//! Module implementing the MPT (Multi-Purpose Transformer) model +//! +//! References: +//! - [MPT Model used by replit-code-v1_5-3b](https://huggingface.co/replit/replit-code-v1_5-3b/blob/main/modeling_mpt.py) +//! - [Configuration](https://huggingface.co/replit/replit-code-v1_5-3b/blob/main/configuration_mpt.py) +//! +//! The model uses grouped query attention and alibi positional embeddings. + +use crate::models::with_tracing::{linear_no_bias, Embedding, Linear}; +/// MPT model used by replit-code-v1_5-3b +/// https://huggingface.co/replit/replit-code-v1_5-3b/blob/main/modeling_mpt.py +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{layer_norm, LayerNorm, VarBuilder}; + +// https://huggingface.co/replit/replit-code-v1_5-3b/blob/main/configuration_mpt.py +#[derive(Debug, Clone, PartialEq)] +pub struct Config { + pub(crate) d_model: usize, + pub(crate) n_heads: usize, + pub(crate) n_layers: usize, + pub(crate) expansion_ratio: usize, + pub(crate) max_seq_len: usize, + pub(crate) vocab_size: usize, + pub(crate) kv_n_heads: usize, + pub(crate) attn_prefix_lm: bool, + pub(crate) attn_alibi: bool, + pub(crate) attn_alibi_bias_max: usize, +} + +impl Config { + pub fn replit_code_v1_5_3b() -> Self { + Self { + d_model: 3072, + n_heads: 24, + n_layers: 32, + expansion_ratio: 4, + max_seq_len: 4096, + vocab_size: 32768, + kv_n_heads: 8, + attn_prefix_lm: false, + attn_alibi: true, + attn_alibi_bias_max: 8, + } + } + + pub fn is_causal(&self) -> bool { + !self.attn_prefix_lm + } +} + +#[derive(Debug, Clone)] +struct GroupedQueryAttention { + wqkv: Linear, + out_proj: Linear, + kv_cache: Option<(Tensor, Tensor)>, + softmax_scale: f64, + head_dim: usize, + d_model: usize, + n_heads: usize, + kv_n_heads: usize, + attn_bias: Tensor, + span: tracing::Span, +} + +impl GroupedQueryAttention { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let head_dim = cfg.d_model / cfg.n_heads; + let wqkv_size = cfg.d_model + 2 * cfg.kv_n_heads * head_dim; + let wqkv = linear_no_bias(cfg.d_model, wqkv_size, vb.pp("Wqkv"))?; + let softmax_scale = 1f64 / (head_dim as f64).sqrt(); + let out_proj = linear_no_bias(cfg.d_model, cfg.d_model, vb.pp("out_proj"))?; + let attn_bias = build_alibi_bias(cfg)?.to_device(vb.device())?; + Ok(Self { + wqkv, + out_proj, + kv_cache: None, + softmax_scale, + head_dim, + d_model: cfg.d_model, + n_heads: cfg.n_heads, + kv_n_heads: cfg.kv_n_heads, + attn_bias, + span: tracing::span!(tracing::Level::TRACE, "gqa"), + }) + } + + fn forward(&mut self, xs: &Tensor, mask: Option<&Tensor>) -> Result { + let _enter = self.span.enter(); + let (b_size, seq_len, _n_embd) = xs.dims3()?; + let qkv = self.wqkv.forward(xs)?; + let query = qkv.narrow(2, 0, self.d_model)?; + let kv_size = self.kv_n_heads * self.head_dim; + let key = qkv.narrow(2, self.d_model, kv_size)?; + let value = qkv.narrow(2, self.d_model + kv_size, kv_size)?; + // scaled_multihead_dot_product_attention + let query = query + .reshape((b_size, seq_len, self.n_heads, ()))? + .transpose(1, 2)?; // b,h,s,d + let key = key + .reshape((b_size, seq_len, self.kv_n_heads, ()))? + .permute((0, 2, 3, 1))?; // b,h,d,s + let value = value + .reshape((b_size, seq_len, self.kv_n_heads, ()))? + .transpose(1, 2)?; // b,h,s,d + let (key, value) = match &self.kv_cache { + None => (key, value), + Some((prev_k, prev_v)) => { + let k = Tensor::cat(&[prev_k, &key], 3)?; + let v = Tensor::cat(&[prev_v, &value], 2)?; + (k, v) + } + }; + self.kv_cache = Some((key.clone(), value.clone())); + let query = query.contiguous()?; + let key = crate::utils::repeat_kv(key, self.n_heads / self.kv_n_heads)?.contiguous()?; + let value = crate::utils::repeat_kv(value, self.n_heads / self.kv_n_heads)?.contiguous()?; + let attn_weights = (query.matmul(&key)? * self.softmax_scale)?; + let attn_bias = { + let s_q = query.dim(D::Minus2)?; + let s_k = key.dim(D::Minus1)?; + let (_, _, a_q, a_k) = self.attn_bias.dims4()?; + let start_q = a_q.saturating_sub(s_q); + let start_k = a_k.saturating_sub(s_k); + self.attn_bias.i((.., .., start_q.., start_k..))? + }; + let attn_weights = attn_weights.broadcast_add(&attn_bias)?; + let attn_weights = match mask { + None => attn_weights, + Some(mask) => masked_fill( + &attn_weights, + &mask.broadcast_as(attn_weights.shape())?, + f32::NEG_INFINITY, + )?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + let attn_output = attn_weights + .matmul(&value)? + .transpose(1, 2)? + .flatten_from(D::Minus2)?; + let out = attn_output.apply(&self.out_proj)?; + Ok(out) + } +} + +#[derive(Debug, Clone)] +struct Ffn { + up_proj: Linear, + down_proj: Linear, +} + +impl Ffn { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden = cfg.d_model * cfg.expansion_ratio; + let up_proj = linear_no_bias(cfg.d_model, hidden, vb.pp("up_proj"))?; + let down_proj = linear_no_bias(hidden, cfg.d_model, vb.pp("down_proj"))?; + Ok(Self { up_proj, down_proj }) + } +} + +impl Module for Ffn { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.up_proj)?.gelu_erf()?.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +struct MPTBlock { + norm1: LayerNorm, // Do we need the low-precision variant? + attn: GroupedQueryAttention, + norm2: LayerNorm, + ffn: Ffn, +} + +impl MPTBlock { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let ln_cfg = candle_nn::LayerNormConfig { + affine: false, + ..Default::default() + }; + let norm1 = layer_norm(cfg.d_model, ln_cfg, vb.pp("norm_1"))?; + let norm2 = layer_norm(cfg.d_model, ln_cfg, vb.pp("norm_2"))?; + let attn = GroupedQueryAttention::new(cfg, vb.pp("attn"))?; + let ffn = Ffn::new(cfg, vb.pp("ffn"))?; + Ok(Self { + norm1, + attn, + norm2, + ffn, + }) + } + + fn forward(&mut self, xs: &Tensor, mask: Option<&Tensor>) -> Result { + let residual = xs; + let xs = xs.apply(&self.norm1)?; + let xs = self.attn.forward(&xs, mask)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.norm2)?.apply(&self.ffn)?; + xs + residual + } +} + +pub(crate) fn build_alibi_bias(cfg: &Config) -> Result { + let full = !cfg.is_causal(); + let seq_len = cfg.max_seq_len; + let alibi_bias = Tensor::arange(1 - seq_len as i64, 1, &Device::Cpu)?; + let alibi_bias = if full { + let a1 = alibi_bias.reshape((1, 1, 1, seq_len))?; + let a2 = alibi_bias.reshape((1, 1, seq_len, 1))?; + a1.broadcast_sub(&a2)?.abs()?.neg()? + } else { + alibi_bias.reshape((1, 1, 1, seq_len))? + }; + let mut n_heads2 = 1; + while n_heads2 < cfg.n_heads { + n_heads2 *= 2 + } + let slopes = (1..=n_heads2) + .map(|v| 1f32 / 2f32.powf((v * cfg.attn_alibi_bias_max) as f32 / n_heads2 as f32)) + .collect::>(); + let slopes = if n_heads2 == cfg.n_heads { + slopes + } else { + slopes + .iter() + .skip(1) + .step_by(2) + .chain(slopes.iter().step_by(2)) + .take(cfg.n_heads) + .cloned() + .collect::>() + }; + let slopes = Tensor::new(slopes, &Device::Cpu)?.reshape((1, (), 1, 1))?; + alibi_bias.to_dtype(DType::F32)?.broadcast_mul(&slopes) +} + +#[derive(Debug, Clone)] +pub struct Model { + wte: Embedding, + blocks: Vec, + norm_f: LayerNorm, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let wte = Embedding::new(cfg.vocab_size, cfg.d_model, vb.pp("wte"))?; + let vb_b = vb.pp("blocks"); + let mut blocks = Vec::with_capacity(cfg.n_layers); + for i in 0..cfg.n_layers { + let block = MPTBlock::new(cfg, vb_b.pp(i))?; + blocks.push(block) + } + let ln_cfg = candle_nn::LayerNormConfig { + affine: false, + ..Default::default() + }; + let norm_f = candle_nn::layer_norm(cfg.d_model, ln_cfg, vb.pp("norm_f"))?; + Ok(Self { + wte, + blocks, + norm_f, + }) + } + + pub fn forward(&mut self, xs: &Tensor) -> Result { + let (_b_size, seq_len) = xs.dims2()?; + let mut xs = xs.apply(&self.wte)?; + let mask = if seq_len <= 1 { + None + } else { + Some(get_mask(seq_len, xs.device())?) + }; + for block in self.blocks.iter_mut() { + xs = block.forward(&xs, mask.as_ref())?; + } + let xs = xs.apply(&self.norm_f)?; + let logits = xs + .narrow(1, seq_len - 1, 1)? + .squeeze(1)? + .matmul(&self.wte.embeddings().t()?)? + .squeeze(1)?; + Ok(logits) + } +} + +pub(crate) fn get_mask(size: usize, device: &Device) -> Result { + let mask: Vec<_> = (0..size) + .flat_map(|i| (0..size).map(move |j| u8::from(j > i))) + .collect(); + Tensor::from_slice(&mask, (size, size), device) +} + +pub(crate) 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) +} diff --git a/patches/candle-transformers/src/models/nvembed_v2/embedding.rs b/patches/candle-transformers/src/models/nvembed_v2/embedding.rs new file mode 100644 index 0000000000..a52192afdf --- /dev/null +++ b/patches/candle-transformers/src/models/nvembed_v2/embedding.rs @@ -0,0 +1,294 @@ +/// Mistral LLM, https://github.com/mistralai/mistral-src +use crate::models::{ + mistral::Config, + with_tracing::{linear_no_bias, Linear, RmsNorm}, +}; +use crate::utils::repeat_kv; +use candle::{DType, Device, Module, Result, Tensor}; +use candle_nn::{Activation, VarBuilder}; +use std::sync::Arc; + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let rope_theta = cfg.rope_theta as f32; + let dim = cfg.hidden_size / cfg.num_attention_heads; + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / rope_theta.powf(i as f32 / dim as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope(q, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(k, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: Activation, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + let gate_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("gate_proj"))?; + let up_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("up_proj"))?; + let down_proj = linear_no_bias(intermediate_sz, hidden_sz, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_act, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = xs.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + hidden_size: usize, + rotary_emb: Arc, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + let head_dim = hidden_sz / num_heads; + let q_proj = linear_no_bias(hidden_sz, num_heads * head_dim, vb.pp("q_proj"))?; + let k_proj = linear_no_bias(hidden_sz, num_kv_heads * head_dim, vb.pp("k_proj"))?; + let v_proj = linear_no_bias(hidden_sz, num_kv_heads * head_dim, vb.pp("v_proj"))?; + let o_proj = linear_no_bias(num_heads * head_dim, hidden_sz, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + hidden_size: hidden_sz, + rotary_emb, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let (query_states, key_states) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states, seqlen_offset)?; + + let key_states = repeat_kv(key_states, self.num_kv_groups)?; + let value_states = repeat_kv(value_states, self.num_kv_groups)?; + + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + let attn_output = attn_weights.matmul(&value_states)?; + + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.hidden_size))? + .apply(&self.o_proj) + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: MLP, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl DecoderLayer { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let self_attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let input_layernorm = + RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let post_attention_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + + let xs = self.self_attn.forward(&xs, attention_mask, seqlen_offset)?; + + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; + residual + xs + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + pub cfg: Config, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(rotary_emb.clone(), cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("norm"))?; + Ok(Self { + embed_tokens, + layers, + norm, + cfg: cfg.clone(), + }) + } + + // Attn mask used to mask out padding tokens + pub fn forward( + &mut self, + attn_mask: &Tensor, + input_ids: &Tensor, + dtype: DType, + ) -> Result { + let mut xs = self.embed_tokens.forward(input_ids)?; + + // Expand to 4d mask for sdpa + let attn_mask = prepare_4d_attention_mask(attn_mask, dtype, None)?; + + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, Some(&attn_mask), 0)?; + } + + // Return hiddens instead of logits + xs.apply(&self.norm) + } +} + +fn prepare_4d_attention_mask( + mask: &Tensor, + dtype: DType, + tgt_len: Option, +) -> Result { + let bsz = mask.dims()[0]; + let src_len = mask.dims()[1]; + let tgt_len = tgt_len.unwrap_or(src_len); + + let expanded_mask = mask + .unsqueeze(1)? + .unsqueeze(2)? + .expand((bsz, 1, tgt_len, src_len))? + .to_dtype(dtype)?; + + let inverted_mask = (1.0 - expanded_mask)?; + + (inverted_mask * get_dtype_min_val(dtype))?.to_dtype(dtype) +} + +fn get_dtype_min_val(dtype: DType) -> f64 { + match dtype { + DType::F32 => f32::MIN as f64, + DType::F64 => f64::MIN, + _ => panic!("Unsupported data type"), + } +} diff --git a/patches/candle-transformers/src/models/nvembed_v2/mod.rs b/patches/candle-transformers/src/models/nvembed_v2/mod.rs new file mode 100644 index 0000000000..8a8f700782 --- /dev/null +++ b/patches/candle-transformers/src/models/nvembed_v2/mod.rs @@ -0,0 +1,18 @@ +//! NV-Embed-v2 +//! +//! NV-Embed-v2 is a text embedding model that combines a Mistral decoder with a latent attention mechanism to produce high-quality text embeddings. +//! +//! This implementation is based on the [paper](https://arxiv.org/pdf/2405.17428) and [weights](https://huggingface.co/nvidia/NV-Embed-v2) +//! +//! # Query-Passage Retrieval Example +//! ```bash +//! cargo run --example nvembed_v2 --release +//! ``` +//! +//! # Sentence Embedding Example +//! ```bash +//! cargo run --example nvembed_v2 --release -- --prompt "Here is a test sentence" +//! ``` + +pub mod embedding; +pub mod model; diff --git a/patches/candle-transformers/src/models/nvembed_v2/model.rs b/patches/candle-transformers/src/models/nvembed_v2/model.rs new file mode 100644 index 0000000000..73ef776e3b --- /dev/null +++ b/patches/candle-transformers/src/models/nvembed_v2/model.rs @@ -0,0 +1,233 @@ +use super::embedding::Model as EmbeddingModel; +use crate::models::{ + mistral::Config, + with_tracing::{layer_norm, linear, linear_no_bias, LayerNorm, Linear}, +}; +use candle::{DType, Device, Result, Tensor, D}; +use candle_nn::{ops::softmax_last_dim, LayerNormConfig, Module, VarBuilder}; + +// Geglu and feedforward from candle-transformers/src/models/stable_diffusion/attention.rs +#[derive(Debug)] +struct GeGlu { + proj: Linear, + span: tracing::Span, +} + +impl GeGlu { + fn new(vs: VarBuilder, dim_in: usize, dim_out: usize) -> Result { + let proj = linear(dim_in, dim_out * 2, vs)?; + let span = tracing::span!(tracing::Level::TRACE, "geglu"); + Ok(Self { proj, span }) + } +} + +impl Module for GeGlu { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let hidden_states_and_gate = self.proj.forward(xs)?.chunk(2, D::Minus1)?; + &hidden_states_and_gate[0] * hidden_states_and_gate[1].gelu()? + } +} + +#[derive(Debug)] +struct FeedForward { + project_in: GeGlu, + linear: Linear, + span: tracing::Span, +} + +impl FeedForward { + fn new(vs: VarBuilder, dim: usize, dim_out: Option, mult: usize) -> Result { + let inner_dim = dim * mult; + let dim_out = dim_out.unwrap_or(dim); + let vs = vs.pp("net"); + let project_in = GeGlu::new(vs.pp("0"), dim, inner_dim)?; + let linear = linear(inner_dim, dim_out, vs.pp("2"))?; + let span = tracing::span!(tracing::Level::TRACE, "ff"); + Ok(Self { + project_in, + linear, + span, + }) + } +} + +impl Module for FeedForward { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let xs = self.project_in.forward(xs)?; + self.linear.forward(&xs) + } +} + +// CrossAttention from candle-transformers/src/models/stable_diffusion/attention.rs +#[derive(Debug)] +struct CrossAttention { + to_q: Linear, + to_kv: Linear, + to_out: Linear, + heads: usize, + scale: f64, + span: tracing::Span, + span_attn: tracing::Span, + span_softmax: tracing::Span, +} + +impl CrossAttention { + fn new( + vs: VarBuilder, + query_dim: usize, + context_dim: Option, + heads: usize, + dim_head: usize, + ) -> Result { + let inner_dim = dim_head * heads; + let context_dim = context_dim.unwrap_or(query_dim); + let scale = 1.0 / f64::sqrt(dim_head as f64); + let to_q = linear_no_bias(query_dim, inner_dim, vs.pp("to_q"))?; + let to_kv = linear_no_bias(context_dim, inner_dim * 2, vs.pp("to_kv"))?; + let to_out = linear_no_bias(inner_dim, query_dim, vs.pp("to_out"))?; + let span = tracing::span!(tracing::Level::TRACE, "xa"); + let span_attn = tracing::span!(tracing::Level::TRACE, "xa-attn"); + let span_softmax = tracing::span!(tracing::Level::TRACE, "xa-softmax"); + Ok(Self { + to_q, + to_kv, + to_out, + heads, + scale, + span, + span_attn, + span_softmax, + }) + } + + fn reshape_heads_to_batch_dim(&self, xs: &Tensor) -> Result { + let (batch_size, seq_len, dim) = xs.dims3()?; + xs.reshape((batch_size, seq_len, self.heads, dim / self.heads))? + .transpose(1, 2)? + .reshape((batch_size * self.heads, seq_len, dim / self.heads)) + } + + fn reshape_batch_dim_to_heads(&self, xs: &Tensor) -> Result { + let (batch_size, seq_len, dim) = xs.dims3()?; + xs.reshape((batch_size / self.heads, self.heads, seq_len, dim))? + .transpose(1, 2)? + .reshape((batch_size / self.heads, seq_len, dim * self.heads)) + } + + fn attention(&self, query: &Tensor, key: &Tensor, value: &Tensor) -> Result { + let _enter = self.span_attn.enter(); + + let in_dtype = query.dtype(); + let query = query.to_dtype(DType::F32)?; + let key = key.to_dtype(DType::F32)?; + let value = value.to_dtype(DType::F32)?; + let xs = query.matmul(&(key.t()? * self.scale)?)?; + let xs = { + let _enter = self.span_softmax.enter(); + softmax_last_dim(&xs)? + }; + let xs = xs.matmul(&value)?.to_dtype(in_dtype)?; + + self.reshape_batch_dim_to_heads(&xs) + } + + fn forward(&self, xs: &Tensor, context: Option<&Tensor>) -> Result { + let _enter = self.span.enter(); + let query = self.to_q.forward(xs)?; + let context = context.unwrap_or(xs).contiguous()?; + let kv_chunks = self + .to_kv + .forward(&context)? + .chunk(2, context.shape().dims().len() - 1)?; + let (key, value) = (kv_chunks[0].clone(), kv_chunks[1].clone()); + let query = self.reshape_heads_to_batch_dim(&query)?; + let key = self.reshape_heads_to_batch_dim(&key)?; + let value = self.reshape_heads_to_batch_dim(&value)?; + + let xs = self.attention(&query, &key, &value)?; + self.to_out.forward(&xs) + } +} + +#[derive(Debug)] +pub struct Model { + embedding_model: EmbeddingModel, + cross_attn: CrossAttention, + cross_attn_norm: LayerNorm, + cross_attn_context_norm: LayerNorm, + ff: FeedForward, + ff_norm: LayerNorm, + latents: Tensor, + pub device: Device, + pub dtype: DType, +} + +impl Model { + pub fn new(vb: VarBuilder) -> Result { + // Embedding model + let cfg = Config::config_7b_v0_1(false); + let embedding_model = EmbeddingModel::new(&cfg, vb.pp("embedding_model"))?; + + // Latent attention + let dim = 4096; + let vb = vb.pp("latent_attention_model"); + let latents = vb.get((512, dim), "latents")?; + + // Cross attend blocks + let vb = vb.pp("cross_attend_blocks"); + let cross_attn_norm = layer_norm(dim, LayerNormConfig::default(), vb.pp("0.norm"))?; + let cross_attn_context_norm = layer_norm( + dim, + candle_nn::LayerNormConfig::default(), + vb.pp("0.norm_context"), + )?; + let cross_attn = CrossAttention::new(vb.pp("0.fn"), dim, None, 8, 4096)?; + + let ff_norm = layer_norm(dim, LayerNormConfig::default(), vb.pp("1.norm"))?; + let ff = FeedForward::new(vb.pp("1.fn"), dim, None, 4)?; + + Ok(Self { + embedding_model, + cross_attn, + cross_attn_norm, + cross_attn_context_norm, + ff, + ff_norm, + latents, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + pub fn forward( + &mut self, + input_ids: &Tensor, + attn_mask: &Tensor, + pool_mask: &Tensor, + ) -> Result { + // Embedding model + let hiddens = self + .embedding_model + .forward(attn_mask, input_ids, self.dtype)?; + + // Latent attention + let b = hiddens.dims()[0]; + let x = self.latents.unsqueeze(0)?.repeat((b, 1, 1))?; + let original_hiddens = &hiddens; + + let hiddens = self.cross_attn_norm.forward(original_hiddens)?; + let x = self.cross_attn_context_norm.forward(&x)?; + let cross_hiddens = (self.cross_attn.forward(&hiddens, Some(&x))? + original_hiddens)?; + + let hiddens = self.ff_norm.forward(&cross_hiddens)?; + let hiddens = (self.ff.forward(&hiddens)? + cross_hiddens)?; + + // Mean pooling + let hiddens_masked = hiddens.broadcast_mul(&pool_mask.unsqueeze(D::Minus1)?)?; + let s = hiddens_masked.sum(1)?; + let d = pool_mask.sum_keepdim(1)?; + s.broadcast_div(&d) + } +} diff --git a/patches/candle-transformers/src/models/olmo.rs b/patches/candle-transformers/src/models/olmo.rs new file mode 100644 index 0000000000..6cf5b1f79d --- /dev/null +++ b/patches/candle-transformers/src/models/olmo.rs @@ -0,0 +1,353 @@ +//! OLMo (Open Language Model) implementation +//! +//! See OLMo model details at: +//! - [Hugging Face](https://huggingface.co/allenai/OLMo) +//! - [OLMo Paper](https://allenai.org/olmo) +//! +//! The model uses: +//! - RoPE embeddings +//! - Sliding window attention +//! - Transformer architecture +//! +//! References: +//! - [Hugging Face Implementation](https://huggingface.co/allenai/OLMo) +//! - [OLMo Paper](https://allenai.org/olmo) +//! + +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::{linear_b, linear_no_bias, Activation, LayerNorm, Linear, VarBuilder}; +use std::sync::Arc; + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub attention_bias: bool, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub hidden_act: candle_nn::Activation, + pub max_position_embeddings: usize, + pub rope_theta: f64, + pub tie_word_embeddings: bool, + pub clip_qkv: Option, +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let dim = cfg.hidden_size / cfg.num_attention_heads; + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: Activation, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + let gate_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("gate_proj"))?; + let up_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("up_proj"))?; + let down_proj = linear_no_bias(intermediate_sz, hidden_sz, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_act, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = xs.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + hidden_size: usize, + rotary_emb: Arc, + qkv_clip: Option, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + let head_dim = hidden_sz / num_heads; + let b = cfg.attention_bias; + let qkv_clip = cfg.clip_qkv; + let q_proj = linear_b(hidden_sz, num_heads * head_dim, b, vb.pp("q_proj"))?; + let k_proj = linear_b(hidden_sz, num_kv_heads * head_dim, b, vb.pp("k_proj"))?; + let v_proj = linear_b(hidden_sz, num_kv_heads * head_dim, b, vb.pp("v_proj"))?; + let o_proj = linear_b(num_heads * head_dim, hidden_sz, b, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + hidden_size: hidden_sz, + rotary_emb, + qkv_clip, + kv_cache: None, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let (query_states, key_states, value_states) = match &self.qkv_clip { + None => (query_states, key_states, value_states), + Some(qkv_clip) => { + let query_states = Tensor::clamp(&query_states, -qkv_clip, *qkv_clip)?; + let key_states = Tensor::clamp(&key_states, -qkv_clip, *qkv_clip)?; + let value_states = Tensor::clamp(&value_states, -qkv_clip, *qkv_clip)?; + (query_states, key_states, value_states) + } + }; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let (query_states, key_states) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states, seqlen_offset)?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?.contiguous()?; + let value_states = + crate::utils::repeat_kv(value_states, self.num_kv_groups)?.contiguous()?; + + let attn_output = { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + }; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.hidden_size))? + .apply(&self.o_proj) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: MLP, + input_layernorm: LayerNorm, + post_attention_layernorm: LayerNorm, +} + +impl DecoderLayer { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let self_attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let ln_weight = Tensor::ones(cfg.hidden_size, vb.dtype(), vb.device())?; + let input_layernorm = LayerNorm::new_no_bias(ln_weight.clone(), 1e-5); + let post_attention_layernorm = LayerNorm::new_no_bias(ln_weight.clone(), 1e-5); + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; + residual + xs + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: LayerNorm, + lm_head: Linear, + device: Device, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb_m.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(rotary_emb.clone(), cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let ln_weight = Tensor::ones(cfg.hidden_size, vb.dtype(), vb.device())?; + let norm = LayerNorm::new_no_bias(ln_weight, 1e-5); + let lm_head = if cfg.tie_word_embeddings { + Linear::new(embed_tokens.embeddings().clone(), None) + } else { + linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))? + }; + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + // Sliding window mask? + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { f32::NEG_INFINITY } else { 0. })) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), self.dtype, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let (b_size, seq_len) = input_ids.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(b_size, seq_len, seqlen_offset)?; + Some(mask) + }; + let mut xs = self.embed_tokens.forward(input_ids)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)? + } + xs.narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +} diff --git a/patches/candle-transformers/src/models/olmo2.rs b/patches/candle-transformers/src/models/olmo2.rs new file mode 100644 index 0000000000..5567cb67f8 --- /dev/null +++ b/patches/candle-transformers/src/models/olmo2.rs @@ -0,0 +1,348 @@ +//! OLMo 2 (Open Language Model) implementation +//! +//! See OLMo 2 model details at: +//! - [Hugging Face Collection](https://huggingface.co/collections/allenai/olmo-2-674117b93ab84e98afc72edc) +//! - [OLMo 2 Paper](https://arxiv.org/abs/2501.00656) +//! +//! +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::{linear_b, linear_no_bias, rms_norm, Activation, Linear, RmsNorm, VarBuilder}; +use std::sync::Arc; + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub attention_bias: bool, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub rms_norm_eps: f64, + pub hidden_act: candle_nn::Activation, + pub max_position_embeddings: usize, + pub rope_theta: f64, + pub tie_word_embeddings: bool, + pub clip_qkv: Option, +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let dim = cfg.hidden_size / cfg.num_attention_heads; + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: Activation, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + let gate_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("gate_proj"))?; + let up_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("up_proj"))?; + let down_proj = linear_no_bias(intermediate_sz, hidden_sz, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_act, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = xs.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + q_norm: RmsNorm, + k_norm: RmsNorm, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + hidden_size: usize, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + let head_dim = hidden_sz / num_heads; + let b = cfg.attention_bias; + let q_proj = linear_b(hidden_sz, num_heads * head_dim, b, vb.pp("q_proj"))?; + let k_proj = linear_b(hidden_sz, num_kv_heads * head_dim, b, vb.pp("k_proj"))?; + let v_proj = linear_b(hidden_sz, num_kv_heads * head_dim, b, vb.pp("v_proj"))?; + let o_proj = linear_b(num_heads * head_dim, hidden_sz, b, vb.pp("o_proj"))?; + let q_norm = rms_norm(hidden_sz, cfg.rms_norm_eps, vb.pp("q_norm"))?; + let k_norm = rms_norm(num_kv_heads * head_dim, cfg.rms_norm_eps, vb.pp("k_norm"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + q_norm, + k_norm, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + hidden_size: hidden_sz, + rotary_emb, + kv_cache: None, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = self.q_norm.forward(&query_states)?; + let key_states = self.k_norm.forward(&key_states)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let (query_states, key_states) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states, seqlen_offset)?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?.contiguous()?; + let value_states = + crate::utils::repeat_kv(value_states, self.num_kv_groups)?.contiguous()?; + + let attn_output = { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + }; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.hidden_size))? + .apply(&self.o_proj) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: MLP, + post_attention_layernorm: RmsNorm, + post_feedforward_layernorm: RmsNorm, +} + +impl DecoderLayer { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let self_attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let post_feedforward_layernorm = rms_norm( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_feedforward_layernorm"), + )?; + let post_attention_layernorm = rms_norm( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + post_attention_layernorm, + post_feedforward_layernorm, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.self_attn.forward(xs, attention_mask, seqlen_offset)?; + let xs = self.post_attention_layernorm.forward(&xs)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = self.mlp.forward(&xs)?; + let xs = self.post_feedforward_layernorm.forward(&xs)?; + residual + xs + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: Linear, + device: Device, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb_m.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(rotary_emb.clone(), cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb_m.pp("norm"))?; + let lm_head = if cfg.tie_word_embeddings { + Linear::new(embed_tokens.embeddings().clone(), None) + } else { + linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))? + }; + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + // Sliding window mask? + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { f32::NEG_INFINITY } else { 0. })) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), self.dtype, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let (b_size, seq_len) = input_ids.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(b_size, seq_len, seqlen_offset)?; + Some(mask) + }; + let mut xs = self.embed_tokens.forward(input_ids)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)? + } + xs.narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +} diff --git a/patches/candle-transformers/src/models/openclip/mod.rs b/patches/candle-transformers/src/models/openclip/mod.rs new file mode 100644 index 0000000000..b3864b815e --- /dev/null +++ b/patches/candle-transformers/src/models/openclip/mod.rs @@ -0,0 +1,13 @@ +//! Open Contrastive Language-Image Pre-Training +//! +//! Open Contrastive Language-Image Pre-Training (OpenCLIP) is an architecture trained on +//! pairs of images with related texts. +//! +//! - 💻 [GH Link](https://github.com/mlfoundations/open_clip) +//! - 📝 [Paper](https://arxiv.org/abs/2212.07143) +//! +//! ## Overview +//! +//! ![](https://raw.githubusercontent.com/mlfoundations/open_clip/main/docs/CLIP.png) + +pub mod text_model; diff --git a/patches/candle-transformers/src/models/openclip/text_model.rs b/patches/candle-transformers/src/models/openclip/text_model.rs new file mode 100644 index 0000000000..7b444e797e --- /dev/null +++ b/patches/candle-transformers/src/models/openclip/text_model.rs @@ -0,0 +1,266 @@ +//! Text encoder as used in most OpenCLIP pretrained models +//! https://github.com/mlfoundations/open_clip + +use candle::{DType, IndexOp, Result, Tensor, D}; +use candle_nn::{ + embedding, layer_norm, linear, ops::softmax_last_dim, Embedding, LayerNorm, Linear, Module, + VarBuilder, +}; + +#[derive(Debug, Clone)] +pub struct Config { + pub vocab_size: usize, + pub embed_dim: usize, + pub intermediate_size: usize, + pub max_position_embeddings: usize, + pub pad_with: Option, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub projection_dim: usize, +} + +impl Config { + pub fn vit_base_patch32() -> Self { + Self { + vocab_size: 49408, + embed_dim: 512, + intermediate_size: 2048, + max_position_embeddings: 77, + pad_with: None, + num_hidden_layers: 12, + num_attention_heads: 8, + projection_dim: 512, + } + } +} + +#[derive(Clone, Debug)] +struct TextEmbeddings { + token_embedding: Embedding, + position_embedding: Tensor, +} + +impl TextEmbeddings { + fn new(vs: VarBuilder, c: &Config) -> Result { + let token_embedding = embedding(c.vocab_size, c.embed_dim, vs.pp("token_embedding"))?; + let position_embedding = vs.get( + (c.max_position_embeddings, c.embed_dim), + "positional_embedding", + )?; + Ok(TextEmbeddings { + token_embedding, + position_embedding, + }) + } +} + +impl Module for TextEmbeddings { + fn forward(&self, input_ids: &Tensor) -> Result { + let seq_length = input_ids.dim(D::Minus1)?; + let inputs_embeds = self.token_embedding.forward(input_ids)?; + + let position_embedding = self.position_embedding.narrow(0, 0, seq_length)?; + + inputs_embeds.broadcast_add(&position_embedding) + } +} + +#[derive(Clone, Debug)] +struct Attention { + k_proj: candle_nn::Linear, + v_proj: candle_nn::Linear, + q_proj: candle_nn::Linear, + out_proj: Linear, + head_dim: usize, + scale: f64, + num_attention_heads: usize, +} + +impl Attention { + fn new(vs: candle_nn::VarBuilder, c: &Config) -> Result { + let embed_dim = c.embed_dim; + let num_attention_heads = c.num_attention_heads; + + let in_proj_weights = vs + .get((embed_dim * 3, embed_dim), "in_proj_weight")? + .chunk(3, 0)?; + let (q_w, k_w, v_w) = ( + &in_proj_weights[0], + &in_proj_weights[1], + &in_proj_weights[2], + ); + let in_proj_biases = vs.get(embed_dim * 3, "in_proj_bias")?.chunk(3, 0)?; + let (q_b, k_b, v_b) = (&in_proj_biases[0], &in_proj_biases[1], &in_proj_biases[2]); + + let q_proj = Linear::new(q_w.clone(), Some(q_b.clone())); + let k_proj = Linear::new(k_w.clone(), Some(k_b.clone())); + let v_proj = Linear::new(v_w.clone(), Some(v_b.clone())); + let out_proj = candle_nn::linear(embed_dim, embed_dim, vs.pp("out_proj"))?; + let head_dim = embed_dim / num_attention_heads; + let scale = (head_dim as f64).powf(-0.5); + + Ok(Attention { + k_proj, + v_proj, + q_proj, + out_proj, + head_dim, + scale, + num_attention_heads, + }) + } + + fn shape_multihead(&self, xs: &Tensor, bsz: usize, seq_len: usize) -> Result { + xs.reshape((bsz, seq_len, self.num_attention_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()? + .to_dtype(DType::F32) + } + + fn forward(&self, xs: &Tensor) -> Result { + let in_dtype = xs.dtype(); + let (bsz, seq_len, embed_dim) = xs.dims3()?; + + let q = self.shape_multihead(&self.q_proj.forward(xs)?, bsz, seq_len)?; + let k = self.shape_multihead(&self.k_proj.forward(xs)?, bsz, seq_len)?; + let v = self.shape_multihead(&self.v_proj.forward(xs)?, bsz, seq_len)?; + let q = (q * self.scale)?; + + let attn_weights = q.matmul(&k.transpose(D::Minus1, D::Minus2)?)?; + + let attn_weights = softmax_last_dim(&attn_weights)?; + + let attn_output = attn_weights.matmul(&v)?.to_dtype(in_dtype)?; + let attn_output = attn_output + .transpose(1, 2)? + .contiguous()? + .reshape((bsz, seq_len, embed_dim))?; + let out = self.out_proj.forward(&attn_output)?; + Ok(out) + } +} + +#[derive(Clone, Debug)] +struct Mlp { + fc1: Linear, + fc2: Linear, +} + +impl Mlp { + fn new(vs: VarBuilder, c: &Config) -> Result { + let fc1 = linear(c.embed_dim, c.intermediate_size, vs.pp("c_fc"))?; + let fc2 = linear(c.intermediate_size, c.embed_dim, vs.pp("c_proj"))?; + + Ok(Mlp { fc1, fc2 }) + } +} + +impl Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.fc1.forward(xs)?; + self.fc2.forward(&xs.gelu_erf()?) + } +} + +#[derive(Clone, Debug)] +struct EncoderLayer { + self_attn: Attention, + layer_norm1: LayerNorm, + mlp: Mlp, + layer_norm2: LayerNorm, +} + +impl EncoderLayer { + fn new(vs: VarBuilder, c: &Config) -> Result { + let self_attn = Attention::new(vs.pp("attn"), c)?; + let layer_norm1 = layer_norm(c.embed_dim, 1e-5, vs.pp("ln_1"))?; + let mlp = Mlp::new(vs.pp("mlp"), c)?; + let layer_norm2 = layer_norm(c.embed_dim, 1e-5, vs.pp("ln_2"))?; + + Ok(EncoderLayer { + self_attn, + layer_norm1, + mlp, + layer_norm2, + }) + } + + fn forward(&self, xs: &Tensor) -> Result { + let residual = xs; + let xs = self.layer_norm1.forward(xs)?; + let xs = self.self_attn.forward(&xs)?; + let xs = (xs + residual)?; + + let residual = &xs; + let xs = self.layer_norm2.forward(&xs)?; + let xs = self.mlp.forward(&xs)?; + let out = (xs + residual)?; + Ok(out) + } +} + +#[derive(Clone, Debug)] +pub struct Encoder { + layers: Vec, +} + +impl Encoder { + pub fn new(vs: VarBuilder, c: &Config) -> Result { + let vs = vs.pp("resblocks"); + let mut layers: Vec = Vec::new(); + for index in 0..c.num_hidden_layers { + let layer = EncoderLayer::new(vs.pp(index.to_string()), c)?; + layers.push(layer) + } + Ok(Encoder { layers }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let mut xs = xs.clone(); + for layer in self.layers.iter() { + xs = layer.forward(&xs)?; + } + Ok(xs) + } +} + +/// A text transformer as used in CLIP variants. +#[derive(Clone, Debug)] +pub struct OpenClipTextTransformer { + embeddings: TextEmbeddings, + encoder: Encoder, + final_layer_norm: LayerNorm, +} + +impl OpenClipTextTransformer { + pub fn new(vs: VarBuilder, c: &Config) -> Result { + let embeddings = TextEmbeddings::new(vs.clone(), c)?; + let final_layer_norm = layer_norm(c.embed_dim, 1e-5, vs.pp("ln_final"))?; + let encoder = Encoder::new(vs.pp("transformer"), c)?; + Ok(OpenClipTextTransformer { + embeddings, + encoder, + final_layer_norm, + }) + } + + pub fn forward(&self, input_ids: &Tensor) -> Result { + let input_ids = self.embeddings.forward(input_ids)?; + let input_ids = self.encoder.forward(&input_ids)?; + self.final_layer_norm.forward(&input_ids) + } +} + +impl Module for OpenClipTextTransformer { + fn forward(&self, input_ids: &Tensor) -> Result { + let output = self.forward(input_ids)?; + let sequence_max_indices = input_ids.argmax(D::Minus1)?.to_dtype(DType::I64)?; + + let mut indices = Vec::new(); + for (batch_idx, &seq_idx) in sequence_max_indices.to_vec1::()?.iter().enumerate() { + let index = output.i((batch_idx, seq_idx as usize))?.unsqueeze(0)?; + indices.push(index); + } + Tensor::cat(&indices, 0) + } +} diff --git a/patches/candle-transformers/src/models/paddleocr_vl/config.rs b/patches/candle-transformers/src/models/paddleocr_vl/config.rs new file mode 100644 index 0000000000..b0b9f75689 --- /dev/null +++ b/patches/candle-transformers/src/models/paddleocr_vl/config.rs @@ -0,0 +1,357 @@ +//! PaddleOCR-VL configuration structures. +//! +//! Defines the configuration for the vision encoder, text decoder, and combined model. + +use candle_nn::Activation; +use serde::Deserialize; + +fn default_vision_hidden_size() -> usize { + 1152 +} + +fn default_vision_intermediate_size() -> usize { + 4304 +} + +fn default_vision_num_hidden_layers() -> usize { + 27 +} + +fn default_vision_num_attention_heads() -> usize { + 16 +} + +fn default_vision_num_channels() -> usize { + 3 +} + +fn default_vision_image_size() -> usize { + 384 +} + +fn default_vision_patch_size() -> usize { + 14 +} + +fn default_vision_hidden_act() -> Activation { + Activation::GeluPytorchTanh +} + +fn default_vision_layer_norm_eps() -> f64 { + 1e-6 +} + +fn default_vision_attention_dropout() -> f64 { + 0.0 +} + +fn default_vision_spatial_merge_size() -> usize { + 2 +} + +/// Vision encoder configuration for PaddleOCR-VL. +/// +/// Uses a NaViT-style dynamic resolution visual encoder with 2D rotary position embeddings. +#[derive(Debug, Clone, Deserialize)] +pub struct VisionConfig { + #[serde(default = "default_vision_hidden_size")] + pub hidden_size: usize, + + #[serde(default = "default_vision_intermediate_size")] + pub intermediate_size: usize, + + #[serde(default = "default_vision_num_hidden_layers")] + pub num_hidden_layers: usize, + + #[serde(default = "default_vision_num_attention_heads")] + pub num_attention_heads: usize, + + #[serde(default = "default_vision_num_channels")] + pub num_channels: usize, + + #[serde(default = "default_vision_image_size")] + pub image_size: usize, + + #[serde(default = "default_vision_patch_size")] + pub patch_size: usize, + + #[serde(default = "default_vision_hidden_act")] + pub hidden_act: Activation, + + #[serde(default = "default_vision_layer_norm_eps")] + pub layer_norm_eps: f64, + + #[serde(default = "default_vision_attention_dropout")] + pub attention_dropout: f64, + + #[serde(default = "default_vision_spatial_merge_size")] + pub spatial_merge_size: usize, +} + +impl Default for VisionConfig { + fn default() -> Self { + Self { + hidden_size: default_vision_hidden_size(), + intermediate_size: default_vision_intermediate_size(), + num_hidden_layers: default_vision_num_hidden_layers(), + num_attention_heads: default_vision_num_attention_heads(), + num_channels: default_vision_num_channels(), + image_size: default_vision_image_size(), + patch_size: default_vision_patch_size(), + hidden_act: default_vision_hidden_act(), + layer_norm_eps: default_vision_layer_norm_eps(), + attention_dropout: default_vision_attention_dropout(), + spatial_merge_size: default_vision_spatial_merge_size(), + } + } +} + +impl VisionConfig { + pub fn head_dim(&self) -> usize { + self.hidden_size / self.num_attention_heads + } +} + +fn default_vocab_size() -> usize { + 103424 +} + +fn default_hidden_size() -> usize { + 1024 +} + +fn default_intermediate_size() -> usize { + 3072 +} + +fn default_num_hidden_layers() -> usize { + 18 +} + +fn default_num_attention_heads() -> usize { + 16 +} + +fn default_num_key_value_heads() -> usize { + 2 +} + +fn default_hidden_act() -> Activation { + Activation::Silu +} + +fn default_max_position_embeddings() -> usize { + 131072 +} + +fn default_rms_norm_eps() -> f64 { + 1e-5 +} + +fn default_rope_theta() -> f64 { + 500000.0 +} + +fn default_head_dim() -> usize { + 128 +} + +fn default_use_bias() -> bool { + false +} + +fn default_tie_word_embeddings() -> bool { + false +} + +fn default_image_token_id() -> u32 { + 100295 +} + +fn default_video_token_id() -> u32 { + 101307 +} + +fn default_vision_start_token_id() -> u32 { + 101305 +} + +fn default_vision_end_token_id() -> u32 { + 101306 +} + +fn default_tokens_per_second() -> usize { + 25 +} + +/// RoPE scaling configuration for multimodal position embeddings. +#[derive(Debug, Clone, Deserialize)] +pub struct RopeScaling { + /// Sections for multimodal RoPE: [temporal, height, width]. + /// Splits head_dim/2 into 3 parts for 3D position encoding. + /// Default: [16, 24, 24] (total = 64 = head_dim/2 for head_dim=128) + #[serde(default = "default_mrope_section")] + pub mrope_section: Vec, + + #[serde(default)] + pub rope_type: Option, +} + +fn default_mrope_section() -> Vec { + vec![16, 24, 24] +} + +impl Default for RopeScaling { + fn default() -> Self { + Self { + mrope_section: default_mrope_section(), + rope_type: Some("default".to_string()), + } + } +} + +/// Combined configuration for PaddleOCR-VL model. +/// +/// The text model parameters are at the top level (not nested in `text_config`), +/// following the HuggingFace format where the main model config contains LLM params directly. +#[derive(Debug, Clone, Deserialize)] +pub struct Config { + // Vision config (nested) + #[serde(default)] + pub vision_config: VisionConfig, + + // Text model parameters (at top level) + #[serde(default = "default_vocab_size")] + pub vocab_size: usize, + + #[serde(default = "default_hidden_size")] + pub hidden_size: usize, + + #[serde(default = "default_intermediate_size")] + pub intermediate_size: usize, + + #[serde(default = "default_num_hidden_layers")] + pub num_hidden_layers: usize, + + #[serde(default = "default_num_attention_heads")] + pub num_attention_heads: usize, + + #[serde(default = "default_num_key_value_heads")] + pub num_key_value_heads: usize, + + #[serde(default = "default_hidden_act")] + pub hidden_act: Activation, + + #[serde(default = "default_max_position_embeddings")] + pub max_position_embeddings: usize, + + #[serde(default = "default_rms_norm_eps", alias = "rms_norm_eps")] + pub layer_norm_eps: f64, + + #[serde(default = "default_rope_theta")] + pub rope_theta: f64, + + #[serde(default = "default_head_dim")] + pub head_dim: usize, + + #[serde(default = "default_use_bias")] + pub use_bias: bool, + + #[serde(default = "default_tie_word_embeddings")] + pub tie_word_embeddings: bool, + + // Special token IDs + #[serde(default = "default_image_token_id")] + pub image_token_id: u32, + + #[serde(default = "default_video_token_id")] + pub video_token_id: u32, + + #[serde(default = "default_vision_start_token_id")] + pub vision_start_token_id: u32, + + #[serde(default = "default_vision_end_token_id")] + pub vision_end_token_id: u32, + + /// RoPE scaling configuration for multimodal position embeddings. + #[serde(default)] + pub rope_scaling: Option, + + /// Tokens per second for video temporal position encoding. + #[serde(default = "default_tokens_per_second")] + pub tokens_per_second: usize, +} + +impl Default for Config { + fn default() -> Self { + Self { + vision_config: VisionConfig::default(), + vocab_size: default_vocab_size(), + hidden_size: default_hidden_size(), + intermediate_size: default_intermediate_size(), + num_hidden_layers: default_num_hidden_layers(), + num_attention_heads: default_num_attention_heads(), + num_key_value_heads: default_num_key_value_heads(), + hidden_act: default_hidden_act(), + max_position_embeddings: default_max_position_embeddings(), + layer_norm_eps: default_rms_norm_eps(), + rope_theta: default_rope_theta(), + head_dim: default_head_dim(), + use_bias: default_use_bias(), + tie_word_embeddings: default_tie_word_embeddings(), + image_token_id: default_image_token_id(), + video_token_id: default_video_token_id(), + vision_start_token_id: default_vision_start_token_id(), + vision_end_token_id: default_vision_end_token_id(), + rope_scaling: Some(RopeScaling::default()), + tokens_per_second: default_tokens_per_second(), + } + } +} + +/// Helper struct for text config (used internally). +/// This provides a view of the text-related config fields. +#[derive(Debug, Clone)] +pub struct TextConfig { + pub vocab_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub hidden_act: Activation, + pub max_position_embeddings: usize, + pub rms_norm_eps: f64, + pub rope_theta: f64, + pub head_dim: usize, + pub use_bias: bool, + pub tie_word_embeddings: bool, + /// Multimodal RoPE sections: [temporal, height, width]. + pub mrope_section: Vec, +} + +impl From<&Config> for TextConfig { + fn from(cfg: &Config) -> Self { + let mrope_section = cfg + .rope_scaling + .as_ref() + .map(|rs| rs.mrope_section.clone()) + .unwrap_or_else(default_mrope_section); + Self { + vocab_size: cfg.vocab_size, + hidden_size: cfg.hidden_size, + intermediate_size: cfg.intermediate_size, + num_hidden_layers: cfg.num_hidden_layers, + num_attention_heads: cfg.num_attention_heads, + num_key_value_heads: cfg.num_key_value_heads, + hidden_act: cfg.hidden_act, + max_position_embeddings: cfg.max_position_embeddings, + rms_norm_eps: cfg.layer_norm_eps, + rope_theta: cfg.rope_theta, + head_dim: cfg.head_dim, + use_bias: cfg.use_bias, + tie_word_embeddings: cfg.tie_word_embeddings, + mrope_section, + } + } +} diff --git a/patches/candle-transformers/src/models/paddleocr_vl/mod.rs b/patches/candle-transformers/src/models/paddleocr_vl/mod.rs new file mode 100644 index 0000000000..88918b81e3 --- /dev/null +++ b/patches/candle-transformers/src/models/paddleocr_vl/mod.rs @@ -0,0 +1,1109 @@ +//! PaddleOCR-VL Vision-Language Model for OCR. +//! +//! PaddleOCR-VL is a state-of-the-art vision-language model for document parsing, +//! combining a NaViT-style visual encoder with the ERNIE-4.5-0.3B language model. +//! +//! Key features: +//! - Dynamic resolution support for variable-sized document images +//! - 2D rotary position embeddings for vision, 1D for text +//! - Grouped Query Attention (GQA) for efficient inference +//! - Supports 109 languages for multilingual OCR +//! - Recognizes text, tables, formulas, and charts +//! +//! Architecture: +//! - Vision Encoder: NaViT-style with 27 layers, 1152 hidden dim, 16 heads +//! - Projector: 2x2 spatial merge + 2-layer MLP (1152*4 → 1024) +//! - Text Decoder: ERNIE-4.5-0.3B with 18 layers, GQA (16 query, 2 KV heads) +//! +//! References: +//! - [Paper](https://arxiv.org/abs/2510.14528) +//! - [HuggingFace Model](https://huggingface.co/PaddlePaddle/PaddleOCR-VL) + +#![allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] + +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::VarBuilder; + +pub mod config; +mod text; +mod vision; + +pub use config::{Config, TextConfig, VisionConfig}; +use text::TextModel; +pub use text::{ + compute_mrope_position_ids, compute_mrope_position_ids_multi, compute_mrope_position_ids_video, + ImageGrid, VideoGrid, +}; +use vision::VisionModel; + +/// Type alias for debug generation output: generated tokens and per-step tensor exports. +pub type GenerateDebugOutput = (Vec, Vec>); + +/// PaddleOCR-VL Model for vision-language OCR tasks. +/// +/// This model combines a NaViT-style vision encoder with an ERNIE-4.5 text decoder +/// for document parsing tasks including OCR, table recognition, formula recognition, +/// and chart recognition. +pub struct PaddleOCRVLModel { + vision: VisionModel, + text: TextModel, + image_token_id: u32, + video_token_id: u32, + dtype: DType, + device: Device, + /// Tracks the M-RoPE position delta for incremental decoding. + /// After prefill with M-RoPE, incremental positions need adjustment. + mrope_position_delta: i64, +} + +impl PaddleOCRVLModel { + /// Create a new PaddleOCR-VL model. + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let text_cfg: TextConfig = cfg.into(); + // Vision model is at "visual.vision_model" + let vision = VisionModel::new( + &cfg.vision_config, + cfg.hidden_size, + vb.pp("visual").pp("vision_model"), + vb.pp("mlp_AR"), // Projector is separate at "mlp_AR" + )?; + // Language model is at "model" (not "language_model.model") + let text = TextModel::new(&text_cfg, vb.clone())?; + + Ok(Self { + vision, + text, + image_token_id: cfg.image_token_id, + video_token_id: cfg.video_token_id, + dtype: vb.dtype(), + device: vb.device().clone(), + mrope_position_delta: 0, + }) + } + + /// Encode image to vision features. + /// + /// # Arguments + /// * `pixel_values` - Image tensor of shape (batch, channels, height, width) + /// * `grid_thw` - Grid dimensions tensor of shape (num_images, 3) with [temporal, height, width] + /// + /// # Returns + /// Vision features projected to text model dimension + pub fn encode_image(&self, pixel_values: &Tensor, grid_thw: &Tensor) -> Result { + self.vision.forward(pixel_values, grid_thw) + } + + /// Encode image with debug output. + pub fn encode_image_debug(&self, pixel_values: &Tensor, grid_thw: &Tensor) -> Result { + self.vision.forward_with_debug(pixel_values, grid_thw, true) + } + + /// Encode image and export intermediate tensors for comparison with PyTorch. + /// + /// Returns vision features and a HashMap of checkpoint tensors. + pub fn encode_image_with_export( + &self, + pixel_values: &Tensor, + grid_thw: &Tensor, + ) -> Result<(Tensor, std::collections::HashMap)> { + self.vision.forward_with_export(pixel_values, grid_thw) + } + + /// Encode multiple images, returning separate embeddings for each. + /// + /// # Arguments + /// * `pixel_values` - Batched image tensor of shape (num_images, channels, height, width) + /// * `grid_thw` - Grid dimensions tensor of shape (num_images, 3) with [temporal, height, width] + /// + /// # Returns + /// Vector of vision feature tensors, one per image + pub fn encode_images_multi( + &self, + pixel_values: &Tensor, + grid_thw: &Tensor, + ) -> Result> { + self.vision.forward_multi(pixel_values, grid_thw) + } + + /// Encode multiple images of different sizes separately. + /// + /// This method handles images with different resolutions by processing + /// each image individually through the vision encoder. + /// + /// # Arguments + /// * `pixel_values_list` - Vector of image tensors, each of shape (1, channels, height, width) + /// * `grid_thw_list` - Vector of grid tensors, each of shape (1, 3) + /// + /// # Returns + /// Vector of vision feature tensors, one per image + pub fn encode_images_separate( + &self, + pixel_values_list: &[Tensor], + grid_thw_list: &[Tensor], + ) -> Result> { + let mut embeddings = Vec::with_capacity(pixel_values_list.len()); + + for (pixel_values, grid_thw) in pixel_values_list.iter().zip(grid_thw_list.iter()) { + let emb = self.vision.forward(pixel_values, grid_thw)?; + embeddings.push(emb); + } + + Ok(embeddings) + } + + /// Forward pass for vision-language generation. + /// + /// # Arguments + /// * `input_ids` - Token IDs of shape (batch, seq_len) + /// * `pixel_values` - Optional image tensor + /// * `grid_thw` - Optional grid dimensions for images + /// * `seqlen_offset` - Sequence length offset for KV cache + /// + /// # Returns + /// Logits for next token prediction + pub fn forward( + &mut self, + input_ids: &Tensor, + pixel_values: Option<&Tensor>, + grid_thw: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (batch_size, seq_len) = input_ids.dims2()?; + + // Get text embeddings + let mut input_embeds = self.text.embed_tokens(input_ids)?; + let hidden_dim = self.text.hidden_size; + + // Track grid dimensions for M-RoPE position computation + let mut merged_grid_h = 0usize; + let mut merged_grid_w = 0usize; + + // If we have images, encode them and inject into embeddings + if let (Some(pixel_values), Some(grid_thw)) = (pixel_values, grid_thw) { + // Encode images + let image_embeds = self.encode_image(pixel_values, grid_thw)?; + let image_embeds = image_embeds.to_dtype(self.dtype)?; + + // Get grid dimensions for M-RoPE (after 2x2 merge) + let grid_thw_vec: Vec = grid_thw.flatten_all()?.to_vec1()?; + if grid_thw_vec.len() >= 3 { + let spatial_merge_size = 2; // 2x2 merge + merged_grid_h = (grid_thw_vec[1] as usize) / spatial_merge_size; + merged_grid_w = (grid_thw_vec[2] as usize) / spatial_merge_size; + } + + // Find image token positions and replace with image embeddings + let input_ids_flat = input_ids.flatten_all()?; + let input_ids_vec = input_ids_flat.to_vec1::()?; + + let mut image_offset = 0usize; + let num_image_tokens = image_embeds.dim(0)?; + + for batch in 0..batch_size { + for pos in 0..seq_len { + let idx = batch * seq_len + pos; + if input_ids_vec[idx] == self.image_token_id && image_offset < num_image_tokens + { + // Replace this token's embedding with image embedding + let img_emb = image_embeds.i(image_offset)?.unsqueeze(0)?.unsqueeze(0)?; + input_embeds = input_embeds.slice_assign( + &[batch..batch + 1, pos..pos + 1, 0..hidden_dim], + &img_emb, + )?; + image_offset += 1; + } + } + } + + // Use M-RoPE with 3D position IDs for prefill with vision tokens + let position_ids = compute_mrope_position_ids( + input_ids, + self.image_token_id, + merged_grid_h, + merged_grid_w, + &self.device, + )?; + + // Compute mrope_position_delta for incremental decoding + // delta = max_position - seq_len + 1, so that position seq_len becomes max_position + 1 + let position_ids_vec: Vec = position_ids.flatten_all()?.to_vec1()?; + let max_pos = position_ids_vec.iter().copied().max().unwrap_or(0); + self.mrope_position_delta = max_pos + 1 - seq_len as i64; + + return self + .text + .forward_embeds_with_mrope(input_embeds, &position_ids); + } + + // Forward through text model with M-RoPE (for incremental decoding) + // + // CRITICAL: We must use M-RoPE during generation, NOT 1D RoPE! + // + // Reason: M-RoPE and 1D RoPE produce DIFFERENT rotations even for the same position + // because M-RoPE splits head_dim by mrope_section [32,48,48] and applies different + // dimension's cos/sin to each section, while 1D RoPE just uses first 64 dims duplicated. + // + // For text tokens, all 3 position dimensions have the same value, but we still need + // to use M-RoPE to maintain consistency with prefill. + // + // Position calculation: seqlen_offset + mrope_position_delta + // This gives the correct sequential position after accounting for the difference + // between sequence index and M-RoPE position caused by 2D vision token positions. + let pos = seqlen_offset as i64 + self.mrope_position_delta; + let (batch_size, seq_len, _) = input_embeds.dims3()?; + + // Create position_ids [3, batch, seq_len] with all dimensions = pos + // For text tokens in generation, all 3 dimensions (temporal, height, width) are identical + let positions: Vec = vec![pos; batch_size * seq_len]; + let pos_tensor = Tensor::from_vec(positions, (batch_size, seq_len), &self.device)?; + let position_ids = Tensor::stack(&[&pos_tensor, &pos_tensor, &pos_tensor], 0)?; + + self.text + .forward_embeds_with_mrope(input_embeds, &position_ids) + } + + /// Forward pass for multi-image vision-language generation. + /// + /// # Arguments + /// * `input_ids` - Token IDs of shape (batch, seq_len) containing multiple image placeholder ranges + /// * `pixel_values` - Batched image tensor of shape (num_images, channels, height, width) + /// * `grid_thw` - Grid dimensions tensor of shape (num_images, 3) with [temporal, height, width] + /// * `seqlen_offset` - Sequence length offset for KV cache (0 for prefill) + /// + /// # Returns + /// Logits for next token prediction + pub fn forward_multi_image( + &mut self, + input_ids: &Tensor, + pixel_values: &Tensor, + grid_thw: &Tensor, + _seqlen_offset: usize, + ) -> Result { + let (batch_size, seq_len) = input_ids.dims2()?; + + // Get text embeddings + let mut input_embeds = self.text.embed_tokens(input_ids)?; + let hidden_dim = self.text.hidden_size; + + // Encode all images, getting separate embeddings for each + let image_embeds_list = self.encode_images_multi(pixel_values, grid_thw)?; + let image_embeds_list: Vec = image_embeds_list + .into_iter() + .map(|t| t.to_dtype(self.dtype)) + .collect::>>()?; + + // Build image grids for M-RoPE position computation + let grid_thw_vec: Vec> = grid_thw.to_vec2()?; + let spatial_merge_size = 2; // 2x2 merge + let image_grids: Vec = grid_thw_vec + .iter() + .map(|g| ImageGrid { + grid_h: (g[1] as usize) / spatial_merge_size, + grid_w: (g[2] as usize) / spatial_merge_size, + }) + .collect(); + + // Find image token ranges and inject embeddings + let input_ids_flat = input_ids.flatten_all()?; + let input_ids_vec = input_ids_flat.to_vec1::()?; + + // Find all image token ranges + let mut image_ranges: Vec<(usize, usize)> = Vec::new(); + let mut in_image = false; + let mut image_start = 0usize; + + for (pos, &token_id) in input_ids_vec.iter().enumerate() { + if token_id == self.image_token_id { + if !in_image { + in_image = true; + image_start = pos; + } + } else if in_image { + image_ranges.push((image_start, pos)); + in_image = false; + } + } + if in_image { + image_ranges.push((image_start, input_ids_vec.len())); + } + + // Verify we have the right number of image ranges + if image_ranges.len() != image_embeds_list.len() { + return Err(candle::Error::Msg(format!( + "Found {} image ranges but have {} encoded images", + image_ranges.len(), + image_embeds_list.len() + ))); + } + + // Inject each image's embeddings at the correct positions + for batch in 0..batch_size { + for (img_idx, ((start, end), embeddings)) in image_ranges + .iter() + .zip(image_embeds_list.iter()) + .enumerate() + { + let num_tokens = end - start; + let num_embeddings = embeddings.dim(0)?; + + if num_tokens != num_embeddings { + return Err(candle::Error::Msg(format!( + "Image {} has {} placeholder tokens but {} embeddings", + img_idx, num_tokens, num_embeddings + ))); + } + + // Replace each placeholder token with the corresponding embedding + for (offset, pos) in (*start..*end).enumerate() { + let img_emb = embeddings.i(offset)?.unsqueeze(0)?.unsqueeze(0)?; + input_embeds = input_embeds + .slice_assign(&[batch..batch + 1, pos..pos + 1, 0..hidden_dim], &img_emb)?; + } + } + } + + // Compute M-RoPE position IDs for multi-image input + let position_ids = compute_mrope_position_ids_multi( + input_ids, + self.image_token_id, + &image_grids, + &self.device, + )?; + + // Compute mrope_position_delta for incremental decoding + let position_ids_vec: Vec = position_ids.flatten_all()?.to_vec1()?; + let max_pos = position_ids_vec.iter().copied().max().unwrap_or(0); + self.mrope_position_delta = max_pos + 1 - seq_len as i64; + + self.text + .forward_embeds_with_mrope(input_embeds, &position_ids) + } + + /// Forward pass for multi-image with variable resolutions. + /// + /// This method handles images of different sizes by processing each + /// image separately through the vision encoder. + /// + /// # Arguments + /// * `input_ids` - Token IDs containing multiple image placeholder ranges + /// * `pixel_values_list` - Vector of image tensors, each (1, C, H, W) + /// * `grid_thw_list` - Vector of grid tensors, each (1, 3) + /// * `_seqlen_offset` - Unused, kept for API consistency + pub fn forward_multi_image_separate( + &mut self, + input_ids: &Tensor, + pixel_values_list: &[Tensor], + grid_thw_list: &[Tensor], + _seqlen_offset: usize, + ) -> Result { + let (batch_size, seq_len) = input_ids.dims2()?; + + // Get text embeddings + let mut input_embeds = self.text.embed_tokens(input_ids)?; + let hidden_dim = self.text.hidden_size; + + // Encode each image separately + let image_embeds_list = self.encode_images_separate(pixel_values_list, grid_thw_list)?; + let image_embeds_list: Vec = image_embeds_list + .into_iter() + .map(|t| t.to_dtype(self.dtype)) + .collect::>>()?; + + // Build image grids for M-RoPE position computation + let spatial_merge_size = 2; // 2x2 merge + let mut image_grids: Vec = Vec::with_capacity(grid_thw_list.len()); + for grid_thw in grid_thw_list { + let grid_vec: Vec> = grid_thw.to_vec2()?; + let g = &grid_vec[0]; + image_grids.push(ImageGrid { + grid_h: (g[1] as usize) / spatial_merge_size, + grid_w: (g[2] as usize) / spatial_merge_size, + }); + } + + // Find image token ranges and inject embeddings + let input_ids_flat = input_ids.flatten_all()?; + let input_ids_vec = input_ids_flat.to_vec1::()?; + + // Find all image token ranges + let mut image_ranges: Vec<(usize, usize)> = Vec::new(); + let mut in_image = false; + let mut image_start = 0usize; + + for (pos, &token_id) in input_ids_vec.iter().enumerate() { + if token_id == self.image_token_id { + if !in_image { + in_image = true; + image_start = pos; + } + } else if in_image { + image_ranges.push((image_start, pos)); + in_image = false; + } + } + if in_image { + image_ranges.push((image_start, input_ids_vec.len())); + } + + // Verify we have the right number of image ranges + if image_ranges.len() != image_embeds_list.len() { + return Err(candle::Error::Msg(format!( + "Found {} image ranges but have {} encoded images", + image_ranges.len(), + image_embeds_list.len() + ))); + } + + // Inject each image's embeddings at the correct positions + for batch in 0..batch_size { + for (img_idx, ((start, end), embeddings)) in image_ranges + .iter() + .zip(image_embeds_list.iter()) + .enumerate() + { + let num_tokens = end - start; + let num_embeddings = embeddings.dim(0)?; + + if num_tokens != num_embeddings { + return Err(candle::Error::Msg(format!( + "Image {} has {} placeholder tokens but {} embeddings", + img_idx, num_tokens, num_embeddings + ))); + } + + // Replace each placeholder token with the corresponding embedding + for (offset, pos) in (*start..*end).enumerate() { + let img_emb = embeddings.i(offset)?.unsqueeze(0)?.unsqueeze(0)?; + input_embeds = input_embeds + .slice_assign(&[batch..batch + 1, pos..pos + 1, 0..hidden_dim], &img_emb)?; + } + } + } + + // Compute M-RoPE position IDs for multi-image input + let position_ids = compute_mrope_position_ids_multi( + input_ids, + self.image_token_id, + &image_grids, + &self.device, + )?; + + // Compute mrope_position_delta for incremental decoding + let position_ids_vec: Vec = position_ids.flatten_all()?.to_vec1()?; + let max_pos = position_ids_vec.iter().copied().max().unwrap_or(0); + self.mrope_position_delta = max_pos + 1 - seq_len as i64; + + self.text + .forward_embeds_with_mrope(input_embeds, &position_ids) + } + + /// Generate text from image using greedy decoding. + /// + /// # Arguments + /// * `input_ids` - Initial token IDs (including image placeholders) + /// * `pixel_values` - Image tensor + /// * `grid_thw` - Grid dimensions for images + /// * `max_new_tokens` - Maximum number of tokens to generate + /// * `eos_token_id` - End of sequence token ID + /// + /// # Returns + /// Generated token IDs + pub fn generate( + &mut self, + input_ids: &Tensor, + pixel_values: &Tensor, + grid_thw: &Tensor, + max_new_tokens: usize, + eos_token_id: u32, + ) -> Result> { + self.clear_kv_cache(); + + let mut generated_tokens = Vec::new(); + let mut current_ids = input_ids.clone(); + + // First forward pass with image + let logits = self.forward(¤t_ids, Some(pixel_values), Some(grid_thw), 0)?; + let next_token = logits + .argmax(D::Minus1)? + .to_dtype(DType::U32)? + .to_vec1::()?[0]; + + generated_tokens.push(next_token); + + if next_token == eos_token_id { + return Ok(generated_tokens); + } + + let mut seqlen_offset = current_ids.dim(1)?; + current_ids = Tensor::new(&[next_token], &self.device)?.unsqueeze(0)?; + + // Subsequent forward passes (text only, using KV cache) + for _ in 1..max_new_tokens { + let logits = self.forward(¤t_ids, None, None, seqlen_offset)?; + let next_token = logits + .argmax(D::Minus1)? + .to_dtype(DType::U32)? + .to_vec1::()?[0]; + + generated_tokens.push(next_token); + + if next_token == eos_token_id { + break; + } + + seqlen_offset += 1; + current_ids = Tensor::new(&[next_token], &self.device)?.unsqueeze(0)?; + } + + Ok(generated_tokens) + } + + /// Generate text from multiple images using greedy decoding. + /// + /// # Arguments + /// * `input_ids` - Initial token IDs (including multiple image placeholder ranges) + /// * `pixel_values` - Batched image tensor of shape (num_images, channels, height, width) + /// * `grid_thw` - Grid dimensions tensor of shape (num_images, 3) + /// * `max_new_tokens` - Maximum number of tokens to generate + /// * `eos_token_id` - End of sequence token ID + /// + /// # Returns + /// Generated token IDs + pub fn generate_multi_image( + &mut self, + input_ids: &Tensor, + pixel_values: &Tensor, + grid_thw: &Tensor, + max_new_tokens: usize, + eos_token_id: u32, + ) -> Result> { + self.clear_kv_cache(); + + let mut generated_tokens = Vec::new(); + let mut current_ids = input_ids.clone(); + + // First forward pass with all images + let logits = self.forward_multi_image(¤t_ids, pixel_values, grid_thw, 0)?; + let next_token = logits + .argmax(D::Minus1)? + .to_dtype(DType::U32)? + .to_vec1::()?[0]; + + generated_tokens.push(next_token); + + if next_token == eos_token_id { + return Ok(generated_tokens); + } + + let mut seqlen_offset = current_ids.dim(1)?; + current_ids = Tensor::new(&[next_token], &self.device)?.unsqueeze(0)?; + + // Subsequent forward passes (text only, using KV cache) + // Uses same incremental decoding as single-image generation + for _ in 1..max_new_tokens { + let logits = self.forward(¤t_ids, None, None, seqlen_offset)?; + let next_token = logits + .argmax(D::Minus1)? + .to_dtype(DType::U32)? + .to_vec1::()?[0]; + + generated_tokens.push(next_token); + + if next_token == eos_token_id { + break; + } + + seqlen_offset += 1; + current_ids = Tensor::new(&[next_token], &self.device)?.unsqueeze(0)?; + } + + Ok(generated_tokens) + } + + /// Generate text from multiple images of different sizes using greedy decoding. + /// + /// This method handles images with different resolutions by processing + /// each image separately through the vision encoder. + /// + /// # Arguments + /// * `input_ids` - Initial token IDs (including multiple image placeholder ranges) + /// * `pixel_values_list` - Vector of image tensors, each (1, C, H, W) + /// * `grid_thw_list` - Vector of grid tensors, each (1, 3) + /// * `max_new_tokens` - Maximum number of tokens to generate + /// * `eos_token_id` - End of sequence token ID + /// + /// # Returns + /// Generated token IDs + pub fn generate_multi_image_separate( + &mut self, + input_ids: &Tensor, + pixel_values_list: &[Tensor], + grid_thw_list: &[Tensor], + max_new_tokens: usize, + eos_token_id: u32, + ) -> Result> { + self.clear_kv_cache(); + + let mut generated_tokens = Vec::new(); + let mut current_ids = input_ids.clone(); + + // First forward pass with all images (processed separately) + let logits = + self.forward_multi_image_separate(¤t_ids, pixel_values_list, grid_thw_list, 0)?; + let next_token = logits + .argmax(D::Minus1)? + .to_dtype(DType::U32)? + .to_vec1::()?[0]; + + generated_tokens.push(next_token); + + if next_token == eos_token_id { + return Ok(generated_tokens); + } + + let mut seqlen_offset = current_ids.dim(1)?; + current_ids = Tensor::new(&[next_token], &self.device)?.unsqueeze(0)?; + + // Subsequent forward passes (text only, using KV cache) + for _ in 1..max_new_tokens { + let logits = self.forward(¤t_ids, None, None, seqlen_offset)?; + let next_token = logits + .argmax(D::Minus1)? + .to_dtype(DType::U32)? + .to_vec1::()?[0]; + + generated_tokens.push(next_token); + + if next_token == eos_token_id { + break; + } + + seqlen_offset += 1; + current_ids = Tensor::new(&[next_token], &self.device)?.unsqueeze(0)?; + } + + Ok(generated_tokens) + } + + /// Forward pass for video input. + /// + /// This method processes video frames with temporal position encoding, + /// where each frame gets sequential temporal positions (t=0, 1, 2, ...) + /// unlike images which all use t=0. + /// + /// # Arguments + /// * `input_ids` - Token IDs containing video placeholder tokens + /// * `pixel_values_video` - Stacked video frames (num_frames * C * H * W flattened) + /// * `video_grid_thw` - Grid dimensions (1, 3) = [temporal, height, width] + /// * `fps` - Frames per second used to extract video frames + /// * `seqlen_offset` - Sequence length offset for KV cache + /// + /// # Returns + /// Logits for next token prediction + pub fn forward_video( + &mut self, + input_ids: &Tensor, + pixel_values_video: &Tensor, + video_grid_thw: &Tensor, + fps: f32, + _seqlen_offset: usize, + ) -> Result { + let (batch_size, seq_len) = input_ids.dims2()?; + + // Get text embeddings + let mut input_embeds = self.text.embed_tokens(input_ids)?; + let hidden_dim = self.text.hidden_size; + + // Encode video frames through vision encoder + // The vision encoder treats video frames similarly to batched images + let video_embeds = self.vision.forward(pixel_values_video, video_grid_thw)?; + let video_embeds = video_embeds.to_dtype(self.dtype)?; + + // Build video grid for M-RoPE position computation + let grid_thw_vec: Vec> = video_grid_thw.to_vec2()?; + let g = &grid_thw_vec[0]; + let spatial_merge_size = 2; // 2x2 merge + let video_grid = VideoGrid { + grid_t: g[0] as usize, + grid_h: (g[1] as usize) / spatial_merge_size, + grid_w: (g[2] as usize) / spatial_merge_size, + }; + // Find video token range and inject embeddings + let input_ids_flat = input_ids.flatten_all()?; + let input_ids_vec = input_ids_flat.to_vec1::()?; + + let mut video_start = None; + let mut video_end = None; + let mut in_video = false; + + for (pos, &token_id) in input_ids_vec.iter().enumerate() { + if token_id == self.video_token_id { + if !in_video { + in_video = true; + video_start = Some(pos); + } + } else if in_video { + video_end = Some(pos); + break; + } + } + if in_video && video_end.is_none() { + video_end = Some(input_ids_vec.len()); + } + + // Inject video embeddings + if let (Some(start), Some(end)) = (video_start, video_end) { + let num_tokens = end - start; + let num_embeddings = video_embeds.dim(0)?; + + if num_tokens != num_embeddings { + return Err(candle::Error::Msg(format!( + "Video has {} placeholder tokens but {} embeddings", + num_tokens, num_embeddings + ))); + } + + for batch in 0..batch_size { + for (offset, pos) in (start..end).enumerate() { + let emb = video_embeds.i(offset)?.unsqueeze(0)?.unsqueeze(0)?; + input_embeds = input_embeds + .slice_assign(&[batch..batch + 1, pos..pos + 1, 0..hidden_dim], &emb)?; + } + } + } + + // Compute temporal scaling parameters for M-RoPE + // HuggingFace Qwen2-VL uses simple sequential temporal indices (0, 1, 2, ...) + // second_per_grid_t * tokens_per_second = 1.0 gives sequential frame indices + // Python shows second_per_grid_ts = 0.5 with tokens_per_second = 2 -> 0.5 * 2 = 1.0 + let second_per_grid_t = 0.5f32; // Match Python processor output + let tokens_per_second = 2usize; + let _ = fps; // fps is used to determine how frames are sampled, not for position encoding + + // Compute M-RoPE position IDs with temporal encoding + let position_ids = compute_mrope_position_ids_video( + input_ids, + self.video_token_id, + &video_grid, + second_per_grid_t, + tokens_per_second, + &self.device, + )?; + + // Compute mrope_position_delta for incremental decoding + let position_ids_vec: Vec = position_ids.flatten_all()?.to_vec1()?; + let max_pos = position_ids_vec.iter().copied().max().unwrap_or(0); + self.mrope_position_delta = max_pos + 1 - seq_len as i64; + + self.text + .forward_embeds_with_mrope(input_embeds, &position_ids) + } + + /// Generate text from video using greedy decoding. + /// + /// # Arguments + /// * `input_ids` - Initial token IDs (including video placeholder tokens) + /// * `pixel_values_video` - Stacked video frames + /// * `video_grid_thw` - Grid dimensions (1, 3) = [temporal, height, width] + /// * `fps` - Frames per second used to extract video frames + /// * `max_new_tokens` - Maximum number of tokens to generate + /// * `eos_token_id` - End of sequence token ID + /// + /// # Returns + /// Generated token IDs + pub fn generate_video( + &mut self, + input_ids: &Tensor, + pixel_values_video: &Tensor, + video_grid_thw: &Tensor, + fps: f32, + max_new_tokens: usize, + eos_token_id: u32, + ) -> Result> { + self.clear_kv_cache(); + + let repetition_penalty = 1.1f32; + let mut generated_tokens = Vec::new(); + let mut current_ids = input_ids.clone(); + + // Helper function to apply repetition penalty + fn apply_repetition_penalty( + logits: &Tensor, + generated: &[u32], + penalty: f32, + ) -> Result { + if generated.is_empty() || penalty == 1.0 { + return Ok(logits.clone()); + } + let device = logits.device(); + let original_shape = logits.dims().to_vec(); + let logits_flat = logits.flatten_all()?; + let mut logits_vec: Vec = logits_flat.to_vec1()?; + for &token in generated { + let idx = token as usize; + if idx < logits_vec.len() { + if logits_vec[idx] > 0.0 { + logits_vec[idx] /= penalty; + } else { + logits_vec[idx] *= penalty; + } + } + } + Tensor::from_vec(logits_vec, original_shape, device) + } + + // First forward pass with video + let logits = + self.forward_video(¤t_ids, pixel_values_video, video_grid_thw, fps, 0)?; + let logits = apply_repetition_penalty(&logits, &generated_tokens, repetition_penalty)?; + let next_token = logits + .argmax(D::Minus1)? + .to_dtype(DType::U32)? + .to_vec1::()?[0]; + + generated_tokens.push(next_token); + + if next_token == eos_token_id { + return Ok(generated_tokens); + } + + let mut seqlen_offset = current_ids.dim(1)?; + current_ids = Tensor::new(&[next_token], &self.device)?.unsqueeze(0)?; + + // Subsequent forward passes (text only, using KV cache) + for _ in 1..max_new_tokens { + let logits = self.forward(¤t_ids, None, None, seqlen_offset)?; + let logits = apply_repetition_penalty(&logits, &generated_tokens, repetition_penalty)?; + let next_token = logits + .argmax(D::Minus1)? + .to_dtype(DType::U32)? + .to_vec1::()?[0]; + + generated_tokens.push(next_token); + + if next_token == eos_token_id { + break; + } + + seqlen_offset += 1; + current_ids = Tensor::new(&[next_token], &self.device)?.unsqueeze(0)?; + } + + Ok(generated_tokens) + } + + /// Clear all KV caches and reset M-RoPE position delta. + pub fn clear_kv_cache(&mut self) { + self.text.clear_kv_cache(); + self.mrope_position_delta = 0; + } + + /// Forward pass with tensor export for decoder comparison. + /// + /// This method captures intermediate tensors at key checkpoints for + /// comparison with PyTorch implementation. + /// + /// # Returns + /// Tuple of (logits, HashMap of checkpoint tensors) + pub fn forward_with_decoder_export( + &mut self, + input_ids: &Tensor, + pixel_values: &Tensor, + grid_thw: &Tensor, + ) -> Result<(Tensor, std::collections::HashMap)> { + use std::collections::HashMap; + + let mut tensors: HashMap = HashMap::new(); + let (batch_size, seq_len) = input_ids.dims2()?; + + // Step 1: Get text embeddings + let mut input_embeds = self.text.embed_tokens(input_ids)?; + tensors.insert( + "input_embeds_before_merge".to_string(), + input_embeds.clone(), + ); + let hidden_dim = self.text.hidden_size; + + // Step 2: Encode images + let image_embeds = self.encode_image(pixel_values, grid_thw)?; + let image_embeds = image_embeds.to_dtype(self.dtype)?; + tensors.insert("vision_embeds".to_string(), image_embeds.clone()); + + // Get grid dimensions for M-RoPE + let grid_thw_vec: Vec = grid_thw.flatten_all()?.to_vec1()?; + let spatial_merge_size = 2; + let merged_grid_h = (grid_thw_vec[1] as usize) / spatial_merge_size; + let merged_grid_w = (grid_thw_vec[2] as usize) / spatial_merge_size; + + // Step 3: Merge vision embeddings into text embeddings + let input_ids_flat = input_ids.flatten_all()?; + let input_ids_vec = input_ids_flat.to_vec1::()?; + let mut image_offset = 0usize; + let num_image_tokens = image_embeds.dim(0)?; + + for batch in 0..batch_size { + for pos in 0..seq_len { + let idx = batch * seq_len + pos; + if input_ids_vec[idx] == self.image_token_id && image_offset < num_image_tokens { + let img_emb = image_embeds.i(image_offset)?.unsqueeze(0)?.unsqueeze(0)?; + input_embeds = input_embeds + .slice_assign(&[batch..batch + 1, pos..pos + 1, 0..hidden_dim], &img_emb)?; + image_offset += 1; + } + } + } + tensors.insert( + "inputs_embeds_after_merge".to_string(), + input_embeds.clone(), + ); + + // Step 4: Compute M-RoPE position IDs + let position_ids = compute_mrope_position_ids( + input_ids, + self.image_token_id, + merged_grid_h, + merged_grid_w, + &self.device, + )?; + tensors.insert("position_ids".to_string(), position_ids.clone()); + + // Compute rope_deltas (max_pos - seq_len + 1) + let position_ids_vec: Vec = position_ids.flatten_all()?.to_vec1()?; + let max_pos = position_ids_vec.iter().copied().max().unwrap_or(0); + let rope_delta = max_pos + 1 - seq_len as i64; + + // CRITICAL: Set mrope_position_delta for incremental decoding + self.mrope_position_delta = rope_delta; + + tensors.insert( + "rope_deltas".to_string(), + Tensor::new(&[rope_delta], &self.device)?, + ); + + // Step 5: Forward through text model with export + let (logits, decoder_tensors) = self + .text + .forward_embeds_with_mrope_export(input_embeds, &position_ids)?; + + // Merge decoder tensors + for (k, v) in decoder_tensors { + tensors.insert(k, v); + } + + // Store last token logits + let last_token_logits = logits.i((.., seq_len - 1, ..))?; + tensors.insert("last_token_logits".to_string(), last_token_logits); + + Ok((logits, tensors)) + } + + /// Generate with debug tensor export at each step. + /// + /// Returns generated tokens and a vector of tensor maps for each step. + pub fn generate_debug( + &mut self, + input_ids: &Tensor, + pixel_values: &Tensor, + grid_thw: &Tensor, + max_steps: usize, + eos_token_id: u32, + ) -> Result { + use std::collections::HashMap; + + self.clear_kv_cache(); + + let mut generated_tokens = Vec::new(); + let mut all_tensors: Vec> = Vec::new(); + + // Step 0: Prefill with image + let (logits, prefill_tensors) = + self.forward_with_decoder_export(input_ids, pixel_values, grid_thw)?; + + let next_token = logits + .i((.., logits.dim(1)? - 1, ..))? + .argmax(D::Minus1)? + .to_dtype(DType::U32)? + .to_vec1::()?[0]; + + let mut step_tensors = prefill_tensors; + step_tensors.insert("step".to_string(), Tensor::new(&[0i64], &self.device)?); + step_tensors.insert( + "predicted_token".to_string(), + Tensor::new(&[next_token as i64], &self.device)?, + ); + step_tensors.insert( + "mrope_delta".to_string(), + Tensor::new(&[self.mrope_position_delta], &self.device)?, + ); + all_tensors.push(step_tensors); + + generated_tokens.push(next_token); + + if next_token == eos_token_id || max_steps <= 1 { + return Ok((generated_tokens, all_tensors)); + } + + // Generation steps + let mut seqlen_offset = input_ids.dim(1)?; + let mut current_ids = Tensor::new(&[next_token], &self.device)?.unsqueeze(0)?; + + for step in 1..max_steps { + // Compute position for M-RoPE + let pos = seqlen_offset as i64 + self.mrope_position_delta; + let (batch_size, seq_len, _) = { + let embeds = self.text.embed_tokens(¤t_ids)?; + embeds.dims3()? + }; + + // Create position_ids [3, batch, seq_len] + let positions: Vec = vec![pos; batch_size * seq_len]; + let pos_tensor = Tensor::from_vec(positions, (batch_size, seq_len), &self.device)?; + let position_ids = Tensor::stack(&[&pos_tensor, &pos_tensor, &pos_tensor], 0)?; + + // Get embeddings + let input_embeds = self.text.embed_tokens(¤t_ids)?; + + // Forward with export + let (logits, decoder_tensors) = self + .text + .forward_embeds_with_mrope_export(input_embeds, &position_ids)?; + + let next_token = logits + .i((.., logits.dim(1)? - 1, ..))? + .argmax(D::Minus1)? + .to_dtype(DType::U32)? + .to_vec1::()?[0]; + + let mut step_tensors: HashMap = decoder_tensors; + step_tensors.insert( + "step".to_string(), + Tensor::new(&[step as i64], &self.device)?, + ); + step_tensors.insert( + "seqlen_offset".to_string(), + Tensor::new(&[seqlen_offset as i64], &self.device)?, + ); + step_tensors.insert( + "mrope_position".to_string(), + Tensor::new(&[pos], &self.device)?, + ); + step_tensors.insert("position_ids".to_string(), position_ids); + step_tensors.insert( + "predicted_token".to_string(), + Tensor::new(&[next_token as i64], &self.device)?, + ); + all_tensors.push(step_tensors); + + generated_tokens.push(next_token); + + if next_token == eos_token_id { + break; + } + + seqlen_offset += 1; + current_ids = Tensor::new(&[next_token], &self.device)?.unsqueeze(0)?; + } + + Ok((generated_tokens, all_tensors)) + } +} diff --git a/patches/candle-transformers/src/models/paddleocr_vl/text.rs b/patches/candle-transformers/src/models/paddleocr_vl/text.rs new file mode 100644 index 0000000000..a1102b0901 --- /dev/null +++ b/patches/candle-transformers/src/models/paddleocr_vl/text.rs @@ -0,0 +1,1260 @@ +//! PaddleOCR-VL Text Model. +//! +//! ERNIE-4.5-0.3B based decoder with RMSNorm, GQA, and M-RoPE (Multimodal RoPE). +//! +//! M-RoPE uses 3D position IDs (temporal, height, width) for vision tokens, +//! allowing the model to encode spatial structure of images. + +use std::sync::Arc; + +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::{embedding, linear_b, rms_norm, Embedding, Linear, Module, RmsNorm, VarBuilder}; + +use super::config::TextConfig; + +/// Multimodal Rotary Position Embedding (M-RoPE). +/// +/// Unlike standard 1D RoPE, M-RoPE supports 3D position IDs for vision tokens: +/// - Temporal position (for video frames, always 0 for images) +/// - Height position (row in the image grid) +/// - Width position (column in the image grid) +/// +/// Text tokens use the same position for all 3 dimensions (equivalent to 1D RoPE). +#[derive(Debug, Clone)] +pub struct RotaryEmbedding { + /// Precomputed cos values for all positions: [max_seq_len, head_dim/2] + cos: Tensor, + /// Precomputed sin values for all positions: [max_seq_len, head_dim/2] + sin: Tensor, + /// M-RoPE section sizes: [temporal, height, width] + mrope_section: Vec, + head_dim: usize, +} + +impl RotaryEmbedding { + pub fn new(cfg: &TextConfig, device: &Device, dtype: DType) -> Result { + let dim = cfg.head_dim; + let max_seq_len = cfg.max_position_embeddings; + + // Compute inverse frequencies + let inv_freq: Vec = (0..dim) + .step_by(2) + .map(|i| 1f32 / (cfg.rope_theta as f32).powf(i as f32 / dim as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), device)?; + + // Compute cos/sin for all positions + let t = Tensor::arange(0u32, max_seq_len as u32, device)? + .to_dtype(DType::F32)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + let sin = freqs.sin()?.to_dtype(dtype)?; + let cos = freqs.cos()?.to_dtype(dtype)?; + + Ok(Self { + cos, + sin, + mrope_section: cfg.mrope_section.clone(), + head_dim: dim, + }) + } + + /// Apply Multimodal RoPE with 3D position IDs. + /// + /// This follows the PyTorch implementation where: + /// 1. Compute cos/sin for each of the 3 position dimensions (temporal, height, width) + /// 2. Split the head_dim into sections based on mrope_section + /// 3. Use temporal positions for first section, height for second, width for third + /// + /// # Arguments + /// * `q` - Query tensor [batch, heads, seq_len, head_dim] + /// * `k` - Key tensor [batch, kv_heads, seq_len, head_dim] + /// * `position_ids` - 3D position IDs [3, batch, seq_len] where dim 0 is [temporal, height, width] + pub fn apply_multimodal_rotary_emb( + &self, + q: &Tensor, + k: &Tensor, + position_ids: &Tensor, + ) -> Result<(Tensor, Tensor)> { + // position_ids: [3, batch, seq_len] + let (three, _batch, _seq_len) = position_ids.dims3()?; + assert_eq!(three, 3, "position_ids must have 3 dimensions"); + + // Compute cos/sin for each position dimension + // Each returns [batch, seq_len, head_dim] with cos/sin of (inv_freq * position) + let (cos_3d, sin_3d) = self.compute_3d_rope_embeddings(position_ids)?; + // cos_3d/sin_3d: [3, batch, seq_len, head_dim] + + // Apply mrope_section to select appropriate bands from each dimension + // mrope_section = [16, 24, 24] splits head_dim=128 into [16, 24, 24, 64] chunks + // where 64 is the remainder. Chunk i uses dimension i % 3. + let (cos, sin) = self.apply_mrope_sections(&cos_3d, &sin_3d)?; + // cos/sin: [batch, seq_len, head_dim] + + // Reshape for broadcasting: [batch, 1, seq_len, head_dim] + let cos = cos.unsqueeze(1)?; + let sin = sin.unsqueeze(1)?; + + // Apply RoPE to q and k + let q_embed = self.apply_rope_to_tensor(q, &cos, &sin)?; + let k_embed = self.apply_rope_to_tensor(k, &cos, &sin)?; + + Ok((q_embed, k_embed)) + } + + /// Compute cos/sin embeddings for 3D position IDs. + /// position_ids: [3, batch, seq_len] + /// Returns: (cos, sin) each with shape [3, batch, seq_len, head_dim] + fn compute_3d_rope_embeddings(&self, position_ids: &Tensor) -> Result<(Tensor, Tensor)> { + let (three, batch, seq_len) = position_ids.dims3()?; + let half_dim = self.head_dim / 2; + + // For each of the 3 dimensions, gather cos/sin based on positions + let mut cos_parts = Vec::new(); + let mut sin_parts = Vec::new(); + + for dim_idx in 0..three { + let pos = position_ids.i(dim_idx)?; // [batch, seq_len] + let pos_flat = pos.flatten_all()?; // [batch * seq_len] + + // Gather from precomputed cos/sin + let cos_gathered = self.cos.index_select(&pos_flat, 0)?; // [batch*seq_len, half_dim] + let sin_gathered = self.sin.index_select(&pos_flat, 0)?; + + // Reshape to [batch, seq_len, half_dim] + let cos_dim = cos_gathered.reshape((batch, seq_len, half_dim))?; + let sin_dim = sin_gathered.reshape((batch, seq_len, half_dim))?; + + // Duplicate to full head_dim: [batch, seq_len, head_dim] + let cos_full = Tensor::cat(&[&cos_dim, &cos_dim], D::Minus1)?; + let sin_full = Tensor::cat(&[&sin_dim, &sin_dim], D::Minus1)?; + + cos_parts.push(cos_full); + sin_parts.push(sin_full); + } + + // Stack to [3, batch, seq_len, head_dim] + let cos_3d = Tensor::stack(&cos_parts, 0)?; + let sin_3d = Tensor::stack(&sin_parts, 0)?; + + Ok((cos_3d, sin_3d)) + } + + /// Apply mrope_section to select bands from each dimension. + /// + /// PyTorch behavior: `cos.split(mrope_section * 2, dim=-1)` where `* 2` is **list repetition**! + /// In Python: `[16, 24, 24] * 2 = [16, 24, 24, 16, 24, 24]` (6 chunks totaling 128) + /// + /// Then `[m[i % 3] for i, m in enumerate(splits)]` selects from the 3D position embeddings: + /// - chunk 0 (dims 0-15): from temporal (i=0, i%3=0) + /// - chunk 1 (dims 16-39): from height (i=1, i%3=1) + /// - chunk 2 (dims 40-63): from width (i=2, i%3=2) + /// - chunk 3 (dims 64-79): from temporal (i=3, i%3=0) + /// - chunk 4 (dims 80-103): from height (i=4, i%3=1) + /// - chunk 5 (dims 104-127): from width (i=5, i%3=2) + /// + /// Final layout: [T:16, H:24, W:24, T:16, H:24, W:24] + fn apply_mrope_sections(&self, cos_3d: &Tensor, sin_3d: &Tensor) -> Result<(Tensor, Tensor)> { + // cos_3d/sin_3d: [3, batch, seq_len, head_dim] + // mrope_section = [16, 24, 24] + // + // In Python: mrope_section * 2 = [16, 24, 24, 16, 24, 24] (list repetition!) + // This creates 6 splits, cycling through temporal/height/width twice + let mut sections_repeated: Vec = Vec::new(); + sections_repeated.extend_from_slice(&self.mrope_section); + sections_repeated.extend_from_slice(&self.mrope_section); + // sections_repeated = [16, 24, 24, 16, 24, 24] + + // Split the head_dim and take from appropriate dimension (i % 3) + let mut cos_parts = Vec::new(); + let mut sin_parts = Vec::new(); + let mut offset = 0; + + for (i, &sec_size) in sections_repeated.iter().enumerate() { + let dim_idx = i % 3; // Cycles: temporal(0), height(1), width(2), temporal(0), ... + // Take slice from dimension dim_idx at the current offset + let cos_slice = cos_3d.i(dim_idx)?.narrow(D::Minus1, offset, sec_size)?; + let sin_slice = sin_3d.i(dim_idx)?.narrow(D::Minus1, offset, sec_size)?; + cos_parts.push(cos_slice); + sin_parts.push(sin_slice); + offset += sec_size; + } + + // Concatenate along head_dim: [batch, seq_len, head_dim] + let cos = Tensor::cat(&cos_parts, D::Minus1)?; + let sin = Tensor::cat(&sin_parts, D::Minus1)?; + + Ok((cos, sin)) + } + + /// Apply rotary embedding to a tensor. + /// x: [batch, heads, seq_len, head_dim] + /// cos/sin: [batch, 1, seq_len, head_dim] + fn apply_rope_to_tensor(&self, x: &Tensor, cos: &Tensor, sin: &Tensor) -> Result { + let x = x.contiguous()?; + + // rotate_half: split x into two halves and rotate + let head_dim = x.dim(D::Minus1)?; + let half_dim = head_dim / 2; + + let x1 = x.narrow(D::Minus1, 0, half_dim)?; + let x2 = x.narrow(D::Minus1, half_dim, half_dim)?; + + // rotate_half gives [-x2, x1] + let x_rotated = Tensor::cat(&[&x2.neg()?, &x1], D::Minus1)?; + + // Apply: x * cos + rotate_half(x) * sin + x.broadcast_mul(cos)? + x_rotated.broadcast_mul(sin)? + } + + /// Apply Multimodal RoPE with export of intermediate tensors for debugging. + pub fn apply_multimodal_rotary_emb_with_export( + &self, + q: &Tensor, + k: &Tensor, + position_ids: &Tensor, + ) -> Result<(Tensor, Tensor, std::collections::HashMap)> { + use std::collections::HashMap; + let mut tensors: HashMap = HashMap::new(); + + let (three, _batch, _seq_len) = position_ids.dims3()?; + assert_eq!(three, 3, "position_ids must have 3 dimensions"); + + // Export position_ids + tensors.insert("position_ids".to_string(), position_ids.clone()); + + // Compute cos/sin for each position dimension + let (cos_3d, sin_3d) = self.compute_3d_rope_embeddings(position_ids)?; + tensors.insert("cos_3d".to_string(), cos_3d.clone()); + tensors.insert("sin_3d".to_string(), sin_3d.clone()); + + // Apply mrope_section to select appropriate bands + let (cos, sin) = self.apply_mrope_sections(&cos_3d, &sin_3d)?; + tensors.insert("cos_after_mrope".to_string(), cos.clone()); + tensors.insert("sin_after_mrope".to_string(), sin.clone()); + + // Export specific position for debugging (position 947 if available) + let seq_len = cos.dim(1)?; + if seq_len > 947 { + tensors.insert("cos_pos947".to_string(), cos.i((.., 947, ..))?.squeeze(1)?); + tensors.insert("sin_pos947".to_string(), sin.i((.., 947, ..))?.squeeze(1)?); + } + + // Reshape for broadcasting: [batch, 1, seq_len, head_dim] + let cos = cos.unsqueeze(1)?; + let sin = sin.unsqueeze(1)?; + + // Apply RoPE to q and k + let q_embed = self.apply_rope_to_tensor(q, &cos, &sin)?; + let k_embed = self.apply_rope_to_tensor(k, &cos, &sin)?; + + Ok((q_embed, k_embed, tensors)) + } +} + +/// Image grid specification for multi-image M-RoPE position computation. +#[derive(Debug, Clone)] +pub struct ImageGrid { + /// Grid height (number of patches in height dimension, after spatial merge) + pub grid_h: usize, + /// Grid width (number of patches in width dimension, after spatial merge) + pub grid_w: usize, +} + +/// Compute 3D M-RoPE position IDs for multi-image multimodal input. +/// +/// This function creates position IDs of shape [3, batch, seq_len] for inputs +/// containing multiple images. Each image's tokens get 2D spatial positions, +/// while text tokens get sequential 1D positions. +/// +/// # Position Layout +/// ```text +/// Text tokens: all 3 dims same (t=h=w=pos) +/// Image tokens: 2D grid positions offset by preceding text +/// - pos_t = offset (temporal = 0 for images) +/// - pos_h = row_in_grid + offset +/// - pos_w = col_in_grid + offset +/// ``` +/// +/// # Arguments +/// * `input_ids` - Token IDs of shape (batch, seq_len) +/// * `image_token_id` - The token ID used for image placeholders +/// * `image_grids` - Grid dimensions for each image (in order of appearance) +/// * `device` - Device to create tensors on +/// +/// # Returns +/// Position IDs tensor of shape [3, batch, seq_len] +pub fn compute_mrope_position_ids_multi( + input_ids: &Tensor, + image_token_id: u32, + image_grids: &[ImageGrid], + device: &Device, +) -> Result { + let (batch, seq_len) = input_ids.dims2()?; + let input_ids_vec: Vec = input_ids.flatten_all()?.to_vec1()?; + + // Create position IDs for all 3 dimensions + let mut pos_t = vec![0i64; batch * seq_len]; + let mut pos_h = vec![0i64; batch * seq_len]; + let mut pos_w = vec![0i64; batch * seq_len]; + + for b in 0..batch { + let batch_start = b * seq_len; + + // Find all image token ranges + let mut image_ranges: Vec<(usize, usize)> = Vec::new(); // (start, end) exclusive + let mut in_image = false; + let mut image_start = 0usize; + + for s in 0..seq_len { + let token_id = input_ids_vec[batch_start + s]; + if token_id == image_token_id { + if !in_image { + in_image = true; + image_start = s; + } + } else if in_image { + image_ranges.push((image_start, s)); + in_image = false; + } + } + // Handle case where image tokens extend to end of sequence + if in_image { + image_ranges.push((image_start, seq_len)); + } + + // Verify we have the right number of image ranges + if image_ranges.len() != image_grids.len() { + return Err(candle::Error::Msg(format!( + "Mismatch: found {} image ranges but {} grids provided", + image_ranges.len(), + image_grids.len() + ))); + } + + // Compute positions + let mut current_pos = 0i64; + let mut range_idx = 0usize; + + for s in 0..seq_len { + let idx = batch_start + s; + + // Check if we're at the start of an image range + if range_idx < image_ranges.len() && s == image_ranges[range_idx].0 { + // Process entire image range + let (img_start, img_end) = image_ranges[range_idx]; + let grid = &image_grids[range_idx]; + let num_vision_tokens = grid.grid_h * grid.grid_w; + + // Verify token count matches grid + let actual_tokens = img_end - img_start; + if actual_tokens != num_vision_tokens { + return Err(candle::Error::Msg(format!( + "Image {} has {} tokens but grid {}x{} = {} expected", + range_idx, actual_tokens, grid.grid_h, grid.grid_w, num_vision_tokens + ))); + } + + // Assign spatial positions to vision tokens + let offset = current_pos; + for vision_idx in 0..num_vision_tokens { + let token_s = img_start + vision_idx; + let token_idx = batch_start + token_s; + + let t_pos = 0i64; // Temporal is 0 for images + let h_pos = (vision_idx / grid.grid_w) as i64; + let w_pos = (vision_idx % grid.grid_w) as i64; + + pos_t[token_idx] = t_pos + offset; + pos_h[token_idx] = h_pos + offset; + pos_w[token_idx] = w_pos + offset; + } + + // Update current_pos to max position in this image + 1 + let max_h = (grid.grid_h - 1) as i64; + let max_w = (grid.grid_w - 1) as i64; + current_pos = offset + max_h.max(max_w) + 1; + + range_idx += 1; + continue; + } + + // Skip if we're inside an image range (already processed) + if range_idx > 0 { + let prev_range = image_ranges[range_idx - 1]; + if s >= prev_range.0 && s < prev_range.1 { + continue; + } + } + if range_idx < image_ranges.len() { + let curr_range = image_ranges[range_idx]; + if s >= curr_range.0 && s < curr_range.1 { + continue; + } + } + + // Text token: all dimensions same + pos_t[idx] = current_pos; + pos_h[idx] = current_pos; + pos_w[idx] = current_pos; + current_pos += 1; + } + } + + // Create tensors and stack + let pos_t = Tensor::from_vec(pos_t, (batch, seq_len), device)?; + let pos_h = Tensor::from_vec(pos_h, (batch, seq_len), device)?; + let pos_w = Tensor::from_vec(pos_w, (batch, seq_len), device)?; + + Tensor::stack(&[pos_t, pos_h, pos_w], 0) +} + +/// Compute 3D M-RoPE position IDs for multimodal input. +/// +/// This function creates position IDs of shape [3, batch, seq_len] following PyTorch's +/// get_rope_index() algorithm: +/// - Text tokens before vision: all 3 dims same, starting from 0 +/// - Vision tokens: (temporal + offset, height + offset, width + offset) +/// - Text tokens after vision: all 3 dims same, continuing from max vision position + 1 +/// +/// For vision tokens, positions encode the 2D spatial structure offset by preceding text. +pub fn compute_mrope_position_ids( + input_ids: &Tensor, + image_token_id: u32, + grid_h: usize, + grid_w: usize, + device: &Device, +) -> Result { + let (batch, seq_len) = input_ids.dims2()?; + let input_ids_vec: Vec = input_ids.flatten_all()?.to_vec1()?; + + // Create position IDs for all 3 dimensions + let mut pos_t = vec![0i64; batch * seq_len]; + let mut pos_h = vec![0i64; batch * seq_len]; + let mut pos_w = vec![0i64; batch * seq_len]; + + for b in 0..batch { + // Find the first image token position + let batch_start = b * seq_len; + let mut first_image_pos = None; + for s in 0..seq_len { + if input_ids_vec[batch_start + s] == image_token_id { + first_image_pos = Some(s); + break; + } + } + + // Compute positions following PyTorch's algorithm + let num_vision_tokens = grid_h * grid_w; + + // Text tokens before vision get sequential positions + let text_before = first_image_pos.unwrap_or(seq_len); + for s in 0..text_before { + let idx = batch_start + s; + pos_t[idx] = s as i64; + pos_h[idx] = s as i64; + pos_w[idx] = s as i64; + } + + // Vision tokens: (temporal, height, width) + text_before offset + let offset = text_before as i64; + let mut vision_idx = 0usize; + let mut max_vision_pos = offset - 1; // Will be updated + + for s in text_before..seq_len { + let idx = batch_start + s; + let token_id = input_ids_vec[idx]; + + if token_id == image_token_id && vision_idx < num_vision_tokens { + // Vision token: spatial position + offset + let t_pos = 0i64; // Temporal is 0 for images + let h_pos = (vision_idx / grid_w) as i64; + let w_pos = (vision_idx % grid_w) as i64; + + pos_t[idx] = t_pos + offset; + pos_h[idx] = h_pos + offset; + pos_w[idx] = w_pos + offset; + + // Track max position for text tokens that follow + max_vision_pos = max_vision_pos + .max(pos_t[idx]) + .max(pos_h[idx]) + .max(pos_w[idx]); + + vision_idx += 1; + } else { + // Text token after vision: continue from max_vision_pos + 1 + max_vision_pos += 1; + pos_t[idx] = max_vision_pos; + pos_h[idx] = max_vision_pos; + pos_w[idx] = max_vision_pos; + } + } + } + + // Create tensors and stack + let pos_t = Tensor::from_vec(pos_t, (batch, seq_len), device)?; + let pos_h = Tensor::from_vec(pos_h, (batch, seq_len), device)?; + let pos_w = Tensor::from_vec(pos_w, (batch, seq_len), device)?; + + Tensor::stack(&[pos_t, pos_h, pos_w], 0) +} + +/// Grid specification for video input. +/// +/// Unlike images which have only spatial dimensions (h, w), +/// video has temporal (t), height (h), and width (w) dimensions. +#[derive(Debug, Clone)] +pub struct VideoGrid { + /// Number of temporal frames (after any temporal patching) + pub grid_t: usize, + /// Number of height patches (after spatial merge) + pub grid_h: usize, + /// Number of width patches (after spatial merge) + pub grid_w: usize, +} + +/// Compute 3D M-RoPE position IDs for video input. +/// +/// Unlike multi-image (where t=0 for all images), video uses sequential +/// temporal positions (t=frame_index) to encode temporal relationships +/// between frames. +/// +/// Position encoding pattern for video with grid_t=3, grid_h=2, grid_w=2: +/// ```text +/// t_index = [0,0,0,0, 1,1,1,1, 2,2,2,2] // Temporal: repeats for h*w per frame +/// h_index = [0,0,1,1, 0,0,1,1, 0,0,1,1] // Height: repeats w times per t +/// w_index = [0,1,0,1, 0,1,0,1, 0,1,0,1] // Width: cycles fastest +/// ``` +/// +/// # Arguments +/// * `input_ids` - Token IDs of shape (batch, seq_len) +/// * `video_token_id` - The token ID used for video placeholders (different from image_token_id!) +/// * `video_grid` - Grid dimensions for the video (temporal, height, width) +/// * `second_per_grid_t` - Time interval per temporal grid unit (= temporal_patch_size / fps) +/// * `tokens_per_second` - Temporal position scaling factor (use 2 for video, matching HuggingFace) +/// * `device` - Device to create tensors on +/// +/// # Returns +/// Position IDs tensor of shape [3, batch, seq_len] +pub fn compute_mrope_position_ids_video( + input_ids: &Tensor, + video_token_id: u32, + video_grid: &VideoGrid, + second_per_grid_t: f32, + tokens_per_second: usize, + device: &Device, +) -> Result { + let (batch, seq_len) = input_ids.dims2()?; + let input_ids_vec: Vec = input_ids.flatten_all()?.to_vec1()?; + + let grid_t = video_grid.grid_t; + let grid_h = video_grid.grid_h; + let grid_w = video_grid.grid_w; + let num_vision_tokens = grid_t * grid_h * grid_w; + + // Create position IDs for all 3 dimensions + let mut pos_t = vec![0i64; batch * seq_len]; + let mut pos_h = vec![0i64; batch * seq_len]; + let mut pos_w = vec![0i64; batch * seq_len]; + + for b in 0..batch { + let batch_start = b * seq_len; + + // Find the video token range + let mut video_start = None; + let mut video_end = None; + let mut in_video = false; + + for s in 0..seq_len { + let token_id = input_ids_vec[batch_start + s]; + if token_id == video_token_id { + if !in_video { + in_video = true; + video_start = Some(s); + } + } else if in_video { + video_end = Some(s); + break; + } + } + // Handle case where video tokens extend to end of sequence + if in_video && video_end.is_none() { + video_end = Some(seq_len); + } + + // Verify video token count matches grid + if let (Some(start), Some(end)) = (video_start, video_end) { + let actual_tokens = end - start; + if actual_tokens != num_vision_tokens { + return Err(candle::Error::Msg(format!( + "Video has {} tokens but grid {}x{}x{} = {} expected", + actual_tokens, grid_t, grid_h, grid_w, num_vision_tokens + ))); + } + } + + // Compute positions + let mut current_pos = 0i64; + let video_range = video_start.zip(video_end); + + for s in 0..seq_len { + let idx = batch_start + s; + + // Check if we're at the start of the video range + if let Some((v_start, v_end)) = video_range { + if s == v_start { + // Process entire video range with 3D positions + let offset = current_pos; + + for vision_idx in 0..num_vision_tokens { + let token_s = v_start + vision_idx; + let token_idx = batch_start + token_s; + + // 3D position: t uses temporal scaling for proper frame spacing + // Formula: t_pos = frame_index * second_per_grid_t * tokens_per_second + // This matches HuggingFace Qwen2-VL processor behavior + let frame_index = vision_idx / (grid_h * grid_w); + let t_pos = (frame_index as f32 + * second_per_grid_t + * tokens_per_second as f32) as i64; + let spatial_idx = vision_idx % (grid_h * grid_w); + let h_pos = (spatial_idx / grid_w) as i64; + let w_pos = (spatial_idx % grid_w) as i64; + + pos_t[token_idx] = t_pos + offset; + pos_h[token_idx] = h_pos + offset; + pos_w[token_idx] = w_pos + offset; + } + + // Update current_pos to max position in video + 1 + // max_t also needs temporal scaling to match the scaled positions + let max_t = + ((grid_t - 1) as f32 * second_per_grid_t * tokens_per_second as f32) as i64; + let max_h = (grid_h - 1) as i64; + let max_w = (grid_w - 1) as i64; + current_pos = offset + max_t.max(max_h).max(max_w) + 1; + + continue; + } + + // Skip if we're inside the video range (already processed) + if s > v_start && s < v_end { + continue; + } + } + + // Text token: all dimensions same + pos_t[idx] = current_pos; + pos_h[idx] = current_pos; + pos_w[idx] = current_pos; + current_pos += 1; + } + } + + // Create tensors and stack + let pos_t = Tensor::from_vec(pos_t, (batch, seq_len), device)?; + let pos_h = Tensor::from_vec(pos_h, (batch, seq_len), device)?; + let pos_w = Tensor::from_vec(pos_w, (batch, seq_len), device)?; + + Tensor::stack(&[pos_t, pos_h, pos_w], 0) +} + +/// Gated MLP block (SwiGLU-style). +struct Mlp { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: candle_nn::Activation, +} + +impl Mlp { + fn new(cfg: &TextConfig, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + let gate_proj = linear_b(hidden_sz, intermediate_sz, cfg.use_bias, vb.pp("gate_proj"))?; + let up_proj = linear_b(hidden_sz, intermediate_sz, cfg.use_bias, vb.pp("up_proj"))?; + let down_proj = linear_b(intermediate_sz, hidden_sz, cfg.use_bias, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_act, + }) + } + + fn forward(&self, xs: &Tensor) -> Result { + let lhs = self.gate_proj.forward(xs)?.apply(&self.act_fn)?; + let rhs = self.up_proj.forward(xs)?; + self.down_proj.forward(&(lhs * rhs)?) + } + + /// Forward with intermediate tensor export for debugging. + fn forward_with_export( + &self, + xs: &Tensor, + ) -> Result<(Tensor, std::collections::HashMap)> { + use std::collections::HashMap; + let mut tensors: HashMap = HashMap::new(); + + // gate_proj: hidden_size -> intermediate_size + let gate_out = self.gate_proj.forward(xs)?; + tensors.insert("gate_proj_out".to_string(), gate_out.clone()); + + // Activation (SiLU) + let gate_act = gate_out.apply(&self.act_fn)?; + tensors.insert("gate_act_out".to_string(), gate_act.clone()); + + // up_proj: hidden_size -> intermediate_size + let up_out = self.up_proj.forward(xs)?; + tensors.insert("up_proj_out".to_string(), up_out.clone()); + + // Element-wise multiplication + let mul_out = (&gate_act * &up_out)?; + tensors.insert("gate_up_mul".to_string(), mul_out.clone()); + + // down_proj: intermediate_size -> hidden_size + let output = self.down_proj.forward(&mul_out)?; + tensors.insert("down_proj_out".to_string(), output.clone()); + + Ok((output, tensors)) + } +} + +/// Multi-head attention with Grouped Query Attention (GQA). +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, + softmax_scale: f64, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &TextConfig, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let head_dim = cfg.head_dim; + let num_kv_groups = num_heads / num_kv_heads; + + let q_proj = linear_b( + hidden_sz, + num_heads * head_dim, + cfg.use_bias, + vb.pp("q_proj"), + )?; + let k_proj = linear_b( + hidden_sz, + num_kv_heads * head_dim, + cfg.use_bias, + vb.pp("k_proj"), + )?; + let v_proj = linear_b( + hidden_sz, + num_kv_heads * head_dim, + cfg.use_bias, + vb.pp("v_proj"), + )?; + let o_proj = linear_b( + num_heads * head_dim, + hidden_sz, + cfg.use_bias, + vb.pp("o_proj"), + )?; + + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + rotary_emb, + kv_cache: None, + softmax_scale: 1.0 / (head_dim as f64).sqrt(), + }) + } + + /// Forward with 3D M-RoPE. + fn forward_with_mrope( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + position_ids: &Tensor, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + // Apply M-RoPE (3D position IDs) + let (query_states, key_states) = self.rotary_emb.apply_multimodal_rotary_emb( + &query_states, + &key_states, + position_ids, + )?; + + self.compute_attention( + query_states, + key_states, + value_states, + attention_mask, + b_sz, + q_len, + ) + } + + /// Shared attention computation. + fn compute_attention( + &mut self, + query_states: Tensor, + key_states: Tensor, + value_states: Tensor, + attention_mask: Option<&Tensor>, + b_sz: usize, + q_len: usize, + ) -> Result { + // KV cache handling + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + // Repeat KV heads for GQA (matches PyTorch's repeat_kv) + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?.contiguous()?; + let value_states = + crate::utils::repeat_kv(value_states, self.num_kv_groups)?.contiguous()?; + + // Compute attention (matches eager_attention_forward_ernie) + let attn_output = { + // attn_weights = query @ key^T * scaling + let attn_weights = + (query_states.matmul(&key_states.transpose(2, 3)?)? * self.softmax_scale)?; + + // Apply causal mask + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + // Softmax in F32 for stability (matches PyTorch's softmax(..., dtype=torch.float32).to(query.dtype)) + let original_dtype = attn_weights.dtype(); + let attn_weights = if original_dtype != DType::F32 { + let attn_weights = attn_weights.to_dtype(DType::F32)?; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.to_dtype(original_dtype)? + } else { + candle_nn::ops::softmax_last_dim(&attn_weights)? + }; + // attn_output = attn_weights @ value + attn_weights.matmul(&value_states)? + }; + + // attn_output.transpose(1, 2).contiguous().reshape(...) + attn_output + .transpose(1, 2)? + .contiguous()? + .reshape((b_sz, q_len, self.num_heads * self.head_dim))? + .apply(&self.o_proj) + } + + /// Forward with 3D M-RoPE and export attention intermediates (for debugging). + /// Matches PyTorch's Ernie4_5Attention.forward + eager_attention_forward_ernie exactly. + pub fn forward_with_mrope_export( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + position_ids: &Tensor, + ) -> Result<(Tensor, std::collections::HashMap)> { + use std::collections::HashMap; + let mut tensors: HashMap = HashMap::new(); + + let (b_sz, q_len, _) = xs.dims3()?; + + // Q, K, V projections (matches: query_states = self.q_proj(hidden_states)) + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + // Reshape to [batch, seq, heads, head_dim] then transpose to [batch, heads, seq, head_dim] + // matches: .view(hidden_shape).transpose(1, 2) + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + tensors.insert("q_pre_rope".to_string(), query_states.clone()); + tensors.insert("k_pre_rope".to_string(), key_states.clone()); + tensors.insert("v".to_string(), value_states.clone()); + + // Apply M-RoPE with export (matches: apply_multimodal_rotary_pos_emb) + let (query_states, key_states, rope_tensors) = self + .rotary_emb + .apply_multimodal_rotary_emb_with_export(&query_states, &key_states, position_ids)?; + + // Merge RoPE tensors with prefix + for (k, v) in rope_tensors { + tensors.insert(format!("rope_{}", k), v); + } + + tensors.insert("q_post_rope".to_string(), query_states.clone()); + tensors.insert("k_post_rope".to_string(), key_states.clone()); + + // No KV cache during prefill + // Repeat KV heads for GQA (matches: repeat_kv in eager_attention_forward_ernie) + let key_states_repeated = + crate::utils::repeat_kv(key_states.clone(), self.num_kv_groups)?.contiguous()?; + let value_states_repeated = + crate::utils::repeat_kv(value_states.clone(), self.num_kv_groups)?.contiguous()?; + + tensors.insert("k_repeated".to_string(), key_states_repeated.clone()); + tensors.insert("v_repeated".to_string(), value_states_repeated.clone()); + + // Attention scores: Q @ K^T * scaling (matches: torch.matmul(query, key_states.transpose(2, 3)) * scaling) + let attn_weights_pre = + (query_states.matmul(&key_states_repeated.transpose(2, 3)?)? * self.softmax_scale)?; + // Skip exporting full attention matrices - too large ([1, 16, 1357, 1357]) + // Just export a slice for verification: last row of attention for each head + let seq_len = attn_weights_pre.dim(2)?; + let attn_last_row = attn_weights_pre.narrow(2, seq_len - 1, 1)?; + tensors.insert("attn_weights_last_row".to_string(), attn_last_row); + + // Apply mask (matches: attn_weights = attn_weights + causal_mask) + let attn_weights_masked = match attention_mask { + None => attn_weights_pre, + Some(mask) => attn_weights_pre.broadcast_add(mask)?, + }; + + // Softmax (matches: softmax(..., dtype=torch.float32).to(query.dtype)) + let original_dtype = attn_weights_masked.dtype(); + let attn_weights = if original_dtype != DType::F32 { + let attn_weights = attn_weights_masked.to_dtype(DType::F32)?; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.to_dtype(original_dtype)? + } else { + candle_nn::ops::softmax_last_dim(&attn_weights_masked)? + }; + // Export last row of softmax attention weights + let attn_softmax_last_row = attn_weights.narrow(2, seq_len - 1, 1)?; + tensors.insert( + "attn_weights_softmax_last_row".to_string(), + attn_softmax_last_row, + ); + + // Attention output (matches: torch.matmul(attn_weights, value_states)) + let attn_output = attn_weights.matmul(&value_states_repeated)?; + tensors.insert("attn_output_pre_transpose".to_string(), attn_output.clone()); + + // Reshape (matches: .transpose(1, 2).contiguous()) + let attn_output = attn_output.transpose(1, 2)?.contiguous()?.reshape(( + b_sz, + q_len, + self.num_heads * self.head_dim, + ))?; + + // Output projection (matches: self.o_proj(attn_output)) + let output = self.o_proj.forward(&attn_output)?; + tensors.insert("attn_output".to_string(), output.clone()); + + Ok((output, tensors)) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None; + } +} + +/// Decoder layer with pre-norm architecture. +struct DecoderLayer { + self_attn: Attention, + mlp: Mlp, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl DecoderLayer { + fn new(rotary_emb: Arc, cfg: &TextConfig, vb: VarBuilder) -> Result { + let self_attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let mlp = Mlp::new(cfg, vb.pp("mlp"))?; + let input_layernorm = + rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let post_attention_layernorm = rms_norm( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + }) + } + + /// Forward with 3D M-RoPE. + fn forward_with_mrope( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + position_ids: &Tensor, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self + .self_attn + .forward_with_mrope(&xs, attention_mask, position_ids)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = self + .mlp + .forward(&xs.apply(&self.post_attention_layernorm)?)?; + residual + xs + } + + /// Forward with 3D M-RoPE and export attention intermediates. + fn forward_with_mrope_export( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + position_ids: &Tensor, + ) -> Result<(Tensor, std::collections::HashMap)> { + use std::collections::HashMap; + let mut tensors: HashMap = HashMap::new(); + + let residual = xs; + tensors.insert("layer_input".to_string(), xs.clone()); + + let xs = self.input_layernorm.forward(xs)?; + tensors.insert("post_input_layernorm".to_string(), xs.clone()); + + let (attn_out, attn_tensors) = + self.self_attn + .forward_with_mrope_export(&xs, attention_mask, position_ids)?; + + // Merge attention tensors with prefix + for (k, v) in attn_tensors { + tensors.insert(format!("attn_{}", k), v); + } + + let xs = (attn_out + residual)?; + tensors.insert("post_attn_residual".to_string(), xs.clone()); + + let residual = &xs; + let post_norm = xs.apply(&self.post_attention_layernorm)?; + tensors.insert("post_attention_layernorm".to_string(), post_norm.clone()); + + // Use MLP forward with export to capture intermediate values + let (mlp_out, mlp_tensors) = self.mlp.forward_with_export(&post_norm)?; + + // Merge MLP tensors with prefix + for (k, v) in mlp_tensors { + tensors.insert(format!("mlp_{}", k), v); + } + + tensors.insert("mlp_output".to_string(), mlp_out.clone()); + + let output = (residual + mlp_out)?; + tensors.insert("layer_output".to_string(), output.clone()); + + Ok((output, tensors)) + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache(); + } +} + +/// PaddleOCR-VL Text Model (ERNIE-4.5 based). +pub struct TextModel { + embed_tokens: Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: Linear, + pub dtype: DType, + pub hidden_size: usize, + device: Device, +} + +impl TextModel { + pub fn new(cfg: &TextConfig, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + + let embed_tokens = embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + + let rotary_emb = Arc::new(RotaryEmbedding::new(cfg, vb.device(), vb.dtype())?); + + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(rotary_emb.clone(), cfg, vb_l.pp(layer_idx))?; + layers.push(layer); + } + + let norm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb_m.pp("norm"))?; + + let lm_head = if cfg.tie_word_embeddings { + Linear::new(embed_tokens.embeddings().clone(), None) + } else { + linear_b(cfg.hidden_size, cfg.vocab_size, false, vb.pp("lm_head"))? + }; + + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + dtype: vb.dtype(), + hidden_size: cfg.hidden_size, + device: vb.device().clone(), + }) + } + + /// Get token embeddings. + pub fn embed_tokens(&self, input_ids: &Tensor) -> Result { + self.embed_tokens.forward(input_ids) + } + + /// Prepare causal attention mask. + fn prepare_causal_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + let mask: Vec = (0..tgt_len) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { f32::NEG_INFINITY } else { 0f32 })) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + /// Forward pass with embeddings using 3D M-RoPE. + /// + /// This method is used for all forward passes (both prefill and generation). + /// M-RoPE must always be used to maintain consistency with the prefill positions. + pub fn forward_embeds_with_mrope( + &mut self, + mut xs: Tensor, + position_ids: &Tensor, + ) -> Result { + let (b_sz, seq_len, _) = xs.dims3()?; + + // Create causal attention mask for prefill + let attention_mask = if seq_len <= 1 { + None + } else { + Some(self.prepare_causal_attention_mask(b_sz, seq_len, 0)?) + }; + + for layer in self.layers.iter_mut() { + xs = layer.forward_with_mrope(&xs, attention_mask.as_ref(), position_ids)?; + } + + xs = xs.apply(&self.norm)?; + + // Only compute logits for last token + self.lm_head + .forward(&xs)? + .i((.., seq_len - 1, ..))? + .contiguous() + } + + /// Clear all KV caches. + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache(); + } + } + + /// Forward pass with M-RoPE and tensor export for debugging. + /// + /// Captures intermediate tensors at key checkpoints for comparison with PyTorch. + /// Layer 1 exports detailed attention intermediates for GQA repeat_kv debugging. + pub fn forward_embeds_with_mrope_export( + &mut self, + mut xs: Tensor, + position_ids: &Tensor, + ) -> Result<(Tensor, std::collections::HashMap)> { + use std::collections::HashMap; + + let mut tensors: HashMap = HashMap::new(); + let (b_sz, seq_len, _) = xs.dims3()?; + + // Causal attention mask + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_causal_attention_mask(b_sz, seq_len, 0)?; + tensors.insert("causal_mask".to_string(), mask.clone()); + Some(mask) + }; + + tensors.insert("layer0_input".to_string(), xs.clone()); + + // Forward through ALL layers, capturing each output + // Layer 1 gets detailed attention export for debugging + for (i, layer) in self.layers.iter_mut().enumerate() { + if i == 1 { + // Layer 1: export all attention intermediates + let (layer_out, layer_tensors) = + layer.forward_with_mrope_export(&xs, attention_mask.as_ref(), position_ids)?; + xs = layer_out; + // Add layer 1 tensors with prefix + for (k, v) in layer_tensors { + tensors.insert(format!("layer1_{}", k), v); + } + } else { + xs = layer.forward_with_mrope(&xs, attention_mask.as_ref(), position_ids)?; + } + // Capture EVERY layer output for detailed comparison + tensors.insert(format!("layer_{}_output", i), xs.clone()); + } + + // Final layer norm + xs = xs.apply(&self.norm)?; + tensors.insert("final_hidden_state".to_string(), xs.clone()); + + // LM head - compute full logits + let logits = self.lm_head.forward(&xs)?; + tensors.insert("logits".to_string(), logits.clone()); + + Ok((logits, tensors)) + } +} diff --git a/patches/candle-transformers/src/models/paddleocr_vl/vision.rs b/patches/candle-transformers/src/models/paddleocr_vl/vision.rs new file mode 100644 index 0000000000..9756cb7d28 --- /dev/null +++ b/patches/candle-transformers/src/models/paddleocr_vl/vision.rs @@ -0,0 +1,1222 @@ +//! PaddleOCR-VL Vision Encoder. +//! +//! NaViT-style dynamic resolution visual encoder with 2D rotary position embeddings. + +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::{layer_norm, linear_b, LayerNorm, LayerNormConfig, Linear, Module, VarBuilder}; +use std::cell::RefCell; +use std::collections::HashMap; + +use super::config::VisionConfig; + +/// Default maximum number of cached position embeddings. +const DEFAULT_POS_EMBED_CACHE_SIZE: usize = 16; + +/// LFU (Least Frequently Used) cache for interpolated position embeddings. +/// +/// Caches interpolated position embeddings keyed by (height, width) grid dimensions. +/// Uses frequency-based eviction when cache is full: the least frequently accessed +/// entry is evicted first. This matches PyTorch's caching behavior. +struct PosEmbedCache { + /// Cached embeddings: (height, width) -> tensor + cache: HashMap<(usize, usize), Tensor>, + /// Access frequency for each key + frequency: HashMap<(usize, usize), usize>, + /// Maximum cache size + max_size: usize, +} + +impl PosEmbedCache { + fn new(max_size: usize) -> Self { + Self { + cache: HashMap::with_capacity(max_size), + frequency: HashMap::with_capacity(max_size), + max_size, + } + } + + /// Get a cached embedding, incrementing its access frequency. + fn get(&mut self, key: (usize, usize)) -> Option { + if let Some(tensor) = self.cache.get(&key) { + *self.frequency.entry(key).or_insert(0) += 1; + Some(tensor.clone()) + } else { + None + } + } + + /// Insert an embedding into the cache, evicting LFU entry if full. + fn insert(&mut self, key: (usize, usize), tensor: Tensor) { + // If already in cache, just update + if let std::collections::hash_map::Entry::Occupied(mut e) = self.cache.entry(key) { + e.insert(tensor); + *self.frequency.entry(key).or_insert(0) += 1; + return; + } + + // Evict LFU entry if at capacity + if self.cache.len() >= self.max_size { + if let Some((&lfu_key, _)) = self.frequency.iter().min_by_key(|(_, &freq)| freq) { + self.cache.remove(&lfu_key); + self.frequency.remove(&lfu_key); + } + } + + // Insert new entry + self.cache.insert(key, tensor); + self.frequency.insert(key, 1); + } + + /// Clear all cached embeddings. + #[allow(dead_code)] + fn clear(&mut self) { + self.cache.clear(); + self.frequency.clear(); + } +} + +/// Patch embedding using Conv2d with interpolated position embedding. +/// +/// Weight names: +/// - embeddings.patch_embedding.{weight,bias} +/// - embeddings.position_embedding.weight (base 27×27 grid for interpolation) +/// - embeddings.packing_position_embedding.weight (fallback, 32768 positions) +/// +/// For dynamic resolution images, the base position embedding grid is bilinearly +/// interpolated to match the actual patch grid size. Interpolated embeddings are +/// cached with LFU eviction to avoid redundant computation. +struct PatchEmbedding { + patch_embedding: candle_nn::Conv2d, + position_embedding: Tensor, // (num_positions, hidden_size) where num_positions = (image_size/patch_size)^2 + #[allow(dead_code)] + packing_position_embedding: candle_nn::Embedding, // Fallback, kept for weight loading + base_grid_size: usize, // sqrt(num_positions), typically 27 for 384/14 + hidden_size: usize, + /// Cache for interpolated position embeddings (LFU eviction) + pos_embed_cache: RefCell, +} + +impl PatchEmbedding { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let conv_cfg = candle_nn::Conv2dConfig { + stride: cfg.patch_size, + ..Default::default() + }; + // Weight: embeddings.patch_embedding (with bias) + let patch_embedding = candle_nn::conv2d( + cfg.num_channels, + cfg.hidden_size, + cfg.patch_size, + conv_cfg, + vb.pp("patch_embedding"), + )?; + + // Weight: embeddings.position_embedding (base grid for interpolation) + // Shape: (num_positions, hidden_size) where num_positions = (image_size/patch_size)^2 + let base_grid_size = cfg.image_size / cfg.patch_size; + let num_positions = base_grid_size * base_grid_size; + let position_embedding = vb + .pp("position_embedding") + .get((num_positions, cfg.hidden_size), "weight")?; + + // Weight: embeddings.packing_position_embedding (32768 positions) - kept for compatibility + let packing_position_embedding = + candle_nn::embedding(32768, cfg.hidden_size, vb.pp("packing_position_embedding"))?; + + Ok(Self { + patch_embedding, + position_embedding, + packing_position_embedding, + base_grid_size, + hidden_size: cfg.hidden_size, + pos_embed_cache: RefCell::new(PosEmbedCache::new(DEFAULT_POS_EMBED_CACHE_SIZE)), + }) + } + + /// Bilinearly interpolate position embeddings to match target grid size. + /// + /// Takes the base position embedding grid (e.g., 27×27) and interpolates it + /// to the target size (e.g., 72×58) using bilinear interpolation. + /// + /// This matches PyTorch's nn.functional.interpolate with mode='bilinear', align_corners=False. + /// Results are cached with LFU eviction to avoid redundant computation. + fn interpolate_pos_encoding(&self, target_h: usize, target_w: usize) -> Result { + let cache_key = (target_h, target_w); + + // Check cache first + if let Some(cached) = self.pos_embed_cache.borrow_mut().get(cache_key) { + return Ok(cached); + } + + let device = self.position_embedding.device(); + let dtype = self.position_embedding.dtype(); + let base_h = self.base_grid_size; + let base_w = self.base_grid_size; + + // If target matches base, just reshape and return (also cache it) + if target_h == base_h && target_w == base_w { + let result = self + .position_embedding + .reshape((1, target_h * target_w, self.hidden_size))? + .to_dtype(dtype)?; + self.pos_embed_cache + .borrow_mut() + .insert(cache_key, result.clone()); + return Ok(result); + } + + // Reshape position embedding to (base_h, base_w, hidden) + let pos_embed = self.position_embedding.to_dtype(DType::F32)?.reshape(( + base_h, + base_w, + self.hidden_size, + ))?; + + // Compute scale factors (align_corners=False style) + let scale_h = base_h as f64 / target_h as f64; + let scale_w = base_w as f64 / target_w as f64; + + // Build interpolated output + let mut output_data = Vec::with_capacity(target_h * target_w * self.hidden_size); + + for ty in 0..target_h { + for tx in 0..target_w { + // Source coordinates (align_corners=False: map center to center) + let sy = (ty as f64 + 0.5) * scale_h - 0.5; + let sx = (tx as f64 + 0.5) * scale_w - 0.5; + + // Clamp to valid range + let sy = sy.max(0.0).min((base_h - 1) as f64); + let sx = sx.max(0.0).min((base_w - 1) as f64); + + // Integer and fractional parts + let sy0 = sy.floor() as usize; + let sx0 = sx.floor() as usize; + let sy1 = (sy0 + 1).min(base_h - 1); + let sx1 = (sx0 + 1).min(base_w - 1); + let fy = (sy - sy0 as f64) as f32; + let fx = (sx - sx0 as f64) as f32; + + // Bilinear weights + let w00 = (1.0 - fy) * (1.0 - fx); + let w01 = (1.0 - fy) * fx; + let w10 = fy * (1.0 - fx); + let w11 = fy * fx; + + // Get the 4 corner embeddings + let e00: Vec = pos_embed.i((sy0, sx0))?.to_vec1()?; + let e01: Vec = pos_embed.i((sy0, sx1))?.to_vec1()?; + let e10: Vec = pos_embed.i((sy1, sx0))?.to_vec1()?; + let e11: Vec = pos_embed.i((sy1, sx1))?.to_vec1()?; + + // Interpolate each dimension + for d in 0..self.hidden_size { + let val = w00 * e00[d] + w01 * e01[d] + w10 * e10[d] + w11 * e11[d]; + output_data.push(val); + } + } + } + + // Create output tensor and cache it + let result = Tensor::from_vec( + output_data, + (1, target_h * target_w, self.hidden_size), + device, + )? + .to_dtype(dtype)?; + self.pos_embed_cache + .borrow_mut() + .insert(cache_key, result.clone()); + Ok(result) + } + + /// Forward pass with interpolated position embeddings for dynamic resolution. + fn forward(&self, xs: &Tensor) -> Result { + // Input: (batch, channels, height, width) + // Output: (batch, num_patches, hidden_size) + let xs = self.patch_embedding.forward(xs)?; + let (batch, hidden, h, w) = xs.dims4()?; + let num_patches = h * w; + + // Reshape to (batch, num_patches, hidden) + let xs = xs.reshape((batch, hidden, num_patches))?.transpose(1, 2)?; + + // Get interpolated position embedding for this grid size + let pos_embed = self.interpolate_pos_encoding(h, w)?; + + // Broadcast add position embedding to each batch + xs.broadcast_add(&pos_embed) + } +} + +/// 2D Rotary Position Embedding for vision. +struct VisionRotaryEmbedding { + inv_freq: Tensor, +} + +impl VisionRotaryEmbedding { + const THETA: f32 = 10000.0; + + fn new(dim: usize, device: &Device) -> Result { + let inv_freq: Vec = (0..dim) + .step_by(2) + .map(|i| 1f32 / Self::THETA.powf(i as f32 / dim as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + Ok(Self { + inv_freq: Tensor::from_vec(inv_freq, (1, inv_freq_len), device)?, + }) + } + + fn make_embeds(&self, seqlen: usize) -> Result { + let seq = + Tensor::arange(0f32, seqlen as f32, self.inv_freq.device())?.unsqueeze(D::Minus1)?; + seq.broadcast_matmul(&self.inv_freq) + } +} + +fn rotate_half(xs: &Tensor) -> Result { + let last_dim = xs.dim(D::Minus1)?; + let xs1 = xs.narrow(D::Minus1, 0, last_dim / 2)?; + let xs2 = xs.narrow(D::Minus1, last_dim / 2, last_dim - last_dim / 2)?; + Tensor::cat(&[&xs2.neg()?, &xs1], D::Minus1) +} + +fn apply_rotary_pos_emb_vision( + q: &Tensor, + k: &Tensor, + cos: &Tensor, + sin: &Tensor, +) -> Result<(Tensor, Tensor)> { + let cos = cos.unsqueeze(D::Minus2)?; + let sin = sin.unsqueeze(D::Minus2)?; + + let q_embed = (q.broadcast_mul(&cos)? + rotate_half(q)?.broadcast_mul(&sin)?)?; + let k_embed = (k.broadcast_mul(&cos)? + rotate_half(k)?.broadcast_mul(&sin)?)?; + Ok((q_embed, k_embed)) +} + +/// Tile size for chunked attention (KV positions per tile). +/// Balances memory usage vs throughput. 512 keeps peak memory under ~500MB per tile. +const ATTENTION_TILE_SIZE: usize = 512; + +/// Chunked attention with online softmax for memory efficiency. +/// +/// For large sequences that would exceed GPU memory limits (e.g., 14K+ patches from +/// high-resolution images), this processes K/V in tiles using the Flash Attention +/// online softmax algorithm. This is mathematically equivalent to standard attention +/// but never materializes the full (seq × seq) attention matrix. +/// +/// # Arguments +/// * `q` - Query tensor, shape (1, heads, q_seq, head_dim) +/// * `k` - Key tensor, shape (1, heads, kv_seq, head_dim) +/// * `v` - Value tensor, shape (1, heads, kv_seq, head_dim) +/// * `scale` - Attention scale factor (typically 1/sqrt(head_dim)) +/// +/// # Returns +/// Output tensor, shape (1, heads, q_seq, head_dim) +fn chunked_attention(q: &Tensor, k: &Tensor, v: &Tensor, scale: f64) -> Result { + let (_, num_heads, q_seq, head_dim) = q.dims4()?; + let kv_seq = k.dim(2)?; + let device = q.device(); + let dtype = q.dtype(); + + // For small sequences, use standard attention (fits in memory) + if kv_seq <= ATTENTION_TILE_SIZE { + let attn_weights = (q.matmul(&k.transpose(2, 3)?)? * scale)?; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + return attn_weights.matmul(v); + } + + // Chunked attention for large sequences using online softmax + let num_tiles = kv_seq.div_ceil(ATTENTION_TILE_SIZE); + + // Initialize accumulators in F32 for numerical stability. + // Use from_vec() to create properly contiguous tensors that work correctly + // with repeated broadcast operations across many loop iterations. + let output_accum_data = vec![0f32; num_heads * q_seq * head_dim]; + let mut output_accum = + Tensor::from_vec(output_accum_data, (1, num_heads, q_seq, head_dim), device)?; + let max_scores_data = vec![-f32::INFINITY; num_heads * q_seq]; + let mut max_scores = Tensor::from_vec(max_scores_data, (1, num_heads, q_seq, 1), device)?; + let sum_exps_data = vec![0f32; num_heads * q_seq]; + let mut sum_exps = Tensor::from_vec(sum_exps_data, (1, num_heads, q_seq, 1), device)?; + + let q_f32 = q.to_dtype(DType::F32)?; + + for tile_idx in 0..num_tiles { + let tile_start = tile_idx * ATTENTION_TILE_SIZE; + let tile_len = (kv_seq - tile_start).min(ATTENTION_TILE_SIZE); + + // Make tiles contiguous before dtype conversion (narrow creates a view) + let k_tile = k + .narrow(2, tile_start, tile_len)? + .contiguous()? + .to_dtype(DType::F32)?; + let v_tile = v + .narrow(2, tile_start, tile_len)? + .contiguous()? + .to_dtype(DType::F32)?; + + // Compute scores for this tile: (1, heads, q_seq, tile_len) + let scores_tile = (q_f32.matmul(&k_tile.transpose(2, 3)?)? * scale)?; + + // Get tile max: (1, heads, q_seq, 1) + let tile_max = scores_tile.max_keepdim(D::Minus1)?; + + // New running max + let new_max = max_scores.maximum(&tile_max)?; + + // Rescale previous accumulator: exp(old_max - new_max) + let rescale = (&max_scores - &new_max)?.exp()?; + output_accum = output_accum.broadcast_mul(&rescale)?; + sum_exps = sum_exps.broadcast_mul(&rescale)?; + + // Compute exp(scores - new_max) for this tile + let exp_scores = scores_tile.broadcast_sub(&new_max)?.exp()?; + + // Update accumulators + output_accum = (output_accum + exp_scores.matmul(&v_tile)?)?; + sum_exps = (sum_exps + exp_scores.sum_keepdim(D::Minus1)?)?; + + max_scores = new_max; + } + + // Final normalization and convert back to original dtype + output_accum.broadcast_div(&sum_exps)?.to_dtype(dtype) +} + +/// Vision MLP block. +struct VisionMlp { + fc1: Linear, + fc2: Linear, + act: candle_nn::Activation, +} + +impl VisionMlp { + fn new( + dim: usize, + hidden_dim: usize, + act: candle_nn::Activation, + vb: VarBuilder, + ) -> Result { + Ok(Self { + fc1: linear_b(dim, hidden_dim, true, vb.pp("fc1"))?, + fc2: linear_b(hidden_dim, dim, true, vb.pp("fc2"))?, + act, + }) + } + + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.fc1.forward(xs)?; + let xs = xs.apply(&self.act)?; + self.fc2.forward(&xs) + } +} + +/// Vision self-attention with 2D RoPE. +/// Weight names: +/// - self_attn.q_proj.{weight,bias} +/// - self_attn.k_proj.{weight,bias} +/// - self_attn.v_proj.{weight,bias} +/// - self_attn.out_proj.{weight,bias} +struct VisionAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + out_proj: Linear, + num_heads: usize, + head_dim: usize, + scale: f64, +} + +impl VisionAttention { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let dim = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let head_dim = dim / num_heads; + Ok(Self { + q_proj: linear_b(dim, dim, true, vb.pp("q_proj"))?, + k_proj: linear_b(dim, dim, true, vb.pp("k_proj"))?, + v_proj: linear_b(dim, dim, true, vb.pp("v_proj"))?, + out_proj: linear_b(dim, dim, true, vb.pp("out_proj"))?, + num_heads, + head_dim, + scale: (head_dim as f64).powf(-0.5), + }) + } + + fn forward( + &self, + xs: &Tensor, + cu_seqlens: &[usize], + cos: &Tensor, + sin: &Tensor, + ) -> Result { + self.forward_impl(xs, cu_seqlens, cos, sin, None) + } + + /// Forward pass with optional debug tensor export. + fn forward_with_debug( + &self, + xs: &Tensor, + cu_seqlens: &[usize], + cos: &Tensor, + sin: &Tensor, + exports: &mut HashMap, + ) -> Result { + self.forward_impl(xs, cu_seqlens, cos, sin, Some(exports)) + } + + fn forward_impl( + &self, + xs: &Tensor, + cu_seqlens: &[usize], + cos: &Tensor, + sin: &Tensor, + mut exports: Option<&mut HashMap>, + ) -> Result { + let seq_len = xs.dim(0)?; + + // Separate Q, K, V projections + let q = self.q_proj.forward(xs)?; + let k = self.k_proj.forward(xs)?; + let v = self.v_proj.forward(xs)?; + + // Export Q, K, V before reshape + if let Some(ref mut exp) = exports { + exp.insert("attn_q_proj".to_string(), q.to_dtype(DType::F32)?); + exp.insert("attn_k_proj".to_string(), k.to_dtype(DType::F32)?); + exp.insert("attn_v_proj".to_string(), v.to_dtype(DType::F32)?); + } + + // Reshape to (seq_len, num_heads, head_dim) + let mut q = q.reshape((seq_len, self.num_heads, self.head_dim))?; + let mut k = k.reshape((seq_len, self.num_heads, self.head_dim))?; + let mut v = v.reshape((seq_len, self.num_heads, self.head_dim))?; + + // Convert to f32 for precision in RoPE + let cos = cos.to_dtype(DType::F32)?; + let sin = sin.to_dtype(DType::F32)?; + q = q.to_dtype(DType::F32)?; + k = k.to_dtype(DType::F32)?; + v = v.to_dtype(DType::F32)?; + + // Export cos/sin and Q/K before RoPE + if let Some(ref mut exp) = exports { + exp.insert("rope_cos".to_string(), cos.clone()); + exp.insert("rope_sin".to_string(), sin.clone()); + exp.insert("q_before_rope".to_string(), q.clone()); + exp.insert("k_before_rope".to_string(), k.clone()); + } + + // Apply 2D RoPE + (q, k) = apply_rotary_pos_emb_vision(&q, &k, &cos, &sin)?; + + // Export Q/K after RoPE + if let Some(ref mut exp) = exports { + exp.insert("q_after_rope".to_string(), q.clone()); + exp.insert("k_after_rope".to_string(), k.clone()); + } + + // Process each image sequence separately (variable length) + let mut outputs = Vec::new(); + + for window in cu_seqlens.windows(2) { + let start = window[0]; + let end = window[1]; + if end <= start { + continue; + } + let len = end - start; + let q_chunk = q.narrow(0, start, len)?.transpose(0, 1)?.contiguous()?; + let k_chunk = k.narrow(0, start, len)?.transpose(0, 1)?.contiguous()?; + let v_chunk = v.narrow(0, start, len)?.transpose(0, 1)?.contiguous()?; + + let mut chunk_out = { + let q = q_chunk.unsqueeze(0)?; + let k = k_chunk.unsqueeze(0)?; + let v = v_chunk.unsqueeze(0)?; + + // Use chunked attention with online softmax for memory efficiency. + // For small sequences (<= 512), falls back to standard attention. + // For large sequences (14K+ patches), uses tiled computation to avoid OOM. + chunked_attention(&q, &k, &v, self.scale)? + }; + + chunk_out = chunk_out.squeeze(0)?.transpose(0, 1)?; + // Synchronize GPU before CPU accesses tensor data (critical for Metal correctness) + chunk_out.device().synchronize()?; + chunk_out = chunk_out.reshape((len, self.num_heads * self.head_dim))?; + outputs.push(chunk_out.to_dtype(xs.dtype())?); + } + + let attn_output = Tensor::cat(&outputs, 0)?; + + // Export before out_proj + if let Some(ref mut exp) = exports { + exp.insert( + "attn_output_before_proj".to_string(), + attn_output.to_dtype(DType::F32)?, + ); + } + + self.out_proj.forward(&attn_output) + } +} + +/// Vision encoder block (pre-norm transformer). +/// Weight names: +/// - layer_norm1.{weight,bias} +/// - layer_norm2.{weight,bias} +/// - self_attn.{q,k,v,out}_proj.{weight,bias} +/// - mlp.fc1.{weight,bias} +/// - mlp.fc2.{weight,bias} +struct VisionBlock { + layer_norm1: LayerNorm, + layer_norm2: LayerNorm, + self_attn: VisionAttention, + mlp: VisionMlp, +} + +impl VisionBlock { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let norm_cfg = LayerNormConfig { + eps: cfg.layer_norm_eps, + ..Default::default() + }; + Ok(Self { + layer_norm1: layer_norm(cfg.hidden_size, norm_cfg, vb.pp("layer_norm1"))?, + layer_norm2: layer_norm(cfg.hidden_size, norm_cfg, vb.pp("layer_norm2"))?, + self_attn: VisionAttention::new(cfg, vb.pp("self_attn"))?, + mlp: VisionMlp::new( + cfg.hidden_size, + cfg.intermediate_size, + cfg.hidden_act, + vb.pp("mlp"), + )?, + }) + } + + fn forward( + &self, + xs: &Tensor, + cu_seqlens: &[usize], + cos: &Tensor, + sin: &Tensor, + ) -> Result { + let normed = self.layer_norm1.forward(xs)?; + let attn_out = self.self_attn.forward(&normed, cu_seqlens, cos, sin)?; + let xs_att = xs.add(&attn_out)?; + let mlp_out = self.mlp.forward(&self.layer_norm2.forward(&xs_att)?)?; + xs_att.add(&mlp_out) + } + + /// Forward pass with debug tensor export for attention internals. + fn forward_with_debug( + &self, + xs: &Tensor, + cu_seqlens: &[usize], + cos: &Tensor, + sin: &Tensor, + exports: &mut HashMap, + ) -> Result { + let normed = self.layer_norm1.forward(xs)?; + exports.insert( + "layer0_after_norm1".to_string(), + normed.to_dtype(DType::F32)?, + ); + + let attn_out = self + .self_attn + .forward_with_debug(&normed, cu_seqlens, cos, sin, exports)?; + exports.insert( + "layer0_attn_output".to_string(), + attn_out.to_dtype(DType::F32)?, + ); + + let xs_att = xs.add(&attn_out)?; + exports.insert( + "layer0_after_attn_residual".to_string(), + xs_att.to_dtype(DType::F32)?, + ); + + let normed2 = self.layer_norm2.forward(&xs_att)?; + exports.insert( + "layer0_after_norm2".to_string(), + normed2.to_dtype(DType::F32)?, + ); + + let mlp_out = self.mlp.forward(&normed2)?; + exports.insert( + "layer0_mlp_output".to_string(), + mlp_out.to_dtype(DType::F32)?, + ); + + xs_att.add(&mlp_out) + } +} + +/// Projector (mlp_AR) - Vision-to-Text bridge. +/// +/// Projects vision features to text model dimension with 2×2 spatial merging. +/// Weight names: mlp_AR.pre_norm, mlp_AR.linear_1, mlp_AR.linear_2 +/// +/// The spatial merge gathers 2×2 patches from the image grid: +/// ```text +/// Input patches (raster order): Merged output: +/// [0, 1, 2, 3] [0+1+4+5, 2+3+6+7] +/// [4, 5, 6, 7] -> [8+9+12+13, 10+11+14+15] +/// [8, 9, 10, 11] +/// [12, 13, 14, 15] +/// ``` +pub struct Projector { + pre_norm: LayerNorm, + linear_1: Linear, + linear_2: Linear, + spatial_merge_size: usize, + hidden_size: usize, +} + +impl Projector { + pub fn new(cfg: &VisionConfig, text_hidden_size: usize, vb: VarBuilder) -> Result { + let merged_hidden_size = cfg.hidden_size * cfg.spatial_merge_size.pow(2); + let norm_cfg = LayerNormConfig { + eps: 1e-5, + ..Default::default() + }; + Ok(Self { + pre_norm: layer_norm(cfg.hidden_size, norm_cfg, vb.pp("pre_norm"))?, + linear_1: linear_b( + merged_hidden_size, + merged_hidden_size, + true, + vb.pp("linear_1"), + )?, + linear_2: linear_b( + merged_hidden_size, + text_hidden_size, + true, + vb.pp("linear_2"), + )?, + spatial_merge_size: cfg.spatial_merge_size, + hidden_size: cfg.hidden_size, + }) + } + + /// Forward pass with proper 2×2 spatial merge. + /// + /// Implements the einops pattern: "(t h m1 w m2) d -> (t h w) (m1 m2 d)" + /// where m1=m2=spatial_merge_size (typically 2). + pub fn forward(&self, xs: &Tensor, grid_thw: &Tensor) -> Result { + let normed = self.pre_norm.forward(xs)?; + + let grid = grid_thw.to_vec2::()?; + let m = self.spatial_merge_size; + + let mut merged_features = Vec::new(); + let mut offset = 0usize; + + for g in &grid { + let t = g[0] as usize; + let h = g[1] as usize; + let w = g[2] as usize; + let seq_len = t * h * w; + + // Extract this image's features + let features = normed.narrow(0, offset, seq_len)?; + offset += seq_len; + + // Reshape to (t, h, w, hidden) + let features = features.reshape((t, h, w, self.hidden_size))?; + + // Merged dimensions + let h_merged = h / m; + let w_merged = w / m; + + // Gather 2×2 blocks: for each merged position, collect m×m patches + // and concatenate their features + let mut blocks = Vec::with_capacity(t * h_merged * w_merged); + + for ti in 0..t { + for hi in 0..h_merged { + for wi in 0..w_merged { + // Collect m×m patches at this merged position + let mut patch_features = Vec::with_capacity(m * m); + for mi in 0..m { + for mj in 0..m { + let patch = features.i((ti, hi * m + mi, wi * m + mj))?; + patch_features.push(patch); + } + } + // Concatenate patch features: (m*m, hidden) -> (m*m * hidden,) + let block = Tensor::cat(&patch_features, 0)?; + blocks.push(block); + } + } + } + + // Stack all blocks: (t * h_merged * w_merged, merged_hidden) + let merged = Tensor::stack(&blocks, 0)?; + merged_features.push(merged); + } + + // Concatenate all images + let merged = Tensor::cat(&merged_features, 0)?; + + // Apply MLP + let xs = self.linear_1.forward(&merged)?; + let xs = xs.gelu()?; + self.linear_2.forward(&xs) + } + + /// Forward pass returning separate embeddings for each image. + /// + /// Unlike `forward()` which concatenates all image features, this method + /// returns a `Vec` where each tensor contains the embeddings for + /// one image. This enables the text model to inject each image's embeddings + /// at the correct positions in multi-image scenarios. + /// + /// # Arguments + /// * `xs` - Vision encoder output of shape (total_patches, hidden_size) + /// * `grid_thw` - Grid dimensions tensor of shape (num_images, 3) + /// + /// # Returns + /// Vector of tensors, one per image, each of shape (num_merged_patches, text_hidden_size) + pub fn forward_multi(&self, xs: &Tensor, grid_thw: &Tensor) -> Result> { + let normed = self.pre_norm.forward(xs)?; + + let grid = grid_thw.to_vec2::()?; + let m = self.spatial_merge_size; + + let mut result = Vec::with_capacity(grid.len()); + let mut offset = 0usize; + + for g in &grid { + let t = g[0] as usize; + let h = g[1] as usize; + let w = g[2] as usize; + let seq_len = t * h * w; + + // Extract this image's features + let features = normed.narrow(0, offset, seq_len)?; + offset += seq_len; + + // Reshape to (t, h, w, hidden) + let features = features.reshape((t, h, w, self.hidden_size))?; + + // Merged dimensions + let h_merged = h / m; + let w_merged = w / m; + + // Gather 2×2 blocks + let mut blocks = Vec::with_capacity(t * h_merged * w_merged); + + for ti in 0..t { + for hi in 0..h_merged { + for wi in 0..w_merged { + let mut patch_features = Vec::with_capacity(m * m); + for mi in 0..m { + for mj in 0..m { + let patch = features.i((ti, hi * m + mi, wi * m + mj))?; + patch_features.push(patch); + } + } + let block = Tensor::cat(&patch_features, 0)?; + blocks.push(block); + } + } + } + + // Stack all blocks: (t * h_merged * w_merged, merged_hidden) + let merged = Tensor::stack(&blocks, 0)?; + + // Apply MLP + let xs = self.linear_1.forward(&merged)?; + let xs = xs.gelu()?; + let projected = self.linear_2.forward(&xs)?; + + result.push(projected); + } + + Ok(result) + } +} + +/// PaddleOCR-VL Vision Model. +/// +/// NaViT-style encoder with 2D RoPE, supporting dynamic image resolutions. +/// Weight structure: +/// - embeddings.patch_embedding, embeddings.position_embedding +/// - encoder.layers.{i}.* +/// - post_layernorm +pub struct VisionModel { + embeddings: PatchEmbedding, + encoder_layers: Vec, + post_layernorm: LayerNorm, + projector: Projector, + rotary_pos_emb: VisionRotaryEmbedding, + hidden_size: usize, + patch_size: usize, +} + +impl VisionModel { + pub fn new( + vision_cfg: &VisionConfig, + text_hidden_size: usize, + vb: VarBuilder, + projector_vb: VarBuilder, + ) -> Result { + // Embeddings: embeddings.patch_embedding, embeddings.position_embedding + let embeddings = PatchEmbedding::new(vision_cfg, vb.pp("embeddings"))?; + + // Encoder layers: encoder.layers.{i}.* + let mut encoder_layers = Vec::with_capacity(vision_cfg.num_hidden_layers); + let vb_encoder = vb.pp("encoder").pp("layers"); + for i in 0..vision_cfg.num_hidden_layers { + encoder_layers.push(VisionBlock::new(vision_cfg, vb_encoder.pp(i))?); + } + + // Post layer norm: post_layernorm + let norm_cfg = LayerNormConfig { + eps: vision_cfg.layer_norm_eps, + ..Default::default() + }; + let post_layernorm = layer_norm(vision_cfg.hidden_size, norm_cfg, vb.pp("post_layernorm"))?; + + // Projector is separate at mlp_AR + let projector = Projector::new(vision_cfg, text_hidden_size, projector_vb)?; + + let head_dim = vision_cfg.head_dim(); + let rotary_pos_emb = VisionRotaryEmbedding::new(head_dim / 2, vb.device())?; + + Ok(Self { + embeddings, + encoder_layers, + post_layernorm, + projector, + rotary_pos_emb, + hidden_size: vision_cfg.hidden_size, + patch_size: vision_cfg.patch_size, + }) + } + + /// Compute 2D rotary position embeddings for variable-size grids. + /// + /// For each patch position, computes (row_embed, col_embed) based on its + /// 2D coordinates in the image grid. Uses raster order: position i has + /// row = i // width, col = i % width. + fn rot_pos_emb(&self, grid_thw: &Tensor) -> Result { + let device = self.rotary_pos_emb.inv_freq.device(); + let grid = grid_thw.to_vec2::()?; + + // Find max grid dimension to build frequency table + let max_hw = grid + .iter() + .flat_map(|v| v[1..3].iter()) + .copied() + .max() + .unwrap_or(0) as usize; + let freq_table = self.rotary_pos_emb.make_embeds(max_hw)?; + + // Build position indices using simple raster order + // Reference: image_pids = arange(t*h*w) % (h*w) + // h_ids = image_pids // w + // w_ids = image_pids % w + let mut rows = Vec::new(); + let mut cols = Vec::new(); + + for g in &grid { + let t = g[0] as usize; + let h = g[1] as usize; + let w = g[2] as usize; + + // For each temporal frame, patches are in raster order + for _ in 0..t { + for pos in 0..(h * w) { + let row = (pos / w) as i64; + let col = (pos % w) as i64; + rows.push(row); + cols.push(col); + } + } + } + + let total_tokens = rows.len(); + let rows = Tensor::from_vec(rows, (total_tokens,), device)?; + let cols = Tensor::from_vec(cols, (total_tokens,), device)?; + + // Get row and column frequency embeddings + let row_embeds = freq_table.index_select(&rows, 0)?; + let col_embeds = freq_table.index_select(&cols, 0)?; + + // Stack and reshape: (tokens, 2, dim/2) -> (tokens, dim) + Tensor::stack(&[row_embeds, col_embeds], D::Minus2)? + .reshape((total_tokens, freq_table.dim(D::Minus1)? * 2)) + } + + /// Build cumulative sequence lengths for variable-length attention. + fn build_cu_seqlens(&self, grid_thw: &Tensor) -> Result> { + let grid = grid_thw.to_vec2::()?; + let mut cu = Vec::with_capacity(grid.iter().map(|v| v[0] as usize).sum::() + 1); + cu.push(0usize); + let mut acc = 0usize; + for g in &grid { + let area = (g[1] * g[2]) as usize; + for _ in 0..(g[0] as usize) { + acc += area; + cu.push(acc); + } + } + Ok(cu) + } + + /// Forward pass for vision encoder. + /// + /// # Arguments + /// * `pixel_values` - Image tensor of shape (batch, channels, height, width) + /// * `grid_thw` - Grid dimensions tensor of shape (num_images, 3) containing [temporal, height, width] + /// + /// # Returns + /// Projected vision features of shape (total_patches / merge_factor, text_hidden_size) + pub fn forward(&self, pixel_values: &Tensor, grid_thw: &Tensor) -> Result { + self.forward_with_debug(pixel_values, grid_thw, false) + } + + /// Forward pass with optional debug output. + pub fn forward_with_debug( + &self, + pixel_values: &Tensor, + grid_thw: &Tensor, + debug: bool, + ) -> Result { + let dtype = pixel_values.dtype(); + + // Get patch embeddings + let hidden_states = self.embeddings.forward(pixel_values)?; + let hidden_states = hidden_states.reshape(((), self.hidden_size))?; + + if debug { + let hs_f32 = hidden_states.to_dtype(DType::F32)?; + let first_10: Vec = hs_f32.i(0)?.narrow(0, 0, 10)?.to_vec1()?; + eprintln!("DEBUG vision encoder:"); + eprintln!( + " patch_embedding+pos output shape: {:?}", + hidden_states.dims() + ); + eprintln!(" embeddings[0,:10]: {:?}", first_10); + let mean = hs_f32.mean_all()?.to_scalar::()?; + eprintln!(" embeddings mean: {:.6}", mean); + } + + // Compute rotary embeddings + let rotary_pos_emb = self.rot_pos_emb(grid_thw)?; + let seq_len = hidden_states.dim(0)?; + let rotary_pos_emb = rotary_pos_emb.reshape((seq_len, ()))?; + let emb = Tensor::cat(&[&rotary_pos_emb, &rotary_pos_emb], D::Minus1)?; + let cos = emb.cos()?.to_dtype(DType::F32)?; + let sin = emb.sin()?.to_dtype(DType::F32)?; + + let cu_seqlens = self.build_cu_seqlens(grid_thw)?; + + // Pass through encoder layers + let mut hidden_states = hidden_states; + for (i, layer) in self.encoder_layers.iter().enumerate() { + hidden_states = layer.forward(&hidden_states, &cu_seqlens, &cos, &sin)?; + + if debug && (i == 0 || i == 13 || i == 26) { + let hs_f32 = hidden_states.to_dtype(DType::F32)?; + let first_10: Vec = hs_f32.i(0)?.narrow(0, 0, 10)?.to_vec1()?; + let mean = hs_f32.mean_all()?.to_scalar::()?; + eprintln!( + " after layer {}: mean={:.6}, [0,:10]={:?}", + i, mean, first_10 + ); + } + } + + // Apply post layer norm + let hidden_states = self.post_layernorm.forward(&hidden_states)?; + + if debug { + let hs_f32 = hidden_states.to_dtype(DType::F32)?; + let first_10: Vec = hs_f32.i(0)?.narrow(0, 0, 10)?.to_vec1()?; + let mean = hs_f32.mean_all()?.to_scalar::()?; + eprintln!( + " after post_layernorm: mean={:.6}, [0,:10]={:?}", + mean, first_10 + ); + } + + // Project to text model dimension with proper 2×2 spatial merging + let output = self.projector.forward(&hidden_states, grid_thw)?; + + if debug { + let out_f32 = output.to_dtype(DType::F32)?; + let first_10: Vec = out_f32.i(0)?.narrow(0, 0, 10)?.to_vec1()?; + let mean = out_f32.mean_all()?.to_scalar::()?; + eprintln!( + " projector output: shape={:?}, mean={:.6}, [0,:10]={:?}", + output.dims(), + mean, + first_10 + ); + } + + output.to_dtype(dtype) + } + + /// Forward pass for multiple images, returning separate embeddings for each. + /// + /// # Arguments + /// * `pixel_values` - Batched image tensor of shape (num_images, channels, height, width) + /// * `grid_thw` - Grid dimensions tensor of shape (num_images, 3) + /// + /// # Returns + /// Vector of tensors, one per image, each of shape (num_merged_patches, text_hidden_size) + pub fn forward_multi(&self, pixel_values: &Tensor, grid_thw: &Tensor) -> Result> { + let dtype = pixel_values.dtype(); + + // Get patch embeddings + let hidden_states = self.embeddings.forward(pixel_values)?; + let hidden_states = hidden_states.reshape(((), self.hidden_size))?; + + // Compute rotary embeddings + let rotary_pos_emb = self.rot_pos_emb(grid_thw)?; + let seq_len = hidden_states.dim(0)?; + let rotary_pos_emb = rotary_pos_emb.reshape((seq_len, ()))?; + let emb = Tensor::cat(&[&rotary_pos_emb, &rotary_pos_emb], D::Minus1)?; + let cos = emb.cos()?.to_dtype(DType::F32)?; + let sin = emb.sin()?.to_dtype(DType::F32)?; + + let cu_seqlens = self.build_cu_seqlens(grid_thw)?; + + // Pass through encoder layers + let mut hidden_states = hidden_states; + for layer in self.encoder_layers.iter() { + hidden_states = layer.forward(&hidden_states, &cu_seqlens, &cos, &sin)?; + } + + // Apply post layer norm + let hidden_states = self.post_layernorm.forward(&hidden_states)?; + + // Project to text model dimension, returning separate tensors per image + let outputs = self.projector.forward_multi(&hidden_states, grid_thw)?; + + // Convert each output to target dtype + outputs.into_iter().map(|t| t.to_dtype(dtype)).collect() + } + + /// Forward pass with tensor export for substitution testing. + /// + /// Returns a HashMap of checkpoint tensors that can be saved for comparison + /// with the PyTorch reference implementation. + pub fn forward_with_export( + &self, + pixel_values: &Tensor, + grid_thw: &Tensor, + ) -> Result<(Tensor, HashMap)> { + let dtype = pixel_values.dtype(); + let mut exports: HashMap = HashMap::new(); + + // Export patchified pixel values to match PyTorch format: (num_patches, 3, 14, 14) + // Input is (batch, channels, height, width), output is (num_patches, channels, patch, patch) + let (batch, channels, height, width) = pixel_values.dims4()?; + let h_patches = height / self.patch_size; + let w_patches = width / self.patch_size; + let patchified = pixel_values + .reshape(( + batch, + channels, + h_patches, + self.patch_size, + w_patches, + self.patch_size, + ))? + .permute((0, 2, 4, 1, 3, 5))? // (batch, h_patches, w_patches, channels, patch_size, patch_size) + .reshape(( + h_patches * w_patches, + channels, + self.patch_size, + self.patch_size, + ))?; + exports.insert("pixel_values".to_string(), patchified.to_dtype(DType::F32)?); + + // 1. Patch embedding (before position embedding) + let patch_out = self.embeddings.patch_embedding.forward(pixel_values)?; + let (batch, hidden, h, w) = patch_out.dims4()?; + let num_patches = h * w; + let patch_out = patch_out + .reshape((batch, hidden, num_patches))? + .transpose(1, 2)?; + exports.insert( + "patch_embedding_output".to_string(), + patch_out.to_dtype(DType::F32)?, + ); + + // 2. Add position embedding (use interpolated 2D position embeddings, same as forward()) + // NOTE: The packing_position_embedding is a fallback; we must use interpolate_pos_encoding + // to match the regular forward path which uses bilinear interpolation of the 27×27 base grid. + let pos_embed = self.embeddings.interpolate_pos_encoding(h, w)?; + let hidden_states = patch_out.broadcast_add(&pos_embed)?; + let hidden_states = hidden_states.reshape(((), self.hidden_size))?; + exports.insert( + "embeddings_output".to_string(), + hidden_states.to_dtype(DType::F32)?, + ); + + // Compute rotary embeddings + let rotary_pos_emb = self.rot_pos_emb(grid_thw)?; + let seq_len = hidden_states.dim(0)?; + let rotary_pos_emb = rotary_pos_emb.reshape((seq_len, ()))?; + let emb = Tensor::cat(&[&rotary_pos_emb, &rotary_pos_emb], D::Minus1)?; + let cos = emb.cos()?.to_dtype(DType::F32)?; + let sin = emb.sin()?.to_dtype(DType::F32)?; + + let cu_seqlens = self.build_cu_seqlens(grid_thw)?; + + // Export RoPE embeddings for comparison + exports.insert("rope_pos_emb_raw".to_string(), rotary_pos_emb.clone()); + + // Pass through encoder layers with checkpoints + // Layer 0 gets detailed debug export + let mut hidden_states = hidden_states; + for (i, layer) in self.encoder_layers.iter().enumerate() { + if i == 0 { + // Use debug forward for layer 0 to capture attention internals + hidden_states = layer.forward_with_debug( + &hidden_states, + &cu_seqlens, + &cos, + &sin, + &mut exports, + )?; + exports.insert( + "layer_0_output".to_string(), + hidden_states.to_dtype(DType::F32)?, + ); + } else { + hidden_states = layer.forward(&hidden_states, &cu_seqlens, &cos, &sin)?; + if i == 13 || i == 26 { + exports.insert( + format!("layer_{}_output", i), + hidden_states.to_dtype(DType::F32)?, + ); + } + } + } + + // Apply post layer norm + let hidden_states = self.post_layernorm.forward(&hidden_states)?; + exports.insert( + "post_layernorm_output".to_string(), + hidden_states.to_dtype(DType::F32)?, + ); + + // Project to text model dimension + let output = self.projector.forward(&hidden_states, grid_thw)?; + exports.insert("projector_output".to_string(), output.to_dtype(DType::F32)?); + + Ok((output.to_dtype(dtype)?, exports)) + } +} diff --git a/patches/candle-transformers/src/models/paligemma.rs b/patches/candle-transformers/src/models/paligemma.rs new file mode 100644 index 0000000000..e992869923 --- /dev/null +++ b/patches/candle-transformers/src/models/paligemma.rs @@ -0,0 +1,170 @@ +//! Multimodal multi-purpose model combining Gemma-based language model with SigLIP image understanding +//! +//! See PaLiGemma details at: +//! - [Paper](https://arxiv.org/abs/2402.05257) +//! - [Google Blog Post](https://blog.research.google/2024/02/paligemma-scaling-language-image.html) +//! +//! The model is a multimodal combination of: +//! - SigLIP vision encoder +//! - Gemma language model +//! - Cross-projection layers +//! +//! References: +//! - [HuggingFace Implementation](https://huggingface.co/google/paligemma-3b) +//! - [Paper: PaLI-3 and Beyond: Scaling Language-Image Learning](https://arxiv.org/abs/2402.05257) +//! + +use crate::models::{gemma, siglip}; +use candle::{Module, Result, Tensor}; +use candle_nn::{linear, Linear, VarBuilder}; + +#[derive(serde::Deserialize, Clone, Debug)] +pub struct Config { + pub vision_config: siglip::VisionConfig, + pub text_config: gemma::Config, + pub projection_dim: usize, +} + +impl Config { + pub fn paligemma_3b_224() -> Self { + // https://huggingface.co/google/paligemma-3b-pt-224/blob/main/config.json + Self { + vision_config: siglip::VisionConfig::paligemma_3b_224(), + text_config: gemma::Config { + hidden_size: 2048, + intermediate_size: 16384, + num_attention_heads: 8, + num_hidden_layers: 18, + num_key_value_heads: 1, + vocab_size: 257216, + // Default values. + rope_theta: 10000., + head_dim: 256, + hidden_act: Some(candle_nn::Activation::GeluPytorchTanh), + hidden_activation: None, + attention_bias: false, + max_position_embeddings: 8192, + rms_norm_eps: 1e-6, + }, + projection_dim: 2048, + } + } + + pub fn paligemma_3b_448() -> Self { + Self { + vision_config: siglip::VisionConfig::paligemma_3b_448(), + text_config: gemma::Config { + hidden_size: 2048, + intermediate_size: 16384, + num_attention_heads: 8, + num_hidden_layers: 18, + num_key_value_heads: 1, + // Default values. + rope_theta: 10000., + head_dim: 256, + hidden_act: Some(candle_nn::Activation::GeluPytorchTanh), + hidden_activation: None, + attention_bias: false, + max_position_embeddings: 8192, + rms_norm_eps: 1e-6, + vocab_size: 257216, + }, + projection_dim: 2048, + } + } +} + +#[derive(Clone, Debug)] +pub struct MultiModalProjector { + linear: Linear, +} + +impl MultiModalProjector { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let linear = linear( + cfg.vision_config.hidden_size, + cfg.projection_dim, + vb.pp("linear"), + )?; + Ok(Self { linear }) + } +} + +impl Module for MultiModalProjector { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.linear) + } +} + +#[derive(Clone, Debug)] +pub struct Model { + pos: usize, + vision_tower: siglip::VisionModel, + multi_modal_projector: MultiModalProjector, + language_model: gemma::Model, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vision_tower = siglip::VisionModel::new( + &cfg.vision_config, + false, + vb.pp("vision_tower.vision_model"), + )?; + let multi_modal_projector = MultiModalProjector::new(cfg, vb.pp("multi_modal_projector"))?; + let language_model = gemma::Model::new(false, &cfg.text_config, vb.pp("language_model"))?; + Ok(Self { + pos: 0, + language_model, + vision_tower, + multi_modal_projector, + }) + } + + pub fn setup(&mut self, pixel_values: &Tensor, input_ids: &Tensor) -> Result { + self.clear_kv_cache(); + let image_features = self + .vision_tower + .forward(pixel_values)? + .apply(&self.multi_modal_projector)?; + let image_features = crate::models::clip::div_l2_norm(&image_features)?; + let text_features = self.language_model.embed_tokens().forward(input_ids)?; + let input_embeds = Tensor::cat(&[image_features, text_features], 1)?; + self.pos = input_embeds.dim(1)?; + self.language_model.forward_embeds(&input_embeds, None, 0) + } + + pub fn forward(&mut self, input_ids: &Tensor) -> Result { + let pos = self.pos; + let seq_len = input_ids.dim(1)?; + self.pos = pos + seq_len; + self.language_model.forward(input_ids, pos) + } + + pub fn forward_without_projection(&mut self, input_ids: &Tensor) -> Result { + self.clear_kv_cache(); + let input_embeds = self.language_model.embed_tokens().forward(input_ids)?; + self.language_model + .forward_embeds_without_projection(&input_embeds, None, 0) + } + pub fn setup_without_projection( + &mut self, + pixel_values: &Tensor, + input_ids: &Tensor, + ) -> Result { + self.clear_kv_cache(); + let image_features = self + .vision_tower + .forward(pixel_values)? + .apply(&self.multi_modal_projector)?; + let image_features = crate::models::clip::div_l2_norm(&image_features)?; + let text_features = self.language_model.embed_tokens().forward(input_ids)?; + let input_embeds = Tensor::cat(&[image_features, text_features], 1)?; + self.language_model + .forward_embeds_without_projection(&input_embeds, None, 0) + } + pub fn clear_kv_cache(&mut self) { + self.pos = 0; + self.language_model.clear_kv_cache() + } +} diff --git a/patches/candle-transformers/src/models/parler_tts.rs b/patches/candle-transformers/src/models/parler_tts.rs new file mode 100644 index 0000000000..b514ee0b28 --- /dev/null +++ b/patches/candle-transformers/src/models/parler_tts.rs @@ -0,0 +1,473 @@ +//! Parler Model implementation for parler_tts text-to-speech synthesis +//! +//! Implements a transformer-based decoder architecture for generating audio tokens +//! from text using discrete tokens. The model converts text into audio segments +//! using multiple codebooks of quantized audio tokens. +//! +//! The model architecture includes: +//! - Multi-head attention layers for text and audio processing +//! - Feed-forward networks +//! - Layer normalization +//! - Positional embeddings +//! - Multiple codebook prediction heads +//! +//! The implementation follows the original parler_tts architecture while focusing +//! on audio token generation for text-to-speech synthesis. +//! + +use crate::generation::LogitsProcessor; +use crate::models::t5; +use candle::{IndexOp, Result, Tensor}; +use candle_nn::{layer_norm, linear_b as linear, Activation, LayerNorm, Linear, VarBuilder}; + +#[derive(serde::Deserialize, Debug, Clone)] +pub struct DecoderConfig { + pub vocab_size: usize, + pub max_position_embeddings: usize, + pub num_hidden_layers: usize, + pub ffn_dim: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: Option, + pub num_cross_attention_key_value_heads: Option, + pub activation_function: Activation, + pub hidden_size: usize, + pub scale_embedding: bool, + pub num_codebooks: usize, + pub pad_token_id: usize, + pub bos_token_id: usize, + pub eos_token_id: usize, + pub tie_word_embeddings: bool, + pub rope_embeddings: bool, + pub rope_theta: f64, +} + +#[derive(serde::Deserialize, Debug, Clone)] +pub struct Config { + pub decoder_start_token_id: u32, + pub pad_token_id: u32, + pub decoder: DecoderConfig, + pub text_encoder: t5::Config, + pub vocab_size: usize, + pub audio_encoder: crate::models::dac::Config, +} + +#[derive(Debug, Clone)] +pub struct Attention { + k_proj: Linear, + v_proj: Linear, + q_proj: Linear, + out_proj: Linear, + is_causal: bool, + kv_cache: Option<(Tensor, Tensor)>, + scaling: f64, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, +} + +impl Attention { + fn new( + num_kv_heads: usize, + is_causal: bool, + cfg: &DecoderConfig, + vb: VarBuilder, + ) -> Result { + if cfg.rope_embeddings { + candle::bail!("rope embeddings are not supported"); + } + let embed_dim = cfg.hidden_size; + let head_dim = embed_dim / cfg.num_attention_heads; + let kv_out_dim = num_kv_heads * head_dim; + let k_proj = linear(embed_dim, kv_out_dim, false, vb.pp("k_proj"))?; + let v_proj = linear(embed_dim, kv_out_dim, false, vb.pp("v_proj"))?; + let q_proj = linear(embed_dim, embed_dim, false, vb.pp("q_proj"))?; + let out_proj = linear(embed_dim, embed_dim, false, vb.pp("out_proj"))?; + Ok(Self { + k_proj, + v_proj, + q_proj, + out_proj, + is_causal, + kv_cache: None, + scaling: (head_dim as f64).powf(-0.5), + num_heads: cfg.num_attention_heads, + num_kv_heads, + num_kv_groups: cfg.num_attention_heads / num_kv_heads, + head_dim, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + key_value_states: Option<&Tensor>, + attention_mask: Option<&Tensor>, + ) -> Result { + let (b_sz, tgt_len, _) = xs.dims3()?; + let query_states = (xs.apply(&self.q_proj)? * self.scaling)? + .reshape((b_sz, tgt_len, self.num_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let key_states = match key_value_states { + Some(states) => states.apply(&self.k_proj)?, + None => xs.apply(&self.k_proj)?, + }; + let key_states = key_states + .reshape((b_sz, (), self.num_kv_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let value_states = match key_value_states { + Some(states) => states.apply(&self.v_proj)?, + None => xs.apply(&self.v_proj)?, + }; + let value_states = value_states + .reshape((b_sz, (), self.num_kv_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + if self.is_causal { + self.kv_cache = Some((key_states.clone(), value_states.clone())); + } + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?.contiguous()?; + let value_states = + crate::utils::repeat_kv(value_states, self.num_kv_groups)?.contiguous()?; + + let attn_weights = query_states.matmul(&key_states.transpose(2, 3)?)?; + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + let attn_output = attn_weights.matmul(&value_states)?; + attn_output + .transpose(1, 2)? + .reshape((b_sz, tgt_len, ()))? + .apply(&self.out_proj) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +pub struct DecoderLayer { + self_attn: Attention, + self_attn_layer_norm: LayerNorm, + encoder_attn: Attention, + encoder_attn_layer_norm: LayerNorm, + fc1: Linear, + fc2: Linear, + final_layer_norm: LayerNorm, + activation: Activation, +} + +impl DecoderLayer { + fn new(cfg: &DecoderConfig, vb: VarBuilder) -> Result { + let kv_heads = cfg.num_key_value_heads.unwrap_or(cfg.num_attention_heads); + let kv_heads_cross = cfg.num_cross_attention_key_value_heads.unwrap_or(kv_heads); + + let self_attn = Attention::new(kv_heads, true, cfg, vb.pp("self_attn"))?; + let encoder_attn = Attention::new(kv_heads_cross, false, cfg, vb.pp("encoder_attn"))?; + let self_attn_layer_norm = + layer_norm(cfg.hidden_size, 1e-5, vb.pp("self_attn_layer_norm"))?; + let encoder_attn_layer_norm = + layer_norm(cfg.hidden_size, 1e-5, vb.pp("encoder_attn_layer_norm"))?; + let fc1 = linear(cfg.hidden_size, cfg.ffn_dim, false, vb.pp("fc1"))?; + let fc2 = linear(cfg.ffn_dim, cfg.hidden_size, false, vb.pp("fc2"))?; + let final_layer_norm = layer_norm(cfg.hidden_size, 1e-5, vb.pp("final_layer_norm"))?; + Ok(Self { + self_attn, + self_attn_layer_norm, + encoder_attn, + encoder_attn_layer_norm, + fc1, + fc2, + final_layer_norm, + activation: cfg.activation_function, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + encoder_xs: &Tensor, + encoder_attention_mask: Option<&Tensor>, + ) -> Result { + // Self attention + let residual = xs; + let xs = xs.apply(&self.self_attn_layer_norm)?; + let xs = self.self_attn.forward(&xs, None, attention_mask)?; + let xs = (residual + xs)?; + + // Cross attention + let residual = &xs; + let xs = xs.apply(&self.encoder_attn_layer_norm)?; + let xs = self + .encoder_attn + .forward(&xs, Some(encoder_xs), encoder_attention_mask)?; + let xs = (residual + xs)?; + + // Fully connected + let residual = &xs; + let xs = xs + .apply(&self.final_layer_norm)? + .apply(&self.fc1)? + .apply(&self.activation)? + .apply(&self.fc2)?; + residual + xs + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache(); + self.encoder_attn.clear_kv_cache(); + } +} + +#[derive(Debug, Clone)] +pub struct Decoder { + embed_tokens: Vec, + embed_positions: Tensor, + layers: Vec, + layer_norm: LayerNorm, + num_codebooks: usize, + hidden_size: usize, + lm_heads: Vec, + dtype: candle::DType, +} + +impl Decoder { + pub fn new(cfg: &DecoderConfig, vb: VarBuilder) -> Result { + let vb_d = vb.pp("model.decoder"); + let mut embed_tokens = Vec::with_capacity(cfg.num_codebooks); + let vb_e = vb_d.pp("embed_tokens"); + for embed_idx in 0..cfg.num_codebooks { + let e = candle_nn::embedding(cfg.vocab_size + 1, cfg.hidden_size, vb_e.pp(embed_idx))?; + embed_tokens.push(e) + } + let embed_positions = vb_d.get( + (cfg.max_position_embeddings, cfg.hidden_size), + "embed_positions.weights", + )?; + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_d.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let layer_norm = layer_norm(cfg.hidden_size, 1e-5, vb_d.pp("layer_norm"))?; + + let mut lm_heads = Vec::with_capacity(cfg.num_codebooks); + let vb_l = vb.pp("lm_heads"); + for lm_idx in 0..cfg.num_codebooks { + let lm_head = linear(cfg.hidden_size, cfg.vocab_size, false, vb_l.pp(lm_idx))?; + lm_heads.push(lm_head) + } + Ok(Self { + embed_tokens, + embed_positions, + layers, + layer_norm, + num_codebooks: cfg.num_codebooks, + lm_heads, + hidden_size: cfg.hidden_size, + dtype: vb.dtype(), + }) + } + + pub fn forward( + &mut self, + input_ids: &Tensor, + prompt_hidden_states: Option<&Tensor>, + attention_mask: Option<&Tensor>, + encoder_xs: &Tensor, + encoder_attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result> { + let (b_sz, num_codebooks, seq_len) = input_ids.dims3()?; + if num_codebooks != self.num_codebooks { + candle::bail!("unexpected num codebooks in input {:?}", input_ids.shape()) + } + let mut inputs_embeds = Tensor::zeros( + (b_sz, seq_len, self.hidden_size), + self.dtype, + input_ids.device(), + )?; + for (idx, embs) in self.embed_tokens.iter().enumerate() { + let e = input_ids.i((.., idx))?.apply(embs)?; + inputs_embeds = (inputs_embeds + e)? + } + let inputs_embeds = match prompt_hidden_states { + None => inputs_embeds, + Some(pis) => Tensor::cat(&[pis, &inputs_embeds], 1)?, + }; + let embed_positions = self + .embed_positions + .i(seqlen_offset..seqlen_offset + inputs_embeds.dim(1)?)?; + let mut xs = (inputs_embeds + embed_positions.unsqueeze(0))?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask, encoder_xs, encoder_attention_mask)?; + } + let xs = xs.apply(&self.layer_norm)?; + let mut lm_logits = Vec::with_capacity(self.num_codebooks); + for lm_head in self.lm_heads.iter() { + let logits = xs.apply(lm_head)?; + lm_logits.push(logits) + } + Ok(lm_logits) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +} + +#[derive(Debug, Clone)] +pub struct Model { + pub embed_prompts: candle_nn::Embedding, + pub enc_to_dec_proj: Option, + pub decoder: Decoder, + pub text_encoder: t5::T5EncoderModel, + pub decoder_start_token_id: u32, + pub pad_token_id: u32, + pub audio_encoder: crate::models::dac::Model, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let text_encoder = t5::T5EncoderModel::load(vb.pp("text_encoder"), &cfg.text_encoder)?; + let decoder = Decoder::new(&cfg.decoder, vb.pp("decoder"))?; + let embed_prompts = candle_nn::embedding( + cfg.vocab_size, + cfg.decoder.hidden_size, + vb.pp("embed_prompts"), + )?; + let enc_to_dec_proj = if cfg.text_encoder.d_model != cfg.decoder.hidden_size { + let proj = linear( + cfg.text_encoder.d_model, + cfg.decoder.hidden_size, + true, + vb.pp("enc_to_dec_proj"), + )?; + Some(proj) + } else { + None + }; + let audio_encoder = + crate::models::dac::Model::new(&cfg.audio_encoder, vb.pp("audio_encoder.model"))?; + Ok(Self { + decoder, + text_encoder, + embed_prompts, + enc_to_dec_proj, + decoder_start_token_id: cfg.decoder_start_token_id, + pad_token_id: cfg.pad_token_id, + audio_encoder, + }) + } + + /// Note that the returned tensor uses the CPU device. + pub fn generate( + &mut self, + prompt_tokens: &Tensor, + description_tokens: &Tensor, + mut lp: LogitsProcessor, + max_steps: usize, + ) -> Result { + self.decoder.clear_kv_cache(); + self.text_encoder.clear_kv_cache(); + let encoded = self.text_encoder.forward(description_tokens)?; + let encoded = match self.enc_to_dec_proj.as_ref() { + None => encoded, + Some(proj) => encoded.apply(proj)?, + }; + let prompt_hidden_states = prompt_tokens.apply(&self.embed_prompts)?; + let num_codebooks = self.decoder.num_codebooks; + let mut audio_tokens = vec![self.decoder_start_token_id; num_codebooks]; + let mut all_audio_tokens = vec![vec![]; num_codebooks]; + let prompt_len = prompt_hidden_states.dim(1)?; + for step in 0..max_steps { + let input_ids = Tensor::from_slice( + audio_tokens.as_slice(), + (1, num_codebooks, 1), + prompt_tokens.device(), + )?; + let (prompt_hidden_states, pos) = if step == 0 { + (Some(&prompt_hidden_states), 0) + } else { + (None, step + prompt_len) + }; + let causal_mask = if pos == 0 { + self.prepare_causal_mask(prompt_len + 1, prompt_len + 1, input_ids.device())? + } else { + self.prepare_causal_mask(1, pos + 1, input_ids.device())? + }; + let logits = self.decoder.forward( + &input_ids, + prompt_hidden_states, + Some(&causal_mask), + &encoded, + None, + pos, + )?; + for (logit_idx, logit) in logits.iter().enumerate() { + if logit_idx > step { + break; + } + if audio_tokens[logit_idx] != self.pad_token_id { + let logit = logit.i((0, logit.dim(1)? - 1))?; + let token = lp.sample(&logit)?; + audio_tokens[logit_idx] = token + } + } + if audio_tokens.iter().all(|v| v == &self.pad_token_id) { + break; + } + for (cb_idx, &token) in audio_tokens.iter().enumerate() { + if token != self.decoder_start_token_id && token != self.pad_token_id { + all_audio_tokens[cb_idx].push(token) + } + } + } + + let min_len = all_audio_tokens.iter().map(|v| v.len()).min().unwrap_or(0); + all_audio_tokens.iter_mut().for_each(|v| { + v.resize(min_len, 0); + }); + let all_audio_tokens = Tensor::new(all_audio_tokens, &candle::Device::Cpu)?; + Ok(all_audio_tokens) + } + + fn prepare_causal_mask( + &self, + q_len: usize, + kv_len: usize, + device: &candle::Device, + ) -> Result { + let mask: Vec<_> = (0..q_len) + .flat_map(|i| { + (0..kv_len).map(move |j| { + if i + kv_len < j + q_len { + f32::NEG_INFINITY + } else { + 0. + } + }) + }) + .collect(); + Tensor::from_slice(&mask, (q_len, kv_len), device) + } +} diff --git a/patches/candle-transformers/src/models/persimmon.rs b/patches/candle-transformers/src/models/persimmon.rs new file mode 100644 index 0000000000..d1e3db316f --- /dev/null +++ b/patches/candle-transformers/src/models/persimmon.rs @@ -0,0 +1,70 @@ +//! Persimmon Model +//! +//! A transformer language model for efficient inference and general-purpose tasks. The model uses a standard transformer architecture with: +//! - Layer normalization for Q/K attention +//! - RoPE embeddings with partial rotary factor +//! - ReLU activation +//! - Separate number of attention heads and KV heads +//! +//! References: +//! - 💻 [Hugging Face Implementation](https://github.com/huggingface/transformers/blob/main/src/transformers/models/persimmon/modeling_persimmon.py) +//! - 💻 [Persimmon Config](https://github.com/huggingface/transformers/blob/main/src/transformers/models/persimmon/configuration_persimmon.py) +//! - 🤗 [Hugging Face](https://huggingface.co/adept/persimmon-8b-base) +//! + +use candle::DType; +use serde::Deserialize; + +pub const DTYPE: DType = DType::F32; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum PositionEmbeddingType { + Absolute, + Alibi, +} + +// https://github.com/huggingface/transformers/blob/main/src/transformers/models/persimmon/configuration_persimmon.py +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub hidden_act: candle_nn::Activation, + pub max_position_embeddings: usize, + pub initializer_range: f64, + pub layer_norm_eps: f64, + pub rms_norm_eps: f64, + pub use_cache: bool, + pub tie_word_embeddings: bool, + pub rope_theta: f64, + pub qk_layernorm: bool, + pub partial_rotary_factor: f64, +} + +impl Config { + pub fn base_8b() -> Self { + // https://huggingface.co/adept/persimmon-8b-base/blob/main/config.json + Self { + hidden_act: candle_nn::Activation::Relu, + hidden_size: 4096, + initializer_range: 0.02, + intermediate_size: 16384, + layer_norm_eps: 1e-05, + max_position_embeddings: 16384, + num_attention_heads: 64, + num_hidden_layers: 36, + num_key_value_heads: 64, + qk_layernorm: true, + rms_norm_eps: 1e-06, + rope_theta: 25000.0, + tie_word_embeddings: false, + use_cache: true, + vocab_size: 262144, + partial_rotary_factor: 0.5, + } + } +} diff --git a/patches/candle-transformers/src/models/phi.rs b/patches/candle-transformers/src/models/phi.rs new file mode 100644 index 0000000000..c94ef6686b --- /dev/null +++ b/patches/candle-transformers/src/models/phi.rs @@ -0,0 +1,365 @@ +//! Microsoft Phi model implementation +//! +//! The Phi series are decoder-only transformers designed for code and language tasks. +//! +//! Key characteristics: +//! - Decoder-only transformer architecture +//! - RoPE embeddings +//! - Layer normalization +//! - QK normalization +//! +//! - ⚡ [Interactive Wasm Example](https://huggingface.co/spaces/radames/Candle-phi1-phi2-wasm-demo) +//! - 🤗 [HF Link](https://huggingface.co/microsoft/phi-2) +//! + +use crate::models::with_tracing::{layer_norm, linear, Embedding, LayerNorm, Linear}; +/// Phi model. +/// https://huggingface.co/microsoft/phi-2 +/// There is an alternative implementation of the phi model in mixformers.rs. +/// This corresponds to the model update made with the following commit: +/// https://huggingface.co/microsoft/phi-2/commit/cb2f4533604d8b67de604e7df03bfe6f3ca22869 +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{Activation, VarBuilder}; +use serde::Deserialize; + +// https://huggingface.co/microsoft/phi-2/blob/main/configuration_phi.py +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Config { + pub(crate) vocab_size: usize, + pub(crate) hidden_size: usize, + pub(crate) intermediate_size: usize, + pub(crate) num_hidden_layers: usize, + pub(crate) num_attention_heads: usize, + pub(crate) num_key_value_heads: Option, + pub(crate) hidden_act: Activation, + pub(crate) max_position_embeddings: usize, + pub(crate) layer_norm_eps: f64, + pub(crate) tie_word_embeddings: bool, + pub(crate) rope_theta: f32, + pub(crate) partial_rotary_factor: f64, + pub(crate) qk_layernorm: bool, +} + +impl Config { + fn num_key_value_heads(&self) -> usize { + self.num_key_value_heads.unwrap_or(self.num_attention_heads) + } + + fn head_dim(&self) -> usize { + self.hidden_size / self.num_attention_heads + } +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + dim: usize, + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(cfg: &Config, dev: &Device) -> Result { + let dim = (cfg.partial_rotary_factor * cfg.head_dim() as f64) as usize; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f32 / dim as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?; + let t = Tensor::arange(0u32, cfg.max_position_embeddings as u32, dev)? + .to_dtype(DType::F32)? + .reshape((cfg.max_position_embeddings, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + dim, + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb(&self, xs: &Tensor, seqlen_offset: usize) -> Result { + let (_b_size, _num_heads, seq_len, _headdim) = xs.dims4()?; + let xs_rot = xs.i((.., .., .., ..self.dim))?.contiguous()?; + let xs_pass = xs.i((.., .., .., self.dim..))?; + let c = self.cos.narrow(0, seqlen_offset, seq_len)?; + let s = self.sin.narrow(0, seqlen_offset, seq_len)?; + let xs_rot = candle_nn::rotary_emb::rope(&xs_rot, &c, &s)?; + Tensor::cat(&[&xs_rot, &xs_pass], D::Minus1) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + fc1: Linear, + fc2: Linear, + act: Activation, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let fc1 = linear(cfg.hidden_size, cfg.intermediate_size, vb.pp("fc1"))?; + let fc2 = linear(cfg.intermediate_size, cfg.hidden_size, vb.pp("fc2"))?; + Ok(Self { + fc1, + fc2, + // This does not match the mixformers implementation where Gelu is used rather than + // GeluNew. + act: cfg.hidden_act, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.fc1)?.apply(&self.act)?.apply(&self.fc2) + } +} + +#[derive(Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + dense: Linear, + kv_cache: Option<(Tensor, Tensor)>, + q_layernorm: Option, + k_layernorm: Option, + rotary_emb: RotaryEmbedding, + softmax_scale: f64, + num_heads: usize, + num_kv_heads: usize, + head_dim: usize, + span: tracing::Span, +} + +fn get_mask(size: usize, device: &Device) -> Result { + let mask: Vec<_> = (0..size) + .flat_map(|i| (0..size).map(move |j| u8::from(j > i))) + .collect(); + Tensor::from_slice(&mask, (size, size), device) +} + +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 Attention { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads(); + let head_dim = cfg.head_dim(); + let q_proj = linear(cfg.hidden_size, num_heads * head_dim, vb.pp("q_proj"))?; + let k_proj = linear(cfg.hidden_size, num_kv_heads * head_dim, vb.pp("k_proj"))?; + let v_proj = linear(cfg.hidden_size, num_kv_heads * head_dim, vb.pp("v_proj"))?; + let dense = linear(num_heads * head_dim, cfg.hidden_size, vb.pp("dense"))?; + // Alternative rope scalings are not supported. + let rotary_emb = RotaryEmbedding::new(cfg, vb.device())?; + let (q_layernorm, k_layernorm) = if cfg.qk_layernorm { + let q_layernorm = layer_norm(head_dim, cfg.layer_norm_eps, vb.pp("q_layernorm"))?; + let k_layernorm = layer_norm(head_dim, cfg.layer_norm_eps, vb.pp("k_layernorm"))?; + (Some(q_layernorm), Some(k_layernorm)) + } else { + (None, None) + }; + let softmax_scale = 1f64 / (head_dim as f64).sqrt(); + Ok(Self { + q_proj, + k_proj, + v_proj, + dense, + kv_cache: None, + q_layernorm, + k_layernorm, + rotary_emb, + softmax_scale, + num_heads, + num_kv_heads, + head_dim, + span: tracing::span!(tracing::Level::TRACE, "attention"), + }) + } + + fn repeat_kv(&self, xs: Tensor) -> Result { + crate::utils::repeat_kv(xs, self.num_heads / self.num_kv_heads) + } + + fn forward(&mut self, xs: &Tensor, mask: Option<&Tensor>) -> Result { + let _enter = self.span.enter(); + let (b_size, seq_len, _n_embd) = xs.dims3()?; + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = match &self.q_layernorm { + None => query_states, + Some(ln) => query_states.apply(ln)?, + }; + let key_states = match &self.k_layernorm { + None => key_states, + Some(ln) => key_states.apply(ln)?, + }; + + let query_states = query_states + .reshape((b_size, seq_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_size, seq_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_size, seq_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + // Rotary embeddings. + let seqlen_offset = match &self.kv_cache { + None => 0, + Some((prev_k, _)) => prev_k.dim(2)?, + }; + let query_states = self + .rotary_emb + .apply_rotary_emb(&query_states, seqlen_offset)?; + let key_states = self + .rotary_emb + .apply_rotary_emb(&key_states, seqlen_offset)?; + + // KV cache. + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let k = Tensor::cat(&[prev_k, &key_states], 2)?; + let v = Tensor::cat(&[prev_v, &value_states], 2)?; + (k, v) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + // Repeat kv. + let key_states = self.repeat_kv(key_states)?.contiguous()?; + let value_states = self.repeat_kv(value_states)?.contiguous()?; + + let attn_weights = (query_states + .to_dtype(DType::F32)? + .contiguous()? + .matmul(&key_states.to_dtype(DType::F32)?.t()?)? + * self.softmax_scale)?; + let attn_weights = match mask { + None => attn_weights, + Some(mask) => masked_fill( + &attn_weights, + &mask.broadcast_left((b_size, self.num_heads))?, + f32::NEG_INFINITY, + )?, + }; + let attn_weights = + candle_nn::ops::softmax_last_dim(&attn_weights)?.to_dtype(value_states.dtype())?; + let attn_output = attn_weights.matmul(&value_states)?; + let attn_output = attn_output + .transpose(1, 2)? + .reshape((b_size, seq_len, ()))?; + attn_output.apply(&self.dense) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: MLP, + input_layernorm: LayerNorm, + span: tracing::Span, +} + +impl DecoderLayer { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let self_attn = Attention::new(cfg, vb.pp("self_attn"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let input_layernorm = layer_norm( + cfg.hidden_size, + cfg.layer_norm_eps, + vb.pp("input_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + span: tracing::span!(tracing::Level::TRACE, "block"), + }) + } + + fn forward(&mut self, xs: &Tensor, mask: Option<&Tensor>) -> Result { + let _enter = self.span.enter(); + let residual = xs; + let xs = xs.apply(&self.input_layernorm)?; + let attn_outputs = self.self_attn.forward(&xs, mask)?; + let feed_forward_hidden_states = self.mlp.forward(&xs)?; + attn_outputs + feed_forward_hidden_states + residual + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } +} + +#[derive(Clone)] +pub struct Model { + embed_tokens: Embedding, + layers: Vec, + final_layernorm: LayerNorm, + lm_head: Linear, + span: tracing::Span, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + Embedding::new(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let final_layernorm = layer_norm( + cfg.hidden_size, + cfg.layer_norm_eps, + vb_m.pp("final_layernorm"), + )?; + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_m = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(cfg, vb_m.pp(layer_idx))?; + layers.push(layer) + } + let lm_head = linear(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))?; + Ok(Self { + embed_tokens, + layers, + final_layernorm, + lm_head, + span: tracing::span!(tracing::Level::TRACE, "model"), + }) + } + + pub fn forward(&mut self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let (_b_size, seq_len) = xs.dims2()?; + let mut xs = xs.apply(&self.embed_tokens)?; + let mask = if seq_len <= 1 { + None + } else { + Some(get_mask(seq_len, xs.device())?) + }; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, mask.as_ref())?; + } + xs.apply(&self.final_layernorm)? + .narrow(1, seq_len - 1, 1)? + .apply(&self.lm_head)? + .squeeze(1) + } + + pub fn clear_kv_cache(&mut self) { + self.layers.iter_mut().for_each(|b| b.clear_kv_cache()) + } +} diff --git a/patches/candle-transformers/src/models/phi3.rs b/patches/candle-transformers/src/models/phi3.rs new file mode 100644 index 0000000000..6535d9a4fd --- /dev/null +++ b/patches/candle-transformers/src/models/phi3.rs @@ -0,0 +1,428 @@ +//! Microsoft Phi-3 model implementation +//! +//! See Phi model details at: +//! - [Phi-3 Model](https://huggingface.co/microsoft/phi-3) +//! +//! The Phi series are decoder-only transformers designed for code and language tasks. +//! Key characteristics: +//! - Decoder-only transformer architecture +//! - RoPE embeddings +//! - Layer normalization +//! - QK normalization +//! - Mixed activation functions +//! - Improved context window handling +//! +//! References: +//! - [Hugging Face Implementation](https://huggingface.co/microsoft/phi-3) +//! - [Alternative Implementation](https://huggingface.co/microsoft/phi-3/tree/main) +//! + +// This implementation is based on: +// https://huggingface.co/microsoft/Phi-3-mini-4k-instruct/blob/main/modeling_phi3.py +use crate::models::with_tracing::{linear_no_bias as linear, Linear, RmsNorm}; +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::VarBuilder; +use std::sync::Arc; + +#[derive(Debug, Clone, serde::Deserialize)] +pub enum RopeScalingType { + #[serde(rename = "longrope")] + LongRope, +} + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct RopeScaling { + pub short_factor: Vec, + pub long_factor: Vec, + #[serde(rename = "type")] + pub type_: RopeScalingType, +} + +// https://huggingface.co/microsoft/Phi-3-mini-4k-instruct/blob/main/config.json +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_act: candle_nn::Activation, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub rms_norm_eps: f64, + pub rope_theta: f64, + pub bos_token_id: Option, + pub eos_token_id: Option, + pub rope_scaling: Option, + pub max_position_embeddings: usize, + pub original_max_position_embeddings: Option, + pub partial_rotary_factor: Option, + #[serde(default)] + pub tie_word_embeddings: bool, +} + +impl Config { + pub fn head_dim(&self) -> usize { + self.hidden_size / self.num_attention_heads + } +} + +#[derive(Debug, Clone)] +pub struct RotaryEmbedding { + partial_dim: Option, + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + pub fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let partial_dim = cfg + .partial_rotary_factor + .as_ref() + .map(|v| (v * cfg.head_dim() as f64) as usize); + let dim = partial_dim.unwrap_or(cfg.head_dim()); + let freqs = match cfg.rope_scaling.as_ref() { + None => { + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq = Tensor::from_vec(inv_freq, (1, ()), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + t.matmul(&inv_freq)? + } + Some(rope_scaling) => { + let inv_freq_s: Vec<_> = (0..dim) + .step_by(2) + .zip(rope_scaling.short_factor.iter()) + .map(|(i, &f)| f / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_s = Tensor::from_vec(inv_freq_s, (1, ()), dev)?.to_dtype(dtype)?; + let max_seq_len = cfg.max_position_embeddings; + match cfg.original_max_position_embeddings { + None => { + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + t.matmul(&inv_freq_s)? + } + Some(original_max_seq_len) => { + let t_s = Tensor::arange(0u32, original_max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((original_max_seq_len, 1))?; + let freq_s = t_s.matmul(&inv_freq_s)?; + let inv_freq_l: Vec<_> = (0..dim) + .step_by(2) + .zip(rope_scaling.long_factor.iter()) + .map(|(i, &f)| f / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_l = + Tensor::from_vec(inv_freq_l, (1, ()), dev)?.to_dtype(dtype)?; + let t_l = + Tensor::arange(original_max_seq_len as u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape(((), 1))?; + let freq_l = t_l.matmul(&inv_freq_l)?; + Tensor::cat(&[&freq_s, &freq_l], 0)? + } + } + } + }; + Ok(Self { + partial_dim, + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn rope(&self, xs: &Tensor, cos: &Tensor, sin: &Tensor) -> Result { + let x = match self.partial_dim { + None => candle_nn::rotary_emb::rope(&xs.contiguous()?, cos, sin)?, + Some(dim) => { + let xs_rot = xs.i((.., .., .., ..dim))?.contiguous()?; + let xs_pass = xs.i((.., .., .., dim..))?; + let xs_rot = candle_nn::rotary_emb::rope(&xs_rot, cos, sin)?; + Tensor::cat(&[&xs_rot, &xs_pass], D::Minus1)?.contiguous()? + } + }; + Ok(x) + } + + pub fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let q_embed = self.rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = self.rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +struct Attention { + qkv_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let head_dim = cfg.head_dim(); + let op_size = num_heads * head_dim + 2 * num_kv_heads * head_dim; + let qkv_proj = linear(cfg.hidden_size, op_size, vb.pp("qkv_proj"))?; + let o_proj = linear(num_heads * head_dim, cfg.hidden_size, vb.pp("o_proj"))?; + Ok(Self { + qkv_proj, + o_proj, + rotary_emb, + kv_cache: None, + num_heads, + num_kv_heads, + num_kv_groups: num_heads / num_kv_heads, + head_dim, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let qkv = self.qkv_proj.forward(xs)?; + let query_pos = self.num_heads * self.head_dim; + let query_states = qkv.narrow(D::Minus1, 0, query_pos)?; + let key_states = qkv.narrow(D::Minus1, query_pos, self.num_kv_heads * self.head_dim)?; + let value_states = qkv.narrow( + D::Minus1, + query_pos + self.num_kv_heads * self.head_dim, + self.num_kv_heads * self.head_dim, + )?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let (query_states, key_states) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states, seqlen_offset)?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?.contiguous()?; + let value_states = + crate::utils::repeat_kv(value_states, self.num_kv_groups)?.contiguous()?; + + let attn_output = { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + }; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, ()))? + .apply(&self.o_proj) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +struct Mlp { + gate_up_proj: Linear, + down_proj: Linear, + act_fn: candle_nn::Activation, + i_size: usize, +} + +impl Mlp { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_size = cfg.hidden_size; + let i_size = cfg.intermediate_size; + let gate_up_proj = linear(hidden_size, 2 * i_size, vb.pp("gate_up_proj"))?; + let down_proj = linear(i_size, hidden_size, vb.pp("down_proj"))?; + Ok(Self { + gate_up_proj, + down_proj, + act_fn: cfg.hidden_act, + i_size, + }) + } +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let up_states = xs.apply(&self.gate_up_proj)?; + let gate = up_states.narrow(D::Minus1, 0, self.i_size)?; + let up_states = up_states.narrow(D::Minus1, self.i_size, self.i_size)?; + let up_states = (up_states * gate.apply(&self.act_fn))?; + up_states.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: Mlp, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl DecoderLayer { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let self_attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let mlp = Mlp::new(cfg, vb.pp("mlp"))?; + let input_layernorm = + RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let post_attention_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; + residual + xs + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: Linear, + device: Device, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb_m.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(rotary_emb.clone(), cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb_m.pp("norm"))?; + let lm_head = if cfg.tie_word_embeddings { + Linear::from_weights(embed_tokens.embeddings().clone(), None) + } else { + linear(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))? + }; + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { f32::NEG_INFINITY } else { 0. })) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let (b_size, seq_len) = input_ids.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(b_size, seq_len, seqlen_offset)?; + Some(mask) + }; + let mut xs = self.embed_tokens.forward(input_ids)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)? + } + xs.narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +} diff --git a/patches/candle-transformers/src/models/pixtral/llava.rs b/patches/candle-transformers/src/models/pixtral/llava.rs new file mode 100644 index 0000000000..4aff26a784 --- /dev/null +++ b/patches/candle-transformers/src/models/pixtral/llava.rs @@ -0,0 +1,98 @@ +use candle::{Module, Result, Tensor}; +use candle_nn::{linear, Linear, VarBuilder}; + +use super::vision_model; +use crate::models::mistral; + +#[derive(serde::Deserialize, Debug, Clone)] +pub struct Config { + pub projector_hidden_act: candle_nn::Activation, + pub text_config: mistral::Config, + pub vision_config: vision_model::Config, + pub image_token_index: usize, + pub image_seq_length: usize, +} + +#[derive(Debug, Clone)] +pub struct MultiModalProjector { + linear_1: Linear, + act: candle_nn::Activation, + linear_2: Linear, +} + +impl MultiModalProjector { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let (hidden_v, hidden_t) = (cfg.vision_config.hidden_size, cfg.text_config.hidden_size); + let linear_1 = linear(hidden_v, hidden_t, vb.pp("linear_1"))?; + let linear_2 = linear(hidden_t, hidden_t, vb.pp("linear_2"))?; + Ok(Self { + linear_1, + act: cfg.projector_hidden_act, + linear_2, + }) + } +} + +impl Module for MultiModalProjector { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.linear_1)? + .apply(&self.act)? + .apply(&self.linear_2) + } +} + +#[derive(Debug, Clone)] +pub struct Model { + pub multi_modal_projector: MultiModalProjector, + pub language_model: mistral::Model, + pub vision_tower: vision_model::Model, + pub patch_size: usize, + pub dtype: candle::DType, + pub pos: usize, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let language_model = mistral::Model::new(&cfg.text_config, vb.pp("language_model"))?; + let vision_tower = vision_model::Model::new( + &cfg.vision_config, + vb.pp("vision_tower").to_dtype(candle::DType::F32), + )?; + let multi_modal_projector = MultiModalProjector::new( + cfg, + vb.pp("multi_modal_projector").to_dtype(candle::DType::F32), + )?; + Ok(Self { + multi_modal_projector, + language_model, + vision_tower, + patch_size: cfg.vision_config.patch_size, + dtype: vb.dtype(), + pos: 0, + }) + } + + pub fn clear_kv_cache(&mut self) { + self.language_model.clear_kv_cache(); + self.pos = 0; + } + + pub fn encode_image(&self, image: &Tensor) -> Result { + let image_embeds = self.vision_tower.forward(image)?; + self.multi_modal_projector.forward(&image_embeds) + } + + pub fn lm_forward(&mut self, input_ids: &Tensor) -> Result { + let (_, seq_len) = input_ids.dims2()?; + let logits = self.language_model.forward(input_ids, self.pos)?; + self.pos += seq_len; + Ok(logits) + } + + pub fn lm_forward_embeds(&mut self, xs: &Tensor) -> Result { + let (_, seq_len, _) = xs.dims3()?; + let logits = self.language_model.forward_embeds(xs, None, self.pos)?; + self.pos += seq_len; + Ok(logits) + } +} diff --git a/patches/candle-transformers/src/models/pixtral/mod.rs b/patches/candle-transformers/src/models/pixtral/mod.rs new file mode 100644 index 0000000000..18bcc5f793 --- /dev/null +++ b/patches/candle-transformers/src/models/pixtral/mod.rs @@ -0,0 +1,43 @@ +//! Pixtral Language-Image Pre-Training +//! +//! Pixtral is an architecture trained for multimodal learning +//! using images paired with text descriptions. +//! +//! - 💻 Transformers Python [reference implementation](https://github.com/huggingface/transformers/tree/main/src/transformers/models/pixtral) +//! - 📝 [Blog Post](https://mistral.ai/news/pixtral-12b/) +//! - 🤗 [HF Model Card](https://huggingface.co/mistralai/Pixtral-12B-2409) +//! - 🤗 [HF Community Model Card](https://huggingface.co/mistral-community/pixtral-12b) +//! +//! # Example +//! +//!
+//! +//!
+//! +//! ```bash +//! cargo run --profile=release-with-debug \ +//! --features cuda \ +//! --example pixtral -- \ +//! --image candle-examples/examples/flux/assets/flux-robot.jpg +//! ``` +//! +//! ```txt +//! Describe the image. +//! +//! The image depicts a charming, rustic robot standing on a sandy beach at sunset. +//! The robot has a vintage, steampunk aesthetic with visible gears and mechanical +//! parts. It is holding a small lantern in one hand, which emits a warm glow, and +//! its other arm is extended forward as if reaching out or guiding the way. The +//! robot's body is adorned with the word "RUST" in bright orange letters, adding to +//! its rustic theme. +//! +//! The background features a dramatic sky filled with clouds, illuminated by the +//! setting sun, casting a golden hue over the scene. Gentle waves lap against the +//! shore, creating a serene and picturesque atmosphere. The overall mood of the +//! image is whimsical and nostalgic, evoking a sense of adventure and tranquility. +//! ``` + +pub mod llava; +pub mod vision_model; + +pub use llava::{Config, Model}; diff --git a/patches/candle-transformers/src/models/pixtral/vision_model.rs b/patches/candle-transformers/src/models/pixtral/vision_model.rs new file mode 100644 index 0000000000..3f884aaf89 --- /dev/null +++ b/patches/candle-transformers/src/models/pixtral/vision_model.rs @@ -0,0 +1,366 @@ +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::{linear_b, rms_norm, Linear, RmsNorm, VarBuilder}; + +fn default_act() -> candle_nn::Activation { + candle_nn::Activation::Silu +} + +fn default_hidden_size() -> usize { + 1024 +} + +fn default_intermediate_size() -> usize { + 4096 +} + +fn default_num_channels() -> usize { + 3 +} + +fn default_num_hidden_layers() -> usize { + 24 +} + +fn default_num_attention_heads() -> usize { + 16 +} + +#[derive(serde::Deserialize, Debug, Clone)] +pub struct Config { + #[serde(default = "default_hidden_size")] + pub hidden_size: usize, + #[serde(default = "default_num_channels")] + pub num_channels: usize, + pub image_size: usize, + pub patch_size: usize, + pub rope_theta: f64, + #[serde(default = "default_intermediate_size")] + pub intermediate_size: usize, + #[serde(default = "default_num_hidden_layers")] + pub num_hidden_layers: usize, + pub head_dim: Option, + #[serde(default = "default_num_attention_heads")] + pub num_attention_heads: usize, + #[serde(default = "default_act")] + pub hidden_act: candle_nn::Activation, +} + +impl Config { + pub fn pixtral_12b_2409() -> Self { + Self { + hidden_size: 1024, + num_channels: 3, + image_size: 1024, + patch_size: 16, + rope_theta: 10000.0, + intermediate_size: 4096, + num_hidden_layers: 24, + num_attention_heads: 16, + head_dim: None, + // Default + hidden_act: candle_nn::Activation::Silu, + } + } + + fn head_dim(&self) -> usize { + self.head_dim + .unwrap_or(self.hidden_size / self.num_attention_heads) + } +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + scale: f64, + num_heads: usize, + head_dim: usize, +} + +impl Attention { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let h = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let head_dim = cfg.head_dim(); + let q_proj = linear_b(h, h, false, vb.pp("q_proj"))?; + let k_proj = linear_b(h, h, false, vb.pp("k_proj"))?; + let v_proj = linear_b(h, h, false, vb.pp("v_proj"))?; + let o_proj = linear_b(h, h, false, vb.pp("o_proj"))?; + let scale = (head_dim as f64).powf(-0.5); + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + scale, + num_heads, + head_dim, + }) + } + + fn forward( + &self, + xs: &Tensor, + emb: &RotaryEmbedding, + subsampled_positions: Option<&Tensor>, + attention_mask: Option<&Tensor>, + ) -> Result { + let (b, patches, _) = xs.dims3()?; + let query_states = xs.apply(&self.q_proj)?; + let key_states = xs.apply(&self.k_proj)?; + let value_states = xs.apply(&self.v_proj)?; + + let shape = (b, patches, self.num_heads, self.head_dim); + let query_states = query_states.reshape(shape)?.transpose(1, 2)?.contiguous()?; + let key_states = key_states.reshape(shape)?.transpose(1, 2)?.contiguous()?; + let value_states = value_states.reshape(shape)?.transpose(1, 2)?.contiguous()?; + + let (query_states, key_states) = + emb.apply_rotary_emb_qkv(&query_states, &key_states, subsampled_positions)?; + let attn_weights = (query_states.matmul(&key_states.t()?)? * self.scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights + .matmul(&value_states)? + .transpose(1, 2)? + .reshape((b, patches, ()))? + .apply(&self.o_proj) + } +} + +#[derive(Debug, Clone)] +struct Mlp { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: candle_nn::Activation, +} + +impl Mlp { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let (h, i) = (cfg.hidden_size, cfg.intermediate_size); + let gate_proj = linear_b(h, i, false, vb.pp("gate_proj"))?; + let up_proj = linear_b(h, i, false, vb.pp("up_proj"))?; + let down_proj = linear_b(i, h, false, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_act, + }) + } +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + (xs.apply(&self.gate_proj)?.apply(&self.act_fn)? * xs.apply(&self.up_proj))? + .apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +struct AttentionLayer { + attention_norm: RmsNorm, + feed_forward: Mlp, + attention: Attention, + ffn_norm: RmsNorm, +} + +impl AttentionLayer { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let attention_norm = rms_norm(cfg.hidden_size, 1e-5, vb.pp("attention_norm"))?; + let feed_forward = Mlp::new(cfg, vb.pp("feed_forward"))?; + let attention = Attention::new(cfg, vb.pp("attention"))?; + let ffn_norm = rms_norm(cfg.hidden_size, 1e-5, vb.pp("ffn_norm"))?; + Ok(Self { + attention_norm, + feed_forward, + attention, + ffn_norm, + }) + } + + fn forward( + &self, + xs: &Tensor, + emb: &RotaryEmbedding, + subsampled_positions: Option<&Tensor>, + attention_mask: Option<&Tensor>, + ) -> Result { + let residual = xs; + let xs = self.attention.forward( + &xs.apply(&self.attention_norm)?, + emb, + subsampled_positions, + attention_mask, + )?; + let xs = (residual + xs)?; + let residual = &xs; + let xs = xs.apply(&self.ffn_norm)?.apply(&self.feed_forward)?; + xs + residual + } +} + +#[derive(Debug, Clone)] +struct Transformer { + layers: Vec, +} + +impl Transformer { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb = vb.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = AttentionLayer::new(cfg, vb.pp(layer_idx))?; + layers.push(layer) + } + Ok(Self { layers }) + } + + fn forward( + &self, + xs: &Tensor, + emb: &RotaryEmbedding, + subsampled_positions: Option<&Tensor>, + attention_mask: Option<&Tensor>, + ) -> Result { + let mut xs = xs.clone(); + for layer in self.layers.iter() { + xs = layer.forward(&xs, emb, subsampled_positions, attention_mask)? + } + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + cos: Tensor, + sin: Tensor, +} + +impl RotaryEmbedding { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dtype = vb.dtype(); + let dev = vb.device(); + let dim = cfg.head_dim(); + let rope_theta = cfg.rope_theta as f32; + let max_patches_per_side = cfg.image_size / cfg.patch_size; + let freqs: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / rope_theta.powf(i as f32 / dim as f32)) + .collect(); + let freqs_h = freqs.iter().step_by(2).copied().collect::>(); + let freqs_h = Tensor::new(freqs_h, dev)?; + let freqs_w = freqs.iter().skip(1).step_by(2).copied().collect::>(); + let freqs_w = Tensor::new(freqs_w, dev)?; + let h = Tensor::arange(0u32, max_patches_per_side as u32, dev)?.to_dtype(DType::F32)?; + let w = Tensor::arange(0u32, max_patches_per_side as u32, dev)?.to_dtype(DType::F32)?; + let freqs_h = h.unsqueeze(1)?.matmul(&freqs_h.unsqueeze(0)?)?; + let freqs_w = w.unsqueeze(1)?.matmul(&freqs_w.unsqueeze(0)?)?; + let inv_freq = Tensor::cat( + &[ + freqs_h.unsqueeze(1)?.repeat((1, max_patches_per_side, 1))?, + freqs_w.unsqueeze(0)?.repeat((max_patches_per_side, 1, 1))?, + ], + D::Minus1, + )? + .reshape(((), dim / 2))?; + let cos = inv_freq.cos()?.to_dtype(dtype)?; + let sin = inv_freq.sin()?.to_dtype(dtype)?; + Ok(Self { cos, sin }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + subsampled_positions: Option<&Tensor>, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, _seq_len, _n_embd) = q.dims4()?; + let (cos, sin) = match subsampled_positions { + None => (&self.cos, &self.sin), + Some(pos) => ( + &self.cos.index_select(pos, 0)?, + &self.sin.index_select(pos, 0)?, + ), + }; + let q_embed = candle_nn::rotary_emb::rope(q, cos, sin)?; + let k_embed = candle_nn::rotary_emb::rope(k, cos, sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +pub struct Model { + patch_conv: candle_nn::Conv2d, + ln_pre: RmsNorm, + transformer: Transformer, + patch_positional_embedding: RotaryEmbedding, + max_image_width: u32, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let conv2d_cfg = candle_nn::Conv2dConfig { + stride: cfg.patch_size, + ..Default::default() + }; + let patch_conv = candle_nn::conv2d_no_bias( + cfg.num_channels, + cfg.hidden_size, + cfg.patch_size, + conv2d_cfg, + vb.pp("patch_conv"), + )?; + let ln_pre = candle_nn::rms_norm(cfg.hidden_size, 1e-5, vb.pp("ln_pre"))?; + let transformer = Transformer::new(cfg, vb.pp("transformer"))?; + let patch_positional_embedding = + RotaryEmbedding::new(cfg, vb.pp("patch_positional_embedding"))?; + let max_image_width = (cfg.image_size / cfg.patch_size) as u32; + Ok(Self { + patch_conv, + ln_pre, + transformer, + patch_positional_embedding, + max_image_width, + }) + } + + pub fn position_ids_in_meshgrid( + &self, + num_patches_h: usize, + num_patches_w: usize, + device: &Device, + ) -> Result { + let idx = Tensor::arange(0, num_patches_h as u32, device)?; + let idy = Tensor::arange(0, num_patches_w as u32, device)?; + let mesh = Tensor::meshgrid(&[idx, idy], false)?; + let ids = (&mesh[0] * (self.max_image_width as f64) + &mesh[1])?.flatten_all()?; + Ok(ids) + } +} + +impl Module for Model { + fn forward(&self, xs: &Tensor) -> Result { + let patch_embeds = xs.apply(&self.patch_conv)?; + let subsampled_positions = Some(self.position_ids_in_meshgrid( + patch_embeds.dim(2)?, + patch_embeds.dim(3)?, + patch_embeds.device(), + )?); + let patch_embeds = patch_embeds.flatten_from(2)?.t()?.apply(&self.ln_pre)?; + self.transformer.forward( + &patch_embeds, + &self.patch_positional_embedding, + subsampled_positions.as_ref(), + None, + ) + } +} diff --git a/patches/candle-transformers/src/models/quantized_blip.rs b/patches/candle-transformers/src/models/quantized_blip.rs new file mode 100644 index 0000000000..acba9ba191 --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_blip.rs @@ -0,0 +1,277 @@ +//! BLIP model implementation with quantization support. +//! +//! BLIP is a vision-language model for image understanding and generation tasks. +//! This implementation provides quantization for reduced memory and compute. +//! +//! Key characteristics: +//! - Vision encoder using ViT architecture +//! - Text decoder using BERT-style transformer +//! - Cross-attention between vision and text features +//! - Support for 8-bit quantization +//! +//! References: +//! - [BLIP Paper](https://arxiv.org/abs/2201.12086) +//! - [Hugging Face Implementation](https://huggingface.co/docs/transformers/model_doc/blip) +//! + +use super::quantized_blip_text as blip_text; +use crate::quantized_nn::{layer_norm, linear, Linear}; +pub use crate::quantized_var_builder::VarBuilder; +use candle::{Module, Result, Tensor, D}; +use candle_nn::{Conv2d, Conv2dConfig, LayerNorm}; + +pub type VisionConfig = super::blip::VisionConfig; +pub type Config = super::blip::Config; + +#[derive(Debug, Clone)] +struct VisionEmbeddings { + class_embedding: Tensor, + patch_embedding: Conv2d, + position_embedding: Tensor, +} + +impl VisionEmbeddings { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let class_embedding = vb + .get((1, 1, cfg.hidden_size), "class_embedding")? + .dequantize(vb.device())?; + let conv_cfg = Conv2dConfig { + stride: cfg.patch_size, + ..Default::default() + }; + let pe_vb = vb.pp("patch_embedding"); + let pe_weight = pe_vb + .get( + (cfg.hidden_size, 3, cfg.patch_size, cfg.patch_size), + "weight", + )? + .dequantize(vb.device())?; + let pe_bias = pe_vb + .get(cfg.hidden_size, "bias")? + .dequantize(vb.device())?; + + let patch_embedding = Conv2d::new(pe_weight, Some(pe_bias), conv_cfg); + let num_patches1 = cfg.image_size / cfg.patch_size; + let num_patches = num_patches1 * num_patches1; + let num_positions = num_patches + 1; + let position_embedding = vb + .get((1, num_positions, cfg.hidden_size), "position_embedding")? + .dequantize(vb.device())?; + Ok(Self { + class_embedding, + patch_embedding, + position_embedding, + }) + } +} + +impl Module for VisionEmbeddings { + fn forward(&self, xs: &Tensor) -> Result { + let target_dtype = xs.dtype(); + let b_size = xs.dim(0)?; + let patch_embeds = xs.apply(&self.patch_embedding)?.flatten_from(2)?.t()?; + let d = self.class_embedding.dim(D::Minus1)?; + let class_embeds = self + .class_embedding + .broadcast_as((b_size, 1, d))? + .to_dtype(target_dtype)?; + let embeddings = Tensor::cat(&[&class_embeds, &patch_embeds], 1)?; + let position_embedding = self.position_embedding.narrow(1, 0, embeddings.dim(1)?)?; + embeddings.broadcast_add(&position_embedding) + } +} + +#[derive(Debug, Clone)] +struct Attention { + qkv: Linear, + projection: Linear, + scale: f64, + num_heads: usize, +} + +impl Attention { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let embed_dim = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let head_dim = embed_dim / num_heads; + let scale = 1f64 / (head_dim as f64).sqrt(); + let qkv = linear(embed_dim, 3 * embed_dim, vb.pp("qkv"))?; + let projection = linear(embed_dim, embed_dim, vb.pp("projection"))?; + Ok(Self { + qkv, + projection, + scale, + num_heads, + }) + } + + fn forward(&self, xs: &Tensor, attn_mask: Option<&Tensor>) -> Result { + let (b_sz, tgt_len, embed_dim) = xs.dims3()?; + let mixed_qkv = xs + .apply(&self.qkv)? + .reshape((b_sz, tgt_len, 3, self.num_heads, embed_dim / self.num_heads))? + .permute((2, 0, 3, 1, 4))?; + let query = mixed_qkv.get(0)?; + let key = mixed_qkv.get(1)?; + let value = mixed_qkv.get(2)?; + let attention_scores = query.matmul(&key.t()?)?; + let attention_scores = (attention_scores * self.scale)?; + let attention_probs = candle_nn::ops::softmax_last_dim(&attention_scores)?; + let attention_probs = match attn_mask { + None => attention_probs, + Some(attn_mask) => (attention_probs * attn_mask)?, + }; + attention_probs + .matmul(&value)? + .permute((0, 2, 1, 3))? + .flatten_from(D::Minus2)? + .apply(&self.projection) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + activation_fn: candle_nn::Activation, + fc1: Linear, + fc2: Linear, +} + +impl MLP { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let fc1 = linear(cfg.hidden_size, cfg.intermediate_size, vb.pp("fc1"))?; + let fc2 = linear(cfg.intermediate_size, cfg.hidden_size, vb.pp("fc2"))?; + Ok(Self { + activation_fn: cfg.hidden_act, + fc1, + fc2, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.fc1)? + .apply(&self.activation_fn)? + .apply(&self.fc2) + } +} + +#[derive(Debug, Clone)] +struct EncoderLayer { + self_attn: Attention, + layer_norm1: LayerNorm, + mlp: MLP, + layer_norm2: LayerNorm, +} + +impl EncoderLayer { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let embed_dim = cfg.hidden_size; + let self_attn = Attention::new(cfg, vb.pp("self_attn"))?; + let layer_norm1 = layer_norm(embed_dim, cfg.layer_norm_eps, vb.pp("layer_norm1"))?; + let layer_norm2 = layer_norm(embed_dim, cfg.layer_norm_eps, vb.pp("layer_norm2"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + Ok(Self { + self_attn, + layer_norm1, + mlp, + layer_norm2, + }) + } + + fn forward(&self, xs: &Tensor, attention_mask: Option<&Tensor>) -> Result { + let residual = xs; + let xs = xs.apply(&self.layer_norm1)?; + let xs = self.self_attn.forward(&xs, attention_mask)?; + let xs = (xs + residual)?; + + let residual = &xs; + let xs = xs.apply(&self.layer_norm2)?.apply(&self.mlp)?; + xs + residual + } +} + +#[derive(Debug, Clone)] +struct Encoder { + layers: Vec, +} + +impl Encoder { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb = vb.pp("layers"); + for i in 0..cfg.num_hidden_layers { + let layer = EncoderLayer::new(cfg, vb.pp(i))?; + layers.push(layer) + } + Ok(Self { layers }) + } + + fn forward(&self, xs: &Tensor, attention_mask: Option<&Tensor>) -> Result { + let mut xs = xs.clone(); + for layer in self.layers.iter() { + xs = layer.forward(&xs, attention_mask)? + } + Ok(xs) + } +} + +#[derive(Debug, Clone)] +pub struct VisionModel { + embeddings: VisionEmbeddings, + encoder: Encoder, + post_layernorm: LayerNorm, +} + +impl VisionModel { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let embeddings = VisionEmbeddings::new(cfg, vb.pp("embeddings"))?; + let encoder = Encoder::new(cfg, vb.pp("encoder"))?; + let post_layernorm = + layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("post_layernorm"))?; + Ok(Self { + embeddings, + encoder, + post_layernorm, + }) + } +} + +impl Module for VisionModel { + fn forward(&self, xs: &Tensor) -> Result { + let xs = xs.apply(&self.embeddings)?; + let encoder_outputs = self.encoder.forward(&xs, None)?; + // Return the last hidden state rather than pooled outputs. + encoder_outputs.apply(&self.post_layernorm) + } +} + +#[derive(Debug, Clone)] +pub struct BlipForConditionalGeneration { + vision_model: VisionModel, + text_decoder: blip_text::TextLMHeadModel, +} + +impl BlipForConditionalGeneration { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vision_model = VisionModel::new(&cfg.vision_config, vb.pp("vision_model"))?; + let text_decoder = + blip_text::TextLMHeadModel::new(&cfg.text_config, vb.pp("text_decoder"))?; + Ok(Self { + vision_model, + text_decoder, + }) + } + + pub fn vision_model(&self) -> &VisionModel { + &self.vision_model + } + + pub fn text_decoder(&mut self) -> &mut blip_text::TextLMHeadModel { + &mut self.text_decoder + } + pub fn reset_kv_cache(&mut self) { + self.text_decoder.reset_kv_cache(); + } +} diff --git a/patches/candle-transformers/src/models/quantized_blip_text.rs b/patches/candle-transformers/src/models/quantized_blip_text.rs new file mode 100644 index 0000000000..7b753fb116 --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_blip_text.rs @@ -0,0 +1,493 @@ +//! Quantized BLIP text module implementation. +//! +//! Provides the text decoder portion of the BLIP model with 8-bit quantization. +//! Uses a BERT-style transformer architecture for text processing. +//! +//! Key components: +//! - Text embeddings layer with position embeddings +//! - Multi-head self attention layers +//! - Cross-attention for vision-text fusion +//! - Layer normalization and feed-forward layers +//! - Quantized linear transformations +//! +//! References: +//! - [BLIP Paper](https://arxiv.org/abs/2201.12086) +//! - [Hugging Face Implementation](https://huggingface.co/docs/transformers/model_doc/blip) +//! + +use crate::models::with_tracing::QMatMul; +use crate::quantized_nn::{layer_norm, linear, Embedding, Linear}; +pub use crate::quantized_var_builder::VarBuilder; +use candle::{Module, Result, Tensor, D}; +use candle_nn::LayerNorm; + +pub type Config = super::blip_text::Config; + +#[derive(Debug, Clone)] +struct TextEmbeddings { + word_embeddings: Embedding, + position_embeddings: Embedding, + layer_norm: LayerNorm, + position_ids: Tensor, +} + +impl TextEmbeddings { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let word_embeddings = + Embedding::new(cfg.vocab_size, cfg.hidden_size, vb.pp("word_embeddings"))?; + let position_embeddings = Embedding::new( + cfg.max_position_embeddings, + cfg.hidden_size, + vb.pp("position_embeddings"), + )?; + let layer_norm = layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("LayerNorm"))?; + let position_ids = + Tensor::arange(0, cfg.max_position_embeddings as u32, vb.device())?.unsqueeze(0)?; + Ok(Self { + word_embeddings, + position_embeddings, + layer_norm, + position_ids, + }) + } + + fn forward(&self, xs: &Tensor, past_kv_len: usize) -> Result { + let seq_len = xs.dim(1)?; + let position_ids = self.position_ids.narrow(1, past_kv_len, seq_len)?; + let embeddings = self.word_embeddings.forward(xs)?; + let position_embeddings = self.position_embeddings.forward(&position_ids)?; + (embeddings + position_embeddings)?.apply(&self.layer_norm) + } +} + +#[derive(Debug, Clone)] +struct TextSelfAttention { + query: Linear, + key: Linear, + value: Linear, + attention_head_size: usize, + num_attention_heads: usize, + attention_scale: f64, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl TextSelfAttention { + fn new(cfg: &Config, is_cross_attention: bool, vb: VarBuilder) -> Result { + let num_attention_heads = cfg.num_attention_heads; + let attention_head_size = cfg.hidden_size / num_attention_heads; + let all_head_size = cfg.num_attention_heads * attention_head_size; + let query = linear(cfg.hidden_size, all_head_size, vb.pp("query"))?; + let in_size = if is_cross_attention { + cfg.encoder_hidden_size + } else { + cfg.hidden_size + }; + let key = linear(in_size, all_head_size, vb.pp("key"))?; + let value = linear(in_size, all_head_size, vb.pp("value"))?; + let attention_scale = 1f64 / (attention_head_size as f64).sqrt(); + Ok(Self { + query, + key, + value, + attention_head_size, + num_attention_heads, + attention_scale, + kv_cache: None, + }) + } + + fn transpose_for_scores(&self, xs: &Tensor) -> Result { + let (b_size, seq_len, _) = xs.dims3()?; + xs.reshape(( + b_size, + seq_len, + self.num_attention_heads, + self.attention_head_size, + ))? + .permute((0, 2, 1, 3)) + } + + fn reset_kv_cache(&mut self) { + self.kv_cache = None + } + + fn forward( + &mut self, + xs: &Tensor, + encoder_hidden_states: Option<&Tensor>, + attention_mask: Option<&Tensor>, + ) -> Result { + let query = self + .transpose_for_scores(&self.query.forward(xs)?)? + .contiguous()?; + let (key, value) = match encoder_hidden_states { + None => { + let key = self.transpose_for_scores(&self.key.forward(xs)?)?; + let value = self.transpose_for_scores(&self.value.forward(xs)?)?; + let (key, value) = match &self.kv_cache { + None => (key, value), + Some((prev_key, prev_value)) => { + let key = Tensor::cat(&[prev_key, &key], 2)?; + let value = Tensor::cat(&[prev_value, &value], 2)?; + (key, value) + } + }; + self.kv_cache = Some((key.clone(), value.clone())); + (key, value) + } + Some(xs) => { + let key = self.transpose_for_scores(&self.key.forward(xs)?)?; + let value = self.transpose_for_scores(&self.value.forward(xs)?)?; + // no kv-cache in this case, but the results could probably be memoized. + (key, value) + } + }; + let key = key.contiguous()?; + let value = value.contiguous()?; + let attention_scores = query.matmul(&key.t()?)?; + let attention_scores = (attention_scores * self.attention_scale)?; + let attention_scores = match attention_mask { + Some(mask) => attention_scores.broadcast_add(mask)?, + None => attention_scores, + }; + let attention_probs = candle_nn::ops::softmax_last_dim(&attention_scores)?; + attention_probs + .matmul(&value)? + .permute((0, 2, 1, 3))? + .flatten_from(D::Minus2) + } +} + +#[derive(Debug, Clone)] +struct TextSelfOutput { + dense: Linear, + layer_norm: LayerNorm, +} + +impl TextSelfOutput { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.hidden_size, cfg.hidden_size, vb.pp("dense"))?; + let layer_norm = layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("LayerNorm"))?; + Ok(Self { dense, layer_norm }) + } + + fn forward(&self, xs: &Tensor, input_tensor: &Tensor) -> Result { + (xs.apply(&self.dense) + input_tensor)?.apply(&self.layer_norm) + } +} + +#[derive(Debug, Clone)] +struct TextAttention { + self_: TextSelfAttention, + output: TextSelfOutput, +} + +impl TextAttention { + fn new(cfg: &Config, is_cross_attention: bool, vb: VarBuilder) -> Result { + let self_ = TextSelfAttention::new(cfg, is_cross_attention, vb.pp("self"))?; + let output = TextSelfOutput::new(cfg, vb.pp("output"))?; + Ok(Self { self_, output }) + } + + fn reset_kv_cache(&mut self) { + self.self_.reset_kv_cache() + } + + fn forward( + &mut self, + xs: &Tensor, + encoder_hidden_states: Option<&Tensor>, + attention_mask: Option<&Tensor>, + ) -> Result { + let self_outputs = self + .self_ + .forward(xs, encoder_hidden_states, attention_mask)?; + self.output.forward(&self_outputs, xs) + } +} + +#[derive(Debug, Clone)] +struct TextIntermediate { + dense: Linear, + intermediate_act_fn: candle_nn::Activation, +} + +impl TextIntermediate { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.hidden_size, cfg.intermediate_size, vb.pp("dense"))?; + Ok(Self { + dense, + intermediate_act_fn: cfg.hidden_act, + }) + } +} + +impl Module for TextIntermediate { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.dense)?.apply(&self.intermediate_act_fn) + } +} + +#[derive(Debug, Clone)] +struct TextOutput { + dense: Linear, + layer_norm: LayerNorm, +} + +impl TextOutput { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.intermediate_size, cfg.hidden_size, vb.pp("dense"))?; + let layer_norm = layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("LayerNorm"))?; + Ok(Self { dense, layer_norm }) + } + + fn forward(&self, xs: &Tensor, input_tensor: &Tensor) -> Result { + (xs.apply(&self.dense)? + input_tensor)?.apply(&self.layer_norm) + } +} + +#[derive(Debug, Clone)] +struct TextLayer { + attention: TextAttention, + cross_attention: Option, + intermediate: TextIntermediate, + output: TextOutput, +} + +impl TextLayer { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let attention = TextAttention::new(cfg, false, vb.pp("attention"))?; + let cross_attention = if cfg.is_decoder { + Some(TextAttention::new(cfg, true, vb.pp("crossattention"))?) + } else { + None + }; + let intermediate = TextIntermediate::new(cfg, vb.pp("intermediate"))?; + let output = TextOutput::new(cfg, vb.pp("output"))?; + Ok(Self { + attention, + cross_attention, + intermediate, + output, + }) + } + + fn reset_kv_cache(&mut self) { + self.attention.reset_kv_cache(); + if let Some(ca) = &mut self.cross_attention { + ca.reset_kv_cache() + } + } + + fn forward( + &mut self, + xs: &Tensor, + encoder_hidden_states: &Tensor, + attention_mask: &Tensor, + ) -> Result { + let attention_output = self.attention.forward(xs, None, Some(attention_mask))?; + let attention_output = match &mut self.cross_attention { + Some(ca) => ca.forward(&attention_output, Some(encoder_hidden_states), None)?, + None => candle::bail!("expected some cross-attn"), + }; + let intermediate_output = self.intermediate.forward(&attention_output)?; + self.output.forward(&intermediate_output, &attention_output) + } +} + +#[derive(Debug, Clone)] +struct TextEncoder { + layers: Vec, +} + +impl TextEncoder { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb = vb.pp("layer"); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + for i in 0..cfg.num_hidden_layers { + let layer = TextLayer::new(cfg, vb.pp(i))?; + layers.push(layer) + } + Ok(Self { layers }) + } + + fn reset_kv_cache(&mut self) { + self.layers.iter_mut().for_each(|l| l.reset_kv_cache()) + } + + fn forward( + &mut self, + xs: &Tensor, + encoder_hidden_states: &Tensor, + attention_mask: &Tensor, + ) -> Result { + let mut xs = xs.clone(); + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, encoder_hidden_states, attention_mask)? + } + Ok(xs) + } +} + +#[derive(Debug, Clone)] +pub struct TextPooler { + dense: Linear, +} + +impl TextPooler { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.hidden_size, cfg.hidden_size, vb.pp("dense"))?; + Ok(Self { dense }) + } +} + +impl Module for TextPooler { + fn forward(&self, xs: &Tensor) -> Result { + xs.narrow(D::Minus1, 0, 1)? + .squeeze(D::Minus1)? + .apply(&self.dense)? + .tanh() + } +} + +#[derive(Debug, Clone)] +struct TextPredictionHeadTransform { + dense: Linear, + transform_act_fn: candle_nn::Activation, + layer_norm: LayerNorm, +} + +impl TextPredictionHeadTransform { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.hidden_size, cfg.hidden_size, vb.pp("dense"))?; + let layer_norm = layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("LayerNorm"))?; + Ok(Self { + dense, + transform_act_fn: cfg.hidden_act, + layer_norm, + }) + } +} + +impl Module for TextPredictionHeadTransform { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.dense)? + .apply(&self.transform_act_fn)? + .apply(&self.layer_norm) + } +} + +#[derive(Debug, Clone)] +struct TextLMPredictionHead { + transform: TextPredictionHeadTransform, + decoder: Linear, +} + +impl TextLMPredictionHead { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let transform = TextPredictionHeadTransform::new(cfg, vb.pp("transform"))?; + let weight = QMatMul::new(cfg.hidden_size, cfg.vocab_size, vb.pp("decoder"))?; + let bias = vb.get(cfg.vocab_size, "bias")?.dequantize(vb.device())?; + let decoder = Linear::from_weights(weight, Some(bias)); + Ok(Self { transform, decoder }) + } +} + +impl Module for TextLMPredictionHead { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.transform)?.apply(&self.decoder) + } +} + +#[derive(Debug, Clone)] +struct TextOnlyMLMHead { + predictions: TextLMPredictionHead, +} + +impl TextOnlyMLMHead { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let predictions = TextLMPredictionHead::new(cfg, vb.pp("predictions"))?; + Ok(Self { predictions }) + } +} + +impl Module for TextOnlyMLMHead { + fn forward(&self, xs: &Tensor) -> Result { + self.predictions.forward(xs) + } +} + +#[derive(Debug, Clone)] +struct TextModel { + embeddings: TextEmbeddings, + encoder: TextEncoder, + past_kv_len: usize, + // We do not need the pooler for caption generation +} + +impl TextModel { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let embeddings = TextEmbeddings::new(cfg, vb.pp("embeddings"))?; + let encoder = TextEncoder::new(cfg, vb.pp("encoder"))?; + Ok(Self { + embeddings, + encoder, + past_kv_len: 0, + }) + } + + fn forward( + &mut self, + input_ids: &Tensor, + encoder_hidden_states: &Tensor, + attention_mask: &Tensor, + ) -> Result { + let (_b_sz, seq_len) = input_ids.dims2()?; + let embedding_output = self.embeddings.forward(input_ids, self.past_kv_len)?; + let sequence_output = + self.encoder + .forward(&embedding_output, encoder_hidden_states, attention_mask)?; + self.past_kv_len += seq_len; + // We're interested in the sequence-output rather than the pooled-output. + Ok(sequence_output) + } + + fn reset_kv_cache(&mut self) { + self.past_kv_len = 0; + self.encoder.reset_kv_cache(); + } +} + +#[derive(Debug, Clone)] +pub struct TextLMHeadModel { + bert: TextModel, + cls: TextOnlyMLMHead, +} + +impl TextLMHeadModel { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let bert = TextModel::new(cfg, vb.pp("bert"))?; + let cls = TextOnlyMLMHead::new(cfg, vb.pp("cls"))?; + Ok(Self { bert, cls }) + } + + pub fn forward( + &mut self, + input_ids: &Tensor, + encoder_hidden_states: &Tensor, + ) -> Result { + let seq_len = input_ids.dim(1)?; + let mask: Vec<_> = (0..seq_len) + .flat_map(|i| (0..seq_len).map(move |j| if j > i { f32::NEG_INFINITY } else { 0f32 })) + .collect(); + let mask = Tensor::from_vec(mask, (seq_len, seq_len), input_ids.device())?; + let sequence_output = self.bert.forward(input_ids, encoder_hidden_states, &mask)?; + let prediction_scores = self.cls.forward(&sequence_output)?; + // return_logits is false so we don't discard the last sequence element. + Ok(prediction_scores) + } + + pub fn reset_kv_cache(&mut self) { + self.bert.reset_kv_cache() + } +} diff --git a/patches/candle-transformers/src/models/quantized_gemma3.rs b/patches/candle-transformers/src/models/quantized_gemma3.rs new file mode 100644 index 0000000000..7af22243c6 --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_gemma3.rs @@ -0,0 +1,480 @@ +//! Gemma 3 model implementation with quantization support. +//! +//! Gemma 3 is a family of multimodal language models developed by Google. +//! This implementation provides quantization for reduced memory usage and faster inference. +//! +//! Key characteristics: +//! - Group-Query Attention (GQA) with specialized key-value heads +//! - RMSNorm for layer normalization +//! - Specialized attention patterns with separate normalization for Q/K/V +//! - Feed-forward network with SwiGLU activation +//! - Support for 2/3/4/8-bit quantization +//! +//! References: +//! - [Gemma 3 Models](https://blog.google/technology/developers/gemma-3/) +//! + +use crate::quantized_nn::RmsNorm; +use candle::quantized::gguf_file; +use candle::quantized::QTensor; +use candle::D; +use candle::{DType, Device, IndexOp, Result, Tensor}; +use candle_nn::{Embedding, Module}; + +pub const MAX_SEQ_LEN: usize = 131072; // Gemma 3 supports 128K context window +pub const DEFAULT_SLIDING_WINDOW_TYPE: usize = 6; +pub const DEFAULT_ROPE_FREQUENCY: f32 = 1_000_000.; +pub const DEFAULT_ROPE_FREQUENCY_SLIDING: f32 = 10_000.; +pub const DEFAULT_ROPE_FREQUENCY_SCALE_FACTOR: f32 = 1.; + +#[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_gate: QMatMul, // ffn_gate in GGUF + feed_forward_up: QMatMul, // ffn_up in GGUF + feed_forward_down: QMatMul, // ffn_down in GGUF +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let gate = self.feed_forward_gate.forward(xs)?; + let up = self.feed_forward_up.forward(xs)?; + let silu = candle_nn::ops::silu(&gate)?; + let gated = (silu * up)?; + self.feed_forward_down.forward(&gated) + } +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(head_dim: usize, rope_frequency: f32, device: &Device) -> Result { + let theta: Vec<_> = (0..head_dim) + .step_by(2) + .map(|i| 1f32 / rope_frequency.powf(i as f32 / head_dim as f32)) + .collect(); + let theta = Tensor::new(theta.as_slice(), device)?; + let idx_theta = Tensor::arange(0, MAX_SEQ_LEN as u32, device)? + .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(Self { sin, cos }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + index_pos: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, index_pos, seq_len)?; + let sin = self.sin.narrow(0, index_pos, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +struct LayerWeights { + // Attention components + attention_wq: QMatMul, + attention_wk: QMatMul, + attention_wv: QMatMul, + attention_wo: QMatMul, + + // Specialized normalization for Q and K + attention_q_norm: RmsNorm, + attention_k_norm: RmsNorm, + + // Layer normalization + attention_norm: RmsNorm, // Applied before attention + post_attention_norm: RmsNorm, // Applied after attention + ffn_norm: RmsNorm, // Applied before feedforward + post_ffn_norm: RmsNorm, // Applied after feedforward + + // Feed-forward network + mlp: Mlp, + + // Attention parameters + n_head: usize, // Number of query heads + n_kv_head: usize, // Number of key-value heads + head_dim: usize, // Dimension of each head + q_dim: usize, // Total dimension for queries + + sliding_window_size: Option, + + rotary_embedding: RotaryEmbedding, + neg_inf: Tensor, + + // Cache + kv_cache: Option<(Tensor, Tensor)>, + + // Tracing + span_attn: tracing::Span, + span_mlp: tracing::Span, +} + +impl LayerWeights { + fn mask( + &self, + b_sz: usize, + seq_len: usize, + index_pos: usize, + dtype: DType, + device: &Device, + ) -> Result { + let mask: Vec<_> = if let Some(sliding_window_size) = self.sliding_window_size { + (0..seq_len) + .flat_map(|i| { + (0..seq_len).map(move |j| { + if i < j || j + sliding_window_size < i { + 0u32 + } else { + 1u32 + } + }) + }) + .collect() + } else { + (0..seq_len) + .flat_map(|i| (0..seq_len).map(move |j| if i < j { 0u32 } else { 1u32 })) + .collect() + }; + let mask = Tensor::from_slice(&mask, (seq_len, seq_len), device)?; + let mask = if index_pos > 0 { + let mask0 = Tensor::zeros((seq_len, index_pos), DType::F32, device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_sz, 1, seq_len, seq_len + index_pos))? + .to_dtype(dtype) + } + + fn forward_attn( + &mut self, + x: &Tensor, + mask: Option<&Tensor>, + index_pos: usize, + ) -> Result { + let _enter = self.span_attn.enter(); + let (b_sz, seq_len, _) = 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.attention_q_norm.forward(&q.contiguous()?)?; + let k = self.attention_k_norm.forward(&k.contiguous()?)?; + + let (q, k) = self + .rotary_embedding + .apply_rotary_emb_qkv(&q, &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)?; // concat on seq dim + let v = Tensor::cat(&[v_cache, &v], 2)?; + (k, v) + } + } + }; + self.kv_cache = Some((k.clone(), v.clone())); // update cache + + // Repeat KV for GQA + let k = crate::utils::repeat_kv(k, self.n_head / self.n_kv_head)?; + let v = crate::utils::repeat_kv(v, self.n_head / self.n_kv_head)?; + + // Scaled Dot-Product Attention + let scale = 1.0 / (self.head_dim as f64).sqrt(); + let mut attn_weights = (q.matmul(&k.transpose(2, 3)?)? * scale)?; + + if let Some(mask) = mask { + let mask = mask.broadcast_as(attn_weights.shape())?; + let neg_inf = self.neg_inf.broadcast_as(attn_weights.dims())?; + attn_weights = mask.eq(0u32)?.where_cond(&neg_inf, &attn_weights)?; + } + + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + let attn_output = attn_weights.matmul(&v)?; + + let attn_output = attn_output + .transpose(1, 2)? + .reshape((b_sz, seq_len, self.q_dim))?; + + self.attention_wo.forward(&attn_output) + } +} + +#[derive(Debug, Clone)] +pub struct ModelWeights { + tok_embeddings: Embedding, + embedding_length: usize, + layers: Vec, + norm: RmsNorm, + output: QMatMul, + span: tracing::Span, + span_output: tracing::Span, +} + +impl ModelWeights { + pub fn from_gguf( + ct: gguf_file::Content, + reader: &mut R, + device: &Device, + ) -> Result { + // Detect architecture prefix by probing which keys exist in metadata. + // This supports gemma3, gemma2, gemma, gemma-embedding, and future variants. + let prefix = ["gemma3", "gemma2", "gemma", "gemma-embedding"] + .iter() + .find(|p| { + ct.metadata + .contains_key(&format!("{}.attention.head_count", p)) + }) + .copied() + .unwrap_or("gemma3"); + + let md_get = |s: &str| { + let key = format!("{prefix}.{s}"); + match ct.metadata.get(&key) { + None => candle::bail!("cannot find {key} in metadata"), + Some(v) => Ok(v), + } + }; + + let head_count = md_get("attention.head_count")?.to_u32()? as usize; + let head_count_kv = md_get("attention.head_count_kv")?.to_u32()? as usize; + let block_count = md_get("block_count")?.to_u32()? as usize; + let embedding_length = md_get("embedding_length")?.to_u32()? as usize; + let key_length = md_get("attention.key_length")?.to_u32()? as usize; + let _value_length = md_get("attention.value_length")?.to_u32()? as usize; + let rms_norm_eps = md_get("attention.layer_norm_rms_epsilon")?.to_f32()? as f64; + let sliding_window_size = md_get("attention.sliding_window")?.to_u32()? as usize; + + let sliding_window_type = md_get("attention.sliding_window_type") + .and_then(|m| Ok(m.to_u32()? as usize)) + .unwrap_or(DEFAULT_SLIDING_WINDOW_TYPE); + + let rope_freq_base = md_get("rope.freq_base") + .and_then(|m| m.to_f32()) + .unwrap_or(DEFAULT_ROPE_FREQUENCY); + + let rope_freq_base_sliding = md_get("rope.local_freq_base") + .and_then(|m| m.to_f32()) + .unwrap_or(DEFAULT_ROPE_FREQUENCY_SLIDING); + + // Unused in Llama.cpp so we aren't using it here. + let _rope_freq_scaling_factor = md_get("rope.scaling.factor") + .and_then(|m| m.to_f32()) + .unwrap_or(DEFAULT_ROPE_FREQUENCY_SCALE_FACTOR); + + // Compute the dimensions for queries, keys, and values + // These are the total dimensions when projected across all heads + let q_dim = head_count * key_length; + + let neg_inf = Tensor::new(f32::NEG_INFINITY, device)?; + + // Load token embeddings and output projection + let tok_embeddings = ct.tensor(reader, "token_embd.weight", device)?; + let tok_embeddings = tok_embeddings.dequantize(device)?; + let norm = RmsNorm::from_qtensor( + ct.tensor(reader, "output_norm.weight", device)?, + rms_norm_eps, + )?; + let output = match ct.tensor(reader, "output.weight", device) { + Ok(tensor) => tensor, + Err(_) => ct.tensor(reader, "token_embd.weight", device)?, // Use tied weights if output.weight doesn't exist + }; + + 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"), device)?; + let attention_wk = ct.tensor(reader, &format!("{prefix}.attn_k.weight"), device)?; + let attention_wv = ct.tensor(reader, &format!("{prefix}.attn_v.weight"), device)?; + let attention_wo = + ct.tensor(reader, &format!("{prefix}.attn_output.weight"), device)?; + + let attention_q_norm = RmsNorm::from_qtensor( + ct.tensor(reader, &format!("{prefix}.attn_q_norm.weight"), device)?, + rms_norm_eps, + )?; + + let attention_k_norm = RmsNorm::from_qtensor( + ct.tensor(reader, &format!("{prefix}.attn_k_norm.weight"), device)?, + rms_norm_eps, + )?; + + let attention_norm = RmsNorm::from_qtensor( + ct.tensor(reader, &format!("{prefix}.attn_norm.weight"), device)?, + rms_norm_eps, + )?; + + let post_attention_norm = RmsNorm::from_qtensor( + ct.tensor( + reader, + &format!("{prefix}.post_attention_norm.weight"), + device, + )?, + rms_norm_eps, + )?; + + let ffn_norm = RmsNorm::from_qtensor( + ct.tensor(reader, &format!("{prefix}.ffn_norm.weight"), device)?, + rms_norm_eps, + )?; + + let post_ffn_norm = RmsNorm::from_qtensor( + ct.tensor(reader, &format!("{prefix}.post_ffw_norm.weight"), device)?, + rms_norm_eps, + )?; + + let feed_forward_gate = + ct.tensor(reader, &format!("{prefix}.ffn_gate.weight"), device)?; + let feed_forward_up = ct.tensor(reader, &format!("{prefix}.ffn_up.weight"), device)?; + let feed_forward_down = + ct.tensor(reader, &format!("{prefix}.ffn_down.weight"), device)?; + + let mlp = Mlp { + feed_forward_gate: QMatMul::from_qtensor(feed_forward_gate)?, + feed_forward_up: QMatMul::from_qtensor(feed_forward_up)?, + feed_forward_down: QMatMul::from_qtensor(feed_forward_down)?, + }; + + // Sliding window pattern hardcoded to 6 because it's not explicitly defined + let is_sliding = (layer_idx + 1) % sliding_window_type > 0; + let sliding_window_size = is_sliding.then_some(sliding_window_size); + let layer_rope_frequency = if is_sliding { + rope_freq_base_sliding + } else { + rope_freq_base + }; + + let rotary_embedding = RotaryEmbedding::new(key_length, layer_rope_frequency, device)?; + + // Tracing spans + let span_attn = tracing::span!(tracing::Level::TRACE, "attn"); + 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_q_norm, + attention_k_norm, + attention_norm, + post_attention_norm, + ffn_norm, + post_ffn_norm, + mlp, + n_head: head_count, + n_kv_head: head_count_kv, + head_dim: key_length, + q_dim, + sliding_window_size, + rotary_embedding, + neg_inf: neg_inf.clone(), + kv_cache: None, + span_attn, + 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), + embedding_length, + layers, + norm, + output: QMatMul::from_qtensor(output)?, + span, + span_output, + }) + } + + pub fn forward(&mut self, x: &Tensor, index_pos: usize) -> Result { + let (b_sz, seq_len) = x.dims2()?; + let _enter = self.span.enter(); + + let mut layer_in = self.tok_embeddings.forward(x)?; + layer_in = (layer_in * (self.embedding_length as f64).sqrt())?; + + for layer in self.layers.iter_mut() { + let attention_mask = if seq_len == 1 { + None + } else { + Some(layer.mask(b_sz, seq_len, index_pos, x.dtype(), x.device())?) + }; + + // Attention block + let residual = &layer_in; + let x = layer.attention_norm.forward(&layer_in)?; + let x = layer.forward_attn(&x, attention_mask.as_ref(), index_pos)?; + let x = layer.post_attention_norm.forward(&x)?; + let x = (x + residual)?; + + // Feed-forward block + let _enter = layer.span_mlp.enter(); + let residual = &x; + let x = layer.ffn_norm.forward(&x)?; + let x = layer.mlp.forward(&x)?; + let x = layer.post_ffn_norm.forward(&x)?; + let x = (x + residual)?; + drop(_enter); + + layer_in = x; + } + + let _enter = self.span_output.enter(); + + let x = layer_in.i((.., seq_len - 1, ..))?; + let x = self.norm.forward(&x)?; + let output = self.output.forward(&x)?; + + Ok(output) + } +} diff --git a/patches/candle-transformers/src/models/quantized_glm4.rs b/patches/candle-transformers/src/models/quantized_glm4.rs new file mode 100644 index 0000000000..51da17f7ee --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_glm4.rs @@ -0,0 +1,477 @@ +//! GLM4 implementation with quantization support. +//! +//! Based on the GLM4 architecture and implemented with quantized weights +//! for reduced memory usage and faster inference on compatible hardware. +//! +//! References: +//! - [GLM4-0414 Models](THUDM/GLM-4-9B-0414) (architecture based on official implementations) +//! +use super::with_tracing::QMatMul; +use crate::{quantized_nn::RmsNorm, utils::repeat_kv}; +use candle::quantized::{gguf_file, QTensor}; +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::{kv_cache::KvCache, Activation, Embedding, Module}; +use std::io::{Read, Seek}; +use std::sync::Arc; + +struct Gguf { + ct: gguf_file::Content, + reader: R, + device: Device, +} + +impl Gguf { + fn new(ct: gguf_file::Content, reader: R, device: Device) -> Self { + Self { ct, reader, device } + } + + fn qmatmul(&mut self, name: &str) -> Result { + let ws = self.ct.tensor(&mut self.reader, name, &self.device)?; + QMatMul::from_weights(ws.into()) + } + + fn rms_norm(&mut self, name: &str, eps: f64) -> Result { + let ws = self.ct.tensor(&mut self.reader, name, &self.device)?; + RmsNorm::from_qtensor(ws, eps) + } + + fn metadata(&self) -> &std::collections::HashMap { + &self.ct.metadata + } + + fn tensor(&mut self, name: &str) -> Result { + self.ct.tensor(&mut self.reader, name, &self.device) + } + + fn unquantized_tensor(&mut self, name: &str, dtype: DType) -> Option { + let t = self.ct.tensor(&mut self.reader, name, &self.device); + if let Ok(t) = &t { + t.dequantize(&self.device).unwrap().to_dtype(dtype).ok() + } else { + None + } + } +} + +#[derive(Debug, Clone)] +struct Mlp { + gate_up_proj: QMatMul, + down_proj: QMatMul, + act_fn: Activation, +} + +impl Mlp { + fn new(gg: &mut Gguf, prefix: &str) -> Result { + //ffn_gate and ffn_up combined into ffn_up + let gate_up_proj = gg.qmatmul(&format!("{prefix}.ffn_up.weight"))?; + let down_proj = gg.qmatmul(&format!("{prefix}.ffn_down.weight"))?; + let act_fn = Activation::Silu; + Ok(Self { + gate_up_proj, + down_proj, + act_fn, + }) + } +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let w = self.gate_up_proj.forward(xs)?; + let dim = w.dims().len() - 1; + let gate = w + .narrow(dim, 0, w.dim(dim)? / 2)? + .contiguous()? + .apply(&self.act_fn)?; + let up_states = w + .narrow(dim, w.dim(dim)? / 2, w.dim(dim)? / 2)? + .contiguous()?; + self.down_proj.forward(&(gate * up_states)?) + } +} + +#[derive(Debug, Clone)] +pub(crate) struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, + rotary_dim: usize, +} + +impl RotaryEmbedding { + pub(crate) fn new( + dtype: DType, + head_dim: usize, + max_position_embeddings: usize, + rope_theta: f64, + partial_rotary_factor: Option, + dev: &Device, + ) -> Result { + let rotary_dim = if let Some(factor) = partial_rotary_factor { + (factor * head_dim as f32) as usize + } else { + head_dim + }; + let max_seq_len = max_position_embeddings; + let inv_freq: Vec<_> = (0..rotary_dim) + .step_by(2) + .map(|i| 1f32 / rope_theta.powf(i as f64 / rotary_dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + rotary_dim, + }) + } + + pub(crate) fn apply(&self, xs: &Tensor, offset: usize) -> Result { + let (_, _, seq_len, _) = xs.dims4()?; + let (s, e) = (offset, offset + seq_len); + let cos = self.cos.i((s..e, ..))?.contiguous()?; + let sin = self.sin.i((s..e, ..))?.contiguous()?; + let xs_rot = xs + .i((0, .., .., ..self.rotary_dim))? + .unsqueeze(0)? + .contiguous()?; + let xs_pass = xs.i((0, .., .., self.rotary_dim..))?.unsqueeze(0)?; + let xs_rot = candle_nn::rotary_emb::rope_i(&xs_rot, &cos, &sin).unwrap(); + Tensor::cat(&[&xs_rot, &xs_pass], D::Minus1)?.contiguous() + } +} + +#[derive(Debug, Clone)] +struct AttentionWeights { + q_proj: QMatMul, + k_proj: QMatMul, + v_proj: QMatMul, + o_proj: QMatMul, + attention_bq: Option, + attention_bk: Option, + attention_bv: Option, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + rotary_emb: Arc, + kv_cache: KvCache, + dtype: DType, + span_attn: tracing::Span, +} + +impl AttentionWeights { + fn new( + gg: &mut Gguf, + num_heads: usize, + num_kv_heads: usize, + head_dim: usize, + rotary_emb: Arc, + prefix: &str, + dtype: DType, + ) -> Result { + let num_kv_groups = num_heads / num_kv_heads; + + let q_proj = gg.qmatmul(&format!("{prefix}.attn_q.weight"))?; + let k_proj = gg.qmatmul(&format!("{prefix}.attn_k.weight"))?; + let v_proj = gg.qmatmul(&format!("{prefix}.attn_v.weight"))?; + let o_proj = gg.qmatmul(&format!("{prefix}.attn_output.weight"))?; + + let attention_bq = gg.unquantized_tensor(&format!("{prefix}.attn_q.bias"), DType::F32); + let attention_bk = gg.unquantized_tensor(&format!("{prefix}.attn_k.bias"), DType::F32); + let attention_bv = gg.unquantized_tensor(&format!("{prefix}.attn_v.bias"), DType::F32); + + // Initialize KV cache with 512 tokens capacity to reduce initial memory allocation. + // The cache will grow in chunks of 512 tokens when needed. + let kv_cache = KvCache::new(2, 512); + + let span_attn = tracing::span!(tracing::Level::TRACE, "attn"); + + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + attention_bq, + attention_bk, + attention_bv, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + rotary_emb, + kv_cache, + dtype, + span_attn, + }) + } + + fn forward(&mut self, x: &Tensor, attn_mask: Option<&Tensor>, offset: usize) -> Result { + let _enter = self.span_attn.enter(); + let (b, l, _) = x.dims3()?; + + let q = self.q_proj.forward(x)?; + let k = self.k_proj.forward(x)?; + let v = self.v_proj.forward(x)?; + let q = if let Some(bq) = &self.attention_bq { + q.broadcast_add(bq)? + } else { + q + }; + + let k = if let Some(bk) = &self.attention_bk { + k.broadcast_add(bk)? + } else { + k + }; + + let v = if let Some(bv) = &self.attention_bv { + v.broadcast_add(bv)? + } else { + v + }; + + let q = q + .reshape((b, l, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let k = k + .reshape((b, l, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let v = v + .reshape((b, l, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let q = self.rotary_emb.apply(&q, offset)?; + let k = self.rotary_emb.apply(&k, offset)?; + + let (q, k, v) = ( + q.to_dtype(self.dtype)?, + k.to_dtype(self.dtype)?, + v.to_dtype(self.dtype)?, + ); + // Reset KV cache if we're at the first position + if offset == 0 { + self.kv_cache.reset(); + } + + let k = k.contiguous()?; + let v = v.contiguous()?; + let (k, v) = self.kv_cache.append(&k.contiguous()?, &v.contiguous()?)?; + + let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?; + let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?; + + let scale = 1.0 / (self.head_dim as f64).sqrt(); + let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?; + if let Some(mask) = attn_mask { + scores = scores.broadcast_add(mask)?; + } + let probs = candle_nn::ops::softmax_last_dim(&scores)?; + let ctx = probs.matmul(&v)?; // (B, H, L, D) + let reshaped_ctx = ctx + .transpose(1, 2)? + .reshape((b, l, self.num_heads * self.head_dim))?; + self.o_proj.forward(&reshaped_ctx.to_dtype(x.dtype())?) + } +} + +#[derive(Debug, Clone)] +struct LayerWeights { + self_attn: AttentionWeights, + mlp: Mlp, + ffn_norm: RmsNorm, + attn_norm: RmsNorm, + post_ffw_norm: RmsNorm, + post_attention_norm: RmsNorm, +} + +impl LayerWeights { + #[allow(clippy::too_many_arguments)] + fn new( + gg: &mut Gguf, + num_attention_heads: usize, + num_key_value_heads: usize, + head_dim: usize, + rms_norm_eps: f64, + rotary: Arc, + layer_idx: usize, + dtype: DType, + ) -> Result { + let prefix = format!("blk.{layer_idx}"); + + let attn_norm = gg.rms_norm(&format!("{prefix}.attn_norm.weight"), rms_norm_eps)?; + let ffn_norm = gg.rms_norm(&format!("{prefix}.ffn_norm.weight"), rms_norm_eps)?; + + let post_ffw_norm = gg.rms_norm(&format!("{prefix}.post_ffw_norm.weight"), rms_norm_eps)?; + let post_attention_norm = gg.rms_norm( + &format!("{prefix}.post_attention_norm.weight"), + rms_norm_eps, + )?; + + let self_attn = AttentionWeights::new( + gg, + num_attention_heads, + num_key_value_heads, + head_dim, + rotary, + &prefix, + dtype, + )?; + let mlp = Mlp::new(gg, &prefix)?; + Ok(Self { + self_attn, + mlp, + attn_norm, + ffn_norm, + post_ffw_norm, + post_attention_norm, + }) + } + + fn forward(&mut self, x: &Tensor, mask: Option<&Tensor>, offset: usize) -> Result { + let residual = x; + let x = self.attn_norm.forward(x)?; + let attn = self.self_attn.forward(&x, mask, offset)?; + let attn = self.post_attention_norm.forward(&attn)?; + let x = (attn + residual)?; + + // MLP + let residual = &x; + let x = self.ffn_norm.forward(&x)?; + let x = self.mlp.forward(&x)?; + let x = self.post_ffw_norm.forward(&x)?; + x + residual + } +} + +#[derive(Debug, Clone)] +pub struct ModelWeights { + embed_tokens: Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: QMatMul, + device: Device, + dtype: DType, + span: tracing::Span, + span_output: tracing::Span, +} + +impl ModelWeights { + pub fn from_gguf( + ct: gguf_file::Content, + reader: &mut R, + device: &Device, + dtype: DType, + ) -> Result { + let mut gg = Gguf::new(ct, reader, device.clone()); + let md_get = |s: &str| match gg.metadata().get(s) { + None => candle::bail!("cannot find {s} in metadata"), + Some(v) => Ok(v), + }; + + let num_attention_heads = md_get("glm4.attention.head_count")?.to_u32()? as usize; + let num_kv_heads = md_get("glm4.attention.head_count_kv")?.to_u32()? as usize; + let head_dim = md_get("glm4.attention.key_length")?.to_u32()? as usize; + let num_layers = md_get("glm4.block_count")?.to_u32()? as usize; + let hidden_size = md_get("glm4.embedding_length")?.to_u32()? as usize; + let max_position_embeddings = md_get("glm4.context_length")?.to_u32()? as usize; + let rms_norm_eps = md_get("glm4.attention.layer_norm_rms_epsilon")?.to_f32()? as f64; + let rope_freq_base = md_get("glm4.rope.freq_base")?.to_f32()? as f64; + + let embed_tensor = gg.tensor("token_embd.weight")?; + let embed_tokens = Embedding::new(embed_tensor.dequantize(device)?, hidden_size); + + let rotary = Arc::new(RotaryEmbedding::new( + DType::F32, + head_dim, + max_position_embeddings, + rope_freq_base, + Some(0.5), //partial rotary factor not embedded in gguf + device, + )?); + + let mut layers = Vec::with_capacity(num_layers); + for i in 0..num_layers { + layers.push(LayerWeights::new( + &mut gg, + num_attention_heads, + num_kv_heads, + head_dim, + rms_norm_eps, + rotary.clone(), + i, + dtype, + )?); + } + + let norm = gg.rms_norm("output_norm.weight", rms_norm_eps)?; + // Load output projection tensor, falling back to tied embeddings like gemma3 + let lm_head_tensor = match gg.tensor("output.weight") { + Ok(tensor) => tensor, + Err(_) => gg.tensor("token_embd.weight")?, + }; + let lm_head = QMatMul::from_weights(lm_head_tensor.into())?; + let span = tracing::span!(tracing::Level::TRACE, "model"); + let span_output = tracing::span!(tracing::Level::TRACE, "output"); + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + device: device.clone(), + dtype, + span, + span_output, + }) + } + + fn causal_mask( + &self, + b: usize, + tgt: usize, + offset: usize, + sw: Option, + ) -> Result { + let minf = f32::NEG_INFINITY; + let mask: Vec<_> = (0..tgt) + .flat_map(|i| { + (0..(tgt + offset)).map(move |j| { + let past_ok = j <= i + offset; + let sw_ok = match sw { + Some(w) => (i + offset) as i64 - j as i64 <= w as i64, + None => true, + }; + if past_ok && sw_ok { + 0. + } else { + minf + } + }) + }) + .collect(); + Tensor::from_slice(&mask, (b, 1, tgt, tgt + offset), &self.device)?.to_dtype(self.dtype) + } + + pub fn forward(&mut self, input: &Tensor, offset: usize) -> Result { + let _enter = self.span.enter(); + let (b, l) = input.dims2()?; + let mut h = self.embed_tokens.forward(input)?; + + let causal_mask = if l == 1 { + None + } else { + Some(self.causal_mask(b, l, offset, None)?) + }; + + for layer in &mut self.layers { + h = layer.forward(&h, causal_mask.as_ref(), offset)?; + } + + let h = self.norm.forward(&h)?; + let _enter = self.span_output.enter(); + let last_hidden = h.narrow(1, l - 1, 1)?; + self.lm_head.forward(&last_hidden)?.squeeze(1) + } +} diff --git a/patches/candle-transformers/src/models/quantized_lfm2.rs b/patches/candle-transformers/src/models/quantized_lfm2.rs new file mode 100644 index 0000000000..d47bfd3317 --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_lfm2.rs @@ -0,0 +1,632 @@ +use crate::quantized_nn::RmsNorm; +use crate::utils::repeat_kv; +use candle::quantized::gguf_file; +use candle::quantized::QMatMul; +use candle::{bail, DType, Device, IndexOp, Result, Tensor}; +use candle_nn::{Conv1d, Conv1dConfig, Embedding, Module}; +use std::collections::HashMap; + +fn get_qtensor( + ct: &gguf_file::Content, + reader: &mut R, + device: &Device, + names: &[String], +) -> Result { + for name in names { + if let Ok(t) = ct.tensor(reader, name, device) { + return Ok(t); + } + } + bail!("cannot find tensor info for {}", names.join(" | ")) +} + +fn get_dequantized( + ct: &gguf_file::Content, + reader: &mut R, + device: &Device, + names: &[String], +) -> Result { + get_qtensor(ct, reader, device, names)?.dequantize(device) +} + +#[derive(Debug, Clone)] +struct Mlp { + w1: QMatMul, + w2: QMatMul, + w3: QMatMul, +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let w1 = self.w1.forward(xs)?; + let w3 = self.w3.forward(xs)?; + self.w2.forward(&(candle_nn::ops::silu(&w1)? * w3)?) + } +} + +#[derive(Debug, Clone)] +struct AttentionLayer { + wq: QMatMul, + wk: QMatMul, + wv: QMatMul, + wo: QMatMul, + q_norm: RmsNorm, + k_norm: RmsNorm, + n_head: usize, + n_kv_head: usize, + head_dim: usize, + cos: Tensor, + sin: Tensor, + neg_inf: Tensor, + kv_cache: Option<(Tensor, Tensor)>, + span_attn: tracing::Span, + span_rot: tracing::Span, +} + +#[derive(Debug, Clone)] +struct ShortConvLayer { + in_proj: QMatMul, + out_proj: QMatMul, + conv: Tensor, + l_cache: usize, + cache: Option, +} + +#[allow(clippy::large_enum_variant)] +#[derive(Debug, Clone)] +enum LayerKind { + Attention(AttentionLayer), + ShortConv(ShortConvLayer), +} + +#[derive(Debug, Clone)] +struct LayerWeights { + operator_norm: RmsNorm, + ffn_norm: RmsNorm, + mlp: Mlp, + kind: LayerKind, + span_mlp: tracing::Span, +} + +fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: &Tensor) -> Result { + let shape = mask.shape(); + let m = mask.where_cond(&on_true.broadcast_as(shape.dims())?, on_false)?; + Ok(m) +} + +fn precomput_freqs_cis( + head_dim: usize, + freq_base: f32, + context_length: usize, + device: &Device, +) -> 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)?; + let idx_theta = Tensor::arange(0, context_length as u32, device)? + .to_dtype(DType::F32)? + .reshape((context_length, 1))? + .matmul(&theta.reshape((1, theta.elem_count()))?)?; + let cos = idx_theta.cos()?; + let sin = idx_theta.sin()?; + Ok((cos, sin)) +} + +impl AttentionLayer { + fn apply_rotary_emb(&self, x: &Tensor, index_pos: usize) -> Result { + let _enter = self.span_rot.enter(); + let (_b, _n, seq_len, _d) = x.dims4()?; + let cos = self.cos.narrow(0, index_pos, seq_len)?; + let sin = self.sin.narrow(0, index_pos, seq_len)?; + candle_nn::rotary_emb::rope(&x.contiguous()?, &cos, &sin) + } + + fn forward(&mut self, xs: &Tensor, mask: Option<&Tensor>, index_pos: usize) -> Result { + let _enter = self.span_attn.enter(); + let (b_sz, seq_len, n_embd) = xs.dims3()?; + + let q = self.wq.forward(xs)?; + let k = self.wk.forward(xs)?; + let v = self.wv.forward(xs)?; + + 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)? + .contiguous()?; + + let q = self.q_norm.forward(&q.contiguous()?)?; + let k = self.k_norm.forward(&k.contiguous()?)?; + + 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)?; + let v = Tensor::cat(&[v_cache, &v], 2)?; + (k, v) + } + } + }; + self.kv_cache = Some((k.clone(), v.clone())); + + let k = repeat_kv(k, self.n_head / self.n_kv_head)?; + let v = repeat_kv(v, self.n_head / self.n_kv_head)?; + + let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?; + let att = match mask { + None => att, + Some(mask) => { + let mask = mask.broadcast_as(att.shape())?; + masked_fill(&att, &mask, &self.neg_inf)? + } + }; + let att = candle_nn::ops::softmax_last_dim(&att)?; + let y = att.matmul(&v.contiguous()?)?; + + let y = y.transpose(1, 2)?.reshape(&[b_sz, seq_len, n_embd])?; + self.wo.forward(&y) + } +} + +impl ShortConvLayer { + fn forward(&mut self, xs: &Tensor, _index_pos: usize) -> Result { + let (b_sz, seq_len, hidden) = xs.dims3()?; + let bcx = self.in_proj.forward(xs)?.transpose(1, 2)?; + let b = bcx.narrow(1, 0, hidden)?; + let c = bcx.narrow(1, hidden, hidden)?; + let x = bcx.narrow(1, 2 * hidden, hidden)?; + let bx = (b * &x)?.contiguous()?; + + // conv_weight shape -> [hidden, l_cache] + let mut conv_weight = self.conv.clone(); + if conv_weight.dims().len() == 3 { + conv_weight = conv_weight.squeeze(1)?; + } else if conv_weight.dims().len() == 2 && conv_weight.dims2()? == (self.l_cache, hidden) { + conv_weight = conv_weight.t()?.contiguous()?; + } + let conv_weight = conv_weight.contiguous()?; + + let mut conv_out = if seq_len == 1 { + let mut state = if let Some(cache) = &self.cache { + cache.clone() + } else { + Tensor::zeros((b_sz, hidden, self.l_cache), bx.dtype(), bx.device())? + }; + + if self.l_cache > 1 { + let tail = state.narrow(2, 1, self.l_cache - 1)?; + state = Tensor::cat(&[tail, bx.clone()], 2)?; + } else { + state = bx.clone(); + } + self.cache = Some(state.clone()); + + (state * &conv_weight.unsqueeze(0)?)? + .sum_keepdim(2)? + .contiguous()? + } else { + let conv = Conv1d::new( + conv_weight + .reshape((hidden, 1, self.l_cache))? + .contiguous()?, + None, + Conv1dConfig { + padding: self.l_cache.saturating_sub(1), + groups: hidden, + ..Default::default() + }, + ); + let mut out = conv.forward(&bx.contiguous()?)?; + out = out.narrow(2, 0, seq_len)?; + + if self.l_cache > 0 { + let (_, _, cur_len) = bx.dims3()?; + let start = cur_len.saturating_sub(self.l_cache); + let mut cache_src = bx.narrow(2, start, cur_len - start)?; + if cache_src.dims3()?.2 < self.l_cache { + let pad = self.l_cache - cache_src.dims3()?.2; + let zeros = + Tensor::zeros((b_sz, hidden, pad), cache_src.dtype(), cache_src.device())?; + cache_src = Tensor::cat(&[zeros, cache_src], 2)?; + } + self.cache = Some(cache_src); + } + + out + }; + + conv_out = (c * &conv_out)?; + let conv_out = conv_out.transpose(1, 2)?.contiguous()?; + self.out_proj.forward(&conv_out) + } +} + +pub struct ModelWeights { + tok_embeddings: Embedding, + layers: Vec, + norm: RmsNorm, + output: QMatMul, + masks: HashMap, + span: tracing::Span, + span_output: tracing::Span, +} + +fn value_to_usize(v: &gguf_file::Value) -> Result { + use gguf_file::Value::*; + match v { + U8(x) => Ok(*x as usize), + I8(x) => Ok(*x as usize), + U16(x) => Ok(*x as usize), + I16(x) => Ok(*x as usize), + U32(x) => Ok(*x as usize), + I32(x) => Ok(*x as usize), + U64(x) => Ok(*x as usize), + I64(x) => Ok(*x as usize), + F32(x) => Ok(*x as usize), + F64(x) => Ok(*x as usize), + Bool(x) => Ok(usize::from(*x)), + String(_) => bail!("unexpected string metadata"), + Array(_) => bail!("array should be handled separately"), + } +} + +fn read_usize_list(v: &gguf_file::Value, len: usize) -> Result> { + use gguf_file::Value::Array; + match v { + Array(arr) => { + let mut out = Vec::with_capacity(arr.len()); + for item in arr { + out.push(value_to_usize(item)?); + } + if out.len() == len { + Ok(out) + } else if out.len() == 1 { + Ok(vec![out[0]; len]) + } else { + bail!( + "unexpected array length in metadata, expected {len} got {}", + out.len() + ) + } + } + _ => Ok(vec![value_to_usize(v)?; len]), + } +} + +impl ModelWeights { + pub fn from_gguf( + ct: gguf_file::Content, + reader: &mut R, + device: &Device, + ) -> Result { + let md_get = |s: &str| match ct.metadata.get(s) { + None => bail!("cannot find {s} in metadata"), + Some(v) => Ok(v), + }; + + let head_count = md_get("lfm2.attention.head_count")?.to_u32()? as usize; + let head_count_kv_meta = md_get("lfm2.attention.head_count_kv")?; + let embedding_length = md_get("lfm2.embedding_length")?.to_u32()? as usize; + let context_length = md_get("lfm2.context_length")?.to_u32()? as usize; + let block_count = md_get("lfm2.block_count")?.to_u32()? as usize; + let rms_norm_eps = md_get("lfm2.attention.layer_norm_rms_epsilon")?.to_f32()? as f64; + let rope_freq_base = md_get("lfm2.rope.freq_base") + .and_then(|m| m.to_f32()) + .unwrap_or(1_000_000f32); + let l_cache = md_get("lfm2.shortconv.l_cache")?.to_u32()? as usize; + + let head_count_kv = read_usize_list(head_count_kv_meta, block_count)?; + let head_dim = embedding_length / head_count; + let (cos, sin) = precomput_freqs_cis(head_dim, rope_freq_base, context_length, device)?; + let neg_inf = Tensor::new(f32::NEG_INFINITY, device)?; + + let tok_embeddings_q = get_qtensor( + &ct, + reader, + device, + &[ + "token_embd.weight", + "tok_embeddings.weight", + "model.embed_tokens.weight", + ] + .iter() + .map(|s| s.to_string()) + .collect::>(), + )?; + let tok_embeddings = tok_embeddings_q.dequantize(device)?; + tracing::debug!( + tok_embd_shape = ?tok_embeddings.shape().dims(), + "loaded lfm2 token embeddings" + ); + + let norm = RmsNorm::from_qtensor( + get_qtensor( + &ct, + reader, + device, + &[ + "output_norm.weight", + "embedding_norm.weight", + "model.embedding_norm.weight", + "model.embedding_norm", + "token_embd_norm.weight", + ] + .iter() + .map(|s| s.to_string()) + .collect::>(), + )?, + rms_norm_eps, + )?; + let output_q = get_qtensor( + &ct, + reader, + device, + &[ + "output.weight", + "lm_head.weight", + "model.output.weight", + "model.lm_head.weight", + ] + .iter() + .map(|s| s.to_string()) + .collect::>(), + ) + .unwrap_or(tok_embeddings_q); + tracing::debug!( + output_shape = ?output_q.shape().dims(), + "loaded lfm2 output weight (using tok_embd if missing)" + ); + + let mut layers = Vec::with_capacity(block_count); + for layer_idx in 0..block_count { + let prefix = format!("blk.{layer_idx}"); + let is_attention = head_count_kv.get(layer_idx).copied().unwrap_or(head_count) > 0; + + let operator_norm = get_qtensor( + &ct, + reader, + device, + &[ + format!("{prefix}.attn_norm.weight"), + format!("{prefix}.operator_norm.weight"), + format!("{prefix}.attention_norm.weight"), + ], + )?; + let ffn_norm = get_qtensor( + &ct, + reader, + device, + &[ + format!("{prefix}.ffn_norm.weight"), + format!("{prefix}.ffn_norm"), + ], + )?; + let mlp = { + let w1 = get_qtensor( + &ct, + reader, + device, + &[ + format!("{prefix}.ffn_gate.weight"), + format!("{prefix}.feed_forward.w1.weight"), + format!("{prefix}.mlp.gate_proj.weight"), + ], + )?; + let w2 = get_qtensor( + &ct, + reader, + device, + &[ + format!("{prefix}.ffn_down.weight"), + format!("{prefix}.feed_forward.w2.weight"), + format!("{prefix}.mlp.down_proj.weight"), + ], + )?; + let w3 = get_qtensor( + &ct, + reader, + device, + &[ + format!("{prefix}.ffn_up.weight"), + format!("{prefix}.feed_forward.w3.weight"), + format!("{prefix}.mlp.up_proj.weight"), + ], + )?; + Mlp { + w1: QMatMul::from_qtensor(w1)?, + w2: QMatMul::from_qtensor(w2)?, + w3: QMatMul::from_qtensor(w3)?, + } + }; + + let kind = if is_attention { + let n_kv_head = head_count_kv[layer_idx]; + let wq = get_qtensor( + &ct, + reader, + device, + &[ + format!("{prefix}.attn_q.weight"), + format!("{prefix}.self_attn.q_proj.weight"), + ], + )?; + let wk = get_qtensor( + &ct, + reader, + device, + &[ + format!("{prefix}.attn_k.weight"), + format!("{prefix}.self_attn.k_proj.weight"), + ], + )?; + let wv = get_qtensor( + &ct, + reader, + device, + &[ + format!("{prefix}.attn_v.weight"), + format!("{prefix}.self_attn.v_proj.weight"), + ], + )?; + let wo = get_qtensor( + &ct, + reader, + device, + &[ + format!("{prefix}.attn_output.weight"), + format!("{prefix}.self_attn.out_proj.weight"), + ], + )?; + let q_norm = get_qtensor( + &ct, + reader, + device, + &[ + format!("{prefix}.attn_q_norm.weight"), + format!("{prefix}.self_attn.q_layernorm.weight"), + format!("{prefix}.attention.q_norm.weight"), + ], + )?; + let k_norm = get_qtensor( + &ct, + reader, + device, + &[ + format!("{prefix}.attn_k_norm.weight"), + format!("{prefix}.self_attn.k_layernorm.weight"), + format!("{prefix}.attention.k_norm.weight"), + ], + )?; + + LayerKind::Attention(AttentionLayer { + wq: QMatMul::from_qtensor(wq)?, + wk: QMatMul::from_qtensor(wk)?, + wv: QMatMul::from_qtensor(wv)?, + wo: QMatMul::from_qtensor(wo)?, + q_norm: RmsNorm::from_qtensor(q_norm, rms_norm_eps)?, + k_norm: RmsNorm::from_qtensor(k_norm, rms_norm_eps)?, + n_head: head_count, + n_kv_head, + head_dim, + cos: cos.clone(), + sin: sin.clone(), + neg_inf: neg_inf.clone(), + kv_cache: None, + span_attn: tracing::span!(tracing::Level::TRACE, "attn"), + span_rot: tracing::span!(tracing::Level::TRACE, "attn-rot"), + }) + } else { + let in_proj = get_qtensor( + &ct, + reader, + device, + &[ + format!("{prefix}.shortconv.in_proj.weight"), + format!("{prefix}.conv.in_proj.weight"), + ], + )?; + let out_proj = get_qtensor( + &ct, + reader, + device, + &[ + format!("{prefix}.shortconv.out_proj.weight"), + format!("{prefix}.conv.out_proj.weight"), + ], + )?; + let conv = get_dequantized( + &ct, + reader, + device, + &[ + format!("{prefix}.shortconv.conv.weight"), + format!("{prefix}.conv.conv.weight"), + format!("{prefix}.shortconv.conv"), + ], + )?; + LayerKind::ShortConv(ShortConvLayer { + in_proj: QMatMul::from_qtensor(in_proj)?, + out_proj: QMatMul::from_qtensor(out_proj)?, + conv, + l_cache, + cache: None, + }) + }; + + layers.push(LayerWeights { + operator_norm: RmsNorm::from_qtensor(operator_norm, rms_norm_eps)?, + ffn_norm: RmsNorm::from_qtensor(ffn_norm, rms_norm_eps)?, + mlp, + kind, + span_mlp: tracing::span!(tracing::Level::TRACE, "ffn"), + }); + } + + Ok(Self { + tok_embeddings: Embedding::new(tok_embeddings, embedding_length), + layers, + norm, + output: QMatMul::from_qtensor(output_q)?, + masks: HashMap::new(), + span: tracing::span!(tracing::Level::TRACE, "model"), + span_output: tracing::span!(tracing::Level::TRACE, "output"), + }) + } + + fn mask(&mut self, t: usize, device: &Device) -> 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)?; + 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 = if seq_len == 1 { + None + } else { + Some(self.mask(seq_len, x.device())?) + }; + + let _enter = self.span.enter(); + let mut hidden = self.tok_embeddings.forward(x)?; + for layer in self.layers.iter_mut() { + let residual = hidden.clone(); + let normed = layer.operator_norm.forward(&hidden)?; + hidden = match &mut layer.kind { + LayerKind::Attention(attn) => attn.forward(&normed, mask.as_ref(), index_pos)?, + LayerKind::ShortConv(conv) => conv.forward(&normed, index_pos)?, + }; + hidden = (hidden + residual)?; + + let residual = hidden.clone(); + let ff = layer.ffn_norm.forward(&hidden)?; + let _enter = layer.span_mlp.enter(); + let ff = layer.mlp.forward(&ff)?; + hidden = (ff + residual)?; + } + let hidden = self.norm.forward(&hidden)?; + let hidden = hidden.i((.., seq_len - 1, ..))?; + let _enter = self.span_output.enter(); + self.output.forward(&hidden) + } +} diff --git a/patches/candle-transformers/src/models/quantized_llama.rs b/patches/candle-transformers/src/models/quantized_llama.rs new file mode 100644 index 0000000000..5b78b8c748 --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_llama.rs @@ -0,0 +1,793 @@ +//! Quantized llama model implementation. +//! +//! This provides a quantized implementation of the llama language model architecture. +//! The model implements parameter efficient quantization for reduced memory usage +//! while maintaining model quality. +//! +//! Key characteristics: +//! - Transformer decoder architecture +//! - Support for 2/3/4/8-bit quantization +//! - Optimized memory usage through quantization +//! - Configurable model sizes and parameter counts +//! +//! - 💻 [GH Link](https://github.com/facebookresearch/llama) +//! - 📝 [Paper](https://arxiv.org/abs/2302.13971) +//! +//! ![](https://raw.githubusercontent.com/huggingface/candle/main/candle-examples/examples/quantized/assets/aoc.gif) +//! + +use std::collections::HashMap; + +use crate::quantized_nn::RmsNorm; +use candle::quantized::QTensor; +use candle::quantized::{ggml_file, gguf_file}; +use candle::{DType, Device, IndexOp, Result, Tensor}; +use candle_nn::{Embedding, Module}; + +pub const MAX_SEQ_LEN: usize = 4096; + +// 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, + neg_inf: 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: &Tensor) -> Result { + let shape = mask.shape(); + let m = mask.where_cond(&on_true.broadcast_as(shape.dims())?, on_false)?; + Ok(m) +} + +/// Causal mask for `t` new query positions attending over `index_pos + t` +/// total key/value positions (the already-cached `index_pos` positions plus +/// the `t` new ones being fed this call). Query row `i` may attend to key +/// column `j` iff `j <= index_pos + i`. +/// +/// `ModelWeights::mask` (used by `forward`) only ever builds a square `[t, +/// t]` mask, which is correct when the cache is empty (`index_pos == 0`, +/// so `index_pos + t == t`) but wrong for a multi-token forward pass +/// continuing an existing cache — attention scores there have shape +/// `[b, heads, t, index_pos + t]` and a `[t, t]` mask fails to broadcast. +/// This case never arose before `forward_all_positions` existed, since +/// every other caller only ever fed a single new token (`t == 1`, masked +/// with `None`) once past the initial prefill. +fn causal_mask_with_offset(t: usize, index_pos: usize, device: &Device) -> Result { + let total = index_pos + t; + let mask: Vec = (0..t) + .flat_map(|i| (0..total).map(move |j| u8::from(j > index_pos + i))) + .collect(); + Tensor::from_slice(&mask, (t, total), device) +} + +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)?; + let sin = self.sin.narrow(0, index_pos, seq_len)?; + // The call to contiguous below is only necessary when processing the prompt. + // When the seq_len is 1 in the inference loop, this is a no-op. + candle_nn::rotary_emb::rope_i(&x.contiguous()?, &cos, &sin) + } + + fn forward_attn( + &mut self, + x: &Tensor, + mask: Option<&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)? + // This call to contiguous ensures that the fast kernel can be called below. It's + // actually a no-op except when processing the initial prompt so has no significant + // impact on performance. + .contiguous()?; + + 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)?; + let v = Tensor::cat(&[v_cache, &v], 2)?; + (k, v) + } + } + }; + self.kv_cache = Some((k.clone(), v.clone())); + + let y = if q.device().is_metal() && seq_len == 1 { + // SDPA will do MQA for us + candle_nn::ops::sdpa( + &q, + &k, + &v, + None, + false, + 1. / (self.head_dim as f32).sqrt(), + 1., + )? + } else { + // Support for MQA, useful for 70B models and mistral. + let k = crate::utils::repeat_kv(k, self.n_head / self.n_kv_head)?; + let v = crate::utils::repeat_kv(v, self.n_head / self.n_kv_head)?; + + let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?; + let att = match mask { + None => att, + Some(mask) => { + let mask = mask.broadcast_as(att.shape())?; + masked_fill(&att, &mask, &self.neg_inf)? + } + }; + let att = candle_nn::ops::softmax_last_dim(&att)?; + // Convert to contiguous as matmul doesn't support strided vs for now. + 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) + } + + /// Same as `forward_attn`, but routes the attention-score computation + /// through `ruvllm_sparse_attention::SubquadraticSparseAttention` + /// instead of dense causal attention, when doing so is valid. + /// + /// `SubquadraticSparseAttention::forward_gqa` requires `q.seq == k.seq + /// == v.seq` (it derives its own causal structure from that shared + /// length) — that only holds for a fresh-cache prompt prefill + /// (`index_pos == 0`), not for continuing an existing KV cache, where + /// queries are new tokens but keys/values also cover cached history. + /// Prefill is also where sparse attention's O(N log N)-vs-O(N^2) win + /// actually matters (long prompts), so this falls back to the ordinary + /// dense/SDPA path from `forward_attn` for every other call. + #[cfg(feature = "sparse-attention")] + fn forward_attn_sparse( + &mut self, + x: &Tensor, + mask: Option<&Tensor>, + index_pos: usize, + sparse: &ruvllm_sparse_attention::SubquadraticSparseAttention, + ) -> 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)? + .contiguous()?; + + 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)?; + let v = Tensor::cat(&[v_cache, &v], 2)?; + (k, v) + } + } + }; + self.kv_cache = Some((k.clone(), v.clone())); + + let y = if index_pos == 0 && b_sz == 1 { + let q3 = tensor_to_tensor3(&q)?; + let k3 = tensor_to_tensor3(&k)?; + let v3 = tensor_to_tensor3(&v)?; + let out3 = sparse + .forward_gqa(&q3, &k3, &v3) + .map_err(|e| candle::Error::Msg(format!("sparse attention failed: {e:?}")))?; + tensor3_to_tensor(&out3, x.device())? + } else if q.device().is_metal() && seq_len == 1 { + candle_nn::ops::sdpa( + &q, + &k, + &v, + None, + false, + 1. / (self.head_dim as f32).sqrt(), + 1., + )? + } else { + let k = crate::utils::repeat_kv(k, self.n_head / self.n_kv_head)?; + let v = crate::utils::repeat_kv(v, self.n_head / self.n_kv_head)?; + + let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?; + let att = match mask { + None => att, + Some(mask) => { + let mask = mask.broadcast_as(att.shape())?; + masked_fill(&att, &mask, &self.neg_inf)? + } + }; + let att = candle_nn::ops::softmax_last_dim(&att)?; + 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) + } +} + +/// Bridge a candle `[1, heads, seq, dim]` tensor to `ruvllm_sparse_attention`'s +/// flat `[seq, heads, dim]`-layout `Tensor3`. +#[cfg(feature = "sparse-attention")] +fn tensor_to_tensor3(t: &Tensor) -> Result { + let (_b, heads, seq, dim) = t.dims4()?; + let t = t.transpose(1, 2)?.contiguous()?.to_dtype(DType::F32)?; + let data: Vec = t.flatten_all()?.to_vec1()?; + Ok(ruvllm_sparse_attention::Tensor3 { + data, + seq, + heads, + dim, + }) +} + +/// Inverse of `tensor_to_tensor3`: back to candle's `[1, heads, seq, dim]`. +#[cfg(feature = "sparse-attention")] +fn tensor3_to_tensor(t3: &ruvllm_sparse_attention::Tensor3, device: &Device) -> Result { + let t = Tensor::from_slice(t3.data.as_slice(), (1, t3.seq, t3.heads, t3.dim), device)?; + t.transpose(1, 2)?.contiguous() +} + +#[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, + device: &Device, +) -> 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)?; + let idx_theta = Tensor::arange(0, MAX_SEQ_LEN as u32, device)? + .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 head_dim = (ct.hparams.n_embd / ct.hparams.n_head) as usize; + let (cos, sin) = precomput_freqs_cis(head_dim, 10000., &ct.device)?; + let neg_inf = Tensor::new(f32::NEG_INFINITY, &ct.device)?; + let tok_embeddings = ct.remove("tok_embeddings.weight")?; + let tok_embeddings = tok_embeddings.dequantize(&ct.device)?; + let norm = RmsNorm::from_qtensor(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::from_qtensor(attention_norm, 1e-5)?, + mlp_or_moe, + ffn_norm: RmsNorm::from_qtensor(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(), + neg_inf: neg_inf.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, + device: &Device, + ) -> Result { + 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()? as f64; + + 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, device)?; + let neg_inf = Tensor::new(f32::NEG_INFINITY, device)?; + + let tok_embeddings_q = ct.tensor(reader, "token_embd.weight", device)?; + let tok_embeddings = tok_embeddings_q.dequantize(device)?; + let norm = RmsNorm::from_qtensor( + ct.tensor(reader, "output_norm.weight", device)?, + rms_norm_eps, + )?; + let output = match ct.tensor(reader, "output.weight", device) { + Ok(tensor) => tensor, + Err(_) => tok_embeddings_q, + }; + 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"), device)?; + let attention_wk = ct.tensor(reader, &format!("{prefix}.attn_k.weight"), device)?; + let attention_wv = ct.tensor(reader, &format!("{prefix}.attn_v.weight"), device)?; + let attention_wo = + ct.tensor(reader, &format!("{prefix}.attn_output.weight"), device)?; + let mlp_or_moe = if n_expert <= 1 { + let feed_forward_w1 = + ct.tensor(reader, &format!("{prefix}.ffn_gate.weight"), device)?; + let feed_forward_w2 = + ct.tensor(reader, &format!("{prefix}.ffn_down.weight"), device)?; + let feed_forward_w3 = + ct.tensor(reader, &format!("{prefix}.ffn_up.weight"), device)?; + 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"), device)?; + 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"), device)?; + let feed_forward_w2 = + ct.tensor(reader, &format!("{prefix}.ffn_down.{i}.weight"), device)?; + let feed_forward_w3 = + ct.tensor(reader, &format!("{prefix}.ffn_up.{i}.weight"), device)?; + 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"), device)?; + let ffn_norm = ct.tensor(reader, &format!("{prefix}.ffn_norm.weight"), device)?; + 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::from_qtensor(attention_norm, rms_norm_eps)?, + mlp_or_moe, + ffn_norm: RmsNorm::from_qtensor(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(), + neg_inf: neg_inf.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, device: &Device) -> 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)?; + 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 = if seq_len == 1 { + None + } else { + Some(self.mask(seq_len, x.device())?) + }; + 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.as_ref(), 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) + } + + /// Like `forward`, but returns logits for every input position instead of + /// only the last one. Needed by callers that verify multiple newly-fed + /// tokens against this model's own predictions in a single forward pass + /// (e.g. speculative decoding), where a per-position distribution is + /// required to check each candidate token, not just a single next-token + /// prediction. Shares the entire layer loop with `forward`; only the + /// final slice-to-last-position step is skipped. + pub fn forward_all_positions(&mut self, x: &Tensor, index_pos: usize) -> Result { + let (_b_sz, seq_len) = x.dims2()?; + let mask = if seq_len == 1 { + None + } else { + Some(causal_mask_with_offset(seq_len, index_pos, x.device())?) + }; + 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.as_ref(), 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 _enter = self.span_output.enter(); + self.output.forward(&x) + } + + /// Like `forward`, but routes each layer's attention through + /// `ruvllm_sparse_attention::SubquadraticSparseAttention` for prompt + /// prefill (falling back to dense attention for incremental decode + /// steps — see `LayerWeights::forward_attn_sparse` for why). This is + /// where sparse attention's subquadratic scaling actually pays off: + /// the prefill's attention cost grows with prompt length squared under + /// dense attention, vs. `O(n log n)` (or better, with landmarks/gating) + /// here. + #[cfg(feature = "sparse-attention")] + pub fn forward_sparse( + &mut self, + x: &Tensor, + index_pos: usize, + sparse: &ruvllm_sparse_attention::SubquadraticSparseAttention, + ) -> Result { + let (_b_sz, seq_len) = x.dims2()?; + let mask = if seq_len == 1 { + None + } else { + Some(causal_mask_with_offset(seq_len, index_pos, x.device())?) + }; + 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_sparse(&x, mask.as_ref(), index_pos, sparse)?; + let x = (attn + residual)?; + + 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) + } + + /// `forward_all_positions` + `forward_sparse` combined: per-position + /// logits (see `forward_all_positions`) with prefill attention routed + /// through sparse attention (see `forward_sparse`). Used by speculative + /// decoding's prompt prefill. + #[cfg(feature = "sparse-attention")] + pub fn forward_all_positions_sparse( + &mut self, + x: &Tensor, + index_pos: usize, + sparse: &ruvllm_sparse_attention::SubquadraticSparseAttention, + ) -> Result { + let (_b_sz, seq_len) = x.dims2()?; + let mask = if seq_len == 1 { + None + } else { + Some(causal_mask_with_offset(seq_len, index_pos, x.device())?) + }; + 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_sparse(&x, mask.as_ref(), index_pos, sparse)?; + let x = (attn + residual)?; + + 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 _enter = self.span_output.enter(); + self.output.forward(&x) + } + + /// Snapshot every layer's KV cache. `Tensor` clones are cheap (shared + /// storage, not a data copy), so this is O(num_layers), not O(context + /// length) — callers that need to walk back a rejected speculative + /// decoding draft can restore an earlier snapshot instead of resetting + /// to position 0 and replaying the entire context. + pub fn snapshot_kv_cache(&self) -> Vec> { + self.layers.iter().map(|l| l.kv_cache.clone()).collect() + } + + /// Restore every layer's KV cache from a snapshot previously returned + /// by `snapshot_kv_cache`. The caller is responsible for also + /// restoring whatever `index_pos` it was passing to `forward`/ + /// `forward_all_positions` alongside this snapshot — the cache + /// tensors alone don't encode position. + pub fn restore_kv_cache(&mut self, snapshot: &[Option<(Tensor, Tensor)>]) { + for (layer, cached) in self.layers.iter_mut().zip(snapshot.iter()) { + layer.kv_cache = cached.clone(); + } + } +} diff --git a/patches/candle-transformers/src/models/quantized_llama2_c.rs b/patches/candle-transformers/src/models/quantized_llama2_c.rs new file mode 100644 index 0000000000..3eb14bb9e6 --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_llama2_c.rs @@ -0,0 +1,260 @@ +//! Quantized Llama2 model implementation. +//! +//! This provides an 8-bit quantized implementation of Meta's LLaMA2 language model +//! for reduced memory usage and faster inference. +//! +//! Key characteristics: +//! - Decoder-only transformer architecture +//! - RoPE position embeddings +//! - Grouped Query Attention +//! - 8-bit quantization of weights +//! +//! References: +//! - [LLaMA2 Paper](https://arxiv.org/abs/2307.09288) +//! - [LLaMA2 Technical Report](https://ai.meta.com/research/publications/llama-2-open-foundation-and-fine-tuned-chat-models/) +//! + +use super::llama2_c::{Cache, Config}; +use crate::quantized_nn::{linear_no_bias as linear, Embedding, Linear, RmsNorm}; +pub use crate::quantized_var_builder::VarBuilder; +use candle::{DType, IndexOp, Module, Result, Tensor, D}; + +fn silu(xs: &Tensor) -> Result { + xs / (xs.neg()?.exp()? + 1.0)? +} + +#[derive(Debug, Clone)] +struct CausalSelfAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + n_head: usize, + n_key_value_head: usize, + head_dim: usize, +} + +impl CausalSelfAttention { + fn apply_rotary_emb(&self, x: &Tensor, index_pos: usize, cache: &Cache) -> Result { + let (b_sz, seq_len, h, n_embd) = x.dims4()?; + let cos = cache.cos.i(index_pos..index_pos + seq_len)?; + let sin = cache.sin.i(index_pos..index_pos + seq_len)?; + let cos = cos.unsqueeze(1)?; + let sin = sin.unsqueeze(1)?; + let cos = cos.broadcast_as((b_sz, seq_len, 1, n_embd / 2, 1))?; + let sin = sin.broadcast_as((b_sz, seq_len, 1, n_embd / 2, 1))?; + let x = x.reshape((b_sz, seq_len, h, n_embd / 2, 2))?; + let x0 = x.narrow(D::Minus1, 0, 1)?; + let x1 = x.narrow(D::Minus1, 1, 1)?; + let dst0 = (x0.broadcast_mul(&cos)? - x1.broadcast_mul(&sin)?)?; + let dst1 = (x0.broadcast_mul(&sin)? + x1.broadcast_mul(&cos)?)?; + let rope = Tensor::cat(&[&dst0, &dst1], D::Minus1)?.reshape((b_sz, seq_len, h, n_embd))?; + Ok(rope) + } + + fn forward( + &self, + x: &Tensor, + index_pos: usize, + block_idx: usize, + cache: &mut Cache, + ) -> Result { + let (b_sz, seq_len, n_embd) = x.dims3()?; + let q = self.q_proj.forward(x)?; + let k = self.k_proj.forward(x)?; + let v = self.v_proj.forward(x)?; + + let q = q.reshape((b_sz, seq_len, self.n_head, self.head_dim))?; + let k = k.reshape((b_sz, seq_len, self.n_key_value_head, self.head_dim))?; + let mut v = v.reshape((b_sz, seq_len, self.n_key_value_head, self.head_dim))?; + + let q = self.apply_rotary_emb(&q, index_pos, cache)?; + let mut k = self.apply_rotary_emb(&k, index_pos, cache)?; + + if cache.use_kv_cache { + if let Some((cache_k, cache_v)) = &cache.kvs[block_idx] { + k = Tensor::cat(&[cache_k, &k], 1)?.contiguous()?; + v = Tensor::cat(&[cache_v, &v], 1)?.contiguous()?; + } + cache.kvs[block_idx] = Some((k.clone(), v.clone())) + } + + let k = self.repeat_kv(k)?; + let v = self.repeat_kv(v)?; + + let q = q.transpose(1, 2)?.contiguous()?; + let k = k.transpose(1, 2)?.contiguous()?; + let v = v.transpose(1, 2)?.contiguous()?; + + let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?; + let att = if seq_len <= 1 { + att + } else { + let mask = cache.mask(seq_len)?.broadcast_as(att.shape())?; + masked_fill(&att, &mask, f32::NEG_INFINITY)? + }; + let att = candle_nn::ops::softmax(&att, D::Minus1)?; + // 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.o_proj.forward(&y)?; + Ok(y) + } + + fn repeat_kv(&self, x: Tensor) -> Result { + let n_rep = self.n_head / self.n_key_value_head; + if n_rep == 1 { + Ok(x) + } else { + let (b_sz, seq_len, n_kv_head, head_dim) = x.dims4()?; + let x = x + .unsqueeze(3)? + .expand((b_sz, seq_len, n_kv_head, n_rep, head_dim))? + .reshape((b_sz, seq_len, n_kv_head * n_rep, head_dim))?; + Ok(x) + } + } + + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let size_in = cfg.dim; + let size_q = (cfg.dim / cfg.n_heads) * cfg.n_heads; + let size_kv = (cfg.dim / cfg.n_heads) * cfg.n_kv_heads; + let q_proj = linear(size_in, size_q, vb.pp("q_proj"))?; + let k_proj = linear(size_in, size_kv, vb.pp("k_proj"))?; + let v_proj = linear(size_in, size_kv, vb.pp("v_proj"))?; + let o_proj = linear(size_q, size_in, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + n_head: cfg.n_heads, + n_key_value_head: cfg.n_kv_heads, + head_dim: cfg.dim / cfg.n_heads, + }) + } +} + +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) +} + +#[derive(Debug, Clone)] +struct Mlp { + c_fc1: Linear, + c_fc2: Linear, + c_proj: Linear, +} + +impl Mlp { + fn new(c_fc1: Linear, c_fc2: Linear, c_proj: Linear) -> Self { + Self { + c_fc1, + c_fc2, + c_proj, + } + } + + fn forward(&self, x: &Tensor) -> Result { + let x = (silu(&self.c_fc1.forward(x)?)? * self.c_fc2.forward(x)?)?; + self.c_proj.forward(&x) + } + + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let h_size = cfg.dim; + let i_size = cfg.hidden_dim; + let c_fc1 = linear(h_size, i_size, vb.pp("gate_proj"))?; + let c_fc2 = linear(h_size, i_size, vb.pp("up_proj"))?; + let c_proj = linear(i_size, h_size, vb.pp("down_proj"))?; + Ok(Self::new(c_fc1, c_fc2, c_proj)) + } +} + +#[derive(Debug, Clone)] +struct Block { + rms_1: RmsNorm, + attn: CausalSelfAttention, + rms_2: RmsNorm, + mlp: Mlp, +} + +impl Block { + fn new(rms_1: RmsNorm, attn: CausalSelfAttention, rms_2: RmsNorm, mlp: Mlp) -> Self { + Self { + rms_1, + attn, + rms_2, + mlp, + } + } + + fn forward( + &self, + x: &Tensor, + index_pos: usize, + block_idx: usize, + cache: &mut Cache, + ) -> Result { + let residual = x; + let x = self.rms_1.forward(x)?; + let x = (self.attn.forward(&x, index_pos, block_idx, cache)? + residual)?; + let residual = &x; + let x = (self.mlp.forward(&self.rms_2.forward(&x)?)? + residual)?; + Ok(x) + } + + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let attn = CausalSelfAttention::load(vb.pp("self_attn"), cfg)?; + let mlp = Mlp::load(vb.pp("mlp"), cfg)?; + let input_layernorm = RmsNorm::new(cfg.dim, cfg.norm_eps, vb.pp("input_layernorm"))?; + let post_attention_layernorm = + RmsNorm::new(cfg.dim, cfg.norm_eps, vb.pp("post_attention_layernorm"))?; + Ok(Self::new( + input_layernorm, + attn, + post_attention_layernorm, + mlp, + )) + } +} + +#[derive(Debug, Clone)] +pub struct QLlama { + wte: Embedding, + blocks: Vec, + ln_f: RmsNorm, + lm_head: Linear, + pub config: Config, +} + +impl QLlama { + pub fn forward(&self, x: &Tensor, index_pos: usize, cache: &mut Cache) -> Result { + let (_b_sz, _seq_len) = x.dims2()?; + let mut x = self.wte.forward(x)?; + for (block_idx, block) in self.blocks.iter().enumerate() { + x = block.forward(&x, index_pos, block_idx, cache)?; + } + let x = self.ln_f.forward(&x)?; + let logits = self.lm_head.forward(&x)?; + logits.to_dtype(DType::F32) + } + + pub fn load(vb: VarBuilder, cfg: Config) -> Result { + let wte = Embedding::new(cfg.vocab_size, cfg.dim, vb.pp("model.embed_tokens"))?; + let lm_head = linear(cfg.dim, cfg.vocab_size, vb.pp("lm_head"))?; + let ln_f = RmsNorm::new(cfg.dim, cfg.norm_eps, vb.pp("model.norm"))?; + let blocks: Vec<_> = (0..cfg.n_layers) + .map(|i| Block::load(vb.pp(format!("model.layers.{i}")), &cfg).unwrap()) + .collect(); + Ok(Self { + wte, + blocks, + ln_f, + lm_head, + config: cfg, + }) + } +} diff --git a/patches/candle-transformers/src/models/quantized_metavoice.rs b/patches/candle-transformers/src/models/quantized_metavoice.rs new file mode 100644 index 0000000000..ac72162715 --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_metavoice.rs @@ -0,0 +1,259 @@ +//! Quantized MetaVoice model implementation. +//! +//! MetaVoice is a conditional text-to-speech model based on a transformer architecture. +//! This implementation provides quantization for reduced memory and compute. +//! +//! Key characteristics: +//! - Transformer-based autoregressive decoder +//! - Speaker conditioning +//! - Support for 8-bit quantization +//! - Key-value caching for efficient inference +//! - RMS normalization layers +//! +//! References: +//! - [MetaVoice Code](https://github.com/metavoiceio/metavoice) +//! + +use crate::quantized_nn::{linear_b, Embedding, Linear, RmsNorm}; +pub use crate::quantized_var_builder::VarBuilder; + +use crate::models::metavoice::repeat_interleave; +use candle::{Module, Result, Tensor, D}; + +pub mod transformer { + use super::*; + + type Config = crate::models::metavoice::transformer::Config; + + #[derive(Debug, Clone)] + struct FeedForward { + w1: Linear, + w2: Linear, + w3: Linear, + span: tracing::Span, + } + + impl FeedForward { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let i_size = cfg.intermediate_size(); + let w1 = linear_b(cfg.dim, i_size, false, vb.pp("swiglu.w1"))?; + let w2 = linear_b(i_size, cfg.dim, false, vb.pp("w2"))?; + let w3 = linear_b(cfg.dim, i_size, false, vb.pp("swiglu.w3"))?; + Ok(Self { + w1, + w2, + w3, + span: tracing::span!(tracing::Level::TRACE, "feed-forward"), + }) + } + } + + impl Module for FeedForward { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let swiglu = (candle_nn::ops::silu(&xs.apply(&self.w1)?)? * xs.apply(&self.w3))?; + swiglu.apply(&self.w2) + } + } + + #[derive(Debug, Clone)] + struct Attention { + wqkv: Linear, + wo: Linear, + dim: usize, + kv_size: usize, + n_local_heads: usize, + head_dim: usize, + n_head: usize, + kv_cache: Option<(Tensor, Tensor)>, + span: tracing::Span, + } + + impl Attention { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let n_local_heads = cfg.n_local_heads(); + let head_dim = cfg.head_dim(); + let total_head_dim = (cfg.n_head + 2 * n_local_heads) * head_dim; + let wqkv = linear_b(cfg.dim, total_head_dim, false, vb.pp("wqkv"))?; + let wo = linear_b(cfg.dim, cfg.dim, false, vb.pp("wo"))?; + Ok(Self { + wqkv, + wo, + dim: cfg.dim, + kv_size: n_local_heads * head_dim, + n_local_heads, + head_dim, + n_head: cfg.n_head, + kv_cache: None, + span: tracing::span!(tracing::Level::TRACE, "attention"), + }) + } + + fn forward(&mut self, xs: &Tensor, _pos: usize, mask: &Tensor) -> Result { + let _enter = self.span.enter(); + let (b_sz, seqlen, _) = xs.dims3()?; + + let qkv = xs.apply(&self.wqkv)?; + let q = qkv.narrow(D::Minus1, 0, self.dim)?; + let k = qkv.narrow(D::Minus1, self.dim, self.kv_size)?; + let v = qkv.narrow(D::Minus1, self.dim + self.kv_size, self.kv_size)?; + let q = q + .reshape((b_sz, seqlen, self.n_head, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let k = k + .reshape((b_sz, seqlen, self.n_local_heads, self.head_dim))? + .transpose(1, 2)?; + let v = v + .reshape((b_sz, seqlen, self.n_local_heads, self.head_dim))? + .transpose(1, 2)?; + + let (k, v) = match &self.kv_cache { + None => (k, v), + Some((prev_k, prev_v)) => { + let k = Tensor::cat(&[prev_k, &k], 2)?; + let v = Tensor::cat(&[prev_v, &v], 2)?; + (k, v) + } + }; + self.kv_cache = Some((k.clone(), v.clone())); + + let k = repeat_interleave(&k, self.n_head / self.n_local_heads, 1)?; + let v = repeat_interleave(&v, self.n_head / self.n_local_heads, 1)?; + + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (q.matmul(&k.transpose(2, 3)?)? * scale)?; + + let attn_weights = attn_weights.broadcast_add(mask)?; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + let attn_output = attn_weights.matmul(&v)?; + attn_output + .transpose(1, 2)? + .reshape((b_sz, seqlen, self.dim))? + .apply(&self.wo) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } + } + + #[derive(Debug, Clone)] + struct Block { + attention: Attention, + feed_forward: FeedForward, + ffn_norm: RmsNorm, + attention_norm: RmsNorm, + span: tracing::Span, + } + + impl Block { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let attention = Attention::new(cfg, vb.pp("attention"))?; + let feed_forward = FeedForward::new(cfg, vb.pp("feed_forward"))?; + let ffn_norm = RmsNorm::new(cfg.dim, cfg.norm_eps, vb.pp("ffn_norm"))?; + let attention_norm = RmsNorm::new(cfg.dim, cfg.norm_eps, vb.pp("attention_norm"))?; + Ok(Self { + attention, + feed_forward, + ffn_norm, + attention_norm, + span: tracing::span!(tracing::Level::TRACE, "block"), + }) + } + + fn forward(&mut self, xs: &Tensor, pos: usize, mask: &Tensor) -> Result { + let _enter = self.span.enter(); + let hs = xs.apply(&self.attention_norm)?; + let hs = (xs + self.attention.forward(&hs, pos, mask))?; + &hs + hs.apply(&self.ffn_norm)?.apply(&self.feed_forward) + } + + fn clear_kv_cache(&mut self) { + self.attention.clear_kv_cache() + } + } + + #[derive(Debug, Clone)] + pub struct Model { + tok_embeddings: Embedding, + pos_embeddings: Embedding, + speaker_cond_pos: Linear, + layers: Vec, + norm: RmsNorm, + output: Linear, + spk_cond_mask: Tensor, + span: tracing::Span, + } + + impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let tok_embeddings = Embedding::new(cfg.vocab_size, cfg.dim, vb.pp("tok_embeddings"))?; + let pos_embeddings = Embedding::new(cfg.block_size, cfg.dim, vb.pp("pos_embeddings"))?; + let speaker_cond_pos = linear_b( + cfg.speaker_emb_dim, + cfg.dim, + false, + vb.pp("speaker_cond_pos"), + )?; + let mut layers = Vec::with_capacity(cfg.n_layer); + let vb_l = vb.pp("layers"); + for layer_idx in 0..cfg.n_layer { + let layer = Block::new(cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = RmsNorm::new(cfg.dim, cfg.norm_eps, vb.pp("norm"))?; + let output = linear_b(cfg.dim, cfg.vocab_size, false, vb.pp("output"))?; + let spk_cond_mask = Tensor::cat( + &[ + Tensor::ones((1, 1, cfg.dim), candle::DType::F32, vb.device())?, + Tensor::zeros((1, 1, cfg.dim), candle::DType::F32, vb.device())?, + ], + 0, + )?; + Ok(Self { + tok_embeddings, + pos_embeddings, + speaker_cond_pos, + layers, + norm, + output, + spk_cond_mask, + span: tracing::span!(tracing::Level::TRACE, "qtransformer"), + }) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } + + pub fn forward(&mut self, xs: &Tensor, spk_emb: &Tensor, pos: usize) -> Result { + let _enter = self.span.enter(); + let (_b_sz, seqlen) = xs.dims2()?; + let mask: Vec<_> = (0..seqlen) + .flat_map(|i| (0..seqlen).map(move |j| if i < j { f32::NEG_INFINITY } else { 0. })) + .collect(); + let mask = Tensor::from_slice(&mask, (1, 1, seqlen, seqlen), xs.device())?; + let input_pos = Tensor::arange(pos as u32, (pos + seqlen) as u32, xs.device())?; + let tok_embeddings = xs.apply(&self.tok_embeddings)?; + let pos_embeddings = input_pos.apply(&self.pos_embeddings)?; + let mut xs = tok_embeddings + .broadcast_add(&pos_embeddings)? + .broadcast_add( + &spk_emb + .apply(&self.speaker_cond_pos)? + .broadcast_mul(&self.spk_cond_mask)?, + )?; + let mask = mask.to_dtype(xs.dtype())?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, pos, &mask)? + } + xs.narrow(1, seqlen - 1, 1)? + .contiguous()? + .apply(&self.norm)? + .apply(&self.output) + } + } +} diff --git a/patches/candle-transformers/src/models/quantized_mistral.rs b/patches/candle-transformers/src/models/quantized_mistral.rs new file mode 100644 index 0000000000..cdb687d573 --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_mistral.rs @@ -0,0 +1,337 @@ +//! Mistral model implementation with quantization support. +//! +//! Mistral is a large language model optimized for efficiency. +//! This implementation provides quantization for reduced memory and compute. +//! +//! Key characteristics: +//! - Sliding window attention mechanism +//! - Grouped query attention (GQA) +//! - RMSNorm for layer normalization +//! - Rotary positional embeddings (RoPE) +//! - Support for 8-bit quantization +//! +//! References: +//! - [Mistral Paper](https://arxiv.org/abs/2310.06825) +//! - [Model Card](https://huggingface.co/mistralai/Mistral-7B-v0.1) +//! + +use crate::quantized_nn::{linear_no_bias, Embedding, Linear, RmsNorm}; +pub use crate::quantized_var_builder::VarBuilder; +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::Activation; +use std::sync::Arc; + +pub use crate::models::mistral::Config; + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(cfg: &Config, dev: &Device) -> Result { + let rope_theta = cfg.rope_theta as f32; + let dim = cfg.hidden_size / cfg.num_attention_heads; + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / rope_theta.powf(i as f32 / dim as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(DType::F32)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope(q, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(k, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: Activation, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + let gate_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("gate_proj"))?; + let up_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("up_proj"))?; + let down_proj = linear_no_bias(intermediate_sz, hidden_sz, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_act, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = xs.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + hidden_size: usize, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + let head_dim = hidden_sz / num_heads; + let q_proj = linear_no_bias(hidden_sz, num_heads * head_dim, vb.pp("q_proj"))?; + let k_proj = linear_no_bias(hidden_sz, num_kv_heads * head_dim, vb.pp("k_proj"))?; + let v_proj = linear_no_bias(hidden_sz, num_kv_heads * head_dim, vb.pp("v_proj"))?; + let o_proj = linear_no_bias(num_heads * head_dim, hidden_sz, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + hidden_size: hidden_sz, + rotary_emb, + kv_cache: None, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let (query_states, key_states) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states, seqlen_offset)?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?; + let value_states = crate::utils::repeat_kv(value_states, self.num_kv_groups)?; + + let attn_output = { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + }; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.hidden_size))? + .apply(&self.o_proj) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: MLP, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl DecoderLayer { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let self_attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let input_layernorm = + RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let post_attention_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; + residual + xs + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: Linear, + sliding_window: Option, + device: Device, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + Embedding::new(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(cfg, vb_m.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(rotary_emb.clone(), cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb_m.pp("norm"))?; + let lm_head = linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))?; + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + sliding_window: cfg.sliding_window, + device: vb.device().clone(), + }) + } + + fn prepare_decoder_attention_mask( + &self, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + let sliding_window = self.sliding_window.unwrap_or(tgt_len + 1); + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| { + (0..tgt_len).map(move |j| { + if i < j || j + sliding_window < i { + f32::NEG_INFINITY + } else { + 0. + } + }) + }) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((1, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(DType::F32) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let (_b_size, seq_len) = input_ids.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(seq_len, seqlen_offset)?; + Some(mask) + }; + let mut xs = self.embed_tokens.forward(input_ids)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)? + } + xs.narrow(1, seq_len - 1, 1)? + .contiguous()? + .apply(&self.norm)? + .apply(&self.lm_head) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +} diff --git a/patches/candle-transformers/src/models/quantized_mixformer.rs b/patches/candle-transformers/src/models/quantized_mixformer.rs new file mode 100644 index 0000000000..8736544625 --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_mixformer.rs @@ -0,0 +1,380 @@ +//! Module containing quantized MixFormer model implementation. +//! +//! MixFormer is an efficient transformer variant for text generation that uses +//! mixture-of-experts and parallel attention/feed-forward blocks. +//! This implementation provides quantization for reduced memory usage. +//! +//! Key features: +//! - Parallel attention and feed-forward computation +//! - Rotary positional embeddings +//! - Optional key-value caching +//! - Support for 8-bit quantization +//! + +use crate::quantized_nn::{layer_norm, linear, Linear}; +pub use crate::quantized_var_builder::VarBuilder; +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::Activation; + +pub use crate::models::mixformer::Config; + +const MAX_SEQ_LEN: usize = 4096; + +#[derive(Debug, Clone)] +struct Embedding { + wte: crate::quantized_nn::Embedding, +} + +impl Embedding { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let wte = crate::quantized_nn::Embedding::new(cfg.vocab_size, cfg.n_embd, vb.pp("wte"))?; + Ok(Self { wte }) + } +} + +impl Module for Embedding { + fn forward(&self, xs: &Tensor) -> Result { + self.wte.forward(xs) + } +} + +fn get_mask(size: usize, device: &Device) -> Result { + let mask: Vec<_> = (0..size) + .flat_map(|i| (0..size).map(move |j| u8::from(j > i))) + .collect(); + Tensor::from_slice(&mask, (size, size), device) +} + +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) +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(dim: usize, max_seq_len: usize, dev: &Device) -> Result { + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / 10000f32.powf(i as f32 / dim as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(DType::F32)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + qkv: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor, Tensor)> { + let (_b_size, seqlen, three, _, _headdim) = qkv.dims5()?; + if three != 3 { + candle::bail!("unexpected shape for qkv {:?}", qkv.shape()) + } + let (_rotary_seqlen, rotary_dim) = self.cos.dims2()?; + let rotary_dim = rotary_dim * 2; + let q_rot = qkv.i((.., .., 0, .., ..rotary_dim))?; + let q_pass = qkv.i((.., .., 0, .., rotary_dim..))?; + let k_rot = qkv.i((.., .., 1, .., ..rotary_dim))?; + let k_pass = qkv.i((.., .., 1, .., rotary_dim..))?; + let q12 = q_rot.chunk(2, D::Minus1)?; + let k12 = k_rot.chunk(2, D::Minus1)?; + let (q1, q2) = (&q12[0], &q12[1]); + let (k1, k2) = (&k12[0], &k12[1]); + let c = self.cos.narrow(0, seqlen_offset, seqlen)?.unsqueeze(1)?; + let s = self.sin.narrow(0, seqlen_offset, seqlen)?.unsqueeze(1)?; + let q_rot = Tensor::cat( + &[ + (q1.broadcast_mul(&c)? - q2.broadcast_mul(&s)?)?, + (q1.broadcast_mul(&s)? + q2.broadcast_mul(&c)?)?, + ], + D::Minus1, + )?; + let k_rot = Tensor::cat( + &[ + (k1.broadcast_mul(&c)? - k2.broadcast_mul(&s)?)?, + (k1.broadcast_mul(&s)? + k2.broadcast_mul(&c)?)?, + ], + D::Minus1, + )?; + let q = Tensor::cat(&[&q_rot, &q_pass], D::Minus1)?; + let k = Tensor::cat(&[&k_rot, &k_pass], D::Minus1)?; + let v = qkv.i((.., .., 2))?; + Ok((q, k, v)) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + fc1: Linear, + fc2: Linear, + act: Activation, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let n_inner = cfg.n_inner.unwrap_or(4 * cfg.n_embd); + let fc1 = linear(cfg.n_embd, n_inner, vb.pp("fc1"))?; + let fc2 = linear(n_inner, cfg.n_embd, vb.pp("fc2"))?; + Ok(Self { + fc1, + fc2, + act: cfg.activation_function, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.fc1)?.apply(&self.act)?.apply(&self.fc2) + } +} + +#[derive(Debug, Clone)] +struct CausalLMHead { + ln: candle_nn::LayerNorm, + linear: Linear, +} + +impl CausalLMHead { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let ln = layer_norm(cfg.n_embd, cfg.layer_norm_epsilon, vb.pp("ln"))?; + let linear = linear(cfg.n_embd, cfg.vocab_size, vb.pp("linear"))?; + Ok(Self { ln, linear }) + } +} + +impl Module for CausalLMHead { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.ln)? + .apply(&self.linear)? + .to_dtype(DType::F32) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MHA { + wqkv: Linear, + out_proj: Linear, + rotary_emb: RotaryEmbedding, + kv_cache: Option<(Tensor, Tensor)>, + head_dim: usize, + n_head: usize, + softmax_scale: f64, + span: tracing::Span, +} + +impl MHA { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let head_dim = cfg.n_embd / cfg.n_head; + let op_size = cfg.n_embd; + let wqkv = linear(cfg.n_embd, 3 * op_size, vb.pp("Wqkv"))?; + let out_proj = linear(op_size, cfg.n_embd, vb.pp("out_proj"))?; + let rotary_emb = RotaryEmbedding::new(cfg.rotary_dim, MAX_SEQ_LEN, vb.device())?; + let softmax_scale = 1f64 / (head_dim as f64).sqrt(); + Ok(Self { + wqkv, + out_proj, + head_dim, + n_head: cfg.n_head, + kv_cache: None, + rotary_emb, + softmax_scale, + span: tracing::span!(tracing::Level::TRACE, "mha"), + }) + } + + fn forward(&mut self, xs: &Tensor, mask: Option<&Tensor>) -> Result { + let _enter = self.span.enter(); + let (b_size, seq_len, _n_embd) = xs.dims3()?; + let qkv = self + .wqkv + .forward(xs)? + .reshape((b_size, seq_len, 3, (), self.head_dim))?; + let seqlen_offset = match &self.kv_cache { + None => 0, + Some((prev_k, _)) => prev_k.dim(1)?, + }; + // In the python implementation, a single tensor is returned with the third axis of size 3. + let (q, k, v) = self.rotary_emb.apply_rotary_emb_qkv(&qkv, seqlen_offset)?; + let (k, v) = match &self.kv_cache { + None => (k, v), + Some((prev_k, prev_v)) => { + let k = Tensor::cat(&[prev_k, &k], 1)?; + let v = Tensor::cat(&[prev_v, &v], 1)?; + (k, v) + } + }; + self.kv_cache = Some((k.clone(), v.clone())); + // scores = torch.einsum('bthd,bshd->bhts', q, k * softmax_scale) + let q = q.transpose(1, 2)?.flatten_to(1)?; // b*h, t, d + let k = k.transpose(1, 2)?.flatten_to(1)?; // b*h, s, d + let v = v.transpose(1, 2)?.flatten_to(1)?; // b*h, s, d + let attn_weights = (q.matmul(&k.t()?)? * self.softmax_scale)?; // b*h, t, s + + // causal_mask = torch.triu(torch.full((seqlen_q, seqlen_k), -10000.0, device=scores.device), 1) + // scores = scores + causal_mask.to(dtype=scores.dtype) + let attn_weights = match mask { + None => attn_weights, + Some(mask) => masked_fill( + &attn_weights, + &mask.broadcast_left(b_size * self.n_head)?, + f32::NEG_INFINITY, + )?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + + // output = torch.einsum('bhts,bshd->bthd', attention_drop, v) + // attn_weights: b*h,t,s, v: b*h,s,d + let attn_output = attn_weights.matmul(&v)?; + // b*h,t,d + let attn_output = attn_output + .reshape((b_size, (), seq_len, self.head_dim))? + .transpose(1, 2)? + .flatten_from(D::Minus2)?; + attn_output.apply(&self.out_proj) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +struct ParallelBlock { + ln: candle_nn::LayerNorm, + mixer: MHA, + mlp: MLP, + span: tracing::Span, +} + +impl ParallelBlock { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let ln = layer_norm(cfg.n_embd, cfg.layer_norm_epsilon, vb.pp("ln"))?; + let mixer = MHA::new(cfg, vb.pp("mixer"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + Ok(Self { + ln, + mixer, + mlp, + span: tracing::span!(tracing::Level::TRACE, "block"), + }) + } + + fn forward(&mut self, xs: &Tensor, mask: Option<&Tensor>) -> Result { + let _enter = self.span.enter(); + let residual = xs; + let xs = xs.apply(&self.ln)?; + let attn_outputs = self.mixer.forward(&xs, mask)?; + let feed_forward_hidden_states = self.mlp.forward(&xs)?; + attn_outputs + feed_forward_hidden_states + residual + } + + fn clear_kv_cache(&mut self) { + self.mixer.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct MixFormerSequentialForCausalLM { + embedding: Embedding, + blocks: Vec, + head: CausalLMHead, + span: tracing::Span, +} + +impl MixFormerSequentialForCausalLM { + pub fn new_v2(cfg: &Config, vb: VarBuilder) -> Result { + let vb_head = vb.pp("lm_head"); + let vb = vb.pp("transformer"); + let embedding = Embedding::new(cfg, vb.pp("embd"))?; + let mut blocks = Vec::new(); + for i in 0..cfg.n_layer { + let block = ParallelBlock::new(cfg, vb.pp("h").pp(i))?; + blocks.push(block) + } + let head = CausalLMHead::new(cfg, vb_head)?; + Ok(Self { + embedding, + blocks, + head, + span: tracing::span!(tracing::Level::TRACE, "mixformer"), + }) + } + + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb = vb.pp("layers"); + let embedding = Embedding::new(cfg, vb.pp(0))?; + let mut blocks = Vec::new(); + for i in 0..cfg.n_layer { + let block = ParallelBlock::new(cfg, vb.pp(i + 1))?; + blocks.push(block); + } + let head = CausalLMHead::new(cfg, vb.pp(cfg.n_layer + 1))?; + Ok(Self { + embedding, + blocks, + head, + span: tracing::span!(tracing::Level::TRACE, "mixformer"), + }) + } + + pub fn forward(&mut self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let (_b_size, seq_len) = xs.dims2()?; + let mut xs = xs.apply(&self.embedding)?; + let mask = if seq_len <= 1 { + None + } else { + Some(get_mask(seq_len, xs.device())?) + }; + for block in self.blocks.iter_mut() { + xs = block.forward(&xs, mask.as_ref())?; + } + xs.narrow(1, seq_len - 1, 1)?.apply(&self.head)?.squeeze(1) + } + + pub fn forward_with_img( + &mut self, + bos_token: &Tensor, + xs: &Tensor, + img_embeds: &Tensor, + ) -> Result { + let _enter = self.span.enter(); + let xs = xs.apply(&self.embedding)?; + let bos_token = bos_token.apply(&self.embedding)?; + // Python implementation sequence order is + // https://github.com/vikhyat/moondream/blob/a9d788a20d1543fb1479edc54106e88cff7759d3/moondream/moondream.py#L43-L56 + let mut xs = Tensor::cat(&[bos_token, img_embeds.clone(), xs], 1)?; + let (_b_size, seq_len, _embds) = xs.dims3()?; + let mask = Some(get_mask(seq_len, xs.device())?); + for block in self.blocks.iter_mut() { + xs = block.forward(&xs, mask.as_ref())? + } + let xs = xs + .narrow(1, seq_len - 1, 1)? + .apply(&self.head)? + .squeeze(1)?; + Ok(xs) + } + + pub fn clear_kv_cache(&mut self) { + self.blocks.iter_mut().for_each(|b| b.clear_kv_cache()) + } +} diff --git a/patches/candle-transformers/src/models/quantized_moondream.rs b/patches/candle-transformers/src/models/quantized_moondream.rs new file mode 100644 index 0000000000..9a49598bcd --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_moondream.rs @@ -0,0 +1,286 @@ +//! Implementation of a quantized Moondream vision language model. +//! +//! Moondream is a lightweight vision-language model for image understanding and generation. +//! This module provides a quantized version for reduced memory usage and faster inference. +//! +//! Key features: +//! - ViT-based vision encoder +//! - Phi-2 text decoder model +//! - Memory efficient 8-bit quantization +//! - Optimized for efficient deployment +//! +//! References: +//! - [Moondream Model](https://github.com/vikhyat/moondream) +//! + +use crate::models::moondream::{Config, VisionConfig}; +use crate::models::quantized_mixformer::MixFormerSequentialForCausalLM as PhiModel; +use crate::quantized_nn::{layer_norm, linear_b, Linear}; +use crate::quantized_var_builder::VarBuilder; +use candle::{IndexOp, Module, Result, Tensor, D}; + +fn scaled_dot_product_attention(q: &Tensor, k: &Tensor, v: &Tensor) -> Result { + let dim = q.dim(D::Minus1)?; + let scale_factor = 1.0 / (dim as f64).sqrt(); + let attn_weights = (q.matmul(&k.t()?)? * scale_factor)?; + candle_nn::ops::softmax_last_dim(&attn_weights)?.matmul(v) +} + +#[derive(Debug, Clone)] +struct LinearPatchEmbedding { + linear: Linear, +} + +impl LinearPatchEmbedding { + fn new(vb: VarBuilder) -> Result { + let linear = linear_b(588, 1152, true, vb.pp("linear"))?; + Ok(Self { linear }) + } +} + +impl Module for LinearPatchEmbedding { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.linear) + } +} + +#[derive(Debug, Clone)] +struct Attention { + num_heads: usize, + head_dim: usize, + qkv: Linear, + proj: Linear, +} + +impl Attention { + pub fn new(vb: VarBuilder, dim: usize, num_heads: usize) -> Result { + let qkv = linear_b(dim, dim * 3, true, vb.pp("qkv"))?; + let proj = linear_b(dim, dim, true, vb.pp("proj"))?; + Ok(Self { + num_heads, + head_dim: dim / num_heads, + qkv, + proj, + }) + } +} + +impl Module for Attention { + fn forward(&self, xs: &Tensor) -> Result { + let (b, n, c) = xs.dims3()?; + let qkv = xs + .apply(&self.qkv)? + .reshape((b, n, 3, self.num_heads, self.head_dim))? + .permute((2, 0, 3, 1, 4))?; + let (q, k, v) = ( + qkv.i(0)?.contiguous()?, + qkv.i(1)?.contiguous()?, + qkv.i(2)?.contiguous()?, + ); + scaled_dot_product_attention(&q, &k, &v)? + .transpose(1, 2)? + .reshape((b, n, c))? + .apply(&self.proj) + } +} + +#[derive(Debug, Clone)] +struct VitBlock { + attn: Attention, + mlp: Mlp, + norm1: candle_nn::LayerNorm, + norm2: candle_nn::LayerNorm, +} + +impl VitBlock { + fn new(vb: VarBuilder, dim: usize, num_heads: usize, cfg: &VisionConfig) -> Result { + let attn = Attention::new(vb.pp("attn"), dim, num_heads)?; + let mlp = Mlp::new(vb.pp("mlp"), dim, cfg.hidden_features, dim, cfg.act)?; + let norm1 = layer_norm(dim, 1e-5, vb.pp("norm1"))?; + let norm2 = layer_norm(dim, 1e-5, vb.pp("norm2"))?; + Ok(Self { + attn, + mlp, + norm1, + norm2, + }) + } +} + +impl Module for VitBlock { + fn forward(&self, xs: &Tensor) -> Result { + let ys = xs.apply(&self.norm1)?.apply(&self.attn)?; + let xs = (xs + &ys)?; + let ys = xs.apply(&self.norm2)?.apply(&self.mlp)?; + let xs = (&xs + &ys)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct VisionTransformer { + patch_embed: LinearPatchEmbedding, + pos_embed: Tensor, + blocks: Vec, + norm: candle_nn::LayerNorm, +} + +impl VisionTransformer { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let patch_embed = LinearPatchEmbedding::new(vb.pp("patch_embed"))?; + let pos_embed = vb + .get((1, cfg.embed_len, cfg.embed_dim), "pos_embed")? + .dequantize(vb.device())?; + let blocks = (0..cfg.num_blocks) + .map(|i| { + VitBlock::new( + vb.pp(format!("blocks.{i}")), + cfg.embed_dim, + cfg.num_heads, + cfg, + ) + }) + .collect::>()?; + let norm = layer_norm(cfg.embed_dim, 1e-5, vb.pp("norm"))?; + Ok(Self { + patch_embed, + pos_embed, + blocks, + norm, + }) + } +} + +impl Module for VisionTransformer { + fn forward(&self, xs: &Tensor) -> Result { + let mut xs = (&xs.apply(&self.patch_embed)? + &self.pos_embed)?; + for block in self.blocks.iter() { + xs = xs.apply(block)?; + } + xs.apply(&self.norm) + } +} + +#[derive(Debug, Clone)] +pub struct Encoder { + model: VisionTransformer, +} + +impl Encoder { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let model = VisionTransformer::new(cfg, vb.pp("model.visual"))?; + Ok(Self { model }) + } +} + +impl Module for Encoder { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.model) + } +} + +#[derive(Debug, Clone)] +struct Mlp { + fc1: Linear, + act: candle_nn::Activation, + fc2: Linear, +} + +impl Mlp { + fn new( + vb: VarBuilder, + in_features: usize, + hidden_features: usize, + out_features: usize, + act: candle_nn::Activation, + ) -> Result { + let fc1 = linear_b(in_features, hidden_features, true, vb.pp("fc1"))?; + let fc2 = linear_b(hidden_features, out_features, true, vb.pp("fc2"))?; + Ok(Self { fc1, act, fc2 }) + } +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.fc1)?.apply(&self.act)?.apply(&self.fc2) + } +} + +#[derive(Debug, Clone)] +struct VisionProjection { + mlp: Mlp, +} + +impl VisionProjection { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let mlp = Mlp::new( + vb.pp("mlp"), + cfg.image_embedding_dim, + cfg.hidden_dim, + cfg.model_dim, + cfg.act, + )?; + Ok(Self { mlp }) + } +} + +impl Module for VisionProjection { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.mlp) + } +} + +#[derive(Debug, Clone)] +pub struct VisionEncoder { + encoder: Encoder, + projection: VisionProjection, +} + +impl VisionEncoder { + pub fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let encoder = Encoder::new(cfg, vb.pp("encoder"))?; + let projection = VisionProjection::new(cfg, vb.pp("projection"))?; + Ok(Self { + encoder, + projection, + }) + } +} + +impl Module for VisionEncoder { + fn forward(&self, xs: &Tensor) -> Result { + let (b, c, hp1, wp2) = xs.dims4()?; + let (p1, p2) = (14, 14); + let h = hp1 / p1; + let w = wp2 / p2; + xs.reshape((b, c, h, p1, h, p2))? + .permute((0, 2, 4, 1, 3, 5))? + .reshape((b, h * w, c * p1 * p2))? + .apply(&self.encoder)? + .apply(&self.projection) + } +} + +pub struct Model { + pub text_model: PhiModel, + pub vision_encoder: VisionEncoder, +} + +impl Model { + pub fn new(config: &Config, vb: VarBuilder) -> Result { + let text_model = PhiModel::new_v2(&config.phi_config, vb.pp("text_model"))?; + let vision_encoder = VisionEncoder::new(&config.vision_config, vb.pp("vision_encoder"))?; + Ok(Self { + text_model, + vision_encoder, + }) + } + + pub fn vision_encoder(&self) -> &VisionEncoder { + &self.vision_encoder + } + + pub fn text_model(&mut self) -> &mut PhiModel { + &mut self.text_model + } +} diff --git a/patches/candle-transformers/src/models/quantized_mpt.rs b/patches/candle-transformers/src/models/quantized_mpt.rs new file mode 100644 index 0000000000..44d8566b7b --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_mpt.rs @@ -0,0 +1,219 @@ +//! Quantized MPT model implementation. +//! +//! MPT (MPT-7B) is a causal transformer model series optimized for code generation. +//! This implementation provides quantization for reduced memory and compute. +//! +//! Key characteristics: +//! - Multi-Query Grouped Attention (MQA) +//! - Support for KV-caching +//! - Pre-computed ALiBi attention biases +//! - Support for 8-bit quantization +//! +//! References: +//! - [Replit Code Models](https://huggingface.co/replit/replit-code-v1_5-3b) +//! - [MPT-7B Implementation](https://github.com/mosaicml/llm-foundry) +//! +/// MPT model used by replit-code-v1_5-3b +/// https://huggingface.co/replit/replit-code-v1_5-3b/blob/main/modeling_mpt.py +/// +use crate::quantized_nn::{layer_norm_no_bias, linear_no_bias, Embedding, Linear}; +pub use crate::quantized_var_builder::VarBuilder; +/// MPT model used by replit-code-v1_5-3b +/// https://huggingface.co/replit/replit-code-v1_5-3b/blob/main/modeling_mpt.py +use candle::{IndexOp, Module, Result, Tensor, D}; +use candle_nn::LayerNorm; + +pub use super::mpt::Config; + +#[derive(Debug, Clone)] +struct GroupedQueryAttention { + wqkv: Linear, + out_proj: Linear, + kv_cache: Option<(Tensor, Tensor)>, + softmax_scale: f64, + head_dim: usize, + d_model: usize, + n_heads: usize, + kv_n_heads: usize, + attn_bias: Tensor, + span: tracing::Span, +} + +impl GroupedQueryAttention { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let head_dim = cfg.d_model / cfg.n_heads; + let wqkv_size = cfg.d_model + 2 * cfg.kv_n_heads * head_dim; + let wqkv = linear_no_bias(cfg.d_model, wqkv_size, vb.pp("Wqkv"))?; + let softmax_scale = 1f64 / (head_dim as f64).sqrt(); + let out_proj = linear_no_bias(cfg.d_model, cfg.d_model, vb.pp("out_proj"))?; + let attn_bias = super::mpt::build_alibi_bias(cfg)?.to_device(vb.device())?; + Ok(Self { + wqkv, + out_proj, + kv_cache: None, + softmax_scale, + head_dim, + d_model: cfg.d_model, + n_heads: cfg.n_heads, + kv_n_heads: cfg.kv_n_heads, + attn_bias, + span: tracing::span!(tracing::Level::TRACE, "gqa"), + }) + } + + fn forward(&mut self, xs: &Tensor, mask: Option<&Tensor>) -> Result { + let _enter = self.span.enter(); + let (b_size, seq_len, _n_embd) = xs.dims3()?; + let qkv = self.wqkv.forward(xs)?; + let query = qkv.narrow(2, 0, self.d_model)?; + let kv_size = self.kv_n_heads * self.head_dim; + let key = qkv.narrow(2, self.d_model, kv_size)?; + let value = qkv.narrow(2, self.d_model + kv_size, kv_size)?; + // scaled_multihead_dot_product_attention + let query = query + .reshape((b_size, seq_len, self.n_heads, ()))? + .transpose(1, 2)?; // b,h,s,d + let key = key + .reshape((b_size, seq_len, self.kv_n_heads, ()))? + .permute((0, 2, 3, 1))?; // b,h,d,s + let value = value + .reshape((b_size, seq_len, self.kv_n_heads, ()))? + .transpose(1, 2)?; // b,h,s,d + let (key, value) = match &self.kv_cache { + None => (key, value), + Some((prev_k, prev_v)) => { + let k = Tensor::cat(&[prev_k, &key], 3)?; + let v = Tensor::cat(&[prev_v, &value], 2)?; + (k, v) + } + }; + self.kv_cache = Some((key.clone(), value.clone())); + let query = query.contiguous()?; + let key = crate::utils::repeat_kv(key, self.n_heads / self.kv_n_heads)?.contiguous()?; + let value = crate::utils::repeat_kv(value, self.n_heads / self.kv_n_heads)?.contiguous()?; + let attn_weights = (query.matmul(&key)? * self.softmax_scale)?; + let attn_bias = { + let s_q = query.dim(D::Minus2)?; + let s_k = key.dim(D::Minus1)?; + let (_, _, a_q, a_k) = self.attn_bias.dims4()?; + let start_q = a_q.saturating_sub(s_q); + let start_k = a_k.saturating_sub(s_k); + self.attn_bias.i((.., .., start_q.., start_k..))? + }; + let attn_weights = attn_weights.broadcast_add(&attn_bias)?; + let attn_weights = match mask { + None => attn_weights, + Some(mask) => super::mpt::masked_fill( + &attn_weights, + &mask.broadcast_as(attn_weights.shape())?, + f32::NEG_INFINITY, + )?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + let attn_output = attn_weights + .matmul(&value)? + .transpose(1, 2)? + .flatten_from(D::Minus2)?; + let out = attn_output.apply(&self.out_proj)?; + Ok(out) + } +} + +#[derive(Debug, Clone)] +struct Ffn { + up_proj: Linear, + down_proj: Linear, +} + +impl Ffn { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden = cfg.d_model * cfg.expansion_ratio; + let up_proj = linear_no_bias(cfg.d_model, hidden, vb.pp("up_proj"))?; + let down_proj = linear_no_bias(hidden, cfg.d_model, vb.pp("down_proj"))?; + Ok(Self { up_proj, down_proj }) + } +} + +impl Module for Ffn { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.up_proj)?.gelu_erf()?.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +struct MPTBlock { + norm1: LayerNorm, // Do we need the low-precision variant? + attn: GroupedQueryAttention, + norm2: LayerNorm, + ffn: Ffn, +} + +impl MPTBlock { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let norm1 = layer_norm_no_bias(cfg.d_model, 1e-5, vb.pp("norm_1"))?; + let norm2 = layer_norm_no_bias(cfg.d_model, 1e-5, vb.pp("norm_2"))?; + let attn = GroupedQueryAttention::new(cfg, vb.pp("attn"))?; + let ffn = Ffn::new(cfg, vb.pp("ffn"))?; + Ok(Self { + norm1, + attn, + norm2, + ffn, + }) + } + + fn forward(&mut self, xs: &Tensor, mask: Option<&Tensor>) -> Result { + let residual = xs; + let xs = xs.apply(&self.norm1)?; + let xs = self.attn.forward(&xs, mask)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.norm2)?.apply(&self.ffn)?; + xs + residual + } +} + +#[derive(Debug, Clone)] +pub struct Model { + wte: Embedding, + blocks: Vec, + norm_f: LayerNorm, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let wte = Embedding::new(cfg.vocab_size, cfg.d_model, vb.pp("wte"))?; + let vb_b = vb.pp("blocks"); + let mut blocks = Vec::with_capacity(cfg.n_layers); + for i in 0..cfg.n_layers { + let block = MPTBlock::new(cfg, vb_b.pp(i))?; + blocks.push(block) + } + let norm_f = layer_norm_no_bias(cfg.d_model, 1e-5, vb.pp("norm_f"))?; + Ok(Self { + wte, + blocks, + norm_f, + }) + } + + pub fn forward(&mut self, xs: &Tensor) -> Result { + let (_b_size, seq_len) = xs.dims2()?; + let mut xs = xs.apply(&self.wte)?; + let mask = if seq_len <= 1 { + None + } else { + Some(super::mpt::get_mask(seq_len, xs.device())?) + }; + for block in self.blocks.iter_mut() { + xs = block.forward(&xs, mask.as_ref())?; + } + let xs = xs.apply(&self.norm_f)?; + let logits = xs + .narrow(1, seq_len - 1, 1)? + .squeeze(1)? + .matmul(&self.wte.embeddings().t()?)? + .squeeze(1)?; + Ok(logits) + } +} diff --git a/patches/candle-transformers/src/models/quantized_phi.rs b/patches/candle-transformers/src/models/quantized_phi.rs new file mode 100644 index 0000000000..b874ad94ea --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_phi.rs @@ -0,0 +1,305 @@ +//! Phi2 model implementation with quantization support. +//! +//! Phi2 is a 2.7B parameter language model using scaled-up Transformer decoder architecture. +//! This implementation provides quantization for reduced memory and compute usage. +//! +//! Key characteristics: +//! - Partial attention with learned mixing to reduce quadratic costs +//! - Layer reuse for improved inference efficiency +//! - Linear transformations with scalar mixing +//! - Rotary positional embeddings (RoPE) +//! - Support for 8-bit quantization +//! +//! References: +//! - [Phi2 Paper](https://arxiv.org/abs/2309.05463) +//! - [Model Card](https://huggingface.co/microsoft/phi-2) +//! + +use std::collections::HashMap; + +use candle::quantized::gguf_file; +use candle::quantized::QTensor; +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{Embedding, LayerNorm}; + +pub const MAX_SEQ_LEN: usize = 4096; + +#[derive(Debug, Clone)] +struct QLinear { + inner: candle::quantized::QMatMul, + bias: Tensor, + span: tracing::Span, +} + +impl QLinear { + fn new( + ct: &gguf_file::Content, + r: &mut R, + name: &str, + device: &Device, + ) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "qmatmul"); + let w = ct.tensor(r, &format!("{name}.weight"), device)?; + let b = ct.tensor(r, &format!("{name}.bias"), device)?; + let inner = candle::quantized::QMatMul::from_qtensor(w)?; + let bias = b.dequantize(device)?; + Ok(Self { inner, bias, span }) + } +} + +impl Module for QLinear { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + self.inner.forward(xs)?.broadcast_add(&self.bias) + } +} + +#[derive(Debug, Clone)] +struct Mlp { + ffn_up: QLinear, + ffn_down: QLinear, +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.ffn_up)?.gelu()?.apply(&self.ffn_down) + } +} + +#[derive(Debug, Clone)] +struct LayerWeights { + attn_qkv: QLinear, + attn_output: QLinear, + attn_norm: LayerNorm, + mlp: Mlp, + n_head: usize, + n_kv_head: usize, + head_dim: usize, + cos: Tensor, + sin: Tensor, + rope_dim: usize, + neg_inf: Tensor, + kv_cache: Option<(Tensor, Tensor)>, + span_attn: tracing::Span, + span_rot: tracing::Span, +} + +fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: &Tensor) -> Result { + let shape = mask.shape(); + let m = mask.where_cond(&on_true.broadcast_as(shape.dims())?, on_false)?; + Ok(m) +} + +impl LayerWeights { + fn apply_rotary_emb(&self, xs: &Tensor, index_pos: usize) -> Result { + let _enter = self.span_rot.enter(); + let (_b_sz, _n_head, seq_len, _n_embd) = xs.dims4()?; + let xs_rot = xs.i((.., .., .., ..self.rope_dim))?; + let xs_pass = xs.i((.., .., .., self.rope_dim..))?; + let cos = self.cos.narrow(0, index_pos, seq_len)?; + let sin = self.sin.narrow(0, index_pos, seq_len)?; + let xs_rot = candle_nn::rotary_emb::rope(&xs_rot.contiguous()?, &cos, &sin)?; + Tensor::cat(&[&xs_rot, &xs_pass], D::Minus1) + } + + fn forward_attn( + &mut self, + x: &Tensor, + mask: Option<&Tensor>, + index_pos: usize, + ) -> Result { + let _enter = self.span_attn.enter(); + let (b_sz, seq_len, n_embd) = x.dims3()?; + let qkv = + self.attn_qkv + .forward(x)? + .reshape((b_sz, seq_len, 3, self.n_head, self.head_dim))?; + + let q = qkv.i((.., .., 0))?.transpose(1, 2)?; + let k = qkv.i((.., .., 1))?.transpose(1, 2)?; + let v = qkv.i((.., .., 2))?.transpose(1, 2)?; + // This call to contiguous ensures that the fast kernel can be called below. It's + // actually a no-op except when processing the initial prompt so has no significant + // impact on performance. + let v = v.contiguous()?; + + let q = self.apply_rotary_emb(&q, index_pos)?.contiguous()?; + let k = self.apply_rotary_emb(&k, index_pos)?; + + let (k, v) = match &self.kv_cache { + None => (k.contiguous()?, v.contiguous()?), + Some((k_cache, v_cache)) => { + if index_pos == 0 { + (k.contiguous()?, v.contiguous()?) + } else { + let k = Tensor::cat(&[k_cache, &k], 2)?; + let v = Tensor::cat(&[v_cache, &v], 2)?; + (k.contiguous()?, v.contiguous()?) + } + } + }; + self.kv_cache = Some((k.clone(), v.clone())); + + let k = crate::utils::repeat_kv(k, self.n_head / self.n_kv_head)?; + let v = crate::utils::repeat_kv(v, self.n_head / self.n_kv_head)?; + + let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?; + let att = match mask { + None => att, + Some(mask) => { + let mask = mask.broadcast_as(att.shape())?; + masked_fill(&att, &mask, &self.neg_inf)? + } + }; + 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.attn_output.forward(&y)?; + Ok(y) + } +} + +#[derive(Debug, Clone)] +pub struct ModelWeights { + tok_embeddings: Embedding, + layers: Vec, + output_norm: LayerNorm, + output: QLinear, + masks: HashMap, + span: tracing::Span, + span_output: tracing::Span, +} + +fn precomput_freqs_cis( + head_dim: usize, + freq_base: f32, + device: &Device, +) -> 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)?; + let idx_theta = Tensor::arange(0, MAX_SEQ_LEN as u32, device)? + .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)) +} + +fn layer_norm(w: QTensor, b: QTensor, eps: f64) -> Result { + let w = w.dequantize(&w.device())?; + let b = b.dequantize(&b.device())?; + let ln = LayerNorm::new(w, b, eps); + Ok(ln) +} + +impl ModelWeights { + pub fn from_gguf( + ct: gguf_file::Content, + reader: &mut R, + device: &Device, + ) -> Result { + 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 head_count = md_get("phi2.attention.head_count")?.to_u32()? as usize; + let head_count_kv = md_get("phi2.attention.head_count_kv")?.to_u32()? as usize; + let block_count = md_get("phi2.block_count")?.to_u32()? as usize; + let embedding_length = md_get("phi2.embedding_length")?.to_u32()? as usize; + let rope_dim = md_get("phi2.rope.dimension_count")?.to_u32()? as usize; + let ln_eps = md_get("phi2.attention.layer_norm_epsilon")?.to_f32()? as f64; + let (cos, sin) = precomput_freqs_cis(rope_dim, 10_000., device)?; + let neg_inf = Tensor::new(f32::NEG_INFINITY, device)?; + + let tok_embeddings = ct.tensor(reader, "token_embd.weight", device)?; + let tok_embeddings = tok_embeddings.dequantize(device)?; + let output_norm = layer_norm( + ct.tensor(reader, "output_norm.weight", device)?, + ct.tensor(reader, "output_norm.bias", device)?, + ln_eps, + )?; + let output = QLinear::new(&ct, reader, "output", device)?; + let mut layers = Vec::with_capacity(block_count); + for layer_idx in 0..block_count { + let prefix = format!("blk.{layer_idx}"); + let ffn_up = QLinear::new(&ct, reader, &format!("{prefix}.ffn_up"), device)?; + let ffn_down = QLinear::new(&ct, reader, &format!("{prefix}.ffn_down"), device)?; + let mlp = Mlp { ffn_up, ffn_down }; + let attn_norm = layer_norm( + ct.tensor(reader, &format!("{prefix}.attn_norm.weight"), device)?, + ct.tensor(reader, &format!("{prefix}.attn_norm.bias"), device)?, + ln_eps, + )?; + let span_attn = tracing::span!(tracing::Level::TRACE, "attn"); + let span_rot = tracing::span!(tracing::Level::TRACE, "attn-rot"); + layers.push(LayerWeights { + attn_qkv: QLinear::new(&ct, reader, &format!("{prefix}.attn_qkv"), device)?, + attn_output: QLinear::new(&ct, reader, &format!("{prefix}.attn_output"), device)?, + attn_norm, + mlp, + n_head: head_count, + n_kv_head: head_count_kv, + head_dim: embedding_length / head_count, + cos: cos.clone(), + sin: sin.clone(), + rope_dim, + neg_inf: neg_inf.clone(), + kv_cache: None, + span_attn, + span_rot, + }) + } + 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, + output_norm, + output, + masks: HashMap::new(), + span, + span_output, + }) + } + + fn mask(&mut self, t: usize, device: &Device) -> 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)?; + self.masks.insert(t, mask.clone()); + Ok(mask) + } + } + + pub fn forward(&mut self, xs: &Tensor, index_pos: usize) -> Result { + let (_b_sz, seq_len) = xs.dims2()?; + let mask = if seq_len == 1 { + None + } else { + Some(self.mask(seq_len, xs.device())?) + }; + let _enter = self.span.enter(); + let mut xs = self.tok_embeddings.forward(xs)?; + for layer in self.layers.iter_mut() { + let residual = &xs; + let xs_norm = xs.apply(&layer.attn_norm)?; + let attn_outputs = layer.forward_attn(&xs_norm, mask.as_ref(), index_pos)?; + let feed_forward_hidden_states = layer.mlp.forward(&xs_norm)?; + xs = (attn_outputs + feed_forward_hidden_states + residual)? + } + let xs = xs.apply(&self.output_norm)?.i((.., seq_len - 1, ..))?; + let _enter = self.span_output.enter(); + self.output.forward(&xs) + } +} diff --git a/patches/candle-transformers/src/models/quantized_phi3.rs b/patches/candle-transformers/src/models/quantized_phi3.rs new file mode 100644 index 0000000000..4a04e43418 --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_phi3.rs @@ -0,0 +1,340 @@ +//! Phi3 model implementation with quantization support. +//! +//! Phi3 is a language model intended for research purposes. +//! This implementation provides quantization for reduced memory usage. +//! +//! Key characteristics: +//! - Multi-head attention +//! - RMSNorm for layer normalization +//! - Rotary positional embeddings (RoPE) +//! - Support for quantization +//! +//! References: +//! - [Model Card](https://huggingface.co/microsoft/phi-3) +//! + +use std::collections::HashMap; + +use candle::quantized::gguf_file; +use candle::quantized::QTensor; +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{kv_cache::KvCache, Embedding, RmsNorm}; + +#[derive(Debug, Clone)] +struct QLinear { + inner: candle::quantized::QMatMul, + span: tracing::Span, +} + +impl QLinear { + fn new( + ct: &gguf_file::Content, + r: &mut R, + name: &str, + device: &Device, + ) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "qmatmul"); + let w = ct.tensor(r, &format!("{name}.weight"), device)?; + let inner = candle::quantized::QMatMul::from_qtensor(w)?; + Ok(Self { inner, span }) + } +} + +impl Module for QLinear { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + self.inner.forward(xs) + } +} + +#[derive(Debug, Clone)] +struct Mlp { + ffn_up: QLinear, + ffn_down: QLinear, + i_size: usize, +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let up_states = xs.apply(&self.ffn_up)?; + let gate = up_states.narrow(D::Minus1, 0, self.i_size)?; + let up_states = up_states.narrow(D::Minus1, self.i_size, self.i_size)?; + let up_states = (up_states * gate.silu()?)?; + up_states.apply(&self.ffn_down) + } +} + +fn rms_norm(w: QTensor, eps: f64) -> Result { + let w = w.dequantize(&w.device())?; + let rms = RmsNorm::new(w, eps); + Ok(rms) +} + +#[derive(Debug, Clone)] +struct LayerWeights { + attn_qkv: QLinear, + attn_output: QLinear, + attn_norm: RmsNorm, + ffn_norm: RmsNorm, + mlp: Mlp, + n_head: usize, + n_kv_head: usize, + head_dim: usize, + cos: Tensor, + sin: Tensor, + neg_inf: Tensor, + kv_cache: KvCache, + use_flash_attn: bool, + span_attn: tracing::Span, + span_rot: tracing::Span, +} + +fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: &Tensor) -> Result { + let shape = mask.shape(); + let m = mask.where_cond(&on_true.broadcast_as(shape.dims())?, on_false)?; + Ok(m) +} + +impl LayerWeights { + fn apply_rotary_emb(&self, xs: &Tensor, index_pos: usize) -> Result { + let _enter = self.span_rot.enter(); + let (_b_sz, _h, seq_len, _n_embd) = xs.dims4()?; + let cos = self.cos.narrow(0, index_pos, seq_len)?; + let sin = self.sin.narrow(0, index_pos, seq_len)?; + candle_nn::rotary_emb::rope(&xs.contiguous()?, &cos, &sin) + } + + fn forward_attn( + &mut self, + x: &Tensor, + mask: Option<&Tensor>, + index_pos: usize, + ) -> Result { + let _enter = self.span_attn.enter(); + let (b_sz, seq_len, n_embd) = x.dims3()?; + let qkv = self.attn_qkv.forward(x)?; + + let query_pos = self.n_head * self.head_dim; + let q = qkv.narrow(D::Minus1, 0, query_pos)?; + let k = qkv.narrow(D::Minus1, query_pos, self.n_kv_head * self.head_dim)?; + let v = qkv.narrow( + D::Minus1, + query_pos + self.n_kv_head * self.head_dim, + self.n_kv_head * self.head_dim, + )?; + + 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)?.contiguous()?; + let k = self.apply_rotary_emb(&k, index_pos)?; + + if index_pos == 0 { + self.kv_cache.reset(); + } + let (k, v) = self.kv_cache.append(&k.contiguous()?, &v.contiguous()?)?; + + let k = crate::utils::repeat_kv(k, self.n_head / self.n_kv_head)?; + let v = crate::utils::repeat_kv(v, self.n_head / self.n_kv_head)?; + + let y = if self.use_flash_attn { + // flash-attn expects (b_sz, seq_len, nheads, head_dim) + let q = q.to_dtype(DType::BF16)?.transpose(1, 2)?; + let k = k.to_dtype(DType::BF16)?.transpose(1, 2)?; + let v = v.to_dtype(DType::BF16)?.transpose(1, 2)?; + let softmax_scale = 1f32 / (self.head_dim as f32).sqrt(); + flash_attn(&q, &k, &v, softmax_scale, seq_len > 1)? + .to_dtype(DType::F32)? + .transpose(1, 2)? + } else { + let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?; + let att = match mask { + None => att, + Some(mask) => { + let mask = mask.broadcast_as(att.shape())?; + masked_fill(&att, &mask, &self.neg_inf)? + } + }; + let att = candle_nn::ops::softmax_last_dim(&att)?; + // Convert to contiguous as matmul doesn't support strided vs for now. + att.matmul(&v)? + }; + let y = y.transpose(1, 2)?.reshape(&[b_sz, seq_len, n_embd])?; + let y = self.attn_output.forward(&y)?; + Ok(y) + } +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +#[derive(Debug, Clone)] +pub struct ModelWeights { + tok_embeddings: Embedding, + layers: Vec, + output_norm: RmsNorm, + output: QLinear, + masks: HashMap, + span: tracing::Span, + span_output: tracing::Span, +} + +fn precomput_freqs_cis( + head_dim: usize, + max_seq_len: usize, + freq_base: f32, + device: &Device, +) -> 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)?; + let idx_theta = Tensor::arange(0, max_seq_len as u32, device)? + .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_gguf( + use_flash_attn: bool, + ct: gguf_file::Content, + reader: &mut R, + device: &Device, + ) -> Result { + 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 head_count = md_get("phi3.attention.head_count")?.to_u32()? as usize; + let head_count_kv = md_get("phi3.attention.head_count_kv")?.to_u32()? as usize; + let block_count = md_get("phi3.block_count")?.to_u32()? as usize; + let embedding_length = md_get("phi3.embedding_length")?.to_u32()? as usize; + let max_seq_len = md_get("phi3.context_length")?.to_u32()? as usize; + let head_dim = embedding_length / head_count; + let i_size = md_get("phi3.feed_forward_length")?.to_u32()? as usize; + let rope_dim = md_get("phi3.rope.dimension_count")?.to_u32()? as usize; + let rms_eps = md_get("phi3.attention.layer_norm_rms_epsilon")?.to_f32()? as f64; + let (cos, sin) = precomput_freqs_cis(rope_dim, max_seq_len, 10_000., device)?; + let neg_inf = Tensor::new(f32::NEG_INFINITY, device)?; + + let tok_embeddings = ct.tensor(reader, "token_embd.weight", device)?; + let tok_embeddings = tok_embeddings.dequantize(device)?; + let output_norm = rms_norm(ct.tensor(reader, "output_norm.weight", device)?, rms_eps)?; + let output = QLinear::new(&ct, reader, "output", device)?; + + let mut layers = Vec::with_capacity(block_count); + for layer_idx in 0..block_count { + let prefix = format!("blk.{layer_idx}"); + let ffn_up = QLinear::new(&ct, reader, &format!("{prefix}.ffn_up"), device)?; + let ffn_down = QLinear::new(&ct, reader, &format!("{prefix}.ffn_down"), device)?; + let mlp = Mlp { + ffn_up, + ffn_down, + i_size, + }; + let attn_norm = rms_norm( + ct.tensor(reader, &format!("{prefix}.attn_norm.weight"), device)?, + rms_eps, + )?; + let ffn_norm = rms_norm( + ct.tensor(reader, &format!("{prefix}.ffn_norm.weight"), device)?, + rms_eps, + )?; + let span_attn = tracing::span!(tracing::Level::TRACE, "attn"); + let span_rot = tracing::span!(tracing::Level::TRACE, "attn-rot"); + let kv_cache = KvCache::new(2, max_seq_len); + layers.push(LayerWeights { + attn_qkv: QLinear::new(&ct, reader, &format!("{prefix}.attn_qkv"), device)?, + attn_output: QLinear::new(&ct, reader, &format!("{prefix}.attn_output"), device)?, + attn_norm, + ffn_norm, + mlp, + n_head: head_count, + n_kv_head: head_count_kv, + head_dim, + cos: cos.clone(), + sin: sin.clone(), + neg_inf: neg_inf.clone(), + kv_cache, + use_flash_attn, + span_attn, + span_rot, + }) + } + 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, + output_norm, + output, + masks: HashMap::new(), + span, + span_output, + }) + } + + fn mask(&mut self, t: usize, device: &Device) -> 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)?; + self.masks.insert(t, mask.clone()); + Ok(mask) + } + } + + pub fn forward(&mut self, xs: &Tensor, index_pos: usize) -> Result { + let (_b_sz, seq_len) = xs.dims2()?; + let mask = if seq_len == 1 { + None + } else { + Some(self.mask(seq_len, xs.device())?) + }; + let _enter = self.span.enter(); + let mut xs = self.tok_embeddings.forward(xs)?; + for layer in self.layers.iter_mut() { + let residual = &xs; + let ys = xs.apply(&layer.attn_norm)?; + let ys = layer.forward_attn(&ys, mask.as_ref(), index_pos)?; + let ys = (ys + residual)?; + let residual = &ys; + let ys = ys.apply(&layer.ffn_norm)?; + let ys = layer.mlp.forward(&ys)?; + xs = (ys + residual)? + } + let xs = xs.apply(&self.output_norm)?.i((.., seq_len - 1, ..))?; + let _enter = self.span_output.enter(); + self.output.forward(&xs) + } +} diff --git a/patches/candle-transformers/src/models/quantized_qwen2.rs b/patches/candle-transformers/src/models/quantized_qwen2.rs new file mode 100644 index 0000000000..c04da56925 --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_qwen2.rs @@ -0,0 +1,338 @@ +//! Qwen2 model implementation with quantization support. +//! +//! Qwen2 is a chat-optimized language model that supports 8-bit quantization +//! for reduced memory usage and faster inference. +//! +//! Key characteristics: +//! - Group Query Attention (GQA) +//! - RMSNorm for layer normalization +//! - Rotary positional embeddings (RoPE) +//! - Support for 8-bit quantization +//! +//! References: +//! - [Model Card](https://huggingface.co/Qwen/Qwen2) +//! + +use crate::{quantized_nn::RmsNorm, utils::repeat_kv}; +use candle::{ + quantized::{gguf_file, QMatMul}, + DType, Device, IndexOp, Result, Tensor, +}; +use candle_nn::{Embedding, Module}; +use std::collections::HashMap; + +#[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)] +struct LayerWeights { + attention_wq: QMatMul, + attention_wk: QMatMul, + attention_wv: QMatMul, + attention_bq: Tensor, + attention_bk: Tensor, + attention_bv: Tensor, + attention_wo: QMatMul, + attention_norm: RmsNorm, + mlp: Mlp, + ffn_norm: RmsNorm, + n_head: usize, + n_kv_head: usize, + head_dim: usize, + cos: Tensor, + sin: Tensor, + neg_inf: 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: &Tensor) -> Result { + let shape = mask.shape(); + let m = mask.where_cond(&on_true.broadcast_as(shape.dims())?, 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)?; + let sin = self.sin.narrow(0, index_pos, seq_len)?; + candle_nn::rotary_emb::rope(&x.contiguous()?, &cos, &sin) + } + + fn forward_attn( + &mut self, + x: &Tensor, + mask: Option<&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.broadcast_add(&self.attention_bq)?; + let k = k.broadcast_add(&self.attention_bk)?; + let v = v.broadcast_add(&self.attention_bv)?; + + let q = q + .reshape((b_sz, seq_len, self.n_head, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let k = k + .reshape((b_sz, seq_len, self.n_kv_head, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let v = v + .reshape((b_sz, seq_len, self.n_kv_head, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + + // let (q, k) = self + // .rotary_embedding + // .apply_rotary_emb_qkv(&q, &k, index_pos)?; + 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)?; + let v = Tensor::cat(&[v_cache, &v], 2)?; + (k, v) + } + } + }; + self.kv_cache = Some((k.clone(), v.clone())); + + // Support for MQA, useful for 70B models and mistral. + let k = repeat_kv(k, self.n_head / self.n_kv_head)?; + let v = repeat_kv(v, self.n_head / self.n_kv_head)?; + + let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?; + let att = match mask { + None => att, + Some(mask) => { + let mask = mask.broadcast_as(att.shape())?; + masked_fill(&att, &mask, &self.neg_inf)? + } + }; + 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) + } +} + +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, + context_length: usize, + device: &Device, +) -> 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)?; + let idx_theta = Tensor::arange(0, context_length as u32, device)? + .to_dtype(DType::F32)? + .reshape((context_length, 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_gguf( + ct: gguf_file::Content, + reader: &mut R, + device: &Device, + ) -> Result { + let md_get = |s: &str| match ct.metadata.get(s) { + None => candle::bail!("cannot find {s} in metadata"), + Some(v) => Ok(v), + }; + + let head_count = md_get("qwen2.attention.head_count")?.to_u32()? as usize; + let head_count_kv = md_get("qwen2.attention.head_count_kv")?.to_u32()? as usize; + let embedding_length = md_get("qwen2.embedding_length")?.to_u32()? as usize; + let context_length = md_get("qwen2.context_length")?.to_u32()? as usize; + let block_count = md_get("qwen2.block_count")?.to_u32()? as usize; + let rms_norm_eps = md_get("qwen2.attention.layer_norm_rms_epsilon")?.to_f32()? as f64; + let rope_freq_base = md_get("qwen2.rope.freq_base") + .and_then(|m| m.to_f32()) + .unwrap_or(10000f32); + + let head_dim = embedding_length / head_count; + + let neg_inf = Tensor::new(f32::NEG_INFINITY, device)?; + + let tok_embeddings = ct.tensor(reader, "token_embd.weight", device)?; + let tok_embeddings = tok_embeddings.dequantize(device)?; + let norm = RmsNorm::from_qtensor( + ct.tensor(reader, "output_norm.weight", device)?, + rms_norm_eps, + )?; + let output = match ct.tensor(reader, "output.weight", device) { + Ok(v) => QMatMul::from_qtensor(v)?, + _ => { + // use tie_word_embeddings + QMatMul::from_qtensor(ct.tensor(reader, "token_embd.weight", device)?)? + } + }; + + let (cos, sin) = precomput_freqs_cis(head_dim, rope_freq_base, context_length, device)?; + + 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"), device)?; + let attention_wk = ct.tensor(reader, &format!("{prefix}.attn_k.weight"), device)?; + let attention_wv = ct.tensor(reader, &format!("{prefix}.attn_v.weight"), device)?; + + let attention_bq = ct.tensor(reader, &format!("{prefix}.attn_q.bias"), device)?; + let attention_bk = ct.tensor(reader, &format!("{prefix}.attn_k.bias"), device)?; + let attention_bv = ct.tensor(reader, &format!("{prefix}.attn_v.bias"), device)?; + + let attention_wo = + ct.tensor(reader, &format!("{prefix}.attn_output.weight"), device)?; + + let mlp = { + let feed_forward_w1 = + ct.tensor(reader, &format!("{prefix}.ffn_gate.weight"), device)?; + let feed_forward_w2 = + ct.tensor(reader, &format!("{prefix}.ffn_down.weight"), device)?; + let feed_forward_w3 = + ct.tensor(reader, &format!("{prefix}.ffn_up.weight"), device)?; + 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.tensor(reader, &format!("{prefix}.attn_norm.weight"), device)?; + let ffn_norm = ct.tensor(reader, &format!("{prefix}.ffn_norm.weight"), device)?; + + 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_bq: attention_bq.dequantize(device)?, + attention_bk: attention_bk.dequantize(device)?, + attention_bv: attention_bv.dequantize(device)?, + attention_wo: QMatMul::from_qtensor(attention_wo)?, + attention_norm: RmsNorm::from_qtensor(attention_norm, rms_norm_eps)?, + cos: cos.clone(), + sin: sin.clone(), + mlp, + ffn_norm: RmsNorm::from_qtensor(ffn_norm, rms_norm_eps)?, + n_head: head_count, + n_kv_head: head_count_kv, + head_dim, + neg_inf: neg_inf.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, + masks: HashMap::new(), + span, + span_output, + }) + } + + fn mask(&mut self, t: usize, device: &Device) -> 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)?; + 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 = if seq_len == 1 { + None + } else { + Some(self.mask(seq_len, x.device())?) + }; + 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.as_ref(), 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.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/patches/candle-transformers/src/models/quantized_qwen3.rs b/patches/candle-transformers/src/models/quantized_qwen3.rs new file mode 100644 index 0000000000..85ccbb0edd --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_qwen3.rs @@ -0,0 +1,433 @@ +//! Qwen3 implementation with quantization support. +//! +//! Based on the Qwen3 architecture and implemented with quantized weights +//! for reduced memory usage and faster inference on compatible hardware. +//! +//! References: +//! - [Qwen3 Models](https://huggingface.co/Qwen/Qwen3-0.6B) (architecture based on official implementations) +//! +use super::with_tracing::QMatMul; +use crate::{quantized_nn::RmsNorm, utils::repeat_kv}; +use candle::quantized::{gguf_file, QTensor}; +use candle::{DType, Device, Result, Tensor}; +use candle_nn::{kv_cache::ConcatKvCache, Activation, Embedding, Module}; +use std::io::{Read, Seek}; +use std::sync::Arc; + +pub struct Gguf { + ct: gguf_file::Content, + reader: R, + device: Device, +} + +impl Gguf { + pub fn new(ct: gguf_file::Content, reader: R, device: Device) -> Self { + Self { ct, reader, device } + } + + pub fn qmatmul(&mut self, name: &str) -> Result { + let ws = self.ct.tensor(&mut self.reader, name, &self.device)?; + QMatMul::from_weights(ws.into()) + } + + pub fn rms_norm(&mut self, name: &str, eps: f64) -> Result { + let ws = self.ct.tensor(&mut self.reader, name, &self.device)?; + RmsNorm::from_qtensor(ws, eps) + } + + pub fn metadata(&self) -> &std::collections::HashMap { + &self.ct.metadata + } + + pub fn tensor(&mut self, name: &str) -> Result { + self.ct.tensor(&mut self.reader, name, &self.device) + } +} + +#[derive(Debug, Clone)] +struct MlpWeights { + gate_proj: QMatMul, + up_proj: QMatMul, + down_proj: QMatMul, + act_fn: Activation, + span: tracing::Span, +} + +impl MlpWeights { + fn new(gg: &mut Gguf, prefix: &str) -> Result { + let gate_proj = gg.qmatmul(&format!("{prefix}.ffn_gate.weight"))?; + let up_proj = gg.qmatmul(&format!("{prefix}.ffn_up.weight"))?; + let down_proj = gg.qmatmul(&format!("{prefix}.ffn_down.weight"))?; + let act_fn = Activation::Silu; + let span = tracing::span!(tracing::Level::TRACE, "mlp"); + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn, + span, + }) + } +} + +impl Module for MlpWeights { + fn forward(&self, x: &Tensor) -> Result { + let _enter = self.span.enter(); + let gate = self.gate_proj.forward(x)?.apply(&self.act_fn)?; + let up = self.up_proj.forward(x)?; + let gated = (gate * up)?; + self.down_proj.forward(&gated) + } +} + +#[derive(Debug, Clone)] +pub struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + pub fn new( + dtype: DType, + head_dim: usize, + max_position_embeddings: usize, + rope_theta: f64, + dev: &Device, + ) -> Result { + let dim = head_dim; + let max_seq_len = max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + /// Apply RoPE (q, k shape: B x H x L x D) + pub fn apply(&self, q: &Tensor, k: &Tensor, offset: usize) -> Result<(Tensor, Tensor)> { + let (_, _, seq_len, _) = q.dims4()?; + let cos = self.cos.narrow(0, offset, seq_len)?.to_dtype(q.dtype())?; + let sin = self.sin.narrow(0, offset, seq_len)?.to_dtype(q.dtype())?; + let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +struct AttentionWeights { + q_proj: QMatMul, + k_proj: QMatMul, + v_proj: QMatMul, + o_proj: QMatMul, + q_norm: RmsNorm, + k_norm: RmsNorm, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + rotary_emb: Arc, + kv_cache: ConcatKvCache, + span_attn: tracing::Span, +} + +impl AttentionWeights { + fn new( + gg: &mut Gguf, + num_heads: usize, + num_kv_heads: usize, + head_dim: usize, + rms_norm_eps: f64, + rotary_emb: Arc, + prefix: &str, + ) -> Result { + let num_kv_groups = num_heads / num_kv_heads; + + let q_proj = gg.qmatmul(&format!("{prefix}.attn_q.weight"))?; + let k_proj = gg.qmatmul(&format!("{prefix}.attn_k.weight"))?; + let v_proj = gg.qmatmul(&format!("{prefix}.attn_v.weight"))?; + let o_proj = gg.qmatmul(&format!("{prefix}.attn_output.weight"))?; + + let q_norm = gg.rms_norm(&format!("{prefix}.attn_q_norm.weight"), rms_norm_eps)?; + let k_norm = gg.rms_norm(&format!("{prefix}.attn_k_norm.weight"), rms_norm_eps)?; + + let kv_cache = ConcatKvCache::new(2); + + let span_attn = tracing::span!(tracing::Level::TRACE, "attn"); + + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + q_norm, + k_norm, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + rotary_emb, + kv_cache, + span_attn, + }) + } + + fn forward(&mut self, x: &Tensor, attn_mask: Option<&Tensor>, offset: usize) -> Result { + let _enter = self.span_attn.enter(); + let (b, l, _) = x.dims3()?; + + let q = self.q_proj.forward(x)?; + let k = self.k_proj.forward(x)?; + let v = self.v_proj.forward(x)?; + + let q = q + .reshape((b, l, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let k = k + .reshape((b, l, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let v = v + .reshape((b, l, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let q_flat = q.flatten(0, 2)?; + let k_flat = k.flatten(0, 2)?; + + let q_flat = self.q_norm.forward(&q_flat)?; + let k_flat = self.k_norm.forward(&k_flat)?; + let q = q_flat.reshape((b, self.num_heads, l, self.head_dim))?; + let k = k_flat.reshape((b, self.num_kv_heads, l, self.head_dim))?; + + let (q, k) = self.rotary_emb.apply(&q, &k, offset)?; + + let (k, v) = self.kv_cache.append(&k, &v)?; + + let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?; + let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?; + + let scale = 1.0 / (self.head_dim as f64).sqrt(); + let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?; + if let Some(m) = attn_mask { + let m_dtype = m.dtype(); + let scores_dtype = scores.dtype(); + let mask = if m_dtype != scores_dtype { + m.to_dtype(scores_dtype)? + } else { + m.clone() + }; + scores = scores.broadcast_add(&mask)?; + } + let probs = candle_nn::ops::softmax_last_dim(&scores)?; + let ctx = probs.matmul(&v)?; // (B, H, L, D) + let reshaped_ctx = ctx + .transpose(1, 2)? + .reshape((b, l, self.num_heads * self.head_dim))?; + self.o_proj.forward(&reshaped_ctx) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache.reset(); + } +} + +#[derive(Debug, Clone)] +struct LayerWeights { + self_attn: AttentionWeights, + mlp: MlpWeights, + ln1: RmsNorm, + ln2: RmsNorm, +} + +impl LayerWeights { + fn new( + gg: &mut Gguf, + num_attention_heads: usize, + num_key_value_heads: usize, + head_dim: usize, + rms_norm_eps: f64, + rotary: Arc, + layer_idx: usize, + ) -> Result { + let prefix = format!("blk.{layer_idx}"); + + let ln1 = gg.rms_norm(&format!("{prefix}.attn_norm.weight"), rms_norm_eps)?; + let ln2 = gg.rms_norm(&format!("{prefix}.ffn_norm.weight"), rms_norm_eps)?; + let self_attn = AttentionWeights::new( + gg, + num_attention_heads, + num_key_value_heads, + head_dim, + rms_norm_eps, + rotary, + &prefix, + )?; + let mlp = MlpWeights::new(gg, &prefix)?; + Ok(Self { + self_attn, + mlp, + ln1, + ln2, + }) + } + + fn forward(&mut self, x: &Tensor, mask: Option<&Tensor>, offset: usize) -> Result { + let h = self.ln1.forward(x)?; + let h = self.self_attn.forward(&h, mask, offset)?; + let x = (x + h)?; + let h2 = self.ln2.forward(&x)?; + let h2 = h2.apply(&self.mlp)?; + x + h2 + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache(); + } +} + +#[derive(Debug, Clone)] +pub struct ModelWeights { + embed_tokens: Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: QMatMul, + device: Device, + dtype: DType, + span: tracing::Span, + span_output: tracing::Span, +} + +impl ModelWeights { + pub fn from_gguf( + ct: gguf_file::Content, + reader: &mut R, + device: &Device, + ) -> Result { + let mut gg = Gguf::new(ct, reader, device.clone()); + let md_get = |s: &str| match gg.metadata().get(s) { + None => candle::bail!("cannot find {s} in metadata"), + Some(v) => Ok(v), + }; + + let num_attention_heads = md_get("qwen3.attention.head_count")?.to_u32()? as usize; + let num_kv_heads = md_get("qwen3.attention.head_count_kv")?.to_u32()? as usize; + let head_dim = md_get("qwen3.attention.key_length")?.to_u32()? as usize; + let num_layers = md_get("qwen3.block_count")?.to_u32()? as usize; + let hidden_size = md_get("qwen3.embedding_length")?.to_u32()? as usize; + let max_position_embeddings = md_get("qwen3.context_length")?.to_u32()? as usize; + let rms_norm_eps = md_get("qwen3.attention.layer_norm_rms_epsilon")?.to_f32()? as f64; + let rope_freq_base = md_get("qwen3.rope.freq_base")?.to_f32()? as f64; + + let dtype = match gg.metadata().get("general.dtype") { + Some(v) => match v.to_u32() { + Ok(0) => DType::F32, + Ok(1) => DType::F16, + _ => DType::F16, + }, + None => DType::F16, + }; + + let embed_tensor = gg.tensor("token_embd.weight")?; + let embed_tokens = Embedding::new(embed_tensor.dequantize(device)?, hidden_size); + + let rotary = Arc::new(RotaryEmbedding::new( + dtype, + head_dim, + max_position_embeddings, + rope_freq_base, + device, + )?); + + let mut layers = Vec::with_capacity(num_layers); + for i in 0..num_layers { + layers.push(LayerWeights::new( + &mut gg, + num_attention_heads, + num_kv_heads, + head_dim, + rms_norm_eps, + rotary.clone(), + i, + )?); + } + + let norm = gg.rms_norm("output_norm.weight", rms_norm_eps)?; + // Load output projection tensor, falling back to tied embeddings like gemma3 + let lm_head_tensor = match gg.tensor("output.weight") { + Ok(tensor) => tensor, + Err(_) => gg.tensor("token_embd.weight")?, + }; + let lm_head = QMatMul::from_weights(lm_head_tensor.into())?; + let span = tracing::span!(tracing::Level::TRACE, "model"); + let span_output = tracing::span!(tracing::Level::TRACE, "output"); + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + device: device.clone(), + dtype, + span, + span_output, + }) + } + + fn causal_mask( + &self, + b: usize, + tgt: usize, + offset: usize, + sw: Option, + ) -> Result { + let minf = f32::NEG_INFINITY; + let mask: Vec<_> = (0..tgt) + .flat_map(|i| { + (0..(tgt + offset)).map(move |j| { + let past_ok = j <= i + offset; + let sw_ok = match sw { + Some(w) => (i + offset) as i64 - j as i64 <= w as i64, + None => true, + }; + if past_ok && sw_ok { + 0. + } else { + minf + } + }) + }) + .collect(); + Tensor::from_slice(&mask, (b, 1, tgt, tgt + offset), &self.device)?.to_dtype(self.dtype) + } + + pub fn forward(&mut self, input: &Tensor, offset: usize) -> Result { + let _enter = self.span.enter(); + let (b, l) = input.dims2()?; + let mut h = self.embed_tokens.forward(input)?; + let causal_mask = if l == 1 { + None + } else { + Some(self.causal_mask(b, l, offset, None)?) + }; + for layer in &mut self.layers { + h = layer.forward(&h, causal_mask.as_ref(), offset)?; + } + let h = self.norm.forward(&h)?; + let _enter = self.span_output.enter(); + let last_hidden = h.narrow(1, l - 1, 1)?; + self.lm_head.forward(&last_hidden)?.squeeze(1) + } + + pub fn clear_kv_cache(&mut self) { + for layer in &mut self.layers { + layer.clear_kv_cache(); + } + } +} diff --git a/patches/candle-transformers/src/models/quantized_qwen3_moe.rs b/patches/candle-transformers/src/models/quantized_qwen3_moe.rs new file mode 100644 index 0000000000..57c3abf599 --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_qwen3_moe.rs @@ -0,0 +1,451 @@ +use super::quantized_qwen3::{Gguf, RotaryEmbedding}; +use super::with_tracing::QMatMul; +use crate::fused_moe::{FusedMoeGGUF, MoeCfg}; +use crate::quantized_nn::RmsNorm; +use crate::utils::repeat_kv; +use candle::quantized::gguf_file; +use candle::{DType, Device, Result, Tensor}; +use candle_nn::kv_cache::ConcatKvCache; +use candle_nn::Linear; +use candle_nn::{Embedding, Module}; +use std::sync::Arc; +#[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)?) + } +} + +enum MoeOrMlp { + FusedMoe(FusedMoeGGUF), + Mlp(Mlp), +} + +impl MoeOrMlp { + fn forward(&self, xs: &Tensor, is_prefill: bool) -> Result { + match self { + Self::Mlp(m) => m.forward(xs), + Self::FusedMoe(m) => m.forward(xs, is_prefill), + } + } +} + +pub struct QuantizedAttention { + attention_wq: QMatMul, + attention_wk: QMatMul, + attention_wv: QMatMul, + attention_bq: Option, + attention_bk: Option, + attention_bv: Option, + attention_wo: QMatMul, + q_norm: Option, + k_norm: Option, + n_head: usize, + n_kv_head: usize, + head_dim: usize, + num_kv_groups: usize, + rotary_emb: Arc, + dtype: DType, + kv_cache: ConcatKvCache, +} + +impl QuantizedAttention { + #[allow(clippy::too_many_arguments)] + pub fn new( + gg: &mut Gguf, + prefix: &str, + dtype: DType, + num_heads: usize, + num_kv_heads: usize, + head_dim: usize, + rms_norm_eps: f64, + device: &Device, + rotary_emb: Arc, + ) -> Result { + let num_kv_groups = num_heads / num_kv_heads; + let attention_wq = gg.qmatmul(&format!("{prefix}.attn_q.weight"))?; + let attention_wk = gg.qmatmul(&format!("{prefix}.attn_k.weight"))?; + let attention_wv = gg.qmatmul(&format!("{prefix}.attn_v.weight"))?; + + let attention_bq = gg.tensor(&format!("{prefix}.attn_q.bias")); + let attention_bk = gg.tensor(&format!("{prefix}.attn_k.bias")); + let attention_bv = gg.tensor(&format!("{prefix}.attn_v.bias")); + + let attention_bq = if let Ok(attention_bq) = attention_bq { + Some(attention_bq.dequantize(device)?.to_dtype(DType::F32)?) + } else { + None + }; + + let attention_bk = if let Ok(attention_bk) = attention_bk { + Some(attention_bk.dequantize(device)?.to_dtype(DType::F32)?) + } else { + None + }; + + let attention_bv = if let Ok(attention_bv) = attention_bv { + Some(attention_bv.dequantize(device)?.to_dtype(DType::F32)?) + } else { + None + }; + + let attention_wo = gg.qmatmul(&format!("{prefix}.attn_output.weight"))?; + let q_norm = Some(gg.rms_norm(&format!("{prefix}.attn_q_norm.weight"), rms_norm_eps)?); + let k_norm = Some(gg.rms_norm(&format!("{prefix}.attn_k_norm.weight"), rms_norm_eps)?); + let kv_cache = ConcatKvCache::new(2); + Ok(QuantizedAttention { + attention_wq, + attention_wk, + attention_wv, + attention_bq, + attention_bk, + attention_bv, + attention_wo, + q_norm, + k_norm, + n_head: num_heads, + n_kv_head: num_kv_heads, + head_dim, + num_kv_groups, + rotary_emb: rotary_emb.clone(), + dtype, + kv_cache, + }) + } + + pub fn forward( + &mut self, + x: &Tensor, + mask: Option<&Tensor>, + input_pos: usize, + ) -> Result { + let (b, seq_len, _) = x.dims3()?; + let in_dtype = x.dtype(); + let q = self.attention_wq.forward(x)?; + let k = self.attention_wk.forward(x)?; + let v = self.attention_wv.forward(x)?; + + let q = if let Some(bq) = &self.attention_bq { + q.broadcast_add(bq)? + } else { + q + }; + + let k = if let Some(bk) = &self.attention_bk { + k.broadcast_add(bk)? + } else { + k + }; + + let v = if let Some(bv) = &self.attention_bv { + v.broadcast_add(bv)? + } else { + v + }; + + let q = q + .reshape((1, seq_len, self.n_head, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let k = k + .reshape((1, seq_len, self.n_kv_head, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let v = v + .reshape((1, seq_len, self.n_kv_head, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + + let (q, k) = if let (Some(q_norm), Some(k_norm)) = (&self.q_norm, &self.k_norm) { + // Per‑head RMSNorm in qwen3 + let q_flat = q.flatten(0, 2)?; // (B*H, L, D) -> (BHL, D) after transpose later + let k_flat = k.flatten(0, 2)?; + + // q_norm and k_norm weights stored in f32 format in qwen3 gguf + let q_flat = q_norm.forward(&q_flat)?; + let k_flat = k_norm.forward(&k_flat)?; + + let q = q_flat.reshape((1, self.n_head, seq_len, self.head_dim))?; + let k = k_flat.reshape((1, self.n_kv_head, seq_len, self.head_dim))?; + + (q, k) + } else { + (q, k) + }; + + let (q, k, v) = ( + q.to_dtype(self.dtype)?, + k.to_dtype(self.dtype)?, + v.to_dtype(self.dtype)?, + ); + + let (q, k) = self.rotary_emb.apply(&q, &k, input_pos)?; + + let (k, v) = self.kv_cache.append(&k, &v)?; + + let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?; + let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?; + + let scale = 1.0 / (self.head_dim as f64).sqrt(); + let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?; + + if let Some(m) = mask { + let m_dtype = m.dtype(); + let scores_dtype = scores.dtype(); + let mask = if m_dtype != scores_dtype { + m.to_dtype(scores_dtype)? + } else { + m.clone() + }; + scores = scores.broadcast_add(&mask)?; + } + + let probs = candle_nn::ops::softmax_last_dim(&scores)?; + let ctx = probs.matmul(&v)?; // (B, H, L, D) + let reshaped_ctx = + ctx.transpose(1, 2)? + .reshape((b, seq_len, self.n_head * self.head_dim))?; + + self.attention_wo.forward(&reshaped_ctx.to_dtype(in_dtype)?) + } +} + +struct LayerWeights { + self_attn: QuantizedAttention, + attention_norm: RmsNorm, + mlp: MoeOrMlp, + ffn_norm: RmsNorm, +} + +impl LayerWeights { + fn forward_attn(&mut self, x: &Tensor, mask: Option<&Tensor>, offset: usize) -> Result { + self.self_attn.forward(x, mask, offset) + } +} + +pub struct GGUFQWenMoE { + tok_embeddings: Embedding, + layers: Vec, + norm: RmsNorm, + output: QMatMul, + dtype: DType, + device: Device, +} + +impl GGUFQWenMoE { + pub fn from_gguf( + ct: gguf_file::Content, + reader: &mut R, + device: &Device, + dtype: DType, + ) -> Result { + let mut gg = Gguf::new(ct, reader, device.clone()); + let md_get = |s: &str| match gg.metadata().get(s) { + None => candle::bail!("cannot find {s} in metadata"), + Some(v) => Ok(v), + }; + let arch = md_get("general.architecture")?.to_string()?; + + let head_count = + md_get(format!("{arch}.attention.head_count").as_str())?.to_u32()? as usize; + let head_count_kv = + md_get(format!("{arch}.attention.head_count_kv").as_str())?.to_u32()? as usize; + + let head_dim = md_get(format!("{arch}.attention.key_length").as_str()); + let embedding_length = + md_get(format!("{arch}.embedding_length").as_str())?.to_u32()? as usize; + let head_dim = if let Ok(head_dim) = head_dim { + head_dim.to_u32()? as usize + } else { + embedding_length / head_count + }; + let context_length = md_get(format!("{arch}.context_length").as_str())?.to_u32()? as usize; + let block_count = md_get(format!("{arch}.block_count").as_str())?.to_u32()? as usize; + let rms_norm_eps = + md_get(format!("{arch}.attention.layer_norm_rms_epsilon").as_str())?.to_f32()? as f64; + let rope_freq_base = md_get(format!("{arch}.rope.freq_base").as_str()) + .and_then(|m| m.to_f32()) + .unwrap_or(10000f32); + let expert_shared_feed_forward_length = + md_get(format!("{arch}.expert_shared_feed_forward_length").as_str()); + let shared_expert_intermediate_size = match expert_shared_feed_forward_length { + Ok(length) => { + if length.to_u32()? > 0 { + Some(length.to_u32()? as usize) + } else { + None + } + } + _ => None, + }; + + let moe_cfg = MoeCfg { + moe_intermediate_size: md_get(format!("{arch}.expert_feed_forward_length").as_str())? + .to_u32()? as usize, + num_experts: md_get(format!("{arch}.expert_count").as_str())?.to_u32()? as usize, + norm_topk_prob: shared_expert_intermediate_size.is_none(), + num_experts_per_tok: md_get(format!("{arch}.expert_used_count").as_str())?.to_u32()? + as usize, + hidden_size: head_dim, + act: candle_nn::Activation::Silu, + decoder_sparse_step: None, + }; + + let tok_embeddings = gg.tensor("token_embd.weight")?; + let tok_embeddings = tok_embeddings.dequantize(device)?; + let norm = gg.rms_norm("output_norm.weight", rms_norm_eps)?; + let output = match gg.qmatmul("output.weight") { + Ok(v) => v, + _ => { + // use tie_word_embeddings + gg.qmatmul("token_embd.weight")? + } + }; + + let rotary_emb = Arc::new(RotaryEmbedding::new( + dtype, + head_dim, + context_length, + rope_freq_base as f64, + device, + )?); + let mut layers = Vec::with_capacity(block_count); + for layer_idx in 0..block_count { + let prefix = format!("blk.{layer_idx}"); + let mlp = if moe_cfg.num_experts > 0 + && (layer_idx + 1) % moe_cfg.decoder_sparse_step.unwrap_or(1) == 0 + { + let gate_ws = gg + .tensor(&format!("{prefix}.ffn_gate_inp.weight"))? + .dequantize(device)? + .to_dtype(DType::F32)?; + let gate = Linear::new(gate_ws, None); + let gate_experts = Arc::new(gg.tensor(&format!("{prefix}.ffn_gate_exps.weight"))?); + let up_experts = Arc::new(gg.tensor(&format!("{prefix}.ffn_up_exps.weight"))?); + let down_experts = Arc::new(gg.tensor(&format!("{prefix}.ffn_down_exps.weight"))?); + let moe = FusedMoeGGUF { + gate, + gate_experts, + up_experts, + down_experts, + act: candle_nn::Activation::Silu, + norm_topk_prob: moe_cfg.norm_topk_prob, + num_experts_per_tok: moe_cfg.num_experts_per_tok, + dtype, + }; + + MoeOrMlp::FusedMoe(moe) + } else { + let mlp = { + let feed_forward_w1 = gg.qmatmul(&format!("{prefix}.ffn_gate.weight"))?; + let feed_forward_w2 = gg.qmatmul(&format!("{prefix}.ffn_down.weight"))?; + let feed_forward_w3 = gg.qmatmul(&format!("{prefix}.ffn_up.weight"))?; + Mlp { + feed_forward_w1, + feed_forward_w2, + feed_forward_w3, + } + }; + MoeOrMlp::Mlp(mlp) + }; + + let attention_norm = + gg.rms_norm(&format!("{prefix}.attn_norm.weight"), rms_norm_eps)?; + let ffn_norm = gg.rms_norm(&format!("{prefix}.ffn_norm.weight"), rms_norm_eps)?; + + let self_attn = QuantizedAttention::new( + &mut gg, + &prefix, + dtype, + head_count, + head_count_kv, + head_dim, + rms_norm_eps, + device, + rotary_emb.clone(), + )?; + layers.push(LayerWeights { + self_attn, + attention_norm, + mlp, + ffn_norm, + }); + } + + Ok(Self { + tok_embeddings: Embedding::new(tok_embeddings, embedding_length), + layers, + norm, + output, + dtype, + device: device.clone(), + }) + } + + fn causal_mask( + &self, + b: usize, + tgt: usize, + offset: usize, + sw: Option, + ) -> Result { + let minf = f32::NEG_INFINITY; + let mask: Vec<_> = (0..tgt) + .flat_map(|i| { + (0..(tgt + offset)).map(move |j| { + let past_ok = j <= i + offset; + let sw_ok = match sw { + Some(w) => (i + offset) as i64 - j as i64 <= w as i64, + None => true, + }; + if past_ok && sw_ok { + 0. + } else { + minf + } + }) + }) + .collect(); + Tensor::from_slice(&mask, (b, 1, tgt, tgt + offset), &self.device)?.to_dtype(self.dtype) + } + + pub fn forward(&mut self, x: &Tensor, offset: usize) -> Result { + let mut xs = self.tok_embeddings.forward(x)?; + let (b, l) = x.dims2()?; + + let causal_mask = if l == 1 { + None + } else { + Some(self.causal_mask(b, l, offset, None)?) + }; + + for layer in self.layers.iter_mut() { + let x = xs; + let residual = &x; + + let x = layer.attention_norm.forward(&x)?; + let attn = layer.forward_attn(&x, causal_mask.as_ref(), offset)?; + let x = (attn + residual)?; + + // MLP + let residual = &x; + let x = layer.ffn_norm.forward(&x)?; + let x = layer.mlp.forward(&x, causal_mask.is_some())?; + let x = (x + residual)?; + xs = x + } + + let xs = xs.narrow(1, l - 1, 1)?; + let xs = self.norm.forward(&xs)?; + self.output.forward(&xs)?.to_dtype(DType::F32)?.squeeze(1) + } +} diff --git a/patches/candle-transformers/src/models/quantized_recurrent_gemma.rs b/patches/candle-transformers/src/models/quantized_recurrent_gemma.rs new file mode 100644 index 0000000000..e40daa1f33 --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_recurrent_gemma.rs @@ -0,0 +1,429 @@ +//! Recurrent Gemma model implementation with quantization support. +//! +//! Gemma is a large language model optimized for efficiency. +//! This implementation provides quantization for reduced memory and compute. +//! +//! Key characteristics: +//! - Recurrent blocks with gated recurrent units +//! - Convolution and attention blocks +//! - RMSNorm for layer normalization +//! - Rotary positional embeddings (RoPE) +//! - Support for 8-bit quantization +//! +//! References: +//! - [Gemma Paper](https://arxiv.org/abs/2401.06751) +//! - [Model Card](https://ai.google.dev/gemma) +//! + +use crate::quantized_nn::{linear_b as linear, Embedding, Linear}; +pub use crate::quantized_var_builder::VarBuilder; +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use std::sync::Arc; + +use crate::models::recurrent_gemma::{Config, Rglru, RmsNorm, RotaryEmbedding, TemporalBlockType}; + +fn rms_norm(size: usize, eps: f64, vb: VarBuilder) -> Result { + let weight = vb.get(size, "weight")?.dequantize(vb.device())?; + Ok(RmsNorm::from_weight(weight, eps)) +} + +#[derive(Debug, Clone)] +struct Mlp { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: candle_nn::Activation, +} + +impl Mlp { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let h = cfg.hidden_size; + let intermediate_size = cfg.intermediate_size / 2; + let gate_proj = linear(h, intermediate_size, true, vb.pp("gate_proj"))?; + let up_proj = linear(h, intermediate_size, true, vb.pp("up_proj"))?; + let down_proj = linear(intermediate_size, h, true, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_activation, + }) + } +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let gate = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; + (gate * xs.apply(&self.up_proj))?.apply(&self.down_proj) + } +} + +fn rglru(cfg: &Config, vb: VarBuilder) -> Result { + let h = cfg.hidden_size; + let lru_width = cfg.lru_width.unwrap_or(h); + let n_heads = cfg.num_attention_heads; + let block_width = lru_width / n_heads; + let recurrent_param = vb.get((lru_width,), "recurrent_param")?; + let input_gate_weight = vb.get((n_heads, block_width, block_width), "input_gate_weight")?; + let input_gate_bias = vb.get((n_heads, block_width), "input_gate_bias")?; + let recurrent_gate_weight = + vb.get((n_heads, block_width, block_width), "recurrent_gate_weight")?; + let recurrent_gate_bias = vb.get((n_heads, block_width), "recurrent_gate_bias")?; + Ok(Rglru { + recurrent_param: recurrent_param.dequantize(vb.device())?, + input_gate_bias: input_gate_bias.dequantize(vb.device())?, + input_gate_weight: input_gate_weight.dequantize(vb.device())?, + recurrent_gate_bias: recurrent_gate_bias.dequantize(vb.device())?, + recurrent_gate_weight: recurrent_gate_weight.dequantize(vb.device())?, + block_width, + n_heads, + recurrent_states: None, + }) +} + +#[derive(Debug, Clone)] +struct RecurrentBlock { + linear_y: Linear, + linear_x: Linear, + linear_out: Linear, + conv_1d: candle_nn::Conv1d, + conv1d_state: Option, + conv1d_width: usize, + rg_lru: Rglru, + act_fn: candle_nn::Activation, +} + +impl RecurrentBlock { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let h = cfg.hidden_size; + let lru_width = cfg.lru_width.unwrap_or(h); + let linear_y = linear(h, lru_width, true, vb.pp("linear_y"))?; + let linear_x = linear(h, lru_width, true, vb.pp("linear_x"))?; + let linear_out = linear(lru_width, h, true, vb.pp("linear_out"))?; + + let conv_1d = { + let ws = vb + .get((lru_width, 1, cfg.conv1d_width), "conv_1d.weight")? + .dequantize(vb.device())?; + let bs = vb.get(lru_width, "conv_1d.bias")?.dequantize(vb.device())?; + let config = candle_nn::Conv1dConfig { + groups: lru_width, + padding: cfg.conv1d_width - 1, + ..Default::default() + }; + candle_nn::Conv1d::new(ws, Some(bs), config) + }; + let rg_lru = rglru(cfg, vb.pp("rg_lru"))?; + Ok(Self { + linear_y, + linear_x, + linear_out, + conv_1d, + conv1d_state: None, + conv1d_width: cfg.conv1d_width, + rg_lru, + act_fn: cfg.hidden_activation, + }) + } + + pub fn forward(&mut self, xs: &Tensor, pos: usize) -> Result { + let (_b_sz, seq_len, _) = xs.dims3()?; + + let y_branch = xs.apply(&self.linear_y)?.apply(&self.act_fn)?; + let x_branch = xs.apply(&self.linear_x)?.transpose(1, 2)?; + let x_branch = if pos == 0 { + let x_len = x_branch.dim(D::Minus1)?; + let pad = self.conv1d_width as i64 - x_len as i64 - 1; + let padded = match pad.cmp(&0) { + std::cmp::Ordering::Equal => x_branch.clone(), + std::cmp::Ordering::Less => { + let rev_pad = (-pad) as usize; + x_branch.narrow(D::Minus1, rev_pad, x_len - rev_pad)? + } + std::cmp::Ordering::Greater => { + x_branch.pad_with_zeros(D::Minus1, pad as usize, 0)? + } + }; + self.conv1d_state = Some(padded); + x_branch + .apply(&self.conv_1d)? + .narrow(D::Minus1, 0, seq_len)? + } else { + let conv_state = match self.conv1d_state.as_ref() { + None => candle::bail!("empty cache despite pos > 0"), + Some(s) => Tensor::cat(&[s, &x_branch], D::Minus1)?, + }; + let w = self.conv_1d.weight().i((.., 0, ..))?; + let x_branch = conv_state.broadcast_mul(&w)?.sum(D::Minus1)?; + let x_branch = match self.conv_1d.bias() { + None => x_branch, + Some(b) => x_branch.broadcast_add(b)?, + }; + let x_branch = x_branch.unsqueeze(D::Minus1)?; + self.conv1d_state = Some(conv_state.i((.., .., 1..))?); + x_branch + }; + let x_branch = x_branch.transpose(1, 2)?; + let x_branch = self.rg_lru.forward(&x_branch, pos)?; + (x_branch * y_branch)?.apply(&self.linear_out) + } +} + +#[derive(Debug, Clone)] +struct SdpaAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + n_heads: usize, + n_kv_heads: usize, + head_dim: usize, + hidden_size: usize, + kv_cache: Option<(Tensor, Tensor)>, + rotary_emb: Arc, +} + +impl SdpaAttention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let h = cfg.hidden_size; + let n_heads = cfg.num_attention_heads; + let n_kv_heads = cfg.num_key_value_heads; + let hd = cfg.head_dim; + let q_proj = linear(h, n_heads * hd, cfg.attention_bias, vb.pp("q_proj"))?; + let k_proj = linear(h, n_kv_heads * hd, cfg.attention_bias, vb.pp("k_proj"))?; + let v_proj = linear(h, n_kv_heads * hd, cfg.attention_bias, vb.pp("v_proj"))?; + let o_proj = linear(n_heads * hd, h, true, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + n_heads, + n_kv_heads, + head_dim: hd, + hidden_size: h, + kv_cache: None, + rotary_emb, + }) + } + + fn repeat_kv(&self, x: Tensor) -> Result { + let n_rep = self.n_heads / self.n_kv_heads; + crate::utils::repeat_kv(x, n_rep) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + pos: usize, + ) -> Result { + let (bsz, q_len, _) = xs.dims3()?; + + let query_states = xs.apply(&self.q_proj)?; + let key_states = xs.apply(&self.k_proj)?; + let value_states = xs.apply(&self.v_proj)?; + + let query_states = query_states + .reshape((bsz, q_len, self.n_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((bsz, q_len, self.n_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((bsz, q_len, self.n_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let query_states = query_states.chunk(2, D::Minus1)?; + let key_states = key_states.chunk(2, D::Minus1)?; + let (query_rot, key_rot) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states[0], &key_states[0], pos)?; + let query_states = Tensor::cat(&[&query_rot, &query_states[1]], D::Minus1)?.contiguous()?; + let key_states = Tensor::cat(&[&key_rot, &key_states[1]], D::Minus1)?.contiguous()?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + let key_states = self.repeat_kv(key_states)?; + let value_states = self.repeat_kv(value_states)?; + let xs = { + let att = (query_states.matmul(&key_states.t()?)? / (self.head_dim as f64).sqrt())?; + let att = if q_len == 1 { + att + } else { + match attention_mask { + None => att, + Some(mask) => att.broadcast_add(mask)?, + } + }; + let att = candle_nn::ops::softmax_last_dim(&att)?; + att.matmul(&value_states.contiguous()?)? + }; + + let xs = xs + .transpose(1, 2)? + .reshape((bsz, q_len, self.hidden_size))?; + self.o_proj.forward(&xs) + } +} + +#[derive(Debug, Clone)] +enum TemporalBlock { + Recurrent(RecurrentBlock), + Attention(SdpaAttention), +} + +impl TemporalBlock { + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + pos: usize, + ) -> Result { + match self { + Self::Recurrent(b) => b.forward(xs, pos), + Self::Attention(b) => b.forward(xs, attention_mask, pos), + } + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + temporal_pre_norm: RmsNorm, + channel_pre_norm: RmsNorm, + temporal_block: TemporalBlock, + mlp_block: Mlp, +} + +impl DecoderLayer { + fn new( + block_idx: usize, + rotary_emb: Arc, + cfg: &Config, + vb: VarBuilder, + ) -> Result { + let h = cfg.hidden_size; + let temporal_pre_norm = rms_norm(h, cfg.rms_norm_eps, vb.pp("temporal_pre_norm"))?; + let channel_pre_norm = rms_norm(h, cfg.rms_norm_eps, vb.pp("channel_pre_norm"))?; + let temporal_block = match cfg.block_types[block_idx % cfg.block_types.len()] { + TemporalBlockType::Recurrent => { + let block = RecurrentBlock::new(cfg, vb.pp("temporal_block"))?; + TemporalBlock::Recurrent(block) + } + TemporalBlockType::Attention => { + let block = SdpaAttention::new(rotary_emb, cfg, vb.pp("temporal_block"))?; + TemporalBlock::Attention(block) + } + }; + let mlp_block = Mlp::new(cfg, vb.pp("mlp_block"))?; + Ok(Self { + temporal_pre_norm, + channel_pre_norm, + temporal_block, + mlp_block, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + pos: usize, + ) -> Result { + let residual = xs; + let xs = xs.apply(&self.temporal_pre_norm)?; + let xs = self.temporal_block.forward(&xs, attention_mask, pos)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.channel_pre_norm)?.apply(&self.mlp_block)?; + xs + residual + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: Embedding, + layers: Vec, + final_norm: RmsNorm, + lm_head: Linear, + hidden_size: usize, + logits_soft_cap: f64, + device: Device, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let embed_tokens = Embedding::new(cfg.vocab_size, cfg.hidden_size, vb.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(DType::F32, cfg, vb.device())?); + let vb_b = vb.pp("layers"); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + for idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(idx, rotary_emb.clone(), cfg, vb_b.pp(idx))?; + layers.push(layer) + } + let final_norm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("final_norm"))?; + let lm_head = linear( + cfg.hidden_size, + cfg.vocab_size, + false, + vb.pp("embed_tokens"), + )?; + Ok(Self { + embed_tokens, + layers, + final_norm, + lm_head, + hidden_size: cfg.hidden_size, + logits_soft_cap: cfg.logits_soft_cap, + device: vb.device().clone(), + }) + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { f32::NEG_INFINITY } else { 0. })) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(DType::F32) + } + + pub fn forward(&mut self, xs: &Tensor, pos: usize) -> Result { + let (b_size, seq_len) = xs.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(b_size, seq_len, pos)?; + Some(mask) + }; + let xs = xs.apply(&self.embed_tokens)?; + let mut xs = (xs * (self.hidden_size as f64).sqrt())?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), pos)?; + } + let logits = xs + .narrow(1, seq_len - 1, 1)? + .apply(&self.final_norm)? + .apply(&self.lm_head)?; + let logits = ((logits / self.logits_soft_cap)?.tanh()? * self.logits_soft_cap)?; + Ok(logits) + } +} diff --git a/patches/candle-transformers/src/models/quantized_rwkv_v5.rs b/patches/candle-transformers/src/models/quantized_rwkv_v5.rs new file mode 100644 index 0000000000..cc5204bf24 --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_rwkv_v5.rs @@ -0,0 +1,303 @@ +//! RWKV v5 model implementation with quantization support. +//! +//! RWKV v5 is an attention-free language model optimized for efficiency. +//! This implementation provides quantization for reduced memory and compute. +//! +//! Key characteristics: +//! - Linear attention mechanism +//! - GroupNorm layer normalization +//! - Time-mixing layers +//! - State-based sequential processing +//! - Support for 8-bit quantization +//! +//! References: +//! - [RWKV Model](https://github.com/BlinkDL/RWKV-LM) +//! - [RWKV v5 Architecture](https://www.rwkv.com/v5) +//! + +use crate::{ + quantized_nn::{layer_norm, linear_no_bias as linear, Embedding, Linear}, + quantized_var_builder::VarBuilder, +}; +use candle::{IndexOp, Result, Tensor}; +use candle_nn::{GroupNorm, LayerNorm, Module}; + +pub use crate::models::rwkv_v5::{Config, State, Tokenizer}; + +#[derive(Debug, Clone)] +struct SelfAttention { + key: Linear, + receptance: Linear, + value: Linear, + gate: Linear, + output: Linear, + ln_x: candle_nn::GroupNorm, + time_mix_key: Tensor, + time_mix_value: Tensor, + time_mix_receptance: Tensor, + time_decay: Tensor, + time_faaaa: Tensor, + time_mix_gate: Tensor, + layer_id: usize, + n_attn_heads: usize, +} + +impl SelfAttention { + fn new(layer_id: usize, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_size = cfg.hidden_size; + let attn_hidden_size = cfg.attention_hidden_size; + let key = linear(hidden_size, attn_hidden_size, vb.pp("key"))?; + let receptance = linear(hidden_size, attn_hidden_size, vb.pp("receptance"))?; + let value = linear(hidden_size, attn_hidden_size, vb.pp("value"))?; + let gate = linear(hidden_size, attn_hidden_size, vb.pp("gate"))?; + let output = linear(attn_hidden_size, hidden_size, vb.pp("output"))?; + + let vb_x = vb.pp("ln_x"); + let ln_x_weight = vb_x.get(hidden_size, "weight")?.dequantize(vb.device())?; + let ln_x_bias = vb_x.get(hidden_size, "bias")?.dequantize(vb.device())?; + + let ln_x = GroupNorm::new( + ln_x_weight, + ln_x_bias, + hidden_size, + hidden_size / cfg.head_size, + 1e-5, + )?; + + let time_mix_key = vb + .get((1, 1, cfg.hidden_size), "time_mix_key")? + .dequantize(vb.device())?; + let time_mix_value = vb + .get((1, 1, cfg.hidden_size), "time_mix_value")? + .dequantize(vb.device())?; + let time_mix_receptance = vb + .get((1, 1, cfg.hidden_size), "time_mix_receptance")? + .dequantize(vb.device())?; + let n_attn_heads = cfg.hidden_size / cfg.head_size; + let time_decay = vb + .get((n_attn_heads, cfg.head_size), "time_decay")? + .dequantize(vb.device())?; + let time_faaaa = vb + .get((n_attn_heads, cfg.head_size), "time_faaaa")? + .dequantize(vb.device())?; + let time_mix_gate = vb + .get((1, 1, cfg.hidden_size), "time_mix_gate")? + .dequantize(vb.device())?; + Ok(Self { + key, + value, + receptance, + gate, + output, + ln_x, + time_mix_key, + time_mix_value, + time_mix_receptance, + time_decay, + time_faaaa, + time_mix_gate, + layer_id, + n_attn_heads, + }) + } + + pub fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let h = self.time_decay.dim(0)?; + let (b, t, s) = xs.dims3()?; + let s = s / h; + let (receptance, key, value, gate) = { + // extract key-value + let shifted = state.per_layer[self.layer_id].extract_key_value.clone(); + let shifted = if shifted.rank() == 2 { + shifted.unsqueeze(1)? + } else { + shifted + }; + let key = ((xs * &self.time_mix_key)? + &shifted * (1.0 - &self.time_mix_key)?)?; + let value = ((xs * &self.time_mix_value)? + &shifted * (1.0 - &self.time_mix_value)?)?; + let receptance = ((xs * &self.time_mix_receptance)? + + &shifted * (1.0 - &self.time_mix_receptance)?)?; + let gate = ((xs * &self.time_mix_gate)? + &shifted * (1.0 - &self.time_mix_gate)?)?; + + let key = self.key.forward(&key)?; + let value = self.value.forward(&value)?; + let receptance = self.receptance.forward(&receptance)?; + let gate = candle_nn::ops::silu(&self.gate.forward(&gate)?)?; + state.per_layer[self.layer_id].extract_key_value = xs.i((.., t - 1))?; + (receptance, key, value, gate) + }; + // linear attention + let mut state_ = state.per_layer[self.layer_id].linear_attention.clone(); + let key = key.reshape((b, t, h, s))?.permute((0, 2, 3, 1))?; + let value = value.reshape((b, t, h, s))?.transpose(1, 2)?; + let receptance = receptance.reshape((b, t, h, s))?.transpose(1, 2)?; + + let time_decay = self + .time_decay + .exp()? + .neg()? + .exp()? + .reshape(((), 1, 1))? + .reshape((self.n_attn_heads, (), 1))?; + let time_faaaa = + self.time_faaaa + .reshape(((), 1, 1))? + .reshape((self.n_attn_heads, (), 1))?; + + let mut out: Vec = Vec::with_capacity(t); + for t_ in 0..t { + let rt = receptance.i((.., .., t_..t_ + 1))?.contiguous()?; + let kt = key.i((.., .., .., t_..t_ + 1))?.contiguous()?; + let vt = value.i((.., .., t_..t_ + 1))?.contiguous()?; + let at = kt.matmul(&vt)?; + let rhs = (time_faaaa.broadcast_mul(&at)? + &state_)?; + let out_ = rt.matmul(&rhs)?.squeeze(2)?; + state_ = (&at + time_decay.broadcast_mul(&state_))?; + out.push(out_) + } + let out = Tensor::cat(&out, 1)?.reshape((b * t, h * s, 1))?; + let out = out.apply(&self.ln_x)?.reshape((b, t, h * s))?; + let out = (out * gate)?.apply(&self.output)?; + state.per_layer[self.layer_id].linear_attention = state_; + Ok(out) + } +} + +#[derive(Debug, Clone)] +struct FeedForward { + time_mix_key: Tensor, + time_mix_receptance: Tensor, + key: Linear, + receptance: Linear, + value: Linear, + layer_id: usize, +} + +impl FeedForward { + fn new(layer_id: usize, cfg: &Config, vb: VarBuilder) -> Result { + let int_size = cfg + .intermediate_size + .unwrap_or(((cfg.hidden_size as f64 * 3.5) as usize) / 32 * 32); + let key = linear(cfg.hidden_size, int_size, vb.pp("key"))?; + let receptance = linear(cfg.hidden_size, cfg.hidden_size, vb.pp("receptance"))?; + let value = linear(int_size, cfg.hidden_size, vb.pp("value"))?; + let time_mix_key = vb + .get((1, 1, cfg.hidden_size), "time_mix_key")? + .dequantize(vb.device())?; + let time_mix_receptance = vb + .get((1, 1, cfg.hidden_size), "time_mix_receptance")? + .dequantize(vb.device())?; + Ok(Self { + key, + receptance, + value, + time_mix_key, + time_mix_receptance, + layer_id, + }) + } + + fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let shifted = &state.per_layer[self.layer_id].feed_forward; + let key = (xs.broadcast_mul(&self.time_mix_key)? + + shifted.broadcast_mul(&(1.0 - &self.time_mix_key)?)?)?; + let receptance = (xs.broadcast_mul(&self.time_mix_receptance)? + + shifted.broadcast_mul(&(1.0 - &self.time_mix_receptance)?)?)?; + let key = key.apply(&self.key)?.relu()?.sqr()?; + let value = key.apply(&self.value)?; + let receptance = candle_nn::ops::sigmoid(&receptance.apply(&self.receptance)?)?; + state.per_layer[self.layer_id].feed_forward = xs.i((.., xs.dim(1)? - 1))?; + let xs = (receptance * value)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct Block { + pre_ln: Option, + ln1: LayerNorm, + ln2: LayerNorm, + attention: SelfAttention, + feed_forward: FeedForward, +} + +impl Block { + fn new(layer_id: usize, cfg: &Config, vb: VarBuilder) -> Result { + let ln1 = layer_norm(cfg.hidden_size, cfg.layer_norm_epsilon, vb.pp("ln1"))?; + let ln2 = layer_norm(cfg.hidden_size, cfg.layer_norm_epsilon, vb.pp("ln2"))?; + let pre_ln = if layer_id == 0 { + let ln = layer_norm(cfg.hidden_size, cfg.layer_norm_epsilon, vb.pp("pre_ln"))?; + Some(ln) + } else { + None + }; + let attention = SelfAttention::new(layer_id, cfg, vb.pp("attention"))?; + let feed_forward = FeedForward::new(layer_id, cfg, vb.pp("feed_forward"))?; + Ok(Self { + pre_ln, + ln1, + ln2, + attention, + feed_forward, + }) + } + + fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let xs = match self.pre_ln.as_ref() { + None => xs.clone(), + Some(pre_ln) => xs.apply(pre_ln)?, + }; + let attention = self.attention.forward(&xs.apply(&self.ln1)?, state)?; + let xs = (xs + attention)?; + let feed_forward = self.feed_forward.forward(&xs.apply(&self.ln2)?, state)?; + let xs = (xs + feed_forward)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embeddings: Embedding, + blocks: Vec, + ln_out: LayerNorm, + head: Linear, + rescale_every: usize, + layers_are_rescaled: bool, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("rwkv"); + let embeddings = Embedding::new(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embeddings"))?; + let mut blocks = Vec::with_capacity(cfg.num_hidden_layers); + let vb_b = vb_m.pp("blocks"); + for block_index in 0..cfg.num_hidden_layers { + let block = Block::new(block_index, cfg, vb_b.pp(block_index))?; + blocks.push(block) + } + let ln_out = layer_norm(cfg.hidden_size, 1e-5, vb_m.pp("ln_out"))?; + let head = linear(cfg.hidden_size, cfg.vocab_size, vb.pp("head"))?; + Ok(Self { + embeddings, + blocks, + ln_out, + head, + rescale_every: cfg.rescale_every, + layers_are_rescaled: false, // This seem to only happen for the f16/bf16 dtypes. + }) + } + + pub fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let (_b_size, _seq_len) = xs.dims2()?; + let mut xs = xs.apply(&self.embeddings)?; + for (block_idx, block) in self.blocks.iter().enumerate() { + xs = block.forward(&xs, state)?; + if self.layers_are_rescaled && (block_idx + 1) % self.rescale_every == 0 { + xs = (xs / 2.)? + } + } + let xs = xs.apply(&self.ln_out)?.apply(&self.head)?; + state.pos += 1; + Ok(xs) + } +} diff --git a/patches/candle-transformers/src/models/quantized_rwkv_v6.rs b/patches/candle-transformers/src/models/quantized_rwkv_v6.rs new file mode 100644 index 0000000000..91288c2e61 --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_rwkv_v6.rs @@ -0,0 +1,350 @@ +//! RWKV v6 model implementation with quantization support. +//! +//! RWKV is a linear attention model that combines the efficiency of RNNs +//! with the parallelizable training of Transformers. Version 6 builds on previous +//! versions with further optimizations. +//! +//! Key characteristics: +//! - Linear attention mechanism +//! - Time mixing layers +//! - Channel mixing layers +//! - RMSNorm for normalization +//! - Support for 8-bit quantization +//! +//! References: +//! - [RWKV Architecture](https://github.com/BlinkDL/RWKV-LM) +//! - [RWKV v6 Release](https://huggingface.co/BlinkDL/rwkv-6) +//! + +use crate::{ + quantized_nn::{layer_norm, linear_no_bias as linear, Embedding, Linear}, + quantized_var_builder::VarBuilder, +}; +use candle::{IndexOp, Result, Tensor}; +use candle_nn::{GroupNorm, LayerNorm, Module}; + +pub use crate::models::rwkv_v5::{Config, State, Tokenizer}; + +#[derive(Debug, Clone)] +struct SelfAttention { + key: Linear, + receptance: Linear, + value: Linear, + gate: Linear, + output: Linear, + ln_x: candle_nn::GroupNorm, + time_mix_x: Tensor, + time_mix_w: Tensor, + time_mix_key: Tensor, + time_mix_value: Tensor, + time_mix_receptance: Tensor, + time_decay: Tensor, + time_faaaa: Tensor, + time_mix_gate: Tensor, + time_decay_w1: Tensor, + time_decay_w2: Tensor, + time_mix_w1: Tensor, + time_mix_w2: Tensor, + layer_id: usize, + n_attn_heads: usize, +} + +impl SelfAttention { + fn new(layer_id: usize, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_size = cfg.hidden_size; + let attn_hidden_size = cfg.attention_hidden_size; + let key = linear(hidden_size, attn_hidden_size, vb.pp("key"))?; + let receptance = linear(hidden_size, attn_hidden_size, vb.pp("receptance"))?; + let value = linear(hidden_size, attn_hidden_size, vb.pp("value"))?; + let gate = linear(hidden_size, attn_hidden_size, vb.pp("gate"))?; + let output = linear(attn_hidden_size, hidden_size, vb.pp("output"))?; + + let vb_x = vb.pp("ln_x"); + let ln_x_weight = vb_x.get(hidden_size, "weight")?.dequantize(vb.device())?; + let ln_x_bias = vb_x.get(hidden_size, "bias")?.dequantize(vb.device())?; + + let ln_x = GroupNorm::new( + ln_x_weight, + ln_x_bias, + hidden_size, + hidden_size / cfg.head_size, + 1e-5, + )?; + + let time_mix_x = vb + .get((1, 1, cfg.hidden_size), "time_mix_x")? + .dequantize(vb.device())?; + let time_mix_w = vb + .get((1, 1, cfg.hidden_size), "time_mix_w")? + .dequantize(vb.device())?; + let time_mix_key = vb + .get((1, 1, cfg.hidden_size), "time_mix_key")? + .dequantize(vb.device())?; + let time_mix_value = vb + .get((1, 1, cfg.hidden_size), "time_mix_value")? + .dequantize(vb.device())?; + let time_mix_receptance = vb + .get((1, 1, cfg.hidden_size), "time_mix_receptance")? + .dequantize(vb.device())?; + let n_attn_heads = cfg.hidden_size / cfg.head_size; + let time_decay = vb + .get((1, 1, cfg.hidden_size), "time_decay")? + .dequantize(vb.device())?; + let time_faaaa = vb + .get((n_attn_heads, cfg.head_size), "time_faaaa")? + .dequantize(vb.device())?; + let time_mix_gate = vb + .get((1, 1, cfg.hidden_size), "time_mix_gate")? + .dequantize(vb.device())?; + let time_decay_w1 = vb + .get((cfg.hidden_size, n_attn_heads * 2), "time_decay_w1")? + .dequantize(vb.device())?; + let time_decay_w2 = vb + .get((n_attn_heads * 2, cfg.hidden_size), "time_decay_w2")? + .dequantize(vb.device())?; + let time_mix_w1 = vb + .get((cfg.hidden_size, n_attn_heads * 5), "time_mix_w1")? + .dequantize(vb.device())?; + let time_mix_w2 = vb + .get((5, n_attn_heads, cfg.hidden_size), "time_mix_w2")? + .dequantize(vb.device())?; + Ok(Self { + key, + value, + receptance, + gate, + output, + ln_x, + time_mix_x, + time_mix_w, + time_mix_key, + time_mix_value, + time_mix_receptance, + time_decay, + time_faaaa, + time_mix_gate, + time_decay_w1, + time_decay_w2, + time_mix_w1, + time_mix_w2, + layer_id, + n_attn_heads, + }) + } + + pub fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let h = self.n_attn_heads; + let (b, t, s) = xs.dims3()?; + let s = s / h; + let (receptance, key, value, gate, w) = { + // extract key-value + let shifted = state.per_layer[self.layer_id].extract_key_value.clone(); + let shifted = if shifted.rank() == 2 { + shifted.unsqueeze(1)? + } else { + shifted + }; + + let sx = (&shifted - xs)?; + let xxx = (xs + &sx * &self.time_mix_x)?; + let xxx = xxx + .broadcast_matmul(&self.time_mix_w1)? + .tanh()? + .reshape((b * t, 5, ()))? + .transpose(0, 1)?; + + let xxx = xxx.matmul(&self.time_mix_w2)?.reshape((5, b, t, ()))?; + + let (mw, mk, mv, mr, mg) = (xxx.i(0)?, xxx.i(1)?, xxx.i(2)?, xxx.i(3)?, xxx.i(4)?); + + let xw = (xs + &sx * (&self.time_mix_w + &mw)?)?; + let xk = (xs + &sx * (&self.time_mix_key + &mk)?)?; + let xv = (xs + &sx * (&self.time_mix_value + &mv)?)?; + let xr = (xs + &sx * (&self.time_mix_receptance + &mr)?)?; + let xg = (xs + &sx * (&self.time_mix_gate + &mg)?)?; + + let w = (&self.time_decay + + xw.broadcast_matmul(&self.time_decay_w1)? + .tanh()? + .broadcast_matmul(&self.time_decay_w2)?)? + .reshape(((), 1, 1))? + .reshape((self.n_attn_heads, (), 1))?; + + let key = self.key.forward(&xk)?; + let value = self.value.forward(&xv)?; + let receptance = self.receptance.forward(&xr)?; + let gate = candle_nn::ops::silu(&self.gate.forward(&xg)?)?; + state.per_layer[self.layer_id].extract_key_value = xs.i((.., t - 1))?; + (receptance, key, value, gate, w) + }; + + // linear attention + let mut state_ = state.per_layer[self.layer_id].linear_attention.clone(); + let key = key.reshape((b, t, h, s))?.permute((0, 2, 3, 1))?; + let value = value.reshape((b, t, h, s))?.transpose(1, 2)?; + let receptance = receptance.reshape((b, t, h, s))?.transpose(1, 2)?; + + let w = w.exp()?.neg()?.exp()?; + + let time_faaaa = + self.time_faaaa + .reshape(((), 1, 1))? + .reshape((self.n_attn_heads, (), 1))?; + + let mut out: Vec = Vec::with_capacity(t); + for t_ in 0..t { + let rt = receptance.i((.., .., t_..t_ + 1))?.contiguous()?; + let kt = key.i((.., .., .., t_..t_ + 1))?.contiguous()?; + let vt = value.i((.., .., t_..t_ + 1))?.contiguous()?; + let at = kt.matmul(&vt)?; + let rhs = (time_faaaa.broadcast_mul(&at)? + &state_)?; + let out_ = rt.matmul(&rhs)?.squeeze(2)?; + state_ = (&at + w.broadcast_mul(&state_))?; + out.push(out_) + } + let out = Tensor::cat(&out, 1)?.reshape((b * t, h * s, 1))?; + let out = out.apply(&self.ln_x)?.reshape((b, t, h * s))?; + let out = (out * gate)?.apply(&self.output)?; + state.per_layer[self.layer_id].linear_attention = state_; + Ok(out) + } +} + +#[derive(Debug, Clone)] +struct FeedForward { + time_mix_key: Tensor, + time_mix_receptance: Tensor, + key: Linear, + receptance: Linear, + value: Linear, + layer_id: usize, +} + +impl FeedForward { + fn new(layer_id: usize, cfg: &Config, vb: VarBuilder) -> Result { + let int_size = cfg + .intermediate_size + .unwrap_or(((cfg.hidden_size as f64 * 3.5) as usize) / 32 * 32); + let key = linear(cfg.hidden_size, int_size, vb.pp("key"))?; + let receptance = linear(cfg.hidden_size, cfg.hidden_size, vb.pp("receptance"))?; + let value = linear(int_size, cfg.hidden_size, vb.pp("value"))?; + let time_mix_key = vb + .get((1, 1, cfg.hidden_size), "time_mix_key")? + .dequantize(vb.device())?; + let time_mix_receptance = vb + .get((1, 1, cfg.hidden_size), "time_mix_receptance")? + .dequantize(vb.device())?; + Ok(Self { + key, + receptance, + value, + time_mix_key, + time_mix_receptance, + layer_id, + }) + } + + fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let shifted = state.per_layer[self.layer_id] + .feed_forward + .broadcast_sub(xs)?; + let key = (xs + shifted.broadcast_mul(&self.time_mix_key)?)?; + let receptance = (xs + shifted.broadcast_mul(&self.time_mix_receptance)?)?; + let key = key.apply(&self.key)?.relu()?.sqr()?; + let value = key.apply(&self.value)?; + let receptance = candle_nn::ops::sigmoid(&receptance.apply(&self.receptance)?)?; + state.per_layer[self.layer_id].feed_forward = xs.i((.., xs.dim(1)? - 1))?; + let xs = (receptance * value)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct Block { + pre_ln: Option, + ln1: LayerNorm, + ln2: LayerNorm, + attention: SelfAttention, + feed_forward: FeedForward, +} + +impl Block { + fn new(layer_id: usize, cfg: &Config, vb: VarBuilder) -> Result { + let ln1 = layer_norm(cfg.hidden_size, cfg.layer_norm_epsilon, vb.pp("ln1"))?; + let ln2 = layer_norm(cfg.hidden_size, cfg.layer_norm_epsilon, vb.pp("ln2"))?; + let pre_ln = if layer_id == 0 { + let ln = layer_norm(cfg.hidden_size, cfg.layer_norm_epsilon, vb.pp("pre_ln"))?; + Some(ln) + } else { + None + }; + let attention = SelfAttention::new(layer_id, cfg, vb.pp("attention"))?; + let feed_forward = FeedForward::new(layer_id, cfg, vb.pp("feed_forward"))?; + Ok(Self { + pre_ln, + ln1, + ln2, + attention, + feed_forward, + }) + } + + fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let xs = match self.pre_ln.as_ref() { + None => xs.clone(), + Some(pre_ln) => xs.apply(pre_ln)?, + }; + let attention = self.attention.forward(&xs.apply(&self.ln1)?, state)?; + let xs = (xs + attention)?; + let feed_forward = self.feed_forward.forward(&xs.apply(&self.ln2)?, state)?; + let xs = (xs + feed_forward)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embeddings: Embedding, + blocks: Vec, + ln_out: LayerNorm, + head: Linear, + rescale_every: usize, + layers_are_rescaled: bool, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("rwkv"); + let embeddings = Embedding::new(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embeddings"))?; + let mut blocks = Vec::with_capacity(cfg.num_hidden_layers); + let vb_b = vb_m.pp("blocks"); + for block_index in 0..cfg.num_hidden_layers { + let block = Block::new(block_index, cfg, vb_b.pp(block_index))?; + blocks.push(block) + } + let ln_out = layer_norm(cfg.hidden_size, 1e-5, vb_m.pp("ln_out"))?; + let head = linear(cfg.hidden_size, cfg.vocab_size, vb.pp("head"))?; + Ok(Self { + embeddings, + blocks, + ln_out, + head, + rescale_every: cfg.rescale_every, + layers_are_rescaled: false, // This seem to only happen for the f16/bf16 dtypes. + }) + } + + pub fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let (_b_size, _seq_len) = xs.dims2()?; + let mut xs = xs.apply(&self.embeddings)?; + for (block_idx, block) in self.blocks.iter().enumerate() { + xs = block.forward(&xs, state)?; + if self.layers_are_rescaled && (block_idx + 1) % self.rescale_every == 0 { + xs = (xs / 2.)? + } + } + let xs = xs.apply(&self.ln_out)?.apply(&self.head)?; + state.pos += 1; + Ok(xs) + } +} diff --git a/patches/candle-transformers/src/models/quantized_stable_lm.rs b/patches/candle-transformers/src/models/quantized_stable_lm.rs new file mode 100644 index 0000000000..d74ed743d8 --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_stable_lm.rs @@ -0,0 +1,301 @@ +//! Module for quantized StableLM implementation. +//! +//! StableLM is a series of open-source large language models +//! optimized for performance and stability. This implementation +//! provides quantization support for efficient model deployment. +//! +//! Key characteristics: +//! - RMSNorm for layer normalization +//! - Rotary positional embeddings (RoPE) +//! - Support for 8-bit quantization +//! +//! References: +//! - [StableLM](https://github.com/Stability-AI/StableLM) +//! + +use crate::quantized_nn::{layer_norm, linear, linear_no_bias, Embedding, Linear}; +pub use crate::quantized_var_builder::VarBuilder; +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::{Activation, LayerNorm}; +use std::sync::Arc; + +pub use crate::models::stable_lm::Config; +use crate::models::stable_lm::RotaryEmbedding; + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: Activation, + span: tracing::Span, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + let gate_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("gate_proj"))?; + let up_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("up_proj"))?; + let down_proj = linear_no_bias(intermediate_sz, hidden_sz, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_act, + span: tracing::span!(tracing::Level::TRACE, "mlp"), + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = xs.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + hidden_size: usize, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, + use_cache: bool, + rotary_ndims: usize, + span: tracing::Span, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let head_dim = cfg.head_dim(); + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let linear_layer = if cfg.use_qkv_bias { + linear + } else { + linear_no_bias + }; + let q_proj = linear_layer(hidden_sz, num_heads * head_dim, vb.pp("q_proj"))?; + let k_proj = linear_layer(hidden_sz, num_kv_heads * head_dim, vb.pp("k_proj"))?; + let v_proj = linear_layer(hidden_sz, num_kv_heads * head_dim, vb.pp("v_proj"))?; + let o_proj = linear_no_bias(num_heads * head_dim, hidden_sz, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups: cfg.num_kv_groups(), + head_dim, + hidden_size: hidden_sz, + rotary_emb, + kv_cache: None, + use_cache: cfg.use_cache, + rotary_ndims: cfg.rotary_ndims(), + span: tracing::span!(tracing::Level::TRACE, "attn"), + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let _enter = self.span.enter(); + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let (rot_ndims, pass_ndims) = (self.rotary_ndims, self.head_dim - self.rotary_ndims); + let query_rot = query_states.narrow(D::Minus1, 0, rot_ndims)?; + let query_pass = query_states.narrow(D::Minus1, rot_ndims, pass_ndims)?; + let key_rot = key_states.narrow(D::Minus1, 0, rot_ndims)?; + let key_pass = key_states.narrow(D::Minus1, rot_ndims, pass_ndims)?; + let (query_rot, key_rot) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_rot, &key_rot, seqlen_offset)?; + let query_states = Tensor::cat(&[query_rot, query_pass], D::Minus1)?.contiguous()?; + let key_states = Tensor::cat(&[key_rot, key_pass], D::Minus1)?.contiguous()?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + if self.use_cache { + self.kv_cache = Some((key_states.clone(), value_states.clone())); + } + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?.contiguous()?; + let value_states = + crate::utils::repeat_kv(value_states, self.num_kv_groups)?.contiguous()?; + + let attn_output = { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + }; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.hidden_size))? + .apply(&self.o_proj) + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: MLP, + input_layernorm: LayerNorm, + post_attention_layernorm: LayerNorm, + span: tracing::Span, +} + +impl DecoderLayer { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let self_attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let input_layernorm = layer_norm( + cfg.hidden_size, + cfg.layer_norm_eps, + vb.pp("input_layernorm"), + )?; + let post_attention_layernorm = layer_norm( + cfg.hidden_size, + cfg.layer_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + span: tracing::span!(tracing::Level::TRACE, "layer"), + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let _enter = self.span.enter(); + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; + residual + xs + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: Embedding, + layers: Vec, + norm: LayerNorm, + lm_head: Linear, + device: Device, + span: tracing::Span, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + Embedding::new(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(DType::F32, cfg, vb_m.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(rotary_emb.clone(), cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb_m.pp("norm"))?; + let lm_head = linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))?; + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + device: vb.device().clone(), + span: tracing::span!(tracing::Level::TRACE, "model"), + }) + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + // Sliding window mask? + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { f32::NEG_INFINITY } else { 0. })) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(DType::F32) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let _enter = self.span.enter(); + let (b_size, seq_len) = input_ids.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(b_size, seq_len, seqlen_offset)?; + Some(mask) + }; + let mut xs = self.embed_tokens.forward(input_ids)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)? + } + xs.narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head) + } +} diff --git a/patches/candle-transformers/src/models/quantized_t5.rs b/patches/candle-transformers/src/models/quantized_t5.rs new file mode 100644 index 0000000000..4fc9c537f8 --- /dev/null +++ b/patches/candle-transformers/src/models/quantized_t5.rs @@ -0,0 +1,799 @@ +//! T5 model implementation with quantization support. +//! +//! T5 is an encoder-decoder model pre-trained on a multi-task mixture of supervised +//! and unsupervised tasks. This implementation provides quantization for reduced +//! memory and compute requirements. +//! +//! Key characteristics: +//! - Encoder-decoder architecture +//! - Layer normalization +//! - Relative positional encodings +//! - Support for 8-bit quantization +//! +//! References: +//! - 📝 [T5 Paper](https://arxiv.org/abs/1910.10683) +//! - 🤗 [Model Card](https://huggingface.co/t5-base) +//! - 🤗 Original model from [T5](https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/modeling_t5.py) + +use crate::models::t5::{deserialize_feed_forward_proj_activation, ActivationWithOptionalGating}; +use crate::models::with_tracing::QMatMul; +use crate::quantized_nn::Embedding; +pub use crate::quantized_var_builder::VarBuilder; +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::Activation; +use serde::Deserialize; +use std::sync::Arc; + +fn default_relative_attention_max_distance() -> usize { + 128 +} + +fn default_is_decoder() -> bool { + false +} + +fn default_use_cache() -> bool { + true +} + +fn default_tie_word_embeddings() -> bool { + true +} + +fn get_mask(size: usize, device: &Device) -> Result { + let mask: Vec<_> = (0..size) + .flat_map(|i| (0..size).map(move |j| u8::from(j > i))) + .collect(); + Tensor::from_slice(&mask, (size, size), device) +} + +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) +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Config { + vocab_size: usize, + d_model: usize, + d_kv: usize, + d_ff: usize, + num_layers: usize, + num_decoder_layers: Option, + num_heads: usize, + relative_attention_num_buckets: usize, + #[serde(default = "default_relative_attention_max_distance")] + relative_attention_max_distance: usize, + dropout_rate: f64, + layer_norm_epsilon: f64, + initializer_factor: f64, + #[serde(default, deserialize_with = "deserialize_feed_forward_proj_activation")] + pub feed_forward_proj: ActivationWithOptionalGating, + #[serde(default = "default_tie_word_embeddings")] + tie_word_embeddings: bool, + #[serde(default = "default_is_decoder")] + is_decoder: bool, + is_encoder_decoder: bool, + #[serde(default = "default_use_cache")] + pub use_cache: bool, + pub pad_token_id: usize, + pub eos_token_id: usize, + pub decoder_start_token_id: Option, +} + +impl Default for Config { + fn default() -> Self { + Self { + vocab_size: 32128, + d_model: 512, + d_kv: 64, + d_ff: 2048, + num_layers: 6, + num_decoder_layers: None, + num_heads: 8, + relative_attention_num_buckets: 32, + relative_attention_max_distance: 128, + dropout_rate: 0.1, + layer_norm_epsilon: 1e-6, + initializer_factor: 1.0, + feed_forward_proj: ActivationWithOptionalGating { + gated: false, + activation: Activation::Relu, + }, + tie_word_embeddings: true, + is_decoder: false, + is_encoder_decoder: true, + use_cache: true, + pad_token_id: 0, + eos_token_id: 1, + decoder_start_token_id: Some(0), + } + } +} + +#[derive(Debug, Clone)] +struct T5LayerNorm { + weight: Tensor, + variance_epsilon: f64, + span: tracing::Span, +} + +impl T5LayerNorm { + fn load(h: usize, eps: f64, vb: VarBuilder) -> Result { + let weight = vb.get(h, "weight")?.dequantize(vb.device())?; + Ok(Self { + weight, + variance_epsilon: eps, + span: tracing::span!(tracing::Level::TRACE, "layer-norm"), + }) + } +} + +impl Module for T5LayerNorm { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let dtype = xs.dtype(); + let xs_f32 = xs.to_dtype(DType::F32)?; + // variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) + let variance = xs_f32.sqr()?.mean_keepdim(D::Minus1)?; + let xs = xs.broadcast_div(&(variance + self.variance_epsilon)?.sqrt()?)?; + let xs = xs.to_dtype(dtype)?; + let xs = xs.broadcast_mul(&self.weight)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct T5DenseActDense { + wi: QMatMul, + wo: QMatMul, + act: Activation, + span: tracing::Span, +} + +impl T5DenseActDense { + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let wi = QMatMul::new(cfg.d_model, cfg.d_ff, vb.pp("wi"))?; + let wo = QMatMul::new(cfg.d_ff, cfg.d_model, vb.pp("wo"))?; + Ok(Self { + wi, + wo, + act: Activation::Relu, + span: tracing::span!(tracing::Level::TRACE, "dense-act-dense"), + }) + } +} + +impl Module for T5DenseActDense { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let xs = self.wi.forward(xs)?; + let xs = self.act.forward(&xs)?; + let xs = self.wo.forward(&xs)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct T5DenseGatedActDense { + wi_0: QMatMul, + wi_1: QMatMul, + wo: QMatMul, + act: Activation, + span: tracing::Span, +} + +impl T5DenseGatedActDense { + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let wi_0 = QMatMul::new(cfg.d_model, cfg.d_ff, vb.pp("wi_0"))?; + let wi_1 = QMatMul::new(cfg.d_model, cfg.d_ff, vb.pp("wi_1"))?; + let wo = QMatMul::new(cfg.d_ff, cfg.d_model, vb.pp("wo"))?; + Ok(Self { + wi_0, + wi_1, + wo, + act: cfg.feed_forward_proj.activation, + span: tracing::span!(tracing::Level::TRACE, "dense-gated-act-dense"), + }) + } +} + +impl Module for T5DenseGatedActDense { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let hidden_gelu = self.act.forward(&self.wi_0.forward(xs)?)?; + let hidden_linear = self.wi_1.forward(xs)?; + let xs = hidden_gelu.broadcast_mul(&hidden_linear)?; + let xs = self.wo.forward(&xs)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct T5LayerFF { + dense_act: Option, + gated_dense_act: Option, + layer_norm: T5LayerNorm, + span: tracing::Span, +} + +impl T5LayerFF { + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let layer_norm = + T5LayerNorm::load(cfg.d_model, cfg.layer_norm_epsilon, vb.pp("layer_norm"))?; + let (dense_act, gated_dense_act) = if cfg.feed_forward_proj.gated { + ( + None, + Some(T5DenseGatedActDense::load(vb.pp("DenseReluDense"), cfg)?), + ) + } else { + ( + Some(T5DenseActDense::load(vb.pp("DenseReluDense"), cfg)?), + None, + ) + }; + Ok(Self { + dense_act, + gated_dense_act, + layer_norm, + span: tracing::span!(tracing::Level::TRACE, "layer-ff"), + }) + } +} + +impl Module for T5LayerFF { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let ys = self.layer_norm.forward(xs)?; + let ys = match &self.dense_act { + Some(dense_act) => dense_act.forward(&ys)?, + None => self.gated_dense_act.as_ref().unwrap().forward(&ys)?, + }; + let xs = (xs + ys)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct T5Attention { + q: QMatMul, + k: QMatMul, + v: QMatMul, + o: QMatMul, + n_heads: usize, + d_kv: usize, + relative_attention_bias: Option, + relative_attention_num_buckets: usize, + relative_attention_max_distance: usize, + inner_dim: usize, + use_cache: bool, + kv_cache: Option<(Tensor, Tensor)>, + span: tracing::Span, + span_cache: tracing::Span, + span_mm: tracing::Span, + span_sm: tracing::Span, +} + +impl T5Attention { + fn load( + has_relative_attention_bias: bool, + decoder: bool, + vb: VarBuilder, + cfg: &Config, + ) -> Result { + let inner_dim = cfg.num_heads * cfg.d_kv; + let q = QMatMul::new(cfg.d_model, inner_dim, vb.pp("q"))?; + let k = QMatMul::new(cfg.d_model, inner_dim, vb.pp("k"))?; + let v = QMatMul::new(cfg.d_model, inner_dim, vb.pp("v"))?; + let o = QMatMul::new(inner_dim, cfg.d_model, vb.pp("o"))?; + let relative_attention_bias = if has_relative_attention_bias { + let emb = Embedding::new( + cfg.relative_attention_num_buckets, + cfg.num_heads, + vb.pp("relative_attention_bias"), + )?; + Some(emb) + } else { + None + }; + Ok(Self { + q, + k, + v, + o, + n_heads: cfg.num_heads, + d_kv: cfg.d_kv, + relative_attention_bias, + relative_attention_num_buckets: cfg.relative_attention_num_buckets, + relative_attention_max_distance: cfg.relative_attention_max_distance, + inner_dim, + use_cache: cfg.use_cache && decoder, + kv_cache: None, + span: tracing::span!(tracing::Level::TRACE, "attention"), + span_cache: tracing::span!(tracing::Level::TRACE, "attention-cache"), + span_mm: tracing::span!(tracing::Level::TRACE, "attention-mm"), + span_sm: tracing::span!(tracing::Level::TRACE, "attention-sm"), + }) + } + + fn forward( + &mut self, + xs: &Tensor, + position_bias: Option<&Tensor>, + key_value_states: Option<&Tensor>, + mask: Option<&Tensor>, + ) -> Result<(Tensor, Option)> { + // Performs Self-attention (if key_value_states is None) or attention + // over source sentence (provided by key_value_states). + let _enter = self.span.enter(); + let kv_input = match key_value_states { + None => xs, + Some(key_value_states) => key_value_states, + }; + let (b_sz, q_len) = (xs.dim(0)?, xs.dim(1)?); + let kv_len = kv_input.dim(1)?; + let q = self.q.forward(xs)?; + let k = self.k.forward(kv_input)?; + let v = self.v.forward(kv_input)?; + let q = q + .reshape((b_sz, q_len, self.n_heads, self.d_kv))? + .transpose(1, 2)? + .contiguous()?; + let mut k = k + .reshape((b_sz, kv_len, self.n_heads, self.d_kv))? + .transpose(1, 2)?; + let mut v = v + .reshape((b_sz, kv_len, self.n_heads, self.d_kv))? + .transpose(1, 2)?; + + if self.use_cache && key_value_states.is_none() { + let _enter = self.span_cache.enter(); + if let Some((kv_cache_k, kv_cache_v)) = &self.kv_cache { + k = Tensor::cat(&[kv_cache_k, &k], 2)?; + v = Tensor::cat(&[kv_cache_v, &v], 2)?; + }; + self.kv_cache = Some((k.clone(), v.clone())); + }; + let k = k.contiguous()?; + let v = v.contiguous()?; + // TODO: Use flash_attn. + let scores = { + let _enter = self.span_mm.enter(); + q.matmul(&k.t()?)? + }; + let scores = match mask { + None => scores, + Some(mask) => masked_fill( + &scores, + &mask + .unsqueeze(0)? + .unsqueeze(0)? + .repeat((b_sz, self.n_heads))?, + f32::NEG_INFINITY, + )?, + }; + + let (scores, position_bias) = match position_bias { + Some(position_bias) => ( + scores.broadcast_add(position_bias)?, + Some(position_bias.clone()), + ), + None => match &self.relative_attention_bias { + None => (scores, None), + Some(relative_attention_bias) => { + // This only handles the bidirectional case. + let kv_len = k.dim(2)?; + let (q_start, q_end) = match self.use_cache { + true => ((kv_len - q_len) as u32, kv_len as u32), + false => (0_u32, kv_len as u32), + }; + let num_buckets = self.relative_attention_num_buckets as u32 / 2; + let max_exact = num_buckets / 2; + let relative_position = (q_start..q_end) + .map(|i| { + (0..kv_len as u32) + .map(|j| { + if i < j { + if j - i < max_exact { + j - i + num_buckets + } else { + let b = f32::log( + (j - i) as f32 / max_exact as f32, + self.relative_attention_max_distance as f32 + / max_exact as f32, + ) * (num_buckets - max_exact) as f32; + u32::min( + max_exact + num_buckets + b as u32, + self.relative_attention_num_buckets as u32 - 1, + ) + } + } else if i - j < max_exact { + i - j + } else { + let b = f32::log( + (i - j) as f32 / max_exact as f32, + self.relative_attention_max_distance as f32 + / max_exact as f32, + ) * (num_buckets - max_exact) as f32; + max_exact + b as u32 + } + }) + .collect::>() + }) + .collect::>>(); + let relative_buckets = Tensor::new(relative_position, q.device())?; + let position_bias = relative_attention_bias + .forward(&relative_buckets)? + .permute((2, 0, 1))? + .unsqueeze(0)?; + (scores.broadcast_add(&position_bias)?, Some(position_bias)) + // TODO: position_bias_masked? + } + }, + }; + + let attn_weights = { + let _enter = self.span_sm.enter(); + candle_nn::ops::softmax_last_dim(&scores)? + }; + let attn_output = attn_weights.matmul(&v)?; + let attn_output = attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.inner_dim))?; + let attn_output = self.o.forward(&attn_output)?; + Ok((attn_output, position_bias)) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +struct T5LayerSelfAttention { + self_attention: T5Attention, + layer_norm: T5LayerNorm, + span: tracing::Span, +} + +impl T5LayerSelfAttention { + fn load(h: bool, d: bool, vb: VarBuilder, cfg: &Config) -> Result { + let self_attention = T5Attention::load(h, d, vb.pp("SelfAttention"), cfg)?; + let layer_norm = + T5LayerNorm::load(cfg.d_model, cfg.layer_norm_epsilon, vb.pp("layer_norm"))?; + Ok(Self { + self_attention, + layer_norm, + span: tracing::span!(tracing::Level::TRACE, "self-attn"), + }) + } + + fn forward( + &mut self, + xs: &Tensor, + position_bias: Option<&Tensor>, + mask: Option<&Tensor>, + ) -> Result<(Tensor, Option)> { + let _enter = self.span.enter(); + let normed_xs = self.layer_norm.forward(xs)?; + let (ys, position_bias) = + self.self_attention + .forward(&normed_xs, position_bias, None, mask)?; + let ys = (xs + ys)?; + Ok((ys, position_bias)) + } + + fn clear_kv_cache(&mut self) { + self.self_attention.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +struct T5LayerCrossAttention { + cross_attention: T5Attention, + layer_norm: T5LayerNorm, + span: tracing::Span, +} + +impl T5LayerCrossAttention { + fn load(decoder: bool, vb: VarBuilder, cfg: &Config) -> Result { + let cross_attention = T5Attention::load(false, decoder, vb.pp("EncDecAttention"), cfg)?; + let layer_norm = + T5LayerNorm::load(cfg.d_model, cfg.layer_norm_epsilon, vb.pp("layer_norm"))?; + Ok(Self { + cross_attention, + layer_norm, + span: tracing::span!(tracing::Level::TRACE, "cross-attn"), + }) + } + + fn forward( + &mut self, + hidden_states: &Tensor, + position_bias: Option<&Tensor>, + key_value_states: &Tensor, + ) -> Result<(Tensor, Option)> { + let _enter = self.span.enter(); + let normed_hidden_states = self.layer_norm.forward(hidden_states)?; + let (ys, position_bias) = self.cross_attention.forward( + &normed_hidden_states, + position_bias, + Some(key_value_states), + None, + )?; + let ys = (hidden_states + ys)?; + Ok((ys, position_bias)) + } + + fn clear_kv_cache(&mut self) { + self.cross_attention.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +struct T5Block { + self_attn: T5LayerSelfAttention, + cross_attn: Option, + ff: T5LayerFF, + span: tracing::Span, +} + +impl T5Block { + fn load( + has_relative_attention_bias: bool, + decoder: bool, + vb: VarBuilder, + cfg: &Config, + ) -> Result { + let vb = vb.pp("layer"); + let self_attn = + T5LayerSelfAttention::load(has_relative_attention_bias, decoder, vb.pp("0"), cfg)?; + let cross_attn = if cfg.is_decoder { + Some(T5LayerCrossAttention::load(decoder, vb.pp("1"), cfg)?) + } else { + None + }; + let ff_i = if cross_attn.is_some() { 2 } else { 1 }; + let ff = T5LayerFF::load(vb.pp(ff_i), cfg)?; + Ok(Self { + self_attn, + cross_attn, + ff, + span: tracing::span!(tracing::Level::TRACE, "block"), + }) + } + + fn forward( + &mut self, + xs: &Tensor, + position_bias: Option<&Tensor>, + encoder_hidden_states: Option<&Tensor>, + ) -> Result<(Tensor, Option)> { + let _enter = self.span.enter(); + // TODO: Cache masks + let mask = match self.cross_attn.is_some() { + true => { + let mask_len = xs.dim(1)?; + // If the input seq length is 1, no need for a mask, this is also helpful to avoid shape + // issues when using the KV cache in the decoder. + if mask_len <= 1 { + None + } else { + Some(get_mask(mask_len, xs.device())?) + } + } + false => None, + }; + let (mut xs, position_bias) = self.self_attn.forward(xs, position_bias, mask.as_ref())?; + // TODO: clamp for f16? + if let Some(cross_attn) = &mut self.cross_attn { + (xs, _) = cross_attn.forward(&xs, None, encoder_hidden_states.unwrap())?; + // TODO: clamp for f16? + } + let xs = self.ff.forward(&xs)?; + // TODO: clamp for f16? + Ok((xs, position_bias)) + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache(); + self.cross_attn.iter_mut().for_each(|c| c.clear_kv_cache()); + } +} + +#[derive(Debug, Clone)] +struct T5Stack { + block: Vec, + shared: Arc, + final_layer_norm: T5LayerNorm, + span: tracing::Span, +} + +impl T5Stack { + fn load(decoder: bool, vb: VarBuilder, shared: &Arc, cfg: &Config) -> Result { + let block = (0..cfg.num_layers) + .map(|i| T5Block::load(i == 0, decoder, vb.pp(format!("block.{i}")), cfg)) + .collect::>>()?; + let final_layer_norm = T5LayerNorm::load( + cfg.d_model, + cfg.layer_norm_epsilon, + vb.pp("final_layer_norm"), + )?; + Ok(Self { + block, + shared: shared.clone(), + final_layer_norm, + span: tracing::span!(tracing::Level::TRACE, "stack"), + }) + } + + fn forward( + &mut self, + input_ids: &Tensor, + encoder_hidden_states: Option<&Tensor>, + ) -> Result { + let _enter = self.span.enter(); + let input_embeds = self.shared.as_ref().forward(input_ids)?; + let mut hidden_states = input_embeds; + let mut position_bias = None; + for block in self.block.iter_mut() { + (hidden_states, position_bias) = block.forward( + &hidden_states, + position_bias.as_ref(), + encoder_hidden_states, + )? + } + self.final_layer_norm.forward(&hidden_states) + } + + fn clear_kv_cache(&mut self) { + self.block.iter_mut().for_each(|b| b.clear_kv_cache()) + } +} + +#[derive(Debug, Clone)] +pub struct T5EncoderModel { + encoder: T5Stack, + device: Device, + span: tracing::Span, +} + +impl T5EncoderModel { + pub fn load(vb: VarBuilder, cfg: &Config) -> Result { + let shared_vb = if vb.contains_key("shared.weight") { + vb.pp("shared") + } else { + vb.pp("decoder").pp("embed_tokens") + }; + let shared = Embedding::new(cfg.vocab_size, cfg.d_model, shared_vb)?; + let shared = Arc::new(shared); + let encoder = T5Stack::load(false, vb.pp("encoder"), &shared, cfg)?; + Ok(Self { + encoder, + device: vb.device().clone(), + span: tracing::span!(tracing::Level::TRACE, "encoder"), + }) + } + + pub fn forward(&mut self, input_ids: &Tensor) -> Result { + let _enter = self.span.enter(); + self.encoder.forward(input_ids, None) + } + + pub fn device(&self) -> &Device { + &self.device + } + + pub fn clear_kv_cache(&mut self) { + self.encoder.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct T5ForConditionalGeneration { + encoder: T5Stack, + decoder: T5Stack, + d_model: usize, + tie_word_embeddings: bool, + lm_head: Option, + shared: Arc, + device: Device, + span_decode: tracing::Span, + span_decode_head: tracing::Span, +} + +impl T5ForConditionalGeneration { + pub fn load(vb: VarBuilder, cfg: &Config) -> Result { + assert!(cfg.is_encoder_decoder); + let d_model = cfg.d_model; + let shared_vb = if vb.contains_key("shared.weight") { + vb.pp("shared") + } else { + vb.pp("decoder").pp("embed_tokens") + }; + let shared = Embedding::new(cfg.vocab_size, cfg.d_model, shared_vb)?; + let shared = Arc::new(shared); + + let mut encoder_cfg = cfg.clone(); + encoder_cfg.is_decoder = false; + encoder_cfg.use_cache = false; + encoder_cfg.is_encoder_decoder = false; + let encoder = T5Stack::load(false, vb.pp("encoder"), &shared, &encoder_cfg)?; + + let mut decoder_cfg = cfg.clone(); + decoder_cfg.is_decoder = true; + decoder_cfg.is_encoder_decoder = false; + decoder_cfg.num_layers = cfg.num_decoder_layers.unwrap_or(cfg.num_layers); + let decoder = T5Stack::load(true, vb.pp("decoder"), &shared, &decoder_cfg)?; + + let tie_word_embeddings = cfg.tie_word_embeddings; + let lm_head = if tie_word_embeddings { + None + } else { + Some(QMatMul::new(cfg.d_model, cfg.vocab_size, vb.pp("lm_head"))?) + }; + + Ok(Self { + encoder, + decoder, + d_model, + tie_word_embeddings, + lm_head, + shared, + device: vb.device().clone(), + span_decode: tracing::span!(tracing::Level::TRACE, "decode"), + span_decode_head: tracing::span!(tracing::Level::TRACE, "decode-head"), + }) + } + + pub fn encode(&mut self, input_ids: &Tensor) -> Result { + self.encoder.forward(input_ids, None) + } + + pub fn decode( + &mut self, + decoder_input_ids: &Tensor, + encoder_output: &Tensor, + ) -> Result { + let _enter = self.span_decode.enter(); + let decoder_output = self + .decoder + .forward(decoder_input_ids, Some(encoder_output))?; + + let scaling_factor = if self.tie_word_embeddings { + // Rescale output before projecting on vocab + // See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586 + (self.d_model as f64).sqrt() + } else { + 1.0 + }; + let sequence_output = ((decoder_output + .narrow(1, decoder_output.dim(1)? - 1, 1)? + .squeeze(1)?) + * scaling_factor)?; + let output = { + let _enter = self.span_decode_head.enter(); + match self.lm_head { + None => sequence_output.matmul(&self.shared.embeddings().t()?)?, + Some(ref lm_head) => lm_head.forward(&sequence_output)?, + } + }; + Ok(output) + } + + pub fn forward(&mut self, input_ids: &Tensor, decoder_input_ids: &Tensor) -> Result { + let encoder_output = self.encode(input_ids)?; + self.decode(decoder_input_ids, &encoder_output) + } + + pub fn device(&self) -> &Device { + &self.device + } + + pub fn clear_kv_cache(&mut self) { + self.encoder.clear_kv_cache(); + self.decoder.clear_kv_cache(); + } +} diff --git a/patches/candle-transformers/src/models/qwen2.rs b/patches/candle-transformers/src/models/qwen2.rs new file mode 100644 index 0000000000..8a29646efe --- /dev/null +++ b/patches/candle-transformers/src/models/qwen2.rs @@ -0,0 +1,402 @@ +//! Qwen2 model implementation with quantization support. +//! +//! Qwen2 is a large language model from Alibaba optimized for efficiency. +//! This implementation provides quantization for reduced memory and compute. +//! +//! Key characteristics: +//! - Streaming decode support +//! - Grouped query attention (GQA) +//! - RMSNorm for layer normalization +//! - Rotary positional embeddings (RoPE) +//! - Support for 8-bit quantization +//! +//! References: +//! - 🤗 [Qwen2 Model](https://huggingface.co/Qwen/Qwen2-7B) +//! + +use crate::models::with_tracing::{linear, linear_no_bias, Linear, RmsNorm}; +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{Activation, VarBuilder}; +use std::sync::Arc; + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub max_position_embeddings: usize, + pub sliding_window: usize, + pub max_window_layers: usize, + pub tie_word_embeddings: bool, + pub rope_theta: f64, + pub rms_norm_eps: f64, + pub use_sliding_window: bool, + pub hidden_act: Activation, +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let dim = cfg.hidden_size / cfg.num_attention_heads; + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: Activation, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + let gate_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("gate_proj"))?; + let up_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("up_proj"))?; + let down_proj = linear_no_bias(intermediate_sz, hidden_sz, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_act, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = xs.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + hidden_size: usize, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + let head_dim = hidden_sz / num_heads; + let q_proj = linear(hidden_sz, num_heads * head_dim, vb.pp("q_proj"))?; + let k_proj = linear(hidden_sz, num_kv_heads * head_dim, vb.pp("k_proj"))?; + let v_proj = linear(hidden_sz, num_kv_heads * head_dim, vb.pp("v_proj"))?; + let o_proj = linear_no_bias(num_heads * head_dim, hidden_sz, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + hidden_size: hidden_sz, + rotary_emb, + kv_cache: None, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let (query_states, key_states) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states, seqlen_offset)?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?.contiguous()?; + let value_states = + crate::utils::repeat_kv(value_states, self.num_kv_groups)?.contiguous()?; + + let attn_output = { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + }; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.hidden_size))? + .apply(&self.o_proj) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: MLP, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl DecoderLayer { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let self_attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let input_layernorm = + RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let post_attention_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; + residual + xs + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + sliding_window: usize, + device: Device, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb_m.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(rotary_emb.clone(), cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb_m.pp("norm"))?; + Ok(Self { + embed_tokens, + layers, + norm, + sliding_window: cfg.sliding_window, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + fn prepare_causal_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + // Sliding window mask? + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| { + (0..tgt_len).map(move |j| { + if i < j || j + self.sliding_window < i { + f32::NEG_INFINITY + } else { + 0. + } + }) + }) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), self.dtype, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + fn prepare_attention_mask(&self, attn_mask: &Tensor) -> Result { + let (b_sz, sql_len) = attn_mask.dims2()?; + let mut mask: Vec = vec![]; + for b in 0..b_sz { + mask.push(attn_mask.i((b, ..))?.expand((1, 1, sql_len, sql_len))?); + } + let mask = Tensor::cat(&mask, 0)?; + let on_true = mask.zeros_like()?.to_dtype(self.dtype)?; + let on_false = Tensor::new(f32::NEG_INFINITY, &self.device)? + .broadcast_as(mask.shape())? + .to_dtype(self.dtype)?; + mask.where_cond(&on_true, &on_false) + } + + pub fn forward( + &mut self, + input_ids: &Tensor, + seqlen_offset: usize, + attn_mask: Option<&Tensor>, + ) -> Result { + let (b_size, seq_len) = input_ids.dims2()?; + let attention_mask: Option = match attn_mask { + Some(mask) => Some(self.prepare_attention_mask(mask)?), + None => { + if seq_len <= 1 { + None + } else { + Some(self.prepare_causal_attention_mask(b_size, seq_len, seqlen_offset)?) + } + } + }; + let mut xs = self.embed_tokens.forward(input_ids)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)? + } + xs.apply(&self.norm) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +} + +#[derive(Debug, Clone)] +pub struct ModelForCausalLM { + base_model: Model, + lm_head: Linear, +} + +impl ModelForCausalLM { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let base_model = Model::new(cfg, vb.clone())?; + let lm_head = if vb.contains_tensor("lm_head.weight") { + linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))? + } else { + Linear::from_weights(base_model.embed_tokens.embeddings().clone(), None) + }; + Ok(Self { + base_model, + lm_head, + }) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let (_b_size, seq_len) = input_ids.dims2()?; + self.base_model + .forward(input_ids, seqlen_offset, None)? + .narrow(1, seq_len - 1, 1)? + .apply(&self.lm_head) + } + + pub fn clear_kv_cache(&mut self) { + self.base_model.clear_kv_cache() + } +} diff --git a/patches/candle-transformers/src/models/qwen2_moe.rs b/patches/candle-transformers/src/models/qwen2_moe.rs new file mode 100644 index 0000000000..0f52008483 --- /dev/null +++ b/patches/candle-transformers/src/models/qwen2_moe.rs @@ -0,0 +1,484 @@ +//! Qwen2 model implementation with Mixture of Experts support. +//! +//! Qwen2 is a large language model using sparse Mixture of Experts (MoE). +//! This implementation provides support for sparsely activated MoE layers. +//! +//! Key characteristics: +//! - Mixture of Experts architecture +//! - Sparse expert activation +//! - Shared expert routing mechanism +//! - Grouped query attention (GQA) +//! - RMSNorm for layer normalization +//! - Rotary positional embeddings (RoPE) +//! +//! References: +//! - [Qwen2 Paper](https://arxiv.org/abs/2401.08985) +//! - [Model Card](https://huggingface.co/Qwen/Qwen2-7B-beta) +//! + +use crate::models::with_tracing::{linear, linear_no_bias, Linear, RmsNorm}; +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::{Activation, VarBuilder}; +use std::sync::Arc; + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub max_position_embeddings: usize, + pub sliding_window: usize, + pub max_window_layers: usize, + pub tie_word_embeddings: bool, + pub rope_theta: f64, + pub rms_norm_eps: f64, + pub use_sliding_window: bool, + pub hidden_act: Activation, + pub decoder_sparse_step: usize, + pub moe_intermediate_size: usize, + pub shared_expert_intermediate_size: usize, + pub num_experts_per_tok: usize, + pub num_experts: usize, + pub norm_topk_prob: bool, +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let dim = cfg.hidden_size / cfg.num_attention_heads; + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: Activation, +} + +impl MLP { + fn new(intermediate_sz: usize, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let gate_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("gate_proj"))?; + let up_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("up_proj"))?; + let down_proj = linear_no_bias(intermediate_sz, hidden_sz, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_act, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = xs.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + hidden_size: usize, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + let head_dim = hidden_sz / num_heads; + let q_proj = linear(hidden_sz, num_heads * head_dim, vb.pp("q_proj"))?; + let k_proj = linear(hidden_sz, num_kv_heads * head_dim, vb.pp("k_proj"))?; + let v_proj = linear(hidden_sz, num_kv_heads * head_dim, vb.pp("v_proj"))?; + let o_proj = linear_no_bias(num_heads * head_dim, hidden_sz, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + hidden_size: hidden_sz, + rotary_emb, + kv_cache: None, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let (query_states, key_states) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states, seqlen_offset)?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?.contiguous()?; + let value_states = + crate::utils::repeat_kv(value_states, self.num_kv_groups)?.contiguous()?; + + let attn_output = { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + }; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.hidden_size))? + .apply(&self.o_proj) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +// https://github.com/huggingface/transformers/blob/536ea2aca234fb48c5c69769431d643b0d93b233/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py#L800 +#[derive(Debug, Clone)] +struct SparseMoeBlock { + gate: Linear, + experts: Vec, + shared_expert: MLP, + shared_expert_gate: Linear, + norm_topk_prob: bool, + num_experts_per_tok: usize, +} + +impl SparseMoeBlock { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let gate = linear_no_bias(cfg.hidden_size, cfg.num_experts, vb.pp("gate"))?; + let mut experts = Vec::with_capacity(cfg.num_experts); + let vb_e = vb.pp("experts"); + for idx in 0..cfg.num_experts { + let expert = MLP::new(cfg.moe_intermediate_size, cfg, vb_e.pp(idx))?; + experts.push(expert) + } + let shared_expert = MLP::new( + cfg.shared_expert_intermediate_size, + cfg, + vb.pp("shared_expert"), + )?; + let shared_expert_gate = linear_no_bias(cfg.hidden_size, 1, vb.pp("shared_expert_gate"))?; + Ok(Self { + gate, + experts, + shared_expert, + shared_expert_gate, + norm_topk_prob: cfg.norm_topk_prob, + num_experts_per_tok: cfg.num_experts_per_tok, + }) + } +} + +impl Module for SparseMoeBlock { + fn forward(&self, xs: &Tensor) -> Result { + let (b_size, seq_len, hidden_dim) = xs.dims3()?; + let xs = xs.reshape(((), hidden_dim))?; + let router_logits = xs.apply(&self.gate)?; + 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 experts_per_tok = routing_weights + .arg_sort_last_dim(false)? + .narrow(D::Minus1, 0, self.num_experts_per_tok)? + .contiguous()?; + let routing_weights = routing_weights.gather(&experts_per_tok, D::Minus1)?; + + // 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 routing_weights = routing_weights.to_dtype(DType::F32)?.to_vec2::()?; + let experts_per_tok = experts_per_tok.to_vec2::()?; + let mut top_x = vec![vec![]; self.experts.len()]; + let mut selected_experts = vec![vec![]; self.experts.len()]; + for (row_idx, (rw, expert_idxs)) in routing_weights + .iter() + .zip(experts_per_tok.iter()) + .enumerate() + { + let sum_rw = rw.iter().sum::(); + for (&rw, &expert_idx) in rw.iter().zip(expert_idxs.iter()) { + top_x[expert_idx as usize].push(row_idx as u32); + let rw = if self.norm_topk_prob { rw / sum_rw } else { rw }; + selected_experts[expert_idx as usize].push(rw) + } + } + + let mut ys = xs.zeros_like()?; + for (expert_idx, expert_layer) in self.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_experts = + Tensor::new(selected_experts[expert_idx].as_slice(), xs.device())? + .reshape(((), 1))? + .to_dtype(xs.dtype())?; + // 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_experts)?; + ys = ys.index_add(&top_x, ¤t_hidden_states, 0)?; + } + let shared_expert_output = xs.apply(&self.shared_expert)?; + let shared_expert_output = shared_expert_output.broadcast_mul(&candle_nn::ops::sigmoid( + &xs.apply(&self.shared_expert_gate)?, + )?)?; + let ys = (ys + shared_expert_output)?; + let ys = ys.reshape((b_size, seq_len, hidden_dim))?; + Ok(ys) + } +} + +#[derive(Debug, Clone)] +enum MlpOrMoeBlock { + Mlp(MLP), + MoeBlock(SparseMoeBlock), +} + +impl Module for MlpOrMoeBlock { + fn forward(&self, xs: &Tensor) -> Result { + match self { + Self::MoeBlock(m) => m.forward(xs), + Self::Mlp(m) => m.forward(xs), + } + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: MlpOrMoeBlock, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl DecoderLayer { + fn new( + layer_idx: usize, + rotary_emb: Arc, + cfg: &Config, + vb: VarBuilder, + ) -> Result { + let self_attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let mlp = if cfg.num_experts > 0 && (layer_idx + 1).is_multiple_of(cfg.decoder_sparse_step) + { + MlpOrMoeBlock::MoeBlock(SparseMoeBlock::new(cfg, vb.pp("mlp"))?) + } else { + MlpOrMoeBlock::Mlp(MLP::new(cfg.intermediate_size, cfg, vb.pp("mlp"))?) + }; + let input_layernorm = + RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let post_attention_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; + residual + xs + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: Linear, + sliding_window: usize, + device: Device, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb_m.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(layer_idx, rotary_emb.clone(), cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb_m.pp("norm"))?; + let lm_head = linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))?; + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + sliding_window: cfg.sliding_window, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + // Sliding window mask? + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| { + (0..tgt_len).map(move |j| { + if i < j || j + self.sliding_window < i { + f32::NEG_INFINITY + } else { + 0. + } + }) + }) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let (b_size, seq_len) = input_ids.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(b_size, seq_len, seqlen_offset)?; + Some(mask) + }; + let mut xs = self.embed_tokens.forward(input_ids)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)? + } + xs.narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +} diff --git a/patches/candle-transformers/src/models/qwen3.rs b/patches/candle-transformers/src/models/qwen3.rs new file mode 100644 index 0000000000..9f018939ae --- /dev/null +++ b/patches/candle-transformers/src/models/qwen3.rs @@ -0,0 +1,389 @@ +use crate::{ + models::with_tracing::{linear_b, linear_no_bias, Linear, RmsNorm}, + utils::repeat_kv, +}; +use candle::{DType, Device, Module, Result, Tensor}; +use candle_nn::{kv_cache::ConcatKvCache, Activation, VarBuilder}; +use std::sync::Arc; + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub head_dim: usize, + pub attention_bias: bool, + pub num_key_value_heads: usize, + pub max_position_embeddings: usize, + pub sliding_window: Option, + pub max_window_layers: usize, + pub tie_word_embeddings: bool, + pub rope_theta: f64, + pub rms_norm_eps: f64, + pub use_sliding_window: bool, + pub hidden_act: Activation, +} + +#[derive(Debug, Clone)] +pub(crate) struct Qwen3RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl Qwen3RotaryEmbedding { + pub(crate) fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let dim = cfg.head_dim; + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(DType::F32)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(DType::F32)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?.to_dtype(dtype)?, + cos: freqs.cos()?.to_dtype(dtype)?, + }) + } + + /// Apply RoPE (q, k shape: B x H x L x D) + pub(crate) fn apply(&self, q: &Tensor, k: &Tensor, offset: usize) -> Result<(Tensor, Tensor)> { + let (_, _, seq_len, _) = q.dims4()?; + let cos = self.cos.narrow(0, offset, seq_len)?; + let sin = self.sin.narrow(0, offset, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +pub(crate) struct Qwen3MLP { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: Activation, +} + +impl Qwen3MLP { + pub(crate) fn new(cfg: &Config, vb: VarBuilder) -> Result { + Ok(Self { + gate_proj: linear_no_bias(cfg.hidden_size, cfg.intermediate_size, vb.pp("gate_proj"))?, + up_proj: linear_no_bias(cfg.hidden_size, cfg.intermediate_size, vb.pp("up_proj"))?, + down_proj: linear_no_bias(cfg.intermediate_size, cfg.hidden_size, vb.pp("down_proj"))?, + act_fn: cfg.hidden_act, + }) + } +} + +impl Module for Qwen3MLP { + fn forward(&self, x: &Tensor) -> Result { + let lhs = x.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = x.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +pub(crate) struct Qwen3Attention { + // projections + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + // norms + q_norm: RmsNorm, + k_norm: RmsNorm, + // hyper params + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + hidden_size: usize, + // utils + rotary_emb: Arc, + kv_cache: ConcatKvCache, +} + +impl Qwen3Attention { + pub(crate) fn new( + cfg: &Config, + rotary_emb: Arc, + vb: VarBuilder, + ) -> Result { + if cfg.use_sliding_window { + candle::bail!("sliding window is not supported") + } + + let head_dim = cfg.head_dim; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + + let q_proj = linear_b( + cfg.hidden_size, + num_heads * head_dim, + cfg.attention_bias, + vb.pp("q_proj"), + )?; + let k_proj = linear_b( + cfg.hidden_size, + num_kv_heads * head_dim, + cfg.attention_bias, + vb.pp("k_proj"), + )?; + let v_proj = linear_b( + cfg.hidden_size, + num_kv_heads * head_dim, + cfg.attention_bias, + vb.pp("v_proj"), + )?; + let o_proj = linear_b( + num_heads * head_dim, + cfg.hidden_size, + cfg.attention_bias, + vb.pp("o_proj"), + )?; + + let q_norm = RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("q_norm"))?; + let k_norm = RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("k_norm"))?; + + // Necessary because the hidden_size in the config isn't always accurate + let hidden_size = head_dim * cfg.num_attention_heads; + + // dim=2 because we concatenate along the sequence dimension + // For tensors of shape [batch, heads, seq, head_dim] + let kv_cache = ConcatKvCache::new(2); + + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + q_norm, + k_norm, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + hidden_size, + rotary_emb, + kv_cache, + }) + } + + pub(crate) fn forward( + &mut self, + x: &Tensor, + attn_mask: Option<&Tensor>, + offset: usize, + ) -> Result { + let (b, l, _) = x.dims3()?; + + // 1. Proj + let q = self.q_proj.forward(x)?; + let k = self.k_proj.forward(x)?; + let v = self.v_proj.forward(x)?; + + // 2. Reshape: (B, L, H, D) -> (B, H, L, D) + let q = q + .reshape((b, l, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let k = k + .reshape((b, l, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let v = v + .reshape((b, l, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + // 3. Per‑head RMSNorm + let q_flat = q.flatten(0, 2)?; // (B*H, L, D) -> (BHL, D) after transpose later + let k_flat = k.flatten(0, 2)?; + let q_flat = self.q_norm.forward(&q_flat)?; + let k_flat = self.k_norm.forward(&k_flat)?; + let q = q_flat.reshape((b, self.num_heads, l, self.head_dim))?; + let k = k_flat.reshape((b, self.num_kv_heads, l, self.head_dim))?; + + // 4. RoPE + let (q, k) = self.rotary_emb.apply(&q, &k, offset)?; + + // 5. Accumulate KV cache + let (k, v) = self.kv_cache.append(&k, &v)?; + + // 6. GQA repeat_kv + let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?; + let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?; + + // 7. Attention score + let scale = 1.0 / (self.head_dim as f64).sqrt(); + let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?; + if let Some(m) = attn_mask { + scores = scores.broadcast_add(m)?; + } + let probs = candle_nn::ops::softmax_last_dim(&scores)?; + let ctx = probs.matmul(&v)?; // (B, H, L, D) + + // 8. Output proj + ctx.transpose(1, 2)? + .reshape((b, l, self.hidden_size))? + .apply(&self.o_proj) + } + + pub(crate) fn clear_kv_cache(&mut self) { + self.kv_cache.reset(); + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Qwen3Attention, + mlp: Qwen3MLP, + ln1: RmsNorm, + ln2: RmsNorm, +} + +impl DecoderLayer { + fn new(cfg: &Config, rotary: Arc, vb: VarBuilder) -> Result { + let self_attn = Qwen3Attention::new(cfg, rotary, vb.pp("self_attn"))?; + let mlp = Qwen3MLP::new(cfg, vb.pp("mlp"))?; + let ln1 = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let ln2 = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + ln1, + ln2, + }) + } + + fn forward(&mut self, x: &Tensor, mask: Option<&Tensor>, offset: usize) -> Result { + let h = self.ln1.forward(x)?; + let h = self.self_attn.forward(&h, mask, offset)?; + let x = (x + h)?; + let h2 = self.ln2.forward(&x)?; + let h2 = h2.apply(&self.mlp)?; + x + h2 + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache(); + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + device: Device, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("model.embed_tokens"))?; + let rotary = Arc::new(Qwen3RotaryEmbedding::new(vb.dtype(), cfg, vb.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb.pp("model.layers"); + for i in 0..cfg.num_hidden_layers { + layers.push(DecoderLayer::new(cfg, rotary.clone(), vb_l.pp(i))?); + } + Ok(Self { + embed_tokens, + layers, + norm: RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("model.norm"))?, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + fn clear_kv_cache(&mut self) { + for l in &mut self.layers { + l.clear_kv_cache(); + } + } + + fn causal_mask( + &self, + b: usize, + tgt: usize, + offset: usize, + sw: Option, + ) -> Result { + let minf = f32::NEG_INFINITY; + let mask: Vec<_> = (0..tgt) + .flat_map(|i| { + (0..(tgt + offset)).map(move |j| { + let past_ok = j <= i + offset; + let sw_ok = match sw { + Some(w) => (i + offset) as i64 - j as i64 <= w as i64, + None => true, + }; + if past_ok && sw_ok { + 0. + } else { + minf + } + }) + }) + .collect(); + Tensor::from_slice(&mask, (b, 1, tgt, tgt + offset), &self.device)?.to_dtype(self.dtype) + } + + pub fn forward(&mut self, input: &Tensor, offset: usize) -> Result { + let (b, l) = input.dims2()?; + let mut h = self.embed_tokens.forward(input)?; + + let causal = if l == 1 { + None + } else { + Some(self.causal_mask(b, l, offset, None)?) + }; + + for layer in &mut self.layers { + h = layer.forward(&h, causal.as_ref(), offset)?; + } + self.norm.forward(&h) + } +} + +#[derive(Debug, Clone)] +pub struct ModelForCausalLM { + base: Model, + lm_head: Linear, +} + +impl ModelForCausalLM { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let base = Model::new(cfg, vb.clone())?; + let lm_head = if cfg.tie_word_embeddings { + Linear::from_weights(base.embed_tokens.embeddings().clone(), None) + } else { + linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))? + }; + Ok(Self { base, lm_head }) + } + + pub fn forward(&mut self, input: &Tensor, offset: usize) -> Result { + let (_, l) = input.dims2()?; + self.base + .forward(input, offset)? + .narrow(1, l - 1, 1)? + .apply(&self.lm_head) + } + + pub fn clear_kv_cache(&mut self) { + self.base.clear_kv_cache(); + } +} diff --git a/patches/candle-transformers/src/models/qwen3_moe.rs b/patches/candle-transformers/src/models/qwen3_moe.rs new file mode 100644 index 0000000000..0576b4c075 --- /dev/null +++ b/patches/candle-transformers/src/models/qwen3_moe.rs @@ -0,0 +1,374 @@ +use crate::{ + fused_moe::{FusedMoe, MoeCfg}, + models::{ + qwen3::{Config as Qwen3Config, Qwen3Attention, Qwen3MLP, Qwen3RotaryEmbedding}, + with_tracing::{linear_no_bias, Linear, RmsNorm}, + }, +}; +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::{Activation, VarBuilder}; +use std::sync::Arc; + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub head_dim: usize, + pub attention_bias: bool, + pub num_key_value_heads: usize, + pub max_position_embeddings: usize, + pub sliding_window: Option, + pub max_window_layers: usize, + pub tie_word_embeddings: bool, + pub rope_theta: f64, + pub rms_norm_eps: f64, + pub use_sliding_window: bool, + pub hidden_act: Activation, + // MoE specific configuration + pub decoder_sparse_step: usize, + pub moe_intermediate_size: usize, + pub num_experts_per_tok: usize, + pub num_experts: usize, + pub norm_topk_prob: bool, +} + +impl From<&Config> for Qwen3Config { + fn from(val: &Config) -> Self { + Qwen3Config { + vocab_size: val.vocab_size, + hidden_size: val.hidden_size, + intermediate_size: val.intermediate_size, + num_hidden_layers: val.num_hidden_layers, + num_attention_heads: val.num_attention_heads, + head_dim: val.head_dim, + attention_bias: val.attention_bias, + num_key_value_heads: val.num_key_value_heads, + max_position_embeddings: val.max_position_embeddings, + sliding_window: val.sliding_window, + max_window_layers: val.max_window_layers, + tie_word_embeddings: val.tie_word_embeddings, + rope_theta: val.rope_theta, + rms_norm_eps: val.rms_norm_eps, + use_sliding_window: val.use_sliding_window, + hidden_act: val.hidden_act, + } + } +} + +#[derive(Debug, Clone)] +struct Qwen3MLPExpert { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: Activation, +} + +impl Qwen3MLPExpert { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + Ok(Self { + gate_proj: linear_no_bias( + cfg.hidden_size, + cfg.moe_intermediate_size, + vb.pp("gate_proj"), + )?, + up_proj: linear_no_bias(cfg.hidden_size, cfg.moe_intermediate_size, vb.pp("up_proj"))?, + down_proj: linear_no_bias( + cfg.moe_intermediate_size, + cfg.hidden_size, + vb.pp("down_proj"), + )?, + act_fn: cfg.hidden_act, + }) + } +} + +impl Module for Qwen3MLPExpert { + fn forward(&self, x: &Tensor) -> Result { + let lhs = x.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = x.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +// Qwen3 Sparse MoE Block implementation +#[derive(Debug, Clone)] +struct Qwen3SparseMoeBlock { + gate: Linear, + experts: Vec, + norm_topk_prob: bool, + num_experts_per_tok: usize, +} + +impl Qwen3SparseMoeBlock { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let gate = linear_no_bias(cfg.hidden_size, cfg.num_experts, vb.pp("gate"))?; + let mut experts = Vec::with_capacity(cfg.num_experts); + let vb_e = vb.pp("experts"); + for idx in 0..cfg.num_experts { + let expert = Qwen3MLPExpert::new(cfg, vb_e.pp(idx))?; + experts.push(expert) + } + Ok(Self { + gate, + experts, + norm_topk_prob: cfg.norm_topk_prob, + num_experts_per_tok: cfg.num_experts_per_tok, + }) + } +} + +impl Module for Qwen3SparseMoeBlock { + fn forward(&self, xs: &Tensor) -> Result { + let (b_size, seq_len, hidden_dim) = xs.dims3()?; + let xs = xs.reshape(((), hidden_dim))?; + let router_logits = xs.apply(&self.gate)?; + let routing_weights = candle_nn::ops::softmax_last_dim(&router_logits)?; + + // Extract topk experts per token + let experts_per_tok = routing_weights + .arg_sort_last_dim(false)? + .narrow(D::Minus1, 0, self.num_experts_per_tok)? + .contiguous()?; + let routing_weights = routing_weights.gather(&experts_per_tok, D::Minus1)?; + + // Extract needed data + let routing_weights = routing_weights.to_dtype(DType::F32)?.to_vec2::()?; + let experts_per_tok = experts_per_tok.to_vec2::()?; + let mut top_x = vec![vec![]; self.experts.len()]; + let mut selected_experts = vec![vec![]; self.experts.len()]; + for (row_idx, (rw, expert_idxs)) in routing_weights + .iter() + .zip(experts_per_tok.iter()) + .enumerate() + { + let sum_rw = rw.iter().sum::(); + for (&rw, &expert_idx) in rw.iter().zip(expert_idxs.iter()) { + top_x[expert_idx as usize].push(row_idx as u32); + let rw = if self.norm_topk_prob { rw / sum_rw } else { rw }; + selected_experts[expert_idx as usize].push(rw) + } + } + + // Process through experts + let mut ys = xs.zeros_like()?; + for (expert_idx, expert_layer) in self.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_experts = + Tensor::new(selected_experts[expert_idx].as_slice(), xs.device())? + .reshape(((), 1))? + .to_dtype(xs.dtype())?; + + let current_state = xs.index_select(&top_x, 0)?.reshape(((), hidden_dim))?; + let current_hidden_states = expert_layer.forward(¤t_state)?; + let current_hidden_states = current_hidden_states.broadcast_mul(&selected_experts)?; + ys = ys.index_add(&top_x, ¤t_hidden_states, 0)?; + } + + ys.reshape((b_size, seq_len, hidden_dim)) + } +} + +// MLP or MoE decision enum +#[derive(Debug, Clone)] +enum Qwen3FeedForward { + Mlp(Qwen3MLP), + NaiveMoE(Qwen3SparseMoeBlock), + FusedMoE(FusedMoe), +} + +impl Qwen3FeedForward { + fn forward(&self, xs: &Tensor, is_prefill: bool) -> Result { + match self { + Self::Mlp(m) => m.forward(xs), + Self::NaiveMoE(m) => m.forward(xs), + Self::FusedMoE(m) => m.forward(xs, is_prefill), + } + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Qwen3Attention, + feed_forward: Qwen3FeedForward, + ln1: RmsNorm, + ln2: RmsNorm, +} + +impl DecoderLayer { + fn new( + layer_idx: usize, + cfg: &Config, + rotary: Arc, + vb: VarBuilder, + ) -> Result { + let self_attn = Qwen3Attention::new(&cfg.into(), rotary, vb.pp("self_attn"))?; + + let moe_cfg = MoeCfg { + hidden_size: cfg.hidden_size, + num_experts: cfg.num_experts, + num_experts_per_tok: cfg.num_experts_per_tok, + moe_intermediate_size: cfg.moe_intermediate_size, + norm_topk_prob: cfg.norm_topk_prob, + act: cfg.hidden_act, + decoder_sparse_step: None, + }; + // Decide whether to use MoE or regular MLP based on layer_idx and decoder_sparse_step + let feed_forward = + if cfg.num_experts > 0 && (layer_idx + 1).is_multiple_of(cfg.decoder_sparse_step) { + if cfg!(feature = "cuda") { + // Use fused MoE kernel on CUDA + Qwen3FeedForward::FusedMoE(FusedMoe::new(&moe_cfg, vb.pp("mlp"), vb.dtype())?) + } else { + Qwen3FeedForward::NaiveMoE(Qwen3SparseMoeBlock::new(cfg, vb.pp("mlp"))?) + } + } else { + Qwen3FeedForward::Mlp(Qwen3MLP::new(&cfg.into(), vb.pp("mlp"))?) + }; + + let ln1 = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let ln2 = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + + Ok(Self { + self_attn, + feed_forward, + ln1, + ln2, + }) + } + + fn forward(&mut self, x: &Tensor, mask: Option<&Tensor>, offset: usize) -> Result { + let h = self.ln1.forward(x)?; + let h = self.self_attn.forward(&h, mask, offset)?; + let x = (x + h)?; + let h2 = self.ln2.forward(&x)?; + let h2 = self.feed_forward.forward(&h2, mask.is_some())?; + x + h2 + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache(); + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + device: Device, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("model.embed_tokens"))?; + let rotary = Arc::new(Qwen3RotaryEmbedding::new( + vb.dtype(), + &cfg.into(), + vb.device(), + )?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb.pp("model.layers"); + for i in 0..cfg.num_hidden_layers { + layers.push(DecoderLayer::new(i, cfg, rotary.clone(), vb_l.pp(i))?); + } + Ok(Self { + embed_tokens, + layers, + norm: RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("model.norm"))?, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + fn clear_kv_cache(&mut self) { + for l in &mut self.layers { + l.clear_kv_cache(); + } + } + + fn causal_mask( + &self, + b: usize, + tgt: usize, + offset: usize, + sw: Option, + ) -> Result { + let minf = f32::NEG_INFINITY; + let mask: Vec<_> = (0..tgt) + .flat_map(|i| { + (0..(tgt + offset)).map(move |j| { + let past_ok = j <= i + offset; + let sw_ok = match sw { + Some(w) => (i + offset) as i64 - j as i64 <= w as i64, + None => true, + }; + if past_ok && sw_ok { + 0. + } else { + minf + } + }) + }) + .collect(); + Tensor::from_slice(&mask, (b, 1, tgt, tgt + offset), &self.device)?.to_dtype(self.dtype) + } + + pub fn forward(&mut self, input: &Tensor, offset: usize) -> Result { + let (b, l) = input.dims2()?; + let mut h = self.embed_tokens.forward(input)?; + + let causal = if l == 1 { + None + } else { + Some(self.causal_mask(b, l, offset, None)?) + }; + + for layer in &mut self.layers { + h = layer.forward(&h, causal.as_ref(), offset)?; + } + self.norm.forward(&h) + } +} + +#[derive(Debug, Clone)] +pub struct ModelForCausalLM { + base: Model, + lm_head: Linear, +} + +impl ModelForCausalLM { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let base = Model::new(cfg, vb.clone())?; + let lm_head = if cfg.tie_word_embeddings { + Linear::from_weights(base.embed_tokens.embeddings().clone(), None) + } else { + linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))? + }; + Ok(Self { base, lm_head }) + } + + pub fn forward(&mut self, input: &Tensor, offset: usize) -> Result { + let (_, l) = input.dims2()?; + self.base + .forward(input, offset)? + .narrow(1, l - 1, 1)? + .apply(&self.lm_head) + } + + pub fn clear_kv_cache(&mut self) { + self.base.clear_kv_cache(); + } +} diff --git a/patches/candle-transformers/src/models/qwen3_vl/config.rs b/patches/candle-transformers/src/models/qwen3_vl/config.rs new file mode 100644 index 0000000000..8cc180d3e9 --- /dev/null +++ b/patches/candle-transformers/src/models/qwen3_vl/config.rs @@ -0,0 +1,71 @@ +use candle_nn::Activation; + +use crate::serde_default_fn; + +serde_default_fn!(Activation, default_vision_hidden_act, Activation::Gelu); +serde_default_fn!(usize, default_in_channels, 3); +serde_default_fn!(usize, default_depth, 32); +serde_default_fn!(usize, default_hidden_size, 3584); +serde_default_fn!(usize, default_out_hidden_size, 3584); +serde_default_fn!(usize, default_intermediate_size, 3420); +serde_default_fn!(usize, default_num_heads, 16); +serde_default_fn!(usize, default_patch_size, 14); +serde_default_fn!(usize, default_spatial_merge_size, 2); +serde_default_fn!(usize, default_temporal_patch_size, 2); +serde_default_fn!(usize, default_num_position_embeddings, 576); +serde_default_fn!(Vec, default_deepstack_visual_indexes, Vec::new()); + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct VisionConfig { + #[serde(default = "default_depth")] + pub depth: usize, + #[serde(default = "default_hidden_size")] + pub hidden_size: usize, + #[serde(default = "default_out_hidden_size")] + pub out_hidden_size: usize, + #[serde(default = "default_vision_hidden_act")] + pub hidden_act: Activation, + #[serde(default = "default_intermediate_size")] + pub intermediate_size: usize, + #[serde(default = "default_num_heads")] + pub num_heads: usize, + #[serde(default = "default_in_channels")] + pub in_chans: usize, + #[serde(default = "default_patch_size")] + pub patch_size: usize, + #[serde(default = "default_spatial_merge_size")] + pub spatial_merge_size: usize, + #[serde(default = "default_temporal_patch_size")] + pub temporal_patch_size: usize, + #[serde(default = "default_num_position_embeddings")] + pub num_position_embeddings: usize, + #[serde(default = "default_deepstack_visual_indexes")] + pub deepstack_visual_indexes: Vec, +} + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct TextConfig { + pub head_dim: usize, + pub vocab_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub hidden_act: Activation, + pub max_position_embeddings: usize, + pub rms_norm_eps: f64, + pub tie_word_embeddings: bool, + pub rope_theta: f64, + pub sliding_window: Option, +} + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + pub text_config: TextConfig, + pub vision_config: VisionConfig, + pub image_token_id: u32, + pub video_token_id: u32, + pub vision_start_token_id: u32, + pub vision_end_token_id: u32, +} diff --git a/patches/candle-transformers/src/models/qwen3_vl/conv3d_temporal_2.rs b/patches/candle-transformers/src/models/qwen3_vl/conv3d_temporal_2.rs new file mode 100644 index 0000000000..f390e3ba4e --- /dev/null +++ b/patches/candle-transformers/src/models/qwen3_vl/conv3d_temporal_2.rs @@ -0,0 +1,77 @@ +//! Conv3dConfig assuming a temporal patch size of 2 + +use candle::{IndexOp, Module, Result, Tensor}; +use candle_nn::{Conv2d, Conv2dConfig, VarBuilder}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Conv3dConfig { + pub padding: usize, + pub stride: usize, + pub dilation: usize, + pub groups: usize, +} + +impl Default for Conv3dConfig { + fn default() -> Self { + Self { + padding: 0, + stride: 1, + dilation: 1, + groups: 1, + } + } +} + +pub struct Conv3dNoBias { + conv2d_1: Conv2d, + conv2d_2: Conv2d, +} + +impl Conv3dNoBias { + pub fn new( + in_channels: usize, + out_channels: usize, + kernel_sizes: [usize; 3], + cfg: Conv3dConfig, + vb: VarBuilder, + ) -> Result { + let ws = vb.get( + ( + out_channels, + in_channels / cfg.groups, + kernel_sizes[0], + kernel_sizes[1], + kernel_sizes[2], + ), + "weight", + )?; + + // Split on temporal dimension + // https://github.com/pytorch/pytorch/issues/139066 + + let w1 = ws.i((.., .., 0, .., ..))?; + let w2 = ws.i((.., .., 1, .., ..))?; + + let cfg = Conv2dConfig { + padding: cfg.padding, + stride: cfg.stride, + dilation: cfg.dilation, + groups: cfg.groups, + cudnn_fwd_algo: None, + }; + + Ok(Self { + conv2d_1: Conv2d::new(w1.contiguous()?, None, cfg), + conv2d_2: Conv2d::new(w2.contiguous()?, None, cfg), + }) + } +} + +impl Module for Conv3dNoBias { + fn forward(&self, xs: &Tensor) -> Result { + let xs1 = xs.i((.., .., 0, .., ..))?; + let xs2 = xs.i((.., .., 1, .., ..))?; + + (self.conv2d_1.forward(&xs1)? + self.conv2d_2.forward(&xs2)?)?.unsqueeze(2) + } +} diff --git a/patches/candle-transformers/src/models/qwen3_vl/mod.rs b/patches/candle-transformers/src/models/qwen3_vl/mod.rs new file mode 100644 index 0000000000..57e78f5082 --- /dev/null +++ b/patches/candle-transformers/src/models/qwen3_vl/mod.rs @@ -0,0 +1,270 @@ +#![allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] + +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::VarBuilder; +use text::Qwen3VLTextModel; +use vision::Qwen3VLVisionModel; + +pub mod config; +mod conv3d_temporal_2; +mod text; +mod vision; + +pub use config::Config; + +use crate::models::deepseek2::NonZeroOp; + +pub struct Qwen3VLModel { + text: Qwen3VLTextModel, + vision: Qwen3VLVisionModel, +} + +impl Qwen3VLModel { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vision = Qwen3VLVisionModel::new(&cfg.vision_config, vb.pp("model").pp("visual"))?; + let text = Qwen3VLTextModel::new(&cfg.text_config, vb.clone())?; + Ok(Self { text, vision }) + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + dtype: DType, + device: &Device, + ) -> Result { + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { f32::NEG_INFINITY } else { 0f32 })) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand(( + b_size, + self.text.num_attn_heads, + tgt_len, + tgt_len + seqlen_offset, + ))? + .to_dtype(dtype) + } + + #[allow(clippy::too_many_arguments)] + pub fn forward( + &self, + input_ids: &Tensor, + pixel_values: Option, + pixel_values_videos: Option, + image_grid_thw: Option, + video_grid_thw: Option, + seqlens: Vec, + continuous_img_pad: Vec>, + continuous_vid_pad: Vec>, + seqlen_offsets: &[usize], + ) -> Result { + let (bs, seqlen) = input_ids.dims2()?; + let attention_mask = if seqlen <= 1 { + Some(self.prepare_decoder_attention_mask( + bs, + seqlen, + seqlen_offsets[0], + self.text.dtype, + input_ids.device(), + )?) + } else { + None + }; + + let mut input_embeds = self.text.embed_tokens(input_ids)?; + let (batch_size, seq_len, hidden_dim) = input_embeds.dims3()?; + let device = input_embeds.device().clone(); + + let mut image_mask_opt: Option = None; + let mut video_mask_opt: Option = None; + let mut deepstack_image_opt: Option> = None; + let mut deepstack_video_opt: Option> = None; + + if let Some(pixel_values) = &pixel_values { + let Some(image_grid_thw_ref) = image_grid_thw.as_ref() else { + candle::bail!("pixel_values require image_grid_thw"); + }; + let mut pixel_values = pixel_values.clone(); + let dims = pixel_values.dims(); + if dims.len() == 3 { + pixel_values = pixel_values.reshape((dims[0] * dims[1], dims[2]))?; + } + let (image_embeds, deepstack_image_embeds) = + self.vision.forward(&pixel_values, image_grid_thw_ref)?; + let image_embeds = image_embeds.to_device(&device)?.to_dtype(self.text.dtype)?; + let mut deepstack_image_embeds = deepstack_image_embeds + .into_iter() + .map(|t| t.to_device(&device)?.to_dtype(self.text.dtype)) + .collect::>>()?; + + let mut offset = 0usize; + let mut image_mask = + Tensor::zeros((batch_size, seq_len), DType::F32, input_ids.device())?; + let total_expected: usize = continuous_img_pad + .iter() + .flat_map(|spans| spans.iter().map(|(s, e)| e - s)) + .sum(); + if image_embeds.dim(0)? != total_expected { + candle::bail!( + "Image embedding length {} does not match placeholder tokens {}", + image_embeds.dim(0)?, + total_expected + ); + } + + for (batch, spans) in continuous_img_pad.iter().enumerate() { + for &(start, end) in spans { + let len = end - start; + let chunk = image_embeds.narrow(0, offset, len)?; + offset += len; + input_embeds = input_embeds.slice_assign( + &[batch..batch + 1, start..end, 0..hidden_dim], + &chunk.unsqueeze(0)?, + )?; + let ones = Tensor::ones((1, len), DType::F32, input_ids.device())?; + image_mask = image_mask.slice_assign(&[batch..batch + 1, start..end], &ones)?; + } + } + image_mask_opt = Some(image_mask.to_dtype(DType::U8)?); + deepstack_image_opt = Some(std::mem::take(&mut deepstack_image_embeds)); + } + + if let Some(pixel_values_videos) = &pixel_values_videos { + let Some(video_grid_thw_ref) = video_grid_thw.as_ref() else { + candle::bail!("pixel_values_videos require video_grid_thw"); + }; + let mut pixel_values = pixel_values_videos.clone(); + let dims = pixel_values.dims(); + if dims.len() == 3 { + pixel_values = pixel_values.reshape((dims[0] * dims[1], dims[2]))?; + } + let (video_embeds, deepstack_video_embeds) = + self.vision.forward(&pixel_values, video_grid_thw_ref)?; + let video_embeds = video_embeds.to_device(&device)?.to_dtype(self.text.dtype)?; + let mut deepstack_video_embeds = deepstack_video_embeds + .into_iter() + .map(|t| t.to_device(&device)?.to_dtype(self.text.dtype)) + .collect::>>()?; + + let mut offset = 0usize; + let mut video_mask = + Tensor::zeros((batch_size, seq_len), DType::F32, input_ids.device())?; + let total_expected: usize = continuous_vid_pad + .iter() + .flat_map(|spans| spans.iter().map(|(s, e)| e - s)) + .sum(); + if video_embeds.dim(0)? != total_expected { + candle::bail!( + "Video embedding length {} does not match placeholder tokens {}", + video_embeds.dim(0)?, + total_expected + ); + } + + for (batch, spans) in continuous_vid_pad.iter().enumerate() { + for &(start, end) in spans { + let len = end - start; + let chunk = video_embeds.narrow(0, offset, len)?; + offset += len; + input_embeds = input_embeds.slice_assign( + &[batch..batch + 1, start..end, 0..hidden_dim], + &chunk.unsqueeze(0)?, + )?; + let ones = Tensor::ones((1, len), DType::F32, input_ids.device())?; + video_mask = video_mask.slice_assign(&[batch..batch + 1, start..end], &ones)?; + } + } + video_mask_opt = Some(video_mask.to_dtype(DType::U8)?); + deepstack_video_opt = Some(std::mem::take(&mut deepstack_video_embeds)); + } + + let (visual_pos_masks, deepstack_visual_embeds) = match ( + image_mask_opt, + deepstack_image_opt, + video_mask_opt, + deepstack_video_opt, + ) { + (Some(image_mask), Some(image_deepstack), Some(video_mask), Some(video_deepstack)) => { + let combined = + (image_mask.to_dtype(DType::F32)? + video_mask.to_dtype(DType::F32)?)?; + let visual_mask = combined.gt(0f32)?.to_dtype(DType::U8)?; + let visual_indices = visual_mask.flatten_all()?.nonzero()?.squeeze(1)?; + let visual_indices_vec = visual_indices.to_vec1::()?; + + let image_flat = image_mask + .flatten_all()? + .to_dtype(DType::U8)? + .to_vec1::()?; + let num_visual = visual_indices_vec.len(); + if image_deepstack.len() != video_deepstack.len() { + candle::bail!( + "DeepStack image layers ({}) do not match video layers ({})", + image_deepstack.len(), + video_deepstack.len() + ); + } + let mut combined_layers = Vec::with_capacity(image_deepstack.len()); + for (img_layer, vid_layer) in image_deepstack.iter().zip(video_deepstack.iter()) { + let mut rows = Vec::with_capacity(num_visual); + let mut img_offset = 0usize; + let mut vid_offset = 0usize; + for &idx in &visual_indices_vec { + let idx = idx as usize; + if image_flat[idx] != 0 { + rows.push(img_layer.i(img_offset)?); + img_offset += 1; + } else { + rows.push(vid_layer.i(vid_offset)?); + vid_offset += 1; + } + } + if img_offset != img_layer.dim(0)? || vid_offset != vid_layer.dim(0)? { + candle::bail!( + "DeepStack feature alignment failed for images ({}/{}) or videos ({}/{})", + img_offset, + img_layer.dim(0)?, + vid_offset, + vid_layer.dim(0)? + ); + } + let row_refs: Vec<&Tensor> = rows.iter().collect(); + combined_layers.push(Tensor::stack(&row_refs, 0)?); + } + (Some(visual_mask), Some(combined_layers)) + } + (Some(image_mask), Some(image_deepstack), _, _) => { + (Some(image_mask), Some(image_deepstack)) + } + (_, _, Some(video_mask), Some(video_deepstack)) => { + (Some(video_mask), Some(video_deepstack)) + } + _ => (None, None), + }; + + let mut ropeidx_attn_mask_bs = Vec::new(); + let max_seqlens = *seqlens.iter().max().unwrap(); + for len in &seqlens { + ropeidx_attn_mask_bs.push(Tensor::new( + [vec![1f32; *len], vec![0f32; max_seqlens - len]].concat(), + input_ids.device(), + )?); + } + + let out = self.text.forward_embeds( + input_embeds, + attention_mask.as_ref(), + seqlen_offsets, + visual_pos_masks.as_ref(), + deepstack_visual_embeds.as_deref(), + )?; + Ok(out) + } +} diff --git a/patches/candle-transformers/src/models/qwen3_vl/text.rs b/patches/candle-transformers/src/models/qwen3_vl/text.rs new file mode 100644 index 0000000000..febe426879 --- /dev/null +++ b/patches/candle-transformers/src/models/qwen3_vl/text.rs @@ -0,0 +1,395 @@ +use std::sync::{Arc, Mutex}; + +use candle::{DType, Device, IndexOp, Result, Tensor}; +use candle_nn::{ + embedding, kv_cache::KvCache, linear, linear_b, rms_norm, Activation, Embedding, Linear, + Module, RmsNorm, VarBuilder, +}; + +use super::config::TextConfig; + +#[derive(Debug, Clone)] +pub struct RotaryEmbedding { + cos: Tensor, + sin: Tensor, +} + +impl RotaryEmbedding { + pub fn new( + base: f32, + head_dim: usize, + max_position_embeddings: usize, + device: &Device, + dtype: DType, + ) -> Result { + let inv_freq: Vec<_> = (0..head_dim) + .step_by(2) + .map(|i| 1f32 / base.powf(i as f32 / head_dim as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), device)?; + let t = Tensor::arange(0u32, max_position_embeddings as u32, device)? + .to_dtype(DType::F32)? + .reshape((max_position_embeddings, 1))?; + let freqs = t.matmul(&inv_freq)?; + let sin = freqs.sin()?.to_dtype(dtype)?; + let cos = freqs.cos()?.to_dtype(dtype)?; + + Ok(Self { cos, sin }) + } + + pub fn forward( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offsets: &[usize], + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _qh, seq_len, _n_embd) = q.dims4()?; + + let rope = candle_nn::rotary_emb::rope; + + let mut q_embeds = Vec::new(); + let mut k_embeds = Vec::new(); + for (i, offset) in seqlen_offsets.iter().enumerate() { + let cos = self.cos.narrow(0, *offset, seq_len)?; + let sin = self.sin.narrow(0, *offset, seq_len)?; + let q_embed = rope(&q.i(i)?.unsqueeze(0)?.contiguous()?, &cos, &sin)?; + let k_embed = rope(&k.i(i)?.unsqueeze(0)?.contiguous()?, &cos, &sin)?; + q_embeds.push(q_embed); + k_embeds.push(k_embed); + } + Ok((Tensor::cat(&q_embeds, 0)?, Tensor::cat(&k_embeds, 0)?)) + } +} + +struct Mlp { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: Activation, +} + +impl Mlp { + fn new(cfg: &TextConfig, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + let gate_proj = linear_b(hidden_sz, intermediate_sz, false, vb.pp("gate_proj"))?; + let up_proj = linear_b(hidden_sz, intermediate_sz, false, vb.pp("up_proj"))?; + let down_proj = linear_b(intermediate_sz, hidden_sz, false, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_act, + }) + } + + fn forward(&self, xs: &Tensor) -> Result { + let lhs = self.gate_proj.forward(xs)?.apply(&self.act_fn)?; + let rhs = self.up_proj.forward(xs)?; + self.down_proj.forward(&(lhs * rhs)?) + } +} + +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + q_norm: RmsNorm, + k_norm: RmsNorm, + num_heads: usize, + num_kv_heads: usize, + head_dim: usize, + rotary_emb: Arc, + n_kv_groups: usize, + softmax_scale: f64, + kv_cache: Arc>, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &TextConfig, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let q_proj = linear_b(hidden_sz, num_heads * cfg.head_dim, false, vb.pp("q_proj"))?; + let k_proj = linear_b( + hidden_sz, + num_kv_heads * cfg.head_dim, + false, + vb.pp("k_proj"), + )?; + let v_proj = linear_b( + hidden_sz, + num_kv_heads * cfg.head_dim, + false, + vb.pp("v_proj"), + )?; + let o_proj = linear_b(num_heads * cfg.head_dim, hidden_sz, false, vb.pp("o_proj"))?; + let q_norm = rms_norm(cfg.head_dim, cfg.rms_norm_eps, vb.pp("q_norm"))?; + let k_norm = rms_norm(cfg.head_dim, cfg.rms_norm_eps, vb.pp("k_norm"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + q_norm, + k_norm, + num_heads, + num_kv_heads, + head_dim: cfg.head_dim, + rotary_emb, + n_kv_groups: cfg.num_attention_heads / cfg.num_key_value_heads, + softmax_scale: 1.0 / (cfg.head_dim as f64).sqrt(), + kv_cache: Arc::new(Mutex::new(KvCache::new(2, cfg.max_position_embeddings))), + }) + } + + #[allow(clippy::too_many_arguments)] + fn forward( + &self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offsets: &[usize], + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + let mut q = self.q_proj.forward(xs)?; + let mut k = self.k_proj.forward(xs)?; + let mut v = self.v_proj.forward(xs)?; + + q = q + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + k = k + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + v = v + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + q = q.apply(&self.q_norm)?; + k = k.apply(&self.k_norm)?; + + (q, k) = self.rotary_emb.forward(&q, &k, seqlen_offsets)?; + + let q = q.contiguous()?; + let k = k.contiguous()?; + let v = v.contiguous()?; + + let (k, v) = self + .kv_cache + .lock() + .expect("Need a lock because of the deepstack injection") + .append(&k, &v)?; + + let k = crate::utils::repeat_kv(k, self.n_kv_groups)?.contiguous()?; + let v = crate::utils::repeat_kv(v, self.n_kv_groups)?.contiguous()?; + + let mut attn_output = { + let attn_weights = (q.matmul(&k.transpose(2, 3)?)? * self.softmax_scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&v)? + }; + + attn_output = attn_output.transpose(1, 2)?.reshape((b_sz, q_len, ()))?; + + self.o_proj.forward(&attn_output) + } +} + +pub struct DecoderLayer { + self_attn: Attention, + mlp: Mlp, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl DecoderLayer { + fn new(rotary_emb: Arc, cfg: &TextConfig, vb: VarBuilder) -> Result { + let self_attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let mlp = Mlp::new(cfg, vb.pp("mlp"))?; + let input_layernorm = + rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let post_attention_layernorm = rms_norm( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + }) + } + + #[allow(clippy::too_many_arguments)] + fn forward( + &self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offsets: &[usize], + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self + .self_attn + .forward(&xs, attention_mask, seqlen_offsets)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = self + .mlp + .forward(&xs.apply(&self.post_attention_layernorm)?)?; + residual + xs + } +} + +pub struct Qwen3VLTextModel { + embed_tokens: Embedding, + pub(super) norm: RmsNorm, + layers: Vec, + lm_head: Linear, + pub(super) dtype: DType, + pub(super) num_attn_heads: usize, +} + +impl Qwen3VLTextModel { + pub fn new(cfg: &TextConfig, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model").pp("language_model"); + + let embed_tokens = embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + + let rotary_emb = Arc::new(RotaryEmbedding::new( + cfg.rope_theta as f32, + cfg.head_dim, + cfg.max_position_embeddings, + vb.device(), + vb_m.dtype(), + )?); + let vb_l = vb_m.pp("layers"); + let mut layers = Vec::new(); + for layer_idx in 0..cfg.num_hidden_layers { + layers.push(DecoderLayer::new( + rotary_emb.clone(), + cfg, + vb_l.pp(layer_idx), + )?); + } + let norm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb_m.pp("norm"))?; + let lm_head = if !cfg.tie_word_embeddings { + linear(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))? + } else { + candle_nn::Linear::new(embed_tokens.embeddings().clone(), None) + }; + Ok(Self { + embed_tokens, + norm, + layers, + lm_head, + dtype: vb.dtype(), + num_attn_heads: cfg.num_attention_heads, + }) + } + + pub fn embed_tokens(&self, input_ids: &Tensor) -> Result { + self.embed_tokens.forward(input_ids) + } + + pub fn forward_embeds( + &self, + mut xs: Tensor, + attention_mask: Option<&Tensor>, + seqlen_offsets: &[usize], + visual_pos_masks: Option<&Tensor>, + deepstack_visual_embeds: Option<&[Tensor]>, + ) -> Result { + let (_, seq_len, _) = xs.dims3()?; + + for (i, layer) in self.layers.iter().enumerate() { + xs = layer.forward( + &xs, + attention_mask + .as_ref() + .map(|m| m.to_device(xs.device()).unwrap()) + .as_ref(), + seqlen_offsets, + )?; + + // Integrate DeepStack visual features when provided. + if let (Some(visual_pos_masks), Some(deepstack)) = + (visual_pos_masks, deepstack_visual_embeds) + { + if i < deepstack.len() { + xs = self.deepstack_process(xs, visual_pos_masks, &deepstack[i])?; + } + } + } + + xs = xs.apply(&self.norm)?; + + self.lm_head + .forward(&xs)? + .i((.., seq_len - 1, ..))? + .contiguous() + } + + fn deepstack_process( + &self, + hidden_states: Tensor, + visual_pos_masks: &Tensor, + visual_embeds: &Tensor, + ) -> Result { + let device = hidden_states.device(); + let dtype = hidden_states.dtype(); + + let mask = visual_pos_masks.to_device(device)?.to_dtype(DType::F32)?; + let mask_flat = mask.flatten_all()?; + + let masked_count = mask_flat.sum_all()?.to_scalar::()? as usize; + let visual_embeds = visual_embeds.to_device(device)?.to_dtype(dtype)?; + + if masked_count == 0 { + if visual_embeds.dim(0)? != 0 { + candle::bail!( + "DeepStack visual embeds ({}) provided but mask is empty", + visual_embeds.dim(0)? + ); + } + return Ok(hidden_states); + } + + if visual_embeds.dim(0)? != masked_count { + candle::bail!( + "Mismatch between DeepStack visual embeds ({}) and mask positions ({})", + visual_embeds.dim(0)?, + masked_count + ); + } + + let (batch, seq, hidden) = hidden_states.dims3()?; + let total_positions = batch * seq; + let mut hidden_flat = hidden_states.reshape((total_positions, hidden))?; + + let prefix = mask_flat.cumsum(0)?; + let rank = (prefix - &mask_flat)?.mul(&mask_flat)?; + let rank_u32 = rank.to_dtype(DType::U32)?; + + let positions = Tensor::arange(0u32, total_positions as u32, device)?; + let positions_f32 = positions.to_dtype(DType::F32)?; + let masked_positions = positions_f32.mul(&mask_flat)?; + + let mut position_per_rank = Tensor::zeros((masked_count,), DType::F32, device)?; + position_per_rank = position_per_rank.scatter_add(&rank_u32, &masked_positions, 0)?; + let position_per_rank = position_per_rank.to_dtype(DType::U32)?; + + let linear_index = position_per_rank.unsqueeze(1)?.repeat((1, hidden))?; + + hidden_flat = hidden_flat.scatter_add(&linear_index, &visual_embeds, 0)?; + hidden_flat.reshape((batch, seq, hidden)) + } +} diff --git a/patches/candle-transformers/src/models/qwen3_vl/vision.rs b/patches/candle-transformers/src/models/qwen3_vl/vision.rs new file mode 100644 index 0000000000..465a7407ff --- /dev/null +++ b/patches/candle-transformers/src/models/qwen3_vl/vision.rs @@ -0,0 +1,585 @@ +use std::f64; + +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::{ + embedding, layer_norm, linear, Activation, Embedding, LayerNorm, LayerNormConfig, Linear, + Module, VarBuilder, +}; + +use crate::models::qwen3_vl::conv3d_temporal_2::{Conv3dConfig, Conv3dNoBias}; + +use super::config::VisionConfig; + +struct PatchEmbed { + proj: Conv3dNoBias, + bias: Tensor, + in_channels: usize, + patch_size: usize, + temporal_patch_size: usize, + hidden_size: usize, +} + +impl PatchEmbed { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let proj_vb = vb.pp("proj"); + let proj = Conv3dNoBias::new( + cfg.in_chans, + cfg.hidden_size, + [cfg.temporal_patch_size, cfg.patch_size, cfg.patch_size], + Conv3dConfig { + stride: cfg.patch_size, + ..Default::default() + }, + proj_vb.clone(), + )?; + let bias = proj_vb.get(cfg.hidden_size, "bias")?; + Ok(Self { + proj, + bias, + in_channels: cfg.in_chans, + patch_size: cfg.patch_size, + temporal_patch_size: cfg.temporal_patch_size, + hidden_size: cfg.hidden_size, + }) + } + + fn forward(&self, xs: &Tensor) -> Result { + let xs = xs.reshape(( + (), + self.in_channels, + self.temporal_patch_size, + self.patch_size, + self.patch_size, + ))?; + let xs = self.proj.forward(&xs)?; + let xs = xs.reshape(((), self.hidden_size))?; + let bias = self.bias.unsqueeze(0)?; + xs.broadcast_add(&bias) + } +} + +struct VisionMlp { + fc1: Linear, + fc2: Linear, + act: Activation, +} + +impl VisionMlp { + fn new(dim: usize, hidden_dim: usize, act: Activation, vb: VarBuilder) -> Result { + Ok(Self { + fc1: linear(dim, hidden_dim, vb.pp("linear_fc1"))?, + fc2: linear(hidden_dim, dim, vb.pp("linear_fc2"))?, + act, + }) + } + + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.fc1.forward(xs)?; + let xs = xs.apply(&self.act)?; + self.fc2.forward(&xs) + } +} + +fn rotate_half(xs: &Tensor) -> Result { + let last_dim = xs.dim(D::Minus1)?; + let xs1 = xs.narrow(D::Minus1, 0, last_dim / 2)?; + let xs2 = xs.narrow(D::Minus1, last_dim / 2, last_dim - last_dim / 2)?; + Tensor::cat(&[&xs2.neg()?, &xs1], D::Minus1) +} + +fn apply_rotary_pos_emb_vision( + q: &Tensor, + k: &Tensor, + cos: &Tensor, + sin: &Tensor, +) -> Result<(Tensor, Tensor)> { + let cos = cos.unsqueeze(D::Minus2)?; + let sin = sin.unsqueeze(D::Minus2)?; + + let q_embed = (q.broadcast_mul(&cos)? + rotate_half(q)?.broadcast_mul(&sin)?)?; + let k_embed = (k.broadcast_mul(&cos)? + rotate_half(k)?.broadcast_mul(&sin)?)?; + Ok((q_embed, k_embed)) +} + +struct VisionAttention { + qkv: Linear, + proj: Linear, + num_heads: usize, + head_dim: usize, +} + +impl VisionAttention { + fn new(dim: usize, num_heads: usize, vb: VarBuilder) -> Result { + Ok(Self { + qkv: linear(dim, dim * 3, vb.pp("qkv"))?, + proj: linear(dim, dim, vb.pp("proj"))?, + num_heads, + head_dim: dim / num_heads, + }) + } + + fn forward( + &self, + xs: &Tensor, + cu_seqlens: &[usize], + cos: &Tensor, + sin: &Tensor, + ) -> Result { + let seq_len = xs.dim(0)?; + let hidden_states = self.qkv.forward(xs)?; + let qkv = hidden_states + .reshape((seq_len, 3, self.num_heads, self.head_dim))? + .permute((1, 0, 2, 3))?; + let mut q = qkv.i(0)?.squeeze(0)?; + let mut k = qkv.i(1)?.squeeze(0)?; + let mut v = qkv.i(2)?.squeeze(0)?; + + let cos = cos.to_dtype(DType::F32)?; + let sin = sin.to_dtype(DType::F32)?; + q = q.to_dtype(DType::F32)?; + k = k.to_dtype(DType::F32)?; + v = v.to_dtype(DType::F32)?; + (q, k) = apply_rotary_pos_emb_vision(&q, &k, &cos, &sin)?; + + let mut outputs = Vec::new(); + for window in cu_seqlens.windows(2) { + let start = window[0]; + let end = window[1]; + if end <= start { + continue; + } + let len = end - start; + let q_chunk = q.narrow(0, start, len)?.transpose(0, 1)?.contiguous()?; + let k_chunk = k.narrow(0, start, len)?.transpose(0, 1)?.contiguous()?; + let v_chunk = v.narrow(0, start, len)?.transpose(0, 1)?.contiguous()?; + + let mut chunk_out = { + let q = q_chunk.unsqueeze(0)?; + let k = k_chunk.unsqueeze(0)?; + let v = v_chunk.unsqueeze(0)?; + + let attn_weights = + (q.matmul(&k.transpose(2, 3)?)? / (self.head_dim as f64).sqrt())?; + + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&v)? + }; + chunk_out = chunk_out.squeeze(0)?.transpose(0, 1)?; + + chunk_out.device().synchronize()?; + chunk_out = chunk_out.reshape((len, self.num_heads * self.head_dim))?; + outputs.push(chunk_out.to_dtype(xs.dtype())?); + } + let attn_output = Tensor::cat(&outputs, 0)?; + self.proj.forward(&attn_output) + } +} + +struct VisionBlock { + norm1: LayerNorm, + norm2: LayerNorm, + attn: VisionAttention, + mlp: VisionMlp, +} + +impl VisionBlock { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let norm_cfg = LayerNormConfig { + eps: 1e-6, + ..Default::default() + }; + let norm1 = layer_norm(cfg.hidden_size, norm_cfg, vb.pp("norm1"))?; + let norm2 = layer_norm(cfg.hidden_size, norm_cfg, vb.pp("norm2"))?; + let attn = VisionAttention::new(cfg.hidden_size, cfg.num_heads, vb.pp("attn"))?; + let mlp = VisionMlp::new( + cfg.hidden_size, + cfg.intermediate_size, + cfg.hidden_act, + vb.pp("mlp"), + )?; + Ok(Self { + norm1, + norm2, + attn, + mlp, + }) + } + + fn forward( + &self, + xs: &Tensor, + cu_seqlens: &[usize], + cos: &Tensor, + sin: &Tensor, + ) -> Result { + let normed = self.norm1.forward(xs)?; + let attn_out = self.attn.forward(&normed, cu_seqlens, cos, sin)?; + let xs_att = xs.add(&attn_out)?; + let mlp_out = self.mlp.forward(&self.norm2.forward(&xs_att)?)?; + xs_att.add(&mlp_out) + } +} + +struct PatchMerger { + norm: LayerNorm, + use_postshuffle_norm: bool, + spatial_merge_unit: usize, + merged_hidden_size: usize, + fc1: Linear, + fc2: Linear, +} + +impl PatchMerger { + fn new(cfg: &VisionConfig, use_postshuffle_norm: bool, vb: VarBuilder) -> Result { + let merged_hidden_size = cfg.hidden_size * cfg.spatial_merge_size.pow(2); + let norm_dim = if use_postshuffle_norm { + merged_hidden_size + } else { + cfg.hidden_size + }; + let norm_cfg = LayerNormConfig { + eps: 1e-6, + ..Default::default() + }; + Ok(Self { + norm: layer_norm(norm_dim, norm_cfg, vb.pp("norm"))?, + use_postshuffle_norm, + spatial_merge_unit: cfg.spatial_merge_size.pow(2), + merged_hidden_size, + fc1: linear(merged_hidden_size, merged_hidden_size, vb.pp("linear_fc1"))?, + fc2: linear(merged_hidden_size, cfg.out_hidden_size, vb.pp("linear_fc2"))?, + }) + } + + fn forward(&self, xs: &Tensor) -> Result { + let seq_len = xs.dim(0)?; + if seq_len % self.spatial_merge_unit != 0 { + candle::bail!( + "Sequence length {} is not divisible by spatial merge unit {}", + seq_len, + self.spatial_merge_unit + ); + } + let grouped = seq_len / self.spatial_merge_unit; + let norm_input = if self.use_postshuffle_norm { + xs.reshape((grouped, self.merged_hidden_size))? + } else { + xs.clone() + }; + let normed = self.norm.forward(&norm_input)?; + let reshaped = if self.use_postshuffle_norm { + normed + } else { + normed.reshape((grouped, self.merged_hidden_size))? + }; + let xs = self.fc1.forward(&reshaped)?; + let xs = xs.gelu()?; + self.fc2.forward(&xs) + } +} + +struct VisionRotaryEmbedding { + inv_freq: Tensor, +} + +impl VisionRotaryEmbedding { + const THETA: f32 = 10000.; + + fn new(dim: usize, device: &Device) -> Result { + let inv_freq = (0..dim) + .step_by(2) + .map(|i| 1f32 / Self::THETA.powf(i as f32 / dim as f32)) + .collect::>(); + let inv_freq_len = inv_freq.len(); + Ok(Self { + inv_freq: Tensor::from_vec(inv_freq, (1, inv_freq_len), device)?, + }) + } + + fn make_embeds(&self, seqlen: usize) -> Result { + let seq = + Tensor::arange(0f32, seqlen as f32, self.inv_freq.device())?.unsqueeze(D::Minus1)?; + seq.broadcast_matmul(&self.inv_freq) + } +} + +pub struct Qwen3VLVisionModel { + patch_embed: PatchEmbed, + pos_embed: Embedding, + blocks: Vec, + merger: PatchMerger, + deepstack_mergers: Vec, + deepstack_lookup: Vec>, + rotary_pos_emb: VisionRotaryEmbedding, + spatial_merge_size: usize, + num_grid_per_side: usize, + hidden_size: usize, +} + +impl Qwen3VLVisionModel { + pub fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let patch_embed = PatchEmbed::new(cfg, vb.pp("patch_embed"))?; + let pos_embed = embedding( + cfg.num_position_embeddings, + cfg.hidden_size, + vb.pp("pos_embed"), + )?; + + let mut blocks = Vec::with_capacity(cfg.depth); + for i in 0..cfg.depth { + blocks.push(VisionBlock::new(cfg, vb.pp(format!("blocks.{i}")))?); + } + + let merger = PatchMerger::new(cfg, false, vb.pp("merger"))?; + let deepstack_mergers = cfg + .deepstack_visual_indexes + .iter() + .enumerate() + .map(|(i, _)| PatchMerger::new(cfg, true, vb.pp(format!("deepstack_merger_list.{i}")))) + .collect::>>()?; + + let mut deepstack_lookup = vec![None; cfg.depth]; + for (idx, &layer_idx) in cfg.deepstack_visual_indexes.iter().enumerate() { + if layer_idx < cfg.depth { + deepstack_lookup[layer_idx] = Some(idx); + } + } + + let head_dim = cfg.hidden_size / cfg.num_heads; + let rotary_pos_emb = VisionRotaryEmbedding::new(head_dim / 2, vb.device())?; + + let num_grid_per_side = (cfg.num_position_embeddings as f64).sqrt().round() as usize; + if num_grid_per_side * num_grid_per_side != cfg.num_position_embeddings { + candle::bail!( + "num_position_embeddings {} is not a perfect square", + cfg.num_position_embeddings + ); + } + + Ok(Self { + patch_embed, + pos_embed, + blocks, + merger, + deepstack_mergers, + deepstack_lookup, + rotary_pos_emb, + spatial_merge_size: cfg.spatial_merge_size, + num_grid_per_side, + hidden_size: cfg.hidden_size, + }) + } + + fn linspace_points(&self, steps: usize) -> Vec { + if steps == 1 { + return vec![0.0]; + } + let max_val = (self.num_grid_per_side - 1) as f32; + let step = max_val / (steps.saturating_sub(1)) as f32; + (0..steps).map(|i| i as f32 * step).collect() + } + + fn fast_pos_embed_interpolate(&self, grid_thw: &Tensor) -> Result { + let device = self.pos_embed.embeddings().device(); + let dtype = self.pos_embed.embeddings().dtype(); + let grid = grid_thw.to_vec2::()?; + + let mut idx_lists: [Vec; 4] = Default::default(); + let mut weight_lists: [Vec; 4] = Default::default(); + let mut hw_lengths = Vec::with_capacity(grid.len()); + + for g in &grid { + let h = g[1] as usize; + let w = g[2] as usize; + hw_lengths.push(h * w); + + let h_vals = self.linspace_points(h); + let w_vals = self.linspace_points(w); + + let h_floor: Vec = h_vals.iter().map(|v| v.floor() as usize).collect(); + let w_floor: Vec = w_vals.iter().map(|v| v.floor() as usize).collect(); + let h_ceil: Vec = h_vals + .iter() + .map(|v| (v.ceil() as usize).min(self.num_grid_per_side - 1)) + .collect(); + let w_ceil: Vec = w_vals + .iter() + .map(|v| (v.ceil() as usize).min(self.num_grid_per_side - 1)) + .collect(); + let dh: Vec = h_vals + .iter() + .zip(&h_floor) + .map(|(v, f)| v - *f as f32) + .collect(); + let dw: Vec = w_vals + .iter() + .zip(&w_floor) + .map(|(v, f)| v - *f as f32) + .collect(); + + for ((&hf, &hc), &dh_val) in h_floor.iter().zip(&h_ceil).zip(&dh) { + for ((&wf, &wc), &dw_val) in w_floor.iter().zip(&w_ceil).zip(&dw) { + let base00 = (hf * self.num_grid_per_side + wf) as i64; + let base01 = (hf * self.num_grid_per_side + wc) as i64; + let base10 = (hc * self.num_grid_per_side + wf) as i64; + let base11 = (hc * self.num_grid_per_side + wc) as i64; + + let w00 = (1.0 - dh_val) * (1.0 - dw_val); + let w01 = (1.0 - dh_val) * dw_val; + let w10 = dh_val * (1.0 - dw_val); + let w11 = dh_val * dw_val; + + idx_lists[0].push(base00); + idx_lists[1].push(base01); + idx_lists[2].push(base10); + idx_lists[3].push(base11); + + weight_lists[0].push(w00); + weight_lists[1].push(w01); + weight_lists[2].push(w10); + weight_lists[3].push(w11); + } + } + } + + let idx_tensors = idx_lists + .iter() + .map(|idxs| Tensor::from_vec(idxs.clone(), (idxs.len(),), device)) + .collect::>>()?; + let idx_tensor = Tensor::stack(&idx_tensors, 0)?; + + let weight_tensors = weight_lists + .iter() + .map(|weights| Tensor::from_vec(weights.clone(), (weights.len(),), device)) + .collect::>>()?; + let weight_tensor = Tensor::stack(&weight_tensors, 0)?.to_dtype(dtype)?; + + let pos_embeds = self.pos_embed.forward(&idx_tensor)?; + let pos_embeds = pos_embeds.broadcast_mul(&weight_tensor.unsqueeze(D::Minus1)?)?; + let pos_embeds = pos_embeds.sum(0)?; + + let mut splits = Vec::with_capacity(hw_lengths.len()); + let mut start = 0; + for len in hw_lengths { + splits.push(pos_embeds.narrow(0, start, len)?); + start += len; + } + + let mut permuted = Vec::with_capacity(grid.len()); + for (pos_embed, g) in splits.into_iter().zip(&grid) { + let t = g[0] as usize; + let h = g[1] as usize; + let w = g[2] as usize; + let pos_embed = pos_embed.repeat((t, 1))?; + let pos_embed = pos_embed.reshape(( + t, + h / self.spatial_merge_size, + self.spatial_merge_size, + w / self.spatial_merge_size, + self.spatial_merge_size, + self.hidden_size, + ))?; + let pos_embed = pos_embed + .permute((0, 1, 3, 2, 4, 5))? + .reshape((t * h * w, self.hidden_size))?; + permuted.push(pos_embed); + } + + Tensor::cat(&permuted, 0) + } + + fn rot_pos_emb(&self, grid_thw: &Tensor) -> Result { + let device = self.rotary_pos_emb.inv_freq.device(); + let grid = grid_thw.to_vec2::()?; + let max_hw = grid + .iter() + .flat_map(|v| v[1..3].iter()) + .copied() + .max() + .unwrap_or(0) as usize; + let freq_table = self.rotary_pos_emb.make_embeds(max_hw)?; + + let mut coords: Vec<(i64, i64)> = Vec::new(); + for g in &grid { + let h = g[1] as usize; + let w = g[2] as usize; + let merged_h = h / self.spatial_merge_size; + let merged_w = w / self.spatial_merge_size; + + let mut base_coords: Vec<(i64, i64)> = Vec::with_capacity(h * w); + for br in 0..merged_h { + for bc in 0..merged_w { + for ir in 0..self.spatial_merge_size { + for ic in 0..self.spatial_merge_size { + base_coords.push(( + (br * self.spatial_merge_size + ir) as i64, + (bc * self.spatial_merge_size + ic) as i64, + )); + } + } + } + } + + for _ in 0..(g[0] as usize) { + coords.extend(base_coords.iter().cloned()); + } + } + + let total_tokens = coords.len(); + let mut rows = Vec::with_capacity(total_tokens); + let mut cols = Vec::with_capacity(total_tokens); + for &(r, c) in &coords { + rows.push(r); + cols.push(c); + } + let rows = Tensor::from_vec(rows, (total_tokens,), device)?; + let cols = Tensor::from_vec(cols, (total_tokens,), device)?; + let row_embeds = freq_table.index_select(&rows, 0)?; + let col_embeds = freq_table.index_select(&cols, 0)?; + Tensor::stack(&[row_embeds, col_embeds], D::Minus2)? + .reshape((total_tokens, freq_table.dim(D::Minus1)? * 2)) + } + + fn build_cu_seqlens(&self, grid_thw: &Tensor) -> Result> { + let grid = grid_thw.to_vec2::()?; + let mut cu = Vec::with_capacity(grid.iter().map(|v| v[0] as usize).sum::() + 1); + cu.push(0usize); + let mut acc = 0usize; + for g in &grid { + let area = (g[1] * g[2]) as usize; + for _ in 0..(g[0] as usize) { + acc += area; + cu.push(acc); + } + } + Ok(cu) + } + + pub fn forward(&self, xs: &Tensor, grid_thw: &Tensor) -> Result<(Tensor, Vec)> { + let dtype = self.pos_embed.embeddings().dtype(); + let xs = self.patch_embed.forward(&xs.to_dtype(dtype)?)?; + let pos_embeds = self.fast_pos_embed_interpolate(grid_thw)?; + let mut hidden_states = xs.add(&pos_embeds)?; + + let rotary_pos_emb = self.rot_pos_emb(grid_thw)?; + let seq_len = hidden_states.dim(0)?; + let rotary_pos_emb = rotary_pos_emb.reshape((seq_len, ()))?; + let emb = Tensor::cat(&[&rotary_pos_emb, &rotary_pos_emb], D::Minus1)?; + let cos = emb.cos()?.to_dtype(DType::F32)?; + let sin = emb.sin()?.to_dtype(DType::F32)?; + + let cu_seqlens = self.build_cu_seqlens(grid_thw)?; + + let mut deepstack_features = Vec::new(); + for (layer_idx, block) in self.blocks.iter().enumerate() { + hidden_states = block.forward(&hidden_states, &cu_seqlens, &cos, &sin)?; + if let Some(merger_idx) = self.deepstack_lookup[layer_idx] { + let feat = self.deepstack_mergers[merger_idx].forward(&hidden_states)?; + deepstack_features.push(feat); + } + } + + let hidden_states = self.merger.forward(&hidden_states)?; + Ok((hidden_states, deepstack_features)) + } +} diff --git a/patches/candle-transformers/src/models/recurrent_gemma.rs b/patches/candle-transformers/src/models/recurrent_gemma.rs new file mode 100644 index 0000000000..d6a029babc --- /dev/null +++ b/patches/candle-transformers/src/models/recurrent_gemma.rs @@ -0,0 +1,660 @@ +//! Recurrent Gemma model implementation +//! +//! Recurrent Gemma is a version of the Gemma language model that incorporates recurrent memory. +//! This allows the model to maintain state between predictions and have longer-range memory. +//! +//! Key characteristics: +//! - Real-gated linear recurrent units (RGLRU) +//! - 1D convolution for local context +//! - RMSNorm for layer normalization +//! - Rotary positional embeddings (RoPE) +//! - Grouped query attention +//! +//! References: +//! - [Gemma: Open Models Based on Gemini Technology](https://blog.google/technology/developers/gemma-open-models/) +//! - [Recurrent Memory model architecture](https://arxiv.org/abs/2402.00441) +//! +//! This implementation is based on the python version from huggingface/transformers. +//! https://github.com/huggingface/transformers/blob/b109257f4fb8b1166e7c53cc5418632014ed53a5/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py#L2 +//! +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{linear_b as linear, Linear, VarBuilder}; +use std::sync::Arc; + +#[derive(serde::Deserialize, Debug, Clone, Copy)] +#[serde(rename_all = "snake_case")] +pub enum TemporalBlockType { + Attention, + Recurrent, +} + +#[derive(serde::Deserialize, Debug, Clone)] +pub struct Config { + pub num_hidden_layers: usize, + pub vocab_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub head_dim: usize, + pub lru_width: Option, + pub attention_window_size: usize, + pub conv1d_width: usize, + pub logits_soft_cap: f64, + pub hidden_activation: candle_nn::Activation, + pub partial_rotary_factor: f64, + pub rms_norm_eps: f64, + pub rope_theta: f64, + #[serde(alias = "_block_types")] + pub block_types: Vec, + pub attention_bias: bool, + #[serde(default = "default_max_seq_len")] + pub max_seq_len: usize, +} + +fn default_max_seq_len() -> usize { + 8192 +} + +#[derive(Debug, Clone)] +pub(crate) struct RmsNorm { + weight: Tensor, + eps: f64, +} + +impl RmsNorm { + pub(crate) fn new(dim: usize, eps: f64, vb: VarBuilder) -> Result { + let weight = vb.get(dim, "weight")?; + Ok(Self { weight, eps }) + } + + pub(crate) fn from_weight(weight: Tensor, eps: f64) -> Self { + Self { weight, eps } + } +} + +impl Module for RmsNorm { + fn forward(&self, x: &Tensor) -> Result { + let x_dtype = x.dtype(); + let internal_dtype = match x_dtype { + DType::F16 | DType::BF16 => DType::F32, + d => d, + }; + let hidden_size = x.dim(D::Minus1)?; + let x = x.to_dtype(internal_dtype)?; + let norm_x = (x.sqr()?.sum_keepdim(D::Minus1)? / hidden_size as f64)?; + let x_normed = x.broadcast_div(&(norm_x + self.eps)?.sqrt()?)?; + x_normed + .to_dtype(x_dtype)? + .broadcast_mul(&(&self.weight + 1.0)?) + } +} + +#[derive(Debug, Clone)] +pub(crate) struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +fn rotate_half(xs: &Tensor) -> Result { + let last_dim = xs.dim(D::Minus1)?; + let xs1 = xs.narrow(D::Minus1, 0, last_dim / 2)?; + let xs2 = xs.narrow(D::Minus1, last_dim / 2, last_dim - last_dim / 2)?; + Tensor::cat(&[&xs2.neg()?, &xs1], D::Minus1) +} + +impl RotaryEmbedding { + pub(crate) fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + if cfg.partial_rotary_factor != 0.5 { + candle::bail!("partial-rotary-factor {} <> 0.5", cfg.partial_rotary_factor) + } + let dim = cfg.head_dim / 2; + let max_seq_len = cfg.max_seq_len; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + let freqs = Tensor::cat(&[&freqs, &freqs], D::Minus1)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + pub(crate) fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let cos = cos.unsqueeze(0)?.unsqueeze(0)?; // (1, 1, seq_len, dim) + let sin = sin.unsqueeze(0)?.unsqueeze(0)?; // (1, 1, seq_len, dim) + let q_embed = (q.broadcast_mul(&cos)? + rotate_half(q)?.broadcast_mul(&sin))?; + let k_embed = (k.broadcast_mul(&cos)? + rotate_half(k)?.broadcast_mul(&sin))?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +struct Mlp { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: candle_nn::Activation, +} + +impl Mlp { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let h = cfg.hidden_size; + let intermediate_size = cfg.intermediate_size / 2; + let gate_proj = linear(h, intermediate_size, true, vb.pp("gate_proj"))?; + let up_proj = linear(h, intermediate_size, true, vb.pp("up_proj"))?; + let down_proj = linear(intermediate_size, h, true, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_activation, + }) + } +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let gate = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; + (gate * xs.apply(&self.up_proj))?.apply(&self.down_proj) + } +} + +// Real-Gated Linear Recurrent Unit +#[derive(Debug, Clone)] +pub(crate) struct Rglru { + pub(crate) recurrent_param: Tensor, + pub(crate) input_gate_weight: Tensor, + pub(crate) input_gate_bias: Tensor, + pub(crate) recurrent_gate_weight: Tensor, + pub(crate) recurrent_gate_bias: Tensor, + pub(crate) block_width: usize, + pub(crate) n_heads: usize, + pub(crate) recurrent_states: Option, +} + +fn baddbmm(a: &Tensor, b: &Tensor, c: &Tensor) -> Result { + a.broadcast_add(&b.matmul(c)?) +} + +fn softplus(xs: &Tensor) -> Result { + (xs.exp()? + 1.0)?.log() +} + +impl Rglru { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let h = cfg.hidden_size; + let lru_width = cfg.lru_width.unwrap_or(h); + let n_heads = cfg.num_attention_heads; + let block_width = lru_width / n_heads; + let recurrent_param = vb.get((lru_width,), "recurrent_param")?; + let input_gate_weight = vb.get((n_heads, block_width, block_width), "input_gate_weight")?; + let input_gate_bias = vb.get((n_heads, block_width), "input_gate_bias")?; + let recurrent_gate_weight = + vb.get((n_heads, block_width, block_width), "recurrent_gate_weight")?; + let recurrent_gate_bias = vb.get((n_heads, block_width), "recurrent_gate_bias")?; + Ok(Self { + recurrent_param, + input_gate_bias, + input_gate_weight, + recurrent_gate_bias, + recurrent_gate_weight, + block_width, + n_heads, + recurrent_states: None, + }) + } + + // https://github.com/huggingface/transformers/blob/0bd58f1ce0573c0e3269de4215a17d318add49b9/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py#L303 + pub(crate) fn forward(&mut self, xs: &Tensor, pos: usize) -> Result { + let (b_sz, seq_len, lru_width) = xs.dims3()?; + let pos = Tensor::arange(pos as u32, (pos + seq_len) as u32, xs.device())?; + let reset = pos.eq(0u32)?.unsqueeze(1)?.unsqueeze(0)?; + let reshape_act = xs + .reshape((b_sz * seq_len, self.n_heads, self.block_width))? + .permute((1, 0, 2))? + .contiguous()?; + + let res = baddbmm( + &self.input_gate_bias.unsqueeze(1)?, + &reshape_act, + &self.input_gate_weight, + )?; + let input_gate = res.transpose(0, 1)?.reshape((b_sz, seq_len, lru_width))?; + let input_gate = candle_nn::ops::sigmoid(&input_gate)?; + let res = baddbmm( + &self.recurrent_gate_bias.unsqueeze(1)?, + &reshape_act, + &self.recurrent_gate_weight, + )?; + let recurrent_gate = res.transpose(0, 1)?.reshape((b_sz, seq_len, lru_width))?; + let recurrent_gate = candle_nn::ops::sigmoid(&recurrent_gate)?; + + let log_recurrent_gate = + (recurrent_gate * (-8.0))?.broadcast_mul(&softplus(&self.recurrent_param)?)?; + let recurrent_gate = log_recurrent_gate.exp()?; + let a_square = (log_recurrent_gate * 2.)?.exp()?; + + // Gate the input. + let gated_inputs = (xs * input_gate)?; + + let reset = reset.to_dtype(a_square.dtype())?; + let multiplier = + reset.broadcast_add(&((1.0 - &reset)?.broadcast_mul(&(1.0 - a_square)?.sqrt()?))?)?; + let normalized_x = (gated_inputs * multiplier.to_dtype(xs.dtype()))?; + + let (hidden_states, recurrent_states) = rnn_scan( + &normalized_x, + &recurrent_gate, + &reset, + self.recurrent_states.as_ref(), + )?; + self.recurrent_states = Some(recurrent_states); + Ok(hidden_states) + } +} + +fn rnn_scan( + hidden_states: &Tensor, + recurrent_gate: &Tensor, + reset: &Tensor, + recurrent_states: Option<&Tensor>, +) -> Result<(Tensor, Tensor)> { + let acc_dtype = DType::F32; + let dev = hidden_states.device(); + let in_dtype = hidden_states.dtype(); + let inv_reset = (1.0 - reset)?.to_dtype(recurrent_gate.dtype())?; + let recurrent_gate = recurrent_gate.broadcast_mul(&inv_reset)?; + let (c, r) = if hidden_states.dim(1)? == 1 { + match recurrent_states { + None => { + let next_state = hidden_states.i((.., 0))?.to_dtype(acc_dtype)?; + (hidden_states.clone(), next_state) + } + Some(recurrent_states) => { + let contextualized_states = + recurrent_gate.to_dtype(acc_dtype)? * recurrent_states.unsqueeze(1)?; + let contextualized_states = + (contextualized_states + hidden_states.to_dtype(acc_dtype)?)?; + let c = contextualized_states.to_dtype(in_dtype)?; + let l = contextualized_states.dim(1)?; + let r = contextualized_states.i((.., l - 1))?; + (c, r) + } + } + } else { + let mut recurrent_states = match recurrent_states { + None => Tensor::zeros(hidden_states.i((.., 0))?.shape(), acc_dtype, dev)?, + Some(r) => r.clone(), + }; + let mut contextualized_states = vec![]; + for t in 0..hidden_states.dim(1)? { + recurrent_states = + (recurrent_gate.i((.., t))?.to_dtype(acc_dtype)? * recurrent_states)?; + recurrent_states = + (recurrent_states + hidden_states.i((.., t))?.to_dtype(acc_dtype)?)?; + contextualized_states.push(recurrent_states.to_dtype(in_dtype)?) + } + let contextualized_states = Tensor::stack(&contextualized_states, 1)?; + (contextualized_states, recurrent_states) + }; + Ok((c, r)) +} + +#[derive(Debug, Clone)] +struct RecurrentBlock { + linear_y: Linear, + linear_x: Linear, + linear_out: Linear, + conv_1d: candle_nn::Conv1d, + conv1d_state: Option, + conv1d_width: usize, + rg_lru: Rglru, + act_fn: candle_nn::Activation, +} + +impl RecurrentBlock { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let h = cfg.hidden_size; + let lru_width = cfg.lru_width.unwrap_or(h); + let linear_y = linear(h, lru_width, true, vb.pp("linear_y"))?; + let linear_x = linear(h, lru_width, true, vb.pp("linear_x"))?; + let linear_out = linear(lru_width, h, true, vb.pp("linear_out"))?; + let conv_1d = candle_nn::conv1d( + lru_width, + lru_width, + cfg.conv1d_width, + candle_nn::Conv1dConfig { + groups: lru_width, + padding: cfg.conv1d_width - 1, + ..Default::default() + }, + vb.pp("conv_1d"), + )?; + let rg_lru = Rglru::new(cfg, vb.pp("rg_lru"))?; + Ok(Self { + linear_y, + linear_x, + linear_out, + conv_1d, + conv1d_state: None, + conv1d_width: cfg.conv1d_width, + rg_lru, + act_fn: cfg.hidden_activation, + }) + } + + pub fn forward(&mut self, xs: &Tensor, pos: usize) -> Result { + let (_b_sz, seq_len, _) = xs.dims3()?; + + let y_branch = xs.apply(&self.linear_y)?.apply(&self.act_fn)?; + let x_branch = xs.apply(&self.linear_x)?.transpose(1, 2)?; + let x_branch = if pos == 0 { + let x_len = x_branch.dim(D::Minus1)?; + let pad = self.conv1d_width as i64 - x_len as i64 - 1; + let padded = match pad.cmp(&0) { + std::cmp::Ordering::Equal => x_branch.clone(), + std::cmp::Ordering::Less => { + let rev_pad = (-pad) as usize; + x_branch.narrow(D::Minus1, rev_pad, x_len - rev_pad)? + } + std::cmp::Ordering::Greater => { + x_branch.pad_with_zeros(D::Minus1, pad as usize, 0)? + } + }; + self.conv1d_state = Some(padded); + x_branch + .apply(&self.conv_1d)? + .narrow(D::Minus1, 0, seq_len)? + } else { + let conv_state = match self.conv1d_state.as_ref() { + None => candle::bail!("empty cache despite pos > 0"), + Some(s) => Tensor::cat(&[s, &x_branch], D::Minus1)?, + }; + let w = self.conv_1d.weight().i((.., 0, ..))?; + let x_branch = conv_state.broadcast_mul(&w)?.sum(D::Minus1)?; + let x_branch = match self.conv_1d.bias() { + None => x_branch, + Some(b) => x_branch.broadcast_add(b)?, + }; + let x_branch = x_branch.unsqueeze(D::Minus1)?; + self.conv1d_state = Some(conv_state.i((.., .., 1..))?); + x_branch + }; + let x_branch = x_branch.transpose(1, 2)?; + let x_branch = self.rg_lru.forward(&x_branch, pos)?; + (x_branch * y_branch)?.apply(&self.linear_out) + } +} + +#[derive(Debug, Clone)] +struct SdpaAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + n_heads: usize, + n_kv_heads: usize, + head_dim: usize, + hidden_size: usize, + kv_cache: Option<(Tensor, Tensor)>, + rotary_emb: Arc, +} + +impl SdpaAttention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let h = cfg.hidden_size; + let n_heads = cfg.num_attention_heads; + let n_kv_heads = cfg.num_key_value_heads; + let hd = cfg.head_dim; + let q_proj = linear(h, n_heads * hd, cfg.attention_bias, vb.pp("q_proj"))?; + let k_proj = linear(h, n_kv_heads * hd, cfg.attention_bias, vb.pp("k_proj"))?; + let v_proj = linear(h, n_kv_heads * hd, cfg.attention_bias, vb.pp("v_proj"))?; + let o_proj = linear(n_heads * hd, h, true, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + n_heads, + n_kv_heads, + head_dim: hd, + hidden_size: h, + kv_cache: None, + rotary_emb, + }) + } + + fn repeat_kv(&self, x: Tensor) -> Result { + let n_rep = self.n_heads / self.n_kv_heads; + crate::utils::repeat_kv(x, n_rep) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + pos: usize, + ) -> Result { + let (bsz, q_len, _) = xs.dims3()?; + + let query_states = xs.apply(&self.q_proj)?; + let key_states = xs.apply(&self.k_proj)?; + let value_states = xs.apply(&self.v_proj)?; + + let query_states = query_states + .reshape((bsz, q_len, self.n_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((bsz, q_len, self.n_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((bsz, q_len, self.n_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let query_states = query_states.chunk(2, D::Minus1)?; + let key_states = key_states.chunk(2, D::Minus1)?; + let (query_rot, key_rot) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states[0], &key_states[0], pos)?; + let query_states = Tensor::cat(&[&query_rot, &query_states[1]], D::Minus1)?.contiguous()?; + let key_states = Tensor::cat(&[&key_rot, &key_states[1]], D::Minus1)?.contiguous()?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + let key_states = self.repeat_kv(key_states)?; + let value_states = self.repeat_kv(value_states)?; + let xs = { + let att = (query_states.matmul(&key_states.t()?)? / (self.head_dim as f64).sqrt())?; + let att = if q_len == 1 { + att + } else { + match attention_mask { + None => att, + Some(mask) => att.broadcast_add(mask)?, + } + }; + let att = candle_nn::ops::softmax_last_dim(&att)?; + att.matmul(&value_states.contiguous()?)? + }; + + let xs = xs + .transpose(1, 2)? + .reshape((bsz, q_len, self.hidden_size))?; + self.o_proj.forward(&xs) + } +} + +#[derive(Debug, Clone)] +enum TemporalBlock { + Recurrent(RecurrentBlock), + Attention(SdpaAttention), +} + +impl TemporalBlock { + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + pos: usize, + ) -> Result { + match self { + Self::Recurrent(b) => b.forward(xs, pos), + Self::Attention(b) => b.forward(xs, attention_mask, pos), + } + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + temporal_pre_norm: RmsNorm, + channel_pre_norm: RmsNorm, + temporal_block: TemporalBlock, + mlp_block: Mlp, +} + +impl DecoderLayer { + fn new( + block_idx: usize, + rotary_emb: Arc, + cfg: &Config, + vb: VarBuilder, + ) -> Result { + let h = cfg.hidden_size; + let temporal_pre_norm = RmsNorm::new(h, cfg.rms_norm_eps, vb.pp("temporal_pre_norm"))?; + let channel_pre_norm = RmsNorm::new(h, cfg.rms_norm_eps, vb.pp("channel_pre_norm"))?; + let temporal_block = match cfg.block_types[block_idx % cfg.block_types.len()] { + TemporalBlockType::Recurrent => { + let block = RecurrentBlock::new(cfg, vb.pp("temporal_block"))?; + TemporalBlock::Recurrent(block) + } + TemporalBlockType::Attention => { + let block = SdpaAttention::new(rotary_emb, cfg, vb.pp("temporal_block"))?; + TemporalBlock::Attention(block) + } + }; + let mlp_block = Mlp::new(cfg, vb.pp("mlp_block"))?; + Ok(Self { + temporal_pre_norm, + channel_pre_norm, + temporal_block, + mlp_block, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + pos: usize, + ) -> Result { + let residual = xs; + let xs = xs.apply(&self.temporal_pre_norm)?; + let xs = self.temporal_block.forward(&xs, attention_mask, pos)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.channel_pre_norm)?.apply(&self.mlp_block)?; + xs + residual + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + final_norm: RmsNorm, + lm_head: Linear, + hidden_size: usize, + logits_soft_cap: f64, + dtype: DType, + device: Device, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb.device())?); + let vb_b = vb.pp("layers"); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + for idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(idx, rotary_emb.clone(), cfg, vb_b.pp(idx))?; + layers.push(layer) + } + let final_norm = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("final_norm"))?; + let lm_head = Linear::new(embed_tokens.embeddings().clone(), None); + Ok(Self { + embed_tokens, + layers, + final_norm, + lm_head, + hidden_size: cfg.hidden_size, + logits_soft_cap: cfg.logits_soft_cap, + dtype: vb.dtype(), + device: vb.device().clone(), + }) + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { f32::NEG_INFINITY } else { 0. })) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + pub fn forward(&mut self, xs: &Tensor, pos: usize) -> Result { + let (b_size, seq_len) = xs.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(b_size, seq_len, pos)?; + Some(mask) + }; + let xs = xs.apply(&self.embed_tokens)?; + let mut xs = (xs * (self.hidden_size as f64).sqrt())?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), pos)?; + } + let logits = xs + .narrow(1, seq_len - 1, 1)? + .apply(&self.final_norm)? + .apply(&self.lm_head)?; + let logits = ((logits / self.logits_soft_cap)?.tanh()? * self.logits_soft_cap)?; + Ok(logits) + } +} diff --git a/patches/candle-transformers/src/models/repvgg.rs b/patches/candle-transformers/src/models/repvgg.rs new file mode 100644 index 0000000000..6e45c2d68c --- /dev/null +++ b/patches/candle-transformers/src/models/repvgg.rs @@ -0,0 +1,314 @@ +//! RepVGG inference implementation +//! +//! Key characteristics: +//! - Efficient inference architecture through structural reparameterization +//! - Single 3x3 conv layer after fusing 3x3 branch, 1x1 branch and identity branch +//! - Different configurations including a0-a2, b0-b3 and variants with group convolutions +//! - High accuracy with VGG-like plain architecture and training +//! +//! References: +//! - [RepVGG Paper](https://arxiv.org/abs/2101.03697). RepVGG: Making VGG-style ConvNets Great Again +//! - [Official Implementation](https://github.com/DingXiaoH/RepVGG) +//! + +use candle::{Result, Tensor, D}; +use candle_nn::{ + batch_norm, conv2d_no_bias, linear, BatchNorm, Conv2d, Conv2dConfig, Func, VarBuilder, +}; + +const CHANNELS_PER_STAGE: [usize; 5] = [64, 64, 128, 256, 512]; + +#[derive(Clone)] +pub struct Config { + a: f32, + b: f32, + groups: usize, + stages: [usize; 4], +} + +impl Config { + pub fn a0() -> Self { + Self { + a: 0.75, + b: 2.5, + groups: 1, + stages: [2, 4, 14, 1], + } + } + + pub fn a1() -> Self { + Self { + a: 1.0, + b: 2.5, + groups: 1, + stages: [2, 4, 14, 1], + } + } + + pub fn a2() -> Self { + Self { + a: 1.5, + b: 2.75, + groups: 1, + stages: [2, 4, 14, 1], + } + } + + pub fn b0() -> Self { + Self { + a: 1.0, + b: 2.5, + groups: 1, + stages: [4, 6, 16, 1], + } + } + + pub fn b1() -> Self { + Self { + a: 2.0, + b: 4.0, + groups: 1, + stages: [4, 6, 16, 1], + } + } + + pub fn b2() -> Self { + Self { + a: 2.5, + b: 5.0, + groups: 1, + stages: [4, 6, 16, 1], + } + } + + pub fn b3() -> Self { + Self { + a: 3.0, + b: 5.0, + groups: 1, + stages: [4, 6, 16, 1], + } + } + + pub fn b1g4() -> Self { + Self { + a: 2.0, + b: 4.0, + groups: 4, + stages: [4, 6, 16, 1], + } + } + + pub fn b2g4() -> Self { + Self { + a: 2.5, + b: 5.0, + groups: 4, + stages: [4, 6, 16, 1], + } + } + + pub fn b3g4() -> Self { + Self { + a: 3.0, + b: 5.0, + groups: 4, + stages: [4, 6, 16, 1], + } + } +} + +// fuses a convolutional kernel and a batchnorm layer into a convolutional layer +// based on the _fuse_bn_tensor method in timm +// see https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/byobnet.py#L602 +fn fuse_conv_bn(weights: &Tensor, bn: BatchNorm) -> Result<(Tensor, Tensor)> { + let (gamma, beta) = bn.weight_and_bias().unwrap(); + let mu = bn.running_mean(); + let sigma = (bn.running_var() + bn.eps())?.sqrt(); + let gps = (gamma / sigma)?; + let bias = (beta - mu * &gps)?; + let weights = weights.broadcast_mul(&gps.reshape(((), 1, 1, 1))?)?; + + Ok((weights, bias)) +} + +// A RepVGG layer has a different training time and inference time architecture. +// The latter is a simple and efficient equivalent transformation of the former +// realized by a structural reparameterization technique, where 3x3 and 1x1 convolutions +// along with identity branches and batchnorm layers are fused into a single 3x3 convolution. +fn repvgg_layer( + has_identity: bool, + dim: usize, + stride: usize, + in_channels: usize, + out_channels: usize, + groups: usize, + vb: VarBuilder, +) -> Result> { + let conv2d_cfg = Conv2dConfig { + stride, + groups, + padding: 1, + ..Default::default() + }; + + // read and reparameterize the 1x1 conv and bn into w1 and b1 + // based on https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/byobnet.py#L543 + + let conv1x1_bn = batch_norm(dim, 1e-5, vb.pp("conv_1x1.bn"))?; + let conv1x1 = conv2d_no_bias( + in_channels, + out_channels, + 1, + conv2d_cfg, + vb.pp("conv_1x1.conv"), + )?; + + let (mut w1, b1) = fuse_conv_bn(conv1x1.weight(), conv1x1_bn)?; + + // resize to 3x3 + w1 = w1.pad_with_zeros(D::Minus1, 1, 1)?; + w1 = w1.pad_with_zeros(D::Minus2, 1, 1)?; + + // read and reparameterize the 3x3 conv and bn into w3 and b3 + let convkxk_bn = batch_norm(dim, 1e-5, vb.pp("conv_kxk.bn"))?; + let conv3x3 = conv2d_no_bias( + in_channels, + out_channels, + 3, + conv2d_cfg, + vb.pp("conv_kxk.conv"), + )?; + + let (w3, b3) = fuse_conv_bn(conv3x3.weight(), convkxk_bn)?; + + let mut w = (w1 + w3)?; + let mut b = (b1 + b3)?; + + // read and reparameterize the identity bn into wi and bi + if has_identity { + let identity_bn = batch_norm(dim, 1e-5, vb.pp("identity"))?; + + // create a 3x3 convolution equivalent to the identity branch + let mut weights: Vec = vec![0.0; conv3x3.weight().elem_count()]; + + // https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/byobnet.py#L620 + let in_dim = in_channels / groups; + for i in 0..in_channels { + weights[i * in_dim * 3 * 3 + (i % in_dim) * 3 * 3 + 4] = 1.0; + } + + let weights = &Tensor::from_vec(weights, w.shape(), w.device())?; + let (wi, bi) = fuse_conv_bn(weights, identity_bn)?; + + w = (w + wi)?; + b = (b + bi)?; + } + + // create the 3x3 conv equivalent to the sum of 3x3, 1x1 and identity branches + let reparam_conv = Conv2d::new(w, Some(b), conv2d_cfg); + + Ok(Func::new(move |xs| { + let xs = xs.apply(&reparam_conv)?.relu()?; + Ok(xs) + })) +} + +// Get the number of output channels per stage taking into account the multipliers +fn output_channels_per_stage(a: f32, b: f32, stage: usize) -> usize { + let channels = CHANNELS_PER_STAGE[stage] as f32; + + match stage { + 0 => std::cmp::min(64, (channels * a) as usize), + 4 => (channels * b) as usize, + _ => (channels * a) as usize, + } +} + +// Each stage is made of layers. The first layer always downsamples with stride 2. +// All but the first layer have a residual connection. +// The G4 variants have a groupwise convolution instead of a dense one on odd layers +// counted across stage boundaries, so we keep track of which layer we are in the +// full model. +fn repvgg_stage(cfg: &Config, idx: usize, vb: VarBuilder) -> Result> { + let nlayers = cfg.stages[idx - 1]; + let mut layers = Vec::with_capacity(nlayers); + let prev_layers: usize = cfg.stages[..idx - 1].iter().sum(); + let out_channels_prev = output_channels_per_stage(cfg.a, cfg.b, idx - 1); + let out_channels = output_channels_per_stage(cfg.a, cfg.b, idx); + + for layer_idx in 0..nlayers { + let (has_identity, stride, in_channels) = if layer_idx == 0 { + (false, 2, out_channels_prev) + } else { + (true, 1, out_channels) + }; + + let groups = if (prev_layers + layer_idx) % 2 == 1 { + cfg.groups + } else { + 1 + }; + + layers.push(repvgg_layer( + has_identity, + out_channels, + stride, + in_channels, + out_channels, + groups, + vb.pp(layer_idx), + )?) + } + + Ok(Func::new(move |xs| { + let mut xs = xs.clone(); + for layer in layers.iter() { + xs = xs.apply(layer)? + } + Ok(xs) + })) +} + +// Build a RepVGG model for a given configuration. +fn repvgg_model(config: &Config, nclasses: Option, vb: VarBuilder) -> Result> { + let cls = match nclasses { + None => None, + Some(nclasses) => { + let outputs = output_channels_per_stage(config.a, config.b, 4); + let linear = linear(outputs, nclasses, vb.pp("head.fc"))?; + Some(linear) + } + }; + + let stem_dim = output_channels_per_stage(config.a, config.b, 0); + let stem = repvgg_layer(false, stem_dim, 2, 3, stem_dim, 1, vb.pp("stem"))?; + let vb = vb.pp("stages"); + let stage1 = repvgg_stage(config, 1, vb.pp(0))?; + let stage2 = repvgg_stage(config, 2, vb.pp(1))?; + let stage3 = repvgg_stage(config, 3, vb.pp(2))?; + let stage4 = repvgg_stage(config, 4, vb.pp(3))?; + + Ok(Func::new(move |xs| { + let xs = xs + .apply(&stem)? + .apply(&stage1)? + .apply(&stage2)? + .apply(&stage3)? + .apply(&stage4)? + .mean(D::Minus1)? + .mean(D::Minus1)?; + match &cls { + None => Ok(xs), + Some(cls) => xs.apply(cls), + } + })) +} + +pub fn repvgg(cfg: &Config, nclasses: usize, vb: VarBuilder) -> Result> { + repvgg_model(cfg, Some(nclasses), vb) +} + +pub fn repvgg_no_final_layer(cfg: &Config, vb: VarBuilder) -> Result> { + repvgg_model(cfg, None, vb) +} diff --git a/patches/candle-transformers/src/models/resnet.rs b/patches/candle-transformers/src/models/resnet.rs new file mode 100644 index 0000000000..31395c8f84 --- /dev/null +++ b/patches/candle-transformers/src/models/resnet.rs @@ -0,0 +1,257 @@ +//! # ResNet Implementation +//! +//! Implementation of ResNet architectures as described in the paper: +//! +//! ## Reference +//! +//! [Deep Residual Learning for Image Recognition](https://arxiv.org/abs/1512.03385) +//! He et al. (2015) +//! +//! This paper introduced ResNet, a deep neural network architecture that utilizes +//! skip connections ("residual connections") to enable training of very deep networks. + +use candle::{Result, D}; +use candle_nn::{batch_norm, Conv2d, Func, VarBuilder}; + +fn conv2d( + c_in: usize, + c_out: usize, + ksize: usize, + padding: usize, + stride: usize, + vb: VarBuilder, +) -> Result { + let conv2d_cfg = candle_nn::Conv2dConfig { + stride, + padding, + ..Default::default() + }; + candle_nn::conv2d_no_bias(c_in, c_out, ksize, conv2d_cfg, vb) +} + +fn downsample(c_in: usize, c_out: usize, stride: usize, vb: VarBuilder) -> Result { + if stride != 1 || c_in != c_out { + let conv = conv2d(c_in, c_out, 1, 0, stride, vb.pp(0))?; + let bn = batch_norm(c_out, 1e-5, vb.pp(1))?; + Ok(Func::new(move |xs| xs.apply(&conv)?.apply_t(&bn, false))) + } else { + Ok(Func::new(|xs| Ok(xs.clone()))) + } +} + +fn basic_block(c_in: usize, c_out: usize, stride: usize, vb: VarBuilder) -> Result { + let conv1 = conv2d(c_in, c_out, 3, 1, stride, vb.pp("conv1"))?; + let bn1 = batch_norm(c_out, 1e-5, vb.pp("bn1"))?; + let conv2 = conv2d(c_out, c_out, 3, 1, 1, vb.pp("conv2"))?; + let bn2 = batch_norm(c_out, 1e-5, vb.pp("bn2"))?; + let downsample = downsample(c_in, c_out, stride, vb.pp("downsample"))?; + Ok(Func::new(move |xs| { + let ys = xs + .apply(&conv1)? + .apply_t(&bn1, false)? + .relu()? + .apply(&conv2)? + .apply_t(&bn2, false)?; + (xs.apply(&downsample)? + ys)?.relu() + })) +} + +fn basic_layer( + c_in: usize, + c_out: usize, + stride: usize, + cnt: usize, + vb: VarBuilder, +) -> Result { + let mut layers = Vec::with_capacity(cnt); + for index in 0..cnt { + let l_in = if index == 0 { c_in } else { c_out }; + let stride = if index == 0 { stride } else { 1 }; + layers.push(basic_block(l_in, c_out, stride, vb.pp(index))?) + } + Ok(Func::new(move |xs| { + let mut xs = xs.clone(); + for layer in layers.iter() { + xs = xs.apply(layer)? + } + Ok(xs) + })) +} + +fn resnet( + nclasses: Option, + c1: usize, + c2: usize, + c3: usize, + c4: usize, + vb: VarBuilder, +) -> Result { + let conv1 = conv2d(3, 64, 7, 3, 2, vb.pp("conv1"))?; + let bn1 = batch_norm(64, 1e-5, vb.pp("bn1"))?; + let layer1 = basic_layer(64, 64, 1, c1, vb.pp("layer1"))?; + let layer2 = basic_layer(64, 128, 2, c2, vb.pp("layer2"))?; + let layer3 = basic_layer(128, 256, 2, c3, vb.pp("layer3"))?; + let layer4 = basic_layer(256, 512, 2, c4, vb.pp("layer4"))?; + let fc = match nclasses { + None => None, + Some(nclasses) => { + let linear = candle_nn::linear(512, nclasses, vb.pp("fc"))?; + Some(linear) + } + }; + Ok(Func::new(move |xs| { + let xs = xs + .apply(&conv1)? + .apply_t(&bn1, false)? + .relu()? + .pad_with_same(D::Minus1, 1, 1)? + .pad_with_same(D::Minus2, 1, 1)? + .max_pool2d_with_stride(3, 2)? + .apply(&layer1)? + .apply(&layer2)? + .apply(&layer3)? + .apply(&layer4)? + .mean(D::Minus1)? + .mean(D::Minus1)?; + match &fc { + None => Ok(xs), + Some(fc) => xs.apply(fc), + } + })) +} + +/// Creates a ResNet-18 model. +pub fn resnet18(num_classes: usize, vb: VarBuilder) -> Result { + resnet(Some(num_classes), 2, 2, 2, 2, vb) +} + +pub fn resnet18_no_final_layer(vb: VarBuilder) -> Result { + resnet(None, 2, 2, 2, 2, vb) +} + +/// Creates a ResNet-34 model. +pub fn resnet34(num_classes: usize, vb: VarBuilder) -> Result { + resnet(Some(num_classes), 3, 4, 6, 3, vb) +} + +pub fn resnet34_no_final_layer(vb: VarBuilder) -> Result { + resnet(None, 3, 4, 6, 3, vb) +} + +// Bottleneck versions for ResNet 50, 101, and 152. +fn bottleneck_block( + c_in: usize, + c_out: usize, + stride: usize, + e: usize, + vb: VarBuilder, +) -> Result { + let e_dim = e * c_out; + let conv1 = conv2d(c_in, c_out, 1, 0, 1, vb.pp("conv1"))?; + let bn1 = batch_norm(c_out, 1e-5, vb.pp("bn1"))?; + let conv2 = conv2d(c_out, c_out, 3, 1, stride, vb.pp("conv2"))?; + let bn2 = batch_norm(c_out, 1e-5, vb.pp("bn2"))?; + let conv3 = conv2d(c_out, e_dim, 1, 0, 1, vb.pp("conv3"))?; + let bn3 = batch_norm(e_dim, 1e-5, vb.pp("bn3"))?; + let downsample = downsample(c_in, e_dim, stride, vb.pp("downsample"))?; + Ok(Func::new(move |xs| { + let ys = xs + .apply(&conv1)? + .apply_t(&bn1, false)? + .relu()? + .apply(&conv2)? + .apply_t(&bn2, false)? + .relu()? + .apply(&conv3)? + .apply_t(&bn3, false)?; + (xs.apply(&downsample)? + ys)?.relu() + })) +} + +fn bottleneck_layer( + c_in: usize, + c_out: usize, + stride: usize, + cnt: usize, + vb: VarBuilder, +) -> Result { + let mut layers = Vec::with_capacity(cnt); + for index in 0..cnt { + let l_in = if index == 0 { c_in } else { 4 * c_out }; + let stride = if index == 0 { stride } else { 1 }; + layers.push(bottleneck_block(l_in, c_out, stride, 4, vb.pp(index))?) + } + Ok(Func::new(move |xs| { + let mut xs = xs.clone(); + for layer in layers.iter() { + xs = xs.apply(layer)? + } + Ok(xs) + })) +} + +fn bottleneck_resnet( + nclasses: Option, + c1: usize, + c2: usize, + c3: usize, + c4: usize, + vb: VarBuilder, +) -> Result { + let conv1 = conv2d(3, 64, 7, 3, 2, vb.pp("conv1"))?; + let bn1 = batch_norm(64, 1e-5, vb.pp("bn1"))?; + let layer1 = bottleneck_layer(64, 64, 1, c1, vb.pp("layer1"))?; + let layer2 = bottleneck_layer(4 * 64, 128, 2, c2, vb.pp("layer2"))?; + let layer3 = bottleneck_layer(4 * 128, 256, 2, c3, vb.pp("layer3"))?; + let layer4 = bottleneck_layer(4 * 256, 512, 2, c4, vb.pp("layer4"))?; + let fc = match nclasses { + None => None, + Some(nclasses) => { + let linear = candle_nn::linear(4 * 512, nclasses, vb.pp("fc"))?; + Some(linear) + } + }; + Ok(Func::new(move |xs| { + let xs = xs + .apply(&conv1)? + .apply_t(&bn1, false)? + .relu()? + .pad_with_same(D::Minus1, 1, 1)? + .pad_with_same(D::Minus2, 1, 1)? + .max_pool2d_with_stride(3, 2)? + .apply(&layer1)? + .apply(&layer2)? + .apply(&layer3)? + .apply(&layer4)? + .mean(D::Minus1)? + .mean(D::Minus1)?; + match &fc { + None => Ok(xs), + Some(fc) => xs.apply(fc), + } + })) +} + +pub fn resnet50(num_classes: usize, vb: VarBuilder) -> Result { + bottleneck_resnet(Some(num_classes), 3, 4, 6, 3, vb) +} + +pub fn resnet50_no_final_layer(vb: VarBuilder) -> Result { + bottleneck_resnet(None, 3, 4, 6, 3, vb) +} + +pub fn resnet101(num_classes: usize, vb: VarBuilder) -> Result { + bottleneck_resnet(Some(num_classes), 3, 4, 23, 3, vb) +} + +pub fn resnet101_no_final_layer(vb: VarBuilder) -> Result { + bottleneck_resnet(None, 3, 4, 23, 3, vb) +} + +pub fn resnet152(num_classes: usize, vb: VarBuilder) -> Result { + bottleneck_resnet(Some(num_classes), 3, 8, 36, 3, vb) +} + +pub fn resnet152_no_final_layer(vb: VarBuilder) -> Result { + bottleneck_resnet(None, 3, 8, 36, 3, vb) +} diff --git a/patches/candle-transformers/src/models/rwkv_v5.rs b/patches/candle-transformers/src/models/rwkv_v5.rs new file mode 100644 index 0000000000..15e386d292 --- /dev/null +++ b/patches/candle-transformers/src/models/rwkv_v5.rs @@ -0,0 +1,441 @@ +//! RWKV v5 model implementation. +//! +//! The [RWKV model](https://wiki.rwkv.com/) is a recurrent neural network model +//! with performance on par with transformer architectures. Several variants are +//! available, candle implements the v5 and v6 versions and can be used with +//! Eagle 7B([blog post](https://blog.rwkv.com/p/eagle-7b-soaring-past-transformers)). +//! +//! Key characteristics: +//! - Time-mix attention mechanism +//! - Channel-mix feed-forward network +//! - Linear attention +//! - Group normalization +//! - Token shift mechanism +//! +//! References: +//! - [RWKV Language Model](https://github.com/BlinkDL/RWKV-LM) +//! - [RWKV v5 Release](https://github.com/BlinkDL/ChatRWKV/tree/main) +//! +//! # Example +//! +//! ```bash +//! cargo run --example rwkv --release -- \ +//! --prompt "The smallest prime is " +//! +//! > avx: true, neon: false, simd128: false, f16c: true +//! > temp: 0.00 repeat-penalty: 1.10 repeat-last-n: 64 +//! > The smallest prime is ϕ(2) = 2. +//! > The smallest composite is ϕ(3) = 3. +//! > The smallest perfect number is ϕ(5) = 5. +//! > The smallest perfect square is ϕ(4) = 4. +//! > The smallest perfect cube is ϕ(6) = 6. +//! ``` + +use super::with_tracing::{layer_norm, linear_no_bias as linear, LayerNorm, Linear}; +use candle::{DType, Device, IndexOp, Result, Tensor}; +use candle_nn::{embedding, Embedding, Module, VarBuilder}; +use std::collections::{HashMap, HashSet}; + +fn default_num_attention_heads() -> usize { + 64 +} + +// https://huggingface.co/RWKV/HF_v5-Eagle-7B/blob/main/configuration_rwkv5.py +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub num_hidden_layers: usize, + pub attention_hidden_size: usize, + #[serde(default = "default_num_attention_heads")] + pub num_attention_heads: usize, + pub head_size: usize, + pub intermediate_size: Option, + pub layer_norm_epsilon: f64, + pub rescale_every: usize, +} + +pub struct StatePerLayer { + pub extract_key_value: Tensor, + pub linear_attention: Tensor, + pub feed_forward: Tensor, +} + +pub struct State { + pub per_layer: Vec, + pub pos: usize, +} + +impl State { + pub fn new(batch_size: usize, cfg: &Config, dev: &Device) -> Result { + let mut per_layer = Vec::with_capacity(cfg.num_hidden_layers); + // Certainly a weird convention but taken from modeling_rwkv5.py + let num_attention_heads = cfg.hidden_size / cfg.num_attention_heads; + for _layer_idx in 0..cfg.num_hidden_layers { + let extract_key_value = Tensor::zeros((batch_size, cfg.hidden_size), DType::F32, dev)?; + let linear_attention = Tensor::zeros( + ( + batch_size, + num_attention_heads, + cfg.hidden_size / num_attention_heads, + cfg.hidden_size / num_attention_heads, + ), + DType::F32, + dev, + )?; + let feed_forward = Tensor::zeros((batch_size, cfg.hidden_size), DType::F32, dev)?; + per_layer.push(StatePerLayer { + extract_key_value, + linear_attention, + feed_forward, + }); + } + Ok(Self { per_layer, pos: 0 }) + } +} + +#[derive(Debug, Clone)] +struct SelfAttention { + key: Linear, + receptance: Linear, + value: Linear, + gate: Linear, + output: Linear, + ln_x: candle_nn::GroupNorm, + time_mix_key: Tensor, + time_mix_value: Tensor, + time_mix_receptance: Tensor, + time_decay: Tensor, + time_faaaa: Tensor, + time_mix_gate: Tensor, + layer_id: usize, + n_attn_heads: usize, +} + +impl SelfAttention { + pub fn new(layer_id: usize, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_size = cfg.hidden_size; + let attn_hidden_size = cfg.attention_hidden_size; + let key = linear(hidden_size, attn_hidden_size, vb.pp("key"))?; + let receptance = linear(hidden_size, attn_hidden_size, vb.pp("receptance"))?; + let value = linear(hidden_size, attn_hidden_size, vb.pp("value"))?; + let gate = linear(hidden_size, attn_hidden_size, vb.pp("gate"))?; + let output = linear(attn_hidden_size, hidden_size, vb.pp("output"))?; + let ln_x = candle_nn::group_norm( + hidden_size / cfg.head_size, + hidden_size, + 1e-5, + vb.pp("ln_x"), + )?; + let time_mix_key = vb.get((1, 1, cfg.hidden_size), "time_mix_key")?; + let time_mix_value = vb.get((1, 1, cfg.hidden_size), "time_mix_value")?; + let time_mix_receptance = vb.get((1, 1, cfg.hidden_size), "time_mix_receptance")?; + let n_attn_heads = cfg.hidden_size / cfg.head_size; + let time_decay = vb.get((n_attn_heads, cfg.head_size), "time_decay")?; + let time_faaaa = vb.get((n_attn_heads, cfg.head_size), "time_faaaa")?; + let time_mix_gate = vb.get((1, 1, cfg.hidden_size), "time_mix_gate")?; + Ok(Self { + key, + value, + receptance, + gate, + output, + ln_x, + time_mix_key, + time_mix_value, + time_mix_receptance, + time_decay, + time_faaaa, + time_mix_gate, + layer_id, + n_attn_heads, + }) + } + + pub fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let h = self.time_decay.dim(0)?; + let (b, t, s) = xs.dims3()?; + let s = s / h; + let (receptance, key, value, gate) = { + // extract key-value + let shifted = state.per_layer[self.layer_id].extract_key_value.clone(); + let shifted = if shifted.rank() == 2 { + shifted.unsqueeze(1)? + } else { + shifted + }; + let key = ((xs * &self.time_mix_key)? + &shifted * (1.0 - &self.time_mix_key)?)?; + let value = ((xs * &self.time_mix_value)? + &shifted * (1.0 - &self.time_mix_value)?)?; + let receptance = ((xs * &self.time_mix_receptance)? + + &shifted * (1.0 - &self.time_mix_receptance)?)?; + let gate = ((xs * &self.time_mix_gate)? + &shifted * (1.0 - &self.time_mix_gate)?)?; + + let key = self.key.forward(&key)?; + let value = self.value.forward(&value)?; + let receptance = self.receptance.forward(&receptance)?; + let gate = candle_nn::ops::silu(&self.gate.forward(&gate)?)?; + state.per_layer[self.layer_id].extract_key_value = xs.i((.., t - 1))?; + (receptance, key, value, gate) + }; + // linear attention + let mut state_ = state.per_layer[self.layer_id].linear_attention.clone(); + let key = key.reshape((b, t, h, s))?.permute((0, 2, 3, 1))?; + let value = value.reshape((b, t, h, s))?.transpose(1, 2)?; + let receptance = receptance.reshape((b, t, h, s))?.transpose(1, 2)?; + + let time_decay = self + .time_decay + .exp()? + .neg()? + .exp()? + .reshape(((), 1, 1))? + .reshape((self.n_attn_heads, (), 1))?; + let time_faaaa = + self.time_faaaa + .reshape(((), 1, 1))? + .reshape((self.n_attn_heads, (), 1))?; + + let mut out: Vec = Vec::with_capacity(t); + for t_ in 0..t { + let rt = receptance.i((.., .., t_..t_ + 1))?.contiguous()?; + let kt = key.i((.., .., .., t_..t_ + 1))?.contiguous()?; + let vt = value.i((.., .., t_..t_ + 1))?.contiguous()?; + let at = kt.matmul(&vt)?; + let rhs = (time_faaaa.broadcast_mul(&at)? + &state_)?; + let out_ = rt.matmul(&rhs)?.squeeze(2)?; + state_ = (&at + time_decay.broadcast_mul(&state_))?; + out.push(out_) + } + let out = Tensor::cat(&out, 1)?.reshape((b * t, h * s, 1))?; + let out = out.apply(&self.ln_x)?.reshape((b, t, h * s))?; + let out = (out * gate)?.apply(&self.output)?; + state.per_layer[self.layer_id].linear_attention = state_; + Ok(out) + } +} + +#[derive(Debug, Clone)] +struct FeedForward { + time_mix_key: Tensor, + time_mix_receptance: Tensor, + key: Linear, + receptance: Linear, + value: Linear, + layer_id: usize, +} + +impl FeedForward { + pub fn new(layer_id: usize, cfg: &Config, vb: VarBuilder) -> Result { + let int_size = cfg + .intermediate_size + .unwrap_or(((cfg.hidden_size as f64 * 3.5) as usize) / 32 * 32); + let key = linear(cfg.hidden_size, int_size, vb.pp("key"))?; + let receptance = linear(cfg.hidden_size, cfg.hidden_size, vb.pp("receptance"))?; + let value = linear(int_size, cfg.hidden_size, vb.pp("value"))?; + let time_mix_key = vb.get((1, 1, cfg.hidden_size), "time_mix_key")?; + let time_mix_receptance = vb.get((1, 1, cfg.hidden_size), "time_mix_receptance")?; + Ok(Self { + key, + receptance, + value, + time_mix_key, + time_mix_receptance, + layer_id, + }) + } + + pub fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let shifted = &state.per_layer[self.layer_id].feed_forward; + let key = (xs.broadcast_mul(&self.time_mix_key)? + + shifted.broadcast_mul(&(1.0 - &self.time_mix_key)?)?)?; + let receptance = (xs.broadcast_mul(&self.time_mix_receptance)? + + shifted.broadcast_mul(&(1.0 - &self.time_mix_receptance)?)?)?; + let key = key.apply(&self.key)?.relu()?.sqr()?; + let value = key.apply(&self.value)?; + let receptance = candle_nn::ops::sigmoid(&receptance.apply(&self.receptance)?)?; + state.per_layer[self.layer_id].feed_forward = xs.i((.., xs.dim(1)? - 1))?; + let xs = (receptance * value)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct Block { + pre_ln: Option, + ln1: LayerNorm, + ln2: LayerNorm, + attention: SelfAttention, + feed_forward: FeedForward, +} + +impl Block { + pub fn new(layer_id: usize, cfg: &Config, vb: VarBuilder) -> Result { + let ln1 = layer_norm(cfg.hidden_size, cfg.layer_norm_epsilon, vb.pp("ln1"))?; + let ln2 = layer_norm(cfg.hidden_size, cfg.layer_norm_epsilon, vb.pp("ln2"))?; + let pre_ln = if layer_id == 0 { + let ln = layer_norm(cfg.hidden_size, cfg.layer_norm_epsilon, vb.pp("pre_ln"))?; + Some(ln) + } else { + None + }; + let attention = SelfAttention::new(layer_id, cfg, vb.pp("attention"))?; + let feed_forward = FeedForward::new(layer_id, cfg, vb.pp("feed_forward"))?; + Ok(Self { + pre_ln, + ln1, + ln2, + attention, + feed_forward, + }) + } + + pub fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let xs = match self.pre_ln.as_ref() { + None => xs.clone(), + Some(pre_ln) => xs.apply(pre_ln)?, + }; + let attention = self.attention.forward(&xs.apply(&self.ln1)?, state)?; + let xs = (xs + attention)?; + let feed_forward = self.feed_forward.forward(&xs.apply(&self.ln2)?, state)?; + let xs = (xs + feed_forward)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embeddings: Embedding, + blocks: Vec, + ln_out: LayerNorm, + head: Linear, + rescale_every: usize, + layers_are_rescaled: bool, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("rwkv"); + let embeddings = embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embeddings"))?; + let mut blocks = Vec::with_capacity(cfg.num_hidden_layers); + let vb_b = vb_m.pp("blocks"); + for block_index in 0..cfg.num_hidden_layers { + let block = Block::new(block_index, cfg, vb_b.pp(block_index))?; + blocks.push(block) + } + let ln_out = layer_norm(cfg.hidden_size, 1e-5, vb_m.pp("ln_out"))?; + let head = linear(cfg.hidden_size, cfg.vocab_size, vb.pp("head"))?; + Ok(Self { + embeddings, + blocks, + ln_out, + head, + rescale_every: cfg.rescale_every, + layers_are_rescaled: false, // This seem to only happen for the f16/bf16 dtypes. + }) + } + + pub fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let (_b_size, _seq_len) = xs.dims2()?; + let mut xs = xs.apply(&self.embeddings)?; + for (block_idx, block) in self.blocks.iter().enumerate() { + xs = block.forward(&xs, state)?; + if self.layers_are_rescaled && (block_idx + 1) % self.rescale_every == 0 { + xs = (xs / 2.)? + } + } + let xs = xs.apply(&self.ln_out)?.apply(&self.head)?; + state.pos += 1; + Ok(xs) + } +} + +type Bytes = Vec; + +// https://github.com/BlinkDL/ChatRWKV/blob/095e812aef15a1f74107f6c39d13578a2412dc46/RWKV_v5_demo.py#L14 +pub struct Tokenizer { + table: Vec>>, + good: Vec>, + idx2token: HashMap>, + token2idx: HashMap, u32>, +} + +impl Tokenizer { + pub fn new>(p: P) -> Result { + let file = std::fs::File::open(p)?; + let token2idx: HashMap = + serde_json::from_reader(file).map_err(candle::Error::wrap)?; + let token2idx = token2idx + .into_iter() + .map(|(key, value)| (key.into_bytes(), value)) + .collect::>(); + let idx2token = token2idx + .iter() + .map(|(key, value)| (*value, key.to_vec())) + .collect::>(); + + let max_idx = token2idx.values().copied().max().unwrap_or(0); + + let mut table = vec![vec![vec![]; 256]; 256]; + let mut good = vec![HashSet::new(); 256]; + for idx in (0..(1 + max_idx)).rev() { + let s = match idx2token.get(&idx) { + None => continue, + Some(s) => s, + }; + if s.len() >= 2 { + let (s0, s1) = (s[0], s[1]); + table[s0 as usize][s1 as usize].push(s.to_vec()); + good[s0 as usize].insert(s1); + } + } + Ok(Self { + table, + good, + idx2token, + token2idx, + }) + } + + pub fn decode_bytes(&self, tokens: &[u32]) -> Vec { + let mut v = Vec::new(); + for token_id in tokens.iter() { + if let Some(token) = self.idx2token.get(token_id) { + v.extend_from_slice(token.as_slice()) + } + } + v + } + + pub fn decode(&self, tokens: &[u32]) -> Result { + let bytes = self.decode_bytes(tokens); + String::from_utf8(bytes).map_err(candle::Error::wrap) + } + + pub fn encode_bytes(&self, bytes: &[u8]) -> Result> { + let mut tokens = Vec::new(); + let mut i = 0; + while i < bytes.len() { + let mut s = vec![bytes[i]]; + if i + 1 < bytes.len() && self.good[bytes[i] as usize].contains(&bytes[i + 1]) { + let table = &self.table[bytes[i] as usize][bytes[i + 1] as usize]; + for table_elem in table.iter() { + if bytes[i..].starts_with(table_elem) { + s = table_elem.to_vec(); + break; + } + } + } + i += s.len(); + let token = match self.token2idx.get(&s) { + None => candle::bail!("unexpected token '{}' {s:?}", String::from_utf8_lossy(&s)), + Some(token) => *token, + }; + tokens.push(token) + } + Ok(tokens) + } + + pub fn encode(&self, str: &str) -> Result> { + self.encode_bytes(str.as_bytes()) + } +} diff --git a/patches/candle-transformers/src/models/rwkv_v6.rs b/patches/candle-transformers/src/models/rwkv_v6.rs new file mode 100644 index 0000000000..5da1c5ce81 --- /dev/null +++ b/patches/candle-transformers/src/models/rwkv_v6.rs @@ -0,0 +1,324 @@ +//! RWKV v6 model implementation. +//! +//! The [RWKV model](https://wiki.rwkv.com/) is a recurrent neural network model +//! with performance on par with transformer architectures. Several variants are +//! available, candle implements the v5 and v6 versions and can be used with +//! Eagle 7B([blog post](https://blog.rwkv.com/p/eagle-7b-soaring-past-transformers)). +//! +//! Key characteristics: +//! - Linear attention mechanism +//! - Time-mixing for temporal dependencies +//! - Group normalization +//! - Feed forward gating +//! - State recycling for efficient inference +//! +//! # Example +//! +//! ```bash +//! cargo run --example rwkv --release -- \ +//! --prompt "The smallest prime is " +//! +//! > avx: true, neon: false, simd128: false, f16c: true +//! > temp: 0.00 repeat-penalty: 1.10 repeat-last-n: 64 +//! > The smallest prime is ϕ(2) = 2. +//! > The smallest composite is ϕ(3) = 3. +//! > The smallest perfect number is ϕ(5) = 5. +//! > The smallest perfect square is ϕ(4) = 4. +//! > The smallest perfect cube is ϕ(6) = 6. +//! ``` + +use super::with_tracing::{layer_norm, linear_no_bias as linear, LayerNorm, Linear}; +use candle::{IndexOp, Result, Tensor}; +use candle_nn::{embedding, Embedding, Module, VarBuilder}; + +pub use crate::models::rwkv_v5::{Config, State, Tokenizer}; + +#[derive(Debug, Clone)] +struct SelfAttention { + key: Linear, + receptance: Linear, + value: Linear, + gate: Linear, + output: Linear, + ln_x: candle_nn::GroupNorm, + time_mix_x: Tensor, + time_mix_w: Tensor, + time_mix_key: Tensor, + time_mix_value: Tensor, + time_mix_receptance: Tensor, + time_decay: Tensor, + time_faaaa: Tensor, + time_mix_gate: Tensor, + time_decay_w1: Tensor, + time_decay_w2: Tensor, + time_mix_w1: Tensor, + time_mix_w2: Tensor, + layer_id: usize, + n_attn_heads: usize, +} + +impl SelfAttention { + fn new(layer_id: usize, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_size = cfg.hidden_size; + let attn_hidden_size = cfg.attention_hidden_size; + let key = linear(hidden_size, attn_hidden_size, vb.pp("key"))?; + let receptance = linear(hidden_size, attn_hidden_size, vb.pp("receptance"))?; + let value = linear(hidden_size, attn_hidden_size, vb.pp("value"))?; + let gate = linear(hidden_size, attn_hidden_size, vb.pp("gate"))?; + let output = linear(attn_hidden_size, hidden_size, vb.pp("output"))?; + let ln_x = candle_nn::group_norm( + hidden_size / cfg.head_size, + hidden_size, + 1e-5, + vb.pp("ln_x"), + )?; + + let time_mix_x = vb.get((1, 1, cfg.hidden_size), "time_mix_x")?; + let time_mix_w = vb.get((1, 1, cfg.hidden_size), "time_mix_w")?; + let time_mix_key = vb.get((1, 1, cfg.hidden_size), "time_mix_key")?; + let time_mix_value = vb.get((1, 1, cfg.hidden_size), "time_mix_value")?; + let time_mix_receptance = vb.get((1, 1, cfg.hidden_size), "time_mix_receptance")?; + let n_attn_heads = cfg.hidden_size / cfg.head_size; + let time_decay = vb.get((1, 1, cfg.hidden_size), "time_decay")?; + let time_faaaa = vb.get((n_attn_heads, cfg.head_size), "time_faaaa")?; + let time_mix_gate = vb.get((1, 1, cfg.hidden_size), "time_mix_gate")?; + let time_decay_w1 = vb.get((cfg.hidden_size, n_attn_heads * 2), "time_decay_w1")?; + let time_decay_w2 = vb.get((n_attn_heads * 2, cfg.hidden_size), "time_decay_w2")?; + let time_mix_w1 = vb.get((cfg.hidden_size, n_attn_heads * 5), "time_mix_w1")?; + let time_mix_w2 = vb.get((5, n_attn_heads, cfg.hidden_size), "time_mix_w2")?; + Ok(Self { + key, + value, + receptance, + gate, + output, + ln_x, + time_mix_x, + time_mix_w, + time_mix_key, + time_mix_value, + time_mix_receptance, + time_decay, + time_faaaa, + time_mix_gate, + time_decay_w1, + time_decay_w2, + time_mix_w1, + time_mix_w2, + layer_id, + n_attn_heads, + }) + } + + pub fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let h = self.n_attn_heads; + let (b, t, s) = xs.dims3()?; + let s = s / h; + let (receptance, key, value, gate, w) = { + // extract key-value + let shifted = state.per_layer[self.layer_id].extract_key_value.clone(); + let shifted = if shifted.rank() == 2 { + shifted.unsqueeze(1)? + } else { + shifted + }; + + let sx = (&shifted - xs)?; + let xxx = (xs + &sx * &self.time_mix_x)?; + let xxx = xxx + .broadcast_matmul(&self.time_mix_w1)? + .tanh()? + .reshape((b * t, 5, ()))? + .transpose(0, 1)?; + + let xxx = xxx.matmul(&self.time_mix_w2)?.reshape((5, b, t, ()))?; + + let (mw, mk, mv, mr, mg) = (xxx.i(0)?, xxx.i(1)?, xxx.i(2)?, xxx.i(3)?, xxx.i(4)?); + + let xw = (xs + &sx * (&self.time_mix_w + &mw)?)?; + let xk = (xs + &sx * (&self.time_mix_key + &mk)?)?; + let xv = (xs + &sx * (&self.time_mix_value + &mv)?)?; + let xr = (xs + &sx * (&self.time_mix_receptance + &mr)?)?; + let xg = (xs + &sx * (&self.time_mix_gate + &mg)?)?; + + let w = (&self.time_decay + + xw.broadcast_matmul(&self.time_decay_w1)? + .tanh()? + .broadcast_matmul(&self.time_decay_w2)?)? + .reshape(((), 1, 1))? + .reshape((self.n_attn_heads, (), 1))?; + + let key = self.key.forward(&xk)?; + let value = self.value.forward(&xv)?; + let receptance = self.receptance.forward(&xr)?; + let gate = candle_nn::ops::silu(&self.gate.forward(&xg)?)?; + state.per_layer[self.layer_id].extract_key_value = xs.i((.., t - 1))?; + (receptance, key, value, gate, w) + }; + + // linear attention + let mut state_ = state.per_layer[self.layer_id].linear_attention.clone(); + let key = key.reshape((b, t, h, s))?.permute((0, 2, 3, 1))?; + let value = value.reshape((b, t, h, s))?.transpose(1, 2)?; + let receptance = receptance.reshape((b, t, h, s))?.transpose(1, 2)?; + + let w = w.exp()?.neg()?.exp()?; + + let time_faaaa = + self.time_faaaa + .reshape(((), 1, 1))? + .reshape((self.n_attn_heads, (), 1))?; + + let mut out: Vec = Vec::with_capacity(t); + for t_ in 0..t { + let rt = receptance.i((.., .., t_..t_ + 1))?.contiguous()?; + let kt = key.i((.., .., .., t_..t_ + 1))?.contiguous()?; + let vt = value.i((.., .., t_..t_ + 1))?.contiguous()?; + let at = kt.matmul(&vt)?; + let rhs = (time_faaaa.broadcast_mul(&at)? + &state_)?; + let out_ = rt.matmul(&rhs)?.squeeze(2)?; + state_ = (&at + w.broadcast_mul(&state_))?; + out.push(out_) + } + let out = Tensor::cat(&out, 1)?.reshape((b * t, h * s, 1))?; + let out = out.apply(&self.ln_x)?.reshape((b, t, h * s))?; + let out = (out * gate)?.apply(&self.output)?; + state.per_layer[self.layer_id].linear_attention = state_; + Ok(out) + } +} + +#[derive(Debug, Clone)] +struct FeedForward { + time_mix_key: Tensor, + time_mix_receptance: Tensor, + key: Linear, + receptance: Linear, + value: Linear, + layer_id: usize, +} + +impl FeedForward { + fn new(layer_id: usize, cfg: &Config, vb: VarBuilder) -> Result { + let int_size = cfg + .intermediate_size + .unwrap_or(((cfg.hidden_size as f64 * 3.5) as usize) / 32 * 32); + let key = linear(cfg.hidden_size, int_size, vb.pp("key"))?; + let receptance = linear(cfg.hidden_size, cfg.hidden_size, vb.pp("receptance"))?; + let value = linear(int_size, cfg.hidden_size, vb.pp("value"))?; + let time_mix_key = vb.get((1, 1, cfg.hidden_size), "time_mix_key")?; + let time_mix_receptance = vb.get((1, 1, cfg.hidden_size), "time_mix_receptance")?; + Ok(Self { + key, + receptance, + value, + time_mix_key, + time_mix_receptance, + layer_id, + }) + } + + fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let shifted = state.per_layer[self.layer_id] + .feed_forward + .broadcast_sub(xs)?; + let key = (xs + shifted.broadcast_mul(&self.time_mix_key)?)?; + let receptance = (xs + shifted.broadcast_mul(&self.time_mix_receptance)?)?; + let key = key.apply(&self.key)?.relu()?.sqr()?; + let value = key.apply(&self.value)?; + let receptance = candle_nn::ops::sigmoid(&receptance.apply(&self.receptance)?)?; + state.per_layer[self.layer_id].feed_forward = xs.i((.., xs.dim(1)? - 1))?; + let xs = (receptance * value)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct Block { + pre_ln: Option, + ln1: LayerNorm, + ln2: LayerNorm, + attention: SelfAttention, + feed_forward: FeedForward, +} + +impl Block { + fn new(layer_id: usize, cfg: &Config, vb: VarBuilder) -> Result { + let ln1 = layer_norm(cfg.hidden_size, cfg.layer_norm_epsilon, vb.pp("ln1"))?; + let ln2 = layer_norm(cfg.hidden_size, cfg.layer_norm_epsilon, vb.pp("ln2"))?; + let pre_ln = if layer_id == 0 { + let ln = layer_norm(cfg.hidden_size, cfg.layer_norm_epsilon, vb.pp("pre_ln"))?; + Some(ln) + } else { + None + }; + let attention = SelfAttention::new(layer_id, cfg, vb.pp("attention"))?; + let feed_forward = FeedForward::new(layer_id, cfg, vb.pp("feed_forward"))?; + Ok(Self { + pre_ln, + ln1, + ln2, + attention, + feed_forward, + }) + } + + fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let xs = match self.pre_ln.as_ref() { + None => xs.clone(), + Some(pre_ln) => xs.apply(pre_ln)?, + }; + let attention = self.attention.forward(&xs.apply(&self.ln1)?, state)?; + let xs = (xs + attention)?; + let feed_forward = self.feed_forward.forward(&xs.apply(&self.ln2)?, state)?; + let xs = (xs + feed_forward)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embeddings: Embedding, + blocks: Vec, + ln_out: LayerNorm, + head: Linear, + rescale_every: usize, + layers_are_rescaled: bool, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("rwkv"); + let embeddings = embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embeddings"))?; + let mut blocks = Vec::with_capacity(cfg.num_hidden_layers); + let vb_b = vb_m.pp("blocks"); + for block_index in 0..cfg.num_hidden_layers { + let block = Block::new(block_index, cfg, vb_b.pp(block_index))?; + blocks.push(block) + } + let ln_out = layer_norm(cfg.hidden_size, 1e-5, vb_m.pp("ln_out"))?; + let head = linear(cfg.hidden_size, cfg.vocab_size, vb.pp("head"))?; + Ok(Self { + embeddings, + blocks, + ln_out, + head, + rescale_every: cfg.rescale_every, + layers_are_rescaled: false, // This seem to only happen for the f16/bf16 dtypes. + }) + } + + pub fn forward(&self, xs: &Tensor, state: &mut State) -> Result { + let (_b_size, _seq_len) = xs.dims2()?; + let mut xs = xs.apply(&self.embeddings)?; + for (block_idx, block) in self.blocks.iter().enumerate() { + xs = block.forward(&xs, state)?; + if self.layers_are_rescaled && (block_idx + 1) % self.rescale_every == 0 { + xs = (xs / 2.)? + } + } + let xs = xs.apply(&self.ln_out)?.apply(&self.head)?; + state.pos += 1; + Ok(xs) + } +} diff --git a/patches/candle-transformers/src/models/segformer.rs b/patches/candle-transformers/src/models/segformer.rs new file mode 100644 index 0000000000..bf72e7c690 --- /dev/null +++ b/patches/candle-transformers/src/models/segformer.rs @@ -0,0 +1,721 @@ +//! Segformer model implementation for semantic segmentation and image classification. +//! +//! Segformer is a transformer-based model designed for vision tasks. It uses a hierarchical +//! structure that progressively generates features at different scales. +//! +//! Key characteristics: +//! - Efficient self-attention with sequence reduction +//! - Hierarchical feature generation +//! - Mix-FFN for local and global feature interaction +//! - Lightweight all-MLP decode head +//! +//! References: +//! - [SegFormer Paper](https://arxiv.org/abs/2105.15203) +//! - [Model Card](https://huggingface.co/nvidia/mit-b0) +//! + +use crate::models::with_tracing::{conv2d, linear, Conv2d, Linear}; +use candle::{Context, Module, ModuleT, Result, Tensor, D}; +use candle_nn::{conv2d_no_bias, layer_norm, Activation, Conv2dConfig, VarBuilder}; +use serde::Deserialize; +use std::collections::HashMap; + +// https://github.com/huggingface/transformers/blob/main/src/transformers/models/segformer/configuration_segformer.py +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Config { + #[serde(default)] + pub id2label: HashMap, + pub num_channels: usize, + pub num_encoder_blocks: usize, + pub depths: Vec, + pub sr_ratios: Vec, + pub hidden_sizes: Vec, + pub patch_sizes: Vec, + pub strides: Vec, + pub num_attention_heads: Vec, + pub mlp_ratios: Vec, + pub hidden_act: candle_nn::Activation, + pub layer_norm_eps: f64, + pub decoder_hidden_size: usize, +} + +#[derive(Debug, Clone)] +struct SegformerOverlapPatchEmbeddings { + projection: Conv2d, + layer_norm: candle_nn::LayerNorm, +} + +impl SegformerOverlapPatchEmbeddings { + fn new( + config: &Config, + patch_size: usize, + stride: usize, + num_channels: usize, + hidden_size: usize, + vb: VarBuilder, + ) -> Result { + let projection = conv2d( + num_channels, + hidden_size, + patch_size, + Conv2dConfig { + stride, + padding: patch_size / 2, + ..Default::default() + }, + vb.pp("proj"), + )?; + let layer_norm = + candle_nn::layer_norm(hidden_size, config.layer_norm_eps, vb.pp("layer_norm"))?; + Ok(Self { + projection, + layer_norm, + }) + } +} + +impl Module for SegformerOverlapPatchEmbeddings { + fn forward(&self, x: &Tensor) -> Result { + let embeddings = self.projection.forward(x)?; + let shape = embeddings.shape(); + // [B, C, H, W] -> [B, H * W, C] + let embeddings = embeddings.flatten_from(2)?.transpose(1, 2)?; + let embeddings = self.layer_norm.forward(&embeddings)?; + // [B, H * W, C] -> [B, C, H, W] + let embeddings = embeddings.transpose(1, 2)?.reshape(shape)?; + Ok(embeddings) + } +} + +#[derive(Debug, Clone)] +struct SegformerEfficientSelfAttention { + num_attention_heads: usize, + attention_head_size: usize, + query: Linear, + key: Linear, + value: Linear, + sr: Option, + layer_norm: Option, +} + +impl SegformerEfficientSelfAttention { + fn new( + config: &Config, + hidden_size: usize, + num_attention_heads: usize, + sequence_reduction_ratio: usize, + vb: VarBuilder, + ) -> Result { + if !hidden_size.is_multiple_of(num_attention_heads) { + candle::bail!( + "The hidden size {} is not a multiple of the number of attention heads {}", + hidden_size, + num_attention_heads + ) + } + let attention_head_size = hidden_size / num_attention_heads; + let all_head_size = num_attention_heads * attention_head_size; + let query = linear(hidden_size, all_head_size, vb.pp("query"))?; + let key = linear(hidden_size, all_head_size, vb.pp("key"))?; + let value = linear(hidden_size, all_head_size, vb.pp("value"))?; + let (sr, layer_norm) = if sequence_reduction_ratio > 1 { + ( + Some(conv2d( + hidden_size, + hidden_size, + sequence_reduction_ratio, + Conv2dConfig { + stride: sequence_reduction_ratio, + ..Default::default() + }, + vb.pp("sr"), + )?), + Some(candle_nn::layer_norm( + hidden_size, + config.layer_norm_eps, + vb.pp("layer_norm"), + )?), + ) + } else { + (None, None) + }; + Ok(Self { + num_attention_heads, + attention_head_size, + query, + key, + value, + sr, + layer_norm, + }) + } + + fn transpose_for_scores(&self, hidden_states: Tensor) -> Result { + let (batch, seq_length, _) = hidden_states.shape().dims3()?; + let new_shape = &[ + batch, + seq_length, + self.num_attention_heads, + self.attention_head_size, + ]; + let hidden_states = hidden_states.reshape(new_shape)?; + let hidden_states = hidden_states.permute((0, 2, 1, 3))?; + Ok(hidden_states) + } +} + +impl Module for SegformerEfficientSelfAttention { + fn forward(&self, x: &Tensor) -> Result { + // [B, C, H, W] -> [B, H * W, C] + let hidden_states = x.flatten_from(2)?.permute((0, 2, 1))?; + let query = self + .transpose_for_scores(self.query.forward(&hidden_states)?)? + .contiguous()?; + let hidden_states = if let (Some(sr), Some(layer_norm)) = (&self.sr, &self.layer_norm) { + let hidden_states = sr.forward(x)?; + // [B, C, H, W] -> [B, H * W, C] + let hidden_states = hidden_states.flatten_from(2)?.permute((0, 2, 1))?; + layer_norm.forward(&hidden_states)? + } else { + // already [B, H * W, C] + hidden_states + }; + // standard self-attention + let key = self + .transpose_for_scores(self.key.forward(&hidden_states)?)? + .contiguous()?; + let value = self + .transpose_for_scores(self.value.forward(&hidden_states)?)? + .contiguous()?; + let attention_scores = + (query.matmul(&key.t()?)? / f64::sqrt(self.attention_head_size as f64))?; + let attention_scores = candle_nn::ops::softmax_last_dim(&attention_scores)?; + let result = attention_scores.matmul(&value)?; + let result = result.permute((0, 2, 1, 3))?.contiguous()?; + result.flatten_from(D::Minus2) + } +} + +#[derive(Debug, Clone)] +struct SegformerSelfOutput { + dense: Linear, +} + +impl SegformerSelfOutput { + fn new(hidden_size: usize, vb: VarBuilder) -> Result { + let dense = linear(hidden_size, hidden_size, vb.pp("dense"))?; + Ok(Self { dense }) + } +} + +impl Module for SegformerSelfOutput { + fn forward(&self, x: &Tensor) -> Result { + self.dense.forward(x) + } +} + +#[derive(Debug, Clone)] +struct SegformerAttention { + attention: SegformerEfficientSelfAttention, + output: SegformerSelfOutput, +} + +impl SegformerAttention { + fn new( + config: &Config, + hidden_size: usize, + num_attention_heads: usize, + sequence_reduction_ratio: usize, + vb: VarBuilder, + ) -> Result { + let attention = SegformerEfficientSelfAttention::new( + config, + hidden_size, + num_attention_heads, + sequence_reduction_ratio, + vb.pp("self"), + )?; + let output = SegformerSelfOutput::new(hidden_size, vb.pp("output"))?; + Ok(Self { attention, output }) + } +} + +impl Module for SegformerAttention { + fn forward(&self, x: &Tensor) -> Result { + let attention_output = self.attention.forward(x)?; + self.output.forward(&attention_output) + } +} + +#[derive(Debug, Clone)] +struct SegformerDWConv { + dw_conv: Conv2d, +} + +impl SegformerDWConv { + fn new(dim: usize, vb: VarBuilder) -> Result { + let dw_conv = conv2d( + dim, + dim, + 3, + Conv2dConfig { + stride: 1, + padding: 1, + groups: dim, + ..Default::default() + }, + vb.pp("dwconv"), + )?; + Ok(Self { dw_conv }) + } +} + +impl Module for SegformerDWConv { + fn forward(&self, x: &Tensor) -> Result { + self.dw_conv.forward(x) + } +} + +#[derive(Debug, Clone)] +struct SegformerMixFFN { + dense1: Linear, + dw_conv: SegformerDWConv, + act: Activation, + dense2: Linear, +} + +impl SegformerMixFFN { + fn new( + config: &Config, + in_features: usize, + hidden_features: usize, + out_features: usize, + vb: VarBuilder, + ) -> Result { + let dense1 = linear(in_features, hidden_features, vb.pp("dense1"))?; + let dw_conv = SegformerDWConv::new(hidden_features, vb.pp("dwconv"))?; + let act = config.hidden_act; + let dense2 = linear(hidden_features, out_features, vb.pp("dense2"))?; + Ok(Self { + dense1, + dw_conv, + act, + dense2, + }) + } +} + +impl Module for SegformerMixFFN { + fn forward(&self, x: &Tensor) -> Result { + let (batch, _, height, width) = x.shape().dims4()?; + let hidden_states = self + .dense1 + .forward(&x.flatten_from(2)?.permute((0, 2, 1))?)?; + let channels = hidden_states.dim(2)?; + let hidden_states = self.dw_conv.forward( + &hidden_states + .permute((0, 2, 1))? + .reshape((batch, channels, height, width))?, + )?; + let hidden_states = self.act.forward(&hidden_states)?; + let hidden_states = self + .dense2 + .forward(&hidden_states.flatten_from(2)?.permute((0, 2, 1))?)?; + let channels = hidden_states.dim(2)?; + hidden_states + .permute((0, 2, 1))? + .reshape((batch, channels, height, width)) + } +} + +#[derive(Debug, Clone)] +struct SegformerLayer { + layer_norm_1: candle_nn::LayerNorm, + attention: SegformerAttention, + layer_norm_2: candle_nn::LayerNorm, + mlp: SegformerMixFFN, +} + +impl SegformerLayer { + fn new( + config: &Config, + hidden_size: usize, + num_attention_heads: usize, + sequence_reduction_ratio: usize, + mlp_ratio: usize, + vb: VarBuilder, + ) -> Result { + let layer_norm_1 = layer_norm(hidden_size, config.layer_norm_eps, vb.pp("layer_norm_1"))?; + let attention = SegformerAttention::new( + config, + hidden_size, + num_attention_heads, + sequence_reduction_ratio, + vb.pp("attention"), + )?; + let layer_norm_2 = layer_norm(hidden_size, config.layer_norm_eps, vb.pp("layer_norm_2"))?; + let mlp = SegformerMixFFN::new( + config, + hidden_size, + hidden_size * mlp_ratio, + hidden_size, + vb.pp("mlp"), + )?; + Ok(Self { + layer_norm_1, + attention, + layer_norm_2, + mlp, + }) + } +} + +impl Module for SegformerLayer { + fn forward(&self, x: &Tensor) -> Result { + let shape = x.shape().dims4()?; + // [B, C, H, W] -> [B, H * W, C] + let hidden_states = x.flatten_from(2)?.permute((0, 2, 1))?; + let layer_norm_output = self.layer_norm_1.forward(&hidden_states)?; + let layer_norm_output = layer_norm_output.permute((0, 2, 1))?.reshape(shape)?; + // attention takes in [B, C, H, W] in order to properly do conv2d (and output [B, H * W, C]) + let attention_output = self.attention.forward(&layer_norm_output)?; + let hidden_states = (attention_output + hidden_states)?; + let layer_norm_output = self.layer_norm_2.forward(&hidden_states)?; + let mlp_output = self + .mlp + .forward(&layer_norm_output.permute((0, 2, 1))?.reshape(shape)?)?; + hidden_states.permute((0, 2, 1))?.reshape(shape)? + mlp_output + } +} + +#[derive(Debug, Clone)] +struct SegformerEncoder { + /// config file + config: Config, + /// a list of embeddings + patch_embeddings: Vec, + /// a list of attention blocks, each consisting of layers + blocks: Vec>, + /// a final list of layer norms + layer_norms: Vec, +} + +impl SegformerEncoder { + fn new(config: Config, vb: VarBuilder) -> Result { + let mut patch_embeddings = Vec::with_capacity(config.num_encoder_blocks); + let mut blocks = Vec::with_capacity(config.num_encoder_blocks); + let mut layer_norms = Vec::with_capacity(config.num_encoder_blocks); + for i in 0..config.num_encoder_blocks { + let patch_size = config.patch_sizes[i]; + let stride = config.strides[i]; + let hidden_size = config.hidden_sizes[i]; + let num_channels = if i == 0 { + config.num_channels + } else { + config.hidden_sizes[i - 1] + }; + patch_embeddings.push(SegformerOverlapPatchEmbeddings::new( + &config, + patch_size, + stride, + num_channels, + hidden_size, + vb.pp(format!("patch_embeddings.{i}")), + )?); + let mut layers = Vec::with_capacity(config.depths[i]); + for j in 0..config.depths[i] { + let sequence_reduction_ratio = config.sr_ratios[i]; + let num_attention_heads = config.num_attention_heads[i]; + let mlp_ratio = config.mlp_ratios[i]; + layers.push(SegformerLayer::new( + &config, + hidden_size, + num_attention_heads, + sequence_reduction_ratio, + mlp_ratio, + vb.pp(format!("block.{i}.{j}")), + )?); + } + blocks.push(layers); + layer_norms.push(layer_norm( + hidden_size, + config.layer_norm_eps, + vb.pp(format!("layer_norm.{i}")), + )?); + } + Ok(Self { + config, + patch_embeddings, + blocks, + layer_norms, + }) + } +} + +impl ModuleWithHiddenStates for SegformerEncoder { + fn forward(&self, x: &Tensor) -> Result> { + let mut all_hidden_states = Vec::with_capacity(self.config.num_encoder_blocks); + let mut hidden_states = x.clone(); + for i in 0..self.config.num_encoder_blocks { + hidden_states = self.patch_embeddings[i].forward(&hidden_states)?; + for layer in &self.blocks[i] { + hidden_states = layer.forward(&hidden_states)?; + } + let shape = hidden_states.shape().dims4()?; + hidden_states = + self.layer_norms[i].forward(&hidden_states.flatten_from(2)?.permute((0, 2, 1))?)?; + hidden_states = hidden_states.permute((0, 2, 1))?.reshape(shape)?; + all_hidden_states.push(hidden_states.clone()); + } + Ok(all_hidden_states) + } +} + +#[derive(Debug, Clone)] +struct SegformerModel { + encoder: SegformerEncoder, +} + +impl SegformerModel { + fn new(config: &Config, vb: VarBuilder) -> Result { + let encoder = SegformerEncoder::new(config.clone(), vb.pp("encoder"))?; + Ok(Self { encoder }) + } +} + +impl ModuleWithHiddenStates for SegformerModel { + fn forward(&self, x: &Tensor) -> Result> { + self.encoder.forward(x) + } +} + +#[derive(Debug, Clone)] +struct SegformerMLP { + proj: Linear, +} + +impl SegformerMLP { + fn new(config: &Config, input_dim: usize, vb: VarBuilder) -> Result { + let proj = linear(input_dim, config.decoder_hidden_size, vb.pp("proj"))?; + Ok(Self { proj }) + } +} + +impl Module for SegformerMLP { + fn forward(&self, x: &Tensor) -> Result { + self.proj.forward(x) + } +} + +#[derive(Debug, Clone)] +struct SegformerDecodeHead { + linear_c: Vec, + linear_fuse: candle_nn::Conv2d, + batch_norm: candle_nn::BatchNorm, + classifier: candle_nn::Conv2d, +} + +impl SegformerDecodeHead { + fn new(config: &Config, num_labels: usize, vb: VarBuilder) -> Result { + let mut linear_c = Vec::with_capacity(config.num_encoder_blocks); + for i in 0..config.num_encoder_blocks { + let hidden_size = config.hidden_sizes[i]; + linear_c.push(SegformerMLP::new( + config, + hidden_size, + vb.pp(format!("linear_c.{i}")), + )?); + } + let linear_fuse = conv2d_no_bias( + config.decoder_hidden_size * config.num_encoder_blocks, + config.decoder_hidden_size, + 1, + Conv2dConfig::default(), + vb.pp("linear_fuse"), + )?; + let batch_norm = candle_nn::batch_norm( + config.decoder_hidden_size, + config.layer_norm_eps, + vb.pp("batch_norm"), + )?; + let classifier = conv2d_no_bias( + config.decoder_hidden_size, + num_labels, + 1, + Conv2dConfig::default(), + vb.pp("classifier"), + )?; + Ok(Self { + linear_c, + linear_fuse, + batch_norm, + classifier, + }) + } + + fn forward(&self, encoder_hidden_states: &[Tensor]) -> Result { + if encoder_hidden_states.len() != self.linear_c.len() { + candle::bail!( + "The number of encoder hidden states {} is not equal to the number of linear layers {}", + encoder_hidden_states.len(), + self.linear_c.len() + ) + } + // most fine layer + let (_, _, upsample_height, upsample_width) = encoder_hidden_states[0].shape().dims4()?; + let mut hidden_states = Vec::with_capacity(self.linear_c.len()); + for (hidden_state, mlp) in encoder_hidden_states.iter().zip(&self.linear_c) { + let (batch, _, height, width) = hidden_state.shape().dims4()?; + let hidden_state = mlp.forward(&hidden_state.flatten_from(2)?.permute((0, 2, 1))?)?; + let hidden_state = hidden_state.permute((0, 2, 1))?.reshape(( + batch, + hidden_state.dim(2)?, + height, + width, + ))?; + let hidden_state = hidden_state.upsample_nearest2d(upsample_height, upsample_width)?; + hidden_states.push(hidden_state); + } + hidden_states.reverse(); + let hidden_states = Tensor::cat(&hidden_states, 1)?; + let hidden_states = self.linear_fuse.forward(&hidden_states)?; + let hidden_states = self.batch_norm.forward_t(&hidden_states, false)?; + let hidden_states = hidden_states.relu()?; + self.classifier.forward(&hidden_states) + } +} + +trait ModuleWithHiddenStates { + fn forward(&self, xs: &Tensor) -> Result>; +} + +#[derive(Debug, Clone)] +pub struct SemanticSegmentationModel { + segformer: SegformerModel, + decode_head: SegformerDecodeHead, +} + +impl SemanticSegmentationModel { + pub fn new(config: &Config, num_labels: usize, vb: VarBuilder) -> Result { + let segformer = SegformerModel::new(config, vb.pp("segformer"))?; + let decode_head = SegformerDecodeHead::new(config, num_labels, vb.pp("decode_head"))?; + Ok(Self { + segformer, + decode_head, + }) + } +} + +impl Module for SemanticSegmentationModel { + fn forward(&self, x: &Tensor) -> Result { + let hidden_states = self.segformer.forward(x)?; + self.decode_head.forward(&hidden_states) + } +} + +#[derive(Debug, Clone)] +pub struct ImageClassificationModel { + segformer: SegformerModel, + classifier: Linear, +} + +impl ImageClassificationModel { + pub fn new(config: &Config, num_labels: usize, vb: VarBuilder) -> Result { + let segformer = SegformerModel::new(config, vb.pp("segformer"))?; + let classifier = linear(config.decoder_hidden_size, num_labels, vb.pp("classifier"))?; + Ok(Self { + segformer, + classifier, + }) + } +} + +impl Module for ImageClassificationModel { + fn forward(&self, x: &Tensor) -> Result { + let all_hidden_states = self.segformer.forward(x)?; + let hidden_states = all_hidden_states.last().context("no last")?; + let hidden_states = hidden_states.flatten_from(2)?.permute((0, 2, 1))?; + let mean = hidden_states.mean(1)?; + self.classifier.forward(&mean) + } +} + +#[cfg(test)] +mod tests { + + use super::*; + + #[test] + fn test_config_json_load() { + let raw_json = r#"{ + "architectures": [ + "SegformerForImageClassification" + ], + "attention_probs_dropout_prob": 0.0, + "classifier_dropout_prob": 0.1, + "decoder_hidden_size": 256, + "depths": [ + 2, + 2, + 2, + 2 + ], + "downsampling_rates": [ + 1, + 4, + 8, + 16 + ], + "drop_path_rate": 0.1, + "hidden_act": "gelu", + "hidden_dropout_prob": 0.0, + "hidden_sizes": [ + 32, + 64, + 160, + 256 + ], + "image_size": 224, + "initializer_range": 0.02, + "layer_norm_eps": 1e-06, + "mlp_ratios": [ + 4, + 4, + 4, + 4 + ], + "model_type": "segformer", + "num_attention_heads": [ + 1, + 2, + 5, + 8 + ], + "num_channels": 3, + "num_encoder_blocks": 4, + "patch_sizes": [ + 7, + 3, + 3, + 3 + ], + "sr_ratios": [ + 8, + 4, + 2, + 1 + ], + "strides": [ + 4, + 2, + 2, + 2 + ], + "torch_dtype": "float32", + "transformers_version": "4.12.0.dev0" + }"#; + let config: Config = serde_json::from_str(raw_json).unwrap(); + assert_eq!(vec![4, 2, 2, 2], config.strides); + assert_eq!(1e-6, config.layer_norm_eps); + } +} diff --git a/patches/candle-transformers/src/models/segment_anything/image_encoder.rs b/patches/candle-transformers/src/models/segment_anything/image_encoder.rs new file mode 100644 index 0000000000..0b3138305f --- /dev/null +++ b/patches/candle-transformers/src/models/segment_anything/image_encoder.rs @@ -0,0 +1,483 @@ +use candle::{DType, IndexOp, Result, Tensor}; +use candle_nn::{layer_norm, LayerNorm, Module, VarBuilder}; + +#[derive(Debug)] +struct PatchEmbed { + proj: candle_nn::Conv2d, + span: tracing::Span, +} + +impl PatchEmbed { + fn new( + in_chans: usize, + embed_dim: usize, + k_size: usize, + stride: usize, + padding: usize, + vb: VarBuilder, + ) -> Result { + let cfg = candle_nn::Conv2dConfig { + stride, + padding, + ..Default::default() + }; + let proj = candle_nn::conv2d(in_chans, embed_dim, k_size, cfg, vb.pp("proj"))?; + let span = tracing::span!(tracing::Level::TRACE, "patch-embed"); + Ok(Self { proj, span }) + } +} + +impl Module for PatchEmbed { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + xs.apply(&self.proj)?.permute((0, 2, 3, 1)) + } +} + +// A custom op to make add_decomposed_rel_pos faster. Most of the time is spent on the final +// addition in the case where b = 12, q_h = q_w = 4096, k_h = k_w = 4096 +// (attn.reshape((b, q_h, q_w, k_h, k_w))? +// + rel_h.unsqueeze(4)?.broadcast_add(&rel_w.unsqueeze(3)?)?)? +// .reshape((b, q_h * q_w, k_h * k_w)) +// Ideally we would perform this operation in place but this is not supported in candle at the +// moment. We should also investigate using f16 rather than f32. +struct Add3(usize, usize, usize, usize, usize); +impl candle::CustomOp3 for Add3 { + fn name(&self) -> &'static str { + "add3" + } + + fn cpu_fwd( + &self, + s1: &candle::CpuStorage, + l1: &candle::Layout, + s2: &candle::CpuStorage, + l2: &candle::Layout, + s3: &candle::CpuStorage, + l3: &candle::Layout, + ) -> Result<(candle::CpuStorage, candle::Shape)> { + use rayon::prelude::*; + + let Add3(b, q_h, q_w, k_h, k_w) = *self; + let s1 = s1.as_slice::()?; + let s1 = match l1.contiguous_offsets() { + None => candle::bail!("input1 has to be contiguous"), + Some((o1, o2)) => &s1[o1..o2], + }; + let s2 = s2.as_slice::()?; + let s2 = match l2.contiguous_offsets() { + None => candle::bail!("input2 has to be contiguous"), + Some((o1, o2)) => &s2[o1..o2], + }; + let s3 = s3.as_slice::()?; + let s3 = match l3.contiguous_offsets() { + None => candle::bail!("input3 has to be contiguous"), + Some((o1, o2)) => &s3[o1..o2], + }; + let mut dst = vec![0f32; b * q_h * q_w * k_h * k_w]; + dst.par_chunks_exact_mut(k_h * k_w) + .enumerate() + .for_each(|(b_idx, dst)| { + let s1_idx = b_idx * k_h * k_w; + let s2_idx = b_idx * k_h; + let s3_idx = b_idx * k_w; + for h_idx in 0..k_h { + let s1_idx = s1_idx + h_idx * k_w; + let s2_idx = s2_idx + h_idx; + let dst_idx = h_idx * k_w; + for w_idx in 0..k_w { + let s1_idx = s1_idx + w_idx; + let s3_idx = s3_idx + w_idx; + let dst_idx = dst_idx + w_idx; + dst[dst_idx] = s1[s1_idx] + s2[s2_idx] + s3[s3_idx] + } + } + }); + let dst = candle::WithDType::to_cpu_storage_owned(dst); + Ok((dst, (b, q_h * q_w, k_h * k_w).into())) + } +} + +#[derive(Debug)] +struct Attention { + qkv: super::Linear, + proj: super::Linear, + num_heads: usize, + scale: f64, + rel_pos_hw: Option<(Tensor, Tensor)>, + span: tracing::Span, + span_matmul: tracing::Span, + span_rel_pos: tracing::Span, + span_softmax: tracing::Span, +} + +impl Attention { + fn new( + dim: usize, + num_heads: usize, + qkv_bias: bool, + use_rel_pos: bool, + input_size: (usize, usize), + vb: VarBuilder, + ) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "attention"); + let span_matmul = tracing::span!(tracing::Level::TRACE, "attn-matmul"); + let span_rel_pos = tracing::span!(tracing::Level::TRACE, "attn-rel-pos"); + let span_softmax = tracing::span!(tracing::Level::TRACE, "attn-sm"); + let qkv = super::linear(vb.pp("qkv"), dim, dim * 3, qkv_bias)?; + let proj = super::linear(vb.pp("proj"), dim, dim, true)?; + let head_dim = dim / num_heads; + let scale = 1. / (head_dim as f64).sqrt(); + let rel_pos_hw = if use_rel_pos { + let h = vb.get((2 * input_size.0 - 1, head_dim), "rel_pos_h")?; + let w = vb.get((2 * input_size.1 - 1, head_dim), "rel_pos_w")?; + Some((h, w)) + } else { + None + }; + Ok(Self { + qkv, + proj, + num_heads, + scale, + rel_pos_hw, + span, + span_matmul, + span_rel_pos, + span_softmax, + }) + } + + fn add_decomposed_rel_pos( + &self, + attn: Tensor, + q: &Tensor, + (q_h, q_w): (usize, usize), + (k_h, k_w): (usize, usize), + ) -> Result { + match &self.rel_pos_hw { + Some((rel_pos_h, rel_pos_w)) => { + let r_h = get_rel_pos(q_h, k_h, rel_pos_h)?; + let r_w = get_rel_pos(q_w, k_w, rel_pos_w)?; + let (b, _, dim) = q.dims3()?; + let r_q = q.reshape((b, q_h, q_w, dim))?; + // rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh) + let rel_h = r_q.matmul(&r_h.broadcast_left(b)?.t()?.contiguous()?)?; + // rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw) + let rel_w = r_q + .transpose(1, 2)? // -> bwhc + .contiguous()? + .matmul(&r_w.broadcast_left(b)?.t()?.contiguous()?)? // bwhc,bwck -> bwhk + .transpose(1, 2)? + .contiguous()?; + if attn.device().is_cpu() { + let op = Add3(b, q_h, q_w, k_h, k_w); + attn.apply_op3_no_bwd(&rel_h, &rel_w, &op) + } else { + (attn.reshape((b, q_h, q_w, k_h, k_w))? + + rel_h.unsqueeze(4)?.broadcast_add(&rel_w.unsqueeze(3)?)?)? + .reshape((b, q_h * q_w, k_h * k_w)) + } + } + None => Ok(attn), + } + } +} + +fn get_rel_pos(q_size: usize, k_size: usize, rel_pos: &Tensor) -> Result { + let max_rel_dist = 2 * usize::max(q_size, k_size) - 1; + let dev = rel_pos.device(); + let rel_pos_resized = if rel_pos.dim(0)? != max_rel_dist { + todo!("interpolation") + } else { + rel_pos + }; + let q_coords = Tensor::arange(0u32, q_size as u32, dev)? + .reshape((q_size, 1))? + .to_dtype(DType::F32)?; + let k_coords = Tensor::arange(0u32, k_size as u32, dev)? + .reshape((1, k_size))? + .to_dtype(DType::F32)?; + let q_coords = (q_coords * f64::max(1f64, k_size as f64 / q_size as f64))?; + let k_coords = (k_coords * f64::max(1f64, q_size as f64 / k_size as f64))?; + let relative_coords = (q_coords.broadcast_sub(&k_coords)? + + (k_size as f64 - 1.) * f64::max(1f64, q_size as f64 / k_size as f64))?; + let (d1, d2) = relative_coords.dims2()?; + let relative_coords = relative_coords.to_dtype(DType::U32)?; + rel_pos_resized + .index_select(&relative_coords.reshape(d1 * d2)?, 0)? + .reshape((d1, d2, ())) +} + +impl Module for Attention { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let (b, h, w, c) = xs.dims4()?; + let qkv = self + .qkv + .forward(&xs.flatten_to(1)?)? + .reshape((b, h * w, 3, self.num_heads, c / self.num_heads))? + .permute((2, 0, 3, 1, 4))? + .reshape((3, b * self.num_heads, h * w, c / self.num_heads))?; + let q = qkv.i(0)?; + let k = qkv.i(1)?; + let v = qkv.i(2)?; + let attn = { + let _enter = self.span_matmul.enter(); + (&q * self.scale)?.matmul(&k.t()?)? + }; + let attn = { + let _enter = self.span_rel_pos.enter(); + self.add_decomposed_rel_pos(attn, &q, (h, w), (h, w))? + }; + let attn = { + let _enter = self.span_softmax.enter(); + candle_nn::ops::softmax_last_dim(&attn)? + }; + let attn = { + let _enter = self.span_matmul.enter(); + attn.matmul(&v)? + }; + let attn = attn + .reshape((b, self.num_heads, h, w, c / self.num_heads))? + .permute((0, 2, 3, 1, 4))? + .reshape((b, h * w, c))?; + self.proj.forward(&attn)?.reshape((b, h, w, c)) + } +} + +#[derive(Debug)] +struct Block { + norm1: LayerNorm, + attn: Attention, + norm2: LayerNorm, + mlp: super::MlpBlock, + window_size: usize, + span: tracing::Span, +} + +impl Block { + fn new( + dim: usize, + num_heads: usize, + qkv_bias: bool, + use_rel_pos: bool, + window_size: usize, + input_size: (usize, usize), + vb: VarBuilder, + ) -> Result { + let norm1 = layer_norm(dim, 1e-6, vb.pp("norm1"))?; + let norm2 = layer_norm(dim, 1e-6, vb.pp("norm2"))?; + let input_size_attn = if window_size == 0 { + input_size + } else { + (window_size, window_size) + }; + let attn = Attention::new( + dim, + num_heads, + qkv_bias, + use_rel_pos, + input_size_attn, + vb.pp("attn"), + )?; + let mlp = super::MlpBlock::new(dim, dim * 4, candle_nn::Activation::Gelu, vb.pp("mlp"))?; + let span = tracing::span!(tracing::Level::TRACE, "ie-block"); + Ok(Self { + norm1, + attn, + norm2, + mlp, + window_size, + span, + }) + } +} + +fn window_partition(xs: Tensor, window_size: usize) -> Result<(Tensor, (usize, usize))> { + let (b, h, w, c) = xs.dims4()?; + let pad_h = (window_size - h % window_size) % window_size; + let pad_w = (window_size - w % window_size) % window_size; + let xs = if pad_h > 0 { + xs.pad_with_zeros(1, 0, pad_h)? + } else { + xs + }; + let xs = if pad_w > 0 { + xs.pad_with_zeros(2, 0, pad_w)? + } else { + xs + }; + let (h_p, w_p) = (h + pad_h, w + pad_w); + let windows = xs + .reshape(( + b, + h_p / window_size, + window_size, + w_p / window_size, + window_size, + c, + ))? + .transpose(2, 3)? + .contiguous()? + .flatten_to(2)?; + Ok((windows, (h_p, w_p))) +} + +fn window_unpartition( + windows: Tensor, + window_size: usize, + (h_p, w_p): (usize, usize), + (h, w): (usize, usize), +) -> Result { + let b = windows.dim(0)? / (h_p * w_p / window_size / window_size); + let xs = windows + .reshape(( + b, + h_p / window_size, + w_p / window_size, + window_size, + window_size, + windows.elem_count() / b / h_p / w_p, + ))? + .transpose(2, 3)? + .contiguous()? + .reshape((b, h_p, w_p, ()))?; + let xs = if h_p > h { xs.narrow(1, 0, h)? } else { xs }; + let xs = if w_p > w { xs.narrow(2, 0, w)? } else { xs }; + Ok(xs) +} + +impl Module for Block { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let shortcut = xs; + let xs = self.norm1.forward(xs)?; + let hw = (xs.dim(1)?, xs.dim(2)?); + let (xs, pad_hw) = if self.window_size > 0 { + window_partition(xs, self.window_size)? + } else { + (xs, (0, 0)) + }; + let xs = self.attn.forward(&xs)?; + let xs = if self.window_size > 0 { + window_unpartition(xs, self.window_size, pad_hw, hw)? + } else { + xs + }; + let xs = (xs + shortcut)?; + &xs + xs.apply(&self.norm2)?.apply(&self.mlp)? + } +} + +#[derive(Debug)] +pub struct ImageEncoderViT { + patch_embed: PatchEmbed, + blocks: Vec, + neck_conv1: candle_nn::Conv2d, + neck_ln1: super::LayerNorm2d, + neck_conv2: candle_nn::Conv2d, + neck_ln2: super::LayerNorm2d, + pos_embed: Option, + span: tracing::Span, +} + +impl ImageEncoderViT { + #[allow(clippy::too_many_arguments)] + pub fn new( + img_size: usize, + patch_size: usize, + in_chans: usize, + embed_dim: usize, + depth: usize, + num_heads: usize, + out_chans: usize, + qkv_bias: bool, + use_rel_pos: bool, + use_abs_pos: bool, + window_size: usize, + global_attn_indexes: &[usize], + vb: VarBuilder, + ) -> Result { + let patch_embed = PatchEmbed::new( + in_chans, + embed_dim, + patch_size, + patch_size, + 0, + vb.pp("patch_embed"), + )?; + let mut blocks = Vec::with_capacity(depth); + let vb_b = vb.pp("blocks"); + for i in 0..depth { + let window_size = if global_attn_indexes.contains(&i) { + 0 + } else { + window_size + }; + let block = Block::new( + embed_dim, + num_heads, + qkv_bias, + use_rel_pos, + window_size, + (img_size / patch_size, img_size / patch_size), + vb_b.pp(i), + )?; + blocks.push(block) + } + let neck_conv1 = candle_nn::conv2d_no_bias( + embed_dim, + out_chans, + 1, + Default::default(), + vb.pp("neck.0"), + )?; + let neck_ln1 = super::LayerNorm2d::new(out_chans, 1e-6, vb.pp("neck.1"))?; + let cfg = candle_nn::Conv2dConfig { + padding: 1, + ..Default::default() + }; + let neck_conv2 = candle_nn::conv2d_no_bias(out_chans, out_chans, 3, cfg, vb.pp("neck.2"))?; + let neck_ln2 = super::LayerNorm2d::new(out_chans, 1e-6, vb.pp("neck.3"))?; + let pos_embed = if use_abs_pos { + let p = vb.get( + (1, img_size / patch_size, img_size / patch_size, embed_dim), + "pos_embed", + )?; + Some(p) + } else { + None + }; + let span = tracing::span!(tracing::Level::TRACE, "image-encoder-vit"); + Ok(Self { + patch_embed, + blocks, + neck_conv1, + neck_ln1, + neck_conv2, + neck_ln2, + pos_embed, + span, + }) + } +} + +impl Module for ImageEncoderViT { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let xs = self.patch_embed.forward(xs)?; + let mut xs = match &self.pos_embed { + Some(pos_embed) => (xs + pos_embed)?, + None => xs, + }; + for block in self.blocks.iter() { + xs = block.forward(&xs)? + } + xs.permute((0, 3, 1, 2))? + .apply(&self.neck_conv1)? + .apply(&self.neck_ln1)? + .apply(&self.neck_conv2)? + .apply(&self.neck_ln2) + } +} diff --git a/patches/candle-transformers/src/models/segment_anything/mask_decoder.rs b/patches/candle-transformers/src/models/segment_anything/mask_decoder.rs new file mode 100644 index 0000000000..1703c80979 --- /dev/null +++ b/patches/candle-transformers/src/models/segment_anything/mask_decoder.rs @@ -0,0 +1,239 @@ +use candle::{IndexOp, Result, Tensor}; +use candle_nn::{Module, VarBuilder}; + +use super::transformer::TwoWayTransformer; + +#[derive(Debug)] +struct MlpMaskDecoder { + layers: Vec, + sigmoid_output: bool, + span: tracing::Span, +} + +impl MlpMaskDecoder { + fn new( + input_dim: usize, + hidden_dim: usize, + output_dim: usize, + num_layers: usize, + sigmoid_output: bool, + vb: VarBuilder, + ) -> Result { + let mut layers = Vec::with_capacity(num_layers); + let vb = vb.pp("layers"); + for i in 0..num_layers { + let in_dim = if i == 0 { input_dim } else { hidden_dim }; + let out_dim = if i + 1 == num_layers { + output_dim + } else { + hidden_dim + }; + let layer = super::linear(vb.pp(i), in_dim, out_dim, true)?; + layers.push(layer) + } + let span = tracing::span!(tracing::Level::TRACE, "mlp-mask-decoder"); + Ok(Self { + layers, + sigmoid_output, + span, + }) + } +} + +impl Module for MlpMaskDecoder { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let mut xs = xs.clone(); + for (i, layer) in self.layers.iter().enumerate() { + xs = layer.forward(&xs)?; + if i + 1 < self.layers.len() { + xs = xs.relu()? + } + } + if self.sigmoid_output { + candle_nn::ops::sigmoid(&xs) + } else { + Ok(xs) + } + } +} + +#[derive(Debug)] +pub struct MaskDecoder { + iou_token: candle_nn::Embedding, + mask_tokens: candle_nn::Embedding, + iou_prediction_head: MlpMaskDecoder, + output_upscaling_conv1: candle_nn::ConvTranspose2d, + output_upscaling_ln: super::LayerNorm2d, + output_upscaling_conv2: candle_nn::ConvTranspose2d, + num_mask_tokens: usize, + output_hypernetworks_mlps: Vec, + transformer: TwoWayTransformer, + span: tracing::Span, +} + +impl MaskDecoder { + pub fn new( + transformer_dim: usize, + num_multimask_outputs: usize, + iou_head_depth: usize, + iou_head_hidden_dim: usize, + vb: VarBuilder, + ) -> Result { + let num_mask_tokens = num_multimask_outputs + 1; + let iou_prediction_head = MlpMaskDecoder::new( + transformer_dim, + iou_head_hidden_dim, + num_mask_tokens, + iou_head_depth, + false, + vb.pp("iou_prediction_head"), + )?; + let iou_token = candle_nn::embedding(1, transformer_dim, vb.pp("iou_token"))?; + let mask_tokens = + candle_nn::embedding(num_mask_tokens, transformer_dim, vb.pp("mask_tokens"))?; + let cfg = candle_nn::ConvTranspose2dConfig { + stride: 2, + ..Default::default() + }; + let output_upscaling_conv1 = candle_nn::conv_transpose2d( + transformer_dim, + transformer_dim / 4, + 2, + cfg, + vb.pp("output_upscaling.0"), + )?; + let output_upscaling_ln = + super::LayerNorm2d::new(transformer_dim / 4, 1e-6, vb.pp("output_upscaling.1"))?; + let output_upscaling_conv2 = candle_nn::conv_transpose2d( + transformer_dim / 4, + transformer_dim / 8, + 2, + cfg, + vb.pp("output_upscaling.3"), + )?; + let mut output_hypernetworks_mlps = Vec::with_capacity(num_mask_tokens); + let vb_o = vb.pp("output_hypernetworks_mlps"); + for i in 0..num_mask_tokens { + let mlp = MlpMaskDecoder::new( + transformer_dim, + transformer_dim, + transformer_dim / 8, + 3, + false, + vb_o.pp(i), + )?; + output_hypernetworks_mlps.push(mlp) + } + let transformer = TwoWayTransformer::new( + /* depth */ 2, + /* embedding_dim */ transformer_dim, + /* num_heads */ 8, + /* mlp_dim */ 2048, + vb.pp("transformer"), + )?; + let span = tracing::span!(tracing::Level::TRACE, "mask-decoder"); + Ok(Self { + iou_token, + mask_tokens, + iou_prediction_head, + output_upscaling_conv1, + output_upscaling_ln, + output_upscaling_conv2, + num_mask_tokens, + output_hypernetworks_mlps, + transformer, + span, + }) + } + + pub fn forward( + &self, + image_embeddings: &Tensor, + image_pe: &Tensor, + sparse_prompt_embeddings: &Tensor, + dense_prompt_embeddings: &Tensor, + multimask_output: bool, + ) -> Result<(Tensor, Tensor)> { + let _enter = self.span.enter(); + let (masks, iou_pred) = self.predict_masks( + image_embeddings, + image_pe, + sparse_prompt_embeddings, + dense_prompt_embeddings, + )?; + let masks = if multimask_output { + masks.i((.., 1..))? + } else { + masks.i((.., 0..1))? + }; + let iou_pred = if multimask_output { + iou_pred.i((.., 1..))? + } else { + iou_pred.i((.., 0..1))? + }; + Ok((masks, iou_pred)) + } + + fn predict_masks( + &self, + image_embeddings: &Tensor, + image_pe: &Tensor, + sparse_prompt_embeddings: &Tensor, + dense_prompt_embeddings: &Tensor, + ) -> Result<(Tensor, Tensor)> { + // Concatenate output tokens. + let output_tokens = Tensor::cat( + &[self.iou_token.embeddings(), self.mask_tokens.embeddings()], + 0, + )?; + let (d1, d2) = output_tokens.dims2()?; + let output_tokens = + output_tokens + .unsqueeze(0)? + .expand((sparse_prompt_embeddings.dim(0)?, d1, d2))?; + let tokens = Tensor::cat(&[&output_tokens, sparse_prompt_embeddings], 1)?; + + // Expand per-image data in batch direction to be per mask + let src = repeat_interleave(image_embeddings, tokens.dim(0)?, 0)?; + let src = src.broadcast_add(dense_prompt_embeddings)?; + let pos_src = repeat_interleave(image_pe, tokens.dim(0)?, 0)?; + let (b, c, h, w) = src.dims4()?; + + // Run the transformer + let (hs, src) = self.transformer.forward(&src, &pos_src, &tokens)?; + let iou_token_out = hs.i((.., 0))?; + let mask_tokens_out = hs.i((.., 1..1 + self.num_mask_tokens))?; + + // Upscale mask embeddings and predict masks using the masks tokens. + let src = src.transpose(1, 2)?.reshape((b, c, h, w))?; + let upscaled_embedding = self + .output_upscaling_conv1 + .forward(&src)? + .apply(&self.output_upscaling_ln)? + .gelu()? + .apply(&self.output_upscaling_conv2)? + .gelu()?; + let mut hyper_in_list = Vec::with_capacity(self.num_mask_tokens); + for (i, mlp) in self.output_hypernetworks_mlps.iter().enumerate() { + let h = mlp.forward(&mask_tokens_out.i((.., i))?)?; + hyper_in_list.push(h) + } + let hyper_in = Tensor::stack(hyper_in_list.as_slice(), 1)?.contiguous()?; + let (b, c, h, w) = upscaled_embedding.dims4()?; + let masks = hyper_in.matmul(&upscaled_embedding.reshape((b, c, h * w))?)?; + let masks = masks.reshape((b, (), h, w))?; + + // Generate mask quality predictions. + let iou_pred = self.iou_prediction_head.forward(&iou_token_out)?; + Ok((masks, iou_pred)) + } +} + +// Equivalent to torch.repeat_interleave +fn repeat_interleave(img: &Tensor, repeats: usize, dim: usize) -> Result { + let img = img.unsqueeze(dim + 1)?; + let mut dims = img.dims().to_vec(); + dims[dim + 1] = repeats; + img.broadcast_as(dims)?.flatten(dim, dim + 1) +} diff --git a/patches/candle-transformers/src/models/segment_anything/mod.rs b/patches/candle-transformers/src/models/segment_anything/mod.rs new file mode 100644 index 0000000000..fe0b099008 --- /dev/null +++ b/patches/candle-transformers/src/models/segment_anything/mod.rs @@ -0,0 +1,117 @@ +//! Segment Anything Model (SAM) +//! +//! SAM is an architecture for image segmentation, capable of segmenting any object +//! in an image based on prompts like points or boxes. //! This model provides a robust and fast image segmentation pipeline that can be tweaked via +//! some prompting (requesting some points to be in the target mask, requesting some +//! points to be part of the background so _not_ in the target mask, specifying some +//! bounding box). +//! +//! - ⚡ [Interactive Wasm Example](https://huggingface.co/spaces/radames/candle-segment-anything-wasm) +//! - 💻 [GH Link](https://github.com/facebookresearch/segment-anything) +//! - 📝 [Paper](https://arxiv.org/abs/2304.02643) +//! - 💡 The default backbone can be replaced by the smaller and faster TinyViT model based on [MobileSAM](https://github.com/ChaoningZhang/MobileSAM). +//! +//! +//! ## Example +//! +//! ```bash +//! cargo run --example segment-anything --release -- \ +//! --image candle-examples/examples/yolo-v8/assets/bike.jpg +//! --use-tiny --point 0.6,0.6 --point 0.6,0.55 +//! ``` +//! +//!
+//! +//! +//! +//!
+//! +//! +//! > Original; Prompt with `--point 0.6,0.55`; Prompt with `--point 0.6,0.6 --point 0.6,0.55` +//! +pub use crate::models::with_tracing::Linear; +use candle::{Result, Tensor}; +use candle_nn::{Module, VarBuilder}; + +pub mod image_encoder; +pub mod mask_decoder; +pub mod prompt_encoder; +pub mod sam; +pub mod tiny_vit; +pub mod transformer; + +pub fn linear(vb: VarBuilder, in_dim: usize, out_dim: usize, bias: bool) -> Result { + if bias { + crate::models::with_tracing::linear(in_dim, out_dim, vb) + } else { + crate::models::with_tracing::linear_no_bias(in_dim, out_dim, vb) + } +} + +#[derive(Debug)] +pub struct LayerNorm2d { + weight: Tensor, + bias: Tensor, + num_channels: usize, + eps: f64, +} + +impl LayerNorm2d { + pub fn new(num_channels: usize, eps: f64, vb: VarBuilder) -> Result { + let weight = vb.get(num_channels, "weight")?; + let bias = vb.get(num_channels, "bias")?; + Ok(Self { + weight, + bias, + num_channels, + eps, + }) + } +} + +impl Module for LayerNorm2d { + fn forward(&self, xs: &Tensor) -> Result { + let u = xs.mean_keepdim(1)?; + let xs = xs.broadcast_sub(&u)?; + let s = xs.sqr()?.mean_keepdim(1)?; + let xs = xs.broadcast_div(&(s + self.eps)?.sqrt()?)?; + xs.broadcast_mul(&self.weight.reshape((1, self.num_channels, 1, 1))?)? + .broadcast_add(&self.bias.reshape((1, self.num_channels, 1, 1))?) + } +} + +#[derive(Debug)] +pub struct MlpBlock { + lin1: Linear, + lin2: Linear, + activation: candle_nn::Activation, + span: tracing::Span, +} + +impl MlpBlock { + pub fn new( + embedding_dim: usize, + mlp_dim: usize, + activation: candle_nn::Activation, + vb: VarBuilder, + ) -> Result { + let lin1 = linear(vb.pp("lin1"), embedding_dim, mlp_dim, true)?; + let lin2 = linear(vb.pp("lin2"), mlp_dim, embedding_dim, true)?; + let span = tracing::span!(tracing::Level::TRACE, "mlp-block"); + Ok(Self { + lin1, + lin2, + activation, + span, + }) + } +} + +impl Module for MlpBlock { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + xs.apply(&self.lin1)? + .apply(&self.activation)? + .apply(&self.lin2) + } +} diff --git a/patches/candle-transformers/src/models/segment_anything/prompt_encoder.rs b/patches/candle-transformers/src/models/segment_anything/prompt_encoder.rs new file mode 100644 index 0000000000..258fb5aa9e --- /dev/null +++ b/patches/candle-transformers/src/models/segment_anything/prompt_encoder.rs @@ -0,0 +1,240 @@ +use candle::{DType, IndexOp, Result, Tensor, D}; +use candle_nn::VarBuilder; + +#[derive(Debug)] +struct PositionEmbeddingRandom { + positional_encoding_gaussian_matrix: Tensor, +} + +impl PositionEmbeddingRandom { + fn new(num_pos_feats: usize, vb: VarBuilder) -> Result { + let positional_encoding_gaussian_matrix = + vb.get((2, num_pos_feats), "positional_encoding_gaussian_matrix")?; + Ok(Self { + positional_encoding_gaussian_matrix, + }) + } + + fn pe_encoding(&self, coords: &Tensor) -> Result { + let coords = coords.affine(2., -1.)?; + let coords = coords.broadcast_matmul(&self.positional_encoding_gaussian_matrix)?; + let coords = (coords * (2. * std::f64::consts::PI))?; + Tensor::cat(&[coords.sin()?, coords.cos()?], D::Minus1) + } + + fn forward(&self, h: usize, w: usize) -> Result { + let device = self.positional_encoding_gaussian_matrix.device(); + let x_embed = (Tensor::arange(0u32, w as u32, device)?.to_dtype(DType::F32)? + 0.5)?; + let y_embed = (Tensor::arange(0u32, h as u32, device)?.to_dtype(DType::F32)? + 0.5)?; + let x_embed = (x_embed / w as f64)? + .reshape((1, ()))? + .broadcast_as((h, w))?; + let y_embed = (y_embed / h as f64)? + .reshape(((), 1))? + .broadcast_as((h, w))?; + let coords = Tensor::stack(&[&x_embed, &y_embed], D::Minus1)?; + self.pe_encoding(&coords)?.permute((2, 0, 1)) + } + + fn forward_with_coords( + &self, + coords_input: &Tensor, + image_size: (usize, usize), + ) -> Result { + let coords0 = (coords_input.narrow(D::Minus1, 0, 1)? / image_size.1 as f64)?; + let coords1 = (coords_input.narrow(D::Minus1, 1, 1)? / image_size.0 as f64)?; + let c = coords_input.dim(D::Minus1)?; + let coords_rest = coords_input.narrow(D::Minus1, 2, c - 2)?; + let coords = Tensor::cat(&[&coords0, &coords1, &coords_rest], D::Minus1)?; + self.pe_encoding(&coords) + } +} + +#[derive(Debug)] +pub struct PromptEncoder { + pe_layer: PositionEmbeddingRandom, + point_embeddings: Vec, + not_a_point_embed: candle_nn::Embedding, + mask_downscaling_conv1: candle_nn::Conv2d, + mask_downscaling_ln1: super::LayerNorm2d, + mask_downscaling_conv2: candle_nn::Conv2d, + mask_downscaling_ln2: super::LayerNorm2d, + mask_downscaling_conv3: candle_nn::Conv2d, + no_mask_embed: candle_nn::Embedding, + image_embedding_size: (usize, usize), + input_image_size: (usize, usize), + embed_dim: usize, + span: tracing::Span, +} + +impl PromptEncoder { + pub fn new( + embed_dim: usize, + image_embedding_size: (usize, usize), + input_image_size: (usize, usize), + mask_in_chans: usize, + vb: VarBuilder, + ) -> Result { + let num_points_embeddings = 4; + let pe_layer = PositionEmbeddingRandom::new(embed_dim / 2, vb.pp("pe_layer"))?; + let not_a_point_embed = candle_nn::embedding(1, embed_dim, vb.pp("not_a_point_embed"))?; + let no_mask_embed = candle_nn::embedding(1, embed_dim, vb.pp("no_mask_embed"))?; + let cfg = candle_nn::Conv2dConfig { + stride: 2, + ..Default::default() + }; + let mask_downscaling_conv1 = + candle_nn::conv2d(1, mask_in_chans / 4, 2, cfg, vb.pp("mask_downscaling.0"))?; + let mask_downscaling_conv2 = candle_nn::conv2d( + mask_in_chans / 4, + mask_in_chans, + 2, + cfg, + vb.pp("mask_downscaling.3"), + )?; + let mask_downscaling_conv3 = candle_nn::conv2d( + mask_in_chans, + embed_dim, + 1, + Default::default(), + vb.pp("mask_downscaling.6"), + )?; + let mask_downscaling_ln1 = + super::LayerNorm2d::new(mask_in_chans / 4, 1e-6, vb.pp("mask_downscaling.1"))?; + let mask_downscaling_ln2 = + super::LayerNorm2d::new(mask_in_chans, 1e-6, vb.pp("mask_downscaling.4"))?; + let mut point_embeddings = Vec::with_capacity(num_points_embeddings); + let vb_e = vb.pp("point_embeddings"); + for i in 0..num_points_embeddings { + let emb = candle_nn::embedding(1, embed_dim, vb_e.pp(i))?; + point_embeddings.push(emb) + } + let span = tracing::span!(tracing::Level::TRACE, "prompt-encoder"); + Ok(Self { + pe_layer, + point_embeddings, + not_a_point_embed, + mask_downscaling_conv1, + mask_downscaling_ln1, + mask_downscaling_conv2, + mask_downscaling_ln2, + mask_downscaling_conv3, + no_mask_embed, + image_embedding_size, + input_image_size, + embed_dim, + span, + }) + } + + pub fn get_dense_pe(&self) -> Result { + self.pe_layer + .forward(self.image_embedding_size.0, self.image_embedding_size.1)? + .unsqueeze(0) + } + + fn embed_masks(&self, masks: &Tensor) -> Result { + masks + .apply(&self.mask_downscaling_conv1)? + .apply(&self.mask_downscaling_ln1)? + .gelu()? + .apply(&self.mask_downscaling_conv2)? + .apply(&self.mask_downscaling_ln2)? + .gelu()? + .apply(&self.mask_downscaling_conv3) + } + + fn embed_points(&self, points: &Tensor, labels: &Tensor, pad: bool) -> Result { + let points = (points + 0.5)?; + let dev = points.device(); + let (points, labels) = if pad { + let padding_point = Tensor::zeros((points.dim(0)?, 1, 2), DType::F32, dev)?; + let padding_label = (Tensor::ones((labels.dim(0)?, 1), DType::F32, dev)? * (-1f64))?; + let points = Tensor::cat(&[&points, &padding_point], 1)?; + let labels = Tensor::cat(&[labels, &padding_label], 1)?; + (points, labels) + } else { + (points, labels.clone()) + }; + let point_embedding = self + .pe_layer + .forward_with_coords(&points, self.input_image_size)?; + let labels = labels.unsqueeze(2)?.broadcast_as(point_embedding.shape())?; + let zeros = point_embedding.zeros_like()?; + let point_embedding = labels.lt(0f32)?.where_cond( + &self + .not_a_point_embed + .embeddings() + .broadcast_as(zeros.shape())?, + &point_embedding, + )?; + let labels0 = labels.eq(0f32)?.where_cond( + &self.point_embeddings[0] + .embeddings() + .broadcast_as(zeros.shape())?, + &zeros, + )?; + let point_embedding = (point_embedding + labels0)?; + let labels1 = labels.eq(1f32)?.where_cond( + &self.point_embeddings[1] + .embeddings() + .broadcast_as(zeros.shape())?, + &zeros, + )?; + let point_embedding = (point_embedding + labels1)?; + Ok(point_embedding) + } + + fn embed_boxes(&self, boxes: &Tensor) -> Result { + let boxes = (boxes + 0.5)?; + let coords = boxes.reshape(((), 2, 2))?; + let corner_embedding = self + .pe_layer + .forward_with_coords(&coords, self.input_image_size)?; + let ce1 = corner_embedding.i((.., 0))?; + let ce2 = corner_embedding.i((.., 1))?; + let ce1 = (ce1 + self.point_embeddings[2].embeddings())?; + let ce2 = (ce2 + self.point_embeddings[3].embeddings())?; + Tensor::cat(&[&ce1, &ce2], 1) + } + + pub fn forward( + &self, + points: Option<(&Tensor, &Tensor)>, + boxes: Option<&Tensor>, + masks: Option<&Tensor>, + ) -> Result<(Tensor, Tensor)> { + let _enter = self.span.enter(); + let se_points = match points { + Some((coords, labels)) => Some(self.embed_points(coords, labels, boxes.is_none())?), + None => None, + }; + let se_boxes = match boxes { + Some(boxes) => Some(self.embed_boxes(boxes)?), + None => None, + }; + let sparse_embeddings = match (se_points, se_boxes) { + (Some(se_points), Some(se_boxes)) => Tensor::cat(&[se_points, se_boxes], 1)?, + (Some(se_points), None) => se_points, + (None, Some(se_boxes)) => se_boxes, + (None, None) => { + let dev = self.no_mask_embed.embeddings().device(); + Tensor::zeros((1, 0, self.embed_dim), DType::F32, dev)? + } + }; + + let dense_embeddings = match masks { + None => { + let emb = self.no_mask_embed.embeddings(); + emb.reshape((1, (), 1, 1))?.expand(( + 1, + emb.elem_count(), + self.image_embedding_size.0, + self.image_embedding_size.1, + ))? + } + Some(masks) => self.embed_masks(masks)?, + }; + Ok((sparse_embeddings, dense_embeddings)) + } +} diff --git a/patches/candle-transformers/src/models/segment_anything/sam.rs b/patches/candle-transformers/src/models/segment_anything/sam.rs new file mode 100644 index 0000000000..7a84eef419 --- /dev/null +++ b/patches/candle-transformers/src/models/segment_anything/sam.rs @@ -0,0 +1,447 @@ +use candle::{DType, IndexOp, Result, Tensor}; +use candle_nn::{Module, VarBuilder}; + +use super::image_encoder::ImageEncoderViT; +use super::mask_decoder::MaskDecoder; +use super::prompt_encoder::PromptEncoder; +use super::tiny_vit::{tiny_vit_5m, TinyViT}; + +const PROMPT_EMBED_DIM: usize = 256; +pub const IMAGE_SIZE: usize = 1024; +const VIT_PATCH_SIZE: usize = 16; +const PRED_IOU_THRESH: f32 = 0.88; +const STABILITY_SCORE_OFFSET: f32 = 1.0; +const STABILITY_SCORE_THRESHOLD: f32 = 0.95; +const MODEL_MASK_THRESHOLD: f32 = 0.0; +const CROP_NMS_THRESH: f32 = 0.7; + +#[derive(Debug)] +enum ImageEncoder { + Original(Box), + TinyViT(Box), +} + +impl Module for ImageEncoder { + fn forward(&self, xs: &Tensor) -> Result { + match self { + Self::Original(vit) => vit.forward(xs), + Self::TinyViT(vit) => vit.forward(xs), + } + } +} + +#[derive(Debug)] +pub struct Sam { + image_encoder: ImageEncoder, + prompt_encoder: PromptEncoder, + mask_decoder: MaskDecoder, + pixel_mean: Tensor, + pixel_std: Tensor, +} + +impl Sam { + pub fn new( + encoder_embed_dim: usize, + encoder_depth: usize, + encoder_num_heads: usize, + encoder_global_attn_indexes: &[usize], + vb: VarBuilder, + ) -> Result { + let image_embedding_size = IMAGE_SIZE / VIT_PATCH_SIZE; + + let image_encoder = ImageEncoderViT::new( + IMAGE_SIZE, + VIT_PATCH_SIZE, + 3, + encoder_embed_dim, + encoder_depth, + encoder_num_heads, + PROMPT_EMBED_DIM, + /* qkv_bias */ true, + /* use_rel_pos */ true, + /* use_abs_pos */ true, + /* window_size */ 14, + /* global_attn_indexes */ encoder_global_attn_indexes, + vb.pp("image_encoder"), + )?; + let prompt_encoder = PromptEncoder::new( + PROMPT_EMBED_DIM, + (image_embedding_size, image_embedding_size), + (IMAGE_SIZE, IMAGE_SIZE), + 16, + vb.pp("prompt_encoder"), + )?; + let mask_decoder = MaskDecoder::new( + PROMPT_EMBED_DIM, + /* num_multitask_outputs */ 3, + /* iou_head_depth */ 3, + /* iou_head_hidden_dim */ 256, + vb.pp("mask_decoder"), + )?; + let pixel_mean = + Tensor::new(&[123.675f32, 116.28, 103.53], vb.device())?.reshape((3, 1, 1))?; + let pixel_std = + Tensor::new(&[58.395f32, 57.12, 57.375], vb.device())?.reshape((3, 1, 1))?; + Ok(Self { + image_encoder: ImageEncoder::Original(image_encoder.into()), + prompt_encoder, + mask_decoder, + pixel_std, + pixel_mean, + }) + } + + pub fn new_tiny(vb: VarBuilder) -> Result { + let image_embedding_size = IMAGE_SIZE / VIT_PATCH_SIZE; + + let image_encoder = tiny_vit_5m(vb.pp("image_encoder"))?; + let prompt_encoder = PromptEncoder::new( + PROMPT_EMBED_DIM, + (image_embedding_size, image_embedding_size), + (IMAGE_SIZE, IMAGE_SIZE), + 16, + vb.pp("prompt_encoder"), + )?; + let mask_decoder = MaskDecoder::new( + PROMPT_EMBED_DIM, + /* num_multitask_outputs */ 3, + /* iou_head_depth */ 3, + /* iou_head_hidden_dim */ 256, + vb.pp("mask_decoder"), + )?; + let pixel_mean = + Tensor::new(&[123.675f32, 116.28, 103.53], vb.device())?.reshape((3, 1, 1))?; + let pixel_std = + Tensor::new(&[58.395f32, 57.12, 57.375], vb.device())?.reshape((3, 1, 1))?; + Ok(Self { + image_encoder: ImageEncoder::TinyViT(image_encoder.into()), + prompt_encoder, + mask_decoder, + pixel_std, + pixel_mean, + }) + } + + pub fn embeddings(&self, img: &Tensor) -> Result { + let img = self.preprocess(img)?.unsqueeze(0)?; + self.image_encoder.forward(&img) + } + + pub fn forward( + &self, + img: &Tensor, + points: &[(f64, f64, bool)], + multimask_output: bool, + ) -> Result<(Tensor, Tensor)> { + let (_c, original_h, original_w) = img.dims3()?; + let img = self.preprocess(img)?.unsqueeze(0)?; + let img_embeddings = self.image_encoder.forward(&img)?; + let (low_res_mask, iou) = self.forward_for_embeddings( + &img_embeddings, + original_h, + original_w, + points, + multimask_output, + )?; + let mask = low_res_mask + .upsample_nearest2d(IMAGE_SIZE, IMAGE_SIZE)? + .get(0)? + .i((.., ..original_h, ..original_w))?; + Ok((mask, iou)) + } + + /// Generate the mask and IOU predictions from some image embeddings and prompt. + /// + /// The prompt is specified as a list of points `(x, y, b)`. `x` and `y` are the point + /// coordinates (between 0 and 1) and `b` is `true` for points that should be part of the mask + /// and `false` for points that should be part of the background and so excluded from the mask. + pub fn forward_for_embeddings( + &self, + img_embeddings: &Tensor, + original_h: usize, + original_w: usize, + points: &[(f64, f64, bool)], + multimask_output: bool, + ) -> Result<(Tensor, Tensor)> { + let image_pe = self.prompt_encoder.get_dense_pe()?; + let points = if points.is_empty() { + None + } else { + let n_points = points.len(); + let xys = points + .iter() + .flat_map(|(x, y, _b)| { + let x = (*x as f32) * (original_w as f32); + let y = (*y as f32) * (original_h as f32); + [x, y] + }) + .collect::>(); + let labels = points + .iter() + .map(|(_x, _y, b)| if *b { 1f32 } else { 0f32 }) + .collect::>(); + let points = Tensor::from_vec(xys, (1, n_points, 2), img_embeddings.device())?; + let labels = Tensor::from_vec(labels, (1, n_points), img_embeddings.device())?; + Some((points, labels)) + }; + let points = points.as_ref().map(|xy| (&xy.0, &xy.1)); + let (sparse_prompt_embeddings, dense_prompt_embeddings) = + self.prompt_encoder.forward(points, None, None)?; + self.mask_decoder.forward( + img_embeddings, + &image_pe, + &sparse_prompt_embeddings, + &dense_prompt_embeddings, + multimask_output, + ) + } + + pub fn unpreprocess(&self, img: &Tensor) -> Result { + let img = img + .broadcast_mul(&self.pixel_std)? + .broadcast_add(&self.pixel_mean)?; + img.maximum(&img.zeros_like()?)? + .minimum(&(img.ones_like()? * 255.)?) + } + + pub fn preprocess(&self, img: &Tensor) -> Result { + let (_c, h, w) = img.dims3()?; + let img = img + .to_dtype(DType::F32)? + .broadcast_sub(&self.pixel_mean)? + .broadcast_div(&self.pixel_std)?; + if h > IMAGE_SIZE || w > IMAGE_SIZE { + candle::bail!("image is too large ({w}, {h}), maximum size {IMAGE_SIZE}") + } + let img = img.pad_with_zeros(1, 0, IMAGE_SIZE - h)?; + img.pad_with_zeros(2, 0, IMAGE_SIZE - w) + } + + fn process_crop( + &self, + img: &Tensor, + cb: CropBox, + point_grids: &[(f64, f64)], + ) -> Result>> { + // Crop the image and calculate embeddings. + let img = img.i((.., cb.y0..cb.y1, cb.x0..cb.x1))?; + let img = self.preprocess(&img)?.unsqueeze(0)?; + let img_embeddings = self.image_encoder.forward(&img)?; + + let crop_w = cb.x1 - cb.x0; + let crop_h = cb.y1 - cb.y0; + + // Generate masks for this crop. + let image_pe = self.prompt_encoder.get_dense_pe()?; + let points = point_grids + .iter() + .map(|&(x, y)| vec![x as f32 * crop_w as f32, y as f32 * crop_h as f32]) + .collect::>(); + + let mut bboxes = Vec::new(); + for points in points.chunks(64) { + // Run the model on this batch. + let points_len = points.len(); + let in_points = Tensor::new(points.to_vec(), img.device())?.unsqueeze(1)?; + let in_labels = Tensor::ones((points_len, 1), DType::F32, img.device())?; + let (sparse_prompt_embeddings, dense_prompt_embeddings) = + self.prompt_encoder + .forward(Some((&in_points, &in_labels)), None, None)?; + + let (low_res_mask, iou_predictions) = self.mask_decoder.forward( + &img_embeddings, + &image_pe, + &sparse_prompt_embeddings, + &dense_prompt_embeddings, + /* multimask_output */ true, + )?; + let low_res_mask = low_res_mask.flatten(0, 1)?; + let iou_predictions = iou_predictions.flatten(0, 1)?.to_vec1::()?; + let dev = low_res_mask.device(); + + for (i, iou) in iou_predictions.iter().enumerate() { + // Filter by predicted IoU. + if *iou < PRED_IOU_THRESH { + continue; + } + let low_res_mask = low_res_mask.get(i)?; + + // Calculate stability score. + let bound = Tensor::new(MODEL_MASK_THRESHOLD + STABILITY_SCORE_OFFSET, dev)? + .broadcast_as(low_res_mask.shape())?; + let intersections = low_res_mask + .ge(&bound)? + .to_dtype(DType::F32)? + .sum_all()? + .to_vec0::()?; + let bound = Tensor::new(MODEL_MASK_THRESHOLD - STABILITY_SCORE_OFFSET, dev)? + .broadcast_as(low_res_mask.shape())?; + let unions = low_res_mask + .ge(&bound)? + .to_dtype(DType::F32)? + .sum_all()? + .to_vec0::()?; + let stability_score = intersections / unions; + if stability_score < STABILITY_SCORE_THRESHOLD { + continue; + } + + // Threshold masks and calculate boxes. + let low_res_mask = low_res_mask + .ge(&Tensor::new(0f32, dev)?.broadcast_as(low_res_mask.shape())?)? + .to_dtype(DType::U32)?; + let low_res_mask_per_x = low_res_mask.sum(0)?.to_vec1::()?; + let low_res_mask_per_y = low_res_mask.sum(1)?.to_vec1::()?; + let min_max_x = min_max_indexes(&low_res_mask_per_x); + let min_max_y = min_max_indexes(&low_res_mask_per_y); + if let Some(((x0, x1), (y0, y1))) = min_max_x.zip(min_max_y) { + let bbox = crate::object_detection::Bbox { + xmin: x0 as f32, + ymin: y0 as f32, + xmax: x1 as f32, + ymax: y1 as f32, + confidence: *iou, + data: low_res_mask, + }; + bboxes.push(bbox); + } + // TODO: + // Filter boxes that touch crop boundaries + // Compress to RLE. + } + } + + let mut bboxes = vec![bboxes]; + // Remove duplicates within this crop. + crate::object_detection::non_maximum_suppression(&mut bboxes, CROP_NMS_THRESH); + + // TODO: Return to the original image frame. + Ok(bboxes.remove(0)) + } + + pub fn generate_masks( + &self, + img: &Tensor, + points_per_side: usize, + crop_n_layer: usize, + crop_overlap_ratio: f64, + crop_n_points_downscale_factor: usize, + ) -> Result>> { + let (_c, h, w) = img.dims3()?; + let point_grids = build_all_layer_point_grids( + points_per_side, + crop_n_layer, + crop_n_points_downscale_factor, + ); + let crop_boxes = generate_crop_boxes((h, w), crop_n_layer, crop_overlap_ratio); + let mut bboxes = Vec::new(); + for crop_box in crop_boxes.into_iter() { + let layer_idx = crop_box.layer_idx; + let b = self.process_crop(img, crop_box, &point_grids[layer_idx])?; + bboxes.extend(b) + } + // TODO: remove duplicates + Ok(bboxes) + } +} + +// Return the first and last indexes i for which values[i] > 0 +fn min_max_indexes(values: &[u32]) -> Option<(usize, usize)> { + let (mut min_i, mut max_i) = (usize::MAX, usize::MIN); + for (i, &s) in values.iter().enumerate() { + if s == 0 { + continue; + } + min_i = usize::min(i, min_i); + max_i = usize::max(i, max_i); + } + if max_i < min_i { + None + } else { + Some((min_i, max_i)) + } +} + +#[derive(Debug)] +struct CropBox { + x0: usize, + y0: usize, + x1: usize, + y1: usize, + layer_idx: usize, +} + +impl CropBox { + fn new(x0: usize, y0: usize, x1: usize, y1: usize, layer_idx: usize) -> Self { + Self { + x0, + y0, + x1, + y1, + layer_idx, + } + } +} + +fn generate_crop_boxes( + (im_h, im_w): (usize, usize), + n_layers: usize, + overlap_ratio: f64, +) -> Vec { + fn crop_len(orig_len: usize, n_crops: usize, overlap: usize) -> usize { + f64::ceil((overlap * (n_crops - 1) + orig_len) as f64 / n_crops as f64) as usize + } + + let short_side = usize::min(im_h, im_w); + + let mut crop_boxes = Vec::new(); + + // Original image. + crop_boxes.push(CropBox::new(0, 0, im_w, im_h, 0)); + + for layer_idx in 1..=n_layers { + let n_crops_per_side = 1 << layer_idx; + let overlap = (overlap_ratio * short_side as f64 * 2. / n_crops_per_side as f64) as usize; + let crop_w = crop_len(im_w, n_crops_per_side, overlap); + let crop_h = crop_len(im_w, n_crops_per_side, overlap); + + for i_x in 0..n_crops_per_side { + let x0 = (crop_w - overlap) * i_x; + for i_y in 0..n_crops_per_side { + let y0 = (crop_h - overlap) * i_y; + let x1 = usize::min(im_w, x0 + crop_w); + let y1 = usize::min(im_h, y0 + crop_h); + crop_boxes.push(CropBox::new(x0, y0, x1, y1, layer_idx)); + } + } + } + + crop_boxes +} + +// Generates a 2D grid of points evenly spaced in [0,1]x[0,1]. +fn build_point_grid(n_per_side: usize) -> Vec<(f64, f64)> { + let offset = 1f64 / (2 * n_per_side) as f64; + let mut points = Vec::with_capacity(n_per_side * n_per_side); + for i_x in 0..n_per_side { + let x = offset + i_x as f64 / n_per_side as f64; + for i_y in 0..n_per_side { + let y = offset + i_y as f64 / n_per_side as f64; + points.push((x, y)) + } + } + points +} + +fn build_all_layer_point_grids( + n_per_side: usize, + n_layers: usize, + scale_per_layer: usize, +) -> Vec> { + let mut points_by_layer = Vec::with_capacity(n_layers + 1); + for i in 0..=n_layers { + let n_points = n_per_side / scale_per_layer.pow(i as u32); + points_by_layer.push(build_point_grid(n_points)) + } + points_by_layer +} diff --git a/patches/candle-transformers/src/models/segment_anything/tiny_vit.rs b/patches/candle-transformers/src/models/segment_anything/tiny_vit.rs new file mode 100644 index 0000000000..d1700cc56b --- /dev/null +++ b/patches/candle-transformers/src/models/segment_anything/tiny_vit.rs @@ -0,0 +1,633 @@ +// Adapted from: +// https://github.com/ChaoningZhang/MobileSAM/blob/master/mobile_sam/modeling/tiny_vit_sam.py +use candle::{IndexOp, Result, Tensor, D}; +use candle_nn::{Conv2dConfig, Module, VarBuilder}; + +const MBCONV_EXPAND_RATIO: usize = 4; +const MLP_RATIO: usize = 4; +const LOCAL_CONV_SIZE: usize = 3; +const IMG_SIZE: usize = 1024; +const IN_CHANNELS: usize = 3; + +#[derive(Debug)] +struct Conv2dBN { + c: candle_nn::Conv2d, + bn: candle_nn::BatchNorm, + span: tracing::Span, +} + +impl Conv2dBN { + fn new(in_: usize, out: usize, ks: usize, cfg: Conv2dConfig, vb: VarBuilder) -> Result { + let c = candle_nn::conv2d_no_bias(in_, out, ks, cfg, vb.pp("c"))?; + let bn = candle_nn::batch_norm(out, 1e-5, vb.pp("bn"))?; + let span = tracing::span!(tracing::Level::TRACE, "conv2d-bn"); + Ok(Self { c, bn, span }) + } +} + +impl Module for Conv2dBN { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + xs.apply(&self.c)?.apply_t(&self.bn, false) + } +} + +#[derive(Debug)] +struct PatchEmbed { + conv1: Conv2dBN, + conv2: Conv2dBN, + span: tracing::Span, +} + +impl PatchEmbed { + fn new(in_chans: usize, embed_dim: usize, vb: VarBuilder) -> Result { + let cfg = candle_nn::Conv2dConfig { + stride: 2, + padding: 1, + ..Default::default() + }; + let conv1 = Conv2dBN::new(in_chans, embed_dim / 2, 3, cfg, vb.pp("seq.0"))?; + let conv2 = Conv2dBN::new(embed_dim / 2, embed_dim, 3, cfg, vb.pp("seq.2"))?; + let span = tracing::span!(tracing::Level::TRACE, "patch-embed"); + Ok(Self { conv1, conv2, span }) + } +} + +impl Module for PatchEmbed { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + xs.apply(&self.conv1)?.gelu()?.apply(&self.conv2) + } +} + +#[derive(Debug)] +struct MBConv { + conv1: Conv2dBN, + conv2: Conv2dBN, + conv3: Conv2dBN, + span: tracing::Span, +} + +impl MBConv { + fn new(in_: usize, out: usize, expand_ratio: usize, vb: VarBuilder) -> Result { + let hidden = in_ * expand_ratio; + let cfg2 = candle_nn::Conv2dConfig { + padding: 1, + groups: hidden, + ..Default::default() + }; + let conv1 = Conv2dBN::new(in_, hidden, 1, Default::default(), vb.pp("conv1"))?; + let conv2 = Conv2dBN::new(hidden, hidden, 3, cfg2, vb.pp("conv2"))?; + let conv3 = Conv2dBN::new(hidden, out, 1, Default::default(), vb.pp("conv3"))?; + let span = tracing::span!(tracing::Level::TRACE, "mb-conv"); + Ok(Self { + conv1, + conv2, + conv3, + span, + }) + } +} + +impl Module for MBConv { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let shortcut = xs; + let xs = xs + .apply(&self.conv1)? + .gelu()? + .apply(&self.conv2)? + .gelu()? + .apply(&self.conv3)?; + (xs + shortcut)?.gelu() + } +} + +#[derive(Debug)] +struct PatchMerging { + conv1: Conv2dBN, + conv2: Conv2dBN, + conv3: Conv2dBN, + input_resolution: (usize, usize), + span: tracing::Span, +} + +impl PatchMerging { + fn new( + input_resolution: (usize, usize), + dim: usize, + out: usize, + vb: VarBuilder, + ) -> Result { + let stride = if [320, 448, 576].contains(&out) { 1 } else { 2 }; + let cfg2 = candle_nn::Conv2dConfig { + padding: 1, + stride, + groups: out, + ..Default::default() + }; + let conv1 = Conv2dBN::new(dim, out, 1, Default::default(), vb.pp("conv1"))?; + let conv2 = Conv2dBN::new(out, out, 3, cfg2, vb.pp("conv2"))?; + let conv3 = Conv2dBN::new(out, out, 1, Default::default(), vb.pp("conv3"))?; + let span = tracing::span!(tracing::Level::TRACE, "patch-merging"); + Ok(Self { + conv1, + conv2, + conv3, + input_resolution, + span, + }) + } +} + +impl Module for PatchMerging { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let xs = if xs.rank() == 3 { + let (h, w) = self.input_resolution; + let b = xs.dim(0)?; + xs.reshape((b, h, w, ()))?.permute((0, 3, 1, 2))? + } else { + xs.clone() + }; + xs.apply(&self.conv1)? + .gelu()? + .apply(&self.conv2)? + .gelu()? + .apply(&self.conv3)? + .flatten_from(2)? + .transpose(1, 2) + } +} + +#[derive(Debug)] +struct ConvLayer { + blocks: Vec, + downsample: Option, + span: tracing::Span, +} + +impl ConvLayer { + fn new( + dim: usize, + out: usize, + input_resolution: (usize, usize), + depth: usize, + downsample: bool, + conv_expand_ratio: usize, + vb: VarBuilder, + ) -> Result { + let vb_b = vb.pp("blocks"); + let mut blocks = Vec::with_capacity(depth); + for index in 0..depth { + let block = MBConv::new(dim, dim, conv_expand_ratio, vb_b.pp(index))?; + blocks.push(block) + } + let downsample = if downsample { + let downsample = PatchMerging::new(input_resolution, dim, out, vb.pp("downsample"))?; + Some(downsample) + } else { + None + }; + let span = tracing::span!(tracing::Level::TRACE, "conv-layer"); + Ok(Self { + blocks, + downsample, + span, + }) + } +} + +impl Module for ConvLayer { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let mut xs = xs.clone(); + for block in self.blocks.iter() { + xs = block.forward(&xs)? + } + match &self.downsample { + None => Ok(xs), + Some(downsample) => downsample.forward(&xs), + } + } +} + +#[derive(Debug)] +struct Mlp { + norm: candle_nn::LayerNorm, + fc1: super::Linear, + fc2: super::Linear, + span: tracing::Span, +} + +impl Mlp { + fn new(in_: usize, hidden: usize, vb: VarBuilder) -> Result { + let norm = candle_nn::layer_norm(in_, 1e-5, vb.pp("norm"))?; + let fc1 = super::linear(vb.pp("fc1"), in_, hidden, true)?; + let fc2 = super::linear(vb.pp("fc2"), hidden, in_, true)?; + let span = tracing::span!(tracing::Level::TRACE, "mlp"); + Ok(Self { + norm, + fc1, + fc2, + span, + }) + } +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + xs.apply(&self.norm)? + .apply(&self.fc1)? + .gelu()? + .apply(&self.fc2) + } +} + +#[derive(Debug)] +struct Attention { + norm: candle_nn::LayerNorm, + qkv: super::Linear, + proj: super::Linear, + ab: Tensor, + key_dim: usize, + num_heads: usize, + d: usize, + dh: usize, + scale: f64, + span: tracing::Span, + span_matmul: tracing::Span, + span_softmax: tracing::Span, +} + +impl Attention { + fn new( + dim: usize, + key_dim: usize, + num_heads: usize, + attn_ratio: usize, + resolution: (usize, usize), + vb: VarBuilder, + ) -> Result { + let d = attn_ratio * key_dim; + let dh = d * num_heads; + let nh_kd = key_dim * num_heads; + let h = dh + nh_kd * 2; + let norm = candle_nn::layer_norm(dim, 1e-5, vb.pp("norm"))?; + let qkv = super::linear(vb.pp("qkv"), dim, h, true)?; + let proj = super::linear(vb.pp("proj"), dh, dim, true)?; + + let points = (0..resolution.0) + .flat_map(|x| (0..resolution.1).map(move |y| (x as i64, y as i64))) + .collect::>(); + let mut idxs = Vec::with_capacity(points.len() * points.len()); + let mut attention_offsets = std::collections::HashMap::new(); + for &(x1, y1) in points.iter() { + for &(x2, y2) in points.iter() { + let offset = ((x2 - x1).abs(), (y2 - y1).abs()); + let l = attention_offsets.len(); + let idx = attention_offsets.entry(offset).or_insert(l); + idxs.push(*idx as u32) + } + } + let attention_biases = vb.get((num_heads, attention_offsets.len()), "attention_biases")?; + let idxs = Tensor::new(idxs, attention_biases.device())?; + let ab = + attention_biases + .index_select(&idxs, 1)? + .reshape(((), points.len(), points.len()))?; + let span = tracing::span!(tracing::Level::TRACE, "attention"); + let span_matmul = tracing::span!(tracing::Level::TRACE, "attn-matmul"); + let span_softmax = tracing::span!(tracing::Level::TRACE, "attn-sm"); + Ok(Self { + norm, + qkv, + proj, + ab, + key_dim, + num_heads, + d, + dh, + scale: 1f64 / (key_dim as f64).sqrt(), + span, + span_matmul, + span_softmax, + }) + } +} + +impl Module for Attention { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let (b, n, _) = xs.dims3()?; + let xs = xs.apply(&self.norm)?; + let qkv = xs.apply(&self.qkv)?.reshape((b, n, self.num_heads, ()))?; + let q = qkv + .narrow(D::Minus1, 0, self.key_dim)? + .permute((0, 2, 1, 3))? + .contiguous()?; + let k = qkv + .narrow(D::Minus1, self.key_dim, self.key_dim)? + .permute((0, 2, 1, 3))? + .contiguous()?; + let v = qkv + .narrow(D::Minus1, 2 * self.key_dim, self.d)? + .permute((0, 2, 1, 3))? + .contiguous()?; + let attn = { + let _enter = self.span_matmul.enter(); + (q.matmul(&k.t()?)? * self.scale)? + }; + let attn = attn.broadcast_add(&self.ab)?; + let attn = { + let _enter = self.span_softmax.enter(); + candle_nn::ops::softmax_last_dim(&attn)? + }; + let attn = { + let _enter = self.span_matmul.enter(); + attn.matmul(&v)? + }; + attn.transpose(1, 2)? + .reshape((b, n, self.dh))? + .apply(&self.proj) + } +} + +#[derive(Debug)] +struct TinyViTBlock { + attn: Attention, + local_conv: Conv2dBN, + mlp: Mlp, + window_size: usize, + input_resolution: (usize, usize), + span: tracing::Span, +} + +impl TinyViTBlock { + fn new( + dim: usize, + input_resolution: (usize, usize), + num_heads: usize, + window_size: usize, + vb: VarBuilder, + ) -> Result { + let head_dim = dim / num_heads; + let attn = Attention::new( + dim, + head_dim, + num_heads, + 1, + (window_size, window_size), + vb.pp("attn"), + )?; + let mlp = Mlp::new(dim, dim * MLP_RATIO, vb.pp("mlp"))?; + let cfg = candle_nn::Conv2dConfig { + padding: LOCAL_CONV_SIZE / 2, + groups: dim, + ..Default::default() + }; + let local_conv = Conv2dBN::new(dim, dim, LOCAL_CONV_SIZE, cfg, vb.pp("local_conv"))?; + let span = tracing::span!(tracing::Level::TRACE, "attention"); + Ok(Self { + attn, + local_conv, + mlp, + window_size, + input_resolution, + span, + }) + } +} + +impl Module for TinyViTBlock { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let (h, w) = self.input_resolution; + let (b, l, c) = xs.dims3()?; + let res_x = xs; + let xs = if h == self.window_size && w == self.window_size { + self.attn.forward(xs)? + } else { + let xs = xs.reshape((b, h, w, c))?; + let pad_b = (self.window_size - h % self.window_size) % self.window_size; + let pad_r = (self.window_size - w % self.window_size) % self.window_size; + + let xs = if pad_b > 0 { + xs.pad_with_zeros(1, 0, pad_b)? + } else { + xs + }; + let xs = if pad_r > 0 { + xs.pad_with_zeros(2, 0, pad_r)? + } else { + xs + }; + let (p_h, p_w) = (h + pad_b, w + pad_r); + let n_h = p_h / self.window_size; + let n_w = p_w / self.window_size; + let xs = xs + .reshape((b, n_h, self.window_size, n_w, self.window_size, c))? + .transpose(2, 3)? + .reshape((b * n_h * n_w, self.window_size * self.window_size, c))?; + let xs = self.attn.forward(&xs)?; + let xs = xs + .reshape((b, n_h, n_w, self.window_size, self.window_size, c))? + .transpose(2, 3)? + .reshape((b, p_h, p_w, c))?; + let xs = if pad_r > 0 { + xs.i((.., .., ..w))?.contiguous()? + } else { + xs + }; + let xs = if pad_b > 0 { + xs.i((.., ..h, ..))?.contiguous()? + } else { + xs + }; + xs.reshape((b, l, c))? + }; + let xs = (xs + res_x)?; + let xs = xs + .transpose(1, 2)? + .reshape((b, c, h, w))? + .apply(&self.local_conv)? + .reshape((b, c, l))? + .transpose(1, 2)?; + &xs + self.mlp.forward(&xs)? + } +} + +#[derive(Debug)] +struct BasicLayer { + blocks: Vec, + downsample: Option, + span: tracing::Span, +} + +impl BasicLayer { + #[allow(clippy::too_many_arguments)] + fn new( + dim: usize, + input_resolution: (usize, usize), + depth: usize, + num_heads: usize, + window_size: usize, + downsample: bool, + out: usize, + vb: VarBuilder, + ) -> Result { + let vb_b = vb.pp("blocks"); + let mut blocks = Vec::with_capacity(depth); + for index in 0..depth { + let block = TinyViTBlock::new( + dim, + input_resolution, + num_heads, + window_size, + vb_b.pp(index), + )?; + blocks.push(block) + } + let downsample = if downsample { + let downsample = PatchMerging::new(input_resolution, dim, out, vb.pp("downsample"))?; + Some(downsample) + } else { + None + }; + let span = tracing::span!(tracing::Level::TRACE, "basic-layer"); + Ok(Self { + blocks, + downsample, + span, + }) + } +} + +impl Module for BasicLayer { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let mut xs = xs.clone(); + for block in self.blocks.iter() { + xs = block.forward(&xs)? + } + match &self.downsample { + None => Ok(xs), + Some(downsample) => downsample.forward(&xs), + } + } +} + +#[derive(Debug)] +pub struct TinyViT { + patch_embed: PatchEmbed, + layer0: ConvLayer, + layers: Vec, + // norm_head: candle_nn::LayerNorm, + // head: candle_nn::Linear, + neck_conv1: candle_nn::Conv2d, + neck_ln1: super::LayerNorm2d, + neck_conv2: candle_nn::Conv2d, + neck_ln2: super::LayerNorm2d, + span: tracing::Span, + span_neck: tracing::Span, +} + +impl TinyViT { + pub fn new( + embed_dims: &[usize], + depths: &[usize], + num_heads: &[usize], + window_sizes: &[usize], + _num_classes: usize, + vb: VarBuilder, + ) -> Result { + let patch_embed = PatchEmbed::new(IN_CHANNELS, embed_dims[0], vb.pp("patch_embed"))?; + let patches_resolution = IMG_SIZE / 4; + + let vb_l = vb.pp("layers"); + let layer0 = ConvLayer::new( + /* dim */ embed_dims[0], + /* out */ embed_dims[1], + /* input_resolution */ (patches_resolution, patches_resolution), + /* depth */ depths[0], + /* downsample */ true, + /* conv_expand_ratio */ MBCONV_EXPAND_RATIO, + vb_l.pp(0), + )?; + + let num_layers = embed_dims.len(); + let mut layers = Vec::with_capacity(num_layers - 1); + for i_layer in 1..num_layers { + let patches_resolution = patches_resolution / (1 << usize::min(i_layer, 2)); + let layer = BasicLayer::new( + /* dim */ embed_dims[i_layer], + /* input_resolution */ (patches_resolution, patches_resolution), + /* depth */ depths[i_layer], + /* num_heads */ num_heads[i_layer], + /* window_size */ window_sizes[i_layer], + /* downsample */ i_layer < num_layers - 1, + /* out */ embed_dims[usize::min(i_layer + 1, num_layers - 1)], + vb_l.pp(i_layer), + )?; + layers.push(layer) + } + + let last_embed_dim = embed_dims[embed_dims.len() - 1]; + // let norm_head = candle_nn::layer_norm(last_embed_dim, 1e-5, vb.pp("norm_head"))?; + // let head = candle_nn::linear(last_embed_dim, num_classes, vb.pp("head"))?; + let neck_conv1 = + candle_nn::conv2d_no_bias(last_embed_dim, 256, 1, Default::default(), vb.pp("neck.0"))?; + let neck_ln1 = super::LayerNorm2d::new(256, 1e-6, vb.pp("neck.1"))?; + let cfg = candle_nn::Conv2dConfig { + padding: 1, + ..Default::default() + }; + let neck_conv2 = candle_nn::conv2d_no_bias(256, 256, 3, cfg, vb.pp("neck.2"))?; + let neck_ln2 = super::LayerNorm2d::new(256, 1e-6, vb.pp("neck.3"))?; + + let span = tracing::span!(tracing::Level::TRACE, "tiny-vit"); + let span_neck = tracing::span!(tracing::Level::TRACE, "neck"); + Ok(Self { + patch_embed, + layer0, + layers, + neck_conv1, + neck_ln1, + neck_conv2, + neck_ln2, + span, + span_neck, + }) + } +} + +impl Module for TinyViT { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let xs = self.patch_embed.forward(xs)?; + let mut xs = self.layer0.forward(&xs)?; + for layer in self.layers.iter() { + xs = layer.forward(&xs)? + } + let (b, _, c) = xs.dims3()?; + let _enter = self.span_neck.enter(); + xs.reshape((b, 64, 64, c))? + .permute((0, 3, 1, 2))? + .apply(&self.neck_conv1)? + .apply(&self.neck_ln1)? + .apply(&self.neck_conv2)? + .apply(&self.neck_ln2) + } +} + +pub fn tiny_vit_5m(vb: VarBuilder) -> Result { + TinyViT::new( + /* embed_dims */ &[64, 128, 160, 320], + /* depths */ &[2, 2, 6, 2], + /* num_heads */ &[2, 4, 5, 10], + /* window_sizes */ &[7, 7, 14, 7], + /* num_classes */ 1000, + vb, + ) +} diff --git a/patches/candle-transformers/src/models/segment_anything/transformer.rs b/patches/candle-transformers/src/models/segment_anything/transformer.rs new file mode 100644 index 0000000000..80efb38c5b --- /dev/null +++ b/patches/candle-transformers/src/models/segment_anything/transformer.rs @@ -0,0 +1,221 @@ +use candle::{Result, Tensor}; +use candle_nn::{layer_norm, LayerNorm, Linear, Module, VarBuilder}; + +#[derive(Debug)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + out_proj: Linear, + num_heads: usize, +} + +impl Attention { + fn new( + embedding_dim: usize, + num_heads: usize, + downsample_rate: usize, + vb: VarBuilder, + ) -> Result { + let internal_dim = embedding_dim / downsample_rate; + let q_proj = candle_nn::linear(embedding_dim, internal_dim, vb.pp("q_proj"))?; + let k_proj = candle_nn::linear(embedding_dim, internal_dim, vb.pp("k_proj"))?; + let v_proj = candle_nn::linear(embedding_dim, internal_dim, vb.pp("v_proj"))?; + let out_proj = candle_nn::linear(internal_dim, embedding_dim, vb.pp("out_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + out_proj, + num_heads, + }) + } + + fn separate_heads(&self, x: &Tensor) -> Result { + let (b, n, c) = x.dims3()?; + x.reshape((b, n, self.num_heads, c / self.num_heads))? + .transpose(1, 2)? + .contiguous() + } + + fn recombine_heads(&self, x: &Tensor) -> Result { + let (b, n_heads, n_tokens, c_per_head) = x.dims4()?; + x.transpose(1, 2)? + .reshape((b, n_tokens, n_heads * c_per_head)) + } + + fn forward(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result { + let q = self.q_proj.forward(&q.contiguous()?)?; + let k = self.k_proj.forward(&k.contiguous()?)?; + let v = self.v_proj.forward(&v.contiguous()?)?; + + let q = self.separate_heads(&q)?; + let k = self.separate_heads(&k)?; + let v = self.separate_heads(&v)?; + + let (_, _, _, c_per_head) = q.dims4()?; + let attn = (q.matmul(&k.t()?)? / (c_per_head as f64).sqrt())?; + let attn = candle_nn::ops::softmax_last_dim(&attn)?; + + let out = attn.matmul(&v)?; + self.recombine_heads(&out)?.apply(&self.out_proj) + } +} + +#[derive(Debug)] +struct TwoWayAttentionBlock { + self_attn: Attention, + norm1: LayerNorm, + cross_attn_token_to_image: Attention, + norm2: LayerNorm, + mlp: super::MlpBlock, + norm3: LayerNorm, + norm4: LayerNorm, + cross_attn_image_to_token: Attention, + skip_first_layer_pe: bool, +} + +impl TwoWayAttentionBlock { + fn new( + embedding_dim: usize, + num_heads: usize, + mlp_dim: usize, + skip_first_layer_pe: bool, + vb: VarBuilder, + ) -> Result { + let norm1 = layer_norm(embedding_dim, 1e-5, vb.pp("norm1"))?; + let norm2 = layer_norm(embedding_dim, 1e-5, vb.pp("norm2"))?; + let norm3 = layer_norm(embedding_dim, 1e-5, vb.pp("norm3"))?; + let norm4 = layer_norm(embedding_dim, 1e-5, vb.pp("norm4"))?; + let self_attn = Attention::new(embedding_dim, num_heads, 1, vb.pp("self_attn"))?; + let cross_attn_token_to_image = Attention::new( + embedding_dim, + num_heads, + 2, + vb.pp("cross_attn_token_to_image"), + )?; + let cross_attn_image_to_token = Attention::new( + embedding_dim, + num_heads, + 2, + vb.pp("cross_attn_image_to_token"), + )?; + let mlp = super::MlpBlock::new( + embedding_dim, + mlp_dim, + candle_nn::Activation::Relu, + vb.pp("mlp"), + )?; + Ok(Self { + self_attn, + norm1, + cross_attn_image_to_token, + norm2, + mlp, + norm3, + norm4, + cross_attn_token_to_image, + skip_first_layer_pe, + }) + } + + fn forward( + &self, + queries: &Tensor, + keys: &Tensor, + query_pe: &Tensor, + key_pe: &Tensor, + ) -> Result<(Tensor, Tensor)> { + // Self attention block + let queries = if self.skip_first_layer_pe { + self.self_attn.forward(queries, queries, queries)? + } else { + let q = (queries + query_pe)?; + let attn_out = self.self_attn.forward(&q, &q, queries)?; + (queries + attn_out)? + }; + let queries = self.norm1.forward(&queries)?; + + // Cross attention block, tokens attending to image embedding + let q = (&queries + query_pe)?; + let k = (keys + key_pe)?; + let attn_out = self.cross_attn_token_to_image.forward(&q, &k, keys)?; + let queries = (&queries + attn_out)?; + let queries = self.norm2.forward(&queries)?; + + // MLP block + let mlp_out = self.mlp.forward(&queries); + let queries = (queries + mlp_out)?; + let queries = self.norm3.forward(&queries)?; + + // Cross attention block, image embedding attending to tokens + let q = (&queries + query_pe)?; + let k = (keys + key_pe)?; + let attn_out = self.cross_attn_image_to_token.forward(&k, &q, &queries)?; + let keys = (keys + attn_out)?; + let keys = self.norm4.forward(&keys)?; + + Ok((queries, keys)) + } +} + +#[derive(Debug)] +pub struct TwoWayTransformer { + layers: Vec, + final_attn_token_to_image: Attention, + norm_final_attn: LayerNorm, +} + +impl TwoWayTransformer { + pub fn new( + depth: usize, + embedding_dim: usize, + num_heads: usize, + mlp_dim: usize, + vb: VarBuilder, + ) -> Result { + let vb_l = vb.pp("layers"); + let mut layers = Vec::with_capacity(depth); + for i in 0..depth { + let layer = + TwoWayAttentionBlock::new(embedding_dim, num_heads, mlp_dim, i == 0, vb_l.pp(i))?; + layers.push(layer) + } + let final_attn_token_to_image = Attention::new( + embedding_dim, + num_heads, + 2, + vb.pp("final_attn_token_to_image"), + )?; + let norm_final_attn = layer_norm(embedding_dim, 1e-5, vb.pp("norm_final_attn"))?; + Ok(Self { + layers, + final_attn_token_to_image, + norm_final_attn, + }) + } + + pub fn forward( + &self, + image_embedding: &Tensor, + image_pe: &Tensor, + point_embedding: &Tensor, + ) -> Result<(Tensor, Tensor)> { + let image_embedding = image_embedding.flatten_from(2)?.permute((0, 2, 1))?; + let image_pe = image_pe.flatten_from(2)?.permute((0, 2, 1))?; + + let mut queries = point_embedding.clone(); + let mut keys = image_embedding; + + for layer in self.layers.iter() { + (queries, keys) = layer.forward(&queries, &keys, point_embedding, &image_pe)? + } + + let q = (&queries + point_embedding)?; + let k = (&keys + image_pe)?; + let attn_out = self.final_attn_token_to_image.forward(&q, &k, &keys)?; + let queries = (queries + attn_out)?.apply(&self.norm_final_attn)?; + + Ok((queries, keys)) + } +} diff --git a/patches/candle-transformers/src/models/siglip.rs b/patches/candle-transformers/src/models/siglip.rs new file mode 100644 index 0000000000..578beea3d8 --- /dev/null +++ b/patches/candle-transformers/src/models/siglip.rs @@ -0,0 +1,798 @@ +//! Siglip model implementation. +//! +//! Siglip architecture combining vision and language for zero-shot tasks. +//! +//! References: +//! - 🤗 [Model Card](https://huggingface.co/google/siglip-base-patch16-224) +//! + +use crate::models::clip::div_l2_norm; +use candle::{IndexOp, Module, Result, Tensor, D}; +use candle_nn::{layer_norm, linear, LayerNorm, Linear, VarBuilder}; + +fn default_text_vocab_size() -> usize { + 32000 +} + +fn default_text_hidden_size() -> usize { + 768 +} + +fn default_text_intermediate_size() -> usize { + 3072 +} + +fn default_text_num_hidden_layers() -> usize { + 12 +} + +fn default_text_num_attention_heads() -> usize { + 12 +} + +fn default_text_max_position_embeddings() -> usize { + 64 +} + +fn default_text_layer_norm_eps() -> f64 { + 1e-6 +} + +fn default_text_pad_token_id() -> u32 { + 1 +} + +fn default_text_bos_token_id() -> u32 { + 49406 +} + +fn default_text_eos_token_id() -> u32 { + 49407 +} + +fn default_text_hidden_act() -> candle_nn::Activation { + candle_nn::Activation::GeluPytorchTanh +} + +// https://github.com/huggingface/transformers/blob/2e24ee4dfa39cc0bc264b89edbccc373c8337086/src/transformers/models/siglip/configuration_siglip.py#L27 +#[derive(serde::Deserialize, Clone, Debug)] +pub struct TextConfig { + #[serde(default = "default_text_vocab_size")] + pub vocab_size: usize, + #[serde(default = "default_text_hidden_size")] + pub hidden_size: usize, + #[serde(default = "default_text_intermediate_size")] + pub intermediate_size: usize, + #[serde(default = "default_text_num_hidden_layers")] + pub num_hidden_layers: usize, + #[serde(default = "default_text_num_attention_heads")] + pub num_attention_heads: usize, + #[serde(default = "default_text_max_position_embeddings")] + pub max_position_embeddings: usize, + #[serde(default = "default_text_hidden_act")] + pub hidden_act: candle_nn::Activation, + #[serde(default = "default_text_layer_norm_eps")] + pub layer_norm_eps: f64, + #[serde(default = "default_text_pad_token_id")] + pub pad_token_id: u32, + #[serde(default = "default_text_bos_token_id")] + pub bos_token_id: u32, + #[serde(default = "default_text_eos_token_id")] + pub eos_token_id: u32, +} + +fn default_vision_hidden_size() -> usize { + 768 +} + +fn default_vision_intermediate_size() -> usize { + 3072 +} + +fn default_vision_num_hidden_layers() -> usize { + 12 +} + +fn default_vision_num_attention_heads() -> usize { + 12 +} + +fn default_vision_num_channels() -> usize { + 3 +} + +fn default_vision_image_size() -> usize { + 224 +} + +fn default_vision_batch_size() -> usize { + 16 +} + +fn default_vision_layer_norm_eps() -> f64 { + 1e-6 +} + +fn default_vision_hidden_act() -> candle_nn::Activation { + candle_nn::Activation::GeluPytorchTanh +} + +// https://github.com/huggingface/transformers/blob/2e24ee4dfa39cc0bc264b89edbccc373c8337086/src/transformers/models/siglip/configuration_siglip.py#L132 +#[derive(serde::Deserialize, Clone, Debug)] +pub struct VisionConfig { + #[serde(default = "default_vision_hidden_size")] + pub hidden_size: usize, + #[serde(default = "default_vision_intermediate_size")] + pub intermediate_size: usize, + #[serde(default = "default_vision_num_hidden_layers")] + pub num_hidden_layers: usize, + #[serde(default = "default_vision_num_attention_heads")] + pub num_attention_heads: usize, + #[serde(default = "default_vision_num_channels")] + pub num_channels: usize, + #[serde(default = "default_vision_image_size")] + pub image_size: usize, + #[serde(default = "default_vision_batch_size")] + pub patch_size: usize, + #[serde(default = "default_vision_hidden_act")] + pub hidden_act: candle_nn::Activation, + #[serde(default = "default_vision_layer_norm_eps")] + pub layer_norm_eps: f64, +} + +trait TransformerConfig { + fn hidden_size(&self) -> usize; + fn intermediate_size(&self) -> usize; + fn num_attention_heads(&self) -> usize; + fn num_hidden_layers(&self) -> usize; + fn layer_norm_eps(&self) -> f64; + fn hidden_act(&self) -> candle_nn::Activation; +} + +impl TransformerConfig for TextConfig { + fn hidden_size(&self) -> usize { + self.hidden_size + } + fn intermediate_size(&self) -> usize { + self.intermediate_size + } + fn num_attention_heads(&self) -> usize { + self.num_attention_heads + } + fn num_hidden_layers(&self) -> usize { + self.num_hidden_layers + } + fn layer_norm_eps(&self) -> f64 { + self.layer_norm_eps + } + fn hidden_act(&self) -> candle_nn::Activation { + self.hidden_act + } +} + +impl TransformerConfig for VisionConfig { + fn hidden_size(&self) -> usize { + self.hidden_size + } + fn intermediate_size(&self) -> usize { + self.intermediate_size + } + fn num_attention_heads(&self) -> usize { + self.num_attention_heads + } + fn num_hidden_layers(&self) -> usize { + self.num_hidden_layers + } + fn layer_norm_eps(&self) -> f64 { + self.layer_norm_eps + } + fn hidden_act(&self) -> candle_nn::Activation { + self.hidden_act + } +} + +impl VisionConfig { + pub fn paligemma_3b_224() -> Self { + Self { + // https://huggingface.co/google/paligemma-3b-pt-224/blob/main/config.json + patch_size: 14, + num_attention_heads: 16, + num_hidden_layers: 27, + hidden_size: 1152, + intermediate_size: 4304, + image_size: 224, // num_image_tokens: (224 / 14)^2 = 256 + // Default values. + num_channels: 3, + hidden_act: candle_nn::Activation::GeluPytorchTanh, + layer_norm_eps: 1e-6, + } + } + + pub fn paligemma_3b_448() -> Self { + Self { + // https://huggingface.co/google/paligemma-3b-pt-448/blob/main/config.json + patch_size: 14, + num_attention_heads: 16, + num_hidden_layers: 27, + hidden_size: 1152, + intermediate_size: 4304, + image_size: 448, // num_image_tokens: (448 / 14)^2 = 1024 + // Default values. + num_channels: 3, + hidden_act: candle_nn::Activation::GeluPytorchTanh, + layer_norm_eps: 1e-6, + } + } + + pub fn paligemma_3b_896() -> Self { + Self { + // https://huggingface.co/google/paligemma-3b-pt-448/blob/main/config.json + patch_size: 14, + num_attention_heads: 16, + num_hidden_layers: 27, + hidden_size: 1152, + intermediate_size: 4304, + image_size: 896, // num_image_tokens: (896 / 14)^2 = 4096 + // Default values. + num_channels: 3, + hidden_act: candle_nn::Activation::GeluPytorchTanh, + layer_norm_eps: 1e-6, + } + } + + pub fn num_patches(&self) -> usize { + (self.image_size / self.patch_size).pow(2) + } +} + +// https://github.com/huggingface/transformers/blob/2e24ee4dfa39cc0bc264b89edbccc373c8337086/src/transformers/models/siglip/configuration_siglip.py#L228 +#[derive(serde::Deserialize, Clone, Debug)] +pub struct Config { + pub text_config: TextConfig, + pub vision_config: VisionConfig, +} + +impl Config { + pub fn base_patch16_224() -> Self { + let text_config = TextConfig { + // https://huggingface.co/google/siglip-base-patch16-224/blob/main/config.json + hidden_size: 768, + intermediate_size: 3072, + num_attention_heads: 12, + vocab_size: 32000, + // Default values. + pad_token_id: 1, + bos_token_id: 49406, + eos_token_id: 49407, + layer_norm_eps: 1e-6, + hidden_act: candle_nn::Activation::GeluPytorchTanh, + max_position_embeddings: 64, + num_hidden_layers: 12, + }; + let vision_config = VisionConfig { + patch_size: 16, + // Default values. + hidden_size: 768, + intermediate_size: 3072, + num_hidden_layers: 12, + num_attention_heads: 12, + num_channels: 3, + image_size: 224, + hidden_act: candle_nn::Activation::GeluPytorchTanh, + layer_norm_eps: 1e-6, + }; + Self { + text_config, + vision_config, + } + } +} + +#[derive(Clone, Debug)] +struct MultiheadAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + out_proj: Linear, + num_heads: usize, +} + +impl MultiheadAttention { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let h = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let w_in_proj = vb.get((3 * h, h), "in_proj_weight")?.chunk(3, 0)?; + let b_in_proj = vb.get(3 * h, "in_proj_bias")?.chunk(3, 0)?; + let q_proj = Linear::new(w_in_proj[0].clone(), Some(b_in_proj[0].clone())); + let k_proj = Linear::new(w_in_proj[1].clone(), Some(b_in_proj[1].clone())); + let v_proj = Linear::new(w_in_proj[2].clone(), Some(b_in_proj[2].clone())); + let out_proj = linear(h, h, vb.pp("out_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + out_proj, + num_heads, + }) + } + + fn separate_heads(&self, x: &Tensor) -> Result { + let (b, n, c) = x.dims3()?; + x.reshape((b, n, self.num_heads, c / self.num_heads))? + .transpose(1, 2)? + .contiguous() + } + + fn recombine_heads(&self, x: &Tensor) -> Result { + let (b, n_heads, n_tokens, c_per_head) = x.dims4()?; + x.transpose(1, 2)? + .reshape((b, n_tokens, n_heads * c_per_head)) + } + + fn forward(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result { + let q = self.q_proj.forward(&q.contiguous()?)?; + let k = self.k_proj.forward(&k.contiguous()?)?; + let v = self.v_proj.forward(&v.contiguous()?)?; + + let q = self.separate_heads(&q)?; + let k = self.separate_heads(&k)?; + let v = self.separate_heads(&v)?; + + let (_, _, _, c_per_head) = q.dims4()?; + let attn = (q.matmul(&k.t()?)? / (c_per_head as f64).sqrt())?; + let attn = candle_nn::ops::softmax_last_dim(&attn)?; + + let out = attn.matmul(&v)?; + self.recombine_heads(&out)?.apply(&self.out_proj) + } +} + +#[derive(Debug, Clone)] +struct MultiheadAttentionPoolingHead { + probe: Tensor, + attention: MultiheadAttention, + layernorm: LayerNorm, + mlp: Mlp, +} + +impl MultiheadAttentionPoolingHead { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let mlp = Mlp::new(cfg, vb.pp("mlp"))?; + let layernorm = layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("layernorm"))?; + let probe = vb.get((1, 1, cfg.hidden_size), "probe")?; + let attention = MultiheadAttention::new(cfg, vb.pp("attention"))?; + Ok(Self { + probe, + attention, + layernorm, + mlp, + }) + } +} + +impl Module for MultiheadAttentionPoolingHead { + fn forward(&self, xs: &Tensor) -> Result { + let batch_size = xs.dim(0)?; + let probe = self.probe.repeat((batch_size, 1, 1))?; + let xs = self.attention.forward(&probe, xs, xs)?; + let residual = &xs; + let xs = xs.apply(&self.layernorm)?.apply(&self.mlp)?; + (xs + residual)?.i((.., 0)) + } +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + out_proj: Linear, + num_heads: usize, + head_dim: usize, + scale: f64, +} + +impl Attention { + fn new(cfg: &C, vb: VarBuilder) -> Result { + let embed_dim = cfg.hidden_size(); + let q_proj = linear(embed_dim, embed_dim, vb.pp("q_proj"))?; + let k_proj = linear(embed_dim, embed_dim, vb.pp("k_proj"))?; + let v_proj = linear(embed_dim, embed_dim, vb.pp("v_proj"))?; + let out_proj = linear(embed_dim, embed_dim, vb.pp("out_proj"))?; + let num_heads = cfg.num_attention_heads(); + let head_dim = embed_dim / num_heads; + Ok(Self { + q_proj, + k_proj, + v_proj, + out_proj, + num_heads, + head_dim, + scale: (head_dim as f64).powf(-0.5), + }) + } + + fn forward(&self, xs: &Tensor, attention_mask: Option<&Tensor>) -> Result { + let (batch_size, q_len, _) = xs.dims3()?; + let query_states = xs.apply(&self.q_proj)?; + let key_states = xs.apply(&self.k_proj)?; + let value_states = xs.apply(&self.v_proj)?; + + let shape = (batch_size, q_len, self.num_heads, self.head_dim); + let query_states = query_states.reshape(shape)?.transpose(1, 2)?.contiguous()?; + let key_states = key_states.reshape(shape)?.transpose(1, 2)?.contiguous()?; + let value_states = value_states.reshape(shape)?.transpose(1, 2)?.contiguous()?; + + let attn_weights = (query_states.matmul(&key_states.t()?)? * self.scale)?; + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + // The original implementation upcasts to f32 but candle_nn::ops::softmax should handle this properly. + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + let attn_outputs = attn_weights + .matmul(&value_states)? + .transpose(1, 2)? + .reshape((batch_size, q_len, ()))? + .apply(&self.out_proj)?; + Ok(attn_outputs) + } +} + +// https://github.com/huggingface/transformers/blob/2e24ee4dfa39cc0bc264b89edbccc373c8337086/src/transformers/models/siglip/modeling_siglip.py#L599 +#[derive(Debug, Clone)] +struct Mlp { + fc1: Linear, + fc2: Linear, + activation_fn: candle_nn::Activation, +} + +impl Mlp { + fn new(cfg: &C, vb: VarBuilder) -> Result { + let hidden_size = cfg.hidden_size(); + let intermediate_size = cfg.intermediate_size(); + let fc1 = candle_nn::linear(hidden_size, intermediate_size, vb.pp("fc1"))?; + let fc2 = candle_nn::linear(intermediate_size, hidden_size, vb.pp("fc2"))?; + Ok(Self { + fc1, + fc2, + activation_fn: cfg.hidden_act(), + }) + } +} + +impl Module for Mlp { + fn forward(&self, xs: &candle::Tensor) -> Result { + xs.apply(&self.fc1)? + .apply(&self.activation_fn)? + .apply(&self.fc2) + } +} + +// https://github.com/huggingface/transformers/blob/2e24ee4dfa39cc0bc264b89edbccc373c8337086/src/transformers/models/siglip/modeling_siglip.py#L614 +#[derive(Debug, Clone)] +struct EncoderLayer { + self_attn: Attention, + layer_norm1: LayerNorm, + mlp: Mlp, + layer_norm2: LayerNorm, +} + +impl EncoderLayer { + fn new(cfg: &C, vb: VarBuilder) -> Result { + let hidden_size = cfg.hidden_size(); + let layer_norm_eps = cfg.layer_norm_eps(); + let self_attn = Attention::new(cfg, vb.pp("self_attn"))?; + let layer_norm1 = layer_norm(hidden_size, layer_norm_eps, vb.pp("layer_norm1"))?; + let mlp = Mlp::new(cfg, vb.pp("mlp"))?; + let layer_norm2 = layer_norm(hidden_size, layer_norm_eps, vb.pp("layer_norm2"))?; + Ok(Self { + self_attn, + layer_norm1, + mlp, + layer_norm2, + }) + } + + fn forward(&self, xs: &Tensor, attention_mask: Option<&Tensor>) -> Result { + let residual = xs; + let xs = xs.apply(&self.layer_norm1)?; + let xs = self.self_attn.forward(&xs, attention_mask)?; + let xs = (residual + xs)?; + let residual = &xs; + let xs = xs.apply(&self.layer_norm2)?.apply(&self.mlp)?; + let xs = (xs + residual)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct Encoder { + layers: Vec, +} + +impl Encoder { + fn new(cfg: &C, vb: VarBuilder) -> Result { + let mut layers = vec![]; + let vb = vb.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers() { + let layer = EncoderLayer::new(cfg, vb.pp(layer_idx))?; + layers.push(layer) + } + Ok(Self { layers }) + } + + fn forward(&self, xs: &Tensor, attention_mask: Option<&Tensor>) -> Result { + let mut xs = xs.clone(); + for layer in self.layers.iter() { + xs = layer.forward(&xs, attention_mask)? + } + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct VisionEmbeddings { + patch_embedding: candle_nn::Conv2d, + position_embedding: Tensor, + patch_size: usize, + base_num_patches_per_side: usize, +} + +impl VisionEmbeddings { + fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result { + let conv2d_cfg = candle_nn::Conv2dConfig { + stride: cfg.patch_size, + ..Default::default() + }; + let patch_embedding = candle_nn::conv2d( + cfg.num_channels, + cfg.hidden_size, + cfg.patch_size, + conv2d_cfg, + vb.pp("patch_embedding"), + )?; + let num_patches_per_side = cfg.image_size / cfg.patch_size; + let embedder = candle_nn::embedding( + num_patches_per_side.pow(2), + cfg.hidden_size(), + vb.pp("position_embedding"), + )?; + let position_embedding = embedder.embeddings(); + let position_embedding = position_embedding + .reshape(( + 1, + num_patches_per_side, + num_patches_per_side, + cfg.hidden_size(), + ))? + .permute((0, 3, 1, 2))?; + Ok(Self { + patch_embedding, + position_embedding, + patch_size: cfg.patch_size, + base_num_patches_per_side: num_patches_per_side, + }) + } +} + +impl Module for VisionEmbeddings { + fn forward(&self, xs: &Tensor) -> Result { + //embed tokens + let (_batch, _channels, _height, _width) = xs.dims4()?; + let embeddings = xs.apply(&self.patch_embedding)?; + // interpolate position embeddings for the current image size (if needed) + let num_patches_h = _height / self.patch_size; + let num_patches_w = _width / self.patch_size; + let resized_position_embedding = if num_patches_w == self.base_num_patches_per_side + && num_patches_h == self.base_num_patches_per_side + { + self.position_embedding.clone() + } else { + self.position_embedding + .interpolate2d(num_patches_h, num_patches_w)? + }; + // Add position embeddings to tokens and flatten from 2D patches to 1D sequence + let embeddings = embeddings + .broadcast_add(&resized_position_embedding)? + .flatten_from(2)? + .transpose(1, 2)?; + Ok(embeddings) + } +} + +#[derive(Debug, Clone)] +struct VisionTransformer { + embeddings: VisionEmbeddings, + encoder: Encoder, + post_layernorm: LayerNorm, + head: Option, +} + +impl VisionTransformer { + fn new(cfg: &VisionConfig, use_head: bool, vb: VarBuilder) -> Result { + let embeddings = VisionEmbeddings::new(cfg, vb.pp("embeddings"))?; + let encoder = Encoder::new(cfg, vb.pp("encoder"))?; + let post_layernorm = + layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("post_layernorm"))?; + let head = if use_head { + Some(MultiheadAttentionPoolingHead::new(cfg, vb.pp("head"))?) + } else { + None + }; + Ok(Self { + embeddings, + encoder, + post_layernorm, + head, + }) + } +} + +impl Module for VisionTransformer { + fn forward(&self, xs: &Tensor) -> Result { + let xs = xs.apply(&self.embeddings)?; + let xs = self.encoder.forward(&xs, None)?; + let xs = xs.apply(&self.post_layernorm)?; + match self.head.as_ref() { + None => Ok(xs), + Some(h) => xs.apply(h), + } + } +} + +#[derive(Debug, Clone)] +pub struct VisionModel { + vision_model: VisionTransformer, +} + +impl VisionModel { + pub fn new(cfg: &VisionConfig, use_head: bool, vb: VarBuilder) -> Result { + let vision_model = VisionTransformer::new(cfg, use_head, vb)?; + Ok(Self { vision_model }) + } +} + +impl Module for VisionModel { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.vision_model) + } +} + +#[derive(Debug, Clone)] +struct TextEmbeddings { + token_embedding: candle_nn::Embedding, + position_embedding: candle_nn::Embedding, + position_ids: Tensor, +} + +impl TextEmbeddings { + fn new(cfg: &TextConfig, vb: VarBuilder) -> Result { + let token_embedding = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("token_embedding"))?; + let position_embedding = candle_nn::embedding( + cfg.max_position_embeddings, + cfg.hidden_size, + vb.pp("position_embedding"), + )?; + let position_ids = + Tensor::arange(0u32, cfg.max_position_embeddings as u32, vb.device())?.unsqueeze(0)?; + Ok(Self { + token_embedding, + position_embedding, + position_ids, + }) + } +} + +impl Module for TextEmbeddings { + fn forward(&self, input_ids: &Tensor) -> Result { + let seq_length = input_ids.dim(D::Minus1)?; + let inputs_embeds = self.token_embedding.forward(input_ids)?; + let position_ids = self.position_ids.narrow(1, 0, seq_length)?; + let position_embedding = self.position_embedding.forward(&position_ids)?; + inputs_embeds.broadcast_add(&position_embedding) + } +} + +#[derive(Debug, Clone)] +pub struct TextTransformer { + embeddings: TextEmbeddings, + encoder: Encoder, + final_layer_norm: LayerNorm, + pub head: Linear, +} + +impl TextTransformer { + fn new(cfg: &TextConfig, vb: VarBuilder) -> Result { + let embeddings = TextEmbeddings::new(cfg, vb.pp("embeddings"))?; + let encoder = Encoder::new(cfg, vb.pp("encoder"))?; + let final_layer_norm = layer_norm( + cfg.hidden_size, + cfg.layer_norm_eps, + vb.pp("final_layer_norm"), + )?; + let head = linear(cfg.hidden_size, cfg.hidden_size, vb.pp("head"))?; + Ok(Self { + embeddings, + encoder, + final_layer_norm, + head, + }) + } +} +impl Module for TextTransformer { + fn forward(&self, input_ids: &Tensor) -> Result { + let (_bsz, seq_len) = input_ids.dims2()?; + let input_ids = self.embeddings.forward(input_ids)?; + let input_ids = self.encoder.forward(&input_ids, None)?; + let last_hidden_state = self.final_layer_norm.forward(&input_ids)?; + last_hidden_state + .i((.., seq_len - 1, ..))? + .contiguous()? + .apply(&self.head) + } +} + +#[derive(Debug, Clone)] +pub struct TextModel { + pub text_model: TextTransformer, +} + +impl TextModel { + pub fn new(cfg: &TextConfig, vb: VarBuilder) -> Result { + let text_model = TextTransformer::new(cfg, vb)?; + Ok(Self { text_model }) + } +} + +impl Module for TextModel { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.text_model) + } +} + +#[derive(Clone, Debug)] +pub struct Model { + text_model: TextModel, + vision_model: VisionModel, + logit_bias: Tensor, + logit_scale: Tensor, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let text_model = TextModel::new(&cfg.text_config, vb.pp("text_model"))?; + let vision_model = VisionModel::new(&cfg.vision_config, true, vb.pp("vision_model"))?; + let logit_scale = vb.get(&[1], "logit_scale")?; + let logit_bias = vb.get(&[1], "logit_bias")?; + Ok(Self { + text_model, + vision_model, + logit_bias, + logit_scale, + }) + } + + pub fn get_text_features(&self, input_ids: &Tensor) -> Result { + input_ids.apply(&self.text_model) + } + + pub fn get_image_features(&self, pixel_values: &Tensor) -> Result { + pixel_values.apply(&self.vision_model) + } + + pub fn forward(&self, pixel_values: &Tensor, input_ids: &Tensor) -> Result<(Tensor, Tensor)> { + let image_features = self.get_image_features(pixel_values)?; + let text_features = self.get_text_features(input_ids)?; + let image_features_normalized = div_l2_norm(&image_features)?; + let text_features_normalized = div_l2_norm(&text_features)?; + let logits_per_text = text_features_normalized.matmul(&image_features_normalized.t()?)?; + let logit_scale = self.logit_scale.exp()?; + let logits_per_text = logits_per_text + .broadcast_mul(&logit_scale)? + .broadcast_add(&self.logit_bias)?; + let logits_per_image = logits_per_text.t()?; + Ok((logits_per_text, logits_per_image)) + } +} diff --git a/patches/candle-transformers/src/models/smol/README.md b/patches/candle-transformers/src/models/smol/README.md new file mode 100644 index 0000000000..5a9e260c9b --- /dev/null +++ b/patches/candle-transformers/src/models/smol/README.md @@ -0,0 +1,259 @@ +# SmolLM Model Family + +This directory contains implementations for the SmolLM family of models +developed by HuggingFace. + +## Models + +### SmolLM2 (see `models/llama`) +SmolLM2 models (135M, 360M, 1.7B) use the standard Llama3 architecture +and are implemented in `models/llama.rs`. No separate implementation +is needed. + +**Variants:** +- HuggingFaceTB/SmolLM2-135M +- HuggingFaceTB/SmolLM2-360M +- HuggingFaceTB/SmolLM2-1.7B + +### SmolLM3 +SmolLM3-3B introduces NoPE (No Positional Encoding) which requires +a custom implementation in `smollm3.rs`. + +**Key innovations:** +- Hybrid RoPE/NoPE (3:1 ratio - every 4th layer uses NoPE) +- GQA with 4 groups (32 attention heads, 8 KV heads) +- Very high rope_theta (5M vs typical 10k-500k) +- Long context support (64k-128k tokens) +- Thinking mode support with `` tags + +**Implementations:** +- `smollm3.rs` - Full precision model (safetensors) +- `quantized_smollm3.rs` - Quantized GGUF model with weight reconstruction + +**Available Models:** +- HuggingFaceTB/SmolLM3-3B (Instruct-tuned) +- HuggingFaceTB/SmolLM3-3B-Base (Base model) +- unsloth/SmolLM3-3B-GGUF (Quantized: Q4_K_M, Q8_0, F16) + +### SmolVLM (planned) +Vision-language model variant, to be implemented. + +## Implementation Details + +### NoPE Architecture +SmolLM3 uses a mixed approach to positional encoding: +```rust +pub fn should_skip_rope(&self, layer_idx: usize) -> bool { + // Method 1: Explicit array from config + if let Some(ref no_rope_layers) = self.no_rope_layers { + if layer_idx < no_rope_layers.len() { + return no_rope_layers[layer_idx] == 0; + } + } + + // Method 2: Interval pattern (SmolLM3-3B default) + // Every 4th layer (indices 3, 7, 11, ...) skips RoPE + if let Some(interval) = self.no_rope_layer_interval { + return (layer_idx + 1) % interval == 0; + } + + false // Default: use RoPE +} +``` + +### Quantized Weight Reconstruction +The quantized implementation includes special handling for Q/K weight +reconstruction to maintain compatibility with the GGUF format's +interleaved weight storage. + +### Thinking Mode +SmolLM3 supports explicit reasoning with thinking tags: +- **Enabled**: `<|im_start|>assistant\n\n` (model generates reasoning) +- **Disabled**: `<|im_start|>assistant\n\n\n\n` (skip to answer) + +## Usage Example + +See `examples/smollm3/main.rs` for a unified implementation that supports +both quantized and full precision models with a single codebase. + +```bash +# Quantized model (recommended) +cargo run --release --example smollm3 -- \ + --model-type quantized \ + --quantization q8_0 \ + --prompt "Explain Rust's ownership system" + +# Full precision model +cargo run --release --example smollm3 -- \ + --model-type full \ + --dtype f16 \ + --prompt "Write a sorting algorithm" + +# Enable thinking mode +cargo run --release --example smollm3 -- \ + --thinking \ + --prompt "Solve this logic puzzle step by step" +``` + +## Performance Characteristics + +| Model Type | Size | Speed | Quality | Use Case | +|------------|-------|-------|---------|----------| +| Q4_K_M | 1.9GB | Fast | Good | Resource-constrained | +| Q8_0 | 3.3GB | Fast | Better | Balanced | +| F16 (GGUF) | 6.2GB | Med | Best | High quality GGUF | +| F16 (Safe) | 6.2GB | Med | Best | Maximum quality | +| F32 (Safe) | 12GB | Slow | Best | Research/debugging | + +# Credits & Attribution + +## SmolLM3 Model + +### Developers +**HuggingFace Team (HuggingFaceTB)** + +The SmolLM family of models represents cutting-edge work in efficient language models, demonstrating that small models can achieve impressive capabilities when trained on high-quality data. + +### Resources +- **Model Card**: https://huggingface.co/HuggingFaceTB/SmolLM3-3B +- **Model Card (Base)**: https://huggingface.co/HuggingFaceTB/SmolLM3-3B-Base +- **Collection**: https://huggingface.co/collections/HuggingFaceTB/smollm3-6723884a9c35673e4f9b74a2 +- **Blog Post**: https://huggingface.co/blog/smollm3 +- **GitHub Repository**: https://github.com/huggingface/smollm +- **License**: Apache 2.0 + +### Key Contributors +The SmolLM project is developed by the HuggingFace team with contributions from researchers focused on efficient LLM architectures and training methods. + +## NoPE Architecture + +### Research Paper +**Title**: "Length Generalization of Causal Transformers without Position Encoding" + +**Authors**: +- Jie Wang (Fudan University) +- Tao Ji (Fudan University) +- Yuanbin Wu (Fudan University) +- Hang Yan (Fudan University) +- Tao Gui (Fudan University) +- Qi Zhang (Fudan University) +- Xuanjing Huang (Fudan University) +- Xiaoling Wang (Fudan University) + +**Published**: NeurIPS 2024 (Thirty-Eighth Annual Conference on Neural Information Processing Systems) + +**Abstract Summary**: The paper demonstrates that removing positional encoding from selected layers (NoPE - No Positional Encoding) can improve length generalization in causal transformers while maintaining or improving performance. SmolLM3 implements this with a 3:1 RoPE/NoPE ratio. + +**Resources**: +- **arXiv**: https://arxiv.org/abs/2410.01926 +- **Conference**: NeurIPS 2024 + +### Key Innovation +The hybrid approach uses: +- **RoPE layers** (75%): Standard rotary positional embeddings for local context +- **NoPE layers** (25%): No positional encoding for improved length generalization +- **Pattern**: Every 4th layer uses NoPE (layers 3, 7, 11, 15, etc.) + +This architecture enables SmolLM3 to handle much longer contexts (64k-128k tokens) while maintaining efficiency. + +## Quantized Models + +### Unsloth +Quantized GGUF models are provided by **Unsloth**, a team focused on making LLM inference and fine-tuning more accessible. + +**Resources**: +- **GGUF Repository**: https://huggingface.co/unsloth/SmolLM3-3B-GGUF +- **Available Quantizations**: Q4_K_M, Q8_0, F16 +- **Website**: https://unsloth.ai/ + +The quantization work enables running SmolLM3 efficiently on consumer hardware with minimal quality loss. + +## Implementation Credits + +### This Candle Implementation +**Implemented for**: Candle ML Framework +**Implementation Date**: Nov 2025 +**Features**: +- Full precision model (F32/F16/BF16) +- Quantized model (Q4_K_M/Q8_0/F16 GGUF) +- Unified example supporting both +- Verified against reference implementations + +**Verification**: +- Full precision: Validated against HuggingFace Transformers Python implementation +- Quantized: Validated against llama.cpp implementation + +### Related Tools & Frameworks + +**Candle**: Minimalist ML framework in Rust by HuggingFace +- GitHub: https://github.com/huggingface/candle + +**llama.cpp**: Efficient LLM inference in C/C++ +- GitHub: https://github.com/ggerganov/llama.cpp +- Used for quantized model verification + +**HuggingFace Transformers**: Reference Python implementation +- GitHub: https://github.com/huggingface/transformers +- Used for full model verification + +## Acknowledgments + +Special thanks to: + +1. **HuggingFace Team** - For developing SmolLM3 and making it openly available under Apache 2.0 license +2. **NoPE Researchers** - For advancing the field with novel positional encoding approaches +3. **Unsloth** - For providing optimized quantized versions +4. **Candle Contributors** - For building an excellent ML framework in Rust +5. **Open Source Community** - For tools like llama.cpp that enable verification and benchmarking + +## Citation + +If you use SmolLM3 in your research or applications, please cite: + +### SmolLM3 Model +```bibtex +@misc{smollm3, + title={SmolLM3}, + author={HuggingFace Team}, + year={2024}, + publisher={HuggingFace}, + howpublished={\url{https://huggingface.co/HuggingFaceTB/SmolLM3-3B}} +} +``` + +### NoPE Paper +```bibtex +@inproceedings{wang2024length, + title={Length Generalization of Causal Transformers without Position Encoding}, + author={Wang, Jie and Ji, Tao and Wu, Yuanbin and Yan, Hang and Gui, Tao and Zhang, Qi and Huang, Xuanjing and Wang, Xiaoling}, + booktitle={Thirty-Eighth Annual Conference on Neural Information Processing Systems}, + year={2024} +} +``` + +### Candle Framework +```bibtex +@software{candle, + title={Candle: Minimalist ML Framework}, + author={HuggingFace}, + year={2024}, + url={https://github.com/huggingface/candle} +} +``` + +## License + +- **SmolLM3 Model**: Apache 2.0 +- **This Implementation**: Follows Candle framework license +- **Candle Framework**: Apache 2.0 and MIT dual-licensed + +## Further Reading + +- **SmolLM Blog Series**: https://huggingface.co/blog/smollm and https://huggingface.co/blog/smollm3 +- **Model Card Details**: https://huggingface.co/HuggingFaceTB/SmolLM3-3B +- **NoPE Paper**: https://arxiv.org/abs/2410.01926 +- **Candle Documentation**: https://huggingface.github.io/candle/ + +--- + +This implementation stands on the shoulders of giants. Thank you to all the researchers, engineers, and open source contributors who make this work possible. diff --git a/patches/candle-transformers/src/models/smol/mod.rs b/patches/candle-transformers/src/models/smol/mod.rs new file mode 100644 index 0000000000..c3900c43ed --- /dev/null +++ b/patches/candle-transformers/src/models/smol/mod.rs @@ -0,0 +1,67 @@ +//! SmolLM model family implementations. +//! +//! The SmolLM family consists of efficient language models developed by HuggingFace: +//! - **SmolLM2** (135M, 360M, 1.7B): Uses standard Llama architecture (see `models::llama`) +//! - **SmolLM3** (3B): Introduces hybrid RoPE/NoPE architecture (implemented here) +//! +//! # SmolLM3 Architecture +//! +//! SmolLM3-3B introduces NoPE (No Positional Encoding) as a key innovation: +//! - 3:1 RoPE/NoPE ratio: every 4th layer skips positional encoding +//! - Grouped Query Attention: 32 attention heads, 8 KV heads (4 groups) +//! - High RoPE theta: 5,000,000 (vs typical 10,000-500,000) +//! - Extended context: 64k-128k tokens +//! +//! # Module Structure +//! +//! - [`smollm3`]: Full precision model implementation (safetensors) +//! - [`quantized_smollm3`]: Quantized model implementation (GGUF) +//! +//! # Example Usage +//! +//! ```ignore +//! use candle_transformers::models::smol::smollm3::{Config, ModelForCausalLM}; +//! use candle_transformers::models::smol::quantized_smollm3::QuantizedModelForCausalLM; +//! use candle::{Device, Tensor}; +//! use candle_nn::VarBuilder; +//! +//! # fn main() -> anyhow::Result<()> { +//! let device = Device::Cpu; +//! +//! // Load full precision model +//! let vb = VarBuilder::zeros(candle::DType::F32, &device); +//! let config = Config::default(); +//! let model = ModelForCausalLM::new(&config, vb)?; +//! +//! // Or load quantized model +//! // let model = QuantizedModelForCausalLM::from_gguf(path, &device)?; +//! +//! // Run inference +//! let input = Tensor::new(&[1u32, 2, 3], &device)?.unsqueeze(0)?; +//! let logits = model.forward(&input, 0)?; +//! # Ok(()) +//! # } +//! ``` +//! +//! # Thinking Mode +//! +//! SmolLM3 supports explicit reasoning via thinking tags in chat templates: +//! - Thinking enabled: `<|im_start|>assistant\n\n` (model generates reasoning) +//! - Thinking disabled: `<|im_start|>assistant\n\n\n\n` (skip to answer) +//! +//! # Performance Considerations +//! +//! | Format | Size | Inference Speed | Quality | +//! |--------|-------|-----------------|---------| +//! | Q4_K_M | 1.9GB | Fastest | Good | +//! | Q8_0 | 3.3GB | Fast | Better | +//! | F16 | 6.2GB | Medium | Best | +//! | F32 | 12GB | Slow | Best | +//! +//! # References +//! +//! - [SmolLM3 Model Card](https://huggingface.co/HuggingFaceTB/SmolLM3-3B) +//! - [NoPE Paper](https://arxiv.org/abs/2410.01926) + +pub mod quantized_smollm3; +pub mod smollm3; diff --git a/patches/candle-transformers/src/models/smol/quantized_smollm3.rs b/patches/candle-transformers/src/models/smol/quantized_smollm3.rs new file mode 100644 index 0000000000..de4f1f318a --- /dev/null +++ b/patches/candle-transformers/src/models/smol/quantized_smollm3.rs @@ -0,0 +1,567 @@ +use crate::models::with_tracing::QMatMul; +use crate::quantized_var_builder::VarBuilder; +use candle::quantized::gguf_file; +use candle::{DType, Device, Module, Result, Tensor}; +use candle_nn::kv_cache::KvCache; +use candle_nn::Activation; +use std::io::Write; +use std::sync::Arc; + +const MAX_SEQ_LEN: usize = 4096; +use candle::IndexOp; + +// ===== RECONSTRUCTION FUNCTION ===== +fn reconstruct_qk_weights(gguf_weight: &Tensor, _num_heads: usize) -> Result { + let total_rows = gguf_weight.dim(0)?; + let half_rows = total_rows / 2; + let chunk_size = 128; + let chunks_per_half = half_rows / chunk_size; + + let mut heads = Vec::new(); + + // First half + for chunk_idx in 0..chunks_per_half { + let chunk_start = chunk_idx * chunk_size; + + // Even rows + let mut head_even = Vec::new(); + for i in (chunk_start..chunk_start + chunk_size).step_by(2) { + head_even.push(gguf_weight.i(i)?); + } + heads.push(Tensor::stack(&head_even, 0)?); + + // Odd rows + let mut head_odd = Vec::new(); + for i in (chunk_start + 1..chunk_start + chunk_size).step_by(2) { + head_odd.push(gguf_weight.i(i)?); + } + heads.push(Tensor::stack(&head_odd, 0)?); + } + + // Second half + for chunk_idx in 0..chunks_per_half { + let chunk_start = half_rows + chunk_idx * chunk_size; + + // Even rows + let mut head_even = Vec::new(); + for i in (chunk_start..chunk_start + chunk_size).step_by(2) { + head_even.push(gguf_weight.i(i)?); + } + heads.push(Tensor::stack(&head_even, 0)?); + + // Odd rows + let mut head_odd = Vec::new(); + for i in (chunk_start + 1..chunk_start + chunk_size).step_by(2) { + head_odd.push(gguf_weight.i(i)?); + } + heads.push(Tensor::stack(&head_odd, 0)?); + } + + Tensor::cat(&heads, 0) +} + +#[derive(Debug, Clone)] +pub struct QuantizedConfig { + pub vocab_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub max_position_embeddings: usize, + pub rope_theta: f64, + pub rms_norm_eps: f64, + pub rope_dimension_count: usize, + pub no_rope_layer_interval: Option, +} + +impl QuantizedConfig { + /// Load config from GGUF metadata + pub fn from_gguf(ct: &gguf_file::Content) -> Result { + let metadata = &ct.metadata; + + // Helper to get required metadata + let get_u32 = |key: &str| -> Result { + metadata + .get(key) + .and_then(|v| v.to_u32().ok()) + .map(|v| v as usize) + .ok_or_else(|| { + candle::Error::Msg(format!("Missing or invalid metadata key: {}", key)) + }) + }; + + let get_f32 = |key: &str| -> Result { + metadata + .get(key) + .and_then(|v| v.to_f32().ok()) + .map(|v| v as f64) + .ok_or_else(|| { + candle::Error::Msg(format!("Missing or invalid metadata key: {}", key)) + }) + }; + + Ok(Self { + vocab_size: get_u32("smollm3.vocab_size")?, + hidden_size: get_u32("smollm3.embedding_length")?, + intermediate_size: get_u32("smollm3.feed_forward_length")?, + num_hidden_layers: get_u32("smollm3.block_count")?, + num_attention_heads: get_u32("smollm3.attention.head_count")?, + num_key_value_heads: get_u32("smollm3.attention.head_count_kv")?, + max_position_embeddings: get_u32("smollm3.context_length").unwrap_or(MAX_SEQ_LEN), + rope_theta: get_f32("smollm3.rope.freq_base")?, + rms_norm_eps: get_f32("smollm3.attention.layer_norm_rms_epsilon")?, + rope_dimension_count: get_u32("smollm3.rope.dimension_count")?, + no_rope_layer_interval: Some(4), + }) + } + + pub fn should_skip_rope(&self, layer_idx: usize) -> bool { + if let Some(interval) = self.no_rope_layer_interval { + return (layer_idx + 1).is_multiple_of(interval); + } + false + } + + pub fn head_dim(&self) -> usize { + self.rope_dimension_count + } +} + +#[derive(Debug, Clone)] +struct RmsNorm { + weight: Tensor, + eps: f64, +} + +impl RmsNorm { + fn new(weight: Tensor, eps: f64) -> Self { + Self { weight, eps } + } + + fn forward(&self, x: &Tensor) -> Result { + let x_dtype = x.dtype(); + let internal_dtype = match x_dtype { + DType::F16 | DType::BF16 => DType::F32, + d => d, + }; + let hidden_size = x.dim(candle::D::Minus1)?; + let x = x.to_dtype(internal_dtype)?; + let norm_x = (x.sqr()?.sum_keepdim(candle::D::Minus1)? / hidden_size as f64)?; + let x_normed = x.broadcast_div(&(norm_x + self.eps)?.sqrt()?)?; + let result = x_normed.broadcast_mul(&self.weight)?; + result.to_dtype(x_dtype) + } +} + +#[derive(Debug, Clone)] +pub struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + pub fn new(dtype: DType, cfg: &QuantizedConfig, dev: &Device) -> Result { + let dim = cfg.head_dim(); + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(DType::F32)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(DType::F32)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?.to_dtype(dtype)?, + cos: freqs.cos()?.to_dtype(dtype)?, + }) + } + + pub fn apply_rotary_emb( + &self, + q: &Tensor, + k: &Tensor, + offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_, _, seq_len, _) = q.dims4()?; + let cos = self.cos.narrow(0, offset, seq_len)?; + let sin = self.sin.narrow(0, offset, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +fn repeat_kv(x: Tensor, n_rep: usize) -> Result { + if n_rep == 1 { + Ok(x) + } else { + let (b, n_kv_heads, seq_len, head_dim) = x.dims4()?; + x.unsqueeze(2)? + .expand(&[b, n_kv_heads, n_rep, seq_len, head_dim])? + .reshape(&[b, n_kv_heads * n_rep, seq_len, head_dim]) + } +} + +#[derive(Debug, Clone)] +struct QuantizedMLP { + gate_proj: QMatMul, + up_proj: QMatMul, + down_proj: QMatMul, +} + +impl QuantizedMLP { + fn new(vb: VarBuilder, _layer_idx: usize) -> Result { + // VarBuilder.get_no_shape() returns Arc which QMatMul::from_weights expects + let gate_proj = QMatMul::from_weights(vb.get_no_shape("ffn_gate.weight")?)?; + let up_proj = QMatMul::from_weights(vb.get_no_shape("ffn_up.weight")?)?; + let down_proj = QMatMul::from_weights(vb.get_no_shape("ffn_down.weight")?)?; + + Ok(Self { + gate_proj, + up_proj, + down_proj, + }) + } + + fn forward(&self, x: &Tensor) -> Result { + let gate = self.gate_proj.forward(x)?.apply(&Activation::Silu)?; + let up = self.up_proj.forward(x)?; + self.down_proj.forward(&(gate * up)?) + } +} + +#[derive(Debug, Clone)] +struct QuantizedAttention { + q_proj: QMatMul, + k_proj: QMatMul, + v_proj: QMatMul, + o_proj: QMatMul, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + rotary_emb: Option>, + skip_rope: bool, + kv_cache: KvCache, +} + +impl QuantizedAttention { + fn new( + vb: VarBuilder, + cfg: &QuantizedConfig, + layer_idx: usize, + rotary_emb: Option>, + ) -> Result { + let head_dim = cfg.head_dim(); + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + + // For v and o weights, use directly from VarBuilder (already quantized) + // VarBuilder.get_no_shape() returns Arc + let v_proj = QMatMul::from_weights(vb.get_no_shape("attn_v.weight")?)?; + let o_proj = QMatMul::from_weights(vb.get_no_shape("attn_output.weight")?)?; + + // For q and k weights, we need to dequantize, reconstruct, then re-quantize + // IMPORTANT: Do reconstruction on CPU to avoid VRAM exhaustion during model loading + let device = vb.device(); + let cpu = Device::Cpu; + + let q_weight_qtensor = vb.get_no_shape("attn_q.weight")?; + let q_weight_raw = q_weight_qtensor.dequantize(&cpu)?; // Dequantize to CPU + let q_weight = reconstruct_qk_weights(&q_weight_raw, num_heads)?; // Reconstruct on CPU + let q_weight = q_weight.to_device(device)?; // Move to GPU + + // Re-quantize (now on GPU) + use candle::quantized::{GgmlDType, QTensor}; + let q_weight_qtensor = QTensor::quantize(&q_weight, GgmlDType::Q8_0)?; + drop(q_weight_raw); // Explicitly free CPU memory + drop(q_weight); + + let k_weight_qtensor = vb.get_no_shape("attn_k.weight")?; + let k_weight_raw = k_weight_qtensor.dequantize(&cpu)?; // Dequantize to CPU + let k_weight = reconstruct_qk_weights(&k_weight_raw, num_kv_heads)?; // Reconstruct on CPU + let k_weight = k_weight.to_device(device)?; // Move to GPU + + // Re-quantize (now on GPU) + let k_weight_qtensor = QTensor::quantize(&k_weight, GgmlDType::Q8_0)?; + drop(k_weight_raw); // Explicitly free CPU memory + drop(k_weight); + + let q_proj = QMatMul::from_weights(Arc::new(q_weight_qtensor))?; + let k_proj = QMatMul::from_weights(Arc::new(k_weight_qtensor))?; + + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups: num_heads / num_kv_heads, + head_dim, + rotary_emb, + skip_rope: cfg.should_skip_rope(layer_idx), + kv_cache: KvCache::new(2, 512), + }) + } + + fn forward(&mut self, x: &Tensor, mask: Option<&Tensor>, offset: usize) -> Result { + let (b, seq_len, _) = x.dims3()?; + + let q = self + .q_proj + .forward(x)? + .reshape((b, seq_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let k = self + .k_proj + .forward(x)? + .reshape((b, seq_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let v = self + .v_proj + .forward(x)? + .reshape((b, seq_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let (q, k) = if self.skip_rope { + (q, k) + } else if let Some(rope) = &self.rotary_emb { + rope.apply_rotary_emb(&q, &k, offset)? + } else { + (q, k) + }; + + // can remove this continguous call if using ConcatKV-Cache https://github.com/huggingface/candle/pull/3143 + let (k, v) = self.kv_cache.append(&k.contiguous()?, &v.contiguous()?)?; + + let k = repeat_kv(k, self.num_kv_groups)?; + let v = repeat_kv(v, self.num_kv_groups)?; + + let scale = 1.0 / (self.head_dim as f64).sqrt(); + // Make q contiguous before matmul to avoid stride mismatch + let q = q.contiguous()?; + let attn_weights = (q.matmul(&k.t()?)? * scale)?; + + let mut attn_weights = match mask { + Some(mask) => attn_weights.broadcast_add(mask)?, + None => attn_weights, + }; + + attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + let attn_output = attn_weights.matmul(&v)?; + + attn_output + .transpose(1, 2)? + .reshape((b, seq_len, self.num_heads * self.head_dim))? + .apply(&self.o_proj) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache.reset(); + } +} + +#[derive(Debug, Clone)] +struct QuantizedDecoderLayer { + self_attn: QuantizedAttention, + mlp: QuantizedMLP, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl QuantizedDecoderLayer { + fn new( + vb: VarBuilder, + cfg: &QuantizedConfig, + layer_idx: usize, + rotary_emb: Option>, + ) -> Result { + let attn_vb = vb.pp(format!("blk.{layer_idx}")); + + Ok(Self { + self_attn: QuantizedAttention::new(attn_vb.clone(), cfg, layer_idx, rotary_emb)?, + mlp: QuantizedMLP::new(attn_vb.clone(), layer_idx)?, + input_layernorm: RmsNorm::new( + attn_vb + .get_no_shape("attn_norm.weight")? + .dequantize(vb.device())?, + cfg.rms_norm_eps, + ), + post_attention_layernorm: RmsNorm::new( + attn_vb + .get_no_shape("ffn_norm.weight")? + .dequantize(vb.device())?, + cfg.rms_norm_eps, + ), + }) + } + + fn forward(&mut self, x: &Tensor, mask: Option<&Tensor>, offset: usize) -> Result { + let residual = x; + let x = self.input_layernorm.forward(x)?; + let x = self.self_attn.forward(&x, mask, offset)?; + let x = (residual + x)?; + + let residual = &x; + let x = self.post_attention_layernorm.forward(&x)?; + let x = self.mlp.forward(&x)?; + residual + x + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache(); + } +} + +#[derive(Debug, Clone)] +pub struct QuantizedModelForCausalLM { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: QMatMul, + device: Device, + config: QuantizedConfig, +} + +impl QuantizedModelForCausalLM { + pub fn from_gguf>(path: P, device: &Device) -> Result { + use candle::quantized::{GgmlDType, QTensor}; + + // Open file once to read metadata + let mut file = std::fs::File::open(path.as_ref())?; + let content = gguf_file::Content::read(&mut file)?; + let config = QuantizedConfig::from_gguf(&content)?; + + // Create VarBuilder for tensor loading + let vb = VarBuilder::from_gguf(path, device)?; + + // Load embedding tensor - dequantize on CPU first to save VRAM + // (will be used for both embed_tokens and lm_head - tied embeddings) + let cpu = Device::Cpu; + let embed_tensor = vb.get_no_shape("token_embd.weight")?.dequantize(&cpu)?; + let embed_tensor_gpu = embed_tensor.to_device(device)?; // Move to GPU for embedding layer + let embed_tokens = candle_nn::Embedding::new(embed_tensor_gpu, config.hidden_size); + + // Create rotary embedding if needed + let needs_rope = (0..config.num_hidden_layers).any(|i| !config.should_skip_rope(i)); + let rotary_emb = if needs_rope { + Some(Arc::new(RotaryEmbedding::new(DType::F32, &config, device)?)) + } else { + None + }; + + // Load decoder layers + let mut layers = Vec::with_capacity(config.num_hidden_layers); + println!("Loading {} decoder layers...", config.num_hidden_layers); + for layer_idx in 0..config.num_hidden_layers { + if layer_idx % 4 == 0 || layer_idx == config.num_hidden_layers - 1 { + print!( + " Layer {}/{}...\r", + layer_idx + 1, + config.num_hidden_layers + ); + std::io::stdout().flush().ok(); + } + layers.push(QuantizedDecoderLayer::new( + vb.clone(), + &config, + layer_idx, + rotary_emb.clone(), + )?); + } + println!( + " Layer {}/{} - Done! ", + config.num_hidden_layers, config.num_hidden_layers + ); + + // Load output norm + let norm = RmsNorm::new( + vb.get_no_shape("output_norm.weight")?.dequantize(device)?, + config.rms_norm_eps, + ); + + // Load LM head - move CPU embedding tensor to GPU, then quantize + let embed_tensor_for_lm = embed_tensor.to_device(device)?; + let embed_qtensor = QTensor::quantize(&embed_tensor_for_lm, GgmlDType::Q8_0)?; + let lm_head = QMatMul::from_weights(Arc::new(embed_qtensor))?; + drop(embed_tensor); // Free CPU memory + drop(embed_tensor_for_lm); + + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + device: device.clone(), + config, + }) + } + + pub fn forward(&mut self, input_ids: &Tensor, offset: usize) -> Result { + let (batch_size, seq_len) = input_ids.dims2()?; + + // Embed tokens + let mut hidden_states = self.embed_tokens.forward(input_ids)?; + + // Create causal mask if needed + let mask = if seq_len > 1 { + Some(self.create_causal_mask(batch_size, seq_len, offset)?) + } else { + None + }; + + // Forward through decoder layers + for layer in &mut self.layers { + hidden_states = layer.forward(&hidden_states, mask.as_ref(), offset)?; + } + + // Final norm + hidden_states = self.norm.forward(&hidden_states)?; + + // LM head (only last token for generation) + let last_hidden = hidden_states.narrow(1, seq_len - 1, 1)?; + let logits = last_hidden.apply(&self.lm_head)?; + + Ok(logits) + } + + fn create_causal_mask( + &self, + batch_size: usize, + tgt_len: usize, + offset: usize, + ) -> Result { + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| { + (0..tgt_len + offset).map(move |j| { + if j <= i + offset { + 0f32 + } else { + f32::NEG_INFINITY + } + }) + }) + .collect(); + + Tensor::from_slice( + &mask, + (batch_size, 1, tgt_len, tgt_len + offset), + &self.device, + ) + } + + pub fn clear_kv_cache(&mut self) { + for layer in &mut self.layers { + layer.clear_kv_cache(); + } + } + + pub fn config(&self) -> &QuantizedConfig { + &self.config + } +} diff --git a/patches/candle-transformers/src/models/smol/smollm3.rs b/patches/candle-transformers/src/models/smol/smollm3.rs new file mode 100644 index 0000000000..e2c7200d5f --- /dev/null +++ b/patches/candle-transformers/src/models/smol/smollm3.rs @@ -0,0 +1,470 @@ +use crate::{ + models::with_tracing::{linear_b, linear_no_bias, Linear, RmsNorm}, + utils::repeat_kv, +}; +use candle::{DType, Device, Module, Result, Tensor}; +use candle_nn::{kv_cache::KvCache, Activation, VarBuilder}; +use std::sync::Arc; + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub max_position_embeddings: usize, + pub tie_word_embeddings: bool, + pub rope_theta: f64, + pub rms_norm_eps: f64, + pub hidden_act: Activation, + // Optional fields + pub attention_bias: Option, + pub attention_dropout: Option, + pub mlp_bias: Option, + pub sliding_window: Option, + pub use_sliding_window: Option, + pub rope_scaling: Option, + pub bos_token_id: Option, + pub eos_token_id: Option, + pub pad_token_id: Option, + pub max_window_layers: Option, + // SmolLM3-specific: NoPE configuration + pub no_rope_layers: Option>, + pub no_rope_layer_interval: Option, +} + +impl Config { + pub fn should_skip_rope(&self, layer_idx: usize) -> bool { + // Method 1: Explicit array (some model variants may provide this) + if let Some(ref no_rope_layers) = self.no_rope_layers { + if layer_idx < no_rope_layers.len() { + // 0 = skip RoPE (NoPE), 1 = use RoPE + return no_rope_layers[layer_idx] == 0; + } + } + + // Method 2: Interval pattern (SmolLM3-3B uses this) + // With interval=4: layers 0,1,2 use RoPE; layer 3 skips RoPE (NoPE) + // Pattern: every 4th layer (3,7,11...) skips RoPE + if let Some(interval) = self.no_rope_layer_interval { + return (layer_idx + 1).is_multiple_of(interval); + } + + // Default: use RoPE on all layers (standard Llama behavior) + false + } + + /// Calculates head_dim from hidden_size and num_attention_heads + pub fn head_dim(&self) -> usize { + self.hidden_size / self.num_attention_heads + } +} + +#[derive(Debug, Clone)] +pub(crate) struct SmolLM3RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl SmolLM3RotaryEmbedding { + pub(crate) fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let dim = cfg.head_dim(); + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(DType::F32)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(DType::F32)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?.to_dtype(dtype)?, + cos: freqs.cos()?.to_dtype(dtype)?, + }) + } + + /// Apply RoPE (q, k shape: B x H x L x D) + pub(crate) fn apply(&self, q: &Tensor, k: &Tensor, offset: usize) -> Result<(Tensor, Tensor)> { + let (_, _, seq_len, _) = q.dims4()?; + let cos = self.cos.narrow(0, offset, seq_len)?; + let sin = self.sin.narrow(0, offset, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +pub(crate) struct SmolLM3MLP { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: Activation, +} + +impl SmolLM3MLP { + pub(crate) fn new(cfg: &Config, vb: VarBuilder) -> Result { + let mlp_bias = cfg.mlp_bias.unwrap_or(false); + Ok(Self { + gate_proj: linear_b( + cfg.hidden_size, + cfg.intermediate_size, + mlp_bias, + vb.pp("gate_proj"), + )?, + up_proj: linear_b( + cfg.hidden_size, + cfg.intermediate_size, + mlp_bias, + vb.pp("up_proj"), + )?, + down_proj: linear_b( + cfg.intermediate_size, + cfg.hidden_size, + mlp_bias, + vb.pp("down_proj"), + )?, + act_fn: cfg.hidden_act, + }) + } +} + +impl Module for SmolLM3MLP { + fn forward(&self, x: &Tensor) -> Result { + let lhs = x.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = x.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +pub(crate) struct SmolLM3Attention { + // projections + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + // hyper params + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + hidden_size: usize, + // utils + rotary_emb: Option>, + kv_cache: KvCache, + // NoPE flag + skip_rope: bool, +} + +impl SmolLM3Attention { + pub(crate) fn new( + cfg: &Config, + layer_idx: usize, + rotary_emb: Option>, + vb: VarBuilder, + ) -> Result { + let use_sliding_window = cfg.use_sliding_window.unwrap_or(false); + if use_sliding_window { + candle::bail!("sliding window is not supported") + } + + let head_dim = cfg.head_dim(); + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + + let attention_bias = cfg.attention_bias.unwrap_or(false); + + let q_proj = linear_b( + cfg.hidden_size, + num_heads * head_dim, + attention_bias, + vb.pp("q_proj"), + )?; + + let k_proj = linear_b( + cfg.hidden_size, + num_kv_heads * head_dim, + attention_bias, + vb.pp("k_proj"), + )?; + + let v_proj = linear_b( + cfg.hidden_size, + num_kv_heads * head_dim, + attention_bias, + vb.pp("v_proj"), + )?; + let o_proj = linear_b( + num_heads * head_dim, + cfg.hidden_size, + attention_bias, + vb.pp("o_proj"), + )?; + + // Necessary because the hidden_size in the config isn't always accurate + let hidden_size = head_dim * cfg.num_attention_heads; + + // Initialize KV cache with 512 tokens capacity to reduce initial memory allocation. + // The cache will grow in chunks of 512 tokens when needed. + let kv_cache = KvCache::new(2, 512); + + // Check if this layer should skip RoPE (NoPE) + let skip_rope = cfg.should_skip_rope(layer_idx); + + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + hidden_size, + rotary_emb, + kv_cache, + skip_rope, + }) + } + + pub(crate) fn forward( + &mut self, + x: &Tensor, + attn_mask: Option<&Tensor>, + offset: usize, + ) -> Result { + let (b, l, _) = x.dims3()?; + + // 1. Proj + let q = self.q_proj.forward(x)?; + let k = self.k_proj.forward(x)?; + let v = self.v_proj.forward(x)?; + + // 2. Reshape: (B, L, H, D) -> (B, H, L, D) + let q = q + .reshape((b, l, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let k = k + .reshape((b, l, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let v = v + .reshape((b, l, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + // 3. RoPE - only apply if this layer should use RoPE (not NoPE) + let (q, k) = if self.skip_rope { + // NoPE: Skip rotary embeddings, but ensure tensors are contiguous + (q.contiguous()?, k.contiguous()?) + } else { + // Apply RoPE + if let Some(ref rope) = self.rotary_emb { + rope.apply(&q, &k, offset)? + } else { + (q, k) + } + }; + + // 4. Accumulate KV cache + // Reset KV cache if we're at the first position + if offset == 0 { + self.kv_cache.reset(); + } + let (k, v) = self.kv_cache.append(&k.contiguous()?, &v.contiguous()?)?; + + // 5. GQA repeat_kv + let k = repeat_kv(k, self.num_kv_groups)?; + let v = repeat_kv(v, self.num_kv_groups)?; + + // 6. Attention score + let scale = 1.0 / (self.head_dim as f64).sqrt(); + let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?; + if let Some(m) = attn_mask { + scores = scores.broadcast_add(m)?; + } + let probs = candle_nn::ops::softmax_last_dim(&scores)?; + let ctx = probs.matmul(&v)?; // (B, H, L, D) + + // 7. Output proj + ctx.transpose(1, 2)? + .reshape((b, l, self.hidden_size))? + .apply(&self.o_proj) + } + + pub fn clear_kv_cache(&mut self) { + self.kv_cache.reset(); + } +} + +#[derive(Debug, Clone)] +pub(crate) struct DecoderLayer { + self_attn: SmolLM3Attention, + mlp: SmolLM3MLP, + ln1: RmsNorm, + ln2: RmsNorm, +} + +impl DecoderLayer { + fn new( + cfg: &Config, + layer_idx: usize, + rotary: Option>, + vb: VarBuilder, + ) -> Result { + let self_attn = SmolLM3Attention::new(cfg, layer_idx, rotary, vb.pp("self_attn"))?; + let mlp = SmolLM3MLP::new(cfg, vb.pp("mlp"))?; + let ln1 = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let ln2 = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + ln1, + ln2, + }) + } + + fn forward(&mut self, x: &Tensor, mask: Option<&Tensor>, offset: usize) -> Result { + let h = self.ln1.forward(x)?; + let h = self.self_attn.forward(&h, mask, offset)?; + let x = (x + h)?; + let h2 = self.ln2.forward(&x)?; + let h2 = h2.apply(&self.mlp)?; + x + h2 + } + + pub fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache(); + } +} + +#[derive(Debug, Clone)] +pub struct Model { + pub(crate) embed_tokens: candle_nn::Embedding, + pub(crate) layers: Vec, + pub(crate) norm: RmsNorm, + device: Device, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("model.embed_tokens"))?; + + // Only create rotary embedding if at least one layer uses RoPE + let needs_rope = (0..cfg.num_hidden_layers).any(|i| !cfg.should_skip_rope(i)); + let rotary = if needs_rope { + Some(Arc::new(SmolLM3RotaryEmbedding::new( + vb.dtype(), + cfg, + vb.device(), + )?)) + } else { + None + }; + + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb.pp("model.layers"); + for i in 0..cfg.num_hidden_layers { + layers.push(DecoderLayer::new(cfg, i, rotary.clone(), vb_l.pp(i))?); + } + Ok(Self { + embed_tokens, + layers, + norm: RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("model.norm"))?, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + pub fn clear_kv_cache(&mut self) { + for l in &mut self.layers { + l.clear_kv_cache(); + } + } + + fn causal_mask( + &self, + b: usize, + tgt: usize, + offset: usize, + sw: Option, + ) -> Result { + let minf = f32::NEG_INFINITY; + let mask: Vec<_> = (0..tgt) + .flat_map(|i| { + (0..(tgt + offset)).map(move |j| { + let past_ok = j <= i + offset; + let sw_ok = match sw { + Some(w) => (i + offset) as i64 - j as i64 <= w as i64, + None => true, + }; + if past_ok && sw_ok { + 0. + } else { + minf + } + }) + }) + .collect(); + Tensor::from_slice(&mask, (b, 1, tgt, tgt + offset), &self.device)?.to_dtype(self.dtype) + } + + pub fn forward(&mut self, input: &Tensor, offset: usize) -> Result { + let (b, l) = input.dims2()?; + + let mut h = self.embed_tokens.forward(input)?; + + let causal = if l == 1 { + None + } else { + Some(self.causal_mask(b, l, offset, None)?) + }; + + for layer in &mut self.layers { + h = layer.forward(&h, causal.as_ref(), offset)?; + } + self.norm.forward(&h) + } +} + +#[derive(Debug, Clone)] +pub struct ModelForCausalLM { + base: Model, + lm_head: Linear, +} + +impl ModelForCausalLM { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let base = Model::new(cfg, vb.clone())?; + let lm_head = if cfg.tie_word_embeddings { + Linear::from_weights(base.embed_tokens.embeddings().clone(), None) + } else { + linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))? + }; + Ok(Self { base, lm_head }) + } + + pub fn forward(&mut self, input: &Tensor, offset: usize) -> Result { + let (_, l) = input.dims2()?; + + self.base + .forward(input, offset)? + .narrow(1, l - 1, 1)? + .apply(&self.lm_head) + } + + pub fn clear_kv_cache(&mut self) { + self.base.clear_kv_cache(); + } +} diff --git a/patches/candle-transformers/src/models/snac.rs b/patches/candle-transformers/src/models/snac.rs new file mode 100644 index 0000000000..65fcb97b41 --- /dev/null +++ b/patches/candle-transformers/src/models/snac.rs @@ -0,0 +1,814 @@ +#![allow(unused)] +//! Implementation of the Multi-Scale Neural Audio Codec (SNAC) +//! +//! See: [SNAC](https://github.com/hubertsiuzdak/snac) +//! +/// Multi-Scale Neural Audio Codec (SNAC) compresses audio into discrete codes at a low bitrate. +/// For more information, read the paper: https://arxiv.org/abs/2410.14411 +/// +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{ + linear_b, Conv1d, Conv1dConfig, ConvTranspose1d, ConvTranspose1dConfig, LayerNorm, Linear, + VarBuilder, +}; + +#[derive(serde::Deserialize, Debug, Clone)] +pub struct Config { + pub sampling_rate: usize, + pub encoder_dim: usize, + pub encoder_rates: Vec, + pub decoder_dim: usize, + pub decoder_rates: Vec, + pub attn_window_size: Option, + pub codebook_size: usize, + pub codebook_dim: usize, + pub vq_strides: Vec, + pub noise: bool, + pub depthwise: bool, +} + +// Equivalent to torch.repeat_interleave +pub fn repeat_interleave( + img: &Tensor, + repeats: usize, + dim: D, +) -> Result { + if repeats == 1 { + return Ok(img.clone()); + } + let dim = dim.to_index(img.shape(), "chunk")?; + let img = img.unsqueeze(dim + 1)?; + let mut dims = img.dims().to_vec(); + dims[dim + 1] = repeats; + img.broadcast_as(dims)?.flatten(dim, dim + 1) +} + +pub fn conv1d_weight_norm( + in_c: usize, + out_c: usize, + kernel_size: usize, + config: candle_nn::Conv1dConfig, + vb: VarBuilder, +) -> Result { + let weight_g = vb.get((out_c, 1, 1), "parametrizations.weight.original0")?; + let weight_v = { + let name = "parametrizations.weight.original1"; + match vb.get((out_c, in_c, kernel_size), name) { + Ok(v) => v, + Err(_) => vb.get((out_c, 1, kernel_size), name)?, + } + }; + let norm_v = weight_v.sqr()?.sum_keepdim((1, 2))?.sqrt()?; + let weight = weight_v.broadcast_mul(&weight_g)?.broadcast_div(&norm_v)?; + let bias = vb.get(out_c, "bias")?; + Ok(Conv1d::new(weight, Some(bias), config)) +} + +pub fn conv1d_weight_norm_no_bias( + in_c: usize, + out_c: usize, + kernel_size: usize, + config: candle_nn::Conv1dConfig, + vb: VarBuilder, +) -> Result { + let weight_g = vb.get((out_c, 1, 1), "parametrizations.weight.original0")?; + let weight_v = { + let name = "parametrizations.weight.original1"; + match vb.get((out_c, in_c, kernel_size), name) { + Ok(v) => v, + Err(_) => vb.get((out_c, 1, kernel_size), name)?, + } + }; + let norm_v = weight_v.sqr()?.sum_keepdim((1, 2))?.sqrt()?; + let weight = weight_v.broadcast_mul(&weight_g)?.broadcast_div(&norm_v)?; + Ok(Conv1d::new(weight, None, config)) +} + +pub fn conv_transpose1d_weight_norm( + in_c: usize, + out_c: usize, + kernel_size: usize, + bias: bool, + config: candle_nn::ConvTranspose1dConfig, + vb: VarBuilder, +) -> Result { + let weight_g = vb.get((in_c, 1, 1), "parametrizations.weight.original0")?; + let weight_v = vb.get( + (in_c, out_c, kernel_size), + "parametrizations.weight.original1", + )?; + let norm_v = weight_v.sqr()?.sum_keepdim((1, 2))?.sqrt()?; + let weight = weight_v.broadcast_mul(&weight_g)?.broadcast_div(&norm_v)?; + let bias = if bias { + Some(vb.get(out_c, "bias")?) + } else { + None + }; + Ok(ConvTranspose1d::new(weight, bias, config)) +} + +// https://github.com/hubertsiuzdak/snac/blob/main/snac/attention.py +#[allow(unused)] +#[derive(Debug, Clone)] +struct SinusoidalEmbeddings { + inv_freq: Tensor, + scale: Tensor, + scale_base: f32, + use_xpos: bool, +} + +impl SinusoidalEmbeddings { + fn new(dim: usize, scale_base: f32, use_xpos: bool, dev: &Device) -> Result { + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / 10_000f32.powf(i as f32 / dim as f32)) + .collect(); + let len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, len, dev)?.to_dtype(DType::F32)?; + let scale: Vec<_> = (0..dim) + .step_by(2) + .map(|i| (i as f32 + 0.4 * dim as f32) / (1.4 * dim as f32)) + .collect(); + let scale = Tensor::from_vec(scale, len, dev)?.to_dtype(DType::F32)?; + Ok(Self { + inv_freq, + scale, + scale_base, + use_xpos, + }) + } +} + +#[allow(unused)] +#[derive(Debug, Clone)] +struct LocalMHA { + norm: LayerNorm, + to_qkv: Linear, + to_out: Linear, + num_heads: usize, + head_dim: usize, + rel_pos: Option, +} + +impl LocalMHA { + fn new( + dim: usize, + window_size: usize, + dim_head: usize, + use_rotary_pos_emb: bool, + vb: VarBuilder, + ) -> Result { + let norm = candle_nn::layer_norm(dim, 1e-5, vb.pp("norm"))?; + let to_qkv = linear_b(dim, dim * 3, false, vb.pp("to_qkv"))?; + let to_out = linear_b(dim, dim, false, vb.pp("to_out"))?; + let rel_pos = if use_rotary_pos_emb { + let rel_pos = + SinusoidalEmbeddings::new(dim_head, window_size as f32 / 2.0, false, vb.device())?; + Some(rel_pos) + } else { + None + }; + Ok(Self { + norm, + to_qkv, + to_out, + rel_pos, + num_heads: dim / dim_head, + head_dim: dim_head, + }) + } +} + +impl Module for LocalMHA { + fn forward(&self, xs: &Tensor) -> Result { + let (b, c, t) = xs.dims3()?; + let residual = xs.clone(); + let xs = xs.transpose(1, 2)?.apply(&self.norm)?; + let qkv = xs.apply(&self.to_qkv)?; + let q = qkv.narrow(D::Minus1, 0, c)?; + let k = qkv.narrow(D::Minus1, c, c)?; + let v = qkv.narrow(D::Minus1, 2 * c, c)?; + let q = q + .reshape((b, t, self.num_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let k = k + .reshape((b, t, self.num_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let v = v + .reshape((b, t, self.num_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let (q, k) = match self.rel_pos { + Some(_) => todo!(), + None => (q, k), + }; + let out = { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (q.matmul(&k.transpose(2, 3)?)? * scale)?; + // Non-causal attention + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&v)? + }; + let out = out + .transpose(1, 2)? + .reshape((b, t, self.num_heads * self.head_dim))? + .apply(&self.to_out)?; + out.transpose(1, 2)? + residual + } +} + +#[derive(Debug, Clone)] +struct Snake1d { + alpha: Tensor, +} + +impl Snake1d { + pub fn new(channels: usize, vb: VarBuilder) -> Result { + let alpha = vb.get((1, channels, 1), "alpha")?; + Ok(Self { alpha }) + } +} + +impl Module for Snake1d { + fn forward(&self, xs: &Tensor) -> Result { + let xs_shape = xs.shape(); + let xs = xs.flatten_from(2)?; + let sin = self.alpha.broadcast_mul(&xs)?.sin()?; + let sin = (&sin * &sin)?; + (xs + (&self.alpha + 1e-9)?.recip()?.broadcast_mul(&sin)?)?.reshape(xs_shape) + } +} + +#[derive(Debug, Clone)] +struct ResidualUnit { + snake1: Snake1d, + conv1: Conv1d, + snake2: Snake1d, + conv2: Conv1d, +} + +impl ResidualUnit { + fn new( + dim: usize, + dilation: usize, + kernel: usize, + groups: usize, + vb: VarBuilder, + ) -> Result { + let pad = ((kernel - 1) * dilation) / 2; + let vb = vb.pp("block"); + let snake1 = Snake1d::new(dim, vb.pp(0))?; + let cfg1 = Conv1dConfig { + dilation, + padding: pad, + groups, + ..Default::default() + }; + let conv1 = conv1d_weight_norm(dim, dim, 7, cfg1, vb.pp(1))?; + let snake2 = Snake1d::new(dim, vb.pp(2))?; + let conv2 = conv1d_weight_norm(dim, dim, 1, Default::default(), vb.pp(3))?; + Ok(Self { + snake1, + conv1, + snake2, + conv2, + }) + } +} + +impl Module for ResidualUnit { + fn forward(&self, xs: &Tensor) -> Result { + let ys = xs + .apply(&self.snake1)? + .apply(&self.conv1)? + .apply(&self.snake2)? + .apply(&self.conv2)?; + let pad = (xs.dim(D::Minus1)? - ys.dim(D::Minus1)?) / 2; + if pad > 0 { + &ys + xs.narrow(D::Minus1, pad, ys.dim(D::Minus1)?) + } else { + ys + xs + } + } +} + +#[derive(Debug, Clone)] +struct NoiseBlock { + linear: Conv1d, +} + +impl NoiseBlock { + fn new(dim: usize, vb: VarBuilder) -> Result { + let linear = conv1d_weight_norm_no_bias(dim, dim, 1, Default::default(), vb.pp("linear"))?; + Ok(Self { linear }) + } +} + +impl Module for NoiseBlock { + fn forward(&self, xs: &Tensor) -> Result { + let (b, _c, t) = xs.dims3()?; + let noise = Tensor::randn(0f32, 1f32, (b, 1, t), xs.device())?; + let h = xs.apply(&self.linear)?; + let n = noise.broadcast_mul(&h)?; + let xs = (xs + n)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct DecoderBlock { + snake1: Snake1d, + conv_tr1: ConvTranspose1d, + noise: Option, + res1: ResidualUnit, + res2: ResidualUnit, + res3: ResidualUnit, +} + +impl DecoderBlock { + fn new( + in_dim: usize, + out_dim: usize, + stride: usize, + noise: bool, + groups: usize, + vb: VarBuilder, + ) -> Result { + let vb = vb.pp("block"); + let snake1 = Snake1d::new(in_dim, vb.pp(0))?; + let cfg = ConvTranspose1dConfig { + stride, + padding: stride.div_ceil(2), + output_padding: stride % 2, + ..Default::default() + }; + let conv_tr1 = + conv_transpose1d_weight_norm(in_dim, out_dim, 2 * stride, true, cfg, vb.pp(1))?; + let (n, noise) = if noise { + let noise = NoiseBlock::new(out_dim, vb.pp(2))?; + (1, Some(noise)) + } else { + (0, None) + }; + let res1 = ResidualUnit::new(out_dim, 1, 7, groups, vb.pp(2 + n))?; + let res2 = ResidualUnit::new(out_dim, 3, 7, groups, vb.pp(3 + n))?; + let res3 = ResidualUnit::new(out_dim, 9, 7, groups, vb.pp(4 + n))?; + Ok(Self { + snake1, + conv_tr1, + noise, + res1, + res2, + res3, + }) + } +} + +impl Module for DecoderBlock { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.snake1)? + .apply(&self.conv_tr1)? + .apply(&self.noise.as_ref())? + .apply(&self.res1)? + .apply(&self.res2)? + .apply(&self.res3) + } +} + +#[derive(Debug, Clone)] +struct EncoderBlock { + res1: ResidualUnit, + res2: ResidualUnit, + res3: ResidualUnit, + snake1: Snake1d, + conv1: Conv1d, +} + +impl EncoderBlock { + fn new( + out_dim: usize, + in_dim: Option, + stride: usize, + groups: usize, + vb: VarBuilder, + ) -> Result { + let vb = vb.pp("block"); + let in_dim = in_dim.unwrap_or(out_dim / 2); + let res1 = ResidualUnit::new(in_dim, 1, 7, groups, vb.pp(0))?; + let res2 = ResidualUnit::new(in_dim, 3, 7, groups, vb.pp(1))?; + let res3 = ResidualUnit::new(in_dim, 9, 7, groups, vb.pp(2))?; + let snake1 = Snake1d::new(in_dim, vb.pp(3))?; + let cfg1 = Conv1dConfig { + stride, + padding: stride.div_ceil(2), + ..Default::default() + }; + let conv1 = conv1d_weight_norm(in_dim, out_dim, 2 * stride, cfg1, vb.pp(4))?; + Ok(Self { + res1, + res2, + res3, + snake1, + conv1, + }) + } +} + +impl candle::Module for EncoderBlock { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.res1)? + .apply(&self.res2)? + .apply(&self.res3)? + .apply(&self.snake1)? + .apply(&self.conv1) + } +} + +#[derive(Debug, Clone)] +pub struct Encoder { + conv1: Conv1d, + blocks: Vec, + local_mha: Option, + conv2: Conv1d, +} + +impl candle::Module for Encoder { + fn forward(&self, xs: &Tensor) -> Result { + let mut xs = xs.apply(&self.conv1)?; + for block in self.blocks.iter() { + xs = xs.apply(block)? + } + xs.apply(&self.conv2) + } +} + +impl Encoder { + fn new( + mut d_model: usize, + strides: &[usize], + depthwise: bool, + attn_window_size: Option, + vb: VarBuilder, + ) -> Result { + let vb = vb.pp("block"); + let mut idx = 0; + let cfg1 = Conv1dConfig { + padding: 3, + ..Default::default() + }; + let conv1 = conv1d_weight_norm(1, d_model, 7, cfg1, vb.pp(idx))?; + idx += 1; + let mut blocks = Vec::with_capacity(strides.len()); + for &stride in strides.iter() { + d_model *= 2; + let groups = if depthwise { d_model / 2 } else { 1 }; + let block = EncoderBlock::new(d_model, None, stride, groups, vb.pp(idx))?; + idx += 1; + blocks.push(block) + } + let local_mha = match attn_window_size { + Some(w) => { + let mha = LocalMHA::new(d_model, w, 64, true, vb.pp(idx))?; + idx += 1; + Some(mha) + } + None => None, + }; + let groups = if depthwise { d_model } else { 1 }; + let cfg2 = Conv1dConfig { + padding: 3, + groups, + ..Default::default() + }; + let conv2 = conv1d_weight_norm(d_model, d_model, 7, cfg2, vb.pp(idx))?; + idx += 1; + Ok(Self { + conv1, + blocks, + local_mha, + conv2, + }) + } +} + +#[derive(Debug, Clone)] +enum ConvInit { + Depthwise(Conv1d, Conv1d), + Standard(Conv1d), +} + +#[derive(Debug, Clone)] +pub struct Decoder { + conv1: ConvInit, + local_mha: Option, + blocks: Vec, + snake1: Snake1d, + conv2: Conv1d, +} + +impl Decoder { + #[allow(clippy::too_many_arguments)] + fn new( + in_c: usize, + mut channels: usize, + rates: &[usize], + noise: bool, + depthwise: bool, + attn_window_size: Option, + d_out: usize, + vb: VarBuilder, + ) -> Result { + let vb = vb.pp("model"); + let mut idx = 0; + let pad3 = Conv1dConfig { + padding: 3, + ..Default::default() + }; + let conv1 = if depthwise { + let cfg1 = Conv1dConfig { + padding: 3, + groups: in_c, + ..Default::default() + }; + let conv1 = conv1d_weight_norm(in_c, in_c, 7, cfg1, vb.pp(idx))?; + idx += 1; + let conv2 = conv1d_weight_norm(in_c, channels, 1, Default::default(), vb.pp(idx))?; + idx += 1; + ConvInit::Depthwise(conv1, conv2) + } else { + let conv1 = conv1d_weight_norm(in_c, channels, 7, pad3, vb.pp(idx))?; + idx += 1; + ConvInit::Standard(conv1) + }; + let mut blocks = Vec::with_capacity(rates.len()); + let local_mha = match attn_window_size { + Some(w) => { + let mha = LocalMHA::new(channels, w, 64, true, vb.pp(idx))?; + idx += 1; + Some(mha) + } + None => None, + }; + for stride in rates.iter() { + let groups = if depthwise { channels / 2 } else { 1 }; + let block = + DecoderBlock::new(channels, channels / 2, *stride, noise, groups, vb.pp(idx))?; + idx += 1; + channels /= 2; + blocks.push(block) + } + let snake1 = Snake1d::new(channels, vb.pp(idx))?; + idx += 1; + let conv2 = conv1d_weight_norm(channels, d_out, 7, pad3, vb.pp(idx))?; + idx += 1; + Ok(Self { + conv1, + local_mha, + blocks, + snake1, + conv2, + }) + } +} + +impl candle::Module for Decoder { + fn forward(&self, xs: &Tensor) -> Result { + let mut xs = match &self.conv1 { + ConvInit::Standard(c) => xs.apply(c)?, + ConvInit::Depthwise(c1, c2) => xs.apply(c1)?.apply(c2)?, + }; + for block in self.blocks.iter() { + xs = xs.apply(block)? + } + xs.apply(&self.snake1)?.apply(&self.conv2) + } +} + +fn normalize(v: &Tensor) -> Result { + v.broadcast_div(&v.sqr()?.sum_keepdim(1)?.sqrt()?) +} + +// https://github.com/hubertsiuzdak/snac/blob/main/snac/vq.py +#[allow(unused)] +#[derive(Clone, Debug)] +struct VectorQuantizer { + in_proj: Conv1d, + out_proj: Conv1d, + codebook: candle_nn::Embedding, + stride: usize, +} + +impl VectorQuantizer { + fn new( + in_dim: usize, + cb_size: usize, + cb_dim: usize, + stride: usize, + vb: VarBuilder, + ) -> Result { + let in_proj = conv1d_weight_norm(in_dim, cb_dim, 1, Default::default(), vb.pp("in_proj"))?; + let out_proj = + conv1d_weight_norm(cb_dim, in_dim, 1, Default::default(), vb.pp("out_proj"))?; + let codebook = candle_nn::embedding(cb_size, cb_dim, vb.pp("codebook"))?; + Ok(Self { + in_proj, + out_proj, + codebook, + stride, + }) + } + + fn decode_latents(&self, latents: &Tensor) -> Result<(Tensor, Tensor)> { + let (b, d, t) = latents.dims3()?; + let encodings = latents.transpose(1, 2)?.reshape((b * t, d))?; + let encodings = normalize(&encodings)?; + let codebook = normalize(self.codebook.embeddings())?; + let dist = (encodings + .sqr()? + .sum_keepdim(1)? + .broadcast_sub(&encodings.matmul(&codebook.t()?)?)? + * 2.0)? + .broadcast_add(&codebook.sqr()?.sum_keepdim(1)?.t()?)?; + let indices = dist.argmin(1)?.reshape((b, ()))?; + let z_q = self.decode_code(&indices)?; + Ok((z_q, indices)) + } + + fn encode(&self, z: &Tensor) -> Result<(Tensor, Tensor)> { + let z = if self.stride > 1 { + let (b, c, t) = z.dims3()?; + z.reshape((b, c, 1, t))? + .avg_pool2d((1, self.stride))? + .squeeze(2)? + } else { + z.clone() + }; + let z_e = z.apply(&self.in_proj)?; + let (z_q, indices) = self.decode_latents(&z_e)?; + let z_q = z_q.apply(&self.out_proj)?; + let z_q = if self.stride > 1 { + repeat_interleave(&z_q, self.stride, D::Minus1)? + } else { + z_q + }; + Ok((z_q, indices)) + } + + fn embed_code(&self, embed_id: &Tensor) -> Result { + embed_id.apply(&self.codebook) + } + + fn decode_code(&self, embed_id: &Tensor) -> Result { + self.embed_code(embed_id)?.transpose(1, 2) + } +} + +#[derive(Clone, Debug)] +pub struct ResidualVectorQuantizer { + quantizers: Vec, +} + +impl ResidualVectorQuantizer { + fn new( + input_dim: usize, + cb_size: usize, + cb_dim: usize, + vq_strides: &[usize], + vb: VarBuilder, + ) -> Result { + let vb = &vb.pp("quantizers"); + let quantizers = vq_strides + .iter() + .enumerate() + .map(|(i, stride)| VectorQuantizer::new(input_dim, cb_size, cb_dim, *stride, vb.pp(i))) + .collect::>>()?; + Ok(Self { quantizers }) + } + + fn encode(&self, z: &Tensor) -> Result<(Tensor, Vec)> { + let mut residual = z.clone(); + let mut z_q = z.zeros_like()?; + let mut codes = Vec::with_capacity(self.quantizers.len()); + for quantizer in self.quantizers.iter() { + let (z_q_i, indices_i) = quantizer.encode(&residual)?; + z_q = (z_q + &z_q_i)?; + residual = (residual - &z_q_i)?; + codes.push(indices_i) + } + Ok((z_q, codes)) + } + + #[allow(clippy::wrong_self_convention)] + fn from_codes(&self, codes: &[&Tensor]) -> Result { + let mut sum = None; + for (quantizer, codes) in self.quantizers.iter().zip(codes.iter()) { + let z_p_i = quantizer.decode_code(codes)?; + let z_q_i = z_p_i.apply(&quantizer.out_proj)?; + let z_q_i = repeat_interleave(&z_q_i, quantizer.stride, D::Minus1)?; + let s = match sum { + None => z_q_i, + Some(s) => (s + z_q_i)?, + }; + sum = Some(s) + } + match sum { + Some(s) => Ok(s), + None => candle::bail!("empty codebooks"), + } + } +} + +fn gcd(mut a: usize, mut b: usize) -> usize { + while b != 0 { + let t = b; + b = a % b; + a = t; + } + a +} + +fn lcm(a: usize, b: usize) -> usize { + a / gcd(a, b) * b +} + +// https://github.com/hubertsiuzdak/snac/blob/main/snac/snac.py +#[derive(Debug, Clone)] +pub struct Model { + pub encoder: Encoder, + pub quantizer: ResidualVectorQuantizer, + pub decoder: Decoder, + pub hop_length: usize, + pub config: Config, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let encoder = Encoder::new( + cfg.encoder_dim, + &cfg.encoder_rates, + cfg.depthwise, + cfg.attn_window_size, + vb.pp("encoder"), + )?; + let latent_dim = cfg.encoder_dim * 2usize.pow(cfg.encoder_rates.len() as u32); + let quantizer = ResidualVectorQuantizer::new( + latent_dim, + cfg.codebook_size, + cfg.codebook_dim, + &cfg.vq_strides, + vb.pp("quantizer"), + )?; + let decoder = Decoder::new( + latent_dim, + cfg.decoder_dim, + &cfg.decoder_rates, + cfg.noise, + cfg.depthwise, + cfg.attn_window_size, + /* d_out */ 1, + vb.pp("decoder"), + )?; + let hop_length = cfg.encoder_rates.iter().product::(); + Ok(Self { + encoder, + decoder, + quantizer, + config: cfg.clone(), + hop_length, + }) + } + + fn preprocess(&self, audio_data: &Tensor) -> Result { + let len = audio_data.dim(D::Minus1)?; + let lcm = lcm( + self.config.vq_strides[0], + self.config.attn_window_size.unwrap_or(1), + ); + let pad_to = self.hop_length * lcm; + let right_pad = len.div_ceil(pad_to) * pad_to - len; + let audio_data = audio_data.pad_with_zeros(D::Minus1, 0, right_pad)?; + Ok(audio_data) + } + + pub fn encode(&self, audio_data: &Tensor) -> Result> { + let audio_data = self.preprocess(audio_data)?; + let z = self.encoder.forward(&audio_data)?; + let (_, codes) = self.quantizer.encode(&z)?; + Ok(codes) + } + + pub fn decode(&self, audio_codes: &[&Tensor]) -> Result { + let audio_values = self.quantizer.from_codes(audio_codes)?; + audio_values.apply(&self.decoder) + } + + pub fn config(&self) -> &Config { + &self.config + } + + pub fn num_codebooks(&self) -> usize { + self.quantizer.quantizers.len() + } +} diff --git a/patches/candle-transformers/src/models/stable_diffusion/attention.rs b/patches/candle-transformers/src/models/stable_diffusion/attention.rs new file mode 100644 index 0000000000..c04e6aa1ff --- /dev/null +++ b/patches/candle-transformers/src/models/stable_diffusion/attention.rs @@ -0,0 +1,567 @@ +//! Attention Based Building Blocks +use candle::{DType, IndexOp, Result, Tensor, D}; +use candle_nn as nn; +use candle_nn::Module; + +#[derive(Debug)] +struct GeGlu { + proj: nn::Linear, + span: tracing::Span, +} + +impl GeGlu { + fn new(vs: nn::VarBuilder, dim_in: usize, dim_out: usize) -> Result { + let proj = nn::linear(dim_in, dim_out * 2, vs.pp("proj"))?; + let span = tracing::span!(tracing::Level::TRACE, "geglu"); + Ok(Self { proj, span }) + } +} + +impl Module for GeGlu { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let hidden_states_and_gate = self.proj.forward(xs)?.chunk(2, D::Minus1)?; + &hidden_states_and_gate[0] * hidden_states_and_gate[1].gelu()? + } +} + +/// A feed-forward layer. +#[derive(Debug)] +struct FeedForward { + project_in: GeGlu, + linear: nn::Linear, + span: tracing::Span, +} + +impl FeedForward { + // The glu parameter in the python code is unused? + // https://github.com/huggingface/diffusers/blob/d3d22ce5a894becb951eec03e663951b28d45135/src/diffusers/models/attention.py#L347 + /// Creates a new feed-forward layer based on some given input dimension, some + /// output dimension, and a multiplier to be used for the intermediary layer. + fn new(vs: nn::VarBuilder, dim: usize, dim_out: Option, mult: usize) -> Result { + let inner_dim = dim * mult; + let dim_out = dim_out.unwrap_or(dim); + let vs = vs.pp("net"); + let project_in = GeGlu::new(vs.pp("0"), dim, inner_dim)?; + let linear = nn::linear(inner_dim, dim_out, vs.pp("2"))?; + let span = tracing::span!(tracing::Level::TRACE, "ff"); + Ok(Self { + project_in, + linear, + span, + }) + } +} + +impl Module for FeedForward { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let xs = self.project_in.forward(xs)?; + self.linear.forward(&xs) + } +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +#[derive(Debug)] +pub struct CrossAttention { + to_q: nn::Linear, + to_k: nn::Linear, + to_v: nn::Linear, + to_out: nn::Linear, + heads: usize, + scale: f64, + slice_size: Option, + span: tracing::Span, + span_attn: tracing::Span, + span_softmax: tracing::Span, + use_flash_attn: bool, +} + +impl CrossAttention { + // Defaults should be heads = 8, dim_head = 64, context_dim = None + pub fn new( + vs: nn::VarBuilder, + query_dim: usize, + context_dim: Option, + heads: usize, + dim_head: usize, + slice_size: Option, + use_flash_attn: bool, + ) -> Result { + let inner_dim = dim_head * heads; + let context_dim = context_dim.unwrap_or(query_dim); + let scale = 1.0 / f64::sqrt(dim_head as f64); + let to_q = nn::linear_no_bias(query_dim, inner_dim, vs.pp("to_q"))?; + let to_k = nn::linear_no_bias(context_dim, inner_dim, vs.pp("to_k"))?; + let to_v = nn::linear_no_bias(context_dim, inner_dim, vs.pp("to_v"))?; + let to_out = nn::linear(inner_dim, query_dim, vs.pp("to_out.0"))?; + let span = tracing::span!(tracing::Level::TRACE, "xa"); + let span_attn = tracing::span!(tracing::Level::TRACE, "xa-attn"); + let span_softmax = tracing::span!(tracing::Level::TRACE, "xa-softmax"); + Ok(Self { + to_q, + to_k, + to_v, + to_out, + heads, + scale, + slice_size, + span, + span_attn, + span_softmax, + use_flash_attn, + }) + } + + fn reshape_heads_to_batch_dim(&self, xs: &Tensor) -> Result { + let (batch_size, seq_len, dim) = xs.dims3()?; + xs.reshape((batch_size, seq_len, self.heads, dim / self.heads))? + .transpose(1, 2)? + .reshape((batch_size * self.heads, seq_len, dim / self.heads)) + } + + fn reshape_batch_dim_to_heads(&self, xs: &Tensor) -> Result { + let (batch_size, seq_len, dim) = xs.dims3()?; + xs.reshape((batch_size / self.heads, self.heads, seq_len, dim))? + .transpose(1, 2)? + .reshape((batch_size / self.heads, seq_len, dim * self.heads)) + } + + fn sliced_attention( + &self, + query: &Tensor, + key: &Tensor, + value: &Tensor, + slice_size: usize, + ) -> Result { + let batch_size_attention = query.dim(0)?; + let mut hidden_states = Vec::with_capacity(batch_size_attention / slice_size); + let in_dtype = query.dtype(); + let query = query.to_dtype(DType::F32)?; + let key = key.to_dtype(DType::F32)?; + let value = value.to_dtype(DType::F32)?; + + for i in 0..batch_size_attention / slice_size { + let start_idx = i * slice_size; + let end_idx = (i + 1) * slice_size; + + let xs = query + .i(start_idx..end_idx)? + .matmul(&(key.i(start_idx..end_idx)?.t()? * self.scale)?)?; + let xs = nn::ops::softmax(&xs, D::Minus1)?.matmul(&value.i(start_idx..end_idx)?)?; + hidden_states.push(xs) + } + let hidden_states = Tensor::stack(&hidden_states, 0)?.to_dtype(in_dtype)?; + self.reshape_batch_dim_to_heads(&hidden_states) + } + + fn attention(&self, query: &Tensor, key: &Tensor, value: &Tensor) -> Result { + let _enter = self.span_attn.enter(); + let xs = if self.use_flash_attn { + let init_dtype = query.dtype(); + let q = query + .to_dtype(candle::DType::F16)? + .unsqueeze(0)? + .transpose(1, 2)?; + let k = key + .to_dtype(candle::DType::F16)? + .unsqueeze(0)? + .transpose(1, 2)?; + let v = value + .to_dtype(candle::DType::F16)? + .unsqueeze(0)? + .transpose(1, 2)?; + flash_attn(&q, &k, &v, self.scale as f32, false)? + .transpose(1, 2)? + .squeeze(0)? + .to_dtype(init_dtype)? + } else { + let in_dtype = query.dtype(); + let query = query.to_dtype(DType::F32)?; + let key = key.to_dtype(DType::F32)?; + let value = value.to_dtype(DType::F32)?; + let xs = query.matmul(&(key.t()? * self.scale)?)?; + let xs = { + let _enter = self.span_softmax.enter(); + nn::ops::softmax_last_dim(&xs)? + }; + xs.matmul(&value)?.to_dtype(in_dtype)? + }; + self.reshape_batch_dim_to_heads(&xs) + } + + pub fn forward(&self, xs: &Tensor, context: Option<&Tensor>) -> Result { + let _enter = self.span.enter(); + let query = self.to_q.forward(xs)?; + let context = context.unwrap_or(xs).contiguous()?; + let key = self.to_k.forward(&context)?; + let value = self.to_v.forward(&context)?; + let query = self.reshape_heads_to_batch_dim(&query)?; + let key = self.reshape_heads_to_batch_dim(&key)?; + let value = self.reshape_heads_to_batch_dim(&value)?; + let dim0 = query.dim(0)?; + let slice_size = self.slice_size.and_then(|slice_size| { + if dim0 < slice_size { + None + } else { + Some(slice_size) + } + }); + let xs = match slice_size { + None => self.attention(&query, &key, &value)?, + Some(slice_size) => self.sliced_attention(&query, &key, &value, slice_size)?, + }; + self.to_out.forward(&xs) + } +} + +/// A basic Transformer block. +#[derive(Debug)] +struct BasicTransformerBlock { + attn1: CrossAttention, + ff: FeedForward, + attn2: CrossAttention, + norm1: nn::LayerNorm, + norm2: nn::LayerNorm, + norm3: nn::LayerNorm, + span: tracing::Span, +} + +impl BasicTransformerBlock { + fn new( + vs: nn::VarBuilder, + dim: usize, + n_heads: usize, + d_head: usize, + context_dim: Option, + sliced_attention_size: Option, + use_flash_attn: bool, + ) -> Result { + let attn1 = CrossAttention::new( + vs.pp("attn1"), + dim, + None, + n_heads, + d_head, + sliced_attention_size, + use_flash_attn, + )?; + let ff = FeedForward::new(vs.pp("ff"), dim, None, 4)?; + let attn2 = CrossAttention::new( + vs.pp("attn2"), + dim, + context_dim, + n_heads, + d_head, + sliced_attention_size, + use_flash_attn, + )?; + let norm1 = nn::layer_norm(dim, 1e-5, vs.pp("norm1"))?; + let norm2 = nn::layer_norm(dim, 1e-5, vs.pp("norm2"))?; + let norm3 = nn::layer_norm(dim, 1e-5, vs.pp("norm3"))?; + let span = tracing::span!(tracing::Level::TRACE, "basic-transformer"); + Ok(Self { + attn1, + ff, + attn2, + norm1, + norm2, + norm3, + span, + }) + } + + fn forward(&self, xs: &Tensor, context: Option<&Tensor>) -> Result { + let _enter = self.span.enter(); + let xs = (self.attn1.forward(&self.norm1.forward(xs)?, None)? + xs)?; + let xs = (self.attn2.forward(&self.norm2.forward(&xs)?, context)? + xs)?; + self.ff.forward(&self.norm3.forward(&xs)?)? + xs + } +} + +#[derive(Debug, Clone, Copy)] +pub struct SpatialTransformerConfig { + pub depth: usize, + pub num_groups: usize, + pub context_dim: Option, + pub sliced_attention_size: Option, + pub use_linear_projection: bool, +} + +impl Default for SpatialTransformerConfig { + fn default() -> Self { + Self { + depth: 1, + num_groups: 32, + context_dim: None, + sliced_attention_size: None, + use_linear_projection: false, + } + } +} + +#[derive(Debug)] +enum Proj { + Conv2d(nn::Conv2d), + Linear(nn::Linear), +} + +// Aka Transformer2DModel +#[derive(Debug)] +pub struct SpatialTransformer { + norm: nn::GroupNorm, + proj_in: Proj, + transformer_blocks: Vec, + proj_out: Proj, + span: tracing::Span, + pub config: SpatialTransformerConfig, +} + +impl SpatialTransformer { + pub fn new( + vs: nn::VarBuilder, + in_channels: usize, + n_heads: usize, + d_head: usize, + use_flash_attn: bool, + config: SpatialTransformerConfig, + ) -> Result { + let inner_dim = n_heads * d_head; + let norm = nn::group_norm(config.num_groups, in_channels, 1e-6, vs.pp("norm"))?; + let proj_in = if config.use_linear_projection { + Proj::Linear(nn::linear(in_channels, inner_dim, vs.pp("proj_in"))?) + } else { + Proj::Conv2d(nn::conv2d( + in_channels, + inner_dim, + 1, + Default::default(), + vs.pp("proj_in"), + )?) + }; + let mut transformer_blocks = vec![]; + let vs_tb = vs.pp("transformer_blocks"); + for index in 0..config.depth { + let tb = BasicTransformerBlock::new( + vs_tb.pp(index.to_string()), + inner_dim, + n_heads, + d_head, + config.context_dim, + config.sliced_attention_size, + use_flash_attn, + )?; + transformer_blocks.push(tb) + } + let proj_out = if config.use_linear_projection { + Proj::Linear(nn::linear(in_channels, inner_dim, vs.pp("proj_out"))?) + } else { + Proj::Conv2d(nn::conv2d( + inner_dim, + in_channels, + 1, + Default::default(), + vs.pp("proj_out"), + )?) + }; + let span = tracing::span!(tracing::Level::TRACE, "spatial-transformer"); + Ok(Self { + norm, + proj_in, + transformer_blocks, + proj_out, + span, + config, + }) + } + + pub fn forward(&self, xs: &Tensor, context: Option<&Tensor>) -> Result { + let _enter = self.span.enter(); + let (batch, _channel, height, weight) = xs.dims4()?; + let residual = xs; + let xs = self.norm.forward(xs)?; + let (inner_dim, xs) = match &self.proj_in { + Proj::Conv2d(p) => { + let xs = p.forward(&xs)?; + let inner_dim = xs.dim(1)?; + let xs = xs + .transpose(1, 2)? + .t()? + .reshape((batch, height * weight, inner_dim))?; + (inner_dim, xs) + } + Proj::Linear(p) => { + let inner_dim = xs.dim(1)?; + let xs = xs + .transpose(1, 2)? + .t()? + .reshape((batch, height * weight, inner_dim))?; + (inner_dim, p.forward(&xs)?) + } + }; + let mut xs = xs; + for block in self.transformer_blocks.iter() { + xs = block.forward(&xs, context)? + } + let xs = match &self.proj_out { + Proj::Conv2d(p) => p.forward( + &xs.reshape((batch, height, weight, inner_dim))? + .t()? + .transpose(1, 2)?, + )?, + Proj::Linear(p) => p + .forward(&xs)? + .reshape((batch, height, weight, inner_dim))? + .t()? + .transpose(1, 2)?, + }; + xs + residual + } +} + +/// Configuration for an attention block. +#[derive(Debug, Clone, Copy)] +pub struct AttentionBlockConfig { + pub num_head_channels: Option, + pub num_groups: usize, + pub rescale_output_factor: f64, + pub eps: f64, +} + +impl Default for AttentionBlockConfig { + fn default() -> Self { + Self { + num_head_channels: None, + num_groups: 32, + rescale_output_factor: 1., + eps: 1e-5, + } + } +} + +#[derive(Debug)] +pub struct AttentionBlock { + group_norm: nn::GroupNorm, + query: nn::Linear, + key: nn::Linear, + value: nn::Linear, + proj_attn: nn::Linear, + channels: usize, + num_heads: usize, + span: tracing::Span, + config: AttentionBlockConfig, +} + +// In the .safetensor weights of official Stable Diffusion 3 Medium Huggingface repo +// https://huggingface.co/stabilityai/stable-diffusion-3-medium +// Linear layer may use a different dimension for the weight in the linear, which is +// incompatible with the current implementation of the nn::linear constructor. +// This is a workaround to handle the different dimensions. +fn get_qkv_linear(channels: usize, vs: nn::VarBuilder) -> Result { + match vs.get((channels, channels), "weight") { + Ok(_) => nn::linear(channels, channels, vs), + Err(_) => { + let weight = vs + .get((channels, channels, 1, 1), "weight")? + .reshape((channels, channels))?; + let bias = vs.get((channels,), "bias")?; + Ok(nn::Linear::new(weight, Some(bias))) + } + } +} + +impl AttentionBlock { + pub fn new(vs: nn::VarBuilder, channels: usize, config: AttentionBlockConfig) -> Result { + let num_head_channels = config.num_head_channels.unwrap_or(channels); + let num_heads = channels / num_head_channels; + let group_norm = + nn::group_norm(config.num_groups, channels, config.eps, vs.pp("group_norm"))?; + let (q_path, k_path, v_path, out_path) = if vs.contains_tensor("to_q.weight") { + ("to_q", "to_k", "to_v", "to_out.0") + } else { + ("query", "key", "value", "proj_attn") + }; + let query = get_qkv_linear(channels, vs.pp(q_path))?; + let key = get_qkv_linear(channels, vs.pp(k_path))?; + let value = get_qkv_linear(channels, vs.pp(v_path))?; + let proj_attn = get_qkv_linear(channels, vs.pp(out_path))?; + let span = tracing::span!(tracing::Level::TRACE, "attn-block"); + Ok(Self { + group_norm, + query, + key, + value, + proj_attn, + channels, + num_heads, + span, + config, + }) + } + + fn transpose_for_scores(&self, xs: Tensor) -> Result { + let (batch, t, h_times_d) = xs.dims3()?; + xs.reshape((batch, t, self.num_heads, h_times_d / self.num_heads))? + .transpose(1, 2) + } +} + +impl Module for AttentionBlock { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let in_dtype = xs.dtype(); + let residual = xs; + let (batch, channel, height, width) = xs.dims4()?; + let xs = self + .group_norm + .forward(xs)? + .reshape((batch, channel, height * width))? + .transpose(1, 2)?; + + let query_proj = self.query.forward(&xs)?; + let key_proj = self.key.forward(&xs)?; + let value_proj = self.value.forward(&xs)?; + + let query_states = self + .transpose_for_scores(query_proj)? + .to_dtype(DType::F32)?; + let key_states = self.transpose_for_scores(key_proj)?.to_dtype(DType::F32)?; + let value_states = self + .transpose_for_scores(value_proj)? + .to_dtype(DType::F32)?; + + // scale is applied twice, hence the -0.25 here rather than -0.5. + // https://github.com/huggingface/diffusers/blob/d3d22ce5a894becb951eec03e663951b28d45135/src/diffusers/models/attention.py#L87 + let scale = f64::powf(self.channels as f64 / self.num_heads as f64, -0.25); + let attention_scores = (query_states * scale)?.matmul(&(key_states.t()? * scale)?)?; + let attention_probs = nn::ops::softmax(&attention_scores, D::Minus1)?; + + // TODO: revert the call to force_contiguous once the three matmul kernels have been + // adapted to handle layout with some dims set to 1. + let xs = attention_probs.matmul(&value_states)?; + let xs = xs.to_dtype(in_dtype)?; + let xs = xs.transpose(1, 2)?.contiguous()?; + let xs = xs.flatten_from(D::Minus2)?; + let xs = self + .proj_attn + .forward(&xs)? + .t()? + .reshape((batch, channel, height, width))?; + (xs + residual)? / self.config.rescale_output_factor + } +} diff --git a/patches/candle-transformers/src/models/stable_diffusion/clip.rs b/patches/candle-transformers/src/models/stable_diffusion/clip.rs new file mode 100644 index 0000000000..4c3f9d512d --- /dev/null +++ b/patches/candle-transformers/src/models/stable_diffusion/clip.rs @@ -0,0 +1,428 @@ +//! Contrastive Language-Image Pre-Training +//! +//! Contrastive Language-Image Pre-Training (CLIP) is an architecture trained on +//! pairs of images with related texts. +//! +//! - [CLIP](https://github.com/openai/CLIP) +use candle::{DType, Device, Result, Tensor, D}; +use candle_nn as nn; +use candle_nn::Module; + +#[derive(Debug, Clone, Copy)] +pub enum Activation { + QuickGelu, + Gelu, + GeluErf, +} + +impl Module for Activation { + fn forward(&self, xs: &Tensor) -> Result { + match self { + Activation::QuickGelu => xs * nn::ops::sigmoid(&(xs * 1.702f64)?)?, + Activation::Gelu => xs.gelu(), + Activation::GeluErf => xs.gelu_erf(), + } + } +} + +#[derive(Debug, Clone)] +pub struct Config { + vocab_size: usize, + embed_dim: usize, // aka config.hidden_size + activation: Activation, // aka config.hidden_act + intermediate_size: usize, + pub max_position_embeddings: usize, + // The character to use for padding, use EOS when not set. + pub pad_with: Option, + num_hidden_layers: usize, + num_attention_heads: usize, + #[allow(dead_code)] + projection_dim: usize, +} + +impl Config { + // The config details can be found in the "text_config" section of this json file: + // https://huggingface.co/openai/clip-vit-large-patch14/blob/main/config.json + pub fn v1_5() -> Self { + Self { + vocab_size: 49408, + embed_dim: 768, + intermediate_size: 3072, + max_position_embeddings: 77, + pad_with: None, + num_hidden_layers: 12, + num_attention_heads: 12, + projection_dim: 768, + activation: Activation::QuickGelu, + } + } + + // https://huggingface.co/stabilityai/stable-diffusion-2-1/blob/main/text_encoder/config.json + pub fn v2_1() -> Self { + Self { + vocab_size: 49408, + embed_dim: 1024, + intermediate_size: 4096, + max_position_embeddings: 77, + pad_with: Some("!".to_string()), + num_hidden_layers: 23, + num_attention_heads: 16, + projection_dim: 512, + activation: Activation::Gelu, + } + } + + // https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/text_encoder/config.json + pub fn sdxl() -> Self { + Self { + vocab_size: 49408, + embed_dim: 768, + intermediate_size: 3072, + max_position_embeddings: 77, + pad_with: Some("!".to_string()), + num_hidden_layers: 12, + num_attention_heads: 12, + projection_dim: 768, + activation: Activation::QuickGelu, + } + } + + // https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/text_encoder_2/config.json + pub fn sdxl2() -> Self { + Self { + vocab_size: 49408, + embed_dim: 1280, + intermediate_size: 5120, + max_position_embeddings: 77, + pad_with: Some("!".to_string()), + num_hidden_layers: 32, + num_attention_heads: 20, + projection_dim: 1280, + activation: Activation::Gelu, + } + } + + pub fn ssd1b() -> Self { + Self::sdxl() + } + + pub fn ssd1b2() -> Self { + Self::sdxl2() + } + + // https://huggingface.co/warp-ai/wuerstchen/blob/main/text_encoder/config.json + pub fn wuerstchen() -> Self { + Self { + vocab_size: 49408, + embed_dim: 1024, + intermediate_size: 4096, + max_position_embeddings: 77, + pad_with: None, + num_hidden_layers: 24, + num_attention_heads: 16, + projection_dim: 1024, + activation: Activation::GeluErf, + } + } + + // https://huggingface.co/warp-ai/wuerstchen-prior/blob/main/text_encoder/config.json + pub fn wuerstchen_prior() -> Self { + Self { + vocab_size: 49408, + embed_dim: 1280, + intermediate_size: 5120, + max_position_embeddings: 77, + pad_with: None, + num_hidden_layers: 32, + num_attention_heads: 20, + projection_dim: 512, + activation: Activation::GeluErf, + } + } +} + +// CLIP Text Model +// https://github.com/huggingface/transformers/blob/674f750a57431222fa2832503a108df3badf1564/src/transformers/models/clip/modeling_clip.py +#[derive(Debug)] +struct ClipTextEmbeddings { + token_embedding: candle_nn::Embedding, + position_embedding: candle_nn::Embedding, + position_ids: Tensor, +} + +impl ClipTextEmbeddings { + fn new(vs: candle_nn::VarBuilder, c: &Config) -> Result { + let token_embedding = + candle_nn::embedding(c.vocab_size, c.embed_dim, vs.pp("token_embedding"))?; + let position_embedding = candle_nn::embedding( + c.max_position_embeddings, + c.embed_dim, + vs.pp("position_embedding"), + )?; + let position_ids = + Tensor::arange(0u32, c.max_position_embeddings as u32, vs.device())?.unsqueeze(0)?; + Ok(ClipTextEmbeddings { + token_embedding, + position_embedding, + position_ids, + }) + } +} + +impl Module for ClipTextEmbeddings { + fn forward(&self, xs: &Tensor) -> Result { + let token_embedding = self.token_embedding.forward(xs)?; + let position_embedding = self.position_embedding.forward(&self.position_ids)?; + token_embedding.broadcast_add(&position_embedding) + } +} + +#[derive(Debug)] +struct ClipAttention { + k_proj: candle_nn::Linear, + v_proj: candle_nn::Linear, + q_proj: candle_nn::Linear, + out_proj: candle_nn::Linear, + head_dim: usize, + scale: f64, + num_attention_heads: usize, +} + +impl ClipAttention { + fn new(vs: candle_nn::VarBuilder, c: &Config) -> Result { + let embed_dim = c.embed_dim; + let num_attention_heads = c.num_attention_heads; + let k_proj = candle_nn::linear(embed_dim, embed_dim, vs.pp("k_proj"))?; + let v_proj = candle_nn::linear(embed_dim, embed_dim, vs.pp("v_proj"))?; + let q_proj = candle_nn::linear(embed_dim, embed_dim, vs.pp("q_proj"))?; + let out_proj = candle_nn::linear(embed_dim, embed_dim, vs.pp("out_proj"))?; + let head_dim = embed_dim / num_attention_heads; + let scale = (head_dim as f64).powf(-0.5); + Ok(ClipAttention { + k_proj, + v_proj, + q_proj, + out_proj, + head_dim, + scale, + num_attention_heads, + }) + } + + fn shape(&self, xs: &Tensor, seq_len: usize, bsz: usize) -> Result { + xs.reshape((bsz, seq_len, self.num_attention_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous() + } + + fn forward(&self, xs: &Tensor, causal_attention_mask: &Tensor) -> Result { + let in_dtype = xs.dtype(); + let (bsz, seq_len, embed_dim) = xs.dims3()?; + let query_states = (self.q_proj.forward(xs)? * self.scale)?; + let proj_shape = (bsz * self.num_attention_heads, seq_len, self.head_dim); + let query_states = self + .shape(&query_states, seq_len, bsz)? + .reshape(proj_shape)? + .to_dtype(DType::F32)?; + let key_states = self + .shape(&self.k_proj.forward(xs)?, seq_len, bsz)? + .reshape(proj_shape)? + .to_dtype(DType::F32)?; + let value_states = self + .shape(&self.v_proj.forward(xs)?, seq_len, bsz)? + .reshape(proj_shape)? + .to_dtype(DType::F32)?; + let attn_weights = query_states.matmul(&key_states.transpose(1, 2)?)?; + + let src_len = key_states.dim(1)?; + let attn_weights = attn_weights + .reshape((bsz, self.num_attention_heads, seq_len, src_len))? + .broadcast_add(causal_attention_mask)?; + let attn_weights = + attn_weights.reshape((bsz * self.num_attention_heads, seq_len, src_len))?; + let attn_weights = candle_nn::ops::softmax(&attn_weights, D::Minus1)?; + + let attn_output = attn_weights.matmul(&value_states)?.to_dtype(in_dtype)?; + let attn_output = attn_output + .reshape((bsz, self.num_attention_heads, seq_len, self.head_dim))? + .transpose(1, 2)? + .reshape((bsz, seq_len, embed_dim))?; + self.out_proj.forward(&attn_output) + } +} + +#[derive(Debug)] +struct ClipMlp { + fc1: candle_nn::Linear, + fc2: candle_nn::Linear, + activation: Activation, +} + +impl ClipMlp { + fn new(vs: candle_nn::VarBuilder, c: &Config) -> Result { + let fc1 = candle_nn::linear(c.embed_dim, c.intermediate_size, vs.pp("fc1"))?; + let fc2 = candle_nn::linear(c.intermediate_size, c.embed_dim, vs.pp("fc2"))?; + Ok(ClipMlp { + fc1, + fc2, + activation: c.activation, + }) + } +} + +impl ClipMlp { + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.fc1.forward(xs)?; + self.fc2.forward(&self.activation.forward(&xs)?) + } +} + +#[derive(Debug)] +struct ClipEncoderLayer { + self_attn: ClipAttention, + layer_norm1: candle_nn::LayerNorm, + mlp: ClipMlp, + layer_norm2: candle_nn::LayerNorm, +} + +impl ClipEncoderLayer { + fn new(vs: candle_nn::VarBuilder, c: &Config) -> Result { + let self_attn = ClipAttention::new(vs.pp("self_attn"), c)?; + let layer_norm1 = candle_nn::layer_norm(c.embed_dim, 1e-5, vs.pp("layer_norm1"))?; + let mlp = ClipMlp::new(vs.pp("mlp"), c)?; + let layer_norm2 = candle_nn::layer_norm(c.embed_dim, 1e-5, vs.pp("layer_norm2"))?; + Ok(ClipEncoderLayer { + self_attn, + layer_norm1, + mlp, + layer_norm2, + }) + } + + fn forward(&self, xs: &Tensor, causal_attention_mask: &Tensor) -> Result { + let residual = xs; + let xs = self.layer_norm1.forward(xs)?; + let xs = self.self_attn.forward(&xs, causal_attention_mask)?; + let xs = (xs + residual)?; + + let residual = &xs; + let xs = self.layer_norm2.forward(&xs)?; + let xs = self.mlp.forward(&xs)?; + xs + residual + } +} + +#[derive(Debug)] +struct ClipEncoder { + layers: Vec, +} + +impl ClipEncoder { + fn new(vs: candle_nn::VarBuilder, c: &Config) -> Result { + let vs = vs.pp("layers"); + let mut layers: Vec = Vec::new(); + for index in 0..c.num_hidden_layers { + let layer = ClipEncoderLayer::new(vs.pp(index.to_string()), c)?; + layers.push(layer) + } + Ok(ClipEncoder { layers }) + } + + fn forward(&self, xs: &Tensor, causal_attention_mask: &Tensor) -> Result { + let mut xs = xs.clone(); + for layer in self.layers.iter() { + xs = layer.forward(&xs, causal_attention_mask)?; + } + Ok(xs) + } +} + +/// A CLIP transformer based model. +#[derive(Debug)] +pub struct ClipTextTransformer { + embeddings: ClipTextEmbeddings, + encoder: ClipEncoder, + final_layer_norm: candle_nn::LayerNorm, +} + +impl ClipTextTransformer { + pub fn new(vs: candle_nn::VarBuilder, c: &Config) -> Result { + let vs = vs.pp("text_model"); + let embeddings = ClipTextEmbeddings::new(vs.pp("embeddings"), c)?; + let encoder = ClipEncoder::new(vs.pp("encoder"), c)?; + let final_layer_norm = candle_nn::layer_norm(c.embed_dim, 1e-5, vs.pp("final_layer_norm"))?; + Ok(ClipTextTransformer { + embeddings, + encoder, + final_layer_norm, + }) + } + + // https://github.com/huggingface/transformers/blob/674f750a57431222fa2832503a108df3badf1564/src/transformers/models/clip/modeling_clip.py#L678 + fn build_causal_attention_mask( + bsz: usize, + seq_len: usize, + mask_after: usize, + device: &Device, + ) -> Result { + let mask: Vec<_> = (0..seq_len) + .flat_map(|i| { + (0..seq_len).map(move |j| { + if j > i || j > mask_after { + f32::MIN + } else { + 0. + } + }) + }) + .collect(); + let mask = Tensor::from_slice(&mask, (seq_len, seq_len), device)?; + mask.broadcast_as((bsz, seq_len, seq_len)) + } + + pub fn forward_with_mask(&self, xs: &Tensor, mask_after: usize) -> Result { + let (bsz, seq_len) = xs.dims2()?; + let xs = self.embeddings.forward(xs)?; + let causal_attention_mask = + Self::build_causal_attention_mask(bsz, seq_len, mask_after, xs.device())?; + let xs = self.encoder.forward(&xs, &causal_attention_mask)?; + self.final_layer_norm.forward(&xs) + } + + pub fn forward_until_encoder_layer( + &self, + xs: &Tensor, + mask_after: usize, + until_layer: isize, + ) -> Result<(Tensor, Tensor)> { + let (bsz, seq_len) = xs.dims2()?; + let xs = self.embeddings.forward(xs)?; + let causal_attention_mask = + Self::build_causal_attention_mask(bsz, seq_len, mask_after, xs.device())?; + + let mut xs = xs.clone(); + let mut intermediate = xs.clone(); + + // Modified encoder.forward that returns the intermediate tensor along with final output. + let until_layer = if until_layer < 0 { + self.encoder.layers.len() as isize + until_layer + } else { + until_layer + } as usize; + + for (layer_id, layer) in self.encoder.layers.iter().enumerate() { + xs = layer.forward(&xs, &causal_attention_mask)?; + if layer_id == until_layer { + intermediate = xs.clone(); + } + } + + Ok((self.final_layer_norm.forward(&xs)?, intermediate)) + } +} + +impl Module for ClipTextTransformer { + fn forward(&self, xs: &Tensor) -> Result { + self.forward_with_mask(xs, usize::MAX) + } +} diff --git a/patches/candle-transformers/src/models/stable_diffusion/ddim.rs b/patches/candle-transformers/src/models/stable_diffusion/ddim.rs new file mode 100644 index 0000000000..d8ef5ec9bb --- /dev/null +++ b/patches/candle-transformers/src/models/stable_diffusion/ddim.rs @@ -0,0 +1,208 @@ +//! # Denoising Diffusion Implicit Models +//! +//! The Denoising Diffusion Implicit Models (DDIM) is a simple scheduler +//! similar to Denoising Diffusion Probabilistic Models (DDPM). The DDPM +//! generative process is the reverse of a Markovian process, DDIM generalizes +//! this to non-Markovian guidance. +//! +//! Denoising Diffusion Implicit Models, J. Song et al, 2020. +//! https://arxiv.org/abs/2010.02502 +use super::schedulers::{ + betas_for_alpha_bar, BetaSchedule, PredictionType, Scheduler, SchedulerConfig, TimestepSpacing, +}; +use candle::{Result, Tensor}; + +/// The configuration for the DDIM scheduler. +#[derive(Debug, Clone, Copy)] +pub struct DDIMSchedulerConfig { + /// The value of beta at the beginning of training. + pub beta_start: f64, + /// The value of beta at the end of training. + pub beta_end: f64, + /// How beta evolved during training. + pub beta_schedule: BetaSchedule, + /// The amount of noise to be added at each step. + pub eta: f64, + /// Adjust the indexes of the inference schedule by this value. + pub steps_offset: usize, + /// prediction type of the scheduler function, one of `epsilon` (predicting + /// the noise of the diffusion process), `sample` (directly predicting the noisy sample`) + /// or `v_prediction` (see section 2.4 https://imagen.research.google/video/paper.pdf) + pub prediction_type: PredictionType, + /// number of diffusion steps used to train the model + pub train_timesteps: usize, + /// time step spacing for the diffusion process + pub timestep_spacing: TimestepSpacing, +} + +impl Default for DDIMSchedulerConfig { + fn default() -> Self { + Self { + beta_start: 0.00085f64, + beta_end: 0.012f64, + beta_schedule: BetaSchedule::ScaledLinear, + eta: 0., + steps_offset: 1, + prediction_type: PredictionType::Epsilon, + train_timesteps: 1000, + timestep_spacing: TimestepSpacing::Leading, + } + } +} + +impl SchedulerConfig for DDIMSchedulerConfig { + fn build(&self, inference_steps: usize) -> Result> { + Ok(Box::new(DDIMScheduler::new(inference_steps, *self)?)) + } +} + +/// The DDIM scheduler. +#[derive(Debug, Clone)] +pub struct DDIMScheduler { + timesteps: Vec, + alphas_cumprod: Vec, + step_ratio: usize, + init_noise_sigma: f64, + pub config: DDIMSchedulerConfig, +} + +// clip_sample: False, set_alpha_to_one: False +impl DDIMScheduler { + /// Creates a new DDIM scheduler given the number of steps to be + /// used for inference as well as the number of steps that was used + /// during training. + fn new(inference_steps: usize, config: DDIMSchedulerConfig) -> Result { + let step_ratio = config.train_timesteps / inference_steps; + let timesteps: Vec = match config.timestep_spacing { + TimestepSpacing::Leading => (0..(inference_steps)) + .map(|s| s * step_ratio + config.steps_offset) + .rev() + .collect(), + TimestepSpacing::Trailing => std::iter::successors(Some(config.train_timesteps), |n| { + if *n > step_ratio { + Some(n - step_ratio) + } else { + None + } + }) + .map(|n| n - 1) + .collect(), + TimestepSpacing::Linspace => { + super::utils::linspace(0.0, (config.train_timesteps - 1) as f64, inference_steps)? + .to_vec1::()? + .iter() + .map(|&f| f as usize) + .rev() + .collect() + } + }; + + let betas = match config.beta_schedule { + BetaSchedule::ScaledLinear => super::utils::linspace( + config.beta_start.sqrt(), + config.beta_end.sqrt(), + config.train_timesteps, + )? + .sqr()?, + BetaSchedule::Linear => { + super::utils::linspace(config.beta_start, config.beta_end, config.train_timesteps)? + } + BetaSchedule::SquaredcosCapV2 => betas_for_alpha_bar(config.train_timesteps, 0.999)?, + }; + let betas = betas.to_vec1::()?; + let mut alphas_cumprod = Vec::with_capacity(betas.len()); + for &beta in betas.iter() { + let alpha = 1.0 - beta; + alphas_cumprod.push(alpha * *alphas_cumprod.last().unwrap_or(&1f64)) + } + Ok(Self { + alphas_cumprod, + timesteps, + step_ratio, + init_noise_sigma: 1., + config, + }) + } +} + +impl Scheduler for DDIMScheduler { + /// Performs a backward step during inference. + fn step(&mut self, model_output: &Tensor, timestep: usize, sample: &Tensor) -> Result { + let timestep = if timestep >= self.alphas_cumprod.len() { + timestep - 1 + } else { + timestep + }; + // https://github.com/huggingface/diffusers/blob/6e099e2c8ce4c4f5c7318e970a8c093dc5c7046e/src/diffusers/schedulers/scheduling_ddim.py#L195 + let prev_timestep = timestep.saturating_sub(self.step_ratio); + let alpha_prod_t = self.alphas_cumprod[timestep]; + let alpha_prod_t_prev = self.alphas_cumprod[prev_timestep]; + let beta_prod_t = 1. - alpha_prod_t; + let beta_prod_t_prev = 1. - alpha_prod_t_prev; + + let (pred_original_sample, pred_epsilon) = match self.config.prediction_type { + PredictionType::Epsilon => { + let pred_original_sample = ((sample - (model_output * beta_prod_t.sqrt())?)? + * (1. / alpha_prod_t.sqrt()))?; + (pred_original_sample, model_output.clone()) + } + PredictionType::VPrediction => { + let pred_original_sample = + ((sample * alpha_prod_t.sqrt())? - (model_output * beta_prod_t.sqrt())?)?; + let pred_epsilon = + ((model_output * alpha_prod_t.sqrt())? + (sample * beta_prod_t.sqrt())?)?; + (pred_original_sample, pred_epsilon) + } + PredictionType::Sample => { + let pred_original_sample = model_output.clone(); + let pred_epsilon = ((sample - &pred_original_sample * alpha_prod_t.sqrt())? + * (1. / beta_prod_t.sqrt()))?; + (pred_original_sample, pred_epsilon) + } + }; + + let variance = (beta_prod_t_prev / beta_prod_t) * (1. - alpha_prod_t / alpha_prod_t_prev); + let std_dev_t = self.config.eta * variance.sqrt(); + + let pred_sample_direction = + (pred_epsilon * (1. - alpha_prod_t_prev - std_dev_t * std_dev_t).sqrt())?; + let prev_sample = + ((pred_original_sample * alpha_prod_t_prev.sqrt())? + pred_sample_direction)?; + if self.config.eta > 0. { + &prev_sample + + Tensor::randn( + 0f32, + std_dev_t as f32, + prev_sample.shape(), + prev_sample.device(), + )? + } else { + Ok(prev_sample) + } + } + + /// Ensures interchangeability with schedulers that need to scale the denoising model input + /// depending on the current timestep. + fn scale_model_input(&self, sample: Tensor, _timestep: usize) -> Result { + Ok(sample) + } + + fn timesteps(&self) -> &[usize] { + self.timesteps.as_slice() + } + + fn add_noise(&self, original: &Tensor, noise: Tensor, timestep: usize) -> Result { + let timestep = if timestep >= self.alphas_cumprod.len() { + timestep - 1 + } else { + timestep + }; + let sqrt_alpha_prod = self.alphas_cumprod[timestep].sqrt(); + let sqrt_one_minus_alpha_prod = (1.0 - self.alphas_cumprod[timestep]).sqrt(); + (original * sqrt_alpha_prod)? + (noise * sqrt_one_minus_alpha_prod)? + } + + fn init_noise_sigma(&self) -> f64 { + self.init_noise_sigma + } +} diff --git a/patches/candle-transformers/src/models/stable_diffusion/ddpm.rs b/patches/candle-transformers/src/models/stable_diffusion/ddpm.rs new file mode 100644 index 0000000000..c7cc7a9a80 --- /dev/null +++ b/patches/candle-transformers/src/models/stable_diffusion/ddpm.rs @@ -0,0 +1,200 @@ +use super::schedulers::{betas_for_alpha_bar, BetaSchedule, PredictionType}; +use candle::{Result, Tensor}; + +#[derive(Debug, Default, Clone, PartialEq, Eq)] +pub enum DDPMVarianceType { + #[default] + FixedSmall, + FixedSmallLog, + FixedLarge, + FixedLargeLog, + Learned, +} + +#[derive(Debug, Clone)] +pub struct DDPMSchedulerConfig { + /// The value of beta at the beginning of training. + pub beta_start: f64, + /// The value of beta at the end of training. + pub beta_end: f64, + /// How beta evolved during training. + pub beta_schedule: BetaSchedule, + /// Option to predicted sample between -1 and 1 for numerical stability. + pub clip_sample: bool, + /// Option to clip the variance used when adding noise to the denoised sample. + pub variance_type: DDPMVarianceType, + /// prediction type of the scheduler function + pub prediction_type: PredictionType, + /// number of diffusion steps used to train the model. + pub train_timesteps: usize, +} + +impl Default for DDPMSchedulerConfig { + fn default() -> Self { + Self { + beta_start: 0.00085, + beta_end: 0.012, + beta_schedule: BetaSchedule::ScaledLinear, + clip_sample: false, + variance_type: DDPMVarianceType::FixedSmall, + prediction_type: PredictionType::Epsilon, + train_timesteps: 1000, + } + } +} + +pub struct DDPMScheduler { + alphas_cumprod: Vec, + init_noise_sigma: f64, + timesteps: Vec, + step_ratio: usize, + pub config: DDPMSchedulerConfig, +} + +impl DDPMScheduler { + pub fn new(inference_steps: usize, config: DDPMSchedulerConfig) -> Result { + let betas = match config.beta_schedule { + BetaSchedule::ScaledLinear => super::utils::linspace( + config.beta_start.sqrt(), + config.beta_end.sqrt(), + config.train_timesteps, + )? + .sqr()?, + BetaSchedule::Linear => { + super::utils::linspace(config.beta_start, config.beta_end, config.train_timesteps)? + } + BetaSchedule::SquaredcosCapV2 => betas_for_alpha_bar(config.train_timesteps, 0.999)?, + }; + + let betas = betas.to_vec1::()?; + let mut alphas_cumprod = Vec::with_capacity(betas.len()); + for &beta in betas.iter() { + let alpha = 1.0 - beta; + alphas_cumprod.push(alpha * *alphas_cumprod.last().unwrap_or(&1f64)) + } + + // min(train_timesteps, inference_steps) + // https://github.com/huggingface/diffusers/blob/8331da46837be40f96fbd24de6a6fb2da28acd11/src/diffusers/schedulers/scheduling_ddpm.py#L187 + let inference_steps = inference_steps.min(config.train_timesteps); + // arange the number of the scheduler's timesteps + let step_ratio = config.train_timesteps / inference_steps; + let timesteps: Vec = (0..inference_steps).map(|s| s * step_ratio).rev().collect(); + + Ok(Self { + alphas_cumprod, + init_noise_sigma: 1.0, + timesteps, + step_ratio, + config, + }) + } + + fn get_variance(&self, timestep: usize) -> f64 { + let prev_t = timestep as isize - self.step_ratio as isize; + let alpha_prod_t = self.alphas_cumprod[timestep]; + let alpha_prod_t_prev = if prev_t >= 0 { + self.alphas_cumprod[prev_t as usize] + } else { + 1.0 + }; + let current_beta_t = 1. - alpha_prod_t / alpha_prod_t_prev; + + // For t > 0, compute predicted variance βt (see formula (6) and (7) from [the pdf](https://arxiv.org/pdf/2006.11239.pdf)) + // and sample from it to get previous sample + // x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample + let variance = (1. - alpha_prod_t_prev) / (1. - alpha_prod_t) * current_beta_t; + + // retrieve variance + match self.config.variance_type { + DDPMVarianceType::FixedSmall => variance.max(1e-20), + // for rl-diffuser https://arxiv.org/abs/2205.09991 + DDPMVarianceType::FixedSmallLog => { + let variance = variance.max(1e-20).ln(); + (variance * 0.5).exp() + } + DDPMVarianceType::FixedLarge => current_beta_t, + DDPMVarianceType::FixedLargeLog => current_beta_t.ln(), + DDPMVarianceType::Learned => variance, + } + } + + pub fn timesteps(&self) -> &[usize] { + self.timesteps.as_slice() + } + + /// Ensures interchangeability with schedulers that need to scale the denoising model input + /// depending on the current timestep. + pub fn scale_model_input(&self, sample: Tensor, _timestep: usize) -> Tensor { + sample + } + + pub fn step(&self, model_output: &Tensor, timestep: usize, sample: &Tensor) -> Result { + let prev_t = timestep as isize - self.step_ratio as isize; + + // https://github.com/huggingface/diffusers/blob/df2b548e893ccb8a888467c2508756680df22821/src/diffusers/schedulers/scheduling_ddpm.py#L272 + // 1. compute alphas, betas + let alpha_prod_t = self.alphas_cumprod[timestep]; + let alpha_prod_t_prev = if prev_t >= 0 { + self.alphas_cumprod[prev_t as usize] + } else { + 1.0 + }; + let beta_prod_t = 1. - alpha_prod_t; + let beta_prod_t_prev = 1. - alpha_prod_t_prev; + let current_alpha_t = alpha_prod_t / alpha_prod_t_prev; + let current_beta_t = 1. - current_alpha_t; + + // 2. compute predicted original sample from predicted noise also called "predicted x_0" of formula (15) + let mut pred_original_sample = match self.config.prediction_type { + PredictionType::Epsilon => { + ((sample - model_output * beta_prod_t.sqrt())? / alpha_prod_t.sqrt())? + } + PredictionType::Sample => model_output.clone(), + PredictionType::VPrediction => { + ((sample * alpha_prod_t.sqrt())? - model_output * beta_prod_t.sqrt())? + } + }; + + // 3. clip predicted x_0 + if self.config.clip_sample { + pred_original_sample = pred_original_sample.clamp(-1f32, 1f32)?; + } + + // 4. Compute coefficients for pred_original_sample x_0 and current sample x_t + // See formula (7) from https://arxiv.org/pdf/2006.11239.pdf + let pred_original_sample_coeff = (alpha_prod_t_prev.sqrt() * current_beta_t) / beta_prod_t; + let current_sample_coeff = current_alpha_t.sqrt() * beta_prod_t_prev / beta_prod_t; + + // 5. Compute predicted previous sample µ_t + // See formula (7) from https://arxiv.org/pdf/2006.11239.pdf + let pred_prev_sample = ((&pred_original_sample * pred_original_sample_coeff)? + + sample * current_sample_coeff)?; + + // https://github.com/huggingface/diffusers/blob/df2b548e893ccb8a888467c2508756680df22821/src/diffusers/schedulers/scheduling_ddpm.py#L305 + // 6. Add noise + let mut variance = model_output.zeros_like()?; + if timestep > 0 { + let variance_noise = model_output.randn_like(0., 1.)?; + if self.config.variance_type == DDPMVarianceType::FixedSmallLog { + variance = (variance_noise * self.get_variance(timestep))?; + } else { + variance = (variance_noise * self.get_variance(timestep).sqrt())?; + } + } + &pred_prev_sample + variance + } + + pub fn add_noise( + &self, + original_samples: &Tensor, + noise: Tensor, + timestep: usize, + ) -> Result { + (original_samples * self.alphas_cumprod[timestep].sqrt())? + + noise * (1. - self.alphas_cumprod[timestep]).sqrt() + } + + pub fn init_noise_sigma(&self) -> f64 { + self.init_noise_sigma + } +} diff --git a/patches/candle-transformers/src/models/stable_diffusion/embeddings.rs b/patches/candle-transformers/src/models/stable_diffusion/embeddings.rs new file mode 100644 index 0000000000..0de5f9a7db --- /dev/null +++ b/patches/candle-transformers/src/models/stable_diffusion/embeddings.rs @@ -0,0 +1,65 @@ +use candle::{Result, Tensor, D}; +use candle_nn as nn; +use candle_nn::Module; + +#[derive(Debug)] +pub struct TimestepEmbedding { + linear_1: nn::Linear, + linear_2: nn::Linear, +} + +impl TimestepEmbedding { + // act_fn: "silu" + pub fn new(vs: nn::VarBuilder, channel: usize, time_embed_dim: usize) -> Result { + let linear_1 = nn::linear(channel, time_embed_dim, vs.pp("linear_1"))?; + let linear_2 = nn::linear(time_embed_dim, time_embed_dim, vs.pp("linear_2"))?; + Ok(Self { linear_1, linear_2 }) + } +} + +impl Module for TimestepEmbedding { + fn forward(&self, xs: &Tensor) -> Result { + let xs = nn::ops::silu(&self.linear_1.forward(xs)?)?; + self.linear_2.forward(&xs) + } +} + +#[derive(Debug)] +pub struct Timesteps { + num_channels: usize, + flip_sin_to_cos: bool, + downscale_freq_shift: f64, +} + +impl Timesteps { + pub fn new(num_channels: usize, flip_sin_to_cos: bool, downscale_freq_shift: f64) -> Self { + Self { + num_channels, + flip_sin_to_cos, + downscale_freq_shift, + } + } +} + +impl Module for Timesteps { + fn forward(&self, xs: &Tensor) -> Result { + let half_dim = (self.num_channels / 2) as u32; + let exponent = (Tensor::arange(0, half_dim, xs.device())?.to_dtype(candle::DType::F32)? + * -f64::ln(10000.))?; + let exponent = (exponent / (half_dim as f64 - self.downscale_freq_shift))?; + let emb = exponent.exp()?.to_dtype(xs.dtype())?; + // emb = timesteps[:, None].float() * emb[None, :] + let emb = xs.unsqueeze(D::Minus1)?.broadcast_mul(&emb.unsqueeze(0)?)?; + let (cos, sin) = (emb.cos()?, emb.sin()?); + let emb = if self.flip_sin_to_cos { + Tensor::cat(&[&cos, &sin], D::Minus1)? + } else { + Tensor::cat(&[&sin, &cos], D::Minus1)? + }; + if self.num_channels % 2 == 1 { + emb.pad_with_zeros(D::Minus2, 0, 1) + } else { + Ok(emb) + } + } +} diff --git a/patches/candle-transformers/src/models/stable_diffusion/euler_ancestral_discrete.rs b/patches/candle-transformers/src/models/stable_diffusion/euler_ancestral_discrete.rs new file mode 100644 index 0000000000..250161ccad --- /dev/null +++ b/patches/candle-transformers/src/models/stable_diffusion/euler_ancestral_discrete.rs @@ -0,0 +1,230 @@ +//! Ancestral sampling with Euler method steps. +//! +//! Based on the original [`k-diffusion` implementation by Katherine Crowson]( https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L72). +//! +use super::{ + schedulers::{ + betas_for_alpha_bar, BetaSchedule, PredictionType, Scheduler, SchedulerConfig, + TimestepSpacing, + }, + utils::interp, +}; +use candle::{bail, Error, Result, Tensor}; + +/// The configuration for the EulerAncestral Discrete scheduler. +#[derive(Debug, Clone, Copy)] +pub struct EulerAncestralDiscreteSchedulerConfig { + /// The value of beta at the beginning of training.n + pub beta_start: f64, + /// The value of beta at the end of training. + pub beta_end: f64, + /// How beta evolved during training. + pub beta_schedule: BetaSchedule, + /// Adjust the indexes of the inference schedule by this value. + pub steps_offset: usize, + /// prediction type of the scheduler function, one of `epsilon` (predicting + /// the noise of the diffusion process), `sample` (directly predicting the noisy sample`) + /// or `v_prediction` (see [section 2.4](https://imagen.research.google/video/paper.pdf)) + pub prediction_type: PredictionType, + /// number of diffusion steps used to train the model + pub train_timesteps: usize, + /// time step spacing for the diffusion process + pub timestep_spacing: TimestepSpacing, +} + +impl Default for EulerAncestralDiscreteSchedulerConfig { + fn default() -> Self { + Self { + beta_start: 0.00085f64, + beta_end: 0.012f64, + beta_schedule: BetaSchedule::ScaledLinear, + steps_offset: 1, + prediction_type: PredictionType::Epsilon, + train_timesteps: 1000, + timestep_spacing: TimestepSpacing::Leading, + } + } +} + +impl SchedulerConfig for EulerAncestralDiscreteSchedulerConfig { + fn build(&self, inference_steps: usize) -> Result> { + Ok(Box::new(EulerAncestralDiscreteScheduler::new( + inference_steps, + *self, + )?)) + } +} + +/// The EulerAncestral Discrete scheduler. +#[derive(Debug, Clone)] +pub struct EulerAncestralDiscreteScheduler { + timesteps: Vec, + sigmas: Vec, + init_noise_sigma: f64, + pub config: EulerAncestralDiscreteSchedulerConfig, +} + +// clip_sample: False, set_alpha_to_one: False +impl EulerAncestralDiscreteScheduler { + /// Creates a new EulerAncestral Discrete scheduler given the number of steps to be + /// used for inference as well as the number of steps that was used + /// during training. + pub fn new( + inference_steps: usize, + config: EulerAncestralDiscreteSchedulerConfig, + ) -> Result { + let step_ratio = config.train_timesteps / inference_steps; + let timesteps: Vec = match config.timestep_spacing { + TimestepSpacing::Leading => (0..(inference_steps)) + .map(|s| s * step_ratio + config.steps_offset) + .rev() + .collect(), + TimestepSpacing::Trailing => std::iter::successors(Some(config.train_timesteps), |n| { + if *n > step_ratio { + Some(n - step_ratio) + } else { + None + } + }) + .map(|n| n - 1) + .collect(), + TimestepSpacing::Linspace => { + super::utils::linspace(0.0, (config.train_timesteps - 1) as f64, inference_steps)? + .to_vec1::()? + .iter() + .map(|&f| f as usize) + .rev() + .collect() + } + }; + + let betas = match config.beta_schedule { + BetaSchedule::ScaledLinear => super::utils::linspace( + config.beta_start.sqrt(), + config.beta_end.sqrt(), + config.train_timesteps, + )? + .sqr()?, + BetaSchedule::Linear => { + super::utils::linspace(config.beta_start, config.beta_end, config.train_timesteps)? + } + BetaSchedule::SquaredcosCapV2 => betas_for_alpha_bar(config.train_timesteps, 0.999)?, + }; + let betas = betas.to_vec1::()?; + let mut alphas_cumprod = Vec::with_capacity(betas.len()); + for &beta in betas.iter() { + let alpha = 1.0 - beta; + alphas_cumprod.push(alpha * *alphas_cumprod.last().unwrap_or(&1f64)) + } + let sigmas: Vec = alphas_cumprod + .iter() + .map(|&f| ((1. - f) / f).sqrt()) + .collect(); + + let sigmas_xa: Vec<_> = (0..sigmas.len()).map(|i| i as f64).collect(); + + let mut sigmas_int = interp( + ×teps.iter().map(|&t| t as f64).collect::>(), + &sigmas_xa, + &sigmas, + ); + sigmas_int.push(0.0); + + // standard deviation of the initial noise distribution + // f64 does not implement Ord such that there is no `max`, so we need to use this workaround + let init_noise_sigma = *sigmas_int + .iter() + .chain(std::iter::once(&0.0)) + .reduce(|a, b| if a > b { a } else { b }) + .expect("init_noise_sigma could not be reduced from sigmas - this should never happen"); + + Ok(Self { + sigmas: sigmas_int, + timesteps, + init_noise_sigma, + config, + }) + } +} + +impl Scheduler for EulerAncestralDiscreteScheduler { + fn timesteps(&self) -> &[usize] { + self.timesteps.as_slice() + } + + /// Ensures interchangeability with schedulers that need to scale the denoising model input + /// depending on the current timestep. + /// + /// Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the K-LMS algorithm + fn scale_model_input(&self, sample: Tensor, timestep: usize) -> Result { + let step_index = match self.timesteps.iter().position(|&t| t == timestep) { + Some(i) => i, + None => bail!("timestep out of this schedulers bounds: {timestep}"), + }; + + let sigma = self + .sigmas + .get(step_index) + .expect("step_index out of sigma bounds - this shouldn't happen"); + + sample / ((sigma.powi(2) + 1.).sqrt()) + } + + /// Performs a backward step during inference. + fn step(&mut self, model_output: &Tensor, timestep: usize, sample: &Tensor) -> Result { + let step_index = self + .timesteps + .iter() + .position(|&p| p == timestep) + .ok_or_else(|| Error::Msg("timestep out of this schedulers bounds".to_string()))?; + + let sigma_from = &self.sigmas[step_index]; + let sigma_to = &self.sigmas[step_index + 1]; + + // 1. compute predicted original sample (x_0) from sigma-scaled predicted noise + let pred_original_sample = match self.config.prediction_type { + PredictionType::Epsilon => (sample - (model_output * *sigma_from))?, + PredictionType::VPrediction => { + ((model_output * (-sigma_from / (sigma_from.powi(2) + 1.0).sqrt()))? + + (sample / (sigma_from.powi(2) + 1.0))?)? + } + PredictionType::Sample => bail!("prediction_type not implemented yet: sample"), + }; + + let sigma_up = (sigma_to.powi(2) * (sigma_from.powi(2) - sigma_to.powi(2)) + / sigma_from.powi(2)) + .sqrt(); + let sigma_down = (sigma_to.powi(2) - sigma_up.powi(2)).sqrt(); + + // 2. convert to a ODE derivative + let derivative = ((sample - pred_original_sample)? / *sigma_from)?; + let dt = sigma_down - *sigma_from; + let prev_sample = (sample + derivative * dt)?; + + let noise = prev_sample.randn_like(0.0, 1.0)?; + + prev_sample + noise * sigma_up + } + + fn add_noise(&self, original: &Tensor, noise: Tensor, timestep: usize) -> Result { + let step_index = self + .timesteps + .iter() + .position(|&p| p == timestep) + .ok_or_else(|| Error::Msg("timestep out of this schedulers bounds".to_string()))?; + + let sigma = self + .sigmas + .get(step_index) + .expect("step_index out of sigma bounds - this shouldn't happen"); + + original + (noise * *sigma)? + } + + fn init_noise_sigma(&self) -> f64 { + match self.config.timestep_spacing { + TimestepSpacing::Trailing | TimestepSpacing::Linspace => self.init_noise_sigma, + TimestepSpacing::Leading => (self.init_noise_sigma.powi(2) + 1.0).sqrt(), + } + } +} diff --git a/patches/candle-transformers/src/models/stable_diffusion/mod.rs b/patches/candle-transformers/src/models/stable_diffusion/mod.rs new file mode 100644 index 0000000000..3c101fc69b --- /dev/null +++ b/patches/candle-transformers/src/models/stable_diffusion/mod.rs @@ -0,0 +1,527 @@ +//! Stable Diffusion +//! +//! Stable Diffusion is a latent text-to-image diffusion model capable of +//! generating photo-realistic images given any text input. +//! +//! - 💻 [Original Repository](https://github.com/CompVis/stable-diffusion) +//! - 🤗 [Hugging Face](https://huggingface.co/runwayml/stable-diffusion-v1-5) +//! - The default scheduler for the v1.5, v2.1 and XL 1.0 version is the Denoising Diffusion Implicit Model scheduler (DDIM). The original paper and some code can be found in the [associated repo](https://github.com/ermongroup/ddim). The default scheduler for the XL Turbo version is the Euler Ancestral scheduler. +//! +//! +//! # Example +//! +//!
+//! rusty robot holding a candle +//!
+//! +//! _"A rusty robot holding a fire torch in its hand."_ Generated by Stable Diffusion XL using Rust and [candle](https://github.com/huggingface/candle). +//! +//! ```bash +//! # example running with cuda +//! # see the candle-examples/examples/stable-diffusion for all options +//! cargo run --example stable-diffusion --release --features=cuda,cudnn \ +//! -- --prompt "a cosmonaut on a horse (hd, realistic, high-def)" +//! +//! # with sd-turbo +//! cargo run --example stable-diffusion --release --features=cuda,cudnn \ +//! -- --prompt "a cosmonaut on a horse (hd, realistic, high-def)" \ +//! --sd-version turbo +//! +//! # with flash attention. +//! # feature flag: `--features flash-attn` +//! # cli flag: `--use-flash-attn`. +//! # flash-attention-v2 is only compatible with Ampere, Ada, \ +//! # or Hopper GPUs (e.g., A100/H100, RTX 3090/4090). +//! cargo run --example stable-diffusion --release --features=cuda,cudnn \ +//! -- --prompt "a cosmonaut on a horse (hd, realistic, high-def)" \ +//! --use-flash-attn +//! ``` + +pub mod attention; +pub mod clip; +pub mod ddim; +pub mod ddpm; +pub mod embeddings; +pub mod euler_ancestral_discrete; +pub mod resnet; +pub mod schedulers; +pub mod unet_2d; +pub mod unet_2d_blocks; +pub mod uni_pc; +pub mod utils; +pub mod vae; + +use std::sync::Arc; + +use candle::{DType, Device, Result}; +use candle_nn as nn; + +use self::schedulers::{Scheduler, SchedulerConfig}; + +#[derive(Clone, Debug)] +pub struct StableDiffusionConfig { + pub width: usize, + pub height: usize, + pub clip: clip::Config, + pub clip2: Option, + autoencoder: vae::AutoEncoderKLConfig, + unet: unet_2d::UNet2DConditionModelConfig, + scheduler: Arc, +} + +impl StableDiffusionConfig { + pub fn v1_5( + sliced_attention_size: Option, + height: Option, + width: Option, + ) -> Self { + let bc = |out_channels, use_cross_attn, attention_head_dim| unet_2d::BlockConfig { + out_channels, + use_cross_attn, + attention_head_dim, + }; + // https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/unet/config.json + let unet = unet_2d::UNet2DConditionModelConfig { + blocks: vec![ + bc(320, Some(1), 8), + bc(640, Some(1), 8), + bc(1280, Some(1), 8), + bc(1280, None, 8), + ], + center_input_sample: false, + cross_attention_dim: 768, + downsample_padding: 1, + flip_sin_to_cos: true, + freq_shift: 0., + layers_per_block: 2, + mid_block_scale_factor: 1., + norm_eps: 1e-5, + norm_num_groups: 32, + sliced_attention_size, + use_linear_projection: false, + }; + let autoencoder = vae::AutoEncoderKLConfig { + block_out_channels: vec![128, 256, 512, 512], + layers_per_block: 2, + latent_channels: 4, + norm_num_groups: 32, + use_quant_conv: true, + use_post_quant_conv: true, + }; + let height = if let Some(height) = height { + assert_eq!(height % 8, 0, "height has to be divisible by 8"); + height + } else { + 512 + }; + + let width = if let Some(width) = width { + assert_eq!(width % 8, 0, "width has to be divisible by 8"); + width + } else { + 512 + }; + + let scheduler = Arc::new(ddim::DDIMSchedulerConfig { + prediction_type: schedulers::PredictionType::Epsilon, + ..Default::default() + }); + + StableDiffusionConfig { + width, + height, + clip: clip::Config::v1_5(), + clip2: None, + autoencoder, + scheduler, + unet, + } + } + + fn v2_1_( + sliced_attention_size: Option, + height: Option, + width: Option, + prediction_type: schedulers::PredictionType, + ) -> Self { + let bc = |out_channels, use_cross_attn, attention_head_dim| unet_2d::BlockConfig { + out_channels, + use_cross_attn, + attention_head_dim, + }; + // https://huggingface.co/stabilityai/stable-diffusion-2-1/blob/main/unet/config.json + let unet = unet_2d::UNet2DConditionModelConfig { + blocks: vec![ + bc(320, Some(1), 5), + bc(640, Some(1), 10), + bc(1280, Some(1), 20), + bc(1280, None, 20), + ], + center_input_sample: false, + cross_attention_dim: 1024, + downsample_padding: 1, + flip_sin_to_cos: true, + freq_shift: 0., + layers_per_block: 2, + mid_block_scale_factor: 1., + norm_eps: 1e-5, + norm_num_groups: 32, + sliced_attention_size, + use_linear_projection: true, + }; + // https://huggingface.co/stabilityai/stable-diffusion-2-1/blob/main/vae/config.json + let autoencoder = vae::AutoEncoderKLConfig { + block_out_channels: vec![128, 256, 512, 512], + layers_per_block: 2, + latent_channels: 4, + norm_num_groups: 32, + use_quant_conv: true, + use_post_quant_conv: true, + }; + let scheduler = Arc::new(ddim::DDIMSchedulerConfig { + prediction_type, + ..Default::default() + }); + + let height = if let Some(height) = height { + assert_eq!(height % 8, 0, "height has to be divisible by 8"); + height + } else { + 768 + }; + + let width = if let Some(width) = width { + assert_eq!(width % 8, 0, "width has to be divisible by 8"); + width + } else { + 768 + }; + + StableDiffusionConfig { + width, + height, + clip: clip::Config::v2_1(), + clip2: None, + autoencoder, + scheduler, + unet, + } + } + + pub fn v2_1( + sliced_attention_size: Option, + height: Option, + width: Option, + ) -> Self { + // https://huggingface.co/stabilityai/stable-diffusion-2-1/blob/main/scheduler/scheduler_config.json + Self::v2_1_( + sliced_attention_size, + height, + width, + schedulers::PredictionType::VPrediction, + ) + } + + fn sdxl_( + sliced_attention_size: Option, + height: Option, + width: Option, + prediction_type: schedulers::PredictionType, + ) -> Self { + let bc = |out_channels, use_cross_attn, attention_head_dim| unet_2d::BlockConfig { + out_channels, + use_cross_attn, + attention_head_dim, + }; + // https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/unet/config.json + let unet = unet_2d::UNet2DConditionModelConfig { + blocks: vec![ + bc(320, None, 5), + bc(640, Some(2), 10), + bc(1280, Some(10), 20), + ], + center_input_sample: false, + cross_attention_dim: 2048, + downsample_padding: 1, + flip_sin_to_cos: true, + freq_shift: 0., + layers_per_block: 2, + mid_block_scale_factor: 1., + norm_eps: 1e-5, + norm_num_groups: 32, + sliced_attention_size, + use_linear_projection: true, + }; + // https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/vae/config.json + let autoencoder = vae::AutoEncoderKLConfig { + block_out_channels: vec![128, 256, 512, 512], + layers_per_block: 2, + latent_channels: 4, + norm_num_groups: 32, + use_quant_conv: true, + use_post_quant_conv: true, + }; + let scheduler = Arc::new(ddim::DDIMSchedulerConfig { + prediction_type, + ..Default::default() + }); + + let height = if let Some(height) = height { + assert_eq!(height % 8, 0, "height has to be divisible by 8"); + height + } else { + 1024 + }; + + let width = if let Some(width) = width { + assert_eq!(width % 8, 0, "width has to be divisible by 8"); + width + } else { + 1024 + }; + + StableDiffusionConfig { + width, + height, + clip: clip::Config::sdxl(), + clip2: Some(clip::Config::sdxl2()), + autoencoder, + scheduler, + unet, + } + } + + fn sdxl_turbo_( + sliced_attention_size: Option, + height: Option, + width: Option, + prediction_type: schedulers::PredictionType, + ) -> Self { + let bc = |out_channels, use_cross_attn, attention_head_dim| unet_2d::BlockConfig { + out_channels, + use_cross_attn, + attention_head_dim, + }; + // https://huggingface.co/stabilityai/sdxl-turbo/blob/main/unet/config.json + let unet = unet_2d::UNet2DConditionModelConfig { + blocks: vec![ + bc(320, None, 5), + bc(640, Some(2), 10), + bc(1280, Some(10), 20), + ], + center_input_sample: false, + cross_attention_dim: 2048, + downsample_padding: 1, + flip_sin_to_cos: true, + freq_shift: 0., + layers_per_block: 2, + mid_block_scale_factor: 1., + norm_eps: 1e-5, + norm_num_groups: 32, + sliced_attention_size, + use_linear_projection: true, + }; + // https://huggingface.co/stabilityai/sdxl-turbo/blob/main/vae/config.json + let autoencoder = vae::AutoEncoderKLConfig { + block_out_channels: vec![128, 256, 512, 512], + layers_per_block: 2, + latent_channels: 4, + norm_num_groups: 32, + use_quant_conv: true, + use_post_quant_conv: true, + }; + let scheduler = Arc::new( + euler_ancestral_discrete::EulerAncestralDiscreteSchedulerConfig { + prediction_type, + timestep_spacing: schedulers::TimestepSpacing::Trailing, + ..Default::default() + }, + ); + + let height = if let Some(height) = height { + assert_eq!(height % 8, 0, "height has to be divisible by 8"); + height + } else { + 512 + }; + + let width = if let Some(width) = width { + assert_eq!(width % 8, 0, "width has to be divisible by 8"); + width + } else { + 512 + }; + + Self { + width, + height, + clip: clip::Config::sdxl(), + clip2: Some(clip::Config::sdxl2()), + autoencoder, + scheduler, + unet, + } + } + + pub fn sdxl( + sliced_attention_size: Option, + height: Option, + width: Option, + ) -> Self { + Self::sdxl_( + sliced_attention_size, + height, + width, + // https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/scheduler/scheduler_config.json + schedulers::PredictionType::Epsilon, + ) + } + + pub fn sdxl_turbo( + sliced_attention_size: Option, + height: Option, + width: Option, + ) -> Self { + Self::sdxl_turbo_( + sliced_attention_size, + height, + width, + // https://huggingface.co/stabilityai/sdxl-turbo/blob/main/scheduler/scheduler_config.json + schedulers::PredictionType::Epsilon, + ) + } + + pub fn ssd1b( + sliced_attention_size: Option, + height: Option, + width: Option, + ) -> Self { + let bc = |out_channels, use_cross_attn, attention_head_dim| unet_2d::BlockConfig { + out_channels, + use_cross_attn, + attention_head_dim, + }; + // https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/unet/config.json + let unet = unet_2d::UNet2DConditionModelConfig { + blocks: vec![ + bc(320, None, 5), + bc(640, Some(2), 10), + bc(1280, Some(10), 20), + ], + center_input_sample: false, + cross_attention_dim: 2048, + downsample_padding: 1, + flip_sin_to_cos: true, + freq_shift: 0., + layers_per_block: 2, + mid_block_scale_factor: 1., + norm_eps: 1e-5, + norm_num_groups: 32, + sliced_attention_size, + use_linear_projection: true, + }; + // https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/vae/config.json + let autoencoder = vae::AutoEncoderKLConfig { + block_out_channels: vec![128, 256, 512, 512], + layers_per_block: 2, + latent_channels: 4, + norm_num_groups: 32, + use_quant_conv: true, + use_post_quant_conv: true, + }; + let scheduler = Arc::new(ddim::DDIMSchedulerConfig { + ..Default::default() + }); + + let height = if let Some(height) = height { + assert_eq!(height % 8, 0, "height has to be divisible by 8"); + height + } else { + 1024 + }; + + let width = if let Some(width) = width { + assert_eq!(width % 8, 0, "width has to be divisible by 8"); + width + } else { + 1024 + }; + + Self { + width, + height, + clip: clip::Config::ssd1b(), + clip2: Some(clip::Config::ssd1b2()), + autoencoder, + scheduler, + unet, + } + } + + pub fn build_vae>( + &self, + vae_weights: P, + device: &Device, + dtype: DType, + ) -> Result { + let vs_ae = + unsafe { nn::VarBuilder::from_mmaped_safetensors(&[vae_weights], dtype, device)? }; + // https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/vae/config.json + let autoencoder = vae::AutoEncoderKL::new(vs_ae, 3, 3, self.autoencoder.clone())?; + Ok(autoencoder) + } + + pub fn build_unet>( + &self, + unet_weights: P, + device: &Device, + in_channels: usize, + use_flash_attn: bool, + dtype: DType, + ) -> Result { + let vs_unet = + unsafe { nn::VarBuilder::from_mmaped_safetensors(&[unet_weights], dtype, device)? }; + let unet = unet_2d::UNet2DConditionModel::new( + vs_unet, + in_channels, + 4, + use_flash_attn, + self.unet.clone(), + )?; + Ok(unet) + } + + pub fn build_unet_sharded>( + &self, + unet_weight_files: &[P], + device: &Device, + in_channels: usize, + use_flash_attn: bool, + dtype: DType, + ) -> Result { + let vs_unet = + unsafe { nn::VarBuilder::from_mmaped_safetensors(unet_weight_files, dtype, device)? }; + unet_2d::UNet2DConditionModel::new( + vs_unet, + in_channels, + 4, + use_flash_attn, + self.unet.clone(), + ) + } + + pub fn build_scheduler(&self, n_steps: usize) -> Result> { + self.scheduler.build(n_steps) + } +} + +pub fn build_clip_transformer>( + clip: &clip::Config, + clip_weights: P, + device: &Device, + dtype: DType, +) -> Result { + let vs = unsafe { nn::VarBuilder::from_mmaped_safetensors(&[clip_weights], dtype, device)? }; + let text_model = clip::ClipTextTransformer::new(vs, clip)?; + Ok(text_model) +} diff --git a/patches/candle-transformers/src/models/stable_diffusion/resnet.rs b/patches/candle-transformers/src/models/stable_diffusion/resnet.rs new file mode 100644 index 0000000000..8a6490c502 --- /dev/null +++ b/patches/candle-transformers/src/models/stable_diffusion/resnet.rs @@ -0,0 +1,141 @@ +//! ResNet Building Blocks +//! +//! Some Residual Network blocks used in UNet models. +//! +//! Denoising Diffusion Implicit Models, K. He and al, 2015. +//! - [Paper](https://arxiv.org/abs/1512.03385) +//! +use crate::models::with_tracing::{conv2d, Conv2d}; +use candle::{Result, Tensor, D}; +use candle_nn as nn; +use candle_nn::Module; + +/// Configuration for a ResNet block. +#[derive(Debug, Clone, Copy)] +pub struct ResnetBlock2DConfig { + /// The number of output channels, defaults to the number of input channels. + pub out_channels: Option, + pub temb_channels: Option, + /// The number of groups to use in group normalization. + pub groups: usize, + pub groups_out: Option, + /// The epsilon to be used in the group normalization operations. + pub eps: f64, + /// Whether to use a 2D convolution in the skip connection. When using None, + /// such a convolution is used if the number of input channels is different from + /// the number of output channels. + pub use_in_shortcut: Option, + // non_linearity: silu + /// The final output is scaled by dividing by this value. + pub output_scale_factor: f64, +} + +impl Default for ResnetBlock2DConfig { + fn default() -> Self { + Self { + out_channels: None, + temb_channels: Some(512), + groups: 32, + groups_out: None, + eps: 1e-6, + use_in_shortcut: None, + output_scale_factor: 1., + } + } +} + +#[derive(Debug)] +pub struct ResnetBlock2D { + norm1: nn::GroupNorm, + conv1: Conv2d, + norm2: nn::GroupNorm, + conv2: Conv2d, + time_emb_proj: Option, + conv_shortcut: Option, + span: tracing::Span, + config: ResnetBlock2DConfig, +} + +impl ResnetBlock2D { + pub fn new( + vs: nn::VarBuilder, + in_channels: usize, + config: ResnetBlock2DConfig, + ) -> Result { + let out_channels = config.out_channels.unwrap_or(in_channels); + let conv_cfg = nn::Conv2dConfig { + stride: 1, + padding: 1, + groups: 1, + dilation: 1, + cudnn_fwd_algo: None, + }; + let norm1 = nn::group_norm(config.groups, in_channels, config.eps, vs.pp("norm1"))?; + let conv1 = conv2d(in_channels, out_channels, 3, conv_cfg, vs.pp("conv1"))?; + let groups_out = config.groups_out.unwrap_or(config.groups); + let norm2 = nn::group_norm(groups_out, out_channels, config.eps, vs.pp("norm2"))?; + let conv2 = conv2d(out_channels, out_channels, 3, conv_cfg, vs.pp("conv2"))?; + let use_in_shortcut = config + .use_in_shortcut + .unwrap_or(in_channels != out_channels); + let conv_shortcut = if use_in_shortcut { + let conv_cfg = nn::Conv2dConfig { + stride: 1, + padding: 0, + groups: 1, + dilation: 1, + cudnn_fwd_algo: None, + }; + Some(conv2d( + in_channels, + out_channels, + 1, + conv_cfg, + vs.pp("conv_shortcut"), + )?) + } else { + None + }; + let time_emb_proj = match config.temb_channels { + None => None, + Some(temb_channels) => Some(nn::linear( + temb_channels, + out_channels, + vs.pp("time_emb_proj"), + )?), + }; + let span = tracing::span!(tracing::Level::TRACE, "resnet2d"); + Ok(Self { + norm1, + conv1, + norm2, + conv2, + time_emb_proj, + span, + config, + conv_shortcut, + }) + } + + pub fn forward(&self, xs: &Tensor, temb: Option<&Tensor>) -> Result { + let _enter = self.span.enter(); + let shortcut_xs = match &self.conv_shortcut { + Some(conv_shortcut) => conv_shortcut.forward(xs)?, + None => xs.clone(), + }; + let xs = self.norm1.forward(xs)?; + let xs = self.conv1.forward(&nn::ops::silu(&xs)?)?; + let xs = match (temb, &self.time_emb_proj) { + (Some(temb), Some(time_emb_proj)) => time_emb_proj + .forward(&nn::ops::silu(temb)?)? + .unsqueeze(D::Minus1)? + .unsqueeze(D::Minus1)? + .broadcast_add(&xs)?, + _ => xs, + }; + let xs = self + .conv2 + .forward(&nn::ops::silu(&self.norm2.forward(&xs)?)?)?; + (shortcut_xs + xs)? / self.config.output_scale_factor + } +} diff --git a/patches/candle-transformers/src/models/stable_diffusion/schedulers.rs b/patches/candle-transformers/src/models/stable_diffusion/schedulers.rs new file mode 100644 index 0000000000..fda592e31c --- /dev/null +++ b/patches/candle-transformers/src/models/stable_diffusion/schedulers.rs @@ -0,0 +1,72 @@ +#![allow(dead_code)] +//! # Diffusion pipelines and models +//! +//! Noise schedulers can be used to set the trade-off between +//! inference speed and quality. +use candle::{Result, Tensor}; + +pub trait SchedulerConfig: std::fmt::Debug + Send + Sync { + fn build(&self, inference_steps: usize) -> Result>; +} + +/// This trait represents a scheduler for the diffusion process. +pub trait Scheduler { + fn timesteps(&self) -> &[usize]; + + fn add_noise(&self, original: &Tensor, noise: Tensor, timestep: usize) -> Result; + + fn init_noise_sigma(&self) -> f64; + + fn scale_model_input(&self, sample: Tensor, _timestep: usize) -> Result; + + fn step(&mut self, model_output: &Tensor, timestep: usize, sample: &Tensor) -> Result; +} + +/// This represents how beta ranges from its minimum value to the maximum +/// during training. +#[derive(Debug, Clone, Copy)] +pub enum BetaSchedule { + /// Linear interpolation. + Linear, + /// Linear interpolation of the square root of beta. + ScaledLinear, + /// Glide cosine schedule + SquaredcosCapV2, +} + +#[derive(Debug, Clone, Copy)] +pub enum PredictionType { + Epsilon, + VPrediction, + Sample, +} + +/// Time step spacing for the diffusion process. +/// +/// "linspace", "leading", "trailing" corresponds to annotation of Table 2. of the [paper](https://arxiv.org/abs/2305.08891) +#[derive(Debug, Default, Clone, Copy)] +pub enum TimestepSpacing { + #[default] + Leading, + Linspace, + Trailing, +} + +/// Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of +/// `(1-beta)` over time from `t = [0,1]`. +/// +/// Contains a function `alpha_bar` that takes an argument `t` and transforms it to the cumulative product of `(1-beta)` +/// up to that part of the diffusion process. +pub(crate) fn betas_for_alpha_bar(num_diffusion_timesteps: usize, max_beta: f64) -> Result { + let alpha_bar = |time_step: usize| { + f64::cos((time_step as f64 + 0.008) / 1.008 * std::f64::consts::FRAC_PI_2).powi(2) + }; + let mut betas = Vec::with_capacity(num_diffusion_timesteps); + for i in 0..num_diffusion_timesteps { + let t1 = i / num_diffusion_timesteps; + let t2 = (i + 1) / num_diffusion_timesteps; + betas.push((1.0 - alpha_bar(t2) / alpha_bar(t1)).min(max_beta)); + } + let betas_len = betas.len(); + Tensor::from_vec(betas, betas_len, &candle::Device::Cpu) +} diff --git a/patches/candle-transformers/src/models/stable_diffusion/unet_2d.rs b/patches/candle-transformers/src/models/stable_diffusion/unet_2d.rs new file mode 100644 index 0000000000..cbef3316dd --- /dev/null +++ b/patches/candle-transformers/src/models/stable_diffusion/unet_2d.rs @@ -0,0 +1,401 @@ +//! 2D UNet Denoising Models +//! +//! The 2D Unet models take as input a noisy sample and the current diffusion +//! timestep and return a denoised version of the input. +use super::embeddings::{TimestepEmbedding, Timesteps}; +use super::unet_2d_blocks::*; +use crate::models::with_tracing::{conv2d, Conv2d}; +use candle::{Result, Tensor}; +use candle_nn as nn; +use candle_nn::Module; + +#[derive(Debug, Clone, Copy)] +pub struct BlockConfig { + pub out_channels: usize, + /// When `None` no cross-attn is used, when `Some(d)` then cross-attn is used and `d` is the + /// number of transformer blocks to be used. + pub use_cross_attn: Option, + pub attention_head_dim: usize, +} + +#[derive(Debug, Clone)] +pub struct UNet2DConditionModelConfig { + pub center_input_sample: bool, + pub flip_sin_to_cos: bool, + pub freq_shift: f64, + pub blocks: Vec, + pub layers_per_block: usize, + pub downsample_padding: usize, + pub mid_block_scale_factor: f64, + pub norm_num_groups: usize, + pub norm_eps: f64, + pub cross_attention_dim: usize, + pub sliced_attention_size: Option, + pub use_linear_projection: bool, +} + +impl Default for UNet2DConditionModelConfig { + fn default() -> Self { + Self { + center_input_sample: false, + flip_sin_to_cos: true, + freq_shift: 0., + blocks: vec![ + BlockConfig { + out_channels: 320, + use_cross_attn: Some(1), + attention_head_dim: 8, + }, + BlockConfig { + out_channels: 640, + use_cross_attn: Some(1), + attention_head_dim: 8, + }, + BlockConfig { + out_channels: 1280, + use_cross_attn: Some(1), + attention_head_dim: 8, + }, + BlockConfig { + out_channels: 1280, + use_cross_attn: None, + attention_head_dim: 8, + }, + ], + layers_per_block: 2, + downsample_padding: 1, + mid_block_scale_factor: 1., + norm_num_groups: 32, + norm_eps: 1e-5, + cross_attention_dim: 1280, + sliced_attention_size: None, + use_linear_projection: false, + } + } +} + +#[derive(Debug)] +pub(crate) enum UNetDownBlock { + Basic(DownBlock2D), + CrossAttn(CrossAttnDownBlock2D), +} + +#[derive(Debug)] +enum UNetUpBlock { + Basic(UpBlock2D), + CrossAttn(CrossAttnUpBlock2D), +} + +#[derive(Debug)] +pub struct UNet2DConditionModel { + conv_in: Conv2d, + time_proj: Timesteps, + time_embedding: TimestepEmbedding, + down_blocks: Vec, + mid_block: UNetMidBlock2DCrossAttn, + up_blocks: Vec, + conv_norm_out: nn::GroupNorm, + conv_out: Conv2d, + span: tracing::Span, + config: UNet2DConditionModelConfig, +} + +impl UNet2DConditionModel { + pub fn new( + vs: nn::VarBuilder, + in_channels: usize, + out_channels: usize, + use_flash_attn: bool, + config: UNet2DConditionModelConfig, + ) -> Result { + let n_blocks = config.blocks.len(); + let b_channels = config.blocks[0].out_channels; + let bl_channels = config.blocks.last().unwrap().out_channels; + let bl_attention_head_dim = config.blocks.last().unwrap().attention_head_dim; + let time_embed_dim = b_channels * 4; + let conv_cfg = nn::Conv2dConfig { + padding: 1, + ..Default::default() + }; + let conv_in = conv2d(in_channels, b_channels, 3, conv_cfg, vs.pp("conv_in"))?; + + let time_proj = Timesteps::new(b_channels, config.flip_sin_to_cos, config.freq_shift); + let time_embedding = + TimestepEmbedding::new(vs.pp("time_embedding"), b_channels, time_embed_dim)?; + + let vs_db = vs.pp("down_blocks"); + let down_blocks = (0..n_blocks) + .map(|i| { + let BlockConfig { + out_channels, + use_cross_attn, + attention_head_dim, + } = config.blocks[i]; + + // Enable automatic attention slicing if the config sliced_attention_size is set to 0. + let sliced_attention_size = match config.sliced_attention_size { + Some(0) => Some(attention_head_dim / 2), + _ => config.sliced_attention_size, + }; + + let in_channels = if i > 0 { + config.blocks[i - 1].out_channels + } else { + b_channels + }; + let db_cfg = DownBlock2DConfig { + num_layers: config.layers_per_block, + resnet_eps: config.norm_eps, + resnet_groups: config.norm_num_groups, + add_downsample: i < n_blocks - 1, + downsample_padding: config.downsample_padding, + ..Default::default() + }; + if let Some(transformer_layers_per_block) = use_cross_attn { + let config = CrossAttnDownBlock2DConfig { + downblock: db_cfg, + attn_num_head_channels: attention_head_dim, + cross_attention_dim: config.cross_attention_dim, + sliced_attention_size, + use_linear_projection: config.use_linear_projection, + transformer_layers_per_block, + }; + let block = CrossAttnDownBlock2D::new( + vs_db.pp(i.to_string()), + in_channels, + out_channels, + Some(time_embed_dim), + use_flash_attn, + config, + )?; + Ok(UNetDownBlock::CrossAttn(block)) + } else { + let block = DownBlock2D::new( + vs_db.pp(i.to_string()), + in_channels, + out_channels, + Some(time_embed_dim), + db_cfg, + )?; + Ok(UNetDownBlock::Basic(block)) + } + }) + .collect::>>()?; + + // https://github.com/huggingface/diffusers/blob/a76f2ad538e73b34d5fe7be08c8eb8ab38c7e90c/src/diffusers/models/unet_2d_condition.py#L462 + let mid_transformer_layers_per_block = match config.blocks.last() { + None => 1, + Some(block) => block.use_cross_attn.unwrap_or(1), + }; + let mid_cfg = UNetMidBlock2DCrossAttnConfig { + resnet_eps: config.norm_eps, + output_scale_factor: config.mid_block_scale_factor, + cross_attn_dim: config.cross_attention_dim, + attn_num_head_channels: bl_attention_head_dim, + resnet_groups: Some(config.norm_num_groups), + use_linear_projection: config.use_linear_projection, + transformer_layers_per_block: mid_transformer_layers_per_block, + ..Default::default() + }; + + let mid_block = UNetMidBlock2DCrossAttn::new( + vs.pp("mid_block"), + bl_channels, + Some(time_embed_dim), + use_flash_attn, + mid_cfg, + )?; + + let vs_ub = vs.pp("up_blocks"); + let up_blocks = (0..n_blocks) + .map(|i| { + let BlockConfig { + out_channels, + use_cross_attn, + attention_head_dim, + } = config.blocks[n_blocks - 1 - i]; + + // Enable automatic attention slicing if the config sliced_attention_size is set to 0. + let sliced_attention_size = match config.sliced_attention_size { + Some(0) => Some(attention_head_dim / 2), + _ => config.sliced_attention_size, + }; + + let prev_out_channels = if i > 0 { + config.blocks[n_blocks - i].out_channels + } else { + bl_channels + }; + let in_channels = { + let index = if i == n_blocks - 1 { + 0 + } else { + n_blocks - i - 2 + }; + config.blocks[index].out_channels + }; + let ub_cfg = UpBlock2DConfig { + num_layers: config.layers_per_block + 1, + resnet_eps: config.norm_eps, + resnet_groups: config.norm_num_groups, + add_upsample: i < n_blocks - 1, + ..Default::default() + }; + if let Some(transformer_layers_per_block) = use_cross_attn { + let config = CrossAttnUpBlock2DConfig { + upblock: ub_cfg, + attn_num_head_channels: attention_head_dim, + cross_attention_dim: config.cross_attention_dim, + sliced_attention_size, + use_linear_projection: config.use_linear_projection, + transformer_layers_per_block, + }; + let block = CrossAttnUpBlock2D::new( + vs_ub.pp(i.to_string()), + in_channels, + prev_out_channels, + out_channels, + Some(time_embed_dim), + use_flash_attn, + config, + )?; + Ok(UNetUpBlock::CrossAttn(block)) + } else { + let block = UpBlock2D::new( + vs_ub.pp(i.to_string()), + in_channels, + prev_out_channels, + out_channels, + Some(time_embed_dim), + ub_cfg, + )?; + Ok(UNetUpBlock::Basic(block)) + } + }) + .collect::>>()?; + + let conv_norm_out = nn::group_norm( + config.norm_num_groups, + b_channels, + config.norm_eps, + vs.pp("conv_norm_out"), + )?; + let conv_out = conv2d(b_channels, out_channels, 3, conv_cfg, vs.pp("conv_out"))?; + let span = tracing::span!(tracing::Level::TRACE, "unet2d"); + Ok(Self { + conv_in, + time_proj, + time_embedding, + down_blocks, + mid_block, + up_blocks, + conv_norm_out, + conv_out, + span, + config, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + timestep: f64, + encoder_hidden_states: &Tensor, + ) -> Result { + let _enter = self.span.enter(); + self.forward_with_additional_residuals(xs, timestep, encoder_hidden_states, None, None) + } + + pub fn forward_with_additional_residuals( + &self, + xs: &Tensor, + timestep: f64, + encoder_hidden_states: &Tensor, + down_block_additional_residuals: Option<&[Tensor]>, + mid_block_additional_residual: Option<&Tensor>, + ) -> Result { + let (bsize, _channels, height, width) = xs.dims4()?; + let device = xs.device(); + let n_blocks = self.config.blocks.len(); + let num_upsamplers = n_blocks - 1; + let default_overall_up_factor = 2usize.pow(num_upsamplers as u32); + let forward_upsample_size = + height % default_overall_up_factor != 0 || width % default_overall_up_factor != 0; + // 0. center input if necessary + let xs = if self.config.center_input_sample { + ((xs * 2.0)? - 1.0)? + } else { + xs.clone() + }; + // 1. time + let emb = (Tensor::ones(bsize, xs.dtype(), device)? * timestep)?; + let emb = self.time_proj.forward(&emb)?; + let emb = self.time_embedding.forward(&emb)?; + // 2. pre-process + let xs = self.conv_in.forward(&xs)?; + // 3. down + let mut down_block_res_xs = vec![xs.clone()]; + let mut xs = xs; + for down_block in self.down_blocks.iter() { + let (_xs, res_xs) = match down_block { + UNetDownBlock::Basic(b) => b.forward(&xs, Some(&emb))?, + UNetDownBlock::CrossAttn(b) => { + b.forward(&xs, Some(&emb), Some(encoder_hidden_states))? + } + }; + down_block_res_xs.extend(res_xs); + xs = _xs; + } + + let new_down_block_res_xs = + if let Some(down_block_additional_residuals) = down_block_additional_residuals { + let mut v = vec![]; + // A previous version of this code had a bug because of the addition being made + // in place via += hence modifying the input of the mid block. + for (i, residuals) in down_block_additional_residuals.iter().enumerate() { + v.push((&down_block_res_xs[i] + residuals)?) + } + v + } else { + down_block_res_xs + }; + let mut down_block_res_xs = new_down_block_res_xs; + + // 4. mid + let xs = self + .mid_block + .forward(&xs, Some(&emb), Some(encoder_hidden_states))?; + let xs = match mid_block_additional_residual { + None => xs, + Some(m) => (m + xs)?, + }; + // 5. up + let mut xs = xs; + let mut upsample_size = None; + for (i, up_block) in self.up_blocks.iter().enumerate() { + let n_resnets = match up_block { + UNetUpBlock::Basic(b) => b.resnets.len(), + UNetUpBlock::CrossAttn(b) => b.upblock.resnets.len(), + }; + let res_xs = down_block_res_xs.split_off(down_block_res_xs.len() - n_resnets); + if i < n_blocks - 1 && forward_upsample_size { + let (_, _, h, w) = down_block_res_xs.last().unwrap().dims4()?; + upsample_size = Some((h, w)) + } + xs = match up_block { + UNetUpBlock::Basic(b) => b.forward(&xs, &res_xs, Some(&emb), upsample_size)?, + UNetUpBlock::CrossAttn(b) => b.forward( + &xs, + &res_xs, + Some(&emb), + upsample_size, + Some(encoder_hidden_states), + )?, + }; + } + // 6. post-process + let xs = self.conv_norm_out.forward(&xs)?; + let xs = nn::ops::silu(&xs)?; + self.conv_out.forward(&xs) + } +} diff --git a/patches/candle-transformers/src/models/stable_diffusion/unet_2d_blocks.rs b/patches/candle-transformers/src/models/stable_diffusion/unet_2d_blocks.rs new file mode 100644 index 0000000000..028c51b744 --- /dev/null +++ b/patches/candle-transformers/src/models/stable_diffusion/unet_2d_blocks.rs @@ -0,0 +1,868 @@ +//! 2D UNet Building Blocks +//! +use super::attention::{ + AttentionBlock, AttentionBlockConfig, SpatialTransformer, SpatialTransformerConfig, +}; +use super::resnet::{ResnetBlock2D, ResnetBlock2DConfig}; +use crate::models::with_tracing::{conv2d, Conv2d}; +use candle::{Module, Result, Tensor, D}; +use candle_nn as nn; + +#[derive(Debug)] +struct Downsample2D { + conv: Option, + padding: usize, + span: tracing::Span, +} + +impl Downsample2D { + fn new( + vs: nn::VarBuilder, + in_channels: usize, + use_conv: bool, + out_channels: usize, + padding: usize, + ) -> Result { + let conv = if use_conv { + let config = nn::Conv2dConfig { + stride: 2, + padding, + ..Default::default() + }; + let conv = conv2d(in_channels, out_channels, 3, config, vs.pp("conv"))?; + Some(conv) + } else { + None + }; + let span = tracing::span!(tracing::Level::TRACE, "downsample2d"); + Ok(Self { + conv, + padding, + span, + }) + } +} + +impl Module for Downsample2D { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + match &self.conv { + None => xs.avg_pool2d(2), + Some(conv) => { + if self.padding == 0 { + let xs = xs + .pad_with_zeros(D::Minus1, 0, 1)? + .pad_with_zeros(D::Minus2, 0, 1)?; + conv.forward(&xs) + } else { + conv.forward(xs) + } + } + } + } +} + +// This does not support the conv-transpose mode. +#[derive(Debug)] +struct Upsample2D { + conv: Conv2d, + span: tracing::Span, +} + +impl Upsample2D { + fn new(vs: nn::VarBuilder, in_channels: usize, out_channels: usize) -> Result { + let config = nn::Conv2dConfig { + padding: 1, + ..Default::default() + }; + let conv = conv2d(in_channels, out_channels, 3, config, vs.pp("conv"))?; + let span = tracing::span!(tracing::Level::TRACE, "upsample2d"); + Ok(Self { conv, span }) + } +} + +impl Upsample2D { + fn forward(&self, xs: &Tensor, size: Option<(usize, usize)>) -> Result { + let _enter = self.span.enter(); + let xs = match size { + None => { + let (_bsize, _channels, h, w) = xs.dims4()?; + xs.upsample_nearest2d(2 * h, 2 * w)? + } + Some((h, w)) => xs.upsample_nearest2d(h, w)?, + }; + self.conv.forward(&xs) + } +} + +#[derive(Debug, Clone, Copy)] +pub struct DownEncoderBlock2DConfig { + pub num_layers: usize, + pub resnet_eps: f64, + pub resnet_groups: usize, + pub output_scale_factor: f64, + pub add_downsample: bool, + pub downsample_padding: usize, +} + +impl Default for DownEncoderBlock2DConfig { + fn default() -> Self { + Self { + num_layers: 1, + resnet_eps: 1e-6, + resnet_groups: 32, + output_scale_factor: 1., + add_downsample: true, + downsample_padding: 1, + } + } +} + +#[derive(Debug)] +pub struct DownEncoderBlock2D { + resnets: Vec, + downsampler: Option, + span: tracing::Span, + pub config: DownEncoderBlock2DConfig, +} + +impl DownEncoderBlock2D { + pub fn new( + vs: nn::VarBuilder, + in_channels: usize, + out_channels: usize, + config: DownEncoderBlock2DConfig, + ) -> Result { + let resnets: Vec<_> = { + let vs = vs.pp("resnets"); + let conv_cfg = ResnetBlock2DConfig { + eps: config.resnet_eps, + out_channels: Some(out_channels), + groups: config.resnet_groups, + output_scale_factor: config.output_scale_factor, + temb_channels: None, + ..Default::default() + }; + (0..(config.num_layers)) + .map(|i| { + let in_channels = if i == 0 { in_channels } else { out_channels }; + ResnetBlock2D::new(vs.pp(i.to_string()), in_channels, conv_cfg) + }) + .collect::>>()? + }; + let downsampler = if config.add_downsample { + let downsample = Downsample2D::new( + vs.pp("downsamplers").pp("0"), + out_channels, + true, + out_channels, + config.downsample_padding, + )?; + Some(downsample) + } else { + None + }; + let span = tracing::span!(tracing::Level::TRACE, "down-enc2d"); + Ok(Self { + resnets, + downsampler, + span, + config, + }) + } +} + +impl Module for DownEncoderBlock2D { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let mut xs = xs.clone(); + for resnet in self.resnets.iter() { + xs = resnet.forward(&xs, None)? + } + match &self.downsampler { + Some(downsampler) => downsampler.forward(&xs), + None => Ok(xs), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct UpDecoderBlock2DConfig { + pub num_layers: usize, + pub resnet_eps: f64, + pub resnet_groups: usize, + pub output_scale_factor: f64, + pub add_upsample: bool, +} + +impl Default for UpDecoderBlock2DConfig { + fn default() -> Self { + Self { + num_layers: 1, + resnet_eps: 1e-6, + resnet_groups: 32, + output_scale_factor: 1., + add_upsample: true, + } + } +} + +#[derive(Debug)] +pub struct UpDecoderBlock2D { + resnets: Vec, + upsampler: Option, + span: tracing::Span, + pub config: UpDecoderBlock2DConfig, +} + +impl UpDecoderBlock2D { + pub fn new( + vs: nn::VarBuilder, + in_channels: usize, + out_channels: usize, + config: UpDecoderBlock2DConfig, + ) -> Result { + let resnets: Vec<_> = { + let vs = vs.pp("resnets"); + let conv_cfg = ResnetBlock2DConfig { + out_channels: Some(out_channels), + eps: config.resnet_eps, + groups: config.resnet_groups, + output_scale_factor: config.output_scale_factor, + temb_channels: None, + ..Default::default() + }; + (0..(config.num_layers)) + .map(|i| { + let in_channels = if i == 0 { in_channels } else { out_channels }; + ResnetBlock2D::new(vs.pp(i.to_string()), in_channels, conv_cfg) + }) + .collect::>>()? + }; + let upsampler = if config.add_upsample { + let upsample = + Upsample2D::new(vs.pp("upsamplers").pp("0"), out_channels, out_channels)?; + Some(upsample) + } else { + None + }; + let span = tracing::span!(tracing::Level::TRACE, "up-dec2d"); + Ok(Self { + resnets, + upsampler, + span, + config, + }) + } +} + +impl Module for UpDecoderBlock2D { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let mut xs = xs.clone(); + for resnet in self.resnets.iter() { + xs = resnet.forward(&xs, None)? + } + match &self.upsampler { + Some(upsampler) => upsampler.forward(&xs, None), + None => Ok(xs), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct UNetMidBlock2DConfig { + pub num_layers: usize, + pub resnet_eps: f64, + pub resnet_groups: Option, + pub attn_num_head_channels: Option, + // attention_type "default" + pub output_scale_factor: f64, +} + +impl Default for UNetMidBlock2DConfig { + fn default() -> Self { + Self { + num_layers: 1, + resnet_eps: 1e-6, + resnet_groups: Some(32), + attn_num_head_channels: Some(1), + output_scale_factor: 1., + } + } +} + +#[derive(Debug)] +pub struct UNetMidBlock2D { + resnet: ResnetBlock2D, + attn_resnets: Vec<(AttentionBlock, ResnetBlock2D)>, + span: tracing::Span, + pub config: UNetMidBlock2DConfig, +} + +impl UNetMidBlock2D { + pub fn new( + vs: nn::VarBuilder, + in_channels: usize, + temb_channels: Option, + config: UNetMidBlock2DConfig, + ) -> Result { + let vs_resnets = vs.pp("resnets"); + let vs_attns = vs.pp("attentions"); + let resnet_groups = config + .resnet_groups + .unwrap_or_else(|| usize::min(in_channels / 4, 32)); + let resnet_cfg = ResnetBlock2DConfig { + eps: config.resnet_eps, + groups: resnet_groups, + output_scale_factor: config.output_scale_factor, + temb_channels, + ..Default::default() + }; + let resnet = ResnetBlock2D::new(vs_resnets.pp("0"), in_channels, resnet_cfg)?; + let attn_cfg = AttentionBlockConfig { + num_head_channels: config.attn_num_head_channels, + num_groups: resnet_groups, + rescale_output_factor: config.output_scale_factor, + eps: config.resnet_eps, + }; + let mut attn_resnets = vec![]; + for index in 0..config.num_layers { + let attn = AttentionBlock::new(vs_attns.pp(index.to_string()), in_channels, attn_cfg)?; + let resnet = ResnetBlock2D::new( + vs_resnets.pp((index + 1).to_string()), + in_channels, + resnet_cfg, + )?; + attn_resnets.push((attn, resnet)) + } + let span = tracing::span!(tracing::Level::TRACE, "mid2d"); + Ok(Self { + resnet, + attn_resnets, + span, + config, + }) + } + + pub fn forward(&self, xs: &Tensor, temb: Option<&Tensor>) -> Result { + let _enter = self.span.enter(); + let mut xs = self.resnet.forward(xs, temb)?; + for (attn, resnet) in self.attn_resnets.iter() { + xs = resnet.forward(&attn.forward(&xs)?, temb)? + } + Ok(xs) + } +} + +#[derive(Debug, Clone, Copy)] +pub struct UNetMidBlock2DCrossAttnConfig { + pub num_layers: usize, + pub resnet_eps: f64, + pub resnet_groups: Option, + pub attn_num_head_channels: usize, + // attention_type "default" + pub output_scale_factor: f64, + pub cross_attn_dim: usize, + pub sliced_attention_size: Option, + pub use_linear_projection: bool, + pub transformer_layers_per_block: usize, +} + +impl Default for UNetMidBlock2DCrossAttnConfig { + fn default() -> Self { + Self { + num_layers: 1, + resnet_eps: 1e-6, + resnet_groups: Some(32), + attn_num_head_channels: 1, + output_scale_factor: 1., + cross_attn_dim: 1280, + sliced_attention_size: None, // Sliced attention disabled + use_linear_projection: false, + transformer_layers_per_block: 1, + } + } +} + +#[derive(Debug)] +pub struct UNetMidBlock2DCrossAttn { + resnet: ResnetBlock2D, + attn_resnets: Vec<(SpatialTransformer, ResnetBlock2D)>, + span: tracing::Span, + pub config: UNetMidBlock2DCrossAttnConfig, +} + +impl UNetMidBlock2DCrossAttn { + pub fn new( + vs: nn::VarBuilder, + in_channels: usize, + temb_channels: Option, + use_flash_attn: bool, + config: UNetMidBlock2DCrossAttnConfig, + ) -> Result { + let vs_resnets = vs.pp("resnets"); + let vs_attns = vs.pp("attentions"); + let resnet_groups = config + .resnet_groups + .unwrap_or_else(|| usize::min(in_channels / 4, 32)); + let resnet_cfg = ResnetBlock2DConfig { + eps: config.resnet_eps, + groups: resnet_groups, + output_scale_factor: config.output_scale_factor, + temb_channels, + ..Default::default() + }; + let resnet = ResnetBlock2D::new(vs_resnets.pp("0"), in_channels, resnet_cfg)?; + let n_heads = config.attn_num_head_channels; + let attn_cfg = SpatialTransformerConfig { + depth: config.transformer_layers_per_block, + num_groups: resnet_groups, + context_dim: Some(config.cross_attn_dim), + sliced_attention_size: config.sliced_attention_size, + use_linear_projection: config.use_linear_projection, + }; + let mut attn_resnets = vec![]; + for index in 0..config.num_layers { + let attn = SpatialTransformer::new( + vs_attns.pp(index.to_string()), + in_channels, + n_heads, + in_channels / n_heads, + use_flash_attn, + attn_cfg, + )?; + let resnet = ResnetBlock2D::new( + vs_resnets.pp((index + 1).to_string()), + in_channels, + resnet_cfg, + )?; + attn_resnets.push((attn, resnet)) + } + let span = tracing::span!(tracing::Level::TRACE, "xa-mid2d"); + Ok(Self { + resnet, + attn_resnets, + span, + config, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + temb: Option<&Tensor>, + encoder_hidden_states: Option<&Tensor>, + ) -> Result { + let _enter = self.span.enter(); + let mut xs = self.resnet.forward(xs, temb)?; + for (attn, resnet) in self.attn_resnets.iter() { + xs = resnet.forward(&attn.forward(&xs, encoder_hidden_states)?, temb)? + } + Ok(xs) + } +} + +#[derive(Debug, Clone, Copy)] +pub struct DownBlock2DConfig { + pub num_layers: usize, + pub resnet_eps: f64, + // resnet_time_scale_shift: "default" + // resnet_act_fn: "swish" + pub resnet_groups: usize, + pub output_scale_factor: f64, + pub add_downsample: bool, + pub downsample_padding: usize, +} + +impl Default for DownBlock2DConfig { + fn default() -> Self { + Self { + num_layers: 1, + resnet_eps: 1e-6, + resnet_groups: 32, + output_scale_factor: 1., + add_downsample: true, + downsample_padding: 1, + } + } +} + +#[derive(Debug)] +pub struct DownBlock2D { + resnets: Vec, + downsampler: Option, + span: tracing::Span, + pub config: DownBlock2DConfig, +} + +impl DownBlock2D { + pub fn new( + vs: nn::VarBuilder, + in_channels: usize, + out_channels: usize, + temb_channels: Option, + config: DownBlock2DConfig, + ) -> Result { + let vs_resnets = vs.pp("resnets"); + let resnet_cfg = ResnetBlock2DConfig { + out_channels: Some(out_channels), + eps: config.resnet_eps, + output_scale_factor: config.output_scale_factor, + temb_channels, + ..Default::default() + }; + let resnets = (0..config.num_layers) + .map(|i| { + let in_channels = if i == 0 { in_channels } else { out_channels }; + ResnetBlock2D::new(vs_resnets.pp(i.to_string()), in_channels, resnet_cfg) + }) + .collect::>>()?; + let downsampler = if config.add_downsample { + let downsampler = Downsample2D::new( + vs.pp("downsamplers").pp("0"), + out_channels, + true, + out_channels, + config.downsample_padding, + )?; + Some(downsampler) + } else { + None + }; + let span = tracing::span!(tracing::Level::TRACE, "down2d"); + Ok(Self { + resnets, + downsampler, + span, + config, + }) + } + + pub fn forward(&self, xs: &Tensor, temb: Option<&Tensor>) -> Result<(Tensor, Vec)> { + let _enter = self.span.enter(); + let mut xs = xs.clone(); + let mut output_states = vec![]; + for resnet in self.resnets.iter() { + xs = resnet.forward(&xs, temb)?; + output_states.push(xs.clone()); + } + let xs = match &self.downsampler { + Some(downsampler) => { + let xs = downsampler.forward(&xs)?; + output_states.push(xs.clone()); + xs + } + None => xs, + }; + Ok((xs, output_states)) + } +} + +#[derive(Debug, Clone, Copy)] +pub struct CrossAttnDownBlock2DConfig { + pub downblock: DownBlock2DConfig, + pub attn_num_head_channels: usize, + pub cross_attention_dim: usize, + // attention_type: "default" + pub sliced_attention_size: Option, + pub use_linear_projection: bool, + pub transformer_layers_per_block: usize, +} + +impl Default for CrossAttnDownBlock2DConfig { + fn default() -> Self { + Self { + downblock: Default::default(), + attn_num_head_channels: 1, + cross_attention_dim: 1280, + sliced_attention_size: None, + use_linear_projection: false, + transformer_layers_per_block: 1, + } + } +} + +#[derive(Debug)] +pub struct CrossAttnDownBlock2D { + downblock: DownBlock2D, + attentions: Vec, + span: tracing::Span, + pub config: CrossAttnDownBlock2DConfig, +} + +impl CrossAttnDownBlock2D { + pub fn new( + vs: nn::VarBuilder, + in_channels: usize, + out_channels: usize, + temb_channels: Option, + use_flash_attn: bool, + config: CrossAttnDownBlock2DConfig, + ) -> Result { + let downblock = DownBlock2D::new( + vs.clone(), + in_channels, + out_channels, + temb_channels, + config.downblock, + )?; + let n_heads = config.attn_num_head_channels; + let cfg = SpatialTransformerConfig { + depth: config.transformer_layers_per_block, + context_dim: Some(config.cross_attention_dim), + num_groups: config.downblock.resnet_groups, + sliced_attention_size: config.sliced_attention_size, + use_linear_projection: config.use_linear_projection, + }; + let vs_attn = vs.pp("attentions"); + let attentions = (0..config.downblock.num_layers) + .map(|i| { + SpatialTransformer::new( + vs_attn.pp(i.to_string()), + out_channels, + n_heads, + out_channels / n_heads, + use_flash_attn, + cfg, + ) + }) + .collect::>>()?; + let span = tracing::span!(tracing::Level::TRACE, "xa-down2d"); + Ok(Self { + downblock, + attentions, + span, + config, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + temb: Option<&Tensor>, + encoder_hidden_states: Option<&Tensor>, + ) -> Result<(Tensor, Vec)> { + let _enter = self.span.enter(); + let mut output_states = vec![]; + let mut xs = xs.clone(); + for (resnet, attn) in self.downblock.resnets.iter().zip(self.attentions.iter()) { + xs = resnet.forward(&xs, temb)?; + xs = attn.forward(&xs, encoder_hidden_states)?; + output_states.push(xs.clone()); + } + let xs = match &self.downblock.downsampler { + Some(downsampler) => { + let xs = downsampler.forward(&xs)?; + output_states.push(xs.clone()); + xs + } + None => xs, + }; + Ok((xs, output_states)) + } +} + +#[derive(Debug, Clone, Copy)] +pub struct UpBlock2DConfig { + pub num_layers: usize, + pub resnet_eps: f64, + // resnet_time_scale_shift: "default" + // resnet_act_fn: "swish" + pub resnet_groups: usize, + pub output_scale_factor: f64, + pub add_upsample: bool, +} + +impl Default for UpBlock2DConfig { + fn default() -> Self { + Self { + num_layers: 1, + resnet_eps: 1e-6, + resnet_groups: 32, + output_scale_factor: 1., + add_upsample: true, + } + } +} + +#[derive(Debug)] +pub struct UpBlock2D { + pub resnets: Vec, + upsampler: Option, + span: tracing::Span, + pub config: UpBlock2DConfig, +} + +impl UpBlock2D { + pub fn new( + vs: nn::VarBuilder, + in_channels: usize, + prev_output_channels: usize, + out_channels: usize, + temb_channels: Option, + config: UpBlock2DConfig, + ) -> Result { + let vs_resnets = vs.pp("resnets"); + let resnet_cfg = ResnetBlock2DConfig { + out_channels: Some(out_channels), + temb_channels, + eps: config.resnet_eps, + output_scale_factor: config.output_scale_factor, + ..Default::default() + }; + let resnets = (0..config.num_layers) + .map(|i| { + let res_skip_channels = if i == config.num_layers - 1 { + in_channels + } else { + out_channels + }; + let resnet_in_channels = if i == 0 { + prev_output_channels + } else { + out_channels + }; + let in_channels = resnet_in_channels + res_skip_channels; + ResnetBlock2D::new(vs_resnets.pp(i.to_string()), in_channels, resnet_cfg) + }) + .collect::>>()?; + let upsampler = if config.add_upsample { + let upsampler = + Upsample2D::new(vs.pp("upsamplers").pp("0"), out_channels, out_channels)?; + Some(upsampler) + } else { + None + }; + let span = tracing::span!(tracing::Level::TRACE, "up2d"); + Ok(Self { + resnets, + upsampler, + span, + config, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + res_xs: &[Tensor], + temb: Option<&Tensor>, + upsample_size: Option<(usize, usize)>, + ) -> Result { + let _enter = self.span.enter(); + let mut xs = xs.clone(); + for (index, resnet) in self.resnets.iter().enumerate() { + xs = Tensor::cat(&[&xs, &res_xs[res_xs.len() - index - 1]], 1)?; + xs = xs.contiguous()?; + xs = resnet.forward(&xs, temb)?; + } + match &self.upsampler { + Some(upsampler) => upsampler.forward(&xs, upsample_size), + None => Ok(xs), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct CrossAttnUpBlock2DConfig { + pub upblock: UpBlock2DConfig, + pub attn_num_head_channels: usize, + pub cross_attention_dim: usize, + // attention_type: "default" + pub sliced_attention_size: Option, + pub use_linear_projection: bool, + pub transformer_layers_per_block: usize, +} + +impl Default for CrossAttnUpBlock2DConfig { + fn default() -> Self { + Self { + upblock: Default::default(), + attn_num_head_channels: 1, + cross_attention_dim: 1280, + sliced_attention_size: None, + use_linear_projection: false, + transformer_layers_per_block: 1, + } + } +} + +#[derive(Debug)] +pub struct CrossAttnUpBlock2D { + pub upblock: UpBlock2D, + pub attentions: Vec, + span: tracing::Span, + pub config: CrossAttnUpBlock2DConfig, +} + +impl CrossAttnUpBlock2D { + pub fn new( + vs: nn::VarBuilder, + in_channels: usize, + prev_output_channels: usize, + out_channels: usize, + temb_channels: Option, + use_flash_attn: bool, + config: CrossAttnUpBlock2DConfig, + ) -> Result { + let upblock = UpBlock2D::new( + vs.clone(), + in_channels, + prev_output_channels, + out_channels, + temb_channels, + config.upblock, + )?; + let n_heads = config.attn_num_head_channels; + let cfg = SpatialTransformerConfig { + depth: config.transformer_layers_per_block, + context_dim: Some(config.cross_attention_dim), + num_groups: config.upblock.resnet_groups, + sliced_attention_size: config.sliced_attention_size, + use_linear_projection: config.use_linear_projection, + }; + let vs_attn = vs.pp("attentions"); + let attentions = (0..config.upblock.num_layers) + .map(|i| { + SpatialTransformer::new( + vs_attn.pp(i.to_string()), + out_channels, + n_heads, + out_channels / n_heads, + use_flash_attn, + cfg, + ) + }) + .collect::>>()?; + let span = tracing::span!(tracing::Level::TRACE, "xa-up2d"); + Ok(Self { + upblock, + attentions, + span, + config, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + res_xs: &[Tensor], + temb: Option<&Tensor>, + upsample_size: Option<(usize, usize)>, + encoder_hidden_states: Option<&Tensor>, + ) -> Result { + let _enter = self.span.enter(); + let mut xs = xs.clone(); + for (index, resnet) in self.upblock.resnets.iter().enumerate() { + xs = Tensor::cat(&[&xs, &res_xs[res_xs.len() - index - 1]], 1)?; + xs = xs.contiguous()?; + xs = resnet.forward(&xs, temb)?; + xs = self.attentions[index].forward(&xs, encoder_hidden_states)?; + } + match &self.upblock.upsampler { + Some(upsampler) => upsampler.forward(&xs, upsample_size), + None => Ok(xs), + } + } +} diff --git a/patches/candle-transformers/src/models/stable_diffusion/uni_pc.rs b/patches/candle-transformers/src/models/stable_diffusion/uni_pc.rs new file mode 100644 index 0000000000..4ac0af3886 --- /dev/null +++ b/patches/candle-transformers/src/models/stable_diffusion/uni_pc.rs @@ -0,0 +1,1005 @@ +//! # UniPC Scheduler +//! +//! UniPC is a training-free framework designed for the fast sampling of diffusion models, which consists of a +//! corrector (UniC) and a predictor (UniP) that share a unified analytical form and support arbitrary orders. +//! +//! UniPC is by design model-agnostic, supporting pixel-space/latent-space DPMs on unconditional/conditional +//! sampling. It can also be applied to both noise prediction and data prediction models. Compared with prior +//! methods, UniPC converges faster thanks to the increased order of accuracy. Both quantitative and qualitative +//! results show UniPC can improve sampling quality, especially at very low step counts (5~10). +//! +//! For more information, see the original publication: +//! UniPC: A Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models, W. Zhao et al, 2023. +//! https://arxiv.org/abs/2302.04867 +//! +//! This work is based largely on UniPC implementation from the diffusers python package: +//! https://raw.githubusercontent.com/huggingface/diffusers/e8aacda762e311505ba05ae340af23b149e37af3/src/diffusers/schedulers/scheduling_unipc_multistep.py +use std::collections::HashSet; +use std::ops::Neg; + +use super::schedulers::PredictionType; +use super::{ + schedulers::{Scheduler, SchedulerConfig}, + utils::{interp, linspace}, +}; +use candle::{Error, IndexOp, Result, Tensor}; + +#[derive(Debug, Clone, Copy)] +pub enum SigmaSchedule { + Karras(KarrasSigmaSchedule), + Exponential(ExponentialSigmaSchedule), +} + +impl SigmaSchedule { + fn sigma_t(&self, t: f64) -> f64 { + match self { + Self::Karras(x) => x.sigma_t(t), + Self::Exponential(x) => x.sigma_t(t), + } + } +} + +impl Default for SigmaSchedule { + fn default() -> Self { + Self::Karras(KarrasSigmaSchedule::default()) + } +} + +#[derive(Debug, Clone, Copy)] +pub struct KarrasSigmaSchedule { + pub sigma_min: f64, + pub sigma_max: f64, + pub rho: f64, +} + +impl KarrasSigmaSchedule { + fn sigma_t(&self, t: f64) -> f64 { + let (min_inv_rho, max_inv_rho) = ( + self.sigma_min.powf(1.0 / self.rho), + self.sigma_max.powf(1.0 / self.rho), + ); + + (max_inv_rho + ((1.0 - t) * (min_inv_rho - max_inv_rho))).powf(self.rho) + } +} + +impl Default for KarrasSigmaSchedule { + fn default() -> Self { + Self { + sigma_max: 10.0, + sigma_min: 0.1, + rho: 4.0, + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct ExponentialSigmaSchedule { + sigma_min: f64, + sigma_max: f64, +} + +impl ExponentialSigmaSchedule { + fn sigma_t(&self, t: f64) -> f64 { + (t * (self.sigma_max.ln() - self.sigma_min.ln()) + self.sigma_min.ln()).exp() + } +} + +impl Default for ExponentialSigmaSchedule { + fn default() -> Self { + Self { + sigma_max: 80.0, + sigma_min: 0.1, + } + } +} + +#[derive(Debug, Default, Clone, Copy)] +pub enum SolverType { + #[default] + Bh1, + Bh2, +} + +#[derive(Debug, Default, Clone, Copy)] +pub enum AlgorithmType { + #[default] + DpmSolverPlusPlus, + SdeDpmSolverPlusPlus, +} + +#[derive(Debug, Default, Clone, Copy)] +pub enum FinalSigmasType { + #[default] + Zero, + SigmaMin, +} + +#[derive(Debug, Clone)] +pub enum TimestepSchedule { + /// Timesteps will be determined by interpolation of sigmas + FromSigmas, + /// Timesteps will be separated by regular intervals + Linspace, +} + +impl TimestepSchedule { + fn timesteps( + &self, + sigma_schedule: &SigmaSchedule, + num_inference_steps: usize, + num_training_steps: usize, + ) -> Result> { + match self { + Self::FromSigmas => { + let sigmas: Tensor = linspace(1., 0., num_inference_steps)? + .to_vec1()? + .into_iter() + .map(|t| sigma_schedule.sigma_t(t)) + .collect::>() + .try_into()?; + let log_sigmas = sigmas.log()?.to_vec1::()?; + let timesteps = interp( + &log_sigmas.iter().copied().rev().collect::>(), + &linspace( + log_sigmas[log_sigmas.len() - 1] - 0.001, + log_sigmas[0] + 0.001, + num_inference_steps, + )? + .to_vec1::()?, + &linspace(0., num_training_steps as f64, num_inference_steps)? + .to_vec1::()?, + ) + .into_iter() + .map(|f| (num_training_steps - 1) - (f as usize)) + .collect::>(); + + Ok(timesteps) + } + + Self::Linspace => { + Ok( + linspace((num_training_steps - 1) as f64, 0., num_inference_steps)? + .to_vec1::()? + .into_iter() + .map(|f| f as usize) + .collect(), + ) + } + } + } +} + +#[derive(Debug, Clone)] +pub enum CorrectorConfiguration { + Disabled, + Enabled { skip_steps: HashSet }, +} + +impl Default for CorrectorConfiguration { + fn default() -> Self { + Self::Enabled { + skip_steps: [0, 1, 2].into_iter().collect(), + } + } +} + +impl CorrectorConfiguration { + pub fn new(disabled_steps: impl IntoIterator) -> Self { + Self::Enabled { + skip_steps: disabled_steps.into_iter().collect(), + } + } +} + +#[derive(Debug, Clone)] +pub struct UniPCSchedulerConfig { + /// Configure the UNIC corrector. By default it is disabled + pub corrector: CorrectorConfiguration, + /// Determines how sigma relates to a given timestep + pub sigma_schedule: SigmaSchedule, + /// Determines the points + pub timestep_schedule: TimestepSchedule, + /// The solver order which can be `1` or higher. It is recommended to use `solver_order=2` for guided + /// sampling, and `solver_order=3` for unconditional sampling. + pub solver_order: usize, + /// Prediction type of the scheduler function + pub prediction_type: PredictionType, + pub num_training_timesteps: usize, + /// Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such + /// as Stable Diffusion. + pub thresholding: bool, + /// The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. + pub dynamic_thresholding_ratio: f64, + /// The threshold value for dynamic thresholding. + pub sample_max_value: f64, + pub solver_type: SolverType, + /// Whether to use lower-order solvers in the final steps. + pub lower_order_final: bool, +} + +impl Default for UniPCSchedulerConfig { + fn default() -> Self { + Self { + corrector: Default::default(), + timestep_schedule: TimestepSchedule::FromSigmas, + sigma_schedule: SigmaSchedule::Karras(Default::default()), + prediction_type: PredictionType::Epsilon, + num_training_timesteps: 1000, + solver_order: 2, + thresholding: false, + dynamic_thresholding_ratio: 0.995, + sample_max_value: 1.0, + solver_type: SolverType::Bh1, + lower_order_final: true, + } + } +} + +impl SchedulerConfig for UniPCSchedulerConfig { + fn build(&self, inference_steps: usize) -> Result> { + Ok(Box::new(EdmDpmMultistepScheduler::new( + self.clone(), + inference_steps, + )?)) + } +} + +struct State { + model_outputs: Vec>, + lower_order_nums: usize, + order: usize, + last_sample: Option, +} + +impl State { + fn new(solver_order: usize) -> Self { + Self { + model_outputs: vec![None; solver_order], + lower_order_nums: 0, + order: 0, + last_sample: None, + } + } + + fn lower_order_nums(&self) -> usize { + self.lower_order_nums + } + + fn update_lower_order_nums(&mut self, n: usize) { + self.lower_order_nums = n; + } + + fn model_outputs(&self) -> &[Option] { + self.model_outputs.as_slice() + } + + fn update_model_output(&mut self, idx: usize, output: Option) { + self.model_outputs[idx] = output; + } + + fn last_sample(&self) -> Option<&Tensor> { + self.last_sample.as_ref() + } + + fn update_last_sample(&mut self, sample: Tensor) { + let _ = self.last_sample.replace(sample); + } + + fn order(&self) -> usize { + self.order + } + + fn update_order(&mut self, order: usize) { + self.order = order; + } +} + +pub struct EdmDpmMultistepScheduler { + schedule: Schedule, + config: UniPCSchedulerConfig, + state: State, +} + +impl EdmDpmMultistepScheduler { + pub fn new(config: UniPCSchedulerConfig, num_inference_steps: usize) -> Result { + let schedule = Schedule::new( + config.timestep_schedule.clone(), + config.sigma_schedule, + num_inference_steps, + config.num_training_timesteps, + )?; + + Ok(Self { + schedule, + state: State::new(config.solver_order), + config, + }) + } + + fn step_index(&self, timestep: usize) -> usize { + let index_candidates = self + .schedule + .timesteps() + .iter() + .enumerate() + .filter(|(_, t)| *t == ×tep) + .map(|(i, _)| i) + .collect::>(); + + match index_candidates.len() { + 0 => 0, + 1 => index_candidates[0], + _ => index_candidates[1], + } + } + + fn timestep(&self, step_idx: usize) -> usize { + self.schedule + .timesteps() + .get(step_idx) + .copied() + .unwrap_or(0) + } + + fn convert_model_output( + &self, + model_output: &Tensor, + sample: &Tensor, + timestep: usize, + ) -> Result { + let (alpha_t, sigma_t) = ( + self.schedule.alpha_t(timestep), + self.schedule.sigma_t(timestep), + ); + + let x0_pred = match self.config.prediction_type { + PredictionType::Epsilon => ((sample - (model_output * sigma_t))? / alpha_t)?, + PredictionType::Sample => model_output.clone(), + PredictionType::VPrediction => ((alpha_t * sample)? - (sigma_t * model_output)?)?, + }; + + if self.config.thresholding { + self.threshold_sample(x0_pred) + } else { + Ok(x0_pred) + } + } + + fn threshold_sample(&self, sample: Tensor) -> Result { + let shape = sample.shape().clone().into_dims(); + let v = sample + .abs()? + .reshape((shape[0], shape[1] * shape[2..].iter().product::()))? + .to_dtype(candle::DType::F64)? + .to_vec2::()?; + let q = stats::Quantile::new(self.config.dynamic_thresholding_ratio) + .with_samples(v.into_iter().flatten()); + let (threshold, max) = (q.quantile().max(self.config.sample_max_value), q.max()); + + sample.clamp(-threshold, threshold)? / (threshold / max).sqrt().min(1.) + } + + fn multistep_uni_p_bh_update(&self, sample: &Tensor, timestep: usize) -> Result { + let step_index = self.step_index(timestep); + let ns = &self.schedule; + let model_outputs = self.state.model_outputs(); + let Some(m0) = &model_outputs[model_outputs.len() - 1] else { + return Err(Error::Msg( + "Expected model output for predictor update".to_string(), + )); + }; + + let (t0, tt) = (timestep, self.timestep(self.step_index(timestep) + 1)); + let (sigma_t, sigma_s0) = (ns.sigma_t(tt), ns.sigma_t(t0)); + let (alpha_t, _alpha_s0) = (ns.alpha_t(tt), ns.alpha_t(t0)); + let (lambda_t, lambda_s0) = (ns.lambda_t(tt), ns.lambda_t(t0)); + + let h = lambda_t - lambda_s0; + let device = sample.device(); + + let (mut rks, mut d1s) = (vec![], vec![]); + for i in 1..self.state.order() { + let ti = self.timestep(step_index.saturating_sub(i + 1)); + let Some(mi) = model_outputs + .get(model_outputs.len().saturating_sub(i + 1)) + .into_iter() + .flatten() + .next() + else { + return Err(Error::Msg( + "Expected model output for predictor update".to_string(), + )); + }; + let (alpha_si, sigma_si) = (ns.alpha_t(ti), ns.sigma_t(ti)); + let lambda_si = alpha_si.ln() - sigma_si.ln(); + let rk = (lambda_si - lambda_s0) / h; + rks.push(rk); + d1s.push(((mi - m0)? / rk)?); + } + rks.push(1.0); + let rks = Tensor::new(rks, device)?; + let (mut r, mut b) = (vec![], vec![]); + + let hh = h.neg(); + let h_phi_1 = hh.exp_m1(); + let mut h_phi_k = h_phi_1 / hh - 1.; + let mut factorial_i = 1.; + + let b_h = match self.config.solver_type { + SolverType::Bh1 => hh, + SolverType::Bh2 => hh.exp_m1(), + }; + + for i in 1..self.state.order() + 1 { + r.push(rks.powf(i as f64 - 1.)?); + b.push(h_phi_k * factorial_i / b_h); + factorial_i = i as f64 + 1.; + h_phi_k = h_phi_k / hh - 1. / factorial_i; + } + + let (r, b) = (Tensor::stack(&r, 0)?, Tensor::new(b, device)?); + let (d1s, rhos_p) = match d1s.len() { + 0 => (None, None), + _ => { + let rhos_p = match self.state.order() { + 2 => Tensor::new(&[0.5f64], m0.device())?.to_dtype(m0.dtype())?, + _ => { + let ((r1, r2), b1) = (r.dims2()?, b.dims1()?); + let inverse = linalg::inverse(&r.i((..(r1 - 1), ..(r2 - 1)))?)?; + let b = b.i(..(b1 - 1))?; + b.broadcast_mul(&inverse)?.sum(1)?.to_dtype(m0.dtype())? + } + }; + + (Some(Tensor::stack(&d1s, 1)?), Some(rhos_p)) + } + }; + + let x_t_ = ((sigma_t / sigma_s0 * sample)? - (alpha_t * h_phi_1 * m0)?)?; + if let (Some(d1s), Some(rhos_p)) = (d1s, rhos_p) { + use linalg::{Permutation, TensordotFixedPosition, TensordotGeneral}; + let output_shape = m0.shape().clone(); + let pred_res = TensordotGeneral { + lhs_permutation: Permutation { dims: vec![0] }, + rhs_permutation: Permutation { + dims: vec![1, 0, 2, 3, 4], + }, + tensordot_fixed_position: TensordotFixedPosition { + len_uncontracted_lhs: 1, + len_uncontracted_rhs: output_shape.dims().iter().product::(), + len_contracted_axes: d1s.dim(1)?, + output_shape, + }, + output_permutation: Permutation { + dims: vec![0, 1, 2, 3], + }, + } + .eval(&rhos_p, &d1s)?; + x_t_ - (alpha_t * b_h * pred_res)? + } else { + Ok(x_t_) + } + } + + fn multistep_uni_c_bh_update( + &self, + model_output: &Tensor, + model_outputs: &[Option], + last_sample: &Tensor, + sample: &Tensor, + timestep: usize, + ) -> Result { + let step_index = self.step_index(timestep); + let Some(m0) = model_outputs.last().into_iter().flatten().next() else { + return Err(Error::Msg( + "Expected model output for corrector update".to_string(), + )); + }; + let model_t = model_output; + let (x, _xt) = (last_sample, sample); + + let (t0, tt, ns) = ( + self.timestep(self.step_index(timestep) - 1), + timestep, + &self.schedule, + ); + let (sigma_t, sigma_s0) = (ns.sigma_t(tt), ns.sigma_t(t0)); + let (alpha_t, _alpha_s0) = (ns.alpha_t(tt), ns.alpha_t(t0)); + let (lambda_t, lambda_s0) = (ns.lambda_t(tt), ns.lambda_t(t0)); + + let h = lambda_t - lambda_s0; + let device = sample.device(); + + let (mut rks, mut d1s) = (vec![], vec![]); + for i in 1..self.state.order() { + let ti = self.timestep(step_index.saturating_sub(i + 1)); + let Some(mi) = model_outputs + .get(model_outputs.len().saturating_sub(i + 1)) + .into_iter() + .flatten() + .next() + else { + return Err(Error::Msg( + "Expected model output for corrector update".to_string(), + )); + }; + let (alpha_si, sigma_si) = (ns.alpha_t(ti), ns.sigma_t(ti)); + let lambda_si = alpha_si.ln() - sigma_si.ln(); + let rk = (lambda_si - lambda_s0) / h; + rks.push(rk); + d1s.push(((mi - m0)? / rk)?); + } + rks.push(1.0); + let rks = Tensor::new(rks, device)?; + let (mut r, mut b) = (vec![], vec![]); + + let hh = h.neg(); + let h_phi_1 = hh.exp_m1(); + let mut h_phi_k = h_phi_1 / hh - 1.; + let mut factorial_i = 1.; + + let b_h = match self.config.solver_type { + SolverType::Bh1 => hh, + SolverType::Bh2 => hh.exp_m1(), + }; + + for i in 1..self.state.order() + 1 { + r.push(rks.powf(i as f64 - 1.)?); + b.push(h_phi_k * factorial_i / b_h); + factorial_i = i as f64 + 1.; + h_phi_k = h_phi_k / hh - 1. / factorial_i; + } + + let (r, b) = (Tensor::stack(&r, 0)?, Tensor::new(b, device)?); + let d1s = match d1s.len() { + 0 => None, + _ => Some(Tensor::stack(&d1s, 1)?), + }; + let rhos_c = match self.state.order() { + 1 => Tensor::new(&[0.5f64], m0.device())?.to_dtype(m0.dtype())?, + _ => { + let inverse = linalg::inverse(&r)?; + b.broadcast_mul(&inverse)?.sum(1)?.to_dtype(m0.dtype())? + } + }; + + let x_t_ = ((sigma_t / sigma_s0 * x)? - (alpha_t * h_phi_1 * m0)?)?; + let corr_res = d1s + .map(|d1s| { + use linalg::{Permutation, TensordotFixedPosition, TensordotGeneral}; + let output_shape = x_t_.shape().clone(); + TensordotGeneral { + lhs_permutation: Permutation { dims: vec![0] }, + rhs_permutation: Permutation { + dims: vec![1, 0, 2, 3, 4], + }, + tensordot_fixed_position: TensordotFixedPosition { + len_uncontracted_lhs: 1, + len_uncontracted_rhs: output_shape.dims().iter().product::(), + len_contracted_axes: d1s.dim(1)?, + output_shape, + }, + output_permutation: Permutation { + dims: vec![0, 1, 2, 3], + }, + } + .eval(&rhos_c.i(..rhos_c.dims()[0] - 1)?, &d1s) + }) + .unwrap_or_else(|| Tensor::zeros_like(m0))?; + + let d1_t = (model_t - m0)?; + let x_t = (x_t_ + - (alpha_t + * b_h + * (corr_res + rhos_c.i(rhos_c.dims()[0] - 1)?.broadcast_mul(&d1_t)?)?)?)?; + + Ok(x_t) + } +} + +impl Scheduler for EdmDpmMultistepScheduler { + fn step(&mut self, model_output: &Tensor, timestep: usize, sample: &Tensor) -> Result { + let step_index = self.step_index(timestep); + let model_output_converted = &self.convert_model_output(model_output, sample, timestep)?; + let sample = match (&self.config.corrector, self.state.last_sample()) { + (CorrectorConfiguration::Enabled { skip_steps: s }, Some(last_sample)) + if !s.contains(&step_index) && step_index > 0 => + { + &self.multistep_uni_c_bh_update( + model_output_converted, + self.state.model_outputs(), + last_sample, + sample, + timestep, + )? + } + (CorrectorConfiguration::Enabled { .. }, _) | (CorrectorConfiguration::Disabled, _) => { + sample + } + }; + + let mut model_outputs = self.state.model_outputs().to_vec(); + for i in 0..self.config.solver_order.saturating_sub(1) { + self.state + .update_model_output(i, model_outputs[i + 1].take()); + } + self.state.update_model_output( + model_outputs.len() - 1, + Some(model_output_converted.clone()), + ); + + let mut this_order = self.config.solver_order; + if self.config.lower_order_final { + this_order = self + .config + .solver_order + .min(self.schedule.timesteps.len() - step_index); + } + self.state + .update_order(this_order.min(self.state.lower_order_nums() + 1)); + + self.state.update_last_sample(sample.clone()); + let prev_sample = self.multistep_uni_p_bh_update(sample, timestep)?; + + let lower_order_nums = self.state.lower_order_nums(); + if lower_order_nums < self.config.solver_order { + self.state.update_lower_order_nums(lower_order_nums + 1); + } + + Ok(prev_sample) + } + + fn scale_model_input(&self, sample: Tensor, _timestep: usize) -> Result { + Ok(sample) + } + + fn timesteps(&self) -> &[usize] { + &self.schedule.timesteps + } + + fn add_noise(&self, original: &Tensor, noise: Tensor, timestep: usize) -> Result { + let (alpha_t, sigma_t) = ( + self.schedule.alpha_t(timestep), + self.schedule.sigma_t(timestep), + ); + + (alpha_t * original)? + (sigma_t * noise)? + } + + fn init_noise_sigma(&self) -> f64 { + self.schedule.sigma_t(self.schedule.num_training_steps()) + } +} + +#[derive(Debug, Clone)] +struct Schedule { + timesteps: Vec, + num_training_steps: usize, + sigma_schedule: SigmaSchedule, + #[allow(unused)] + timestep_schedule: TimestepSchedule, +} + +impl Schedule { + fn new( + timestep_schedule: TimestepSchedule, + sigma_schedule: SigmaSchedule, + num_inference_steps: usize, + num_training_steps: usize, + ) -> Result { + Ok(Self { + timesteps: timestep_schedule.timesteps( + &sigma_schedule, + num_inference_steps, + num_training_steps, + )?, + timestep_schedule, + sigma_schedule, + num_training_steps, + }) + } + + fn timesteps(&self) -> &[usize] { + &self.timesteps + } + + fn num_training_steps(&self) -> usize { + self.num_training_steps + } + + fn t(&self, step: usize) -> f64 { + (step as f64 + 1.) / self.num_training_steps as f64 + } + + fn alpha_t(&self, t: usize) -> f64 { + (1. / (self.sigma_schedule.sigma_t(self.t(t)).powi(2) + 1.)).sqrt() + } + + fn sigma_t(&self, t: usize) -> f64 { + self.sigma_schedule.sigma_t(self.t(t)) * self.alpha_t(t) + } + + fn lambda_t(&self, t: usize) -> f64 { + self.alpha_t(t).ln() - self.sigma_t(t).ln() + } +} + +mod stats { + //! This is a slightly modified form of the P² quantile implementation from https://github.com/vks/average. + //! Also see: http://www.cs.wustl.edu/~jain/papers/ftp/psqr.pdf + use num_traits::{Float, ToPrimitive}; + + #[derive(Debug, Clone)] + pub struct Quantile { + q: [f64; 5], + n: [i64; 5], + m: [f64; 5], + dm: [f64; 5], + max: Option, + } + + impl Quantile { + pub fn new(p: f64) -> Quantile { + assert!((0. ..=1.).contains(&p)); + Quantile { + q: [0.; 5], + n: [1, 2, 3, 4, 0], + m: [1., 1. + 2. * p, 1. + 4. * p, 3. + 2. * p, 5.], + dm: [0., p / 2., p, (1. + p) / 2., 1.], + max: None, + } + } + + pub fn max(&self) -> f64 { + self.max.unwrap_or(f64::NAN) + } + + fn p(&self) -> f64 { + self.dm[2] + } + + fn parabolic(&self, i: usize, d: f64) -> f64 { + let s = d.round() as i64; + self.q[i] + + d / (self.n[i + 1] - self.n[i - 1]).to_f64().unwrap() + * ((self.n[i] - self.n[i - 1] + s).to_f64().unwrap() + * (self.q[i + 1] - self.q[i]) + / (self.n[i + 1] - self.n[i]).to_f64().unwrap() + + (self.n[i + 1] - self.n[i] - s).to_f64().unwrap() + * (self.q[i] - self.q[i - 1]) + / (self.n[i] - self.n[i - 1]).to_f64().unwrap()) + } + + fn linear(&self, i: usize, d: f64) -> f64 { + let sum = if d < 0. { i - 1 } else { i + 1 }; + self.q[i] + d * (self.q[sum] - self.q[i]) / (self.n[sum] - self.n[i]).to_f64().unwrap() + } + + pub fn quantile(&self) -> f64 { + if self.len() >= 5 { + return self.q[2]; + } + + if self.is_empty() { + return f64::NAN; + } + let mut heights: [f64; 4] = [self.q[0], self.q[1], self.q[2], self.q[3]]; + let len = self.len() as usize; + debug_assert!(len < 5); + sort_floats(&mut heights[..len]); + let desired_index = (len as f64) * self.p() - 1.; + let mut index = desired_index.ceil(); + if desired_index == index && index >= 0. { + let index = index.round() as usize; + debug_assert!(index < 5); + if index < len - 1 { + return 0.5 * self.q[index] + 0.5 * self.q[index + 1]; + } + } + index = index.max(0.); + let mut index = index.round() as usize; + debug_assert!(index < 5); + index = index.min(len - 1); + self.q[index] + } + + fn len(&self) -> u64 { + self.n[4] as u64 + } + + fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub fn add(&mut self, x: f64) { + self.max = self.max.map(|y| y.max(x)).or(Some(x)); + + if self.n[4] < 5 { + self.q[self.n[4] as usize] = x; + self.n[4] += 1; + if self.n[4] == 5 { + sort_floats(&mut self.q); + } + return; + } + + let mut k: usize; + if x < self.q[0] { + self.q[0] = x; + k = 0; + } else { + k = 4; + for i in 1..5 { + if x < self.q[i] { + k = i; + break; + } + } + if self.q[4] < x { + self.q[4] = x; + } + }; + + for i in k..5 { + self.n[i] += 1; + } + for i in 0..5 { + self.m[i] += self.dm[i]; + } + + for i in 1..4 { + let d = self.m[i] - self.n[i].to_f64().unwrap(); + if d >= 1. && self.n[i + 1] - self.n[i] > 1 + || d <= -1. && self.n[i - 1] - self.n[i] < -1 + { + let d = Float::signum(d); + let q_new = self.parabolic(i, d); + if self.q[i - 1] < q_new && q_new < self.q[i + 1] { + self.q[i] = q_new; + } else { + self.q[i] = self.linear(i, d); + } + let delta = d.round() as i64; + debug_assert_eq!(delta.abs(), 1); + self.n[i] += delta; + } + } + } + + pub fn with_samples(mut self, samples: impl IntoIterator) -> Self { + for sample in samples { + self.add(sample); + } + + self + } + } + + fn sort_floats(v: &mut [f64]) { + v.sort_unstable_by(|a, b| a.total_cmp(b)); + } +} + +mod linalg { + use candle::{IndexOp, Result, Shape, Tensor}; + + pub fn inverse(m: &Tensor) -> Result { + adjoint(m)? / determinant(m)?.to_scalar::()? + } + + pub fn adjoint(m: &Tensor) -> Result { + cofactor(m)?.transpose(0, 1) + } + + pub fn cofactor(m: &Tensor) -> Result { + let s = m.shape().dim(0)?; + if s == 2 { + let mut v = vec![]; + for i in 0..2 { + let mut x = vec![]; + for j in 0..2 { + x.push((m.i((i, j))? * (-1.0f64).powi(i as i32 + j as i32))?) + } + v.push(Tensor::stack(&x, 0)?.unsqueeze(0)?); + } + return Tensor::stack(&v, 1)?.squeeze(0); + } + + let minors = minors(m)?; + let mut v = vec![]; + for i in 0..s { + let mut x = vec![]; + for j in 0..s { + let det = (determinant(&minors.i((i, j))?)? + * ((-1.0f64).powi(i as i32) * (-1.0f64).powi(j as i32)))?; + x.push(det); + } + v.push(Tensor::stack(&x, 0)?.unsqueeze(0)?); + } + + Tensor::stack(&v, 1)?.squeeze(0) + } + + pub fn determinant(m: &Tensor) -> Result { + let s = m.shape().dim(0)?; + if s == 2 { + return (m.i((0, 0))? * m.i((1, 1))?)? - (m.i((0, 1))? * m.i((1, 0))?); + } + + let cofactor = cofactor(m)?; + let m0 = m.i((0, 0))?; + let det = (0..s) + .map(|i| m.i((0, i))? * cofactor.i((0, i))?) + .try_fold(m0.zeros_like()?, |acc, cur| acc + cur?)?; + + Ok(det) + } + + pub fn minors(m: &Tensor) -> Result { + let s = m.shape().dim(0)?; + if s == 1 { + return m.i((0, 0)); + } + + let mut v = vec![]; + for i in 0..s { + let msub = Tensor::cat(&[m.i((..i, ..))?, m.i(((i + 1).., ..))?], 0)?; + let mut x = vec![]; + for j in 0..s { + let t = Tensor::cat(&[msub.i((.., ..j))?, msub.i((.., (j + 1)..))?], 1)?; + x.push(t); + } + v.push(Tensor::stack(&x, 0)?.unsqueeze(0)?); + } + + Tensor::stack(&v, 1)?.squeeze(0) + } + + #[derive(Debug)] + pub struct TensordotGeneral { + pub lhs_permutation: Permutation, + pub rhs_permutation: Permutation, + pub tensordot_fixed_position: TensordotFixedPosition, + pub output_permutation: Permutation, + } + + impl TensordotGeneral { + pub fn eval(&self, lhs: &Tensor, rhs: &Tensor) -> Result { + let permuted_lhs = self.lhs_permutation.eval(lhs)?; + let permuted_rhs = self.rhs_permutation.eval(rhs)?; + let tensordotted = self + .tensordot_fixed_position + .eval(&permuted_lhs, &permuted_rhs)?; + self.output_permutation.eval(&tensordotted) + } + } + + #[derive(Debug)] + pub struct TensordotFixedPosition { + pub len_uncontracted_lhs: usize, + pub len_uncontracted_rhs: usize, + pub len_contracted_axes: usize, + pub output_shape: Shape, + } + + impl TensordotFixedPosition { + fn eval(&self, lhs: &Tensor, rhs: &Tensor) -> Result { + let lhs_view = lhs.reshape((self.len_uncontracted_lhs, self.len_contracted_axes))?; + let rhs_view = rhs.reshape((self.len_contracted_axes, self.len_uncontracted_rhs))?; + + lhs_view.matmul(&rhs_view)?.reshape(&self.output_shape) + } + } + + #[derive(Debug)] + pub struct Permutation { + pub dims: Vec, + } + + impl Permutation { + fn eval(&self, tensor: &Tensor) -> Result { + tensor.permute(self.dims.as_slice()) + } + } +} diff --git a/patches/candle-transformers/src/models/stable_diffusion/utils.rs b/patches/candle-transformers/src/models/stable_diffusion/utils.rs new file mode 100644 index 0000000000..0118bafc54 --- /dev/null +++ b/patches/candle-transformers/src/models/stable_diffusion/utils.rs @@ -0,0 +1,61 @@ +use candle::{Device, Result, Tensor}; + +pub fn linspace(start: f64, stop: f64, steps: usize) -> Result { + if steps == 0 { + Tensor::from_vec(Vec::::new(), steps, &Device::Cpu) + } else if steps == 1 { + Tensor::from_vec(vec![start], steps, &Device::Cpu) + } else { + let delta = (stop - start) / (steps - 1) as f64; + let vs = (0..steps) + .map(|step| start + step as f64 * delta) + .collect::>(); + Tensor::from_vec(vs, steps, &Device::Cpu) + } +} + +/// A linear interpolator for a sorted array of x and y values. +struct LinearInterpolator<'x, 'y> { + xp: &'x [f64], + fp: &'y [f64], + cache: usize, +} + +impl LinearInterpolator<'_, '_> { + fn accel_find(&mut self, x: f64) -> usize { + let xidx = self.cache; + if x < self.xp[xidx] { + self.cache = self.xp[0..xidx].partition_point(|o| *o < x); + self.cache = self.cache.saturating_sub(1); + } else if x >= self.xp[xidx + 1] { + self.cache = self.xp[xidx..self.xp.len()].partition_point(|o| *o < x) + xidx; + self.cache = self.cache.saturating_sub(1); + } + + self.cache + } + + fn eval(&mut self, x: f64) -> f64 { + if x < self.xp[0] || x > self.xp[self.xp.len() - 1] { + return f64::NAN; + } + + let idx = self.accel_find(x); + + let x_l = self.xp[idx]; + let x_h = self.xp[idx + 1]; + let y_l = self.fp[idx]; + let y_h = self.fp[idx + 1]; + let dx = x_h - x_l; + if dx > 0.0 { + y_l + (x - x_l) / dx * (y_h - y_l) + } else { + f64::NAN + } + } +} + +pub fn interp(x: &[f64], xp: &[f64], fp: &[f64]) -> Vec { + let mut interpolator = LinearInterpolator { xp, fp, cache: 0 }; + x.iter().map(|&x| interpolator.eval(x)).collect() +} diff --git a/patches/candle-transformers/src/models/stable_diffusion/vae.rs b/patches/candle-transformers/src/models/stable_diffusion/vae.rs new file mode 100644 index 0000000000..b3aba80277 --- /dev/null +++ b/patches/candle-transformers/src/models/stable_diffusion/vae.rs @@ -0,0 +1,403 @@ +#![allow(dead_code)] +//! # Variational Auto-Encoder (VAE) Models. +//! +//! Auto-encoder models compress their input to a usually smaller latent space +//! before expanding it back to its original shape. This results in the latent values +//! compressing the original information. +use super::unet_2d_blocks::{ + DownEncoderBlock2D, DownEncoderBlock2DConfig, UNetMidBlock2D, UNetMidBlock2DConfig, + UpDecoderBlock2D, UpDecoderBlock2DConfig, +}; +use candle::{Result, Tensor}; +use candle_nn as nn; +use candle_nn::Module; + +#[derive(Debug, Clone)] +struct EncoderConfig { + // down_block_types: DownEncoderBlock2D + block_out_channels: Vec, + layers_per_block: usize, + norm_num_groups: usize, + double_z: bool, +} + +impl Default for EncoderConfig { + fn default() -> Self { + Self { + block_out_channels: vec![64], + layers_per_block: 2, + norm_num_groups: 32, + double_z: true, + } + } +} + +#[derive(Debug)] +struct Encoder { + conv_in: nn::Conv2d, + down_blocks: Vec, + mid_block: UNetMidBlock2D, + conv_norm_out: nn::GroupNorm, + conv_out: nn::Conv2d, + #[allow(dead_code)] + config: EncoderConfig, +} + +impl Encoder { + fn new( + vs: nn::VarBuilder, + in_channels: usize, + out_channels: usize, + config: EncoderConfig, + ) -> Result { + let conv_cfg = nn::Conv2dConfig { + padding: 1, + ..Default::default() + }; + let conv_in = nn::conv2d( + in_channels, + config.block_out_channels[0], + 3, + conv_cfg, + vs.pp("conv_in"), + )?; + let mut down_blocks = vec![]; + let vs_down_blocks = vs.pp("down_blocks"); + for index in 0..config.block_out_channels.len() { + let out_channels = config.block_out_channels[index]; + let in_channels = if index > 0 { + config.block_out_channels[index - 1] + } else { + config.block_out_channels[0] + }; + let is_final = index + 1 == config.block_out_channels.len(); + let cfg = DownEncoderBlock2DConfig { + num_layers: config.layers_per_block, + resnet_eps: 1e-6, + resnet_groups: config.norm_num_groups, + add_downsample: !is_final, + downsample_padding: 0, + ..Default::default() + }; + let down_block = DownEncoderBlock2D::new( + vs_down_blocks.pp(index.to_string()), + in_channels, + out_channels, + cfg, + )?; + down_blocks.push(down_block) + } + let last_block_out_channels = *config.block_out_channels.last().unwrap(); + let mid_cfg = UNetMidBlock2DConfig { + resnet_eps: 1e-6, + output_scale_factor: 1., + attn_num_head_channels: None, + resnet_groups: Some(config.norm_num_groups), + ..Default::default() + }; + let mid_block = + UNetMidBlock2D::new(vs.pp("mid_block"), last_block_out_channels, None, mid_cfg)?; + let conv_norm_out = nn::group_norm( + config.norm_num_groups, + last_block_out_channels, + 1e-6, + vs.pp("conv_norm_out"), + )?; + let conv_out_channels = if config.double_z { + 2 * out_channels + } else { + out_channels + }; + let conv_cfg = nn::Conv2dConfig { + padding: 1, + ..Default::default() + }; + let conv_out = nn::conv2d( + last_block_out_channels, + conv_out_channels, + 3, + conv_cfg, + vs.pp("conv_out"), + )?; + Ok(Self { + conv_in, + down_blocks, + mid_block, + conv_norm_out, + conv_out, + config, + }) + } +} + +impl Encoder { + fn forward(&self, xs: &Tensor) -> Result { + let mut xs = xs.apply(&self.conv_in)?; + for down_block in self.down_blocks.iter() { + xs = xs.apply(down_block)? + } + let xs = self + .mid_block + .forward(&xs, None)? + .apply(&self.conv_norm_out)?; + nn::ops::silu(&xs)?.apply(&self.conv_out) + } +} + +#[derive(Debug, Clone)] +struct DecoderConfig { + // up_block_types: UpDecoderBlock2D + block_out_channels: Vec, + layers_per_block: usize, + norm_num_groups: usize, +} + +impl Default for DecoderConfig { + fn default() -> Self { + Self { + block_out_channels: vec![64], + layers_per_block: 2, + norm_num_groups: 32, + } + } +} + +#[derive(Debug)] +struct Decoder { + conv_in: nn::Conv2d, + up_blocks: Vec, + mid_block: UNetMidBlock2D, + conv_norm_out: nn::GroupNorm, + conv_out: nn::Conv2d, + #[allow(dead_code)] + config: DecoderConfig, +} + +impl Decoder { + fn new( + vs: nn::VarBuilder, + in_channels: usize, + out_channels: usize, + config: DecoderConfig, + ) -> Result { + let n_block_out_channels = config.block_out_channels.len(); + let last_block_out_channels = *config.block_out_channels.last().unwrap(); + let conv_cfg = nn::Conv2dConfig { + padding: 1, + ..Default::default() + }; + let conv_in = nn::conv2d( + in_channels, + last_block_out_channels, + 3, + conv_cfg, + vs.pp("conv_in"), + )?; + let mid_cfg = UNetMidBlock2DConfig { + resnet_eps: 1e-6, + output_scale_factor: 1., + attn_num_head_channels: None, + resnet_groups: Some(config.norm_num_groups), + ..Default::default() + }; + let mid_block = + UNetMidBlock2D::new(vs.pp("mid_block"), last_block_out_channels, None, mid_cfg)?; + let mut up_blocks = vec![]; + let vs_up_blocks = vs.pp("up_blocks"); + let reversed_block_out_channels: Vec<_> = + config.block_out_channels.iter().copied().rev().collect(); + for index in 0..n_block_out_channels { + let out_channels = reversed_block_out_channels[index]; + let in_channels = if index > 0 { + reversed_block_out_channels[index - 1] + } else { + reversed_block_out_channels[0] + }; + let is_final = index + 1 == n_block_out_channels; + let cfg = UpDecoderBlock2DConfig { + num_layers: config.layers_per_block + 1, + resnet_eps: 1e-6, + resnet_groups: config.norm_num_groups, + add_upsample: !is_final, + ..Default::default() + }; + let up_block = UpDecoderBlock2D::new( + vs_up_blocks.pp(index.to_string()), + in_channels, + out_channels, + cfg, + )?; + up_blocks.push(up_block) + } + let conv_norm_out = nn::group_norm( + config.norm_num_groups, + config.block_out_channels[0], + 1e-6, + vs.pp("conv_norm_out"), + )?; + let conv_cfg = nn::Conv2dConfig { + padding: 1, + ..Default::default() + }; + let conv_out = nn::conv2d( + config.block_out_channels[0], + out_channels, + 3, + conv_cfg, + vs.pp("conv_out"), + )?; + Ok(Self { + conv_in, + up_blocks, + mid_block, + conv_norm_out, + conv_out, + config, + }) + } +} + +impl Decoder { + fn forward(&self, xs: &Tensor) -> Result { + let mut xs = self.mid_block.forward(&self.conv_in.forward(xs)?, None)?; + for up_block in self.up_blocks.iter() { + xs = up_block.forward(&xs)? + } + let xs = self.conv_norm_out.forward(&xs)?; + let xs = nn::ops::silu(&xs)?; + self.conv_out.forward(&xs) + } +} + +#[derive(Debug, Clone)] +pub struct AutoEncoderKLConfig { + pub block_out_channels: Vec, + pub layers_per_block: usize, + pub latent_channels: usize, + pub norm_num_groups: usize, + pub use_quant_conv: bool, + pub use_post_quant_conv: bool, +} + +impl Default for AutoEncoderKLConfig { + fn default() -> Self { + Self { + block_out_channels: vec![64], + layers_per_block: 1, + latent_channels: 4, + norm_num_groups: 32, + use_quant_conv: true, + use_post_quant_conv: true, + } + } +} + +pub struct DiagonalGaussianDistribution { + mean: Tensor, + std: Tensor, +} + +impl DiagonalGaussianDistribution { + pub fn new(parameters: &Tensor) -> Result { + let mut parameters = parameters.chunk(2, 1)?.into_iter(); + let mean = parameters.next().unwrap(); + let logvar = parameters.next().unwrap(); + let std = (logvar * 0.5)?.exp()?; + Ok(DiagonalGaussianDistribution { mean, std }) + } + + pub fn sample(&self) -> Result { + let sample = self.mean.randn_like(0., 1.); + &self.mean + &self.std * sample + } +} + +// https://github.com/huggingface/diffusers/blob/970e30606c2944e3286f56e8eb6d3dc6d1eb85f7/src/diffusers/models/vae.py#L485 +// This implementation is specific to the config used in stable-diffusion-v1-5 +// https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/vae/config.json +#[derive(Debug)] +pub struct AutoEncoderKL { + encoder: Encoder, + decoder: Decoder, + quant_conv: Option, + post_quant_conv: Option, + pub config: AutoEncoderKLConfig, +} + +impl AutoEncoderKL { + pub fn new( + vs: nn::VarBuilder, + in_channels: usize, + out_channels: usize, + config: AutoEncoderKLConfig, + ) -> Result { + let latent_channels = config.latent_channels; + let encoder_cfg = EncoderConfig { + block_out_channels: config.block_out_channels.clone(), + layers_per_block: config.layers_per_block, + norm_num_groups: config.norm_num_groups, + double_z: true, + }; + let encoder = Encoder::new(vs.pp("encoder"), in_channels, latent_channels, encoder_cfg)?; + let decoder_cfg = DecoderConfig { + block_out_channels: config.block_out_channels.clone(), + layers_per_block: config.layers_per_block, + norm_num_groups: config.norm_num_groups, + }; + let decoder = Decoder::new(vs.pp("decoder"), latent_channels, out_channels, decoder_cfg)?; + let conv_cfg = Default::default(); + + let quant_conv = { + if config.use_quant_conv { + Some(nn::conv2d( + 2 * latent_channels, + 2 * latent_channels, + 1, + conv_cfg, + vs.pp("quant_conv"), + )?) + } else { + None + } + }; + let post_quant_conv = { + if config.use_post_quant_conv { + Some(nn::conv2d( + latent_channels, + latent_channels, + 1, + conv_cfg, + vs.pp("post_quant_conv"), + )?) + } else { + None + } + }; + Ok(Self { + encoder, + decoder, + quant_conv, + post_quant_conv, + config, + }) + } + + /// Returns the distribution in the latent space. + pub fn encode(&self, xs: &Tensor) -> Result { + let xs = self.encoder.forward(xs)?; + let parameters = match &self.quant_conv { + None => xs, + Some(quant_conv) => quant_conv.forward(&xs)?, + }; + DiagonalGaussianDistribution::new(¶meters) + } + + /// Takes as input some sampled values. + pub fn decode(&self, xs: &Tensor) -> Result { + let xs = match &self.post_quant_conv { + None => xs, + Some(post_quant_conv) => &post_quant_conv.forward(xs)?, + }; + self.decoder.forward(xs) + } +} diff --git a/patches/candle-transformers/src/models/stable_lm.rs b/patches/candle-transformers/src/models/stable_lm.rs new file mode 100644 index 0000000000..536f7727e4 --- /dev/null +++ b/patches/candle-transformers/src/models/stable_lm.rs @@ -0,0 +1,433 @@ +//! StableLM model implementation. +//! +//! StableLM is a family of language models trained by Stability AI. +//! This implementation supports the StableLM architecture. +//! +//! Key characteristics: +//! - Grouped query attention (GQA) +//! - Layer normalization +//! - Rotary positional embeddings (RoPE) +//! - Support for different model sizes (3B, 7B) +//! +//! References: +//! - 🤗 [Model Card](https://huggingface.co/stabilityai/stablelm-3b-4e1t) +//! + +use crate::models::with_tracing::{linear, linear_no_bias, Linear}; +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::{Activation, LayerNorm, VarBuilder}; +use serde::Deserialize; +use std::sync::Arc; + +// https://huggingface.co/stabilityai/stablelm-3b-4e1t/blob/main/configuration_stablelm.py +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Config { + pub(crate) vocab_size: usize, + pub(crate) intermediate_size: usize, + pub(crate) hidden_size: usize, + pub(crate) num_hidden_layers: usize, + pub(crate) num_attention_heads: usize, + pub(crate) num_key_value_heads: usize, + pub(crate) hidden_act: Activation, + pub(crate) partial_rotary_factor: f64, + pub(crate) rope_theta: f64, + pub(crate) max_position_embeddings: usize, + pub(crate) layer_norm_eps: f64, + pub(crate) use_cache: bool, + #[serde(default)] + pub(crate) use_qkv_bias: bool, // Used in StableLM-2 + #[serde(default)] + pub(crate) use_flash_attn: bool, // Not in config.json +} + +impl Config { + pub fn stablelm_3b_4e1t(use_flash_attn: bool) -> Self { + Self { + vocab_size: 50304, + intermediate_size: 6912, + hidden_size: 2560, + num_hidden_layers: 32, + num_attention_heads: 32, + num_key_value_heads: 32, + hidden_act: Activation::Silu, + partial_rotary_factor: 0.25, + rope_theta: 10_000., + max_position_embeddings: 4096, + layer_norm_eps: 1e-5, + use_qkv_bias: false, + use_cache: true, + use_flash_attn, + } + } + + pub fn head_dim(&self) -> usize { + self.hidden_size / self.num_attention_heads + } + + pub fn rotary_ndims(&self) -> usize { + (self.head_dim() as f64 * self.partial_rotary_factor) as usize + } + + pub fn num_kv_groups(&self) -> usize { + self.num_attention_heads / self.num_key_value_heads + } + + pub fn set_use_flash_attn(&mut self, use_flash_attn: bool) { + self.use_flash_attn = use_flash_attn + } +} + +#[derive(Debug)] +pub(crate) struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +fn rotate_half(xs: &Tensor) -> Result { + let xs = xs.chunk(2, D::Minus1)?; + Tensor::cat(&[&xs[1].neg()?, &xs[0]], D::Minus1) +} + +impl RotaryEmbedding { + pub(crate) fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let dim = cfg.rotary_ndims(); + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + let freqs = Tensor::cat(&[&freqs, &freqs], D::Minus1)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + pub(crate) fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let cos = cos.unsqueeze(0)?.unsqueeze(0)?; // (1, 1, seq_len, dim) + let sin = sin.unsqueeze(0)?.unsqueeze(0)?; // (1, 1, seq_len, dim) + let q_embed = (q.broadcast_mul(&cos)? + rotate_half(q)?.broadcast_mul(&sin))?; + let k_embed = (k.broadcast_mul(&cos)? + rotate_half(k)?.broadcast_mul(&sin))?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: Activation, + span: tracing::Span, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + let gate_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("gate_proj"))?; + let up_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("up_proj"))?; + let down_proj = linear_no_bias(intermediate_sz, hidden_sz, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_act, + span: tracing::span!(tracing::Level::TRACE, "mlp"), + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = xs.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +#[derive(Debug)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + hidden_size: usize, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, + use_cache: bool, + rotary_ndims: usize, + use_flash_attn: bool, + span: tracing::Span, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let head_dim = cfg.head_dim(); + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let linear_layer = if cfg.use_qkv_bias { + linear + } else { + linear_no_bias + }; + + let q_proj = linear_layer(hidden_sz, num_heads * head_dim, vb.pp("q_proj"))?; + let k_proj = linear_layer(hidden_sz, num_kv_heads * head_dim, vb.pp("k_proj"))?; + let v_proj = linear_layer(hidden_sz, num_kv_heads * head_dim, vb.pp("v_proj"))?; + let o_proj = linear_no_bias(num_heads * head_dim, hidden_sz, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups: cfg.num_kv_groups(), + head_dim, + hidden_size: hidden_sz, + rotary_emb, + kv_cache: None, + use_cache: cfg.use_cache, + rotary_ndims: cfg.rotary_ndims(), + use_flash_attn: cfg.use_flash_attn, + span: tracing::span!(tracing::Level::TRACE, "attn"), + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let _enter = self.span.enter(); + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let (rot_ndims, pass_ndims) = (self.rotary_ndims, self.head_dim - self.rotary_ndims); + let query_rot = query_states.narrow(D::Minus1, 0, rot_ndims)?; + let query_pass = query_states.narrow(D::Minus1, rot_ndims, pass_ndims)?; + let key_rot = key_states.narrow(D::Minus1, 0, rot_ndims)?; + let key_pass = key_states.narrow(D::Minus1, rot_ndims, pass_ndims)?; + let (query_rot, key_rot) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_rot, &key_rot, seqlen_offset)?; + let query_states = Tensor::cat(&[query_rot, query_pass], D::Minus1)?.contiguous()?; + let key_states = Tensor::cat(&[key_rot, key_pass], D::Minus1)?.contiguous()?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + if self.use_cache { + self.kv_cache = Some((key_states.clone(), value_states.clone())); + } + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?.contiguous()?; + let value_states = + crate::utils::repeat_kv(value_states, self.num_kv_groups)?.contiguous()?; + + let attn_output = if self.use_flash_attn { + // flash-attn expects (b_sz, seq_len, nheads, head_dim) + let q = query_states.transpose(1, 2)?; + let k = key_states.transpose(1, 2)?; + let v = value_states.transpose(1, 2)?; + let softmax_scale = 1f32 / (self.head_dim as f32).sqrt(); + flash_attn(&q, &k, &v, softmax_scale, q_len > 1)?.transpose(1, 2)? + } else { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + }; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.hidden_size))? + .apply(&self.o_proj) + } +} + +#[derive(Debug)] +struct DecoderLayer { + self_attn: Attention, + mlp: MLP, + input_layernorm: LayerNorm, + post_attention_layernorm: LayerNorm, + span: tracing::Span, +} + +impl DecoderLayer { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let self_attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let input_layernorm = candle_nn::layer_norm( + cfg.hidden_size, + cfg.layer_norm_eps, + vb.pp("input_layernorm"), + )?; + let post_attention_layernorm = candle_nn::layer_norm( + cfg.hidden_size, + cfg.layer_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + span: tracing::span!(tracing::Level::TRACE, "layer"), + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let _enter = self.span.enter(); + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; + residual + xs + } +} + +#[derive(Debug)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: LayerNorm, + lm_head: Linear, + device: Device, + dtype: DType, + span: tracing::Span, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb_m.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(rotary_emb.clone(), cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = candle_nn::layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb_m.pp("norm"))?; + let lm_head = linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))?; + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + device: vb.device().clone(), + dtype: vb.dtype(), + span: tracing::span!(tracing::Level::TRACE, "model"), + }) + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + // Sliding window mask? + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { f32::NEG_INFINITY } else { 0. })) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let _enter = self.span.enter(); + let (b_size, seq_len) = input_ids.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(b_size, seq_len, seqlen_offset)?; + Some(mask) + }; + let mut xs = self.embed_tokens.forward(input_ids)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)? + } + xs.narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head) + } +} diff --git a/patches/candle-transformers/src/models/starcoder2.rs b/patches/candle-transformers/src/models/starcoder2.rs new file mode 100644 index 0000000000..266221e5c8 --- /dev/null +++ b/patches/candle-transformers/src/models/starcoder2.rs @@ -0,0 +1,351 @@ +//! StarCoder model implementation with quantization support. +//! +//! StarCoder is a large language model optimized for code generation. +//! This implementation provides quantization for reduced memory and compute. +//! +//! Key characteristics: +//! - Causal self-attention mechanism +//! - Multi-query attention (MQA) +//! - LayerNorm for normalization +//! - Absolute positional embeddings +//! - Support for 8-bit quantization +//! +//! References: +//! - 📝 [StarCoder Paper](https://arxiv.org/abs/2305.06161) +//! - 🤗 [Model Card](https://huggingface.co/bigcode/starcoder) +//! + +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::{layer_norm, linear_b, LayerNorm, Linear, VarBuilder}; +use std::sync::Arc; + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + vocab_size: usize, + hidden_size: usize, + intermediate_size: usize, + num_hidden_layers: usize, + num_attention_heads: usize, + num_key_value_heads: usize, + hidden_act: candle_nn::Activation, + max_position_embeddings: usize, + norm_epsilon: f64, + rope_theta: f64, + use_bias: bool, + sliding_window: Option, +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +fn rotate_half(xs: &Tensor) -> Result { + let last_dim = xs.dim(D::Minus1)?; + let xs1 = xs.narrow(D::Minus1, 0, last_dim / 2)?; + let xs2 = xs.narrow(D::Minus1, last_dim / 2, last_dim - last_dim / 2)?; + Tensor::cat(&[&xs2.neg()?, &xs1], D::Minus1) +} + +impl RotaryEmbedding { + fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let dim = cfg.hidden_size / cfg.num_attention_heads; + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + let freqs = Tensor::cat(&[&freqs, &freqs], D::Minus1)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let cos = cos.unsqueeze(0)?.unsqueeze(0)?; // (1, 1, seq_len, dim) + let sin = sin.unsqueeze(0)?.unsqueeze(0)?; // (1, 1, seq_len, dim) + let q_embed = (q.broadcast_mul(&cos)? + rotate_half(q)?.broadcast_mul(&sin))?; + let k_embed = (k.broadcast_mul(&cos)? + rotate_half(k)?.broadcast_mul(&sin))?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + c_fc: Linear, + c_proj: Linear, + act: candle_nn::Activation, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let (h_size, i_size) = (cfg.hidden_size, cfg.intermediate_size); + let c_fc = linear_b(h_size, i_size, cfg.use_bias, vb.pp("c_fc"))?; + let c_proj = linear_b(i_size, h_size, cfg.use_bias, vb.pp("c_proj"))?; + Ok(Self { + c_fc, + c_proj, + act: cfg.hidden_act, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.c_fc)?.apply(&self.act)?.apply(&self.c_proj) + } +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + hidden_size: usize, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + let head_dim = hidden_sz / num_heads; + let b = cfg.use_bias; + let q_proj = linear_b(hidden_sz, num_heads * head_dim, b, vb.pp("q_proj"))?; + let k_proj = linear_b(hidden_sz, num_kv_heads * head_dim, b, vb.pp("k_proj"))?; + let v_proj = linear_b(hidden_sz, num_kv_heads * head_dim, b, vb.pp("v_proj"))?; + let o_proj = linear_b(num_heads * head_dim, hidden_sz, b, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + hidden_size: hidden_sz, + rotary_emb, + kv_cache: None, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let (query_states, key_states) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states, seqlen_offset)?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?; + let value_states = crate::utils::repeat_kv(value_states, self.num_kv_groups)?; + + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + let attn_output = attn_weights.matmul(&value_states)?; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.hidden_size))? + .apply(&self.o_proj) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: MLP, + input_layernorm: LayerNorm, + post_attention_layernorm: LayerNorm, +} + +impl DecoderLayer { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let self_attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let input_layernorm = + layer_norm(cfg.hidden_size, cfg.norm_epsilon, vb.pp("input_layernorm"))?; + let post_attention_layernorm = layer_norm( + cfg.hidden_size, + cfg.norm_epsilon, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; + residual + xs + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: LayerNorm, + lm_head: Linear, + sliding_window: Option, + device: Device, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb_m.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(rotary_emb.clone(), cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = layer_norm(cfg.hidden_size, cfg.norm_epsilon, vb_m.pp("norm"))?; + let lm_head = candle_nn::Linear::new(embed_tokens.embeddings().clone(), None); + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + sliding_window: cfg.sliding_window, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + let sliding_window = self.sliding_window.unwrap_or(tgt_len + 42); + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| { + (0..tgt_len).map(move |j| { + if i < j || j + sliding_window < i { + f32::NEG_INFINITY + } else { + 0. + } + }) + }) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let (b_size, seq_len) = input_ids.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(b_size, seq_len, seqlen_offset)?; + Some(mask) + }; + let mut xs = self.embed_tokens.forward(input_ids)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)? + } + xs.narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +} diff --git a/patches/candle-transformers/src/models/stella_en_v5.rs b/patches/candle-transformers/src/models/stella_en_v5.rs new file mode 100644 index 0000000000..4e98791daa --- /dev/null +++ b/patches/candle-transformers/src/models/stella_en_v5.rs @@ -0,0 +1,801 @@ +//! Stella v5 model implementation. +//! +//! Stella is a dense text embedding model optimized for retrieval and similarity tasks. +//! This implementation provides support for multiple embedding dimensions. +//! +//! Key characteristics: +//! - Dense text embeddings optimized for similarity search +//! - Multiple output dimension support (256 to 8192) +//! - Grouped query attention (GQA) +//! - RMSNorm for layer normalization +//! - Rotary positional embeddings (RoPE) +//! +//! References: +//! - [MRL Framework](https://arxiv.org/abs/2205.13147) +//! - [Model Card](https://huggingface.co/dunzhang/stella_en_1.5B_v5) +//! + +use crate::models::with_tracing::{linear, linear_no_bias, Linear, RmsNorm}; +use candle::{DType, Device, Error, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{layer_norm, Activation, LayerNorm, VarBuilder}; +use std::sync::Arc; + +// internal representation for identifying which model is being used +#[derive(Debug, Default, Copy, Clone, PartialEq, serde::Deserialize)] +pub enum ModelVariant { + #[default] + Large, // 1.5B + Small, // 400M +} + +// Same as `qwen2` family of models with the exception being the `embed_head` +// The final `output` causal modelling head is swapped with a learned `dense` layer, `embed_head` +#[derive(Debug, Default, Clone, PartialEq, serde::Deserialize)] +pub struct Config { + pub variant: ModelVariant, + pub vocab_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub max_position_embeddings: usize, + pub rope_theta: f64, + pub embed_head: EmbedHead, + pub norm_eps: f64, // RMSNorm for 1.5B || LayerNorm for 400M + pub activation_fn: Activation, // Silu for 1.5B || Gelu for 400M + // Unique to 1.5B + pub num_key_value_heads: usize, + // Unique to 400M + pub type_vocab_size: usize, + pub scaling_factor: f64, +} + +// Excerpt from `stella` model card: +// `Stella_en_1.5B_v5` models have been trained on [MRL](https://arxiv.org/abs/2205.13147) enabling multiple output dimensions +// Embed head represents the config for various embedding dims supported +#[derive(Debug, Default, Clone, PartialEq, serde::Deserialize)] +pub struct EmbedHead { + pub in_features: usize, + pub out_features: usize, +} + +/// An enum variant representing the Embedding head dimensions `stella` is trained on +/// As the [model-card](https://huggingface.co/dunzhang/stella_en_1.5B_v5#introduction) suggests, D1024 is good enough for most cases +#[derive(Debug, Default, Clone, Copy)] +pub enum EmbedDim { + Dim256, + Dim768, + #[default] + Dim1024, + Dim2048, + Dim4096, + Dim6144, + Dim8192, +} + +impl EmbedDim { + pub fn config(&self, in_features: usize) -> EmbedHead { + EmbedHead { + in_features, + out_features: match &self { + Self::Dim256 => 256, + Self::Dim768 => 768, + Self::Dim1024 => 1024, + Self::Dim2048 => 2048, + Self::Dim4096 => 4096, + Self::Dim6144 => 6144, + Self::Dim8192 => 8192, + }, + } + } +} + +// Initialize a new `stella_en` model - with 400M variant or 1.5B variant +impl Config { + /// Initialize a new `stella_en_1.5B_v5`` model with given embedding dim + pub fn new_1_5_b_v5(embed_dim: EmbedDim) -> Self { + // Representing config.json at https://huggingface.co/dunzhang/stella_en_1.5B_v5/blob/main/config.json + // Removed `sliding_window` related config which is basically being carried forward from `qwen2` but not used here + Self { + variant: ModelVariant::Large, + activation_fn: candle_nn::Activation::Silu, + vocab_size: 151646, + hidden_size: 1536, + intermediate_size: 8960, + num_hidden_layers: 28, + num_attention_heads: 12, + num_key_value_heads: 2, + max_position_embeddings: 131072, + rope_theta: 1000000., + norm_eps: 1e-06, + embed_head: embed_dim.config(1536), + ..Default::default() + } + } + + /// Initialize new `stella_en_400M_v5` + pub fn new_400_m_v5(embed_dim: EmbedDim) -> Self { + Self { + variant: ModelVariant::Small, + vocab_size: 30528, + hidden_size: 1024, + intermediate_size: 4096, + num_hidden_layers: 24, + num_attention_heads: 16, + max_position_embeddings: 8192, + type_vocab_size: 2, + norm_eps: 1e-12, + scaling_factor: 2.0, + rope_theta: 160000.0, + activation_fn: Activation::Gelu, + embed_head: embed_dim.config(1024), + ..Default::default() + } + } +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let dim = cfg.hidden_size / cfg.num_attention_heads; + // Factoring in `scaling factor` for `400M` variant + let max_seq_len = if cfg.scaling_factor == 0. { + cfg.max_position_embeddings + } else { + ((cfg.max_position_embeddings as f64) * cfg.scaling_factor) as usize + }; + + // let rot_dim = if cfg.variant == ModelVariant::Small { dim / 2 } else { dim }; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| { + // Scaled rope_theta for 400M variant + let rope_theta = if cfg.scaling_factor == 0. { + cfg.rope_theta + } else { + cfg.rope_theta * cfg.scaling_factor + }; + let mut freq = 1. / rope_theta.powf(i as f64 / dim as f64); + + if cfg.scaling_factor != 0. { + freq /= cfg.scaling_factor.powf(2.0 / (dim as f64)) + } + + freq as f32 + }) + .collect(); + + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + + // Calculate position embeddings with scaled sequence length + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + // if cfg.variant == ModelVariant::Small { + // freqs = Tensor::cat(&[&freqs, &freqs], 1)? + // } + + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + // TODO: re-visit this + fn apply_rotary_emb_qkv(&self, q: &Tensor, k: &Tensor) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, 0, seq_len)?; + let sin = self.sin.narrow(0, 0, seq_len)?; + + let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + variant: ModelVariant, + gate_proj: Linear, + up_proj: Option, // `up_proj` only for 1.5B variant + down_proj: Linear, + act_fn: Activation, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + + let (gate_proj, up_proj, down_proj) = match cfg.variant { + ModelVariant::Large => ( + linear_no_bias(hidden_sz, intermediate_sz, vb.pp("gate_proj"))?, + Some(linear_no_bias( + hidden_sz, + intermediate_sz, + vb.pp("up_proj"), + )?), + linear_no_bias(intermediate_sz, hidden_sz, vb.pp("down_proj"))?, + ), + ModelVariant::Small => ( + linear_no_bias(hidden_sz, intermediate_sz * 2, vb.pp("up_gate_proj"))?, + None, + linear(intermediate_sz, hidden_sz, vb.pp("down_proj"))?, + ), + }; + + Ok(Self { + variant: cfg.variant, + gate_proj, + up_proj, + down_proj, + act_fn: cfg.activation_fn, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + let up = self.gate_proj.forward(xs)?; + + let (lhs, rhs) = match self.variant { + ModelVariant::Large => { + let lhs = up.apply(&self.act_fn)?; + let rhs = xs.apply(self.up_proj.as_ref().unwrap())?; + + (lhs, rhs) + } + ModelVariant::Small => { + // Get the dimensions + let (_batch_size, _seq_len, hidden_dim) = up.dims3()?; + let split_size = hidden_dim / 2; + + // Split along the last dimension (hidden_dim) + let up_states = up.narrow(2, 0, split_size)?; + let gate = up.narrow(2, split_size, split_size)?.apply(&self.act_fn)?; + + (up_states, gate) + } + }; + + (lhs * rhs)?.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +struct Attention { + qkv_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + hidden_size: usize, + rotary_emb: Arc, + variant: ModelVariant, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = if num_kv_heads > 0 { + num_heads / num_kv_heads + } else { + 0 + }; + let head_dim = hidden_sz / num_heads; + + let (qkv_proj, o_proj) = match cfg.variant { + ModelVariant::Large => { + // The 1.5B variant comes with separate `q, k, v` layers, let's merge it and standardize + // Weights + let q_w = vb + .pp("q_proj") + .get((num_heads * head_dim, hidden_sz), "weight")?; + let k_w = vb + .pp("k_proj") + .get((num_kv_heads * head_dim, hidden_sz), "weight")?; + let v_w = vb + .pp("v_proj") + .get((num_kv_heads * head_dim, hidden_sz), "weight")?; + // Biases + let q_b = vb.pp("q_proj").get(num_heads * head_dim, "bias")?; + let k_b = vb.pp("k_proj").get(num_kv_heads * head_dim, "bias")?; + let v_b = vb.pp("v_proj").get(num_kv_heads * head_dim, "bias")?; + + let qkv_w = Tensor::cat(&[&q_w, &k_w, &v_w], 0)?; + let qkv_b = Tensor::cat(&[&q_b, &k_b, &v_b], 0)?; + + ( + Linear::from_weights(qkv_w, Some(qkv_b)), + linear_no_bias(num_heads * head_dim, hidden_sz, vb.pp("o_proj"))?, + ) + } + ModelVariant::Small => ( + linear(hidden_sz, 3 * num_heads * head_dim, vb.pp("qkv_proj"))?, + linear(num_heads * head_dim, hidden_sz, vb.pp("o_proj"))?, + ), + }; + + Ok(Self { + qkv_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + hidden_size: hidden_sz, + rotary_emb, + variant: cfg.variant, + }) + } + + fn forward(&mut self, xs: &Tensor, attention_mask: Option<&Tensor>) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let qkv = self.qkv_proj.forward(xs)?; + + let n_kv_heads = match self.variant { + ModelVariant::Large => self.num_kv_heads, + ModelVariant::Small => self.num_heads, + }; + + let (query_states, key_states, value_states) = match self.variant { + ModelVariant::Large => { + let q_sz = self.num_heads * self.head_dim; + let kv_sz = n_kv_heads * self.head_dim; + + let q = qkv.narrow(D::Minus1, 0, q_sz)?.reshape(( + b_sz, + q_len, + self.num_heads, + self.head_dim, + ))?; + let k = qkv.narrow(D::Minus1, q_sz, kv_sz)?.reshape(( + b_sz, + q_len, + n_kv_heads, + self.head_dim, + ))?; + let v = qkv.narrow(D::Minus1, q_sz + kv_sz, kv_sz)?.reshape(( + b_sz, + q_len, + n_kv_heads, + self.head_dim, + ))?; + + (q, k, v) + } + ModelVariant::Small => { + // Split into Q, K, V and reshape to match PyTorch shapes + let qkv = qkv.reshape((b_sz, q_len, 3, self.num_heads, self.head_dim))?; + + ( + qkv.i((.., .., 0, .., ..))?, + qkv.i((.., .., 1, .., ..))?, + qkv.i((.., .., 2, .., ..))?, + ) + } + }; + + let query_states = query_states.transpose(1, 2)?.contiguous()?; + let key_states = key_states.transpose(1, 2)?.contiguous()?; + let value_states = value_states.transpose(1, 2)?.contiguous()?; + + let (query_states, key_states) = self + .rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states)?; + + // The 1.5B is expected to have grouped query attention + let (key_states, value_states) = if self.variant == ModelVariant::Large { + ( + crate::utils::repeat_kv(key_states, self.num_kv_groups)?.contiguous()?, + crate::utils::repeat_kv(value_states, self.num_kv_groups)?.contiguous()?, + ) + } else { + (key_states, value_states) + }; + + let attn_output = { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = query_states.matmul(&key_states.transpose(2, 3)?)?; + let attn_weights = (attn_weights * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + + attn_weights.matmul(&value_states)? + }; + + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.hidden_size))? + .apply(&self.o_proj) + } +} + +#[derive(Debug, Clone)] +enum NormType { + Layer(LayerNorm), + Rms(RmsNorm), +} + +#[derive(Debug, Clone)] +struct Layer { + variant: ModelVariant, + attention: Attention, + mlp: MLP, + // For 1.5B: this is `input_layernorm` + // For 400M: this is `output_layernorm` + layernorm: NormType, + post_attention_layernorm: NormType, +} + +impl Layer { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let attention = Attention::new( + rotary_emb, + cfg, + vb.pp(if cfg.variant == ModelVariant::Large { + "self_attn" + } else { + "attention" + }), + )?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let (layernorm, post_attention_layernorm) = match cfg.variant { + ModelVariant::Large => ( + NormType::Rms(RmsNorm::new( + cfg.hidden_size, + cfg.norm_eps, + vb.pp("input_layernorm"), + )?), + NormType::Rms(RmsNorm::new( + cfg.hidden_size, + cfg.norm_eps, + vb.pp("post_attention_layernorm"), + )?), + ), + ModelVariant::Small => ( + NormType::Layer(layer_norm( + cfg.hidden_size, + candle_nn::LayerNormConfig { + eps: cfg.norm_eps, + ..Default::default() + }, + vb.pp("mlp_ln"), + )?), + NormType::Layer(layer_norm( + cfg.hidden_size, + candle_nn::LayerNormConfig { + eps: cfg.norm_eps, + ..Default::default() + }, + vb.pp("attn_ln"), + )?), + ), + }; + + Ok(Self { + variant: cfg.variant, + attention, + mlp, + layernorm, + post_attention_layernorm, + }) + } + + fn forward(&mut self, xs: &Tensor, attention_mask: Option<&Tensor>) -> Result { + // Here, the application of normalizations and activation calculations differ + // For Large [1.5B]: + // residual = x + // state = other_layernorm(xs) + // state = attention(state) + // state += residual + // residual = state + // state = mlp(attention_layernorm(state)) + // -> residual + state + // For Small [400M]: + // residual = x; + // state = attention(x) + // state += residual + // state = attention_layernorm(state) + // residual = state + // state = mlp(state) + // state += residual + // -> other_layernorm(state) + let residual = xs; + + match self.variant { + ModelVariant::Large => { + let (attn_ln, input_ln) = if let (NormType::Rms(attn_ln), NormType::Rms(input_ln)) = + (&self.post_attention_layernorm, &self.layernorm) + { + (attn_ln, input_ln) + } else { + return Err(candle::error::Error::Msg( + "Stella 1.5B expects RMSNorm".to_string(), + )); + }; + + let xs = input_ln.forward(xs)?; + let xs = (self.attention.forward(&xs, attention_mask)? + residual)?; + + let residual = &xs; + let xs = xs.apply(attn_ln)?.apply(&self.mlp)?; + + residual + xs + } + ModelVariant::Small => { + let (attn_ln, output_ln) = + if let (NormType::Layer(attn_ln), NormType::Layer(input_ln)) = + (&self.post_attention_layernorm, &self.layernorm) + { + (attn_ln, input_ln) + } else { + return Err(candle::error::Error::Msg( + "Stella 400M expects RMSNorm".to_string(), + )); + }; + + let xs = (self.attention.forward(xs, attention_mask)? + residual)?; + let xs = attn_ln.forward(&xs)?; + + let residual = &xs; + let xs = (self.mlp.forward(&xs)? + residual)?; + + output_ln.forward(&xs) + } + } + } +} + +#[derive(Debug, Clone)] +pub struct Embeddings { + variant: ModelVariant, + // For 1.5B: this is the `embed_tokens` + // For 400M: this is the `word_embeddings` + embeddings: candle_nn::Embedding, + // following are specifically for 400M + token_type_embeddings: Option, + layer_norm: Option, + position_ids: Option, +} + +impl Embeddings { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let (embeddings, token_type_embeddings, layer_norm, position_ids) = match cfg.variant { + ModelVariant::Large => ( + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("embed_tokens"))?, + None, + None, + None, + ), + ModelVariant::Small => { + let vb = vb.pp("embeddings"); + let weight = vb.pp("LayerNorm").get_with_hints( + cfg.hidden_size, + "weight", + candle_nn::Init::Const(1.0), + )?; + let bias = vb.pp("LayerNorm").get_with_hints( + cfg.hidden_size, + "bias", + candle_nn::Init::Const(0.0), + )?; + let dev = bias.device().clone(); + + let layer_norm = candle_nn::LayerNorm::new(weight, bias, cfg.norm_eps); + + ( + candle_nn::embedding( + cfg.vocab_size, + cfg.hidden_size, + vb.pp("word_embeddings"), + )?, + Some(candle_nn::embedding( + cfg.type_vocab_size, + cfg.hidden_size, + vb.pp("token_type_embeddings"), + )?), + Some(layer_norm), + Some(Tensor::arange( + 0u32, + cfg.max_position_embeddings as u32, + &dev, + )?), + ) + } + }; + + Ok(Self { + variant: cfg.variant, + embeddings, + token_type_embeddings, + layer_norm, + position_ids, + }) + } +} + +impl Module for Embeddings { + fn forward(&self, xs: &Tensor) -> Result { + let embd = self.embeddings.forward(xs)?; + // For 1.5B just forward the embeddings + if self.variant == ModelVariant::Large { + return Ok(embd); + } + + let (token_type_embed, layer_norm, pos_ids) = + if let (Some(token_type_embd), Some(layer_norm), Some(position_ids)) = ( + &self.token_type_embeddings, + &self.layer_norm, + &self.position_ids, + ) { + (token_type_embd, layer_norm, position_ids) + } else { + return Err(Error::Msg( + "Stella 400M requires `token_type_embeddings`, `layer_norm` and `position_ids`" + .to_string(), + )); + }; + + let (batch_size, seq_length) = xs.dims2()?; + + let pos_ids = pos_ids + .as_ref() + .narrow(0, 0, seq_length)? + .expand((batch_size, seq_length))?; + + layer_norm.forward(&embd.add(&token_type_embed.forward(&pos_ids.zeros_like()?)?)?) + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embeddings: Embeddings, + layers: Vec, + norm: Option, + device: Device, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = match cfg.variant { + ModelVariant::Large => vb.pp("model"), + ModelVariant::Small => vb.pp("new"), + }; + // let embed_tokens = + // candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let embeddings = Embeddings::new(cfg, vb_m.clone())?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb_m.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = match cfg.variant { + ModelVariant::Large => vb_m.pp("layers"), + ModelVariant::Small => vb_m.pp("encoder").pp("layer"), + }; + for layer_idx in 0..cfg.num_hidden_layers { + let layer = Layer::new(rotary_emb.clone(), cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = match cfg.variant { + ModelVariant::Large => Some(RmsNorm::new( + cfg.hidden_size, + cfg.norm_eps, + vb_m.pp("norm"), + )?), + ModelVariant::Small => None, + }; + Ok(Self { + embeddings, + layers, + norm, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + fn prepare_attention_mask(&self, attn_mask: &Tensor) -> Result { + let (b_sz, sql_len) = attn_mask.dims2()?; + let mut mask: Vec = vec![]; + for b in 0..b_sz { + mask.push(attn_mask.i((b, ..))?.expand((1, 1, sql_len, sql_len))?); + } + let mask = Tensor::cat(&mask, 0)?; + let on_true = mask.zeros_like()?.to_dtype(self.dtype)?; + let on_false = Tensor::new(f32::NEG_INFINITY, &self.device)? + .broadcast_as(mask.shape())? + .to_dtype(self.dtype)?; + mask.where_cond(&on_true, &on_false) + } + + pub fn forward(&mut self, input_ids: &Tensor, mask: &Tensor) -> Result { + let (_, seq_len) = input_ids.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + // This is not a `causal language modelling` task, we'll need to prepare a `non-causal` attention + Some(self.prepare_attention_mask(mask)?) + }; + + let mut xs = self.embeddings.forward(input_ids)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref())? + } + + if let Some(n) = &self.norm { + xs.apply(n) + } else { + Ok(xs) + } + } +} + +#[derive(Debug)] +pub struct EmbeddingModel { + base_model: Model, + lm_head: Linear, +} + +impl EmbeddingModel { + pub fn new(cfg: &Config, base_vb: VarBuilder, embed_vb: VarBuilder) -> Result { + let base_model = Model::new(cfg, base_vb.clone())?; + let lm_head = linear( + cfg.embed_head.in_features, + cfg.embed_head.out_features, + embed_vb.pp("linear"), + )?; + + Ok(Self { + base_model, + lm_head, + }) + } + + pub fn forward(&mut self, input_ids: &Tensor, mask: &Tensor) -> Result { + let x = self.base_model.forward(input_ids, mask)?; + let x = self.pool(&x, mask)?; + + // No matter what keeping the final activations as F32 helps with the accuracy + self.lm_head.forward(&x.to_dtype(DType::F32)?) // [B_sz, dim_size] + } + + /// Same as forward pass but normalizes the output + pub fn forward_norm(&mut self, input_ids: &Tensor, mask: &Tensor) -> Result { + let x = self.forward(input_ids, mask)?; + // Normalize + x.broadcast_div(&x.sqr()?.sum_keepdim(1)?.sqrt()?) + } + + fn pool(&self, x: &Tensor, mask: &Tensor) -> Result { + let mask = mask.to_dtype(x.dtype())?; // [B_Sz, Seq_len] + let (batch_size, seq_len, hidden_dim) = x.dims3()?; + // expanding the shape of the mask from [B_Sz, Seq_len] -> [B_Sz, Seq_len, Hidden_size] + let mask_expanded = mask + .unsqueeze(2)? + .broadcast_as((batch_size, seq_len, hidden_dim))?; // [B_Sz, Seq_len, Hidden_dim] + + let x = (x * &mask_expanded)?; + + // Sum + let sum_mask = mask + .sum(1)? + .unsqueeze(1)? + .expand((batch_size, hidden_dim))?; + x.sum(1)? / sum_mask + } +} diff --git a/patches/candle-transformers/src/models/t5.rs b/patches/candle-transformers/src/models/t5.rs new file mode 100644 index 0000000000..5d23549f21 --- /dev/null +++ b/patches/candle-transformers/src/models/t5.rs @@ -0,0 +1,953 @@ +//! T5 model implementation. +//! +//! T5 (Text-to-Text Transfer Transformer) is a unified text-to-text transformer model. +//! This implementation follows the original model architecture. +//! +//! Key characteristics: +//! - Text-to-text framework +//! - Relative positional embeddings +//! - T5-specific layer normalization +//! - Encoder-decoder architecture +//! - Support for sequence-to-sequence tasks +//! +//! References: +//! - ⚡ [Interactive Wasm Example](https://huggingface.co/spaces/radames/Candle-T5-Generation-Wasm) +//! - 💻[GH Model](https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/modeling_t5.py) +//! - 🤗 [HF Link](https://huggingface.co/docs/transformers/model_doc/t5) +//! - 📝 [T5 Paper](https://arxiv.org/abs/1910.10683) +//! +//! # Encoder-decoder example: +//! +//! ```bash +//! cargo run --example t5 --release -- \ +//! --model-id "t5-small" \ +//! --prompt "translate to German: A beautiful candle." \ +//! --decode +//! > ... +//! > Eine schöne Kerze. +//! > 9 tokens generated (2.42 token/s) +//! ``` +//! +//! Variants such as [flan-t5](https://huggingface.co/google/flan-t5-small), [flan-ul2](https://huggingface.co/google/flan-ul2) (with `--revision "refs/pr/25"`), and [Co-EdIT](https://huggingface.co/grammarly/coedit-large) are also supported. +//! +//! # Translation with MADLAD +//! +//! +//! [MADLAD-400](https://arxiv.org/abs/2309.04662) is a series of multilingual machine translation T5 models trained on 250 billion tokens covering over 450 languages using publicly available data. These models are competitive with significantly larger models. +//! +//! ```bash +//! cargo run --example t5 --release -- \ +//! --model-id "jbochi/madlad400-3b-mt" \ +//! --prompt "<2de> How are you, my friend?" \ +//! --decode --temperature 0 +//! ... +//! Wie geht es dir, mein Freund? +//! ``` +//! +//! ## Sentence embedding example +//! +//! ```bash +//! cargo run --example t5 --release -- \ +//! --model-id "t5-small" --prompt "A beautiful candle." +//! ... +//! [[[ 0.0515, -0.0541, -0.0761, ..., -0.0392, 0.1511, -0.0265], +//! [-0.0974, 0.0998, -0.1659, ..., -0.2450, 0.1738, -0.0164], +//! [ 0.0624, -0.1024, 0.0430, ..., -0.1388, 0.0564, -0.2962], +//! [-0.0389, -0.1173, 0.0026, ..., 0.1064, -0.1065, 0.0990], +//! [ 0.1300, 0.0027, -0.0326, ..., 0.0026, -0.0317, 0.0851]]] +//! Tensor[[1, 5, 512], f32] +//! Took 303.766583ms +//! ``` + +use crate::models::with_tracing::Embedding; +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::{Activation, VarBuilder}; +use serde::Deserialize; +use std::sync::Arc; + +#[derive(Debug, Clone)] +pub struct Linear { + weight: Tensor, + span: tracing::Span, +} + +pub fn linear_no_bias(d1: usize, d2: usize, vb: VarBuilder) -> Result { + let init_ws = candle_nn::init::DEFAULT_KAIMING_NORMAL; + let weight = vb.get_with_hints((d2, d1), "weight", init_ws)?; + let span = tracing::span!(tracing::Level::TRACE, "linear"); + Ok(Linear { weight, span }) +} + +impl Module for Linear { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let weight = self.weight.to_dtype(xs.dtype())?; + let w = match *xs.dims() { + [b1, b2, _, _] => weight.broadcast_left((b1, b2))?.t()?, + [bsize, _, _] => weight.broadcast_left(bsize)?.t()?, + _ => weight.t()?, + }; + xs.matmul(&w) + } +} + +fn default_relative_attention_max_distance() -> usize { + 128 +} + +fn default_is_decoder() -> bool { + false +} + +fn default_use_cache() -> bool { + true +} + +fn default_tie_word_embeddings() -> bool { + true +} + +fn get_mask(size: usize, device: &Device) -> Result { + let mask: Vec<_> = (0..size) + .flat_map(|i| (0..size).map(move |j| u8::from(j > i))) + .collect(); + Tensor::from_slice(&mask, (size, size), device) +} + +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) +} + +#[derive(Debug, Deserialize, Default, Clone, PartialEq)] +pub struct ActivationWithOptionalGating { + pub gated: bool, + pub activation: candle_nn::Activation, +} + +pub fn deserialize_feed_forward_proj_activation<'de, D>( + deserializer: D, +) -> std::result::Result +where + D: serde::de::Deserializer<'de>, +{ + match String::deserialize(deserializer)?.as_str() { + "gated-gelu" => Ok(ActivationWithOptionalGating { + gated: true, + activation: candle_nn::Activation::NewGelu, + }), + "gated-silu" => Ok(ActivationWithOptionalGating { + gated: true, + activation: candle_nn::Activation::Silu, + }), + buf => { + let activation = serde_plain::from_str(buf).map_err(serde::de::Error::custom)?; + Ok(ActivationWithOptionalGating { + gated: false, + activation, + }) + } + } +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub d_model: usize, + pub d_kv: usize, + pub d_ff: usize, + pub num_layers: usize, + pub num_decoder_layers: Option, + pub num_heads: usize, + pub relative_attention_num_buckets: usize, + #[serde(default = "default_relative_attention_max_distance")] + pub relative_attention_max_distance: usize, + pub dropout_rate: f64, + pub layer_norm_epsilon: f64, + pub initializer_factor: f64, + #[serde(default, deserialize_with = "deserialize_feed_forward_proj_activation")] + pub feed_forward_proj: ActivationWithOptionalGating, + #[serde(default = "default_tie_word_embeddings")] + pub tie_word_embeddings: bool, + #[serde(default = "default_is_decoder")] + pub is_decoder: bool, + pub is_encoder_decoder: bool, + #[serde(default = "default_use_cache")] + pub use_cache: bool, + pub pad_token_id: usize, + pub eos_token_id: usize, + pub decoder_start_token_id: Option, +} + +impl Default for Config { + fn default() -> Self { + Self { + vocab_size: 32128, + d_model: 512, + d_kv: 64, + d_ff: 2048, + num_layers: 6, + num_decoder_layers: None, + num_heads: 8, + relative_attention_num_buckets: 32, + relative_attention_max_distance: 128, + dropout_rate: 0.1, + layer_norm_epsilon: 1e-6, + initializer_factor: 1.0, + feed_forward_proj: ActivationWithOptionalGating { + gated: false, + activation: Activation::Relu, + }, + tie_word_embeddings: true, + is_decoder: false, + is_encoder_decoder: true, + use_cache: true, + pad_token_id: 0, + eos_token_id: 1, + decoder_start_token_id: Some(0), + } + } +} + +impl Config { + // https://huggingface.co/facebook/musicgen-small/blob/495da4ad086b3416a27c6187f9239f9fd96f3962/config.json#L184 + pub fn musicgen_small() -> Self { + Self { + d_ff: 3072, + d_kv: 64, + d_model: 768, + dropout_rate: 0.1, + eos_token_id: 1, + feed_forward_proj: ActivationWithOptionalGating { + gated: false, + activation: Activation::Relu, + }, + tie_word_embeddings: true, + initializer_factor: 1.0, + is_decoder: false, + is_encoder_decoder: true, + layer_norm_epsilon: 1e-6, + num_decoder_layers: Some(12), + num_heads: 12, + num_layers: 12, + pad_token_id: 0, + decoder_start_token_id: Some(0), + relative_attention_max_distance: 128, + relative_attention_num_buckets: 32, + use_cache: true, + vocab_size: 32128, + } + } +} + +#[derive(Debug, Clone)] +struct T5LayerNorm { + weight: Tensor, + variance_epsilon: f64, + span: tracing::Span, +} + +impl T5LayerNorm { + fn load(h: usize, eps: f64, vb: VarBuilder) -> Result { + let weight = vb.get(h, "weight")?; + Ok(Self { + weight, + variance_epsilon: eps, + span: tracing::span!(tracing::Level::TRACE, "layer-norm"), + }) + } +} + +impl Module for T5LayerNorm { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let dtype = xs.dtype(); + let xs_f32 = xs.to_dtype(DType::F32)?; + // variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) + let variance = xs_f32.sqr()?.mean_keepdim(D::Minus1)?; + let xs = xs_f32.broadcast_div(&(variance + self.variance_epsilon)?.sqrt()?)?; + let xs = xs.to_dtype(dtype)?; + let xs = xs.broadcast_mul(&self.weight.to_dtype(dtype)?)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct T5DenseActDense { + wi: Linear, + wo: Linear, + act: Activation, + span: tracing::Span, +} + +impl T5DenseActDense { + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let wi = linear_no_bias(cfg.d_model, cfg.d_ff, vb.pp("wi"))?; + let wo = linear_no_bias(cfg.d_ff, cfg.d_model, vb.pp("wo"))?; + Ok(Self { + wi, + wo, + act: Activation::Relu, + span: tracing::span!(tracing::Level::TRACE, "dense-act-dense"), + }) + } +} + +impl Module for T5DenseActDense { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let xs = self.wi.forward(xs)?; + let xs = self.act.forward(&xs)?; + let xs = self.wo.forward(&xs)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct T5DenseGatedActDense { + wi_0: Linear, + wi_1: Linear, + wo: Linear, + act: Activation, + span: tracing::Span, +} + +impl T5DenseGatedActDense { + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let wi_0 = linear_no_bias(cfg.d_model, cfg.d_ff, vb.pp("wi_0"))?; + let wi_1 = linear_no_bias(cfg.d_model, cfg.d_ff, vb.pp("wi_1"))?; + let wo = linear_no_bias(cfg.d_ff, cfg.d_model, vb.pp("wo"))?; + Ok(Self { + wi_0, + wi_1, + wo, + act: cfg.feed_forward_proj.activation, + span: tracing::span!(tracing::Level::TRACE, "dense-gated-act-dense"), + }) + } +} + +impl Module for T5DenseGatedActDense { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let hidden_gelu = self.act.forward(&self.wi_0.forward(xs)?)?; + let hidden_linear = self.wi_1.forward(xs)?; + let xs = hidden_gelu.broadcast_mul(&hidden_linear)?; + let xs = self.wo.forward(&xs)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct T5LayerFF { + dense_act: Option, + gated_dense_act: Option, + layer_norm: T5LayerNorm, + span: tracing::Span, +} + +impl T5LayerFF { + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let layer_norm = + T5LayerNorm::load(cfg.d_model, cfg.layer_norm_epsilon, vb.pp("layer_norm"))?; + let (dense_act, gated_dense_act) = if cfg.feed_forward_proj.gated { + ( + None, + Some(T5DenseGatedActDense::load(vb.pp("DenseReluDense"), cfg)?), + ) + } else { + ( + Some(T5DenseActDense::load(vb.pp("DenseReluDense"), cfg)?), + None, + ) + }; + Ok(Self { + dense_act, + gated_dense_act, + layer_norm, + span: tracing::span!(tracing::Level::TRACE, "layer-ff"), + }) + } +} + +impl Module for T5LayerFF { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + let ys = self.layer_norm.forward(xs)?; + let ys = match &self.dense_act { + Some(dense_act) => dense_act.forward(&ys)?, + None => self.gated_dense_act.as_ref().unwrap().forward(&ys)?, + }; + let xs = (xs + ys)?; + Ok(xs) + } +} + +#[derive(Debug, Clone)] +struct T5Attention { + q: Linear, + k: Linear, + v: Linear, + o: Linear, + n_heads: usize, + d_kv: usize, + relative_attention_bias: Option, + relative_attention_num_buckets: usize, + relative_attention_max_distance: usize, + inner_dim: usize, + use_cache: bool, + kv_cache: Option<(Tensor, Tensor)>, + span: tracing::Span, + span_cache: tracing::Span, + span_mm: tracing::Span, + span_sm: tracing::Span, +} + +impl T5Attention { + fn load( + has_relative_attention_bias: bool, + decoder: bool, + vb: VarBuilder, + cfg: &Config, + ) -> Result { + let inner_dim = cfg.num_heads * cfg.d_kv; + let q = linear_no_bias(cfg.d_model, inner_dim, vb.pp("q"))?; + let k = linear_no_bias(cfg.d_model, inner_dim, vb.pp("k"))?; + let v = linear_no_bias(cfg.d_model, inner_dim, vb.pp("v"))?; + let o = linear_no_bias(inner_dim, cfg.d_model, vb.pp("o"))?; + let relative_attention_bias = if has_relative_attention_bias { + let emb = Embedding::new( + cfg.relative_attention_num_buckets, + cfg.num_heads, + vb.pp("relative_attention_bias"), + )?; + Some(emb) + } else { + None + }; + Ok(Self { + q, + k, + v, + o, + n_heads: cfg.num_heads, + d_kv: cfg.d_kv, + relative_attention_bias, + relative_attention_num_buckets: cfg.relative_attention_num_buckets, + relative_attention_max_distance: cfg.relative_attention_max_distance, + inner_dim, + use_cache: cfg.use_cache && decoder, + kv_cache: None, + span: tracing::span!(tracing::Level::TRACE, "attention"), + span_cache: tracing::span!(tracing::Level::TRACE, "attention-cache"), + span_mm: tracing::span!(tracing::Level::TRACE, "attention-mm"), + span_sm: tracing::span!(tracing::Level::TRACE, "attention-sm"), + }) + } + + fn forward( + &mut self, + xs: &Tensor, + position_bias: Option<&Tensor>, + key_value_states: Option<&Tensor>, + mask: Option<&Tensor>, + ) -> Result<(Tensor, Option)> { + // Performs Self-attention (if key_value_states is None) or attention + // over source sentence (provided by key_value_states). + let _enter = self.span.enter(); + let kv_input = match key_value_states { + None => xs, + Some(key_value_states) => key_value_states, + }; + let (b_sz, q_len) = (xs.dim(0)?, xs.dim(1)?); + let kv_len = kv_input.dim(1)?; + let q = self.q.forward(xs)?; + let k = self.k.forward(kv_input)?; + let v = self.v.forward(kv_input)?; + let q = q + .reshape((b_sz, q_len, self.n_heads, self.d_kv))? + .transpose(1, 2)? + .contiguous()?; + let mut k = k + .reshape((b_sz, kv_len, self.n_heads, self.d_kv))? + .transpose(1, 2)?; + let mut v = v + .reshape((b_sz, kv_len, self.n_heads, self.d_kv))? + .transpose(1, 2)?; + + if self.use_cache && key_value_states.is_none() { + let _enter = self.span_cache.enter(); + if let Some((kv_cache_k, kv_cache_v)) = &self.kv_cache { + k = Tensor::cat(&[kv_cache_k, &k], 2)?; + v = Tensor::cat(&[kv_cache_v, &v], 2)?; + }; + self.kv_cache = Some((k.clone(), v.clone())); + }; + let k = k.contiguous()?; + let v = v.contiguous()?; + // TODO: Use flash_attn. + let scores = { + let _enter = self.span_mm.enter(); + q.matmul(&k.t()?)? + }; + let scores = match mask { + None => scores, + Some(mask) => masked_fill( + &scores, + &mask + .unsqueeze(0)? + .unsqueeze(0)? + .repeat((b_sz, self.n_heads))?, + f32::NEG_INFINITY, + )?, + }; + + let (scores, position_bias) = match position_bias { + Some(position_bias) => ( + scores.broadcast_add(position_bias)?, + Some(position_bias.clone()), + ), + None => match &self.relative_attention_bias { + None => (scores, None), + Some(relative_attention_bias) => { + // This only handles the bidirectional case. + let kv_len = k.dim(2)?; + let (q_start, q_end) = match self.use_cache { + true => ((kv_len - q_len) as u32, kv_len as u32), + false => (0_u32, kv_len as u32), + }; + let num_buckets = self.relative_attention_num_buckets as u32 / 2; + let max_exact = num_buckets / 2; + let relative_position = (q_start..q_end) + .map(|i| { + (0..kv_len as u32) + .map(|j| { + if i < j { + if j - i < max_exact { + j - i + num_buckets + } else { + let b = f32::log( + (j - i) as f32 / max_exact as f32, + self.relative_attention_max_distance as f32 + / max_exact as f32, + ) * (num_buckets - max_exact) as f32; + u32::min( + max_exact + num_buckets + b as u32, + self.relative_attention_num_buckets as u32 - 1, + ) + } + } else if i - j < max_exact { + i - j + } else { + let b = f32::log( + (i - j) as f32 / max_exact as f32, + self.relative_attention_max_distance as f32 + / max_exact as f32, + ) * (num_buckets - max_exact) as f32; + u32::min(max_exact + b as u32, num_buckets - 1) + } + }) + .collect::>() + }) + .collect::>>(); + let relative_buckets = Tensor::new(relative_position, q.device())?; + let position_bias = relative_attention_bias + .forward(&relative_buckets)? + .permute((2, 0, 1))? + .unsqueeze(0)? + .to_dtype(scores.dtype())?; + (scores.broadcast_add(&position_bias)?, Some(position_bias)) + // TODO: position_bias_masked? + } + }, + }; + + let attn_weights = { + let _enter = self.span_sm.enter(); + candle_nn::ops::softmax_last_dim(&scores)? + }; + let attn_output = attn_weights.matmul(&v)?; + let attn_output = attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.inner_dim))?; + let attn_output = self.o.forward(&attn_output)?; + Ok((attn_output, position_bias)) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +struct T5LayerSelfAttention { + self_attention: T5Attention, + layer_norm: T5LayerNorm, + span: tracing::Span, +} + +impl T5LayerSelfAttention { + fn load(h: bool, d: bool, vb: VarBuilder, cfg: &Config) -> Result { + let self_attention = T5Attention::load(h, d, vb.pp("SelfAttention"), cfg)?; + let layer_norm = + T5LayerNorm::load(cfg.d_model, cfg.layer_norm_epsilon, vb.pp("layer_norm"))?; + Ok(Self { + self_attention, + layer_norm, + span: tracing::span!(tracing::Level::TRACE, "self-attn"), + }) + } + + fn forward( + &mut self, + xs: &Tensor, + position_bias: Option<&Tensor>, + mask: Option<&Tensor>, + ) -> Result<(Tensor, Option)> { + let _enter = self.span.enter(); + let normed_xs = self.layer_norm.forward(xs)?; + let (ys, position_bias) = + self.self_attention + .forward(&normed_xs, position_bias, None, mask)?; + let ys = (xs + ys)?; + Ok((ys, position_bias)) + } + + fn clear_kv_cache(&mut self) { + self.self_attention.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +struct T5LayerCrossAttention { + cross_attention: T5Attention, + layer_norm: T5LayerNorm, + span: tracing::Span, +} + +impl T5LayerCrossAttention { + fn load(decoder: bool, vb: VarBuilder, cfg: &Config) -> Result { + let cross_attention = T5Attention::load(false, decoder, vb.pp("EncDecAttention"), cfg)?; + let layer_norm = + T5LayerNorm::load(cfg.d_model, cfg.layer_norm_epsilon, vb.pp("layer_norm"))?; + Ok(Self { + cross_attention, + layer_norm, + span: tracing::span!(tracing::Level::TRACE, "cross-attn"), + }) + } + + fn forward( + &mut self, + hidden_states: &Tensor, + position_bias: Option<&Tensor>, + key_value_states: &Tensor, + ) -> Result<(Tensor, Option)> { + let _enter = self.span.enter(); + let normed_hidden_states = self.layer_norm.forward(hidden_states)?; + let (ys, position_bias) = self.cross_attention.forward( + &normed_hidden_states, + position_bias, + Some(key_value_states), + None, + )?; + let ys = (hidden_states + ys)?; + Ok((ys, position_bias)) + } + + fn clear_kv_cache(&mut self) { + self.cross_attention.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +struct T5Block { + self_attn: T5LayerSelfAttention, + cross_attn: Option, + ff: T5LayerFF, + span: tracing::Span, +} + +impl T5Block { + fn load( + has_relative_attention_bias: bool, + decoder: bool, + vb: VarBuilder, + cfg: &Config, + ) -> Result { + let vb = vb.pp("layer"); + let self_attn = + T5LayerSelfAttention::load(has_relative_attention_bias, decoder, vb.pp("0"), cfg)?; + let cross_attn = if cfg.is_decoder { + Some(T5LayerCrossAttention::load(decoder, vb.pp("1"), cfg)?) + } else { + None + }; + let ff_i = if cross_attn.is_some() { 2 } else { 1 }; + let ff = T5LayerFF::load(vb.pp(ff_i.to_string()), cfg)?; + Ok(Self { + self_attn, + cross_attn, + ff, + span: tracing::span!(tracing::Level::TRACE, "block"), + }) + } + + fn forward( + &mut self, + xs: &Tensor, + position_bias: Option<&Tensor>, + encoder_hidden_states: Option<&Tensor>, + ) -> Result<(Tensor, Option)> { + let _enter = self.span.enter(); + // TODO: Cache masks + let mask = match self.cross_attn.is_some() { + true => { + let mask_len = xs.dim(1)?; + // If the input seq length is 1, no need for a mask, this is also helpful to avoid shape + // issues when using the KV cache in the decoder. + if mask_len <= 1 { + None + } else { + Some(get_mask(mask_len, xs.device())?) + } + } + false => None, + }; + let (mut xs, position_bias) = self.self_attn.forward(xs, position_bias, mask.as_ref())?; + // TODO: clamp for f16? + if let Some(cross_attn) = &mut self.cross_attn { + (xs, _) = cross_attn.forward(&xs, None, encoder_hidden_states.unwrap())?; + // TODO: clamp for f16? + } + let xs = self.ff.forward(&xs)?; + // TODO: clamp for f16? + Ok((xs, position_bias)) + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache(); + self.cross_attn.iter_mut().for_each(|c| c.clear_kv_cache()); + } +} + +#[derive(Debug, Clone)] +struct T5Stack { + block: Vec, + shared: Arc, + final_layer_norm: T5LayerNorm, + span: tracing::Span, +} + +impl T5Stack { + fn load(decoder: bool, vb: VarBuilder, shared: &Arc, cfg: &Config) -> Result { + let block = (0..cfg.num_layers) + .map(|i| T5Block::load(i == 0, decoder, vb.pp(format!("block.{i}")), cfg)) + .collect::>>()?; + let final_layer_norm = T5LayerNorm::load( + cfg.d_model, + cfg.layer_norm_epsilon, + vb.pp("final_layer_norm"), + )?; + Ok(Self { + block, + shared: shared.clone(), + final_layer_norm, + span: tracing::span!(tracing::Level::TRACE, "stack"), + }) + } + + fn forward( + &mut self, + input_ids: &Tensor, + encoder_hidden_states: Option<&Tensor>, + ) -> Result { + self.forward_dt(input_ids, encoder_hidden_states, None) + } + + fn forward_dt( + &mut self, + input_ids: &Tensor, + encoder_hidden_states: Option<&Tensor>, + dtype: Option, + ) -> Result { + let _enter = self.span.enter(); + let input_embeds = self.shared.as_ref().forward(input_ids)?; + let input_embeds = match dtype { + None => input_embeds, + Some(dtype) => input_embeds.to_dtype(dtype)?, + }; + let mut hidden_states = input_embeds; + let mut position_bias = None; + for block in self.block.iter_mut() { + (hidden_states, position_bias) = block.forward( + &hidden_states, + position_bias.as_ref(), + encoder_hidden_states, + )? + } + self.final_layer_norm.forward(&hidden_states) + } + + fn clear_kv_cache(&mut self) { + self.block.iter_mut().for_each(|b| b.clear_kv_cache()) + } +} + +#[derive(Debug, Clone)] +pub struct T5EncoderModel { + encoder: T5Stack, + device: Device, + span: tracing::Span, +} + +impl T5EncoderModel { + pub fn load(vb: VarBuilder, cfg: &Config) -> Result { + let shared_vb = if vb.contains_tensor("shared.weight") { + vb.pp("shared") + } else if vb.contains_tensor("decoder.embed_tokens") { + vb.pp("decoder").pp("embed_tokens") + } else { + vb.pp("encoder").pp("embed_tokens") + }; + let shared = Embedding::new(cfg.vocab_size, cfg.d_model, shared_vb)?; + let shared = Arc::new(shared); + let encoder = T5Stack::load(false, vb.pp("encoder"), &shared, cfg)?; + Ok(Self { + encoder, + device: vb.device().clone(), + span: tracing::span!(tracing::Level::TRACE, "encoder"), + }) + } + + pub fn forward(&mut self, input_ids: &Tensor) -> Result { + let _enter = self.span.enter(); + self.encoder.forward(input_ids, None) + } + + pub fn forward_dt(&mut self, input_ids: &Tensor, dtype: Option) -> Result { + let _enter = self.span.enter(); + self.encoder.forward_dt(input_ids, None, dtype) + } + + pub fn device(&self) -> &Device { + &self.device + } + + pub fn clear_kv_cache(&mut self) { + self.encoder.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct T5ForConditionalGeneration { + encoder: T5Stack, + decoder: T5Stack, + d_model: usize, + tie_word_embeddings: bool, + lm_head: Option, + shared: Arc, + device: Device, + span_decode: tracing::Span, + span_decode_head: tracing::Span, +} + +impl T5ForConditionalGeneration { + pub fn load(vb: VarBuilder, cfg: &Config) -> Result { + assert!(cfg.is_encoder_decoder); + let d_model = cfg.d_model; + let shared_vb = if vb.contains_tensor("shared.weight") { + vb.pp("shared") + } else { + vb.pp("decoder").pp("embed_tokens") + }; + let shared = Embedding::new(cfg.vocab_size, cfg.d_model, shared_vb)?; + let shared = Arc::new(shared); + + let mut encoder_cfg = cfg.clone(); + encoder_cfg.is_decoder = false; + encoder_cfg.use_cache = false; + encoder_cfg.is_encoder_decoder = false; + let encoder = T5Stack::load(false, vb.pp("encoder"), &shared, &encoder_cfg)?; + + let mut decoder_cfg = cfg.clone(); + decoder_cfg.is_decoder = true; + decoder_cfg.is_encoder_decoder = false; + decoder_cfg.num_layers = cfg.num_decoder_layers.unwrap_or(cfg.num_layers); + let decoder = T5Stack::load(true, vb.pp("decoder"), &shared, &decoder_cfg)?; + + let tie_word_embeddings = cfg.tie_word_embeddings; + let lm_head = if tie_word_embeddings { + None + } else { + Some(linear_no_bias( + cfg.d_model, + cfg.vocab_size, + vb.pp("lm_head"), + )?) + }; + + Ok(Self { + encoder, + decoder, + d_model, + tie_word_embeddings, + lm_head, + shared, + device: vb.device().clone(), + span_decode: tracing::span!(tracing::Level::TRACE, "decode"), + span_decode_head: tracing::span!(tracing::Level::TRACE, "decode-head"), + }) + } + + pub fn encode(&mut self, input_ids: &Tensor) -> Result { + self.encoder.forward(input_ids, None) + } + + pub fn decode( + &mut self, + decoder_input_ids: &Tensor, + encoder_output: &Tensor, + ) -> Result { + let _enter = self.span_decode.enter(); + let decoder_output = self + .decoder + .forward(decoder_input_ids, Some(encoder_output))?; + + let scaling_factor = if self.tie_word_embeddings { + // Rescale output before projecting on vocab + // See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586 + (self.d_model as f64).sqrt() + } else { + 1.0 + }; + let sequence_output = ((decoder_output + .narrow(1, decoder_output.dim(1)? - 1, 1)? + .squeeze(1)?) + * scaling_factor)?; + let output = { + let _enter = self.span_decode_head.enter(); + match self.lm_head { + None => sequence_output.matmul(&self.shared.embeddings().t()?)?, + Some(ref lm_head) => lm_head.forward(&sequence_output)?, + } + }; + Ok(output) + } + + pub fn forward(&mut self, input_ids: &Tensor, decoder_input_ids: &Tensor) -> Result { + let encoder_output = self.encode(input_ids)?; + self.decode(decoder_input_ids, &encoder_output) + } + + pub fn device(&self) -> &Device { + &self.device + } + + pub fn clear_kv_cache(&mut self) { + self.encoder.clear_kv_cache(); + self.decoder.clear_kv_cache(); + } +} diff --git a/patches/candle-transformers/src/models/trocr.rs b/patches/candle-transformers/src/models/trocr.rs new file mode 100644 index 0000000000..88418dd3ca --- /dev/null +++ b/patches/candle-transformers/src/models/trocr.rs @@ -0,0 +1,517 @@ +//! TrOCR model implementation. +//! +//! TrOCR is a Transformer-based OCR model that uses a Vision Transformer encoder +//! and a BART-like decoder for optical character recognition. +//! +//! Key characteristics: +//! - Vision Transformer encoder for image processing +//! - BART-style decoder for text generation +//! - Learned positional embeddings +//! - Layer normalization and self-attention +//! +//! References: +//! - [Paper](https://arxiv.org/abs/2109.10282) +//! - [Model Card](https://huggingface.co/microsoft/trocr-base-handwritten) +//! + +use crate::models::vit::{Config, Embeddings, Encoder}; +use candle::{DType, Result, Tensor}; +use candle_nn::{ + embedding, layer_norm, linear_no_bias, Embedding, LayerNorm, Linear, Module, VarBuilder, +}; + +fn default_tie_word_embeddings() -> bool { + true +} +fn default_use_learned_position_embeddings() -> bool { + true +} + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct TrOCRConfig { + pub vocab_size: usize, + pub d_model: usize, + pub cross_attention_hidden_size: usize, + pub decoder_layers: usize, + pub decoder_attention_heads: usize, + pub decoder_ffn_dim: usize, + pub activation_function: candle_nn::Activation, + pub max_position_embeddings: usize, + pub dropout: f64, + pub attention_dropout: f64, + pub activation_dropout: f64, + pub decoder_start_token_id: u32, + pub init_std: f64, + pub decoder_layerdrop: f64, + pub use_cache: bool, + pub scale_embedding: bool, + pub pad_token_id: usize, + pub bos_token_id: usize, + pub eos_token_id: u32, + pub decoder_vocab_size: Option, + #[serde(default = "default_use_learned_position_embeddings")] + pub use_learned_position_embeddings: bool, + #[serde(default = "default_tie_word_embeddings")] + pub tie_word_embeddings: bool, +} + +impl Default for TrOCRConfig { + fn default() -> Self { + Self { + vocab_size: 50265, + d_model: 1024, + cross_attention_hidden_size: 768, + decoder_layers: 12, + decoder_attention_heads: 16, + decoder_ffn_dim: 4096, + activation_function: candle_nn::Activation::Gelu, + max_position_embeddings: 512, + dropout: 0.1, + attention_dropout: 0.0, + activation_dropout: 0.0, + decoder_start_token_id: 2, + init_std: 0.02, + decoder_layerdrop: 0.0, + use_cache: true, + scale_embedding: false, + pad_token_id: 1, + bos_token_id: 0, + eos_token_id: 2, + decoder_vocab_size: Some(50265), + use_learned_position_embeddings: true, + tie_word_embeddings: true, + } + } +} + +#[derive(Debug, Clone)] +struct TrOCRLearnedPositionalEmbedding { + offset: usize, + weights: Embedding, +} + +impl TrOCRLearnedPositionalEmbedding { + fn load(vb: VarBuilder, cfg: &TrOCRConfig) -> Result { + let offset: usize = 2; + let num_embeddings = cfg.max_position_embeddings; + let embedding_dim = cfg.d_model; + let weights = embedding(num_embeddings + offset, embedding_dim, vb)?; + + Ok(Self { offset, weights }) + } + + fn new_sinusoidal(vb: VarBuilder, cfg: &TrOCRConfig) -> Result { + // https://github.com/huggingface/transformers/blob/58e3d23e97078f361a533b9ec4a6a2de674ea52a/src/transformers/models/trocr/modeling_trocr.py#L81 + let embedding_dim = cfg.d_model; + let half_dim = embedding_dim / 2; + let num_positions = cfg.max_position_embeddings + cfg.pad_token_id + 1; + let dev = vb.device(); + let inv_freq: Vec<_> = (0..half_dim) + .map(|i| 1f32 / 10000f32.powf(i as f32 / (half_dim - 1) as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?; + let t = Tensor::arange(0u32, num_positions as u32, dev)? + .to_dtype(DType::F32)? + .reshape((num_positions, 1))?; + let freqs = t.matmul(&inv_freq)?; + let emb = Tensor::cat(&[freqs.sin()?, freqs.cos()?], 1)?; + let emb = Tensor::cat( + &[ + emb.narrow(0, 0, cfg.pad_token_id)?, + Tensor::zeros((1, embedding_dim), DType::F32, dev)?, + emb.narrow(0, cfg.pad_token_id + 1, cfg.max_position_embeddings)?, + ], + 0, + )? + .contiguous()?; + let emb = Embedding::new(emb, embedding_dim); + Ok(Self { + offset: cfg.pad_token_id + 1, + weights: emb, + }) + } + + fn forward(&mut self, input_ids: &Tensor, past_key_values_length: u32) -> Result { + let (b_sz, seq_len) = input_ids.dims2()?; + + let positions = Tensor::arange( + past_key_values_length, + seq_len as u32 + past_key_values_length, + input_ids.device(), + )? + .expand((b_sz, seq_len))?; + + let positions = + positions.broadcast_add(&Tensor::new(self.offset as u32, input_ids.device())?)?; + self.weights.forward(&positions) + } +} + +#[derive(Debug, Clone)] +struct TrOCRAttention { + head_dim: usize, + num_heads: usize, + is_decoder: bool, + scaling: f64, + k_proj: Linear, + v_proj: Linear, + q_proj: Linear, + out_proj: Linear, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl TrOCRAttention { + fn load( + vb: VarBuilder, + cfg: &TrOCRConfig, + kdim: Option, + vdim: Option, + ) -> Result { + let embed_dim = cfg.d_model; + let num_heads = cfg.decoder_attention_heads; + let head_dim = embed_dim / num_heads; + let kdim = kdim.unwrap_or(embed_dim); + let vdim = vdim.unwrap_or(embed_dim); + + let k_proj = linear_no_bias(kdim, embed_dim, vb.pp("k_proj"))?; + let v_proj = linear_no_bias(vdim, embed_dim, vb.pp("v_proj"))?; + let q_proj = linear_no_bias(embed_dim, embed_dim, vb.pp("q_proj"))?; + + let out_proj = linear_no_bias(embed_dim, embed_dim, vb.pp("out_proj"))?; + Ok(Self { + head_dim, + num_heads, + is_decoder: true, + scaling: 1. / (head_dim as f64).sqrt(), + k_proj, + v_proj, + q_proj, + out_proj, + kv_cache: None, + }) + } + + fn reset_kv_cache(&mut self) { + self.kv_cache = None + } + + fn _shape(&self, tensor: &Tensor, bsz: usize) -> Result { + tensor + .reshape((bsz, (), self.num_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous() + } + + fn forward( + &mut self, + xs: &Tensor, + kv_states: Option<&Tensor>, + attn_mask: Option<&Tensor>, + ) -> Result { + let (b_sz, tgt_len, _) = xs.dims3()?; + let query_states = (xs.apply(&self.q_proj)? * self.scaling)?; + let (key_states, value_states) = match kv_states { + None => { + let key_states = self._shape(&xs.apply(&self.k_proj)?, b_sz)?; + let value_states = self._shape(&xs.apply(&self.v_proj)?, b_sz)?; + if self.is_decoder { + let kv_states = match &self.kv_cache { + None => (key_states, value_states), + Some((p_key_states, p_value_states)) => { + let key_states = Tensor::cat(&[p_key_states, &key_states], 2)?; + let value_states = Tensor::cat(&[p_value_states, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some(kv_states.clone()); + kv_states + } else { + (key_states, value_states) + } + } + Some(kv_states) => { + let key_states = self._shape(&kv_states.apply(&self.k_proj)?, b_sz)?; + let value_states = self._shape(&kv_states.apply(&self.v_proj)?, b_sz)?; + (key_states, value_states) + } + }; + let proj_shape = (b_sz * self.num_heads, (), self.head_dim); + let query_states = self._shape(&query_states, b_sz)?.reshape(proj_shape)?; + let key_states = key_states.reshape(proj_shape)?; + let value_states = value_states.reshape(proj_shape)?; + let attn_weights = query_states.matmul(&key_states.transpose(1, 2)?)?; + let attn_weights = match attn_mask { + None => attn_weights, + Some(attn_mask) => attn_weights.broadcast_add(attn_mask)?, + }; + let attn_probs = candle_nn::ops::softmax_last_dim(&attn_weights)?; + let attn_output = attn_probs.matmul(&value_states)?; + attn_output + .reshape((b_sz, self.num_heads, tgt_len, self.head_dim))? + .transpose(1, 2)? + .reshape((b_sz, tgt_len, self.head_dim * self.num_heads))? + .apply(&self.out_proj) + } +} + +#[derive(Debug, Clone)] +struct TrOCRDecoderLayer { + self_attn: TrOCRAttention, + activation_fn: candle_nn::Activation, + self_attn_layer_norm: LayerNorm, + encoder_attn: TrOCRAttention, + encoder_attn_layer_norm: LayerNorm, + fc1: Linear, + fc2: Linear, + final_layer_norm: LayerNorm, +} + +impl TrOCRDecoderLayer { + fn load(vb: VarBuilder, cfg: &TrOCRConfig) -> Result { + let embed_dim = cfg.d_model; + let self_attn = TrOCRAttention::load(vb.pp("self_attn"), cfg, None, None)?; + let self_attn_layer_norm = layer_norm(embed_dim, 1e-5, vb.pp("self_attn_layer_norm"))?; + let encoder_attn = TrOCRAttention::load( + vb.pp("encoder_attn"), + cfg, + Some(cfg.cross_attention_hidden_size), + Some(cfg.cross_attention_hidden_size), + )?; + let encoder_attn_layer_norm = + layer_norm(embed_dim, 1e-5, vb.pp("encoder_attn_layer_norm"))?; + let fc1 = linear_no_bias(embed_dim, cfg.decoder_ffn_dim, vb.pp("fc1"))?; + let fc2 = linear_no_bias(cfg.decoder_ffn_dim, embed_dim, vb.pp("fc2"))?; + let final_layer_norm = layer_norm(embed_dim, 1e-5, vb.pp("final_layer_norm"))?; + Ok(Self { + self_attn, + activation_fn: cfg.activation_function, + self_attn_layer_norm, + encoder_attn, + encoder_attn_layer_norm, + fc1, + fc2, + final_layer_norm, + }) + } + + fn reset_kv_cache(&mut self) { + self.self_attn.reset_kv_cache(); + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: &Tensor, + encoder_hidden_states: Option<&Tensor>, + ) -> Result { + let residual = xs.clone(); + let xs = self.self_attn.forward(xs, None, Some(attention_mask))?; + let xs = (xs + residual)?; + let mut xs = self.self_attn_layer_norm.forward(&xs)?; + + if let Some(encoder_hidden_states) = &encoder_hidden_states { + let residual = xs.clone(); + let encoder_attention_mask = attention_mask.clone(); // TODO + xs = self.encoder_attn.forward( + &xs, + Some(encoder_hidden_states), + Some(&encoder_attention_mask), + )?; + xs = (xs + residual)?; + xs = self.encoder_attn_layer_norm.forward(&xs)? + } + + let residual = xs.clone(); + let xs = self.fc1.forward(&xs)?; + let xs = self.activation_fn.forward(&xs)?; + let xs = self.fc2.forward(&xs)?; + let xs = (xs + residual)?; + let xs = self.final_layer_norm.forward(&xs)?; + + Ok(xs) + } +} + +#[derive(Debug, Clone)] +pub struct TrOCRDecoder { + layers: Vec, + embed_scale: Option, + embed_tokens: Embedding, + embed_positions: TrOCRLearnedPositionalEmbedding, +} + +impl TrOCRDecoder { + fn new(cfg: &TrOCRConfig, vb: VarBuilder) -> Result { + let vb = vb.pp("decoder.model.decoder"); + + let embed_tokens = embedding(cfg.vocab_size, cfg.d_model, vb.pp("embed_tokens"))?; + let embed_positions = if cfg.use_learned_position_embeddings { + TrOCRLearnedPositionalEmbedding::load(vb.pp("embed_positions"), cfg)? + } else { + TrOCRLearnedPositionalEmbedding::new_sinusoidal(vb.pp("embed_positions"), cfg)? + }; + let mut layers = Vec::with_capacity(cfg.decoder_layers); + let vb_l = vb.pp("layers"); + for idx in 0..cfg.decoder_layers { + let layer = TrOCRDecoderLayer::load(vb_l.pp(idx), cfg)?; + layers.push(layer) + } + let embed_scale = if cfg.scale_embedding { + Some((cfg.d_model as f64).sqrt()) + } else { + None + }; + + Ok(Self { + layers, + embed_scale, + embed_tokens, + embed_positions, + }) + } + + fn reset_kv_cache(&mut self) { + self.layers.iter_mut().for_each(|l| l.reset_kv_cache()) + } + + pub fn forward( + &mut self, + xs: &Tensor, + encoder_xs: Option<&Tensor>, + past_kv_len: usize, + attn_mask: &Tensor, + ) -> Result { + let embed_pos = self.embed_positions.forward(xs, past_kv_len as u32)?; + let xs = xs.apply(&self.embed_tokens)?; + + let xs = match self.embed_scale { + None => xs, + Some(scale) => (xs * scale)?, + }; + + let mut xs = xs.broadcast_add(&embed_pos)?; + + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attn_mask, encoder_xs)?; + } + Ok(xs) + } +} + +#[derive(Debug, Clone)] +pub struct TrOCREncoder { + embeddings: Embeddings, + encoder: Encoder, + layernorm: LayerNorm, +} + +impl TrOCREncoder { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_v = vb.pp("encoder"); + + let embeddings = Embeddings::new(cfg, false, vb_v.pp("embeddings"))?; + + let encoder = Encoder::new(cfg, vb_v.pp("encoder"))?; + let layernorm = layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb_v.pp("layernorm"))?; + + Ok(Self { + embeddings, + encoder, + layernorm, + }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let embedding_output = self.embeddings.forward(xs, None, false)?; + let encoder_outputs = self.encoder.forward(&embedding_output)?; + + self.layernorm.forward(&encoder_outputs) + } +} + +#[derive(Debug, Clone)] +pub struct TrOCRForCausalLM { + decoder: TrOCRDecoder, + output_projection: Linear, +} + +impl TrOCRForCausalLM { + pub fn new(decoder_cfg: &TrOCRConfig, vb: VarBuilder) -> Result { + let decoder = TrOCRDecoder::new(decoder_cfg, vb.clone())?; + let output_projection = if decoder_cfg.tie_word_embeddings { + candle_nn::Linear::new(decoder.embed_tokens.embeddings().clone(), None) + } else { + candle_nn::linear_no_bias( + decoder_cfg.d_model, + decoder_cfg.vocab_size, + vb.pp("decoder.output_projection"), + )? + }; + Ok(Self { + decoder, + output_projection, + }) + } + + pub fn forward( + &mut self, + xs: &Tensor, + encoder_xs: Option<&Tensor>, + past_kv_len: usize, + attn_mask: &Tensor, + ) -> Result { + let xs = self + .decoder + .forward(xs, encoder_xs, past_kv_len, attn_mask)?; + let xs = xs.apply(&self.output_projection)?; + + Ok(xs) + } + + fn reset_kv_cache(&mut self) { + self.decoder.reset_kv_cache(); + } +} + +#[derive(Debug, Clone)] +pub struct TrOCRModel { + encoder: TrOCREncoder, + decoder: TrOCRForCausalLM, +} + +impl TrOCRModel { + pub fn new(encoder_cfg: &Config, decoder_cfg: &TrOCRConfig, vb: VarBuilder) -> Result { + let encoder = TrOCREncoder::new(encoder_cfg, vb.clone())?; + let decoder = TrOCRForCausalLM::new(decoder_cfg, vb)?; + Ok(Self { encoder, decoder }) + } + + pub fn encoder(&mut self) -> &mut TrOCREncoder { + &mut self.encoder + } + + pub fn decoder(&mut self) -> &mut TrOCRForCausalLM { + &mut self.decoder + } + + pub fn decode( + &mut self, + xs: &Tensor, + encoder_xs: &Tensor, + past_kv_len: usize, + ) -> Result { + let seq_len = xs.dim(1)?; + let mask: Vec<_> = (0..seq_len) + .flat_map(|i| (0..seq_len).map(move |j| if j > i { f32::NEG_INFINITY } else { 0f32 })) + .collect(); + let mask = Tensor::from_vec(mask, (seq_len, seq_len), xs.device())?; + + self.decoder + .forward(xs, Some(encoder_xs), past_kv_len, &mask) + } + + pub fn reset_kv_cache(&mut self) { + self.decoder.reset_kv_cache(); + } +} diff --git a/patches/candle-transformers/src/models/vgg.rs b/patches/candle-transformers/src/models/vgg.rs new file mode 100644 index 0000000000..57f9ae67bb --- /dev/null +++ b/patches/candle-transformers/src/models/vgg.rs @@ -0,0 +1,267 @@ +//! VGG-16 model implementation. +//! +//! VGG-16 is a convolutional neural network architecture. It consists of 13 +//! convolutional layers followed by 3 fully connected layers. +//! +//! Key characteristics: +//! - Conv layers with 3x3 filters +//! - Max pooling after every 2-3 conv layers +//! - Three fully connected layers of 4096, 4096, 1000 units +//! - ReLU activation and dropout +//! +//! References: +//! - [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://arxiv.org/abs/1409.1556) +//! + +use candle::{ModuleT, Result, Tensor}; +use candle_nn::{FuncT, VarBuilder}; + +// Enum representing the different VGG models +pub enum Models { + Vgg13, + Vgg16, + Vgg19, +} + +// Struct representing a VGG model +#[derive(Debug)] +pub struct Vgg<'a> { + blocks: Vec>, +} + +// Struct representing the configuration for the pre-logit layer +struct PreLogitConfig { + in_dim: (usize, usize, usize, usize), + target_in: usize, + target_out: usize, +} + +// Implementation of the VGG model +impl<'a> Vgg<'a> { + // Function to create a new VGG model + pub fn new(vb: VarBuilder<'a>, model: Models) -> Result { + let blocks = match model { + Models::Vgg13 => vgg13_blocks(vb)?, + Models::Vgg16 => vgg16_blocks(vb)?, + Models::Vgg19 => vgg19_blocks(vb)?, + }; + Ok(Self { blocks }) + } +} + +// Implementation of the forward pass for the VGG model +impl ModuleT for Vgg<'_> { + fn forward_t(&self, xs: &Tensor, train: bool) -> Result { + let mut xs = xs.unsqueeze(0)?; + for block in self.blocks.iter() { + xs = xs.apply_t(block, train)?; + } + Ok(xs) + } +} + +// Function to create a conv2d block +// The block is composed of two conv2d layers followed by a max pool layer +fn conv2d_block(convs: &[(usize, usize, &str)], vb: &VarBuilder) -> Result> { + let layers = convs + .iter() + .map(|&(in_c, out_c, name)| { + candle_nn::conv2d( + in_c, + out_c, + 3, + candle_nn::Conv2dConfig { + stride: 1, + padding: 1, + ..Default::default() + }, + vb.pp(name), + ) + }) + .collect::>>()?; + + Ok(FuncT::new(move |xs, _train| { + let mut xs = xs.clone(); + for layer in layers.iter() { + xs = xs.apply(layer)?.relu()? + } + xs = xs.max_pool2d_with_stride(2, 2)?; + Ok(xs) + })) +} + +// Function to create a fully connected layer +// The layer is composed of two linear layers followed by a dropout layer +fn fully_connected( + num_classes: usize, + pre_logit_1: PreLogitConfig, + pre_logit_2: PreLogitConfig, + vb: VarBuilder, +) -> Result { + let lin = get_weights_and_biases( + &vb.pp("pre_logits.fc1"), + pre_logit_1.in_dim, + pre_logit_1.target_in, + pre_logit_1.target_out, + )?; + let lin2 = get_weights_and_biases( + &vb.pp("pre_logits.fc2"), + pre_logit_2.in_dim, + pre_logit_2.target_in, + pre_logit_2.target_out, + )?; + let dropout1 = candle_nn::Dropout::new(0.5); + let dropout2 = candle_nn::Dropout::new(0.5); + let dropout3 = candle_nn::Dropout::new(0.5); + Ok(FuncT::new(move |xs, train| { + let xs = xs.reshape((1, pre_logit_1.target_out))?; + let xs = xs.apply_t(&dropout1, train)?.apply(&lin)?.relu()?; + let xs = xs.apply_t(&dropout2, train)?.apply(&lin2)?.relu()?; + let lin3 = candle_nn::linear(4096, num_classes, vb.pp("head.fc"))?; + let xs = xs.apply_t(&dropout3, train)?.apply(&lin3)?.relu()?; + Ok(xs) + })) +} + +// Function to get the weights and biases for a layer +// This is required because the weights and biases are stored in different format than our linear layer expects +fn get_weights_and_biases( + vs: &VarBuilder, + in_dim: (usize, usize, usize, usize), + target_in: usize, + target_out: usize, +) -> Result { + let init_ws = candle_nn::init::DEFAULT_KAIMING_NORMAL; + let ws = vs.get_with_hints(in_dim, "weight", init_ws)?; + let ws = ws.reshape((target_in, target_out))?; + let bound = 1. / (target_out as f64).sqrt(); + let init_bs = candle_nn::Init::Uniform { + lo: -bound, + up: bound, + }; + let bs = vs.get_with_hints(target_in, "bias", init_bs)?; + Ok(candle_nn::Linear::new(ws, Some(bs))) +} + +fn vgg13_blocks(vb: VarBuilder) -> Result> { + let num_classes = 1000; + let blocks = vec![ + conv2d_block(&[(3, 64, "features.0"), (64, 64, "features.2")], &vb)?, + conv2d_block(&[(64, 128, "features.5"), (128, 128, "features.7")], &vb)?, + conv2d_block(&[(128, 256, "features.10"), (256, 256, "features.12")], &vb)?, + conv2d_block(&[(256, 512, "features.15"), (512, 512, "features.17")], &vb)?, + conv2d_block(&[(512, 512, "features.20"), (512, 512, "features.22")], &vb)?, + fully_connected( + num_classes, + PreLogitConfig { + in_dim: (4096, 512, 7, 7), + target_in: 4096, + target_out: 512 * 7 * 7, + }, + PreLogitConfig { + in_dim: (4096, 4096, 1, 1), + target_in: 4096, + target_out: 4096, + }, + vb.clone(), + )?, + ]; + Ok(blocks) +} + +fn vgg16_blocks(vb: VarBuilder) -> Result> { + let num_classes = 1000; + let blocks = vec![ + conv2d_block(&[(3, 64, "features.0"), (64, 64, "features.2")], &vb)?, + conv2d_block(&[(64, 128, "features.5"), (128, 128, "features.7")], &vb)?, + conv2d_block( + &[ + (128, 256, "features.10"), + (256, 256, "features.12"), + (256, 256, "features.14"), + ], + &vb, + )?, + conv2d_block( + &[ + (256, 512, "features.17"), + (512, 512, "features.19"), + (512, 512, "features.21"), + ], + &vb, + )?, + conv2d_block( + &[ + (512, 512, "features.24"), + (512, 512, "features.26"), + (512, 512, "features.28"), + ], + &vb, + )?, + fully_connected( + num_classes, + PreLogitConfig { + in_dim: (4096, 512, 7, 7), + target_in: 4096, + target_out: 512 * 7 * 7, + }, + PreLogitConfig { + in_dim: (4096, 4096, 1, 1), + target_in: 4096, + target_out: 4096, + }, + vb.clone(), + )?, + ]; + Ok(blocks) +} + +fn vgg19_blocks(vb: VarBuilder) -> Result> { + let num_classes = 1000; + let blocks = vec![ + conv2d_block(&[(3, 64, "features.0"), (64, 64, "features.2")], &vb)?, + conv2d_block(&[(64, 128, "features.5"), (128, 128, "features.7")], &vb)?, + conv2d_block( + &[ + (128, 256, "features.10"), + (256, 256, "features.12"), + (256, 256, "features.14"), + (256, 256, "features.16"), + ], + &vb, + )?, + conv2d_block( + &[ + (256, 512, "features.19"), + (512, 512, "features.21"), + (512, 512, "features.23"), + (512, 512, "features.25"), + ], + &vb, + )?, + conv2d_block( + &[ + (512, 512, "features.28"), + (512, 512, "features.30"), + (512, 512, "features.32"), + (512, 512, "features.34"), + ], + &vb, + )?, + fully_connected( + num_classes, + PreLogitConfig { + in_dim: (4096, 512, 7, 7), + target_in: 4096, + target_out: 512 * 7 * 7, + }, + PreLogitConfig { + in_dim: (4096, 4096, 1, 1), + target_in: 4096, + target_out: 4096, + }, + vb.clone(), + )?, + ]; + Ok(blocks) +} diff --git a/patches/candle-transformers/src/models/vit.rs b/patches/candle-transformers/src/models/vit.rs new file mode 100644 index 0000000000..49ab463017 --- /dev/null +++ b/patches/candle-transformers/src/models/vit.rs @@ -0,0 +1,416 @@ +//! Vision Transformer (ViT) implementation. +//! +//! Vision Transformer applies transformer architecture to image classification +//! by splitting images into patches and processing them as a sequence. +//! +//! Key characteristics: +//! - Image patches as sequence tokens +//! - Self-attention between patches +//! - Position embeddings +//! - CLS token for classification +//! - Layer normalization +//! +//! References: +//! - [ViT Paper](https://arxiv.org/abs/2010.11929) +//! - [Model Card](https://huggingface.co/google/vit-base-patch16-224) +//! + +use crate::models::with_tracing::{conv2d, linear, linear_no_bias, Conv2d, Linear}; +use candle::{IndexOp, Module, Result, Tensor, D}; +use candle_nn::{layer_norm, LayerNorm, VarBuilder}; + +// https://github.com/huggingface/transformers/blob/main/src/transformers/models/vit/configuration_vit.py +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + pub hidden_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub intermediate_size: usize, + pub hidden_act: candle_nn::Activation, + pub layer_norm_eps: f64, + pub image_size: usize, + pub patch_size: usize, + pub num_channels: usize, + pub qkv_bias: bool, +} + +impl Config { + // https://huggingface.co/google/vit-base-patch16-224/blob/main/config.json + pub fn vit_base_patch16_224() -> Self { + Self { + hidden_size: 768, + num_hidden_layers: 12, + num_attention_heads: 12, + intermediate_size: 3072, + hidden_act: candle_nn::Activation::Gelu, + layer_norm_eps: 1e-12, + image_size: 224, + patch_size: 16, + num_channels: 3, + qkv_bias: true, + } + } + + pub fn microsoft_trocr_base_handwritten() -> Self { + Self { + hidden_size: 768, + num_hidden_layers: 12, + num_attention_heads: 12, + intermediate_size: 3072, + hidden_act: candle_nn::Activation::Gelu, + layer_norm_eps: 1e-12, + image_size: 384, + patch_size: 16, + num_channels: 3, + qkv_bias: false, + } + } +} + +#[derive(Debug, Clone)] +struct PatchEmbeddings { + num_patches: usize, + projection: Conv2d, +} + +impl PatchEmbeddings { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let image_size = cfg.image_size; + let patch_size = cfg.patch_size; + let num_patches = (image_size / patch_size) * (image_size / patch_size); + let conv_cfg = candle_nn::Conv2dConfig { + stride: patch_size, + ..Default::default() + }; + let projection = conv2d( + cfg.num_channels, + cfg.hidden_size, + patch_size, + conv_cfg, + vb.pp("projection"), + )?; + Ok(Self { + num_patches, + projection, + }) + } +} + +impl Module for PatchEmbeddings { + fn forward(&self, pixel_values: &Tensor) -> Result { + let (_b_size, _num_channels, _height, _width) = pixel_values.dims4()?; + self.projection + .forward(pixel_values)? + .flatten_from(2)? + .transpose(1, 2) + } +} + +#[derive(Debug, Clone)] +pub struct Embeddings { + cls_token: Tensor, + mask_token: Option, + patch_embeddings: PatchEmbeddings, + position_embeddings: Tensor, + hidden_size: usize, +} + +impl Embeddings { + pub fn new(cfg: &Config, use_mask_token: bool, vb: VarBuilder) -> Result { + let hidden_size = cfg.hidden_size; + let cls_token = vb.get((1, 1, hidden_size), "cls_token")?; + let mask_token = if use_mask_token { + Some(vb.get((1, 1, hidden_size), "mask_token")?) + } else { + None + }; + let patch_embeddings = PatchEmbeddings::new(cfg, vb.pp("patch_embeddings"))?; + let num_patches = patch_embeddings.num_patches; + let position_embeddings = + vb.get((1, num_patches + 1, hidden_size), "position_embeddings")?; + Ok(Self { + cls_token, + mask_token, + patch_embeddings, + position_embeddings, + hidden_size, + }) + } + + fn interpolate_pos_encoding( + &self, + _embeddings: &Tensor, + _height: usize, + _width: usize, + ) -> Result { + todo!() + } + + pub fn forward( + &self, + pixel_values: &Tensor, + bool_masked_pos: Option<&Tensor>, + interpolate_pos_encoding: bool, + ) -> Result { + let (b_size, _num_channels, height, width) = pixel_values.dims4()?; + let embeddings = self.patch_embeddings.forward(pixel_values)?; + let embeddings = match (bool_masked_pos, &self.mask_token) { + (None, _) => embeddings, + (Some(_), None) => candle::bail!("bool_masked_pos set without mask_token"), + (Some(bool_masked_pos), Some(mask_tokens)) => { + let seq_len = embeddings.dim(1)?; + let mask_tokens = mask_tokens.broadcast_as((b_size, seq_len, self.hidden_size))?; + let mask = bool_masked_pos + .unsqueeze(D::Minus1)? + .to_dtype(mask_tokens.dtype())?; + ((mask_tokens * &mask)? - (embeddings * (mask - 1.)?)?)? + } + }; + let cls_tokens = self.cls_token.broadcast_as((b_size, 1, self.hidden_size))?; + let embeddings = Tensor::cat(&[&cls_tokens, &embeddings], 1)?; + if interpolate_pos_encoding { + let pos = self.interpolate_pos_encoding(&embeddings, height, width)?; + embeddings.broadcast_add(&pos) + } else { + embeddings.broadcast_add(&self.position_embeddings) + } + } +} + +#[derive(Debug, Clone)] +struct SelfAttention { + query: Linear, + key: Linear, + value: Linear, + num_attention_heads: usize, + attention_head_size: usize, +} + +impl SelfAttention { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let attention_head_size = cfg.hidden_size / cfg.num_attention_heads; + let num_attention_heads = cfg.num_attention_heads; + let all_head_size = num_attention_heads * attention_head_size; + let linear = |name| { + if cfg.qkv_bias { + linear(cfg.hidden_size, all_head_size, vb.pp(name)) + } else { + linear_no_bias(cfg.hidden_size, all_head_size, vb.pp(name)) + } + }; + let query = linear("query")?; + let key = linear("key")?; + let value = linear("value")?; + Ok(Self { + query, + key, + value, + num_attention_heads, + attention_head_size, + }) + } + + fn transpose_for_scores(&self, xs: &Tensor) -> Result { + let (b_size, seq_len, _) = xs.dims3()?; + xs.reshape(( + b_size, + seq_len, + self.num_attention_heads, + self.attention_head_size, + ))? + .permute((0, 2, 1, 3)) + } +} + +impl Module for SelfAttention { + fn forward(&self, xs: &Tensor) -> Result { + let query = self.query.forward(xs)?; + let key = self.key.forward(xs)?; + let value = self.value.forward(xs)?; + + let query = self.transpose_for_scores(&query)?.contiguous()?; + let key = self.transpose_for_scores(&key)?.contiguous()?; + let value = self.transpose_for_scores(&value)?.contiguous()?; + + let attention_scores = + (query.matmul(&key.t()?)? / f64::sqrt(self.attention_head_size as f64))?; + let attention_probs = candle_nn::ops::softmax_last_dim(&attention_scores)?; + attention_probs + .matmul(&value)? + .permute((0, 2, 1, 3))? + .contiguous()? + .flatten_from(D::Minus2) + } +} + +#[derive(Debug, Clone)] +struct SelfOutput { + dense: Linear, +} + +impl SelfOutput { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.hidden_size, cfg.hidden_size, vb.pp("dense"))?; + Ok(Self { dense }) + } +} + +impl Module for SelfOutput { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.dense) + } +} + +#[derive(Debug, Clone)] +struct Attention { + attention: SelfAttention, + output: SelfOutput, +} + +impl Attention { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let attention = SelfAttention::new(cfg, vb.pp("attention"))?; + let output = SelfOutput::new(cfg, vb.pp("output"))?; + Ok(Self { attention, output }) + } +} + +impl Module for Attention { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.attention)?.apply(&self.output) + } +} + +#[derive(Debug, Clone)] +struct Intermediate { + dense: Linear, + intermediate_act_fn: candle_nn::Activation, +} + +impl Intermediate { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.hidden_size, cfg.intermediate_size, vb.pp("dense"))?; + Ok(Self { + dense, + intermediate_act_fn: cfg.hidden_act, + }) + } +} + +impl Module for Intermediate { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.dense)?.apply(&self.intermediate_act_fn) + } +} + +#[derive(Debug, Clone)] +struct Output { + dense: Linear, +} + +impl Output { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.intermediate_size, cfg.hidden_size, vb.pp("dense"))?; + Ok(Self { dense }) + } + + fn forward(&self, xs: &Tensor, input_tensor: &Tensor) -> Result { + xs.apply(&self.dense)? + input_tensor + } +} + +#[derive(Debug, Clone)] +struct Layer { + attention: Attention, + intermediate: Intermediate, + output: Output, + layernorm_before: LayerNorm, + layernorm_after: LayerNorm, +} + +impl Layer { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let attention = Attention::new(cfg, vb.pp("attention"))?; + let intermediate = Intermediate::new(cfg, vb.pp("intermediate"))?; + let output = Output::new(cfg, vb.pp("output"))?; + let h_sz = cfg.hidden_size; + let layernorm_before = layer_norm(h_sz, cfg.layer_norm_eps, vb.pp("layernorm_before"))?; + let layernorm_after = layer_norm(h_sz, cfg.layer_norm_eps, vb.pp("layernorm_after"))?; + Ok(Self { + attention, + intermediate, + output, + layernorm_after, + layernorm_before, + }) + } +} + +impl Module for Layer { + fn forward(&self, xs: &Tensor) -> Result { + let xs = (xs.apply(&self.layernorm_before)?.apply(&self.attention)? + xs)?; + let ys = xs.apply(&self.layernorm_after)?.apply(&self.intermediate)?; + self.output.forward(&ys, &xs) + } +} + +#[derive(Debug, Clone)] +pub struct Encoder { + layers: Vec, +} + +impl Encoder { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb = vb.pp("layer"); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + for i in 0..cfg.num_hidden_layers { + let layer = Layer::new(cfg, vb.pp(i))?; + layers.push(layer) + } + Ok(Self { layers }) + } +} + +impl Module for Encoder { + fn forward(&self, xs: &Tensor) -> Result { + let mut xs = xs.clone(); + for layer in self.layers.iter() { + xs = xs.apply(layer)? + } + Ok(xs) + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embeddings: Embeddings, + encoder: Encoder, + layernorm: LayerNorm, + // no need for pooling layer for image classification + classifier: Linear, +} + +impl Model { + pub fn new(cfg: &Config, num_labels: usize, vb: VarBuilder) -> Result { + let vb_v = vb.pp("vit"); + let embeddings = Embeddings::new(cfg, false, vb_v.pp("embeddings"))?; + let encoder = Encoder::new(cfg, vb_v.pp("encoder"))?; + let layernorm = layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb_v.pp("layernorm"))?; + let classifier = linear(cfg.hidden_size, num_labels, vb.pp("classifier"))?; + Ok(Self { + embeddings, + encoder, + layernorm, + classifier, + }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let embedding_output = self.embeddings.forward(xs, None, false)?; + let encoder_outputs = self.encoder.forward(&embedding_output)?; + encoder_outputs + .i((.., 0, ..))? + .apply(&self.layernorm)? + .apply(&self.classifier) + } +} diff --git a/patches/candle-transformers/src/models/voxtral/audio.rs b/patches/candle-transformers/src/models/voxtral/audio.rs new file mode 100644 index 0000000000..8d577a2b8f --- /dev/null +++ b/patches/candle-transformers/src/models/voxtral/audio.rs @@ -0,0 +1,67 @@ +use candle::{DType, Device, Error, Tensor}; + +use crate::models::whisper::audio::{log_mel_spectrogram_, Float}; + +pub fn pcm_to_mel(samples: &[T], filters: &[T]) -> Vec { + log_mel_spectrogram_( + samples, + filters, + super::N_FFT, + super::HOP_LENGTH, + super::N_MELS, + false, + ) +} + +/// Process audio using exact WhisperFeatureExtractor algorithm then apply VoxtralProcessor chunking +pub fn extract_features(audio: &[f32], filters: &[f32], device: &Device) -> Result { + const N_MELS: usize = super::N_MELS; + + // Use the exact WhisperFeatureExtractor algorithm + // Use the whisper implementation from the parent module + let mel_vec = pcm_to_mel(audio, filters); + + // The whisper implementation returns Vec in shape (n_mel * n_len) + // We need to reshape it to match the expected tensor format + let n_mel = super::N_MELS; + let n_len = mel_vec.len() / n_mel; + + // Create tensor with shape (n_mel, n_len) then add batch dimension + let mel_tensor = Tensor::from_vec(mel_vec, (n_mel, n_len), device)?; + let mel_tensor = mel_tensor.unsqueeze(0)?; // Add batch dimension -> (1, n_mel, n_len) + + // Convert tensor back to Vec for compatibility with existing code + let mel = mel_tensor.flatten_all()?.to_vec1::()?; + let mel_len = mel.len(); + + // Apply VoxtralProcessor chunking exactly like Python + let total_frames = mel_len / N_MELS; + let max_source_positions = 3000; // From VoxtralProcessor defaults + + // Python approach: reshape (feature_size, total_frames) -> (feature_size, -1, max_source_positions) + // First, create mel tensor with shape (N_MELS, total_frames) + let mel_tensor = Tensor::from_vec(mel, (N_MELS, total_frames), device) + .map_err(|e| Error::Msg(format!("Failed to create mel tensor: {e}")))?; + + // Calculate number of chunks (equivalent to Python's -1 dimension in reshape) + let num_chunks = total_frames.div_ceil(max_source_positions); + + // Pad the mel tensor to be divisible by max_source_positions + let padded_frames = num_chunks * max_source_positions; + let padding_needed = padded_frames - total_frames; + + let mel_padded = if padding_needed > 0 { + let padding = Tensor::zeros((N_MELS, padding_needed), DType::F32, device)?; + Tensor::cat(&[&mel_tensor, &padding], 1)? + } else { + mel_tensor + }; + + // Reshape to (N_MELS, num_chunks, max_source_positions) + let reshaped = mel_padded.reshape((N_MELS, num_chunks, max_source_positions))?; + + // Transpose to (num_chunks, N_MELS, max_source_positions) - matching Python's transpose(0,1) + let audio_features = reshaped.transpose(0, 1)?; + + Ok(audio_features) +} diff --git a/patches/candle-transformers/src/models/voxtral/mod.rs b/patches/candle-transformers/src/models/voxtral/mod.rs new file mode 100644 index 0000000000..e2e747511b --- /dev/null +++ b/patches/candle-transformers/src/models/voxtral/mod.rs @@ -0,0 +1,14 @@ +pub mod audio; +pub mod model; +pub mod voxtral_llama; + +pub use audio::extract_features; +pub use model::{ + VoxtralCache, VoxtralConfig, VoxtralEncoder, VoxtralEncoderConfig, + VoxtralForConditionalGeneration, VoxtralGenerationConfig, VoxtralMultiModalProjector, +}; +pub use voxtral_llama::{VoxtralLlama, VoxtralLlamaCache, VoxtralLlamaConfig}; + +pub const N_FFT: usize = 400; +pub const HOP_LENGTH: usize = 160; +pub const N_MELS: usize = 128; diff --git a/patches/candle-transformers/src/models/voxtral/model.rs b/patches/candle-transformers/src/models/voxtral/model.rs new file mode 100644 index 0000000000..09535f7804 --- /dev/null +++ b/patches/candle-transformers/src/models/voxtral/model.rs @@ -0,0 +1,1074 @@ +use super::voxtral_llama::{VoxtralLlama, VoxtralLlamaCache, VoxtralLlamaConfig}; +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{ + layer_norm, linear, linear_no_bias, Conv1d, Dropout, LayerNorm, Linear, VarBuilder, +}; +use rand::Rng; + +#[derive(Debug, Clone)] +pub struct VoxtralEncoderConfig { + pub vocab_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub head_dim: usize, + pub scale_embedding: bool, + pub activation_function: String, + pub num_mel_bins: usize, + pub max_source_positions: usize, + pub initializer_range: f64, + pub attention_dropout: f64, + // These are set to 0.0 for compatibility with Whisper modular architecture + pub dropout: f64, + pub layerdrop: f64, + pub activation_dropout: f64, +} + +#[derive(Debug, Clone)] +pub struct VoxtralConfig { + pub audio_config: VoxtralEncoderConfig, + pub text_config: VoxtralLlamaConfig, + pub audio_token_id: usize, + pub projector_hidden_act: String, +} + +impl Default for VoxtralConfig { + fn default() -> Self { + Self { + audio_config: VoxtralEncoderConfig::default(), + text_config: VoxtralLlamaConfig::voxtral_3b(), + audio_token_id: 24, + projector_hidden_act: "gelu".to_string(), + } + } +} + +impl Default for VoxtralEncoderConfig { + fn default() -> Self { + Self { + vocab_size: 51866, + hidden_size: 1280, + intermediate_size: 5120, + num_hidden_layers: 32, + num_attention_heads: 20, + num_key_value_heads: 20, + head_dim: 64, + scale_embedding: false, + activation_function: "gelu".to_string(), + num_mel_bins: 128, + max_source_positions: 1500, + initializer_range: 0.02, + attention_dropout: 0.0, + // Set for Whisper compatibility + dropout: 0.0, + layerdrop: 0.0, + activation_dropout: 0.0, + } + } +} + +impl VoxtralEncoderConfig { + /// Ensures dropout values are properly set for Whisper compatibility + pub fn with_whisper_compatibility(mut self) -> Self { + self.dropout = 0.0; + self.layerdrop = 0.0; + self.activation_dropout = 0.0; + self + } +} + +/// Custom cache for multimodal inputs +#[derive(Debug, Clone)] +pub struct VoxtralCache { + cache: VoxtralLlamaCache, + audio_processed: bool, + cached_audio_embeds: Option, + cached_audio_positions: Option>, +} + +#[derive(Debug, Clone)] +pub struct VoxtralGenerationConfig { + pub max_new_tokens: usize, + pub temperature: f64, + pub top_p: Option, + pub device: Device, + /// If cache is None, the model will create a new cache. + pub cache: Option, +} + +impl VoxtralGenerationConfig { + pub fn new(device: Device) -> Self { + Self { + max_new_tokens: 500, + temperature: 0.0, + top_p: None, + device, + cache: None, + } + } +} + +impl VoxtralCache { + pub fn new( + use_kv_cache: bool, + dtype: DType, + config: &VoxtralLlamaConfig, + device: &Device, + ) -> Result { + Ok(Self { + cache: VoxtralLlamaCache::new(use_kv_cache, dtype, config, device)?, + audio_processed: false, + cached_audio_embeds: None, + cached_audio_positions: None, + }) + } + + pub fn reset(&mut self) { + // Reset the audio cache state + self.audio_processed = false; + self.cached_audio_embeds = None; + self.cached_audio_positions = None; + // Note: LlamaCache reset needs to be handled at a higher level + // as it requires device access + } +} + +/// Safely clamp tensor values for different dtypes +fn safe_clamp(x: &Tensor) -> Result { + match x.dtype() { + DType::F16 => { + // Match PyTorch exactly: torch.finfo(torch.float16).max - 1000 = 64504.0 + let max_val = 64504.0; + x.clamp(-max_val, max_val) + } + DType::BF16 => { + // BF16 has larger range, typically doesn't need clamping + Ok(x.clone()) + } + _ => Ok(x.clone()), + } +} + +/// Replace audio tokens in embeddings with projected audio features +pub fn replace_audio_tokens( + inputs_embeds: &Tensor, + audio_embeds: &Tensor, + audio_positions: &[(usize, usize)], + device: &Device, +) -> Result { + if audio_positions.is_empty() { + return Ok(inputs_embeds.clone()); + } + + let (batch_size, seq_len, hidden_size) = inputs_embeds.dims3()?; + let num_audio_tokens = audio_positions.len(); + + // HF-style: audio_embeds shape is (total_audio_seq_len, hidden_size) + let audio_embeds_dims = audio_embeds.dims2()?; + let total_audio_embeds = audio_embeds_dims.0; + + // HF-style: Use audio embeddings one-to-one with audio tokens + // We should now have the right number of audio tokens in the input sequence + let audio_embeds = if total_audio_embeds >= num_audio_tokens { + // Take the first num_audio_tokens embeddings to match the audio tokens + if num_audio_tokens == total_audio_embeds { + audio_embeds.clone() + } else { + audio_embeds.i(0..num_audio_tokens)? + } + } else { + candle::bail!( + "Not enough audio embeddings: need {}, got {}. Input sequence should have {} audio tokens.", + num_audio_tokens, + total_audio_embeds, + total_audio_embeds + ); + }; + + // Create result tensor starting with text embeddings + let mut result = inputs_embeds.clone(); + + // Replace audio tokens with audio embeddings + // Since we don't have scatter operations, we'll do this manually + for (idx, &(batch_idx, seq_idx)) in audio_positions.iter().enumerate() { + if batch_idx >= batch_size || seq_idx >= seq_len { + candle::bail!( + "Invalid audio position: ({}, {}) for tensor shape ({}, {}, {})", + batch_idx, + seq_idx, + batch_size, + seq_len, + hidden_size + ); + } + + // Get the audio embedding for this position + let audio_embed = audio_embeds.i(idx)?; + + // Create a mask for this specific position + let mut position_mask = vec![0f32; batch_size * seq_len]; + position_mask[batch_idx * seq_len + seq_idx] = 1.0; + let position_mask = Tensor::new(position_mask.as_slice(), device)? + .reshape((batch_size, seq_len, 1))? + .to_dtype(inputs_embeds.dtype())?; + + // Broadcast audio embedding to full tensor shape + let audio_embed_broadcast = audio_embed.unsqueeze(0)?.unsqueeze(0)?.broadcast_as(( + batch_size, + seq_len, + hidden_size, + ))?; + + // Update result: keep original where mask is 0, use audio where mask is 1 + let inverse_mask = (1.0 - &position_mask)?; + result = (result.broadcast_mul(&inverse_mask)? + + audio_embed_broadcast.broadcast_mul(&position_mask)?)?; + } + + Ok(result) +} + +/// Find positions of audio tokens in input sequences +pub fn find_audio_token_positions( + input_ids: &Tensor, + audio_token_id: usize, +) -> Result> { + // Handle both i64 and u32 token types by converting to i64 first if needed + let input_ids = if input_ids.dtype() == candle::DType::U32 { + input_ids.to_dtype(candle::DType::I64)? + } else { + input_ids.clone() + }; + + let input_ids = input_ids.to_vec2::()?; + let mut positions = Vec::new(); + + for (batch_idx, sequence) in input_ids.iter().enumerate() { + for (seq_idx, &token_id) in sequence.iter().enumerate() { + if token_id as usize == audio_token_id { + positions.push((batch_idx, seq_idx)); + } + } + } + + Ok(positions) +} + +#[derive(Debug, Clone)] +struct VoxtralAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + out_proj: Linear, + num_heads: usize, + head_dim: usize, + scaling: f64, + attention_dropout: Dropout, +} + +impl VoxtralAttention { + fn new(cfg: &VoxtralEncoderConfig, vb: VarBuilder) -> Result { + let embed_dim = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let head_dim = embed_dim / num_heads; + + if head_dim * num_heads != embed_dim { + candle::bail!( + "embed_dim must be divisible by num_heads ({} % {} != 0)", + embed_dim, + num_heads + ); + } + + let scaling = (head_dim as f64).powf(-0.5); + + let q_proj = linear(embed_dim, embed_dim, vb.pp("q_proj"))?; + let k_proj = linear_no_bias(embed_dim, embed_dim, vb.pp("k_proj"))?; + let v_proj = linear(embed_dim, embed_dim, vb.pp("v_proj"))?; + let out_proj = linear(embed_dim, embed_dim, vb.pp("out_proj"))?; + + let attention_dropout = Dropout::new(cfg.attention_dropout as f32); + + Ok(Self { + q_proj, + k_proj, + v_proj, + out_proj, + num_heads, + head_dim, + scaling, + attention_dropout, + }) + } + + fn reshape_for_scores(&self, x: &Tensor, seq_len: usize, bsz: usize) -> Result { + x.reshape((bsz, seq_len, self.num_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous() + } +} + +impl Module for VoxtralAttention { + fn forward(&self, x: &Tensor) -> Result { + let (bsz, seq_len, _) = x.dims3()?; + + // Project queries, keys, and values - apply scaling to queries to match PyTorch SDPA + let q = (self.q_proj.forward(x)? * self.scaling)?; + let k = self.k_proj.forward(x)?; + let v = self.v_proj.forward(x)?; + + // Reshape for multi-head attention: (batch, seq_len, num_heads, head_dim) -> (batch, num_heads, seq_len, head_dim) + let q = self.reshape_for_scores(&q, seq_len, bsz)?; + let k = self.reshape_for_scores(&k, seq_len, bsz)?; + let v = self.reshape_for_scores(&v, seq_len, bsz)?; + + // Manual SDPA-like implementation to match Python's numerical behavior exactly + // Use F16 precision throughout to match PyTorch's F16 model + let scores = q.matmul(&k.transpose(D::Minus2, D::Minus1)?)?; + + // Apply softmax in same precision as input (F16) to match Python + let attn_weights = candle_nn::ops::softmax_last_dim(&scores)?; + + // Apply attention dropout (disabled during inference) + let attn_weights = self.attention_dropout.forward(&attn_weights, false)?; + + // Apply attention to values + let attn_output = attn_weights.matmul(&v)?; + + // Reshape back to (batch, seq_len, embed_dim) + let attn_output = attn_output.transpose(1, 2)?.contiguous()?.reshape(( + bsz, + seq_len, + self.num_heads * self.head_dim, + ))?; + + self.out_proj.forward(&attn_output) + } +} + +#[derive(Debug, Clone)] +struct VoxtralEncoderLayer { + self_attn: VoxtralAttention, + self_attn_layer_norm: LayerNorm, + fc1: Linear, + fc2: Linear, + final_layer_norm: LayerNorm, + activation: candle_nn::Activation, + dropout: Dropout, + activation_dropout: Dropout, +} + +impl VoxtralEncoderLayer { + fn new(cfg: &VoxtralEncoderConfig, vb: VarBuilder) -> Result { + let embed_dim = cfg.hidden_size; + + let self_attn = VoxtralAttention::new(cfg, vb.pp("self_attn"))?; + let self_attn_layer_norm = layer_norm(embed_dim, 1e-5, vb.pp("self_attn_layer_norm"))?; + let fc1 = linear(embed_dim, cfg.intermediate_size, vb.pp("fc1"))?; + let fc2 = linear(cfg.intermediate_size, embed_dim, vb.pp("fc2"))?; + let final_layer_norm = layer_norm(embed_dim, 1e-5, vb.pp("final_layer_norm"))?; + + let activation = match cfg.activation_function.as_str() { + "gelu" => candle_nn::Activation::Gelu, + "relu" => candle_nn::Activation::Relu, + _ => candle::bail!( + "Unsupported activation function: {}", + cfg.activation_function + ), + }; + + let dropout = Dropout::new(cfg.dropout as f32); + let activation_dropout = Dropout::new(cfg.activation_dropout as f32); + + Ok(Self { + self_attn, + self_attn_layer_norm, + fc1, + fc2, + final_layer_norm, + activation, + dropout, + activation_dropout, + }) + } + + pub fn get_fc1_out_dim(&self) -> usize { + // Return the intermediate size from the config + // Since Linear doesn't expose out_dim + self.fc1.weight().dims()[0] + } + + fn forward(&self, x: &Tensor, training: bool) -> Result { + // Self-attention with residual connection + let residual = x; + let x = self.self_attn_layer_norm.forward(x)?; + let x = self.self_attn.forward(&x)?; + let x = self.dropout.forward(&x, training)?; + let x = (x + residual)?; + + // Feed-forward network with residual connection + let residual = &x; + let x = self.final_layer_norm.forward(&x)?; + let x = self.fc1.forward(&x)?; + let x = x.apply(&self.activation)?; + let x = self.activation_dropout.forward(&x, training)?; + let x = self.fc2.forward(&x)?; + let x = self.dropout.forward(&x, training)?; + let x = (x + residual)?; + + // Safe clamping for numerical stability + safe_clamp(&x) + } +} + +#[derive(Debug, Clone)] +pub struct VoxtralEncoder { + conv1: Conv1d, + conv2: Conv1d, + embed_positions: Tensor, + layers: Vec, + layer_norm: LayerNorm, + dropout: Dropout, + layerdrop: f64, +} + +impl VoxtralEncoder { + pub fn new(cfg: &VoxtralEncoderConfig, vb: VarBuilder) -> Result { + // Ensure Whisper compatibility + let cfg = cfg.clone().with_whisper_compatibility(); + + let embed_dim = cfg.hidden_size; + + // Convolutional layers for processing mel features + let conv1 = candle_nn::conv1d( + cfg.num_mel_bins, + embed_dim, + 3, + candle_nn::Conv1dConfig { + padding: 1, + ..Default::default() + }, + vb.pp("conv1"), + )?; + + let conv2 = candle_nn::conv1d( + embed_dim, + embed_dim, + 3, + candle_nn::Conv1dConfig { + stride: 2, + padding: 1, + ..Default::default() + }, + vb.pp("conv2"), + )?; + + // Position embeddings + let embed_positions = vb.get( + (cfg.max_source_positions, embed_dim), + "embed_positions.weight", + )?; + + // Transformer layers + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + for i in 0..cfg.num_hidden_layers { + layers.push(VoxtralEncoderLayer::new( + &cfg, + vb.pp(format!("layers.{i}")), + )?); + } + + let layer_norm = layer_norm(embed_dim, 1e-5, vb.pp("layer_norm"))?; + let dropout = Dropout::new(cfg.dropout as f32); + + Ok(Self { + conv1, + conv2, + embed_positions, + layers, + layer_norm, + dropout, + layerdrop: cfg.layerdrop, + }) + } + + pub fn forward(&self, input_features: &Tensor) -> Result { + self.forward_with_training(input_features, false) + } + + pub fn forward_with_training(&self, input_features: &Tensor, training: bool) -> Result { + // Keep conv layers in F16 to avoid shape issues + let expected_dtype = self.conv1.weight().dtype(); + let input_features = if input_features.dtype() != expected_dtype { + input_features.to_dtype(expected_dtype)? + } else { + input_features.clone() + }; + + // Apply convolutional layers with GELU activation + let x = if false { + // Keep conv layers in F16 + // Convert conv1 weights to F32 for computation + let conv1_weight_f32 = self.conv1.weight().to_dtype(DType::F32)?; + let conv1_bias_f32 = if let Some(bias) = self.conv1.bias() { + Some(bias.to_dtype(DType::F32)?) + } else { + None + }; + + // Manual conv1d operation with F32 precision - conv1 has stride=1, padding=1 + let mut conv_result = input_features.conv1d(&conv1_weight_f32, 1, 1, 1, 1)?; + if let Some(bias) = conv1_bias_f32 { + conv_result = conv_result.broadcast_add(&bias.unsqueeze(0)?.unsqueeze(2)?)?; + } + conv_result + } else { + self.conv1.forward(&input_features)? + }; + + // Apply GELU activation after conv1 (matches Python: conv1 -> GELU) + let x = x.gelu()?; + + // Apply conv2 (matches Python: conv2) + let x = if false { + // Keep conv layers in F16 + // Convert conv2 weights to F32 for computation + let conv2_weight_f32 = self.conv2.weight().to_dtype(DType::F32)?; + let conv2_bias_f32 = if let Some(bias) = self.conv2.bias() { + Some(bias.to_dtype(DType::F32)?) + } else { + None + }; + + // Manual conv1d operation with F32 precision - conv2 has stride=2, padding=1 + let mut conv_result = x.conv1d(&conv2_weight_f32, 2, 1, 1, 1)?; + if let Some(bias) = conv2_bias_f32 { + conv_result = conv_result.broadcast_add(&bias.unsqueeze(0)?.unsqueeze(2)?)?; + } + conv_result + } else { + self.conv2.forward(&x)? + }; + + // Apply GELU activation after conv2 (FIX: matches Python: conv2 -> GELU) + let x = x.gelu()?; + + // Reshape: (batch, embed_dim, seq_len) -> (batch, seq_len, embed_dim) + let x = x.transpose(1, 2)?; + + // Add position embeddings - handle F32 position embeddings + F16 hidden states like PyTorch + let seq_len = x.dim(1)?; + let positions = self.embed_positions.i(..seq_len)?; + + // PyTorch automatically promotes F16 + F32 -> F32, then converts back to original dtype + // We need to match this behavior exactly + let x = if false { + // Keep position embeddings in mixed precision + // Force F32 computation for position embeddings + let x_f32 = x.to_dtype(candle::DType::F32)?; + let positions_f32 = positions.to_dtype(candle::DType::F32)?; + x_f32.broadcast_add(&positions_f32)? // Keep result in F32 + } else if x.dtype() != positions.dtype() { + // Convert hidden states to F32 for addition (positions are already F32) + let x_f32 = x.to_dtype(candle::DType::F32)?; + let result_f32 = x_f32.broadcast_add(&positions)?; + // Convert back to original hidden states dtype (F16) + result_f32.to_dtype(x.dtype())? + } else { + x.broadcast_add(&positions)? + }; + + // Apply dropout + let mut x = self.dropout.forward(&x, training)?; + + for (idx, layer) in self.layers.iter().enumerate() { + // Keep all computation in F16 + x = self.forward_layer_with_dropout(&x, layer, idx, training)?; + } + + // Apply final layer normalization (critical for proper output values!) + let x = self.layer_norm.forward(&x)?; + + Ok(x) + } + + /// Forward a single layer with stochastic depth (layer dropout) + fn forward_layer_with_dropout( + &self, + x: &Tensor, + layer: &VoxtralEncoderLayer, + _layer_idx: usize, + training: bool, + ) -> Result { + if training && self.layerdrop > 0.0 { + // Apply stochastic depth with proper randomization + let mut rng = rand::rng(); + let keep_prob = 1.0 - self.layerdrop; + let keep: bool = rng.random::() < keep_prob; + + if !keep { + // Skip layer entirely (identity mapping) + return Ok(x.clone()); + } + } + + layer.forward(x, training) + } + + /// Get the output dimension of the first FC layer (needed for projector) + pub fn get_intermediate_size(&self) -> usize { + if !self.layers.is_empty() { + self.layers[0].get_fc1_out_dim() + } else { + // Fallback to config value + 5120 // Default intermediate size + } + } + + /// Process long audio sequences in chunks to save memory + pub fn process_long_audio( + &self, + input_features: &Tensor, + chunk_size: usize, + overlap: usize, + ) -> Result { + let (_batch_size, _num_mel, seq_len) = input_features.dims3()?; + + if seq_len <= chunk_size { + return self.forward(input_features); + } + + let mut outputs = Vec::new(); + let step = chunk_size - overlap; + + for start in (0..seq_len).step_by(step) { + let end = (start + chunk_size).min(seq_len); + let chunk = input_features.i((.., .., start..end))?; + + // Process chunk + let output = self.forward(&chunk)?; + + // Handle overlap by averaging + if !outputs.is_empty() && overlap > 0 { + let overlap_frames = overlap / 2; // Account for conv2 stride + let last_output: &mut Tensor = outputs.last_mut().unwrap(); + let last_len = last_output.dim(1)?; + + // Average overlapping regions + let overlap_start = last_len.saturating_sub(overlap_frames); + let overlap_new = output.i((.., ..overlap_frames, ..))?; + let overlap_old = last_output.i((.., overlap_start.., ..))?; + let averaged = ((overlap_old + overlap_new)? * 0.5)?; + + // Update last output + *last_output = + Tensor::cat(&[&last_output.i((.., ..overlap_start, ..))?, &averaged], 1)?; + + // Add non-overlapping part of current chunk + outputs.push(output.i((.., overlap_frames.., ..))?); + } else { + outputs.push(output); + } + } + + // Concatenate all outputs + let outputs_ref: Vec<&Tensor> = outputs.iter().collect(); + Tensor::cat(&outputs_ref, 1) + } +} + +#[derive(Debug, Clone)] +pub struct VoxtralMultiModalProjector { + linear_1: Linear, + linear_2: Linear, + activation: candle_nn::Activation, +} + +impl VoxtralMultiModalProjector { + pub fn new(cfg: &VoxtralConfig, vb: VarBuilder) -> Result { + let linear_1 = linear_no_bias( + cfg.audio_config.intermediate_size, + cfg.text_config.hidden_size, + vb.pp("linear_1"), + )?; + + let linear_2 = linear_no_bias( + cfg.text_config.hidden_size, + cfg.text_config.hidden_size, + vb.pp("linear_2"), + )?; + + let activation = match cfg.projector_hidden_act.as_str() { + "gelu" => candle_nn::Activation::Gelu, + "relu" => candle_nn::Activation::Relu, + _ => candle::bail!( + "Unsupported projector activation: {}", + cfg.projector_hidden_act + ), + }; + + Ok(Self { + linear_1, + linear_2, + activation, + }) + } + + pub fn forward(&self, audio_features: &Tensor) -> Result { + let x = self.linear_1.forward(audio_features)?; + let x = x.apply(&self.activation)?; + self.linear_2.forward(&x) + } +} + +#[derive(Debug, Clone)] +pub struct VoxtralForConditionalGeneration { + audio_tower: VoxtralEncoder, + language_model: VoxtralLlama, + multi_modal_projector: VoxtralMultiModalProjector, + audio_token_id: usize, + audio_config: VoxtralEncoderConfig, + text_config: VoxtralLlamaConfig, +} + +impl VoxtralForConditionalGeneration { + pub fn new(cfg: &VoxtralConfig, vb: VarBuilder) -> Result { + let audio_tower = VoxtralEncoder::new(&cfg.audio_config, vb.pp("audio_tower"))?; + let language_model = VoxtralLlama::load(vb.pp("language_model"), &cfg.text_config)?; + let multi_modal_projector = + VoxtralMultiModalProjector::new(cfg, vb.pp("multi_modal_projector"))?; + + Ok(Self { + audio_tower, + language_model, + multi_modal_projector, + audio_token_id: cfg.audio_token_id, + audio_config: cfg.audio_config.clone(), + text_config: cfg.text_config.clone(), + }) + } + + /// Get the audio token ID used for this model + pub fn audio_token_id(&self) -> usize { + self.audio_token_id + } + + /// Get the text model configuration + pub fn text_config(&self) -> &VoxtralLlamaConfig { + &self.text_config + } + + /// Get the audio encoder configuration + pub fn audio_config(&self) -> &VoxtralEncoderConfig { + &self.audio_config + } + + /// Process audio features through encoder and projector + pub fn get_audio_embeds(&self, input_features: &Tensor) -> Result { + let audio_outputs = self.audio_tower.forward(input_features)?; + + // Following HF implementation: reshape to (-1, config.intermediate_size) before projection + // Python: audio_hidden_states.reshape(-1, self.config.audio_config.intermediate_size) + // This transforms [1, 1500, 1280] -> [375, 5120] using intermediate_size from config + let (batch_size, seq_len, hidden_size) = audio_outputs.dims3()?; + + // The key insight: Python reshapes from [1, 1500, 1280] to [375, 5120] + // This means 1500 * 1280 = 375 * 5120 (1920000 elements) + // So we need: new_batch_size = (batch_size * seq_len * hidden_size) / intermediate_size + let total_elements = batch_size * seq_len * hidden_size; + let new_batch_size = total_elements / self.audio_config.intermediate_size; + + // Verify the division is exact + if total_elements % self.audio_config.intermediate_size != 0 { + return Err(candle::Error::DimOutOfRange { + shape: candle::Shape::from_dims(&[batch_size, seq_len, hidden_size]), + dim: 0, + op: "reshape", + }); + } + + let audio_hidden = + audio_outputs.reshape((new_batch_size, self.audio_config.intermediate_size))?; + + // Project to text space - this gives us embeddings for each audio position + let projected = self.multi_modal_projector.forward(&audio_hidden)?; + + // Return shape: (batch_size * seq_len, text_hidden_size) + // This matches HF implementation - no pooling, keep all audio token embeddings + Ok(projected) + } + + /// Process long audio sequences efficiently + pub fn get_audio_embeds_chunked( + &self, + input_features: &Tensor, + chunk_size: usize, + overlap: usize, + ) -> Result { + let audio_outputs = + self.audio_tower + .process_long_audio(input_features, chunk_size, overlap)?; + + // Reshape and project (now outputs hidden_size, needs reshape to intermediate_size) + let (batch_size, seq_len, hidden_size) = audio_outputs.dims3()?; + // Apply same reshape logic as get_audio_embeds + let total_elements = batch_size * seq_len * hidden_size; + let new_batch_size = total_elements / self.audio_config.intermediate_size; + let audio_hidden = + audio_outputs.reshape((new_batch_size, self.audio_config.intermediate_size))?; + + let projected = self.multi_modal_projector.forward(&audio_hidden)?; + + // Reshape back to (batch_size, seq_len, text_hidden_size) for pooling + let text_hidden_size = self.text_config.hidden_size; + let projected = projected.reshape((batch_size, seq_len, text_hidden_size))?; + + // Apply mean pooling to reduce to single audio embedding per batch + let pooled = projected.mean(1)?; // Mean across sequence dimension + + // Return shape: (batch_size, text_hidden_size) + Ok(pooled) + } + + /// Forward pass with audio features and text input + pub fn forward( + &self, + input_ids: &Tensor, + input_features: Option<&Tensor>, + cache: &mut VoxtralCache, + index_pos: usize, + ) -> Result { + // Get text embeddings + let mut inputs_embeds = self.language_model.embed(input_ids)?; + + // If audio features are provided and not yet processed + if let Some(features) = input_features { + if !cache.audio_processed { + let audio_embeds = self.get_audio_embeds(features)?; + + let audio_positions = find_audio_token_positions(input_ids, self.audio_token_id)?; + + // Cache for future use + cache.cached_audio_embeds = Some(audio_embeds.clone()); + cache.cached_audio_positions = Some(audio_positions.clone()); + cache.audio_processed = true; + + inputs_embeds = replace_audio_tokens( + &inputs_embeds, + &audio_embeds, + &audio_positions, + input_ids.device(), + )?; + } + } + + // Forward through language model using forward_input_embed + self.language_model + .forward_input_embed(&inputs_embeds, index_pos, &mut cache.cache) + } + + /// Generate text given audio input + pub fn generate( + &self, + input_ids: &Tensor, + input_features: Option<&Tensor>, + config: VoxtralGenerationConfig, + ) -> Result> { + // Validate inputs + if config.max_new_tokens == 0 { + return input_ids.i(0)?.to_vec1::(); // Get first batch + } + + if config.temperature < 0.0 { + candle::bail!( + "Temperature must be non-negative, got {}", + config.temperature + ); + } + + if let Some(p) = config.top_p { + if !(0.0..=1.0).contains(&p) { + candle::bail!("top_p must be between 0 and 1, got {}", p); + } + } + + let mut final_cache = if let Some(cache) = config.cache { + cache + } else { + // Get the dtype from the language model by creating a small embedding + let dummy_token = Tensor::new(&[1u32], &config.device)?; + let dummy_embed = self.language_model.embed(&dummy_token)?; + let model_dtype = dummy_embed.dtype(); + VoxtralCache::new(true, model_dtype, &self.text_config, &config.device)? + }; + let mut tokens = input_ids.i(0)?.to_vec1::()?; // Get first batch + let initial_len = tokens.len(); + + for idx in 0..config.max_new_tokens { + let (input, index_pos) = if idx == 0 { + (input_ids.clone(), 0) + } else { + // For subsequent generation steps, use only the last token + let last_token = tokens[tokens.len() - 1]; + let calculated_pos = initial_len + idx - 1; + ( + Tensor::new(&[last_token], &config.device)?.unsqueeze(0)?, + calculated_pos, + ) + }; + + let logits = if idx == 0 { + // First pass - include audio features + match self.forward(&input, input_features, &mut final_cache, index_pos) { + Ok(logits) => logits, + Err(e) => { + return Err(candle::Error::Msg(format!( + "Failed to generate tokens: {e}" + ))); + } + } + } else { + // Subsequent passes - text only + match self.forward(&input, None, &mut final_cache, index_pos) { + Ok(logits) => logits, + Err(e) => { + return Err(candle::Error::Msg(format!( + "Failed to generate tokens: {e}" + ))); + } + } + }; + + // Handle both 2D [batch, vocab] and 3D [batch, seq_len, vocab] logits + let logits = if logits.dims().len() == 3 { + // 3D case: [batch, seq_len, vocab] -> get last token + logits.i((.., logits.dim(1)? - 1, ..))? + } else { + // 2D case: [batch, vocab] -> already the right shape + logits + }; + + let next_token = if config.temperature > 0.0 { + // Sample with temperature + let prs = (logits / config.temperature)?; + let prs = candle_nn::ops::softmax_last_dim(&prs)?; + + if let Some(top_p_val) = config.top_p { + // Apply top-p sampling + sample_top_p(&prs.squeeze(0)?, top_p_val, &config.device)? + } else { + // Sample from full distribution + let probs_vec = prs.squeeze(0)?.to_vec1::()?; + let mut rng = rand::rng(); + let mut cumsum = 0.0; + let rand_val: f32 = rng.random(); + let mut sampled = 0u32; + + for (idx, &prob) in probs_vec.iter().enumerate() { + cumsum += prob; + if cumsum > rand_val { + sampled = idx as u32; + break; + } + } + sampled + } + } else { + // Greedy decoding - find the token with highest probability + let argmax_result = match logits.argmax(D::Minus1) { + Ok(result) => result, + Err(e) => { + return Err(candle::Error::Msg(format!("Argmax failed: {e}"))); + } + }; + + // Handle the case where argmax returns [1] instead of scalar + + if argmax_result.dims().is_empty() { + // Already a scalar + match argmax_result.to_scalar::() { + Ok(token) => token, + Err(e) => { + return Err(candle::Error::Msg(format!("to_scalar failed: {e}"))); + } + } + } else if argmax_result.dims() == [1] { + // Shape [1] - extract the single element + match argmax_result.i(0) { + Ok(scalar_tensor) => match scalar_tensor.to_scalar::() { + Ok(token) => token, + Err(e) => { + return Err(candle::Error::Msg(format!( + "to_scalar on extracted element failed: {e}" + ))); + } + }, + Err(e) => { + return Err(candle::Error::Msg(format!( + "indexing argmax result failed: {e}" + ))); + } + } + } else { + return Err(candle::Error::Msg(format!( + "Unexpected argmax result shape: {:?}", + argmax_result.shape() + ))); + } + }; + + tokens.push(next_token); + + // Check for EOS tokens - Voxtral uses different EOS tokens than hardcoded 2 + // Based on the Mistral/Voxtral tokenizer, common EOS tokens are: + // 2 = , 0 = , 128001, 128009 from various chat formats + let eos_tokens = [2u32, 128001, 128009, 128256]; // Don't include 0 as it might be valid generation + + // Check for EOS tokens only if not ignoring them + if eos_tokens.contains(&next_token) { + break; + } + + // Also break if we get repeated pad tokens (might indicate the model is stuck) + if next_token == 0 && tokens.len() > 5 { + let last_5_tokens = &tokens[tokens.len() - 5..]; + if last_5_tokens.iter().all(|&t| t == 0) { + break; + } + } + } + + Ok(tokens) + } +} + +/// Sample from top-p probability distribution +fn sample_top_p(probs: &Tensor, top_p: f64, _device: &Device) -> Result { + let (sorted_probs, sorted_indices) = probs.sort_last_dim(false)?; + let cumsum = sorted_probs.cumsum(D::Minus1)?; + let mask = cumsum.le(top_p)?; + + // Apply mask and renormalize + let filtered_probs = sorted_probs.where_cond(&mask, &Tensor::zeros_like(&sorted_probs)?)?; + let filtered_probs = (&filtered_probs / filtered_probs.sum_keepdim(D::Minus1)?)?; + + // Sample from filtered distribution + // Since multinomial is not available, we'll use a simple sampling approach + let probs_vec = filtered_probs.to_vec1::()?; + let mut cumsum = 0.0; + let mut rng = rand::rng(); + let rand_val: f32 = rng.random(); + let mut sample_idx = 0; + + for (idx, &prob) in probs_vec.iter().enumerate() { + cumsum += prob; + if cumsum > rand_val { + sample_idx = idx; + break; + } + } + + sorted_indices.i(sample_idx)?.to_scalar::() +} diff --git a/patches/candle-transformers/src/models/voxtral/voxtral_llama.rs b/patches/candle-transformers/src/models/voxtral/voxtral_llama.rs new file mode 100644 index 0000000000..bca9c99ddf --- /dev/null +++ b/patches/candle-transformers/src/models/voxtral/voxtral_llama.rs @@ -0,0 +1,471 @@ +use crate::models::with_tracing::{linear_no_bias as linear, Linear, RmsNorm}; +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::{embedding, Embedding, Module, VarBuilder}; +use serde::Deserialize; +use std::collections::HashMap; + +pub const DEFAULT_MAX_SEQ_LEN: usize = 4096; + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct VoxtralLlamaConfig { + pub hidden_size: usize, + pub intermediate_size: usize, + pub vocab_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub head_dim: Option, // explicit head_dim from config + pub use_flash_attn: bool, + pub rms_norm_eps: f64, + pub rope_theta: f32, + pub max_position_embeddings: usize, + pub tie_word_embeddings: bool, +} + +impl VoxtralLlamaConfig { + /// Voxtral 3B text model configuration + pub fn voxtral_3b() -> Self { + Self { + hidden_size: 3072, + intermediate_size: 8192, + vocab_size: 131072, + num_hidden_layers: 30, + num_attention_heads: 32, + num_key_value_heads: 8, + head_dim: Some(128), // Voxtral uses explicit head_dim=128 + use_flash_attn: true, + rms_norm_eps: 1e-5, + rope_theta: 100_000_000.0, + max_position_embeddings: 131072, + tie_word_embeddings: false, + } + } + + /// Voxtral 24B text model configuration + pub fn voxtral_24b() -> Self { + Self { + hidden_size: 5120, + intermediate_size: 32768, + vocab_size: 131072, + num_hidden_layers: 40, + num_attention_heads: 32, + num_key_value_heads: 8, + head_dim: Some(128), // Voxtral uses explicit head_dim=128 + use_flash_attn: true, + rms_norm_eps: 1e-5, + rope_theta: 100_000_000.0, + max_position_embeddings: 131072, + tie_word_embeddings: false, + } + } +} + +#[derive(Debug, Clone)] +pub struct VoxtralLlamaCache { + masks: HashMap, + pub use_kv_cache: bool, + kvs: Vec>, + cos: Tensor, + sin: Tensor, + device: Device, +} + +fn calculate_default_inv_freq(cfg: &VoxtralLlamaConfig) -> Vec { + let head_dim = cfg + .head_dim + .unwrap_or(cfg.hidden_size / cfg.num_attention_heads); + (0..head_dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f32 / head_dim as f32)) + .collect() +} + +impl VoxtralLlamaCache { + pub fn new( + use_kv_cache: bool, + dtype: DType, + config: &VoxtralLlamaConfig, + device: &Device, + ) -> Result { + // precompute freqs_cis + let theta = calculate_default_inv_freq(config); + + let theta = Tensor::new(theta, device)?; + + let idx_theta = Tensor::arange(0, config.max_position_embeddings as u32, device)? + .to_dtype(DType::F32)? + .reshape((config.max_position_embeddings, 1))? + .matmul(&theta.reshape((1, theta.elem_count()))?)?; + // This is different from the paper, see: + // https://github.com/huggingface/transformers/blob/6112b1c6442aaf7affd2b0676a1cd4eee30c45cf/src/transformers/models/llama/modeling_llama.py#L112 # trufflehog:ignore + let cos = idx_theta.cos()?.to_dtype(dtype)?; + let sin = idx_theta.sin()?.to_dtype(dtype)?; + Ok(Self { + masks: HashMap::new(), + use_kv_cache, + kvs: vec![None; config.num_hidden_layers], + device: device.clone(), + cos, + sin, + }) + } + + 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), &self.device)?; + self.masks.insert(t, mask.clone()); + Ok(mask) + } + } +} + +#[derive(Debug, Clone)] +struct CausalSelfAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_attention_heads: usize, + num_key_value_heads: usize, + head_dim: usize, + use_flash_attn: bool, + span: tracing::Span, + span_rot: tracing::Span, + max_position_embeddings: usize, +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +impl CausalSelfAttention { + fn apply_rotary_emb( + &self, + x: &Tensor, + index_pos: usize, + cache: &VoxtralLlamaCache, + ) -> Result { + let _enter = self.span_rot.enter(); + let (_b_sz, _, seq_len, _hidden_size) = x.dims4()?; + let cos = cache.cos.narrow(0, index_pos, seq_len)?; + let sin = cache.sin.narrow(0, index_pos, seq_len)?; + + // Ensure dtype consistency between input tensor and position embeddings + let x_dtype = x.dtype(); + let cos = if cos.dtype() != x_dtype { + cos.to_dtype(x_dtype)? + } else { + cos + }; + let sin = if sin.dtype() != x_dtype { + sin.to_dtype(x_dtype)? + } else { + sin + }; + + candle_nn::rotary_emb::rope(x, &cos, &sin) + } + + fn forward( + &self, + x: &Tensor, + index_pos: usize, + block_idx: usize, + cache: &mut VoxtralLlamaCache, + ) -> Result { + let _enter = self.span.enter(); + let (b_sz, seq_len, _hidden_size) = x.dims3()?; + let q = self.q_proj.forward(x)?; + let k = self.k_proj.forward(x)?; + let v = self.v_proj.forward(x)?; + + let q = q + .reshape((b_sz, seq_len, self.num_attention_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let k = k + .reshape((b_sz, seq_len, self.num_key_value_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let mut v = v + .reshape((b_sz, seq_len, self.num_key_value_heads, self.head_dim))? + .transpose(1, 2)?; + + let q = self.apply_rotary_emb(&q, index_pos, cache)?; + let mut k = self.apply_rotary_emb(&k, index_pos, cache)?; + + if cache.use_kv_cache { + if let Some((cache_k, cache_v)) = &cache.kvs[block_idx] { + k = Tensor::cat(&[cache_k, &k], 2)?.contiguous()?; + v = Tensor::cat(&[cache_v, &v], 2)?.contiguous()?; + let k_seq_len = k.dims()[1]; + if k_seq_len > self.max_position_embeddings { + k = k + .narrow( + D::Minus1, + k_seq_len - self.max_position_embeddings, + self.max_position_embeddings, + )? + .contiguous()? + } + let v_seq_len = v.dims()[1]; + if v_seq_len > 2 * self.max_position_embeddings { + v = v + .narrow( + D::Minus1, + v_seq_len - self.max_position_embeddings, + self.max_position_embeddings, + )? + .contiguous()? + } + } + cache.kvs[block_idx] = Some((k.clone(), v.clone())) + } + + let k = self.repeat_kv(k)?; + let v = self.repeat_kv(v)?; + + let y = if self.use_flash_attn { + // flash-attn expects (b_sz, seq_len, nheads, head_dim) + let q = q.transpose(1, 2)?; + let k = k.transpose(1, 2)?; + let v = v.transpose(1, 2)?; + let softmax_scale = 1f32 / (self.head_dim as f32).sqrt(); + flash_attn(&q, &k, &v, softmax_scale, seq_len > 1)?.transpose(1, 2)? + } else { + let in_dtype = q.dtype(); + let q = q.to_dtype(DType::F32)?; + let k = k.to_dtype(DType::F32)?; + let v = v.to_dtype(DType::F32)?; + let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?; + let att = if seq_len == 1 { + att + } else { + let mask = cache.mask(seq_len)?.broadcast_as(att.shape())?; + 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. + att.matmul(&v.contiguous()?)?.to_dtype(in_dtype)? + }; + // Use the actual tensor dimensions from attention computation + let actual_hidden_size = self.num_attention_heads * self.head_dim; + let y = y + .transpose(1, 2)? + .reshape(&[b_sz, seq_len, actual_hidden_size])?; + let y = self.o_proj.forward(&y)?; + Ok(y) + } + + fn repeat_kv(&self, x: Tensor) -> Result { + crate::utils::repeat_kv(x, self.num_attention_heads / self.num_key_value_heads) + } + + fn load(vb: VarBuilder, cfg: &VoxtralLlamaConfig) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "attn"); + let span_rot = tracing::span!(tracing::Level::TRACE, "attn-rot"); + let size_in = cfg.hidden_size; + + // Use explicit head_dim if provided, otherwise calculate from hidden_size + let head_dim = cfg + .head_dim + .unwrap_or(cfg.hidden_size / cfg.num_attention_heads); + let size_q = head_dim * cfg.num_attention_heads; + let size_kv = head_dim * cfg.num_key_value_heads; + + let q_proj = linear(size_in, size_q, vb.pp("q_proj"))?; + let k_proj = linear(size_in, size_kv, vb.pp("k_proj"))?; + let v_proj = linear(size_in, size_kv, vb.pp("v_proj"))?; + let o_proj = linear(size_q, size_in, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_attention_heads: cfg.num_attention_heads, + num_key_value_heads: cfg.num_key_value_heads, + head_dim, // use the calculated head_dim from above + use_flash_attn: cfg.use_flash_attn, + span, + span_rot, + max_position_embeddings: cfg.max_position_embeddings, + }) + } +} + +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) +} + +#[derive(Debug, Clone)] +struct Mlp { + c_fc1: Linear, + c_fc2: Linear, + c_proj: Linear, + span: tracing::Span, +} + +impl Mlp { + fn forward(&self, x: &Tensor) -> Result { + let _enter = self.span.enter(); + let x = (candle_nn::ops::silu(&self.c_fc1.forward(x)?)? * self.c_fc2.forward(x)?)?; + self.c_proj.forward(&x) + } + + fn load(vb: VarBuilder, cfg: &VoxtralLlamaConfig) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "mlp"); + let h_size = cfg.hidden_size; + let i_size = cfg.intermediate_size; + let c_fc1 = linear(h_size, i_size, vb.pp("gate_proj"))?; + let c_fc2 = linear(h_size, i_size, vb.pp("up_proj"))?; + let c_proj = linear(i_size, h_size, vb.pp("down_proj"))?; + Ok(Self { + c_fc1, + c_fc2, + c_proj, + span, + }) + } +} + +#[derive(Debug, Clone)] +struct Block { + rms_1: RmsNorm, + attn: CausalSelfAttention, + rms_2: RmsNorm, + mlp: Mlp, + span: tracing::Span, +} + +impl Block { + fn forward( + &self, + x: &Tensor, + index_pos: usize, + block_idx: usize, + cache: &mut VoxtralLlamaCache, + ) -> Result { + let _enter = self.span.enter(); + let residual = x; + let x = self.rms_1.forward(x)?; + let x = (self.attn.forward(&x, index_pos, block_idx, cache)? + residual)?; + let residual = &x; + let x = (self.mlp.forward(&self.rms_2.forward(&x)?)? + residual)?; + Ok(x) + } + + fn load(vb: VarBuilder, cfg: &VoxtralLlamaConfig) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "block"); + let attn = CausalSelfAttention::load(vb.pp("self_attn"), cfg)?; + let mlp = Mlp::load(vb.pp("mlp"), cfg)?; + let rms_1 = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let rms_2 = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + rms_1, + attn, + rms_2, + mlp, + span, + }) + } +} + +#[derive(Debug, Clone)] +pub struct VoxtralLlama { + wte: Embedding, + blocks: Vec, + ln_f: RmsNorm, + lm_head: Linear, +} + +impl VoxtralLlama { + // required by LLaVA + pub fn embed(&self, x: &Tensor) -> Result { + self.wte.forward(x) + } + // required by LLaVA + pub fn forward_input_embed( + &self, + input_embed: &Tensor, + index_pos: usize, + cache: &mut VoxtralLlamaCache, + ) -> Result { + let (_, seq_len, _) = input_embed.dims3()?; + let mut x = input_embed.clone(); + for (block_idx, block) in self.blocks.iter().enumerate() { + x = block.forward(&x, index_pos, block_idx, cache)?; + } + let x = self.ln_f.forward(&x)?; + // Handle both single token and multi-token sequences properly + let x = if seq_len == 1 { + x.i((.., 0, ..))? + } else { + x.i((.., seq_len - 1, ..))? + } + .contiguous()?; + let logits = self.lm_head.forward(&x)?; + logits.to_dtype(DType::F32) + } + + pub fn forward( + &self, + x: &Tensor, + index_pos: usize, + cache: &mut VoxtralLlamaCache, + ) -> Result { + let (_b_sz, seq_len) = x.dims2()?; + let mut x = self.wte.forward(x)?; + for (block_idx, block) in self.blocks.iter().enumerate() { + x = block.forward(&x, index_pos, block_idx, cache)?; + } + let x = self.ln_f.forward(&x)?; + let x = x.i((.., seq_len - 1, ..))?.contiguous()?; + let logits = self.lm_head.forward(&x)?; + logits.to_dtype(DType::F32) + } + + pub fn load(vb: VarBuilder, cfg: &VoxtralLlamaConfig) -> Result { + let wte = embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("model.embed_tokens"))?; + let lm_head = if cfg.tie_word_embeddings { + Linear::from_weights(wte.embeddings().clone(), None) + } else { + linear(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))? + }; + let ln_f = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("model.norm"))?; + let blocks: Vec<_> = (0..cfg.num_hidden_layers) + .map(|i| Block::load(vb.pp(format!("model.layers.{i}")), cfg).unwrap()) + .collect(); + + Ok(Self { + wte, + blocks, + ln_f, + lm_head, + }) + } +} diff --git a/patches/candle-transformers/src/models/whisper/audio.rs b/patches/candle-transformers/src/models/whisper/audio.rs new file mode 100644 index 0000000000..1206fdf081 --- /dev/null +++ b/patches/candle-transformers/src/models/whisper/audio.rs @@ -0,0 +1,336 @@ +// Audio processing code, adapted from whisper.cpp +// https://github.com/ggerganov/whisper.cpp + +use candle::utils::get_num_threads; +use std::sync::Arc; +use std::thread; + +pub trait Float: + num_traits::Float + num_traits::FloatConst + num_traits::NumAssign + Send + Sync +{ +} + +impl Float for f32 {} +impl Float for f64 {} + +// https://github.com/ggerganov/whisper.cpp/blob/4774d2feb01a772a15de81ffc34b34a1f294f020/whisper.cpp#L2357 +fn fft(inp: &[T]) -> Vec { + let n = inp.len(); + let zero = T::zero(); + if n == 1 { + return vec![inp[0], zero]; + } + if n % 2 == 1 { + return dft(inp); + } + let mut out = vec![zero; n * 2]; + + let mut even = Vec::with_capacity(n / 2); + let mut odd = Vec::with_capacity(n / 2); + + for (i, &inp) in inp.iter().enumerate() { + if i % 2 == 0 { + even.push(inp) + } else { + odd.push(inp); + } + } + + let even_fft = fft(&even); + let odd_fft = fft(&odd); + + let two_pi = T::PI() + T::PI(); + let n_t = T::from(n).unwrap(); + for k in 0..n / 2 { + let k_t = T::from(k).unwrap(); + let theta = two_pi * k_t / n_t; + let re = theta.cos(); + let im = -theta.sin(); + + let re_odd = odd_fft[2 * k]; + let im_odd = odd_fft[2 * k + 1]; + + out[2 * k] = even_fft[2 * k] + re * re_odd - im * im_odd; + out[2 * k + 1] = even_fft[2 * k + 1] + re * im_odd + im * re_odd; + + out[2 * (k + n / 2)] = even_fft[2 * k] - re * re_odd + im * im_odd; + out[2 * (k + n / 2) + 1] = even_fft[2 * k + 1] - re * im_odd - im * re_odd; + } + out +} + +// https://github.com/ggerganov/whisper.cpp/blob/4774d2feb01a772a15de81ffc34b34a1f294f020/whisper.cpp#L2337 +fn dft(inp: &[T]) -> Vec { + let zero = T::zero(); + let n = inp.len(); + let two_pi = T::PI() + T::PI(); + + let mut out = Vec::with_capacity(2 * n); + let n_t = T::from(n).unwrap(); + for k in 0..n { + let k_t = T::from(k).unwrap(); + let mut re = zero; + let mut im = zero; + + for (j, &inp) in inp.iter().enumerate() { + let j_t = T::from(j).unwrap(); + let angle = two_pi * k_t * j_t / n_t; + re += inp * angle.cos(); + im -= inp * angle.sin(); + } + + out.push(re); + out.push(im); + } + out +} + +#[allow(clippy::too_many_arguments)] +// https://github.com/ggerganov/whisper.cpp/blob/4774d2feb01a772a15de81ffc34b34a1f294f020/whisper.cpp#L2414 +fn log_mel_spectrogram_w( + ith: usize, + hann: &[T], + samples: &[T], + filters: &[T], + fft_size: usize, + fft_step: usize, + speed_up: bool, + n_len: usize, + n_mel: usize, + n_threads: usize, +) -> Vec { + let n_fft = if speed_up { + 1 + fft_size / 4 + } else { + 1 + fft_size / 2 + }; + + let zero = T::zero(); + let half = T::from(0.5).unwrap(); + let mut fft_in = vec![zero; fft_size]; + let mut mel = vec![zero; n_len * n_mel]; + let n_samples = samples.len(); + let end = std::cmp::min(n_samples / fft_step + 1, n_len); + + for i in (ith..end).step_by(n_threads) { + let offset = i * fft_step; + + // apply Hanning window + for j in 0..std::cmp::min(fft_size, n_samples - offset) { + fft_in[j] = hann[j] * samples[offset + j]; + } + + // fill the rest with zeros + if n_samples - offset < fft_size { + fft_in[n_samples - offset..].fill(zero); + } + + // FFT + let mut fft_out: Vec = fft(&fft_in); + + // Calculate modulus^2 of complex numbers + for j in 0..fft_size { + fft_out[j] = fft_out[2 * j] * fft_out[2 * j] + fft_out[2 * j + 1] * fft_out[2 * j + 1]; + } + for j in 1..fft_size / 2 { + let v = fft_out[fft_size - j]; + fft_out[j] += v; + } + + if speed_up { + // scale down in the frequency domain results in a speed up in the time domain + for j in 0..n_fft { + fft_out[j] = half * (fft_out[2 * j] + fft_out[2 * j + 1]); + } + } + + // mel spectrogram + for j in 0..n_mel { + let mut sum = zero; + let mut k = 0; + // Unroll loop + while k < n_fft.saturating_sub(3) { + sum += fft_out[k] * filters[j * n_fft + k] + + fft_out[k + 1] * filters[j * n_fft + k + 1] + + fft_out[k + 2] * filters[j * n_fft + k + 2] + + fft_out[k + 3] * filters[j * n_fft + k + 3]; + k += 4; + } + // Handle remainder + while k < n_fft { + sum += fft_out[k] * filters[j * n_fft + k]; + k += 1; + } + mel[j * n_len + i] = T::max(sum, T::from(1e-10).unwrap()).log10(); + } + } + mel +} + +pub fn log_mel_spectrogram_( + samples: &[T], + filters: &[T], + fft_size: usize, + fft_step: usize, + n_mel: usize, + speed_up: bool, +) -> Vec { + let zero = T::zero(); + let two_pi = T::PI() + T::PI(); + let half = T::from(0.5).unwrap(); + let one = T::from(1.0).unwrap(); + let four = T::from(4.0).unwrap(); + let fft_size_t = T::from(fft_size).unwrap(); + + let hann: Vec = (0..fft_size) + .map(|i| half * (one - ((two_pi * T::from(i).unwrap()) / fft_size_t).cos())) + .collect(); + let n_len = samples.len() / fft_step; + + // pad audio with at least one extra chunk of zeros + let pad = 100 * super::CHUNK_LENGTH / 2; + let n_len = if !n_len.is_multiple_of(pad) { + (n_len / pad + 1) * pad + } else { + n_len + }; + let n_len = n_len + pad; + let samples = { + let mut samples_padded = samples.to_vec(); + let to_add = n_len * fft_step - samples.len(); + samples_padded.extend(std::iter::repeat_n(zero, to_add)); + samples_padded + }; + + // ensure that the number of threads is even and less than 12 + let n_threads = std::cmp::min(get_num_threads() - get_num_threads() % 2, 12); + let n_threads = std::cmp::max(n_threads, 2); + + let hann = Arc::new(hann); + let samples = Arc::new(samples); + let filters = Arc::new(filters); + + // use scope to allow for non static references to be passed to the threads + // and directly collect the results into a single vector + let all_outputs = thread::scope(|s| { + (0..n_threads) + // create threads and return their handles + .map(|thread_id| { + let hann = Arc::clone(&hann); + let samples = Arc::clone(&samples); + let filters = Arc::clone(&filters); + // spawn new thread and start work + s.spawn(move || { + log_mel_spectrogram_w( + thread_id, &hann, &samples, &filters, fft_size, fft_step, speed_up, n_len, + n_mel, n_threads, + ) + }) + }) + .collect::>() + .into_iter() + // wait for each thread to finish and collect their results + .map(|handle| handle.join().expect("Thread failed")) + .collect::>() + }); + + let l = all_outputs[0].len(); + let mut mel = vec![zero; l]; + + // iterate over mel spectrogram segments, dividing work by threads. + for segment_start in (0..l).step_by(n_threads) { + // go through each thread's output. + for thread_output in all_outputs.iter() { + // add each thread's piece to our mel spectrogram. + for offset in 0..n_threads { + let mel_index = segment_start + offset; // find location in mel. + if mel_index < mel.len() { + // Make sure we don't go out of bounds. + mel[mel_index] += thread_output[mel_index]; + } + } + } + } + + let mmax = mel + .iter() + .max_by(|&u, &v| u.partial_cmp(v).unwrap_or(std::cmp::Ordering::Greater)) + .copied() + .unwrap_or(zero) + - T::from(8).unwrap(); + for m in mel.iter_mut() { + let v = T::max(*m, mmax); + *m = v / four + one + } + mel +} + +pub fn pcm_to_mel(cfg: &super::Config, samples: &[T], filters: &[T]) -> Vec { + log_mel_spectrogram_( + samples, + filters, + super::N_FFT, + super::HOP_LENGTH, + cfg.num_mel_bins, + false, + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_fft() { + let input = vec![0.0, 1.0, 0.0, 0.0]; + let output = fft(&input); + assert_eq!( + output, + vec![ + 1.0, + 0.0, + 6.123233995736766e-17, + -1.0, + -1.0, + 0.0, + -6.123233995736766e-17, + 1.0 + ] + ); + } + + #[test] + fn test_dft() { + let input = vec![0.0, 1.0, 0.0, 0.0]; + let output = dft(&input); + assert_eq!( + output, + vec![ + 1.0, + 0.0, + 6.123233995736766e-17, + -1.0, + -1.0, + -1.2246467991473532e-16, + -1.8369701987210297e-16, + 1.0 + ] + ); + } + + #[test] + fn test_log_mel_spectrogram() { + let samples = vec![0.0; 1000]; + let filters = vec![0.0; 1000]; + let output = log_mel_spectrogram_(&samples, &filters, 100, 10, 10, false); + assert_eq!(output.len(), 30_000); + } + + #[test] + fn test_tiny_log_mel_spectrogram() { + let samples = vec![0.0; 100]; + let filters = vec![0.0; 100]; + let output = log_mel_spectrogram_(&samples, &filters, 20, 2, 2, false); + assert_eq!(output.len(), 6_000); + } +} diff --git a/patches/candle-transformers/src/models/whisper/mod.rs b/patches/candle-transformers/src/models/whisper/mod.rs new file mode 100644 index 0000000000..d7082ea6d8 --- /dev/null +++ b/patches/candle-transformers/src/models/whisper/mod.rs @@ -0,0 +1,58 @@ +//! Whisper Model Implementation +//! +//! Whisper is an automatic speech recognition (ASR) system trained on large amounts +//! of multilingual and multitask supervised data collected from the web. It can be used to +//! convert audio files (in the `.wav` format) to text. Supported features include +//! language detection as well as multilingual speech recognition. +//! +//! - ⚡ [Interactive Wasm Example](https://huggingface.co/spaces/lmz/candle-whisper) +//! - 💻 [GH Link](https://github.com/openai/whisper) +//! - 💻 Transformers Python [reference implementation](https://github.com/huggingface/transformers/blob/main/src/transformers/models/whisper/modeling_whisper.py) +//! +//! +pub mod audio; +pub mod model; +pub mod quantized_model; + +use serde::Deserialize; + +// The names in comments correspond to the original implementation: +// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L17 +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Config { + pub num_mel_bins: usize, // n_mels + pub max_source_positions: usize, // n_audio_ctx + pub d_model: usize, // n_audio_state + pub encoder_attention_heads: usize, // n_audio_head + pub encoder_layers: usize, // n_audio_layer + pub vocab_size: usize, // n_vocab + pub max_target_positions: usize, // n_text_ctx + // pub n_text_state: usize, + pub decoder_attention_heads: usize, // n_text_head + pub decoder_layers: usize, // n_text_layer + #[serde(default)] + pub suppress_tokens: Vec, +} + +pub const DTYPE: candle::DType = candle::DType::F32; + +// Audio parameters. +pub const SAMPLE_RATE: usize = 16000; +pub const N_FFT: usize = 400; +pub const HOP_LENGTH: usize = 160; +pub const CHUNK_LENGTH: usize = 30; +pub const N_SAMPLES: usize = CHUNK_LENGTH * SAMPLE_RATE; // 480000 samples in a 30-second chunk +pub const N_FRAMES: usize = N_SAMPLES / HOP_LENGTH; // 3000 frames in a mel spectrogram input + +pub const NO_SPEECH_THRESHOLD: f64 = 0.6; +pub const LOGPROB_THRESHOLD: f64 = -1.0; +pub const TEMPERATURES: [f64; 6] = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]; +pub const COMPRESSION_RATIO_THRESHOLD: f64 = 2.4; + +// Tokenizer dependent bits. +pub const SOT_TOKEN: &str = "<|startoftranscript|>"; +pub const TRANSCRIBE_TOKEN: &str = "<|transcribe|>"; +pub const TRANSLATE_TOKEN: &str = "<|translate|>"; +pub const NO_TIMESTAMPS_TOKEN: &str = "<|notimestamps|>"; +pub const EOT_TOKEN: &str = "<|endoftext|>"; +pub const NO_SPEECH_TOKENS: [&str; 2] = ["<|nocaptions|>", "<|nospeech|>"]; diff --git a/patches/candle-transformers/src/models/whisper/model.rs b/patches/candle-transformers/src/models/whisper/model.rs new file mode 100644 index 0000000000..2f34b1800f --- /dev/null +++ b/patches/candle-transformers/src/models/whisper/model.rs @@ -0,0 +1,400 @@ +use super::Config; +use crate::models::with_tracing::{linear, linear_no_bias, Linear}; +use candle::{Device, IndexOp, Result, Tensor, D}; +use candle_nn::{embedding, Conv1d, Conv1dConfig, Embedding, LayerNorm, Module, VarBuilder}; + +fn conv1d( + in_channels: usize, + out_channels: usize, + kernel_size: usize, + config: Conv1dConfig, + vb: VarBuilder, +) -> Result { + let weight = vb.get((out_channels, in_channels, kernel_size), "weight")?; + let bias = vb.get(out_channels, "bias")?; + Ok(Conv1d::new(weight, Some(bias), config)) +} + +fn layer_norm(size: usize, vb: VarBuilder) -> Result { + let weight = vb.get(size, "weight")?; + let bias = vb.get(size, "bias")?; + Ok(LayerNorm::new(weight, bias, 1e-5)) +} + +// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L62 +#[derive(Debug, Clone)] +struct MultiHeadAttention { + query: Linear, + key: Linear, + value: Linear, + out: Linear, + n_head: usize, + span: tracing::Span, + softmax_span: tracing::Span, + matmul_span: tracing::Span, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl MultiHeadAttention { + fn load(n_state: usize, n_head: usize, vb: VarBuilder) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "multi-head-attn"); + let softmax_span = tracing::span!(tracing::Level::TRACE, "multi-head-attn-softmax"); + let matmul_span = tracing::span!(tracing::Level::TRACE, "multi-head-attn-matmul"); + let query = linear(n_state, n_state, vb.pp("q_proj"))?; + let value = linear(n_state, n_state, vb.pp("v_proj"))?; + let key = linear_no_bias(n_state, n_state, vb.pp("k_proj"))?; + let out = linear(n_state, n_state, vb.pp("out_proj"))?; + Ok(Self { + query, + key, + value, + out, + n_head, + span, + softmax_span, + matmul_span, + kv_cache: None, + }) + } + + fn forward( + &mut self, + x: &Tensor, + xa: Option<&Tensor>, + mask: Option<&Tensor>, + flush_cache: bool, + ) -> Result { + let _enter = self.span.enter(); + let q = self.query.forward(x)?; + let (k, v) = match xa { + None => { + let k = self.key.forward(x)?; + let v = self.value.forward(x)?; + (k, v) + } + Some(x) => { + if flush_cache { + self.kv_cache = None; + } + if let Some((k, v)) = &self.kv_cache { + (k.clone(), v.clone()) + } else { + let k = self.key.forward(x)?; + let v = self.value.forward(x)?; + self.kv_cache = Some((k.clone(), v.clone())); + (k, v) + } + } + }; + let wv = self.qkv_attention(&q, &k, &v, mask)?; + let out = self.out.forward(&wv)?; + Ok(out) + } + + fn reshape_head(&self, x: &Tensor) -> Result { + let (n_batch, n_ctx, n_state) = x.dims3()?; + let target_dims = &[n_batch, n_ctx, self.n_head, n_state / self.n_head]; + x.reshape(target_dims)?.transpose(1, 2) + } + + fn qkv_attention( + &self, + q: &Tensor, + k: &Tensor, + v: &Tensor, + mask: Option<&Tensor>, + ) -> Result { + let (_, n_ctx, n_state) = q.dims3()?; + let scale = ((n_state / self.n_head) as f64).powf(-0.25); + let q = (self.reshape_head(q)? * scale)?; + let k = (self.reshape_head(k)?.transpose(2, 3)? * scale)?; + let v = self.reshape_head(v)?.contiguous()?; + let mut qk = { + let _enter = self.matmul_span.enter(); + q.matmul(&k)? + }; + if let Some(mask) = mask { + let mask = mask.i((0..n_ctx, 0..n_ctx))?; + qk = qk.broadcast_add(&mask)? + } + let w = { + let _enter = self.softmax_span.enter(); + candle_nn::ops::softmax_last_dim(&qk)? + }; + let wv = { + let _enter = self.matmul_span.enter(); + w.matmul(&v)? + } + .transpose(1, 2)? + .flatten_from(2)?; + Ok(wv) + } + + fn reset_kv_cache(&mut self) { + self.kv_cache = None; + } +} + +// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L111 +#[derive(Debug, Clone)] +struct ResidualAttentionBlock { + attn: MultiHeadAttention, + attn_ln: LayerNorm, + cross_attn: Option<(MultiHeadAttention, LayerNorm)>, + mlp_linear1: Linear, + mlp_linear2: Linear, + mlp_ln: LayerNorm, + span: tracing::Span, +} + +impl ResidualAttentionBlock { + fn load(n_state: usize, n_head: usize, ca: bool, vb: VarBuilder) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "residual-attn"); + let attn = MultiHeadAttention::load(n_state, n_head, vb.pp("self_attn"))?; + let attn_ln = layer_norm(n_state, vb.pp("self_attn_layer_norm"))?; + let cross_attn = if ca { + let cross_attn = MultiHeadAttention::load(n_state, n_head, vb.pp("encoder_attn"))?; + let cross_attn_ln = layer_norm(n_state, vb.pp("encoder_attn_layer_norm"))?; + Some((cross_attn, cross_attn_ln)) + } else { + None + }; + let n_mlp = n_state * 4; + let mlp_linear1 = linear(n_state, n_mlp, vb.pp("fc1"))?; + let mlp_linear2 = linear(n_mlp, n_state, vb.pp("fc2"))?; + let mlp_ln = layer_norm(n_state, vb.pp("final_layer_norm"))?; + Ok(Self { + attn, + attn_ln, + cross_attn, + mlp_linear1, + mlp_linear2, + mlp_ln, + span, + }) + } + + fn forward( + &mut self, + x: &Tensor, + xa: Option<&Tensor>, + mask: Option<&Tensor>, + flush_kv_cache: bool, + ) -> Result { + let _enter = self.span.enter(); + let attn = self + .attn + .forward(&self.attn_ln.forward(x)?, None, mask, flush_kv_cache)?; + let mut x = (x + attn)?; + if let Some((attn, ln)) = &mut self.cross_attn { + x = (&x + attn.forward(&ln.forward(&x)?, xa, None, flush_kv_cache)?)?; + } + let mlp = self.mlp_linear2.forward( + &self + .mlp_linear1 + .forward(&self.mlp_ln.forward(&x)?)? + .gelu()?, + )?; + x + mlp + } + + fn reset_kv_cache(&mut self) { + self.attn.reset_kv_cache(); + if let Some((attn, _)) = &mut self.cross_attn { + attn.reset_kv_cache(); + } + } +} + +fn sinusoids(length: usize, channels: usize, device: &Device) -> Result { + let max_timescale = 10000f32; + let log_timescale_increment = max_timescale.ln() / (channels / 2 - 1) as f32; + let inv_timescales: Vec<_> = (0..channels / 2) + .map(|i| (i as f32 * (-log_timescale_increment)).exp()) + .collect(); + let inv_timescales = Tensor::new(inv_timescales.as_slice(), device)?.unsqueeze(0)?; + let arange = Tensor::arange(0, length as u32, device)? + .to_dtype(candle::DType::F32)? + .unsqueeze(1)?; + let sh = (length, channels / 2); + let scaled_time = (arange.broadcast_as(sh)? * inv_timescales.broadcast_as(sh)?)?; + let sincos = Tensor::cat(&[scaled_time.sin()?, scaled_time.cos()?], 1)?; + Ok(sincos) +} + +// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L143 +#[derive(Debug, Clone)] +pub struct AudioEncoder { + conv1: Conv1d, + conv2: Conv1d, + positional_embedding: Tensor, + blocks: Vec, + ln_post: LayerNorm, + span: tracing::Span, + conv1_span: tracing::Span, + conv2_span: tracing::Span, +} + +impl AudioEncoder { + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "audio-encoder"); + let conv1_span = tracing::span!(tracing::Level::TRACE, "conv1"); + let conv2_span = tracing::span!(tracing::Level::TRACE, "conv2"); + let n_state = cfg.d_model; + let n_head = cfg.encoder_attention_heads; + let n_ctx = cfg.max_source_positions; + let cfg1 = Conv1dConfig { + padding: 1, + stride: 1, + groups: 1, + dilation: 1, + cudnn_fwd_algo: None, + }; + let cfg2 = Conv1dConfig { + padding: 1, + stride: 2, + groups: 1, + dilation: 1, + cudnn_fwd_algo: None, + }; + let conv1 = conv1d(cfg.num_mel_bins, n_state, 3, cfg1, vb.pp("conv1"))?; + let conv2 = conv1d(n_state, n_state, 3, cfg2, vb.pp("conv2"))?; + let positional_embedding = sinusoids(n_ctx, n_state, vb.device())?; + let blocks = (0..cfg.encoder_layers) + .map(|i| { + ResidualAttentionBlock::load(n_state, n_head, false, vb.pp(format!("layers.{i}"))) + }) + .collect::>>()?; + let ln_post = layer_norm(n_state, vb.pp("layer_norm"))?; + Ok(Self { + conv1, + conv2, + positional_embedding, + blocks, + ln_post, + conv1_span, + conv2_span, + span, + }) + } + + pub fn forward(&mut self, x: &Tensor, flush_kv_cache: bool) -> Result { + let _enter = self.span.enter(); + let x = { + let _enter = self.conv1_span.enter(); + self.conv1.forward(x)?.gelu()? + }; + let x = { + let _enter = self.conv2_span.enter(); + self.conv2.forward(&x)?.gelu()? + }; + let x = x.transpose(1, 2)?; + let (_bsize, seq_len, _hidden) = x.dims3()?; + let positional_embedding = self.positional_embedding.narrow(0, 0, seq_len)?; + let mut x = x.broadcast_add(&positional_embedding)?; + for block in self.blocks.iter_mut() { + x = block.forward(&x, None, None, flush_kv_cache)? + } + let x = self.ln_post.forward(&x)?; + Ok(x) + } +} + +// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L176 +#[derive(Debug, Clone)] +pub struct TextDecoder { + token_embedding: Embedding, + positional_embedding: Tensor, + blocks: Vec, + ln: LayerNorm, + mask: Tensor, + span: tracing::Span, + span_final: tracing::Span, +} + +impl TextDecoder { + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "text-decoder"); + let span_final = tracing::span!(tracing::Level::TRACE, "text-decoder-final"); + let n_state = cfg.d_model; + let n_head = cfg.decoder_attention_heads; + let n_ctx = cfg.max_target_positions; + let token_embedding = embedding(cfg.vocab_size, n_state, vb.pp("embed_tokens"))?; + let positional_embedding = vb.get((n_ctx, n_state), "embed_positions.weight")?; + let blocks = (0..cfg.decoder_layers) + .map(|i| { + ResidualAttentionBlock::load(n_state, n_head, true, vb.pp(format!("layers.{i}"))) + }) + .collect::>>()?; + let ln = layer_norm(n_state, vb.pp("layer_norm"))?; + let mask: Vec<_> = (0..n_ctx) + .flat_map(|i| (0..n_ctx).map(move |j| if j > i { f32::NEG_INFINITY } else { 0f32 })) + .collect(); + let mask = Tensor::from_vec(mask, (n_ctx, n_ctx), vb.device())?; + Ok(Self { + token_embedding, + positional_embedding, + blocks, + ln, + mask, + span, + span_final, + }) + } + + pub fn forward(&mut self, x: &Tensor, xa: &Tensor, flush_kv_cache: bool) -> Result { + let _enter = self.span.enter(); + let last = x.dim(D::Minus1)?; + let token_embedding = self.token_embedding.forward(x)?; + let positional_embedding = self.positional_embedding.narrow(0, 0, last)?; + let mut x = token_embedding.broadcast_add(&positional_embedding)?; + for block in self.blocks.iter_mut() { + x = block.forward(&x, Some(xa), Some(&self.mask), flush_kv_cache)?; + } + self.ln.forward(&x) + } + + pub fn final_linear(&self, x: &Tensor) -> Result { + let b_size = x.dim(0)?; + let w = self.token_embedding.embeddings().broadcast_left(b_size)?; + let logits = { + let _enter = self.span_final.enter(); + x.matmul(&w.t()?)? + }; + Ok(logits) + } + + pub fn reset_kv_cache(&mut self) { + for block in self.blocks.iter_mut() { + block.reset_kv_cache(); + } + } +} + +// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L221 +#[derive(Debug, Clone)] +pub struct Whisper { + pub encoder: AudioEncoder, + pub decoder: TextDecoder, + pub config: Config, +} + +impl Whisper { + pub fn load(vb: &VarBuilder, config: Config) -> Result { + let encoder = AudioEncoder::load(vb.pp("model.encoder"), &config)?; + let decoder = TextDecoder::load(vb.pp("model.decoder"), &config)?; + Ok(Self { + encoder, + decoder, + config, + }) + } + + pub fn reset_kv_cache(&mut self) { + self.encoder + .blocks + .iter_mut() + .for_each(|b| b.reset_kv_cache()); + self.decoder.reset_kv_cache(); + } +} diff --git a/patches/candle-transformers/src/models/whisper/quantized_model.rs b/patches/candle-transformers/src/models/whisper/quantized_model.rs new file mode 100644 index 0000000000..15130fbdaa --- /dev/null +++ b/patches/candle-transformers/src/models/whisper/quantized_model.rs @@ -0,0 +1,401 @@ +use super::Config; +use crate::quantized_nn::{layer_norm, linear, linear_no_bias, Embedding, Linear}; +pub use crate::quantized_var_builder::VarBuilder; +use candle::{Device, IndexOp, Result, Tensor, D}; +use candle_nn::{Conv1d, Conv1dConfig, LayerNorm, Module}; + +fn conv1d( + in_channels: usize, + out_channels: usize, + kernel_size: usize, + config: Conv1dConfig, + vb: VarBuilder, +) -> Result { + let weight = vb + .get((out_channels, in_channels, kernel_size), "weight")? + .dequantize(vb.device())?; + let bias = vb.get(out_channels, "bias")?.dequantize(vb.device())?; + Ok(Conv1d::new(weight, Some(bias), config)) +} + +// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L62 +#[derive(Debug, Clone)] +struct MultiHeadAttention { + query: Linear, + key: Linear, + value: Linear, + out: Linear, + n_head: usize, + span: tracing::Span, + softmax_span: tracing::Span, + matmul_span: tracing::Span, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl MultiHeadAttention { + fn load(n_state: usize, n_head: usize, vb: VarBuilder) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "multi-head-attn"); + let softmax_span = tracing::span!(tracing::Level::TRACE, "multi-head-attn-softmax"); + let matmul_span = tracing::span!(tracing::Level::TRACE, "multi-head-attn-matmul"); + let query = linear(n_state, n_state, vb.pp("q_proj"))?; + let value = linear(n_state, n_state, vb.pp("v_proj"))?; + let key = linear_no_bias(n_state, n_state, vb.pp("k_proj"))?; + let out = linear(n_state, n_state, vb.pp("out_proj"))?; + Ok(Self { + query, + key, + value, + out, + n_head, + span, + softmax_span, + matmul_span, + kv_cache: None, + }) + } + + fn forward( + &mut self, + x: &Tensor, + xa: Option<&Tensor>, + mask: Option<&Tensor>, + flush_cache: bool, + ) -> Result { + let _enter = self.span.enter(); + let q = self.query.forward(x)?; + let (k, v) = match xa { + None => { + let k = self.key.forward(x)?; + let v = self.value.forward(x)?; + (k, v) + } + Some(x) => { + if flush_cache { + self.kv_cache = None; + } + if let Some((k, v)) = &self.kv_cache { + (k.clone(), v.clone()) + } else { + let k = self.key.forward(x)?; + let v = self.value.forward(x)?; + self.kv_cache = Some((k.clone(), v.clone())); + (k, v) + } + } + }; + let wv = self.qkv_attention(&q, &k, &v, mask)?; + let out = self.out.forward(&wv)?; + Ok(out) + } + + fn reshape_head(&self, x: &Tensor) -> Result { + let (n_batch, n_ctx, n_state) = x.dims3()?; + let target_dims = &[n_batch, n_ctx, self.n_head, n_state / self.n_head]; + x.reshape(target_dims)?.transpose(1, 2) + } + + fn qkv_attention( + &self, + q: &Tensor, + k: &Tensor, + v: &Tensor, + mask: Option<&Tensor>, + ) -> Result { + let (_, n_ctx, n_state) = q.dims3()?; + let scale = ((n_state / self.n_head) as f64).powf(-0.25); + let q = (self.reshape_head(q)? * scale)?; + let k = (self.reshape_head(k)?.transpose(2, 3)? * scale)?; + let v = self.reshape_head(v)?.contiguous()?; + let mut qk = { + let _enter = self.matmul_span.enter(); + q.matmul(&k)? + }; + if let Some(mask) = mask { + let mask = mask.i((0..n_ctx, 0..n_ctx))?; + qk = qk.broadcast_add(&mask)? + } + let w = { + let _enter = self.softmax_span.enter(); + candle_nn::ops::softmax_last_dim(&qk)? + }; + let wv = { + let _enter = self.matmul_span.enter(); + w.matmul(&v)? + } + .transpose(1, 2)? + .flatten_from(2)?; + Ok(wv) + } + + fn reset_kv_cache(&mut self) { + self.kv_cache = None; + } +} + +// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L111 +#[derive(Debug, Clone)] +struct ResidualAttentionBlock { + attn: MultiHeadAttention, + attn_ln: LayerNorm, + cross_attn: Option<(MultiHeadAttention, LayerNorm)>, + mlp_linear1: Linear, + mlp_linear2: Linear, + mlp_ln: LayerNorm, + span: tracing::Span, +} + +impl ResidualAttentionBlock { + fn load(n_state: usize, n_head: usize, ca: bool, vb: VarBuilder) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "residual-attn"); + let attn = MultiHeadAttention::load(n_state, n_head, vb.pp("self_attn"))?; + let attn_ln = layer_norm(n_state, 1e-5, vb.pp("self_attn_layer_norm"))?; + let cross_attn = if ca { + let cross_attn = MultiHeadAttention::load(n_state, n_head, vb.pp("encoder_attn"))?; + let cross_attn_ln = layer_norm(n_state, 1e-5, vb.pp("encoder_attn_layer_norm"))?; + Some((cross_attn, cross_attn_ln)) + } else { + None + }; + let n_mlp = n_state * 4; + let mlp_linear1 = linear(n_state, n_mlp, vb.pp("fc1"))?; + let mlp_linear2 = linear(n_mlp, n_state, vb.pp("fc2"))?; + let mlp_ln = layer_norm(n_state, 1e-5, vb.pp("final_layer_norm"))?; + Ok(Self { + attn, + attn_ln, + cross_attn, + mlp_linear1, + mlp_linear2, + mlp_ln, + span, + }) + } + + fn forward( + &mut self, + x: &Tensor, + xa: Option<&Tensor>, + mask: Option<&Tensor>, + flush_kv_cache: bool, + ) -> Result { + let _enter = self.span.enter(); + let attn = self + .attn + .forward(&self.attn_ln.forward(x)?, None, mask, flush_kv_cache)?; + let mut x = (x + attn)?; + if let Some((attn, ln)) = &mut self.cross_attn { + x = (&x + attn.forward(&ln.forward(&x)?, xa, None, flush_kv_cache)?)?; + } + let mlp = x + .apply(&self.mlp_ln)? + .apply(&self.mlp_linear1)? + .gelu()? + .apply(&self.mlp_linear2)?; + x + mlp + } + + fn reset_kv_cache(&mut self) { + self.attn.reset_kv_cache(); + if let Some((attn, _)) = &mut self.cross_attn { + attn.reset_kv_cache(); + } + } +} + +fn sinusoids(length: usize, channels: usize, device: &Device) -> Result { + let max_timescale = 10000f32; + let log_timescale_increment = max_timescale.ln() / (channels / 2 - 1) as f32; + let inv_timescales: Vec<_> = (0..channels / 2) + .map(|i| (i as f32 * (-log_timescale_increment)).exp()) + .collect(); + let inv_timescales = Tensor::new(inv_timescales.as_slice(), device)?.unsqueeze(0)?; + let arange = Tensor::arange(0, length as u32, device)? + .to_dtype(candle::DType::F32)? + .unsqueeze(1)?; + let sh = (length, channels / 2); + let scaled_time = (arange.broadcast_as(sh)? * inv_timescales.broadcast_as(sh)?)?; + let sincos = Tensor::cat(&[scaled_time.sin()?, scaled_time.cos()?], 1)?; + Ok(sincos) +} + +// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L143 +#[derive(Debug, Clone)] +pub struct AudioEncoder { + conv1: Conv1d, + conv2: Conv1d, + positional_embedding: Tensor, + blocks: Vec, + ln_post: LayerNorm, + span: tracing::Span, + conv1_span: tracing::Span, + conv2_span: tracing::Span, +} + +impl AudioEncoder { + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "audio-encoder"); + let conv1_span = tracing::span!(tracing::Level::TRACE, "conv1"); + let conv2_span = tracing::span!(tracing::Level::TRACE, "conv2"); + let n_state = cfg.d_model; + let n_head = cfg.encoder_attention_heads; + let n_ctx = cfg.max_source_positions; + let cfg1 = Conv1dConfig { + padding: 1, + stride: 1, + groups: 1, + dilation: 1, + cudnn_fwd_algo: None, + }; + let cfg2 = Conv1dConfig { + padding: 1, + stride: 2, + groups: 1, + dilation: 1, + cudnn_fwd_algo: None, + }; + let conv1 = conv1d(cfg.num_mel_bins, n_state, 3, cfg1, vb.pp("conv1"))?; + let conv2 = conv1d(n_state, n_state, 3, cfg2, vb.pp("conv2"))?; + let positional_embedding = sinusoids(n_ctx, n_state, vb.device())?; + let blocks = (0..cfg.encoder_layers) + .map(|i| { + ResidualAttentionBlock::load(n_state, n_head, false, vb.pp(format!("layers.{i}"))) + }) + .collect::>>()?; + let ln_post = layer_norm(n_state, 1e-5, vb.pp("layer_norm"))?; + Ok(Self { + conv1, + conv2, + positional_embedding, + blocks, + ln_post, + conv1_span, + conv2_span, + span, + }) + } + + pub fn forward(&mut self, x: &Tensor, flush_kv_cache: bool) -> Result { + let _enter = self.span.enter(); + let x = { + let _enter = self.conv1_span.enter(); + self.conv1.forward(x)?.gelu()? + }; + let x = { + let _enter = self.conv2_span.enter(); + self.conv2.forward(&x)?.gelu()? + }; + let x = x.transpose(1, 2)?; + let (_bsize, seq_len, _hidden) = x.dims3()?; + let positional_embedding = self.positional_embedding.narrow(0, 0, seq_len)?; + let mut x = x.broadcast_add(&positional_embedding)?; + for block in self.blocks.iter_mut() { + x = block.forward(&x, None, None, flush_kv_cache)? + } + let x = self.ln_post.forward(&x)?; + Ok(x) + } + + pub fn reset_kv_cache(&mut self) { + for block in self.blocks.iter_mut() { + block.reset_kv_cache(); + } + } +} + +// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L176 +#[derive(Debug, Clone)] +pub struct TextDecoder { + token_embedding: Embedding, + positional_embedding: Tensor, + blocks: Vec, + ln: LayerNorm, + mask: Tensor, + span: tracing::Span, + span_final: tracing::Span, +} + +impl TextDecoder { + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "text-decoder"); + let span_final = tracing::span!(tracing::Level::TRACE, "text-decoder-final"); + let n_state = cfg.d_model; + let n_head = cfg.decoder_attention_heads; + let n_ctx = cfg.max_target_positions; + let token_embedding = Embedding::new(cfg.vocab_size, n_state, vb.pp("embed_tokens"))?; + let positional_embedding = vb + .get((n_ctx, n_state), "embed_positions.weight")? + .dequantize(vb.device())?; + let blocks = (0..cfg.decoder_layers) + .map(|i| { + ResidualAttentionBlock::load(n_state, n_head, true, vb.pp(format!("layers.{i}"))) + }) + .collect::>>()?; + let ln = layer_norm(n_state, 1e-5, vb.pp("layer_norm"))?; + let mask: Vec<_> = (0..n_ctx) + .flat_map(|i| (0..n_ctx).map(move |j| if j > i { f32::NEG_INFINITY } else { 0f32 })) + .collect(); + let mask = Tensor::from_vec(mask, (n_ctx, n_ctx), vb.device())?; + Ok(Self { + token_embedding, + positional_embedding, + blocks, + ln, + mask, + span, + span_final, + }) + } + + pub fn forward(&mut self, x: &Tensor, xa: &Tensor, flush_kv_cache: bool) -> Result { + let _enter = self.span.enter(); + let last = x.dim(D::Minus1)?; + let token_embedding = self.token_embedding.forward(x)?; + let positional_embedding = self.positional_embedding.narrow(0, 0, last)?; + let mut x = token_embedding.broadcast_add(&positional_embedding)?; + for block in self.blocks.iter_mut() { + x = block.forward(&x, Some(xa), Some(&self.mask), flush_kv_cache)?; + } + self.ln.forward(&x) + } + + pub fn final_linear(&self, x: &Tensor) -> Result { + let b_size = x.dim(0)?; + let w = self.token_embedding.embeddings().broadcast_left(b_size)?; + let logits = { + let _enter = self.span_final.enter(); + x.matmul(&w.t()?)? + }; + Ok(logits) + } + + pub fn reset_kv_cache(&mut self) { + for block in self.blocks.iter_mut() { + block.reset_kv_cache(); + } + } +} + +// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L221 +#[derive(Debug, Clone)] +pub struct Whisper { + pub encoder: AudioEncoder, + pub decoder: TextDecoder, + pub config: Config, +} + +impl Whisper { + pub fn load(vb: &VarBuilder, config: Config) -> Result { + let encoder = AudioEncoder::load(vb.pp("model.encoder"), &config)?; + let decoder = TextDecoder::load(vb.pp("model.decoder"), &config)?; + Ok(Self { + encoder, + decoder, + config, + }) + } + + pub fn reset_kv_cache(&mut self) { + self.encoder.reset_kv_cache(); + self.decoder.reset_kv_cache(); + } +} diff --git a/patches/candle-transformers/src/models/with_tracing.rs b/patches/candle-transformers/src/models/with_tracing.rs new file mode 100644 index 0000000000..f4706c7e95 --- /dev/null +++ b/patches/candle-transformers/src/models/with_tracing.rs @@ -0,0 +1,195 @@ +use candle::{Module, Result, Tensor}; +use candle_nn::VarBuilder; + +#[derive(Debug, Clone)] +pub struct Embedding { + inner: candle_nn::Embedding, + span: tracing::Span, +} + +impl Embedding { + pub fn new(d1: usize, d2: usize, vb: VarBuilder) -> Result { + let inner = candle_nn::embedding(d1, d2, vb)?; + let span = tracing::span!(tracing::Level::TRACE, "embedding"); + Ok(Self { inner, span }) + } + + pub fn from_weights(weights: Tensor) -> Result { + let (_in_size, out_size) = weights.dims2()?; + let inner = candle_nn::Embedding::new(weights, out_size); + let span = tracing::span!(tracing::Level::TRACE, "embedding"); + Ok(Self { inner, span }) + } + + pub fn embeddings(&self) -> &Tensor { + self.inner.embeddings() + } +} + +impl Module for Embedding { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + self.inner.forward(xs) + } +} + +#[derive(Debug, Clone)] +pub struct Linear { + inner: candle_nn::Linear, + span: tracing::Span, +} + +impl Linear { + pub fn from_weights(weights: Tensor, bias: Option) -> Self { + let inner = candle_nn::Linear::new(weights, bias); + let span = tracing::span!(tracing::Level::TRACE, "linear"); + Self { inner, span } + } +} + +pub fn linear_b(d1: usize, d2: usize, b: bool, vb: VarBuilder) -> Result { + let inner = candle_nn::linear_b(d1, d2, b, vb)?; + let span = tracing::span!(tracing::Level::TRACE, "linear"); + Ok(Linear { inner, span }) +} + +pub fn linear(d1: usize, d2: usize, vb: VarBuilder) -> Result { + let inner = candle_nn::linear(d1, d2, vb)?; + let span = tracing::span!(tracing::Level::TRACE, "linear"); + Ok(Linear { inner, span }) +} + +pub fn linear_no_bias(d1: usize, d2: usize, vb: VarBuilder) -> Result { + let inner = candle_nn::linear_no_bias(d1, d2, vb)?; + let span = tracing::span!(tracing::Level::TRACE, "linear"); + Ok(Linear { inner, span }) +} + +impl Module for Linear { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + self.inner.forward(xs) + } +} + +// Wrap the conv2d op to provide some tracing. +#[derive(Debug, Clone)] +pub struct Conv2d { + inner: candle_nn::Conv2d, + span: tracing::Span, +} + +impl Module for Conv2d { + fn forward(&self, x: &Tensor) -> Result { + let _enter = self.span.enter(); + self.inner.forward(x) + } +} + +pub fn conv2d( + in_channels: usize, + out_channels: usize, + kernel_size: usize, + cfg: candle_nn::Conv2dConfig, + vs: candle_nn::VarBuilder, +) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "conv2d"); + let inner = candle_nn::conv2d(in_channels, out_channels, kernel_size, cfg, vs)?; + Ok(Conv2d { inner, span }) +} + +// QMatMul wrapper adding some tracing. +#[derive(Clone)] +pub struct QMatMul { + inner: candle::quantized::QMatMul, + span: tracing::Span, +} + +impl QMatMul { + pub fn new( + out_dim: usize, + in_dim: usize, + vb: crate::quantized_var_builder::VarBuilder, + ) -> Result { + let ws = vb.get((in_dim, out_dim), "weight")?; + let inner = candle::quantized::QMatMul::from_arc(ws)?; + let span = tracing::span!(tracing::Level::TRACE, "qmatmul"); + Ok(Self { inner, span }) + } + + pub fn from_weights(ws: std::sync::Arc) -> Result { + let inner = candle::quantized::QMatMul::from_arc(ws)?; + let span = tracing::span!(tracing::Level::TRACE, "qmatmul"); + Ok(Self { inner, span }) + } +} + +impl Module for QMatMul { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + self.inner.forward(xs) + } +} + +impl std::fmt::Debug for QMatMul { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "QMatMul") + } +} + +#[derive(Clone, Debug)] +pub struct LayerNorm { + inner: candle_nn::LayerNorm, + span: tracing::Span, +} + +impl LayerNorm { + pub fn new(weight: Tensor, bias: Tensor, eps: f64) -> Self { + let inner = candle_nn::LayerNorm::new(weight, bias, eps); + let span = tracing::span!(tracing::Level::TRACE, "layer-norm"); + Self { inner, span } + } +} + +impl Module for LayerNorm { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + self.inner.forward(xs) + } +} + +pub fn layer_norm>( + size: usize, + c: C, + vb: VarBuilder, +) -> Result { + let inner = candle_nn::layer_norm(size, c, vb)?; + let span = tracing::span!(tracing::Level::TRACE, "layer-norm"); + Ok(LayerNorm { inner, span }) +} + +#[derive(Debug, Clone)] +pub struct RmsNorm { + inner: candle_nn::RmsNorm, + span: tracing::Span, +} + +impl RmsNorm { + pub fn new(size: usize, eps: f64, vb: VarBuilder) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "rms-norm"); + let inner = candle_nn::rms_norm(size, eps, vb)?; + Ok(Self { inner, span }) + } + + pub fn forward_diff(&self, x: &Tensor) -> Result { + let _enter = self.span.enter(); + self.inner.forward_diff(x) + } +} + +impl Module for RmsNorm { + fn forward(&self, x: &Tensor) -> Result { + let _enter = self.span.enter(); + self.inner.forward(x) + } +} diff --git a/patches/candle-transformers/src/models/wuerstchen/attention_processor.rs b/patches/candle-transformers/src/models/wuerstchen/attention_processor.rs new file mode 100644 index 0000000000..0b90cb9d6a --- /dev/null +++ b/patches/candle-transformers/src/models/wuerstchen/attention_processor.rs @@ -0,0 +1,118 @@ +use candle::{Module, Result, Tensor}; +use candle_nn::{linear, Linear, VarBuilder}; + +// A simplified version of: +// https://github.com/huggingface/diffusers/blob/119ad2c3dc8a8fb8446a83f4bf6f20929487b47f/src/diffusers/models/attention_processor.py#L38 +#[derive(Debug)] +pub struct Attention { + to_q: Linear, + to_k: Linear, + to_v: Linear, + to_out: Linear, + heads: usize, + scale: f64, + use_flash_attn: bool, +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +impl Attention { + pub fn new( + query_dim: usize, + heads: usize, + dim_head: usize, + use_flash_attn: bool, + vb: VarBuilder, + ) -> Result { + let inner_dim = dim_head * heads; + let scale = 1.0 / f64::sqrt(dim_head as f64); + let to_q = linear(query_dim, inner_dim, vb.pp("to_q"))?; + let to_k = linear(query_dim, inner_dim, vb.pp("to_k"))?; + let to_v = linear(query_dim, inner_dim, vb.pp("to_v"))?; + let to_out = linear(inner_dim, query_dim, vb.pp("to_out.0"))?; + Ok(Self { + to_q, + to_k, + to_v, + to_out, + scale, + heads, + use_flash_attn, + }) + } + + fn batch_to_head_dim(&self, xs: &Tensor) -> Result { + let (b_size, seq_len, dim) = xs.dims3()?; + xs.reshape((b_size / self.heads, self.heads, seq_len, dim))? + .permute((0, 2, 1, 3))? + .reshape((b_size / self.heads, seq_len, dim * self.heads)) + } + + fn head_to_batch_dim(&self, xs: &Tensor) -> Result { + let (b_size, seq_len, dim) = xs.dims3()?; + xs.reshape((b_size, seq_len, self.heads, dim / self.heads))? + .permute((0, 2, 1, 3))? + .reshape((b_size * self.heads, seq_len, dim / self.heads)) + } + + fn get_attention_scores(&self, query: &Tensor, key: &Tensor) -> Result { + let attn_probs = (query.matmul(&key.t()?)? * self.scale)?; + candle_nn::ops::softmax_last_dim(&attn_probs) + } + + pub fn forward(&self, xs: &Tensor, encoder_hidden_states: &Tensor) -> Result { + let (b_size, channel, h, w) = xs.dims4()?; + let xs = xs.reshape((b_size, channel, h * w))?.t()?; + + let query = self.to_q.forward(&xs)?; + let key = self.to_k.forward(encoder_hidden_states)?; + let value = self.to_v.forward(encoder_hidden_states)?; + + let query = self.head_to_batch_dim(&query)?; + let key = self.head_to_batch_dim(&key)?; + let value = self.head_to_batch_dim(&value)?; + + let xs = if self.use_flash_attn { + let init_dtype = query.dtype(); + let q = query + .to_dtype(candle::DType::F16)? + .unsqueeze(0)? + .transpose(1, 2)?; + let k = key + .to_dtype(candle::DType::F16)? + .unsqueeze(0)? + .transpose(1, 2)?; + let v = value + .to_dtype(candle::DType::F16)? + .unsqueeze(0)? + .transpose(1, 2)?; + flash_attn(&q, &k, &v, self.scale as f32, false)? + .transpose(1, 2)? + .squeeze(0)? + .to_dtype(init_dtype)? + } else { + let attn_prs = self.get_attention_scores(&query, &key)?; + attn_prs.matmul(&value)? + }; + let xs = self.batch_to_head_dim(&xs)?; + + self.to_out + .forward(&xs)? + .t()? + .reshape((b_size, channel, h, w)) + } +} diff --git a/patches/candle-transformers/src/models/wuerstchen/common.rs b/patches/candle-transformers/src/models/wuerstchen/common.rs new file mode 100644 index 0000000000..c89ec919a3 --- /dev/null +++ b/patches/candle-transformers/src/models/wuerstchen/common.rs @@ -0,0 +1,203 @@ +use candle::{DType, Module, Result, Tensor, D}; +use candle_nn::VarBuilder; + +// https://github.com/huggingface/diffusers/blob/19edca82f1ff194c07317369a92b470dbae97f34/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_common.py#L22 +#[derive(Debug)] +pub struct WLayerNorm { + eps: f64, +} + +impl WLayerNorm { + pub fn new(_size: usize) -> Result { + Ok(Self { eps: 1e-6 }) + } +} + +impl Module for WLayerNorm { + fn forward(&self, xs: &Tensor) -> Result { + let xs = xs.permute((0, 2, 3, 1))?; + + let x_dtype = xs.dtype(); + let internal_dtype = match x_dtype { + DType::F16 | DType::BF16 => DType::F32, + d => d, + }; + + let hidden_size = xs.dim(D::Minus1)?; + let xs = xs.to_dtype(internal_dtype)?; + let mean_x = (xs.sum_keepdim(D::Minus1)? / hidden_size as f64)?; + let xs = xs.broadcast_sub(&mean_x)?; + let norm_x = (xs.sqr()?.sum_keepdim(D::Minus1)? / hidden_size as f64)?; + xs.broadcast_div(&(norm_x + self.eps)?.sqrt()?)? + .to_dtype(x_dtype)? + .permute((0, 3, 1, 2)) + } +} + +#[derive(Debug)] +pub struct LayerNormNoWeights { + eps: f64, +} + +impl LayerNormNoWeights { + pub fn new(_size: usize) -> Result { + Ok(Self { eps: 1e-6 }) + } +} + +impl Module for LayerNormNoWeights { + fn forward(&self, xs: &Tensor) -> Result { + let x_dtype = xs.dtype(); + let internal_dtype = match x_dtype { + DType::F16 | DType::BF16 => DType::F32, + d => d, + }; + let hidden_size = xs.dim(D::Minus1)?; + let xs = xs.to_dtype(internal_dtype)?; + let mean_x = (xs.sum_keepdim(D::Minus1)? / hidden_size as f64)?; + let xs = xs.broadcast_sub(&mean_x)?; + let norm_x = (xs.sqr()?.sum_keepdim(D::Minus1)? / hidden_size as f64)?; + xs.broadcast_div(&(norm_x + self.eps)?.sqrt()?)? + .to_dtype(x_dtype) + } +} + +#[derive(Debug)] +pub struct TimestepBlock { + mapper: candle_nn::Linear, +} + +impl TimestepBlock { + pub fn new(c: usize, c_timestep: usize, vb: VarBuilder) -> Result { + let mapper = candle_nn::linear(c_timestep, c * 2, vb.pp("mapper"))?; + Ok(Self { mapper }) + } + + pub fn forward(&self, xs: &Tensor, t: &Tensor) -> Result { + let ab = self + .mapper + .forward(t)? + .unsqueeze(2)? + .unsqueeze(3)? + .chunk(2, 1)?; + xs.broadcast_mul(&(&ab[0] + 1.)?)?.broadcast_add(&ab[1]) + } +} + +#[derive(Debug)] +pub struct GlobalResponseNorm { + gamma: Tensor, + beta: Tensor, +} + +impl GlobalResponseNorm { + pub fn new(dim: usize, vb: VarBuilder) -> Result { + let gamma = vb.get((1, 1, 1, dim), "gamma")?; + let beta = vb.get((1, 1, 1, dim), "beta")?; + Ok(Self { gamma, beta }) + } +} + +impl Module for GlobalResponseNorm { + fn forward(&self, xs: &Tensor) -> Result { + let agg_norm = xs.sqr()?.sum_keepdim((1, 2))?.sqrt()?; + let stand_div_norm = + agg_norm.broadcast_div(&(agg_norm.mean_keepdim(D::Minus1)? + 1e-6)?)?; + xs.broadcast_mul(&stand_div_norm)? + .broadcast_mul(&self.gamma)? + .broadcast_add(&self.beta)? + + xs + } +} + +#[derive(Debug)] +pub struct ResBlock { + depthwise: candle_nn::Conv2d, + norm: WLayerNorm, + channelwise_lin1: candle_nn::Linear, + channelwise_grn: GlobalResponseNorm, + channelwise_lin2: candle_nn::Linear, +} + +impl ResBlock { + pub fn new(c: usize, c_skip: usize, ksize: usize, vb: VarBuilder) -> Result { + let cfg = candle_nn::Conv2dConfig { + padding: ksize / 2, + groups: c, + ..Default::default() + }; + let depthwise = candle_nn::conv2d(c + c_skip, c, ksize, cfg, vb.pp("depthwise"))?; + let norm = WLayerNorm::new(c)?; + let channelwise_lin1 = candle_nn::linear(c, c * 4, vb.pp("channelwise.0"))?; + let channelwise_grn = GlobalResponseNorm::new(c * 4, vb.pp("channelwise.2"))?; + let channelwise_lin2 = candle_nn::linear(c * 4, c, vb.pp("channelwise.4"))?; + Ok(Self { + depthwise, + norm, + channelwise_lin1, + channelwise_grn, + channelwise_lin2, + }) + } + + pub fn forward(&self, xs: &Tensor, x_skip: Option<&Tensor>) -> Result { + let x_res = xs; + let xs = match x_skip { + None => xs.clone(), + Some(x_skip) => Tensor::cat(&[xs, x_skip], 1)?, + }; + let xs = xs + .apply(&self.depthwise)? + .apply(&self.norm)? + .permute((0, 2, 3, 1))?; + let xs = xs + .apply(&self.channelwise_lin1)? + .gelu_erf()? + .apply(&self.channelwise_grn)? + .apply(&self.channelwise_lin2)? + .permute((0, 3, 1, 2))?; + xs + x_res + } +} +use super::attention_processor::Attention; +#[derive(Debug)] +pub struct AttnBlock { + self_attn: bool, + norm: WLayerNorm, + attention: Attention, + kv_mapper_lin: candle_nn::Linear, +} + +impl AttnBlock { + pub fn new( + c: usize, + c_cond: usize, + nhead: usize, + self_attn: bool, + use_flash_attn: bool, + vb: VarBuilder, + ) -> Result { + let norm = WLayerNorm::new(c)?; + let attention = Attention::new(c, nhead, c / nhead, use_flash_attn, vb.pp("attention"))?; + let kv_mapper_lin = candle_nn::linear(c_cond, c, vb.pp("kv_mapper.1"))?; + Ok(Self { + self_attn, + norm, + attention, + kv_mapper_lin, + }) + } + + pub fn forward(&self, xs: &Tensor, kv: &Tensor) -> Result { + let kv = candle_nn::ops::silu(kv)?.apply(&self.kv_mapper_lin)?; + let norm_xs = self.norm.forward(xs)?; + let kv = if self.self_attn { + let (b_size, channel, _, _) = xs.dims4()?; + let norm_xs = norm_xs.reshape((b_size, channel, ()))?.transpose(1, 2)?; + Tensor::cat(&[&norm_xs, &kv], 1)?.contiguous()? + } else { + kv + }; + xs + self.attention.forward(&norm_xs, &kv) + } +} diff --git a/patches/candle-transformers/src/models/wuerstchen/ddpm.rs b/patches/candle-transformers/src/models/wuerstchen/ddpm.rs new file mode 100644 index 0000000000..9e69b86839 --- /dev/null +++ b/patches/candle-transformers/src/models/wuerstchen/ddpm.rs @@ -0,0 +1,103 @@ +use candle::{Result, Tensor}; + +#[derive(Debug, Clone)] +pub struct DDPMWSchedulerConfig { + scaler: f64, + s: f64, +} + +impl Default for DDPMWSchedulerConfig { + fn default() -> Self { + Self { + scaler: 1f64, + s: 0.008f64, + } + } +} + +pub struct DDPMWScheduler { + init_alpha_cumprod: f64, + init_noise_sigma: f64, + timesteps: Vec, + pub config: DDPMWSchedulerConfig, +} + +impl DDPMWScheduler { + pub fn new(inference_steps: usize, config: DDPMWSchedulerConfig) -> Result { + let init_alpha_cumprod = (config.s / (1. + config.s) * std::f64::consts::PI) + .cos() + .powi(2); + let timesteps = (0..=inference_steps) + .map(|i| 1. - i as f64 / inference_steps as f64) + .collect::>(); + Ok(Self { + init_alpha_cumprod, + init_noise_sigma: 1.0, + timesteps, + config, + }) + } + + pub fn timesteps(&self) -> &[f64] { + &self.timesteps + } + + fn alpha_cumprod(&self, t: f64) -> f64 { + let scaler = self.config.scaler; + let s = self.config.s; + let t = if scaler > 1. { + 1. - (1. - t).powf(scaler) + } else if scaler < 1. { + t.powf(scaler) + } else { + t + }; + let alpha_cumprod = ((t + s) / (1. + s) * std::f64::consts::PI * 0.5) + .cos() + .powi(2) + / self.init_alpha_cumprod; + alpha_cumprod.clamp(0.0001, 0.9999) + } + + fn previous_timestep(&self, ts: f64) -> f64 { + let index = self + .timesteps + .iter() + .enumerate() + .map(|(idx, v)| (idx, (v - ts).abs())) + .min_by(|x, y| x.1.total_cmp(&y.1)) + .unwrap() + .0; + self.timesteps[index + 1] + } + + /// Ensures interchangeability with schedulers that need to scale the denoising model input + /// depending on the current timestep. + pub fn scale_model_input(&self, sample: Tensor, _timestep: usize) -> Tensor { + sample + } + + pub fn step(&self, model_output: &Tensor, ts: f64, sample: &Tensor) -> Result { + let prev_t = self.previous_timestep(ts); + + let alpha_cumprod = self.alpha_cumprod(ts); + let alpha_cumprod_prev = self.alpha_cumprod(prev_t); + let alpha = alpha_cumprod / alpha_cumprod_prev; + + let mu = (sample - model_output * ((1. - alpha) / (1. - alpha_cumprod).sqrt()))?; + let mu = (mu * (1. / alpha).sqrt())?; + + let std_noise = mu.randn_like(0., 1.)?; + let std = + std_noise * ((1. - alpha) * (1. - alpha_cumprod_prev) / (1. - alpha_cumprod)).sqrt(); + if prev_t == 0. { + Ok(mu) + } else { + mu + std + } + } + + pub fn init_noise_sigma(&self) -> f64 { + self.init_noise_sigma + } +} diff --git a/patches/candle-transformers/src/models/wuerstchen/diffnext.rs b/patches/candle-transformers/src/models/wuerstchen/diffnext.rs new file mode 100644 index 0000000000..64a48c8a8a --- /dev/null +++ b/patches/candle-transformers/src/models/wuerstchen/diffnext.rs @@ -0,0 +1,396 @@ +use super::common::{AttnBlock, GlobalResponseNorm, LayerNormNoWeights, TimestepBlock, WLayerNorm}; +use candle::{DType, Module, Result, Tensor, D}; +use candle_nn::VarBuilder; + +#[derive(Debug)] +pub struct ResBlockStageB { + depthwise: candle_nn::Conv2d, + norm: WLayerNorm, + channelwise_lin1: candle_nn::Linear, + channelwise_grn: GlobalResponseNorm, + channelwise_lin2: candle_nn::Linear, +} + +impl ResBlockStageB { + pub fn new(c: usize, c_skip: usize, ksize: usize, vb: VarBuilder) -> Result { + let cfg = candle_nn::Conv2dConfig { + groups: c, + padding: ksize / 2, + ..Default::default() + }; + let depthwise = candle_nn::conv2d(c, c, ksize, cfg, vb.pp("depthwise"))?; + let norm = WLayerNorm::new(c)?; + let channelwise_lin1 = candle_nn::linear(c + c_skip, c * 4, vb.pp("channelwise.0"))?; + let channelwise_grn = GlobalResponseNorm::new(4 * c, vb.pp("channelwise.2"))?; + let channelwise_lin2 = candle_nn::linear(c * 4, c, vb.pp("channelwise.4"))?; + Ok(Self { + depthwise, + norm, + channelwise_lin1, + channelwise_grn, + channelwise_lin2, + }) + } + + pub fn forward(&self, xs: &Tensor, x_skip: Option<&Tensor>) -> Result { + let x_res = xs; + let xs = xs.apply(&self.depthwise)?.apply(&self.norm)?; + let xs = match x_skip { + None => xs.clone(), + Some(x_skip) => Tensor::cat(&[&xs, x_skip], 1)?, + }; + let xs = xs + .permute((0, 2, 3, 1))? + .contiguous()? + .apply(&self.channelwise_lin1)? + .gelu()? + .apply(&self.channelwise_grn)? + .apply(&self.channelwise_lin2)? + .permute((0, 3, 1, 2))?; + xs + x_res + } +} + +#[derive(Debug)] +struct SubBlock { + res_block: ResBlockStageB, + ts_block: TimestepBlock, + attn_block: Option, +} + +#[derive(Debug)] +struct DownBlock { + layer_norm: Option, + conv: Option, + sub_blocks: Vec, +} + +#[derive(Debug)] +struct UpBlock { + sub_blocks: Vec, + layer_norm: Option, + conv: Option, +} + +#[derive(Debug)] +pub struct WDiffNeXt { + clip_mapper: candle_nn::Linear, + effnet_mappers: Vec>, + seq_norm: LayerNormNoWeights, + embedding_conv: candle_nn::Conv2d, + embedding_ln: WLayerNorm, + down_blocks: Vec, + up_blocks: Vec, + clf_ln: WLayerNorm, + clf_conv: candle_nn::Conv2d, + c_r: usize, + patch_size: usize, +} + +impl WDiffNeXt { + #[allow(clippy::too_many_arguments)] + pub fn new( + c_in: usize, + c_out: usize, + c_r: usize, + c_cond: usize, + clip_embd: usize, + patch_size: usize, + use_flash_attn: bool, + vb: VarBuilder, + ) -> Result { + const C_HIDDEN: [usize; 4] = [320, 640, 1280, 1280]; + const BLOCKS: [usize; 4] = [4, 4, 14, 4]; + const NHEAD: [usize; 4] = [1, 10, 20, 20]; + const INJECT_EFFNET: [bool; 4] = [false, true, true, true]; + const EFFNET_EMBD: usize = 16; + + let clip_mapper = candle_nn::linear(clip_embd, c_cond, vb.pp("clip_mapper"))?; + let mut effnet_mappers = Vec::with_capacity(2 * INJECT_EFFNET.len()); + let vb_e = vb.pp("effnet_mappers"); + for (i, &inject) in INJECT_EFFNET.iter().enumerate() { + let c = if inject { + Some(candle_nn::conv2d( + EFFNET_EMBD, + c_cond, + 1, + Default::default(), + vb_e.pp(i), + )?) + } else { + None + }; + effnet_mappers.push(c) + } + for (i, &inject) in INJECT_EFFNET.iter().rev().enumerate() { + let c = if inject { + Some(candle_nn::conv2d( + EFFNET_EMBD, + c_cond, + 1, + Default::default(), + vb_e.pp(i + INJECT_EFFNET.len()), + )?) + } else { + None + }; + effnet_mappers.push(c) + } + let seq_norm = LayerNormNoWeights::new(c_cond)?; + let embedding_ln = WLayerNorm::new(C_HIDDEN[0])?; + let embedding_conv = candle_nn::conv2d( + c_in * patch_size * patch_size, + C_HIDDEN[0], + 1, + Default::default(), + vb.pp("embedding.1"), + )?; + + let mut down_blocks = Vec::with_capacity(C_HIDDEN.len()); + for (i, &c_hidden) in C_HIDDEN.iter().enumerate() { + let vb = vb.pp("down_blocks").pp(i); + let (layer_norm, conv, start_layer_i) = if i > 0 { + let layer_norm = WLayerNorm::new(C_HIDDEN[i - 1])?; + let cfg = candle_nn::Conv2dConfig { + stride: 2, + ..Default::default() + }; + let conv = candle_nn::conv2d(C_HIDDEN[i - 1], c_hidden, 2, cfg, vb.pp("0.1"))?; + (Some(layer_norm), Some(conv), 1) + } else { + (None, None, 0) + }; + let mut sub_blocks = Vec::with_capacity(BLOCKS[i]); + let mut layer_i = start_layer_i; + for _j in 0..BLOCKS[i] { + let c_skip = if INJECT_EFFNET[i] { c_cond } else { 0 }; + let res_block = ResBlockStageB::new(c_hidden, c_skip, 3, vb.pp(layer_i))?; + layer_i += 1; + let ts_block = TimestepBlock::new(c_hidden, c_r, vb.pp(layer_i))?; + layer_i += 1; + let attn_block = if i == 0 { + None + } else { + let attn_block = AttnBlock::new( + c_hidden, + c_cond, + NHEAD[i], + true, + use_flash_attn, + vb.pp(layer_i), + )?; + layer_i += 1; + Some(attn_block) + }; + let sub_block = SubBlock { + res_block, + ts_block, + attn_block, + }; + sub_blocks.push(sub_block) + } + let down_block = DownBlock { + layer_norm, + conv, + sub_blocks, + }; + down_blocks.push(down_block) + } + + let mut up_blocks = Vec::with_capacity(C_HIDDEN.len()); + for (i, &c_hidden) in C_HIDDEN.iter().enumerate().rev() { + let vb = vb.pp("up_blocks").pp(C_HIDDEN.len() - 1 - i); + let mut sub_blocks = Vec::with_capacity(BLOCKS[i]); + let mut layer_i = 0; + for j in 0..BLOCKS[i] { + let c_skip = if INJECT_EFFNET[i] { c_cond } else { 0 }; + let c_skip_res = if i < BLOCKS.len() - 1 && j == 0 { + c_hidden + c_skip + } else { + c_skip + }; + let res_block = ResBlockStageB::new(c_hidden, c_skip_res, 3, vb.pp(layer_i))?; + layer_i += 1; + let ts_block = TimestepBlock::new(c_hidden, c_r, vb.pp(layer_i))?; + layer_i += 1; + let attn_block = if i == 0 { + None + } else { + let attn_block = AttnBlock::new( + c_hidden, + c_cond, + NHEAD[i], + true, + use_flash_attn, + vb.pp(layer_i), + )?; + layer_i += 1; + Some(attn_block) + }; + let sub_block = SubBlock { + res_block, + ts_block, + attn_block, + }; + sub_blocks.push(sub_block) + } + let (layer_norm, conv) = if i > 0 { + let layer_norm = WLayerNorm::new(C_HIDDEN[i - 1])?; + let cfg = candle_nn::ConvTranspose2dConfig { + stride: 2, + ..Default::default() + }; + let conv = candle_nn::conv_transpose2d( + c_hidden, + C_HIDDEN[i - 1], + 2, + cfg, + vb.pp(layer_i).pp(1), + )?; + (Some(layer_norm), Some(conv)) + } else { + (None, None) + }; + let up_block = UpBlock { + layer_norm, + conv, + sub_blocks, + }; + up_blocks.push(up_block) + } + + let clf_ln = WLayerNorm::new(C_HIDDEN[0])?; + let clf_conv = candle_nn::conv2d( + C_HIDDEN[0], + 2 * c_out * patch_size * patch_size, + 1, + Default::default(), + vb.pp("clf.1"), + )?; + Ok(Self { + clip_mapper, + effnet_mappers, + seq_norm, + embedding_conv, + embedding_ln, + down_blocks, + up_blocks, + clf_ln, + clf_conv, + c_r, + patch_size, + }) + } + + fn gen_r_embedding(&self, r: &Tensor) -> Result { + const MAX_POSITIONS: usize = 10000; + let r = (r * MAX_POSITIONS as f64)?; + let half_dim = self.c_r / 2; + let emb = (MAX_POSITIONS as f64).ln() / (half_dim - 1) as f64; + let emb = (Tensor::arange(0u32, half_dim as u32, r.device())?.to_dtype(DType::F32)? + * -emb)? + .exp()?; + let emb = r.unsqueeze(1)?.broadcast_mul(&emb.unsqueeze(0)?)?; + let emb = Tensor::cat(&[emb.sin()?, emb.cos()?], 1)?; + let emb = if self.c_r % 2 == 1 { + emb.pad_with_zeros(D::Minus1, 0, 1)? + } else { + emb + }; + emb.to_dtype(r.dtype()) + } + + fn gen_c_embeddings(&self, clip: &Tensor) -> Result { + clip.apply(&self.clip_mapper)?.apply(&self.seq_norm) + } + + pub fn forward( + &self, + xs: &Tensor, + r: &Tensor, + effnet: &Tensor, + clip: Option<&Tensor>, + ) -> Result { + const EPS: f64 = 1e-3; + + let r_embed = self.gen_r_embedding(r)?; + let clip = match clip { + None => None, + Some(clip) => Some(self.gen_c_embeddings(clip)?), + }; + let x_in = xs; + + let mut xs = xs + .apply(&|xs: &_| candle_nn::ops::pixel_unshuffle(xs, self.patch_size))? + .apply(&self.embedding_conv)? + .apply(&self.embedding_ln)?; + + let mut level_outputs = Vec::new(); + for (i, down_block) in self.down_blocks.iter().enumerate() { + if let Some(ln) = &down_block.layer_norm { + xs = xs.apply(ln)? + } + if let Some(conv) = &down_block.conv { + xs = xs.apply(conv)? + } + let skip = match &self.effnet_mappers[i] { + None => None, + Some(m) => { + let effnet = effnet.interpolate2d(xs.dim(D::Minus2)?, xs.dim(D::Minus1)?)?; + Some(m.forward(&effnet)?) + } + }; + for block in down_block.sub_blocks.iter() { + xs = block.res_block.forward(&xs, skip.as_ref())?; + xs = block.ts_block.forward(&xs, &r_embed)?; + if let Some(attn_block) = &block.attn_block { + xs = attn_block.forward(&xs, clip.as_ref().unwrap())?; + } + } + level_outputs.push(xs.clone()) + } + level_outputs.reverse(); + let mut xs = level_outputs[0].clone(); + + for (i, up_block) in self.up_blocks.iter().enumerate() { + let effnet_c = match &self.effnet_mappers[self.down_blocks.len() + i] { + None => None, + Some(m) => { + let effnet = effnet.interpolate2d(xs.dim(D::Minus2)?, xs.dim(D::Minus1)?)?; + Some(m.forward(&effnet)?) + } + }; + for (j, block) in up_block.sub_blocks.iter().enumerate() { + let skip = if j == 0 && i > 0 { + Some(&level_outputs[i]) + } else { + None + }; + let skip = match (skip, effnet_c.as_ref()) { + (Some(skip), Some(effnet_c)) => Some(Tensor::cat(&[skip, effnet_c], 1)?), + (None, Some(skip)) | (Some(skip), None) => Some(skip.clone()), + (None, None) => None, + }; + xs = block.res_block.forward(&xs, skip.as_ref())?; + xs = block.ts_block.forward(&xs, &r_embed)?; + if let Some(attn_block) = &block.attn_block { + xs = attn_block.forward(&xs, clip.as_ref().unwrap())?; + } + } + if let Some(ln) = &up_block.layer_norm { + xs = xs.apply(ln)? + } + if let Some(conv) = &up_block.conv { + xs = xs.apply(conv)? + } + } + + let ab = xs + .apply(&self.clf_ln)? + .apply(&self.clf_conv)? + .apply(&|xs: &_| candle_nn::ops::pixel_shuffle(xs, self.patch_size))? + .chunk(2, 1)?; + let b = ((candle_nn::ops::sigmoid(&ab[1])? * (1. - EPS * 2.))? + EPS)?; + (x_in - &ab[0])? / b + } +} diff --git a/patches/candle-transformers/src/models/wuerstchen/mod.rs b/patches/candle-transformers/src/models/wuerstchen/mod.rs new file mode 100644 index 0000000000..ae42c4a884 --- /dev/null +++ b/patches/candle-transformers/src/models/wuerstchen/mod.rs @@ -0,0 +1,22 @@ +//! Würstchen Efficient Diffusion Model +//! +//! Würstchen is an efficient diffusion model architecture for generating images using +//! a two-stage approach with a small decoder and prior network. +//! +//! - 💻 [GH Link](https://github.com/dome272/Wuerstchen) +//! - 🤗 [HF Link](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py) +//! - 📝 [Paper](https://openreview.net/pdf?id=gU58AyJlYz) +//! +//! ## Example +//! +//!
+//! +//!

"Anthropomorphic cat dressed as a fire fighter"

+//!
+ +pub mod attention_processor; +pub mod common; +pub mod ddpm; +pub mod diffnext; +pub mod paella_vq; +pub mod prior; diff --git a/patches/candle-transformers/src/models/wuerstchen/paella_vq.rs b/patches/candle-transformers/src/models/wuerstchen/paella_vq.rs new file mode 100644 index 0000000000..58f795bbea --- /dev/null +++ b/patches/candle-transformers/src/models/wuerstchen/paella_vq.rs @@ -0,0 +1,211 @@ +use super::common::LayerNormNoWeights; +use candle::{Module, Result, Tensor}; +use candle_nn::VarBuilder; + +#[derive(Debug)] +pub struct MixingResidualBlock { + norm1: LayerNormNoWeights, + depthwise_conv: candle_nn::Conv2d, + norm2: LayerNormNoWeights, + channelwise_lin1: candle_nn::Linear, + channelwise_lin2: candle_nn::Linear, + gammas: Vec, +} + +impl MixingResidualBlock { + pub fn new(inp: usize, embed_dim: usize, vb: VarBuilder) -> Result { + let norm1 = LayerNormNoWeights::new(inp)?; + let norm2 = LayerNormNoWeights::new(inp)?; + let cfg = candle_nn::Conv2dConfig { + groups: inp, + ..Default::default() + }; + let depthwise_conv = candle_nn::conv2d(inp, inp, 3, cfg, vb.pp("depthwise.1"))?; + let channelwise_lin1 = candle_nn::linear(inp, embed_dim, vb.pp("channelwise.0"))?; + let channelwise_lin2 = candle_nn::linear(embed_dim, inp, vb.pp("channelwise.2"))?; + let gammas = vb.get(6, "gammas")?.to_vec1::()?; + Ok(Self { + norm1, + depthwise_conv, + norm2, + channelwise_lin1, + channelwise_lin2, + gammas, + }) + } +} + +impl Module for MixingResidualBlock { + fn forward(&self, xs: &Tensor) -> Result { + let mods = &self.gammas; + let x_temp = xs + .permute((0, 2, 3, 1))? + .apply(&self.norm1)? + .permute((0, 3, 1, 2))? + .affine(1. + mods[0] as f64, mods[1] as f64)?; + let x_temp = candle_nn::ops::replication_pad2d(&x_temp, 1)?; + let xs = (xs + x_temp.apply(&self.depthwise_conv)? * mods[2] as f64)?; + let x_temp = xs + .permute((0, 2, 3, 1))? + .apply(&self.norm2)? + .permute((0, 3, 1, 2))? + .affine(1. + mods[3] as f64, mods[4] as f64)?; + let x_temp = x_temp + .permute((0, 2, 3, 1))? + .contiguous()? + .apply(&self.channelwise_lin1)? + .gelu()? + .apply(&self.channelwise_lin2)? + .permute((0, 3, 1, 2))?; + xs + x_temp * mods[5] as f64 + } +} + +#[derive(Debug)] +pub struct PaellaVQ { + in_block_conv: candle_nn::Conv2d, + out_block_conv: candle_nn::Conv2d, + down_blocks: Vec<(Option, MixingResidualBlock)>, + down_blocks_conv: candle_nn::Conv2d, + down_blocks_bn: candle_nn::BatchNorm, + up_blocks_conv: candle_nn::Conv2d, + up_blocks: Vec<(Vec, Option)>, +} + +impl PaellaVQ { + pub fn new(vb: VarBuilder) -> Result { + const IN_CHANNELS: usize = 3; + const OUT_CHANNELS: usize = 3; + const LATENT_CHANNELS: usize = 4; + const EMBED_DIM: usize = 384; + const BOTTLENECK_BLOCKS: usize = 12; + const C_LEVELS: [usize; 2] = [EMBED_DIM / 2, EMBED_DIM]; + + let in_block_conv = candle_nn::conv2d( + IN_CHANNELS * 4, + C_LEVELS[0], + 1, + Default::default(), + vb.pp("in_block.1"), + )?; + let out_block_conv = candle_nn::conv2d( + C_LEVELS[0], + OUT_CHANNELS * 4, + 1, + Default::default(), + vb.pp("out_block.0"), + )?; + + let mut down_blocks = Vec::new(); + let vb_d = vb.pp("down_blocks"); + let mut d_idx = 0; + for (i, &c_level) in C_LEVELS.iter().enumerate() { + let conv_block = if i > 0 { + let cfg = candle_nn::Conv2dConfig { + padding: 1, + stride: 2, + ..Default::default() + }; + let block = candle_nn::conv2d(C_LEVELS[i - 1], c_level, 4, cfg, vb_d.pp(d_idx))?; + d_idx += 1; + Some(block) + } else { + None + }; + let res_block = MixingResidualBlock::new(c_level, c_level * 4, vb_d.pp(d_idx))?; + d_idx += 1; + down_blocks.push((conv_block, res_block)) + } + let vb_d = vb_d.pp(d_idx); + let down_blocks_conv = candle_nn::conv2d_no_bias( + C_LEVELS[1], + LATENT_CHANNELS, + 1, + Default::default(), + vb_d.pp(0), + )?; + let down_blocks_bn = candle_nn::batch_norm(LATENT_CHANNELS, 1e-5, vb_d.pp(1))?; + + let mut up_blocks = Vec::new(); + let vb_u = vb.pp("up_blocks"); + let mut u_idx = 0; + let up_blocks_conv = candle_nn::conv2d( + LATENT_CHANNELS, + C_LEVELS[1], + 1, + Default::default(), + vb_u.pp(u_idx).pp(0), + )?; + u_idx += 1; + for (i, &c_level) in C_LEVELS.iter().rev().enumerate() { + let mut res_blocks = Vec::new(); + let n_bottleneck_blocks = if i == 0 { BOTTLENECK_BLOCKS } else { 1 }; + for _j in 0..n_bottleneck_blocks { + let res_block = MixingResidualBlock::new(c_level, c_level * 4, vb_u.pp(u_idx))?; + u_idx += 1; + res_blocks.push(res_block) + } + let conv_block = if i < C_LEVELS.len() - 1 { + let cfg = candle_nn::ConvTranspose2dConfig { + padding: 1, + stride: 2, + ..Default::default() + }; + let block = candle_nn::conv_transpose2d( + c_level, + C_LEVELS[C_LEVELS.len() - i - 2], + 4, + cfg, + vb_u.pp(u_idx), + )?; + u_idx += 1; + Some(block) + } else { + None + }; + up_blocks.push((res_blocks, conv_block)) + } + Ok(Self { + in_block_conv, + down_blocks, + down_blocks_conv, + down_blocks_bn, + up_blocks, + up_blocks_conv, + out_block_conv, + }) + } + + pub fn encode(&self, xs: &Tensor) -> Result { + let mut xs = candle_nn::ops::pixel_unshuffle(xs, 2)?.apply(&self.in_block_conv)?; + for down_block in self.down_blocks.iter() { + if let Some(conv) = &down_block.0 { + xs = xs.apply(conv)? + } + xs = xs.apply(&down_block.1)? + } + xs.apply(&self.down_blocks_conv)? + .apply_t(&self.down_blocks_bn, false) + } + + pub fn decode(&self, xs: &Tensor) -> Result { + // TODO: quantizer if we want to support `force_not_quantize=False`. + let mut xs = xs.apply(&self.up_blocks_conv)?; + for up_block in self.up_blocks.iter() { + for b in up_block.0.iter() { + xs = xs.apply(b)?; + } + if let Some(conv) = &up_block.1 { + xs = xs.apply(conv)? + } + } + xs.apply(&self.out_block_conv)? + .apply(&|xs: &_| candle_nn::ops::pixel_shuffle(xs, 2)) + } +} + +impl Module for PaellaVQ { + fn forward(&self, xs: &Tensor) -> Result { + self.decode(&self.encode(xs)?) + } +} diff --git a/patches/candle-transformers/src/models/wuerstchen/prior.rs b/patches/candle-transformers/src/models/wuerstchen/prior.rs new file mode 100644 index 0000000000..97ccf0e21d --- /dev/null +++ b/patches/candle-transformers/src/models/wuerstchen/prior.rs @@ -0,0 +1,103 @@ +use super::common::{AttnBlock, ResBlock, TimestepBlock}; +use candle::{DType, Result, Tensor, D}; +use candle_nn::VarBuilder; + +#[derive(Debug)] +struct Block { + res_block: ResBlock, + ts_block: TimestepBlock, + attn_block: AttnBlock, +} + +#[derive(Debug)] +pub struct WPrior { + projection: candle_nn::Conv2d, + cond_mapper_lin1: candle_nn::Linear, + cond_mapper_lin2: candle_nn::Linear, + blocks: Vec, + out_ln: super::common::WLayerNorm, + out_conv: candle_nn::Conv2d, + c_r: usize, +} + +impl WPrior { + #[allow(clippy::too_many_arguments)] + pub fn new( + c_in: usize, + c: usize, + c_cond: usize, + c_r: usize, + depth: usize, + nhead: usize, + use_flash_attn: bool, + vb: VarBuilder, + ) -> Result { + let projection = candle_nn::conv2d(c_in, c, 1, Default::default(), vb.pp("projection"))?; + let cond_mapper_lin1 = candle_nn::linear(c_cond, c, vb.pp("cond_mapper.0"))?; + let cond_mapper_lin2 = candle_nn::linear(c, c, vb.pp("cond_mapper.2"))?; + let out_ln = super::common::WLayerNorm::new(c)?; + let out_conv = candle_nn::conv2d(c, c_in * 2, 1, Default::default(), vb.pp("out.1"))?; + let mut blocks = Vec::with_capacity(depth); + for index in 0..depth { + let res_block = ResBlock::new(c, 0, 3, vb.pp(format!("blocks.{}", 3 * index)))?; + let ts_block = TimestepBlock::new(c, c_r, vb.pp(format!("blocks.{}", 3 * index + 1)))?; + let attn_block = AttnBlock::new( + c, + c, + nhead, + true, + use_flash_attn, + vb.pp(format!("blocks.{}", 3 * index + 2)), + )?; + blocks.push(Block { + res_block, + ts_block, + attn_block, + }) + } + Ok(Self { + projection, + cond_mapper_lin1, + cond_mapper_lin2, + blocks, + out_ln, + out_conv, + c_r, + }) + } + + pub fn gen_r_embedding(&self, r: &Tensor) -> Result { + const MAX_POSITIONS: usize = 10000; + let r = (r * MAX_POSITIONS as f64)?; + let half_dim = self.c_r / 2; + let emb = (MAX_POSITIONS as f64).ln() / (half_dim - 1) as f64; + let emb = (Tensor::arange(0u32, half_dim as u32, r.device())?.to_dtype(DType::F32)? + * -emb)? + .exp()?; + let emb = r.unsqueeze(1)?.broadcast_mul(&emb.unsqueeze(0)?)?; + let emb = Tensor::cat(&[emb.sin()?, emb.cos()?], 1)?; + let emb = if self.c_r % 2 == 1 { + emb.pad_with_zeros(D::Minus1, 0, 1)? + } else { + emb + }; + emb.to_dtype(r.dtype()) + } + + pub fn forward(&self, xs: &Tensor, r: &Tensor, c: &Tensor) -> Result { + let x_in = xs; + let mut xs = xs.apply(&self.projection)?; + let c_embed = c + .apply(&self.cond_mapper_lin1)? + .apply(&|xs: &_| candle_nn::ops::leaky_relu(xs, 0.2))? + .apply(&self.cond_mapper_lin2)?; + let r_embed = self.gen_r_embedding(r)?; + for block in self.blocks.iter() { + xs = block.res_block.forward(&xs, None)?; + xs = block.ts_block.forward(&xs, &r_embed)?; + xs = block.attn_block.forward(&xs, &c_embed)?; + } + let ab = xs.apply(&self.out_ln)?.apply(&self.out_conv)?.chunk(2, 1)?; + (x_in - &ab[0])? / ((&ab[1] - 1.)?.abs()? + 1e-5) + } +} diff --git a/patches/candle-transformers/src/models/xlm_roberta.rs b/patches/candle-transformers/src/models/xlm_roberta.rs new file mode 100644 index 0000000000..ee94f0687d --- /dev/null +++ b/patches/candle-transformers/src/models/xlm_roberta.rs @@ -0,0 +1,538 @@ +use crate::models::with_tracing::{linear, Linear}; +use candle::{DType, Module, Result, Tensor}; +use candle_nn::{ + embedding, layer_norm, ops::softmax_last_dim, Activation, Embedding, LayerNorm, VarBuilder, +}; + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + pub hidden_size: usize, + pub layer_norm_eps: f64, + pub attention_probs_dropout_prob: f32, + pub hidden_dropout_prob: f32, + pub num_attention_heads: usize, + pub position_embedding_type: String, + pub intermediate_size: usize, + pub hidden_act: Activation, + pub num_hidden_layers: usize, + pub vocab_size: usize, + pub max_position_embeddings: usize, + pub type_vocab_size: usize, + pub pad_token_id: u32, +} + +struct XLMRobertaEmbeddings { + word_embeddings: Embedding, + position_embeddings: Option, + token_type_embeddings: Embedding, + layer_norm: LayerNorm, + padding_idx: u32, + span: tracing::Span, +} + +impl XLMRobertaEmbeddings { + fn load(vb: VarBuilder, config: &Config) -> Result { + let word_embeddings = embedding( + config.vocab_size, + config.hidden_size, + vb.pp("word_embeddings"), + )?; + let position_embeddings = embedding( + config.max_position_embeddings, + config.hidden_size, + vb.pp("position_embeddings"), + )?; + let token_type_embeddings = embedding( + config.type_vocab_size, + config.hidden_size, + vb.pp("token_type_embeddings"), + )?; + let layer_norm = layer_norm( + config.hidden_size, + config.layer_norm_eps, + vb.pp("LayerNorm"), + )?; + Ok(Self { + word_embeddings, + position_embeddings: Some(position_embeddings), + token_type_embeddings, + layer_norm, + padding_idx: config.pad_token_id, + span: tracing::span!(tracing::Level::TRACE, "embeddings"), + }) + } + + fn forward(&self, input_ids: &Tensor, token_type_ids: &Tensor) -> Result { + let _enter = self.span.enter(); + let (_bsize, _) = input_ids.dims2()?; + let input_embeddings = self.word_embeddings.forward(input_ids)?; + let token_type_embeddings = self.token_type_embeddings.forward(token_type_ids)?; + let mut embeddings = (&input_embeddings + token_type_embeddings)?; + if let Some(position_embeddings) = &self.position_embeddings { + let mask = input_ids + .ne(self.padding_idx)? + .to_dtype(input_embeddings.dtype())?; + let cumsum = mask.cumsum(1)?; + let position_ids = (cumsum * mask)? + .broadcast_add( + &Tensor::try_from(self.padding_idx)? + .to_dtype(input_embeddings.dtype())? + .to_device(input_embeddings.device())?, + )? + .to_dtype(candle::DType::U32)?; + embeddings = embeddings.broadcast_add(&position_embeddings.forward(&position_ids)?)?; + } + let embeddings = self.layer_norm.forward(&embeddings)?; + Ok(embeddings) + } +} + +struct XLMRobertaSelfAttention { + num_attention_heads: usize, + attention_head_size: usize, + all_head_size: usize, + query: Linear, + key: Linear, + value: Linear, +} + +impl XLMRobertaSelfAttention { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let attention_head_size = cfg.hidden_size / cfg.num_attention_heads; + let all_head_size = cfg.num_attention_heads * attention_head_size; + Ok(Self { + num_attention_heads: cfg.num_attention_heads, + attention_head_size, + all_head_size, + query: linear(cfg.hidden_size, all_head_size, vb.pp("query"))?, + key: linear(cfg.hidden_size, all_head_size, vb.pp("key"))?, + value: linear(cfg.hidden_size, all_head_size, vb.pp("value"))?, + }) + } + + fn transpose_for_scores(&self, x: &Tensor) -> Result { + let mut new_x_shape = x.dims().to_vec(); + new_x_shape[2] = self.num_attention_heads; + new_x_shape.push(self.attention_head_size); + let x = x.reshape(new_x_shape)?; + x.permute((0, 2, 1, 3))?.contiguous() + } + + fn forward( + &self, + hidden_states: &Tensor, + encoder_hidden_states: Option<&Tensor>, + attention_mask: &Tensor, + past_key_value: Option<(&Tensor, &Tensor)>, + encoder_attention_mask: Option<&Tensor>, + ) -> Result { + let mixed_query_layer = self.query.forward(hidden_states)?; + let is_cross_attention = encoder_hidden_states.is_some(); + let (key_layer, value_layer, attention_mask) = if is_cross_attention { + if let Some((past_key, past_value)) = past_key_value { + let key_layer = past_key.clone(); + let value_layer = past_value.clone(); + let attention_mask = encoder_attention_mask.unwrap().clone(); + (key_layer, value_layer, Some(attention_mask)) + } else { + let key_layer = + self.transpose_for_scores(&self.key.forward(encoder_hidden_states.unwrap())?)?; + let value_layer = self + .transpose_for_scores(&self.value.forward(encoder_hidden_states.unwrap())?)?; + let attention_mask = encoder_attention_mask.unwrap(); + (key_layer, value_layer, Some(attention_mask.clone())) + } + } else if let Some((past_key, past_value)) = past_key_value { + let mut key_layer = self.transpose_for_scores(&self.key.forward(hidden_states)?)?; + let mut value_layer = self.transpose_for_scores(&self.value.forward(hidden_states)?)?; + key_layer = Tensor::cat(&[past_key.clone(), key_layer], 2)?; + value_layer = Tensor::cat(&[past_value.clone(), value_layer], 2)?; + (key_layer, value_layer, Some(attention_mask.clone())) + } else { + let key_layer = self.transpose_for_scores(&self.key.forward(hidden_states)?)?; + let value_layer = self.transpose_for_scores(&self.value.forward(hidden_states)?)?; + (key_layer, value_layer, Some(attention_mask.clone())) + }; + + let query_layer = self.transpose_for_scores(&mixed_query_layer)?; + let mut attention_scores = query_layer.matmul(&key_layer.transpose(2, 3)?)?; + let scale = 1f64 / f64::sqrt(self.attention_head_size as f64); + + attention_scores = (attention_scores * scale)?; + attention_scores = match attention_mask { + None => attention_scores, + Some(mask) => { + attention_scores.broadcast_add(&mask.to_dtype(attention_scores.dtype())?)? + } + }; + let attention_probs = softmax_last_dim(&attention_scores)?; + + let context_layer = attention_probs + .matmul(&value_layer)? + .permute((0, 2, 1, 3))? + .contiguous()?; + let mut new_context_layer_shape = + context_layer.dims()[..context_layer.dims().len() - 2].to_vec(); + new_context_layer_shape.push(self.all_head_size); + let context_layer = context_layer.reshape(new_context_layer_shape)?; + + Ok(context_layer) + } +} + +struct XLMRobertaSelfOutput { + dense: Linear, + layernorm: LayerNorm, +} + +impl XLMRobertaSelfOutput { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.hidden_size, cfg.hidden_size, vb.pp("dense"))?; + let layernorm = + candle_nn::layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("LayerNorm"))?; + Ok(Self { dense, layernorm }) + } + + fn forward(&self, hidden_states: &Tensor, input_tensor: &Tensor) -> Result { + let hidden_states = self.dense.forward(hidden_states)?; + let hidden_states = self.layernorm.forward(&(hidden_states + input_tensor)?)?; + Ok(hidden_states) + } +} + +struct XLMRobertaAttention { + output: XLMRobertaSelfOutput, + self_attention: XLMRobertaSelfAttention, +} + +impl XLMRobertaAttention { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let output = XLMRobertaSelfOutput::new(cfg, vb.pp("output"))?; + let self_attention = XLMRobertaSelfAttention::new(cfg, vb.pp("self"))?; + Ok(Self { + output, + self_attention, + }) + } + + fn forward( + &self, + hidden_states: &Tensor, + attention_mask: &Tensor, + encoder_hidden_states: Option<&Tensor>, + encoder_attention_mask: Option<&Tensor>, + past_key_value: Option<(&Tensor, &Tensor)>, + ) -> Result<(Tensor, Tensor)> { + let self_outputs = self.self_attention.forward( + hidden_states, + encoder_hidden_states, + attention_mask, + past_key_value, + encoder_attention_mask, + )?; + let attention_output = self.output.forward(&self_outputs, hidden_states)?; + Ok((attention_output, self_outputs)) + } +} + +struct XLMRobertaOutput { + dense: Linear, + layernorm: LayerNorm, +} + +impl XLMRobertaOutput { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.intermediate_size, cfg.hidden_size, vb.pp("dense"))?; + let layernorm = + candle_nn::layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("LayerNorm"))?; + Ok(Self { dense, layernorm }) + } + + fn forward(&self, hidden_states: &Tensor, input_tensor: &Tensor) -> Result { + let hidden_states = self.dense.forward(hidden_states)?; + let hidden_states = self.layernorm.forward(&(hidden_states + input_tensor)?)?; + Ok(hidden_states) + } +} + +struct XLMRobertaIntermediate { + dense: Linear, + intermediate_act_fn: Activation, +} + +impl XLMRobertaIntermediate { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.hidden_size, cfg.intermediate_size, vb.pp("dense"))?; + let intermediate_act_fn = cfg.hidden_act; + Ok(Self { + dense, + intermediate_act_fn, + }) + } + + fn forward(&self, hidden_states: &Tensor) -> Result { + let hidden_states = self.dense.forward(hidden_states)?; + let hidden_states = self.intermediate_act_fn.forward(&hidden_states)?; + Ok(hidden_states) + } +} + +struct XLMRobertaLayer { + attention: XLMRobertaAttention, + intermediate: XLMRobertaIntermediate, + output: XLMRobertaOutput, +} + +impl XLMRobertaLayer { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let attention = XLMRobertaAttention::new(cfg, vb.pp("attention"))?; + let intermediate = XLMRobertaIntermediate::new(cfg, vb.pp("intermediate"))?; + let output = XLMRobertaOutput::new(cfg, vb.pp("output"))?; + Ok(Self { + attention, + intermediate, + output, + }) + } + + fn forward( + &self, + hidden_states: &Tensor, + attention_mask: &Tensor, + encoder_hidden_states: Option<&Tensor>, + encoder_attention_mask: Option<&Tensor>, + past_key_value: Option<(&Tensor, &Tensor)>, + ) -> Result<(Tensor, Tensor)> { + let self_attention_outputs = self.attention.forward( + hidden_states, + attention_mask, + encoder_hidden_states, + encoder_attention_mask, + past_key_value, + )?; + let attention_output = self_attention_outputs.0; + let outputs = self_attention_outputs.1; + let intermediate_output = self.intermediate.forward(&attention_output)?; + let layer_output = self + .output + .forward(&intermediate_output, &attention_output)?; + Ok((layer_output, outputs)) + } +} + +struct XLMRobertaEncoder { + layers: Vec, +} + +impl XLMRobertaEncoder { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let layers = (0..cfg.num_hidden_layers) + .map(|i| XLMRobertaLayer::new(cfg, vb.pp(format!("layer.{i}")))) + .collect::>>()?; + Ok(Self { layers }) + } + + fn forward( + &self, + hidden_states: &Tensor, + attention_mask: &Tensor, + encoder_hidden_states: Option<&Tensor>, + encoder_attention_mask: Option<&Tensor>, + past_key_value: Option<(&Tensor, &Tensor)>, + ) -> Result { + let mut hidden_states = hidden_states.clone(); + for layer_module in self.layers.iter() { + let layer_outputs = layer_module.forward( + &hidden_states, + attention_mask, + encoder_hidden_states, + encoder_attention_mask, + past_key_value, + )?; + hidden_states = layer_outputs.0; + } + Ok(hidden_states) + } +} + +pub struct XLMRobertaModel { + encoder: XLMRobertaEncoder, + embeddings: XLMRobertaEmbeddings, +} + +impl XLMRobertaModel { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let encoder = XLMRobertaEncoder::new(cfg, vb.pp("encoder"))?; + let embeddings = XLMRobertaEmbeddings::load(vb.pp("embeddings"), cfg)?; + Ok(Self { + encoder, + embeddings, + }) + } + + pub fn forward( + &self, + input_ids: &Tensor, + attention_mask: &Tensor, + token_type_ids: &Tensor, + past_key_value: Option<(&Tensor, &Tensor)>, + encoder_hidden_states: Option<&Tensor>, + encoder_attention_mask: Option<&Tensor>, + ) -> Result { + let hidden_states = self.embeddings.forward(input_ids, token_type_ids)?; + let attention_mask = prepare_4d_attention_mask(attention_mask, DType::F32, None)? + .to_device(hidden_states.device())?; + let hidden_states = self.encoder.forward( + &hidden_states, + &attention_mask, + encoder_hidden_states, + encoder_attention_mask, + past_key_value, + )?; + Ok(hidden_states) + } +} + +struct XLMRobertaLMHead { + dense: Linear, + layer_norm: LayerNorm, +} + +impl XLMRobertaLMHead { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.hidden_size, cfg.hidden_size, vb.pp("dense"))?; + let layer_norm = + candle_nn::layer_norm(cfg.hidden_size, cfg.layer_norm_eps, vb.pp("layer_norm"))?; + Ok(Self { dense, layer_norm }) + } + + fn forward(&self, hidden_states: &Tensor, shared_embeddings: &Tensor) -> Result { + let hidden_states = self.dense.forward(hidden_states)?; + let hidden_states = candle_nn::Activation::Gelu.forward(&hidden_states)?; + let hidden_states = self.layer_norm.forward(&hidden_states)?; + let hidden_states = hidden_states.broadcast_matmul(shared_embeddings)?; + Ok(hidden_states) + } +} + +pub struct XLMRobertaForMaskedLM { + roberta: XLMRobertaModel, + lm_head: XLMRobertaLMHead, +} + +impl XLMRobertaForMaskedLM { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let roberta = XLMRobertaModel::new(cfg, vb.pp("roberta"))?; + let lm_head = XLMRobertaLMHead::new(cfg, vb.pp("lm_head"))?; + Ok(Self { roberta, lm_head }) + } + + pub fn forward( + &self, + input_ids: &Tensor, + attention_mask: &Tensor, + token_type_ids: &Tensor, + past_key_value: Option<(&Tensor, &Tensor)>, + encoder_hidden_states: Option<&Tensor>, + encoder_attention_mask: Option<&Tensor>, + ) -> Result { + let hidden_states = self.roberta.forward( + input_ids, + attention_mask, + token_type_ids, + past_key_value, + encoder_hidden_states, + encoder_attention_mask, + )?; + let lm_logits = self.lm_head.forward( + &hidden_states, + &self + .roberta + .embeddings + .word_embeddings + .embeddings() + .t()? + .unsqueeze(0)?, + )?; + Ok(lm_logits) + } +} + +struct XLMRobertaClassificationHead { + dense: Linear, + out_proj: Linear, +} + +impl XLMRobertaClassificationHead { + fn new(num_labels: usize, cfg: &Config, vb: VarBuilder) -> Result { + let dense = linear(cfg.hidden_size, cfg.hidden_size, vb.pp("dense"))?; + let out_proj = linear(cfg.hidden_size, num_labels, vb.pp("out_proj"))?; + Ok(Self { dense, out_proj }) + } + + fn forward(&self, hidden_states: &Tensor) -> Result { + let cls_states = hidden_states.get_on_dim(1, 0)?.contiguous()?; + let hidden_states = self.dense.forward(&cls_states)?; + // The activation used in the classification head is tanh, as per the original + // implementation. + // https://github.com/huggingface/transformers/blob/6e3063422c4b1c014aa60c32b9254fd2902f0f28/src/transformers/models/xlm_roberta/modeling_xlm_roberta.py#L1454 + let hidden_states = self.out_proj.forward(&hidden_states.tanh()?)?; + Ok(hidden_states) + } +} + +pub struct XLMRobertaForSequenceClassification { + roberta: XLMRobertaModel, + classifier: XLMRobertaClassificationHead, +} + +impl XLMRobertaForSequenceClassification { + pub fn new(num_labels: usize, cfg: &Config, vb: VarBuilder) -> Result { + let roberta = XLMRobertaModel::new(cfg, vb.pp("roberta"))?; + let classifier = XLMRobertaClassificationHead::new(num_labels, cfg, vb.pp("classifier"))?; + Ok(Self { + roberta, + classifier, + }) + } + + pub fn forward( + &self, + input_ids: &Tensor, + attention_mask: &Tensor, + token_type_ids: &Tensor, + ) -> Result { + let hidden_states = + self.roberta + .forward(input_ids, attention_mask, token_type_ids, None, None, None)?; + self.classifier.forward(&hidden_states) + } +} + +fn prepare_4d_attention_mask( + mask: &Tensor, + dtype: DType, + tgt_len: Option, +) -> Result { + let bsz = mask.dim(0)?; + let src_len = mask.dim(1)?; + let tgt_len = tgt_len.unwrap_or(src_len); + + let expanded_mask = mask + .unsqueeze(1)? + .unsqueeze(2)? + .expand((bsz, 1, tgt_len, src_len))? + .to_dtype(dtype)?; + + let inverted_mask = (1.0 - expanded_mask)?; + + (inverted_mask * get_dtype_min_val(dtype))?.to_dtype(dtype) +} + +fn get_dtype_min_val(dtype: DType) -> f64 { + match dtype { + DType::F32 => f32::MIN as f64, + DType::F64 => f64::MIN, + _ => panic!("Unsupported data type"), + } +} diff --git a/patches/candle-transformers/src/models/yi.rs b/patches/candle-transformers/src/models/yi.rs new file mode 100644 index 0000000000..8a2fb111be --- /dev/null +++ b/patches/candle-transformers/src/models/yi.rs @@ -0,0 +1,364 @@ +//! Yi model implementation. +//! +//! This candle implementation uses a pre-trained Yi decoder-only large language model for inference. +//! The model was trained by 01.AI and follows a standard transformer architecture similar to LLaMA. +//! +//! Original code: +//! - 💻 [Yi Model](https://huggingface.co/01-ai/Yi-6B) +//! - 💻 [Yi Modeling Code](https://huggingface.co/01-ai/Yi-6B/blob/main/modeling_yi.py) +//! - 📝 [Technical Report](https://arxiv.org/abs/2403.04652) Yi: Open Foundation Models by 01.AI +//! +//! Key characteristics: +//! - Multi-head attention with rotary positional embeddings +//! - RMS normalization +//! - SwiGLU activation in feed-forward layers +//! - Grouped-query attention for efficient inference +//! + +use crate::models::with_tracing::{linear_no_bias, Linear, RmsNorm}; +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::{Activation, VarBuilder}; +use std::sync::Arc; + +#[derive(Debug, Clone, PartialEq)] +pub struct Config { + pub(crate) vocab_size: usize, + pub(crate) hidden_size: usize, + pub(crate) intermediate_size: usize, + pub(crate) num_hidden_layers: usize, + pub(crate) num_attention_heads: usize, + pub(crate) num_key_value_heads: usize, + pub(crate) hidden_act: Activation, + pub(crate) max_position_embeddings: usize, + pub(crate) rms_norm_eps: f64, + pub(crate) rope_theta: f64, +} + +impl Config { + pub fn config_6b() -> Self { + Self { + vocab_size: 64000, + hidden_size: 4096, + intermediate_size: 11008, + num_hidden_layers: 32, + num_attention_heads: 32, + num_key_value_heads: 4, + hidden_act: Activation::Silu, + max_position_embeddings: 4096, + rms_norm_eps: 1e-5, + rope_theta: 5_000_000., + } + } + + pub fn config_34b() -> Self { + Self { + vocab_size: 64000, + hidden_size: 7168, + intermediate_size: 20480, + num_hidden_layers: 60, + num_attention_heads: 56, + num_key_value_heads: 8, + hidden_act: Activation::Silu, + max_position_embeddings: 4096, + rms_norm_eps: 1e-5, + rope_theta: 5_000_000., + } + } +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +fn rotate_half(xs: &Tensor) -> Result { + let last_dim = xs.dim(D::Minus1)?; + let xs1 = xs.narrow(D::Minus1, 0, last_dim / 2)?; + let xs2 = xs.narrow(D::Minus1, last_dim / 2, last_dim - last_dim / 2)?; + Tensor::cat(&[&xs2.neg()?, &xs1], D::Minus1) +} + +impl RotaryEmbedding { + fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let dim = cfg.hidden_size / cfg.num_attention_heads; + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / 10000f32.powf(i as f32 / dim as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + let freqs = Tensor::cat(&[&freqs, &freqs], D::Minus1)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let cos = cos.unsqueeze(0)?.unsqueeze(0)?; // (1, 1, seq_len, dim) + let sin = sin.unsqueeze(0)?.unsqueeze(0)?; // (1, 1, seq_len, dim) + let q_embed = (q.broadcast_mul(&cos)? + rotate_half(q)?.broadcast_mul(&sin))?; + let k_embed = (k.broadcast_mul(&cos)? + rotate_half(k)?.broadcast_mul(&sin))?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: Activation, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + let gate_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("gate_proj"))?; + let up_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("up_proj"))?; + let down_proj = linear_no_bias(intermediate_sz, hidden_sz, vb.pp("down_proj"))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + act_fn: cfg.hidden_act, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result { + let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = xs.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + hidden_size: usize, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + let head_dim = hidden_sz / num_heads; + let q_proj = linear_no_bias(hidden_sz, num_heads * head_dim, vb.pp("q_proj"))?; + let k_proj = linear_no_bias(hidden_sz, num_kv_heads * head_dim, vb.pp("k_proj"))?; + let v_proj = linear_no_bias(hidden_sz, num_kv_heads * head_dim, vb.pp("v_proj"))?; + let o_proj = linear_no_bias(num_heads * head_dim, hidden_sz, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + hidden_size: hidden_sz, + rotary_emb, + kv_cache: None, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let (query_states, key_states) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states, seqlen_offset)?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?; + let value_states = crate::utils::repeat_kv(value_states, self.num_kv_groups)?; + + let attn_output = { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + }; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.hidden_size))? + .apply(&self.o_proj) + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: MLP, + ln1: RmsNorm, + ln2: RmsNorm, +} + +impl DecoderLayer { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let self_attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let ln1 = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let ln2 = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + ln1, + ln2, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.ln1.forward(xs)?; + let xs = self.self_attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.ln2)?.apply(&self.mlp)?; + residual + xs + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: Linear, + device: Device, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb_m.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(rotary_emb.clone(), cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb_m.pp("norm"))?; + let lm_head = linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))?; + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + // Sliding window mask? + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { f32::NEG_INFINITY } else { 0. })) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let (b_size, seq_len) = input_ids.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(b_size, seq_len, seqlen_offset)?; + Some(mask) + }; + let mut xs = self.embed_tokens.forward(input_ids)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)? + } + xs.narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head) + } +} diff --git a/patches/candle-transformers/src/models/z_image/mod.rs b/patches/candle-transformers/src/models/z_image/mod.rs new file mode 100644 index 0000000000..ddb454721d --- /dev/null +++ b/patches/candle-transformers/src/models/z_image/mod.rs @@ -0,0 +1,43 @@ +/* + * @Author: SpenserCai + * @Date: 2026-01-02 11:35:48 + * @version: + * @LastEditors: SpenserCai + * @LastEditTime: 2026-01-02 11:48:26 + * @Description: file content + */ +//! Z-Image Model +//! +//! Z-Image is a text-to-image generation model from Alibaba using Flow Matching. +//! +//! - 🤗 [Hugging Face Model](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) +//! - [Official Website](https://z-image-turbo.org/) +//! +//! # Example +//! +//! ```bash +//! cargo run --features metal --example z_image --release -- \ +//! --prompt "A beautiful landscape" --height 1024 --width 1024 +//! ``` +//! +//! # Architecture +//! +//! - Transformer: ~24B parameters, 30 main layers + 2 noise_refiner + 2 context_refiner +//! - Text Encoder: Qwen3 (hidden_size=2560, 36 layers) +//! - VAE: AutoencoderKL (diffusers format) +//! - Scheduler: FlowMatchEulerDiscreteScheduler (shift=3.0) + +pub mod preprocess; +pub mod sampling; +pub mod scheduler; +pub mod text_encoder; +pub mod transformer; +pub mod vae; + +// Re-export main types +pub use preprocess::{prepare_inputs, PreparedInputs}; +pub use sampling::{get_noise, get_schedule, postprocess_image}; +pub use scheduler::{calculate_shift, FlowMatchEulerDiscreteScheduler, SchedulerConfig}; +pub use text_encoder::{TextEncoderConfig, ZImageTextEncoder}; +pub use transformer::{Config, ZImageTransformer2DModel}; +pub use vae::{AutoEncoderKL, VaeConfig}; diff --git a/patches/candle-transformers/src/models/z_image/preprocess.rs b/patches/candle-transformers/src/models/z_image/preprocess.rs new file mode 100644 index 0000000000..7b7f8755c8 --- /dev/null +++ b/patches/candle-transformers/src/models/z_image/preprocess.rs @@ -0,0 +1,169 @@ +//! Input preprocessing utilities for Z-Image +//! +//! Provides padding and mask construction to convert variable-length inputs +//! into fixed-shape batch tensors. + +use candle::{DType, Device, Result, Tensor}; + +use super::transformer::SEQ_MULTI_OF; + +/// Preprocessed inputs structure +#[derive(Debug, Clone)] +pub struct PreparedInputs { + /// Latent tensor (B, C, 1, H, W) + pub latents: Tensor, + /// Padded caption features (B, max_text_len, dim) + pub cap_feats: Tensor, + /// Caption attention mask (B, max_text_len), 1=valid, 0=padding + pub cap_mask: Tensor, + /// Original text lengths for each sample + pub text_lengths: Vec, +} + +/// Compute padding length to align to SEQ_MULTI_OF +#[inline] +pub fn compute_padding_len(ori_len: usize) -> usize { + (SEQ_MULTI_OF - (ori_len % SEQ_MULTI_OF)) % SEQ_MULTI_OF +} + +/// Pad variable-length text embeddings to uniform length +/// +/// # Arguments +/// * `text_embeddings` - Variable-length text embeddings, each of shape (seq_len, dim) +/// * `pad_value` - Padding value (typically 0.0) +/// * `device` - Device +/// +/// # Returns +/// * Padded tensor (B, max_len, dim) +/// * Attention mask (B, max_len), 1=valid, 0=padding +/// * Original lengths +pub fn pad_text_embeddings( + text_embeddings: &[Tensor], + pad_value: f32, + device: &Device, +) -> Result<(Tensor, Tensor, Vec)> { + if text_embeddings.is_empty() { + candle::bail!("text_embeddings cannot be empty"); + } + + let batch_size = text_embeddings.len(); + let dim = text_embeddings[0].dim(1)?; + let dtype = text_embeddings[0].dtype(); + + // Compute max length and align to SEQ_MULTI_OF + let lengths: Vec = text_embeddings + .iter() + .map(|t| t.dim(0)) + .collect::>>()?; + let max_len = *lengths.iter().max().unwrap(); + let padded_len = max_len + compute_padding_len(max_len); + + // Build padded tensor and mask + let mut padded_list = Vec::with_capacity(batch_size); + let mut mask_list = Vec::with_capacity(batch_size); + + for (i, emb) in text_embeddings.iter().enumerate() { + let seq_len = lengths[i]; + let pad_len = padded_len - seq_len; + + // Pad embedding + let padded = if pad_len > 0 { + let padding = Tensor::full(pad_value, (pad_len, dim), device)?.to_dtype(dtype)?; + Tensor::cat(&[emb, &padding], 0)? + } else { + emb.clone() + }; + padded_list.push(padded); + + // Create mask: 1 for valid, 0 for padding + let valid = Tensor::ones((seq_len,), DType::U8, device)?; + let mask = if pad_len > 0 { + let invalid = Tensor::zeros((pad_len,), DType::U8, device)?; + Tensor::cat(&[&valid, &invalid], 0)? + } else { + valid + }; + mask_list.push(mask); + } + + // Stack into batch + let cap_feats = Tensor::stack(&padded_list, 0)?; + let cap_mask = Tensor::stack(&mask_list, 0)?; + + Ok((cap_feats, cap_mask, lengths)) +} + +/// Prepare all inputs, converting variable-length inputs to fixed-shape batch tensors +/// +/// # Arguments +/// * `latents` - Latent tensor (B, C, H, W) +/// * `text_embeddings` - Variable-length text embeddings, each of shape (seq_len, cap_feat_dim) +/// * `device` - Device +/// +/// # Returns +/// PreparedInputs containing all preprocessed tensors +pub fn prepare_inputs( + latents: &Tensor, + text_embeddings: &[Tensor], + device: &Device, +) -> Result { + // Latents: (B, C, H, W) -> (B, C, 1, H, W) add frame dimension + let latents = latents.unsqueeze(2)?; + + // Pad text embeddings + let (cap_feats, cap_mask, text_lengths) = pad_text_embeddings(text_embeddings, 0.0, device)?; + + Ok(PreparedInputs { + latents, + cap_feats, + cap_mask, + text_lengths, + }) +} + +/// Create attention mask for a single sample +/// Useful for testing or simplified scenarios +pub fn create_attention_mask( + valid_len: usize, + total_len: usize, + device: &Device, +) -> Result { + let valid = Tensor::ones((valid_len,), DType::U8, device)?; + if valid_len < total_len { + let invalid = Tensor::zeros((total_len - valid_len,), DType::U8, device)?; + Tensor::cat(&[&valid, &invalid], 0) + } else { + Ok(valid) + } +} + +/// Create a batch of uniform text embeddings +/// +/// # Arguments +/// * `text_embedding` - Single text embedding (seq_len, dim) +/// * `batch_size` - Number of copies to create +/// +/// # Returns +/// Batched text embeddings (batch_size, seq_len, dim) +pub fn batch_text_embedding(text_embedding: &Tensor, batch_size: usize) -> Result { + let (seq_len, dim) = text_embedding.dims2()?; + text_embedding + .unsqueeze(0)? + .broadcast_as((batch_size, seq_len, dim))? + .contiguous() +} + +/// Create a batch of uniform masks +/// +/// # Arguments +/// * `mask` - Single mask (seq_len,) +/// * `batch_size` - Number of copies to create +/// +/// # Returns +/// Batched masks (batch_size, seq_len) +pub fn batch_mask(mask: &Tensor, batch_size: usize) -> Result { + let seq_len = mask.dim(0)?; + mask.unsqueeze(0)? + .broadcast_as((batch_size, seq_len))? + .contiguous() +} diff --git a/patches/candle-transformers/src/models/z_image/sampling.rs b/patches/candle-transformers/src/models/z_image/sampling.rs new file mode 100644 index 0000000000..8d035a34fa --- /dev/null +++ b/patches/candle-transformers/src/models/z_image/sampling.rs @@ -0,0 +1,133 @@ +//! Sampling utilities for Z-Image model. + +use candle::{DType, Device, Result, Tensor}; + +/// Generate initial Gaussian noise +/// +/// # Arguments +/// * `batch_size` - Batch size +/// * `channels` - Number of channels (typically 16, VAE latent channels) +/// * `height` - Height (latent space, i.e., image_height / 16) +/// * `width` - Width (latent space) +/// * `device` - Compute device +/// +/// # Returns +/// Noise tensor of shape (batch_size, channels, height, width) +pub fn get_noise( + batch_size: usize, + channels: usize, + height: usize, + width: usize, + device: &Device, +) -> Result { + Tensor::randn(0f32, 1.0, (batch_size, channels, height, width), device) +} + +/// Get linear time schedule with shift +/// +/// # Arguments +/// * `num_steps` - Number of inference steps +/// * `mu` - Time shift parameter (from calculate_shift) +/// +/// # Returns +/// Time points from 1.0 to 0.0 (num_steps+1 points) +pub fn get_schedule(num_steps: usize, mu: f64) -> Vec { + let timesteps: Vec = (0..=num_steps) + .map(|v| v as f64 / num_steps as f64) + .rev() + .collect(); + + // Apply time shift (for Flow Matching) + timesteps + .into_iter() + .map(|t| { + if t <= 0.0 || t >= 1.0 { + t // boundary case + } else { + let e = mu.exp(); + e / (e + (1.0 / t - 1.0)) + } + }) + .collect() +} + +/// Post-process image from VAE output +/// Converts from [-1, 1] to [0, 255] u8 image +pub fn postprocess_image(image: &Tensor) -> Result { + let image = image.clamp(-1.0, 1.0)?; + let image = ((image + 1.0)? * 127.5)?; + image.to_dtype(DType::U8) +} + +/// CFG configuration +#[derive(Debug, Clone)] +pub struct CfgConfig { + /// Guidance scale (typically 5.0) + pub guidance_scale: f64, + /// CFG truncation threshold (1.0 = full CFG, 0.0 = no CFG) + pub cfg_truncation: f64, + /// Whether to normalize CFG output + pub cfg_normalization: bool, +} + +impl Default for CfgConfig { + fn default() -> Self { + Self { + guidance_scale: 5.0, + cfg_truncation: 1.0, + cfg_normalization: false, + } + } +} + +/// Apply Classifier-Free Guidance +/// +/// # Arguments +/// * `pos_pred` - Positive (conditional) prediction +/// * `neg_pred` - Negative (unconditional) prediction +/// * `cfg` - CFG configuration +/// * `t_norm` - Normalized time [0, 1] +pub fn apply_cfg( + pos_pred: &Tensor, + neg_pred: &Tensor, + cfg: &CfgConfig, + t_norm: f64, +) -> Result { + // CFG truncation: disable CFG in late sampling + let current_scale = if t_norm > cfg.cfg_truncation { + 0.0 + } else { + cfg.guidance_scale + }; + + if current_scale <= 0.0 { + return Ok(pos_pred.clone()); + } + + // CFG formula: pred = pos + scale * (pos - neg) + let diff = (pos_pred - neg_pred)?; + let pred = (pos_pred + (diff * current_scale)?)?; + + // Optional: CFG normalization (limit output norm) + if cfg.cfg_normalization { + let ori_norm = pos_pred.sqr()?.sum_all()?.sqrt()?; + let new_norm = pred.sqr()?.sum_all()?.sqrt()?; + let ori_norm_val = ori_norm.to_scalar::()?; + let new_norm_val = new_norm.to_scalar::()?; + + if new_norm_val > ori_norm_val { + let scale = ori_norm_val / new_norm_val; + return pred * scale as f64; + } + } + + Ok(pred) +} + +/// Scale latents to initial noise level +/// +/// For flow matching, the initial sample should be pure noise. +/// This function scales the noise by the initial sigma. +pub fn scale_noise(noise: &Tensor, sigma: f64) -> Result { + noise * sigma +} diff --git a/patches/candle-transformers/src/models/z_image/scheduler.rs b/patches/candle-transformers/src/models/z_image/scheduler.rs new file mode 100644 index 0000000000..e5aaff2b6a --- /dev/null +++ b/patches/candle-transformers/src/models/z_image/scheduler.rs @@ -0,0 +1,237 @@ +//! FlowMatch Euler Discrete Scheduler for Z-Image +//! +//! Implements the flow matching scheduler used in Z-Image generation. + +use candle::{Result, Tensor}; + +/// FlowMatchEulerDiscreteScheduler configuration +#[derive(Debug, Clone, serde::Deserialize)] +pub struct SchedulerConfig { + #[serde(default = "default_num_train_timesteps")] + pub num_train_timesteps: usize, + #[serde(default = "default_shift")] + pub shift: f64, + #[serde(default)] + pub use_dynamic_shifting: bool, +} + +fn default_num_train_timesteps() -> usize { + 1000 +} +fn default_shift() -> f64 { + 3.0 +} + +impl Default for SchedulerConfig { + fn default() -> Self { + Self { + num_train_timesteps: default_num_train_timesteps(), + shift: default_shift(), + use_dynamic_shifting: false, + } + } +} + +impl SchedulerConfig { + /// Create configuration for Z-Image Turbo + pub fn z_image_turbo() -> Self { + Self { + num_train_timesteps: 1000, + shift: 3.0, + use_dynamic_shifting: false, + } + } +} + +/// FlowMatch Euler Discrete Scheduler +#[derive(Debug, Clone)] +pub struct FlowMatchEulerDiscreteScheduler { + /// Configuration + pub config: SchedulerConfig, + /// Timesteps for inference + pub timesteps: Vec, + /// Sigma values + pub sigmas: Vec, + /// Minimum sigma + pub sigma_min: f64, + /// Maximum sigma + pub sigma_max: f64, + /// Current step index + step_index: usize, +} + +impl FlowMatchEulerDiscreteScheduler { + pub fn new(config: SchedulerConfig) -> Self { + let num_train_timesteps = config.num_train_timesteps; + let shift = config.shift; + + // Generate initial sigmas + let timesteps: Vec = (1..=num_train_timesteps).rev().map(|t| t as f64).collect(); + + let sigmas: Vec = timesteps + .iter() + .map(|&t| t / num_train_timesteps as f64) + .collect(); + + // Apply shift + let sigmas: Vec = if !config.use_dynamic_shifting { + sigmas + .iter() + .map(|&s| shift * s / (1.0 + (shift - 1.0) * s)) + .collect() + } else { + sigmas + }; + + let timesteps: Vec = sigmas + .iter() + .map(|&s| s * num_train_timesteps as f64) + .collect(); + + let sigma_max = sigmas[0]; + let sigma_min = *sigmas.last().unwrap_or(&0.0); + + Self { + config, + timesteps, + sigmas, + sigma_min, + sigma_max, + step_index: 0, + } + } + + /// Set timesteps for inference + /// + /// # Arguments + /// * `num_inference_steps` - Number of denoising steps + /// * `mu` - Optional time shift parameter (from calculate_shift) + pub fn set_timesteps(&mut self, num_inference_steps: usize, mu: Option) { + let sigma_max = self.sigmas[0]; + let sigma_min = *self.sigmas.last().unwrap_or(&0.0); + + // Linear interpolation to generate timesteps + let timesteps: Vec = (0..num_inference_steps) + .map(|i| { + let t = i as f64 / num_inference_steps as f64; + sigma_max * (1.0 - t) + sigma_min * t + }) + .map(|s| s * self.config.num_train_timesteps as f64) + .collect(); + + let mut sigmas: Vec = timesteps + .iter() + .map(|&t| t / self.config.num_train_timesteps as f64) + .collect(); + + // Apply shift + if let Some(mu) = mu { + if self.config.use_dynamic_shifting { + // time_shift: exp(mu) / (exp(mu) + (1/t - 1)) + sigmas = sigmas + .iter() + .map(|&t| { + if t <= 0.0 { + 0.0 + } else { + let e_mu = mu.exp(); + e_mu / (e_mu + (1.0 / t - 1.0)) + } + }) + .collect(); + } + } else if !self.config.use_dynamic_shifting { + let shift = self.config.shift; + sigmas = sigmas + .iter() + .map(|&s| shift * s / (1.0 + (shift - 1.0) * s)) + .collect(); + } + + // Add terminal sigma = 0 + sigmas.push(0.0); + + self.timesteps = timesteps; + self.sigmas = sigmas; + self.step_index = 0; + } + + /// Get current sigma value + pub fn current_sigma(&self) -> f64 { + self.sigmas[self.step_index] + } + + /// Get current timestep (for model input) + /// Converts scheduler timestep to model input format: (1000 - t) / 1000 + pub fn current_timestep_normalized(&self) -> f64 { + let t = self.timesteps.get(self.step_index).copied().unwrap_or(0.0); + (1000.0 - t) / 1000.0 + } + + /// Euler step + /// + /// # Arguments + /// * `model_output` - Model predicted velocity field + /// * `sample` - Current sample x_t + /// + /// # Returns + /// Next sample x_{t-1} + pub fn step(&mut self, model_output: &Tensor, sample: &Tensor) -> Result { + let sigma = self.sigmas[self.step_index]; + let sigma_next = self.sigmas[self.step_index + 1]; + + let dt = sigma_next - sigma; + + // prev_sample = sample + dt * model_output + let prev_sample = (sample + (model_output * dt)?)?; + + self.step_index += 1; + Ok(prev_sample) + } + + /// Reset scheduler state + pub fn reset(&mut self) { + self.step_index = 0; + } + + /// Get number of inference steps + pub fn num_inference_steps(&self) -> usize { + self.timesteps.len() + } + + /// Get current step index + pub fn step_index(&self) -> usize { + self.step_index + } + + /// Check if denoising is complete + pub fn is_complete(&self) -> bool { + self.step_index >= self.timesteps.len() + } +} + +/// Calculate timestep shift parameter mu +/// +/// # Arguments +/// * `image_seq_len` - Image sequence length (after patchify) +/// * `base_seq_len` - Base sequence length (typically 256) +/// * `max_seq_len` - Maximum sequence length (typically 4096) +/// * `base_shift` - Base shift value (typically 0.5) +/// * `max_shift` - Maximum shift value (typically 1.15) +pub fn calculate_shift( + image_seq_len: usize, + base_seq_len: usize, + max_seq_len: usize, + base_shift: f64, + max_shift: f64, +) -> f64 { + let m = (max_shift - base_shift) / (max_seq_len - base_seq_len) as f64; + let b = base_shift - m * base_seq_len as f64; + image_seq_len as f64 * m + b +} + +/// Constants for shift calculation +pub const BASE_IMAGE_SEQ_LEN: usize = 256; +pub const MAX_IMAGE_SEQ_LEN: usize = 4096; +pub const BASE_SHIFT: f64 = 0.5; +pub const MAX_SHIFT: f64 = 1.15; diff --git a/patches/candle-transformers/src/models/z_image/text_encoder.rs b/patches/candle-transformers/src/models/z_image/text_encoder.rs new file mode 100644 index 0000000000..de4ad7f640 --- /dev/null +++ b/patches/candle-transformers/src/models/z_image/text_encoder.rs @@ -0,0 +1,453 @@ +//! Z-Image Text Encoder (Qwen3 Adapter) +//! +//! This module provides a Qwen3-based text encoder for Z-Image. +//! Key difference from the standard Qwen3 model: +//! - Returns the **second-to-last layer** hidden states (hidden_states[-2]) +//! - Does NOT apply the final RMSNorm + +use crate::models::with_tracing::{linear_b, Linear, RmsNorm}; +use candle::{DType, Device, Module, Result, Tensor}; +use candle_nn::{Activation, VarBuilder}; +use std::sync::Arc; + +/// Text Encoder configuration (Qwen3-based) +#[derive(Debug, Clone, serde::Deserialize)] +pub struct TextEncoderConfig { + #[serde(default = "default_vocab_size")] + pub vocab_size: usize, + #[serde(default = "default_hidden_size")] + pub hidden_size: usize, + #[serde(default = "default_intermediate_size")] + pub intermediate_size: usize, + #[serde(default = "default_num_hidden_layers")] + pub num_hidden_layers: usize, + #[serde(default = "default_num_attention_heads")] + pub num_attention_heads: usize, + #[serde(default = "default_num_key_value_heads")] + pub num_key_value_heads: usize, + #[serde(default = "default_head_dim")] + pub head_dim: usize, + #[serde(default = "default_rms_norm_eps")] + pub rms_norm_eps: f64, + #[serde(default = "default_rope_theta")] + pub rope_theta: f64, + #[serde(default = "default_attention_bias")] + pub attention_bias: bool, + #[serde(default = "default_hidden_act")] + pub hidden_act: Activation, + #[serde(default = "default_max_position_embeddings")] + pub max_position_embeddings: usize, +} + +fn default_vocab_size() -> usize { + 151936 +} +fn default_hidden_size() -> usize { + 2560 +} +fn default_intermediate_size() -> usize { + 9728 +} +fn default_num_hidden_layers() -> usize { + 36 +} +fn default_num_attention_heads() -> usize { + 32 +} +fn default_num_key_value_heads() -> usize { + 8 +} +fn default_head_dim() -> usize { + 128 +} +fn default_rms_norm_eps() -> f64 { + 1e-6 +} +fn default_rope_theta() -> f64 { + 1_000_000.0 +} +fn default_attention_bias() -> bool { + false +} +fn default_hidden_act() -> Activation { + Activation::Silu +} +fn default_max_position_embeddings() -> usize { + 40960 +} + +impl Default for TextEncoderConfig { + fn default() -> Self { + Self::z_image() + } +} + +impl TextEncoderConfig { + /// Create configuration for Z-Image Text Encoder + pub fn z_image() -> Self { + Self { + vocab_size: 151936, + hidden_size: 2560, + intermediate_size: 9728, + num_hidden_layers: 36, + num_attention_heads: 32, + num_key_value_heads: 8, + head_dim: 128, + rms_norm_eps: 1e-6, + rope_theta: 1_000_000.0, + attention_bias: false, + hidden_act: Activation::Silu, + max_position_embeddings: 40960, + } + } +} + +// ==================== Rotary Embedding ==================== + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(dtype: DType, cfg: &TextEncoderConfig, dev: &Device) -> Result { + let dim = cfg.head_dim; + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(DType::F32)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(DType::F32)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?.to_dtype(dtype)?, + cos: freqs.cos()?.to_dtype(dtype)?, + }) + } + + /// Apply RoPE (q, k shape: B x H x L x D) + fn apply(&self, q: &Tensor, k: &Tensor, offset: usize) -> Result<(Tensor, Tensor)> { + let (_, _, seq_len, _) = q.dims4()?; + let cos = self.cos.narrow(0, offset, seq_len)?; + let sin = self.sin.narrow(0, offset, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +// ==================== MLP ==================== + +#[derive(Debug, Clone)] +struct Mlp { + gate_proj: candle_nn::Linear, + up_proj: candle_nn::Linear, + down_proj: candle_nn::Linear, + act_fn: Activation, +} + +impl Mlp { + fn new(cfg: &TextEncoderConfig, vb: VarBuilder) -> Result { + Ok(Self { + gate_proj: candle_nn::linear_no_bias( + cfg.hidden_size, + cfg.intermediate_size, + vb.pp("gate_proj"), + )?, + up_proj: candle_nn::linear_no_bias( + cfg.hidden_size, + cfg.intermediate_size, + vb.pp("up_proj"), + )?, + down_proj: candle_nn::linear_no_bias( + cfg.intermediate_size, + cfg.hidden_size, + vb.pp("down_proj"), + )?, + act_fn: cfg.hidden_act, + }) + } +} + +impl Module for Mlp { + fn forward(&self, x: &Tensor) -> Result { + let lhs = x.apply(&self.gate_proj)?.apply(&self.act_fn)?; + let rhs = x.apply(&self.up_proj)?; + (lhs * rhs)?.apply(&self.down_proj) + } +} + +// ==================== Attention ==================== + +fn repeat_kv(x: Tensor, n_rep: usize) -> Result { + if n_rep == 1 { + Ok(x) + } else { + let (b_sz, n_kv_head, seq_len, head_dim) = x.dims4()?; + x.unsqueeze(2)? + .broadcast_as((b_sz, n_kv_head, n_rep, seq_len, head_dim))? + .reshape((b_sz, n_kv_head * n_rep, seq_len, head_dim)) + } +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + q_norm: RmsNorm, + k_norm: RmsNorm, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + hidden_size: usize, + rotary_emb: Arc, +} + +impl Attention { + fn new( + cfg: &TextEncoderConfig, + rotary_emb: Arc, + vb: VarBuilder, + ) -> Result { + let head_dim = cfg.head_dim; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + + let q_proj = linear_b( + cfg.hidden_size, + num_heads * head_dim, + cfg.attention_bias, + vb.pp("q_proj"), + )?; + let k_proj = linear_b( + cfg.hidden_size, + num_kv_heads * head_dim, + cfg.attention_bias, + vb.pp("k_proj"), + )?; + let v_proj = linear_b( + cfg.hidden_size, + num_kv_heads * head_dim, + cfg.attention_bias, + vb.pp("v_proj"), + )?; + let o_proj = linear_b( + num_heads * head_dim, + cfg.hidden_size, + cfg.attention_bias, + vb.pp("o_proj"), + )?; + + let q_norm = RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("q_norm"))?; + let k_norm = RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("k_norm"))?; + + let hidden_size = head_dim * cfg.num_attention_heads; + + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + q_norm, + k_norm, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + hidden_size, + rotary_emb, + }) + } + + fn forward(&self, x: &Tensor, attn_mask: Option<&Tensor>, offset: usize) -> Result { + let (b, l, _) = x.dims3()?; + + // 1. Proj + let q = self.q_proj.forward(x)?; + let k = self.k_proj.forward(x)?; + let v = self.v_proj.forward(x)?; + + // 2. Reshape: (B, L, H, D) -> (B, H, L, D) + let q = q + .reshape((b, l, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let k = k + .reshape((b, l, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let v = v + .reshape((b, l, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + // 3. Per-head RMSNorm + let q_flat = q.flatten(0, 2)?; + let k_flat = k.flatten(0, 2)?; + let q_flat = self.q_norm.forward(&q_flat)?; + let k_flat = self.k_norm.forward(&k_flat)?; + let q = q_flat.reshape((b, self.num_heads, l, self.head_dim))?; + let k = k_flat.reshape((b, self.num_kv_heads, l, self.head_dim))?; + + // 4. RoPE + let (q, k) = self.rotary_emb.apply(&q, &k, offset)?; + + // 5. GQA repeat_kv + let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?; + let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?; + + // 6. Attention score + let scale = 1.0 / (self.head_dim as f64).sqrt(); + let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?; + if let Some(m) = attn_mask { + scores = scores.broadcast_add(m)?; + } + let probs = candle_nn::ops::softmax_last_dim(&scores)?; + let ctx = probs.matmul(&v)?; // (B, H, L, D) + + // 7. Output proj + ctx.transpose(1, 2)? + .reshape((b, l, self.hidden_size))? + .apply(&self.o_proj) + } +} + +// ==================== Decoder Layer ==================== + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: Mlp, + ln1: RmsNorm, + ln2: RmsNorm, +} + +impl DecoderLayer { + fn new(cfg: &TextEncoderConfig, rotary: Arc, vb: VarBuilder) -> Result { + let self_attn = Attention::new(cfg, rotary, vb.pp("self_attn"))?; + let mlp = Mlp::new(cfg, vb.pp("mlp"))?; + let ln1 = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let ln2 = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + ln1, + ln2, + }) + } + + fn forward(&self, x: &Tensor, mask: Option<&Tensor>, offset: usize) -> Result { + let h = self.ln1.forward(x)?; + let h = self.self_attn.forward(&h, mask, offset)?; + let x = (x + h)?; + let h2 = self.ln2.forward(&x)?; + let h2 = h2.apply(&self.mlp)?; + x + h2 + } +} + +// ==================== ZImageTextEncoder ==================== + +/// Z-Image Text Encoder (Qwen3-based) +/// +/// Returns the second-to-last layer hidden states (hidden_states[-2]) +/// without applying the final RMSNorm. +#[derive(Debug, Clone)] +pub struct ZImageTextEncoder { + embed_tokens: candle_nn::Embedding, + layers: Vec, + num_hidden_layers: usize, + device: Device, + dtype: DType, +} + +impl ZImageTextEncoder { + pub fn new(cfg: &TextEncoderConfig, vb: VarBuilder) -> Result { + // Note: weights have "model." prefix + let vb_model = vb.pp("model"); + + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb_model.pp("embed_tokens"))?; + + let rotary = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb.device())?); + + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_layers = vb_model.pp("layers"); + for i in 0..cfg.num_hidden_layers { + layers.push(DecoderLayer::new(cfg, rotary.clone(), vb_layers.pp(i))?); + } + + // NOTE: We do NOT load the final norm (model.norm.weight) + // because we return the second-to-last layer output without final norm + + Ok(Self { + embed_tokens, + layers, + num_hidden_layers: cfg.num_hidden_layers, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + /// Create causal attention mask + fn causal_mask(&self, b: usize, tgt: usize, offset: usize) -> Result { + let minf = f32::NEG_INFINITY; + let mask: Vec<_> = (0..tgt) + .flat_map(|i| { + (0..(tgt + offset)).map(move |j| if j <= i + offset { 0.0 } else { minf }) + }) + .collect(); + Tensor::from_slice(&mask, (b, 1, tgt, tgt + offset), &self.device)?.to_dtype(self.dtype) + } + + /// Encode text, returning second-to-last layer hidden states + /// + /// # Arguments + /// * `input_ids` - Token IDs (B, seq_len) + /// + /// # Returns + /// Hidden states (B, seq_len, hidden_size) from layer[-2] + /// + /// **Important**: Returns raw output from layer[-2] WITHOUT final RMSNorm + pub fn forward(&self, input_ids: &Tensor) -> Result { + let (b, l) = input_ids.dims2()?; + let mut hidden_states = self.embed_tokens.forward(input_ids)?; + + let causal = if l == 1 { + None + } else { + Some(self.causal_mask(b, l, 0)?) + }; + + // num_hidden_layers = 36, second-to-last layer index = 34 + let target_layer = self.num_hidden_layers - 2; + + for (i, layer) in self.layers.iter().enumerate() { + hidden_states = layer.forward(&hidden_states, causal.as_ref(), 0)?; + + // Return after second-to-last layer, do NOT apply final norm + if i == target_layer { + return Ok(hidden_states); + } + } + + // Should not reach here + candle::bail!("Layer index out of bounds") + } + + /// Get the output dimension (hidden_size) + pub fn hidden_size(&self) -> usize { + // This is derived from embed_tokens weight shape + self.embed_tokens.embeddings().dim(1).unwrap_or(2560) + } +} diff --git a/patches/candle-transformers/src/models/z_image/transformer.rs b/patches/candle-transformers/src/models/z_image/transformer.rs new file mode 100644 index 0000000000..1b810fe431 --- /dev/null +++ b/patches/candle-transformers/src/models/z_image/transformer.rs @@ -0,0 +1,1087 @@ +//! Z-Image Transformer (ZImageTransformer2DModel) +//! +//! Core transformer implementation for Z-Image text-to-image generation. + +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{linear, linear_no_bias, VarBuilder}; + +use crate::models::with_tracing::RmsNorm; + +// ==================== Flash Attention Wrapper ==================== + +/// Flash Attention wrapper for CUDA platform +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +#[allow(dead_code)] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + candle::bail!("flash-attn feature not enabled, compile with '--features flash-attn'") +} + +// ==================== Constants ==================== + +/// AdaLN embedding dimension (256) +pub const ADALN_EMBED_DIM: usize = 256; +/// Sequence padding alignment (32) +pub const SEQ_MULTI_OF: usize = 32; +/// Frequency embedding size for timestep encoding +pub const FREQUENCY_EMBEDDING_SIZE: usize = 256; +/// Max period for sinusoidal encoding +pub const MAX_PERIOD: f64 = 10000.0; + +// ==================== Config ==================== + +/// Z-Image Transformer configuration +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + #[serde(default = "default_patch_size")] + pub all_patch_size: Vec, + #[serde(default = "default_f_patch_size")] + pub all_f_patch_size: Vec, + #[serde(default = "default_in_channels")] + pub in_channels: usize, + #[serde(default = "default_dim")] + pub dim: usize, + #[serde(default = "default_n_layers")] + pub n_layers: usize, + #[serde(default = "default_n_refiner_layers")] + pub n_refiner_layers: usize, + #[serde(default = "default_n_heads")] + pub n_heads: usize, + #[serde(default = "default_n_kv_heads")] + pub n_kv_heads: usize, + #[serde(default = "default_norm_eps")] + pub norm_eps: f64, + #[serde(default = "default_qk_norm")] + pub qk_norm: bool, + #[serde(default = "default_cap_feat_dim")] + pub cap_feat_dim: usize, + #[serde(default = "default_rope_theta")] + pub rope_theta: f64, + #[serde(default = "default_t_scale")] + pub t_scale: f64, + #[serde(default = "default_axes_dims")] + pub axes_dims: Vec, + #[serde(default = "default_axes_lens")] + pub axes_lens: Vec, + /// Whether to use accelerated attention (CUDA flash-attn / Metal SDPA) + /// Default is true, automatically selects optimal implementation per platform + #[serde(default = "default_use_accelerated_attn")] + pub use_accelerated_attn: bool, +} + +fn default_use_accelerated_attn() -> bool { + true +} + +fn default_patch_size() -> Vec { + vec![2] +} +fn default_f_patch_size() -> Vec { + vec![1] +} +fn default_in_channels() -> usize { + 16 +} +fn default_dim() -> usize { + 3840 +} +fn default_n_layers() -> usize { + 30 +} +fn default_n_refiner_layers() -> usize { + 2 +} +fn default_n_heads() -> usize { + 30 +} +fn default_n_kv_heads() -> usize { + 30 +} +fn default_norm_eps() -> f64 { + 1e-5 +} +fn default_qk_norm() -> bool { + true +} +fn default_cap_feat_dim() -> usize { + 2560 +} +fn default_rope_theta() -> f64 { + 256.0 +} +fn default_t_scale() -> f64 { + 1000.0 +} +fn default_axes_dims() -> Vec { + vec![32, 48, 48] +} +fn default_axes_lens() -> Vec { + vec![1536, 512, 512] +} + +impl Config { + /// Create configuration for Z-Image Turbo model + pub fn z_image_turbo() -> Self { + Self { + all_patch_size: vec![2], + all_f_patch_size: vec![1], + in_channels: 16, + dim: 3840, + n_layers: 30, + n_refiner_layers: 2, + n_heads: 30, + n_kv_heads: 30, + norm_eps: 1e-5, + qk_norm: true, + cap_feat_dim: 2560, + rope_theta: 256.0, + t_scale: 1000.0, + axes_dims: vec![32, 48, 48], + axes_lens: vec![1536, 512, 512], + use_accelerated_attn: true, + } + } + + /// Set whether to use accelerated attention (for debugging) + pub fn set_use_accelerated_attn(&mut self, enabled: bool) { + self.use_accelerated_attn = enabled; + } + + /// Get head dimension + pub fn head_dim(&self) -> usize { + self.dim / self.n_heads + } + + /// Get hidden dimension for FFN + /// Matches Python: int(dim / 3 * 8) = 10240 for dim=3840 + pub fn hidden_dim(&self) -> usize { + (self.dim / 3) * 8 + } +} + +// ==================== TimestepEmbedder ==================== + +/// Timestep embedding using sinusoidal encoding + MLP +#[derive(Debug, Clone)] +pub struct TimestepEmbedder { + linear1: candle_nn::Linear, + linear2: candle_nn::Linear, + frequency_embedding_size: usize, +} + +impl TimestepEmbedder { + pub fn new(out_size: usize, mid_size: usize, vb: VarBuilder) -> Result { + let linear1 = linear(FREQUENCY_EMBEDDING_SIZE, mid_size, vb.pp("mlp").pp("0"))?; + let linear2 = linear(mid_size, out_size, vb.pp("mlp").pp("2"))?; + Ok(Self { + linear1, + linear2, + frequency_embedding_size: FREQUENCY_EMBEDDING_SIZE, + }) + } + + fn timestep_embedding(&self, t: &Tensor, device: &Device, dtype: DType) -> Result { + let half = self.frequency_embedding_size / 2; + let freqs = Tensor::arange(0u32, half as u32, device)?.to_dtype(DType::F32)?; + let freqs = (freqs * (-MAX_PERIOD.ln() / half as f64))?.exp()?; + let args = t + .unsqueeze(1)? + .to_dtype(DType::F32)? + .broadcast_mul(&freqs.unsqueeze(0)?)?; + let embedding = Tensor::cat(&[args.cos()?, args.sin()?], D::Minus1)?; + embedding.to_dtype(dtype) + } + + pub fn forward(&self, t: &Tensor) -> Result { + let device = t.device(); + let dtype = self.linear1.weight().dtype(); + let t_freq = self.timestep_embedding(t, device, dtype)?; + t_freq.apply(&self.linear1)?.silu()?.apply(&self.linear2) + } +} + +// ==================== FeedForward (SwiGLU) ==================== + +/// SwiGLU feedforward network +#[derive(Debug, Clone)] +pub struct FeedForward { + w1: candle_nn::Linear, + w2: candle_nn::Linear, + w3: candle_nn::Linear, +} + +impl FeedForward { + pub fn new(dim: usize, hidden_dim: usize, vb: VarBuilder) -> Result { + let w1 = linear_no_bias(dim, hidden_dim, vb.pp("w1"))?; + let w2 = linear_no_bias(hidden_dim, dim, vb.pp("w2"))?; + let w3 = linear_no_bias(dim, hidden_dim, vb.pp("w3"))?; + Ok(Self { w1, w2, w3 }) + } +} + +impl Module for FeedForward { + fn forward(&self, x: &Tensor) -> Result { + let x1 = x.apply(&self.w1)?.silu()?; + let x3 = x.apply(&self.w3)?; + (x1 * x3)?.apply(&self.w2) + } +} + +// ==================== QkNorm ==================== + +/// QK normalization using RMSNorm +#[derive(Debug, Clone)] +pub struct QkNorm { + norm_q: RmsNorm, + norm_k: RmsNorm, +} + +impl QkNorm { + pub fn new(head_dim: usize, eps: f64, vb: VarBuilder) -> Result { + let norm_q = RmsNorm::new(head_dim, eps, vb.pp("norm_q"))?; + let norm_k = RmsNorm::new(head_dim, eps, vb.pp("norm_k"))?; + Ok(Self { norm_q, norm_k }) + } + + pub fn forward(&self, q: &Tensor, k: &Tensor) -> Result<(Tensor, Tensor)> { + // q, k shape: (B, seq_len, n_heads, head_dim) + let q = self.norm_q.forward(q)?; + let k = self.norm_k.forward(k)?; + Ok((q, k)) + } +} + +// ==================== RopeEmbedder (3D) ==================== + +/// 3D Rotary Position Embedding for video/image generation +#[derive(Debug, Clone)] +pub struct RopeEmbedder { + #[allow(dead_code)] + theta: f64, + axes_dims: Vec, + #[allow(dead_code)] + axes_lens: Vec, + /// Pre-computed cos cache per axis + cos_cached: Vec, + /// Pre-computed sin cache per axis + sin_cached: Vec, +} + +impl RopeEmbedder { + pub fn new( + theta: f64, + axes_dims: Vec, + axes_lens: Vec, + device: &Device, + dtype: DType, + ) -> Result { + assert_eq!(axes_dims.len(), axes_lens.len()); + let mut cos_cached = Vec::with_capacity(axes_dims.len()); + let mut sin_cached = Vec::with_capacity(axes_dims.len()); + + for (d, e) in axes_dims.iter().zip(axes_lens.iter()) { + let half_d = d / 2; + let inv_freq: Vec = (0..half_d) + .map(|i| 1.0 / (theta as f32).powf((2 * i) as f32 / *d as f32)) + .collect(); + let inv_freq = Tensor::from_vec(inv_freq, half_d, device)?; + + let positions = Tensor::arange(0u32, *e as u32, device)?.to_dtype(DType::F32)?; + let freqs = positions + .unsqueeze(1)? + .broadcast_mul(&inv_freq.unsqueeze(0)?)?; + + cos_cached.push(freqs.cos()?.to_dtype(dtype)?); + sin_cached.push(freqs.sin()?.to_dtype(dtype)?); + } + + Ok(Self { + theta, + axes_dims, + axes_lens, + cos_cached, + sin_cached, + }) + } + + /// Get RoPE cos/sin from position IDs + /// ids: (seq_len, 3) - [frame_id, height_id, width_id] + pub fn forward(&self, ids: &Tensor) -> Result<(Tensor, Tensor)> { + let mut cos_parts = Vec::with_capacity(self.axes_dims.len()); + let mut sin_parts = Vec::with_capacity(self.axes_dims.len()); + + for (i, _) in self.axes_dims.iter().enumerate() { + let axis_ids = ids.i((.., i))?.contiguous()?; // (seq_len,) - must be contiguous for Metal + let cos_i = self.cos_cached[i].index_select(&axis_ids, 0)?; + let sin_i = self.sin_cached[i].index_select(&axis_ids, 0)?; + cos_parts.push(cos_i); + sin_parts.push(sin_i); + } + + let cos = Tensor::cat(&cos_parts, D::Minus1)?; // (seq_len, head_dim/2) + let sin = Tensor::cat(&sin_parts, D::Minus1)?; + Ok((cos, sin)) + } +} + +/// Apply RoPE (real-number form, equivalent to PyTorch complex multiplication) +/// +/// x: (B, seq_len, n_heads, head_dim) +/// cos, sin: (seq_len, head_dim/2) +pub fn apply_rotary_emb(x: &Tensor, cos: &Tensor, sin: &Tensor) -> Result { + let (b, seq_len, n_heads, head_dim) = x.dims4()?; + let half_dim = head_dim / 2; + + // Reshape x to interleaved real/imag form: (B, seq_len, n_heads, half_dim, 2) + let x = x.reshape((b, seq_len, n_heads, half_dim, 2))?; + + // Extract real and imag parts + let x_real = x.i((.., .., .., .., 0))?; // (B, seq_len, n_heads, half_dim) + let x_imag = x.i((.., .., .., .., 1))?; + + // Expand cos/sin for broadcasting: (seq_len, half_dim) -> (1, seq_len, 1, half_dim) + let cos = cos.unsqueeze(0)?.unsqueeze(2)?; + let sin = sin.unsqueeze(0)?.unsqueeze(2)?; + + // Complex multiplication: (a + bi)(c + di) = (ac - bd) + (ad + bc)i + let y_real = (x_real.broadcast_mul(&cos)? - x_imag.broadcast_mul(&sin)?)?; + let y_imag = (x_real.broadcast_mul(&sin)? + x_imag.broadcast_mul(&cos)?)?; + + // Interleave back + Tensor::stack(&[y_real, y_imag], D::Minus1)?.reshape((b, seq_len, n_heads, head_dim)) +} + +// ==================== ZImageAttention ==================== + +/// Z-Image attention with QK normalization and 3D RoPE +#[derive(Debug, Clone)] +pub struct ZImageAttention { + to_q: candle_nn::Linear, + to_k: candle_nn::Linear, + to_v: candle_nn::Linear, + to_out: candle_nn::Linear, + qk_norm: Option, + n_heads: usize, + head_dim: usize, + use_accelerated_attn: bool, +} + +impl ZImageAttention { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let dim = cfg.dim; + let n_heads = cfg.n_heads; + let head_dim = cfg.head_dim(); + + let to_q = linear_no_bias(dim, n_heads * head_dim, vb.pp("to_q"))?; + let to_k = linear_no_bias(dim, cfg.n_kv_heads * head_dim, vb.pp("to_k"))?; + let to_v = linear_no_bias(dim, cfg.n_kv_heads * head_dim, vb.pp("to_v"))?; + let to_out = linear_no_bias(n_heads * head_dim, dim, vb.pp("to_out").pp("0"))?; + + let qk_norm = if cfg.qk_norm { + Some(QkNorm::new(head_dim, 1e-5, vb.clone())?) + } else { + None + }; + + Ok(Self { + to_q, + to_k, + to_v, + to_out, + qk_norm, + n_heads, + head_dim, + use_accelerated_attn: cfg.use_accelerated_attn, + }) + } + + pub fn forward( + &self, + hidden_states: &Tensor, + attention_mask: Option<&Tensor>, + cos: &Tensor, + sin: &Tensor, + ) -> Result { + let (b, seq_len, _) = hidden_states.dims3()?; + + // Project to Q, K, V + let q = hidden_states.apply(&self.to_q)?; + let k = hidden_states.apply(&self.to_k)?; + let v = hidden_states.apply(&self.to_v)?; + + // Reshape: (B, seq_len, n_heads * head_dim) -> (B, seq_len, n_heads, head_dim) + let q = q.reshape((b, seq_len, self.n_heads, self.head_dim))?; + let k = k.reshape((b, seq_len, self.n_heads, self.head_dim))?; + let v = v.reshape((b, seq_len, self.n_heads, self.head_dim))?; + + // Apply QK norm + let (q, k) = if let Some(ref norm) = self.qk_norm { + norm.forward(&q, &k)? + } else { + (q, k) + }; + + // Apply RoPE + let q = apply_rotary_emb(&q, cos, sin)?; + let k = apply_rotary_emb(&k, cos, sin)?; + + // Transpose for attention: (B, n_heads, seq_len, head_dim) + let q = q.transpose(1, 2)?.contiguous()?; + let k = k.transpose(1, 2)?.contiguous()?; + let v = v.transpose(1, 2)?.contiguous()?; + + let scale = 1.0 / (self.head_dim as f64).sqrt(); + let device = hidden_states.device(); + + // Cross-platform attention dispatch + let context = self.attention_dispatch(&q, &k, &v, attention_mask, scale, device)?; + + // Reshape back: (B, n_heads, seq_len, head_dim) -> (B, seq_len, dim) + let context = context.transpose(1, 2)?.reshape((b, seq_len, ()))?; + + context.apply(&self.to_out) + } + + /// Cross-platform attention dispatch + fn attention_dispatch( + &self, + q: &Tensor, + k: &Tensor, + v: &Tensor, + mask: Option<&Tensor>, + scale: f64, + device: &Device, + ) -> Result { + // If acceleration disabled, use basic implementation + if !self.use_accelerated_attn { + return self.attention_basic(q, k, v, mask, scale); + } + + // Platform dispatch: prefer optimal implementation per platform + if device.is_cuda() { + self.attention_cuda(q, k, v, mask, scale) + } else if device.is_metal() { + self.attention_metal(q, k, v, mask, scale) + } else { + // CPU fallback + self.attention_basic(q, k, v, mask, scale) + } + } + + /// CUDA: Use Flash Attention + #[allow(unused_variables)] + fn attention_cuda( + &self, + q: &Tensor, + k: &Tensor, + v: &Tensor, + mask: Option<&Tensor>, + scale: f64, + ) -> Result { + #[cfg(feature = "flash-attn")] + { + // flash_attn does not directly support custom mask + // Fallback to basic implementation when mask is present + if mask.is_some() { + return self.attention_basic(q, k, v, mask, scale); + } + + // flash_attn input format: (batch, seq_len, num_heads, head_size) + // Current format: (batch, num_heads, seq_len, head_size) + let q = q.transpose(1, 2)?; + let k = k.transpose(1, 2)?; + let v = v.transpose(1, 2)?; + + let result = flash_attn(&q, &k, &v, scale as f32, false)?; + result.transpose(1, 2) + } + + #[cfg(not(feature = "flash-attn"))] + { + // flash-attn not compiled, fallback to basic + self.attention_basic(q, k, v, mask, scale) + } + } + + /// Metal: Use fused SDPA kernel + fn attention_metal( + &self, + q: &Tensor, + k: &Tensor, + v: &Tensor, + mask: Option<&Tensor>, + scale: f64, + ) -> Result { + // Prepare SDPA format mask + let sdpa_mask = self.prepare_sdpa_mask(mask, q)?; + + // candle_nn::ops::sdpa + // Input format: (bs, qhead, seq, hidden) - matches current format + // Supports: BF16/F16/F32, head_dim=128 + candle_nn::ops::sdpa(q, k, v, sdpa_mask.as_ref(), false, scale as f32, 1.0) + } + + /// Fallback implementation + fn attention_basic( + &self, + q: &Tensor, + k: &Tensor, + v: &Tensor, + mask: Option<&Tensor>, + scale: f64, + ) -> Result { + let mut attn_weights = (q.matmul(&k.transpose(2, 3)?)? * scale)?; + + if let Some(m) = mask { + // mask: (B, seq_len) -> (B, 1, 1, seq_len) + let m = m.unsqueeze(1)?.unsqueeze(2)?; + let m = m.to_dtype(attn_weights.dtype())?; + // 1=valid, 0=padding -> 0=valid, -inf=padding + let m = ((m - 1.0)? * 1e9)?; + attn_weights = attn_weights.broadcast_add(&m)?; + } + + let attn_probs = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_probs.matmul(v) + } + + /// Prepare SDPA format mask + fn prepare_sdpa_mask(&self, mask: Option<&Tensor>, q: &Tensor) -> Result> { + match mask { + Some(m) => { + // mask: (B, seq_len) -> (B, n_heads, seq_len, seq_len) + let (b, _, seq_len, _) = q.dims4()?; + let m = m.unsqueeze(1)?.unsqueeze(2)?; + let m = m.to_dtype(q.dtype())?; + // SDPA uses additive mask: 0=valid, -inf=masked + let m = ((m - 1.0)? * 1e9)?; + // broadcast to (B, n_heads, seq_len, seq_len) + let m = m.broadcast_as((b, self.n_heads, seq_len, seq_len))?; + Ok(Some(m)) + } + None => Ok(None), + } + } +} + +// ==================== ZImageTransformerBlock ==================== + +/// Z-Image transformer block with optional AdaLN modulation +#[derive(Debug, Clone)] +pub struct ZImageTransformerBlock { + attention: ZImageAttention, + feed_forward: FeedForward, + attention_norm1: RmsNorm, + attention_norm2: RmsNorm, + ffn_norm1: RmsNorm, + ffn_norm2: RmsNorm, + adaln_modulation: Option, +} + +impl ZImageTransformerBlock { + pub fn new(cfg: &Config, modulation: bool, vb: VarBuilder) -> Result { + let dim = cfg.dim; + let hidden_dim = cfg.hidden_dim(); + + let attention = ZImageAttention::new(cfg, vb.pp("attention"))?; + let feed_forward = FeedForward::new(dim, hidden_dim, vb.pp("feed_forward"))?; + + let attention_norm1 = RmsNorm::new(dim, cfg.norm_eps, vb.pp("attention_norm1"))?; + let attention_norm2 = RmsNorm::new(dim, cfg.norm_eps, vb.pp("attention_norm2"))?; + let ffn_norm1 = RmsNorm::new(dim, cfg.norm_eps, vb.pp("ffn_norm1"))?; + let ffn_norm2 = RmsNorm::new(dim, cfg.norm_eps, vb.pp("ffn_norm2"))?; + + let adaln_modulation = if modulation { + let adaln_dim = dim.min(ADALN_EMBED_DIM); + Some(linear( + adaln_dim, + 4 * dim, + vb.pp("adaLN_modulation").pp("0"), + )?) + } else { + None + }; + + Ok(Self { + attention, + feed_forward, + attention_norm1, + attention_norm2, + ffn_norm1, + ffn_norm2, + adaln_modulation, + }) + } + + pub fn forward( + &self, + x: &Tensor, + attn_mask: Option<&Tensor>, + cos: &Tensor, + sin: &Tensor, + adaln_input: Option<&Tensor>, + ) -> Result { + if let Some(ref adaln) = self.adaln_modulation { + let adaln_input = adaln_input.expect("adaln_input required when modulation=true"); + // (B, 256) -> (B, 4*dim) -> (B, 1, 4*dim) -> chunk into 4 + let modulation = adaln_input.apply(adaln)?.unsqueeze(1)?; + let chunks = modulation.chunk(4, D::Minus1)?; + let (scale_msa, gate_msa, scale_mlp, gate_mlp) = + (&chunks[0], &chunks[1], &chunks[2], &chunks[3]); + + // Apply tanh gate + let gate_msa = gate_msa.tanh()?; + let gate_mlp = gate_mlp.tanh()?; + let scale_msa = (scale_msa + 1.0)?; + let scale_mlp = (scale_mlp + 1.0)?; + + // Attention block + let normed = self.attention_norm1.forward(x)?; + let scaled = normed.broadcast_mul(&scale_msa)?; + let attn_out = self.attention.forward(&scaled, attn_mask, cos, sin)?; + let attn_out = self.attention_norm2.forward(&attn_out)?; + let x = (x + gate_msa.broadcast_mul(&attn_out)?)?; + + // FFN block + let normed = self.ffn_norm1.forward(&x)?; + let scaled = normed.broadcast_mul(&scale_mlp)?; + let ffn_out = self.feed_forward.forward(&scaled)?; + let ffn_out = self.ffn_norm2.forward(&ffn_out)?; + x + gate_mlp.broadcast_mul(&ffn_out)? + } else { + // Without modulation + let normed = self.attention_norm1.forward(x)?; + let attn_out = self.attention.forward(&normed, attn_mask, cos, sin)?; + let attn_out = self.attention_norm2.forward(&attn_out)?; + let x = (x + attn_out)?; + + let normed = self.ffn_norm1.forward(&x)?; + let ffn_out = self.feed_forward.forward(&normed)?; + let ffn_out = self.ffn_norm2.forward(&ffn_out)?; + x + ffn_out + } + } +} + +// ==================== FinalLayer ==================== + +/// LayerNorm without learnable parameters (elementwise_affine=False) +#[derive(Debug, Clone)] +pub struct LayerNormNoParams { + eps: f64, +} + +impl LayerNormNoParams { + pub fn new(eps: f64) -> Self { + Self { eps } + } +} + +impl Module for LayerNormNoParams { + fn forward(&self, x: &Tensor) -> Result { + let x_dtype = x.dtype(); + let internal_dtype = match x_dtype { + DType::F16 | DType::BF16 => DType::F32, + d => d, + }; + let hidden_size = x.dim(D::Minus1)?; + let x = x.to_dtype(internal_dtype)?; + // Subtract mean + let mean_x = (x.sum_keepdim(D::Minus1)? / hidden_size as f64)?; + let x = x.broadcast_sub(&mean_x)?; + // Divide by std + let norm_x = (x.sqr()?.sum_keepdim(D::Minus1)? / hidden_size as f64)?; + let x_normed = x.broadcast_div(&(norm_x + self.eps)?.sqrt()?)?; + x_normed.to_dtype(x_dtype) + } +} + +/// Final layer for output projection +#[derive(Debug, Clone)] +pub struct FinalLayer { + norm_final: LayerNormNoParams, + linear: candle_nn::Linear, + adaln_silu: candle_nn::Linear, +} + +impl FinalLayer { + pub fn new(hidden_size: usize, out_channels: usize, vb: VarBuilder) -> Result { + let norm_final = LayerNormNoParams::new(1e-6); + let linear = candle_nn::linear(hidden_size, out_channels, vb.pp("linear"))?; + let adaln_dim = hidden_size.min(ADALN_EMBED_DIM); + let adaln_silu = + candle_nn::linear(adaln_dim, hidden_size, vb.pp("adaLN_modulation").pp("1"))?; + + Ok(Self { + norm_final, + linear, + adaln_silu, + }) + } + + pub fn forward(&self, x: &Tensor, c: &Tensor) -> Result { + let scale = c.silu()?.apply(&self.adaln_silu)?; + let scale = (scale + 1.0)?.unsqueeze(1)?; + let x = self.norm_final.forward(x)?.broadcast_mul(&scale)?; + x.apply(&self.linear) + } +} + +// ==================== Patchify / Unpatchify ==================== + +/// Convert image to patch sequence +/// Matches Python: image.view(C, F_t, pF, H_t, pH, W_t, pW).permute(1,3,5,2,4,6,0) +/// +/// For Z-Image with F=1, pF=1, we optimize to use 6D operations. +/// input: (B, C, 1, H, W) +/// output: (B, num_patches, patch_dim), (F, H, W) original size +pub fn patchify( + x: &Tensor, + patch_size: usize, + f_patch_size: usize, +) -> Result<(Tensor, (usize, usize, usize))> { + let (b, c, f, h, w) = x.dims5()?; + let ph = patch_size; + let pw = patch_size; + let pf = f_patch_size; + + let f_tokens = f / pf; + let h_tokens = h / ph; + let w_tokens = w / pw; + let num_patches = f_tokens * h_tokens * w_tokens; + let patch_dim = pf * ph * pw * c; + + // For F=1, pF=1 case (image generation), use optimized 6D path + if f == 1 && pf == 1 { + // Step 1: Squeeze F dimension: (B, C, 1, H, W) -> (B, C, H, W) + let x = x.squeeze(2)?; + + // Step 2: Reshape H into (H_tokens, pH): (B, C, H, W) -> (B, C, H_t, pH, W) + let x = x.reshape((b, c, h_tokens, ph, w))?; + + // Step 3: Reshape W into (W_tokens, pW): (B, C, H_t, pH, W) -> (B, C, H_t, pH, W_t, pW) + let x = x.reshape((b, c, h_tokens, ph, w_tokens, pw))?; + + // Step 4: Permute to match Python: (C, H_t, pH, W_t, pW) -> (H_t, W_t, pH, pW, C) + // For batch: (B, C, H_t, pH, W_t, pW) -> (B, H_t, W_t, pH, pW, C) + // Permutation: (0, 2, 4, 3, 5, 1) + let x = x.permute((0, 2, 4, 3, 5, 1))?; + + // Step 5: Reshape to patches: (B, H_t, W_t, pH, pW, C) -> (B, H_t*W_t, pH*pW*C) + let x = x.reshape((b, num_patches, patch_dim))?; + + Ok((x, (f, h, w))) + } else { + // General case: use contiguous + reshape approach + // This is less common for Z-Image image generation + let x = x.permute((0, 2, 3, 4, 1))?.contiguous()?; // (B, F, H, W, C) + let x = x.reshape((b, f_tokens, pf, h_tokens, ph, w_tokens * pw * c))?; + let x = x.permute((0, 1, 3, 5, 2, 4))?.contiguous()?; + let x = x.reshape((b, num_patches, patch_dim))?; + Ok((x, (f, h, w))) + } +} + +/// Convert patch sequence back to image +/// Matches Python: x.view(F_t, H_t, W_t, pF, pH, pW, C).permute(6,0,3,1,4,2,5) +/// +/// For Z-Image with F=1, pF=1, we optimize to use 6D operations. +/// input: (B, seq_len, patch_dim) +/// output: (B, C, F, H, W) +pub fn unpatchify( + x: &Tensor, + size: (usize, usize, usize), + patch_size: usize, + f_patch_size: usize, + out_channels: usize, +) -> Result { + let (f, h, w) = size; + let ph = patch_size; + let pw = patch_size; + let pf = f_patch_size; + + let f_tokens = f / pf; + let h_tokens = h / ph; + let w_tokens = w / pw; + let ori_len = f_tokens * h_tokens * w_tokens; + + let (b, _, _) = x.dims3()?; + let x = x.narrow(1, 0, ori_len)?; // Remove padding + + // For F=1, pF=1 case (image generation), use optimized 6D path + if f == 1 && pf == 1 { + // Step 1: Reshape to (B, H_t, W_t, pH, pW, C) + let x = x.reshape((b, h_tokens, w_tokens, ph, pw, out_channels))?; + + // Step 2: Permute to match Python: (H_t, W_t, pH, pW, C) -> (C, H_t, pH, W_t, pW) + // For batch: (B, H_t, W_t, pH, pW, C) -> (B, C, H_t, pH, W_t, pW) + // Permutation: (0, 5, 1, 3, 2, 4) + let x = x.permute((0, 5, 1, 3, 2, 4))?; + + // Step 3: Reshape to combine H and W: (B, C, H_t, pH, W_t, pW) -> (B, C, H, W) + let x = x.reshape((b, out_channels, h, w))?; + + // Step 4: Add back F dimension: (B, C, H, W) -> (B, C, 1, H, W) + let x = x.unsqueeze(2)?; + + Ok(x) + } else { + // General case + let x = x.reshape((b, f_tokens, h_tokens, w_tokens, pf * ph * pw * out_channels))?; + let x = x.reshape((b, f_tokens, h_tokens, w_tokens * pf, ph, pw * out_channels))?; + let x = x.permute((0, 5, 1, 3, 2, 4))?.contiguous()?; + let x = x.reshape((b, out_channels, f, h, w))?; + Ok(x) + } +} + +/// Create 3D coordinate grid for RoPE position IDs +/// size: (F, H, W) +/// start: (f0, h0, w0) +/// output: (F*H*W, 3) +pub fn create_coordinate_grid( + size: (usize, usize, usize), + start: (usize, usize, usize), + device: &Device, +) -> Result { + let (f, h, w) = size; + let (f0, h0, w0) = start; + + let mut coords = Vec::with_capacity(f * h * w * 3); + for fi in 0..f { + for hi in 0..h { + for wi in 0..w { + coords.push((f0 + fi) as u32); + coords.push((h0 + hi) as u32); + coords.push((w0 + wi) as u32); + } + } + } + + Tensor::from_vec(coords, (f * h * w, 3), device) +} + +// ==================== ZImageTransformer2DModel ==================== + +/// Z-Image Transformer 2D Model +#[derive(Debug, Clone)] +pub struct ZImageTransformer2DModel { + t_embedder: TimestepEmbedder, + cap_embedder_norm: RmsNorm, + cap_embedder_linear: candle_nn::Linear, + x_embedder: candle_nn::Linear, + final_layer: FinalLayer, + #[allow(dead_code)] + x_pad_token: Tensor, + #[allow(dead_code)] + cap_pad_token: Tensor, + noise_refiner: Vec, + context_refiner: Vec, + layers: Vec, + rope_embedder: RopeEmbedder, + cfg: Config, +} + +impl ZImageTransformer2DModel { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let device = vb.device(); + let dtype = vb.dtype(); + + // TimestepEmbedder + let adaln_dim = cfg.dim.min(ADALN_EMBED_DIM); + let t_embedder = TimestepEmbedder::new(adaln_dim, 1024, vb.pp("t_embedder"))?; + + // Caption embedder + let cap_embedder_norm = RmsNorm::new( + cfg.cap_feat_dim, + cfg.norm_eps, + vb.pp("cap_embedder").pp("0"), + )?; + let cap_embedder_linear = linear(cfg.cap_feat_dim, cfg.dim, vb.pp("cap_embedder").pp("1"))?; + + // Patch embedder (assuming patch_size=2, f_patch_size=1) + let patch_dim = cfg.all_f_patch_size[0] + * cfg.all_patch_size[0] + * cfg.all_patch_size[0] + * cfg.in_channels; + let x_embedder = linear(patch_dim, cfg.dim, vb.pp("all_x_embedder").pp("2-1"))?; + + // Final layer + let out_channels = cfg.all_patch_size[0] + * cfg.all_patch_size[0] + * cfg.all_f_patch_size[0] + * cfg.in_channels; + let final_layer = + FinalLayer::new(cfg.dim, out_channels, vb.pp("all_final_layer").pp("2-1"))?; + + // Pad tokens + let x_pad_token = vb.get((1, cfg.dim), "x_pad_token")?; + let cap_pad_token = vb.get((1, cfg.dim), "cap_pad_token")?; + + // Noise refiner (with modulation) + let mut noise_refiner = Vec::with_capacity(cfg.n_refiner_layers); + for i in 0..cfg.n_refiner_layers { + noise_refiner.push(ZImageTransformerBlock::new( + cfg, + true, + vb.pp("noise_refiner").pp(i), + )?); + } + + // Context refiner (without modulation) + let mut context_refiner = Vec::with_capacity(cfg.n_refiner_layers); + for i in 0..cfg.n_refiner_layers { + context_refiner.push(ZImageTransformerBlock::new( + cfg, + false, + vb.pp("context_refiner").pp(i), + )?); + } + + // Main layers (with modulation) + let mut layers = Vec::with_capacity(cfg.n_layers); + for i in 0..cfg.n_layers { + layers.push(ZImageTransformerBlock::new( + cfg, + true, + vb.pp("layers").pp(i), + )?); + } + + // RoPE embedder + let rope_embedder = RopeEmbedder::new( + cfg.rope_theta, + cfg.axes_dims.clone(), + cfg.axes_lens.clone(), + device, + dtype, + )?; + + Ok(Self { + t_embedder, + cap_embedder_norm, + cap_embedder_linear, + x_embedder, + final_layer, + x_pad_token, + cap_pad_token, + noise_refiner, + context_refiner, + layers, + rope_embedder, + cfg: cfg.clone(), + }) + } + + /// Forward pass + /// + /// # Arguments + /// * `x` - Latent tensor (B, C, F, H, W) + /// * `t` - Timesteps [0, 1] (B,) + /// * `cap_feats` - Caption features (B, text_len, cap_feat_dim) + /// * `cap_mask` - Caption attention mask (B, text_len), 1=valid, 0=padding + pub fn forward( + &self, + x: &Tensor, + t: &Tensor, + cap_feats: &Tensor, + cap_mask: &Tensor, + ) -> Result { + let device = x.device(); + let (b, _c, f, h, w) = x.dims5()?; + let patch_size = self.cfg.all_patch_size[0]; + let f_patch_size = self.cfg.all_f_patch_size[0]; + + // 1. Timestep embedding + let t_scaled = (t * self.cfg.t_scale)?; + let adaln_input = self.t_embedder.forward(&t_scaled)?; // (B, 256) + + // 2. Patchify and embed image + let (x_patches, orig_size) = patchify(x, patch_size, f_patch_size)?; + let mut x = x_patches.apply(&self.x_embedder)?; // (B, img_seq, dim) + let img_seq_len = x.dim(1)?; + + // 3. Create image position IDs + let f_tokens = f / f_patch_size; + let h_tokens = h / patch_size; + let w_tokens = w / patch_size; + let text_len = cap_feats.dim(1)?; + + let x_pos_ids = create_coordinate_grid( + (f_tokens, h_tokens, w_tokens), + (text_len + 1, 0, 0), // offset for text + device, + )?; + let (x_cos, x_sin) = self.rope_embedder.forward(&x_pos_ids)?; + + // 4. Caption embedding + let cap_normed = self.cap_embedder_norm.forward(cap_feats)?; + let mut cap = cap_normed.apply(&self.cap_embedder_linear)?; // (B, text_len, dim) + + // 5. Create caption position IDs + let cap_pos_ids = create_coordinate_grid((text_len, 1, 1), (1, 0, 0), device)?; + let (cap_cos, cap_sin) = self.rope_embedder.forward(&cap_pos_ids)?; + + // 6. Create attention masks + let x_attn_mask = Tensor::ones((b, img_seq_len), DType::U8, device)?; + let cap_attn_mask = cap_mask.to_dtype(DType::U8)?; + + // 7. Noise refiner (process image with modulation) + for layer in &self.noise_refiner { + x = layer.forward(&x, Some(&x_attn_mask), &x_cos, &x_sin, Some(&adaln_input))?; + } + + // 8. Context refiner (process text without modulation) + for layer in &self.context_refiner { + cap = layer.forward(&cap, Some(&cap_attn_mask), &cap_cos, &cap_sin, None)?; + } + + // 9. Concatenate image and text: [image_tokens, text_tokens] + let unified = Tensor::cat(&[&x, &cap], 1)?; // (B, img_seq + text_len, dim) + + // 10. Create unified position IDs and attention mask + let unified_pos_ids = Tensor::cat(&[&x_pos_ids, &cap_pos_ids], 0)?; + let (unified_cos, unified_sin) = self.rope_embedder.forward(&unified_pos_ids)?; + let unified_attn_mask = Tensor::cat(&[&x_attn_mask, &cap_attn_mask], 1)?; + + // 11. Main transformer layers + let mut unified = unified; + for layer in &self.layers { + unified = layer.forward( + &unified, + Some(&unified_attn_mask), + &unified_cos, + &unified_sin, + Some(&adaln_input), + )?; + } + + // 12. Final layer (only on image portion) + let x_out = unified.narrow(1, 0, img_seq_len)?; + let x_out = self.final_layer.forward(&x_out, &adaln_input)?; + + // 13. Unpatchify + unpatchify( + &x_out, + orig_size, + patch_size, + f_patch_size, + self.cfg.in_channels, + ) + } + + /// Get model configuration + pub fn config(&self) -> &Config { + &self.cfg + } +} diff --git a/patches/candle-transformers/src/models/z_image/vae.rs b/patches/candle-transformers/src/models/z_image/vae.rs new file mode 100644 index 0000000000..c78ee3123b --- /dev/null +++ b/patches/candle-transformers/src/models/z_image/vae.rs @@ -0,0 +1,684 @@ +//! Z-Image VAE (AutoEncoderKL) - Diffusers Format +//! +//! This VAE implementation uses the diffusers weight naming format, +//! which is different from the Flux autoencoder original format. +//! +//! Key differences from Flux autoencoder: +//! 1. Weight paths: `encoder.down_blocks.{i}.resnets.{j}.*` vs `encoder.down.{i}.block.{j}.*` +//! 2. Attention naming: `to_q/to_k/to_v/to_out.0.*` vs `q/k/v/proj_out.*` +//! 3. Shortcut naming: `conv_shortcut.*` vs `nin_shortcut.*` + +use candle::{Module, Result, Tensor, D}; +use candle_nn::{conv2d, group_norm, Conv2d, Conv2dConfig, GroupNorm, VarBuilder}; + +// ==================== Config ==================== + +/// VAE configuration +#[derive(Debug, Clone, serde::Deserialize)] +pub struct VaeConfig { + #[serde(default = "default_in_channels")] + pub in_channels: usize, + #[serde(default = "default_out_channels")] + pub out_channels: usize, + #[serde(default = "default_latent_channels")] + pub latent_channels: usize, + #[serde(default = "default_block_out_channels")] + pub block_out_channels: Vec, + #[serde(default = "default_layers_per_block")] + pub layers_per_block: usize, + #[serde(default = "default_scaling_factor")] + pub scaling_factor: f64, + #[serde(default = "default_shift_factor")] + pub shift_factor: f64, + #[serde(default = "default_norm_num_groups")] + pub norm_num_groups: usize, +} + +fn default_in_channels() -> usize { + 3 +} +fn default_out_channels() -> usize { + 3 +} +fn default_latent_channels() -> usize { + 16 +} +fn default_block_out_channels() -> Vec { + vec![128, 256, 512, 512] +} +fn default_layers_per_block() -> usize { + 2 +} +fn default_scaling_factor() -> f64 { + 0.3611 +} +fn default_shift_factor() -> f64 { + 0.1159 +} +fn default_norm_num_groups() -> usize { + 32 +} + +impl Default for VaeConfig { + fn default() -> Self { + Self::z_image() + } +} + +impl VaeConfig { + /// Create configuration for Z-Image VAE + pub fn z_image() -> Self { + Self { + in_channels: 3, + out_channels: 3, + latent_channels: 16, + block_out_channels: vec![128, 256, 512, 512], + layers_per_block: 2, + scaling_factor: 0.3611, + shift_factor: 0.1159, + norm_num_groups: 32, + } + } +} + +// ==================== Attention ==================== + +fn scaled_dot_product_attention(q: &Tensor, k: &Tensor, v: &Tensor) -> Result { + let dim = q.dim(D::Minus1)?; + let scale_factor = 1.0 / (dim as f64).sqrt(); + let attn_weights = (q.matmul(&k.t()?)? * scale_factor)?; + candle_nn::ops::softmax_last_dim(&attn_weights)?.matmul(v) +} + +/// VAE Attention block (diffusers format) +/// +/// Note: VAE attention uses Linear with bias (2D weight shape) +/// Unlike Transformer attention which uses linear_no_bias +#[derive(Debug, Clone)] +struct Attention { + group_norm: GroupNorm, + to_q: candle_nn::Linear, + to_k: candle_nn::Linear, + to_v: candle_nn::Linear, + to_out: candle_nn::Linear, +} + +impl Attention { + fn new(channels: usize, num_groups: usize, vb: VarBuilder) -> Result { + let group_norm = group_norm(num_groups, channels, 1e-6, vb.pp("group_norm"))?; + // VAE attention uses Linear with bias + let to_q = candle_nn::linear(channels, channels, vb.pp("to_q"))?; + let to_k = candle_nn::linear(channels, channels, vb.pp("to_k"))?; + let to_v = candle_nn::linear(channels, channels, vb.pp("to_v"))?; + let to_out = candle_nn::linear(channels, channels, vb.pp("to_out").pp("0"))?; + Ok(Self { + group_norm, + to_q, + to_k, + to_v, + to_out, + }) + } +} + +impl Module for Attention { + fn forward(&self, xs: &Tensor) -> Result { + let residual = xs; + let (b, c, h, w) = xs.dims4()?; + + // GroupNorm + let xs = xs.apply(&self.group_norm)?; + + // (B, C, H, W) -> (B, H, W, C) -> (B*H*W, C) + let xs = xs.permute((0, 2, 3, 1))?.reshape((b * h * w, c))?; + + // Linear projections + let q = xs.apply(&self.to_q)?; // (B*H*W, C) + let k = xs.apply(&self.to_k)?; + let v = xs.apply(&self.to_v)?; + + // Reshape for attention: (B*H*W, C) -> (B, H*W, C) -> (B, 1, H*W, C) + let q = q.reshape((b, h * w, c))?.unsqueeze(1)?; + let k = k.reshape((b, h * w, c))?.unsqueeze(1)?; + let v = v.reshape((b, h * w, c))?.unsqueeze(1)?; + + // Scaled dot-product attention + let xs = scaled_dot_product_attention(&q, &k, &v)?; + + // (B, 1, H*W, C) -> (B*H*W, C) + let xs = xs.squeeze(1)?.reshape((b * h * w, c))?; + + // Output projection + let xs = xs.apply(&self.to_out)?; + + // (B*H*W, C) -> (B, H, W, C) -> (B, C, H, W) + let xs = xs.reshape((b, h, w, c))?.permute((0, 3, 1, 2))?; + + // Residual connection + xs + residual + } +} + +// ==================== ResnetBlock2D ==================== + +/// ResNet block (diffusers format) +#[derive(Debug, Clone)] +struct ResnetBlock2D { + norm1: GroupNorm, + conv1: Conv2d, + norm2: GroupNorm, + conv2: Conv2d, + conv_shortcut: Option, +} + +impl ResnetBlock2D { + fn new( + in_channels: usize, + out_channels: usize, + num_groups: usize, + vb: VarBuilder, + ) -> Result { + let conv_cfg = Conv2dConfig { + padding: 1, + ..Default::default() + }; + + let norm1 = group_norm(num_groups, in_channels, 1e-6, vb.pp("norm1"))?; + let conv1 = conv2d(in_channels, out_channels, 3, conv_cfg, vb.pp("conv1"))?; + let norm2 = group_norm(num_groups, out_channels, 1e-6, vb.pp("norm2"))?; + let conv2 = conv2d(out_channels, out_channels, 3, conv_cfg, vb.pp("conv2"))?; + + let conv_shortcut = if in_channels != out_channels { + Some(conv2d( + in_channels, + out_channels, + 1, + Default::default(), + vb.pp("conv_shortcut"), + )?) + } else { + None + }; + + Ok(Self { + norm1, + conv1, + norm2, + conv2, + conv_shortcut, + }) + } +} + +impl Module for ResnetBlock2D { + fn forward(&self, xs: &Tensor) -> Result { + let h = xs + .apply(&self.norm1)? + .apply(&candle_nn::Activation::Swish)? + .apply(&self.conv1)? + .apply(&self.norm2)? + .apply(&candle_nn::Activation::Swish)? + .apply(&self.conv2)?; + + match &self.conv_shortcut { + Some(conv) => xs.apply(conv)? + h, + None => xs + h, + } + } +} + +// ==================== DownEncoderBlock2D ==================== + +#[derive(Debug, Clone)] +struct Downsample2D { + conv: Conv2d, +} + +impl Downsample2D { + fn new(channels: usize, vb: VarBuilder) -> Result { + let conv_cfg = Conv2dConfig { + stride: 2, + padding: 0, + ..Default::default() + }; + let conv = conv2d(channels, channels, 3, conv_cfg, vb.pp("conv"))?; + Ok(Self { conv }) + } +} + +impl Module for Downsample2D { + fn forward(&self, xs: &Tensor) -> Result { + // Manual padding: (0, 1, 0, 1) for right=1, bottom=1 + let xs = xs.pad_with_zeros(D::Minus1, 0, 1)?; // width: right + let xs = xs.pad_with_zeros(D::Minus2, 0, 1)?; // height: bottom + xs.apply(&self.conv) + } +} + +#[derive(Debug, Clone)] +struct DownEncoderBlock2D { + resnets: Vec, + downsampler: Option, +} + +impl DownEncoderBlock2D { + fn new( + in_channels: usize, + out_channels: usize, + num_layers: usize, + num_groups: usize, + add_downsample: bool, + vb: VarBuilder, + ) -> Result { + let mut resnets = Vec::with_capacity(num_layers); + let vb_resnets = vb.pp("resnets"); + + for i in 0..num_layers { + let in_c = if i == 0 { in_channels } else { out_channels }; + resnets.push(ResnetBlock2D::new( + in_c, + out_channels, + num_groups, + vb_resnets.pp(i), + )?); + } + + let downsampler = if add_downsample { + Some(Downsample2D::new( + out_channels, + vb.pp("downsamplers").pp("0"), + )?) + } else { + None + }; + + Ok(Self { + resnets, + downsampler, + }) + } +} + +impl Module for DownEncoderBlock2D { + fn forward(&self, xs: &Tensor) -> Result { + let mut h = xs.clone(); + for resnet in &self.resnets { + h = h.apply(resnet)?; + } + if let Some(ds) = &self.downsampler { + h = h.apply(ds)?; + } + Ok(h) + } +} + +// ==================== UpDecoderBlock2D ==================== + +#[derive(Debug, Clone)] +struct Upsample2D { + conv: Conv2d, +} + +impl Upsample2D { + fn new(channels: usize, vb: VarBuilder) -> Result { + let conv_cfg = Conv2dConfig { + padding: 1, + ..Default::default() + }; + let conv = conv2d(channels, channels, 3, conv_cfg, vb.pp("conv"))?; + Ok(Self { conv }) + } +} + +impl Module for Upsample2D { + fn forward(&self, xs: &Tensor) -> Result { + let (_, _, h, w) = xs.dims4()?; + xs.upsample_nearest2d(h * 2, w * 2)?.apply(&self.conv) + } +} + +#[derive(Debug, Clone)] +struct UpDecoderBlock2D { + resnets: Vec, + upsampler: Option, +} + +impl UpDecoderBlock2D { + fn new( + in_channels: usize, + out_channels: usize, + num_layers: usize, // decoder has num_layers + 1 resnets per block + num_groups: usize, + add_upsample: bool, + vb: VarBuilder, + ) -> Result { + let mut resnets = Vec::with_capacity(num_layers + 1); + let vb_resnets = vb.pp("resnets"); + + for i in 0..=num_layers { + let in_c = if i == 0 { in_channels } else { out_channels }; + resnets.push(ResnetBlock2D::new( + in_c, + out_channels, + num_groups, + vb_resnets.pp(i), + )?); + } + + let upsampler = if add_upsample { + Some(Upsample2D::new(out_channels, vb.pp("upsamplers").pp("0"))?) + } else { + None + }; + + Ok(Self { resnets, upsampler }) + } +} + +impl Module for UpDecoderBlock2D { + fn forward(&self, xs: &Tensor) -> Result { + let mut h = xs.clone(); + for resnet in &self.resnets { + h = h.apply(resnet)?; + } + if let Some(us) = &self.upsampler { + h = h.apply(us)?; + } + Ok(h) + } +} + +// ==================== UNetMidBlock2D ==================== + +#[derive(Debug, Clone)] +struct UNetMidBlock2D { + resnet_0: ResnetBlock2D, + attention: Attention, + resnet_1: ResnetBlock2D, +} + +impl UNetMidBlock2D { + fn new(channels: usize, num_groups: usize, vb: VarBuilder) -> Result { + let resnet_0 = + ResnetBlock2D::new(channels, channels, num_groups, vb.pp("resnets").pp("0"))?; + let attention = Attention::new(channels, num_groups, vb.pp("attentions").pp("0"))?; + let resnet_1 = + ResnetBlock2D::new(channels, channels, num_groups, vb.pp("resnets").pp("1"))?; + Ok(Self { + resnet_0, + attention, + resnet_1, + }) + } +} + +impl Module for UNetMidBlock2D { + fn forward(&self, xs: &Tensor) -> Result { + xs.apply(&self.resnet_0)? + .apply(&self.attention)? + .apply(&self.resnet_1) + } +} + +// ==================== Encoder ==================== + +/// VAE Encoder +#[derive(Debug, Clone)] +pub struct Encoder { + conv_in: Conv2d, + down_blocks: Vec, + mid_block: UNetMidBlock2D, + conv_norm_out: GroupNorm, + conv_out: Conv2d, +} + +impl Encoder { + pub fn new(cfg: &VaeConfig, vb: VarBuilder) -> Result { + let conv_cfg = Conv2dConfig { + padding: 1, + ..Default::default() + }; + let conv_in = conv2d( + cfg.in_channels, + cfg.block_out_channels[0], + 3, + conv_cfg, + vb.pp("conv_in"), + )?; + + let mut down_blocks = Vec::with_capacity(cfg.block_out_channels.len()); + let vb_down = vb.pp("down_blocks"); + + for (i, &out_channels) in cfg.block_out_channels.iter().enumerate() { + let in_channels = if i == 0 { + cfg.block_out_channels[0] + } else { + cfg.block_out_channels[i - 1] + }; + let add_downsample = i < cfg.block_out_channels.len() - 1; + down_blocks.push(DownEncoderBlock2D::new( + in_channels, + out_channels, + cfg.layers_per_block, + cfg.norm_num_groups, + add_downsample, + vb_down.pp(i), + )?); + } + + let mid_channels = *cfg.block_out_channels.last().unwrap(); + let mid_block = UNetMidBlock2D::new(mid_channels, cfg.norm_num_groups, vb.pp("mid_block"))?; + + let conv_norm_out = group_norm( + cfg.norm_num_groups, + mid_channels, + 1e-6, + vb.pp("conv_norm_out"), + )?; + let conv_out = conv2d( + mid_channels, + 2 * cfg.latent_channels, + 3, + conv_cfg, + vb.pp("conv_out"), + )?; + + Ok(Self { + conv_in, + down_blocks, + mid_block, + conv_norm_out, + conv_out, + }) + } +} + +impl Module for Encoder { + fn forward(&self, xs: &Tensor) -> Result { + let mut h = xs.apply(&self.conv_in)?; + for block in &self.down_blocks { + h = h.apply(block)?; + } + h.apply(&self.mid_block)? + .apply(&self.conv_norm_out)? + .apply(&candle_nn::Activation::Swish)? + .apply(&self.conv_out) + } +} + +// ==================== Decoder ==================== + +/// VAE Decoder +#[derive(Debug, Clone)] +pub struct Decoder { + conv_in: Conv2d, + mid_block: UNetMidBlock2D, + up_blocks: Vec, + conv_norm_out: GroupNorm, + conv_out: Conv2d, +} + +impl Decoder { + pub fn new(cfg: &VaeConfig, vb: VarBuilder) -> Result { + let conv_cfg = Conv2dConfig { + padding: 1, + ..Default::default() + }; + let mid_channels = *cfg.block_out_channels.last().unwrap(); + + let conv_in = conv2d( + cfg.latent_channels, + mid_channels, + 3, + conv_cfg, + vb.pp("conv_in"), + )?; + let mid_block = UNetMidBlock2D::new(mid_channels, cfg.norm_num_groups, vb.pp("mid_block"))?; + + // Decoder up_blocks order is reversed from encoder down_blocks + let reversed_channels: Vec = cfg.block_out_channels.iter().rev().cloned().collect(); + let mut up_blocks = Vec::with_capacity(reversed_channels.len()); + let vb_up = vb.pp("up_blocks"); + + for (i, &out_channels) in reversed_channels.iter().enumerate() { + let in_channels = if i == 0 { + mid_channels + } else { + reversed_channels[i - 1] + }; + let add_upsample = i < reversed_channels.len() - 1; + up_blocks.push(UpDecoderBlock2D::new( + in_channels, + out_channels, + cfg.layers_per_block, + cfg.norm_num_groups, + add_upsample, + vb_up.pp(i), + )?); + } + + let final_channels = *reversed_channels.last().unwrap(); + let conv_norm_out = group_norm( + cfg.norm_num_groups, + final_channels, + 1e-6, + vb.pp("conv_norm_out"), + )?; + let conv_out = conv2d( + final_channels, + cfg.out_channels, + 3, + conv_cfg, + vb.pp("conv_out"), + )?; + + Ok(Self { + conv_in, + mid_block, + up_blocks, + conv_norm_out, + conv_out, + }) + } +} + +impl Module for Decoder { + fn forward(&self, xs: &Tensor) -> Result { + let mut h = xs.apply(&self.conv_in)?.apply(&self.mid_block)?; + for block in &self.up_blocks { + h = h.apply(block)?; + } + h.apply(&self.conv_norm_out)? + .apply(&candle_nn::Activation::Swish)? + .apply(&self.conv_out) + } +} + +// ==================== DiagonalGaussian ==================== + +/// Diagonal Gaussian distribution sampling (VAE reparameterization trick) +#[derive(Debug, Clone)] +pub struct DiagonalGaussian { + sample: bool, +} + +impl DiagonalGaussian { + pub fn new(sample: bool) -> Self { + Self { sample } + } +} + +impl Module for DiagonalGaussian { + fn forward(&self, xs: &Tensor) -> Result { + let chunks = xs.chunk(2, 1)?; // Split along channel dimension + let mean = &chunks[0]; + let logvar = &chunks[1]; + + if self.sample { + let std = (logvar * 0.5)?.exp()?; + mean + (std * mean.randn_like(0., 1.)?)? + } else { + Ok(mean.clone()) + } + } +} + +// ==================== AutoEncoderKL ==================== + +/// Z-Image VAE (AutoEncoderKL) - Diffusers Format +#[derive(Debug, Clone)] +pub struct AutoEncoderKL { + encoder: Encoder, + decoder: Decoder, + reg: DiagonalGaussian, + scale_factor: f64, + shift_factor: f64, +} + +impl AutoEncoderKL { + pub fn new(cfg: &VaeConfig, vb: VarBuilder) -> Result { + let encoder = Encoder::new(cfg, vb.pp("encoder"))?; + let decoder = Decoder::new(cfg, vb.pp("decoder"))?; + let reg = DiagonalGaussian::new(true); + + Ok(Self { + encoder, + decoder, + reg, + scale_factor: cfg.scaling_factor, + shift_factor: cfg.shift_factor, + }) + } + + /// Encode image to latent space + /// xs: (B, 3, H, W) RGB image, range [-1, 1] + /// Returns: (B, latent_channels, H/8, W/8) + pub fn encode(&self, xs: &Tensor) -> Result { + let z = xs.apply(&self.encoder)?.apply(&self.reg)?; + (z - self.shift_factor)? * self.scale_factor + } + + /// Decode latent to image + /// xs: (B, latent_channels, H/8, W/8) + /// Returns: (B, 3, H, W) RGB image, range [-1, 1] + pub fn decode(&self, xs: &Tensor) -> Result { + let xs = ((xs / self.scale_factor)? + self.shift_factor)?; + xs.apply(&self.decoder) + } + + /// Get scaling factor + pub fn scale_factor(&self) -> f64 { + self.scale_factor + } + + /// Get shift factor + pub fn shift_factor(&self) -> f64 { + self.shift_factor + } +} + +impl Module for AutoEncoderKL { + fn forward(&self, xs: &Tensor) -> Result { + self.decode(&self.encode(xs)?) + } +} diff --git a/patches/candle-transformers/src/object_detection.rs b/patches/candle-transformers/src/object_detection.rs new file mode 100644 index 0000000000..d1b78cfa25 --- /dev/null +++ b/patches/candle-transformers/src/object_detection.rs @@ -0,0 +1,116 @@ +//! Bounding Boxes and Intersection +//! +//! This module provides functionality for handling bounding boxes and their manipulation, +//! particularly in the context of object detection. It includes tools for calculating +//! intersection over union (IoU) and non-maximum suppression (NMS). + +/// A bounding box around an object. +#[derive(Debug, Clone)] +pub struct Bbox { + pub xmin: f32, + pub ymin: f32, + pub xmax: f32, + pub ymax: f32, + pub confidence: f32, + pub data: D, +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct KeyPoint { + pub x: f32, + pub y: f32, + pub mask: f32, +} + +/// Intersection over union of two bounding boxes. +pub fn iou(b1: &Bbox, b2: &Bbox) -> f32 { + let b1_area = (b1.xmax - b1.xmin + 1.) * (b1.ymax - b1.ymin + 1.); + let b2_area = (b2.xmax - b2.xmin + 1.) * (b2.ymax - b2.ymin + 1.); + let i_xmin = b1.xmin.max(b2.xmin); + let i_xmax = b1.xmax.min(b2.xmax); + let i_ymin = b1.ymin.max(b2.ymin); + let i_ymax = b1.ymax.min(b2.ymax); + let i_area = (i_xmax - i_xmin + 1.).max(0.) * (i_ymax - i_ymin + 1.).max(0.); + i_area / (b1_area + b2_area - i_area) +} + +pub fn non_maximum_suppression(bboxes: &mut [Vec>], threshold: f32) { + // Perform non-maximum suppression. + for bboxes_for_class in bboxes.iter_mut() { + bboxes_for_class.sort_by(|b1, b2| b2.confidence.partial_cmp(&b1.confidence).unwrap()); + let mut current_index = 0; + for index in 0..bboxes_for_class.len() { + let mut drop = false; + for prev_index in 0..current_index { + let iou = iou(&bboxes_for_class[prev_index], &bboxes_for_class[index]); + if iou > threshold { + drop = true; + break; + } + } + if !drop { + bboxes_for_class.swap(current_index, index); + current_index += 1; + } + } + bboxes_for_class.truncate(current_index); + } +} + +// Updates confidences starting at highest and comparing subsequent boxes. +fn update_confidences( + bboxes_for_class: &[Bbox], + updated_confidences: &mut [f32], + iou_threshold: f32, + sigma: f32, +) { + let len = bboxes_for_class.len(); + for current_index in 0..len { + let current_bbox = &bboxes_for_class[current_index]; + for index in (current_index + 1)..len { + let iou_val = iou(current_bbox, &bboxes_for_class[index]); + if iou_val > iou_threshold { + // Decay calculation from page 4 of: https://arxiv.org/pdf/1704.04503 + let decay = (-iou_val * iou_val / sigma).exp(); + let updated_confidence = bboxes_for_class[index].confidence * decay; + updated_confidences[index] = updated_confidence; + } + } + } +} + +// Sorts the bounding boxes by confidence and applies soft non-maximum suppression. +// This function is based on the algorithm described in https://arxiv.org/pdf/1704.04503 +pub fn soft_non_maximum_suppression( + bboxes: &mut [Vec>], + iou_threshold: Option, + confidence_threshold: Option, + sigma: Option, +) { + let iou_threshold = iou_threshold.unwrap_or(0.5); + let confidence_threshold = confidence_threshold.unwrap_or(0.1); + let sigma = sigma.unwrap_or(0.5); + + for bboxes_for_class in bboxes.iter_mut() { + // Sort boxes by confidence in descending order + bboxes_for_class.sort_by(|b1, b2| b2.confidence.partial_cmp(&b1.confidence).unwrap()); + let mut updated_confidences = bboxes_for_class + .iter() + .map(|bbox| bbox.confidence) + .collect::>(); + update_confidences( + bboxes_for_class, + &mut updated_confidences, + iou_threshold, + sigma, + ); + // Update confidences, set to 0.0 if below threshold + for (i, &confidence) in updated_confidences.iter().enumerate() { + bboxes_for_class[i].confidence = if confidence < confidence_threshold { + 0.0 + } else { + confidence + }; + } + } +} diff --git a/patches/candle-transformers/src/pipelines/mod.rs b/patches/candle-transformers/src/pipelines/mod.rs new file mode 100644 index 0000000000..d1bc14e2ac --- /dev/null +++ b/patches/candle-transformers/src/pipelines/mod.rs @@ -0,0 +1 @@ +pub mod text_generation; diff --git a/patches/candle-transformers/src/pipelines/text_generation.rs b/patches/candle-transformers/src/pipelines/text_generation.rs new file mode 100644 index 0000000000..8b13789179 --- /dev/null +++ b/patches/candle-transformers/src/pipelines/text_generation.rs @@ -0,0 +1 @@ + diff --git a/patches/candle-transformers/src/quantized_nn.rs b/patches/candle-transformers/src/quantized_nn.rs new file mode 100644 index 0000000000..4a83253d2e --- /dev/null +++ b/patches/candle-transformers/src/quantized_nn.rs @@ -0,0 +1,126 @@ +//! Utilities for quanitized network layers +//! +//! This module contains various implementations of standard neural network layers, modules and +//! utilities including embedding, linear layers, and various normalization techniques. +//! Most implementations provide quantized weights support. + +use crate::models::with_tracing::QMatMul; +use crate::quantized_var_builder::VarBuilder; +use candle::quantized::QTensor; +use candle::{Module, Result, Tensor}; + +#[derive(Debug, Clone)] +pub struct Embedding { + inner: candle_nn::Embedding, + span: tracing::Span, +} + +impl Embedding { + pub fn new(d1: usize, d2: usize, vb: VarBuilder) -> Result { + let embeddings = vb.get((d1, d2), "weight")?.dequantize(vb.device())?; + let inner = candle_nn::Embedding::new(embeddings, d2); + let span = tracing::span!(tracing::Level::TRACE, "embedding"); + Ok(Self { inner, span }) + } + + pub fn embeddings(&self) -> &Tensor { + self.inner.embeddings() + } +} + +impl Module for Embedding { + fn forward(&self, xs: &Tensor) -> Result { + let _enter = self.span.enter(); + self.inner.forward(xs) + } +} + +#[derive(Debug, Clone)] +pub struct Linear { + weight: QMatMul, + bias: Option, +} + +impl Linear { + pub fn from_arc(weight: std::sync::Arc, bias: Option) -> Result { + let weight = QMatMul::from_weights(weight)?; + Ok(Self { weight, bias }) + } + + pub fn from_weights(weight: QMatMul, bias: Option) -> Self { + Self { weight, bias } + } +} + +impl Module for Linear { + fn forward(&self, x: &Tensor) -> candle::Result { + let x = x.apply(&self.weight)?; + match &self.bias { + None => Ok(x), + Some(bias) => x.broadcast_add(bias), + } + } +} + +pub fn linear_b(in_dim: usize, out_dim: usize, bias: bool, vb: VarBuilder) -> Result { + let bias = if bias { + Some(vb.get(out_dim, "bias")?.dequantize(vb.device())?) + } else { + None + }; + let weight = QMatMul::new(in_dim, out_dim, vb)?; + Ok(Linear { weight, bias }) +} + +pub fn linear(in_dim: usize, out_dim: usize, vb: VarBuilder) -> Result { + let bias = vb.get(out_dim, "bias")?.dequantize(vb.device())?; + let weight = QMatMul::new(in_dim, out_dim, vb)?; + Ok(Linear { + weight, + bias: Some(bias), + }) +} + +pub fn layer_norm(size: usize, eps: f64, vb: VarBuilder) -> Result { + let weight = vb.get(size, "weight")?.dequantize(vb.device())?; + let bias = vb.get(size, "bias")?.dequantize(vb.device())?; + Ok(candle_nn::LayerNorm::new(weight, bias, eps)) +} + +pub fn layer_norm_no_bias(size: usize, eps: f64, vb: VarBuilder) -> Result { + let weight = vb.get(size, "weight")?.dequantize(vb.device())?; + Ok(candle_nn::LayerNorm::new_no_bias(weight, eps)) +} + +pub fn linear_no_bias(in_dim: usize, out_dim: usize, vb: VarBuilder) -> Result { + let weight = QMatMul::new(in_dim, out_dim, vb)?; + Ok(Linear { weight, bias: None }) +} + +#[derive(Debug, Clone)] +pub struct RmsNorm { + weight: Tensor, + eps: f64, + span: tracing::Span, +} + +impl RmsNorm { + pub fn new(size: usize, eps: f64, vb: VarBuilder) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "rms-norm"); + let weight = vb.get(size, "weight")?.dequantize(vb.device())?; + Ok(Self { weight, eps, span }) + } + + pub fn from_qtensor(weight: QTensor, eps: f64) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "rms-norm"); + let weight = weight.dequantize(&weight.device())?; + Ok(Self { weight, eps, span }) + } +} + +impl Module for RmsNorm { + fn forward(&self, x: &Tensor) -> Result { + let _enter = self.span.enter(); + candle_nn::ops::rms_norm(x, &self.weight, self.eps as f32) + } +} diff --git a/patches/candle-transformers/src/quantized_var_builder.rs b/patches/candle-transformers/src/quantized_var_builder.rs new file mode 100644 index 0000000000..2ac64aa5e7 --- /dev/null +++ b/patches/candle-transformers/src/quantized_var_builder.rs @@ -0,0 +1,104 @@ +//! Varbuilder for Loading gguf files +//! +//! VarBuilder is a utility to store quantized tensors from a [GGUF model file](https://huggingface.co/docs/hub/gguf). +//! These tensors can be loaded from disk using `from_gguf` or from an in-memory +//! buffer using `from_gguf_buffer`. + +use candle::quantized::QTensor; +use candle::{Device, Result, Shape}; +use std::sync::Arc; + +// VarBuilder specialized for QTensors +#[derive(Clone)] +pub struct VarBuilder { + data: Arc>>, + path: Vec, + device: Device, +} + +impl VarBuilder { + pub fn from_gguf>(p: P, device: &Device) -> Result { + let mut file = std::fs::File::open(p)?; + let content = candle::quantized::gguf_file::Content::read(&mut file)?; + let mut data = std::collections::HashMap::new(); + for tensor_name in content.tensor_infos.keys() { + let tensor = content.tensor(&mut file, tensor_name, device)?; + data.insert(tensor_name.to_string(), Arc::new(tensor)); + } + Ok(Self { + data: Arc::new(data), + path: Vec::new(), + device: device.clone(), + }) + } + + pub fn from_gguf_buffer(buffer: &[u8], device: &Device) -> Result { + let mut cursor = std::io::Cursor::new(buffer); + let content = candle::quantized::gguf_file::Content::read(&mut cursor)?; + let mut data = std::collections::HashMap::new(); + for tensor_name in content.tensor_infos.keys() { + let tensor = content.tensor(&mut cursor, tensor_name, device)?; + data.insert(tensor_name.to_string(), Arc::new(tensor)); + } + Ok(Self { + data: Arc::new(data), + path: Vec::new(), + device: device.clone(), + }) + } + + pub fn pp(&self, s: S) -> Self { + let mut path = self.path.clone(); + path.push(s.to_string()); + Self { + data: self.data.clone(), + path, + device: self.device.clone(), + } + } + + fn path(&self, tensor_name: &str) -> String { + if self.path.is_empty() { + tensor_name.to_string() + } else { + [&self.path.join("."), tensor_name].join(".") + } + } + + pub fn get>(&self, s: S, name: &str) -> Result> { + let path = self.path(name); + match self.data.get(&path) { + None => { + candle::bail!("cannot find tensor {path}") + } + Some(qtensor) => { + let shape = s.into(); + if qtensor.shape() != &shape { + candle::bail!( + "shape mismatch for {name}, got {:?}, expected {shape:?}", + qtensor.shape() + ) + } + Ok(qtensor.clone()) + } + } + } + + pub fn get_no_shape(&self, name: &str) -> Result> { + let path = self.path(name); + match self.data.get(&path) { + None => { + candle::bail!("cannot find tensor {name}") + } + Some(qtensor) => Ok(qtensor.clone()), + } + } + + pub fn device(&self) -> &Device { + &self.device + } + + pub fn contains_key(&self, key: &str) -> bool { + self.data.contains_key(key) + } +} diff --git a/patches/candle-transformers/src/utils.rs b/patches/candle-transformers/src/utils.rs new file mode 100644 index 0000000000..884d4f378a --- /dev/null +++ b/patches/candle-transformers/src/utils.rs @@ -0,0 +1,38 @@ +//! Apply penalty and repeat_kv + +use candle::{Result, Tensor}; + +pub fn apply_repeat_penalty(logits: &Tensor, penalty: f32, context: &[u32]) -> Result { + let device = logits.device(); + let mut logits = logits.to_dtype(candle::DType::F32)?.to_vec1::()?; + let mut already_seen = std::collections::HashSet::new(); + for token_id in context { + if already_seen.contains(token_id) { + continue; + } + already_seen.insert(token_id); + if let Some(logit) = logits.get_mut(*token_id as usize) { + if *logit >= 0. { + *logit /= penalty + } else { + *logit *= penalty + } + } + } + let logits_len = logits.len(); + Tensor::from_vec(logits, logits_len, device) +} + +/// Repeats a key or value tensor for grouped query attention +/// The input tensor should have a shape `(batch, num_kv_heads, seq_len, head_dim)`, +pub fn repeat_kv(xs: Tensor, n_rep: usize) -> Result { + if n_rep == 1 { + Ok(xs) + } else { + let (b_sz, n_kv_head, seq_len, head_dim) = xs.dims4()?; + // Using cat is faster than a broadcast as it avoids going through a potentially + // strided copy. + // https://github.com/huggingface/candle/pull/2043 + Tensor::cat(&vec![&xs; n_rep], 2)?.reshape((b_sz, n_kv_head * n_rep, seq_len, head_dim)) + } +} diff --git a/patches/candle-transformers/tests/generation_tests.rs b/patches/candle-transformers/tests/generation_tests.rs new file mode 100644 index 0000000000..ee7df16999 --- /dev/null +++ b/patches/candle-transformers/tests/generation_tests.rs @@ -0,0 +1,78 @@ +use candle::{Device, Result, Tensor}; +use candle_transformers::generation::LogitsProcessor; + +#[test] +fn sample_with_zero_temperature() -> Result<()> { + let mut logits_process = LogitsProcessor::new(1337, None, None); + let logits = Tensor::new(&[0.1, 0.2, 0.3, 0.4], &Device::Cpu)?; + let token = logits_process.sample(&logits)?; + assert_eq!(token, 3); + Ok(()) +} + +#[test] +fn sample_with_temperature() -> Result<()> { + let mut logits_process = LogitsProcessor::new(42, Some(0.9), None); + let logits = Tensor::new(&[0.1, 0.2, 0.3, 0.4], &Device::Cpu)?; + let token = logits_process.sample(&logits)?; + assert_eq!(token, 0); + Ok(()) +} + +#[test] +fn sample_with_top_p() -> Result<()> { + let mut logits_process = LogitsProcessor::new(42, Some(1.0), Some(0.5)); + let logits = Tensor::new(&[0.1, 0.2, 0.3, 0.4], &Device::Cpu)?; + let token = logits_process.sample(&logits)?; + assert_eq!(token, 2); + Ok(()) +} + +#[test] +fn sample_with_top_k() -> Result<()> { + let mut logits_process = LogitsProcessor::from_sampling( + 42, + candle_transformers::generation::Sampling::TopK { + k: 1, + temperature: 1.0, + }, + ); + let logits = Tensor::new(&[0.1, 0.2, 0.3, 0.4], &Device::Cpu)?; + let token = logits_process.sample(&logits)?; + assert_eq!(token, 3); + let mut logits_process = LogitsProcessor::from_sampling( + 42, + candle_transformers::generation::Sampling::TopK { + k: 2, + temperature: 1.0, + }, + ); + let logits = Tensor::new(&[0.1, 0.2, 0.3, 0.4], &Device::Cpu)?; + let token = logits_process.sample(&logits)?; + assert_eq!(token, 3); + let token = logits_process.sample(&logits)?; + assert_eq!(token, 2); + Ok(()) +} + +#[test] +fn sample_gumbel() -> Result<()> { + let mut logits_process = LogitsProcessor::from_sampling( + 42, + candle_transformers::generation::Sampling::GumbelSoftmax { temperature: 1.0 }, + ); + let logits = Tensor::new(&[-1.0, 0.0, 0.2, 1.0], &Device::Cpu)?; + let sm = candle_nn::ops::softmax(&logits, 0)?.to_vec1::()?; + let mut counts = vec![0f64; 4]; + let samples = 100000; + for _ in 0..samples { + let token = logits_process.sample(&logits)?; + counts[token as usize] += 1f64 / samples as f64; + } + for i in 0..4 { + if (counts[i] - sm[i]).abs() > 0.05 { + panic!("pr mismatch {counts:?} {sm:?}"); + } + } + Ok(()) +} diff --git a/patches/candle-transformers/tests/nms_tests.rs b/patches/candle-transformers/tests/nms_tests.rs new file mode 100644 index 0000000000..d70f6fdf32 --- /dev/null +++ b/patches/candle-transformers/tests/nms_tests.rs @@ -0,0 +1,222 @@ +use candle::Result; +use candle_transformers::object_detection::{ + non_maximum_suppression, soft_non_maximum_suppression, Bbox, +}; + +#[test] +fn nms_basic() -> Result<()> { + // Boxes based upon https://thepythoncode.com/article/non-maximum-suppression-using-opencv-in-python + let mut bboxes = vec![vec![ + Bbox { + xmin: 245.0, + ymin: 305.0, + xmax: 575.0, + ymax: 490.0, + confidence: 0.9, + data: (), + }, // Box 1 + Bbox { + xmin: 235.0, + ymin: 300.0, + xmax: 485.0, + ymax: 515.0, + confidence: 0.8, + data: (), + }, // Box 2 + Bbox { + xmin: 305.0, + ymin: 270.0, + xmax: 540.0, + ymax: 500.0, + confidence: 0.6, + data: (), + }, // Box 3 + ]]; + + non_maximum_suppression(&mut bboxes, 0.5); + let bboxes = bboxes.into_iter().next().unwrap(); + assert_eq!(bboxes.len(), 1); + assert_eq!(bboxes[0].confidence, 0.9); + + Ok(()) +} + +#[test] +fn softnms_basic_functionality() -> Result<()> { + let mut bboxes = vec![vec![ + Bbox { + xmin: 0.0, + ymin: 0.0, + xmax: 1.0, + ymax: 1.0, + confidence: 0.5, + data: (), + }, + Bbox { + xmin: 0.1, + ymin: 0.1, + xmax: 1.1, + ymax: 1.1, + confidence: 0.9, + data: (), + }, + Bbox { + xmin: 0.2, + ymin: 0.2, + xmax: 1.2, + ymax: 1.2, + confidence: 0.6, + data: (), + }, + ]]; + + soft_non_maximum_suppression(&mut bboxes, Some(0.5), Some(0.1), Some(0.5)); + + // Should decay boxes following highest confidence box + assert!(bboxes[0][0].confidence == 0.9); + assert!(bboxes[0][1].confidence < 0.5); + assert!(bboxes[0][2].confidence < 0.6); + Ok(()) +} + +#[test] +fn softnms_confidence_decay() -> Result<()> { + let mut bboxes = vec![vec![ + Bbox { + xmin: 0.0, + ymin: 0.0, + xmax: 1.0, + ymax: 1.0, + confidence: 0.9, + data: (), + }, // Reference box + Bbox { + xmin: 0.1, + ymin: 0.1, + xmax: 1.1, + ymax: 1.1, + confidence: 0.8, + data: (), + }, // Overlapping box + ]]; + + soft_non_maximum_suppression(&mut bboxes, Some(0.5), Some(0.1), Some(0.5)); + + // Check that confidence of the overlapping box is decayed + assert!(bboxes[0][0].confidence == 0.9); + assert!(bboxes[0][1].confidence < 0.8); + Ok(()) +} + +#[test] +fn softnms_confidence_threshold() -> Result<()> { + let mut bboxes = vec![vec![ + Bbox { + xmin: 0.0, + ymin: 0.0, + xmax: 1.0, + ymax: 1.0, + confidence: 0.9, + data: (), + }, + Bbox { + xmin: 0.1, + ymin: 0.1, + xmax: 1.1, + ymax: 1.1, + confidence: 0.05, + data: (), + }, + ]]; + + soft_non_maximum_suppression(&mut bboxes, Some(0.5), Some(0.1), Some(0.5)); + + // Box with confidence below the threshold should be removed + assert_eq!(bboxes[0].len(), 2); + assert_eq!(bboxes[0][0].confidence, 0.9); + assert_eq!(bboxes[0][1].confidence, 0.00); + Ok(()) +} + +#[test] +fn softnms_no_overlap() -> Result<()> { + let mut bboxes = vec![vec![ + Bbox { + xmin: 0.0, + ymin: 0.0, + xmax: 1.0, + ymax: 1.0, + confidence: 0.9, + data: (), + }, + Bbox { + xmin: 2.0, + ymin: 2.0, + xmax: 3.0, + ymax: 3.0, + confidence: 0.8, + data: (), + }, + ]]; + + soft_non_maximum_suppression(&mut bboxes, Some(0.5), Some(0.1), Some(0.5)); + + // Both boxes should remain as they do not significantly overlap + assert_eq!(bboxes[0].len(), 2); + assert_eq!(bboxes[0][0].confidence, 0.9); + assert_eq!(bboxes[0][1].confidence, 0.8); + Ok(()) +} +#[test] +fn softnms_no_bbox() -> Result<()> { + let mut bboxes: Vec>> = vec![]; + soft_non_maximum_suppression(&mut bboxes, Some(0.5), Some(0.1), Some(0.5)); + assert!(bboxes.is_empty()); + Ok(()) +} + +#[test] +fn softnms_single_bbox() -> Result<()> { + let mut bboxes = vec![vec![Bbox { + xmin: 0.0, + ymin: 0.0, + xmax: 1.0, + ymax: 1.0, + confidence: 0.9, + data: (), + }]]; + soft_non_maximum_suppression(&mut bboxes, Some(0.5), Some(0.1), Some(0.5)); + assert_eq!(bboxes[0].len(), 1); + Ok(()) +} + +#[test] +fn softnms_equal_confidence_overlap() -> Result<()> { + let mut bboxes = vec![vec![ + Bbox { + xmin: 0.0, + ymin: 0.0, + xmax: 1.0, + ymax: 1.0, + confidence: 0.5, + data: (), + }, + Bbox { + xmin: 0.1, + ymin: 0.1, + xmax: 1.1, + ymax: 1.1, + confidence: 0.5, + data: (), + }, + ]]; + + soft_non_maximum_suppression(&mut bboxes, Some(0.5), Some(0.1), Some(0.5)); + + // First box will be reference box, second box should be decayed + // Implementation must change to have both be decayed + assert_eq!(bboxes[0].len(), 2); + assert!(bboxes[0][0].confidence == 0.5); + assert!(bboxes[0][1].confidence < 0.5); + Ok(()) +}