runtime/models_v2/
internlm.rs

1//! InternLM Model V2 - Clean implementation using solid abstractions
2//!
3//! This implements the InternLM architecture including:
4//! - InternLM-7B, InternLM-20B, InternLM2-7B, InternLM2-20B
5//!
6//! InternLM-specific characteristics:
7//! - Uses packed QKV attention (wqkv weight) - split after projection
8//! - Uses w1, w2, w3 for MLP instead of gate_proj, up_proj, down_proj
9//! - Weight path uses `.attention.` instead of `.self_attn.`
10//! - Weight path uses `.feed_forward.` instead of `.mlp.`
11
12use crate::model_config;
13use super::traits::*;
14use anyhow::Result;
15use serde::{Serialize, Deserialize};
16
17/// InternLM model configuration using the model_config macro
18model_config!(InternLMConfig {
19    vocab_size: usize = 103168,
20    hidden_size: usize = 4096,
21    intermediate_size: usize = 11008,
22    num_hidden_layers: usize = 32,
23    num_attention_heads: usize = 32,
24    num_key_value_heads: usize = 8,
25    hidden_act: String = "silu".to_string(),
26    max_position_embeddings: usize = 2048,
27    initializer_range: f32 = 0.02,
28    rms_norm_eps: f32 = 1e-6,
29    use_cache: bool = true,
30    pad_token_id: i64 = 0,
31    bos_token_id: i64 = 1,
32    eos_token_id: i64 = 2,
33    tie_word_embeddings: bool = false,
34    rope_theta: f32 = 10000.0,
35    attention_dropout: f32 = 0.0,
36    bias: bool = false,
37});
38
39impl InternLMConfig {
40    /// Create InternLMConfig from GGUF model configuration
41    pub fn from_gguf_config(gguf: &crate::weight_loader_core::GGUFModelConfig) -> Self {
42        Self {
43            vocab_size: gguf.vocab_size,
44            hidden_size: gguf.hidden_size,
45            intermediate_size: gguf.intermediate_size,
46            num_hidden_layers: gguf.num_hidden_layers,
47            num_attention_heads: gguf.num_attention_heads,
48            num_key_value_heads: gguf.num_key_value_heads,
49            rms_norm_eps: gguf.rms_norm_eps,
50            rope_theta: gguf.rope_theta,
51            max_position_embeddings: gguf.max_position_embeddings,
52            ..Default::default()
53        }
54    }
55}
56
57/// Main InternLM model implementation
58pub struct InternLMModelV2 {
59    config: InternLMConfig,
60    device: Device,
61
62    // Model components using unified Tensor type
63    embed_tokens: Tensor,
64    layers: Vec<InternLMLayer>,
65    norm: Tensor,
66    lm_head: Tensor,
67}
68
69/// InternLM transformer layer
70pub struct InternLMLayer {
71    self_attn: InternLMAttention,
72    mlp: InternLMMLP,
73    input_layernorm: Tensor,
74    post_attention_layernorm: Tensor,
75}
76
77/// InternLM attention mechanism
78/// Supports both packed QKV (wqkv) and separate Q/K/V projections
79pub struct InternLMAttention {
80    // Packed QKV projection (InternLM style)
81    wqkv: Option<Tensor>,
82    // Separate projections (standard style)
83    q_proj: Option<Tensor>,
84    k_proj: Option<Tensor>,
85    v_proj: Option<Tensor>,
86    // Output projection
87    wo: Tensor,
88    num_heads: usize,
89    num_key_value_heads: usize,
90    head_dim: usize,
91    scale: f32,
92}
93
94/// InternLM MLP (feed-forward network)
95/// Uses w1, w2, w3 naming convention
96pub struct InternLMMLP {
97    w1: Tensor, // gate_proj
98    w2: Tensor, // down_proj
99    w3: Tensor, // up_proj
100    hidden_act: String,
101}
102
103impl Model for InternLMModelV2 {
104    type Config = InternLMConfig;
105
106    fn new(config: InternLMConfig) -> Result<Self> {
107        let device = Device::CPU;
108
109        let embed_tokens = ops_fn::zeros(
110            &[config.vocab_size, config.hidden_size],
111            DataType::Float32,
112            &device
113        )?;
114
115        let norm = ops_fn::zeros(
116            &[config.hidden_size],
117            DataType::Float32,
118            &device
119        )?;
120
121        let lm_head = if config.tie_word_embeddings {
122            embed_tokens.clone()
123        } else {
124            ops_fn::zeros(
125                &[config.hidden_size, config.vocab_size],
126                DataType::Float32,
127                &device
128            )?
129        };
130
131        // Create transformer layers
132        let mut layers = Vec::with_capacity(config.num_hidden_layers);
133        for _ in 0..config.num_hidden_layers {
134            layers.push(InternLMLayer::new(&config, &device)?);
135        }
136
137        Ok(Self {
138            config,
139            device,
140            embed_tokens,
141            layers,
142            norm,
143            lm_head,
144        })
145    }
146
147    fn from_weights(config: InternLMConfig, weights: ModelWeights) -> Result<Self> {
148        let mut model = Self::new(config)?;
149
150        // Load weights from the unified weight container
151        // Note: embedding weight is not transposed (used for index lookup)
152        if let Some(embed_weights) = weights.get("model.embed_tokens.weight") {
153            model.embed_tokens = embed_weights.clone();
154        }
155
156        if let Some(norm_weights) = weights.get("model.norm.weight") {
157            model.norm = norm_weights.clone();
158        }
159
160        // lm_head weight needs transpose: [vocab, hidden] -> [hidden, vocab] for matmul
161        if let Some(lm_head_weights) = weights.get("lm_head.weight") {
162            model.lm_head = ops_fn::transpose(lm_head_weights)?;
163        }
164
165        // Load layer weights
166        for (i, layer) in model.layers.iter_mut().enumerate() {
167            layer.load_weights(&weights, i)?;
168        }
169
170        Ok(model)
171    }
172
173    fn forward(&self, inputs: &ModelInputs) -> Result<ModelOutputs> {
174        match inputs {
175            ModelInputs::Text { input_ids, attention_mask, .. } => {
176                // 1. Token embedding
177                let mut hidden_states = ops_fn::embedding(input_ids, &self.embed_tokens)?;
178
179                // 2. Apply transformer layers (with RoPE)
180                for layer in &self.layers {
181                    hidden_states = layer.forward(&hidden_states, attention_mask.as_ref(), self.config.rope_theta)?;
182                }
183
184                // 3. Final layer norm
185                hidden_states = ops_fn::layer_norm(&hidden_states, &self.norm, None, self.config.rms_norm_eps)?;
186
187                // 4. Language modeling head
188                let logits = ops_fn::matmul(&hidden_states, &self.lm_head)?;
189
190                Ok(ModelOutputs::Logits {
191                    logits,
192                    hidden_states: None,
193                })
194            }
195            ModelInputs::Multimodal { input_ids, .. } => {
196                let text_inputs = ModelInputs::Text {
197                    input_ids: input_ids.clone(),
198                    attention_mask: None,
199                    position_ids: None,
200                };
201                self.forward(&text_inputs)
202            }
203            _ => Err(anyhow::anyhow!("InternLM model only supports text and multimodal inputs")),
204        }
205    }
206
207    fn generate(&self, prompt: &str, config: &GenerationConfig) -> Result<String> {
208        use crate::tokenizer::Tokenizer;
209        use rand::Rng;
210
211        // 1. Tokenize prompt
212        let tokenizer = Tokenizer::new();
213        let mut tokens: Vec<u32> = tokenizer.encode(prompt);
214
215        // 2. Generation loop
216        for _ in 0..config.max_new_tokens {
217            // Create input tensor from current tokens
218            let tokens_i64: Vec<i64> = tokens.iter().map(|&t| t as i64).collect();
219            let input_tensor = Tensor::from_i64_slice(&tokens_i64, &[1, tokens.len()], &self.device)?;
220
221            let inputs = ModelInputs::Text {
222                input_ids: input_tensor,
223                attention_mask: None,
224                position_ids: None,
225            };
226
227            // 3. Forward pass
228            let outputs = self.forward(&inputs)?;
229
230            // 4. Get logits and sample next token
231            let logits = match outputs {
232                ModelOutputs::Logits { logits, .. } => logits,
233                _ => return Err(anyhow::anyhow!("Expected logits output")),
234            };
235
236            // Get last token logits
237            let logits_candle = logits.to_candle()?;
238            let shape = logits_candle.dims();
239
240            // Extract last position logits [batch, seq, vocab] -> [vocab]
241            let last_logits = if shape.len() == 3 {
242                let seq_len = shape[1];
243                logits_candle
244                    .narrow(1, seq_len - 1, 1)?
245                    .squeeze(1)?
246                    .squeeze(0)?
247            } else {
248                let seq_len = shape[0];
249                logits_candle
250                    .narrow(0, seq_len - 1, 1)?
251                    .squeeze(0)?
252            };
253
254            // Convert to probabilities and sample
255            let logits_vec: Vec<f32> = last_logits.to_vec1()?;
256
257            let next_token = if config.do_sample && config.temperature > 0.0 {
258                // Temperature sampling
259                let scaled: Vec<f32> = logits_vec.iter()
260                    .map(|&x| x / config.temperature)
261                    .collect();
262
263                // Softmax
264                let max_val = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
265                let exp_sum: f32 = scaled.iter().map(|&x| (x - max_val).exp()).sum();
266                let probs: Vec<f32> = scaled.iter()
267                    .map(|&x| (x - max_val).exp() / exp_sum)
268                    .collect();
269
270                // Sample from distribution
271                let mut rng = rand::thread_rng();
272                let random_val: f32 = rng.gen();
273                let mut cumulative = 0.0;
274                let mut sampled = 0u32;
275
276                for (idx, &prob) in probs.iter().enumerate() {
277                    cumulative += prob;
278                    if random_val <= cumulative {
279                        sampled = idx as u32;
280                        break;
281                    }
282                }
283                sampled
284            } else {
285                // Greedy sampling
286                let mut max_idx = 0;
287                let mut max_val = logits_vec[0];
288                for (idx, &val) in logits_vec.iter().enumerate() {
289                    if val > max_val {
290                        max_val = val;
291                        max_idx = idx;
292                    }
293                }
294                max_idx as u32
295            };
296
297            // 5. Check for EOS
298            if next_token == config.eos_token_id {
299                break;
300            }
301
302            // 6. Append token
303            tokens.push(next_token);
304        }
305
306        // 7. Decode and return
307        Ok(tokenizer.decode(&tokens))
308    }
309
310    fn config(&self) -> &Self::Config {
311        &self.config
312    }
313
314    fn memory_requirements(&self) -> MemoryRequirements {
315        let param_size = self.config.vocab_size * self.config.hidden_size +
316                        self.config.num_hidden_layers * (
317                            4 * self.config.hidden_size * self.config.hidden_size +
318                            3 * self.config.hidden_size * self.config.intermediate_size
319                        );
320
321        let param_bytes = param_size * 4;
322        let kv_cache_bytes = 2 * self.config.num_hidden_layers *
323                           self.config.max_position_embeddings *
324                           self.config.hidden_size * 4;
325
326        MemoryRequirements {
327            gpu_memory: param_bytes,
328            cpu_memory: param_bytes / 4,
329            kv_cache_memory: kv_cache_bytes,
330            peak_memory: param_bytes + kv_cache_bytes,
331        }
332    }
333
334    fn to_device(&mut self, device: &Device) -> Result<()> {
335        self.embed_tokens = self.embed_tokens.to_device(device)?;
336        self.norm = self.norm.to_device(device)?;
337        self.lm_head = self.lm_head.to_device(device)?;
338
339        for layer in &mut self.layers {
340            layer.to_device(device)?;
341        }
342
343        self.device = device.clone();
344        Ok(())
345    }
346}
347
348impl InternLMLayer {
349    fn new(config: &InternLMConfig, device: &Device) -> Result<Self> {
350        let self_attn = InternLMAttention::new(config, device)?;
351        let mlp = InternLMMLP::new(config, device)?;
352
353        let input_layernorm = ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?;
354        let post_attention_layernorm = ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?;
355
356        Ok(Self {
357            self_attn,
358            mlp,
359            input_layernorm,
360            post_attention_layernorm,
361        })
362    }
363
364    fn forward(&self, hidden_states: &Tensor, attention_mask: Option<&Tensor>, rope_theta: f32) -> Result<Tensor> {
365        // 1. Pre-attention layer norm
366        let normed = ops_fn::layer_norm(hidden_states, &self.input_layernorm, None, 1e-6)?;
367
368        // 2. Self attention (with RoPE)
369        let attn_output = self.self_attn.forward(&normed, attention_mask, rope_theta)?;
370
371        // 3. Residual connection
372        let hidden_states = ops_fn::add(hidden_states, &attn_output)?;
373
374        // 4. Pre-MLP layer norm
375        let normed = ops_fn::layer_norm(&hidden_states, &self.post_attention_layernorm, None, 1e-6)?;
376
377        // 5. MLP
378        let mlp_output = self.mlp.forward(&normed)?;
379
380        // 6. Residual connection
381        let output = ops_fn::add(&hidden_states, &mlp_output)?;
382
383        Ok(output)
384    }
385
386    fn load_weights(&mut self, weights: &ModelWeights, layer_idx: usize) -> Result<()> {
387        let prefix = format!("model.layers.{}", layer_idx);
388
389        // Load attention weights - InternLM uses .attention. path
390        // Check for packed QKV first (InternLM style)
391        if let Some(wqkv) = weights.get(&format!("{}.attention.wqkv.weight", prefix)) {
392            self.self_attn.wqkv = Some(ops_fn::transpose(wqkv)?);
393            self.self_attn.q_proj = None;
394            self.self_attn.k_proj = None;
395            self.self_attn.v_proj = None;
396        } else {
397            // Fall back to separate Q/K/V projections
398            if let Some(q_proj) = weights.get(&format!("{}.attention.q_proj.weight", prefix))
399                .or_else(|| weights.get(&format!("{}.self_attn.q_proj.weight", prefix))) {
400                self.self_attn.q_proj = Some(ops_fn::transpose(q_proj)?);
401            }
402            if let Some(k_proj) = weights.get(&format!("{}.attention.k_proj.weight", prefix))
403                .or_else(|| weights.get(&format!("{}.self_attn.k_proj.weight", prefix))) {
404                self.self_attn.k_proj = Some(ops_fn::transpose(k_proj)?);
405            }
406            if let Some(v_proj) = weights.get(&format!("{}.attention.v_proj.weight", prefix))
407                .or_else(|| weights.get(&format!("{}.self_attn.v_proj.weight", prefix))) {
408                self.self_attn.v_proj = Some(ops_fn::transpose(v_proj)?);
409            }
410            self.self_attn.wqkv = None;
411        }
412
413        // Output projection - InternLM uses wo
414        if let Some(wo) = weights.get(&format!("{}.attention.wo.weight", prefix))
415            .or_else(|| weights.get(&format!("{}.attention.o_proj.weight", prefix)))
416            .or_else(|| weights.get(&format!("{}.self_attn.o_proj.weight", prefix))) {
417            self.self_attn.wo = ops_fn::transpose(wo)?;
418        }
419
420        // Load MLP weights - InternLM uses .feed_forward. with w1, w2, w3
421        if let Some(w1) = weights.get(&format!("{}.feed_forward.w1.weight", prefix))
422            .or_else(|| weights.get(&format!("{}.mlp.gate_proj.weight", prefix))) {
423            self.mlp.w1 = ops_fn::transpose(w1)?;
424        }
425        if let Some(w2) = weights.get(&format!("{}.feed_forward.w2.weight", prefix))
426            .or_else(|| weights.get(&format!("{}.mlp.down_proj.weight", prefix))) {
427            self.mlp.w2 = ops_fn::transpose(w2)?;
428        }
429        if let Some(w3) = weights.get(&format!("{}.feed_forward.w3.weight", prefix))
430            .or_else(|| weights.get(&format!("{}.mlp.up_proj.weight", prefix))) {
431            self.mlp.w3 = ops_fn::transpose(w3)?;
432        }
433
434        // Load layer norm weights (no transpose needed - 1D tensors)
435        if let Some(input_ln) = weights.get(&format!("{}.input_layernorm.weight", prefix))
436            .or_else(|| weights.get(&format!("{}.attention_norm.weight", prefix))) {
437            self.input_layernorm = input_ln.clone();
438        }
439        if let Some(post_ln) = weights.get(&format!("{}.post_attention_layernorm.weight", prefix))
440            .or_else(|| weights.get(&format!("{}.ffn_norm.weight", prefix))) {
441            self.post_attention_layernorm = post_ln.clone();
442        }
443
444        Ok(())
445    }
446
447    fn to_device(&mut self, device: &Device) -> Result<()> {
448        self.self_attn.to_device(device)?;
449        self.mlp.to_device(device)?;
450        self.input_layernorm = self.input_layernorm.to_device(device)?;
451        self.post_attention_layernorm = self.post_attention_layernorm.to_device(device)?;
452        Ok(())
453    }
454}
455
456/// Apply Rotary Position Embedding (RoPE) to Q and K tensors
457fn apply_rope(
458    q: &candle_core::Tensor,
459    k: &candle_core::Tensor,
460    seq_len: usize,
461    head_dim: usize,
462    rope_theta: f32,
463) -> Result<(candle_core::Tensor, candle_core::Tensor)> {
464    let device = q.device();
465
466    // Compute inverse frequencies: 1 / (theta^(2i/d)) for i in [0, d/2)
467    let half_dim = head_dim / 2;
468    let inv_freq: Vec<f32> = (0..half_dim)
469        .map(|i| 1.0 / rope_theta.powf((2 * i) as f32 / head_dim as f32))
470        .collect();
471
472    // Create position indices [0, 1, 2, ..., seq_len-1]
473    let positions: Vec<f32> = (0..seq_len).map(|p| p as f32).collect();
474
475    // Compute angles: pos * inv_freq -> [seq_len, half_dim]
476    let mut angles = Vec::with_capacity(seq_len * half_dim);
477    for pos in &positions {
478        for freq in &inv_freq {
479            angles.push(pos * freq);
480        }
481    }
482
483    let angles_tensor = candle_core::Tensor::from_vec(angles, &[seq_len, half_dim], device)?;
484
485    // Compute cos and sin
486    let cos = angles_tensor.cos()?;
487    let sin = angles_tensor.sin()?;
488
489    // Reshape for broadcasting: [1, 1, seq_len, half_dim]
490    let cos = cos.unsqueeze(0)?.unsqueeze(0)?;
491    let sin = sin.unsqueeze(0)?.unsqueeze(0)?;
492
493    // Apply RoPE rotation
494    let q_half1 = q.narrow(3, 0, half_dim)?;
495    let q_half2 = q.narrow(3, half_dim, half_dim)?;
496    let k_half1 = k.narrow(3, 0, half_dim)?;
497    let k_half2 = k.narrow(3, half_dim, half_dim)?;
498
499    // Apply rotation
500    let q_rot1 = (q_half1.broadcast_mul(&cos)? - q_half2.broadcast_mul(&sin)?)?;
501    let q_rot2 = (q_half1.broadcast_mul(&sin)? + q_half2.broadcast_mul(&cos)?)?;
502    let k_rot1 = (k_half1.broadcast_mul(&cos)? - k_half2.broadcast_mul(&sin)?)?;
503    let k_rot2 = (k_half1.broadcast_mul(&sin)? + k_half2.broadcast_mul(&cos)?)?;
504
505    // Concatenate rotated halves
506    let q_rotated = candle_core::Tensor::cat(&[&q_rot1, &q_rot2], 3)?;
507    let k_rotated = candle_core::Tensor::cat(&[&k_rot1, &k_rot2], 3)?;
508
509    Ok((q_rotated, k_rotated))
510}
511
512impl InternLMAttention {
513    fn new(config: &InternLMConfig, device: &Device) -> Result<Self> {
514        let num_heads = config.num_attention_heads;
515        let num_key_value_heads = config.num_key_value_heads;
516        let head_dim = config.hidden_size / num_heads;
517        let scale = 1.0 / (head_dim as f32).sqrt();
518
519        // Initialize packed QKV projection
520        let wqkv_size = (num_heads + 2 * num_key_value_heads) * head_dim;
521        let wqkv = ops_fn::zeros(&[config.hidden_size, wqkv_size], DataType::Float32, device)?;
522        let wo = ops_fn::zeros(&[num_heads * head_dim, config.hidden_size], DataType::Float32, device)?;
523
524        Ok(Self {
525            wqkv: Some(wqkv),
526            q_proj: None,
527            k_proj: None,
528            v_proj: None,
529            wo,
530            num_heads,
531            num_key_value_heads,
532            head_dim,
533            scale,
534        })
535    }
536
537    fn forward(&self, hidden_states: &Tensor, _attention_mask: Option<&Tensor>, rope_theta: f32) -> Result<Tensor> {
538        // Get batch and sequence length from hidden_states shape
539        let shape = hidden_states.shape();
540        let (batch_size, seq_len, _hidden_size) = if shape.len() == 3 {
541            (shape[0], shape[1], shape[2])
542        } else if shape.len() == 2 {
543            (1, shape[0], shape[1])
544        } else {
545            return Err(anyhow::anyhow!("Invalid hidden_states shape: {:?}", shape));
546        };
547
548        // 1. Project to Q, K, V - handle both packed and separate projections
549        let (query_states, key_states, value_states) = if let Some(wqkv) = &self.wqkv {
550            // Packed QKV projection (InternLM style)
551            let qkv_out = ops_fn::matmul(hidden_states, wqkv)?;
552            let qkv_candle = qkv_out.to_candle()?;
553
554            // Split QKV: [batch, seq, (num_heads + 2*num_kv_heads) * head_dim]
555            let q_size = self.num_heads * self.head_dim;
556            let kv_size = self.num_key_value_heads * self.head_dim;
557
558            let q = qkv_candle.narrow(2, 0, q_size)?;
559            let k = qkv_candle.narrow(2, q_size, kv_size)?;
560            let v = qkv_candle.narrow(2, q_size + kv_size, kv_size)?;
561
562            (Tensor::from_candle(q), Tensor::from_candle(k), Tensor::from_candle(v))
563        } else {
564            // Separate Q/K/V projections
565            let q_proj = self.q_proj.as_ref().ok_or_else(|| anyhow::anyhow!("Missing q_proj"))?;
566            let k_proj = self.k_proj.as_ref().ok_or_else(|| anyhow::anyhow!("Missing k_proj"))?;
567            let v_proj = self.v_proj.as_ref().ok_or_else(|| anyhow::anyhow!("Missing v_proj"))?;
568
569            let query_states = ops_fn::matmul(hidden_states, q_proj)?;
570            let key_states = ops_fn::matmul(hidden_states, k_proj)?;
571            let value_states = ops_fn::matmul(hidden_states, v_proj)?;
572
573            (query_states, key_states, value_states)
574        };
575
576        // 2. Reshape for multi-head attention
577        let q_candle = query_states.to_candle()?;
578        let k_candle = key_states.to_candle()?;
579        let v_candle = value_states.to_candle()?;
580
581        let q_reshaped = q_candle
582            .reshape(&[batch_size, seq_len, self.num_heads, self.head_dim])?
583            .transpose(1, 2)?;
584
585        let k_reshaped = k_candle
586            .reshape(&[batch_size, seq_len, self.num_key_value_heads, self.head_dim])?
587            .transpose(1, 2)?;
588
589        let v_reshaped = v_candle
590            .reshape(&[batch_size, seq_len, self.num_key_value_heads, self.head_dim])?
591            .transpose(1, 2)?;
592
593        // 3. Apply RoPE (Rotary Position Embedding) to Q and K
594        let (q_with_rope, k_with_rope) = apply_rope(&q_reshaped, &k_reshaped, seq_len, self.head_dim, rope_theta)?;
595
596        // 4. Handle GQA (Grouped Query Attention) - repeat K/V heads to match Q heads
597        let num_groups = self.num_heads / self.num_key_value_heads;
598        let (k_expanded, v_expanded) = if num_groups > 1 {
599            let k_rep = k_with_rope
600                .unsqueeze(2)?
601                .broadcast_as(&[batch_size, self.num_key_value_heads, num_groups, seq_len, self.head_dim])?
602                .reshape(&[batch_size, self.num_heads, seq_len, self.head_dim])?;
603            let v_rep = v_reshaped
604                .unsqueeze(2)?
605                .broadcast_as(&[batch_size, self.num_key_value_heads, num_groups, seq_len, self.head_dim])?
606                .reshape(&[batch_size, self.num_heads, seq_len, self.head_dim])?;
607            (k_rep, v_rep)
608        } else {
609            (k_with_rope, v_reshaped)
610        };
611
612        // 5. Scaled dot-product attention
613        let k_t = k_expanded.transpose(2, 3)?;
614
615        let q_contiguous = q_with_rope.contiguous()?;
616        let k_contiguous = k_t.contiguous()?;
617
618        let scores = q_contiguous.matmul(&k_contiguous)?;
619        let scaled_scores = (scores * (self.scale as f64))?;
620
621        // Apply causal mask
622        let device = scaled_scores.device();
623        let causal_mask = {
624            let mut mask_data = vec![0.0f32; seq_len * seq_len];
625            for i in 0..seq_len {
626                for j in 0..seq_len {
627                    if j > i {
628                        mask_data[i * seq_len + j] = f32::NEG_INFINITY;
629                    }
630                }
631            }
632            candle_core::Tensor::from_vec(mask_data, &[1, 1, seq_len, seq_len], device)?
633        };
634
635        let masked_scores = scaled_scores.broadcast_add(&causal_mask)?;
636
637        // Softmax over last dimension
638        let attention_weights = candle_nn::ops::softmax_last_dim(&masked_scores)?;
639
640        // Apply attention to values
641        let v_contiguous = v_expanded.contiguous()?;
642        let attn_output = attention_weights.matmul(&v_contiguous)?;
643
644        // 6. Reshape back
645        let attn_output = attn_output
646            .transpose(1, 2)?
647            .reshape(&[batch_size, seq_len, self.num_heads * self.head_dim])?;
648
649        let attn_output = Tensor::from_candle(attn_output);
650
651        // 7. Output projection
652        let output = ops_fn::matmul(&attn_output, &self.wo)?;
653
654        Ok(output)
655    }
656
657    fn to_device(&mut self, device: &Device) -> Result<()> {
658        if let Some(wqkv) = &self.wqkv {
659            self.wqkv = Some(wqkv.to_device(device)?);
660        }
661        if let Some(q_proj) = &self.q_proj {
662            self.q_proj = Some(q_proj.to_device(device)?);
663        }
664        if let Some(k_proj) = &self.k_proj {
665            self.k_proj = Some(k_proj.to_device(device)?);
666        }
667        if let Some(v_proj) = &self.v_proj {
668            self.v_proj = Some(v_proj.to_device(device)?);
669        }
670        self.wo = self.wo.to_device(device)?;
671        Ok(())
672    }
673}
674
675impl InternLMMLP {
676    fn new(config: &InternLMConfig, device: &Device) -> Result<Self> {
677        let w1 = ops_fn::zeros(&[config.hidden_size, config.intermediate_size], DataType::Float32, device)?;
678        let w2 = ops_fn::zeros(&[config.intermediate_size, config.hidden_size], DataType::Float32, device)?;
679        let w3 = ops_fn::zeros(&[config.hidden_size, config.intermediate_size], DataType::Float32, device)?;
680
681        Ok(Self {
682            w1,
683            w2,
684            w3,
685            hidden_act: config.hidden_act.clone(),
686        })
687    }
688
689    fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
690        // 1. Gate and up projections
691        let gate_output = ops_fn::matmul(hidden_states, &self.w1)?;
692        let up_output = ops_fn::matmul(hidden_states, &self.w3)?;
693
694        // 2. Apply activation (SiLU for InternLM)
695        let gate_activated = match self.hidden_act.as_str() {
696            "silu" | "swish" => ops_fn::silu(&gate_output)?,
697            "gelu" => ops_fn::gelu(&gate_output)?,
698            _ => return Err(anyhow::anyhow!("Unsupported activation: {}", self.hidden_act)),
699        };
700
701        // 3. Element-wise multiplication (gating)
702        let gated = ops_fn::mul(&gate_activated, &up_output)?;
703
704        // 4. Down projection
705        let output = ops_fn::matmul(&gated, &self.w2)?;
706
707        Ok(output)
708    }
709
710    fn to_device(&mut self, device: &Device) -> Result<()> {
711        self.w1 = self.w1.to_device(device)?;
712        self.w2 = self.w2.to_device(device)?;
713        self.w3 = self.w3.to_device(device)?;
714        Ok(())
715    }
716}
717
718#[cfg(test)]
719mod tests {
720    use super::*;
721
722    #[test]
723    fn test_internlm_model_creation() {
724        let config = InternLMConfig {
725            vocab_size: 1000,
726            hidden_size: 128,
727            intermediate_size: 512,
728            num_hidden_layers: 2,
729            num_attention_heads: 8,
730            num_key_value_heads: 8,
731            ..Default::default()
732        };
733
734        let model = InternLMModelV2::new(config).unwrap();
735        assert_eq!(model.config().vocab_size(), 1000);
736        assert_eq!(model.config().hidden_size(), 128);
737        assert_eq!(model.config().num_layers(), 2);
738    }
739
740    #[test]
741    fn test_internlm_forward_pass() {
742        let config = InternLMConfig {
743            vocab_size: 100,
744            hidden_size: 64,
745            intermediate_size: 256,
746            num_hidden_layers: 1,
747            num_attention_heads: 4,
748            num_key_value_heads: 4,
749            ..Default::default()
750        };
751
752        let model = InternLMModelV2::new(config).unwrap();
753        let input_ids = ops_fn::zeros(&[2, 8], DataType::Int64, &Device::CPU).unwrap();
754        let inputs = ModelInputs::text(input_ids);
755
756        let outputs = model.forward(&inputs).unwrap();
757        match outputs {
758            ModelOutputs::Logits { logits, .. } => {
759                assert_eq!(logits.shape(), &[2, 8, 100]);
760            }
761            _ => panic!("Expected logits output"),
762        }
763    }
764
765    #[test]
766    fn test_internlm_generation() {
767        let config = InternLMConfig {
768            vocab_size: 256,
769            hidden_size: 64,
770            intermediate_size: 256,
771            num_hidden_layers: 1,
772            num_attention_heads: 4,
773            num_key_value_heads: 4,
774            ..Default::default()
775        };
776        let model = InternLMModelV2::new(config).unwrap();
777        let gen_config = GenerationConfig {
778            max_new_tokens: 5,
779            ..Default::default()
780        };
781
782        let output = model.generate("Hello", &gen_config).unwrap();
783        assert!(!output.is_empty());
784    }
785
786    #[test]
787    fn test_internlm_config_from_gguf() {
788        let gguf_config = crate::weight_loader_core::GGUFModelConfig {
789            architecture: "internlm".to_string(),
790            vocab_size: 65536,
791            hidden_size: 4096,
792            intermediate_size: 11008,
793            num_hidden_layers: 32,
794            num_attention_heads: 32,
795            num_key_value_heads: 8,
796            head_dim: 128,
797            rms_norm_eps: 1e-5,
798            rope_theta: 10000.0,
799            max_position_embeddings: 4096,
800        };
801
802        let config = InternLMConfig::from_gguf_config(&gguf_config);
803        assert_eq!(config.vocab_size, 65536);
804        assert_eq!(config.hidden_size, 4096);
805        assert_eq!(config.num_key_value_heads, 8);
806    }
807}
runtime/models_v2/internlm.rs

runtime/models_v2/
internlm.rs
