Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,508 @@
|
||||
//! Comprehensive tests for RWKV (Receptance Weighted Key Value) model
|
||||
//!
|
||||
//! Tests for linear complexity RNN with transformer-level performance,
|
||||
//! covering time-mixing, channel-mixing, WKV computation, and state caching.
|
||||
|
||||
#[cfg(all(test, feature = "disabled_tests"))]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::layers::*;
|
||||
use rtx_tensor::{Tensor, Device};
|
||||
use approx::assert_abs_diff_eq;
|
||||
|
||||
fn setup_device() -> Device {
|
||||
Device::cuda(0).unwrap_or(Device::default())
|
||||
}
|
||||
|
||||
// Test configuration structures
|
||||
#[test]
|
||||
fn test_rwkv_config_creation() {
|
||||
let config = RwkvConfig::new(768, 16);
|
||||
assert_eq!(config.d_model, 768);
|
||||
assert_eq!(config.n_layer, 16);
|
||||
assert_eq!(config.version, RwkvVersion::V6);
|
||||
assert!(config.use_layer_norm);
|
||||
assert!(config.prenorm);
|
||||
assert_eq!(config.time_mix_extra_dim, 32);
|
||||
assert_eq!(config.time_decay_extra_dim, 64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rwkv_config_variants() {
|
||||
let config_v4 = RwkvConfig::new(512, 12).with_version(RwkvVersion::V4);
|
||||
let config_v5 = RwkvConfig::new(1024, 24).with_version(RwkvVersion::V5);
|
||||
let config_v6 = RwkvConfig::new(2048, 32).with_version(RwkvVersion::V6);
|
||||
|
||||
assert_eq!(config_v4.version, RwkvVersion::V4);
|
||||
assert_eq!(config_v5.version, RwkvVersion::V5);
|
||||
assert_eq!(config_v6.version, RwkvVersion::V6);
|
||||
}
|
||||
|
||||
// Test WKV (Weighted Key-Value) computation core
|
||||
#[test]
|
||||
fn test_wkv_computation_basic() {
|
||||
let device = setup_device();
|
||||
let batch_size = 2;
|
||||
let seq_len = 10;
|
||||
let d_model = 64;
|
||||
|
||||
// Create test tensors
|
||||
let k = Tensor::randn(&[batch_size, seq_len, d_model], &device).unwrap();
|
||||
let v = Tensor::randn(&[batch_size, seq_len, d_model], &device).unwrap();
|
||||
let w = Tensor::randn(&[batch_size, seq_len, d_model], &device).unwrap(); // time decay
|
||||
let u = Tensor::randn(&[d_model], &device).unwrap(); // time first
|
||||
|
||||
let wkv_op = WkvComputation::new();
|
||||
let result = wkv_op.forward(&k, &v, &w, &u).unwrap();
|
||||
|
||||
assert_eq!(result.shape().dims(), &[batch_size, seq_len, d_model]);
|
||||
assert!(!result.isnan().any().unwrap());
|
||||
assert!(!result.isinf().any().unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wkv_computation_with_state() {
|
||||
let device = setup_device();
|
||||
let batch_size = 1;
|
||||
let seq_len = 5;
|
||||
let d_model = 32;
|
||||
|
||||
let k = Tensor::randn(&[batch_size, seq_len, d_model], &device).unwrap();
|
||||
let v = Tensor::randn(&[batch_size, seq_len, d_model], &device).unwrap();
|
||||
let w = Tensor::randn(&[batch_size, seq_len, d_model], &device).unwrap();
|
||||
let u = Tensor::randn(&[d_model], &device).unwrap();
|
||||
|
||||
// Test with initial state
|
||||
let mut state = RwkvState::new(batch_size, d_model, &device).unwrap();
|
||||
|
||||
let wkv_op = WkvComputation::new();
|
||||
let result = wkv_op.forward_with_state(&k, &v, &w, &u, Some(&mut state)).unwrap();
|
||||
|
||||
assert_eq!(result.shape().dims(), &[batch_size, seq_len, d_model]);
|
||||
assert!(!state.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wkv_exponential_decay_properties() {
|
||||
let device = setup_device();
|
||||
let batch_size = 1;
|
||||
let seq_len = 8;
|
||||
let d_model = 16;
|
||||
|
||||
// Create decay weights with known values
|
||||
let w = Tensor::full([batch_size, seq_len, d_model], -2.0, &device).unwrap();
|
||||
let k = Tensor::ones([batch_size, seq_len, d_model], &device).unwrap();
|
||||
let v = Tensor::ones([batch_size, seq_len, d_model], &device).unwrap();
|
||||
let u = Tensor::zeros([d_model], &device).unwrap();
|
||||
|
||||
let wkv_op = WkvComputation::new();
|
||||
let result = wkv_op.forward(&k, &v, &w, &u).unwrap();
|
||||
|
||||
// Check that later positions have less influence (exponential decay)
|
||||
let first_pos = result.narrow(1, 0, 1).unwrap();
|
||||
let last_pos = result.narrow(1, seq_len - 1, 1).unwrap();
|
||||
|
||||
// The exact comparison depends on WKV implementation
|
||||
assert_eq!(result.shape().dims(), &[batch_size, seq_len, d_model]);
|
||||
}
|
||||
|
||||
// Test time-mixing block
|
||||
#[test]
|
||||
fn test_time_mixing_creation() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(256, 8);
|
||||
|
||||
let time_mix = TimeMixing::new(&config, 0, &device).unwrap();
|
||||
assert_eq!(time_mix.layer_id(), 0);
|
||||
assert_eq!(time_mix.device(), &device);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_time_mixing_forward() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(128, 4);
|
||||
let batch_size = 2;
|
||||
let seq_len = 16;
|
||||
|
||||
let time_mix = TimeMixing::new(&config, 0, &device).unwrap();
|
||||
let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device).unwrap();
|
||||
|
||||
let output = time_mix.forward(&input).unwrap();
|
||||
|
||||
assert_eq!(output.shape().dims(), input.shape().dims());
|
||||
assert!(!output.isnan().any().unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_time_mixing_with_state_caching() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(64, 2);
|
||||
let batch_size = 1;
|
||||
let seq_len = 8;
|
||||
|
||||
let time_mix = TimeMixing::new(&config, 0, &device).unwrap();
|
||||
let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device).unwrap();
|
||||
|
||||
// Forward pass with state caching
|
||||
let mut state = RwkvState::new(batch_size, config.d_model, &device).unwrap();
|
||||
let output = time_mix.forward_with_state(&input, Some(&mut state)).unwrap();
|
||||
|
||||
assert_eq!(output.shape().dims(), input.shape().dims());
|
||||
assert!(!state.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_time_mixing_receptance_key_value_gates() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(96, 3);
|
||||
let batch_size = 1;
|
||||
let seq_len = 4;
|
||||
|
||||
let time_mix = TimeMixing::new(&config, 0, &device).unwrap();
|
||||
let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device).unwrap();
|
||||
|
||||
// Test that gates produce expected shapes
|
||||
let gates = time_mix.compute_gates(&input).unwrap();
|
||||
|
||||
assert_eq!(gates.receptance.shape().dims(), &[batch_size, seq_len, config.d_model]);
|
||||
assert_eq!(gates.key.shape().dims(), &[batch_size, seq_len, config.d_model]);
|
||||
assert_eq!(gates.value.shape().dims(), &[batch_size, seq_len, config.d_model]);
|
||||
assert_eq!(gates.time_decay.shape().dims(), &[batch_size, seq_len, config.d_model]);
|
||||
assert_eq!(gates.time_first.shape().dims(), &[config.d_model]);
|
||||
}
|
||||
|
||||
// Test channel-mixing block
|
||||
#[test]
|
||||
fn test_channel_mixing_creation() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(192, 6);
|
||||
|
||||
let channel_mix = ChannelMixing::new(&config, 1, &device).unwrap();
|
||||
assert_eq!(channel_mix.layer_id(), 1);
|
||||
assert_eq!(channel_mix.device(), &device);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_channel_mixing_forward() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(160, 5);
|
||||
let batch_size = 2;
|
||||
let seq_len = 12;
|
||||
|
||||
let channel_mix = ChannelMixing::new(&config, 0, &device).unwrap();
|
||||
let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device).unwrap();
|
||||
|
||||
let output = channel_mix.forward(&input).unwrap();
|
||||
|
||||
assert_eq!(output.shape().dims(), input.shape().dims());
|
||||
assert!(!output.isnan().any().unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_channel_mixing_feed_forward_structure() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(80, 2).with_ffn_dim(320); // 4x expansion
|
||||
let batch_size = 1;
|
||||
let seq_len = 6;
|
||||
|
||||
let channel_mix = ChannelMixing::new(&config, 0, &device).unwrap();
|
||||
let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device).unwrap();
|
||||
|
||||
// Test feed-forward expansion and contraction
|
||||
let ff_output = channel_mix.feed_forward(&input).unwrap();
|
||||
|
||||
assert_eq!(ff_output.shape().dims(), &[batch_size, seq_len, config.d_model]);
|
||||
}
|
||||
|
||||
// Test complete RWKV block
|
||||
#[test]
|
||||
fn test_rwkv_block_creation() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(384, 8);
|
||||
|
||||
let rwkv_block = RwkvBlock::new(&config, 2, &device).unwrap();
|
||||
assert_eq!(rwkv_block.layer_id(), 2);
|
||||
assert_eq!(rwkv_block.device(), &device);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rwkv_block_forward() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(256, 6);
|
||||
let batch_size = 2;
|
||||
let seq_len = 20;
|
||||
|
||||
let rwkv_block = RwkvBlock::new(&config, 0, &device).unwrap();
|
||||
let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device).unwrap();
|
||||
|
||||
let output = rwkv_block.forward(&input).unwrap();
|
||||
|
||||
assert_eq!(output.shape().dims(), input.shape().dims());
|
||||
assert!(!output.isnan().any().unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rwkv_block_residual_connections() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(128, 4);
|
||||
let batch_size = 1;
|
||||
let seq_len = 8;
|
||||
|
||||
let rwkv_block = RwkvBlock::new(&config, 0, &device).unwrap();
|
||||
let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device).unwrap();
|
||||
|
||||
let output = rwkv_block.forward(&input).unwrap();
|
||||
|
||||
// Check that output magnitude is reasonable (residual connections)
|
||||
let input_norm = input.norm(None, None, false).unwrap();
|
||||
let output_norm = output.norm(None, None, false).unwrap();
|
||||
|
||||
// Output should have similar magnitude due to residual connections
|
||||
let norm_ratio = output_norm.div(&input_norm).unwrap().get_item([]).unwrap();
|
||||
assert!(norm_ratio > 0.5 && norm_ratio < 2.0, "Norm ratio: {}", norm_ratio);
|
||||
}
|
||||
|
||||
// Test layer normalization integration
|
||||
#[test]
|
||||
fn test_rwkv_block_with_layer_norm() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(144, 3).with_layer_norm(true);
|
||||
let batch_size = 1;
|
||||
let seq_len = 10;
|
||||
|
||||
let rwkv_block = RwkvBlock::new(&config, 0, &device).unwrap();
|
||||
let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device).unwrap();
|
||||
|
||||
let output = rwkv_block.forward(&input).unwrap();
|
||||
|
||||
assert_eq!(output.shape().dims(), input.shape().dims());
|
||||
|
||||
// Check that layer norm is applied (output should have unit variance approximately)
|
||||
let output_var = output.var(Some(vec![2]), false, false).unwrap().mean(None, false).unwrap();
|
||||
let var_value = output_var.get_item([]).unwrap();
|
||||
assert!(var_value > 0.5 && var_value < 2.0, "Output variance: {}", var_value);
|
||||
}
|
||||
|
||||
// Test state caching for inference
|
||||
#[test]
|
||||
fn test_rwkv_state_creation() {
|
||||
let device = setup_device();
|
||||
let batch_size = 2;
|
||||
let d_model = 128;
|
||||
|
||||
let state = RwkvState::new(batch_size, d_model, &device).unwrap();
|
||||
|
||||
assert_eq!(state.batch_size(), batch_size);
|
||||
assert_eq!(state.d_model(), d_model);
|
||||
assert!(state.is_empty());
|
||||
assert_eq!(state.device(), &device);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rwkv_state_operations() {
|
||||
let device = setup_device();
|
||||
let batch_size = 1;
|
||||
let d_model = 64;
|
||||
|
||||
let mut state = RwkvState::new(batch_size, d_model, &device).unwrap();
|
||||
|
||||
// Set state
|
||||
let test_state = Tensor::randn(&[batch_size, d_model], &device).unwrap();
|
||||
state.set_layer_state(0, "time_mix", test_state.clone()).unwrap();
|
||||
|
||||
assert!(!state.is_empty());
|
||||
|
||||
// Get state
|
||||
let retrieved_state = state.get_layer_state(0, "time_mix").unwrap();
|
||||
assert_eq!(retrieved_state.shape().dims(), test_state.shape().dims());
|
||||
|
||||
// Clear state
|
||||
state.clear();
|
||||
assert!(state.is_empty());
|
||||
}
|
||||
|
||||
// Test version-specific features
|
||||
#[test]
|
||||
fn test_rwkv_v4_specific_features() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(128, 6).with_version(RwkvVersion::V4);
|
||||
let batch_size = 1;
|
||||
let seq_len = 8;
|
||||
|
||||
let rwkv_block = RwkvBlock::new(&config, 0, &device).unwrap();
|
||||
let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device).unwrap();
|
||||
|
||||
let output = rwkv_block.forward(&input).unwrap();
|
||||
assert_eq!(output.shape().dims(), input.shape().dims());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rwkv_v5_specific_features() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(160, 8).with_version(RwkvVersion::V5);
|
||||
let batch_size = 1;
|
||||
let seq_len = 12;
|
||||
|
||||
let rwkv_block = RwkvBlock::new(&config, 0, &device).unwrap();
|
||||
let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device).unwrap();
|
||||
|
||||
let output = rwkv_block.forward(&input).unwrap();
|
||||
assert_eq!(output.shape().dims(), input.shape().dims());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rwkv_v6_specific_features() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(192, 10).with_version(RwkvVersion::V6);
|
||||
let batch_size = 1;
|
||||
let seq_len = 16;
|
||||
|
||||
let rwkv_block = RwkvBlock::new(&config, 0, &device).unwrap();
|
||||
let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device).unwrap();
|
||||
|
||||
let output = rwkv_block.forward(&input).unwrap();
|
||||
assert_eq!(output.shape().dims(), input.shape().dims());
|
||||
}
|
||||
|
||||
// Test RNN mode vs parallel mode
|
||||
#[test]
|
||||
fn test_rwkv_rnn_mode() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(96, 4);
|
||||
let batch_size = 1;
|
||||
let seq_len = 1; // Single token for RNN mode
|
||||
|
||||
let rwkv_block = RwkvBlock::new(&config, 0, &device).unwrap();
|
||||
let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device).unwrap();
|
||||
|
||||
let mut state = RwkvState::new(batch_size, config.d_model, &device).unwrap();
|
||||
let output = rwkv_block.forward_with_state(&input, Some(&mut state)).unwrap();
|
||||
|
||||
assert_eq!(output.shape().dims(), &[batch_size, seq_len, config.d_model]);
|
||||
assert!(!state.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rwkv_parallel_mode() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(128, 6);
|
||||
let batch_size = 2;
|
||||
let seq_len = 32; // Long sequence for parallel mode
|
||||
|
||||
let rwkv_block = RwkvBlock::new(&config, 0, &device).unwrap();
|
||||
let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device).unwrap();
|
||||
|
||||
let output = rwkv_block.forward(&input).unwrap();
|
||||
|
||||
assert_eq!(output.shape().dims(), &[batch_size, seq_len, config.d_model]);
|
||||
assert!(!output.isnan().any().unwrap());
|
||||
}
|
||||
|
||||
// Test numerical stability
|
||||
#[test]
|
||||
fn test_rwkv_numerical_stability() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(64, 2);
|
||||
let batch_size = 1;
|
||||
let seq_len = 100; // Long sequence to test stability
|
||||
|
||||
let rwkv_block = RwkvBlock::new(&config, 0, &device).unwrap();
|
||||
|
||||
// Test with large values
|
||||
let large_input = Tensor::full([batch_size, seq_len, config.d_model], 10.0, &device).unwrap();
|
||||
let large_output = rwkv_block.forward(&large_input).unwrap();
|
||||
|
||||
assert!(!large_output.isnan().any().unwrap());
|
||||
assert!(!large_output.isinf().any().unwrap());
|
||||
|
||||
// Test with small values
|
||||
let small_input = Tensor::full([batch_size, seq_len, config.d_model], 0.001, &device).unwrap();
|
||||
let small_output = rwkv_block.forward(&small_input).unwrap();
|
||||
|
||||
assert!(!small_output.isnan().any().unwrap());
|
||||
assert!(!small_output.isinf().any().unwrap());
|
||||
}
|
||||
|
||||
// Test parameter counting
|
||||
#[test]
|
||||
fn test_rwkv_parameter_count() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(256, 8);
|
||||
|
||||
let rwkv_block = RwkvBlock::new(&config, 0, &device).unwrap();
|
||||
let params = rwkv_block.parameters();
|
||||
|
||||
assert!(!params.is_empty());
|
||||
|
||||
// Calculate expected parameter count
|
||||
let d_model = config.d_model;
|
||||
let ffn_dim = config.ffn_dim.unwrap_or(4 * d_model);
|
||||
|
||||
let time_mix_params = 4 * d_model * d_model + 2 * d_model; // R,K,V,O + time_decay, time_first
|
||||
let channel_mix_params = 2 * d_model * ffn_dim; // key, value projections
|
||||
let norm_params = if config.use_layer_norm { 4 * d_model } else { 0 }; // 2 layer norms
|
||||
|
||||
let expected_params = time_mix_params + channel_mix_params + norm_params;
|
||||
let actual_param_count: usize = params.iter().map(|p| p.numel()).sum();
|
||||
|
||||
// Allow some variance in parameter count due to implementation details
|
||||
assert!(actual_param_count > expected_params / 2);
|
||||
assert!(actual_param_count < expected_params * 2);
|
||||
}
|
||||
|
||||
// Test integration with Layer trait
|
||||
#[test]
|
||||
fn test_rwkv_layer_trait_integration() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(128, 4);
|
||||
|
||||
let rwkv_block = RwkvBlock::new(&config, 0, &device).unwrap();
|
||||
let batch_size = 1;
|
||||
let seq_len = 8;
|
||||
let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device).unwrap();
|
||||
|
||||
// Test Layer trait methods
|
||||
assert_eq!(rwkv_block.layer_type(), "RwkvBlock");
|
||||
assert_eq!(rwkv_block.device(), &device);
|
||||
|
||||
let output = rwkv_block.forward(&input).unwrap();
|
||||
assert_eq!(output.shape().dims(), input.shape().dims());
|
||||
|
||||
let params = rwkv_block.parameters();
|
||||
let params_mut = rwkv_block.parameters_mut();
|
||||
assert_eq!(params.len(), params_mut.len());
|
||||
}
|
||||
|
||||
// Performance and memory tests
|
||||
#[test]
|
||||
fn test_rwkv_linear_complexity() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(64, 2);
|
||||
let batch_size = 1;
|
||||
|
||||
let rwkv_block = RwkvBlock::new(&config, 0, &device).unwrap();
|
||||
|
||||
// Test different sequence lengths
|
||||
for seq_len in [16, 32, 64, 128].iter() {
|
||||
let input = Tensor::randn(&[batch_size, *seq_len, config.d_model], &device).unwrap();
|
||||
let output = rwkv_block.forward(&input).unwrap();
|
||||
|
||||
assert_eq!(output.shape().dims(), &[batch_size, *seq_len, config.d_model]);
|
||||
assert!(!output.isnan().any().unwrap());
|
||||
}
|
||||
}
|
||||
|
||||
// Test error conditions
|
||||
#[test]
|
||||
fn test_rwkv_error_conditions() {
|
||||
let device = setup_device();
|
||||
let config = RwkvConfig::new(128, 4);
|
||||
|
||||
// Test with invalid dimensions
|
||||
let rwkv_block = RwkvBlock::new(&config, 0, &device).unwrap();
|
||||
let invalid_input = Tensor::randn(&[2, 8, 64], &device).unwrap(); // Wrong d_model
|
||||
|
||||
let result = rwkv_block.forward(&invalid_input);
|
||||
assert!(result.is_err()); // Should fail due to dimension mismatch
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user