From b616d024ad7c34392ba12f5d96178fbe59dc1409 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 8 Mar 2026 22:31:49 +0100 Subject: [PATCH] fix(dqn): IQN GPU PER weights, staged GPU buffer, CUDA default in all ML crates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three fixes for GPU PER hot path: 1. IQN quantile loss used empty CPU weights Vec instead of GPU-resident weights_tensor_cached — caused CUDA_ERROR_ILLEGAL_ADDRESS from uninitialized GPU memory. Now uses cached GPU tensor matching C51 and standard DQN paths. 2. GpuPrioritized add()/add_batch() replaced with StagedGpuBuffer: add() stages on CPU (Vec::push, zero GPU ops), sample() batch-flushes staging→GPU in one DMA before sampling. Production path (insert_batch_tensors) bypasses staging entirely — GPU→GPU with zero CPU. 3. All 9 ML sub-crates default to cuda feature so `cargo test -p ml-dqn` exercises GPU code paths on CUDA workstations. CI service crates use default-features=false, unaffected. Test results: 350 passed (was 343+7 failed), 0 failed, 1 ignored. Co-Authored-By: Claude Opus 4.6 --- crates/ml-core/Cargo.toml | 2 +- crates/ml-dqn/Cargo.toml | 2 +- crates/ml-dqn/src/dqn.rs | 15 ++-- crates/ml-dqn/src/replay_buffer_type.rs | 103 +++++++++++++++++++----- crates/ml-ensemble/Cargo.toml | 2 +- crates/ml-explainability/Cargo.toml | 2 +- crates/ml-hyperopt/Cargo.toml | 2 +- crates/ml-labeling/Cargo.toml | 2 +- crates/ml-ppo/Cargo.toml | 2 +- crates/ml-supervised/Cargo.toml | 2 +- crates/ml/Cargo.toml | 2 +- crates/ml/src/trainers/dqn/config.rs | 2 +- 12 files changed, 105 insertions(+), 33 deletions(-) diff --git a/crates/ml-core/Cargo.toml b/crates/ml-core/Cargo.toml index 184161e6b..659e69b48 100644 --- a/crates/ml-core/Cargo.toml +++ b/crates/ml-core/Cargo.toml @@ -14,7 +14,7 @@ categories.workspace = true description = "Shared ML types, traits, and infrastructure for Foxhunt" [features] -default = [] +default = ["cuda"] cuda = ["candle-core/cuda", "candle-core/cudnn", "candle-nn/cuda", "candle-nn/cudnn"] s3-storage = ["aws-config", "aws-sdk-s3", "aws-types", "aws-credential-types", "urlencoding"] high-precision = ["rust_decimal/serde-float"] diff --git a/crates/ml-dqn/Cargo.toml b/crates/ml-dqn/Cargo.toml index 88389411d..bcd249d76 100644 --- a/crates/ml-dqn/Cargo.toml +++ b/crates/ml-dqn/Cargo.toml @@ -14,7 +14,7 @@ categories.workspace = true description = "DQN reinforcement learning for Foxhunt trading" [features] -default = [] +default = ["cuda"] cuda = ["candle-core/cuda", "candle-core/cudnn", "candle-nn/cuda", "candle-nn/cudnn"] [dependencies] diff --git a/crates/ml-dqn/src/dqn.rs b/crates/ml-dqn/src/dqn.rs index a98b2dbb8..9a45955f5 100644 --- a/crates/ml-dqn/src/dqn.rs +++ b/crates/ml-dqn/src/dqn.rs @@ -2493,11 +2493,16 @@ impl DQN { self.config.iqn_kappa, ).map_err(|e| MLError::TrainingError(format!("Quantile Huber loss failed: {}", e)))?; - // Apply PER importance-sampling weights (detached — not learnable) - let weights_tensor = Tensor::from_vec(weights.clone(), batch_size, &device)? - .to_dtype(dtype) - .map_err(|e| MLError::TrainingError(format!("Failed to cast weights tensor: {}", e)))? - .detach(); + // IS weights: use cached GPU tensor or create from Vec + let weights_tensor = if let Some(ref wt) = weights_tensor_cached { + wt.clone() + } else { + Tensor::from_vec(weights.clone(), batch_size, device) + .map_err(|e| MLError::TrainingError(format!("Failed to create weights tensor: {}", e)))? + .to_dtype(dtype) + .map_err(|e| MLError::TrainingError(format!("Failed to cast weights tensor: {}", e)))? + .detach() + }; let weighted_loss = (per_sample_loss * weights_tensor)?.mean_all()?; weighted_loss diff --git a/crates/ml-dqn/src/replay_buffer_type.rs b/crates/ml-dqn/src/replay_buffer_type.rs index 0d9490316..a3f8dc951 100644 --- a/crates/ml-dqn/src/replay_buffer_type.rs +++ b/crates/ml-dqn/src/replay_buffer_type.rs @@ -50,6 +50,59 @@ pub struct GpuBatch { pub indices: candle_core::Tensor, // [batch_size] u32 on GPU (buffer indices) } +/// GPU buffer + CPU staging: `add()` stages on CPU (zero GPU ops), +/// `sample()` batch-flushes staging → GPU in one DMA before sampling. +#[cfg(feature = "cuda")] +pub struct StagedGpuBuffer { + pub gpu: crate::gpu_replay_buffer::GpuReplayBuffer, + staging: Vec, +} + +#[cfg(feature = "cuda")] +impl StagedGpuBuffer { + fn flush(&mut self) -> Result<(), MLError> { + if self.staging.is_empty() { + return Ok(()); + } + let n = self.staging.len(); + let sd = self.staging.first().map_or(0, |e| e.state.len()); + let device = self.gpu.device().clone(); + + let mut states_flat = Vec::with_capacity(n * sd); + let mut next_flat = Vec::with_capacity(n * sd); + let mut actions = Vec::with_capacity(n); + let mut rewards = Vec::with_capacity(n); + let mut dones = Vec::with_capacity(n); + + for exp in self.staging.drain(..) { + states_flat.extend_from_slice(&exp.state); + if exp.state.len() < sd { + states_flat.resize(states_flat.len() + sd - exp.state.len(), 0.0); + } + next_flat.extend_from_slice(&exp.next_state); + if exp.next_state.len() < sd { + next_flat.resize(next_flat.len() + sd - exp.next_state.len(), 0.0); + } + actions.push(exp.action as u32); + rewards.push(exp.reward_f32()); + dones.push(if exp.done { 1.0_f32 } else { 0.0 }); + } + + let states_t = candle_core::Tensor::from_vec(states_flat, (n, sd), &device) + .map_err(|e| MLError::TrainingError(format!("GPU flush states: {e}")))?; + let next_t = candle_core::Tensor::from_vec(next_flat, (n, sd), &device) + .map_err(|e| MLError::TrainingError(format!("GPU flush next_states: {e}")))?; + let actions_t = candle_core::Tensor::from_vec(actions, n, &device) + .map_err(|e| MLError::TrainingError(format!("GPU flush actions: {e}")))?; + let rewards_t = candle_core::Tensor::from_vec(rewards, n, &device) + .map_err(|e| MLError::TrainingError(format!("GPU flush rewards: {e}")))?; + let dones_t = candle_core::Tensor::from_vec(dones, n, &device) + .map_err(|e| MLError::TrainingError(format!("GPU flush dones: {e}")))?; + + self.gpu.insert_batch(&states_t, &next_t, &actions_t, &rewards_t, &dones_t) + } +} + /// Enum wrapper for runtime buffer type switching. /// /// Clone is cheap (Arc inner types) and enables sharing the buffer between @@ -60,9 +113,11 @@ pub enum ReplayBufferType { Uniform(Arc>), /// Prioritized sampling based on TD errors (Rainbow DQN) Prioritized(Arc), - /// GPU-resident prioritized replay — all sampling on GPU (zero CPU round-trips) + /// GPU-resident prioritized replay with CPU staging buffer. + /// `add()` stages on CPU (zero GPU ops). `sample()` flushes staging → GPU + /// in one batch DMA, then samples entirely on GPU. #[cfg(feature = "cuda")] - GpuPrioritized(Arc>), + GpuPrioritized(Arc>), } impl std::fmt::Debug for ReplayBufferType { @@ -134,7 +189,10 @@ impl ReplayBufferType { }; let buffer = GpuReplayBuffer::new(config, device)?; - Ok(Self::GpuPrioritized(Arc::new(Mutex::new(buffer)))) + Ok(Self::GpuPrioritized(Arc::new(Mutex::new(StagedGpuBuffer { + gpu: buffer, + staging: Vec::new(), + })))) } /// Try GPU-resident PER, fall back to CPU PER on allocation failure. @@ -222,8 +280,9 @@ impl ReplayBufferType { } #[cfg(feature = "cuda")] Self::GpuPrioritized(buffer) => { - let buf = buffer.lock(); - let gpu_batch = buf.sample_proportional(batch_size)?; + let mut buf = buffer.lock(); + buf.flush()?; + let gpu_batch = buf.gpu.sample_proportional(batch_size)?; Ok(BatchSample { experiences: vec![], // Empty — data is on GPU weights: vec![], // Empty — weights on GPU @@ -246,8 +305,8 @@ impl ReplayBufferType { buffer.push(experience) } #[cfg(feature = "cuda")] - Self::GpuPrioritized(_) => { - // No-op: GPU buffer uses insert_batch() from GPU experience collector + Self::GpuPrioritized(buf) => { + buf.lock().staging.push(experience); Ok(()) } } @@ -275,8 +334,8 @@ impl ReplayBufferType { Ok(()) } #[cfg(feature = "cuda")] - Self::GpuPrioritized(_) => { - // No-op: GPU buffer uses insert_batch() with tensors directly + Self::GpuPrioritized(buf) => { + buf.lock().staging.extend(experiences); Ok(()) } } @@ -311,7 +370,7 @@ impl ReplayBufferType { match self { Self::GpuPrioritized(buffer) => { let mut buf = buffer.lock(); - buf.update_priorities_gpu(indices, td_errors) + buf.gpu.update_priorities_gpu(indices, td_errors) } Self::Uniform(_) | Self::Prioritized(_) => Ok(()), // No-op for non-GPU buffers } @@ -322,7 +381,7 @@ impl ReplayBufferType { match self { Self::Prioritized(buffer) => buffer.step(), #[cfg(feature = "cuda")] - Self::GpuPrioritized(buffer) => buffer.lock().step(), + Self::GpuPrioritized(buffer) => buffer.lock().gpu.step(), _ => {} } } @@ -353,7 +412,7 @@ impl ReplayBufferType { Self::Uniform(_) => 1.0, Self::Prioritized(buffer) => buffer.current_beta(), #[cfg(feature = "cuda")] - Self::GpuPrioritized(buffer) => buffer.lock().current_beta(), + Self::GpuPrioritized(buffer) => buffer.lock().gpu.current_beta(), } } @@ -363,7 +422,10 @@ impl ReplayBufferType { Self::Uniform(buffer) => buffer.lock().len(), Self::Prioritized(buffer) => buffer.len(), #[cfg(feature = "cuda")] - Self::GpuPrioritized(buffer) => buffer.lock().len(), + Self::GpuPrioritized(buffer) => { + let buf = buffer.lock(); + buf.gpu.len() + buf.staging.len() + } } } @@ -384,7 +446,9 @@ impl ReplayBufferType { } #[cfg(feature = "cuda")] Self::GpuPrioritized(buffer) => { - buffer.lock().clear(); + let mut buf = buffer.lock(); + buf.gpu.clear(); + buf.staging.clear(); } } } @@ -395,7 +459,10 @@ impl ReplayBufferType { Self::Uniform(buffer) => buffer.lock().len() >= batch_size, Self::Prioritized(buffer) => buffer.can_sample(batch_size), #[cfg(feature = "cuda")] - Self::GpuPrioritized(buffer) => buffer.lock().can_sample(batch_size), + Self::GpuPrioritized(buffer) => { + let buf = buffer.lock(); + buf.gpu.len() + buf.staging.len() >= batch_size + } } } @@ -495,13 +562,13 @@ impl ReplayBufferType { Self::Uniform(buffer) => buffer.lock().capacity(), Self::Prioritized(buffer) => buffer.capacity(), #[cfg(feature = "cuda")] - Self::GpuPrioritized(buffer) => buffer.lock().capacity(), + Self::GpuPrioritized(buffer) => buffer.lock().gpu.capacity(), } } - /// Get a locked reference to the inner GPU replay buffer (GpuPrioritized only). + /// Get a locked reference to the staged GPU buffer (GpuPrioritized only). #[cfg(feature = "cuda")] - pub fn as_gpu_buffer(&self) -> Option> { + pub fn as_gpu_buffer(&self) -> Option> { match self { Self::GpuPrioritized(buffer) => Some(buffer.lock()), Self::Uniform(_) | Self::Prioritized(_) => None, diff --git a/crates/ml-ensemble/Cargo.toml b/crates/ml-ensemble/Cargo.toml index 9ff49574d..e168eb5b1 100644 --- a/crates/ml-ensemble/Cargo.toml +++ b/crates/ml-ensemble/Cargo.toml @@ -14,7 +14,7 @@ categories.workspace = true description = "ML ensemble coordination — voting, confidence, gating, hot-swap, inference pipeline" [features] -default = [] +default = ["cuda"] cuda = ["candle-core/cuda", "candle-core/cudnn"] [dependencies] diff --git a/crates/ml-explainability/Cargo.toml b/crates/ml-explainability/Cargo.toml index 2f8ce147c..a7783e885 100644 --- a/crates/ml-explainability/Cargo.toml +++ b/crates/ml-explainability/Cargo.toml @@ -14,7 +14,7 @@ categories.workspace = true description = "Model explainability (integrated gradients) for Foxhunt ML" [features] -default = [] +default = ["cuda"] cuda = ["candle-core/cuda", "candle-core/cudnn", "candle-nn/cuda", "candle-nn/cudnn"] [dependencies] diff --git a/crates/ml-hyperopt/Cargo.toml b/crates/ml-hyperopt/Cargo.toml index dd948f905..f059682d0 100644 --- a/crates/ml-hyperopt/Cargo.toml +++ b/crates/ml-hyperopt/Cargo.toml @@ -14,7 +14,7 @@ categories.workspace = true description = "ML hyperparameter optimization — PSO, TPE, campaigns, sensitivity analysis" [features] -default = [] +default = ["cuda"] cuda = ["candle-core/cuda", "candle-core/cudnn"] [dependencies] diff --git a/crates/ml-labeling/Cargo.toml b/crates/ml-labeling/Cargo.toml index 84e4151e5..abe081710 100644 --- a/crates/ml-labeling/Cargo.toml +++ b/crates/ml-labeling/Cargo.toml @@ -14,7 +14,7 @@ categories.workspace = true description = "ML labeling algorithms for Foxhunt HFT training data" [features] -default = [] +default = ["cuda"] cuda = ["candle-core/cuda", "candle-core/cudnn"] [dependencies] diff --git a/crates/ml-ppo/Cargo.toml b/crates/ml-ppo/Cargo.toml index dc6cf868a..ee37f782f 100644 --- a/crates/ml-ppo/Cargo.toml +++ b/crates/ml-ppo/Cargo.toml @@ -14,7 +14,7 @@ categories.workspace = true description = "PPO reinforcement learning for Foxhunt trading" [features] -default = [] +default = ["cuda"] cuda = ["candle-core/cuda", "candle-core/cudnn", "candle-nn/cuda", "candle-nn/cudnn"] [dependencies] diff --git a/crates/ml-supervised/Cargo.toml b/crates/ml-supervised/Cargo.toml index dda42407a..4f7a34b63 100644 --- a/crates/ml-supervised/Cargo.toml +++ b/crates/ml-supervised/Cargo.toml @@ -14,7 +14,7 @@ categories.workspace = true description = "Supervised models (TFT, Mamba, Liquid, TGGN, TLOB, KAN, xLSTM, Diffusion)" [features] -default = [] +default = ["cuda"] cuda = ["candle-core/cuda", "candle-core/cudnn", "candle-nn/cuda", "candle-nn/cudnn"] [dependencies] diff --git a/crates/ml/Cargo.toml b/crates/ml/Cargo.toml index 672809bfc..df46301ca 100644 --- a/crates/ml/Cargo.toml +++ b/crates/ml/Cargo.toml @@ -15,7 +15,7 @@ categories.workspace = true [features] # MINIMAL features for HFT inference only - ALL HEAVY ML REMOVED # CUDA opt-in: service crates build on CPU nodes without nvcc -default = ["minimal-inference"] +default = ["minimal-inference", "cuda"] # PRODUCTION FEATURES - LIGHTWEIGHT ONLY minimal-inference = [] # Minimal inference with no optional deps diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 083907a9b..d8874d206 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -298,7 +298,7 @@ impl DQNAgentType { match self { Self::Standard(agent) => { if let Some(mut gpu_buf) = agent.memory.as_gpu_buffer() { - gpu_buf.insert_batch(states, next_states, actions, rewards, dones) + gpu_buf.gpu.insert_batch(states, next_states, actions, rewards, dones) } else { // CPU PER fallback: convert tensors back to experiences let batch_size = states.dim(0).map_err(|e| {