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

fix: OOB read in compute_expected_q — tile per_sample_support [N,3] instead of 2-float v_range ptr

Browse Source 
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<f32> 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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-15 23:43:38 +02:00
parent a58f71f6e1
commit 8a54c8a32c
2 changed files with 48 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -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<f32>,
/// 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::<f32>(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::<f32>(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::<i32>(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<f32> = (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;

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

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