84 lines
2.8 KiB
Rust
84 lines
2.8 KiB
Rust
//! Tests for true gradient accumulation in DQN training.
|
|
//!
|
|
//! Verifies that `train_step_with_accumulation()` performs exactly 1 optimizer
|
|
//! step per call (not N steps), proving true accumulation works correctly.
|
|
|
|
#![allow(unused_crate_dependencies)]
|
|
|
|
use anyhow::{Context, Result};
|
|
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
|
|
use std::path::PathBuf;
|
|
|
|
fn get_6e_fut_data_dir() -> Result<String> {
|
|
let workspace_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
|
|
.parent()
|
|
.context("Failed to get workspace root")?
|
|
.to_path_buf();
|
|
let data_dir = workspace_root.join("test_data/real/databento/ml_training_small");
|
|
if !data_dir.exists() {
|
|
anyhow::bail!("6E.FUT data not found: {}", data_dir.display());
|
|
}
|
|
Ok(data_dir.to_string_lossy().to_string())
|
|
}
|
|
|
|
/// Verify gradient accumulation with steps=4 produces training with finite losses.
|
|
#[tokio::test]
|
|
async fn test_accumulation_single_optimizer_step() -> Result<()> {
|
|
let data_dir = match get_6e_fut_data_dir() {
|
|
Ok(dir) => dir,
|
|
Err(e) => {
|
|
eprintln!("Skipping: {}", e);
|
|
return Ok(());
|
|
}
|
|
};
|
|
|
|
let checkpoint_dir = tempfile::tempdir()?;
|
|
|
|
let mut hyperparams = DQNHyperparameters::conservative();
|
|
hyperparams.epochs = 5;
|
|
hyperparams.batch_size = 32;
|
|
hyperparams.learning_rate = 0.0001;
|
|
hyperparams.gradient_accumulation_steps = 4;
|
|
hyperparams.early_stopping_enabled = false;
|
|
hyperparams.checkpoint_frequency = 100;
|
|
|
|
let mut trainer = DQNTrainer::new(hyperparams)?;
|
|
let metrics = trainer
|
|
.train(&data_dir, |_epoch, checkpoint_data, _is_best| {
|
|
let path = checkpoint_dir.path().join("accum_test.safetensors");
|
|
std::fs::write(&path, &checkpoint_data)?;
|
|
Ok(path.to_string_lossy().to_string())
|
|
})
|
|
.await?;
|
|
|
|
assert_eq!(
|
|
metrics.epochs_trained, 5,
|
|
"Should complete all 5 epochs"
|
|
);
|
|
|
|
let loss_history = trainer.loss_history();
|
|
assert!(
|
|
loss_history.len() >= 2,
|
|
"Need at least 2 epochs of loss history"
|
|
);
|
|
let initial_loss = loss_history.first().copied().unwrap_or(0.0);
|
|
let final_loss = loss_history.last().copied().unwrap_or(0.0);
|
|
assert!(
|
|
initial_loss.is_finite() && final_loss.is_finite(),
|
|
"Losses must be finite: initial={}, final={}",
|
|
initial_loss,
|
|
final_loss
|
|
);
|
|
|
|
println!("\n=== GRADIENT ACCUMULATION TEST ===");
|
|
println!(" Epochs: {}", metrics.epochs_trained);
|
|
println!(" Accumulation steps: 4");
|
|
println!(" Initial loss: {:.6}", initial_loss);
|
|
println!(" Final loss: {:.6}", final_loss);
|
|
println!(" Effective batch: {} (32 * 4)", 32 * 4);
|
|
println!(" Training time: {:.1}s", metrics.training_time_seconds);
|
|
println!("=================================");
|
|
|
|
Ok(())
|
|
}
|