Files
rustytorch/archive/legacy_files_backup/llama_attention_broken.rs
T
2026-03-04 00:08:42 +00:00

619 lines
20 KiB
Rust

//! LLaMA Attention Mechanisms and Components
//!
//! Specialized attention implementations for LLaMA architecture including
//! rotary position embeddings, grouped-query attention, and SwiGLU.
use crate::architectures::TransformerConfig;
use crate::{Result, TransformerError};
use rtx_tensor::{Tensor, Device, DType};
use serde::{Deserialize, Serialize};
use tracing::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<RopeScaling>,
/// 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
pub dim: usize,
/// Maximum sequence length
pub max_seq_len: usize,
/// Theta parameter
pub theta: f64,
/// Precomputed cosine values
pub cos_cached: Tensor,
/// Precomputed sine values
pub sin_cached: Tensor,
}
impl RotaryPositionEmbedding {
/// Create new RoPE embeddings
pub fn new(dim: usize, max_seq_len: usize, theta: f64, device: &Device) -> Result<Self> {
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, vec![half_dim], &Device::Cpu)?;
// 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, vec![max_seq_len], &Device::Cpu)?;
// 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(0, 0, seq_len)?; // TODO: Fix slice API
let sin = self.sin_cached.slice(0, 0, seq_len)?; // TODO: Fix slice API
// 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<Tensor> {
// Split x into two halves
let half_dim = self.dim / 2;
// TODO: Implement proper tensor slicing for rotary embeddings
let x1 = x.clone(); // Placeholder
let x2 = x.clone(); // Placeholder
// 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
pub hidden_size: usize,
/// Intermediate size
pub intermediate_size: usize,
}
impl SwiGLU {
/// Create new SwiGLU layer
pub fn new(
hidden_size: usize,
intermediate_size: usize,
bias: bool,
device: &Device,
) -> Result<Self> {
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<Tensor> {
// 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<Tensor> {
// 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 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
pub config: LLaMAConfig,
}
impl LLaMAAttention {
/// Create new LLaMA attention
pub fn new(config: &LLaMAConfig, device: &Device) -> Result<Self> {
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<Tensor> {
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<Tensor> {
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<Tensor> {
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<Tensor> {
// 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)
}
}
#[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_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());
}
}