Mamba2 implementation (huggingface#3264)

Author: Anri-Lombard

candle-examples/examples/mamba2/README.md

@@ -0,0 +1,56 @@
+# candle-mamba2: Mamba2 implementation
+
+Candle implementation of _Mamba2_ [1] inference. Mamba2 introduces the State Space
+Duality (SSD) framework which unifies structured SSMs and attention variants.
+
+- [1]. [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://arxiv.org/abs/2405.21060)
+
+## Running the example
+
+```bash
+cargo run --example mamba2 --release -- --prompt "Mamba is the"
+```
+
+## Supported models
+
+| Model | HuggingFace ID |
+|-------|----------------|
+| Mamba2-130m | `AntonV/mamba2-130m-hf` |
+| Mamba2-370m | `AntonV/mamba2-370m-hf` |
+| Mamba2-780m | `AntonV/mamba2-780m-hf` |
+| Mamba2-1.3b | `AntonV/mamba2-1.3b-hf` |
+| Mamba2-2.7b | `AntonV/mamba2-2.7b-hf` |
+
+## Verification
+
+Outputs match the PyTorch transformers `Mamba2ForCausalLM` reference implementation.
+
+### mamba2-130m
+
+```bash
+cargo run --example mamba2 --release -- \
+  --prompt "Mamba is the" \
+  --which mamba2-130m \
+  --sample-len 20 \
+  --repeat-penalty 1.0
+```
+
+Expected output:
+```
+Mamba is the most popular and popular game in the world. It is a game where you can play with your friends
+```
+
+### mamba2-370m
+
+```bash
+cargo run --example mamba2 --release -- \
+  --prompt "Mamba is the" \
+  --which mamba2-370m \
+  --sample-len 20 \
+  --repeat-penalty 1.0
+```
+
+Expected output:
+```
+Mamba is the first game in the series to feature a new character, the Mamba, who is a female version
+```

candle-examples/examples/mamba2/main.rs

@@ -0,0 +1,326 @@
+#[cfg(feature = "mkl")]
+extern crate intel_mkl_src;
+
+#[cfg(feature = "accelerate")]
+extern crate accelerate_src;
+
+use anyhow::{Error as E, Result};
+use clap::{Parser, ValueEnum};
+
+use candle_transformers::models::mamba2::{Config, Model, State};
+
+use candle::{DType, Device, Tensor};
+use candle_examples::token_output_stream::TokenOutputStream;
+use candle_nn::VarBuilder;
+use candle_transformers::generation::LogitsProcessor;
+use hf_hub::{api::sync::Api, Repo, RepoType};
+use tokenizers::Tokenizer;
+
+struct TextGeneration {
+    model: Model,
+    config: Config,
+    device: Device,
+    tokenizer: TokenOutputStream,
+    logits_processor: LogitsProcessor,
+    repeat_penalty: f32,
+    repeat_last_n: usize,
+    use_prefill: bool,
+    chunk_size: usize,
+}
+
+impl TextGeneration {
+    #[allow(clippy::too_many_arguments)]
+    fn new(
+        model: Model,
+        config: Config,
+        tokenizer: Tokenizer,
+        seed: u64,
+        temp: Option<f64>,
+        top_p: Option<f64>,
+        repeat_penalty: f32,
+        repeat_last_n: usize,
+        use_prefill: bool,
+        chunk_size: usize,
+        device: &Device,
+    ) -> Self {
+        let logits_processor = LogitsProcessor::new(seed, temp, top_p);
+        Self {
+            model,
+            config,
+            tokenizer: TokenOutputStream::new(tokenizer),
+            logits_processor,
+            repeat_penalty,
+            repeat_last_n,
+            use_prefill,
+            chunk_size,
+            device: device.clone(),
+        }
+    }
+
+    fn run(&mut self, prompt: &str, sample_len: usize) -> Result<()> {
+        use std::io::Write;
+        self.tokenizer.clear();
+        let dtype = self.model.dtype();
+        let mut tokens = self
+            .tokenizer
+            .tokenizer()
+            .encode(prompt, true)
+            .map_err(E::msg)?
+            .get_ids()
+            .to_vec();
+        let mut generated_tokens = 0usize;
+        let eos_token = match self.tokenizer.get_token("<|endoftext|>") {
+            Some(token) => token,
+            None => anyhow::bail!("cannot find the <|endoftext|> token"),
+        };
+        let mut state = State::new(1, &self.config, dtype, &self.device)?;
+        let mut next_logits = None;
+
+        if self.use_prefill && tokens.len() > 1 {
+            let prefill_start = std::time::Instant::now();
+            // Prefill mode: process all tokens at once
+            let input = Tensor::new(&tokens[..], &self.device)?.unsqueeze(0)?;
+            let logits = self
+                .model
+                .forward_prefill(&input, &mut state, self.chunk_size)?;
+            // Get logits for last position
+            next_logits = Some(logits.narrow(1, tokens.len() - 1, 1)?.squeeze(1)?);
+            for &t in tokens.iter() {
+                if let Some(t) = self.tokenizer.next_token(t)? {
+                    print!("{t}")
+                }
+            }
+            println!(
+                "\n[Prefill {} tokens in {:.2}ms]",
+                tokens.len(),
+                prefill_start.elapsed().as_secs_f64() * 1000.0
+            );
+        } else {
+            // Step-by-step mode
+            for &t in tokens.iter() {
+                let input = Tensor::new(&[t], &self.device)?;
+                let logits = self.model.forward(&input, &mut state)?;
+                next_logits = Some(logits);
+                if let Some(t) = self.tokenizer.next_token(t)? {
+                    print!("{t}")
+                }
+            }
+        }
+        std::io::stdout().flush()?;
+
+        let start_gen = std::time::Instant::now();
+        for _ in 0..sample_len {
+            let logits = match next_logits.as_ref() {
+                Some(logits) => logits,
+                None => anyhow::bail!("cannot work on an empty prompt"),
+            };
+            let logits = logits.squeeze(0)?.to_dtype(dtype)?;
+            let logits = if self.repeat_penalty == 1. {
+                logits
+            } else {
+                let start_at = tokens.len().saturating_sub(self.repeat_last_n);
+                candle_transformers::utils::apply_repeat_penalty(
+                    &logits,
+                    self.repeat_penalty,
+                    &tokens[start_at..],
+                )?
+            };
+            let next_token = self.logits_processor.sample(&logits)?;
+            tokens.push(next_token);
+            generated_tokens += 1;
+            if next_token == eos_token {
+                break;
+            }
+            if let Some(t) = self.tokenizer.next_token(next_token)? {
+                print!("{t}");
+                std::io::stdout().flush()?;
+            }
+
+            let input = Tensor::new(&[next_token], &self.device)?;
+            next_logits = Some(self.model.forward(&input, &mut state)?)
+        }
+        let dt = start_gen.elapsed();
+        if let Some(rest) = self.tokenizer.decode_rest().map_err(E::msg)? {
+            print!("{rest}");
+        }
+        std::io::stdout().flush()?;
+        println!(
+            "\n{generated_tokens} tokens generated ({:.2} token/s)",
+            generated_tokens as f64 / dt.as_secs_f64(),
+        );
+        Ok(())
+    }
+}
+
+#[derive(Parser, ValueEnum, Clone, Copy, PartialEq, Eq, Debug)]
+enum Which {
+    Mamba2_130m,
+    Mamba2_370m,
+    Mamba2_780m,
+    Mamba2_1_3b,
+    Mamba2_2_7b,
+}
+
+impl std::fmt::Display for Which {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{self:?}")
+    }
+}
+
+impl Which {
+    fn model_id(&self) -> &'static str {
+        match self {
+            Self::Mamba2_130m => "AntonV/mamba2-130m-hf",
+            Self::Mamba2_370m => "AntonV/mamba2-370m-hf",
+            Self::Mamba2_780m => "AntonV/mamba2-780m-hf",
+            Self::Mamba2_1_3b => "AntonV/mamba2-1.3b-hf",
+            Self::Mamba2_2_7b => "AntonV/mamba2-2.7b-hf",
+        }
+    }
+}
+
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+struct Args {
+    /// Run on CPU rather than on GPU.
+    #[arg(long)]
+    cpu: bool,
+
+    /// Enable tracing (generates a trace-timestamp.json file).
+    #[arg(long)]
+    tracing: bool,
+
+    #[arg(long)]
+    prompt: String,
+
+    /// The temperature used to generate samples.
+    #[arg(long)]
+    temperature: Option<f64>,
+
+    /// Nucleus sampling probability cutoff.
+    #[arg(long)]
+    top_p: Option<f64>,
+
+    /// The seed to use when generating random samples.
+    #[arg(long, default_value_t = 299792458)]
+    seed: u64,
+
+    /// The length of the sample to generate (in tokens).
+    #[arg(long, short = 'n', default_value_t = 5000)]
+    sample_len: usize,
+
+    #[arg(long, default_value = "mamba2-130m")]
+    which: Which,
+
+    #[arg(long)]
+    model_id: Option<String>,
+
+    #[arg(long)]
+    tokenizer_file: Option<String>,
+
+    #[arg(long)]
+    weight_files: Option<String>,
+
+    #[arg(long)]
+    config_file: Option<String>,
+
+    #[arg(long, default_value = "f32")]
+    dtype: String,
+
+    /// Penalty to be applied for repeating tokens, 1. means no penalty.
+    #[arg(long, default_value_t = 1.1)]
+    repeat_penalty: f32,
+
+    /// The context size to consider for the repeat penalty.
+    #[arg(long, default_value_t = 64)]
+    repeat_last_n: usize,
+
+    /// Use chunked prefill for processing the initial prompt.
+    #[arg(long)]
+    use_prefill: bool,
+
+    /// Chunk size for prefill (default 256).
+    #[arg(long, default_value_t = 256)]
+    chunk_size: usize,
+}
+
+fn main() -> Result<()> {
+    use std::str::FromStr;
+    use tracing_chrome::ChromeLayerBuilder;
+    use tracing_subscriber::prelude::*;
+
+    let args = Args::parse();
+    let _guard = if args.tracing {
+        let (chrome_layer, guard) = ChromeLayerBuilder::new().build();
+        tracing_subscriber::registry().with(chrome_layer).init();
+        Some(guard)
+    } else {
+        None
+    };
+    println!(
+        "avx: {}, neon: {}, simd128: {}, f16c: {}",
+        candle::utils::with_avx(),
+        candle::utils::with_neon(),
+        candle::utils::with_simd128(),
+        candle::utils::with_f16c()
+    );
+    println!(
+        "temp: {:.2} repeat-penalty: {:.2} repeat-last-n: {}",
+        args.temperature.unwrap_or(0.),
+        args.repeat_penalty,
+        args.repeat_last_n
+    );
+
+    let start = std::time::Instant::now();
+    let api = Api::new()?;
+    let model_id = args
+        .model_id
+        .unwrap_or_else(|| args.which.model_id().to_string());
+    let repo = api.repo(Repo::new(model_id.clone(), RepoType::Model));
+    let tokenizer_filename = match args.tokenizer_file {
+        Some(file) => std::path::PathBuf::from(file),
+        None => repo.get("tokenizer.json")?,
+    };
+    let config_filename = match args.config_file {
+        Some(file) => std::path::PathBuf::from(file),
+        None => repo.get("config.json")?,
+    };
+    let filenames = match args.weight_files {
+        Some(files) => files
+            .split(',')
+            .map(std::path::PathBuf::from)
+            .collect::<Vec<_>>(),
+        None => {
+            vec![repo.get("model.safetensors")?]
+        }
+    };
+    println!("retrieved the files in {:?}", start.elapsed());
+    let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
+
+    let start = std::time::Instant::now();
+    // Config contains `Infinity` which is not valid JSON, replace with a large number
+    let config_str = std::fs::read_to_string(config_filename)?;
+    let config_str = config_str.replace("Infinity", "1e30");
+    let config: Config = serde_json::from_str(&config_str)?;
+    let device = candle_examples::device(args.cpu)?;
+    let dtype = DType::from_str(&args.dtype)?;
+    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&filenames, dtype, &device)? };
+    let model = Model::new(&config, vb.pp("backbone"))?;
+    println!("loaded the model in {:?}", start.elapsed());
+
+    let mut pipeline = TextGeneration::new(
+        model,
+        config,
+        tokenizer,
+        args.seed,
+        args.temperature,
+        args.top_p,
+        args.repeat_penalty,
+        args.repeat_last_n,
+        args.use_prefill,
+        args.chunk_size,
+        &device,
+    );
+    pipeline.run(&args.prompt, args.sample_len)?;
+    Ok(())
+}