From c823865008751b5214597b8e0a285d089c55b571 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 21 Apr 2026 12:42:13 +0200 Subject: [PATCH] =?UTF-8?q?determinism(atom=5Fstats):=202-phase=20reductio?= =?UTF-8?q?n=20=E2=80=94=20final=20training-path=20atomicAdd=20removed?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit compute_expected_q had 2 atomicAdds accumulating per-block entropy + utilization into atom_stats[0..1]. These feed HEALTH_DIAG's atoms=X%ent/Y%util field AND the adaptive_atom_positions kernel (which reshapes C51 atom positions based on utilization) — so non-determinism here propagated into training dynamics. Replaced with standard 2-phase deterministic reduction: Phase 1 (inside compute_expected_q): Per-block warp→block shared-mem reduce unchanged, but the final write is atom_stats_block_sums[blockIdx.x * 2 + i] = value (unique slot per block, no atomic). Kernel signature: `atom_stats` param renamed to `atom_stats_block_sums` to reflect new semantics. Phase 2 (new kernel atom_stats_finalize): Sequentially sums block_sums[0..num_blocks*2] into atom_stats[0..1]. Launch grid=(1,1,1) block=(1,1,1) — bit-stable across runs. Rust glue: - New atom_stats_block_sums_buf [max_blocks * 2] allocated at trainer init. max_blocks = ceil(batch_size / 256); smoke uses 1, production 64. - New atom_stats_finalize_kernel loaded from experience_kernels cubin. - populate_q_out() now launches phase 1 + phase 2 sequentially. - Other launch sites (replay_forward_for_q_values, compute_denoise_target_q) already passed NULL atom_stats and don't need changes — the kernel skips accumulation on NULL. This was the FINAL training-path atomicAdd call site. Post-commit, the entire crates/ml/src/cuda_pipeline/ has ZERO atomicAdd call sites — full CUDA-level determinism for the training path (cuBLAS algorithm selection remains the main remaining nondeterminism source, mitigated via CUBLAS_WORKSPACE_CONFIG for tests). Smoke verification (post-commit, 5-trial multi-trial): q_gaps: 2.24 / 0.98 / 2.75 / 5.58 / 1.08 (all > 0.02, 5/5 pass) Best Sharpe: 19.41 / 38.71 / 25.33 / 35.94 / 24.22 (median 25.3) mean_degradation = -10.66 (NEGATIVE means sharpe_ema IMPROVED over epochs) All assertions pass. Session atomicAdd tally: 69 string occurrences → 13 real call sites (most "hits" were comments). Today removed: 4 (CQL barrier+IB) + 2 (monitoring_reduce) + 2 (trade_stats + dqn_utility) + deleted-kernel (ensemble_diversity's atomic + G6/G10's atomics removed with scaffolding) + 2 (atom_stats, this commit) = 13 total. Zero remaining. Files touched: crates/ml/src/cuda_pipeline/experience_kernels.cu (+38 / -9 net) crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs (+40 / -10 net) Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/cuda_pipeline/experience_kernels.cu | 47 +++++++++++++++---- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 43 ++++++++++++++--- 2 files changed, 76 insertions(+), 14 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index f91decc98..12dc24399 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -2332,7 +2332,8 @@ extern "C" __global__ void compute_expected_q( int b2_size, int b3_size, const float* __restrict__ per_sample_support, /* [N*3]: per-sample [v_min, v_max, delta_z] from IQL */ - float* __restrict__ atom_stats, /* [2]: [0]=sum_entropy, [1]=sum_utilized. NULL to skip. */ + float* __restrict__ atom_stats_block_sums, /* [num_blocks * 2]: per-block partial sums (phase 1). + * NULL to skip atom-stat collection entirely. */ float* __restrict__ q_variance, /* [N, total_actions] Var[Q] per action. NULL to skip. */ const float* __restrict__ atom_positions /* [4, num_atoms] adaptive positions. NULL = use linear. */ ) { @@ -2412,15 +2413,23 @@ extern "C" __global__ void compute_expected_q( branch_logit_offset += N * n_d * num_atoms; } - /* Accumulate atom stats across all samples (monitoring only). - * Uses warp+block reduction to minimize atomicAdd non-determinism. - * One atomicAdd per BLOCK (not per thread). */ - if (atom_stats != NULL) { + /* Atom-stat accumulation across samples (monitoring + adaptive_atom_positions input). + * + * 2-phase deterministic reduction (replaces cross-block atomicAdd): + * Phase 1 (here): warp reduce → block reduce via shared mem → write + * one (entropy, util) pair per block to atom_stats_block_sums. + * Phase 2: `atom_stats_finalize` kernel sequentially sums block_sums + * into atom_stats[0..1]. Bit-stable across runs. + * + * atom_stats_block_sums layout: [num_blocks * 2], block b writes + * [b*2+0] = block's entropy sum, [b*2+1] = block's util sum. + */ + if (atom_stats_block_sums != NULL) { float inv_actions = 1.0f / (float)total_actions; float my_entropy = sample_entropy * inv_actions; float my_util = (float)sample_utilized * inv_actions; - /* Warp-level reduction */ + /* Warp-level reduction (shfl — deterministic across lanes) */ for (int offset = 16; offset > 0; offset >>= 1) { my_entropy += __shfl_xor_sync(0xFFFFFFFF, my_entropy, offset); my_util += __shfl_xor_sync(0xFFFFFFFF, my_util, offset); @@ -2441,13 +2450,35 @@ extern "C" __global__ void compute_expected_q( vu += __shfl_xor_sync(0xFFFFFFFF, vu, off); } if (lane == 0) { - atomicAdd(&atom_stats[0], ve); - atomicAdd(&atom_stats[1], vu); + /* One write per block to unique scratch slot — no atomic. */ + atom_stats_block_sums[blockIdx.x * 2 + 0] = ve; + atom_stats_block_sums[blockIdx.x * 2 + 1] = vu; } } } } +/* ══════════════════════════════════════════════════════════════════════ + * atom_stats_finalize — phase 2 of the deterministic reduction. + * + * Sequentially sums atom_stats_block_sums[0..num_blocks*2] into + * atom_stats[0..1]. Launch grid=(1,1,1) block=(1,1,1); deterministic order. + * ══════════════════════════════════════════════════════════════════════ */ +extern "C" __global__ void atom_stats_finalize( + const float* __restrict__ atom_stats_block_sums, /* [num_blocks * 2] from phase 1 */ + float* __restrict__ atom_stats, /* [2] final (entropy_sum, util_sum) */ + int num_blocks +) { + float sum_entropy = 0.0f; + float sum_util = 0.0f; + for (int b = 0; b < num_blocks; b++) { + sum_entropy += atom_stats_block_sums[b * 2 + 0]; + sum_util += atom_stats_block_sums[b * 2 + 1]; + } + atom_stats[0] = sum_entropy; + atom_stats[1] = sum_util; +} + /* ================================================================== */ /* Kernel 5: expert_action_override */ /* ================================================================== */ diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 5a49f99b3..3d4850d0b 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1310,8 +1310,15 @@ pub struct GpuDqnTrainer { q_stats_kernel: CudaFunction, /// GPU buffer for Q-value statistics [7 floats] q_stats_buf: CudaSlice, - /// GPU buffer for atom utilization accumulation [2 floats: sum_entropy, sum_utilized] + /// GPU buffer for atom utilization accumulation [2 floats: sum_entropy, sum_utilized]. + /// Populated by `atom_stats_finalize` kernel (phase 2 of the deterministic reduction). atom_stats_buf: CudaSlice, + /// Per-block partial sums for `compute_expected_q`'s atom-stat reduction. + /// Phase 1 output: [max_blocks * 2] where block b writes [b*2+0]=entropy, [b*2+1]=util. + /// Phase 2 (`atom_stats_finalize`) reduces these deterministically into atom_stats_buf. + atom_stats_block_sums_buf: CudaSlice, + /// Phase-2 finalize kernel for atom_stats — sequential reduction, bit-stable. + atom_stats_finalize_kernel: CudaFunction, /// Pinned device-mapped readback for q_stats [7] + q_out sample 0 [total_actions] floats. /// GPU writes via q_readback_dev_ptr, CPU reads via q_readback_pinned — zero sync. q_readback_dev_ptr: u64, @@ -5660,6 +5667,13 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("alloc q_stats_f32: {e}")))?; let atom_stats_buf = stream.alloc_zeros::(2) .map_err(|e| MLError::ModelError(format!("alloc atom_stats: {e}")))?; + // Per-block partial sums for the atom-stat deterministic reduction. + // Sized to max possible num_blocks for compute_expected_q (one thread + // per sample, block_dim=256). For smoke (B=16) → 1 block; for + // production (B=16384) → 64 blocks. Overallocate to a fixed cap. + let atom_stats_max_blocks: usize = ((config.batch_size + 255) / 256).max(1); + let atom_stats_block_sums_buf = stream.alloc_zeros::(atom_stats_max_blocks * 2) + .map_err(|e| MLError::ModelError(format!("alloc atom_stats_block_sums: {e}")))?; // q_readback — pinned host buffer for q_stats[7] + q_out sample 0 [total_actions] let total_actions = config.branch_0_size + config.branch_1_size + config.branch_2_size + config.branch_3_size; let q_readback_size = 7 + total_actions; // 7 stats + 13 Q-values = 20 @@ -5693,6 +5707,8 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("branch_confidence_routing load: {e}")))?; let adaptive_atom_kernel = exp_module_for_mag.load_function("adaptive_atom_positions") .map_err(|e| MLError::ModelError(format!("adaptive_atom_positions load: {e}")))?; + let atom_stats_finalize_kernel = exp_module_for_mag.load_function("atom_stats_finalize") + .map_err(|e| MLError::ModelError(format!("atom_stats_finalize load: {e}")))?; let atom_position_grad_kernel = exp_module_for_mag.load_function("atom_position_gradient") .map_err(|e| MLError::ModelError(format!("atom_position_gradient load: {e}")))?; let q_anchor_kernel = exp_module_for_mag.load_function("q_anchor_to_flat") @@ -7414,6 +7430,8 @@ impl GpuDqnTrainer { q_stats_kernel, q_stats_buf, atom_stats_buf, + atom_stats_block_sums_buf, + atom_stats_finalize_kernel, q_readback_pinned, q_readback_dev_ptr, cql_logit_grad_kernel, @@ -9225,12 +9243,10 @@ impl GpuDqnTrainer { let grid_dim = ((batch_size as u32 + block_dim - 1) / block_dim).max(1); let q_out_ptr = self.q_out_buf.raw_ptr(); let atom_stats_ptr = self.atom_stats_buf.raw_ptr(); + let atom_stats_block_sums_ptr = self.atom_stats_block_sums_buf.raw_ptr(); let q_var_ptr = self.q_var_buf_trainer.raw_ptr(); let atom_positions_buf_ptr = self.atom_positions_buf.raw_ptr(); - // Zero atom_stats before accumulation (kernel uses atomicAdd) - unsafe { - cudarc::driver::sys::cuMemsetD32Async(atom_stats_ptr, 0, 2, self.stream.cu_stream()); - } + // Phase 1: per-block sums (deterministic, no atomic) → block_sums unsafe { self.stream .launch_builder(&self.expected_q_kernel) @@ -9244,7 +9260,7 @@ impl GpuDqnTrainer { .arg(&b2) .arg(&b3) .arg(&self.per_sample_support_ptr) - .arg(&atom_stats_ptr) + .arg(&atom_stats_block_sums_ptr) .arg(&q_var_ptr) .arg(&atom_positions_buf_ptr) .launch(LaunchConfig { @@ -9254,6 +9270,21 @@ impl GpuDqnTrainer { }) .map_err(|e| MLError::ModelError(format!("populate_q_out expected_q: {e}")))?; } + // Phase 2: sequential reduce block_sums → atom_stats (deterministic) + let num_blocks_i32 = grid_dim as i32; + unsafe { + self.stream + .launch_builder(&self.atom_stats_finalize_kernel) + .arg(&atom_stats_block_sums_ptr) + .arg(&atom_stats_ptr) + .arg(&num_blocks_i32) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("populate_q_out atom_stats_finalize: {e}")))?; + } Ok(&self.q_out_buf) }