1use crate::model_config;
6use super::traits::*;
7use anyhow::Result;
8use serde::{Serialize, Deserialize};
9
10model_config!(InternVLConfig {
11 vocab_size: usize = 92544,
12 hidden_size: usize = 4096,
13 intermediate_size: usize = 14336,
14 num_hidden_layers: usize = 32,
15 num_attention_heads: usize = 32,
16 num_key_value_heads: usize = 8,
17 max_position_embeddings: usize = 8192,
18 rms_norm_eps: f32 = 1e-5,
19 rope_theta: f32 = 10000.0,
20 vision_hidden_size: usize = 3200,
21 vision_intermediate_size: usize = 12800,
22 vision_num_hidden_layers: usize = 48,
23 vision_num_attention_heads: usize = 25,
24 vision_patch_size: usize = 14,
25 vision_image_size: usize = 448,
26 pad_token_id: i64 = 0,
27 bos_token_id: i64 = 1,
28 eos_token_id: i64 = 2,
29});
30
31impl InternVLConfig {
32 pub fn from_gguf_config(gguf: &crate::weight_loader_core::GGUFModelConfig) -> Self {
33 Self {
34 vocab_size: gguf.vocab_size,
35 hidden_size: gguf.hidden_size,
36 intermediate_size: gguf.intermediate_size,
37 num_hidden_layers: gguf.num_hidden_layers,
38 num_attention_heads: gguf.num_attention_heads,
39 num_key_value_heads: gguf.num_key_value_heads,
40 rms_norm_eps: gguf.rms_norm_eps,
41 rope_theta: gguf.rope_theta,
42 ..Default::default()
43 }
44 }
45}
46
47pub struct InternVLModelV2 {
48 config: InternVLConfig,
49 device: Device,
50 vision_encoder: InternViTEncoder,
51 mlp_projector: Tensor,
52 embed_tokens: Tensor,
53 layers: Vec<InternLMLayer>,
54 norm: Tensor,
55 lm_head: Tensor,
56}
57
58pub struct InternViTEncoder {
59 patch_embed: Tensor,
60 blocks: Vec<InternViTBlock>,
61 norm: Tensor,
62 hidden_size: usize,
63}
64
65pub struct InternViTBlock {
66 norm1: Tensor,
67 attn_qkv: Tensor,
68 attn_proj: Tensor,
69 norm2: Tensor,
70 mlp_fc1: Tensor,
71 mlp_fc2: Tensor,
72 num_heads: usize,
73}
74
75pub struct InternLMLayer {
76 self_attn: InternLMAttention,
77 mlp: InternLMMLP,
78 input_layernorm: Tensor,
79 post_attention_layernorm: Tensor,
80}
81
82pub struct InternLMAttention {
83 q_proj: Tensor,
84 k_proj: Tensor,
85 v_proj: Tensor,
86 o_proj: Tensor,
87 num_heads: usize,
88 num_kv_heads: usize,
89 head_dim: usize,
90}
91
92pub struct InternLMMLP {
93 gate_proj: Tensor,
94 up_proj: Tensor,
95 down_proj: Tensor,
96}
97
98impl Model for InternVLModelV2 {
99 type Config = InternVLConfig;
100
101 fn new(config: InternVLConfig) -> Result<Self> {
102 let device = Device::CPU;
103
104 let vision_encoder = InternViTEncoder::new(&config, &device)?;
105 let mlp_projector = ops_fn::zeros(&[config.vision_hidden_size, config.hidden_size], DataType::Float32, &device)?;
106 let embed_tokens = ops_fn::zeros(&[config.vocab_size, config.hidden_size], DataType::Float32, &device)?;
107 let norm = ops_fn::zeros(&[config.hidden_size], DataType::Float32, &device)?;
108 let lm_head = ops_fn::zeros(&[config.hidden_size, config.vocab_size], DataType::Float32, &device)?;
109
110 let mut layers = Vec::with_capacity(config.num_hidden_layers);
111 for _ in 0..config.num_hidden_layers {
112 layers.push(InternLMLayer::new(&config, &device)?);
113 }
114
115 Ok(Self { config, device, vision_encoder, mlp_projector, embed_tokens, layers, norm, lm_head })
116 }
117
118 fn from_weights(config: InternVLConfig, weights: ModelWeights) -> Result<Self> {
119 let mut model = Self::new(config)?;
120 if let Some(w) = weights.get("model.embed_tokens.weight") {
121 model.embed_tokens = w.clone();
122 }
123 if let Some(w) = weights.get("model.norm.weight") {
124 model.norm = w.clone();
125 }
126 if let Some(w) = weights.get("lm_head.weight") {
127 model.lm_head = ops_fn::transpose(w)?;
128 }
129 Ok(model)
130 }
131
132 fn forward(&self, inputs: &ModelInputs) -> Result<ModelOutputs> {
133 match inputs {
134 ModelInputs::Text { input_ids, .. } => {
135 let mut hidden = ops_fn::embedding(input_ids, &self.embed_tokens)?;
136 for layer in &self.layers {
137 hidden = layer.forward(&hidden)?;
138 }
139 hidden = ops_fn::rms_norm(&hidden, &self.norm, self.config.rms_norm_eps)?;
140 let logits = ops_fn::matmul(&hidden, &self.lm_head)?;
141 Ok(ModelOutputs::Logits { logits, hidden_states: None })
142 }
143 _ => Err(anyhow::anyhow!("InternVL requires text input")),
144 }
145 }
146
147 fn generate(&self, prompt: &str, config: &GenerationConfig) -> Result<String> {
148 use crate::tokenizer::Tokenizer;
149 let tokenizer = Tokenizer::new();
150 let mut tokens: Vec<u32> = tokenizer.encode(prompt);
151
152 for _ in 0..config.max_new_tokens {
153 let input_ids = Tensor::from_i64_slice(
154 &tokens.iter().map(|&t| t as i64).collect::<Vec<_>>(),
155 &[1, tokens.len()],
156 &self.device
157 )?;
158 let outputs = self.forward(&ModelInputs::text(input_ids))?;
159 let logits = match outputs {
160 ModelOutputs::Logits { logits, .. } => logits,
161 _ => return Err(anyhow::anyhow!("Expected logits")),
162 };
163
164 let logits_vec: Vec<f32> = logits.to_candle()?.flatten_all()?.to_vec1()?;
165 let start = (tokens.len() - 1) * self.config.vocab_size;
166 let next_token = logits_vec[start..start + self.config.vocab_size].iter()
167 .enumerate()
168 .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
169 .map(|(idx, _)| idx as u32)
170 .unwrap_or(0);
171
172 if next_token == config.eos_token_id { break; }
173 tokens.push(next_token);
174 }
175
176 Ok(tokenizer.decode(&tokens))
177 }
178
179 fn config(&self) -> &Self::Config { &self.config }
180
181 fn memory_requirements(&self) -> MemoryRequirements {
182 let param_size = (self.config.vocab_size * self.config.hidden_size) * 4;
183 MemoryRequirements { gpu_memory: param_size, cpu_memory: param_size / 4, kv_cache_memory: param_size / 8, peak_memory: param_size * 2 }
184 }
185
186 fn to_device(&mut self, device: &Device) -> Result<()> {
187 self.device = device.clone();
188 self.embed_tokens = self.embed_tokens.to_device(device)?;
189 self.norm = self.norm.to_device(device)?;
190 self.lm_head = self.lm_head.to_device(device)?;
191 Ok(())
192 }
193}
194
195impl InternViTEncoder {
196 fn new(config: &InternVLConfig, device: &Device) -> Result<Self> {
197 let patch_dim = 3 * config.vision_patch_size * config.vision_patch_size;
198 let patch_embed = ops_fn::zeros(&[patch_dim, config.vision_hidden_size], DataType::Float32, device)?;
199 let norm = ops_fn::zeros(&[config.vision_hidden_size], DataType::Float32, device)?;
200
201 let mut blocks = Vec::with_capacity(config.vision_num_hidden_layers);
202 for _ in 0..config.vision_num_hidden_layers {
203 blocks.push(InternViTBlock::new(config, device)?);
204 }
205
206 Ok(Self { patch_embed, blocks, norm, hidden_size: config.vision_hidden_size })
207 }
208}
209
210impl InternViTBlock {
211 fn new(config: &InternVLConfig, device: &Device) -> Result<Self> {
212 Ok(Self {
213 norm1: ops_fn::zeros(&[config.vision_hidden_size], DataType::Float32, device)?,
214 attn_qkv: ops_fn::zeros(&[config.vision_hidden_size, config.vision_hidden_size * 3], DataType::Float32, device)?,
215 attn_proj: ops_fn::zeros(&[config.vision_hidden_size, config.vision_hidden_size], DataType::Float32, device)?,
216 norm2: ops_fn::zeros(&[config.vision_hidden_size], DataType::Float32, device)?,
217 mlp_fc1: ops_fn::zeros(&[config.vision_hidden_size, config.vision_intermediate_size], DataType::Float32, device)?,
218 mlp_fc2: ops_fn::zeros(&[config.vision_intermediate_size, config.vision_hidden_size], DataType::Float32, device)?,
219 num_heads: config.vision_num_attention_heads,
220 })
221 }
222}
223
224impl InternLMLayer {
225 fn new(config: &InternVLConfig, device: &Device) -> Result<Self> {
226 Ok(Self {
227 self_attn: InternLMAttention::new(config, device)?,
228 mlp: InternLMMLP::new(config, device)?,
229 input_layernorm: ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?,
230 post_attention_layernorm: ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?,
231 })
232 }
233
234 fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
235 let residual = hidden_states.clone();
236 let hidden = ops_fn::rms_norm(hidden_states, &self.input_layernorm, 1e-5)?;
237 let hidden = self.self_attn.forward(&hidden)?;
238 let hidden = ops_fn::add(&residual, &hidden)?;
239
240 let residual = hidden.clone();
241 let hidden = ops_fn::rms_norm(&hidden, &self.post_attention_layernorm, 1e-5)?;
242 let hidden = self.mlp.forward(&hidden)?;
243 ops_fn::add(&residual, &hidden)
244 }
245}
246
247impl InternLMAttention {
248 fn new(config: &InternVLConfig, device: &Device) -> Result<Self> {
249 let head_dim = config.hidden_size / config.num_attention_heads;
250 Ok(Self {
251 q_proj: ops_fn::zeros(&[config.hidden_size, config.num_attention_heads * head_dim], DataType::Float32, device)?,
252 k_proj: ops_fn::zeros(&[config.hidden_size, config.num_key_value_heads * head_dim], DataType::Float32, device)?,
253 v_proj: ops_fn::zeros(&[config.hidden_size, config.num_key_value_heads * head_dim], DataType::Float32, device)?,
254 o_proj: ops_fn::zeros(&[config.num_attention_heads * head_dim, config.hidden_size], DataType::Float32, device)?,
255 num_heads: config.num_attention_heads,
256 num_kv_heads: config.num_key_value_heads,
257 head_dim,
258 })
259 }
260
261 fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
262 let shape = hidden_states.shape();
263 let (batch_size, seq_len, _) = (shape[0], shape[1], shape[2]);
264
265 let q = ops_fn::matmul(hidden_states, &self.q_proj)?;
266 let k = ops_fn::matmul(hidden_states, &self.k_proj)?;
267 let v = ops_fn::matmul(hidden_states, &self.v_proj)?;
268
269 let q = q.to_candle()?.reshape(&[batch_size, seq_len, self.num_heads, self.head_dim])?.transpose(1, 2)?;
270 let k = k.to_candle()?.reshape(&[batch_size, seq_len, self.num_kv_heads, self.head_dim])?.transpose(1, 2)?;
271 let v = v.to_candle()?.reshape(&[batch_size, seq_len, self.num_kv_heads, self.head_dim])?.transpose(1, 2)?;
272
273 let num_groups = self.num_heads / self.num_kv_heads;
274 let (k, v) = if num_groups > 1 {
275 (k.unsqueeze(2)?.broadcast_as(&[batch_size, self.num_kv_heads, num_groups, seq_len, self.head_dim])?.reshape(&[batch_size, self.num_heads, seq_len, self.head_dim])?,
276 v.unsqueeze(2)?.broadcast_as(&[batch_size, self.num_kv_heads, num_groups, seq_len, self.head_dim])?.reshape(&[batch_size, self.num_heads, seq_len, self.head_dim])?)
277 } else { (k, v) };
278
279 let scale = (self.head_dim as f32).powf(-0.5);
280 let scores = q.contiguous()?.matmul(&k.transpose(2, 3)?.contiguous()?)?;
281 let scores = (scores * (scale as f64))?;
282
283 let device = scores.device();
284 let mask = {
285 let mut m = vec![0.0f32; seq_len * seq_len];
286 for i in 0..seq_len { for j in (i + 1)..seq_len { m[i * seq_len + j] = f32::NEG_INFINITY; } }
287 candle_core::Tensor::from_vec(m, &[1, 1, seq_len, seq_len], device)?
288 };
289
290 let scores = scores.broadcast_add(&mask)?;
291 let attn_weights = candle_nn::ops::softmax_last_dim(&scores)?;
292 let attn_output = attn_weights.matmul(&v.contiguous()?)?
293 .transpose(1, 2)?
294 .reshape(&[batch_size, seq_len, self.num_heads * self.head_dim])?;
295
296 ops_fn::matmul(&Tensor::from_candle(attn_output), &self.o_proj)
297 }
298}
299
300impl InternLMMLP {
301 fn new(config: &InternVLConfig, device: &Device) -> Result<Self> {
302 Ok(Self {
303 gate_proj: ops_fn::zeros(&[config.hidden_size, config.intermediate_size], DataType::Float32, device)?,
304 up_proj: ops_fn::zeros(&[config.hidden_size, config.intermediate_size], DataType::Float32, device)?,
305 down_proj: ops_fn::zeros(&[config.intermediate_size, config.hidden_size], DataType::Float32, device)?,
306 })
307 }
308
309 fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
310 let gate = ops_fn::silu(&ops_fn::matmul(hidden_states, &self.gate_proj)?)?;
311 let up = ops_fn::matmul(hidden_states, &self.up_proj)?;
312 ops_fn::matmul(&ops_fn::mul(&gate, &up)?, &self.down_proj)
313 }
314}
315
316#[cfg(test)]
317mod tests {
318 use super::*;
319
320 #[test]
321 fn test_internvl_config() {
322 let config = InternVLConfig::default();
323 assert_eq!(config.vocab_size, 92544);
324 }
325
326 #[test]
327 fn test_internvl_forward() {
328 let config = InternVLConfig {
329 vocab_size: 100, hidden_size: 32, intermediate_size: 128,
330 num_hidden_layers: 1, num_attention_heads: 2, num_key_value_heads: 2,
331 vision_hidden_size: 16, vision_intermediate_size: 64,
332 vision_num_hidden_layers: 1, vision_num_attention_heads: 2,
333 ..Default::default()
334 };
335
336 let model = InternVLModelV2::new(config).unwrap();
337 let input_ids = ops_fn::zeros(&[1, 4], DataType::Int64, &Device::CPU).unwrap();
338 let outputs = model.forward(&ModelInputs::text(input_ids)).unwrap();
339
340 match outputs {
341 ModelOutputs::Logits { logits, .. } => assert_eq!(logits.shape(), &[1, 4, 100]),
342 _ => panic!("Expected logits"),
343 }
344 }
345}