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

perf(cuda): eliminate Candle VarMap/Tensor from DQN GPU training hot path

Browse Source 
Replace all Candle tensor operations in the fused training loop with
pure cudarc CudaSlice operations to prevent Candle tensor Drops from
recording events on the default stream (which conflicts with our forked
training stream via cudarc event tracking).

Key changes:
- Re-enable disable_event_tracking() in GpuDqnTrainer::new() — safe now
  that no Candle tensors are created during training
- Rewrite upload_batch_gpu(): BF16 states/next_states use DtoD + bf16→f32
  kernel instead of Candle to_dtype(F32); F32 rewards/dones/weights use
  direct DtoD copy with layout offset handling
- Load bf16_to_f32_kernel from training module (was defined in CUH but
  not loaded)
- Add train_value_step_raw() to GpuIqlTrainer that takes CudaSlice<f32>
  directly, bypassing Candle tensor manipulation
- IQL in fused training now reuses DQN trainer's already-converted F32
  states_buf/rewards_buf instead of creating Candle temporaries
- Fix dtod_from_candle_f32/u32 to respect Candle layout start_offset

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-17 11:33:14 +01:00
parent 44ddd38112
commit 3b4ba67fc5
3 changed files with 321 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

232
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -236,6 +236,7 @@ pub struct GpuDqnTrainer {
ema_kernel: CudaFunction,
per_update_kernel: CudaFunction,
f32_to_bf16_kernel: CudaFunction,
bf16_to_f32_kernel: CudaFunction,
// ── BF16 weight mirrors (forward kernel reads BF16 for tensor core throughput) ──
// Allocated lazily on first train_step when weight sets are available.
@@ -317,6 +318,14 @@ pub struct GpuDqnTrainer {
/// Used by `execute_train_scalars_only()` to avoid reading back td_errors.
scalars_readback_buf: CudaSlice<f32>,
scalars_readback_host: [f32; 2],
// ── BF16 batch staging buffers ──────────────────────────────────
// Pre-allocated CudaSlice<u16> buffers for DtoD copy of BF16 tensors
// from GpuBatch (states, next_states only — stored in BF16 ring buffer).
// The bf16_to_f32_kernel converts these to the F32 trainer input buffers.
// Rewards, dones, weights are F32 in GpuBatch and use direct DtoD to F32 bufs.
bf16_states_buf: CudaSlice<u16>, // [B * STATE_DIM]
bf16_next_states_buf: CudaSlice<u16>, // [B * STATE_DIM]
}
impl GpuDqnTrainer {
@@ -334,6 +343,14 @@ impl GpuDqnTrainer {
&self.save_h_s2
}
/// Reference to the states buffer on GPU.
///
/// Shape: `[B, STATE_DIM]` — contains the batch's states after `upload_batch_gpu()`.
/// Used by IQL value network to read states without Candle tensor intermediaries.
pub fn states_buf(&self) -> &CudaSlice<f32> {
&self.states_buf
}
/// Reference to the next_states buffer on GPU.
///
/// Shape: `[B, STATE_DIM]` — contains the batch's next states after `train_step()`.
@@ -383,18 +400,16 @@ impl GpuDqnTrainer {
let b = config.batch_size;
let total_params = compute_total_params(&config);
// All GPU components now share this single forked CudaStream (stored in
// DQNTrainer). Disable cudarc's automatic event tracking: it injects
// cuStreamWaitEvent/cuEventRecord into every buffer access, creating
// cross-stream event references that break CUDA Graph capture. Safe
// because all GPU work runs on this single stream — no cross-stream
// synchronization needed.
//
// SAFETY: single-stream architecture, all work synchronized on `stream`.
// All GPU components share this single forked CudaStream. Disable cudarc's
// automatic event tracking: it injects cuStreamWaitEvent/cuEventRecord into
// every CudaSlice access, creating cross-stream event references that
// invalidate CUDA Graph capture and conflict with Candle tensor Drops on
// the default stream. Safe: single forked stream, all work serialized.
// SAFETY: single-stream, all GPU work is serialized on `stream`.
unsafe { stream.context().disable_event_tracking(); }
// ── Compile all 5 training kernels from same module ──────────
let (forward_loss_kernel, backward_kernel, grad_norm_kernel, adam_update_kernel, f32_to_bf16_kernel) =
// ── Compile all 6 training kernels from same module ──────────
let (forward_loss_kernel, backward_kernel, grad_norm_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel) =
compile_training_kernels(&stream, &config)?;
// ── Compile EMA kernel (standalone — not captured in CUDA Graph) ──
@@ -463,6 +478,13 @@ impl GpuDqnTrainer {
let scalars_readback_buf = alloc_f32(&stream, 2, "scalars_readback")?;
let scalars_readback_host = [0.0_f32; 2];
// ── Allocate BF16 batch staging buffers ──────────────────────
// DtoD from GpuBatch's BF16 Candle tensors (states, next_states only)
// → these staging buffers, then bf16_to_f32_kernel → F32 trainer bufs.
// Rewards, dones, weights are F32 in GpuBatch and skip BF16 staging.
let bf16_states_buf = alloc_u16(&stream, b * config.state_dim, "bf16_states")?;
let bf16_next_states_buf = alloc_u16(&stream, b * config.state_dim, "bf16_next_states")?;
// ── Shared memory sizing ────────────────────────────────────
let shmem_bytes = compute_shmem_bytes(&config);
@@ -497,6 +519,7 @@ impl GpuDqnTrainer {
ema_kernel,
per_update_kernel,
f32_to_bf16_kernel,
bf16_to_f32_kernel,
online_dueling_bf16: None,
online_branching_bf16: None,
target_dueling_bf16: None,
@@ -537,9 +560,25 @@ impl GpuDqnTrainer {
batch_max_buf,
scalars_readback_buf,
scalars_readback_host,
bf16_states_buf,
bf16_next_states_buf,
})
}
/// Reference to the bf16→f32 conversion kernel.
///
/// Exposed so that callers (e.g., `FusedTrainingCtx`) can convert BF16
/// GpuBatch tensors to F32 CudaSlice buffers without creating Candle
/// temporary tensors.
pub fn bf16_to_f32_kernel(&self) -> &CudaFunction {
&self.bf16_to_f32_kernel
}
/// Reference to the trainer's forked CudaStream.
pub fn stream(&self) -> &Arc<CudaStream> {
&self.stream
}
// ═══════════════════════════════════════════════════════════════════
// BF16 weight mirror management
// ═══════════════════════════════════════════════════════════════════
@@ -1223,55 +1262,38 @@ impl GpuDqnTrainer {
Ok(())
}
/// Upload batch data from GPU tensors directly — zero CPU roundtrip.
/// Upload batch data from GPU tensors — pure CudaSlice, zero Candle temporaries.
///
/// Takes Candle tensors from `GpuBatch`, casts BF16→F32 on GPU if needed,
/// and copies to the trainer's CudaSlice buffers via DtoD memcpy.
/// Replaces the GPU→CPU `.to_vec1()` → CPU→GPU `memcpy_htod` roundtrip.
/// GpuBatch layout from GpuReplayBuffer:
/// - states/next_states: BF16 (stored in BF16 ring buffer)
/// - rewards/dones/weights: F32 (stored in F32 ring buffer / computed as F32)
/// - actions: U32 (stored in U32 ring buffer)
///
/// ~6 DtoD copies vs old path's 1 HtoD + 6 DtoD + GPU→CPU readback.
/// Net saving: eliminates PCIe round-trip (~10-30μs on H100 Gen5).
/// For BF16 tensors (states, next_states): DtoD → BF16 staging → bf16_to_f32 kernel → F32 buf.
/// For F32 tensors (rewards, dones, weights): DtoD → F32 buf directly.
/// For U32 actions: DtoD → I32 buf (bit-compatible for values in [0, 44]).
///
/// No Candle `to_dtype()` temporaries — eliminates Drop/event conflicts with forked stream.
fn upload_batch_gpu(
&mut self,
gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch,
) -> Result<(), MLError> {
use candle_core::DType;
let b = self.config.batch_size;
let sd = self.config.state_dim;
// ── Cast to F32 + contiguous on GPU (no CPU touch) ───────────
let states_f32 = gpu_batch.states.to_dtype(DType::F32)
.and_then(|t| t.contiguous())
.map_err(|e| MLError::ModelError(format!("GPU batch states F32: {e}")))?;
let next_f32 = gpu_batch.next_states.to_dtype(DType::F32)
.and_then(|t| t.contiguous())
.map_err(|e| MLError::ModelError(format!("GPU batch next_states F32: {e}")))?;
let rewards_f32 = gpu_batch.rewards.to_dtype(DType::F32)
.and_then(|t| t.contiguous())
.map_err(|e| MLError::ModelError(format!("GPU batch rewards F32: {e}")))?;
let dones_f32 = gpu_batch.dones.to_dtype(DType::F32)
.and_then(|t| t.contiguous())
.map_err(|e| MLError::ModelError(format!("GPU batch dones F32: {e}")))?;
let weights_f32 = gpu_batch.weights.to_dtype(DType::F32)
.and_then(|t| t.contiguous())
.map_err(|e| MLError::ModelError(format!("GPU batch weights F32: {e}")))?;
// Actions: U32 tensor (bit-compatible with I32 for values in [0, 44])
let actions_cont = gpu_batch.actions.contiguous()
.map_err(|e| MLError::ModelError(format!("GPU batch actions contiguous: {e}")))?;
// ── States & next_states: BF16 → staging → bf16_to_f32 kernel → F32 ──
dtod_from_candle_bf16(&gpu_batch.states, &self.bf16_states_buf, b * sd, &self.stream, "states")?;
dtod_from_candle_bf16(&gpu_batch.next_states, &self.bf16_next_states_buf, b * sd, &self.stream, "next_states")?;
launch_bf16_to_f32(&self.bf16_to_f32_kernel, &self.bf16_states_buf, &self.states_buf, b * sd, &self.stream, "states")?;
launch_bf16_to_f32(&self.bf16_to_f32_kernel, &self.bf16_next_states_buf, &self.next_states_buf, b * sd, &self.stream, "next_states")?;
// ── DtoD copy from Candle tensor storage → trainer CudaSlice buffers ──
dtod_from_candle_f32(&states_f32, &self.states_buf, b * sd, &self.stream, "states")?;
dtod_from_candle_f32(&next_f32, &self.next_states_buf, b * sd, &self.stream, "next_states")?;
dtod_from_candle_u32_to_i32(&actions_cont, &self.actions_buf, b, &self.stream, "actions")?;
dtod_from_candle_f32(&rewards_f32, &self.rewards_buf, b, &self.stream, "rewards")?;
dtod_from_candle_f32(&dones_f32, &self.dones_buf, b, &self.stream, "dones")?;
dtod_from_candle_f32(&weights_f32, &self.is_weights_buf, b, &self.stream, "is_weights")?;
// ── Rewards, dones, IS-weights: already F32 — direct DtoD ──
dtod_from_candle_f32(&gpu_batch.rewards, &self.rewards_buf, b, &self.stream, "rewards")?;
dtod_from_candle_f32(&gpu_batch.dones, &self.dones_buf, b, &self.stream, "dones")?;
dtod_from_candle_f32(&gpu_batch.weights, &self.is_weights_buf, b, &self.stream, "is_weights")?;
// ── Sync stream: ensure DtoD copies complete before temp tensors drop ──
// One stream sync vs old path's GPU→CPU→GPU PCIe roundtrip.
self.stream.synchronize()
.map_err(|e| MLError::ModelError(format!("GPU upload sync: {e}")))?;
// ── Actions: U32 → I32 DtoD (bit-compatible for values in [0, 44]) ──
dtod_from_candle_u32_to_i32(&gpu_batch.actions, &self.actions_buf, b, &self.stream, "actions")?;
Ok(())
}
@@ -1736,13 +1758,13 @@ impl GpuDqnTrainer {
// ── Compilation ─────────────────────────────────────────────────────────────
/// Compile all 5 training kernels from a single NVRTC compilation.
/// Compile all 6 training kernels from a single NVRTC compilation.
///
/// Returns (forward_loss, backward, grad_norm, adam_update, f32_to_bf16) kernel functions.
/// Returns (forward_loss, backward, grad_norm, adam_update, f32_to_bf16, bf16_to_f32) kernel functions.
fn compile_training_kernels(
stream: &Arc<CudaStream>,
config: &GpuDqnTrainConfig,
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
// shmem_max_in_dim must cover ALL weight matrix row widths used by
// cooperative_load_tile_bf16 in the kernel — including value/advantage heads.
// Missing value_h/adv_h caused shmem tile overflow → CUDA_ERROR_ILLEGAL_ADDRESS
@@ -1847,6 +1869,11 @@ fn compile_training_kernels(
.map_err(|e| {
MLError::ModelError(format!("f32_to_bf16_kernel load: {e}"))
})?;
let bf16_to_f32 = module
.load_function("bf16_to_f32_kernel")
.map_err(|e| {
MLError::ModelError(format!("bf16_to_f32_kernel load: {e}"))
})?;
// Opt in to larger dynamic shared memory for forward_loss and backward
// kernels. CUDA default is 48 KB per block; our tile sizes on ≥100 KB GPUs
@@ -1873,8 +1900,8 @@ fn compile_training_kernels(
);
}
info!("GpuDqnTrainer: 5 kernels compiled and loaded (incl. BF16 converter)");
Ok((forward_loss, backward, grad_norm, adam_update, f32_to_bf16))
info!("GpuDqnTrainer: 6 kernels compiled and loaded (incl. BF16 converters)");
Ok((forward_loss, backward, grad_norm, adam_update, f32_to_bf16, bf16_to_f32))
}
/// Compile the standalone Polyak EMA kernel.
@@ -2049,7 +2076,7 @@ fn dtod_from_candle_f32(
stream: &Arc<CudaStream>,
ctx: &str,
) -> Result<(), MLError> {
let (storage_guard, _layout) = tensor.storage_and_layout();
let (storage_guard, layout) = tensor.storage_and_layout();
let cuda_storage = match &*storage_guard {
candle_core::Storage::Cuda(cs) => cs,
_ => {
@@ -2061,7 +2088,10 @@ fn dtod_from_candle_f32(
let src_slice = cuda_storage.as_cuda_slice::<f32>().map_err(|e| {
MLError::ModelError(format!("dtod_from_candle_f32 {ctx}: as_cuda_slice: {e}"))
})?;
let src_ptr = raw_device_ptr(src_slice, stream);
let offset = layout.start_offset();
let (ptr, guard) = src_slice.device_ptr(stream);
let _no_drop = std::mem::ManuallyDrop::new(guard);
let src_ptr = ptr + (offset * std::mem::size_of::<f32>()) as u64;
let dst_ptr = raw_device_ptr(dst, stream);
dtod_copy(dst_ptr, src_ptr, num_elements * std::mem::size_of::<f32>(), stream, 0, ctx)
}
@@ -2076,7 +2106,7 @@ fn dtod_from_candle_u32_to_i32(
stream: &Arc<CudaStream>,
ctx: &str,
) -> Result<(), MLError> {
let (storage_guard, _layout) = tensor.storage_and_layout();
let (storage_guard, layout) = tensor.storage_and_layout();
let cuda_storage = match &*storage_guard {
candle_core::Storage::Cuda(cs) => cs,
_ => {
@@ -2088,10 +2118,11 @@ fn dtod_from_candle_u32_to_i32(
let src_slice = cuda_storage.as_cuda_slice::<u32>().map_err(|e| {
MLError::ModelError(format!("dtod_from_candle_u32 {ctx}: as_cuda_slice: {e}"))
})?;
let offset = layout.start_offset();
let src_ptr = {
let (ptr, guard) = src_slice.device_ptr(stream);
let _no_drop = std::mem::ManuallyDrop::new(guard);
ptr
ptr + (offset * std::mem::size_of::<u32>()) as u64
};
let dst_ptr = raw_device_ptr_i32(dst, stream);
dtod_copy(dst_ptr, src_ptr, num_elements * std::mem::size_of::<u32>(), stream, 0, ctx)
@@ -2118,3 +2149,90 @@ fn alloc_i32(
MLError::ModelError(format!("GpuDqnTrainer alloc {name} ({n} i32): {e}"))
})
}
fn alloc_u16(
stream: &Arc<CudaStream>,
n: usize,
name: &str,
) -> Result<CudaSlice<u16>, MLError> {
stream.alloc_zeros::<u16>(n).map_err(|e| {
MLError::ModelError(format!("GpuDqnTrainer alloc {name} ({n} u16): {e}"))
})
}
/// DtoD copy from a BF16 Candle tensor to a pre-allocated CudaSlice<u16>.
///
/// Extracts the raw device pointer from the tensor's CudaStorage (BF16 = u16 bit pattern)
/// and issues an async DtoD memcpy. No Candle temp tensors are created — only raw pointer
/// extraction via `storage_and_layout()`.
fn dtod_from_candle_bf16(
tensor: &candle_core::Tensor,
dst: &CudaSlice<u16>,
num_elements: usize,
stream: &Arc<CudaStream>,
ctx: &str,
) -> Result<(), MLError> {
// If tensor is already F32, cast path not needed — extract as f32 and error.
// This function expects BF16 tensors (GpuReplayBuffer stores BF16).
let (storage_guard, layout) = tensor.storage_and_layout();
let cuda_storage = match &*storage_guard {
candle_core::Storage::Cuda(cs) => cs,
_ => {
return Err(MLError::ModelError(format!(
"dtod_from_candle_bf16 {ctx}: tensor must be on CUDA device"
)));
}
};
// BF16 tensors are stored as half::bf16 in Candle. The raw GPU storage
// is u16-compatible (same bit width). We use as_cuda_slice::<half::bf16>()
// to get the slice, then extract the raw device pointer for DtoD copy.
let src_slice = cuda_storage.as_cuda_slice::<half::bf16>().map_err(|e| {
MLError::ModelError(format!("dtod_from_candle_bf16 {ctx}: as_cuda_slice<bf16>: {e}"))
})?;
let offset = layout.start_offset();
let (ptr, guard) = src_slice.device_ptr(stream);
let _no_drop = std::mem::ManuallyDrop::new(guard);
// Offset the source pointer by the layout's start_offset (in elements)
let src_ptr = ptr + (offset * std::mem::size_of::<u16>()) as u64;
let (dst_ptr, dst_guard) = dst.device_ptr(stream);
let _no_drop2 = std::mem::ManuallyDrop::new(dst_guard);
dtod_copy(dst_ptr, src_ptr, num_elements * std::mem::size_of::<u16>(), stream, 0, ctx)
}
/// Launch the bf16_to_f32 CUDA kernel: converts BF16 CudaSlice<u16> → F32 CudaSlice<f32>.
///
/// Grid: ceil(n/256), Block: 256. Kernel signature: `bf16_to_f32_kernel(src, dst, n)`.
/// Zero Candle tensor involvement — pure cudarc kernel launch.
pub(crate) fn launch_bf16_to_f32(
kernel: &CudaFunction,
src: &CudaSlice<u16>,
dst: &CudaSlice<f32>,
n: usize,
stream: &Arc<CudaStream>,
ctx: &str,
) -> Result<(), MLError> {
let n_i32 = n as i32;
let launch_cfg = LaunchConfig {
grid_dim: (((n + 255) / 256) as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
// Safety: src is CudaSlice<u16> holding BF16 bit patterns, dst is CudaSlice<f32>,
// both are valid GPU allocations on the same context. Kernel reads src[0..n] and
// writes dst[0..n].
unsafe {
stream
.launch_builder(kernel)
.arg(src)
.arg(dst)
.arg(&n_i32)
.launch(launch_cfg)
.map_err(|e| MLError::ModelError(format!("bf16_to_f32_kernel {ctx}: {e}")))?;
}
Ok(())
}

140
crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs
View File

@@ -373,6 +373,146 @@ impl GpuIqlTrainer {
Ok(loss_host[0])
}
/// Raw CudaSlice variant of `train_value_step` — zero Candle tensor involvement.
///
/// Takes pre-converted F32 CudaSlice buffers directly, skipping the Candle
/// tensor `to_dtype(F32)` / `contiguous()` / `as_cuda_slice()` chain that
/// creates temporary Candle tensors whose Drops conflict with the forked stream.
pub fn train_value_step_raw(
&mut self,
states_f32: &CudaSlice<f32>,
q_values_f32: &CudaSlice<f32>,
) -> Result<f32, MLError> {
let b = self.config.batch_size;
// Zero total_loss and grad_norm before kernel launches
let zero_f32 = [0.0_f32];
self.stream.memcpy_htod(&zero_f32, &mut self.total_loss_buf)
.map_err(|e| MLError::ModelError(format!("IQL zero total_loss: {e}")))?;
self.stream.memcpy_htod(&zero_f32, &mut self.grad_norm_buf)
.map_err(|e| MLError::ModelError(format!("IQL zero grad_norm: {e}")))?;
let batch_size_i32 = b as i32;
let total_params_i32 = self.total_params as i32;
// 1. Forward + loss kernel
let fwd_config = LaunchConfig {
grid_dim: (b as u32, 1, 1),
block_dim: (32, 1, 1),
shared_mem_bytes: 0,
};
// Safety: states_f32 and q_values_f32 are valid F32 CudaSlice buffers
// on the same CUDA context. All internal buffers are pre-allocated.
unsafe {
self.stream
.launch_builder(&self.forward_loss_kernel)
.arg(states_f32)
.arg(q_values_f32)
.arg(&self.params_buf)
.arg(&mut self.v_out_buf)
.arg(&mut self.loss_buf)
.arg(&mut self.total_loss_buf)
.arg(&mut self.save_pre1)
.arg(&mut self.save_pre2)
.arg(&mut self.save_h1)
.arg(&mut self.save_h2)
.arg(&batch_size_i32)
.launch(fwd_config)
.map_err(|e| MLError::ModelError(format!("IQL forward+loss kernel: {e}")))?;
}
// 2. Backward kernel
let bwd_config = LaunchConfig {
grid_dim: (b as u32, 1, 1),
block_dim: (32, 1, 1),
shared_mem_bytes: 0,
};
// Safety: all pointers are valid GPU allocations on the same stream.
unsafe {
self.stream
.launch_builder(&self.backward_kernel)
.arg(states_f32)
.arg(q_values_f32)
.arg(&self.v_out_buf)
.arg(&self.params_buf)
.arg(&self.save_pre1)
.arg(&self.save_pre2)
.arg(&self.save_h1)
.arg(&self.save_h2)
.arg(&mut self.grad_buf)
.arg(&batch_size_i32)
.launch(bwd_config)
.map_err(|e| MLError::ModelError(format!("IQL backward kernel: {e}")))?;
}
// 3. Grad norm kernel
let norm_blocks = (self.total_params + 255) / 256;
let norm_config = LaunchConfig {
grid_dim: (norm_blocks as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
// Safety: grad_buf and grad_norm_buf are valid allocations.
unsafe {
self.stream
.launch_builder(&self.grad_norm_kernel)
.arg(&self.grad_buf)
.arg(&mut self.grad_norm_buf)
.arg(&total_params_i32)
.launch(norm_config)
.map_err(|e| MLError::ModelError(format!("IQL grad_norm kernel: {e}")))?;
}
// 4. Adam update kernel
self.adam_step += 1;
let adam_blocks = (self.total_params + 255) / 256;
let adam_config = LaunchConfig {
grid_dim: (adam_blocks as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
let lr = self.config.lr;
let beta1 = self.config.beta1;
let beta2 = self.config.beta2;
let eps = self.config.epsilon;
let wd = self.config.weight_decay;
let mgn = self.config.max_grad_norm;
let t = self.adam_step;
// Safety: all parameter and optimizer buffers are valid allocations.
unsafe {
self.stream
.launch_builder(&self.adam_kernel)
.arg(&mut self.params_buf)
.arg(&mut self.grad_buf)
.arg(&mut self.m_buf)
.arg(&mut self.v_buf)
.arg(&self.grad_norm_buf)
.arg(&lr)
.arg(&beta1)
.arg(&beta2)
.arg(&eps)
.arg(&wd)
.arg(&mgn)
.arg(&t)
.arg(&total_params_i32)
.launch(adam_config)
.map_err(|e| MLError::ModelError(format!("IQL adam kernel: {e}")))?;
}
// Read back total loss
let mut loss_host = [0.0_f32];
self.stream
.memcpy_dtoh(&self.total_loss_buf, &mut loss_host)
.map_err(|e| MLError::ModelError(format!("IQL loss readback: {e}")))?;
Ok(loss_host[0])
}
/// Compute advantages A(s,a) = Q(s,a) - V(s) for a batch.
///
/// Runs V(s) forward pass only (no loss, no backward) and subtracts

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

@@ -348,14 +348,14 @@ impl FusedTrainingCtx {
}
// ── Step 4: IQL value network (if enabled) ───────────────────────
// Uses GpuBatch tensors directly — GPU dtype cast, no CPU touch.
// Reuses the DQN trainer's F32 states_buf and rewards_buf that were
// already populated by train_step_gpu's bf16→f32 upload path.
// Zero Candle tensor temporaries — pure CudaSlice references.
if let Some(ref mut iql) = self.gpu_iql {
let states_tensor = effective_gpu.states.to_dtype(candle_core::DType::F32)
.map_err(|e| anyhow::anyhow!("IQL states F32 cast: {e}"))?;
let rewards_tensor = effective_gpu.rewards.to_dtype(candle_core::DType::F32)
.map_err(|e| anyhow::anyhow!("IQL rewards F32 cast: {e}"))?;
let states_f32 = self.trainer.states_buf();
let rewards_f32 = self.trainer.rewards_buf();
match iql.train_value_step(&states_tensor, &rewards_tensor) {
match iql.train_value_step_raw(states_f32, rewards_f32) {
Ok(value_loss) => {
tracing::debug!(
iql_value_loss = value_loss,