Initial commit

crates/training/rtx-transformers/src/optimizers/novograd_tests.rs.disabled

@@ -0,0 +1,573 @@
+//! Comprehensive tests for NovoGrad optimizer
+//!
+//! NovoGrad: Layer-wise gradient normalization with adaptive second moment averaging
+//! Key features tested:
+//! - Layer-wise gradient normalization: g_norm = g / ||g||
+//! - Adaptive second moment: v_t = β2 * v_{t-1} + (1-β2) * ||g||²
+//! - Momentum with normalized gradients: m_t = β1 * m_{t-1} + (lr / √(v_t + ε)) * g_norm
+//! - Weight update: θ_t+1 = θ_t - m_t - λ * θ_t (with weight decay)
+//! - Bias correction for early iterations
+//! - Gradient clipping support
+
+#[cfg(test)]
+mod tests {
+    use super::super::{NovoGradOptimizer, NovoGradConfig, Optimizer};
+    use crate::{Result, TransformerError};
+    use rtx_tensor::{Tensor, Device, DType};
+    use std::collections::HashMap;
+
+    // Test configuration defaults
+    fn default_novograd_config() -> NovoGradConfig {
+        NovoGradConfig {
+            learning_rate: 1e-3,
+            beta1: 0.95,
+            beta2: 0.98,
+            eps: 1e-8,
+            weight_decay: 0.0,
+            grad_averaging: false,
+        }
+    }
+
+    #[test]
+    fn test_novograd_config_validation() {
+        // Valid config should pass
+        let config = default_novograd_config();
+        assert!(NovoGradOptimizer::new(config).is_ok());
+
+        // Invalid learning rate
+        let mut invalid_config = default_novograd_config();
+        invalid_config.learning_rate = -0.1;
+        assert!(NovoGradOptimizer::new(invalid_config).is_err());
+
+        // Invalid beta1
+        let mut invalid_config = default_novograd_config();
+        invalid_config.beta1 = 1.1;
+        assert!(NovoGradOptimizer::new(invalid_config).is_err());
+
+        // Invalid beta2
+        let mut invalid_config = default_novograd_config();
+        invalid_config.beta2 = -0.1;
+        assert!(NovoGradOptimizer::new(invalid_config).is_err());
+
+        // Invalid eps
+        let mut invalid_config = default_novograd_config();
+        invalid_config.eps = -1e-8;
+        assert!(NovoGradOptimizer::new(invalid_config).is_err());
+
+        // Invalid weight decay
+        let mut invalid_config = default_novograd_config();
+        invalid_config.weight_decay = -0.5;
+        assert!(NovoGradOptimizer::new(invalid_config).is_err());
+    }
+
+    #[test]
+    fn test_novograd_parameter_creation() -> Result<()> {
+        let config = default_novograd_config();
+        let optimizer = NovoGradOptimizer::new(config)?;
+
+        assert_eq!(optimizer.learning_rate(), 1e-3);
+        assert_eq!(optimizer.beta1(), 0.95);
+        assert_eq!(optimizer.beta2(), 0.98);
+        assert_eq!(optimizer.eps(), 1e-8);
+        assert_eq!(optimizer.weight_decay(), 0.0);
+        assert_eq!(optimizer.grad_averaging(), false);
+        assert_eq!(optimizer.optimizer_type(), "NovoGrad");
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_novograd_gradient_normalization() -> Result<()> {
+        let config = NovoGradConfig {
+            learning_rate: 0.01,
+            beta1: 0.95,
+            beta2: 0.98,
+            eps: 1e-8,
+            weight_decay: 0.0,
+            grad_averaging: false,
+        };
+
+        let mut optimizer = NovoGradOptimizer::new(config)?;
+        let device = Device::cpu();
+
+        // Create test parameter and gradient with known norm
+        let param = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], &[4], &device)?;
+        let grad = Tensor::from_data(vec![3.0, 4.0, 0.0, 0.0], &[4], &device)?; // ||grad|| = 5.0
+
+        // First step - should normalize gradient by its norm
+        let updated_param1 = optimizer.step_param("test_param", &param, &grad)?;
+        
+        // Verify state was created
+        assert!(optimizer.has_state("test_param"));
+        assert_eq!(optimizer.get_step_count("test_param")?, 1);
+
+        // For first step with normalized gradient: g_norm = [0.6, 0.8, 0.0, 0.0]
+        // Second moment: v_1 = (1-0.98) * 25 = 0.5
+        // Momentum: m_1 = (0.01 / √(0.5 + 1e-8)) * [0.6, 0.8, 0.0, 0.0]
+        let param_data = param.to_cpu()?;
+        let updated_data = updated_param1.to_cpu()?;
+        
+        // Parameters should be updated in proportion to normalized gradient
+        let update_ratio = (param_data[1] - updated_data[1]) / (param_data[0] - updated_data[0]);
+        assert!((update_ratio - 4.0/3.0).abs() < 1e-6, 
+                "Update should be proportional to normalized gradient: expected 4/3, got {}", update_ratio);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_novograd_momentum_accumulation() -> Result<()> {
+        let config = NovoGradConfig {
+            learning_rate: 0.01,
+            beta1: 0.95,
+            beta2: 0.98,
+            eps: 1e-8,
+            weight_decay: 0.0,
+            grad_averaging: false,
+        };
+
+        let mut optimizer = NovoGradOptimizer::new(config)?;
+        let device = Device::cpu();
+
+        let param = Tensor::from_data(vec![1.0], &[1], &device)?;
+        let grad = Tensor::from_data(vec![2.0], &[1], &device)?; // ||grad|| = 2.0, normalized = 1.0
+
+        // Take multiple steps with same gradient
+        let mut current_param = param.clone();
+        let mut momentum_effects = Vec::new();
+        
+        for step in 0..5 {
+            let prev_data = current_param.to_cpu()?;
+            current_param = optimizer.step_param("param", &current_param, &grad)?;
+            let new_data = current_param.to_cpu()?;
+            
+            let update_magnitude = (prev_data[0] - new_data[0]).abs();
+            momentum_effects.push(update_magnitude);
+            
+            assert_eq!(optimizer.get_step_count("param")?, step + 1);
+        }
+
+        // Momentum should build up over iterations for consistent gradient direction
+        // Later updates should be influenced by accumulated momentum
+        assert!(momentum_effects.len() == 5, "Should have 5 momentum measurements");
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_novograd_second_moment_averaging() -> Result<()> {
+        let config = NovoGradConfig {
+            learning_rate: 0.01,
+            beta1: 0.95,
+            beta2: 0.98,
+            eps: 1e-8,
+            weight_decay: 0.0,
+            grad_averaging: false,
+        };
+
+        let mut optimizer = NovoGradOptimizer::new(config)?;
+        let device = Device::cpu();
+
+        let param = Tensor::from_data(vec![1.0], &[1], &device)?;
+        
+        // Use different gradient magnitudes to test second moment averaging
+        let gradients = vec![
+            Tensor::from_data(vec![1.0], &[1], &device)?, // ||g|| = 1.0
+            Tensor::from_data(vec![2.0], &[1], &device)?, // ||g|| = 2.0
+            Tensor::from_data(vec![0.5], &[1], &device)?, // ||g|| = 0.5
+        ];
+
+        let mut current_param = param;
+        for (step, grad) in gradients.iter().enumerate() {
+            current_param = optimizer.step_param("param", &current_param, grad)?;
+            assert_eq!(optimizer.get_step_count("param")?, step as i64 + 1);
+        }
+
+        // The optimizer should have adapted to different gradient norms
+        assert!(optimizer.has_state("param"));
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_novograd_bias_correction() -> Result<()> {
+        let config = NovoGradConfig {
+            learning_rate: 0.1,
+            beta1: 0.95,
+            beta2: 0.98,
+            eps: 1e-8,
+            weight_decay: 0.0,
+            grad_averaging: false,
+        };
+
+        let mut optimizer = NovoGradOptimizer::new(config)?;
+        let device = Device::cpu();
+
+        let param = Tensor::from_data(vec![1.0], &[1], &device)?;
+        let grad = Tensor::from_data(vec![1.0], &[1], &device)?;
+
+        // First few steps should show bias correction effects
+        let mut update_magnitudes = Vec::new();
+        let mut current_param = param;
+
+        for _step in 0..10 {
+            let prev_data = current_param.to_cpu()?;
+            current_param = optimizer.step_param("param", &current_param, &grad)?;
+            let new_data = current_param.to_cpu()?;
+            
+            let update_magnitude = (prev_data[0] - new_data[0]).abs();
+            update_magnitudes.push(update_magnitude);
+        }
+
+        // Early steps should have larger updates due to bias correction
+        // Later steps should stabilize
+        assert!(update_magnitudes[0] > update_magnitudes[9] * 0.5,
+                "Early steps should benefit from bias correction");
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_novograd_multi_parameter() -> Result<()> {
+        let config = default_novograd_config();
+        let mut optimizer = NovoGradOptimizer::new(config)?;
+        let device = Device::cpu();
+
+        // Create multiple parameters with different layer characteristics
+        let weight1 = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], &[2, 2], &device)?;
+        let weight2 = Tensor::from_data(vec![0.1, 0.2], &[2], &device)?;
+        let bias = Tensor::from_data(vec![0.01], &[1], &device)?;
+
+        // Create gradients with different norms (simulating different layers)
+        let weight1_grad = Tensor::from_data(vec![0.1, -0.1, 0.2, -0.2], &[2, 2], &device)?; // norm ≈ 0.316
+        let weight2_grad = Tensor::from_data(vec![1.0, -1.0], &[2], &device)?; // norm ≈ 1.414
+        let bias_grad = Tensor::from_data(vec![0.01], &[1], &device)?; // norm = 0.01
+
+        // Store gradients and process
+        let mut gradients = HashMap::new();
+        gradients.insert("weight1".to_string(), weight1_grad);
+        gradients.insert("weight2".to_string(), weight2_grad);
+        gradients.insert("bias".to_string(), bias_grad);
+
+        optimizer.set_gradients(gradients)?;
+        let parameter_updates = optimizer.step(0.01)?;
+
+        // Verify updates exist for all parameters
+        assert!(parameter_updates.contains_key("weight1"));
+        assert!(parameter_updates.contains_key("weight2"));
+        assert!(parameter_updates.contains_key("bias"));
+
+        // Verify shapes match
+        assert_eq!(parameter_updates["weight1"].shape().dims(), &[2, 2]);
+        assert_eq!(parameter_updates["weight2"].shape().dims(), &[2]);
+        assert_eq!(parameter_updates["bias"].shape().dims(), &[1]);
+
+        // Verify state was created for all parameters
+        assert!(optimizer.has_state("weight1"));
+        assert!(optimizer.has_state("weight2"));
+        assert!(optimizer.has_state("bias"));
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_novograd_weight_decay() -> Result<()> {
+        let config = NovoGradConfig {
+            learning_rate: 0.1,
+            beta1: 0.95,
+            beta2: 0.98,
+            eps: 1e-8,
+            weight_decay: 0.01,
+            grad_averaging: false,
+        };
+
+        let mut optimizer = NovoGradOptimizer::new(config)?;
+        let device = Device::cpu();
+
+        let param = Tensor::from_data(vec![2.0], &[1], &device)?; // Non-zero parameter
+        let grad = Tensor::from_data(vec![0.0], &[1], &device)?; // Zero gradient
+
+        // With zero gradient, only weight decay should affect parameter
+        let updated_param = optimizer.step_param("param", &param, &grad)?;
+        let param_data = param.to_cpu()?;
+        let updated_data = updated_param.to_cpu()?;
+
+        // Weight decay should reduce parameter magnitude: θ_new = θ_old - λ * θ_old
+        let expected = param_data[0] * (1.0 - 0.01);
+        assert!((updated_data[0] - expected).abs() < 1e-6, 
+                "Weight decay should reduce parameter: expected {}, got {}", expected, updated_data[0]);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_novograd_gradient_averaging() -> Result<()> {
+        let config = NovoGradConfig {
+            learning_rate: 0.01,
+            beta1: 0.95,
+            beta2: 0.98,
+            eps: 1e-8,
+            weight_decay: 0.0,
+            grad_averaging: true,
+        };
+
+        let mut optimizer = NovoGradOptimizer::new(config)?;
+        let device = Device::cpu();
+
+        let param = Tensor::from_data(vec![1.0, 2.0], &[2], &device)?;
+        let grad = Tensor::from_data(vec![3.0, 4.0], &[2], &device)?; // ||grad|| = 5.0
+
+        let updated_param = optimizer.step_param("param", &param, &grad)?;
+        
+        // With gradient averaging enabled, behavior should differ from standard NovoGrad
+        assert!(optimizer.has_state("param"));
+        assert_eq!(optimizer.get_step_count("param")?, 1);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_novograd_learning_rate_updates() -> Result<()> {
+        let config = default_novograd_config();
+        let mut optimizer = NovoGradOptimizer::new(config)?;
+
+        assert_eq!(optimizer.learning_rate(), 1e-3);
+
+        // Update learning rate
+        optimizer.set_learning_rate(0.01)?;
+        assert_eq!(optimizer.learning_rate(), 0.01);
+
+        // Invalid learning rate should error
+        assert!(optimizer.set_learning_rate(-0.1).is_err());
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_novograd_gradient_clipping_compatibility() -> Result<()> {
+        let config = default_novograd_config();
+        let mut optimizer = NovoGradOptimizer::new(config)?;
+        let device = Device::cpu();
+
+        // Create large gradients that would need clipping
+        let large_grad1 = Tensor::from_data(vec![10.0, 20.0], &[2], &device)?;
+        let large_grad2 = Tensor::from_data(vec![30.0], &[1], &device)?;
+
+        let mut gradients = HashMap::new();
+        gradients.insert("param1".to_string(), large_grad1);
+        gradients.insert("param2".to_string(), large_grad2);
+
+        optimizer.set_gradients(gradients)?;
+
+        // Should be able to compute gradient norm before clipping
+        let total_norm_squared = 10.0*10.0 + 20.0*20.0 + 30.0*30.0; // 100 + 400 + 900 = 1400
+        let expected_norm = total_norm_squared.sqrt(); // ≈ 37.42
+
+        // The optimizer should handle large gradients without issues
+        let parameter_updates = optimizer.step(0.01)?;
+        assert_eq!(parameter_updates.len(), 2);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_novograd_zero_gradient_handling() -> Result<()> {
+        let config = default_novograd_config();
+        let mut optimizer = NovoGradOptimizer::new(config)?;
+        let device = Device::cpu();
+
+        let param = Tensor::from_data(vec![1.0], &[1], &device)?;
+        let zero_grad = Tensor::from_data(vec![0.0], &[1], &device)?;
+
+        // Zero gradient should not cause errors (gradient norm will be 0)
+        let updated_param = optimizer.step_param("param", &param, &zero_grad)?;
+        
+        let param_data = param.to_cpu()?;
+        let updated_data = updated_param.to_cpu()?;
+        
+        // With zero gradient and no weight decay, parameter should be unchanged
+        // (normalized gradient is 0/0 which should be handled gracefully)
+        assert!((updated_data[0] - param_data[0]).abs() < 1e-6, 
+                "Zero gradient should not change parameter significantly");
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_novograd_convergence_quadratic() -> Result<()> {
+        let config = NovoGradConfig {
+            learning_rate: 0.1,
+            beta1: 0.95,
+            beta2: 0.98,
+            eps: 1e-8,
+            weight_decay: 0.0,
+            grad_averaging: false,
+        };
+
+        let mut optimizer = NovoGradOptimizer::new(config)?;
+        let device = Device::cpu();
+
+        // Optimize f(x) = (x - 3)^2, minimum at x = 3
+        let target = 3.0;
+        let mut param = Tensor::from_data(vec![0.0], &[1], &device)?;
+        let mut loss_history = Vec::new();
+
+        for _step in 0..50 {
+            let param_data = param.to_cpu()?;
+            let current_value = param_data[0];
+            let loss = (current_value - target).powi(2);
+            loss_history.push(loss);
+
+            // Gradient of (x - 3)^2 is 2(x - 3)
+            let gradient_value = 2.0 * (current_value - target);
+            let grad = Tensor::from_data(vec![gradient_value], &[1], &device)?;
+
+            param = optimizer.step_param("param", &param, &grad)?;
+        }
+
+        let final_data = param.to_cpu()?;
+        let final_value = final_data[0];
+        let final_loss = (final_value - target).powi(2);
+
+        // Should converge reasonably close to target
+        assert!(
+            (final_value - target).abs() < 0.5,
+            "Should converge close to target, got {} vs {}",
+            final_value, target
+        );
+
+        // Loss should decrease significantly
+        assert!(
+            final_loss < loss_history[0] * 0.1,
+            "Loss should decrease significantly from {} to {}",
+            loss_history[0], final_loss
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_novograd_state_management() -> Result<()> {
+        let config = default_novograd_config();
+        let mut optimizer = NovoGradOptimizer::new(config)?;
+        let device = Device::cpu();
+
+        let param = Tensor::from_data(vec![1.0, 2.0, 3.0], &[3], &device)?;
+        let grad = Tensor::from_data(vec![0.1, 0.2, 0.3], &[3], &device)?;
+
+        // Initialize state
+        optimizer.step_param("param", &param, &grad)?;
+        
+        // NovoGrad should store momentum and second moment (per-layer)
+        assert!(optimizer.has_state("param"));
+
+        // Reset state and verify cleanup
+        optimizer.reset_state("param")?;
+        assert!(!optimizer.has_state("param"));
+
+        // Reset all state
+        optimizer.step_param("param1", &param, &grad)?;
+        optimizer.step_param("param2", &param, &grad)?;
+        optimizer.reset_all_state();
+        assert!(!optimizer.has_state("param1"));
+        assert!(!optimizer.has_state("param2"));
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_novograd_gradient_shapes() -> Result<()> {
+        let config = default_novograd_config();
+        let mut optimizer = NovoGradOptimizer::new(config)?;
+        let device = Device::cpu();
+
+        let param = Tensor::from_data(vec![1.0, 2.0], &[2], &device)?;
+        let wrong_grad = Tensor::from_data(vec![0.1], &[1], &device)?;
+
+        // Shape mismatch should error
+        let result = optimizer.step_param("param", &param, &wrong_grad);
+        assert!(result.is_err());
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_novograd_deterministic_behavior() -> Result<()> {
+        let config = default_novograd_config();
+        let device = Device::cpu();
+
+        // Run same optimization twice
+        let mut results1 = Vec::new();
+        let mut results2 = Vec::new();
+
+        for trial in 0..2 {
+            let mut optimizer = NovoGradOptimizer::new(config.clone())?;
+            let param = Tensor::from_data(vec![1.0], &[1], &device)?;
+            let grad = Tensor::from_data(vec![0.5], &[1], &device)?;
+
+            let mut current_param = param;
+            for _step in 0..5 {
+                current_param = optimizer.step_param("param", &current_param, &grad)?;
+                let value = current_param.to_cpu()?[0];
+                
+                if trial == 0 {
+                    results1.push(value);
+                } else {
+                    results2.push(value);
+                }
+            }
+        }
+
+        // Results should be identical (deterministic)
+        assert_eq!(results1.len(), results2.len());
+        for (r1, r2) in results1.iter().zip(results2.iter()) {
+            assert!(
+                (r1 - r2).abs() < 1e-6,
+                "Results should be deterministic: {} vs {}",
+                r1, r2
+            );
+        }
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_novograd_large_batch_stability() -> Result<()> {
+        // NovoGrad is particularly designed for large batch training
+        let config = NovoGradConfig {
+            learning_rate: 0.01,
+            beta1: 0.95,
+            beta2: 0.98,
+            eps: 1e-8,
+            weight_decay: 0.0,
+            grad_averaging: false,
+        };
+
+        let mut optimizer = NovoGradOptimizer::new(config)?;
+        let device = Device::cpu();
+
+        // Simulate large batch gradients (high variance)
+        let param = Tensor::from_data(vec![1.0, 2.0, 3.0], &[3], &device)?;
+        
+        let batch_gradients = vec![
+            Tensor::from_data(vec![1.0, 0.1, 0.01], &[3], &device)?,     // Different scales
+            Tensor::from_data(vec![0.01, 1.0, 0.1], &[3], &device)?,
+            Tensor::from_data(vec![0.1, 0.01, 1.0], &[3], &device)?,
+        ];
+
+        let mut current_param = param;
+        for (step, grad) in batch_gradients.iter().enumerate() {
+            let prev_param = current_param.clone();
+            current_param = optimizer.step_param("param", &current_param, grad)?;
+            
+            // Should remain stable and not explode
+            let param_data = current_param.to_cpu()?;
+            for &value in param_data.iter() {
+                assert!(value.is_finite(), "Parameter should remain finite at step {}", step);
+                assert!(value.abs() < 100.0, "Parameter should not explode at step {}", step);
+            }
+        }
+
+        Ok(())
+    }
+}