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

fix(gpu): resolve 3 pre-existing issues flagged by zen code review

Browse Source 
1. PER duplicate index accumulation (HIGH): update_priorities_gpu()
   used index_add delta trick which accumulates deltas for duplicate
   indices, overshooting target priority. Now deduplicates the small
   indices tensor (batch_size=256, ~1KB) via HashMap before delta
   computation. Fast path (no dupes) reuses original tensors.

2. RegimeConditional silent experience drop (MEDIUM):
   insert_batch_tensors() returned Ok(()) for CPU replay buffers,
   silently discarding all GPU-collected experiences. Now converts
   tensors→Vec<Experience> via extracted helper (lazy, only on first
   CPU head encountered) and inserts into CPU buffer.

3. Unsafe direct indexing (LOW): gpu_experience_collector.rs used
   &states[s..e] in fallback paths, violating deny(indexing_slicing).
   Replaced with .get() safe bounds checking.

392/392 ml-dqn + 874/874 ml tests pass, 0 clippy warnings.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-10 00:53:16 +01:00
parent 5149c71444
commit ff7d79b03a
3 changed files with 117 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
crates/ml-dqn/src/gpu_replay_buffer.rs
View File

@@ -594,15 +594,48 @@ impl GpuReplayBuffer {
// Delta trick: gather old priorities at indices, compute delta = new - old,
// then index_add the deltas back. Single batched GPU kernel, no CPU loop.
//
// For rare duplicate indices (batch_size << capacity), the deltas may
// stack — index_add adds ALL deltas, so duplicates double-apply.
// Post-clamp prevents priorities from going to zero/negative (→ NaN weights).
// Dedup: index_add accumulates ALL deltas at each position, so duplicate
// indices would double-apply. Pull the small indices tensor (batch_size
// = 256 → 1 KB) to CPU, deduplicate keeping the last td_error per index,
// then rebuild GPU tensors with unique entries only.
let indices_i64 = indices.to_dtype(DType::I64)?;
let old_prios = self.priorities.index_select(&indices_i64, 0)?;
let delta = (&clamped - &old_prios)?;
let (unique_idx, unique_prios) = {
let idx_vec: Vec<i64> = indices_i64.to_vec1().map_err(|e| {
MLError::ModelError(format!("indices to_vec1: {e}"))
})?;
let prio_vec: Vec<f32> = clamped.to_vec1().map_err(|e| {
MLError::ModelError(format!("clamped to_vec1: {e}"))
})?;
// Last-write-wins: later occurrences overwrite earlier ones
let mut seen = std::collections::HashMap::with_capacity(idx_vec.len());
for (i, &idx) in idx_vec.iter().enumerate() {
seen.insert(idx, i);
}
if seen.len() < idx_vec.len() {
// Duplicates found — rebuild with unique entries
let mut u_idx = Vec::with_capacity(seen.len());
let mut u_prio = Vec::with_capacity(seen.len());
for (&idx, &pos) in &seen {
u_idx.push(idx);
if let Some(&p) = prio_vec.get(pos) {
u_prio.push(p);
}
}
let len = u_idx.len();
(
Tensor::from_vec(u_idx, &[len], &self.device)?,
Tensor::from_vec(u_prio, &[len], &self.device)?,
)
} else {
// No duplicates — fast path, reuse existing tensors
(indices_i64, clamped)
}
};
let old_prios = self.priorities.index_select(&unique_idx, 0)?;
let delta = (&unique_prios - &old_prios)?;
self.priorities = self
.priorities
.index_add(&indices_i64, &delta, 0)?
.index_add(&unique_idx, &delta, 0)?
.clamp(self.config.epsilon, 1e6)?;
Ok(())

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

@@ -812,13 +812,17 @@ impl GpuExperienceCollector {
} else {
// Out of bounds — copy current state as fallback
let s_start = idx * sd;
next_states.extend_from_slice(&states[s_start..s_start + sd]);
if let Some(s) = states.get(s_start..s_start + sd) {
next_states.extend_from_slice(s);
}
}
} else {
// Terminal or last timestep — next_state = current state
// (masked by done flag in Bellman equation)
let s_start = idx * sd;
next_states.extend_from_slice(&states[s_start..s_start + sd]);
if let Some(s) = states.get(s_start..s_start + sd) {
next_states.extend_from_slice(s);
}
}
}
}

123
crates/ml/src/trainers/dqn/config.rs
View File

@@ -281,6 +281,50 @@ impl DQNAgentType {
}
}
/// Convert GPU tensors to `Vec<Experience>` for CPU replay buffer insertion.
///
/// Used as fallback when a replay buffer is CPU-based (Prioritized or Uniform)
/// but experiences were collected on GPU.
#[cfg(feature = "cuda")]
fn tensors_to_experiences(
states: &candle_core::Tensor,
next_states: &candle_core::Tensor,
actions: &candle_core::Tensor,
rewards: &candle_core::Tensor,
dones: &candle_core::Tensor,
) -> Result<Vec<Experience>, MLError> {
let batch_size = states.dim(0).map_err(|e| {
MLError::TrainingError(format!("states dim: {e}"))
})?;
let state_dim = states.dim(1).map_err(|e| {
MLError::TrainingError(format!("states dim1: {e}"))
})?;
let s_flat: Vec<f32> = states.to_dtype(candle_core::DType::F32)?.flatten_all()?.to_vec1()
.map_err(|e| MLError::TrainingError(format!("states to_vec1: {e}")))?;
let ns_flat: Vec<f32> = next_states.to_dtype(candle_core::DType::F32)?.flatten_all()?.to_vec1()
.map_err(|e| MLError::TrainingError(format!("next_states to_vec1: {e}")))?;
let a_vec: Vec<u32> = actions.to_dtype(candle_core::DType::U32)?.to_vec1()
.map_err(|e| MLError::TrainingError(format!("actions to_vec1: {e}")))?;
let r_vec: Vec<f32> = rewards.to_dtype(candle_core::DType::F32)?.to_vec1()
.map_err(|e| MLError::TrainingError(format!("rewards to_vec1: {e}")))?;
let d_vec: Vec<f32> = dones.to_dtype(candle_core::DType::F32)?.to_vec1()
.map_err(|e| MLError::TrainingError(format!("dones to_vec1: {e}")))?;
let mut experiences = Vec::with_capacity(batch_size);
for i in 0..batch_size {
let start = i * state_dim;
let end = start + state_dim;
experiences.push(Experience {
state: s_flat.get(start..end).unwrap_or(&[]).to_vec(),
action: a_vec.get(i).copied().unwrap_or(0) as u8,
reward: (r_vec.get(i).copied().unwrap_or(0.0) * 1_000_000.0) as i32,
next_state: ns_flat.get(start..end).unwrap_or(&[]).to_vec(),
done: d_vec.get(i).map_or(false, |&v| v > 0.5),
timestamp: 0, // fallback path — timestamp not meaningful for replay
});
}
Ok(experiences)
}
/// Insert a batch of experience tensors directly into the GPU replay buffer.
///
/// When the replay buffer is `GpuPrioritized`, this feeds tensors straight to
@@ -300,63 +344,40 @@ impl DQNAgentType {
if let Some(mut gpu_buf) = agent.memory.as_gpu_buffer() {
gpu_buf.gpu.insert_batch(states, next_states, actions, rewards, dones)
} else {
// CPU PER fallback: convert tensors back to experiences
let batch_size = states.dim(0).map_err(|e| {
MLError::TrainingError(format!("states dim: {e}"))
})?;
let state_dim = states.dim(1).map_err(|e| {
MLError::TrainingError(format!("states dim1: {e}"))
})?;
let s_flat: Vec<f32> = states.to_dtype(candle_core::DType::F32)?.flatten_all()?.to_vec1()
.map_err(|e| MLError::TrainingError(format!("states to_vec1: {e}")))?;
let ns_flat: Vec<f32> = next_states.to_dtype(candle_core::DType::F32)?.flatten_all()?.to_vec1()
.map_err(|e| MLError::TrainingError(format!("next_states to_vec1: {e}")))?;
let a_vec: Vec<u32> = actions.to_dtype(candle_core::DType::U32)?.to_vec1()
.map_err(|e| MLError::TrainingError(format!("actions to_vec1: {e}")))?;
let r_vec: Vec<f32> = rewards.to_dtype(candle_core::DType::F32)?.to_vec1()
.map_err(|e| MLError::TrainingError(format!("rewards to_vec1: {e}")))?;
let d_vec: Vec<f32> = dones.to_dtype(candle_core::DType::F32)?.to_vec1()
.map_err(|e| MLError::TrainingError(format!("dones to_vec1: {e}")))?;
let mut experiences = Vec::with_capacity(batch_size);
for i in 0..batch_size {
let start = i * state_dim;
let end = start + state_dim;
experiences.push(crate::dqn::Experience {
state: s_flat.get(start..end).unwrap_or(&[]).to_vec(),
action: a_vec.get(i).copied().unwrap_or(0) as u8,
reward: (r_vec.get(i).copied().unwrap_or(0.0) * 1_000_000.0) as i32,
next_state: ns_flat.get(start..end).unwrap_or(&[]).to_vec(),
done: d_vec.get(i).map_or(false, |&v| v > 0.5),
timestamp: 0, // fallback path — timestamp not meaningful for replay
});
}
let experiences = Self::tensors_to_experiences(
states, next_states, actions, rewards, dones,
)?;
agent.memory.add_batch(experiences)
}
}
Self::RegimeConditional(agent) => {
// Insert into all head buffers — regime routing happens at train time
// via GPU tensor mask operations (zero CPU roundtrip).
agent.get_trending_head().map(|h| &h.memory).map_or(Ok(()), |m| {
if let Some(mut gpu_buf) = m.as_gpu_buffer() {
gpu_buf.gpu.insert_batch(states, next_states, actions, rewards, dones)
} else {
Ok(())
}
})?;
agent.get_ranging_head().map(|h| &h.memory).map_or(Ok(()), |m| {
if let Some(mut gpu_buf) = m.as_gpu_buffer() {
gpu_buf.gpu.insert_batch(states, next_states, actions, rewards, dones)
} else {
Ok(())
}
})?;
agent.get_volatile_head().map(|h| &h.memory).map_or(Ok(()), |m| {
if let Some(mut gpu_buf) = m.as_gpu_buffer() {
gpu_buf.gpu.insert_batch(states, next_states, actions, rewards, dones)
} else {
Ok(())
}
})?;
// CPU fallback converts tensors→experiences once, shared across heads.
let mut cpu_experiences: Option<Vec<Experience>> = None;
macro_rules! insert_head {
($getter:ident) => {
if let Some(head) = agent.$getter() {
if let Some(mut gpu_buf) = head.memory.as_gpu_buffer() {
gpu_buf.gpu.insert_batch(states, next_states, actions, rewards, dones)?;
} else {
if cpu_experiences.is_none() {
cpu_experiences = Some(Self::tensors_to_experiences(
states, next_states, actions, rewards, dones,
)?);
}
if let Some(ref exps) = cpu_experiences {
head.memory.add_batch(exps.clone())?;
}
}
}
};
}
insert_head!(get_trending_head);
insert_head!(get_ranging_head);
insert_head!(get_volatile_head);
Ok(())
}
}