tckm.yan/aha
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add voxcpm with some bug

Browse Source
This commit is contained in:
jhqxxx
2025-10-03 22:25:58 +08:00
parent ae36194a4b
commit f33eaaee0d
30 changed files with 2411 additions and 165 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/models/minicpm4/model.rs
+79 -40
View File
@@ -7,38 +7,41 @@ use crate::{
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};
use candle_core::{D, Device, Tensor};
use candle_nn::{Embedding, Linear, Module, RmsNorm, VarBuilder, embedding, rms_norm};
pub struct MiniCPMLongRoPE {
head_dim: usize,
rope_theta: f32,
max_position_embeddings: usize,
short_factor: Vec<f32>,
long_factor: Vec<f32>,
original_max_position_embeddings: usize,
max_seq_len_cached: usize,
scaling_factor: f64,
inv_freq: Tensor,
cos_cached: Tensor,
sin_cached: Tensor,
device: Device,
}
impl MiniCPMLongRoPE {
pub fn new(cfg: &MiniCPM4Config, device: &Device) -> Result<Self> {
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 max_position_embeddings = cfg.max_position_embeddings;
let scale = max_position_embeddings as f64 / original_max_position_embeddings as f64;
let scaling_factor =
(1.0 + (scale as f64).ln() + (original_max_position_embeddings as f64).ln()).sqrt();
(1.0 + scale.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 inv_freq =
Tensor::from_slice(&inv_freq, (1, inv_freq.len()), device)?;
let max_seq_len_cached = max_position_embeddings;
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::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)?;
@@ -47,41 +50,46 @@ impl MiniCPMLongRoPE {
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,
max_seq_len_cached,
scaling_factor,
inv_freq,
cos_cached,
sin_cached,
device: device.clone(),
})
}
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)?;
pub fn update_cos_sin_cache(&mut self, seqlen: usize) -> Result<()> {
self.max_seq_len_cached = seqlen;
let t = Tensor::arange(0.0_f32, seqlen as f32, &self.device)?
.reshape((seqlen, 1))?;
let mut ext_factors = Tensor::from_slice(
&self.short_factor,
(1, self.short_factor.len()),
&self.device,
)?;
if seqlen > self.original_max_position_embeddings {
ext_factors =
Tensor::from_slice(&self.long_factor, (1, self.long_factor.len()), device)?;
Tensor::from_slice(&self.long_factor, (1, self.long_factor.len()), &self.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)?;
let cos_cached = emb.cos()?.affine(self.scaling_factor, 0.0)?;
let sin_cached = emb.sin()?.affine(self.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)> {
pub fn forward(&mut self, pos_offset: usize, seqlen: usize) -> Result<(Tensor, Tensor)> {
if pos_offset + seqlen > self.max_seq_len_cached {
let _ = self.update_cos_sin_cache(pos_offset + seqlen)?;
}
let cos = self.cos_cached.narrow(0, pos_offset, seqlen)?;
let sin = self.sin_cached.narrow(0, pos_offset, seqlen)?;
Ok((cos, sin))
}
}
@@ -133,13 +141,21 @@ impl MiniCPMDecoderLayer {
sin: &Tensor,
attention_mask: Option<&Tensor>,
) -> Result<Tensor> {
let residual = xs;
let residual = xs.clone();
let xs = self.input_layernorm.forward(xs)?;
let xs = self.self_attn.forward(&xs, cos, sin, attention_mask)?;
let xs = (xs + residual)?;
let xs = self.self_attn.forward(&xs, cos, sin, attention_mask, true)?;
let xs = (residual
+ xs.affine(
self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(),
0.0,
))?;
let residual = &xs;
let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?;
let xs = (residual + xs)?;
let xs = (residual
+ xs.affine(
self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(),
0.0,
))?;
Ok(xs)
}
@@ -150,13 +166,21 @@ impl MiniCPMDecoderLayer {
sin: &Tensor,
attention_mask: Option<&Tensor>,
) -> Result<Tensor> {
let residual = xs;
let residual = xs.clone();
let xs = self.input_layernorm.forward(xs)?;
let xs = self.self_attn.forward_step(&xs, cos, sin, attention_mask)?;
let xs = (xs + residual)?;
let xs = self.self_attn.forward_step(&xs, cos, sin, attention_mask, true)?;
let xs = (residual
+ xs.affine(
self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(),
0.0,
))?;
let residual = &xs;
let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?;
let xs = (residual + xs)?;
let xs = (residual
+ xs.affine(
self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(),
0.0,
))?;
Ok(xs)
}
pub fn clear_kv_cache(&mut self) {
@@ -175,6 +199,7 @@ pub struct MiniCPMModel {
impl MiniCPMModel {
pub fn new(vb: VarBuilder, cfg: MiniCPM4Config) -> Result<Self> {
let vb = vb.pp("model");
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");
@@ -191,13 +216,16 @@ impl MiniCPMModel {
layers,
norm,
rope_emb,
lm_head
lm_head,
})
}
pub fn forward(&self, input_ids: &Tensor, position_id: usize) -> Result<Tensor> {
pub fn forward(&mut self, input_ids: &Tensor, position_id: usize) -> Result<Tensor> {
let (bs, seq_len) = input_ids.dims2()?;
let input_embeds = self.embed_tokens.forward(&input_ids)?;
let input_embeds = self
.embed_tokens
.forward(&input_ids)?
.affine(self.cfg.scale_emb, 0.0)?;
let attention_mask: Option<&Tensor> = {
if seq_len <= 1 {
None
@@ -210,7 +238,7 @@ impl MiniCPMModel {
)?)
}
};
let (cos, sin) = self.rope_emb.forward(position_id, seq_len)?;
let mut hidden_states = input_embeds;
for decode_layer in &self.layers {
@@ -218,13 +246,20 @@ impl MiniCPMModel {
}
hidden_states = self.norm.forward(&hidden_states)?;
let hidden_state = hidden_states.narrow(1, seq_len - 1, 1)?;
let hidden_state = hidden_state.affine(
1.0 / (self.cfg.hidden_size / self.cfg.dim_model_base) as f64,
0.0,
)?;
let logits = self.lm_head.forward(&hidden_state)?;
Ok(logits)
}
pub fn forward_step(&mut self, input_ids: &Tensor, position_id: usize) -> Result<Tensor> {
let (bs, seq_len) = input_ids.dims2()?;
let input_embeds = self.embed_tokens.forward(&input_ids)?;
let input_embeds = self
.embed_tokens
.forward(&input_ids)?
.affine(self.cfg.scale_emb, 0.0)?;
let attention_mask: Option<&Tensor> = {
if seq_len <= 1 {
None
@@ -237,14 +272,18 @@ impl MiniCPMModel {
)?)
}
};
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 =
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 hidden_state = hidden_state.affine(
1.0 / (self.cfg.hidden_size / self.cfg.dim_model_base) as f64,
0.0,
)?;
let logits = self.lm_head.forward(&hidden_state)?;
Ok(logits)
}