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

fix: atomicExch for seg_tree leaf writes prevents duplicate-index race

Browse Source 
PER samples with replacement, so a batch can contain duplicate indices.
With non-atomic leaf writes, two threads processing the same index both
read the same old_pa, but only one write survives — internal nodes get
double the delta while the leaf changed once, permanently corrupting
the tree sum.

atomicExch returns the exact replaced value so each thread propagates
the correct delta regardless of concurrent duplicate writes.

Also fix readback_pinned layout doc: slot 9 is used for diversity loss.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-06 12:18:40 +02:00
parent 57c6fdcdd3
commit 22f30fe37f
2 changed files with 4 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
crates/ml-dqn/src/seg_tree_kernel.cu
View File

@@ -56,8 +56,7 @@ extern "C" __global__ void seg_tree_update(
// ancestor. atomicAdd is commutative+associative, so concurrent threads
// produce correct sums without synchronization barriers.
int leaf = capacity + (int)idx;
float old_pa = tree[leaf];
tree[leaf] = pa;
float old_pa = atomicExch(&tree[leaf], pa);
float delta = pa - old_pa;
int node = leaf >> 1;
@@ -90,8 +89,7 @@ extern "C" __global__ void seg_tree_insert(
// Write leaf and propagate delta to root via atomicAdd (race-free).
int leaf = capacity + (int)idx;
float old_pa = tree[leaf];
tree[leaf] = pa;
float old_pa = atomicExch(&tree[leaf], pa);
float delta = pa - old_pa;
int node = leaf >> 1;

3
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -636,7 +636,8 @@ pub struct GpuDqnTrainer {
/// [0]=loss, [1]=mse_loss, [2]=grad_norm_sq (per-step, from graph_adam)
/// [3]=avg_max_q, [4]=q_min, [5]=q_max, [6]=q_mean, [7]=q_var (every 50 steps)
/// [8]=causal_mean_sens (every N steps)
/// [9..16]=reserved (nan_flags, future use)
/// [9]=ensemble_diversity_loss (epoch boundary)
/// [10..16]=reserved (future use)
/// Pinned memory enables true async cuMemcpyDtoHAsync without CPU blocking.
readback_pinned: *mut f32,
/// Whether there's an in-flight readback to collect