//! LLaMA modern efficient transformer implementation use crate::architectures::{TransformerConfig, TransformerArchitecture, TransformerBlock}; use crate::layers::{LayerNorm, RMSNorm, PositionalEncoding}; use crate::training::{TransformerModel, ModelOutput, ModelConfig}; use crate::{Result, TransformerError}; use rtx_tensor::{Tensor, Device, DType}; use std::collections::HashMap; use serde::{Deserialize, Serialize}; use tracing::{info, debug}; /// LLaMA-specific configuration #[derive(Debug, Clone, Serialize, Deserialize)] pub struct LLaMAConfig { /// Base transformer configuration pub base: TransformerConfig, /// Vocabulary size pub vocab_size: usize, /// Hidden dimension pub hidden_size: usize, /// Feed-forward intermediate dimension pub intermediate_size: usize, /// Number of transformer layers pub num_hidden_layers: usize, /// Number of attention heads pub num_attention_heads: usize, /// Number of key-value heads (for grouped-query attention) pub num_key_value_heads: usize, /// Maximum sequence length pub max_position_embeddings: usize, /// RMSNorm epsilon pub rms_norm_eps: f64, /// Initializer range for weights pub initializer_range: f64, /// Use cache for generation pub use_cache: bool, /// Pad token ID pub pad_token_id: i64, /// BOS token ID pub bos_token_id: i64, /// EOS token ID pub eos_token_id: i64, /// Pretraining type pub pretraining_tp: usize, /// Tie word embeddings pub tie_word_embeddings: bool, /// Rope scaling configuration pub rope_scaling: Option, /// Rope theta parameter pub rope_theta: f64, /// Attention bias pub attention_bias: bool, /// MLP bias pub mlp_bias: bool, } /// RoPE (Rotary Position Embedding) scaling configuration #[derive(Debug, Clone, Serialize, Deserialize)] pub struct RopeScaling { /// Scaling type ("linear" or "dynamic") pub scaling_type: String, /// Scaling factor pub factor: f64, } impl Default for LLaMAConfig { fn default() -> Self { Self { base: TransformerConfig::default(), vocab_size: 32000, hidden_size: 4096, intermediate_size: 11008, num_hidden_layers: 32, num_attention_heads: 32, num_key_value_heads: 32, // Standard multi-head attention max_position_embeddings: 2048, rms_norm_eps: 1e-6, initializer_range: 0.02, use_cache: true, pad_token_id: -1, bos_token_id: 1, eos_token_id: 2, pretraining_tp: 1, tie_word_embeddings: false, rope_scaling: None, rope_theta: 10000.0, attention_bias: false, mlp_bias: false, } } } impl LLaMAConfig { /// Create LLaMA 7B configuration pub fn llama_7b() -> Self { Self::default() } /// Create LLaMA 13B configuration pub fn llama_13b() -> Self { Self { hidden_size: 5120, intermediate_size: 13824, num_hidden_layers: 40, num_attention_heads: 40, num_key_value_heads: 40, ..Self::default() } } /// Create LLaMA 30B configuration pub fn llama_30b() -> Self { Self { hidden_size: 6656, intermediate_size: 17920, num_hidden_layers: 60, num_attention_heads: 52, num_key_value_heads: 52, ..Self::default() } } /// Create LLaMA 65B configuration pub fn llama_65b() -> Self { Self { hidden_size: 8_192, intermediate_size: 22016, num_hidden_layers: 80, num_attention_heads: 64, num_key_value_heads: 64, ..Self::default() } } /// Create LLaMA 2 7B configuration pub fn llama2_7b() -> Self { Self { vocab_size: 32000, max_position_embeddings: 4096, ..Self::llama_7b() } } /// Create LLaMA 2 13B configuration pub fn llama2_13b() -> Self { Self { vocab_size: 32000, max_position_embeddings: 4096, ..Self::llama_13b() } } /// Create LLaMA 2 70B configuration with grouped-query attention pub fn llama2_70b() -> Self { Self { vocab_size: 32000, hidden_size: 8_192, intermediate_size: 28672, num_hidden_layers: 80, num_attention_heads: 64, num_key_value_heads: 8, // Grouped-query attention max_position_embeddings: 4096, ..Self::default() } } /// Create Code Llama configuration pub fn code_llama() -> Self { Self { vocab_size: 32016, max_position_embeddings: 16_384, // Longer context for code rope_theta: 1_000_000.0, // Higher rope theta for longer sequences ..Self::llama2_7b() } } /// Validate configuration parameters pub fn validate(&self) -> Result<()> { if self.hidden_size % self.num_attention_heads != 0 { return Err(TransformerError::config( "hidden_size must be divisible by num_attention_heads" )); } if self.num_key_value_heads > self.num_attention_heads { return Err(TransformerError::config( "num_key_value_heads cannot be greater than num_attention_heads" )); } if self.num_attention_heads % self.num_key_value_heads != 0 { return Err(TransformerError::config( "num_attention_heads must be divisible by num_key_value_heads" )); } if self.vocab_size == 0 { return Err(TransformerError::config("vocab_size must be greater than 0")); } if self.num_hidden_layers == 0 { return Err(TransformerError::config("num_hidden_layers must be greater than 0")); } Ok(()) } } /// Rotary Position Embedding (RoPE) implementation #[derive(Debug)] pub struct RotaryPositionEmbedding { /// Dimension per head dim: usize, /// Maximum sequence length max_seq_len: usize, /// Theta parameter theta: f64, /// Precomputed cosine values cos_cached: Tensor, /// Precomputed sine values sin_cached: Tensor, } impl RotaryPositionEmbedding { /// Create new RoPE embeddings pub fn new(dim: usize, max_seq_len: usize, theta: f64, device: &Device) -> Result { let half_dim = dim / 2; // Create frequency inverse (1/theta^(2i/d)) let mut freqs = Vec::with_capacity(half_dim); for i in 0..half_dim { let freq = 1.0 / theta.powf(2.0 * i as f64 / dim as f64); freqs.push(freq as f32); } let freqs_tensor = Tensor::from_data(freqs, &[half_dim], DType::F32)?; // Create position indices let mut positions = Vec::with_capacity(max_seq_len); for i in 0..max_seq_len { positions.push(i as f32); } let positions_tensor = Tensor::from_data(positions, &[max_seq_len], DType::F32)?; // Compute angles: position * freq let angles = positions_tensor.unsqueeze(-1)?.matmul(&freqs_tensor.unsqueeze(0)?)?; // Precompute cosine and sine values let cos_cached = angles.cos()?; let sin_cached = angles.sin()?; Ok(Self { dim, max_seq_len, theta, cos_cached, sin_cached, }) } /// Apply rotary position embedding pub fn forward(&self, q: &Tensor, k: &Tensor, position_ids: &Tensor) -> Result<(Tensor, Tensor)> { let seq_len = position_ids.shape()[1]; if seq_len > self.max_seq_len { return Err(TransformerError::architecture( format!("Sequence length {} exceeds maximum {}", seq_len, self.max_seq_len) )); } // Get cos and sin for current positions let cos = self.cos_cached.slice(&[..seq_len, ..])?; let sin = self.sin_cached.slice(&[..seq_len, ..])?; // Apply rotation to queries and keys let q_rotated = self.apply_rotation(q, &cos, &sin)?; let k_rotated = self.apply_rotation(k, &cos, &sin)?; Ok((q_rotated, k_rotated)) } /// Apply rotation transformation fn apply_rotation(&self, x: &Tensor, cos: &Tensor, sin: &Tensor) -> Result { // Split x into two halves let half_dim = self.dim / 2; let x1 = x.slice(&[.., .., .., ..half_dim])?; let x2 = x.slice(&[.., .., .., half_dim..])?; // Apply rotation: x1 * cos - x2 * sin, x1 * sin + x2 * cos let rotated_x1 = (x1.clone() * cos.clone()) - (x2.clone() * sin.clone()); let rotated_x2 = (x1 * sin.clone()) + (x2 * cos.clone()); // Concatenate back Tensor::cat(&[rotated_x1, rotated_x2], -1) } } /// SwiGLU activation function used in LLaMA #[derive(Debug)] pub struct SwiGLU { /// Gate projection pub gate_proj: Tensor, /// Up projection pub up_proj: Tensor, /// Down projection pub down_proj: Tensor, /// Hidden size hidden_size: usize, /// Intermediate size intermediate_size: usize, } impl SwiGLU { /// Create new SwiGLU layer pub fn new( hidden_size: usize, intermediate_size: usize, bias: bool, device: &Device, ) -> Result { let gate_proj = Tensor::randn(&[intermediate_size, hidden_size], DType::F32, device)? * 0.02; let up_proj = Tensor::randn(&[intermediate_size, hidden_size], DType::F32, device)? * 0.02; let down_proj = Tensor::randn(&[hidden_size, intermediate_size], DType::F32, device)? * 0.02; Ok(Self { gate_proj, up_proj, down_proj, hidden_size, intermediate_size, }) } /// Forward pass: SwiGLU(x) = Swish(gate(x)) * up(x) @ down pub fn forward(&self, x: &Tensor) -> Result { // Gate projection let gate = x.matmul(&self.gate_proj.transpose(-1, -2)?)?; // Up projection let up = x.matmul(&self.up_proj.transpose(-1, -2)?)?; // SwiGLU: Swish(gate) * up let swish_gate = self.swish(&gate)?; let gated = swish_gate * up; // Down projection gated.matmul(&self.down_proj.transpose(-1, -2)?) } /// Swish activation: x * sigmoid(x) fn swish(&self, x: &Tensor) -> Result { // Swish(x) = x * sigmoid(x) = x * (1 / (1 + exp(-x))) // This is also known as SiLU (Sigmoid Linear Unit) // Compute sigmoid(x) = 1 / (1 + exp(-x)) let neg_x = (-x.clone())?; let exp_neg_x = neg_x.exp()?; let one_plus_exp = (exp_neg_x + 1.0)?; let sigmoid_x = (1.0 / one_plus_exp)?; // Compute x * sigmoid(x) (x.clone() * sigmoid_x) .map_err(|e| crate::TransformerError::ArchitectureError( format!("Failed to compute swish activation: {}", e) )) } } /// LLaMA transformer block with RMSNorm and SwiGLU #[derive(Debug)] pub struct LLaMABlock { /// Input RMSNorm pub input_layernorm: RMSNorm, /// Self-attention layer pub self_attn: LLaMAAttention, /// Post-attention RMSNorm pub post_attention_layernorm: RMSNorm, /// MLP layer (SwiGLU) pub mlp: SwiGLU, } impl LLaMABlock { /// Create new LLaMA block pub fn new(config: &LLaMAConfig, device: &Device) -> Result { let input_layernorm = RMSNorm::new(config.hidden_size, config.rms_norm_eps, device)?; let self_attn = LLaMAAttention::new(config, device)?; let post_attention_layernorm = RMSNorm::new(config.hidden_size, config.rms_norm_eps, device)?; let mlp = SwiGLU::new( config.hidden_size, config.intermediate_size, config.mlp_bias, device, )?; Ok(Self { input_layernorm, self_attn, post_attention_layernorm, mlp, }) } /// Forward pass with residual connections pub fn forward(&self, hidden_states: &Tensor, position_ids: &Tensor) -> Result { // Pre-attention RMSNorm let normed_input = self.input_layernorm.forward(hidden_states)?; // Self-attention with residual let attn_output = self.self_attn.forward(&normed_input, position_ids)?; let hidden_states = (hidden_states.clone() + attn_output)?; // Pre-MLP RMSNorm let normed_hidden = self.post_attention_layernorm.forward(&hidden_states)?; // MLP with residual let mlp_output = self.mlp.forward(&normed_hidden)?; let output = (hidden_states + mlp_output)?; Ok(output) } } /// LLaMA attention with grouped-query attention and RoPE #[derive(Debug)] pub struct LLaMAAttention { /// Query projection pub q_proj: Tensor, /// Key projection pub k_proj: Tensor, /// Value projection pub v_proj: Tensor, /// Output projection pub o_proj: Tensor, /// Rotary position embedding pub rotary_emb: RotaryPositionEmbedding, /// Configuration config: LLaMAConfig, } impl LLaMAAttention { /// Create new LLaMA attention pub fn new(config: &LLaMAConfig, device: &Device) -> Result { let head_dim = config.hidden_size / config.num_attention_heads; let q_proj = Tensor::randn( &[config.num_attention_heads * head_dim, config.hidden_size], DType::F32, device, )? * config.initializer_range as f32; let k_proj = Tensor::randn( &[config.num_key_value_heads * head_dim, config.hidden_size], DType::F32, device, )? * config.initializer_range as f32; let v_proj = Tensor::randn( &[config.num_key_value_heads * head_dim, config.hidden_size], DType::F32, device, )? * config.initializer_range as f32; let o_proj = Tensor::randn( &[config.hidden_size, config.num_attention_heads * head_dim], DType::F32, device, )? * config.initializer_range as f32; let rotary_emb = RotaryPositionEmbedding::new( head_dim, config.max_position_embeddings, config.rope_theta, device, )?; Ok(Self { q_proj, k_proj, v_proj, o_proj, rotary_emb, config: config.clone(), }) } /// Forward pass with grouped-query attention pub fn forward(&self, hidden_states: &Tensor, position_ids: &Tensor) -> Result { let batch_size = hidden_states.shape()[0]; let seq_len = hidden_states.shape()[1]; let head_dim = self.config.hidden_size / self.config.num_attention_heads; // Project to Q, K, V let q = hidden_states.matmul(&self.q_proj.transpose(-1, -2)?)?; let k = hidden_states.matmul(&self.k_proj.transpose(-1, -2)?)?; let v = hidden_states.matmul(&self.v_proj.transpose(-1, -2)?)?; // Reshape for multi-head attention let q = q.reshape(&[batch_size, seq_len, self.config.num_attention_heads, head_dim])? .transpose(1, 2)?; let k = k.reshape(&[batch_size, seq_len, self.config.num_key_value_heads, head_dim])? .transpose(1, 2)?; let v = v.reshape(&[batch_size, seq_len, self.config.num_key_value_heads, head_dim])? .transpose(1, 2)?; // Apply rotary position embedding let (q, k) = self.rotary_emb.forward(&q, &k, position_ids)?; // Grouped-query attention: repeat K, V if needed let (k, v) = if self.config.num_key_value_heads < self.config.num_attention_heads { let repeat_factor = self.config.num_attention_heads / self.config.num_key_value_heads; let k = self.repeat_kv(&k, repeat_factor)?; let v = self.repeat_kv(&v, repeat_factor)?; (k, v) } else { (k, v) }; // Scaled dot-product attention let attn_output = self.scaled_dot_product_attention(&q, &k, &v)?; // Reshape and project output let attn_output = attn_output.transpose(1, 2)? .reshape(&[batch_size, seq_len, self.config.hidden_size])?; attn_output.matmul(&self.o_proj.transpose(-1, -2)?) } /// Repeat key-value tensors for grouped-query attention fn repeat_kv(&self, tensor: &Tensor, repeat_factor: usize) -> Result { if repeat_factor == 1 { return Ok(tensor.clone()); } // Efficient key-value repetition for grouped-query attention // Repeat key/value tensors to match the number of query heads let [batch_size, seq_len, n_kv_heads, head_dim] = tensor.shape(); // Calculate repetition factor let n_heads = self.n_heads; let rep_factor = n_heads / n_kv_heads; if rep_factor == 1 { // No repetition needed return Ok(tensor.clone()); } // Reshape to [batch, seq, n_kv_heads, 1, head_dim] let reshaped = tensor.reshape(&[batch_size, seq_len, n_kv_heads, 1, head_dim])?; // Repeat along the new dimension: [batch, seq, n_kv_heads, rep_factor, head_dim] let repeated = reshaped.repeat(&[1, 1, 1, rep_factor, 1])?; // Reshape to [batch, seq, n_heads, head_dim] let output_shape = [batch_size, seq_len, n_heads, head_dim]; repeated.reshape(&output_shape) .map_err(|e| crate::TransformerError::ArchitectureError( format!("Failed to repeat key-value tensor: {}", e) )) } /// Scaled dot-product attention fn scaled_dot_product_attention(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result { let head_dim = self.config.hidden_size / self.config.num_attention_heads; let scale = 1.0 / (head_dim as f32).sqrt(); // QK^T / sqrt(d_k) let scores = q.matmul(&k.transpose(-1, -2)?)? * scale; // Apply causal mask let seq_len = scores.shape()[2]; let causal_mask = self.create_causal_mask(seq_len, scores.device())?; let masked_scores = scores + causal_mask; // Softmax let attn_weights = masked_scores.softmax(-1)?; // Apply to values attn_weights.matmul(v) } /// Create causal attention mask fn create_causal_mask(&self, seq_len: usize, device: &Device) -> Result { // Create lower triangular mask: mask[i,j] = 0 if i >= j, -inf if i < j // This prevents attention to future positions // Create a matrix filled with negative infinity let mut mask_data = vec![-f32::INFINITY; seq_len * seq_len]; // Set lower triangle (including diagonal) to 0 for i in 0..seq_len { for j in 0..=i { // j <= i for lower triangle mask_data[i * seq_len + j] = 0.0; } } // Create tensor from the mask data let mask = Tensor::from_vec( mask_data, &[1, 1, seq_len, seq_len], DType::F32, device, ).map_err(|e| crate::TransformerError::ArchitectureError( format!("Failed to create causal mask: {}", e) ))?; Ok(mask) } } /// Complete LLaMA model #[derive(Debug)] pub struct LLaMAModel { /// Model configuration config: LLaMAConfig, /// Token embeddings embed_tokens: Tensor, /// Stack of LLaMA blocks layers: Vec, /// Final RMSNorm norm: RMSNorm, /// Language modeling head lm_head: Tensor, /// Device device: Device, /// Training mode training: bool, } impl LLaMAModel { /// Create a new LLaMA model pub fn new(config: LLaMAConfig, device: &Device) -> Result { config.validate()?; info!("Creating LLaMA model with config: {:?}", config); // Token embeddings let embed_tokens = Tensor::randn( &[config.vocab_size, config.hidden_size], DType::F32, device, )? * config.initializer_range as f32; // LLaMA blocks let mut layers = Vec::with_capacity(config.num_hidden_layers); for i in 0..config.num_hidden_layers { debug!("Creating LLaMA layer {}/{}", i + 1, config.num_hidden_layers); let layer = LLaMABlock::new(&config, device)?; layers.push(layer); } // Final RMSNorm let norm = RMSNorm::new(config.hidden_size, config.rms_norm_eps, device)?; // Language modeling head let lm_head = if config.tie_word_embeddings { embed_tokens.clone() } else { Tensor::randn( &[config.vocab_size, config.hidden_size], DType::F32, device, )? * config.initializer_range as f32 }; info!("LLaMA model created successfully with {} parameters", Self::count_parameters(&embed_tokens, &layers, &norm, &lm_head)); Ok(Self { config, embed_tokens, layers, norm, lm_head, device: device.clone(), training: false, }) } /// Count total parameters in the model fn count_parameters( embed_tokens: &Tensor, layers: &[LLaMABlock], norm: &RMSNorm, lm_head: &Tensor, ) -> usize { let mut total = 0; // Token embeddings total += embed_tokens.numel(); // LLaMA layers (approximation) total += layers.len() * 2000000; // Placeholder // Final norm total += norm.weight.numel(); // LM head total += lm_head.numel(); total } /// Forward pass pub fn forward(&mut self, input_ids: &Tensor, labels: Option<&Tensor>) -> Result { debug!("LLaMA forward pass: input shape {:?}", input_ids.shape()); let seq_len = input_ids.shape()[1]; // Create position IDs let position_ids = self.create_position_ids(input_ids)?; // Token embeddings let mut hidden_states = self.embed_lookup(&input_ids)?; // Apply LLaMA layers for (i, layer) in self.layers.iter().enumerate() { debug!("Applying LLaMA layer {}", i); hidden_states = layer.forward(&hidden_states, &position_ids)?; } // Final normalization hidden_states = self.norm.forward(&hidden_states)?; // Language modeling head let logits = hidden_states.matmul(&self.lm_head.transpose(-1, -2)?)?; // Compute loss if labels are provided let loss = if let Some(labels) = labels { self.compute_loss(&logits, labels)? } else { None }; Ok(ModelOutput { loss, logits, additional_outputs: HashMap::new(), }) } /// Create position IDs fn create_position_ids(&self, input_ids: &Tensor) -> Result { let batch_size = input_ids.shape()[0]; let seq_len = input_ids.shape()[1]; let mut position_ids = Vec::with_capacity(batch_size * seq_len); for _ in 0..batch_size { for pos in 0..seq_len { position_ids.push(pos as i64); } } Tensor::from_data(position_ids, &[batch_size, seq_len], DType::I64) } /// Embedding lookup fn embed_lookup(&self, input_ids: &Tensor) -> Result { // Efficient embedding lookup using gather operation let batch_size = input_ids.shape()[0]; let seq_len = input_ids.shape()[1]; // In a real implementation, this would use efficient gathering: // 1. Use optimized embedding lookup kernels // 2. Handle out-of-bounds indices gracefully // 3. Support gradient computation for training // For now, create a simplified embedding lookup // Clamp input_ids to valid range let vocab_size = self.embed_tokens.shape()[0]; let clamped_ids = input_ids.clamp(0, vocab_size as i64 - 1)?; // Use indexing to gather embeddings // This is a simplified version - real implementation would be more efficient let mut output_data = Vec::with_capacity(batch_size * seq_len * self.config.hidden_size); // For each position in the input for batch_idx in 0..batch_size { for seq_idx in 0..seq_len { // Get the token ID at this position let token_id = clamped_ids.get_scalar([batch_idx, seq_idx])? as usize; // Get the embedding vector for this token let embedding = self.embed_tokens.slice(&[token_id, ..])?; // Add to output for embed_dim in 0..self.config.hidden_size { let val = embedding.get_scalar([embed_dim])?; output_data.push(val); } } } // Create output tensor Tensor::from_vec( output_data, &[batch_size, seq_len, self.config.hidden_size], DType::F32, input_ids.device(), ).map_err(|e| crate::TransformerError::ArchitectureError( format!("Failed to create embedding lookup result: {}", e) )) } /// Compute cross-entropy loss for language modeling fn compute_loss(&self, logits: &Tensor, labels: &Tensor) -> Result { debug!("Computing cross-entropy loss"); // Cross-entropy loss for language modeling // logits: [batch_size, seq_len, vocab_size] // labels: [batch_size, seq_len] let batch_size = logits.shape()[0]; let seq_len = logits.shape()[1]; let vocab_size = logits.shape()[2]; // Flatten tensors for loss computation let logits_flat = logits.reshape(&[batch_size * seq_len, vocab_size])?; let labels_flat = labels.reshape(&[batch_size * seq_len])?; // Apply softmax to get probabilities let log_probs = self.log_softmax(&logits_flat)?; // Compute negative log likelihood let mut total_loss = 0.0f32; let mut num_tokens = 0; for i in 0..(batch_size * seq_len) { let target_id = labels_flat.get_scalar([i])? as usize; // Skip padding tokens (assume -100 or similar sentinel value) if target_id == usize::MAX || target_id >= vocab_size { continue; } let log_prob = log_probs.get_scalar([i, target_id])?; total_loss -= log_prob; num_tokens += 1; } // Average loss over valid tokens let avg_loss = if num_tokens > 0 { total_loss / num_tokens as f32 } else { 0.0 }; Tensor::scalar(avg_loss, logits.dtype(), logits.device()) .map_err(|e| crate::TransformerError::ArchitectureError( format!("Failed to compute cross-entropy loss: {}", e) )) } /// Compute log softmax for numerical stability fn log_softmax(&self, logits: &Tensor) -> Result { // log_softmax(x) = x - log(sum(exp(x))) // For numerical stability: log_softmax(x) = x - max(x) - log(sum(exp(x - max(x)))) let max_logits = logits.max_keepdim(-1)?; let shifted_logits = (logits.clone() - max_logits.clone())?; let exp_shifted = shifted_logits.exp()?; let sum_exp = exp_shifted.sum_keepdim(-1)?; let log_sum_exp = sum_exp.log()?; (logits.clone() - max_logits - log_sum_exp) .map_err(|e| crate::TransformerError::ArchitectureError( format!("Failed to compute log softmax: {}", e) )) } /// Generate text using the model pub fn generate( &mut self, input_ids: &Tensor, max_length: usize, temperature: f32, do_sample: bool, top_k: Option, top_p: Option, ) -> Result { self.set_training(false); let mut current_ids = input_ids.clone(); let batch_size = input_ids.shape()[0]; let initial_length = input_ids.shape()[1]; for step in 0..(max_length - initial_length) { debug!("Generation step {}/{}", step + 1, max_length - initial_length); // Forward pass let output = self.forward(¤t_ids, None)?; let mut logits = output.logits; // Get logits for the last token logits = logits.slice(&[.., -1, ..])?; // Apply temperature if temperature != 1.0 { logits = logits / temperature; } // Apply top-k filtering if let Some(k) = top_k { logits = self.top_k_filtering(&logits, k)?; } // Apply top-p (nucleus) filtering if let Some(p) = top_p { logits = self.top_p_filtering(&logits, p)?; } // Sample next token let next_token = if do_sample { self.sample_from_logits(&logits)? } else { self.greedy_from_logits(&logits)? }; // Append next token current_ids = Tensor::cat(&[current_ids, next_token.unsqueeze(-1)], -1)?; // Check for EOS token // TODO: Implement proper stopping criteria } Ok(current_ids) } /// Top-k filtering: Keep only the k largest logits, set others to -inf fn top_k_filtering(&self, logits: &Tensor, k: usize) -> Result { if k >= logits.shape().last().unwrap() { // k is larger than vocab size, no filtering needed return Ok(logits.clone()); } let vocab_size = logits.shape().last().unwrap(); let batch_size = logits.numel() / vocab_size; // Get flattened view for processing let flat_logits = logits.flatten()?; let mut filtered_data = vec![-f32::INFINITY; flat_logits.numel()]; // Process each sequence in the batch for batch_idx in 0..batch_size { let start_idx = batch_idx * vocab_size; let end_idx = start_idx + vocab_size; // Extract logits for this sequence let mut seq_logits: Vec<(f32, usize)> = Vec::with_capacity(vocab_size); for i in start_idx..end_idx { let value = flat_logits.get_scalar([i])?; seq_logits.push((value, i)); } // Sort by logits value (descending) seq_logits.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap()); // Keep only top-k values for (logit_val, orig_idx) in seq_logits.into_iter().take(k) { filtered_data[orig_idx] = logit_val; } } // Create filtered tensor Tensor::from_vec( filtered_data, logits.shape(), logits.dtype(), logits.device(), ).map_err(|e| crate::TransformerError::ArchitectureError( format!("Failed to apply top-k filtering: {}", e) )) } /// Top-p (nucleus) filtering: Keep tokens with cumulative probability <= p fn top_p_filtering(&self, logits: &Tensor, p: f32) -> Result { if p >= 1.0 { // No filtering needed return Ok(logits.clone()); } let vocab_size = logits.shape().last().unwrap(); let batch_size = logits.numel() / vocab_size; // Get flattened view for processing let flat_logits = logits.flatten()?; let mut filtered_data = vec![-f32::INFINITY; flat_logits.numel()]; // Process each sequence in the batch for batch_idx in 0..batch_size { let start_idx = batch_idx * vocab_size; let end_idx = start_idx + vocab_size; // Extract logits for this sequence let mut seq_logits: Vec<(f32, usize)> = Vec::with_capacity(vocab_size); for i in start_idx..end_idx { let value = flat_logits.get_scalar([i])?; seq_logits.push((value, i)); } // Sort by logits value (descending) seq_logits.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap()); // Convert to probabilities using softmax let max_logit = seq_logits[0].0; let mut probs: Vec<(f32, usize)> = seq_logits .iter() .map(|(logit, idx)| { let exp_logit = (logit - max_logit).exp(); (exp_logit, *idx) }) .collect(); // Normalize probabilities let sum_probs: f32 = probs.iter().map(|(prob, _)| prob).sum(); for (prob, _) in probs.iter_mut() { *prob /= sum_probs; } // Calculate cumulative probabilities and filter let mut cumulative_prob = 0.0; for (prob, orig_idx) in probs { cumulative_prob += prob; if cumulative_prob <= p { // Keep this token filtered_data[orig_idx] = flat_logits.get_scalar([orig_idx])?; } else { // Stop adding tokens break; } } } // Create filtered tensor Tensor::from_vec( filtered_data, logits.shape(), logits.dtype(), logits.device(), ).map_err(|e| crate::TransformerError::ArchitectureError( format!("Failed to apply top-p filtering: {}", e) )) } /// Sample from logits distribution using multinomial sampling fn sample_from_logits(&self, logits: &Tensor) -> Result { // Convert logits to probabilities using softmax let probs = self.softmax(logits)?; // Sample from multinomial distribution self.multinomial_sample(&probs) } /// Apply softmax to convert logits to probabilities fn softmax(&self, logits: &Tensor) -> Result { // Subtract max for numerical stability: softmax(x - max(x)) = softmax(x) let max_logits = logits.max_keepdim(-1)?; let shifted_logits = (logits.clone() - max_logits)?; // Compute exponentials let exp_logits = shifted_logits.exp()?; // Sum along last dimension let sum_exp = exp_logits.sum_keepdim(-1)?; // Normalize (exp_logits / sum_exp) .map_err(|e| crate::TransformerError::ArchitectureError( format!("Failed to compute softmax: {}", e) )) } /// Sample from multinomial distribution fn multinomial_sample(&self, probs: &Tensor) -> Result { use rand::Rng; let vocab_size = probs.shape().last().unwrap(); let batch_size = probs.numel() / vocab_size; let mut rng = rand::thread_rng(); let flat_probs = probs.flatten()?; let mut sampled_tokens = Vec::with_capacity(batch_size); // Sample for each sequence in the batch for batch_idx in 0..batch_size { let start_idx = batch_idx * vocab_size; let end_idx = start_idx + vocab_size; // Extract probabilities for this sequence let mut seq_probs = Vec::with_capacity(vocab_size); for i in start_idx..end_idx { let prob = flat_probs.get_scalar([i])?; seq_probs.push(prob); } // Ensure probabilities sum to 1 (numerical precision) let prob_sum: f32 = seq_probs.iter().sum(); if prob_sum > 0.0 { for prob in seq_probs.iter_mut() { *prob /= prob_sum; } } else { // All probabilities are zero, fallback to uniform for prob in seq_probs.iter_mut() { *prob = 1.0 / vocab_size as f32; } } // Sample using cumulative distribution let random_val: f32 = rng.gen(); let mut cumulative_prob = 0.0; let mut selected_token = 0; for (token_id, &prob) in seq_probs.iter().enumerate() { cumulative_prob += prob; if random_val <= cumulative_prob { selected_token = token_id; break; } } sampled_tokens.push(selected_token as f32); } // Create output tensor let output_shape = if batch_size == 1 { vec![1] } else { vec![batch_size, 1] }; Tensor::from_vec( sampled_tokens, &output_shape, DType::F32, probs.device(), ).map_err(|e| crate::TransformerError::ArchitectureError( format!("Failed to sample from multinomial distribution: {}", e) )) } /// Greedy selection from logits fn greedy_from_logits(&self, logits: &Tensor) -> Result { // TODO: Implement argmax operation Tensor::zeros_typed(&[logits.shape()[0]], DType::I64, logits.device()) } } impl TransformerModel for LLaMAModel { fn forward(&mut self, input_ids: &Tensor, labels: Option<&Tensor>) -> Result { self.forward(input_ids, labels) } fn parameters(&self) -> HashMap { let mut params = HashMap::new(); // Embeddings params.insert("model.embed_tokens.weight".to_string(), self.embed_tokens.clone()); // Layers (simplified) for (i, _layer) in self.layers.iter().enumerate() { // TODO: Add actual layer parameters params.insert(format!("model.layers.{}.placeholder", i), self.embed_tokens.clone()); // Placeholder } // Final norm params.insert("model.norm.weight".to_string(), self.norm.weight.clone()); // LM head params.insert("lm_head.weight".to_string(), self.lm_head.clone()); params } fn update_parameters(&mut self, updates: &HashMap) -> Result<()> { for (name, update) in updates { match name.as_str() { "model.embed_tokens.weight" => { self.embed_tokens = update.clone(); } "model.norm.weight" => { self.norm.weight = update.clone(); } "lm_head.weight" => { self.lm_head = update.clone(); } _ => { debug!("Updating parameter: {}", name); } } } Ok(()) } fn config(&self) -> ModelConfig { ModelConfig { model_type: "LLaMA".to_string(), num_parameters: Self::count_parameters( &self.embed_tokens, &self.layers, &self.norm, &self.lm_head, ), dtype: DType::F32, config: HashMap::new(), } } fn set_training(&mut self, training: bool) { self.training = training; debug!("Set LLaMA training mode: {}", training); } fn memory_stats(&self) -> HashMap { let mut stats = HashMap::new(); stats.insert("num_hidden_layers".to_string(), self.config.num_hidden_layers); stats.insert("hidden_size".to_string(), self.config.hidden_size); stats.insert("vocab_size".to_string(), self.config.vocab_size); stats.insert("num_attention_heads".to_string(), self.config.num_attention_heads); stats.insert("num_key_value_heads".to_string(), self.config.num_key_value_heads); stats } } impl TransformerArchitecture for LLaMAModel { fn forward(&self, input: &Tensor) -> Result { Ok(input.clone()) } fn architecture_type(&self) -> &'static str { "LLaMA" } fn device(&self) -> &Device { &self.device } fn parameters(&self) -> Vec<&Tensor> { vec![&self.embed_tokens, &self.lm_head] } fn parameters_mut(&mut self) -> Vec<&mut Tensor> { vec![&mut self.embed_tokens, &mut self.lm_head] } fn config(&self) -> &TransformerConfig { &self.config.base } } #[cfg(test)] mod tests { use super::*; use rtx_tensor::Device; #[test] fn test_llama_config_validation() { let mut config = LLaMAConfig::default(); assert!(config.validate().is_ok()); // Test invalid configuration config.hidden_size = 100; config.num_attention_heads = 7; // 100 is not divisible by 7 assert!(config.validate().is_err()); } #[test] fn test_llama_config_presets() { let llama_7b = LLaMAConfig::llama_7b(); assert_eq!(llama_7b.hidden_size, 4096); assert_eq!(llama_7b.num_hidden_layers, 32); let llama2_70b = LLaMAConfig::llama2_70b(); assert_eq!(llama2_70b.num_key_value_heads, 8); // Grouped-query attention let code_llama = LLaMAConfig::code_llama(); assert_eq!(code_llama.max_position_embeddings, 16_384); } #[test] fn test_rope_creation() { let device = Device::Cpu; let rope = RotaryPositionEmbedding::new(128, 2048, 10000.0, &device); assert!(rope.is_ok()); let rope = rope.unwrap(); assert_eq!(rope.dim, 128); assert_eq!(rope.max_seq_len, 2048); } #[test] fn test_swiglu_creation() { let device = Device::Cpu; let swiglu = SwiGLU::new(4096, 11008, false, &device); assert!(swiglu.is_ok()); let swiglu = swiglu.unwrap(); assert_eq!(swiglu.hidden_size, 4096); assert_eq!(swiglu.intermediate_size, 11008); } #[test] fn test_llama_model_creation() { let config = LLaMAConfig::llama_7b(); let device = Device::Cpu; let model = LLaMAModel::new(config, &device); assert!(model.is_ok()); let model = model.unwrap(); assert_eq!(model.architecture_type(), "LLaMA"); assert_eq!(model.config().model_type, "LLaMA"); } #[test] fn test_grouped_query_attention() { let mut config = LLaMAConfig::llama2_70b(); assert_eq!(config.num_attention_heads, 64); assert_eq!(config.num_key_value_heads, 8); assert!(config.validate().is_ok()); // Test invalid GQA configuration config.num_key_value_heads = 65; // Greater than attention heads assert!(config.validate().is_err()); } }