Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,761 @@
|
||||
//! `LLaMA` Attention Mechanisms and Components (Simplified Stub)
|
||||
//!
|
||||
//! This is a simplified stub version to get compilation working.
|
||||
//! TODO: Implement full `LLaMA` attention with rotary embeddings and grouped-query attention.
|
||||
|
||||
use crate::architectures::TransformerConfig;
|
||||
use crate::tensor_bridge::TensorBridge;
|
||||
use crate::{Result, TransformerError};
|
||||
use rtx_tensor::{Device, Tensor};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// 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
|
||||
#[must_use]
|
||||
pub fn llama_7b() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Create `LLaMA` 13B configuration
|
||||
#[must_use]
|
||||
pub fn llama_13b() -> Self {
|
||||
let mut config = Self::default();
|
||||
config.hidden_size = 5120;
|
||||
config.intermediate_size = 13824;
|
||||
config.num_hidden_layers = 40;
|
||||
config.num_attention_heads = 40;
|
||||
config.num_key_value_heads = 40;
|
||||
config
|
||||
}
|
||||
|
||||
/// Create `LLaMA` 30B configuration
|
||||
#[must_use]
|
||||
pub fn llama_30b() -> Self {
|
||||
let mut config = Self::default();
|
||||
config.hidden_size = 6656;
|
||||
config.intermediate_size = 17920;
|
||||
config.num_hidden_layers = 60;
|
||||
config.num_attention_heads = 52;
|
||||
config.num_key_value_heads = 52;
|
||||
config
|
||||
}
|
||||
|
||||
/// Create `LLaMA` 65B configuration
|
||||
#[must_use]
|
||||
pub fn llama_65b() -> Self {
|
||||
let mut config = Self::default();
|
||||
config.hidden_size = 8192;
|
||||
config.intermediate_size = 22016;
|
||||
config.num_hidden_layers = 80;
|
||||
config.num_attention_heads = 64;
|
||||
config.num_key_value_heads = 64;
|
||||
config
|
||||
}
|
||||
|
||||
/// Create Code Llama configuration
|
||||
pub fn code_llama(base_size: &str) -> Result<Self> {
|
||||
let mut config = match base_size {
|
||||
"7b" => Self::llama_7b(),
|
||||
"13b" => Self::llama_13b(),
|
||||
"34b" => {
|
||||
let mut config = Self::default();
|
||||
config.hidden_size = 8192;
|
||||
config.intermediate_size = 22016;
|
||||
config.num_hidden_layers = 48;
|
||||
config.num_attention_heads = 64;
|
||||
config.num_key_value_heads = 8; // Grouped-query attention
|
||||
config.max_position_embeddings = 16384; // Extended context
|
||||
config
|
||||
}
|
||||
_ => {
|
||||
return Err(TransformerError::config(format!(
|
||||
"Unsupported Code Llama size: {base_size}"
|
||||
)));
|
||||
}
|
||||
};
|
||||
|
||||
// Code Llama specific adjustments
|
||||
config.vocab_size = 32016; // Extended vocabulary
|
||||
config.rope_theta = 1000000.0; // Extended RoPE
|
||||
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
/// Validate configuration parameters
|
||||
pub fn validate(&self) -> Result<()> {
|
||||
if self.vocab_size == 0 {
|
||||
return Err(TransformerError::Generic(
|
||||
"vocab_size must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.hidden_size == 0 {
|
||||
return Err(TransformerError::Generic(
|
||||
"hidden_size must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.num_hidden_layers == 0 {
|
||||
return Err(TransformerError::Generic(
|
||||
"num_hidden_layers must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.num_attention_heads == 0 {
|
||||
return Err(TransformerError::Generic(
|
||||
"num_attention_heads must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if !self.hidden_size.is_multiple_of(self.num_attention_heads) {
|
||||
return Err(TransformerError::Generic(
|
||||
"hidden_size must be divisible by num_attention_heads".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Rotary Position Embedding (`RoPE`) implementation
|
||||
#[derive(Debug)]
|
||||
pub struct RotaryPositionEmbedding {
|
||||
dim: usize,
|
||||
max_seq_len: usize,
|
||||
theta: f64,
|
||||
device: Device,
|
||||
// Precomputed frequency tensors for efficiency
|
||||
cos_cached: Option<Tensor>,
|
||||
sin_cached: Option<Tensor>,
|
||||
}
|
||||
|
||||
impl RotaryPositionEmbedding {
|
||||
/// Create new `RoPE` layer
|
||||
pub fn new(dim: usize, max_seq_len: usize, theta: f64, device: &Device) -> Result<Self> {
|
||||
let mut rope = Self {
|
||||
dim,
|
||||
max_seq_len,
|
||||
theta,
|
||||
device: device.clone(),
|
||||
cos_cached: None,
|
||||
sin_cached: None,
|
||||
};
|
||||
|
||||
// Precompute cos/sin values for efficiency
|
||||
rope.precompute_freqs(max_seq_len)?;
|
||||
Ok(rope)
|
||||
}
|
||||
|
||||
/// Precompute frequency tensors
|
||||
fn precompute_freqs(&mut self, seq_len: usize) -> Result<()> {
|
||||
// Create frequency tensor: theta^(-2k/dim) for k = 0, 1, ..., dim//2-1
|
||||
let half_dim = self.dim / 2;
|
||||
let freqs: Vec<f32> = (0..half_dim)
|
||||
.map(|i| 1.0 / self.theta.powf(2.0 * i as f64 / self.dim as f64) as f32)
|
||||
.collect();
|
||||
|
||||
let freq_tensor = Tensor::from_vec(freqs, &[half_dim], &self.device)?;
|
||||
|
||||
// Create position tensor: [0, 1, 2, ..., seq_len-1]
|
||||
let positions: Vec<f32> = (0..seq_len).map(|i| i as f32).collect();
|
||||
let pos_tensor = Tensor::from_vec(positions, &[seq_len], &self.device)?;
|
||||
|
||||
// Compute outer product: pos[i] * freq[j] for all i, j
|
||||
let pos_expanded = pos_tensor.unsqueeze(1)?; // [seq_len, 1]
|
||||
let freq_expanded = freq_tensor.unsqueeze(0)?; // [1, half_dim]
|
||||
let angles = pos_expanded.matmul(&freq_expanded)?; // [seq_len, half_dim]
|
||||
|
||||
// Compute cos and sin
|
||||
self.cos_cached = Some(angles.cos()?);
|
||||
self.sin_cached = Some(angles.sin()?);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Apply rotary position encoding to queries and keys
|
||||
pub fn forward(&self, q: &Tensor, k: &Tensor, seq_len: usize) -> Result<(Tensor, Tensor)> {
|
||||
if seq_len > self.max_seq_len {
|
||||
return Err(TransformerError::architecture(format!(
|
||||
"Sequence length {} exceeds maximum {}",
|
||||
seq_len, self.max_seq_len
|
||||
)));
|
||||
}
|
||||
|
||||
// Get cached cos/sin values
|
||||
let cos = self.cos_cached.as_ref().ok_or_else(|| {
|
||||
TransformerError::architecture("RoPE frequencies not precomputed".to_string())
|
||||
})?;
|
||||
let sin = self.sin_cached.as_ref().ok_or_else(|| {
|
||||
TransformerError::architecture("RoPE frequencies not precomputed".to_string())
|
||||
})?;
|
||||
|
||||
// Apply rotary position encoding
|
||||
let q_rot = self.apply_rotary_pos_emb(q, cos, sin, seq_len)?;
|
||||
let k_rot = self.apply_rotary_pos_emb(k, cos, sin, seq_len)?;
|
||||
|
||||
Ok((q_rot, k_rot))
|
||||
}
|
||||
|
||||
/// Apply rotary position embedding to a tensor
|
||||
fn apply_rotary_pos_emb(
|
||||
&self,
|
||||
x: &Tensor,
|
||||
cos: &Tensor,
|
||||
sin: &Tensor,
|
||||
seq_len: usize,
|
||||
) -> Result<Tensor> {
|
||||
// Extract shape: [batch_size, seq_len, num_heads, head_dim]
|
||||
let shape = x.shape();
|
||||
let batch_size = shape[0];
|
||||
let input_seq_len = shape[1];
|
||||
let num_heads = shape[2];
|
||||
let head_dim = shape[3];
|
||||
|
||||
let actual_seq_len = input_seq_len.min(seq_len);
|
||||
|
||||
// Reshape x to [batch_size, seq_len, num_heads, head_dim/2, 2]
|
||||
let half_dim = head_dim / 2;
|
||||
let x_reshaped = x.view([batch_size, actual_seq_len, num_heads, half_dim, 2])?;
|
||||
|
||||
// Split into x1 and x2 (real and imaginary parts)
|
||||
let x1 = x_reshaped.slice(4, 0, 1)?; // [..., 0]
|
||||
let x2 = x_reshaped.slice(4, 1, 2)?; // [..., 1]
|
||||
|
||||
// Get cos/sin for the sequence length
|
||||
let cos_seq = cos.slice(0, 0, actual_seq_len)?; // [actual_seq_len, half_dim]
|
||||
let sin_seq = sin.slice(0, 0, actual_seq_len)?; // [actual_seq_len, half_dim]
|
||||
|
||||
// Expand cos/sin to match tensor dimensions
|
||||
// [actual_seq_len, half_dim] -> [1, actual_seq_len, 1, half_dim, 1]
|
||||
let cos_expanded = cos_seq.view([1, actual_seq_len, 1, half_dim, 1])?;
|
||||
let sin_expanded = sin_seq.view([1, actual_seq_len, 1, half_dim, 1])?;
|
||||
|
||||
// Apply rotation:
|
||||
// x_rot1 = x1 * cos - x2 * sin
|
||||
// x_rot2 = x1 * sin + x2 * cos
|
||||
let x_rot1 = x1.mul(&cos_expanded)?.sub(&x2.mul(&sin_expanded)?)?;
|
||||
let x_rot2 = x1.mul(&sin_expanded)?.add(&x2.mul(&cos_expanded)?)?;
|
||||
|
||||
// Concatenate rotated parts
|
||||
let x_rot = Tensor::cat(&[x_rot1, x_rot2], 4)?;
|
||||
|
||||
// Reshape back to original shape
|
||||
x_rot
|
||||
.view([batch_size, actual_seq_len, num_heads, head_dim])
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
/// `LLaMA` Multi-Head Attention with `RoPE` and GQA
|
||||
#[derive(Debug)]
|
||||
pub struct LLaMAAttention {
|
||||
num_heads: usize,
|
||||
num_key_value_heads: usize,
|
||||
head_dim: usize,
|
||||
hidden_size: usize,
|
||||
rope: RotaryPositionEmbedding,
|
||||
// Linear projection weights
|
||||
q_proj: Tensor,
|
||||
k_proj: Tensor,
|
||||
v_proj: Tensor,
|
||||
o_proj: Tensor,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl LLaMAAttention {
|
||||
/// Create new `LLaMA` attention layer
|
||||
pub fn new(config: &LLaMAConfig, device: &Device) -> Result<Self> {
|
||||
let head_dim = config.hidden_size / config.num_attention_heads;
|
||||
|
||||
let rope = RotaryPositionEmbedding::new(
|
||||
head_dim,
|
||||
config.max_position_embeddings,
|
||||
config.rope_theta,
|
||||
device,
|
||||
)?;
|
||||
|
||||
// Initialize projection weights
|
||||
let scale = (config.hidden_size as f32).recip().sqrt();
|
||||
|
||||
// Query projection: [hidden_size, num_heads * head_dim]
|
||||
let q_proj = Tensor::randn(
|
||||
&[config.hidden_size, config.num_attention_heads * head_dim],
|
||||
device,
|
||||
)?
|
||||
.mul_scalar(scale)?;
|
||||
|
||||
// Key and Value projections for GQA: [hidden_size, num_kv_heads * head_dim]
|
||||
let k_proj = Tensor::randn(
|
||||
&[config.hidden_size, config.num_key_value_heads * head_dim],
|
||||
device,
|
||||
)?
|
||||
.mul_scalar(scale)?;
|
||||
let v_proj = Tensor::randn(
|
||||
&[config.hidden_size, config.num_key_value_heads * head_dim],
|
||||
device,
|
||||
)?
|
||||
.mul_scalar(scale)?;
|
||||
|
||||
// Output projection: [num_heads * head_dim, hidden_size]
|
||||
let o_proj = Tensor::randn(
|
||||
&[config.num_attention_heads * head_dim, config.hidden_size],
|
||||
device,
|
||||
)?
|
||||
.mul_scalar(scale)?;
|
||||
|
||||
Ok(Self {
|
||||
num_heads: config.num_attention_heads,
|
||||
num_key_value_heads: config.num_key_value_heads,
|
||||
head_dim,
|
||||
hidden_size: config.hidden_size,
|
||||
rope,
|
||||
q_proj,
|
||||
k_proj,
|
||||
v_proj,
|
||||
o_proj,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Forward pass of attention
|
||||
pub fn forward(&self, hidden_states: &Tensor, position_ids: Option<&Tensor>) -> Result<Tensor> {
|
||||
let batch_size = hidden_states.shape()[0];
|
||||
let seq_len = hidden_states.shape()[1];
|
||||
|
||||
// 1. Linear projections
|
||||
let q = self.compute_queries(hidden_states)?;
|
||||
let k = self.compute_keys(hidden_states)?;
|
||||
let v = self.compute_values(hidden_states)?;
|
||||
|
||||
// Reshape for multi-head attention
|
||||
// Q: [batch, seq, num_heads, head_dim]
|
||||
let q = q.view([batch_size, seq_len, self.num_heads, self.head_dim])?;
|
||||
|
||||
// K, V: [batch, seq, num_kv_heads, head_dim]
|
||||
let k = k.view([batch_size, seq_len, self.num_key_value_heads, self.head_dim])?;
|
||||
let v = v.view([batch_size, seq_len, self.num_key_value_heads, self.head_dim])?;
|
||||
|
||||
// 2. Apply rotary position encoding
|
||||
let (q, k) = self.rope.forward(&q, &k, seq_len)?;
|
||||
|
||||
// 3. Expand K, V for grouped-query attention if needed
|
||||
let (k, v) = if self.num_key_value_heads < self.num_heads {
|
||||
let repeat_factor = self.num_heads / self.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)
|
||||
};
|
||||
|
||||
// 4. Compute scaled dot-product attention
|
||||
let attention_output = self.scaled_dot_product_attention(&q, &k, &v)?;
|
||||
|
||||
// 5. Reshape and apply output projection
|
||||
let attention_output =
|
||||
attention_output.view([batch_size, seq_len, self.num_heads * self.head_dim])?;
|
||||
let output = attention_output.matmul(&self.o_proj)?;
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
/// Compute query projections
|
||||
fn compute_queries(&self, hidden_states: &Tensor) -> Result<Tensor> {
|
||||
hidden_states
|
||||
.matmul(&self.q_proj)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||||
}
|
||||
|
||||
/// Compute key projections
|
||||
fn compute_keys(&self, hidden_states: &Tensor) -> Result<Tensor> {
|
||||
hidden_states
|
||||
.matmul(&self.k_proj)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||||
}
|
||||
|
||||
/// Compute value projections
|
||||
fn compute_values(&self, hidden_states: &Tensor) -> Result<Tensor> {
|
||||
hidden_states
|
||||
.matmul(&self.v_proj)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||||
}
|
||||
|
||||
/// 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());
|
||||
}
|
||||
|
||||
let shape = tensor.shape();
|
||||
let batch_size = shape[0];
|
||||
let seq_len = shape[1];
|
||||
let num_kv_heads = shape[2];
|
||||
let head_dim = shape[3];
|
||||
|
||||
// Expand tensor by repeating along the head dimension
|
||||
// [batch, seq, num_kv_heads, head_dim] -> [batch, seq, num_kv_heads, repeat_factor, head_dim]
|
||||
let expanded = tensor.unsqueeze(3)?; // [batch, seq, num_kv_heads, 1, head_dim]
|
||||
let repeated =
|
||||
expanded.expand(&[batch_size, seq_len, num_kv_heads, repeat_factor, head_dim])?;
|
||||
|
||||
// Reshape to [batch, seq, num_kv_heads * repeat_factor, head_dim]
|
||||
repeated
|
||||
.view([batch_size, seq_len, num_kv_heads * repeat_factor, head_dim])
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||||
}
|
||||
|
||||
/// Compute scaled dot-product attention
|
||||
fn scaled_dot_product_attention(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result<Tensor> {
|
||||
let scale = (self.head_dim as f32).recip().sqrt();
|
||||
|
||||
// Transpose for attention computation: [batch, seq, heads, dim] -> [batch, heads, seq, dim]
|
||||
let q = q.transpose(1, 2)?;
|
||||
let k = k.transpose(1, 2)?;
|
||||
let v = v.transpose(1, 2)?;
|
||||
|
||||
// Compute attention scores: Q @ K^T / sqrt(d_k)
|
||||
let k_transposed = k.transpose(-2, -1)?;
|
||||
let attention_scores = q.matmul(&k_transposed)?.mul_scalar(scale)?;
|
||||
|
||||
// Apply causal mask (lower triangular)
|
||||
let seq_len = attention_scores.shape()[2];
|
||||
let causal_mask = self.create_causal_mask(seq_len)?;
|
||||
let masked_scores = attention_scores.add(&causal_mask)?;
|
||||
|
||||
// Apply softmax
|
||||
let attention_probs = masked_scores.softmax(-1)?;
|
||||
|
||||
// Apply attention to values: attention_probs @ V
|
||||
let attention_output = attention_probs.matmul(&v)?;
|
||||
|
||||
// Transpose back: [batch, heads, seq, dim] -> [batch, seq, heads, dim]
|
||||
attention_output
|
||||
.transpose(1, 2)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||||
}
|
||||
|
||||
/// Create causal attention mask
|
||||
fn create_causal_mask(&self, seq_len: usize) -> Result<Tensor> {
|
||||
// Create a lower triangular mask filled with 0s and -inf
|
||||
let mut mask_data = Vec::with_capacity(seq_len * seq_len);
|
||||
|
||||
for i in 0..seq_len {
|
||||
for j in 0..seq_len {
|
||||
if j <= i {
|
||||
mask_data.push(0.0f32); // Allow attention
|
||||
} else {
|
||||
mask_data.push(f32::NEG_INFINITY); // Mask future positions
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Tensor::from_vec(mask_data, &[seq_len, seq_len], &self.device)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
/// `RMSNorm` layer for `LLaMA`
|
||||
#[derive(Debug)]
|
||||
pub struct RMSNorm {
|
||||
hidden_size: usize,
|
||||
eps: f64,
|
||||
weight: Tensor,
|
||||
}
|
||||
|
||||
impl RMSNorm {
|
||||
/// Create new `RMSNorm` layer
|
||||
pub fn new(hidden_size: usize, eps: f64, device: &Device) -> Result<Self> {
|
||||
// Initialize weight to ones
|
||||
let weight = Tensor::ones([hidden_size], device)?;
|
||||
|
||||
Ok(Self {
|
||||
hidden_size,
|
||||
eps,
|
||||
weight,
|
||||
})
|
||||
}
|
||||
|
||||
/// Get the weight tensor (for parameter access)
|
||||
#[must_use]
|
||||
pub fn weight(&self) -> &Tensor {
|
||||
&self.weight
|
||||
}
|
||||
|
||||
/// Forward pass of `RMSNorm`
|
||||
pub fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
|
||||
// RMSNorm: x / sqrt(mean(x^2) + eps) * weight
|
||||
|
||||
// Compute x^2
|
||||
let x_squared = hidden_states.mul(hidden_states)?;
|
||||
|
||||
// Compute mean along last dimension (hidden_size)
|
||||
let mean_squared = x_squared.mean_dim(-1, true)?; // Keep dim for broadcasting
|
||||
|
||||
// Add epsilon for numerical stability
|
||||
let variance = mean_squared.add_scalar(self.eps as f32)?;
|
||||
|
||||
// Compute 1 / sqrt(variance)
|
||||
let inv_std = variance.rsqrt()?;
|
||||
|
||||
// Normalize
|
||||
let normalized = hidden_states.mul(&inv_std)?;
|
||||
|
||||
// Apply learnable weight
|
||||
normalized
|
||||
.mul(&self.weight)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
/// `SwiGLU` activation function for `LLaMA` MLP
|
||||
#[derive(Debug)]
|
||||
pub struct SwiGLU {
|
||||
gate_proj: Tensor, // W_gate: [hidden_size, intermediate_size]
|
||||
up_proj: Tensor, // W_up: [hidden_size, intermediate_size]
|
||||
down_proj: Tensor, // W_down: [intermediate_size, hidden_size]
|
||||
}
|
||||
|
||||
impl SwiGLU {
|
||||
/// Create new `SwiGLU` layer
|
||||
pub fn new(config: &LLaMAConfig, device: &Device) -> Result<Self> {
|
||||
let scale = (config.hidden_size as f32).recip().sqrt();
|
||||
|
||||
// Initialize projection weights
|
||||
let gate_proj = Tensor::randn(&[config.hidden_size, config.intermediate_size], device)?
|
||||
.mul_scalar(scale)?;
|
||||
let up_proj = Tensor::randn(&[config.hidden_size, config.intermediate_size], device)?
|
||||
.mul_scalar(scale)?;
|
||||
let down_proj = Tensor::randn(&[config.intermediate_size, config.hidden_size], device)?
|
||||
.mul_scalar(scale)?;
|
||||
|
||||
Ok(Self {
|
||||
gate_proj,
|
||||
up_proj,
|
||||
down_proj,
|
||||
})
|
||||
}
|
||||
|
||||
/// Forward pass of `SwiGLU`
|
||||
pub fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
|
||||
// SwiGLU(x) = Swish(xW_gate) * (xW_up) @ W_down
|
||||
// where Swish(x) = x * sigmoid(x)
|
||||
|
||||
// Gate projection and activation
|
||||
let gate_output = hidden_states.matmul(&self.gate_proj)?;
|
||||
let gate_activated = self.swish(&gate_output)?;
|
||||
|
||||
// Up projection (no activation)
|
||||
let up_output = hidden_states.matmul(&self.up_proj)?;
|
||||
|
||||
// Element-wise multiplication (gating)
|
||||
let gated = gate_activated.mul(&up_output)?;
|
||||
|
||||
// Down projection
|
||||
gated
|
||||
.matmul(&self.down_proj)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||||
}
|
||||
|
||||
/// Swish activation function: x * sigmoid(x)
|
||||
fn swish(&self, x: &Tensor) -> Result<Tensor> {
|
||||
// Swish(x) = x * sigmoid(x) = x * (1 / (1 + exp(-x)))
|
||||
let sigmoid = x.sigmoid()?;
|
||||
x.mul(&sigmoid)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
/// `LLaMA` MLP layer
|
||||
#[derive(Debug)]
|
||||
pub struct LLaMAMLP {
|
||||
swiglu: SwiGLU,
|
||||
// TODO: Add down projection
|
||||
}
|
||||
|
||||
impl LLaMAMLP {
|
||||
/// Create new `LLaMA` MLP layer
|
||||
pub fn new(config: &LLaMAConfig, device: &Device) -> Result<Self> {
|
||||
let swiglu = SwiGLU::new(config, device)?;
|
||||
|
||||
Ok(Self { swiglu })
|
||||
}
|
||||
|
||||
/// Forward pass of MLP
|
||||
pub fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
|
||||
// TODO: Implement proper MLP computation
|
||||
// For now, just pass through SwiGLU
|
||||
self.swiglu.forward(hidden_states)
|
||||
}
|
||||
}
|
||||
|
||||
/// Complete `LLaMA` transformer layer
|
||||
#[derive(Debug)]
|
||||
pub struct LLaMADecoderLayer {
|
||||
self_attn: LLaMAAttention,
|
||||
mlp: LLaMAMLP,
|
||||
input_layernorm: RMSNorm,
|
||||
post_attention_layernorm: RMSNorm,
|
||||
}
|
||||
|
||||
impl LLaMADecoderLayer {
|
||||
/// Create new `LLaMA` decoder layer
|
||||
pub fn new(config: &LLaMAConfig, device: &Device) -> Result<Self> {
|
||||
let self_attn = LLaMAAttention::new(config, device)?;
|
||||
let mlp = LLaMAMLP::new(config, device)?;
|
||||
let input_layernorm = RMSNorm::new(config.hidden_size, config.rms_norm_eps, device)?;
|
||||
let post_attention_layernorm =
|
||||
RMSNorm::new(config.hidden_size, config.rms_norm_eps, device)?;
|
||||
|
||||
Ok(Self {
|
||||
self_attn,
|
||||
mlp,
|
||||
input_layernorm,
|
||||
post_attention_layernorm,
|
||||
})
|
||||
}
|
||||
|
||||
/// Forward pass of decoder layer
|
||||
pub fn forward(
|
||||
&self,
|
||||
hidden_states: &Tensor,
|
||||
attention_mask: Option<&Tensor>,
|
||||
) -> Result<Tensor> {
|
||||
// Pre-attention norm
|
||||
let normed_hidden_states = self.input_layernorm.forward(hidden_states)?;
|
||||
|
||||
// Self-attention with residual connection
|
||||
let attn_output = self
|
||||
.self_attn
|
||||
.forward(&normed_hidden_states, attention_mask)?;
|
||||
let hidden_states = hidden_states.add(&attn_output)?;
|
||||
|
||||
// Pre-MLP norm
|
||||
let normed_hidden_states = self.post_attention_layernorm.forward(&hidden_states)?;
|
||||
|
||||
// MLP with residual connection
|
||||
let mlp_output = self.mlp.forward(&normed_hidden_states)?;
|
||||
let hidden_states = hidden_states.add(&mlp_output)?;
|
||||
|
||||
Ok(hidden_states)
|
||||
}
|
||||
}
|
||||
|
||||
/// Complete `LLaMA` model
|
||||
#[derive(Debug)]
|
||||
pub struct LLaMAModel {
|
||||
config: LLaMAConfig,
|
||||
layers: Vec<LLaMADecoderLayer>,
|
||||
norm: RMSNorm,
|
||||
// TODO: Add embed_tokens and lm_head
|
||||
}
|
||||
|
||||
impl LLaMAModel {
|
||||
/// Create new `LLaMA` model
|
||||
pub fn new(config: LLaMAConfig, device: &Device) -> Result<Self> {
|
||||
let mut layers = Vec::new();
|
||||
for _ in 0..config.num_hidden_layers {
|
||||
layers.push(LLaMADecoderLayer::new(&config, device)?);
|
||||
}
|
||||
|
||||
let norm = RMSNorm::new(config.hidden_size, config.rms_norm_eps, device)?;
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
layers,
|
||||
norm,
|
||||
})
|
||||
}
|
||||
|
||||
/// Forward pass of `LLaMA` model
|
||||
pub fn forward(&self, input_ids: &Tensor, _attention_mask: Option<&Tensor>) -> Result<Tensor> {
|
||||
// TODO: Implement proper forward pass with:
|
||||
// 1. Token embeddings
|
||||
// 2. All transformer layers
|
||||
// 3. Final normalization
|
||||
// 4. Language modeling head
|
||||
|
||||
// For now, just return a dummy tensor
|
||||
let device = Device::cuda(0).unwrap_or_default();
|
||||
Tensor::zeros(
|
||||
vec![1, input_ids.shape().dims()[1], self.config.hidden_size],
|
||||
&device,
|
||||
)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user