From 8a54c8a32ce4c80709e06fdc8cd2551160bfaf8b Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 15 Apr 2026 23:43:38 +0200 Subject: [PATCH] =?UTF-8?q?fix:=20OOB=20read=20in=20compute=5Fexpected=5Fq?= =?UTF-8?q?=20=E2=80=94=20tile=20per=5Fsample=5Fsupport=20[N,3]=20instead?= =?UTF-8?q?=20of=202-float=20v=5Frange=20ptr?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The compute_expected_q and quantile_q_select kernels read per_sample_support[i*3+0/1/2] (3 floats per sample), but the experience collector was passing eval_v_range_ptr which is only 2 floats (v_min, v_max). Every sample after sample 0 read out of bounds. Replace the u64 pointer field with a proper CudaSlice buffer [alloc_episodes, 3] that is tiled with [v_min, v_max, delta_z] once per epoch via update_per_sample_support(). Co-Authored-By: Claude Opus 4.6 (1M context) --- .../cuda_pipeline/gpu_experience_collector.rs | 34 +++++++++++++------ .../src/trainers/dqn/trainer/training_loop.rs | 27 +++++++++++++-- 2 files changed, 48 insertions(+), 13 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 12cf1f5ac..72eb206f8 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -461,8 +461,9 @@ pub struct GpuExperienceCollector { network_dims: (usize, usize, usize, usize), /// Number of C51 atoms (1 for non-distributional). num_atoms: usize, - /// Per-sample C51 support pointer — points to IQL's per_sample_support_buf [B, 3]. - eval_v_range_ptr: u64, + /// Per-sample C51 support tiled [alloc_episodes, 3]: [v_min, v_max, delta_z] per episode. + /// Filled once per epoch via `update_per_sample_support()` — matches kernel expectation. + per_sample_support_buf: CudaSlice, /// Per-sample epsilon from IQL expectile gap (0 = use cosine schedule). per_sample_epsilon_ptr: u64, /// Branch sizes [3, 3, 3, 3] for 4-branch hierarchical DQN. @@ -1076,6 +1077,12 @@ impl GpuExperienceCollector { let exp_bn_concat = stream.alloc_zeros::(alloc_episodes * (bn_alloc + portfolio_dim_bn) + 128) .map_err(|e| MLError::ModelError(format!("alloc exp_bn_concat: {e}")))?; + // Per-sample C51 support buffer [alloc_episodes, 3] — tiled v_min/v_max/delta_z. + // Filled once per epoch in update_per_sample_support() — replaces the 2-float + // eval_v_range pointer that caused OOB reads in compute_expected_q/quantile_q_select. + let per_sample_support_buf = stream.alloc_zeros::(alloc_episodes * 3) + .map_err(|e| MLError::ModelError(format!("alloc per_sample_support_buf: {e}")))?; + // Task 8: GPU-resident step counter for experience collection loop let step_counter_gpu = stream.alloc_zeros::(1) .map_err(|e| MLError::ModelError(format!("alloc step_counter_gpu: {e}")))?; @@ -1100,7 +1107,7 @@ impl GpuExperienceCollector { market_dim, network_dims, num_atoms, - eval_v_range_ptr: 0, // Set via set_eval_v_range_ptr() after fused ctx init + per_sample_support_buf, // Tiled [alloc_episodes, 3]; filled via update_per_sample_support() per_sample_epsilon_ptr: 0, // Set via set_per_sample_epsilon_ptr() after IQL init branch_sizes, alloc_episodes, @@ -2023,7 +2030,6 @@ impl GpuExperienceCollector { // handles num_atoms=1 correctly, and logits are now f32. { let na = self.num_atoms as i32; - let eval_v_range_ptr = self.eval_v_range_ptr; unsafe { let null_atom_stats = 0u64; let null_atom_positions = 0u64; // NULL = use linear atom spacing @@ -2038,7 +2044,7 @@ impl GpuExperienceCollector { .arg(&b1) .arg(&b2) .arg(&b3) - .arg(&eval_v_range_ptr) + .arg(&self.per_sample_support_buf) // [N, 3] per-sample support .arg(&null_atom_stats) .arg(&mut self.q_var_buf) .arg(&null_atom_positions) // atom_positions (NULL = linear) @@ -2144,7 +2150,6 @@ impl GpuExperienceCollector { // q_values (E[Q]) is preserved for monitoring, Q-gap computation, and loss. { let na = self.num_atoms as i32; - let eval_v_range_ptr = self.eval_v_range_ptr; let iqn_r = self.iqn_readiness; let util_e = self.util_ema; unsafe { @@ -2159,7 +2164,7 @@ impl GpuExperienceCollector { .arg(&b1) .arg(&b2) .arg(&b3) - .arg(&eval_v_range_ptr) + .arg(&self.per_sample_support_buf) // [N, 3] per-sample support .arg(&iqn_r) .arg(&util_e) .launch(launch_cfg) @@ -2380,9 +2385,18 @@ impl GpuExperienceCollector { info!(trainer_params_ptr = ptr, "Experience collector: zero-copy trainer params pointer set"); } - /// Set the raw device pointer to the trainer's v_range_buf [v_min, v_max]. - pub fn set_eval_v_range_ptr(&mut self, ptr: u64) { - self.eval_v_range_ptr = ptr; + /// Fill per_sample_support_buf [alloc_episodes, 3] with uniform v_min/v_max/delta_z. + /// + /// Called once per epoch (not hot-path) when the eval v_range changes. + /// Replaces the old `set_eval_v_range_ptr` which passed a 2-float pointer + /// to kernels that expected [N, 3], causing OOB reads for every sample > 0. + pub fn update_per_sample_support(&mut self, v_min: f32, v_max: f32) -> Result<(), MLError> { + let n = self.alloc_episodes; + let delta_z = (v_max - v_min) / (self.num_atoms as f32 - 1.0).max(1.0); + let data: Vec = (0..n).flat_map(|_| [v_min, v_max, delta_z]).collect(); + super::htod_f32(&self.stream, &data, &mut self.per_sample_support_buf)?; + debug!(v_min, v_max, delta_z, n, "per_sample_support tiled for experience collector"); + Ok(()) } pub fn set_per_sample_epsilon_ptr(&mut self, ptr: u64) { self.per_sample_epsilon_ptr = ptr; diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 3d5e993bf..85a4112ff 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -913,11 +913,14 @@ impl DQNTrainer { .collect(); collector.upload_ofi_features(&flat)?; } - // Wire params + v_range pointers from fused trainer (zero-copy weight access) + // Wire params + epsilon pointers from fused trainer (zero-copy weight access) if let Some(ref fused_ctx) = self.fused_ctx { collector.set_trainer_params_ptr(fused_ctx.params_flat_ptr()); - collector.set_eval_v_range_ptr(fused_ctx.eval_v_range_ptr()); collector.set_per_sample_epsilon_ptr(fused_ctx.per_sample_epsilon_ptr()); + // Tile per_sample_support with initial v_range (will be re-tiled each epoch) + let vr = fused_ctx.eval_v_range(); + collector.update_per_sample_support(vr[0], vr[1]) + .map_err(|e| anyhow::anyhow!("init per_sample_support: {e}"))?; } // Wire curiosity weights to fused trainer for Q-penalty if let Some(ref mut fused_ctx) = self.fused_ctx { @@ -1101,6 +1104,20 @@ impl DQNTrainer { return Err(anyhow::anyhow!("GPU episode reset FAILED (no CPU fallback): {e}")); } + // Tile per-sample support [N, 3] from current eval v_range before experience collection. + // Reads v_min/v_max from fused trainer (which may have been updated last epoch), + // falls back to config defaults for epoch 0 / pre-fused-init. + { + let (v_min, v_max) = if let Some(ref fused_ctx) = self.fused_ctx { + let vr = fused_ctx.eval_v_range(); + (vr[0], vr[1]) + } else { + (self.hyperparams.v_min as f32, self.hyperparams.v_max as f32) + }; + collector.update_per_sample_support(v_min, v_max) + .map_err(|e| anyhow::anyhow!("update_per_sample_support: {e}"))?; + } + let agent = self.agent.read().await; let epsilon = agent.get_epsilon(); drop(agent); @@ -1343,7 +1360,11 @@ impl DQNTrainer { // Zero-copy: wire collector to read weights directly from trainer's params_buf. if let Some(ref mut collector) = self.gpu_experience_collector { collector.set_trainer_params_ptr(ctx.params_flat_ptr()); - collector.set_eval_v_range_ptr(ctx.eval_v_range_ptr()); + // Tile per_sample_support with current v_range + let vr = ctx.eval_v_range(); + if let Err(e) = collector.update_per_sample_support(vr[0], vr[1]) { + tracing::error!("update_per_sample_support on fused re-init: {e}"); + } } self.fused_ctx = Some(ctx); }