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

refactor(cuda): replace Tensor→CudaSlice conversions with direct accessors in trainers/

Browse Source 
Add CudaSlice<f32> accessor methods to GpuTrainResult and GradientResult
(loss_cuda_slice, grad_norm_cuda_slice) in ml-dqn, eliminating redundant
tensor_to_cuda_slice_f32 calls at every training guard check site.

- GpuTrainResult: add from_fused_scalars() constructor and CudaSlice extractors
- GradientResult: add CudaSlice extractors that delegate to GpuTrainResult
- fused_training.rs: use from_fused_scalars() instead of Tensor::new() wrapping
- train_step.rs: use result.loss_cuda_slice() instead of tensor_to_cuda_slice_f32
- training_loop.rs: same CudaSlice accessor pattern for both fused and accum paths

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-17 19:01:02 +01:00
parent 31a7c1b171
commit 8a6bca399c
4 changed files with 159 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -833,6 +833,26 @@ pub struct GradientResult {
pub grad_norm_gpu: Option<Tensor>,
}
impl GradientResult {
/// Extract the GPU loss tensor as `CudaSlice<f32>` for the training guard.
///
/// Returns `None` if `loss_tensor_gpu` is `None`.
pub fn loss_cuda_slice(
&self,
) -> Option<Result<candle_core::cuda_backend::cudarc::driver::CudaSlice<f32>, MLError>> {
self.loss_tensor_gpu.as_ref().map(GpuTrainResult::tensor_scalar_to_cuda_slice)
}
/// Extract the GPU grad norm tensor as `CudaSlice<f32>` for the training guard.
///
/// Returns `None` if `grad_norm_gpu` is `None`.
pub fn grad_norm_cuda_slice(
&self,
) -> Option<Result<candle_core::cuda_backend::cudarc::driver::CudaSlice<f32>, 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<f32>` — 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<candle_core::cuda_backend::cudarc::driver::CudaSlice<f32>, MLError> {
Self::tensor_scalar_to_cuda_slice(&self.loss_gpu)
}
/// Extract grad_norm as a `CudaSlice<f32>` — same as [`loss_cuda_slice`].
pub fn grad_norm_cuda_slice(
&self,
) -> Result<candle_core::cuda_backend::cudarc::driver::CudaSlice<f32>, 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<Self, MLError> {
Ok(Self {
loss_gpu: Tensor::new(loss, device)?,
grad_norm_gpu: Tensor::new(grad_norm, device)?,
})
}
/// Extract a contiguous F32 scalar tensor to CudaSlice<f32>.
///
/// Also used by [`GradientResult`] accessors.
pub fn tensor_scalar_to_cuda_slice(
tensor: &Tensor,
) -> Result<candle_core::cuda_backend::cudarc::driver::CudaSlice<f32>, 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::<f32>().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::<f32>(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::<f32>();
#[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).

16
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -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).

50
crates/ml/src/trainers/dqn/trainer/train_step.rs
View File

@@ -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<candle_core::Var> = 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) ===

33
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -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);
}
}