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

fix: null pointer crash in increment_step_counters + wire PER rng_step

Browse Source 
Root cause: rng_step was passed as 0u64 (null) to increment_step_counters
kernel, causing atomicAdd on address 0 → CUDA_ERROR_ILLEGAL_ADDRESS →
cascading CUBLAS_STATUS_EXECUTION_FAILED on all subsequent operations.

Fixes:
- Replace all atomicAdd with plain writes (single-thread kernel)
- Add null guards for optional pointers (iqn_t, attn_t, rng_step)
- Allocate pinned device-mapped rng_step in GpuReplayBuffer
- Wire rng_step_dev_ptr from replay buffer to FusedTrainingCtx
- Remove stale atomicAdd comment

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-19 01:36:37 +02:00
parent 1d3a533f99
commit 56c5458008
5 changed files with 43 additions and 4755 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -174,6 +174,10 @@ pub struct GpuReplayBuffer {
insert_scratch_cap: usize, // current allocation size
scratch_f32: CudaSlice<f32>, // [1] reusable temp for flush/apply_max
rng_step: u32,
/// Pinned device-mapped RNG step counter — GPU writes, host reads.
/// Wrapped in usize for Send/Sync (raw ptr cast at usage site).
rng_step_pinned: usize,
rng_step_dev_ptr: u64,
/// Size of the most recent `sample_proportional` call. Used by `sample_weights_ref` /
/// `sample_indices_ref` to return correctly-sized views (pre-allocated buffers are
/// `max_batch_size` large; only the first `last_batch_size` elements are valid).
@@ -263,6 +267,16 @@ impl GpuReplayBuffer {
let insert_ep_buf = a32i(stream, isc, "ins_ep")?;
let scratch_f32 = a32f(stream, 1, "scratch")?;
// Pinned device-mapped RNG step counter for GPU-side increment
let (rng_step_pinned, rng_step_dev_ptr) = unsafe {
let mut host_ptr: *mut std::ffi::c_void = std::ptr::null_mut();
let mut dev_ptr: u64 = 0;
cudarc::driver::sys::cuMemAllocHost_v2(&mut host_ptr, std::mem::size_of::<i32>());
cudarc::driver::sys::cuMemHostGetDevicePointer_v2(&mut dev_ptr, host_ptr, 0);
*(host_ptr as *mut i32) = 0;
(host_ptr as usize, dev_ptr) // usize for Send/Sync
};
Ok(Self {
config, stream: Arc::clone(stream), kernels: k,
states: s, next_states: ns, actions: a, rewards: r, dones: d, priorities: p,
@@ -282,6 +296,8 @@ impl GpuReplayBuffer {
insert_prio_buf, insert_idx_buf, insert_ep_buf, insert_scratch_cap: isc,
scratch_f32,
rng_step: 0,
rng_step_pinned,
rng_step_dev_ptr,
last_batch_size: 0,
trainer_states_ptr: 0,
trainer_next_states_ptr: 0,
@@ -338,6 +354,9 @@ impl GpuReplayBuffer {
self.trainer_state_dim_padded = state_dim_padded;
}
/// Pinned device-mapped RNG step counter pointer for GPU-side increment.
pub fn rng_step_dev_ptr(&self) -> u64 { self.rng_step_dev_ptr }
/// Grow insert scratch buffers if `eff` exceeds current allocation.
/// Only triggers cuMemAlloc on the first large insert — subsequent calls reuse.
fn ensure_insert_scratch(&mut self, eff: usize) -> Result<(), MLError> {

17
crates/ml/src/cuda_pipeline/graph_utility_kernels.cu
View File

@@ -131,14 +131,15 @@ extern "C" __global__ void increment_step_counters(
{
/* adam_t is the master step — its post-increment value drives the
* cosine schedule so that step 1 produces the first non-initial tau. */
int step = atomicAdd(adam_t, 1) + 1;
atomicAdd(sel_t, 1);
atomicAdd(denoise_t, 1);
atomicAdd(iql_t, 1);
atomicAdd(iql_low_t, 1);
atomicAdd(iqn_t, 1);
atomicAdd(attn_t, 1);
atomicAdd(rng_step, 1);
// Single thread — plain writes, no atomicAdd needed
int step = (*adam_t += 1);
*sel_t += 1;
*denoise_t += 1;
*iql_t += 1;
*iql_low_t += 1;
if (iqn_t) *iqn_t += 1;
if (attn_t) *attn_t += 1;
if (rng_step) *rng_step += 1;
float progress = fminf((float)step / (float)anneal_steps, 1.0f);
float cosine = 0.5f * (1.0f + cosf(progress * 3.14159265f));

10
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -292,6 +292,8 @@ pub(crate) struct FusedTrainingCtx {
maintenance_child: Option<ChildGraph>,
iql_modulate_child: Option<ChildGraph>,
per_priority_child: Option<ChildGraph>,
/// PER RNG step counter (pinned device-mapped, GPU-side increment).
per_rng_step_dev_ptr: u64,
/// Composed parent -- single launch replays all children sequentially.
parent_graph: Option<ParentGraph>,
/// Eval-only forward graph exec for deterministic Q-value replay.
@@ -768,6 +770,7 @@ impl FusedTrainingCtx {
maintenance_child: None,
iql_modulate_child: None,
per_priority_child: None,
per_rng_step_dev_ptr: 0,
parent_graph: None,
eval_forward_exec: None,
tau_host: 0.0,
@@ -1437,7 +1440,7 @@ impl FusedTrainingCtx {
let iql_low_t = self.gpu_iql_low.t_dev_ptr();
let iqn_t = self.gpu_iqn.as_ref().map(|i| i.t_dev_ptr()).unwrap_or(0);
let attn_t = self.gpu_attention.as_ref().map(|a| a.t_dev_ptr()).unwrap_or(0);
let rng_step = 0u64; // PER rng_step dev_ptr — 0 means no-op inside kernel
let rng_step = self.per_rng_step_dev_ptr;
self.trainer.submit_counter_increments(iql_t, iql_low_t, iqn_t, attn_t, rng_step)
.map_err(|e| anyhow::anyhow!("counters: {e}"))?;
// Stochastic depth mask (already a GPU kernel)
@@ -2237,6 +2240,11 @@ impl FusedTrainingCtx {
self.trainer.state_dim_padded()
}
/// Set the PER RNG step counter pointer (pinned device-mapped).
pub(crate) fn set_per_rng_step_dev_ptr(&mut self, ptr: u64) {
self.per_rng_step_dev_ptr = ptr;
}
/// Run cross-branch Q attention: q_out_buf → q_coord_buf [batch_size, 12].
pub(crate) fn launch_q_attention(&self, batch_size: usize) -> Result<()> {
self.trainer.launch_q_attention(batch_size)

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

@@ -328,7 +328,7 @@ impl DQNTrainer {
}
drop(agent);
// Wire direct-to-trainer gather: PER gather writes directly to trainer buffers.
if let Some(ref fused) = self.fused_ctx {
if let Some(ref mut fused) = self.fused_ctx {
let mut agent_w = self.agent.write().await;
agent_w.memory_mut().gpu.set_trainer_buffers(
fused.trainer_states_buf_ptr(),
@@ -339,6 +339,7 @@ impl DQNTrainer {
fused.trainer_is_weights_buf_ptr(),
fused.trainer_state_dim_padded(),
);
fused.set_per_rng_step_dev_ptr(agent_w.memory().gpu.rng_step_dev_ptr());
}
}
}
@@ -689,7 +690,7 @@ impl DQNTrainer {
}
// Apply E5: ensemble agreement threshold
if let Some(ref fused) = self.fused_ctx {
if let Some(ref mut fused) = self.fused_ctx {
fused.set_var_ema(result.agreement_threshold);
}
}
@@ -1448,7 +1449,7 @@ impl DQNTrainer {
}
drop(agent);
// Wire direct-to-trainer gather: PER gather writes directly to trainer buffers.
if let Some(ref fused) = self.fused_ctx {
if let Some(ref mut fused) = self.fused_ctx {
let mut agent_w = self.agent.write().await;
agent_w.memory_mut().gpu.set_trainer_buffers(
fused.trainer_states_buf_ptr(),
@@ -1459,6 +1460,7 @@ impl DQNTrainer {
fused.trainer_is_weights_buf_ptr(),
fused.trainer_state_dim_padded(),
);
fused.set_per_rng_step_dev_ptr(agent_w.memory().gpu.rng_step_dev_ptr());
}
}
}
@@ -1559,7 +1561,7 @@ impl DQNTrainer {
!guard_past_warmup,
).map_err(|e| anyhow::anyhow!("guard check: {e}"))?;
if gr.halt_nan {
if let Some(ref fused) = self.fused_ctx {
if let Some(ref mut fused) = self.fused_ctx {
if let Ok(flags) = fused.read_nan_flags() {
let names = ["STATES_bf16", "on_v_logits", "on_b_logits", "mse_loss", "f32_params_PRE", "bf16_params_PRE", "grad_buf", "save_probs_bf16"];
let flagged: Vec<_> = flags.iter().enumerate()

4742
pal_precommit.changeset
View File

File diff suppressed because it is too large Load Diff