Add the Mamba 2 architecture

mistralrs-core/src/layers.rs

@@ -77,6 +77,25 @@
  }
 }
 
+pub struct GatedRmsNorm(RmsNorm);
+
+impl GatedRmsNorm {
+ pub fn new(size: usize, eps: f64, vb: VarBuilder) -> Result<Self> {
+ let inner = candle_nn::rms_norm_non_quant(size, eps, vb)?;
+ let w = inner.inner().weight().clone();
+ Ok(Self(RmsNorm { eps, weight: w }))
+ }
+
+ fn forward(&self, x: &Tensor, gate: Option<&Tensor>) -> Result<Tensor> {
+ let x = if let Some(gate) = gate {
+ x.broadcast_mul(&gate.to_dtype(x.dtype())?)?.contiguous()?
+ } else {
+ x.contiguous()?
+ };
+ candle_nn::ops::rms_norm(&x, &self.0.weight, self.0.eps as f32)
+ }
+}
+
 /// RoPE supporting LongRope
 #[derive(Debug, Clone)]
 pub struct PhiRotaryEmbedding {

mistralrs-core/src/models/mamba2.rs

@@ -0,0 +1,235 @@
+#![allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
+
+use std::sync::Arc;
+
+use candle_core::{Result, Tensor};
+use candle_nn::{
+ conv1d, conv1d_no_bias, Activation, Conv1d, Conv1dConfig, Embedding, Linear, VarBuilder,
+};
+use mistralrs_quant::{linear_b, QuantMethod, QuantMethodConfig, QuantizedConfig, UnquantLinear};
+use serde::Deserialize;
+
+use crate::{
+ layers::{GatedRmsNorm, RmsNorm},
+ paged_attention::AttentionImplementation,
+ pipeline::NormalLoadingMetadata,
+ serde_default_fn,
+};
+
+serde_default_fn!(usize, num_heads_default, 128);
+serde_default_fn!(usize, head_dim_default, 64);
+serde_default_fn!(usize, vocab_size_default, 32768);
+serde_default_fn!(usize, hidden_size_default, 4096);
+serde_default_fn!(usize, state_size_default, 128);
+serde_default_fn!(usize, num_hidden_layers_default, 64);
+serde_default_fn!(f64, layer_norm_epsilon_default, 1e-5);
+serde_default_fn!(f64, expand_default, 2.0);
+serde_default_fn!(usize, conv_kernel_default, 4);
+serde_default_fn!(usize, n_groups_default, 2);
+serde_default_fn!(bool, use_bias_default, false);
+serde_default_fn!(bool, use_conv_bias_default, true);
+serde_default_fn!(Activation, hidden_act_default, Activation::Silu);
+serde_default_fn!(bool, residual_in_fp32_default, true);
+serde_default_fn!(f64, time_step_min_default, 0.001);
+serde_default_fn!(f64, time_step_max_default, 0.1);
+serde_default_fn!(f64, time_step_floor_default, 0.0001);
+serde_default_fn!((f64, f64), time_step_limit_default, (0.0, f64::INFINITY));
+serde_default_fn!(bool, rescale_prenorm_residual_default, false);
+serde_default_fn!(bool, norm_before_gate_default, true);
+serde_default_fn!(bool, rms_norm_default, true);
+serde_default_fn!(usize, chunk_size_default, 256);
+
+#[derive(Debug, Clone, Deserialize, Default)]
+pub enum TimeStepRank {
+ #[default]
+ Auto,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct Mamba2Config {
+ #[serde(default = "num_heads_default")]
+ num_heads: usize,
+ #[serde(default = "head_dim_default")]
+ head_dim: usize,
+ #[serde(default = "vocab_size_default")]
+ vocab_size: usize,
+ #[serde(default = "hidden_size_default")]
+ hidden_size: usize,
+ #[serde(default = "state_size_default")]
+ state_size: usize,
+ #[serde(default = "num_hidden_layers_default")]
+ num_hidden_layers: usize,
+ #[serde(default = "layer_norm_epsilon_default")]
+ layer_norm_epsilon: f64,
+ #[serde(default = "expand_default")]
+ expand: f64,
+ #[serde(default = "conv_kernel_default")]
+ conv_kernel: usize,
+ #[serde(default = "n_groups_default")]
+ n_groups: usize,
+ #[serde(default = "use_bias_default")]
+ use_bias: bool,
+ #[serde(default = "use_conv_bias_default")]
+ use_conv_bias: bool,
+ #[serde(default = "hidden_act_default")]
+ hidden_act: Activation,
+ #[serde(default = "residual_in_fp32_default")]
+ residual_in_fp32: bool,
+ #[serde(default = "Default::default")]
+ time_step_rank: TimeStepRank,
+ #[serde(default = "time_step_min_default")]
+ time_step_min: f64,
+ #[serde(default = "time_step_max_default")]
+ time_step_max: f64,
+ #[serde(default = "time_step_floor_default")]
+ time_step_floor: f64,
+ #[serde(default = "time_step_limit_default")]
+ time_step_limit: (f64, f64),
+ #[serde(default = "rescale_prenorm_residual_default")]
+ rescale_prenorm_residual: bool,
+ #[serde(default = "norm_before_gate_default")]
+ norm_before_gate: bool,
+ #[serde(default = "rms_norm_default")]
+ rms_norm: bool,
+ #[serde(default = "chunk_size_default")]
+ chunk_size: usize,
+ quantization_config: Option<QuantizedConfig>,
+}
+
+// https://github.com/huggingface/transformers/blob/main/src/transformers/models/mamba2/modeling_mamba2.py#L406
+struct Mixer {
+ conv1d: Conv1d,
+ in_proj: Arc<dyn QuantMethod>,
+ dt_bias: Tensor,
+ a_log: Tensor,
+ d: Tensor,
+ norm: GatedRmsNorm,
+ out_proj: Arc<dyn QuantMethod>,
+}
+
+impl Mixer {
+ fn new(cfg: &Mamba2Config, vb: VarBuilder) -> Result<Self> {
+ let intermediate_size = (cfg.expand * cfg.hidden_size as f64) as usize;
+ let conv_dim = intermediate_size + 2 * cfg.n_groups * cfg.state_size;
+ let projection_size = intermediate_size + conv_dim + cfg.num_heads;
+
+ let conv1d_fn = if cfg.use_conv_bias {
+ conv1d
+ } else {
+ conv1d_no_bias
+ };
+
+ let conv1d = conv1d_fn(
+ conv_dim,
+ conv_dim,
+ cfg.conv_kernel,
+ Conv1dConfig {
+ padding: cfg.conv_kernel - 1,
+ groups: conv_dim,
+ stride: 1,
+ dilation: 1,
+ },
+ vb.pp("conv1d"),
+ )?;
+
+ let in_proj = linear_b(
+ cfg.hidden_size,
+ projection_size,
+ cfg.use_bias,
+ &cfg.quantization_config,
+ vb.pp("in_proj"),
+ )?;
+
+ let out_proj = linear_b(
+ intermediate_size,
+ cfg.hidden_size,
+ cfg.use_bias,
+ &cfg.quantization_config,
+ vb.pp("out_proj"),
+ )?;
+
+ // Time step proj, discretization
+ let dt_bias = vb.get((cfg.num_heads,), "dt_bias")?;
+
+ // S4D real init, not discretized
+ let a_log = vb.get((1, cfg.num_heads + 1), "A_log")?;
+ let d = vb.get((cfg.num_heads,), "D")?;
+
+ let norm = GatedRmsNorm::new(intermediate_size, cfg.layer_norm_epsilon, vb.pp("norm"))?;
+
+ Ok(Self {
+ conv1d,
+ in_proj,
+ out_proj,
+ dt_bias,
+ a_log,
+ d,
+ norm,
+ })
+ }
+}
+
+struct Layer {
+ norm: RmsNorm,
+ mixer: Mixer,
+ res_in_f32: bool,
+}
+
+impl Layer {
+ fn new(cfg: &Mamba2Config, vb: VarBuilder) -> Result<Self> {
+ let norm = RmsNorm::new(cfg.hidden_size, cfg.layer_norm_epsilon, vb.pp("norm"))?;
+ let mixer = Mixer::new(cfg, vb.pp("mixer"))?;
+ Ok(Self {
+ norm,
+ mixer,
+ res_in_f32: cfg.residual_in_fp32,
+ })
+ }
+}
+
+pub struct Model {
+ lm_head: Arc<dyn QuantMethod>,
+ embeddings: Embedding,
+ norm_f: RmsNorm,
+ layers: Vec<Layer>,
+}
+
+impl Model {
+ pub fn new(
+ cfg: &Mamba2Config,
+ vb: VarBuilder,
+ _is_gptx: bool,
+ _normal_loading_metadata: NormalLoadingMetadata,
+ _attention_mechanism: AttentionImplementation,
+ ) -> Result<Self> {
+ if let Some(ref quant_cfg) = &cfg.quantization_config {
+ tracing::info!(
+ "Using {} quantization in {} bits.",
+ quant_cfg.quant_method.to_string(),
+ quant_cfg.bits
+ );
+ }
+
+ let vb = vb.pp("backbone");
+
+ let embeddings =
+ candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("embeddings"))?;
+ // Tied to lm_head...
+ let lm_head = Arc::new(UnquantLinear::new(QuantMethodConfig::Unquantized(
+ Linear::new(embeddings.embeddings().clone(), None),
+ ))?);
+ let norm_f = RmsNorm::new(cfg.hidden_size, cfg.layer_norm_epsilon, vb.pp("norm_f"))?;
+
+ let mut layers = Vec::with_capacity(cfg.num_hidden_layers);
+ for idx in 0..cfg.num_hidden_layers {
+ layers.push(Layer::new(cfg, vb.pp(idx))?);
+ }
+
+ Ok(Self {
+ lm_head,
+ embeddings,
+ norm_f,
+ layers,
+ })
+ }
+}

mistralrs-core/src/models/mod.rs

@@ -1,6 +1,7 @@
 pub(crate) mod gemma;
 pub(crate) mod gemma2;
 pub(crate) mod llama;
+pub(crate) mod mamba2;
 pub(crate) mod mistral;
 pub(crate) mod mixtral;
 pub(crate) mod phi2;