add qwen3 and fun-asr-nano
This commit is contained in:
jhqxxx
@@ -0,0 +1,277 @@
|
||||
use anyhow::Result;
|
||||
use candle_core::Tensor;
|
||||
use candle_nn::{
|
||||
Embedding, Linear, Module, RmsNorm, VarBuilder, embedding, linear, linear_no_bias, rms_norm,
|
||||
};
|
||||
|
||||
use crate::{
|
||||
models::{
|
||||
common::{GateUpDownMLP, eager_attention_forward},
|
||||
qwen3::config::Qwen3Config,
|
||||
},
|
||||
position_embed::rope::{RoPE, apply_rotary_pos_emb},
|
||||
utils::tensor_utils::prepare_causal_attention_mask,
|
||||
};
|
||||
|
||||
pub struct Qwen3Attention {
|
||||
q_proj: Linear,
|
||||
k_proj: Linear,
|
||||
v_proj: Linear,
|
||||
o_proj: Linear,
|
||||
q_norm: RmsNorm,
|
||||
k_norm: RmsNorm,
|
||||
num_attention_heads: usize,
|
||||
num_key_value_heads: usize,
|
||||
num_kv_groups: usize,
|
||||
head_dim: usize,
|
||||
scaling: f64,
|
||||
kv_cache: Option<(Tensor, Tensor)>,
|
||||
}
|
||||
|
||||
impl Qwen3Attention {
|
||||
pub fn new(config: &Qwen3Config, vb: VarBuilder) -> Result<Self> {
|
||||
let hidden_size = config.hidden_size;
|
||||
let num_attention_heads = config.num_attention_heads;
|
||||
let head_dim = config.head_dim;
|
||||
let num_key_value_heads = config.num_key_value_heads;
|
||||
let num_kv_groups = num_attention_heads / num_key_value_heads;
|
||||
let scaling = 1f64 / f64::sqrt(head_dim as f64);
|
||||
let (q_proj, k_proj, v_proj, o_proj) = if config.attention_bias {
|
||||
let q_proj = linear(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?;
|
||||
let k_proj = linear(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?;
|
||||
let v_proj = linear(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?;
|
||||
let o_proj = linear(num_attention_heads * head_dim, hidden_size, vb.pp("o_proj"))?;
|
||||
(q_proj, k_proj, v_proj, o_proj)
|
||||
} else {
|
||||
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(num_attention_heads * head_dim, hidden_size, vb.pp("o_proj"))?;
|
||||
(q_proj, k_proj, v_proj, o_proj)
|
||||
};
|
||||
let q_norm = rms_norm(head_dim, config.rms_norm_eps, vb.pp("q_norm"))?;
|
||||
let k_norm = rms_norm(head_dim, config.rms_norm_eps, vb.pp("k_norm"))?;
|
||||
Ok(Self {
|
||||
q_proj,
|
||||
k_proj,
|
||||
v_proj,
|
||||
o_proj,
|
||||
q_norm,
|
||||
k_norm,
|
||||
num_attention_heads,
|
||||
num_key_value_heads,
|
||||
num_kv_groups,
|
||||
head_dim,
|
||||
scaling,
|
||||
kv_cache: None,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn forward(
|
||||
&mut self,
|
||||
xs: &Tensor,
|
||||
cos: &Tensor,
|
||||
sin: &Tensor,
|
||||
attention_mask: Option<&Tensor>,
|
||||
) -> Result<Tensor> {
|
||||
let (b_sz, q_len, _) = xs.dims3()?;
|
||||
let query_states = self.q_proj.forward(xs)?.reshape((
|
||||
b_sz,
|
||||
q_len,
|
||||
self.num_attention_heads,
|
||||
self.head_dim,
|
||||
))?;
|
||||
let query_states = self.q_norm.forward(&query_states)?.transpose(1, 2)?;
|
||||
let key_states = self.k_proj.forward(xs)?.reshape((
|
||||
b_sz,
|
||||
q_len,
|
||||
self.num_key_value_heads,
|
||||
self.head_dim,
|
||||
))?;
|
||||
let key_states = self.k_norm.forward(&key_states)?.transpose(1, 2)?;
|
||||
let value_states = self.v_proj.forward(xs)?;
|
||||
let value_states = value_states
|
||||
.reshape((b_sz, q_len, self.num_key_value_heads, self.head_dim))?
|
||||
.transpose(1, 2)?;
|
||||
let (query_states, key_states) =
|
||||
apply_rotary_pos_emb(&query_states, &key_states, cos, sin, false)?;
|
||||
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 attn_output = eager_attention_forward(
|
||||
&query_states,
|
||||
&key_states,
|
||||
&value_states,
|
||||
Some(self.num_kv_groups),
|
||||
attention_mask,
|
||||
self.scaling,
|
||||
)?;
|
||||
let attn_output =
|
||||
attn_output.reshape((b_sz, q_len, self.num_attention_heads * self.head_dim))?;
|
||||
let attn_output = attn_output.apply(&self.o_proj)?;
|
||||
Ok(attn_output)
|
||||
}
|
||||
|
||||
pub fn clear_kv_cache(&mut self) {
|
||||
self.kv_cache = None
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Qwen3DecoderLayer {
|
||||
self_attn: Qwen3Attention,
|
||||
mlp: GateUpDownMLP,
|
||||
input_layernorm: RmsNorm,
|
||||
post_attention_layernorm: RmsNorm,
|
||||
}
|
||||
|
||||
impl Qwen3DecoderLayer {
|
||||
pub fn new(config: &Qwen3Config, vb: VarBuilder) -> Result<Self> {
|
||||
let self_attn = Qwen3Attention::new(config, vb.pp("self_attn"))?;
|
||||
let mlp = GateUpDownMLP::new(
|
||||
vb.pp("mlp"),
|
||||
config.hidden_size,
|
||||
config.intermediate_size,
|
||||
config.hidden_act,
|
||||
false,
|
||||
)?;
|
||||
let input_layernorm = rms_norm(
|
||||
config.hidden_size,
|
||||
config.rms_norm_eps,
|
||||
vb.pp("input_layernorm"),
|
||||
)?;
|
||||
let post_attention_layernorm = rms_norm(
|
||||
config.hidden_size,
|
||||
config.rms_norm_eps,
|
||||
vb.pp("post_attention_layernorm"),
|
||||
)?;
|
||||
Ok(Self {
|
||||
self_attn,
|
||||
mlp,
|
||||
input_layernorm,
|
||||
post_attention_layernorm,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn forward(
|
||||
&mut self,
|
||||
xs: &Tensor,
|
||||
cos: &Tensor,
|
||||
sin: &Tensor,
|
||||
attention_mask: Option<&Tensor>,
|
||||
) -> Result<Tensor> {
|
||||
let residual = xs.clone();
|
||||
let xs = self.input_layernorm.forward(xs)?;
|
||||
let xs = self.self_attn.forward(&xs, cos, sin, attention_mask)?;
|
||||
let xs = residual.add(&xs)?;
|
||||
let residual = xs.clone();
|
||||
let xs = self.post_attention_layernorm.forward(&xs)?;
|
||||
let xs = self.mlp.forward(&xs)?;
|
||||
let xs = residual.add(&xs)?;
|
||||
Ok(xs)
|
||||
}
|
||||
|
||||
pub fn clear_kv_cache(&mut self) {
|
||||
self.self_attn.clear_kv_cache();
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Qwen3Model {
|
||||
embed_tokens: Embedding,
|
||||
layers: Vec<Qwen3DecoderLayer>,
|
||||
norm: RmsNorm,
|
||||
rotary_emb: RoPE,
|
||||
lm_head: Linear,
|
||||
}
|
||||
|
||||
impl Qwen3Model {
|
||||
pub fn new(config: &Qwen3Config, vb: VarBuilder) -> Result<Self> {
|
||||
let vb = vb.pp("model");
|
||||
let vocab_size = config.vocab_size;
|
||||
let embed_tokens = embedding(vocab_size, config.hidden_size, vb.pp("embed_tokens"))?;
|
||||
let mut layers = vec![];
|
||||
let vb_l = vb.pp("layers");
|
||||
for layer_idx in 0..config.num_hidden_layers {
|
||||
let layer = Qwen3DecoderLayer::new(config, vb_l.pp(layer_idx))?;
|
||||
layers.push(layer)
|
||||
}
|
||||
let norm = rms_norm(config.hidden_size, config.rms_norm_eps, vb.pp("norm"))?;
|
||||
let head_dim = config.head_dim;
|
||||
let rotary_emb = RoPE::new(head_dim, config.rope_theta, vb.device())?;
|
||||
let lm_head = if config.tie_word_embeddings {
|
||||
Linear::new(embed_tokens.embeddings().clone(), None)
|
||||
} else {
|
||||
linear_no_bias(config.hidden_size, config.vocab_size, vb.pp("lm_head"))?
|
||||
};
|
||||
Ok(Self {
|
||||
embed_tokens,
|
||||
layers,
|
||||
norm,
|
||||
rotary_emb,
|
||||
lm_head,
|
||||
})
|
||||
}
|
||||
pub fn forward(
|
||||
&mut self,
|
||||
input_ids: Option<&Tensor>,
|
||||
inputs_embeds: Option<&Tensor>,
|
||||
seqlen_offset: usize,
|
||||
) -> Result<Tensor> {
|
||||
if input_ids.is_none() && inputs_embeds.is_none() {
|
||||
return Err(anyhow::anyhow!(
|
||||
"You must specify exactly one of input_ids or inputs_embeds"
|
||||
));
|
||||
}
|
||||
let inputs_embeds = if let Some(inputs_embeds) = inputs_embeds {
|
||||
inputs_embeds.clone()
|
||||
} else {
|
||||
let input_ids = input_ids.unwrap();
|
||||
self.embedding_token_id(input_ids)?
|
||||
};
|
||||
let (bs, seq_len, _) = inputs_embeds.dims3()?;
|
||||
let attention_mask: Option<Tensor> = {
|
||||
if seq_len <= 1 {
|
||||
None
|
||||
} else {
|
||||
Some(prepare_causal_attention_mask(
|
||||
bs,
|
||||
seq_len,
|
||||
0,
|
||||
inputs_embeds.device(),
|
||||
)?)
|
||||
}
|
||||
};
|
||||
|
||||
let (cos, sin) = self
|
||||
.rotary_emb
|
||||
.forward(seqlen_offset, seq_len, inputs_embeds.device())?;
|
||||
|
||||
let mut hidden_states = inputs_embeds;
|
||||
for decode_layer in &mut self.layers {
|
||||
hidden_states =
|
||||
decode_layer.forward(&hidden_states, &cos, &sin, attention_mask.as_ref())?;
|
||||
}
|
||||
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 embedding_token_id(&self, input_ids: &Tensor) -> Result<Tensor> {
|
||||
Ok(self.embed_tokens.forward(input_ids)?)
|
||||
}
|
||||
|
||||
pub fn clear_kv_cache(&mut self) {
|
||||
for layer in self.layers.iter_mut() {
|
||||
layer.clear_kv_cache()
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user