diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 1aca3461c..7339d02b8 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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, scalars_readback_host: [f32; 2], + + // ── BF16 batch staging buffers ────────────────────────────────── + // Pre-allocated CudaSlice 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, // [B * STATE_DIM] + bf16_next_states_buf: CudaSlice, // [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 { + &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 { + &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, 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, 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::().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::()) as u64; let dst_ptr = raw_device_ptr(dst, stream); dtod_copy(dst_ptr, src_ptr, num_elements * std::mem::size_of::(), stream, 0, ctx) } @@ -2076,7 +2106,7 @@ fn dtod_from_candle_u32_to_i32( stream: &Arc, 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::().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::()) as u64 }; let dst_ptr = raw_device_ptr_i32(dst, stream); dtod_copy(dst_ptr, src_ptr, num_elements * std::mem::size_of::(), stream, 0, ctx) @@ -2118,3 +2149,90 @@ fn alloc_i32( MLError::ModelError(format!("GpuDqnTrainer alloc {name} ({n} i32): {e}")) }) } + +fn alloc_u16( + stream: &Arc, + n: usize, + name: &str, +) -> Result, MLError> { + stream.alloc_zeros::(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. +/// +/// 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, + num_elements: usize, + stream: &Arc, + 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::() + // to get the slice, then extract the raw device pointer for DtoD copy. + let src_slice = cuda_storage.as_cuda_slice::().map_err(|e| { + MLError::ModelError(format!("dtod_from_candle_bf16 {ctx}: as_cuda_slice: {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::()) 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::(), stream, 0, ctx) +} + +/// Launch the bf16_to_f32 CUDA kernel: converts BF16 CudaSlice → F32 CudaSlice. +/// +/// 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, + dst: &CudaSlice, + n: usize, + stream: &Arc, + 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 holding BF16 bit patterns, dst is CudaSlice, + // 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(()) +} diff --git a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs index 187b67109..986c94206 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs @@ -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, + q_values_f32: &CudaSlice, + ) -> Result { + 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 diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 5cdf0ea3b..2409765bb 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -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,