gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(dqn): GPU PER monitoring funcs used empty CPU experiences, causing cuBLAS crash

Browse Source 
Four monitoring functions (estimate_avg_q_value_with_early_stopping,
collect_qvalue_statistics, compute_q_gap_for_epoch, compute_per_action_q_values)
iterated over batch_sample.experiences to build CPU tensors for forward passes.
When GPU PER (GpuPrioritized) is active, experiences is always vec![] — all data
lives on GPU tensors in gpu_batch. This created zero-element tensors with non-zero
shapes, triggering CUBLAS_STATUS_INVALID_VALUE on the next forward pass.

Fix: all four functions now check for gpu_batch.states and use it directly,
falling back to CPU experiences only for non-GPU buffers. Also removes debug
eprintln probes, wires new_on_device for DQN/RegimeConditionalDQN construction,
and adds GPU-aware smoke tests (33 pass, 874 total, 0 failures).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-09 10:47:46 +01:00
parent 0f45537e6a
commit d12273e164
17 changed files with 1148 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
crates/ml-core/src/mixed_precision.rs
View File

@@ -233,16 +233,11 @@ pub fn detect_from_gpu_name_auto() -> Option<MixedPrecisionConfig> {
pub fn training_dtype(device: &candle_core::Device) -> candle_core::DType {
match device {
candle_core::Device::Cuda(_) => {
if let Some(config) = detect_from_gpu_name_auto() {
match config.dtype {
DTypeSelection::BF16 => candle_core::DType::BF16,
// F16 mixed precision uses F32 weights + F16 forward;
// for *training dtype* we stay F32 (only Ampere+ gets BF16).
DTypeSelection::F16 => candle_core::DType::F32,
}
} else {
candle_core::DType::F32
}
// BF16 mixed precision is disabled for now: Candle's linear.forward()
// requires matching dtypes between weights and input, and several codepaths
// (GPU experience collector, portfolio simulator) assume F32 tensors.
// Re-enable once all GPU kernels support BF16. See comment at line 207.
candle_core::DType::F32
}
candle_core::Device::Cpu | candle_core::Device::Metal(_) => candle_core::DType::F32,
}

13
crates/ml-dqn/src/dqn.rs
View File

@@ -1093,8 +1093,17 @@ pub struct DQN {
}
impl DQN {
/// Create new `DQN`
/// Create new `DQN` with auto-detected device (GPU if available, CPU otherwise).
pub fn new(config: DQNConfig) -> Result<Self, MLError> {
let device = Device::cuda_if_available(0)?;
Self::new_on_device(config, device)
}
/// Create new `DQN` on a specific device.
///
/// Use this when the caller needs to control the device (e.g. trainer passes
/// its own device to keep network and data on the same device).
pub fn new_on_device(config: DQNConfig, device: Device) -> Result<Self, MLError> {
if config.state_dim == 0 {
return Err(MLError::ConfigError("DQN requires state_dim > 0".to_owned()));
}
@@ -1102,8 +1111,6 @@ impl DQN {
return Err(MLError::ConfigError("DQN requires num_actions > 0".to_owned()));
}
let device = Device::cuda_if_available(0)?; // Use GPU if available, fallback to CPU
// Create main Q-network
let q_network = Sequential::new(
config.state_dim,

12
crates/ml-dqn/src/regime_conditional.rs
View File

@@ -204,16 +204,20 @@ impl RegimeConditionalDQN {
/// Returns error if head creation fails
pub fn new(config: DQNConfig) -> Result<Self, MLError> {
let device = Device::cuda_if_available(0)?;
Self::new_on_device(config, device)
}
/// Create regime-conditional DQN on a specific device.
pub fn new_on_device(config: DQNConfig, device: Device) -> Result<Self, MLError> {
// Create 3 independent heads with shared memory
let trending_config = config.clone();
let ranging_config = config.clone();
let volatile_config = config.clone();
// Create heads - they'll each create their own memory based on config.use_per
let trending_head = DQN::new(trending_config)?;
let ranging_head = DQN::new(ranging_config)?;
let volatile_head = DQN::new(volatile_config)?;
// Create heads on the same device as the trainer
let trending_head = DQN::new_on_device(trending_config, device.clone())?;
let ranging_head = DQN::new_on_device(ranging_config, device.clone())?;
let volatile_head = DQN::new_on_device(volatile_config, device.clone())?;
// Note: Each head now has its own replay buffer (uniform or prioritized based on config)
// Shared memory across heads is not currently supported with ReplayBufferType enum

8
crates/ml/src/cuda_pipeline/common_device_functions.cuh
View File

@@ -11,12 +11,12 @@
/* State layout — overridable via NVRTC #define injection.
* Must match the actual feature buffer width (MARKET_DIM), network
* input dimension (STATE_DIM), and portfolio feature count (PORTFOLIO_DIM).
* Default: 40 basic market + 3 portfolio = 43, no alignment padding. */
* Default: 40 basic market + 3 portfolio = 43, tensor-core-aligned to 48. */
#ifndef STATE_DIM
#define STATE_DIM 54
#define STATE_DIM 48
#endif
#ifndef MARKET_DIM
#define MARKET_DIM 51
#define MARKET_DIM 40
#endif
#ifndef PORTFOLIO_DIM
#define PORTFOLIO_DIM 3
@@ -376,7 +376,7 @@ __device__ float diversity_entropy(
* Returns MSE prediction error clamped to [0, max_reward].
*
* Architecture: [CUR_INPUT=35] -> [CUR_HIDDEN=64] LReLU -> [CUR_OUTPUT=32]
* Input: 32 state features (first 32 of state[54]) + 3-class one-hot.
* Input: 32 state features (first 32 of state[STATE_DIM]) + 3-class one-hot.
* Action mapping: exposure/9: <=1 -> Short(0), ==2 -> Flat(1), >=3 -> Long(2).
*/
__device__ float curiosity_inference(

9
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -30,8 +30,9 @@ use super::gpu_weights::{
// Constants
// ---------------------------------------------------------------------------
/// Fallback STATE_DIM when not dynamically configured (original pre-OFI value).
const DEFAULT_STATE_DIM: usize = 54;
/// Fallback STATE_DIM when not dynamically configured.
/// 40 market + 3 portfolio = 43, tensor-core-aligned to 48.
const DEFAULT_STATE_DIM: usize = 48;
const MAX_EPISODES: usize = 256;
const MAX_TIMESTEPS: usize = 1000;
const PORTFOLIO_STATE_SIZE: usize = 8;
@@ -310,7 +311,7 @@ impl GpuExperienceCollector {
set
}
None => {
RmsNormWeightSet::ones(&stream)?
RmsNormWeightSet::ones(&stream, network_dims)?
}
};
let target_rmsnorm = match extract_rmsnorm_weights(target_vars, &stream)? {
@@ -319,7 +320,7 @@ impl GpuExperienceCollector {
set
}
None => {
RmsNormWeightSet::ones(&stream)?
RmsNormWeightSet::ones(&stream, network_dims)?
}
};

25
crates/ml/src/cuda_pipeline/gpu_weights.rs
View File

@@ -200,15 +200,22 @@ const RMSNORM_WEIGHT_NAMES: [&str; 4] = [
/// - advantage_rmsnorm:[128] (after advantage FC)
#[allow(missing_debug_implementations)]
pub struct RmsNormWeightSet {
pub gamma_s0: CudaSlice<f32>, // shared_rmsnorm_0.weight [256]
pub gamma_s1: CudaSlice<f32>, // shared_rmsnorm_1.weight [256]
pub gamma_v: CudaSlice<f32>, // value_rmsnorm.weight [128]
pub gamma_a: CudaSlice<f32>, // advantage_rmsnorm.weight [128]
pub gamma_s0: CudaSlice<f32>, // shared_rmsnorm_0.weight [shared_h1]
pub gamma_s1: CudaSlice<f32>, // shared_rmsnorm_1.weight [shared_h2]
pub gamma_v: CudaSlice<f32>, // value_rmsnorm.weight [value_h]
pub gamma_a: CudaSlice<f32>, // advantage_rmsnorm.weight [adv_h]
}
impl RmsNormWeightSet {
/// Create all-ones RMSNorm weights (no-op scaling, used when distributional is disabled).
pub fn ones(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
///
/// Dimensions must match the actual network layer sizes to avoid GPU OOB reads.
/// `network_dims`: `(shared_h1, shared_h2, value_h, adv_h)`.
pub fn ones(
stream: &Arc<CudaStream>,
network_dims: (usize, usize, usize, usize),
) -> Result<Self, MLError> {
let (shared_h1, shared_h2, value_h, adv_h) = network_dims;
let alloc_ones = |size: usize| -> Result<CudaSlice<f32>, MLError> {
let data = vec![1.0_f32; size];
let mut buf = stream
@@ -220,10 +227,10 @@ impl RmsNormWeightSet {
Ok(buf)
};
Ok(Self {
gamma_s0: alloc_ones(256)?,
gamma_s1: alloc_ones(256)?,
gamma_v: alloc_ones(128)?,
gamma_a: alloc_ones(128)?,
gamma_s0: alloc_ones(shared_h1)?,
gamma_s1: alloc_ones(shared_h2)?,
gamma_v: alloc_ones(value_h)?,
gamma_a: alloc_ones(adv_h)?,
})
}
}

35
crates/ml/src/cuda_pipeline/mod.rs
View File

@@ -922,8 +922,26 @@ mod tests {
let target_q_values = vec![0.5_f32; total];
let td_errors = vec![0.1_f32; total];
// next_states: shift states by one timestep within each episode,
// terminal steps reuse current state
let next_states: Vec<f32> = (0..total)
.flat_map(|idx| {
let t = idx % timesteps;
let ep = idx / timesteps;
let src = if t + 1 < timesteps {
(ep * timesteps + t + 1) * state_dim
} else {
idx * state_dim
};
states.get(src..src + state_dim)
.unwrap_or(&states[idx * state_dim..idx * state_dim + state_dim])
.to_vec()
})
.collect();
let batch = ExperienceBatch {
states,
next_states,
actions,
rewards,
done_flags,
@@ -1036,8 +1054,25 @@ mod tests {
.map(|i| if i % timesteps == timesteps - 1 { 1 } else { 0 })
.collect();
// next_states: shift states by one timestep, terminal reuses current
let next_states: Vec<f32> = (0..total)
.flat_map(|idx| {
let t = idx % timesteps;
let ep = idx / timesteps;
let src = if t + 1 < timesteps {
(ep * timesteps + t + 1) * state_dim
} else {
idx * state_dim
};
states.get(src..src + state_dim)
.unwrap_or(&states[idx * state_dim..idx * state_dim + state_dim])
.to_vec()
})
.collect();
let batch = ExperienceBatch {
states,
next_states,
actions: vec![0; total],
rewards: vec![0.0; total],
done_flags,

34
crates/ml/src/hyperopt/adapters/dqn.rs
View File

@@ -458,10 +458,11 @@ impl ParameterSpace for DQNParams {
(0.2, 0.6), // 9: per_beta_start (linear)
// Rainbow DQN extensions (6D)
// v_range: with EMA-normalized rewards in [-0.45, +0.32] and gamma~0.93/n_steps~4,
// max discounted return ≈ 0.32 * (1-0.93^4)/(1-0.93) ≈ 1.2. Range [2, 10] covers this.
(2.0, 10.0), // 10: v_range (symmetric: v_min=-v_range, v_max=+v_range)
(2.0, 10.0), // 11: v_range_mirror (UNUSED — kept for 28D compat, mirrors index 10)
// v_range must cover Q-value range, not reward range.
// Q_max = max_reward / (1-gamma): gamma=0.88→3.75, gamma=0.99→45.
// Range [10, 50] covers the full gamma search space [0.88, 0.99].
(10.0, 50.0), // 10: v_range (symmetric: v_min=-v_range, v_max=+v_range)
(10.0, 50.0), // 11: v_range_mirror (UNUSED — kept for 28D compat, mirrors index 10)
(0.1_f64.ln(), 1.0_f64.ln()), // 12: noisy_sigma_init (log scale)
(128.0, 512.0), // 13: dueling_hidden_dim (linear, step=128)
(3.0, 5.0), // 14: n_steps (raised min: n=3 is Rainbow standard)
@@ -519,7 +520,7 @@ impl ParameterSpace for DQNParams {
// Rainbow DQN extensions — SYMMETRIC support (prevents init bias)
// Asymmetric v_min/v_max causes Q-values to initialize at midpoint ≠ 0,
// creating a self-reinforcing equilibrium that freezes learning.
let v_range = x[10].clamp(2.0, 10.0);
let v_range = x[10].clamp(10.0, 50.0);
let v_min = -v_range;
let v_max = v_range;
let noisy_sigma_init = x[12].exp().clamp(0.1, 1.0);
@@ -2466,6 +2467,7 @@ impl HyperparameterOptimizable for DQNTrainer {
minimum_profit_factor: params.minimum_profit_factor,
// Phase 3: GPU experience collection (defaults, not in search space)
enable_gpu_experience_collector: true,
gpu_n_episodes: 128,
gpu_timesteps_per_episode: 500,
avg_spread: 0.0001,
@@ -3597,8 +3599,8 @@ mod tests {
assert_eq!(bounds[7], (0.5, 2.0)); // transaction_cost_multiplier
assert_eq!(bounds[8], (0.4, 0.8)); // per_alpha
assert_eq!(bounds[9], (0.2, 0.6)); // per_beta_start
assert_eq!(bounds[10], (2.0, 10.0)); // v_range (symmetric support)
assert_eq!(bounds[11], (2.0, 10.0)); // v_range_mirror (unused, same as 10)
assert_eq!(bounds[10], (10.0, 50.0)); // v_range (symmetric support, covers gamma 0.88-0.99)
assert_eq!(bounds[11], (10.0, 50.0)); // v_range_mirror (unused, same as 10)
assert_eq!(bounds[13], (128.0, 512.0)); // dueling_hidden_dim
assert_eq!(bounds[14], (3.0, 5.0)); // n_steps (raised min: Rainbow standard)
assert_eq!(bounds[15], (51.0, 201.0)); // num_atoms
@@ -3654,7 +3656,7 @@ mod tests {
let continuous = vec![
3e-5_f64.ln(), 92.0, 0.92, 97_273_f64.ln(), 4.0, 24.77_f64.ln(), 0.03, 1.0,
0.6, 0.4, // 8-9: per_alpha, per_beta_start
7.0, 7.0, 0.5_f64.ln(), // 10-12: v_range, v_range_mirror, noisy_sigma_init
25.0, 25.0, 0.5_f64.ln(), // 10-12: v_range, v_range_mirror, noisy_sigma_init
256.0, 3.0, 101.0, // 13-15: dueling_hidden_dim, n_steps, num_atoms
1e-4_f64.ln(), // 16: weight_decay
0.5, 0.25, // 17-18: kelly_fractional, kelly_max_fraction
@@ -3680,14 +3682,14 @@ mod tests {
assert!((params.curiosity_weight - 0.0).abs() < 1e-6); // C2: fixed
assert!((params.noisy_epsilon_floor - 0.0).abs() < 1e-6); // C2: fixed
// Verify symmetric support
assert!((params.v_min - (-7.0)).abs() < 1e-6, "v_min should be -v_range");
assert!((params.v_max - 7.0).abs() < 1e-6, "v_max should be +v_range");
assert!((params.v_min - (-25.0)).abs() < 1e-6, "v_min should be -v_range");
assert!((params.v_max - 25.0).abs() < 1e-6, "v_max should be +v_range");
// Test PER parameter bounds (min values) -- 28D
let continuous_min = vec![
2e-5_f64.ln(), 64.0, 0.88, 50_000_f64.ln(), 1.0, 10.0_f64.ln(), 0.05, 0.5,
0.4, 0.2, // per_alpha min, per_beta_start min
2.0, 2.0, 0.1_f64.ln(), // v_range min, v_range_mirror, noisy_sigma_init
10.0, 10.0, 0.1_f64.ln(), // v_range min, v_range_mirror, noisy_sigma_init
128.0, 3.0, 51.0, // dueling_hidden_dim, n_steps (min=3), num_atoms
1e-4_f64.ln(), // weight_decay min
0.25, 0.1, // kelly_fractional, kelly_max_fraction
@@ -3707,14 +3709,14 @@ mod tests {
assert!(params_min.use_dueling);
assert!(params_min.use_distributional);
assert!(params_min.use_noisy_nets);
assert!((params_min.v_min - (-2.0)).abs() < 1e-6, "v_min should be -v_range at min");
assert!((params_min.v_max - 2.0).abs() < 1e-6, "v_max should be +v_range at min");
assert!((params_min.v_min - (-10.0)).abs() < 1e-6, "v_min should be -v_range at min");
assert!((params_min.v_max - 10.0).abs() < 1e-6, "v_max should be +v_range at min");
// Test PER parameter bounds (max values) -- 28D
let continuous_max = vec![
8e-5_f64.ln(), 512.0, 0.99, 100_000_f64.ln(), 4.0, 40.0_f64.ln(), 0.5, 2.0,
0.8, 0.6, // per_alpha max, per_beta_start max
10.0, 10.0, 1.0_f64.ln(), // v_range max, v_range_mirror, noisy_sigma_init
50.0, 50.0, 1.0_f64.ln(), // v_range max, v_range_mirror, noisy_sigma_init
512.0, 5.0, 201.0, // dueling_hidden_dim, n_steps, num_atoms
1e-2_f64.ln(), // weight_decay max
0.75, 0.5, // kelly_fractional, kelly_max_fraction
@@ -3734,8 +3736,8 @@ mod tests {
assert!(params_max.use_dueling);
assert!(params_max.use_distributional);
assert!(params_max.use_noisy_nets);
assert!((params_max.v_min - (-10.0)).abs() < 1e-6, "v_min should be -v_range at max");
assert!((params_max.v_max - 10.0).abs() < 1e-6, "v_max should be +v_range at max");
assert!((params_max.v_min - (-50.0)).abs() < 1e-6, "v_min should be -v_range at max");
assert!((params_max.v_max - 50.0).abs() < 1e-6, "v_max should be +v_range at max");
}
#[test]

13
crates/ml/src/trainers/dqn/config.rs
View File

@@ -835,6 +835,10 @@ pub struct DQNHyperparameters {
pub qr_kappa: f64,
// Phase 3: GPU experience collection kernel configuration
/// Enable the NVRTC-compiled GPU experience collection kernel.
/// When false, experience collection uses the batched CPU path while
/// the training step (forward/backward/optimizer) still runs on CUDA.
pub enable_gpu_experience_collector: bool,
/// Number of parallel episodes per GPU kernel launch (default: 128, max: 256)
/// Higher values improve GPU utilization on larger GPUs (e.g. H100: 256)
pub gpu_n_episodes: usize,
@@ -1001,8 +1005,8 @@ impl DQNHyperparameters {
// Wave 2.3: Distributional RL (BUG #36 FIXED — scatter_add gradient flow verified)
use_distributional: true, // Default: ENABLED (C51 Rainbow DQN standard, BUG #36 fixed)
num_atoms: 51, // Rainbow DQN standard: 51 atoms
v_min: -2.0, // BUG #5 FIX: Align with reward range ±2 (was -1000.0, 500x too large!)
v_max: 2.0, // BUG #5 FIX: Align with reward range ±2 (was +1000.0, 500x too large!)
v_min: -25.0, // Must cover Q-value range, not reward range. With gamma=0.92, reward±2 → Q≈±25
v_max: 25.0, // Matches DQNConfig default. Wider is safe — just lowers per-atom resolution
// Wave 2.4: Noisy Networks (WAVE 6.4: ENABLED BY DEFAULT)
use_noisy_nets: true, // Default: enabled (replaces epsilon-greedy)
@@ -1076,6 +1080,7 @@ impl DQNHyperparameters {
qr_kappa: 1.0, // Default: 1.0 (standard quantile Huber loss)
// Phase 3: GPU experience collection
enable_gpu_experience_collector: true, // Default: enabled (zero-roundtrip CUDA kernel)
gpu_n_episodes: 128, // Default: 128 (good for 4-8GB VRAM GPUs)
gpu_timesteps_per_episode: 500, // Default: 500 timesteps per episode
avg_spread: 0.0001, // Default: 1bp (ES/NQ futures)
@@ -1200,8 +1205,8 @@ pub(crate) fn dqn_config_2025() -> DQNConfig {
dueling_hidden_dim: 256,
use_distributional: true,
num_atoms: 51,
v_min: -2.0,
v_max: 2.0,
v_min: -25.0,
v_max: 25.0,
n_steps: 3,
// Stability & Safety

1
crates/ml/src/trainers/dqn/mod.rs
View File

@@ -28,6 +28,7 @@ mod monitoring;
mod risk;
mod statistics;
mod trainer;
mod smoke_tests;
// Re-export all public items for backward compatibility
pub use config::{DQNAgentType, DQNHyperparameters};

156
crates/ml/src/trainers/dqn/smoke_tests/feature_coverage.rs Normal file
View File

@@ -0,0 +1,156 @@
//! Feature coverage smoke tests.
//!
//! Verify that every major DQN feature flag can be enabled and trained
//! without panics or NaN. CPU-only for CI safety.
use super::helpers::{assert_finite, cpu_trainer_with, smoke_params, synthetic_data};
/// Helper: train with modified params, assert loss is finite.
async fn train_and_check(
mut params: crate::trainers::dqn::DQNHyperparameters,
label: &str,
) -> anyhow::Result<()> {
// Keep tests fast
params.epochs = 3;
params.batch_size = 16;
params.buffer_size = 500;
params.min_replay_size = 32;
params.warmup_steps = 0;
params.checkpoint_frequency = 100;
params.early_stopping_enabled = false;
params.enable_stress_testing = false;
params.enable_compliance = false;
let mut trainer = cpu_trainer_with(params)?;
let data = synthetic_data(200);
let val = synthetic_data(50);
let metrics = trainer
.train_with_preloaded_data(data, val, |_epoch, _bytes, _best| {
Ok("skip".to_owned())
})
.await?;
assert_finite(metrics.loss, label);
Ok(())
}
/// Dueling DQN architecture (separate value/advantage streams).
#[tokio::test]
async fn test_dueling_network_trains() -> anyhow::Result<()> {
let mut p = smoke_params();
p.use_dueling = true;
p.use_distributional = false;
p.use_noisy_nets = false;
train_and_check(p, "dueling_loss").await
}
/// Distributional RL (C51) trains without NaN.
#[tokio::test]
async fn test_distributional_c51_trains() -> anyhow::Result<()> {
let mut p = smoke_params();
p.use_distributional = true;
p.use_dueling = false;
p.use_noisy_nets = false;
train_and_check(p, "distributional_c51_loss").await
}
/// Noisy Networks exploration (replaces epsilon-greedy).
#[tokio::test]
async fn test_noisy_nets_trains() -> anyhow::Result<()> {
let mut p = smoke_params();
p.use_noisy_nets = true;
p.use_dueling = false;
p.use_distributional = false;
train_and_check(p, "noisy_nets_loss").await
}
/// Full Rainbow DQN: dueling + distributional + noisy + PER + n-step.
#[tokio::test]
async fn test_rainbow_full_trains() -> anyhow::Result<()> {
let mut p = smoke_params();
p.use_dueling = true;
p.use_distributional = true;
p.use_noisy_nets = true;
p.use_per = true;
p.n_steps = 3;
train_and_check(p, "rainbow_loss").await
}
/// QR-DQN (IQN quantile regression) replaces C51.
#[tokio::test]
async fn test_iqn_qr_dqn_trains() -> anyhow::Result<()> {
let mut p = smoke_params();
p.use_qr_dqn = true;
p.num_quantiles = 32;
p.qr_kappa = 1.0;
p.use_distributional = false;
train_and_check(p, "qr_dqn_loss").await
}
/// Conservative Q-Learning (CQL) regularization.
#[tokio::test]
async fn test_cql_regularization_trains() -> anyhow::Result<()> {
let mut p = smoke_params();
p.use_cql = true;
p.cql_alpha = 0.1;
train_and_check(p, "cql_loss").await
}
/// Curiosity-driven exploration module.
#[tokio::test]
async fn test_curiosity_module_trains() -> anyhow::Result<()> {
let mut p = smoke_params();
p.curiosity_weight = 0.1;
train_and_check(p, "curiosity_loss").await
}
/// Kelly criterion position sizing.
#[tokio::test]
async fn test_kelly_sizing_enabled() -> anyhow::Result<()> {
let mut p = smoke_params();
p.enable_kelly_sizing = true;
p.kelly_fractional = 0.5;
p.kelly_max_fraction = 0.25;
train_and_check(p, "kelly_loss").await
}
/// Circuit breaker (5-failure trip mechanism).
#[tokio::test]
async fn test_circuit_breaker_enabled() -> anyhow::Result<()> {
let mut p = smoke_params();
p.enable_circuit_breaker = true;
train_and_check(p, "circuit_breaker_loss").await
}
/// Action masking (filters invalid actions by position limits).
#[tokio::test]
async fn test_action_masking_enabled() -> anyhow::Result<()> {
let mut p = smoke_params();
p.enable_action_masking = true;
p.max_position_absolute = 2.0;
train_and_check(p, "action_masking_loss").await
}
/// Gradient accumulation (effective batch = batch_size * steps).
#[tokio::test]
async fn test_gradient_accumulation_trains() -> anyhow::Result<()> {
let mut p = smoke_params();
p.gradient_accumulation_steps = 2;
p.batch_size = 8; // effective batch = 16
train_and_check(p, "gradient_accumulation_loss").await
}
/// Multi-step returns (n-step TD targets).
#[tokio::test]
async fn test_n_step_returns_trains() -> anyhow::Result<()> {
let mut p = smoke_params();
p.n_steps = 3;
train_and_check(p, "n_step_loss").await
}
/// Double DQN (uses target net for action selection to reduce overestimation).
#[tokio::test]
async fn test_double_dqn_enabled() -> anyhow::Result<()> {
let mut p = smoke_params();
p.use_double_dqn = true;
train_and_check(p, "double_dqn_loss").await
}

272
crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs Normal file
View File

@@ -0,0 +1,272 @@
use super::helpers::*;
use candle_core::{DType, Device, Tensor};
use candle_nn::Module;
use crate::dqn::mixed_precision::training_dtype;
// GPU replay buffer is only available with the cuda feature (the module is
// gated on #[cfg(feature = "cuda")] in ml-dqn). The buffer itself works on
// Device::Cpu too -- the gate is compile-time only.
#[cfg(feature = "cuda")]
use crate::dqn::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};
#[cfg(feature = "cuda")]
use crate::dqn::replay_buffer_type::GpuBatch;
// ---------------------------------------------------------------------------
// GpuReplayBuffer tests (compiled only with cuda feature)
// ---------------------------------------------------------------------------
#[cfg(feature = "cuda")]
fn test_buffer_config(capacity: usize, state_dim: usize) -> GpuReplayBufferConfig {
GpuReplayBufferConfig {
capacity,
state_dim,
alpha: 0.6,
beta_start: 0.4,
beta_max: 1.0,
beta_annealing_steps: 1000,
epsilon: 1e-6,
}
}
/// Insert random experiences into a GPU replay buffer on CPU device.
#[cfg(feature = "cuda")]
fn insert_random_batch(
buf: &mut GpuReplayBuffer,
n: usize,
state_dim: usize,
) -> anyhow::Result<()> {
let device = Device::Cpu;
let states = Tensor::randn(0.0_f32, 1.0, &[n, state_dim], &device)?;
let next_states = Tensor::randn(0.0_f32, 1.0, &[n, state_dim], &device)?;
let actions = Tensor::zeros(&[n], DType::U32, &device)?;
let rewards = Tensor::randn(0.0_f32, 1.0, &[n], &device)?;
let dones = Tensor::zeros(&[n], DType::F32, &device)?;
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones)?;
Ok(())
}
#[cfg(feature = "cuda")]
#[tokio::test]
async fn test_gpu_replay_buffer_insert_and_device() -> anyhow::Result<()> {
let config = test_buffer_config(100, 48);
let mut buf = GpuReplayBuffer::new(config, &Device::Cpu)?;
insert_random_batch(&mut buf, 20, 48)?;
assert_eq!(buf.len(), 20);
assert!(matches!(buf.device(), Device::Cpu));
Ok(())
}
#[cfg(feature = "cuda")]
#[tokio::test]
async fn test_gpu_replay_buffer_proportional_sample_valid() -> anyhow::Result<()> {
let config = test_buffer_config(100, 48);
let mut buf = GpuReplayBuffer::new(config, &Device::Cpu)?;
insert_random_batch(&mut buf, 50, 48)?;
let batch: GpuBatch = buf.sample_proportional(16)?;
assert_eq!(batch.states.dims(), &[16, 48]);
assert_eq!(batch.weights.dims(), &[16]);
// All weights must be positive
let weights: Vec<f32> = batch.weights.to_vec1()?;
for (i, w) in weights.iter().enumerate() {
assert!(*w > 0.0, "weight at index {i} is not positive: {w}");
}
// All indices must be in range [0, 50)
let indices: Vec<u32> = batch.indices.to_vec1()?;
for (i, &idx) in indices.iter().enumerate() {
assert!(
(idx as usize) < 50,
"index at position {i} out of range: {idx}"
);
}
Ok(())
}
#[cfg(feature = "cuda")]
#[tokio::test]
async fn test_gpu_replay_buffer_rank_based_sample_valid() -> anyhow::Result<()> {
let config = test_buffer_config(100, 48);
let mut buf = GpuReplayBuffer::new(config, &Device::Cpu)?;
insert_random_batch(&mut buf, 50, 48)?;
let batch: GpuBatch = buf.sample_rank_based(16)?;
assert_eq!(batch.states.dims(), &[16, 48]);
assert_eq!(batch.weights.dims(), &[16]);
// All weights must be positive
let weights: Vec<f32> = batch.weights.to_vec1()?;
for (i, w) in weights.iter().enumerate() {
assert!(*w > 0.0, "weight at index {i} is not positive: {w}");
}
// All indices must be in range [0, 50)
let indices: Vec<u32> = batch.indices.to_vec1()?;
for (i, &idx) in indices.iter().enumerate() {
assert!(
(idx as usize) < 50,
"index at position {i} out of range: {idx}"
);
}
Ok(())
}
#[cfg(feature = "cuda")]
#[tokio::test]
async fn test_gpu_replay_buffer_priority_update_valid() -> anyhow::Result<()> {
let config = test_buffer_config(100, 48);
let epsilon = config.epsilon;
let mut buf = GpuReplayBuffer::new(config, &Device::Cpu)?;
insert_random_batch(&mut buf, 50, 48)?;
// Sample and get indices
let batch = buf.sample_proportional(16)?;
let td_errors = Tensor::randn(0.0_f32, 1.0, &[16], &Device::Cpu)?;
buf.update_priorities_gpu(&batch.indices, &td_errors)?;
// Verify sampling still works with updated priorities (proves update succeeded)
let batch2 = buf.sample_proportional(16)?;
assert_eq!(batch2.weights.dims(), &[16]);
// Weights should still be positive after priority update
let weights: Vec<f32> = batch2.weights.to_vec1()?;
for (i, w) in weights.iter().enumerate() {
assert!(*w > 0.0, "weight at index {i} not positive after update: {w}");
}
// Epsilon is the priority floor -- verify it's reasonable
assert!(epsilon > 0.0, "epsilon should be positive");
Ok(())
}
#[cfg(feature = "cuda")]
#[tokio::test]
async fn test_searchsorted_cpu_correctness() -> anyhow::Result<()> {
// Test indirectly: create a buffer with known size, sample, and verify
// that sampled indices are always in valid range. The CPU fallback of
// SearchSorted uses binary_search_by internally.
let config = test_buffer_config(50, 4);
let mut buf = GpuReplayBuffer::new(config, &Device::Cpu)?;
// Fill with exactly 50 experiences
insert_random_batch(&mut buf, 50, 4)?;
assert_eq!(buf.len(), 50);
// Sample multiple times to exercise the searchsorted codepath
for trial in 0..5 {
let batch = buf.sample_proportional(10)?;
let indices: Vec<u32> = batch.indices.to_vec1()?;
for (i, &idx) in indices.iter().enumerate() {
assert!(
(idx as usize) < 50,
"trial {trial}, position {i}: index {idx} out of range [0, 50)"
);
}
}
Ok(())
}
#[tokio::test]
async fn test_train_step_produces_finite_metrics() -> anyhow::Result<()> {
let mut trainer = cpu_trainer()?;
let train_data = synthetic_data(200);
let val_data = synthetic_data(50);
let metrics = trainer
.train_with_preloaded_data(train_data, val_data, |_epoch, _bytes, _is_best| {
Ok(String::new())
})
.await?;
// Final loss must be finite (not NaN, not Inf)
assert_finite(metrics.loss, "final_loss");
Ok(())
}
#[tokio::test]
async fn test_mixed_precision_dtype_selection() -> anyhow::Result<()> {
let cpu_dtype = training_dtype(&Device::Cpu);
assert_eq!(cpu_dtype, DType::F32);
Ok(())
}
/// Diagnose: what training_dtype does the GPU get?
#[tokio::test]
async fn test_gpu_training_dtype_diagnosis() -> anyhow::Result<()> {
let dev = safe_device();
let dtype = training_dtype(&dev);
eprintln!(" Device: {:?}, training_dtype: {:?}", dev, dtype);
// Verify a linear layer forward pass works with training_dtype weights
let varmap = candle_nn::VarMap::new();
let vs = candle_nn::VarBuilder::from_varmap(&varmap, dtype, &dev);
let layer = candle_nn::linear(48, 32, vs.pp("test"))?;
let input = Tensor::zeros(&[2, 48], dtype, &dev)?;
let out = layer.forward(&input)?;
eprintln!(" Linear forward OK: {:?}, dtype: {:?}", out.dims(), out.dtype());
Ok(())
}
/// Diagnose: does training work on GPU when the NVRTC experience collector
/// is disabled? Isolates cuBLAS-level issues from NVRTC side effects.
#[tokio::test]
async fn test_training_without_gpu_collector() -> anyhow::Result<()> {
let mut params = smoke_params();
params.enable_gpu_experience_collector = false; // no NVRTC, pure cuBLAS training
let mut trainer = cpu_trainer_with(params)?;
let train_data = synthetic_data(200);
let val_data = synthetic_data(50);
let result = trainer
.train_with_preloaded_data(train_data, val_data, |_epoch, _bytes, _is_best| {
Ok(String::new())
})
.await;
match &result {
Ok(m) => eprintln!(" Training OK, loss={}", m.loss),
Err(e) => eprintln!(" Training FAILED: {}", e),
}
let metrics = result?;
assert_finite(metrics.loss, "final_loss_no_collector");
Ok(())
}
/// Diagnose: can we create + step an AdamW optimizer on GPU with training_dtype?
#[tokio::test]
async fn test_gpu_adamw_creation() -> anyhow::Result<()> {
use candle_nn::Optimizer;
let dev = safe_device();
let dtype = training_dtype(&dev);
eprintln!(" Device: {:?}, dtype: {:?}", dev, dtype);
let varmap = candle_nn::VarMap::new();
let vs = candle_nn::VarBuilder::from_varmap(&varmap, dtype, &dev);
let layer = candle_nn::linear(48, 32, vs.pp("test"))?;
// Forward + backward
let input = Tensor::zeros(&[2, 48], dtype, &dev)?;
let out = layer.forward(&input)?;
let loss = out.sqr()?.mean_all()?;
let grads = loss.backward()?;
// Create optimizer (this is where CUDA_ERROR_ILLEGAL_ADDRESS happened)
let params = candle_nn::ParamsAdamW {
lr: 1e-4,
..Default::default()
};
let mut opt = candle_nn::AdamW::new(varmap.all_vars(), params)?;
opt.step(&grads)?;
eprintln!(" AdamW step OK");
Ok(())
}

85
crates/ml/src/trainers/dqn/smoke_tests/helpers.rs Normal file
View File

@@ -0,0 +1,85 @@
use crate::trainers::dqn::{DQNHyperparameters, DQNTrainer};
use crate::features::extraction::FeatureVector;
use candle_core::Device;
/// Create a fast test config with small sizes
pub(super) fn smoke_params() -> DQNHyperparameters {
let mut p = DQNHyperparameters::conservative();
p.batch_size = 16;
p.buffer_size = 500;
p.min_replay_size = 32;
p.epochs = 3;
p.warmup_steps = 0;
p.checkpoint_frequency = 100; // never checkpoint in tests
p.early_stopping_enabled = false;
// Disable stress testing, compliance, and regime Q-network to speed up tests
// and avoid CUBLAS init errors on machines without GPU runtime
p.enable_stress_testing = false;
p.enable_compliance = false;
p.enable_regime_qnetwork = false;
// Set explicit hidden_dim_base to skip GPU capability probe
// (cached_capabilities() touches CUDA even when device is CPU)
p.hidden_dim_base = Some(32);
// Shrink GPU experience collection to fit the small test buffer.
// Default 128×500 = 63,500 experiences per collection exceeds buffer_size=500.
p.gpu_n_episodes = 2;
p.gpu_timesteps_per_episode = 50;
p
}
/// Generate synthetic training data (feature vectors + OHLC targets)
pub(super) fn synthetic_data(n: usize) -> Vec<(FeatureVector, Vec<f64>)> {
use rand::Rng;
let mut rng = rand::thread_rng();
(0..n)
.map(|_| {
let mut fv = [0.0_f64; 40];
for v in &mut fv {
*v = rng.gen_range(-1.0..1.0);
}
let targets = vec![
4000.0 + rng.gen_range(-50.0..50.0), // open
4010.0 + rng.gen_range(-50.0..50.0), // high
3990.0 + rng.gen_range(-50.0..50.0), // low
4005.0 + rng.gen_range(-50.0..50.0), // close
];
(fv, targets)
})
.collect()
}
/// Select safe device: CPU when CUDA compiled but no GPU available.
/// Tests a small matmul to detect broken CUDA runtime before trainer creation.
pub(super) fn safe_device() -> Device {
#[cfg(feature = "cuda")]
{
if let Ok(dev) = Device::new_cuda(0) {
let probe = candle_core::Tensor::zeros(&[2, 2], candle_core::DType::F32, &dev)
.and_then(|t| t.matmul(&t));
if probe.is_ok() {
return dev;
}
}
}
Device::Cpu
}
/// Create trainer on a safe device (CPU fallback if CUDA broken)
pub(super) fn cpu_trainer() -> anyhow::Result<DQNTrainer> {
DQNTrainer::new_with_device(smoke_params(), safe_device())
}
/// Create trainer with custom params on safe device
pub(super) fn cpu_trainer_with(params: DQNHyperparameters) -> anyhow::Result<DQNTrainer> {
DQNTrainer::new_with_device(params, safe_device())
}
/// Assert a value is finite (not NaN or Inf)
pub(super) fn assert_finite(val: f64, name: &str) {
assert!(val.is_finite(), "{name} is not finite: {val}");
}
/// Assert a value is finite f32
pub(super) fn assert_finite_f32(val: f32, name: &str) {
assert!(val.is_finite(), "{name} is not finite: {val}");
}

10
crates/ml/src/trainers/dqn/smoke_tests/mod.rs Normal file
View File

@@ -0,0 +1,10 @@
#[cfg(test)]
mod helpers;
#[cfg(test)]
mod gpu_residency;
#[cfg(test)]
mod training_stability;
#[cfg(test)]
mod feature_coverage;
#[cfg(test)]
mod performance;

146
crates/ml/src/trainers/dqn/smoke_tests/performance.rs Normal file
View File

@@ -0,0 +1,146 @@
//! Performance smoke tests.
//!
//! Speed benchmarks for the DQN training pipeline. All tests are marked
//! `#[ignore]` so they never run in CI -- invoke manually on a GPU box:
//!
//! ```sh
//! cargo test -p ml --lib smoke_tests::performance -- --ignored --nocapture
//! ```
use super::helpers::{assert_finite, smoke_params, synthetic_data};
use crate::trainers::dqn::DQNTrainer;
use std::time::Instant;
/// Measure end-to-end training throughput (3 epochs, 1000 samples).
#[tokio::test]
#[ignore] // Run manually: cargo test -p ml --lib smoke_tests::performance -- --ignored --nocapture
async fn test_training_throughput_measurement() -> anyhow::Result<()> {
let mut trainer = DQNTrainer::new(smoke_params())?; // auto-detect GPU
let data = synthetic_data(1000);
let val = synthetic_data(100);
let start = Instant::now();
let metrics = trainer
.train_with_preloaded_data(data, val, |_epoch, _bytes, _best| {
Ok("skip".to_owned())
})
.await?;
let elapsed = start.elapsed();
assert_finite(metrics.loss, "loss");
eprintln!(
" Training throughput: {:.1}s for {} epochs",
elapsed.as_secs_f64(),
metrics.epochs_trained
);
eprintln!(" Avg loss: {:.6}", metrics.loss);
if let Some(&avg_q) = metrics.additional_metrics.get("avg_q_value") {
eprintln!(" Avg Q-value: {:.4}", avg_q);
}
if let Some(&final_eps) = metrics.additional_metrics.get("final_epsilon") {
eprintln!(" Final epsilon: {:.4}", final_eps);
}
Ok(())
}
/// Measure per-sample latency of the GPU PER replay buffer.
#[cfg(feature = "cuda")]
#[tokio::test]
#[ignore] // Run manually on GPU
async fn test_per_sample_latency() -> anyhow::Result<()> {
use candle_core::{DType, Device, Tensor};
use crate::dqn::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};
let device = Device::new_cuda(0)?;
let state_dim: usize = 48;
let capacity: usize = 10_000;
let config = GpuReplayBufferConfig {
capacity,
state_dim,
alpha: 0.6,
beta_start: 0.4,
beta_max: 1.0,
beta_annealing_steps: 1000,
epsilon: 1e-6,
};
let mut buf = GpuReplayBuffer::new(config, &device)?;
// Fill with 5000 experiences
let fill_count: usize = 5000;
let batch_insert: usize = 100;
for _ in 0..(fill_count / batch_insert) {
let states = Tensor::randn(0.0_f32, 1.0, &[batch_insert, state_dim], &device)?;
let next_states = Tensor::randn(0.0_f32, 1.0, &[batch_insert, state_dim], &device)?;
let actions = Tensor::zeros(&[batch_insert], DType::U32, &device)?;
let rewards = Tensor::randn(0.0_f32, 1.0, &[batch_insert], &device)?;
let dones = Tensor::zeros(&[batch_insert], DType::F32, &device)?;
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones)?;
}
assert_eq!(buf.len(), fill_count);
// Warmup
for _ in 0..5 {
let _ = buf.sample_proportional(64)?;
}
// Timed: 100 proportional samples of batch_size=64
let sample_rounds: usize = 100;
let start = Instant::now();
for _ in 0..sample_rounds {
let _ = buf.sample_proportional(64)?;
}
let elapsed = start.elapsed();
let us_per_sample = elapsed.as_micros() as f64 / sample_rounds as f64;
eprintln!(
" PER proportional sample latency: {:.1} us/sample ({} rounds, batch_size=64)",
us_per_sample, sample_rounds
);
assert!(
us_per_sample < 50_000.0,
"sample latency {us_per_sample:.1} us is unreasonably high"
);
Ok(())
}
/// Single-epoch training on real data (if available).
#[tokio::test]
#[ignore] // Run manually: FOXHUNT_TEST_DATA=./test_data/futures-baseline cargo test ...
async fn test_real_data_single_epoch() -> anyhow::Result<()> {
let data_dir = std::env::var("FOXHUNT_TEST_DATA")
.unwrap_or_else(|_| "../../test_data/futures-baseline".to_owned());
let data_path = std::path::Path::new(&data_dir);
if !data_path.exists() {
eprintln!(
" SKIP: test data dir does not exist: {}",
data_path.display()
);
return Ok(());
}
let mut params = smoke_params();
params.epochs = 1;
let mut trainer = DQNTrainer::new(params)?;
let start = Instant::now();
let metrics = trainer
.train(&data_dir, |_epoch, _bytes, _best| {
Ok("skip".to_owned())
})
.await?;
let elapsed = start.elapsed();
assert_finite(metrics.loss, "real_data_loss");
eprintln!(
" Real-data single epoch: {:.1}s, loss={:.6}",
elapsed.as_secs_f64(),
metrics.loss
);
if let Some(&avg_q) = metrics.additional_metrics.get("avg_q_value") {
eprintln!(" Avg Q-value: {:.4}", avg_q);
}
Ok(())
}

254
crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs Normal file
View File

@@ -0,0 +1,254 @@
//! Training stability smoke tests.
//!
//! Verify that short DQN training runs produce sane, finite metrics.
//! All tests use CPU device for CI compatibility.
use super::helpers::{assert_finite, cpu_trainer, cpu_trainer_with, smoke_params, synthetic_data};
/// Train a few epochs and verify the returned loss is finite and non-negative.
#[tokio::test]
async fn test_training_produces_finite_loss() -> anyhow::Result<()> {
let mut trainer = cpu_trainer()?;
let data = synthetic_data(200);
let val = synthetic_data(50);
let metrics = trainer
.train_with_preloaded_data(data, val, |_epoch, _bytes, _best| {
Ok("skip".to_owned())
})
.await?;
assert_finite(metrics.loss, "loss");
assert!(metrics.loss >= 0.0, "loss should be non-negative");
Ok(())
}
/// Gradient norms must be finite and positive (gradients are flowing).
#[tokio::test]
async fn test_gradient_norms_finite() -> anyhow::Result<()> {
let mut params = smoke_params();
params.epochs = 3;
let mut trainer = cpu_trainer_with(params)?;
let data = synthetic_data(200);
let val = synthetic_data(50);
let metrics = trainer
.train_with_preloaded_data(data, val, |_epoch, _bytes, _best| {
Ok("skip".to_owned())
})
.await?;
let grad_norm = metrics
.additional_metrics
.get("avg_gradient_norm")
.copied()
.unwrap_or(f64::NAN);
assert_finite(grad_norm, "avg_gradient_norm");
assert!(
grad_norm > 0.0,
"avg_gradient_norm should be positive, got {grad_norm}"
);
Ok(())
}
/// Q-values should stay bounded during short training (no explosion).
#[tokio::test]
async fn test_q_values_bounded() -> anyhow::Result<()> {
let mut trainer = cpu_trainer()?;
let data = synthetic_data(200);
let val = synthetic_data(50);
let metrics = trainer
.train_with_preloaded_data(data, val, |_epoch, _bytes, _best| {
Ok("skip".to_owned())
})
.await?;
let avg_q = metrics
.additional_metrics
.get("avg_q_value")
.copied()
.unwrap_or(f64::NAN);
assert_finite(avg_q, "avg_q_value");
assert!(
(-1000.0..=1000.0).contains(&avg_q),
"avg_q_value {avg_q} out of sane range [-1000, 1000]"
);
Ok(())
}
/// Epsilon should decay below its starting value after training.
#[tokio::test]
async fn test_epsilon_decays() -> anyhow::Result<()> {
let mut params = smoke_params();
params.epsilon_start = 1.0;
params.epsilon_decay = 0.9;
let mut trainer = cpu_trainer_with(params)?;
let data = synthetic_data(200);
let val = synthetic_data(50);
let metrics = trainer
.train_with_preloaded_data(data, val, |_epoch, _bytes, _best| {
Ok("skip".to_owned())
})
.await?;
let final_eps = metrics
.additional_metrics
.get("final_epsilon")
.copied()
.unwrap_or(f64::NAN);
assert_finite(final_eps, "final_epsilon");
assert!(
final_eps < 1.0,
"epsilon should have decayed below 1.0, got {final_eps}"
);
Ok(())
}
/// Short training must not produce NaN loss.
#[tokio::test]
async fn test_no_nan_loss_short_training() -> anyhow::Result<()> {
let mut params = smoke_params();
params.epochs = 2;
let mut trainer = cpu_trainer_with(params)?;
let data = synthetic_data(100);
let val = synthetic_data(30);
let metrics = trainer
.train_with_preloaded_data(data, val, |_epoch, _bytes, _best| {
Ok("skip".to_owned())
})
.await?;
assert!(
metrics.loss.is_finite(),
"loss should be finite after 2-epoch training, got {}",
metrics.loss
);
Ok(())
}
/// PER importance-sampling weights must be finite and positive.
#[tokio::test]
async fn test_per_weights_in_valid_range() -> anyhow::Result<()> {
use candle_core::{DType, Device, Tensor};
use crate::dqn::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};
let device = Device::Cpu;
let state_dim = 43;
let config = GpuReplayBufferConfig {
capacity: 128,
state_dim,
alpha: 0.6,
beta_start: 0.4,
beta_max: 1.0,
beta_annealing_steps: 1000,
epsilon: 1e-6,
};
let mut buf = GpuReplayBuffer::new(config, &device)?;
// Fill with 50 experiences
for _ in 0..50 {
let s = Tensor::randn(0.0_f32, 1.0, &[1, state_dim], &device)?;
let ns = Tensor::randn(0.0_f32, 1.0, &[1, state_dim], &device)?;
let a = Tensor::zeros(&[1], DType::U32, &device)?;
let r = Tensor::randn(0.0_f32, 1.0, &[1], &device)?;
let d = Tensor::zeros(&[1], DType::F32, &device)?;
buf.insert_batch(&s, &ns, &a, &r, &d)?;
}
let batch = buf.sample_proportional(16)?;
let weights: Vec<f32> = batch.weights.to_vec1()?;
for (i, &w) in weights.iter().enumerate() {
assert!(
w.is_finite() && w > 0.0,
"weight[{i}] = {w} is not finite-positive"
);
}
Ok(())
}
/// PER sampled indices must be within buffer bounds.
#[tokio::test]
async fn test_per_indices_in_valid_range() -> anyhow::Result<()> {
use candle_core::{DType, Device, Tensor};
use crate::dqn::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};
let device = Device::Cpu;
let state_dim = 43;
let config = GpuReplayBufferConfig {
capacity: 128,
state_dim,
alpha: 0.6,
beta_start: 0.4,
beta_max: 1.0,
beta_annealing_steps: 1000,
epsilon: 1e-6,
};
let mut buf = GpuReplayBuffer::new(config, &device)?;
for _ in 0..50 {
let s = Tensor::randn(0.0_f32, 1.0, &[1, state_dim], &device)?;
let ns = Tensor::randn(0.0_f32, 1.0, &[1, state_dim], &device)?;
let a = Tensor::zeros(&[1], DType::U32, &device)?;
let r = Tensor::randn(0.0_f32, 1.0, &[1], &device)?;
let d = Tensor::zeros(&[1], DType::F32, &device)?;
buf.insert_batch(&s, &ns, &a, &r, &d)?;
}
let batch = buf.sample_proportional(16)?;
let indices: Vec<u32> = batch.indices.to_vec1()?;
for (i, &idx) in indices.iter().enumerate() {
assert!(
idx < 50,
"index[{i}] = {idx} is out of bounds (buffer has 50 experiences)"
);
}
Ok(())
}
/// Training must record at least one episode.
#[tokio::test]
async fn test_total_episodes_nonzero() -> anyhow::Result<()> {
let mut params = smoke_params();
params.epochs = 3;
let mut trainer = cpu_trainer_with(params)?;
let data = synthetic_data(200);
let val = synthetic_data(50);
let metrics = trainer
.train_with_preloaded_data(data, val, |_epoch, _bytes, _best| {
Ok("skip".to_owned())
})
.await?;
assert!(
metrics.epochs_trained > 0,
"epochs_trained should be > 0, got {}",
metrics.epochs_trained
);
Ok(())
}
/// Dueling DQN architecture should still produce finite loss.
#[tokio::test]
async fn test_training_with_dueling_produces_finite() -> anyhow::Result<()> {
let mut params = smoke_params();
params.use_dueling = true;
let mut trainer = cpu_trainer_with(params)?;
let data = synthetic_data(200);
let val = synthetic_data(50);
let metrics = trainer
.train_with_preloaded_data(data, val, |_epoch, _bytes, _best| {
Ok("skip".to_owned())
})
.await?;
assert_finite(metrics.loss, "dueling_loss");
Ok(())
}
/// Distributional (C51) DQN should produce finite loss.
#[tokio::test]
async fn test_training_with_distributional_produces_finite() -> anyhow::Result<()> {
let mut params = smoke_params();
params.use_distributional = true;
let mut trainer = cpu_trainer_with(params)?;
let data = synthetic_data(200);
let val = synthetic_data(50);
let metrics = trainer
.train_with_preloaded_data(data, val, |_epoch, _bytes, _best| {
Ok("skip".to_owned())
})
.await?;
assert_finite(metrics.loss, "distributional_loss");
Ok(())
}

175
crates/ml/src/trainers/dqn/trainer.rs
View File

@@ -494,12 +494,12 @@ impl DQNTrainer {
info!("Creating regime-conditional DQN with 3 heads (Trending, Ranging, Volatile)");
info!(" - Regime detection: ADX (index 211) + Entropy (index 219)");
info!(" - Classification: Trending (ADX>25), Volatile (ADX≤25 & Entropy>0.7), Ranging (ADX≤25 & Entropy≤0.7)");
let regime_agent = RegimeConditionalDQN::new(config)
let regime_agent = RegimeConditionalDQN::new_on_device(config, device.clone())
.map_err(|e| anyhow::anyhow!("Failed to create regime-conditional DQN: {}", e))?;
DQNAgentType::RegimeConditional(regime_agent)
} else {
info!("Creating standard DQN with single Q-network head");
let standard_agent = DQN::new(config)
let standard_agent = DQN::new_on_device(config, device.clone())
.map_err(|e| anyhow::anyhow!("Failed to create DQN agent: {}", e))?;
DQNAgentType::Standard(standard_agent)
};
@@ -1074,22 +1074,32 @@ impl DQNTrainer {
// Sample experiences from replay buffer
let batch_sample = agent.memory().sample(sample_size)?;
let experiences = batch_sample.experiences;
// Extract states and create batch tensor
let states: Vec<f32> = experiences
.iter()
.flat_map(|exp| exp.state.iter().copied())
.collect();
let state_dim = agent.get_state_dim();
let batch_tensor = Tensor::from_vec(
states,
(sample_size, state_dim),
agent.device()
).map_err(|e| crate::MLError::ModelError(format!("Failed to create batch tensor: {}", e)))?
.to_dtype(training_dtype(agent.device()))
.map_err(|e| crate::MLError::ModelError(format!("Failed to cast batch tensor to training dtype: {}", e)))?;
// GPU PER path: use gpu_batch.states directly (experiences vec is empty)
#[allow(unused_assignments)]
let mut batch_tensor_opt: Option<Tensor> = None;
#[cfg(feature = "cuda")]
{
if let Some(ref gpu) = batch_sample.gpu_batch {
batch_tensor_opt = Some(
gpu.states.to_dtype(training_dtype(agent.device()))
.map_err(|e| crate::MLError::ModelError(format!("GPU Q-stat states dtype cast: {}", e)))?
);
}
}
let batch_tensor = if let Some(t) = batch_tensor_opt {
t
} else {
let states: Vec<f32> = batch_sample.experiences
.iter()
.flat_map(|exp| exp.state.iter().copied())
.collect();
Tensor::from_vec(states, (sample_size, state_dim), agent.device())
.map_err(|e| crate::MLError::ModelError(format!("Failed to create batch tensor: {}", e)))?
.to_dtype(training_dtype(agent.device()))
.map_err(|e| crate::MLError::ModelError(format!("Failed to cast batch tensor to training dtype: {}", e)))?
};
// Forward pass to get Q-values [batch_size, num_actions]
let q_values = agent.forward(&batch_tensor)?;
@@ -1720,7 +1730,9 @@ impl DQNTrainer {
// Phase 1c: Initialize zero-roundtrip GPU experience collector (once)
#[cfg(feature = "cuda")]
if self.gpu_experience_collector.is_none() {
if self.hyperparams.enable_gpu_experience_collector
&& self.gpu_experience_collector.is_none()
{
if let candle_core::Device::Cuda(ref cuda_dev) = self.device {
use crate::cuda_pipeline::gpu_experience_collector::GpuExperienceCollector;
let stream = cuda_dev.cuda_stream();
@@ -1806,11 +1818,10 @@ impl DQNTrainer {
if let Some(result) = init_result {
match result {
Ok(collector) => {
info!("GPU experience collector initialized (zero-roundtrip CUDA kernel)");
self.gpu_experience_collector = Some(collector);
}
Err(e) => {
warn!("GPU experience collector init failed (CPU fallback): {}", e);
warn!("GPU experience collector init failed, falling back to CPU: {e}");
}
}
}
@@ -2570,6 +2581,7 @@ impl DQNTrainer {
// **PHASE 2: Batched Training from Replay Buffer**
// Now that buffer is populated, perform batched training
// This reduces train_step() calls from 1000×/epoch to ~8×/epoch (125× reduction)
let batch_size = self.hyperparams.batch_size;
let num_training_steps = if self.can_train().await? {
// Calculate number of training steps based on dataset size and batch size
@@ -4168,25 +4180,36 @@ impl DQNTrainer {
let batch_sample = buffer
.sample(sample_size)
.map_err(|e| anyhow::anyhow!("Failed to sample experiences: {}", e))?;
let samples = batch_sample.experiences;
// OPTIMIZATION: Batch all states into single tensor for parallel GPU processing
// WAVE 10.4: Get state dimension from agent configuration (fixes hardcoded STATE_DIM bug)
let state_dim = agent.get_state_dim();
// Pad replay buffer states (raw dim) to aligned state_dim for tensor core alignment
let batched_states: Vec<f32> = samples.iter().flat_map(|exp| {
let mut s = exp.state.clone();
s.resize(state_dim, 0.0); // zero-pad if shorter than aligned dim
s
}).collect();
// Create batched tensor [batch_size, state_dim]
let batch_tensor =
// GPU PER path: use gpu_batch.states directly (experiences vec is empty)
// CPU path: build tensor from experiences
#[allow(unused_assignments)]
let mut batch_tensor = None;
#[cfg(feature = "cuda")]
{
if let Some(ref gpu) = batch_sample.gpu_batch {
batch_tensor = Some(
gpu.states.to_dtype(training_dtype(agent.device()))
.map_err(|e| anyhow::anyhow!("GPU Q-est states dtype cast: {}", e))?
);
}
}
let batch_tensor = if let Some(t) = batch_tensor {
t
} else {
let samples = &batch_sample.experiences;
let batched_states: Vec<f32> = samples.iter().flat_map(|exp| {
let mut s = exp.state.clone();
s.resize(state_dim, 0.0);
s
}).collect();
Tensor::from_vec(batched_states, (sample_size, state_dim), agent.device())
.map_err(|e| anyhow::anyhow!("Failed to create batched state tensor: {}", e))?
.to_dtype(training_dtype(agent.device()))
.map_err(|e| anyhow::anyhow!("Failed to cast batched state tensor to training dtype: {}", e))?;
.map_err(|e| anyhow::anyhow!("Failed to cast batched state tensor to training dtype: {}", e))?
};
// WAVE 23 P0 Fix: Check for Q-value divergence (early stopping)
// This calls log_q_values() which returns Err if divergence detected for consecutive checks
@@ -4239,22 +4262,37 @@ impl DQNTrainer {
};
let state_dim = agent.get_state_dim();
// Pad replay buffer states to aligned dim for tensor core alignment
let batched_states: Vec<f32> = batch_sample.experiences
.iter()
.flat_map(|exp| {
let mut s = exp.state.clone();
s.resize(state_dim, 0.0);
s
})
.collect();
let batch_tensor = match Tensor::from_vec(batched_states, (sample_size, state_dim), agent.device()) {
Ok(t) => match t.to_dtype(training_dtype(agent.device())) {
// GPU PER path: use gpu_batch.states directly (experiences vec is empty)
#[allow(unused_mut)]
let mut batch_tensor_opt: Option<Tensor> = None;
#[cfg(feature = "cuda")]
{
if let Some(ref gpu) = batch_sample.gpu_batch {
batch_tensor_opt = gpu.states
.to_dtype(training_dtype(agent.device()))
.ok();
}
}
let batch_tensor = if let Some(t) = batch_tensor_opt {
t
} else {
let batched_states: Vec<f32> = batch_sample.experiences
.iter()
.flat_map(|exp| {
let mut s = exp.state.clone();
s.resize(state_dim, 0.0);
s
})
.collect();
let t = match Tensor::from_vec(batched_states, (sample_size, state_dim), agent.device()) {
Ok(t) => t,
Err(_) => return None,
},
Err(_) => return None,
};
match t.to_dtype(training_dtype(agent.device())) {
Ok(t) => t,
Err(_) => return None,
}
};
let batch_q_values = match agent.forward(&batch_tensor) {
@@ -4318,24 +4356,39 @@ impl DQNTrainer {
};
let state_dim = agent.get_state_dim();
let batched_states: Vec<f32> = batch_sample
.experiences
.iter()
.flat_map(|exp| {
let mut s = exp.state.clone();
s.resize(state_dim, 0.0);
s
})
.collect();
let batch_tensor =
match Tensor::from_vec(batched_states, (sample_size, state_dim), agent.device()) {
Ok(t) => match t.to_dtype(training_dtype(agent.device())) {
Ok(t) => t,
Err(_) => return None,
},
// GPU PER path: use gpu_batch.states directly (experiences vec is empty)
#[allow(unused_mut)]
let mut batch_tensor_opt: Option<Tensor> = None;
#[cfg(feature = "cuda")]
{
if let Some(ref gpu) = batch_sample.gpu_batch {
batch_tensor_opt = gpu.states
.to_dtype(training_dtype(agent.device()))
.ok();
}
}
let batch_tensor = if let Some(t) = batch_tensor_opt {
t
} else {
let batched_states: Vec<f32> = batch_sample
.experiences
.iter()
.flat_map(|exp| {
let mut s = exp.state.clone();
s.resize(state_dim, 0.0);
s
})
.collect();
let t = match Tensor::from_vec(batched_states, (sample_size, state_dim), agent.device()) {
Ok(t) => t,
Err(_) => return None,
};
match t.to_dtype(training_dtype(agent.device())) {
Ok(t) => t,
Err(_) => return None,
}
};
let batch_q_values = match agent.forward(&batch_tensor) {
Ok(q) => q,