255 lines
8.3 KiB
Rust
255 lines
8.3 KiB
Rust
//! Integration tests for Mamba with transformer architecture
|
|
//!
|
|
//! This module demonstrates how Mamba blocks can be integrated with existing
|
|
//! transformer components to create hybrid architectures.
|
|
|
|
use super::{MambaBlock, MambaConfig, LayerNorm, PositionalEncoding};
|
|
use crate::Result;
|
|
use rtx_tensor::{Tensor, Device};
|
|
|
|
/// Hybrid transformer block using Mamba instead of attention
|
|
#[derive(Debug)]
|
|
pub struct MambaTransformerBlock {
|
|
/// Mamba state space model
|
|
mamba: MambaBlock,
|
|
/// Layer normalization before Mamba
|
|
norm1: LayerNorm,
|
|
/// Layer normalization after feed-forward
|
|
norm2: LayerNorm,
|
|
/// Feed-forward network weights (simplified)
|
|
ff_weight1: Tensor,
|
|
ff_weight2: Tensor,
|
|
/// Dropout probability
|
|
dropout: f32,
|
|
/// Device
|
|
device: Device,
|
|
}
|
|
|
|
impl MambaTransformerBlock {
|
|
/// Create new hybrid transformer block with Mamba
|
|
pub fn new(d_model: usize, ff_dim: usize, mamba_config: MambaConfig, device: &Device) -> Result<Self> {
|
|
let mamba = MambaBlock::new(mamba_config, device)?;
|
|
let norm1 = LayerNorm::new(d_model, 1e-5, device)?;
|
|
let norm2 = LayerNorm::new(d_model, 1e-5, device)?;
|
|
|
|
// Initialize feed-forward weights
|
|
let ff_weight1 = Tensor::randn(&[d_model, ff_dim], device)?;
|
|
let ff_weight2 = Tensor::randn(&[ff_dim, d_model], device)?;
|
|
|
|
Ok(Self {
|
|
mamba,
|
|
norm1,
|
|
norm2,
|
|
ff_weight1,
|
|
ff_weight2,
|
|
dropout: 0.1,
|
|
device: device.clone(),
|
|
})
|
|
}
|
|
|
|
/// Forward pass through hybrid block
|
|
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
|
// Pre-norm: LayerNorm -> Mamba -> Residual
|
|
let normed = self.norm1.forward(x)?;
|
|
let mamba_out = self.mamba.forward(&normed)?;
|
|
let x = x.add(&mamba_out)?; // Residual connection
|
|
|
|
// Feed-forward with residual connection
|
|
let normed = self.norm2.forward(&x)?;
|
|
let ff1 = normed.matmul(&self.ff_weight1)?.relu()?;
|
|
let ff2 = ff1.matmul(&self.ff_weight2)?;
|
|
let output = x.add(&ff2)?; // Residual connection
|
|
|
|
Ok(output)
|
|
}
|
|
}
|
|
|
|
/// Complete Mamba-based transformer model
|
|
#[derive(Debug)]
|
|
pub struct MambaTransformer {
|
|
/// Token embeddings
|
|
embeddings: Tensor,
|
|
/// Positional encoding
|
|
pos_encoding: PositionalEncoding,
|
|
/// Stack of Mamba transformer blocks
|
|
blocks: Vec<MambaTransformerBlock>,
|
|
/// Final layer normalization
|
|
final_norm: LayerNorm,
|
|
/// Output projection head
|
|
output_proj: Tensor,
|
|
/// Device
|
|
device: Device,
|
|
}
|
|
|
|
impl MambaTransformer {
|
|
/// Create new Mamba transformer
|
|
pub fn new(
|
|
vocab_size: usize,
|
|
d_model: usize,
|
|
num_layers: usize,
|
|
max_seq_len: usize,
|
|
device: &Device,
|
|
) -> Result<Self> {
|
|
// Initialize embeddings
|
|
let embeddings = Tensor::randn(&[vocab_size, d_model], device)?;
|
|
|
|
// Positional encoding
|
|
let pos_encoding = PositionalEncoding::new(d_model, max_seq_len, device)?;
|
|
|
|
// Create Mamba blocks
|
|
let mut blocks = Vec::new();
|
|
for _ in 0..num_layers {
|
|
let mamba_config = MambaConfig::new(d_model, 16, 4);
|
|
let block = MambaTransformerBlock::new(d_model, d_model * 4, mamba_config, device)?;
|
|
blocks.push(block);
|
|
}
|
|
|
|
// Final components
|
|
let final_norm = LayerNorm::new(d_model, 1e-5, device)?;
|
|
let output_proj = Tensor::randn(&[d_model, vocab_size], device)?;
|
|
|
|
Ok(Self {
|
|
embeddings,
|
|
pos_encoding,
|
|
blocks,
|
|
final_norm,
|
|
output_proj,
|
|
device: device.clone(),
|
|
})
|
|
}
|
|
|
|
/// Forward pass through full transformer
|
|
pub fn forward(&self, input_ids: &Tensor) -> Result<Tensor> {
|
|
let batch_size = input_ids.shape().dims()[0];
|
|
let seq_len = input_ids.shape().dims()[1];
|
|
|
|
// Get embeddings (simplified - assumes input_ids are indices)
|
|
let mut x = self.embeddings.clone(); // Simplified embedding lookup
|
|
|
|
// Add positional encoding
|
|
x = self.pos_encoding.forward(&x)?;
|
|
|
|
// Pass through Mamba blocks
|
|
for block in &self.blocks {
|
|
x = block.forward(&x)?;
|
|
}
|
|
|
|
// Final normalization and projection
|
|
x = self.final_norm.forward(&x)?;
|
|
let logits = x.matmul(&self.output_proj)?;
|
|
|
|
Ok(logits)
|
|
}
|
|
|
|
/// Generate text using the model (simplified)
|
|
pub fn generate(&self, prompt: &Tensor, max_length: usize) -> Result<Tensor> {
|
|
let mut current_seq = prompt.clone();
|
|
|
|
for _ in 0..max_length {
|
|
let logits = self.forward(¤t_seq)?;
|
|
|
|
// Simple greedy decoding (take argmax of last token)
|
|
// In practice, would implement proper sampling
|
|
let next_token = logits.clone(); // Simplified
|
|
|
|
// Would concatenate next_token to current_seq
|
|
// For now, just return the logits
|
|
return Ok(logits);
|
|
}
|
|
|
|
Ok(current_seq)
|
|
}
|
|
}
|
|
|
|
#[cfg(all(test, feature = "disabled_tests"))]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[tokio::test]
|
|
async fn test_mamba_transformer_block() -> Result<()> {
|
|
let device = Device::cpu();
|
|
let d_model = 256;
|
|
let ff_dim = 1024;
|
|
let mamba_config = MambaConfig::new(d_model, 16, 4);
|
|
|
|
let block = MambaTransformerBlock::new(d_model, ff_dim, mamba_config, &device)?;
|
|
|
|
let batch_size = 2;
|
|
let seq_len = 10;
|
|
let input = Tensor::randn(&[batch_size, seq_len, d_model], &device)?;
|
|
|
|
let output = block.forward(&input)?;
|
|
|
|
// Output should have same shape as input
|
|
assert_eq!(output.shape().dims(), input.shape().dims());
|
|
|
|
// Verify output is not just zeros
|
|
let output_data = output.to_vec()?;
|
|
assert!(output_data.iter().any(|&x| x != 0.0));
|
|
assert!(output_data.iter().all(|&x| x.is_finite()));
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_mamba_transformer_full() -> Result<()> {
|
|
let device = Device::cpu();
|
|
let vocab_size = 1000;
|
|
let d_model = 128;
|
|
let num_layers = 2;
|
|
let max_seq_len = 50;
|
|
|
|
let transformer = MambaTransformer::new(
|
|
vocab_size, d_model, num_layers, max_seq_len, &device
|
|
)?;
|
|
|
|
let batch_size = 1;
|
|
let seq_len = 10;
|
|
let input_ids = Tensor::randint(0, vocab_size as i32, &[batch_size, seq_len], &device)?;
|
|
|
|
let logits = transformer.forward(&input_ids)?;
|
|
|
|
// Check output shape
|
|
assert_eq!(logits.shape().dims(), &[batch_size, seq_len, vocab_size]);
|
|
|
|
// Verify meaningful output
|
|
let logits_data = logits.to_vec()?;
|
|
assert!(logits_data.iter().all(|&x| x.is_finite()));
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_mamba_vs_attention_efficiency() -> Result<()> {
|
|
let device = Device::cpu();
|
|
let d_model = 256;
|
|
let mamba_config = MambaConfig::new(d_model, 16, 4);
|
|
let mamba_block = MambaBlock::new(mamba_config, &device)?;
|
|
|
|
// Test with different sequence lengths to verify linear scaling
|
|
let sequence_lengths = vec![50, 100, 200];
|
|
let mut processing_times = Vec::new();
|
|
|
|
for seq_len in sequence_lengths {
|
|
let input = Tensor::randn(&[1, seq_len, d_model], &device)?;
|
|
|
|
let start = std::time::Instant::now();
|
|
let _output = mamba_block.forward(&input)?;
|
|
let elapsed = start.elapsed();
|
|
|
|
processing_times.push(elapsed.as_nanos() as f64);
|
|
}
|
|
|
|
// Mamba should scale approximately linearly
|
|
let ratio_1_2 = processing_times[1] / processing_times[0];
|
|
let ratio_2_3 = processing_times[2] / processing_times[1];
|
|
|
|
// Allow for some variance in timing
|
|
assert!(ratio_1_2 < 3.0, "Scaling should be roughly linear");
|
|
assert!(ratio_2_3 < 3.0, "Scaling should be roughly linear");
|
|
|
|
println!("Mamba scaling ratios: {:.2}, {:.2}", ratio_1_2, ratio_2_3);
|
|
|
|
Ok(())
|
|
}
|
|
} |