add readme

src/models/common/mod.rs

@@ -0,0 +1,264 @@
+use anyhow::Result;
+use candle_core::{Tensor, D};
+use candle_nn::{Activation, Linear, Module, VarBuilder, linear, linear_no_bias};
+
+use crate::{position_embed::rope::apply_rotary_pos_emb, utils::tensor_utils::repeat_kv};
+
+#[derive(Debug, Clone)]
+pub struct MLPWithBias {
+    gate_proj: Linear,
+    up_proj: Linear,
+    down_proj: Linear,
+    act_fn: Activation,
+}
+
+impl MLPWithBias {
+    pub fn new(
+        vb: VarBuilder,
+        hidden_size: usize,
+        intermediate_size: usize,
+        act_fn: Activation,
+    ) -> Result<Self> {
+        let gate_proj = linear(hidden_size, intermediate_size, vb.pp("gate_proj"))?;
+        let up_proj = linear(hidden_size, intermediate_size, vb.pp("up_proj"))?;
+        let down_proj = linear(intermediate_size, hidden_size, vb.pp("down_proj"))?;
+        Ok(Self {
+            gate_proj,
+            up_proj,
+            down_proj,
+            act_fn,
+        })
+    }
+}
+
+impl Module for MLPWithBias {
+    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
+        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)]
+pub struct MLPNoBias {
+    gate_proj: Linear,
+    up_proj: Linear,
+    down_proj: Linear,
+    act_fn: Activation,
+}
+
+impl MLPNoBias {
+    pub fn new(
+        vb: VarBuilder,
+        hidden_size: usize,
+        intermediate_size: usize,
+        act_fn: Activation,
+    ) -> Result<Self> {
+        let gate_proj = linear_no_bias(hidden_size, intermediate_size, vb.pp("gate_proj"))?;
+        let up_proj = linear_no_bias(hidden_size, intermediate_size, vb.pp("up_proj"))?;
+        let down_proj = linear_no_bias(intermediate_size, hidden_size, vb.pp("down_proj"))?;
+        Ok(Self {
+            gate_proj,
+            up_proj,
+            down_proj,
+            act_fn,
+        })
+    }
+}
+
+impl Module for MLPNoBias {
+    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
+        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)]
+pub struct AttentionNobias {
+    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,
+    kv_cache: Option<(Tensor, Tensor)>,
+}
+
+impl AttentionNobias {
+    pub fn new(vb: VarBuilder, hidden_size: usize, num_attention_heads: usize, num_key_value_heads: usize) -> Result<Self> {
+        let num_kv_groups = num_attention_heads / num_key_value_heads;
+        let head_dim = hidden_size / num_attention_heads;
+        let q_proj = linear_no_bias(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?;
+        let k_proj = linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?;
+        let v_proj = linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?;
+        let o_proj = linear_no_bias(hidden_size, hidden_size, vb.pp("o_proj"))?;
+        Ok(Self {
+            q_proj,
+            k_proj,
+            v_proj,
+            o_proj,
+            num_heads: num_attention_heads,
+            num_kv_heads: num_key_value_heads,
+            num_kv_groups,
+            head_dim,
+            hidden_size,
+            kv_cache: None,
+        })
+    }
+
+    pub fn forward(
+        &self,
+        xs: &Tensor,
+        cos: &Tensor,
+        sin: &Tensor,
+        attention_mask: Option<&Tensor>,
+        tof32: bool,
+    ) -> Result<Tensor> {
+        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) =
+            apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)?;
+
+        let key_states = repeat_kv(key_states, self.num_kv_groups)?.contiguous()?;
+        let value_states = repeat_kv(value_states, self.num_kv_groups)?.contiguous()?;
+        let query_states = query_states.contiguous()?;
+        let attn_output = {
+            let scale = 1f64 / f64::sqrt(self.head_dim as f64);
+            #[cfg(not(feature = "flash-attn"))]
+            {
+                let attn_weights =
+                    query_states.matmul(&key_states.transpose(D::Minus2, D::Minus1)?)?;
+                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)?;
+                let attn_weights = attn_weights.matmul(&value_states)?;
+                attn_weights
+            }
+            #[cfg(feature = "flash-attn")]
+            {
+                // use flash-attn,
+                // flash-attn shape: (bs, seq_len, num_head, head_dim)
+                let query_states = query_states.transpose(1, 2)?;
+                let key_states = key_states.transpose(1, 2)?;
+                let value_states = value_states.transpose(1, 2)?;
+                let attn_output = candle_flash_attn::flash_attn(
+                    &query_states,
+                    &key_states,
+                    &value_states,
+                    scale as f32,
+                    attention_mask.is_some(),
+                )?
+                .transpose(1, 2)?;
+                attn_output
+            }
+        };
+        let attn_output =
+            attn_output
+                .transpose(1, 2)?
+                .contiguous()?
+                .reshape((b_sz, q_len, self.hidden_size))?;
+        let attn_output = attn_output.apply(&self.o_proj)?;
+        Ok(attn_output)
+    }
+
+    pub fn forward_with_cache(
+        &mut self,
+        xs: &Tensor,
+        cos: &Tensor,
+        sin: &Tensor,
+        attention_mask: Option<&Tensor>,
+        tof32: bool,
+    ) -> Result<Tensor> {
+        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) =
+            apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)?;
+        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 = repeat_kv(key_states, self.num_kv_groups)?.contiguous()?;
+        let value_states = repeat_kv(value_states, self.num_kv_groups)?.contiguous()?;
+        let query_states = query_states.contiguous()?;
+        let attn_output = {
+            let scale = 1f64 / f64::sqrt(self.head_dim as f64);
+            #[cfg(not(feature = "flash-attn"))]
+            {
+                let attn_weights =
+                    query_states.matmul(&key_states.transpose(D::Minus2, D::Minus1)?)?;
+                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)?;
+                let attn_weights = attn_weights.matmul(&value_states)?;
+                attn_weights
+            }
+            #[cfg(feature = "flash-attn")]
+            {
+                // use flash-attn,
+                // flash-attn shape: (bs, seq_len, num_head, head_dim)
+                let query_states = query_states.transpose(1, 2)?;
+                let key_states = key_states.transpose(1, 2)?;
+                let value_states = value_states.transpose(1, 2)?;
+                let attn_output = candle_flash_attn::flash_attn(
+                    &query_states,
+                    &key_states,
+                    &value_states,
+                    scale as f32,
+                    attention_mask.is_some(),
+                )?
+                .transpose(1, 2)?;
+                attn_output
+            }
+        };
+        let attn_output =
+            attn_output
+                .transpose(1, 2)?
+                .contiguous()?
+                .reshape((b_sz, q_len, self.hidden_size))?;
+        let attn_output = attn_output.apply(&self.o_proj)?;
+        Ok(attn_output)
+    }
+
+    pub fn clear_kv_cache(&mut self) {
+        self.kv_cache = None
+    }
+}