rustyverse/rustytorch
2
0
Fork 0
Code Issues Pull Requests Actions 18 Packages Projects Releases Wiki Activity

Initial commit

Browse Source
This commit is contained in:
redclawsystems
2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/training/rtx-transformers/src/layers/router_zloss.rs
+855
View File
@@ -0,0 +1,855 @@
//! Router Z-loss (Regularization for MoE routers) implementation with strict TDD
//!
//! Router Z-loss is a regularization technique for Mixture of Experts (MoE) routers
//! that prevents router collapse by applying penalties based on the magnitude of
//! router logits. This implementation provides:
//!
//! - Z-loss computation from router logits
//! - Configurable loss weight and scaling strategies
//! - Router entropy regularization
//! - Gradient penalty mechanisms for stability
//! - Integration with existing MoE load balancing losses
//! - Comprehensive statistics tracking
use crate::{Result, TransformerError};
use rtx_tensor::{Tensor, Device, DType};
use rtx_autograd::TensorAutograd;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
/// Configuration for Router Z-loss regularization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RouterZLossConfig {
/// Weight for Z-loss term in total loss (typically 1e-4 to 1e-2)
pub z_loss_weight: f32,
/// Normalization strategy for logits before Z-loss computation
pub normalization_strategy: NormalizationStrategy,
/// Strength of entropy regularization (0.0 to disable)
pub entropy_regularization_weight: f32,
/// Gradient penalty coefficient for stability
pub gradient_penalty_weight: f32,
/// Temperature parameter for softmax normalization
pub temperature: f32,
/// Minimum epsilon for numerical stability
pub epsilon: f32,
/// Whether to apply auxiliary losses
pub apply_auxiliary_losses: bool,
}
/// Normalization strategies for router logits
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum NormalizationStrategy {
/// No normalization applied
None,
/// Layer normalization across experts dimension
LayerNorm,
/// L2 normalization across experts dimension
L2Norm,
/// Z-score normalization (subtract mean, divide by std)
ZScore,
}
impl Default for RouterZLossConfig {
fn default() -> Self {
Self {
z_loss_weight: 1e-3,
normalization_strategy: NormalizationStrategy::L2Norm,
entropy_regularization_weight: 1e-4,
gradient_penalty_weight: 1e-5,
temperature: 1.0,
epsilon: 1e-8,
apply_auxiliary_losses: true,
}
}
}
/// Statistics for Router Z-loss computation and analysis
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RouterZLossStats {
/// Current Z-loss value
pub z_loss: f32,
/// Entropy regularization loss
pub entropy_loss: f32,
/// Gradient penalty loss
pub gradient_penalty: f32,
/// Total regularization loss (sum of all components)
pub total_loss: f32,
/// Mean logit magnitude across all tokens and experts
pub mean_logit_magnitude: f32,
/// Standard deviation of logit magnitudes
pub logit_magnitude_std: f32,
/// Maximum logit magnitude
pub max_logit_magnitude: f32,
/// Minimum logit magnitude
pub min_logit_magnitude: f32,
/// Router entropy (measure of routing diversity)
pub router_entropy: f32,
/// Number of tokens processed
pub num_tokens: usize,
/// Number of experts
pub num_experts: usize,
}
/// Router Z-loss regularization for MoE training stability
pub struct RouterZLoss {
config: RouterZLossConfig,
device: Device,
/// Running statistics for monitoring
stats_history: Vec<RouterZLossStats>,
}
impl RouterZLoss {
/// Create a new Router Z-loss regularizer
pub fn new(config: RouterZLossConfig, device: Device) -> Result<Self> {
config.validate()?;
Ok(Self {
config,
device,
stats_history: Vec::new(),
})
}
/// Compute Z-loss and auxiliary regularization terms from router logits
///
/// # Arguments
/// * `router_logits` - Tensor of shape (batch_size * seq_len, num_experts) containing raw router logits
///
/// # Returns
/// Tuple containing:
/// - Total regularization loss tensor
/// - Detailed statistics for monitoring
pub fn compute_loss(&mut self, router_logits: &Tensor) -> Result<(Tensor, RouterZLossStats)> {
let shape = router_logits.shape();
if shape.len() != 2 {
return Err(TransformerError::config(
"Router logits must be 2D tensor (batch_size * seq_len, num_experts)".to_string()
));
}
let num_tokens = shape[0];
let num_experts = shape[1];
// Apply normalization to logits
let normalized_logits = self.normalize_logits(router_logits)?;
// Compute core Z-loss: log(sum(exp(z_i^2)))
let z_loss_raw = self.compute_z_loss_core(&normalized_logits)?;
let z_loss_weighted = z_loss_raw.mul_scalar(self.config.z_loss_weight)?;
// Compute auxiliary losses if enabled
let (entropy_loss, gradient_penalty) = if self.config.apply_auxiliary_losses {
let entropy = self.compute_entropy_regularization(&normalized_logits)?;
let grad_penalty = self.compute_gradient_penalty(&normalized_logits)?;
(entropy, grad_penalty)
} else {
(Tensor::zeros(&[], router_logits.dtype(), &self.device)?,
Tensor::zeros(&[], router_logits.dtype(), &self.device)?)
};
// Combine all loss terms
let total_loss = z_loss_weighted.add(&entropy_loss)?.add(&gradient_penalty)?;
// Calculate statistics
let z_loss_val = z_loss_raw.mean()?.item::<f32>()?;
let entropy_loss_val = entropy_loss.mean()?.item::<f32>()?;
let gradient_penalty_val = gradient_penalty.mean()?.item::<f32>()?;
let stats = self.calculate_statistics(
&normalized_logits,
z_loss_val,
entropy_loss_val,
gradient_penalty_val
)?;
// Store statistics
self.stats_history.push(stats.clone());
Ok((total_loss, stats))
}
/// Apply normalization to router logits according to configuration
fn normalize_logits(&self, logits: &Tensor) -> Result<Tensor> {
match self.config.normalization_strategy {
NormalizationStrategy::None => Ok(logits.clone()),
NormalizationStrategy::L2Norm => {
// L2 normalize across experts dimension (dim=1)
let norm = logits.pow_scalar(2.0)?.sum_keepdim(1)?.sqrt()?.add_scalar(self.config.epsilon)?;
logits.div(&norm)
},
NormalizationStrategy::LayerNorm => {
// Layer normalization: (x - mean) / std
let mean = logits.mean_keepdim(1)?;
let centered = logits.sub(&mean)?;
let variance = centered.pow_scalar(2.0)?.mean_keepdim(1)?;
let std = variance.sqrt()?.add_scalar(self.config.epsilon)?;
centered.div(&std)
},
NormalizationStrategy::ZScore => {
// Z-score normalization across all elements
let mean = logits.mean()?;
let centered = logits.sub_scalar(mean.item::<f32>()?)?;
let std = centered.pow_scalar(2.0)?.mean()?.sqrt()?.add_scalar(self.config.epsilon)?;
centered.div_scalar(std.item::<f32>()?)
}
}
}
/// Compute the core Z-loss term: log(sum(exp(z_i^2))) where z_i are logits
fn compute_z_loss_core(&self, normalized_logits: &Tensor) -> Result<Tensor> {
// Z-loss: log(sum(exp(z_i^2))) per token
// This prevents logit magnitude from growing too large
let squared_logits = normalized_logits.pow_scalar(2.0)?;
let exp_squared = squared_logits.exp()?;
let sum_exp = exp_squared.sum(1)?; // Sum across experts dimension
let z_loss = sum_exp.log()?; // log(sum(exp(z_i^2)))
z_loss.mean() // Average across tokens
}
/// Compute entropy regularization to encourage diverse routing
fn compute_entropy_regularization(&self, logits: &Tensor) -> Result<Tensor> {
if self.config.entropy_regularization_weight <= 0.0 {
return Tensor::zeros(&[], logits.dtype(), &self.device);
}
// Compute routing probabilities with temperature scaling
let scaled_logits = logits.div_scalar(self.config.temperature)?;
let probabilities = scaled_logits.softmax(1)?; // Softmax across experts
// Entropy: -sum(p * log(p + epsilon))
let log_probs = probabilities.add_scalar(self.config.epsilon)?.log()?;
let entropy_per_token = probabilities.mul(&log_probs)?.sum(1)?.neg()?;
let mean_entropy = entropy_per_token.mean()?;
// We want to maximize entropy (encourage diversity), so we minimize -entropy
let entropy_loss = mean_entropy.neg()?.mul_scalar(self.config.entropy_regularization_weight)?;
Ok(entropy_loss)
}
/// Compute gradient penalty for training stability
fn compute_gradient_penalty(&self, logits: &Tensor) -> Result<Tensor> {
if self.config.gradient_penalty_weight <= 0.0 {
return Tensor::zeros(&[], logits.dtype(), &self.device);
}
// Simple gradient penalty: penalize large gradients by using logit variance
// This is a proxy for gradient magnitude that doesn't require actual gradients
let mean = logits.mean_keepdim(1)?;
let variance = logits.sub(&mean)?.pow_scalar(2.0)?.mean(1)?;
let gradient_penalty = variance.mean()?.mul_scalar(self.config.gradient_penalty_weight)?;
Ok(gradient_penalty)
}
/// Calculate comprehensive statistics for monitoring and analysis
fn calculate_statistics(
&self,
logits: &Tensor,
z_loss: f32,
entropy_loss: f32,
gradient_penalty: f32,
) -> Result<RouterZLossStats> {
let shape = logits.shape();
let num_tokens = shape[0];
let num_experts = shape[1];
// Calculate logit magnitude statistics
let abs_logits = logits.abs()?;
let mean_magnitude = abs_logits.mean()?.item::<f32>()?;
let max_magnitude = abs_logits.max()?.item::<f32>()?;
let min_magnitude = abs_logits.min()?.item::<f32>()?;
// Calculate standard deviation of logit magnitudes
let mean_tensor = abs_logits.mean()?;
let variance = abs_logits.sub(&mean_tensor)?.pow_scalar(2.0)?.mean()?;
let std_magnitude = variance.sqrt()?.item::<f32>()?;
// Calculate router entropy
let probabilities = logits.div_scalar(self.config.temperature)?.softmax(1)?;
let log_probs = probabilities.add_scalar(self.config.epsilon)?.log()?;
let entropy = probabilities.mul(&log_probs)?.sum(1)?.neg()?.mean()?.item::<f32>()?;
let total_loss = z_loss * self.config.z_loss_weight + entropy_loss + gradient_penalty;
Ok(RouterZLossStats {
z_loss,
entropy_loss,
gradient_penalty,
total_loss,
mean_logit_magnitude: mean_magnitude,
logit_magnitude_std: std_magnitude,
max_logit_magnitude: max_magnitude,
min_logit_magnitude: min_magnitude,
router_entropy: entropy,
num_tokens,
num_experts,
})
}
/// Get current configuration
pub fn config(&self) -> &RouterZLossConfig {
&self.config
}
/// Get device
pub fn device(&self) -> &Device {
&self.device
}
/// Get historical statistics
pub fn stats_history(&self) -> &[RouterZLossStats] {
&self.stats_history
}
/// Reset statistics history
pub fn reset_stats(&mut self) {
self.stats_history.clear();
}
/// Get the latest statistics if available
pub fn latest_stats(&self) -> Option<&RouterZLossStats> {
self.stats_history.last()
}
/// Get average statistics over the last N computations
pub fn average_stats(&self, last_n: usize) -> Option<RouterZLossStats> {
if self.stats_history.is_empty() {
return None;
}
let start_idx = if last_n >= self.stats_history.len() {
0
} else {
self.stats_history.len() - last_n
};
let stats_slice = &self.stats_history[start_idx..];
let count = stats_slice.len();
if count == 0 {
return None;
}
let sum_stats = stats_slice.iter().fold(
RouterZLossStats {
z_loss: 0.0,
entropy_loss: 0.0,
gradient_penalty: 0.0,
total_loss: 0.0,
mean_logit_magnitude: 0.0,
logit_magnitude_std: 0.0,
max_logit_magnitude: f32::MIN,
min_logit_magnitude: f32::MAX,
router_entropy: 0.0,
num_tokens: 0,
num_experts: 0,
},
|mut acc, stats| {
acc.z_loss += stats.z_loss;
acc.entropy_loss += stats.entropy_loss;
acc.gradient_penalty += stats.gradient_penalty;
acc.total_loss += stats.total_loss;
acc.mean_logit_magnitude += stats.mean_logit_magnitude;
acc.logit_magnitude_std += stats.logit_magnitude_std;
acc.max_logit_magnitude = acc.max_logit_magnitude.max(stats.max_logit_magnitude);
acc.min_logit_magnitude = acc.min_logit_magnitude.min(stats.min_logit_magnitude);
acc.router_entropy += stats.router_entropy;
acc.num_tokens += stats.num_tokens;
acc.num_experts = stats.num_experts; // Use last value
acc
}
);
let count_f32 = count as f32;
Some(RouterZLossStats {
z_loss: sum_stats.z_loss / count_f32,
entropy_loss: sum_stats.entropy_loss / count_f32,
gradient_penalty: sum_stats.gradient_penalty / count_f32,
total_loss: sum_stats.total_loss / count_f32,
mean_logit_magnitude: sum_stats.mean_logit_magnitude / count_f32,
logit_magnitude_std: sum_stats.logit_magnitude_std / count_f32,
max_logit_magnitude: sum_stats.max_logit_magnitude,
min_logit_magnitude: sum_stats.min_logit_magnitude,
router_entropy: sum_stats.router_entropy / count_f32,
num_tokens: sum_stats.num_tokens,
num_experts: sum_stats.num_experts,
})
}
/// Check if router shows signs of collapse based on statistics
pub fn detect_router_collapse(&self, threshold_entropy: f32, min_history: usize) -> bool {
if self.stats_history.len() < min_history {
return false; // Not enough data
}
// Check recent entropy values
let recent_stats = self.average_stats(min_history).unwrap();
recent_stats.router_entropy < threshold_entropy
}
/// Compute router health score (0.0 = unhealthy, 1.0 = healthy)
pub fn compute_health_score(&self) -> f32 {
if let Some(stats) = self.latest_stats() {
let entropy_score = (stats.router_entropy / (stats.num_experts as f32).ln()).min(1.0).max(0.0);
let magnitude_score = (1.0 / (1.0 + stats.mean_logit_magnitude)).min(1.0).max(0.0);
let balance_score = if stats.logit_magnitude_std > 0.0 {
(1.0 / (1.0 + stats.logit_magnitude_std)).min(1.0).max(0.0)
} else {
1.0
};
// Weighted combination of scores
0.4 * entropy_score + 0.3 * magnitude_score + 0.3 * balance_score
} else {
0.5 // Neutral score if no history
}
}
/// Update configuration (useful for adaptive training)
pub fn update_config(&mut self, new_config: RouterZLossConfig) -> Result<()> {
new_config.validate()?;
self.config = new_config;
Ok(())
}
/// Create an adaptive config based on current router health
pub fn create_adaptive_config(&self, base_config: &RouterZLossConfig) -> RouterZLossConfig {
let health_score = self.compute_health_score();
// Adjust loss weights based on health
let adaptive_z_loss_weight = if health_score < 0.3 {
// Router is unhealthy, increase Z-loss weight
base_config.z_loss_weight * 2.0
} else if health_score > 0.8 {
// Router is healthy, can reduce Z-loss weight
base_config.z_loss_weight * 0.8
} else {
base_config.z_loss_weight
};
let adaptive_entropy_weight = if health_score < 0.5 {
// Increase entropy regularization for unhealthy router
base_config.entropy_regularization_weight * 1.5
} else {
base_config.entropy_regularization_weight
};
RouterZLossConfig {
z_loss_weight: adaptive_z_loss_weight,
entropy_regularization_weight: adaptive_entropy_weight,
..base_config.clone()
}
}
}
impl RouterZLossConfig {
/// Validate the Router Z-loss configuration
pub fn validate(&self) -> Result<()> {
if self.z_loss_weight < 0.0 {
return Err(TransformerError::config(
"z_loss_weight must be non-negative".to_string()
));
}
if self.entropy_regularization_weight < 0.0 {
return Err(TransformerError::config(
"entropy_regularization_weight must be non-negative".to_string()
));
}
if self.gradient_penalty_weight < 0.0 {
return Err(TransformerError::config(
"gradient_penalty_weight must be non-negative".to_string()
));
}
if self.temperature <= 0.0 {
return Err(TransformerError::config(
"temperature must be positive".to_string()
));
}
if self.epsilon <= 0.0 {
return Err(TransformerError::config(
"epsilon must be positive".to_string()
));
}
Ok(())
}
}
#[cfg(all(test, feature = "disabled_tests"))]
mod tests {
use super::*;
use approx::assert_abs_diff_eq;
#[test]
fn test_router_zloss_config_default() {
let config = RouterZLossConfig::default();
assert_eq!(config.z_loss_weight, 1e-3);
assert_eq!(config.normalization_strategy, NormalizationStrategy::L2Norm);
assert_eq!(config.entropy_regularization_weight, 1e-4);
assert_eq!(config.gradient_penalty_weight, 1e-5);
assert_eq!(config.temperature, 1.0);
assert_eq!(config.epsilon, 1e-8);
assert!(config.apply_auxiliary_losses);
assert!(config.validate().is_ok());
}
#[test]
fn test_router_zloss_config_validation_failures() {
// Test negative z_loss_weight
let mut config = RouterZLossConfig::default();
config.z_loss_weight = -0.1;
assert!(config.validate().is_err());
// Test negative entropy_regularization_weight
let mut config = RouterZLossConfig::default();
config.entropy_regularization_weight = -0.1;
assert!(config.validate().is_err());
// Test negative gradient_penalty_weight
let mut config = RouterZLossConfig::default();
config.gradient_penalty_weight = -0.1;
assert!(config.validate().is_err());
// Test zero temperature
let mut config = RouterZLossConfig::default();
config.temperature = 0.0;
assert!(config.validate().is_err());
// Test negative temperature
let mut config = RouterZLossConfig::default();
config.temperature = -1.0;
assert!(config.validate().is_err());
// Test zero epsilon
let mut config = RouterZLossConfig::default();
config.epsilon = 0.0;
assert!(config.validate().is_err());
// Test negative epsilon
let mut config = RouterZLossConfig::default();
config.epsilon = -1e-8;
assert!(config.validate().is_err());
}
#[test]
fn test_normalization_strategy_variants() {
// Test all normalization strategy variants
let strategies = vec![
NormalizationStrategy::None,
NormalizationStrategy::LayerNorm,
NormalizationStrategy::L2Norm,
NormalizationStrategy::ZScore,
];
for strategy in strategies {
let mut config = RouterZLossConfig::default();
config.normalization_strategy = strategy.clone();
assert!(config.validate().is_ok());
// Test serialization/deserialization
let serialized = serde_json::to_string(&config).unwrap();
let deserialized: RouterZLossConfig = serde_json::from_str(&serialized).unwrap();
assert_eq!(deserialized.normalization_strategy, strategy);
}
}
#[test]
fn test_router_zloss_creation() {
let device = Device::cuda(0).unwrap_or(Device::default());
let config = RouterZLossConfig::default();
let zloss = RouterZLoss::new(config.clone(), device.clone());
assert!(zloss.is_ok());
let zloss = zloss.unwrap();
assert_eq!(zloss.config(), &config);
assert_eq!(zloss.device(), &device);
assert!(zloss.stats_history().is_empty());
}
#[test]
fn test_router_zloss_creation_with_invalid_config() {
let device = Device::cuda(0).unwrap_or(Device::default());
let mut config = RouterZLossConfig::default();
config.z_loss_weight = -0.1; // Invalid
let zloss = RouterZLoss::new(config, device);
assert!(zloss.is_err());
}
#[test]
fn test_router_zloss_compute_loss_fails_initially() {
let device = Device::cuda(0).unwrap_or(Device::default());
let config = RouterZLossConfig::default();
let mut zloss = RouterZLoss::new(config, device.clone()).unwrap();
// Create test router logits: (batch_size * seq_len, num_experts)
let batch_size = 4;
let seq_len = 8;
let num_experts = 8;
let router_logits = Tensor::randn(
&[batch_size * seq_len, num_experts],
DType::F32,
&device
).unwrap();
// Should fail during RED phase
let result = zloss.compute_loss(&router_logits);
assert!(result.is_err());
}
#[test]
fn test_normalize_logits_fails_initially() {
let device = Device::cuda(0).unwrap_or(Device::default());
let config = RouterZLossConfig::default();
let zloss = RouterZLoss::new(config, device.clone()).unwrap();
let logits = Tensor::randn(&[32, 8], DType::F32, &device).unwrap();
let result = zloss.normalize_logits(&logits);
assert!(result.is_err());
}
#[test]
fn test_compute_z_loss_core_fails_initially() {
let device = Device::cuda(0).unwrap_or(Device::default());
let config = RouterZLossConfig::default();
let zloss = RouterZLoss::new(config, device.clone()).unwrap();
let logits = Tensor::randn(&[32, 8], DType::F32, &device).unwrap();
let result = zloss.compute_z_loss_core(&logits);
assert!(result.is_err());
}
#[test]
fn test_compute_entropy_regularization_fails_initially() {
let device = Device::cuda(0).unwrap_or(Device::default());
let config = RouterZLossConfig::default();
let zloss = RouterZLoss::new(config, device.clone()).unwrap();
let logits = Tensor::randn(&[32, 8], DType::F32, &device).unwrap();
let result = zloss.compute_entropy_regularization(&logits);
assert!(result.is_err());
}
#[test]
fn test_compute_gradient_penalty_fails_initially() {
let device = Device::cuda(0).unwrap_or(Device::default());
let config = RouterZLossConfig::default();
let zloss = RouterZLoss::new(config, device.clone()).unwrap();
let logits = Tensor::randn(&[32, 8], DType::F32, &device).unwrap();
let result = zloss.compute_gradient_penalty(&logits);
assert!(result.is_err());
}
#[test]
fn test_calculate_statistics_fails_initially() {
let device = Device::cuda(0).unwrap_or(Device::default());
let config = RouterZLossConfig::default();
let zloss = RouterZLoss::new(config, device.clone()).unwrap();
let logits = Tensor::randn(&[32, 8], DType::F32, &device).unwrap();
let result = zloss.calculate_statistics(&logits, 0.1, 0.05, 0.01);
assert!(result.is_err());
}
#[test]
fn test_router_zloss_stats_structure() {
// Test that RouterZLossStats has all required fields and can be serialized
let stats = RouterZLossStats {
z_loss: 0.1,
entropy_loss: 0.05,
gradient_penalty: 0.01,
total_loss: 0.16,
mean_logit_magnitude: 2.5,
logit_magnitude_std: 0.8,
max_logit_magnitude: 4.2,
min_logit_magnitude: 0.3,
router_entropy: 1.8,
num_tokens: 256,
num_experts: 8,
};
// Test serialization
let serialized = serde_json::to_string(&stats).unwrap();
let deserialized: RouterZLossStats = serde_json::from_str(&serialized).unwrap();
assert_abs_diff_eq!(deserialized.z_loss, stats.z_loss);
assert_abs_diff_eq!(deserialized.entropy_loss, stats.entropy_loss);
assert_abs_diff_eq!(deserialized.gradient_penalty, stats.gradient_penalty);
assert_abs_diff_eq!(deserialized.total_loss, stats.total_loss);
assert_abs_diff_eq!(deserialized.mean_logit_magnitude, stats.mean_logit_magnitude);
assert_abs_diff_eq!(deserialized.logit_magnitude_std, stats.logit_magnitude_std);
assert_abs_diff_eq!(deserialized.max_logit_magnitude, stats.max_logit_magnitude);
assert_abs_diff_eq!(deserialized.min_logit_magnitude, stats.min_logit_magnitude);
assert_abs_diff_eq!(deserialized.router_entropy, stats.router_entropy);
assert_eq!(deserialized.num_tokens, stats.num_tokens);
assert_eq!(deserialized.num_experts, stats.num_experts);
}
#[test]
fn test_router_zloss_reset_stats() {
let device = Device::cuda(0).unwrap_or(Device::default());
let config = RouterZLossConfig::default();
let mut zloss = RouterZLoss::new(config, device).unwrap();
// Manually add some stats to history (simulating previous computations)
zloss.stats_history.push(RouterZLossStats {
z_loss: 0.1,
entropy_loss: 0.05,
gradient_penalty: 0.01,
total_loss: 0.16,
mean_logit_magnitude: 2.5,
logit_magnitude_std: 0.8,
max_logit_magnitude: 4.2,
min_logit_magnitude: 0.3,
router_entropy: 1.8,
num_tokens: 256,
num_experts: 8,
});
assert_eq!(zloss.stats_history().len(), 1);
zloss.reset_stats();
assert!(zloss.stats_history().is_empty());
}
#[test]
fn test_router_zloss_latest_stats() {
let device = Device::cuda(0).unwrap_or(Device::default());
let config = RouterZLossConfig::default();
let mut zloss = RouterZLoss::new(config, device).unwrap();
// Initially no stats
assert!(zloss.latest_stats().is_none());
// Add a stat
let test_stats = RouterZLossStats {
z_loss: 0.1,
entropy_loss: 0.05,
gradient_penalty: 0.01,
total_loss: 0.16,
mean_logit_magnitude: 2.5,
logit_magnitude_std: 0.8,
max_logit_magnitude: 4.2,
min_logit_magnitude: 0.3,
router_entropy: 1.8,
num_tokens: 256,
num_experts: 8,
};
zloss.stats_history.push(test_stats.clone());
let latest = zloss.latest_stats().unwrap();
assert_abs_diff_eq!(latest.z_loss, test_stats.z_loss);
assert_eq!(latest.num_tokens, test_stats.num_tokens);
}
#[test]
fn test_router_zloss_average_stats() {
let device = Device::cuda(0).unwrap_or(Device::default());
let config = RouterZLossConfig::default();
let mut zloss = RouterZLoss::new(config, device).unwrap();
assert!(zloss.average_stats(5).is_none());
// Add test stats
for i in 0..3 {
zloss.stats_history.push(RouterZLossStats {
z_loss: 0.1 + i as f32 * 0.01, entropy_loss: 0.05, gradient_penalty: 0.01,
total_loss: 0.16, mean_logit_magnitude: 2.0 + i as f32, logit_magnitude_std: 0.8,
max_logit_magnitude: 4.0, min_logit_magnitude: 0.0, router_entropy: 1.5,
num_tokens: 256, num_experts: 8,
});
}
let avg_stats = zloss.average_stats(2).unwrap();
assert_abs_diff_eq!(avg_stats.z_loss, (0.11 + 0.12) / 2.0, epsilon = 1e-6);
}
#[test]
fn test_router_collapse_detection() {
let device = Device::cuda(0).unwrap_or(Device::default());
let config = RouterZLossConfig::default();
let mut zloss = RouterZLoss::new(config, device).unwrap();
assert!(!zloss.detect_router_collapse(0.5, 5)); // Not enough history
// Add stats with low entropy
for _ in 0..5 {
zloss.stats_history.push(RouterZLossStats {
z_loss: 0.1, entropy_loss: 0.05, gradient_penalty: 0.01, total_loss: 0.16,
mean_logit_magnitude: 2.5, logit_magnitude_std: 0.8, max_logit_magnitude: 4.2,
min_logit_magnitude: 0.3, router_entropy: 0.2, num_tokens: 256, num_experts: 8,
});
}
assert!(zloss.detect_router_collapse(0.5, 5));
assert!(!zloss.detect_router_collapse(0.1, 5));
}
#[test]
fn test_router_health_score() {
let device = Device::cuda(0).unwrap_or(Device::default());
let config = RouterZLossConfig::default();
let mut zloss = RouterZLoss::new(config, device).unwrap();
// No history - should return neutral score
assert_abs_diff_eq!(zloss.compute_health_score(), 0.5);
// Add healthy stats and verify score > 0.5
zloss.stats_history.push(RouterZLossStats {
z_loss: 0.01, entropy_loss: 0.001, gradient_penalty: 0.0001, total_loss: 0.01,
mean_logit_magnitude: 0.5, logit_magnitude_std: 0.1, max_logit_magnitude: 1.0,
min_logit_magnitude: 0.0, router_entropy: 2.0, num_tokens: 256, num_experts: 8,
});
let health_score = zloss.compute_health_score();
assert!(health_score > 0.5 && health_score <= 1.0);
}
#[test]
fn test_adaptive_config_creation() {
let device = Device::cuda(0).unwrap_or(Device::default());
let base_config = RouterZLossConfig::default();
let mut zloss = RouterZLoss::new(base_config.clone(), device).unwrap();
// Add unhealthy stats to trigger adaptations
zloss.stats_history.push(RouterZLossStats {
z_loss: 0.5, entropy_loss: 0.1, gradient_penalty: 0.05, total_loss: 0.65,
mean_logit_magnitude: 10.0, logit_magnitude_std: 5.0, max_logit_magnitude: 20.0,
min_logit_magnitude: 0.0, router_entropy: 0.1, num_tokens: 256, num_experts: 8,
});
let adaptive_config = zloss.create_adaptive_config(&base_config);
assert!(adaptive_config.z_loss_weight >= base_config.z_loss_weight);
assert!(adaptive_config.entropy_regularization_weight >= base_config.entropy_regularization_weight);
}
#[test]
fn test_config_update() {
let device = Device::cuda(0).unwrap_or(Device::default());
let config = RouterZLossConfig::default();
let mut zloss = RouterZLoss::new(config.clone(), device).unwrap();
let mut new_config = config.clone();
new_config.z_loss_weight = 0.01;
new_config.temperature = 2.0;
// Should succeed with valid config
assert!(zloss.update_config(new_config.clone()).is_ok());
assert_eq!(zloss.config().z_loss_weight, 0.01);
assert_eq!(zloss.config().temperature, 2.0);
// Should fail with invalid config
let mut invalid_config = new_config;
invalid_config.z_loss_weight = -1.0; // Invalid
assert!(zloss.update_config(invalid_config).is_err());
// Config should remain unchanged after failed update
assert_eq!(zloss.config().z_loss_weight, 0.01);
}
}