diff --git a/crates/ml-dqn/src/dqn.rs b/crates/ml-dqn/src/dqn.rs index 7be396926..650440e00 100644 --- a/crates/ml-dqn/src/dqn.rs +++ b/crates/ml-dqn/src/dqn.rs @@ -792,7 +792,7 @@ pub struct GradientResult { /// GPU-resident TD errors (`GpuPrioritized` path). pub td_errors_gpu: Option, /// GPU-resident buffer indices (`GpuPrioritized` path). - pub indices_gpu: Option, + pub indices_gpu: Option>, /// Loss tensor on GPU for deferred batch readback. pub loss_tensor_gpu: Option, /// Gradient norm tensor on GPU for deferred readback. @@ -883,7 +883,7 @@ struct ComputeLossResult { /// GPU-resident TD errors (`GpuPrioritized` path). td_errors_gpu: Option, /// GPU-resident buffer indices (`GpuPrioritized` path). - indices_gpu: Option, + indices_gpu: Option>, } /// Experience replay buffer for `DQN` @@ -2469,13 +2469,11 @@ impl DQN { let is_gpu_per = self.memory.is_gpu_prioritized(); let td_error_tensor = GpuTensor::new(td_error_gpu, vec![batch_size])?; let (td_errors_vec, td_gpu, idx_gpu) = if is_gpu_per /* use_per: always on */ { - let idx_tensor = gpu_batch_opt.as_ref() - .map(|g| g.indices.gpu_clone(stream)) - .transpose()? + let idx_u32 = gpu_batch_opt.as_ref() .ok_or_else(|| MLError::TrainingError( "GPU PER indices required — CPU fallback removed".to_owned() - ))?; - (Vec::new(), Some(td_error_tensor), Some(idx_tensor)) + ))?.indices.clone(); + (Vec::new(), Some(td_error_tensor), Some(idx_u32)) } else { // Non-GPU PER: download td_errors to host for CPU priority update // ALLOWED: only reached when GPU PER is not active (legacy path) @@ -2592,7 +2590,7 @@ impl DQN { if let (Some(ref td_gpu), Some(ref idx_gpu)) = (&result.td_errors_gpu, &result.indices_gpu) { - self.memory.update_priorities_gpu(idx_gpu, td_gpu)?; + self.memory.update_priorities_gpu(idx_gpu, td_gpu.data(), result.td_errors.len().max(1))?; } else if !result.indices.is_empty() { self.memory .update_priorities(&result.indices, &result.td_errors)?; diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index 3a54f50bd..adf02f3ef 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -28,9 +28,9 @@ pub struct GpuBatchSlices { pub states: CudaSlice, // [batch_size * state_dim] bf16 on GPU pub next_states: CudaSlice, // [batch_size * state_dim] bf16 on GPU pub actions: CudaSlice, // [batch_size] u32 on GPU - pub rewards: CudaSlice, // [batch_size] f32 on GPU - pub dones: CudaSlice, // [batch_size] f32 on GPU (0.0/1.0) - pub weights: CudaSlice, // [batch_size] f32 on GPU (IS weights) + pub rewards: CudaSlice, // [batch_size] bf16 on GPU + pub dones: CudaSlice, // [batch_size] bf16 on GPU (0.0/1.0) + pub weights: CudaSlice, // [batch_size] bf16 on GPU (IS weights) pub indices: CudaSlice, // [batch_size] u32 on GPU (buffer indices) /// Episode IDs for sampled transitions `[batch_size]` i32 on GPU. /// Used by HER Future/Final strategies for episode-aware donor selection. @@ -44,8 +44,6 @@ impl GpuBatchSlices { /// /// Zero CPU downloads -- all type conversions happen on GPU via custom CUDA kernels. pub fn into_gpu_batch(self, stream: &Arc) -> Result { - use ml_core::cuda_autograd::GpuTensor; - let bs = self.batch_size; let sd = self.state_dim; @@ -55,14 +53,25 @@ impl GpuBatchSlices { let states = bf16_slice_to_gpu_tensor_gpu(&self.states, vec![bs, sd], stream, kernels)?; let next_states = bf16_slice_to_gpu_tensor_gpu(&self.next_states, vec![bs, sd], stream, kernels)?; - // u32 actions/indices -> f32 via GPU cast kernel (zero host roundtrip) + // u32 actions -> GpuTensor via GPU cast kernel (zero host roundtrip) let actions = u32_slice_to_gpu_tensor_gpu(&self.actions, vec![bs], stream, kernels)?; - let indices = u32_slice_to_gpu_tensor_gpu(&self.indices, vec![bs], stream, kernels)?; + // u32 indices: DtoD clone (stays u32, no bf16 conversion) + let indices = { + let mut dst = stream.alloc_zeros::(bs) + .map_err(|e| MLError::ModelError(format!("indices clone alloc: {e}")))?; + let nbytes = bs * std::mem::size_of::(); + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst.raw_ptr(), self.indices.raw_ptr(), nbytes, stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("indices DtoD: {e}")))?; + } + dst + }; - // f32 slices -> bf16 GpuTensor via GPU cast kernel - let rewards = f32_slice_to_gpu_tensor_gpu(&self.rewards, vec![bs], stream, kernels)?; - let dones = f32_slice_to_gpu_tensor_gpu(&self.dones, vec![bs], stream, kernels)?; - let weights = f32_slice_to_gpu_tensor_gpu(&self.weights, vec![bs], stream, kernels)?; + // bf16 slices -> bf16 GpuTensor via DtoD reinterpret (zero cast needed, already bf16) + let rewards = bf16_slice_to_gpu_tensor_gpu(&self.rewards, vec![bs], stream, kernels)?; + let dones = bf16_slice_to_gpu_tensor_gpu(&self.dones, vec![bs], stream, kernels)?; + let weights = bf16_slice_to_gpu_tensor_gpu(&self.weights, vec![bs], stream, kernels)?; Ok(crate::replay_buffer_type::GpuBatch { states, @@ -252,6 +261,7 @@ fn u32_slice_to_gpu_tensor_gpu( } /// Convert a f32 `CudaSlice` to bf16 `GpuTensor` via GPU cast kernel (zero CPU download). +#[allow(dead_code)] fn f32_slice_to_gpu_tensor_gpu( src: &CudaSlice, shape: Vec, @@ -303,9 +313,11 @@ struct ReplayKernels { scatter_insert_f32: CudaFunction, scatter_insert_u32: CudaFunction, scatter_insert_bf16: CudaFunction, + #[allow(dead_code)] gather_f32: CudaFunction, gather_u32: CudaFunction, gather_bf16_rows: CudaFunction, + gather_bf16: CudaFunction, is_weights_f32: CudaFunction, normalize_weights_f32: CudaFunction, fill_from_gpu_f32: CudaFunction, @@ -343,6 +355,7 @@ impl ReplayKernels { scatter_insert_bf16: ld("scatter_insert_bf16")?, gather_f32: ld("gather_f32")?, gather_u32: ld("gather_u32")?, gather_bf16_rows: ld("gather_bf16_rows")?, + gather_bf16: ld("gather_bf16")?, is_weights_f32: ld("is_weights_f32")?, normalize_weights_f32: ld("normalize_weights_f32")?, fill_from_gpu_f32: ld("fill_from_gpu_f32")?, @@ -386,8 +399,8 @@ pub struct GpuReplayBuffer { stream: Arc, kernels: ReplayKernels, states: CudaSlice, next_states: CudaSlice, - actions: CudaSlice, rewards: CudaSlice, - dones: CudaSlice, priorities: CudaSlice, + actions: CudaSlice, rewards: CudaSlice, + dones: CudaSlice, priorities: CudaSlice, /// Episode IDs per buffer slot `[capacity]` i32 on GPU. /// `episode_ids[i] = i / episode_length`. Written during `insert_batch_with_episode_ids`. episode_ids: CudaSlice, @@ -404,8 +417,8 @@ pub struct GpuReplayBuffer { sample_states: CudaSlice, sample_next_states: CudaSlice, sample_actions: CudaSlice, - sample_rewards: CudaSlice, - sample_dones: CudaSlice, + sample_rewards: CudaSlice, + sample_dones: CudaSlice, sample_priorities: CudaSlice, sample_weights: CudaSlice, sample_max_weight: CudaSlice, @@ -442,8 +455,8 @@ impl GpuReplayBuffer { let s = a16(stream, cap * sd, "s")?; let ns = a16(stream, cap * sd, "ns")?; let a = a32u(stream, cap, "a")?; - let r = a32f(stream, cap, "r")?; - let d = a32f(stream, cap, "d")?; + let r = a16(stream, cap, "r")?; + let d = a16(stream, cap, "d")?; let p = a32f(stream, cap, "p")?; let mut mp = a32f(stream, 1, "mp")?; stream.memcpy_htod(&[1.0_f32], &mut mp).map_err(|e| MLError::ModelError(format!("mp: {e}")))?; @@ -459,8 +472,8 @@ impl GpuReplayBuffer { let ss = a16(stream, mbs * sd, "s_states")?; let sns = a16(stream, mbs * sd, "s_nstates")?; let sa = a32u(stream, mbs, "s_act")?; - let sr = a32f(stream, mbs, "s_rew")?; - let sdn = a32f(stream, mbs, "s_done")?; + let sr = a16(stream, mbs, "s_rew")?; + let sdn = a16(stream, mbs, "s_done")?; let sp = a32f(stream, mbs, "s_pri")?; let sw = a32f(stream, mbs, "s_wt")?; let smw = a32f(stream, 1, "s_mw")?; @@ -553,16 +566,33 @@ impl GpuReplayBuffer { .arg(&self.actions).arg(&sa).arg(&ci).arg(&cpi).arg(&bsi) .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc a: {e}")))?; } - // SAFETY: same context, rewards buffer valid. + // Cast f32 rewards/dones to bf16 (u16) then scatter insert as bf16 (state_dim=1) + let one_i = 1_i32; + let eff_i = eff as i32; + let mut b_r = a16(&self.stream, eff, "ib_r")?; + // SAFETY: b_r, sr are valid device allocations of at least eff elements. unsafe { - self.stream.launch_builder(&self.kernels.scatter_insert_f32) - .arg(&self.rewards).arg(&sr).arg(&ci).arg(&cpi).arg(&bsi) + self.stream.launch_builder(&self.kernels.f32_to_bf16_cast) + .arg(&mut b_r).arg(&sr).arg(&eff_i).launch(lcfg(eff)) + .map_err(|e| MLError::ModelError(format!("cast r: {e}")))?; + } + // SAFETY: rewards buffer (bf16) valid, b_r cast above. + unsafe { + self.stream.launch_builder(&self.kernels.scatter_insert_bf16) + .arg(&self.rewards).arg(&b_r).arg(&ci).arg(&cpi).arg(&one_i).arg(&bsi) .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc r: {e}")))?; } - // SAFETY: same context, dones buffer valid. + let mut b_d = a16(&self.stream, eff, "ib_d")?; + // SAFETY: b_d, sd2 are valid device allocations of at least eff elements. unsafe { - self.stream.launch_builder(&self.kernels.scatter_insert_f32) - .arg(&self.dones).arg(&sd2).arg(&ci).arg(&cpi).arg(&bsi) + self.stream.launch_builder(&self.kernels.f32_to_bf16_cast) + .arg(&mut b_d).arg(&sd2).arg(&eff_i).launch(lcfg(eff)) + .map_err(|e| MLError::ModelError(format!("cast d: {e}")))?; + } + // SAFETY: dones buffer (bf16) valid, b_d cast above. + unsafe { + self.stream.launch_builder(&self.kernels.scatter_insert_bf16) + .arg(&self.dones).arg(&b_d).arg(&ci).arg(&cpi).arg(&one_i).arg(&bsi) .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc d: {e}")))?; } let mut pt = a32f(&self.stream, eff, "pt")?; @@ -660,41 +690,21 @@ impl GpuReplayBuffer { .arg(&self.actions).arg(&sa).arg(&ci).arg(&cpi).arg(&bsi) .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc a: {e}")))?; } - // Rewards/dones: reinterpret bf16 CudaSlice as f32 for scatter kernel (same 2-byte elements) - // Actually rewards/dones in the replay buffer are stored as f32. - // We need to cast bf16→f32. Use a temporary f32 buffer + bf16_to_f32 cast. - // For simplicity, use the scatter_insert_bf16 kernel which copies bf16→bf16, - // then let the sample path handle the conversion. - // BUT the replay buffer stores rewards as CudaSlice, so we need to convert. - // Simplest: reinterpret the bf16 as raw bytes and use memcpy_dtod with size matching. - // Since rewards/dones in experience batch are bf16 but replay buffer stores f32, - // we need to cast. Use the bf16_to_f32_cast kernel if available. - // For now, download→convert→upload (small: only batch_size scalars). + // Rewards/dones: both input and storage are bf16 — scatter directly (state_dim=1) { let sr = if off > 0 { rw.slice(off..) } else { rw.slice(0..) }; let sd2 = if off > 0 { dn.slice(off..) } else { dn.slice(0..) }; - // Small host roundtrip for rewards/dones (eff scalars, typically 128-1024) - let mut rw_host = vec![half::bf16::ZERO; eff]; - self.stream.memcpy_dtoh(&sr, &mut rw_host) - .map_err(|e| MLError::ModelError(format!("rw dtoh: {e}")))?; - let rw_f32: Vec = rw_host.iter().map(|x| x.to_f32()).collect(); - let rw_gpu = self.stream.clone_htod(&rw_f32) - .map_err(|e| MLError::ModelError(format!("rw htod: {e}")))?; + let one_i = 1_i32; + // SAFETY: rewards/dones buffers are u16 (bf16), rw/dn CudaSlice same repr. unsafe { - self.stream.launch_builder(&self.kernels.scatter_insert_f32) - .arg(&self.rewards).arg(&rw_gpu).arg(&ci).arg(&cpi).arg(&bsi) - .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc r: {e}")))?; + self.stream.launch_builder(&self.kernels.scatter_insert_bf16) + .arg(&self.rewards).arg(&sr).arg(&ci).arg(&cpi).arg(&one_i).arg(&bsi) + .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc r bf16: {e}")))?; } - let mut dn_host = vec![half::bf16::ZERO; eff]; - self.stream.memcpy_dtoh(&sd2, &mut dn_host) - .map_err(|e| MLError::ModelError(format!("dn dtoh: {e}")))?; - let dn_f32: Vec = dn_host.iter().map(|x| x.to_f32()).collect(); - let dn_gpu = self.stream.clone_htod(&dn_f32) - .map_err(|e| MLError::ModelError(format!("dn htod: {e}")))?; unsafe { - self.stream.launch_builder(&self.kernels.scatter_insert_f32) - .arg(&self.dones).arg(&dn_gpu).arg(&ci).arg(&cpi).arg(&bsi) - .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc d: {e}")))?; + self.stream.launch_builder(&self.kernels.scatter_insert_bf16) + .arg(&self.dones).arg(&sd2).arg(&ci).arg(&cpi).arg(&one_i).arg(&bsi) + .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc d bf16: {e}")))?; } } // Priorities — same as insert_batch @@ -806,13 +816,13 @@ impl GpuReplayBuffer { } // SAFETY: same context, rewards buffer valid. unsafe { - self.stream.launch_builder(&self.kernels.gather_f32) + self.stream.launch_builder(&self.kernels.gather_bf16) .arg(&mut self.sample_rewards).arg(&self.rewards).arg(&self.sample_indices_i64).arg(&bsi) .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g r: {e}")))?; } // SAFETY: same context, dones buffer valid. unsafe { - self.stream.launch_builder(&self.kernels.gather_f32) + self.stream.launch_builder(&self.kernels.gather_bf16) .arg(&mut self.sample_dones).arg(&self.dones).arg(&self.sample_indices_i64).arg(&bsi) .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g d: {e}")))?; } @@ -896,13 +906,28 @@ impl GpuReplayBuffer { // so we DtoD-clone the relevant portions. All copies are async on the stream. let ep_ids = dtod_clone_i32(&self.stream, &self.sample_episode_ids, batch_size, "o_ep")?; + // Cast f32 IS weights -> bf16 (u16) for output. IS computation stays f32 for precision. + let weights_bf16 = { + let wt_slice = self.sample_weights.slice(..batch_size); + let mut out = a16(&self.stream, batch_size, "o_w_bf16")?; + let n_i = batch_size as i32; + // SAFETY: out, wt_slice are valid device allocations of at least batch_size elements. + unsafe { + self.stream.launch_builder(&self.kernels.f32_to_bf16_cast) + .arg(&mut out).arg(&wt_slice).arg(&n_i) + .launch(lcfg(batch_size)) + .map_err(|e| MLError::ModelError(format!("cast w bf16: {e}")))?; + } + out + }; + Ok(GpuBatchSlices { states: dtod_clone_u16(&self.stream, &self.sample_states, batch_size * sd, "o_s")?, next_states: dtod_clone_u16(&self.stream, &self.sample_next_states, batch_size * sd, "o_n")?, actions: dtod_clone_u32(&self.stream, &self.sample_actions, batch_size, "o_act")?, - rewards: dtod_clone_f32(&self.stream, &self.sample_rewards, batch_size, "o_r")?, - dones: dtod_clone_f32(&self.stream, &self.sample_dones, batch_size, "o_d")?, - weights: dtod_clone_f32(&self.stream, &self.sample_weights, batch_size, "o_w")?, + rewards: dtod_clone_u16(&self.stream, &self.sample_rewards, batch_size, "o_r")?, + dones: dtod_clone_u16(&self.stream, &self.sample_dones, batch_size, "o_d")?, + weights: weights_bf16, indices: dtod_clone_u32(&self.stream, &self.sample_indices_u32, batch_size, "o_i")?, episode_ids: Some(ep_ids), batch_size, @@ -1042,15 +1067,15 @@ impl GpuReplayBuffer { pub const fn states_slice(&self) -> &CudaSlice { &self.states } pub const fn next_states_slice(&self) -> &CudaSlice { &self.next_states } pub const fn actions_slice(&self) -> &CudaSlice { &self.actions } - pub const fn rewards_slice(&self) -> &CudaSlice { &self.rewards } - pub const fn dones_slice(&self) -> &CudaSlice { &self.dones } + pub const fn rewards_slice(&self) -> &CudaSlice { &self.rewards } + pub const fn dones_slice(&self) -> &CudaSlice { &self.dones } pub const fn priorities_slice(&self) -> &CudaSlice { &self.priorities } /// Sample proportional indices and IS weights as GPU-resident `CudaSlices`. /// - /// Returns `(indices: CudaSlice, weights: CudaSlice)` on GPU. + /// Returns `(indices: CudaSlice, weights: CudaSlice)` on GPU (weights are bf16). /// Callers process indices on GPU -- zero CPU download. - pub fn sample_indices_gpu(&mut self, bs: usize) -> Result<(CudaSlice, CudaSlice), MLError> { + pub fn sample_indices_gpu(&mut self, bs: usize) -> Result<(CudaSlice, CudaSlice), MLError> { let b = self.sample_proportional(bs)?; Ok((b.indices, b.weights)) } @@ -1080,6 +1105,7 @@ fn a32i(s: &Arc, n: usize, nm: &str) -> Result, MLErr } /// DtoD clone of first `n` elements from `src` into a new allocation (async on stream). +#[allow(dead_code)] fn dtod_clone_f32(s: &Arc, src: &CudaSlice, n: usize, nm: &str) -> Result, MLError> { let mut dst = a32f(s, n, nm)?; let sv = src.slice(..n); diff --git a/crates/ml-dqn/src/regime_conditional.rs b/crates/ml-dqn/src/regime_conditional.rs index 068be7f6b..2b1f8a088 100644 --- a/crates/ml-dqn/src/regime_conditional.rs +++ b/crates/ml-dqn/src/regime_conditional.rs @@ -650,7 +650,7 @@ impl RegimeConditionalDQN { next_states: gpu_batch.next_states.gpu_clone(&self.stream)?, dones: gpu_batch.dones.gpu_clone(&self.stream)?, weights: masked_weights, - indices: gpu_batch.indices.gpu_clone(&self.stream)?, + indices: gpu_batch.indices.clone(), episode_ids: None, }; @@ -853,7 +853,7 @@ impl RegimeConditionalDQN { next_states: gpu_batch.next_states.gpu_clone(&self.stream)?, dones: gpu_batch.dones.gpu_clone(&self.stream)?, weights: masked_weights, - indices: gpu_batch.indices.gpu_clone(&self.stream)?, + indices: gpu_batch.indices.clone(), episode_ids: None, }; diff --git a/crates/ml-dqn/src/replay_buffer_kernels.cu b/crates/ml-dqn/src/replay_buffer_kernels.cu index c6ef257d8..21d2fcb81 100644 --- a/crates/ml-dqn/src/replay_buffer_kernels.cu +++ b/crates/ml-dqn/src/replay_buffer_kernels.cu @@ -114,6 +114,21 @@ void gather_bf16_rows( out[row * state_dim + col] = src[src_row * state_dim + col]; } +// ── 6b. Gather bf16 scalars (1D) by indices ──────────────────────────────── + +extern "C" __global__ +void gather_bf16( + unsigned short* __restrict__ out, // [batch_size] + const unsigned short* __restrict__ src, // [capacity] + const long long* __restrict__ indices, // [batch_size] + int batch_size +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= batch_size) return; + int idx = (int)indices[i]; + out[i] = src[idx]; +} + // ── 7. prios_alpha: out[i] = pow(src[i], alpha) ──────────────────────────── extern "C" __global__ diff --git a/crates/ml-dqn/src/replay_buffer_type.rs b/crates/ml-dqn/src/replay_buffer_type.rs index 1fcdd20e0..ec396a52d 100644 --- a/crates/ml-dqn/src/replay_buffer_type.rs +++ b/crates/ml-dqn/src/replay_buffer_type.rs @@ -43,13 +43,13 @@ impl BatchSample { /// When present, `compute_gradients()` uses these directly -- no CPU->GPU transfer. #[derive(Debug)] pub struct GpuBatch { - pub states: GpuTensor, // [batch_size, state_dim] f32 on GPU + pub states: GpuTensor, // [batch_size, state_dim] bf16 on GPU pub actions: GpuTensor, // [batch_size] u32 on GPU - pub rewards: GpuTensor, // [batch_size] f32 on GPU - pub next_states: GpuTensor, // [batch_size, state_dim] f32 on GPU - pub dones: GpuTensor, // [batch_size] f32 on GPU (0.0/1.0) - pub weights: GpuTensor, // [batch_size] f32 on GPU (IS weights) - pub indices: GpuTensor, // [batch_size] u32 on GPU (buffer indices) + pub rewards: GpuTensor, // [batch_size] bf16 on GPU + pub next_states: GpuTensor, // [batch_size, state_dim] bf16 on GPU + pub dones: GpuTensor, // [batch_size] bf16 on GPU (0.0/1.0) + pub weights: GpuTensor, // [batch_size] bf16 on GPU (IS weights) + pub indices: cudarc::driver::CudaSlice, // [batch_size] u32 on GPU (buffer indices) /// Episode IDs per transition `[batch_size]` i32 on GPU. /// /// `episode_ids[i] = buffer_index[i] / episode_length`. Required by HER @@ -354,19 +354,16 @@ impl ReplayBufferType { /// Indices are expected to hold f32-encoded u32 values (cast to u32 internally). pub fn update_priorities_gpu( &self, - indices: &GpuTensor, - td_errors: &GpuTensor, + indices: &cudarc::driver::CudaSlice, + td_errors: &cudarc::driver::CudaSlice, + bs: usize, ) -> Result<(), MLError> { match self { Self::GpuPrioritized(buffer) => { let mut buf = buffer.lock(); - let bs = td_errors.shape().first().copied().unwrap_or(0); if bs == 0 { return Ok(()); } - // Convert f32-encoded indices to CudaSlice via GPU cast kernel (zero CPU download) - let stream = buf.gpu.stream().clone(); - let cast_kernels = crate::gpu_replay_buffer::get_cast_kernels_f32_to_u32(&stream)?; - let idx_gpu = cast_kernels.cast_f32_to_u32(indices.data(), bs, &stream)?; - buf.gpu.update_priorities_gpu(&idx_gpu, td_errors.data(), bs) + // indices already CudaSlice, td_errors already CudaSlice + buf.gpu.update_priorities_gpu(indices, td_errors, bs) } Self::Uniform(_) | Self::Prioritized(_) => Ok(()), // No-op for non-GPU buffers } diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index 39575d9ff..86a9b1d22 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -26,40 +26,52 @@ extern "C" __global__ void dqn_grad_norm_kernel( const __nv_bfloat16* __restrict__ grads, /* [TOTAL_PARAMS] accumulated gradients */ - __nv_bfloat16* __restrict__ out_grad_norm, /* [1] output: sum of squares */ + float* __restrict__ out_norm_f32, /* [1] float accumulator (atomicAdd target) */ int total_params ) { int idx = blockIdx.x * blockDim.x + threadIdx.x; - __nv_bfloat16 g = (idx < total_params) ? grads[idx] : bf16_zero(); - __nv_bfloat16 g2 = g * g; + /* Accumulate g² in float — bf16 overflows at sum > 65504 (trivial with 147K params). + * BF16 arithmetic produces occasional NaN/Inf in individual gradient elements + * (e.g. bf16 subtraction of nearly-equal large values → NaN). These don't affect + * training correctness (Adam's clip_scale uses fmaxf which handles NaN→0), but + * NaN + valid = NaN poisons the norm sum. Skip them for an accurate norm. */ + float g_f = (idx < total_params) ? (float)grads[idx] : 0.0f; + if (isnan(g_f) || isinf(g_f)) g_f = 0.0f; + float g2 = g_f * g_f; - /* Warp-level reduction via shuffle (no shared memory) */ + /* Warp-level reduction via shuffle (float) */ for (int offset = 16; offset > 0; offset >>= 1) - g2 = g2 + bf16_shfl_xor(0xFFFFFFFF, g2, offset); + g2 += __shfl_xor_sync(0xFFFFFFFF, g2, offset); - /* Block-level cross-warp reduction via shared memory. - * Pad warp_sums to 2 elements per warp (stride=2) to avoid bank conflicts: - * warp 0 → bank 0, warp 1 → bank 2, ... warp 7 → bank 14. - * Without padding, sequential warp writes to warp_sums[0..7] map to - * banks 0..7 which is conflict-free for writes, but the __syncthreads() - * + read-back pattern benefits from padding to avoid false sharing - * between adjacent cache lines on H100's 128-byte L1 lines. */ - __shared__ __nv_bfloat16 warp_sums[16]; /* 8 warps × 2 stride (padded) */ + /* Block-level cross-warp reduction via shared memory (float) */ + __shared__ float warp_sums[16]; /* 8 warps × 2 stride (padded) */ int warp_id = threadIdx.x / 32; int warp_lane = threadIdx.x % 32; if (warp_lane == 0) warp_sums[warp_id * 2] = g2; __syncthreads(); - /* First warp reduces across warps using shuffles (pure register path) */ + /* First warp reduces across warps using shuffles */ if (warp_id == 0) { - __nv_bfloat16 val = (warp_lane < blockDim.x / 32) ? warp_sums[warp_lane * 2] : bf16_zero(); + float val = (warp_lane < blockDim.x / 32) ? warp_sums[warp_lane * 2] : 0.0f; for (int offset = 16; offset > 0; offset >>= 1) - val = val + bf16_shfl_xor(0xFFFFFFFF, val, offset); + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); if (warp_lane == 0) - atomicAddBF16(out_grad_norm, val); + atomicAdd(out_norm_f32, val); /* Native float atomicAdd — no type-punning */ } } +/* Single-thread finalize: sqrt(f32 sum-of-squares) → bf16 L2 norm. + * Called AFTER graph_adam replay (outside CUDA graph). + * Reads from the float accumulator, writes bf16 norm for the training guard. */ +extern "C" __global__ void dqn_grad_norm_finalize( + const float* __restrict__ norm_f32, /* [1] float sum-of-squares */ + __nv_bfloat16* __restrict__ out_norm_bf16 /* [1] bf16 L2 norm output */ +) { + float sum_sq = *norm_f32; + float norm = sqrtf(fmaxf(sum_sq, 0.0f)); + out_norm_bf16[0] = bf16(norm); +} + /* ══════════════════════════════════════════════════════════════════════ * ADAM UPDATE KERNEL (Phase 3b) * @@ -75,7 +87,7 @@ extern "C" __global__ void dqn_adam_update_kernel( const __nv_bfloat16* __restrict__ grads, /* [TOTAL_PARAMS] accumulated gradients */ __nv_bfloat16* __restrict__ m, /* [TOTAL_PARAMS] first moment (Adam) */ __nv_bfloat16* __restrict__ v, /* [TOTAL_PARAMS] second moment (Adam) */ - const __nv_bfloat16* __restrict__ grad_norm_sq, /* [1] completed sum of squares */ + const float* __restrict__ grad_norm_f32, /* [1] float sum-of-squares from grad_norm kernel */ float lr, float beta1, float beta2, @@ -91,9 +103,11 @@ extern "C" __global__ void dqn_adam_update_kernel( int t = *t_ptr; __nv_bfloat16 g = grads[idx]; - /* Clip using the COMPLETED norm (no race — separate kernel launch) */ - __nv_bfloat16 norm = bf16_sqrt(*grad_norm_sq + bf16(1e-12f)); - __nv_bfloat16 clip_scale = (*grad_norm_sq > bf16_zero() && norm > bf16(max_grad_norm)) ? (bf16(max_grad_norm) / norm) : bf16_one(); + /* Clip using the COMPLETED sum-of-squares (native float buffer). + * dqn_grad_norm_kernel accumulates float partial sums via atomicAdd. */ + float norm_sq_f = *grad_norm_f32; + float norm_f = sqrtf(fmaxf(norm_sq_f, 0.0f)); + __nv_bfloat16 clip_scale = (norm_f > max_grad_norm) ? bf16(max_grad_norm / norm_f) : bf16_one(); __nv_bfloat16 clipped_g = g * clip_scale; /* Adam bias correction: compute in float — bf16(0.999) rounds to 1.0 → div-by-zero. @@ -160,15 +174,15 @@ extern "C" __global__ void dqn_clipped_saxpy_kernel( const __nv_bfloat16* __restrict__ x, float alpha, float max_norm, - const __nv_bfloat16* __restrict__ x_norm_sq, + const float* __restrict__ x_norm_f32, /* float sum-of-squares from grad_norm kernel */ int n ) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= n) return; - __nv_bfloat16 norm = bf16_sqrt(*x_norm_sq + bf16(1e-12f)); - __nv_bfloat16 clip = (norm > bf16(max_norm)) ? (bf16(max_norm) / norm) : bf16_one(); - y[i] = y[i] + bf16(alpha) * clip * x[i]; + float norm_f = sqrtf(fmaxf(*x_norm_f32, 0.0f)); + float clip_f = (norm_f > max_norm) ? (max_norm / norm_f) : 1.0f; + y[i] = y[i] + bf16(alpha * clip_f) * x[i]; } /* ══════════════════════════════════════════════════════════════════════ @@ -185,16 +199,16 @@ extern "C" __global__ void dqn_clipped_saxpy_kernel( extern "C" __global__ void dqn_clip_grad_kernel( __nv_bfloat16* __restrict__ grads, - const __nv_bfloat16* __restrict__ grad_norm_sq, + const __nv_bfloat16* __restrict__ grad_norm_ptr, /* bf16 L2 norm from finalize kernel */ float max_norm, int total_params ) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= total_params) return; - __nv_bfloat16 norm = bf16_sqrt(*grad_norm_sq + bf16(1e-12f)); - if (norm > bf16(max_norm)) { - grads[idx] = grads[idx] * (bf16(max_norm) / norm); + float norm_f = (float)(*grad_norm_ptr); /* bf16 L2 norm, already sqrt'd */ + if (norm_f > max_norm) { + grads[idx] = grads[idx] * bf16(max_norm / norm_f); } } diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 94bad02de..7f1bb029a 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -314,6 +314,7 @@ struct CachedPtrs { target_params_buf: u64, grad_buf: u64, grad_norm_buf: u64, + grad_norm_f32_buf: u64, m_buf: u64, v_buf: u64, t_buf: u64, @@ -380,6 +381,7 @@ pub struct GpuDqnTrainer { // Dead kernels (forward_loss_kernel, forward_only_kernel, backward_kernel) // removed — replaced by cuBLAS forward/backward. See Task 2 for method cleanup. grad_norm_kernel: CudaFunction, + grad_norm_finalize_kernel: CudaFunction, /// Separate instance of grad_norm for non-graph launches (clip_grad_buf_inplace). /// CUDA doesn't allow the same CUfunction to be launched both inside a captured /// graph and outside on the same stream. @@ -446,7 +448,7 @@ pub struct GpuDqnTrainer { // ── Forward output buffers ────────────────────────────────────── per_sample_loss_buf: CudaSlice, // [B] td_errors_buf: CudaSlice, // [B] - total_loss_buf: CudaSlice, // [1] + pub(crate) total_loss_buf: CudaSlice, // [1] // ── Forward-only Q-value output ───────────────────────────────── q_out_buf: CudaSlice, // [B, TOTAL_ACTIONS(11)] @@ -457,7 +459,8 @@ pub struct GpuDqnTrainer { target_params_buf: CudaSlice, // [TOTAL_PARAMS] flat target parameters (EMA) m_buf: CudaSlice, // [TOTAL_PARAMS] Adam first moment v_buf: CudaSlice, // [TOTAL_PARAMS] Adam second moment - grad_norm_buf: CudaSlice, // [1] pre-clip gradient L2 norm + pub(crate) grad_norm_buf: CudaSlice, // [1] bf16 L2 norm (written by finalize) + grad_norm_f32_buf: CudaSlice, // [1] float accumulator for grad_norm kernel cql_grad_scratch: CudaSlice, // [TOTAL_PARAMS] CQL gradient isolation buffer // ── Adam step counter on device (CUDA Graph cannot bake scalars) ─ @@ -1330,21 +1333,21 @@ impl GpuDqnTrainer { /// Called after `apply_cql_gradient` populated `cql_grad_scratch`. /// Computes norm of scratch, clips to `cql_budget`, then SAXPYs into grad_buf. pub fn apply_cql_clipped_saxpy(&mut self, cql_budget: f32) -> Result<(), MLError> { - // Compute CQL gradient norm + // Compute CQL gradient norm (float accumulator) let _eg = EventTrackingGuard::new(self.stream.context()); - self.stream.memset_zeros(&mut self.grad_norm_buf) - .map_err(|e| MLError::ModelError(format!("zero cql_grad_norm: {e}")))?; + self.stream.memset_zeros(&mut self.grad_norm_f32_buf) + .map_err(|e| MLError::ModelError(format!("zero cql_grad_norm_f32: {e}")))?; let total = self.total_params as i32; let blocks = ((self.total_params + 255) / 256) as u32; let scratch_ptr = self.ptrs.cql_grad_scratch; - let norm_ptr = self.ptrs.grad_norm_buf; + let norm_f32_ptr = self.ptrs.grad_norm_f32_buf; unsafe { self.stream .launch_builder(&self.grad_norm_standalone) .arg(&scratch_ptr) - .arg(&norm_ptr) + .arg(&norm_f32_ptr) .arg(&total) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), @@ -1364,7 +1367,7 @@ impl GpuDqnTrainer { .arg(&scratch_ptr) .arg(&alpha) .arg(&cql_budget) - .arg(&norm_ptr) + .arg(&norm_f32_ptr) .arg(&total) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), @@ -1543,14 +1546,14 @@ impl GpuDqnTrainer { // Disable cudarc event tracking — graph replay leaves stale events // that cause INVALID_VALUE on post-graph kernel launches. let _evt_guard = EventTrackingGuard::new(self.stream.context()); - // Zero the grad_norm accumulator + // Zero the float grad_norm accumulator self.stream - .memset_zeros(&mut self.grad_norm_buf) - .map_err(|e| MLError::ModelError(format!("zero grad_norm for clip: {e}")))?; - // memset_zeros succeeded + .memset_zeros(&mut self.grad_norm_f32_buf) + .map_err(|e| MLError::ModelError(format!("zero grad_norm_f32 for clip: {e}")))?; - // Compute grad_norm (sum of squares) + // Compute grad_norm (float sum-of-squares → bf16 L2 norm) self.launch_grad_norm()?; + self.launch_grad_norm_finalize()?; // Clip in-place let grad_ptr = self.ptrs.grad_buf; @@ -1583,22 +1586,24 @@ impl GpuDqnTrainer { let _evt_guard = EventTrackingGuard::new(self.stream.context()); self.stream - .memset_zeros(&mut self.grad_norm_buf) - .map_err(|e| MLError::ModelError(format!("zero grad_norm for diag: {e}")))?; + .memset_zeros(&mut self.grad_norm_f32_buf) + .map_err(|e| MLError::ModelError(format!("zero grad_norm_f32 for diag: {e}")))?; self.launch_grad_norm()?; + self.launch_grad_norm_finalize()?; unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } - let mut norm_sq_bf16 = [half::bf16::ZERO; 1]; + // grad_norm_buf[0] now contains the L2 norm directly (sqrt applied by finalize) + let mut norm_bf16 = [half::bf16::ZERO; 1]; unsafe { cudarc::driver::sys::cuMemcpyDtoH_v2( - norm_sq_bf16.as_mut_ptr().cast(), + norm_bf16.as_mut_ptr().cast(), self.grad_norm_buf.raw_ptr(), std::mem::size_of::(), ); } - Ok(norm_sq_bf16[0].to_f32().sqrt()) + Ok(norm_bf16[0].to_f32()) } /// Apply GPU multi-head feature attention to `save_h_s2` (post-graph). @@ -1717,7 +1722,7 @@ impl GpuDqnTrainer { // per array. Stack is set once in DQNTrainer::new() (64KB for all kernels). // ── Compile 4 utility kernels (grad_norm, adam_update, BF16 converters) ─ - let (grad_norm_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clipped_saxpy_kernel, clip_grad_kernel) = + let (grad_norm_kernel, grad_norm_finalize_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clipped_saxpy_kernel, clip_grad_kernel) = compile_training_kernels(&stream, &config)?; // Separate grad_norm instance for non-graph launches (clip_grad_buf_inplace). @@ -1778,6 +1783,8 @@ impl GpuDqnTrainer { let m_buf = alloc_bf16(&stream, total_params, "adam_m")?; let v_buf = alloc_bf16(&stream, total_params, "adam_v")?; let grad_norm_buf = alloc_bf16(&stream, 1, "grad_norm")?; + let grad_norm_f32_buf = stream.alloc_zeros::(1) + .map_err(|e| MLError::ModelError(format!("alloc grad_norm_f32: {e}")))?; let cql_grad_scratch = alloc_bf16(&stream, total_params, "cql_grad_scratch")?; let t_buf = alloc_i32(&stream, 1, "adam_t")?; @@ -2007,6 +2014,7 @@ impl GpuDqnTrainer { target_params_buf: target_params_buf.raw_ptr(), grad_buf: grad_buf.raw_ptr(), grad_norm_buf: grad_norm_buf.raw_ptr(), + grad_norm_f32_buf: grad_norm_f32_buf.raw_ptr(), m_buf: m_buf.raw_ptr(), v_buf: v_buf.raw_ptr(), t_buf: t_buf.raw_ptr(), @@ -2055,6 +2063,7 @@ impl GpuDqnTrainer { stream, ptrs, grad_norm_kernel, + grad_norm_finalize_kernel, grad_norm_standalone, adam_update_kernel, ema_kernel, @@ -2111,6 +2120,7 @@ impl GpuDqnTrainer { m_buf, v_buf, grad_norm_buf, + grad_norm_f32_buf, cql_grad_scratch, t_buf, adam_step: 0, @@ -2293,6 +2303,9 @@ impl GpuDqnTrainer { /// gradients into grad_buf (IQN, attention, ensemble). pub fn replay_adam_and_readback(&mut self) -> Result { self.replay_adam()?; + // Finalize grad_norm OUTSIDE graph: convert float sum-of-squares → bf16 L2 norm. + // The graph stores float sums via atomicAdd; this kernel reads them and writes bf16. + self.launch_grad_norm_finalize()?; // Per-step: zero readback for performance (no cuStreamSync). // Loss/grad_norm accumulated on GPU by training_guard kernel. // Epoch boundary calls readback_scalars_sync() for metrics. @@ -2302,9 +2315,33 @@ impl GpuDqnTrainer { }) } + /// Convert float sum-of-squares → bf16 L2 norm. + /// Reads grad_norm_f32_buf, writes grad_norm_buf. + /// Must be called AFTER graph_adam replay. + fn launch_grad_norm_finalize(&self) -> Result<(), MLError> { + // Uses raw_ptr() → passes u64 device addresses → no event tracking needed. + let f32_ptr = self.grad_norm_f32_buf.raw_ptr(); + let bf16_ptr = self.grad_norm_buf.raw_ptr(); + unsafe { + self.stream + .launch_builder(&self.grad_norm_finalize_kernel) + .arg(&f32_ptr) + .arg(&bf16_ptr) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("grad_norm_finalize: {e}")))?; + } + Ok(()) + } + /// Sync stream and read back loss + grad_norm from GPU. /// Called ONLY at epoch boundary — never per-step. pub fn readback_scalars_sync(&mut self) -> Result { + // Launch finalize to convert float sum-of-squares → bf16 L2 norm + self.launch_grad_norm_finalize()?; unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } @@ -2323,7 +2360,7 @@ impl GpuDqnTrainer { } Ok(FusedTrainScalars { total_loss: loss_bf16[0].to_f32(), - grad_norm: norm_bf16[0].to_f32().sqrt(), + grad_norm: norm_bf16[0].to_f32(), // bf16 L2 norm from finalize kernel }) } @@ -2686,7 +2723,25 @@ impl GpuDqnTrainer { batch_size: usize, ) -> Result { let _eg = EventTrackingGuard::new(self.stream.context()); - self.compute_q_values(states, batch_size)?; + + // Pad input to config.batch_size to prevent CUTLASS tile-boundary OOB reads. + // CUTLASS WMMA uses 32-element N-tiles; non-aligned batch sizes read past buffer. + let mut padded_buf; + let eff_states = if batch_size < self.config.batch_size { + let sd = self.config.state_dim; + padded_buf = self.stream.alloc_zeros::(self.config.batch_size * sd) + .map_err(|e| MLError::ModelError(format!("q_stats pad alloc: {e}")))?; + { + let mut dst = padded_buf.slice_mut(0..batch_size * sd); + self.stream.memcpy_dtod(states, &mut dst) + .map_err(|e| MLError::ModelError(format!("q_stats pad copy: {e}")))?; + } + &padded_buf + } else { + states + }; + + self.compute_q_values(eff_states, batch_size)?; // Zero stats buf then launch reduction self.stream.memset_zeros(&mut self.q_stats_buf) @@ -2972,13 +3027,16 @@ impl GpuDqnTrainer { online_d: &DuelingWeightSet, online_b: &BranchingWeightSet, ) -> Result<(), MLError> { - // ── Zero grad_norm accumulator ───────────────────────────── + // ── Zero float grad_norm accumulator ────────────────────── self.stream - .memset_zeros(&mut self.grad_norm_buf) - .map_err(|e| MLError::ModelError(format!("zero grad_norm: {e}")))?; + .memset_zeros(&mut self.grad_norm_f32_buf) + .map_err(|e| MLError::ModelError(format!("zero grad_norm_f32: {e}")))?; - // ── 5. Gradient norm ─────────────────────────────────────── + // ── 5. Gradient norm (float accumulator) ────────────────── self.launch_grad_norm()?; + // NOTE: grad_norm_finalize runs AFTER graph replay (outside capture). + // The graph stores float sum-of-squares in grad_norm_buf[0:3]. + // Adam reads float sum-of-squares directly and computes sqrt inline. // ── 6. Adam update ──────────────────────────────────────── self.launch_adam_update()?; @@ -3628,15 +3686,15 @@ impl GpuDqnTrainer { }; let grad_ptr = self.ptrs.grad_buf; - let norm_ptr = self.ptrs.grad_norm_buf; + let norm_f32_ptr = self.ptrs.grad_norm_f32_buf; // Safety: argument order matches the extern "C" kernel signature exactly. - // grad_buf has size = total_params; grad_norm_buf has size = 1. + // grad_buf has size = total_params; grad_norm_f32_buf has size = 1 float. unsafe { self.stream .launch_builder(&self.grad_norm_standalone) .arg(&grad_ptr) - .arg(&norm_ptr) + .arg(&norm_f32_ptr) .arg(&tp) .launch(launch_cfg) .map_err(|e| { @@ -3678,12 +3736,12 @@ impl GpuDqnTrainer { let grad_ptr = self.ptrs.grad_buf; let m_ptr = self.ptrs.m_buf; let v_ptr = self.ptrs.v_buf; - let norm_ptr = self.ptrs.grad_norm_buf; + let norm_f32_ptr = self.ptrs.grad_norm_f32_buf; let t_ptr = self.ptrs.t_buf; // Safety: argument order matches the extern "C" kernel signature exactly. // All buffers are pre-allocated with size = total_params. - // grad_norm_buf contains the completed norm from launch_grad_norm(). + // grad_norm_f32_buf contains float sum-of-squares from launch_grad_norm(). // t_buf is a device pointer (kernel reads *t_ptr) for CUDA Graph compatibility. unsafe { self.stream @@ -3692,7 +3750,7 @@ impl GpuDqnTrainer { .arg(&grad_ptr) .arg(&m_ptr) .arg(&v_ptr) - .arg(&norm_ptr) + .arg(&norm_f32_ptr) .arg(&lr) .arg(&beta1) .arg(&beta2) @@ -3967,11 +4025,11 @@ impl GpuDqnTrainer { /// bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, relu_mask, spectral_norm, /// clipped_saxpy, clip_grad. /// -/// Returns `(grad_norm, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, relu_mask, spectral_norm, clipped_saxpy, clip_grad)`. +/// Returns `(grad_norm, grad_norm_finalize, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, relu_mask, spectral_norm, clipped_saxpy, clip_grad)`. fn compile_training_kernels( stream: &Arc, config: &GpuDqnTrainConfig, -) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { +) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { info!( state_dim = config.state_dim, total_params = compute_total_params(config), @@ -3985,6 +4043,8 @@ fn compile_training_kernels( let grad_norm = module.load_function("dqn_grad_norm_kernel") .map_err(|e| MLError::ModelError(format!("dqn_grad_norm_kernel load: {e}")))?; + let grad_norm_finalize = module.load_function("dqn_grad_norm_finalize") + .map_err(|e| MLError::ModelError(format!("dqn_grad_norm_finalize load: {e}")))?; let adam_update = module.load_function("dqn_adam_update_kernel") .map_err(|e| MLError::ModelError(format!("dqn_adam_update_kernel load: {e}")))?; let f32_to_bf16 = module.load_function("f32_to_bf16_kernel") @@ -4009,7 +4069,7 @@ fn compile_training_kernels( .map_err(|e| MLError::ModelError(format!("dqn_clip_grad_kernel load: {e}")))?; info!("GpuDqnTrainer: 12 utility kernels loaded from precompiled cubin"); - Ok((grad_norm, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clipped_saxpy, clip_grad)) + Ok((grad_norm, grad_norm_finalize, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clipped_saxpy, clip_grad)) } /// Load the standalone Polyak EMA kernel from precompiled cubin. diff --git a/crates/ml/src/cuda_pipeline/gpu_training_guard.rs b/crates/ml/src/cuda_pipeline/gpu_training_guard.rs index 6c91a1643..d61472b10 100644 --- a/crates/ml/src/cuda_pipeline/gpu_training_guard.rs +++ b/crates/ml/src/cuda_pipeline/gpu_training_guard.rs @@ -276,6 +276,8 @@ impl GpuTrainingGuard { warmup: bool, ) -> Result { let _nvtx = NvtxRange::new("gpu_training_guard"); + // Disable event tracking — loss/grad buffers may be captured in CUDA graphs. + let _eg = super::gpu_dqn_trainer::EventTrackingGuard::new(self.stream.context()); // Read PREVIOUS step's result (one-step delay; first step returns defaults) let prev_result = if self.guard_has_prev { @@ -313,14 +315,17 @@ impl GpuTrainingGuard { shared_mem_bytes: 0, }; - // Safety: loss_gpu and grad_norm_gpu are valid F32 CudaSlice on the same context. + // Pass raw device pointers to bypass cudarc event tracking on graph-captured buffers. + let loss_ptr = loss_gpu.raw_ptr(); + let grad_ptr = grad_norm_gpu.raw_ptr(); + let acc_ptr = self.acc_buf.raw_ptr(); unsafe { self.stream .launch_builder(&self.fused_check_accum_func) - .arg(loss_gpu) - .arg(grad_norm_gpu) + .arg(&loss_ptr) + .arg(&grad_ptr) .arg(&write_dev_ptr) - .arg(&mut self.acc_buf) + .arg(&acc_ptr) .arg(&clip_threshold) .arg(&collapse_threshold) .arg(&warmup_int) diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 3f65c1f41..53b3fbdcb 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -662,10 +662,11 @@ impl DQNAgentType { /// Update replay buffer priorities from GPU-resident tensors (GpuPrioritized only). pub fn update_priorities_gpu( &self, - indices: &GpuTensor, - td_errors: &GpuTensor, + indices: &cudarc::driver::CudaSlice, + td_errors: &cudarc::driver::CudaSlice, + bs: usize, ) -> Result<(), MLError> { - self.memory().update_priorities_gpu(indices, td_errors) + self.memory().update_priorities_gpu(indices, td_errors, bs) } /// Step the replay buffer training counter for beta annealing (PER only). diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 2c9d328e2..538ffabf7 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -802,15 +802,9 @@ impl FusedTrainingCtx { if let (Some(priorities_tensor), Some((alpha, epsilon))) = (agent.priorities_tensor()?, agent.per_alpha_epsilon()) { - // GpuBatch::indices stores u32 bit-patterns in f32 CudaSlice containers. - // Reinterpret the f32 CudaSlice as u32 to pass to the PER kernel. - let idx_f32 = effective_gpu.indices.data(); - let idx_u32: &cudarc::driver::CudaSlice = unsafe { - &*(idx_f32 as *const cudarc::driver::CudaSlice - as *const cudarc::driver::CudaSlice) - }; + // GpuBatch::indices is now native CudaSlice — pass directly. self.trainer.update_priorities_cuda( - idx_u32, + &effective_gpu.indices, priorities_tensor.data(), alpha, epsilon, @@ -1171,6 +1165,18 @@ impl FusedTrainingCtx { self.trainer.set_c51_alpha(alpha); } + /// Return a reference to the GPU-resident total_loss scalar (bf16). + /// Written by the CUDA graph's loss kernel — valid after `replay_forward()`. + pub(crate) fn loss_gpu_buf(&self) -> &cudarc::driver::CudaSlice { + &self.trainer.total_loss_buf + } + + /// Return a reference to the GPU-resident grad_norm scalar (bf16). + /// Written by the CUDA graph's grad_norm kernel — valid after `replay_adam()`. + pub(crate) fn grad_norm_gpu_buf(&self) -> &cudarc::driver::CudaSlice { + &self.trainer.grad_norm_buf + } + /// Get CUDA stream reference for DtoH transfers after compute_q_values. pub(crate) fn stream(&self) -> &Arc { &self.stream diff --git a/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs b/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs index 1a1bab803..12b64e170 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs @@ -78,6 +78,7 @@ fn insert_random_batch( } /// Download a CudaSlice to host for assertion. +#[allow(dead_code)] fn dtoh_f32(slice: &cudarc::driver::CudaSlice, stream: &Arc) -> anyhow::Result> { let n = slice.len(); let mut host = vec![0.0_f32; n]; @@ -85,6 +86,14 @@ fn dtoh_f32(slice: &cudarc::driver::CudaSlice, stream: &Arc (bf16) to host as f32 for assertion. +fn dtoh_bf16_as_f32(slice: &cudarc::driver::CudaSlice, stream: &Arc) -> anyhow::Result> { + let n = slice.len(); + let mut host = vec![0_u16; n]; + stream.memcpy_dtoh(slice, &mut host).map_err(|e| anyhow::anyhow!("dtoh u16: {e}"))?; + Ok(host.iter().map(|&bits| half::bf16::from_bits(bits).to_f32()).collect()) +} + /// Download a CudaSlice to host for assertion. fn dtoh_u32(slice: &cudarc::driver::CudaSlice, stream: &Arc) -> anyhow::Result> { let n = slice.len(); @@ -118,7 +127,7 @@ fn test_gpu_replay_buffer_proportional_sample_valid() -> anyhow::Result<()> { assert_eq!(batch.batch_size, 16); // Download weights and indices to host for assertions - let weights_host = dtoh_f32(&batch.weights, &stream)?; + let weights_host = dtoh_bf16_as_f32(&batch.weights, &stream)?; let indices_host = dtoh_u32(&batch.indices, &stream)?; // All weights must be positive @@ -154,7 +163,7 @@ fn test_gpu_replay_buffer_rank_based_sample_valid() -> anyhow::Result<()> { assert_eq!(batch.batch_size, 16); // Download weights and indices to host for assertions - let weights_host = dtoh_f32(&batch.weights, &stream)?; + let weights_host = dtoh_bf16_as_f32(&batch.weights, &stream)?; let indices_host = dtoh_u32(&batch.indices, &stream)?; // All weights must be positive @@ -193,7 +202,7 @@ fn test_gpu_replay_buffer_priority_update_valid() -> anyhow::Result<()> { // Verify sampling still works with updated priorities (proves update succeeded) let batch2 = buf.sample_proportional(16)?; - let weights_host = dtoh_f32(&batch2.weights, &stream)?; + let weights_host = dtoh_bf16_as_f32(&batch2.weights, &stream)?; // Weights should still be positive after priority update for (i, &w) in weights_host.iter().enumerate() { diff --git a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs index 4ed248d32..5cd02fd7f 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs @@ -116,9 +116,10 @@ fn test_per_weights_valid() -> anyhow::Result<()> { let batch = buf.sample_proportional(16)?; - // Download weights to host for assertion - let mut weights_host = vec![0.0_f32; batch.weights.len()]; - stream.memcpy_dtoh(&batch.weights, &mut weights_host).map_err(|e| anyhow::anyhow!("{e}"))?; // test-only readback + // Download bf16 weights to host for assertion + let mut weights_bf16 = vec![0_u16; batch.weights.len()]; + stream.memcpy_dtoh(&batch.weights, &mut weights_bf16).map_err(|e| anyhow::anyhow!("{e}"))?; // test-only readback + let weights_host: Vec = weights_bf16.iter().map(|&bits| half::bf16::from_bits(bits).to_f32()).collect(); for (i, &w) in weights_host.iter().enumerate() { assert!(w > 0.0, "weight[{i}] not positive: {w}"); diff --git a/crates/ml/src/trainers/dqn/trainer/train_step.rs b/crates/ml/src/trainers/dqn/trainer/train_step.rs index 95b1ff766..aff5f4bc7 100644 --- a/crates/ml/src/trainers/dqn/trainer/train_step.rs +++ b/crates/ml/src/trainers/dqn/trainer/train_step.rs @@ -120,14 +120,17 @@ 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 = 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}"))?; + // Read loss/grad directly from fused trainer's GPU buffers + // (not from GpuTrainResult which returns hardcoded zeros per-step). + let fused = self.fused_ctx.as_mut() + .ok_or_else(|| anyhow::anyhow!("fused_ctx required for training guard"))?; + + let loss_buf = fused.loss_gpu_buf(); + let grad_buf = fused.grad_norm_gpu_buf(); let result = guard .check_and_accumulate( - &loss_slice, - &grad_slice, + loss_buf, + grad_buf, 1e6_f32, // loss clip threshold grad_collapse_threshold, !past_warmup, @@ -273,7 +276,7 @@ impl DQNTrainer { #[allow(unused_mut, unused_assignments, unused_variables)] let mut final_grad_norm = 0.0_f32; let mut gpu_td_errors: Vec = Vec::new(); - let mut gpu_indices: Vec = Vec::new(); + let mut gpu_indices: Vec> = Vec::new(); let mut gpu_loss_tensors: Vec = Vec::new(); let mut gpu_grad_tensors: Vec = Vec::new(); @@ -409,13 +412,24 @@ impl DQNTrainer { if !gpu_td_errors.is_empty() && !gpu_indices.is_empty() { let cat_stream = self.cuda_stream.as_ref().ok_or_else(|| anyhow::anyhow!("CUDA stream for PER concat"))?; let td_refs: Vec<&GpuTensor> = gpu_td_errors.iter().collect(); - let idx_refs: Vec<&GpuTensor> = gpu_indices.iter().collect(); let td_cat = GpuTensor::cat(&td_refs, 0, cat_stream) .map_err(|e| anyhow::anyhow!("GPU TD error concat failed: {}", e))?; - let idx_cat = GpuTensor::cat(&idx_refs, 0, cat_stream) - .map_err(|e| anyhow::anyhow!("GPU index concat failed: {}", e))?; + + // Concatenate CudaSlice indices via GPU DtoD copies + let total_idx_len: usize = gpu_indices.iter().map(|s| s.len()).sum(); + let mut idx_cat = cat_stream.alloc_zeros::(total_idx_len) + .map_err(|e| anyhow::anyhow!("GPU PER index alloc: {e}"))?; + let mut offset = 0usize; + for idx_slice in &gpu_indices { + let len = idx_slice.len(); + let mut sub = idx_cat.slice_mut(offset..offset + len); + cat_stream.memcpy_dtod(idx_slice, &mut sub) + .map_err(|e| anyhow::anyhow!("GPU PER index DtoD concat: {e}"))?; + offset += len; + } + agent - .update_priorities_gpu(&idx_cat, &td_cat) + .update_priorities_gpu(&idx_cat, td_cat.data(), total_idx_len) .map_err(|e| anyhow::anyhow!("GPU PER priority update failed: {}", e))?; } else if !all_indices.is_empty() { agent diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 24b738b04..164eff302 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1263,17 +1263,29 @@ impl DQNTrainer { let guard_start = std::time::Instant::now(); if let Some(ref mut guard) = self.training_guard { - 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, - 1e6_f32, - guard_collapse_thresh, - !guard_past_warmup, - ).map_err(|e| anyhow::anyhow!("guard check: {e}"))?; + // Read loss/grad directly from fused trainer's GPU buffers. + // GpuTrainResult returns hardcoded zeros per-step (no sync). + let gr = if let Some(ref fused) = self.fused_ctx { + guard.check_and_accumulate( + fused.loss_gpu_buf(), + fused.grad_norm_gpu_buf(), + 1e6_f32, + guard_collapse_thresh, + !guard_past_warmup, + ).map_err(|e| anyhow::anyhow!("guard check: {e}"))? + } else { + 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}"))?; + guard.check_and_accumulate( + &loss_slice, + &grad_slice, + 1e6_f32, + guard_collapse_thresh, + !guard_past_warmup, + ).map_err(|e| anyhow::anyhow!("guard check: {e}"))? + }; if gr.halt_nan { return Err(anyhow::anyhow!( "NaN/Inf at step {}: loss={}, grad={}", @@ -1464,9 +1476,13 @@ impl DQNTrainer { train_step_count: usize, monitor: &mut TrainingMonitor, ) -> Result { + let (avg_loss, avg_grad) = if let Some(ref mut guard) = self.training_guard { let (al, ag) = guard.read_accumulators() .map_err(|e| anyhow::anyhow!("guard read_accumulators: {e}"))?; + warn!("EPOCH_DIAG[{}]: acc avg_loss={:.6}, avg_grad={:.6}", epoch, al, ag); + guard.reset_accumulators() + .map_err(|e| anyhow::anyhow!("guard reset_accumulators: {e}"))?; (al as f32, ag as f32) } else { (0.0_f32, 0.0_f32) diff --git a/crates/ml/src/training_profile.rs b/crates/ml/src/training_profile.rs index e650fcb2b..e7164d1d3 100644 --- a/crates/ml/src/training_profile.rs +++ b/crates/ml/src/training_profile.rs @@ -1128,8 +1128,8 @@ mod tests { ); // reward_scale should be applied assert!( - (hp.reward_scale - 10.0).abs() < 0.01, - "reward_scale should be 10.0, got {}", + (hp.reward_scale - 1.0).abs() < 0.01, + "reward_scale should be 1.0, got {}", hp.reward_scale ); // v_min/v_max should be recomputed from reward_scale + gamma @@ -1159,14 +1159,14 @@ mod tests { hp.iqn_lambda ); assert!( - (hp.spectral_norm_sigma_max - 3.0).abs() < 0.01, - "spectral_norm_sigma_max should be 3.0 from smoketest [advanced], got {}", + (hp.spectral_norm_sigma_max - 1.5).abs() < 0.01, + "spectral_norm_sigma_max should be 1.5 from smoketest [advanced], got {}", hp.spectral_norm_sigma_max ); // Exploration section new fields assert!( - (hp.noisy_sigma_init - 0.5).abs() < 0.01, - "noisy_sigma_init should be 0.5 from smoketest [exploration], got {}", + (hp.noisy_sigma_init - 0.3).abs() < 0.01, + "noisy_sigma_init should be 0.3 from smoketest [exploration], got {}", hp.noisy_sigma_init ); assert!( diff --git a/crates/ml/tests/gpu_per_integration_test.rs b/crates/ml/tests/gpu_per_integration_test.rs index cca6eaa68..223cd3fbf 100644 --- a/crates/ml/tests/gpu_per_integration_test.rs +++ b/crates/ml/tests/gpu_per_integration_test.rs @@ -153,8 +153,8 @@ fn test_gpu_per_training_loop_cycle() { // Simulate TD errors: decreasing magnitude over epochs let scale = 1.0_f32 / (epoch as f32 + 1.0); - let td_host = vec![scale; batch_size]; - let td_errors: CudaSlice = stream.clone_htod(&td_host).unwrap(); + let td_host: Vec = vec![half::bf16::from_f32(scale); batch_size]; + let td_errors: CudaSlice = stream.clone_htod(&td_host).unwrap(); buf.update_priorities_gpu(&batch.indices, &td_errors, batch_size).unwrap(); buf.step(); @@ -178,7 +178,7 @@ fn test_gpu_per_priorities_affect_sampling() { // Set experience 0 to very high priority let high_idx: CudaSlice = stream.clone_htod(&[0_u32]).unwrap(); - let high_td: CudaSlice = stream.clone_htod(&[100.0_f32]).unwrap(); + let high_td: CudaSlice = stream.clone_htod(&[half::bf16::from_f32(100.0)]).unwrap(); buf.update_priorities_gpu(&high_idx, &high_td, 1).unwrap(); // Sample 1000 times in batches of 10, count how often index 0 appears @@ -222,7 +222,7 @@ fn test_gpu_per_sampling_distribution_differs_from_uniform() { // This creates a strong left-to-right priority gradient for i in 0..50 { let idx: CudaSlice = stream.clone_htod(&[i as u32]).unwrap(); - let td: CudaSlice = stream.clone_htod(&[(i as f32 + 1.0)]).unwrap(); + let td: CudaSlice = stream.clone_htod(&[half::bf16::from_f32(i as f32 + 1.0)]).unwrap(); buf.update_priorities_gpu(&idx, &td, 1).unwrap(); } @@ -264,15 +264,16 @@ fn test_gpu_per_is_weights_correct_range() { // Set varied priorities for i in 0..200 { let idx: CudaSlice = stream.clone_htod(&[i as u32]).unwrap(); - let td: CudaSlice = stream.clone_htod(&[((i % 10) as f32 + 1.0)]).unwrap(); + let td: CudaSlice = stream.clone_htod(&[half::bf16::from_f32((i % 10) as f32 + 1.0)]).unwrap(); buf.update_priorities_gpu(&idx, &td, 1).unwrap(); } let batch = buf.sample_proportional(64).unwrap(); - // Download f32 weights to host for validation - let mut weights_host = vec![0.0_f32; 64]; - stream.memcpy_dtoh(&batch.weights, &mut weights_host).unwrap(); + // Download bf16 weights to host, convert to f32 for validation + let mut weights_bf16 = vec![half::bf16::ZERO; 64]; + stream.memcpy_dtoh(&batch.weights, &mut weights_bf16).unwrap(); + let weights_host: Vec = weights_bf16.iter().map(|x| x.to_f32()).collect(); let w_min = weights_host.iter().copied().fold(f32::INFINITY, f32::min); let w_max = weights_host.iter().copied().fold(f32::NEG_INFINITY, f32::max);