1327 lines
51 KiB
Rust
1327 lines
51 KiB
Rust
//! 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<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
|
|
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<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, &[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<Tensor> {
|
|
// 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<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 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<Self> {
|
|
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<Tensor> {
|
|
// 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<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)
|
|
}
|
|
}
|
|
|
|
/// Complete LLaMA model
|
|
#[derive(Debug)]
|
|
pub struct LLaMAModel {
|
|
/// Model configuration
|
|
config: LLaMAConfig,
|
|
/// Token embeddings
|
|
embed_tokens: Tensor,
|
|
/// Stack of LLaMA blocks
|
|
layers: Vec<LLaMABlock>,
|
|
/// 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<Self> {
|
|
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<ModelOutput> {
|
|
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<Tensor> {
|
|
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<Tensor> {
|
|
// 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<Tensor> {
|
|
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<Tensor> {
|
|
// 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<usize>,
|
|
top_p: Option<f32>,
|
|
) -> Result<Tensor> {
|
|
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<Tensor> {
|
|
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<Tensor> {
|
|
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<Tensor> {
|
|
// 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<Tensor> {
|
|
// 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<Tensor> {
|
|
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<Tensor> {
|
|
// 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<ModelOutput> {
|
|
self.forward(input_ids, labels)
|
|
}
|
|
|
|
fn parameters(&self) -> HashMap<String, Tensor> {
|
|
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<String, Tensor>) -> 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<String, usize> {
|
|
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<Tensor> {
|
|
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());
|
|
}
|
|
}
|