From 8c86325ff81d01774faef08f70cca19c3bb111c5 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 19 Mar 2026 23:50:33 +0100 Subject: [PATCH] fix: fused training supports RegimeConditional + GpuTensor::cat dim>0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three fixes: 1. FusedTrainingCtx::new now accepts RegimeConditionalDQN by using primary_head() — the old code rejected it, causing silent fallback to non-fused train_step() which has action-space mismatch with branching DQN (5-action Q-table vs 45-action factored indices → CUDA_ERROR_ILLEGAL_ADDRESS buffer overrun) 2. ensure_fused_ctx() returns Result instead of () — fused init failure is now a hard error (no silent CPU fallback) 3. GpuTensor::cat now supports dim>0 concatenation (was unimplemented, caused "dim=1 > 0 not yet implemented" error in validation path) Root cause chain: RegimeConditional rejected by fused init → silent fallback to DQN::train_step() → compute_loss_internal() uses num_actions=5 but actions are 0-44 → gather with out-of-bounds offsets → CUDA_ERROR_ILLEGAL_ADDRESS → async error poisons CUDA context → next stream.synchronize() deadlocks forever Remaining: curiosity kernel crash also poisons the stream. The curiosity training runs inside collect_gpu_experiences() and its async error blocks the PER insert. This needs separate investigation of curiosity_training_kernel.cu. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml-core/src/cuda_autograd/gpu_tensor.rs | 25 ++++++++++++++++--- crates/ml/src/trainers/dqn/fused_training.rs | 6 +---- .../ml/src/trainers/dqn/trainer/train_step.rs | 13 ++++++---- .../src/trainers/dqn/trainer/training_loop.rs | 2 +- 4 files changed, 31 insertions(+), 15 deletions(-) diff --git a/crates/ml-core/src/cuda_autograd/gpu_tensor.rs b/crates/ml-core/src/cuda_autograd/gpu_tensor.rs index 1a208d374..7df51166a 100644 --- a/crates/ml-core/src/cuda_autograd/gpu_tensor.rs +++ b/crates/ml-core/src/cuda_autograd/gpu_tensor.rs @@ -248,10 +248,27 @@ impl GpuTensor { } return Self::from_host(&flat, new_shape, stream); } - // General case: not yet optimized - Err(MLError::ModelError(format!( - "cat: dim={dim} > 0 not yet implemented for GpuTensor" - ))) + // General dim>0: interleave rows along the concat dimension + let ref_shape = tensors.first().map_or(&[] as &[usize], |t| t.shape()); + // outer = product of dims before `dim`, inner = product of dims after `dim` + let outer: usize = ref_shape[..dim].iter().product(); + let inner: usize = if dim + 1 < ref_shape.len() { ref_shape[dim + 1..].iter().product() } else { 1 }; + let mut flat = Vec::with_capacity(outer * total_dim * inner); + for o in 0..outer { + for (ti, t) in tensors.iter().enumerate() { + let d = t.shape().get(dim).copied().unwrap_or(0); + let src = &all_data[ti]; + let src_stride = d * inner; + let start = o * src_stride; + let end = start + src_stride; + if end <= src.len() { + flat.extend_from_slice(&src[start..end]); + } + } + } + let mut new_shape = ref_shape.to_vec(); + new_shape[dim] = total_dim; + Self::from_host(&flat, new_shape, stream) } /// Stack tensors along a new dimension 0. diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 099a7604f..42503f5df 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -98,11 +98,7 @@ impl FusedTrainingCtx { ) -> Result { let dqn = match agent { DQNAgentType::Standard(d) => d, - DQNAgentType::RegimeConditional(_) => { - return Err(anyhow::anyhow!( - "Fused CUDA training requires Standard DQN agent (not RegimeConditional)" - )); - } + DQNAgentType::RegimeConditional(regime) => regime.primary_head(), }; let branching_net = dqn.branching_q_network.as_ref().ok_or_else(|| { diff --git a/crates/ml/src/trainers/dqn/trainer/train_step.rs b/crates/ml/src/trainers/dqn/trainer/train_step.rs index e6045bbd6..376c21902 100644 --- a/crates/ml/src/trainers/dqn/trainer/train_step.rs +++ b/crates/ml/src/trainers/dqn/trainer/train_step.rs @@ -634,18 +634,17 @@ impl DQNTrainer { /// Lazy-init fused CUDA training context (Standard DQN only). /// Recreate if batch_size changed (OOM recovery). - pub(crate) async fn ensure_fused_ctx(&mut self) { + pub(crate) async fn ensure_fused_ctx(&mut self) -> anyhow::Result<()> { let needs_init = match &self.fused_ctx { None => self.device.is_cuda(), Some(ctx) => ctx.batch_size() != self.current_batch_size, }; if !needs_init { - return; + return Ok(()); } - // Fused context requires the unified forked CudaStream let stream = match self.cuda_stream { Some(ref s) => std::sync::Arc::clone(s), - None => return, + None => return Ok(()), }; if self.fused_ctx.is_some() { info!("Fused CUDA context: batch_size changed, recreating"); @@ -660,8 +659,12 @@ impl DQNTrainer { self.fused_ctx = Some(ctx); } Err(e) => { - warn!("Fused CUDA training init failed, using Candle path: {e}"); + return Err(anyhow::anyhow!( + "Fused CUDA training init FAILED (no fallback — non-fused path \ + has action-space mismatch with branching DQN): {e}" + )); } } + Ok(()) } } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 5ad363026..ca9e417ca 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -885,7 +885,7 @@ impl DQNTrainer { // Lazy-init fused CUDA training context (Standard DQN only). // Recreate if batch_size changed (OOM recovery). if num_training_steps > 0 { - self.ensure_fused_ctx().await; + self.ensure_fused_ctx().await?; } // Lazy-init + reset guard accumulators for this epoch