48 lines
1.8 KiB
Rust
48 lines
1.8 KiB
Rust
//! Expert dropout standalone test
|
|
|
|
use rtx_transformers::layers::{
|
|
ExpertDropoutConfig, DropoutStrategy, ExpertDropoutLayer,
|
|
MoEConfig, ExpertOutputs, DropoutScheduler, ExpertImportanceScorer
|
|
};
|
|
use rtx_tensor::{Tensor, Device, DType};
|
|
|
|
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
// Test basic functionality
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let moe_config = MoEConfig::new(8, 2, 768, 3072);
|
|
let dropout_config = ExpertDropoutConfig::new(0.2, DropoutStrategy::Random);
|
|
|
|
let dropout_layer = ExpertDropoutLayer::new(dropout_config, moe_config, &device)?;
|
|
println!("✓ Expert dropout layer created successfully");
|
|
println!(" Number of experts: {}", dropout_layer.num_experts());
|
|
println!(" Training mode: {}", dropout_layer.is_training());
|
|
|
|
// Test different strategies
|
|
let strategies = vec![
|
|
DropoutStrategy::Random,
|
|
DropoutStrategy::Block,
|
|
DropoutStrategy::Progressive,
|
|
DropoutStrategy::LoadAware,
|
|
];
|
|
|
|
for strategy in strategies {
|
|
let config = ExpertDropoutConfig::new(0.1, strategy.clone());
|
|
assert!(config.validate().is_ok());
|
|
println!("✓ Strategy {:?} validation passed", strategy);
|
|
}
|
|
|
|
// Test dropout scheduler
|
|
let scheduler = DropoutScheduler::new(0.5, 0.1, 1000);
|
|
println!("✓ Dropout scheduler created");
|
|
println!(" Initial rate: {}", scheduler.get_dropout_rate(0));
|
|
println!(" Mid rate: {}", scheduler.get_dropout_rate(500));
|
|
println!(" Final rate: {}", scheduler.get_dropout_rate(1000));
|
|
|
|
// Test expert importance scorer
|
|
let scorer = ExpertImportanceScorer::new(8);
|
|
println!("✓ Expert importance scorer created");
|
|
println!(" Initial scores: {:?}", scorer.get_importance_scores());
|
|
|
|
println!("\n🎉 All expert dropout tests passed!");
|
|
Ok(())
|
|
} |