diff --git a/src/lib.rs b/src/lib.rs index fb19131..264715b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,4 @@ -use crate::models::{GenerateModel, qwen2_5vl::generate::Qwen2_5VLGenerateModel}; +use crate::models::{minicpm4::generate::MiniCPMGenerateModel, qwen2_5vl::generate::Qwen2_5VLGenerateModel, GenerateModel}; use anyhow::{Ok, Result}; use candle_core::{DType, Device}; pub mod chat_template; @@ -9,6 +9,7 @@ pub mod utils; pub enum ModelType { Qwen2_5VL, + MiniCPM4, } impl ModelType { @@ -22,6 +23,10 @@ impl ModelType { ModelType::Qwen2_5VL => { let model = Qwen2_5VLGenerateModel::init(model_path, device, dtype)?; Ok(Box::new(model)) + }, + ModelType::MiniCPM4 => { + let model = MiniCPMGenerateModel::init(model_path, device, dtype)?; + Ok(Box::new(model)) } } } diff --git a/src/models/base_modules/mod.rs b/src/models/base_modules/mod.rs new file mode 100644 index 0000000..be470d1 --- /dev/null +++ b/src/models/base_modules/mod.rs @@ -0,0 +1,262 @@ +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 { + 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 { + 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 { + 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 { + 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 { + 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>, + ) -> 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) = + apply_rotary_pos_emb(&query_states, &key_states, cos, sin)?; + + 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_step( + &mut self, + xs: &Tensor, + cos: &Tensor, + sin: &Tensor, + attention_mask: Option<&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)?; + let (query_states, key_states) = + apply_rotary_pos_emb(&query_states, &key_states, cos, sin)?; + 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 + } +} diff --git a/src/models/minicpm4/config.rs b/src/models/minicpm4/config.rs new file mode 100644 index 0000000..af4d7d4 --- /dev/null +++ b/src/models/minicpm4/config.rs @@ -0,0 +1,32 @@ +use candle_nn::Activation; + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct RopeScalingConfig { + pub rope_type: String, + pub long_factor: Vec, + pub short_factor: Vec, + pub original_max_position_embeddings: usize, +} + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct MiniCPM4Config { + pub bos_token_id: u32, + pub eos_token_id: Vec, + pub hidden_act: Activation, + pub hidden_size: usize, + pub intermediate_size: usize, + pub max_position_embeddings: usize, + pub num_attention_heads: usize, + pub num_hidden_layers: usize, + pub num_key_value_heads: usize, + pub rms_norm_eps: f64, + pub rope_scaling: RopeScalingConfig, + pub torch_dtype: String, + pub vocab_size: usize, + // pub use_mup: bool, + pub scale_emb:f32, + pub dim_model_base: usize, + pub scale_depth: f32, + // pub rope_theta: f32, + // pub kv_channels: i32, +} \ No newline at end of file diff --git a/src/models/minicpm4/generate.rs b/src/models/minicpm4/generate.rs new file mode 100644 index 0000000..813f28e --- /dev/null +++ b/src/models/minicpm4/generate.rs @@ -0,0 +1,136 @@ +use crate::models::minicpm4::config::MiniCPM4Config; +use crate::models::minicpm4::model::MiniCPMModel; +// use crate::models::GenerateStream; +use crate::utils::utils::{ + build_completion_chunk_response, build_completion_response, find_safetensors_files, get_device, + get_dtype, get_logit_processor, +}; +use crate::{ + chat_template::chat_template::ChatTemplate, models::GenerateModel, + tokenizer::tokenizer::TokenizerModel, +}; +use anyhow::{Result, anyhow}; +use candle_core::{D, DType, Device, IndexOp, Tensor}; +use candle_nn::VarBuilder; +use openai_dive::v1::resources::chat::{ + ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, +}; +use rocket::async_stream::stream; +use rocket::futures::Stream; + +pub struct MiniCPMGenerateModel<'a> { + chat_template: ChatTemplate<'a>, + tokenizer: TokenizerModel, + minicpm: MiniCPMModel, + device: Device, + endoftext_id: u32, + im_end_id: u32, +} + +impl<'a> GenerateModel for MiniCPMGenerateModel<'a> { + fn init(path: &str, device: Option<&Device>, dtype: Option) -> Result { + let chat_template = ChatTemplate::init(path)?; + let tokenizer = TokenizerModel::init(path)?; + let config_path = path.to_string() + "/config.json"; + let cfg: MiniCPM4Config = serde_json::from_slice(&std::fs::read(config_path)?)?; + let device = &get_device(device); + let cfg_dtype = cfg.torch_dtype.as_str(); + let dtype = get_dtype(dtype, cfg_dtype); + let endoftext_id = cfg.eos_token_id[0]; + let im_end_id = cfg.eos_token_id[1]; + let model_list = find_safetensors_files(&path)?; + let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, device)? }; + let minicpm = MiniCPMModel::new(vb, cfg)?; + + Ok(MiniCPMGenerateModel { + chat_template, + tokenizer, + minicpm, + device: device.clone(), + endoftext_id, + im_end_id, + }) + } + + fn generate(&mut self, mes: ChatCompletionParameters) -> Result { + let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p); + let mes_render = self.chat_template.apply_chat_template(&mes)?; + let mut input_ids = self.tokenizer.text_encode(mes_render, &self.device)?; + let mut seq_len = input_ids.dim(1)?; + let mut seqlen_offset = 0; + let mut generate = Vec::new(); + let sample_len = match mes.max_tokens { + Some(max) => max, + None => 512, + }; + for _ in 0..sample_len { + let logits = self.minicpm.forward_step(&input_ids, seqlen_offset)?; + let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?; + let next_token = logit_processor.sample(&logits)?; + generate.push(next_token); + if next_token == self.endoftext_id || next_token == self.im_end_id { + break; + } + seqlen_offset += seq_len; + seq_len = 1; + input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; + } + let res = self.tokenizer.token_decode(generate)?; + self.minicpm.clear_kv_cache(); + let response = build_completion_response(res, "minicpm"); + Ok(response) + } + fn generate_stream( + &mut self, + mes: ChatCompletionParameters, + ) -> Result>> { + let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p); + let mes_render = self.chat_template.apply_chat_template(&mes)?; + let mut input_ids = self.tokenizer.text_encode(mes_render, &self.device)?; + let mut seq_len = input_ids.dim(1)?; + let mut seqlen_offset = 0; + let sample_len = match mes.max_tokens { + Some(max) => max, + None => 512, + }; + let stream = stream! { + let mut error_tokens = Vec::new(); + for _ in 0..sample_len { + let logits = self.minicpm.forward_step( + &input_ids, + seqlen_offset, + )?; + let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?; + let next_token = logit_processor.sample(&logits)?; + let mut decode_ids = Vec::new(); + if error_tokens.len() > 0 { + decode_ids.extend_from_slice(&error_tokens); + } + decode_ids.push(next_token); + let decoded_token = self.tokenizer.token_decode(decode_ids).map_err(|e| anyhow!(format!("stream decode error{}", e)))?; + if decoded_token.contains("�") { + error_tokens.push(next_token); + if error_tokens.len() > 3 { + error_tokens.clear(); + } + seqlen_offset += seq_len; + seq_len = 1; + input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; + continue; + } + error_tokens.clear(); + let chunk = build_completion_chunk_response(decoded_token, "minicpm", None, None); + yield Ok(chunk); + if next_token == self.endoftext_id || next_token == self.im_end_id { + break; + } + seqlen_offset += seq_len; + seq_len = 1; + input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; + + } + self.minicpm.clear_kv_cache(); + }; + Ok(stream) + } +} diff --git a/src/models/minicpm4/mod.rs b/src/models/minicpm4/mod.rs new file mode 100644 index 0000000..de0a4df --- /dev/null +++ b/src/models/minicpm4/mod.rs @@ -0,0 +1,3 @@ +pub mod config; +pub mod model; +pub mod generate; \ No newline at end of file diff --git a/src/models/minicpm4/model.rs b/src/models/minicpm4/model.rs new file mode 100644 index 0000000..9e17c91 --- /dev/null +++ b/src/models/minicpm4/model.rs @@ -0,0 +1,257 @@ +use crate::{ + models::{ + base_modules::{AttentionNobias, MLPNoBias}, + minicpm4::config::MiniCPM4Config, + }, + position_embed::rope::compute_default_rope_parameters, + utils::tensor_utils::prepare_causal_attention_mask, +}; +use anyhow::{Ok, Result}; +use candle_core::{D, DType, Device, Tensor, Var}; +use candle_nn::{embedding, rms_norm, Embedding, Linear, Module, RmsNorm, VarBuilder}; + +pub struct MiniCPMLongRoPE { + head_dim: usize, + rope_theta: f32, + max_position_embeddings: usize, + short_factor: Vec, + long_factor: Vec, + original_max_position_embeddings: usize, + inv_freq: Tensor, + cos_cached: Tensor, + sin_cached: Tensor, +} +impl MiniCPMLongRoPE { + pub fn new(cfg: &MiniCPM4Config, device: &Device) -> Result { + let head_dim = cfg.hidden_size / cfg.num_attention_heads; + let rope_theta = 10000.0; + let max_position_embeddings = cfg.max_position_embeddings; + let short_factor = cfg.rope_scaling.short_factor.clone(); + let long_factor = cfg.rope_scaling.short_factor.clone(); + let original_max_position_embeddings = cfg.rope_scaling.original_max_position_embeddings; + let scale = max_position_embeddings / original_max_position_embeddings; + let scaling_factor = + (1.0 + (scale as f64).ln() + (original_max_position_embeddings as f64).ln()).sqrt(); + let inv_freq = compute_default_rope_parameters(head_dim, rope_theta); + let inv_freq = Tensor::from_slice(&inv_freq, (1, inv_freq.len()), device)?; + let t = Tensor::arange(0.0_f32, max_position_embeddings as f32, device)? + .reshape((max_position_embeddings, 1))?; + // short_factor.len() = 32 + // head_dim = 1024 / 16 = 64, inv_freq.len() = 32 + let ext_factors = Tensor::from_slice(&short_factor, (1, short_factor.len()), device)?; + let ext_factors = Tensor::ones_like(&ext_factors)?.div(&ext_factors)?; + // (seq_len, 1) matmul (1, 32) -> (seq_len, 32) * (1, 32)-> (seq_len, 32) + let freqs = t.matmul(&ext_factors)?.broadcast_mul(&inv_freq)?; + + let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?; + let cos_cached = emb.cos()?.affine(scaling_factor, 0.0)?; + let sin_cached = emb.sin()?.affine(scaling_factor, 0.0)?; + Ok(Self { + head_dim, + rope_theta, + max_position_embeddings, + short_factor, + long_factor, + original_max_position_embeddings, + inv_freq, + cos_cached, + sin_cached, + }) + } + pub fn update_cos_sin_cache(&mut self, seqlen: usize, device: &Device) -> Result<()> { + let t = Tensor::arange(0.0_f32, seqlen as f32, device)?.reshape((seqlen, 1))?; + let mut ext_factors = + Tensor::from_slice(&self.short_factor, (1, self.short_factor.len()), device)?; + if seqlen > self.original_max_position_embeddings { + ext_factors = + Tensor::from_slice(&self.long_factor, (1, self.long_factor.len()), device)?; + } + let ext_factors = Tensor::ones_like(&ext_factors)?.div(&ext_factors)?; + let freqs = t.matmul(&ext_factors)?.broadcast_mul(&self.inv_freq)?; + let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?; + let scale = seqlen / self.original_max_position_embeddings; + let scaling_factor = + (1.0 + (scale as f64).ln() + (self.original_max_position_embeddings as f64).ln()) + .sqrt(); + let cos_cached = emb.cos()?.affine(scaling_factor, 0.0)?; + let sin_cached = emb.sin()?.affine(scaling_factor, 0.0)?; + self.cos_cached = cos_cached; + self.sin_cached = sin_cached; + Ok(()) + } + pub fn forward(&self, pos_offset: usize, seqlen: usize) -> Result<(Tensor, Tensor)> { + let cos = self.cos_cached.narrow(0, pos_offset, seqlen)?; + let sin = self.sin_cached.narrow(0, pos_offset, seqlen)?; + Ok((cos, sin)) + } +} + +pub struct MiniCPMDecoderLayer { + self_attn: AttentionNobias, + mlp: MLPNoBias, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, + scale_depth: f32, + num_hidden_layers: usize, +} + +impl MiniCPMDecoderLayer { + pub fn new(vb: VarBuilder, cfg: &MiniCPM4Config) -> Result { + let self_attn = AttentionNobias::new( + vb.pp("self_attn"), + cfg.hidden_size, + cfg.num_attention_heads, + cfg.num_key_value_heads, + )?; + let mlp = MLPNoBias::new( + vb.pp("mlp"), + cfg.hidden_size, + cfg.intermediate_size, + cfg.hidden_act, + )?; + 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, + scale_depth: cfg.scale_depth, + num_hidden_layers: cfg.num_hidden_layers, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + cos: &Tensor, + sin: &Tensor, + attention_mask: Option<&Tensor>, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, cos, sin, attention_mask)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; + let xs = (residual + xs)?; + Ok(xs) + } + + pub fn forward_step( + &mut self, + xs: &Tensor, + cos: &Tensor, + sin: &Tensor, + attention_mask: Option<&Tensor>, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward_step(&xs, cos, sin, attention_mask)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; + let xs = (residual + xs)?; + Ok(xs) + } + pub fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache(); + } +} + +pub struct MiniCPMModel { + cfg: MiniCPM4Config, + embed_tokens: Embedding, + layers: Vec, + norm: RmsNorm, + rope_emb: MiniCPMLongRoPE, + lm_head: Linear, +} + +impl MiniCPMModel { + pub fn new(vb: VarBuilder, cfg: MiniCPM4Config) -> Result { + let embed_tokens = embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("embed_tokens"))?; + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_layers = vb.pp("layers"); + for i in 0..cfg.num_hidden_layers { + let layer = MiniCPMDecoderLayer::new(vb_layers.pp(i), &cfg)?; + layers.push(layer); + } + let norm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("norm"))?; + let rope_emb = MiniCPMLongRoPE::new(&cfg, vb.device())?; + let lm_head = Linear::new(embed_tokens.embeddings().clone(), None); + Ok(Self { + cfg, + embed_tokens, + layers, + norm, + rope_emb, + lm_head + }) + } + + pub fn forward(&self, input_ids: &Tensor, position_id: usize) -> Result { + let (bs, seq_len) = input_ids.dims2()?; + let input_embeds = self.embed_tokens.forward(&input_ids)?; + let attention_mask: Option<&Tensor> = { + if seq_len <= 1 { + None + } else { + Some(&prepare_causal_attention_mask( + bs, + seq_len, + position_id, + input_ids.device(), + )?) + } + }; + + let (cos, sin) = self.rope_emb.forward(position_id, seq_len)?; + let mut hidden_states = input_embeds; + for decode_layer in &self.layers { + hidden_states = decode_layer.forward(&hidden_states, &cos, &sin, attention_mask)?; + } + hidden_states = self.norm.forward(&hidden_states)?; + let hidden_state = hidden_states.narrow(1, seq_len - 1, 1)?; + let logits = self.lm_head.forward(&hidden_state)?; + Ok(logits) + } + + pub fn forward_step(&mut self, input_ids: &Tensor, position_id: usize) -> Result { + let (bs, seq_len) = input_ids.dims2()?; + let input_embeds = self.embed_tokens.forward(&input_ids)?; + let attention_mask: Option<&Tensor> = { + if seq_len <= 1 { + None + } else { + Some(&prepare_causal_attention_mask( + bs, + seq_len, + position_id, + input_ids.device(), + )?) + } + }; + + let (cos, sin) = self.rope_emb.forward(position_id, seq_len)?; + let mut hidden_states = input_embeds; + for decode_layer in &mut self.layers { + hidden_states = decode_layer.forward_step(&hidden_states, &cos, &sin, attention_mask)?; + } + hidden_states = self.norm.forward(&hidden_states)?; + let hidden_state = hidden_states.narrow(1, seq_len - 1, 1)?; + let logits = self.lm_head.forward(&hidden_state)?; + Ok(logits) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +} diff --git a/src/models/mod.rs b/src/models/mod.rs index cd3f008..c84b1dd 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,4 +1,7 @@ pub mod qwen2_5vl; +pub mod minicpm4; +pub mod base_modules; + use anyhow::Result; use candle_core::{DType, Device}; use openai_dive::v1::resources::chat::{ diff --git a/src/models/qwen2_5vl/config.rs b/src/models/qwen2_5vl/config.rs index a3f3a9d..f6a62f6 100644 --- a/src/models/qwen2_5vl/config.rs +++ b/src/models/qwen2_5vl/config.rs @@ -24,10 +24,10 @@ pub struct RopeScaling { } #[derive(Debug, Clone, PartialEq, serde::Deserialize)] -pub struct Config { +pub struct Qwen2_5VLConfig { pub attention_dropout: f32, - pub bos_token_id: usize, - pub eos_token_id: usize, + pub bos_token_id: u32, + pub eos_token_id: u32, pub vision_start_token_id: usize, pub vision_end_token_id: usize, pub vision_token_id: usize, diff --git a/src/models/qwen2_5vl/generate.rs b/src/models/qwen2_5vl/generate.rs index 42c9020..e0f5ac5 100644 --- a/src/models/qwen2_5vl/generate.rs +++ b/src/models/qwen2_5vl/generate.rs @@ -1,5 +1,5 @@ // use crate::models::GenerateStream; -use crate::models::qwen2_5vl::config::Config; +use crate::models::qwen2_5vl::config::Qwen2_5VLConfig; use crate::utils::utils::{ build_completion_chunk_response, build_completion_response, find_safetensors_files, get_device, get_dtype, get_logit_processor, @@ -15,7 +15,6 @@ use crate::{ use anyhow::{Result, anyhow}; use candle_core::{D, DType, Device, IndexOp, Tensor}; use candle_nn::VarBuilder; -use candle_transformers::generation::LogitsProcessor; use openai_dive::v1::resources::chat::{ ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, }; @@ -37,13 +36,13 @@ impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> { let chat_template = ChatTemplate::init(path)?; let tokenizer = TokenizerModel::init(path)?; let config_path = path.to_string() + "/config.json"; - let cfg: Config = serde_json::from_slice(&std::fs::read(config_path)?)?; + let cfg: Qwen2_5VLConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; let device = &get_device(device); let cfg_dtype = cfg.torch_dtype.as_str(); let dtype = get_dtype(dtype, cfg_dtype); let pre_processor = Qwen2_5VLProcessor::new(device, dtype)?; - let endoftext_id = cfg.bos_token_id as u32; - let im_end_id = cfg.eos_token_id as u32; + let endoftext_id = cfg.bos_token_id; + let im_end_id = cfg.eos_token_id; let model_list = find_safetensors_files(&path)?; let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, device)? }; let qwen2_5_vl = Qwen2_5VLModel::new(cfg, vb)?; diff --git a/src/models/qwen2_5vl/model.rs b/src/models/qwen2_5vl/model.rs index 0a32f00..1b4c1f7 100644 --- a/src/models/qwen2_5vl/model.rs +++ b/src/models/qwen2_5vl/model.rs @@ -1,12 +1,10 @@ use crate::{ - models::qwen2_5vl::config::{Config, RopeScaling}, + models::qwen2_5vl::config::{Qwen2_5VLConfig, RopeScaling}, position_embed::rope::{ - Qwen2_5VLTextRotaryEmbedding, Qwen2_5VisionRotaryEmbedding, apply_rotary_pos_emb, - apply_rotary_pos_emb_vision, + apply_rotary_pos_emb, apply_rotary_pos_emb_vision, Qwen2_5VLTextRotaryEmbedding, Qwen2_5VisionRotaryEmbedding }, utils::tensor_utils::{ - get_equal_mask, get_vision_next_indices, masked_scatter_dim0, nonzero_index, - safe_arg_sort_last_dim, zero_index, + get_equal_mask, get_vision_next_indices, masked_scatter_dim0, nonzero_index, repeat_kv, safe_arg_sort_last_dim, zero_index }, }; use anyhow::{Result, anyhow}; @@ -20,7 +18,7 @@ pub struct Qwen2_5VisionPatchEmbed { } impl Qwen2_5VisionPatchEmbed { - pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + pub fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result { let patch_size = cfg.vision_config.patch_size; let temporal_patch_size = cfg.vision_config.temporal_patch_size; let in_channels = cfg.vision_config.in_chans; @@ -60,7 +58,7 @@ pub struct Qwen2_5VLPatchMerger { } impl Qwen2_5VLPatchMerger { - pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + pub fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result { let hidden_size = cfg.vision_config.hidden_size * (cfg.vision_config.spatial_merge_size.pow(2)); let ln_q = rms_norm( @@ -99,7 +97,7 @@ struct Qwen2_5VLVisionMLP { } impl Qwen2_5VLVisionMLP { - fn new(cfg: &Config, vb: VarBuilder) -> Result { + fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result { let hidden_sz = cfg.vision_config.hidden_size; let intermediate_sz = cfg.vision_config.intermediate_size; let gate_proj = linear(hidden_sz, intermediate_sz, vb.pp("gate_proj"))?; @@ -131,7 +129,7 @@ struct Qwen2_5VLVisionAttention { } impl Qwen2_5VLVisionAttention { - fn new(cfg: &Config, vb: VarBuilder) -> Result { + fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result { let hidden_size = cfg.vision_config.hidden_size; let num_heads = cfg.vision_config.num_heads; let head_dim = hidden_size / num_heads; @@ -200,7 +198,7 @@ struct Qwen2_5VLVisionBlock { } impl Qwen2_5VLVisionBlock { - fn new(cfg: &Config, vb: VarBuilder) -> Result { + fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result { let attn = Qwen2_5VLVisionAttention::new(cfg, vb.pp("attn"))?; let mlp = Qwen2_5VLVisionMLP::new(cfg, vb.pp("mlp"))?; let norm1 = rms_norm( @@ -254,7 +252,7 @@ pub struct Qwen2_5VLVisionModel { } impl Qwen2_5VLVisionModel { - pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + pub fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result { let spatial_merge_size = cfg.vision_config.spatial_merge_size; let patch_size = cfg.vision_config.patch_size; let fullatt_block_indexes = cfg.vision_config.fullatt_block_indexes.clone(); @@ -539,23 +537,6 @@ impl Qwen2_5VLVisionModel { } } -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 - let kv = Tensor::cat(&vec![&xs; n_rep], 2)?.reshape(( - b_sz, - n_kv_head * n_rep, - seq_len, - head_dim, - ))?; - Ok(kv) - } -} #[derive(Debug, Clone)] struct Qwen2_5VLTextMLP { @@ -566,7 +547,7 @@ struct Qwen2_5VLTextMLP { } impl Qwen2_5VLTextMLP { - fn new(cfg: &Config, vb: VarBuilder) -> Result { + fn new(cfg: &Qwen2_5VLConfig, 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"))?; @@ -605,7 +586,7 @@ struct Qwen2_5VLTextAttention { } impl Qwen2_5VLTextAttention { - fn new(cfg: &Config, vb: VarBuilder) -> Result { + fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result { let hidden_size = cfg.hidden_size; let num_heads = cfg.num_attention_heads; let num_kv_heads = cfg.num_key_value_heads; @@ -721,7 +702,7 @@ struct Qwen2_5VLTextDecoderLayer { } impl Qwen2_5VLTextDecoderLayer { - fn new(cfg: &Config, vb: VarBuilder) -> Result { + fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result { let self_attn = Qwen2_5VLTextAttention::new(cfg, vb.pp("self_attn"))?; let mlp = Qwen2_5VLTextMLP::new(cfg, vb.pp("mlp"))?; let input_layernorm = @@ -774,7 +755,7 @@ pub struct Qwen2_5VLTextModel { } impl Qwen2_5VLTextModel { - pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + pub fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result { let embed_tokens = candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("embed_tokens"))?; let head_dim = cfg.hidden_size / cfg.num_attention_heads; @@ -879,13 +860,13 @@ impl Qwen2_5VLTextModel { pub struct Qwen2_5VLModel { visual: Qwen2_5VLVisionModel, model: Qwen2_5VLTextModel, - pub cfg: Config, + pub cfg: Qwen2_5VLConfig, lm_head: Linear, rope_deltas: Option, } impl Qwen2_5VLModel { - pub fn new(cfg: Config, vb: VarBuilder) -> Result { + pub fn new(cfg: Qwen2_5VLConfig, vb: VarBuilder) -> Result { let visual = Qwen2_5VLVisionModel::new(&cfg, vb.pp("visual"))?; let model = Qwen2_5VLTextModel::new(&cfg, vb.pp("model"))?; let vocab_size = cfg.vocab_size; diff --git a/src/position_embed/rope.rs b/src/position_embed/rope.rs index 165bcf0..cc70efc 100644 --- a/src/position_embed/rope.rs +++ b/src/position_embed/rope.rs @@ -2,8 +2,6 @@ use anyhow::Result; use candle_core::{D, DType, Device, IndexOp, Tensor}; use candle_transformers::models::deepseek2::SplitOp; -use crate::models::qwen2_5vl::config::RopeScaling; - pub fn compute_default_rope_parameters(dim: usize, base: f32) -> Vec { let inv_freq: Vec = (0..dim) .step_by(2) diff --git a/src/utils/tensor_utils.rs b/src/utils/tensor_utils.rs index 9426025..41a851c 100644 --- a/src/utils/tensor_utils.rs +++ b/src/utils/tensor_utils.rs @@ -1,6 +1,55 @@ use anyhow::{Result, anyhow}; use candle_core::{D, DType, Device, IndexOp, Tensor, shape::Dim}; +pub fn prepare_causal_attention_mask( + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + device: &Device +) -> 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), device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::U32, device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + let mask = mask + .expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(DType::U32)?; + Ok(mask) +} + +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 + let kv = Tensor::cat(&vec![&xs; n_rep], 2)?.reshape(( + b_sz, + n_kv_head * n_rep, + seq_len, + head_dim, + ))?; + Ok(kv) + } +} + pub fn split(t: &Tensor, splits: &[usize], dim: D) -> Result> { let dim = dim.to_index(t.shape(), "split")?; let mut split_res = Vec::new(); diff --git a/tests/config_tests.rs b/tests/config_tests.rs index 445e86d..3c30f03 100644 --- a/tests/config_tests.rs +++ b/tests/config_tests.rs @@ -1,12 +1,23 @@ -use aha::models::qwen2_5vl::config::Config; +use aha::models::{minicpm4::config::MiniCPM4Config, qwen2_5vl::config::Qwen2_5VLConfig}; use anyhow::Result; #[test] -fn qwen2_5vl_config() -> Result<()> { - // cargo test qwen2_5vl_config -- --nocapture +fn qwen2_5_vl_config() -> Result<()> { + // cargo test -F cuda,flash-attn qwen2_5vl_config -- --nocapture let model_path = "/home/jhq/huggingface_model/Qwen/Qwen2.5-VL-3B-Instruct/"; let config_path = model_path.to_string() + "/config.json"; - let config: Config = serde_json::from_slice(&std::fs::read(config_path)?)?; + let config: Qwen2_5VLConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; println!("{:?}", config); Ok(()) } + + +#[test] +fn minicpm4_config() -> Result<()> { + // cargo test -F cuda,flash-attn minicpm4_config -- --nocapture + let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/"; + let config_path = model_path.to_string() + "/config.json"; + let config: MiniCPM4Config = serde_json::from_slice(&std::fs::read(config_path)?)?; + println!("{:?}", config); + Ok(()) +} \ No newline at end of file diff --git a/tests/test_minicpm4.rs b/tests/test_minicpm4.rs new file mode 100644 index 0000000..74bebd4 --- /dev/null +++ b/tests/test_minicpm4.rs @@ -0,0 +1,97 @@ +use std::time::Instant; + +use anyhow::Result; +use candle_core::{DType, Device}; +use openai_dive::v1::resources::chat::ChatCompletionParameters; + +#[test] +fn qwen2_5vl_generate() -> Result<()> { + // test with cpu :(太慢了, : RUST_BACKTRACE=1 cargo test qwen2_5vl_generate -- --nocapture + // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda qwen2_5vl_generate -- --nocapture + // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn qwen2_5vl_generate -- --nocapture + let device = Device::cuda_if_available(0)?; + let dtype = DType::BF16; + + let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/"; + + let message = r#" + { + "model": "minicpm4", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "你是谁" + } + ] + } + ] + } + "#; + let mes: ChatCompletionParameters = serde_json::from_str(message)?; + let i_start = Instant::now(); + // let mut model = Qwen2_5VLGenerateModel::init(model_path, &device, dtype)?; + let mut model = ModelType::init(ModelType::Qwen2_5VL, model_path, None, None)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in load model is: {:?}", i_duration); + + let i_start = Instant::now(); + let result = model.generate(mes)?; + println!("generate: \n {:?}", result); + let i_duration = i_start.elapsed(); + println!("Time elapsed in generate is: {:?}", i_duration); + + Ok(()) +} + +#[tokio::test] +async fn qwen2_5vl_stream() -> Result<()> { + // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn qwen2_5vl_generate -- --nocapture + let device = Device::cuda_if_available(0)?; + let dtype = DType::BF16; + + let model_path = "/home/jhq/huggingface_model/Qwen/Qwen2.5-VL-3B-Instruct/"; + + let message = r#" + { + "model": "qwen2.5vl", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image", + "image_url": + { + "url": "file://./assets/img/ocr_test.png" + } + }, + { + "type": "text", + "text": "请分析图片并提取所有可见文本内容，按从左到右、从上到下的布局，返回纯文本" + } + ] + } + ] + } + "#; + let mes: ChatCompletionParameters = serde_json::from_str(message)?; + let i_start = Instant::now(); + // let mut model = Qwen2_5VLGenerateModel::init(model_path, &device, dtype)?; + let mut model = ModelType::init(ModelType::Qwen2_5VL, model_path, None, None)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in load model is: {:?}", i_duration); + + let i_start = Instant::now(); + let mut stream = pin!(model.generate_stream(mes)?); + while let Some(item) = stream.next().await { + println!("generate: \n {:?}", item); + } + + let i_duration = i_start.elapsed(); + println!("Time elapsed in generate is: {:?}", i_duration); + + Ok(()) +} diff --git a/tests/weight_test.rs b/tests/weight_test.rs new file mode 100644 index 0000000..dd8761b --- /dev/null +++ b/tests/weight_test.rs @@ -0,0 +1,18 @@ +use aha::utils::utils::find_safetensors_files; +use anyhow::Result; +use candle_core::{safetensors, Device}; +#[test] +fn minicpm4_weight() -> Result<()> { + let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/"; + let model_list = find_safetensors_files(&model_path)?; + let device = Device::Cpu; + for m in model_list { + let weights = safetensors::load(m, &device)?; + for (key, tensor) in weights.iter() { + println!("=== {} ===", key); + println!("Shape: {:?}", tensor.shape()); + println!("DType: {:?}", tensor.dtype()); + } + } + Ok(()) +} \ No newline at end of file