diff --git a/crates/ml-dqn/src/dqn.rs b/crates/ml-dqn/src/dqn.rs index 93cb6730f..0c9cc4605 100644 --- a/crates/ml-dqn/src/dqn.rs +++ b/crates/ml-dqn/src/dqn.rs @@ -833,6 +833,26 @@ pub struct GradientResult { pub grad_norm_gpu: Option, } +impl GradientResult { + /// Extract the GPU loss tensor as `CudaSlice` for the training guard. + /// + /// Returns `None` if `loss_tensor_gpu` is `None`. + pub fn loss_cuda_slice( + &self, + ) -> Option, MLError>> { + self.loss_tensor_gpu.as_ref().map(GpuTrainResult::tensor_scalar_to_cuda_slice) + } + + /// Extract the GPU grad norm tensor as `CudaSlice` for the training guard. + /// + /// Returns `None` if `grad_norm_gpu` is `None`. + pub fn grad_norm_cuda_slice( + &self, + ) -> Option, MLError>> { + self.grad_norm_gpu.as_ref().map(GpuTrainResult::tensor_scalar_to_cuda_slice) + } +} + /// GPU-resident training step result — **zero CPU readback**. /// /// Loss and gradient norm stay as GPU scalar tensors. The caller (trainer) @@ -846,6 +866,87 @@ pub struct GpuTrainResult { pub grad_norm_gpu: Tensor, } +impl GpuTrainResult { + /// Extract loss as a `CudaSlice` — zero-alloc view into the underlying + /// Candle storage (scalar F32 tensor guaranteed contiguous). + /// + /// Returns an owned `CudaSlice` via DtoD copy (4 bytes for a scalar). + /// Use this instead of `tensor_to_cuda_slice_f32(&self.loss_gpu)` to avoid + /// re-importing the converter at every call site. + pub fn loss_cuda_slice( + &self, + ) -> Result, MLError> { + Self::tensor_scalar_to_cuda_slice(&self.loss_gpu) + } + + /// Extract grad_norm as a `CudaSlice` — same as [`loss_cuda_slice`]. + pub fn grad_norm_cuda_slice( + &self, + ) -> Result, MLError> { + Self::tensor_scalar_to_cuda_slice(&self.grad_norm_gpu) + } + + /// Construct from raw f32 scalars (fused CUDA training path). + /// + /// Uploads two f32 scalars to the GPU as rank-0 Tensors, avoiding the + /// need for the caller to import `candle_core::Tensor`. + pub fn from_fused_scalars(loss: f32, grad_norm: f32, device: &Device) -> Result { + Ok(Self { + loss_gpu: Tensor::new(loss, device)?, + grad_norm_gpu: Tensor::new(grad_norm, device)?, + }) + } + + /// Extract a contiguous F32 scalar tensor to CudaSlice. + /// + /// Also used by [`GradientResult`] accessors. + pub fn tensor_scalar_to_cuda_slice( + tensor: &Tensor, + ) -> Result, MLError> { + let tensor = if tensor.dtype() != DType::F32 { + tensor.to_dtype(DType::F32).map_err(|e| { + MLError::ModelError(format!("GpuTrainResult scalar dtype cast: {e}")) + })? + } else { + tensor.clone() + }; + let tensor = tensor.contiguous().map_err(|e| { + MLError::ModelError(format!("GpuTrainResult scalar contiguous: {e}")) + })?; + let (storage, layout) = tensor.storage_and_layout(); + match &*storage { + candle_core::Storage::Cuda(cs) => { + let slice = cs.as_cuda_slice::().map_err(|e| { + MLError::ModelError(format!("GpuTrainResult as_cuda_slice: {e}")) + })?; + let view = slice.slice(layout.start_offset()..); + let n = view.len(); + let stream = cs.device.cuda_stream(); + let dst = stream.alloc_zeros::(n).map_err(|e| { + MLError::ModelError(format!("GpuTrainResult alloc: {e}")) + })?; + { + use candle_core::cuda_backend::cudarc::driver::DevicePtr; + let (src_ptr, _src_guard) = view.device_ptr(&stream); + let (dst_ptr, _dst_guard) = dst.device_ptr(&stream); + let num_bytes = n * std::mem::size_of::(); + #[allow(unsafe_code)] + unsafe { + candle_core::cuda_backend::cudarc::driver::result::memcpy_dtod_async( + dst_ptr, src_ptr, num_bytes, stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("GpuTrainResult DtoD: {e}")))?; + } + } + drop(storage); + Ok(dst) + } + _ => Err(MLError::ModelError( + "GpuTrainResult: tensor must be on CUDA device".to_owned(), + )), + } + } +} + /// Internal result from forward pass + loss computation (no backward pass). struct ComputeLossResult { /// The loss tensor (still in the computation graph for backward pass). diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index c24895d83..d270b8142 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -25,7 +25,7 @@ use std::sync::Arc; use anyhow::Result; -use candle_core::{Device, Tensor}; +use candle_core::{Device, Tensor}; // Tensor: HER batch relabeling ops use tracing::info; use crate::cuda_pipeline::gpu_dqn_trainer::{GpuDqnTrainConfig, GpuDqnTrainer}; @@ -441,13 +441,13 @@ impl FusedTrainingCtx { self.steps_since_varmap_sync += 1; - // ── Step 7: Create GPU scalar tensors for monitoring ───────────── - Ok(GpuTrainResult { - loss_gpu: Tensor::new(fused_result.total_loss, device) - .map_err(|e| anyhow::anyhow!("Fused loss->Tensor: {e}"))?, - grad_norm_gpu: Tensor::new(fused_result.grad_norm, device) - .map_err(|e| anyhow::anyhow!("Fused grad_norm->Tensor: {e}"))?, - }) + // ── Step 7: Wrap raw scalars into GpuTrainResult ───────────────── + // Uses from_fused_scalars() to avoid importing Tensor at this call site. + GpuTrainResult::from_fused_scalars( + fused_result.total_loss, + fused_result.grad_norm, + device, + ).map_err(|e| anyhow::anyhow!("Fused scalars->GpuTrainResult: {e}")) } /// Sync CudaSlice weights back to VarMap (deferred -- called at epoch boundary). diff --git a/crates/ml/src/trainers/dqn/trainer/train_step.rs b/crates/ml/src/trainers/dqn/trainer/train_step.rs index 2f4989f8b..936d71e57 100644 --- a/crates/ml/src/trainers/dqn/trainer/train_step.rs +++ b/crates/ml/src/trainers/dqn/trainer/train_step.rs @@ -120,10 +120,10 @@ impl DQNTrainer { let warmup_steps = (self.collapse_warmup_buffer_size as f64 * 0.2) as u64; let past_warmup = self.gradient_logging_step as u64 > warmup_steps; - let loss_slice = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&gpu_result.loss_gpu) - .map_err(|e| anyhow::anyhow!("GPU guard loss->CudaSlice: {e}"))?; - let grad_slice = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&gpu_result.grad_norm_gpu) - .map_err(|e| anyhow::anyhow!("GPU guard grad->CudaSlice: {e}"))?; + let loss_slice = gpu_result.loss_cuda_slice() + .map_err(|e| anyhow::anyhow!("GPU guard loss CudaSlice: {e}"))?; + let grad_slice = gpu_result.grad_norm_cuda_slice() + .map_err(|e| anyhow::anyhow!("GPU guard grad CudaSlice: {e}"))?; let result = guard .check_and_accumulate( &loss_slice, @@ -345,6 +345,17 @@ impl DQNTrainer { .compute_gradients(batch) .map_err(|e| anyhow::anyhow!("Gradient computation step {} failed: {}", step, e))?; + // Extract CudaSlice handles and destructure BEFORE moving fields out. + let loss_sl_opt = result.loss_cuda_slice(); + let grad_sl_opt = result.grad_norm_cuda_slice(); + let r_grads = result.grads; + let r_td_errors = result.td_errors; + let r_indices = result.indices.clone(); + let r_td_gpu = result.td_errors_gpu; + let r_idx_gpu = result.indices_gpu; + let r_loss_gpu = result.loss_tensor_gpu; + let r_gn_gpu = result.grad_norm_gpu; + // Get vars for accumulation. Var is an Arc wrapper so cloning is cheap. let vars: Vec = agent .optimizer_vars() @@ -352,18 +363,18 @@ impl DQNTrainer { crate::gradient_accumulation::accumulate_grads( &mut accumulated_grads, - result.grads, + r_grads, &vars, ) .map_err(|e| anyhow::anyhow!("Gradient accumulation step {} failed: {}", step, e))?; - all_td_errors.extend(result.td_errors); - all_indices.extend(result.indices); - // GPU guard: check + accumulate loss/grad for this sub-step (borrows - // tensors before the move into gpu_*_tensors below). + all_td_errors.extend(r_td_errors); + all_indices.extend(r_indices); + // GPU guard: check + accumulate loss/grad for this sub-step. + // Uses CudaSlice accessors on GradientResult to avoid tensor_to_cuda_slice_f32. { - if let (Some(ref loss_gpu), Some(ref gn_gpu)) = - (&result.loss_tensor_gpu, &result.grad_norm_gpu) + if let (Some(loss_sl_r), Some(grad_sl_r)) = + (loss_sl_opt, grad_sl_opt) { if let Some(ref mut guard) = self.training_guard { let grad_collapse_threshold = @@ -373,10 +384,10 @@ impl DQNTrainer { let warmup_steps = (self.collapse_warmup_buffer_size as f64 * 0.2) as u64; let past_warmup = self.gradient_logging_step as u64 > warmup_steps; - let loss_slice = crate::cuda_pipeline::tensor_to_cuda_slice_f32(loss_gpu) - .map_err(|e| anyhow::anyhow!("GPU guard accum loss->CudaSlice: {e}"))?; - let grad_slice = crate::cuda_pipeline::tensor_to_cuda_slice_f32(gn_gpu) - .map_err(|e| anyhow::anyhow!("GPU guard accum grad->CudaSlice: {e}"))?; + let loss_slice = loss_sl_r + .map_err(|e| anyhow::anyhow!("GPU guard accum loss CudaSlice: {e}"))?; + let grad_slice = grad_sl_r + .map_err(|e| anyhow::anyhow!("GPU guard accum grad CudaSlice: {e}"))?; let guard_result = guard.check_and_accumulate( &loss_slice, &grad_slice, @@ -401,20 +412,19 @@ impl DQNTrainer { } } { - if let Some(td_gpu) = result.td_errors_gpu { + if let Some(td_gpu) = r_td_gpu { gpu_td_errors.push(td_gpu); } - if let Some(idx_gpu) = result.indices_gpu { + if let Some(idx_gpu) = r_idx_gpu { gpu_indices.push(idx_gpu); } - if let Some(loss_gpu) = result.loss_tensor_gpu { + if let Some(loss_gpu) = r_loss_gpu { gpu_loss_tensors.push(loss_gpu); } - if let Some(gn_gpu) = result.grad_norm_gpu { + if let Some(gn_gpu) = r_gn_gpu { gpu_grad_tensors.push(gn_gpu); } } - // CPU sentinel fallback (non-CUDA only) } // === Phase 2: Average and apply gradients (single optimizer step) === diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 863bb7197..7f50ad2fc 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -938,10 +938,10 @@ impl DQNTrainer { }; if let Some(ref mut guard) = self.training_guard { - let loss_slice = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&_gpu_result.loss_gpu) - .map_err(|e| anyhow::anyhow!("guard loss->CudaSlice: {e}"))?; - let grad_slice = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&_gpu_result.grad_norm_gpu) - .map_err(|e| anyhow::anyhow!("guard grad->CudaSlice: {e}"))?; + let loss_slice = _gpu_result.loss_cuda_slice() + .map_err(|e| anyhow::anyhow!("guard loss CudaSlice: {e}"))?; + let grad_slice = _gpu_result.grad_norm_cuda_slice() + .map_err(|e| anyhow::anyhow!("guard grad CudaSlice: {e}"))?; let gr = guard.check_and_accumulate( &loss_slice, &grad_slice, @@ -992,23 +992,30 @@ impl DQNTrainer { } }; + // Extract CudaSlice handles and destructure result BEFORE moving grads. + let loss_sl_opt = result.loss_cuda_slice(); + let grad_sl_opt = result.grad_norm_cuda_slice(); + let r_td_gpu = result.td_errors_gpu; + let r_idx_gpu = result.indices_gpu; + let r_grads = result.grads; + let vars = agent.optimizer_vars() .map_err(|e| anyhow::anyhow!("optimizer vars: {e}"))?; crate::gradient_accumulation::accumulate_grads( &mut accumulated_grads, - result.grads, + r_grads, &vars, ).map_err(|e| anyhow::anyhow!("grad accum: {e}"))?; { - if let (Some(ref loss_t), Some(ref gn_t)) = - (&result.loss_tensor_gpu, &result.grad_norm_gpu) + if let (Some(loss_sl_r), Some(grad_sl_r)) = + (loss_sl_opt, grad_sl_opt) { if let Some(ref mut guard) = self.training_guard { - let loss_sl = crate::cuda_pipeline::tensor_to_cuda_slice_f32(loss_t) - .map_err(|e| anyhow::anyhow!("accum loss->CudaSlice: {e}"))?; - let grad_sl = crate::cuda_pipeline::tensor_to_cuda_slice_f32(gn_t) - .map_err(|e| anyhow::anyhow!("accum grad->CudaSlice: {e}"))?; + let loss_sl = loss_sl_r + .map_err(|e| anyhow::anyhow!("accum loss CudaSlice: {e}"))?; + let grad_sl = grad_sl_r + .map_err(|e| anyhow::anyhow!("accum grad CudaSlice: {e}"))?; let gr = guard.check_and_accumulate( &loss_sl, &grad_sl, 1e6_f32, guard_collapse_thresh, !guard_past_warmup, @@ -1027,10 +1034,10 @@ impl DQNTrainer { } } } - if let Some(td) = result.td_errors_gpu { + if let Some(td) = r_td_gpu { group_td_gpu.push(td); } - if let Some(idx) = result.indices_gpu { + if let Some(idx) = r_idx_gpu { group_idx_gpu.push(idx); } }