jhqxxx
1211 lines
48 KiB
Rust
1211 lines
48 KiB
Rust
use crate::{
|
||
models::qwen2_5vl::config::{Qwen2_5VLConfig, RopeScaling},
|
||
position_embed::rope::{
|
||
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, repeat_kv, safe_arg_sort_last_dim, zero_index
|
||
},
|
||
};
|
||
use anyhow::{Result, anyhow};
|
||
use candle_core::{D, DType, Device, IndexOp, Tensor};
|
||
use candle_nn::{
|
||
Activation, Init, Linear, Module, RmsNorm, VarBuilder, linear, linear_no_bias, rms_norm,
|
||
};
|
||
|
||
pub struct Qwen2_5VisionPatchEmbed {
|
||
conv3d_weight: Tensor,
|
||
}
|
||
|
||
impl Qwen2_5VisionPatchEmbed {
|
||
pub fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
|
||
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;
|
||
let embed_dim = cfg.vision_config.hidden_size;
|
||
// conv3d weight key: visual.patch_embed.proj.weight, value: Tensor[dims 1280, 3, 2, 14, 14; bf16, cuda:0]
|
||
// (1280, 3, 2, 14, 14) -> (1280, 1176) -> (1176, 1280)
|
||
let conv3d_weight = vb
|
||
.get_with_hints(
|
||
(
|
||
embed_dim,
|
||
in_channels,
|
||
temporal_patch_size,
|
||
patch_size,
|
||
patch_size,
|
||
),
|
||
"proj.weight",
|
||
Init::Const(1.),
|
||
)?
|
||
.flatten(1, 4)?
|
||
.t()?;
|
||
Ok(Self { conv3d_weight })
|
||
}
|
||
|
||
pub fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
|
||
// hidden_states shape: (grid_t*grid_h*grid_w, c*temporal_patch_size*patch_size*patch_size)
|
||
// ((), 1176) matmul (1176, 1280) -> ((), 1280)
|
||
let hidden_states = hidden_states.matmul(&self.conv3d_weight)?;
|
||
Ok(hidden_states)
|
||
}
|
||
}
|
||
|
||
pub struct Qwen2_5VLPatchMerger {
|
||
hidden_size: usize,
|
||
ln_q: RmsNorm,
|
||
mlp_0: Linear,
|
||
mlp_2: Linear,
|
||
}
|
||
|
||
impl Qwen2_5VLPatchMerger {
|
||
pub fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
|
||
let hidden_size =
|
||
cfg.vision_config.hidden_size * (cfg.vision_config.spatial_merge_size.pow(2));
|
||
let ln_q = rms_norm(
|
||
cfg.vision_config.hidden_size,
|
||
cfg.rms_norm_eps,
|
||
vb.pp("ln_q"),
|
||
)?;
|
||
let mlp_0 = linear(hidden_size, hidden_size, vb.pp("mlp.0"))?;
|
||
let mlp_2 = linear(
|
||
hidden_size,
|
||
cfg.vision_config.out_hidden_size,
|
||
vb.pp("mlp.2"),
|
||
)?;
|
||
Ok(Self {
|
||
hidden_size,
|
||
ln_q,
|
||
mlp_0,
|
||
mlp_2,
|
||
})
|
||
}
|
||
}
|
||
impl Module for Qwen2_5VLPatchMerger {
|
||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||
let xs = xs.apply(&self.ln_q)?.reshape(((), self.hidden_size))?;
|
||
let xs = xs.apply(&self.mlp_0)?.gelu()?.apply(&self.mlp_2)?;
|
||
Ok(xs)
|
||
}
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
struct Qwen2_5VLVisionMLP {
|
||
gate_proj: Linear,
|
||
up_proj: Linear,
|
||
down_proj: Linear,
|
||
act_fn: Activation,
|
||
}
|
||
|
||
impl Qwen2_5VLVisionMLP {
|
||
fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
|
||
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"))?;
|
||
let up_proj = linear(hidden_sz, intermediate_sz, vb.pp("up_proj"))?;
|
||
let down_proj = linear(intermediate_sz, hidden_sz, vb.pp("down_proj"))?;
|
||
Ok(Self {
|
||
gate_proj,
|
||
up_proj,
|
||
down_proj,
|
||
act_fn: cfg.hidden_act,
|
||
})
|
||
}
|
||
}
|
||
|
||
impl Module for Qwen2_5VLVisionMLP {
|
||
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)]
|
||
struct Qwen2_5VLVisionAttention {
|
||
qkv: Linear,
|
||
proj: Linear,
|
||
num_heads: usize,
|
||
scale: Tensor,
|
||
}
|
||
|
||
impl Qwen2_5VLVisionAttention {
|
||
fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
|
||
let hidden_size = cfg.vision_config.hidden_size;
|
||
let num_heads = cfg.vision_config.num_heads;
|
||
let head_dim = hidden_size / num_heads;
|
||
let qkv = linear(hidden_size, hidden_size * 3, vb.pp("qkv"))?;
|
||
let proj = linear(hidden_size, hidden_size, vb.pp("proj"))?;
|
||
let scale = Tensor::new(vec![1f32 / (head_dim as f32).sqrt()], vb.device())?
|
||
.to_dtype(vb.dtype())?;
|
||
Ok(Self {
|
||
qkv,
|
||
proj,
|
||
num_heads,
|
||
scale,
|
||
})
|
||
}
|
||
|
||
fn forward(
|
||
&self,
|
||
xs: &Tensor,
|
||
cos: &Tensor,
|
||
sin: &Tensor,
|
||
attention_mask: &Tensor,
|
||
) -> Result<Tensor> {
|
||
// xs: (seq_len, hidden_size)
|
||
let seq_length = xs.dim(0)?;
|
||
// (seq_len, hidden_size) -> (seq_len, hidden_size*3)
|
||
// -> (seq_len, 3, num_heads, head_dim) -> (3, seq_len, num_heads, head_dim)
|
||
let qkv_states = xs
|
||
.apply(&self.qkv)?
|
||
.reshape((seq_length, 3, self.num_heads, ()))?
|
||
.permute((1, 0, 2, 3))?;
|
||
// (seq_len, num_heads, head_dim)
|
||
let query_states = qkv_states.i(0)?.contiguous()?;
|
||
let key_states = qkv_states.i(1)?.contiguous()?;
|
||
let value_states = qkv_states.i(2)?.contiguous()?;
|
||
let (query_states, key_states) =
|
||
apply_rotary_pos_emb_vision(&query_states, &key_states, cos, sin)?;
|
||
// (seq_len, num_heads, head_dim) -> (num_heads, seq_len, head_dim)
|
||
let query_states = query_states.transpose(0, 1)?.contiguous()?;
|
||
let key_states = key_states.transpose(0, 1)?.contiguous()?;
|
||
let value_states = value_states.transpose(0, 1)?.contiguous()?;
|
||
|
||
let attn_output = {
|
||
let attn_weights = query_states
|
||
.matmul(&key_states.transpose(D::Minus2, D::Minus1)?)?
|
||
.broadcast_mul(&self.scale)?;
|
||
let attn_weights = attn_weights.broadcast_add(&attention_mask)?;
|
||
let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?;
|
||
attn_weights.matmul(&value_states)?
|
||
};
|
||
// (num_heads, seq_len, head_dim) -> (seq_len, num_heads, head_dim) -> (seq_len, hidden_size)
|
||
let attn_output = attn_output
|
||
.transpose(0, 1)?
|
||
.reshape((seq_length, ()))?
|
||
.contiguous()?;
|
||
let attn_ouput = attn_output.apply(&self.proj)?;
|
||
Ok(attn_ouput)
|
||
}
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
struct Qwen2_5VLVisionBlock {
|
||
attn: Qwen2_5VLVisionAttention,
|
||
mlp: Qwen2_5VLVisionMLP,
|
||
norm1: RmsNorm,
|
||
norm2: RmsNorm,
|
||
}
|
||
|
||
impl Qwen2_5VLVisionBlock {
|
||
fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
|
||
let attn = Qwen2_5VLVisionAttention::new(cfg, vb.pp("attn"))?;
|
||
let mlp = Qwen2_5VLVisionMLP::new(cfg, vb.pp("mlp"))?;
|
||
let norm1 = rms_norm(
|
||
cfg.vision_config.hidden_size,
|
||
cfg.rms_norm_eps,
|
||
vb.pp("norm1"),
|
||
)?;
|
||
let norm2 = rms_norm(
|
||
cfg.vision_config.hidden_size,
|
||
cfg.rms_norm_eps,
|
||
vb.pp("norm2"),
|
||
)?;
|
||
|
||
Ok(Self {
|
||
attn,
|
||
mlp,
|
||
norm1,
|
||
norm2,
|
||
})
|
||
}
|
||
|
||
fn forward(
|
||
&self,
|
||
xs: &Tensor,
|
||
cos: &Tensor,
|
||
sin: &Tensor,
|
||
attention_mask: &Tensor,
|
||
) -> Result<Tensor> {
|
||
let residual = xs;
|
||
let xs = self.norm1.forward(xs)?;
|
||
let xs = self.attn.forward(&xs, cos, sin, attention_mask)?;
|
||
let xs = (xs + residual)?;
|
||
let residual = &xs;
|
||
let xs = xs.apply(&self.norm2)?.apply(&self.mlp)?;
|
||
let xs = (residual + xs)?;
|
||
Ok(xs)
|
||
}
|
||
}
|
||
|
||
pub struct Qwen2_5VLVisionModel {
|
||
spatial_merge_size: usize,
|
||
patch_size: usize,
|
||
fullatt_block_indexes: Vec<usize>,
|
||
window_size: usize,
|
||
spatial_merge_unit: usize,
|
||
patch_embed: Qwen2_5VisionPatchEmbed,
|
||
rotary_pos_emb: Qwen2_5VisionRotaryEmbedding,
|
||
blocks: Vec<Qwen2_5VLVisionBlock>,
|
||
merger: Qwen2_5VLPatchMerger,
|
||
dtype: DType,
|
||
}
|
||
|
||
impl Qwen2_5VLVisionModel {
|
||
pub fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
|
||
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();
|
||
let window_size = cfg.vision_config.window_size;
|
||
let spatial_merge_unit = spatial_merge_size * spatial_merge_size;
|
||
let head_dim = cfg.vision_config.hidden_size / cfg.vision_config.num_heads;
|
||
let patch_embed = Qwen2_5VisionPatchEmbed::new(cfg, vb.pp("patch_embed"))?;
|
||
let rotary_pos_emb = Qwen2_5VisionRotaryEmbedding::new(head_dim / 2, None);
|
||
let mut blocks = Vec::new();
|
||
let vb_blocks = vb.pp("blocks");
|
||
for i in 0..cfg.vision_config.depth {
|
||
let block = Qwen2_5VLVisionBlock::new(cfg, vb_blocks.pp(i))?;
|
||
blocks.push(block);
|
||
}
|
||
let merger = Qwen2_5VLPatchMerger::new(cfg, vb.pp("merger"))?;
|
||
let dtype = vb.dtype();
|
||
Ok(Self {
|
||
spatial_merge_size,
|
||
patch_size,
|
||
fullatt_block_indexes,
|
||
window_size,
|
||
spatial_merge_unit,
|
||
patch_embed,
|
||
rotary_pos_emb,
|
||
blocks,
|
||
merger,
|
||
dtype,
|
||
})
|
||
}
|
||
|
||
pub fn rot_pos_emb(&self, grid_thw: &Tensor) -> Result<Tensor> {
|
||
let mut pos_ids = Vec::new();
|
||
for i in 0..grid_thw.dim(0)? {
|
||
let [t, h, w] = grid_thw.i(i)?.to_vec1::<u32>()?[..] else {
|
||
return Err(anyhow!(format!("grid_thw Expected exactly 3 elements")));
|
||
};
|
||
// hpos_ids shape (h, w)
|
||
let hpos_ids = Tensor::arange(0, h, grid_thw.device())?
|
||
.unsqueeze(1)?
|
||
.expand((h as usize, w as usize))?;
|
||
let hpos_ids = hpos_ids.reshape((
|
||
h as usize / self.spatial_merge_size,
|
||
self.spatial_merge_size,
|
||
w as usize / self.spatial_merge_size,
|
||
self.spatial_merge_size,
|
||
))?;
|
||
let hpos_ids = hpos_ids.permute((0, 2, 1, 3))?.flatten_all()?;
|
||
let wpos_ids = Tensor::arange(0, w, grid_thw.device())?
|
||
.unsqueeze(0)?
|
||
.expand((h as usize, w as usize))?;
|
||
let wpos_ids = wpos_ids.reshape((
|
||
h as usize / self.spatial_merge_size,
|
||
self.spatial_merge_size,
|
||
w as usize / self.spatial_merge_size,
|
||
self.spatial_merge_size,
|
||
))?;
|
||
let wpos_ids = wpos_ids.permute((0, 2, 1, 3))?.flatten_all()?;
|
||
// thw_pos_ids shape (h*w, 2)
|
||
let thw_pos_ids =
|
||
Tensor::stack(&[&hpos_ids, &wpos_ids], D::Minus1)?.repeat((t as usize, 1))?;
|
||
pos_ids.push(thw_pos_ids);
|
||
}
|
||
let pos_ids = Tensor::cat(&pos_ids, 0)?.contiguous()?;
|
||
let max_grid_size = grid_thw.i((.., 1..))?.max_all()?.to_scalar::<u32>()?;
|
||
let rotary_pos_emb_full = self
|
||
.rotary_pos_emb
|
||
.forward(max_grid_size as usize, grid_thw.device())?;
|
||
|
||
// contiguous()一定要加！！！很重要！！！！，不然index_select出来的是错的
|
||
// 找错找了半天，都是泪啊，做维度索引操作后contiguous顺手写上总没错
|
||
// 第一列是h维度的索引
|
||
let pos_ids_h = pos_ids.i((.., 0))?.contiguous()?;
|
||
// 第二列是w维度的索引
|
||
let pos_ids_w = pos_ids.i((.., 1))?.contiguous()?;
|
||
let rotary_pos_emb_h = rotary_pos_emb_full.index_select(&pos_ids_h, 0)?;
|
||
let rotary_pos_emb_w = rotary_pos_emb_full.index_select(&pos_ids_w, 0)?;
|
||
// 每个patch融合h索引和w索引两个的位置编码信息
|
||
let rotary_pos_emb = Tensor::cat(&[rotary_pos_emb_h, rotary_pos_emb_w], 1)?.contiguous()?;
|
||
Ok(rotary_pos_emb)
|
||
}
|
||
|
||
pub fn get_window_index(&self, grid_thw: &Tensor) -> Result<(Tensor, Tensor)> {
|
||
let mut window_index = Vec::new();
|
||
let mut cu_window_seqlens = vec![0];
|
||
let mut window_index_id = 0_i64;
|
||
|
||
let vit_merger_window_size =
|
||
(self.window_size / self.spatial_merge_size / self.patch_size) as u32;
|
||
for i in 0..grid_thw.dim(0)? {
|
||
let [grid_t, grid_h, grid_w] = grid_thw.i(i)?.to_vec1::<u32>()?[..] else {
|
||
return Err(anyhow!(format!("grid_thw Expected exactly 3 elements")));
|
||
};
|
||
let llm_grid_h = grid_h / self.spatial_merge_size as u32;
|
||
let llm_grid_w = grid_w / self.spatial_merge_size as u32;
|
||
// 因为后续需要使用-100来做填充，所以需要int类型
|
||
// candle好像不支持i32， DType里面都没有定义i32, 所以这里使用i64
|
||
let mut index = Tensor::arange(
|
||
window_index_id,
|
||
window_index_id + (grid_t * llm_grid_h * llm_grid_w) as i64,
|
||
grid_thw.device(),
|
||
)?
|
||
.reshape((grid_t as usize, llm_grid_h as usize, llm_grid_w as usize))?
|
||
.contiguous()?;
|
||
// python transformers 中实现如下
|
||
// let pad_h = (vit_merger_window_size - llm_grid_h % vit_merger_window_size);
|
||
// 后面加上 % vit_merger_window_size，保证llm_grid_h能整除vit_merger_window_size时不需要pad
|
||
// 按理说能整除应该是不需要pad的，transformers中这样实现不知道是不是有什么其他原因
|
||
let pad_h = (vit_merger_window_size - llm_grid_h % vit_merger_window_size)
|
||
% vit_merger_window_size;
|
||
let pad_w = (vit_merger_window_size - llm_grid_w % vit_merger_window_size)
|
||
% vit_merger_window_size;
|
||
let num_window_h = (llm_grid_h + pad_h) / vit_merger_window_size;
|
||
let num_window_w = (llm_grid_w + pad_w) / vit_merger_window_size;
|
||
if pad_h > 0 {
|
||
let pad_h_t = Tensor::new(vec![-100_i64], grid_thw.device())?
|
||
.broadcast_as((grid_t as usize, pad_h as usize, llm_grid_w as usize))?
|
||
.contiguous()?;
|
||
index = Tensor::cat(&[&index, &pad_h_t], 1)?;
|
||
}
|
||
if pad_w > 0 {
|
||
let pad_w_t = Tensor::new(vec![-100_i64], grid_thw.device())?
|
||
.broadcast_as((
|
||
grid_t as usize,
|
||
(llm_grid_h + pad_h) as usize,
|
||
pad_w as usize,
|
||
))?
|
||
.contiguous()?;
|
||
index = Tensor::cat(&[&index, &pad_w_t], 2)?;
|
||
}
|
||
let index_padded = index
|
||
.reshape((
|
||
grid_t as usize,
|
||
num_window_h as usize,
|
||
vit_merger_window_size as usize,
|
||
num_window_w as usize,
|
||
vit_merger_window_size as usize,
|
||
))?
|
||
.permute((0, 1, 3, 2, 4))?;
|
||
let index_padded = index_padded
|
||
.reshape((
|
||
grid_t as usize,
|
||
(num_window_h * num_window_w) as usize,
|
||
vit_merger_window_size as usize,
|
||
vit_merger_window_size as usize,
|
||
))?
|
||
.contiguous()?;
|
||
let is_pad = Tensor::new(vec![-100_i64], grid_thw.device())?;
|
||
let seqlens = index_padded
|
||
.broadcast_ne(&is_pad)?
|
||
.sum((2, 3))?
|
||
.flatten_all()?;
|
||
let index_padded = index_padded.flatten_all()?;
|
||
let not_pad = index_padded.broadcast_ne(&is_pad)?.to_vec1::<u8>()?;
|
||
let indices: Vec<u32> = not_pad
|
||
.iter()
|
||
.enumerate()
|
||
.filter_map(|(idx, &val)| if val != 0 { Some(idx as u32) } else { None })
|
||
.collect();
|
||
let indices_tensor = Tensor::from_slice(&indices, indices.len(), grid_thw.device())?;
|
||
let index_new = index_padded.gather(&indices_tensor, 0)?;
|
||
window_index.push(index_new);
|
||
|
||
let seq_len_last = cu_window_seqlens[cu_window_seqlens.len() - 1];
|
||
// cumsum方法i64类型执行会报错，先转成F64计算后再转回i64
|
||
let cu_seqlens_tmp = seqlens
|
||
.to_dtype(candle_core::DType::F64)?
|
||
.cumsum(0)?
|
||
.to_dtype(candle_core::DType::I64)?
|
||
.broadcast_mul(&Tensor::new(
|
||
vec![self.spatial_merge_unit as i64],
|
||
grid_thw.device(),
|
||
)?)?
|
||
.broadcast_add(&Tensor::new(vec![seq_len_last], grid_thw.device())?)?;
|
||
cu_window_seqlens.extend_from_slice(&cu_seqlens_tmp.to_vec1::<i64>()?);
|
||
window_index_id += (grid_t * llm_grid_h * llm_grid_w) as i64;
|
||
}
|
||
let window_index_tensor = Tensor::cat(&window_index, 0)?;
|
||
let cu_window_seqlens_tensor = Tensor::from_slice(
|
||
&cu_window_seqlens,
|
||
cu_window_seqlens.len(),
|
||
grid_thw.device(),
|
||
)?
|
||
.to_dtype(candle_core::DType::U32)?;
|
||
Ok((window_index_tensor, cu_window_seqlens_tensor))
|
||
}
|
||
|
||
pub fn get_attention_mask(
|
||
&self,
|
||
cu_seqlens: &Tensor,
|
||
seq_len: usize,
|
||
device: &Device,
|
||
dtype: DType,
|
||
) -> Result<Tensor> {
|
||
let mut attention_mask =
|
||
Tensor::new(f32::NEG_INFINITY, device)?.broadcast_as((1, seq_len, seq_len))?;
|
||
for i in 1..cu_seqlens.dim(0)? {
|
||
let start = cu_seqlens.i(i - 1)?.to_scalar::<u32>()? as usize;
|
||
let end = cu_seqlens.i(i)?.to_scalar::<u32>()? as usize;
|
||
let block_size = end - start;
|
||
let zeros =
|
||
Tensor::zeros((1, block_size, block_size), candle_core::DType::F32, device)?;
|
||
attention_mask =
|
||
attention_mask.slice_assign(&[(0..1), (start..end), (start..end)], &zeros)?;
|
||
}
|
||
let attention_mask = attention_mask.to_dtype(dtype)?.contiguous()?;
|
||
Ok(attention_mask)
|
||
}
|
||
|
||
pub fn forward(&self, hidden_states: &Tensor, grid_thw: &Tensor) -> Result<Tensor> {
|
||
// hidden_states: (seq_len, hidden_size)
|
||
// grid_thw: (num_images_or_videos, 3), temporal, height, width
|
||
let hidden_states = hidden_states.to_dtype(self.dtype)?;
|
||
let hidden_states = self.patch_embed.forward(&hidden_states)?;
|
||
let rotary_pos_emb = self.rot_pos_emb(grid_thw)?;
|
||
let (window_index, cu_window_seqlens) = self.get_window_index(grid_thw)?;
|
||
let seq_len = hidden_states.dim(0)?;
|
||
let hidden_states = hidden_states
|
||
.reshape((
|
||
seq_len / self.spatial_merge_unit,
|
||
self.spatial_merge_unit,
|
||
(),
|
||
))?
|
||
.contiguous()?;
|
||
let hidden_states = hidden_states.index_select(&window_index, 0)?;
|
||
let mut hidden_states = hidden_states.reshape((seq_len, ()))?;
|
||
let rotary_pos_emb = rotary_pos_emb.reshape((
|
||
seq_len / self.spatial_merge_unit,
|
||
self.spatial_merge_unit,
|
||
(),
|
||
))?;
|
||
let rotary_pos_emb = rotary_pos_emb.index_select(&window_index, 0)?;
|
||
let rotary_pos_emb = rotary_pos_emb.reshape((seq_len, ()))?;
|
||
let emb = Tensor::cat(&[&rotary_pos_emb, &rotary_pos_emb], D::Minus1)?;
|
||
let cos = emb.cos()?.to_dtype(hidden_states.dtype())?;
|
||
let sin = emb.sin()?.to_dtype(hidden_states.dtype())?;
|
||
let cu_seqlens = grid_thw.i((.., 1))?.mul(&grid_thw.i((.., 2))?)?;
|
||
let grid_t = grid_thw.i((.., 0))?.to_vec1::<u32>()?;
|
||
let cu_seqlens_full = match cu_seqlens.rank() {
|
||
1 => cu_seqlens.repeat(grid_t[0] as usize)?,
|
||
2 => {
|
||
let mut cu_seqlens_repeat = Vec::new();
|
||
for (index, t) in grid_t.iter().enumerate() {
|
||
cu_seqlens_repeat.push(cu_seqlens.i(index)?.repeat(t.clone() as usize)?);
|
||
}
|
||
Tensor::cat(&cu_seqlens_repeat, 0)?.flatten_all()?
|
||
}
|
||
_ => {
|
||
return Err(anyhow!(format!("create cu_seqlens error")));
|
||
}
|
||
};
|
||
let cu_seqlens = cu_seqlens_full
|
||
.to_dtype(DType::F64)?
|
||
.cumsum(0)?
|
||
.to_dtype(DType::U32)?;
|
||
let pad_zero = Tensor::from_vec(vec![0_u32], 1, hidden_states.device())?;
|
||
let cu_seqlens = Tensor::cat(&[&pad_zero, &cu_seqlens], D::Minus1)?;
|
||
let attention_mask_window = self.get_attention_mask(
|
||
&cu_window_seqlens,
|
||
seq_len,
|
||
hidden_states.device(),
|
||
hidden_states.dtype(),
|
||
)?;
|
||
let attention_mask_full = self.get_attention_mask(
|
||
&cu_seqlens,
|
||
seq_len,
|
||
hidden_states.device(),
|
||
hidden_states.dtype(),
|
||
)?;
|
||
let mut attention_mask = attention_mask_window.clone();
|
||
for (layer_num, block) in self.blocks.iter().enumerate() {
|
||
if self.fullatt_block_indexes.contains(&layer_num) {
|
||
attention_mask = attention_mask_full.clone();
|
||
} else {
|
||
attention_mask = attention_mask_window.clone();
|
||
}
|
||
hidden_states = block.forward(&hidden_states, &cos, &sin, &attention_mask)?;
|
||
}
|
||
let hidden_states = self.merger.forward(&hidden_states)?;
|
||
let reverse_indices = safe_arg_sort_last_dim(&window_index, true)?;
|
||
let hidden_states = hidden_states.index_select(&reverse_indices, 0)?;
|
||
Ok(hidden_states)
|
||
}
|
||
}
|
||
|
||
|
||
#[derive(Debug, Clone)]
|
||
struct Qwen2_5VLTextMLP {
|
||
gate_proj: Linear,
|
||
up_proj: Linear,
|
||
down_proj: Linear,
|
||
act_fn: Activation,
|
||
}
|
||
|
||
impl Qwen2_5VLTextMLP {
|
||
fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
|
||
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"))?;
|
||
let up_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("up_proj"))?;
|
||
let down_proj = linear_no_bias(intermediate_sz, hidden_sz, vb.pp("down_proj"))?;
|
||
|
||
Ok(Self {
|
||
gate_proj,
|
||
up_proj,
|
||
down_proj,
|
||
act_fn: cfg.hidden_act,
|
||
})
|
||
}
|
||
}
|
||
|
||
impl Module for Qwen2_5VLTextMLP {
|
||
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)]
|
||
struct Qwen2_5VLTextAttention {
|
||
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 Qwen2_5VLTextAttention {
|
||
fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
|
||
let hidden_size = cfg.hidden_size;
|
||
let num_heads = cfg.num_attention_heads;
|
||
let num_kv_heads = cfg.num_key_value_heads;
|
||
let num_kv_groups = num_heads / num_kv_heads;
|
||
let head_dim = hidden_size / num_heads;
|
||
let q_proj = linear(hidden_size, num_heads * head_dim, vb.pp("q_proj"))?;
|
||
let k_proj = linear(hidden_size, num_kv_heads * head_dim, vb.pp("k_proj"))?;
|
||
let v_proj = linear(hidden_size, num_kv_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_kv_heads,
|
||
num_kv_groups,
|
||
head_dim,
|
||
hidden_size,
|
||
kv_cache: None,
|
||
})
|
||
}
|
||
|
||
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)?;
|
||
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)
|
||
}
|
||
|
||
fn clear_kv_cache(&mut self) {
|
||
self.kv_cache = None
|
||
}
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
struct Qwen2_5VLTextDecoderLayer {
|
||
self_attn: Qwen2_5VLTextAttention,
|
||
mlp: Qwen2_5VLTextMLP,
|
||
input_layernorm: RmsNorm,
|
||
post_attention_layernorm: RmsNorm,
|
||
}
|
||
|
||
impl Qwen2_5VLTextDecoderLayer {
|
||
fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
|
||
let self_attn = Qwen2_5VLTextAttention::new(cfg, vb.pp("self_attn"))?;
|
||
let mlp = Qwen2_5VLTextMLP::new(cfg, vb.pp("mlp"))?;
|
||
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,
|
||
})
|
||
}
|
||
|
||
fn forward(
|
||
&mut self,
|
||
xs: &Tensor,
|
||
cos: &Tensor,
|
||
sin: &Tensor,
|
||
attention_mask: Option<&Tensor>,
|
||
) -> Result<Tensor> {
|
||
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)
|
||
}
|
||
|
||
fn clear_kv_cache(&mut self) {
|
||
self.self_attn.clear_kv_cache()
|
||
}
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
pub struct Qwen2_5VLTextModel {
|
||
pub embed_tokens: candle_nn::Embedding,
|
||
layers: Vec<Qwen2_5VLTextDecoderLayer>,
|
||
norm: RmsNorm,
|
||
rotary_emb: Qwen2_5VLTextRotaryEmbedding,
|
||
dtype: DType,
|
||
sliding_window: usize,
|
||
device: Device,
|
||
rope_scaling: RopeScaling,
|
||
}
|
||
|
||
impl Qwen2_5VLTextModel {
|
||
pub fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
|
||
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;
|
||
let rotary_emb = Qwen2_5VLTextRotaryEmbedding::new(head_dim, cfg.rope_theta);
|
||
let mut layers = Vec::with_capacity(cfg.num_hidden_layers);
|
||
let vb_l = vb.pp("layers");
|
||
for layer_idx in 0..cfg.num_hidden_layers {
|
||
let layer = Qwen2_5VLTextDecoderLayer::new(cfg, vb_l.pp(layer_idx))?;
|
||
layers.push(layer)
|
||
}
|
||
let norm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("norm"))?;
|
||
let sliding_window = cfg.sliding_window;
|
||
let rope_scaling = cfg.rope_scaling.clone();
|
||
Ok(Self {
|
||
embed_tokens,
|
||
layers,
|
||
norm,
|
||
rotary_emb,
|
||
dtype: vb.dtype(),
|
||
sliding_window,
|
||
device: vb.device().clone(),
|
||
rope_scaling,
|
||
})
|
||
}
|
||
|
||
fn prepare_causal_attention_mask(
|
||
&self,
|
||
b_size: usize,
|
||
tgt_len: usize,
|
||
seqlen_offset: usize,
|
||
) -> Result<Tensor> {
|
||
// Sliding window mask?
|
||
let mask: Vec<_> = (0..tgt_len)
|
||
.flat_map(|i| {
|
||
(0..tgt_len).map(move |j| {
|
||
if i < j || j + self.sliding_window < i {
|
||
f32::NEG_INFINITY
|
||
} else {
|
||
0.
|
||
}
|
||
})
|
||
})
|
||
.collect();
|
||
let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?;
|
||
let mask = if seqlen_offset > 0 {
|
||
let mask0 = Tensor::zeros((tgt_len, seqlen_offset), self.dtype, &self.device)?;
|
||
Tensor::cat(&[&mask0, &mask], D::Minus1)?
|
||
} else {
|
||
mask
|
||
};
|
||
let mask = mask
|
||
.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))?
|
||
.to_dtype(self.dtype)?;
|
||
Ok(mask)
|
||
}
|
||
pub fn forward(
|
||
&mut self,
|
||
inputs_embeds: &Tensor,
|
||
seqlen_offset: usize,
|
||
position_ids: Option<&Tensor>,
|
||
) -> Result<Tensor> {
|
||
let (b_size, seq_len, _) = inputs_embeds.dims3()?;
|
||
|
||
let position_ids = match position_ids {
|
||
Some(ids) => ids.clone(),
|
||
None => Tensor::arange(
|
||
seqlen_offset as u32,
|
||
(seq_len + seqlen_offset) as u32,
|
||
inputs_embeds.device(),
|
||
)?
|
||
.unsqueeze(0)?
|
||
.unsqueeze(0)?
|
||
.broadcast_as((3, b_size, seq_len))?,
|
||
};
|
||
let (cos, sin) = self.rotary_emb.forward(
|
||
&position_ids,
|
||
self.dtype,
|
||
self.rope_scaling.mrope_section.clone(),
|
||
)?;
|
||
let mut xs = inputs_embeds.clone();
|
||
let attention_mask: Option<&Tensor> = {
|
||
if seq_len <= 1 {
|
||
None
|
||
} else {
|
||
Some(&self.prepare_causal_attention_mask(b_size, seq_len, seqlen_offset)?)
|
||
}
|
||
};
|
||
for layer in self.layers.iter_mut() {
|
||
xs = layer.forward(&xs, &cos, &sin, attention_mask)?;
|
||
}
|
||
let xs = xs.apply(&self.norm)?;
|
||
Ok(xs)
|
||
}
|
||
|
||
pub fn clear_kv_cache(&mut self) {
|
||
for layer in self.layers.iter_mut() {
|
||
layer.clear_kv_cache()
|
||
}
|
||
}
|
||
}
|
||
|
||
pub struct Qwen2_5VLModel {
|
||
visual: Qwen2_5VLVisionModel,
|
||
model: Qwen2_5VLTextModel,
|
||
pub cfg: Qwen2_5VLConfig,
|
||
lm_head: Linear,
|
||
rope_deltas: Option<Tensor>,
|
||
}
|
||
|
||
impl Qwen2_5VLModel {
|
||
pub fn new(cfg: Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
|
||
let visual = Qwen2_5VLVisionModel::new(&cfg, vb.pp("visual"))?;
|
||
let model = Qwen2_5VLTextModel::new(&cfg, vb.pp("model"))?;
|
||
let vocab_size = cfg.vocab_size;
|
||
let lm_head = if cfg.tie_word_embeddings {
|
||
Linear::new(model.embed_tokens.embeddings().clone(), None)
|
||
} else {
|
||
linear_no_bias(cfg.hidden_size, vocab_size, vb.pp("lm_head"))?
|
||
};
|
||
|
||
Ok(Self {
|
||
visual,
|
||
model,
|
||
cfg,
|
||
lm_head,
|
||
rope_deltas: None,
|
||
})
|
||
}
|
||
|
||
pub fn get_rope_index(
|
||
&self,
|
||
input_ids: &Tensor,
|
||
image_grid_thw: Option<&Tensor>,
|
||
video_grid_thw: Option<&Tensor>,
|
||
mask: Option<&Tensor>,
|
||
second_per_grid_ts: Option<Vec<f32>>,
|
||
) -> Result<(Tensor, Tensor)> {
|
||
let spatial_merge_size = self.cfg.vision_config.spatial_merge_size;
|
||
let mut mrope_position_deltas: Vec<i64> = Vec::new();
|
||
if image_grid_thw.is_some() || video_grid_thw.is_some() {
|
||
let total_input_ids = input_ids.clone();
|
||
let mut mask_;
|
||
if mask.is_none() {
|
||
mask_ = Tensor::ones_like(&total_input_ids)?;
|
||
} else {
|
||
mask_ = mask.unwrap().clone();
|
||
}
|
||
let mut position_ids = Tensor::ones(
|
||
(3, input_ids.dim(0)?, input_ids.dim(1)?),
|
||
input_ids.dtype(),
|
||
input_ids.device(),
|
||
)?;
|
||
let mut image_index = 0;
|
||
let mut video_index = 0;
|
||
for i in 0..total_input_ids.dim(0)? {
|
||
let mut input_ids_i = total_input_ids.i(i)?;
|
||
let mask_i = mask_.i(i)?;
|
||
// 推理时, attention_mask如果是全1向量,取非0索引的操作没必要
|
||
if mask_i.sum_all()?.to_scalar::<u32>()? != mask_i.dim(0)? as u32 {
|
||
let nonzero_idx = nonzero_index(&mask_i)?;
|
||
input_ids_i = input_ids_i.gather(&nonzero_idx, 0)?;
|
||
}
|
||
let mut text_start = 0;
|
||
let mut text_end = 0;
|
||
let mut thw = vec![];
|
||
let mut second_per_grid_t = 0_f32;
|
||
let mut llm_pos_ids_list: Vec<Tensor> = Vec::new();
|
||
// vision start的下一个索引
|
||
let vision_indices =
|
||
get_vision_next_indices(&input_ids_i, self.cfg.vision_start_token_id as u32);
|
||
match vision_indices {
|
||
Ok(indeices) => {
|
||
let vision_tokens = input_ids_i.gather(&indeices, 0)?.to_vec1::<u32>()?;
|
||
let vision_indices_vec = indeices.to_vec1::<u32>()?;
|
||
for (j, &token) in vision_tokens.iter().enumerate() {
|
||
if token == self.cfg.image_token_id as u32 {
|
||
thw = image_grid_thw.unwrap().i(image_index)?.to_vec1::<u32>()?;
|
||
image_index += 1;
|
||
text_end = vision_indices_vec[j];
|
||
second_per_grid_t = 0.0;
|
||
}
|
||
if token == self.cfg.video_token_id as u32 {
|
||
thw = video_grid_thw.unwrap().i(video_index)?.to_vec1::<u32>()?;
|
||
text_end = vision_indices_vec[j];
|
||
second_per_grid_t = match second_per_grid_ts {
|
||
None => 1.0,
|
||
Some(ref vec) => vec[video_index],
|
||
};
|
||
video_index += 1;
|
||
}
|
||
let llm_grid_t = thw[0];
|
||
let llm_grid_h = thw[1] / spatial_merge_size as u32;
|
||
let llm_grid_w = thw[2] / spatial_merge_size as u32;
|
||
let text_len = text_end - text_start;
|
||
let start_idx = if llm_pos_ids_list.len() > 0 {
|
||
llm_pos_ids_list[llm_pos_ids_list.len() - 1]
|
||
.max_all()?
|
||
.to_scalar::<u32>()?
|
||
+ 1
|
||
} else {
|
||
0
|
||
};
|
||
let pos_ids = Tensor::arange(
|
||
start_idx,
|
||
start_idx + text_len,
|
||
input_ids_i.device(),
|
||
)?
|
||
.unsqueeze(0)?
|
||
.broadcast_as((3usize, text_len as usize))?;
|
||
llm_pos_ids_list.push(pos_ids);
|
||
let range_tensor = Tensor::arange(0, llm_grid_t, input_ids_i.device())?
|
||
.unsqueeze(D::Minus1)?;
|
||
let expanded_range = range_tensor.broadcast_as((
|
||
llm_grid_t as usize,
|
||
(llm_grid_h * llm_grid_w) as usize,
|
||
))?;
|
||
let time_tensor = expanded_range
|
||
.broadcast_mul(&Tensor::new(
|
||
vec![
|
||
(second_per_grid_t
|
||
* self.cfg.vision_config.tokens_per_second as f32)
|
||
as u32,
|
||
],
|
||
input_ids_i.device(),
|
||
)?)?
|
||
.broadcast_add(&Tensor::new(
|
||
vec![start_idx + text_len],
|
||
input_ids_i.device(),
|
||
)?)?;
|
||
let t_index = time_tensor.flatten_all()?;
|
||
let h_index = Tensor::arange(
|
||
start_idx + text_len,
|
||
start_idx + text_len + llm_grid_h,
|
||
input_ids_i.device(),
|
||
)?
|
||
.unsqueeze(0)?
|
||
.unsqueeze(D::Minus1)?
|
||
.broadcast_as((
|
||
llm_grid_t as usize,
|
||
llm_grid_h as usize,
|
||
llm_grid_w as usize,
|
||
))?
|
||
.flatten_all()?;
|
||
let w_index = Tensor::arange(
|
||
start_idx + text_len,
|
||
start_idx + text_len + llm_grid_w,
|
||
input_ids_i.device(),
|
||
)?
|
||
.unsqueeze(0)?
|
||
.unsqueeze(0)?
|
||
.broadcast_as((
|
||
llm_grid_t as usize,
|
||
llm_grid_h as usize,
|
||
llm_grid_w as usize,
|
||
))?
|
||
.flatten_all()?;
|
||
|
||
let thw_index = Tensor::stack(&[t_index, h_index, w_index], 0)?;
|
||
llm_pos_ids_list.push(thw_index);
|
||
text_start = text_end + llm_grid_t * llm_grid_h * llm_grid_w;
|
||
}
|
||
}
|
||
Err(e) => {
|
||
println!("get vision_indices err: {}", e);
|
||
}
|
||
};
|
||
|
||
if text_start < input_ids_i.dim(0)? as u32 {
|
||
let start_idx = if llm_pos_ids_list.len() > 0 {
|
||
llm_pos_ids_list[llm_pos_ids_list.len() - 1]
|
||
.max_all()?
|
||
.to_scalar::<u32>()?
|
||
+ 1
|
||
} else {
|
||
0
|
||
};
|
||
let text_len = input_ids_i.dim(0)? as u32 - text_start;
|
||
let pos_ids =
|
||
Tensor::arange(start_idx, start_idx + text_len, input_ids_i.device())?
|
||
.unsqueeze(0)?
|
||
.broadcast_as((3usize, text_len as usize))?;
|
||
llm_pos_ids_list.push(pos_ids);
|
||
}
|
||
let llm_position = Tensor::cat(&llm_pos_ids_list, 1)?.reshape((3, 1, ()))?;
|
||
position_ids = position_ids
|
||
.slice_assign(&[(0..3), (i..i + 1), (0..input_ids.dim(1)?)], &llm_position)?;
|
||
let position_deltas = llm_position.max_all()?.to_scalar::<u32>()? as i64 + 1
|
||
- input_ids_i.dim(0)? as i64;
|
||
mrope_position_deltas.push(position_deltas);
|
||
}
|
||
|
||
let mut mrope_position_deltas = Tensor::new(mrope_position_deltas, input_ids.device())?;
|
||
if mrope_position_deltas.rank() == 1 {
|
||
mrope_position_deltas = mrope_position_deltas.unsqueeze(0)?;
|
||
}
|
||
return Ok((position_ids.contiguous()?, mrope_position_deltas));
|
||
} else {
|
||
if mask.is_some() {
|
||
let mut position_ids = mask
|
||
.unwrap()
|
||
.to_dtype(candle_core::DType::F64)?
|
||
.cumsum(D::Minus1)?
|
||
.to_dtype(candle_core::DType::U32)?
|
||
.broadcast_sub(&Tensor::new(vec![1_u32], input_ids.device())?)?;
|
||
for i in 0..position_ids.dim(0)? {
|
||
let mut position_ids_i = position_ids.i(i)?;
|
||
let mask_i = mask.unwrap().i(i)?;
|
||
// 如果有pad, 将填充位置置为1
|
||
// 当bs>1, 可能存在不同序列长度，需要添加pad使seq_len长度一致
|
||
if mask_i.sum_all()?.to_scalar::<u32>()? != mask_i.dim(0)? as u32 {
|
||
let zero_indices = zero_index(&mask_i)?;
|
||
let replace_1 = Tensor::ones(
|
||
zero_indices.dim(0)?,
|
||
candle_core::DType::U32,
|
||
input_ids.device(),
|
||
)?;
|
||
position_ids_i = position_ids_i
|
||
.scatter(&zero_indices, &replace_1, 0)?
|
||
.unsqueeze(0)?;
|
||
position_ids = position_ids.slice_assign(
|
||
&[(i..i + 1), (0..position_ids.dim(1)?)],
|
||
&position_ids_i,
|
||
)?;
|
||
}
|
||
}
|
||
position_ids = position_ids
|
||
.unsqueeze(0)?
|
||
.broadcast_as((3, input_ids.dim(0)?, input_ids.dim(1)?))?
|
||
.contiguous()?;
|
||
let mut mrope_position_deltas = position_ids
|
||
.max(0)?
|
||
.max(D::Minus1)?
|
||
.broadcast_sub(&Tensor::new(
|
||
vec![mask.unwrap().dim(D::Minus1)? as u32 - 1],
|
||
input_ids.device(),
|
||
)?)?
|
||
.contiguous()?;
|
||
if mrope_position_deltas.rank() == 1 {
|
||
mrope_position_deltas = mrope_position_deltas.unsqueeze(0)?;
|
||
}
|
||
return Ok((position_ids, mrope_position_deltas));
|
||
} else {
|
||
let position_ids =
|
||
Tensor::arange(0_u32, input_ids.dim(D::Minus1)? as u32, input_ids.device())?
|
||
.unsqueeze(0)?
|
||
.unsqueeze(0)?
|
||
.broadcast_as((3, input_ids.dim(0)?, input_ids.dim(D::Minus1)?))?
|
||
.contiguous()?;
|
||
let mrope_position_deltas = Tensor::zeros(
|
||
(input_ids.dim(0)?, 1),
|
||
input_ids.dtype(),
|
||
input_ids.device(),
|
||
)?;
|
||
Ok((position_ids, mrope_position_deltas))
|
||
}
|
||
}
|
||
}
|
||
|
||
pub fn forward(
|
||
&mut self,
|
||
input_ids: &Tensor,
|
||
pixel_values: Option<&Tensor>,
|
||
image_grid_thw: Option<&Tensor>,
|
||
pixel_values_video: Option<&Tensor>,
|
||
video_grid_thw: Option<&Tensor>,
|
||
mask: &Tensor,
|
||
cache_position: Option<&Tensor>,
|
||
seqlen_offset: usize,
|
||
second_per_grid_ts: Option<Vec<f32>>,
|
||
) -> Result<Tensor> {
|
||
// input_ids shape: (bs, seq_len)
|
||
let mut inputs_embeds = self.model.embed_tokens.forward(&input_ids)?;
|
||
// inputs_embeds shape: (bs, seq_len, hidden_dim)
|
||
if pixel_values.is_some() && image_grid_thw.is_some() {
|
||
// image_embed shape: (seq_len, hidden_dim)
|
||
let image_embed = self
|
||
.visual
|
||
.forward(pixel_values.unwrap(), image_grid_thw.unwrap())?;
|
||
let vision_mask = get_equal_mask(&input_ids, self.cfg.image_token_id as u32)?;
|
||
|
||
let n_image_tokens = vision_mask.sum_all()?.to_scalar::<u32>()?;
|
||
if n_image_tokens as usize != image_embed.dim(0)? {
|
||
return Err(anyhow!(format!(
|
||
"n_image_token num: {} not equal to image_embed len: {}",
|
||
n_image_tokens,
|
||
image_embed.dim(0)?
|
||
)));
|
||
}
|
||
inputs_embeds = masked_scatter_dim0(&inputs_embeds, &image_embed, &vision_mask)?;
|
||
}
|
||
if pixel_values_video.is_some() && video_grid_thw.is_some() {
|
||
let video_embed = self
|
||
.visual
|
||
.forward(pixel_values_video.unwrap(), video_grid_thw.unwrap())?;
|
||
|
||
let vision_mask = get_equal_mask(&input_ids, self.cfg.video_token_id as u32)?;
|
||
let n_video_tokens = vision_mask.sum_all()?.to_scalar::<u32>()?;
|
||
if n_video_tokens as usize != video_embed.dim(0)? {
|
||
return Err(anyhow!(format!(
|
||
"n_image_token num: {} not equal to image_embed len: {}",
|
||
n_video_tokens,
|
||
video_embed.dim(0)?
|
||
)));
|
||
}
|
||
inputs_embeds = masked_scatter_dim0(&inputs_embeds, &video_embed, &vision_mask)?;
|
||
}
|
||
let mut position_ids;
|
||
let mut rope_deltas;
|
||
if (cache_position.is_some() && cache_position.unwrap().i(0)?.to_scalar::<u32>()? == 0)
|
||
|| self.rope_deltas.is_none()
|
||
{
|
||
(position_ids, rope_deltas) = self.get_rope_index(
|
||
input_ids,
|
||
image_grid_thw,
|
||
video_grid_thw,
|
||
Some(mask),
|
||
second_per_grid_ts,
|
||
)?;
|
||
self.rope_deltas = Some(rope_deltas);
|
||
} else {
|
||
let (bs, seq_len, _) = inputs_embeds.dims3()?;
|
||
let delta = if cache_position.is_some() {
|
||
cache_position
|
||
.unwrap()
|
||
.i(0)?
|
||
.to_dtype(self.rope_deltas.as_ref().unwrap().dtype())?
|
||
.broadcast_add(&self.rope_deltas.as_ref().unwrap())?
|
||
.contiguous()?
|
||
.to_dtype(candle_core::DType::U32)?
|
||
} else {
|
||
Tensor::zeros(1, inputs_embeds.dtype(), inputs_embeds.device())?
|
||
};
|
||
position_ids = Tensor::arange(0u32, seq_len as u32, input_ids.device())?
|
||
.unsqueeze(0)?
|
||
.broadcast_as((bs, seq_len))?
|
||
.broadcast_add(&delta)?
|
||
.unsqueeze(0)?
|
||
.broadcast_as((3, bs, seq_len))?
|
||
.contiguous()?;
|
||
}
|
||
let outputs = self
|
||
.model
|
||
.forward(&inputs_embeds, seqlen_offset, Some(&position_ids))?;
|
||
let seq_len = outputs.dim(1)?;
|
||
let hidden_state = outputs.narrow(1, seq_len - 1, 1)?;
|
||
let logits = self.lm_head.forward(&hidden_state)?;
|
||
Ok(logits)
|
||
}
|
||
pub fn clear_kv_cache(&mut self) {
|
||
self.model.clear_kv_cache();
|
||
}
|
||
}
|