diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 44a477743..b1fe3a42f 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -42,6 +42,7 @@ const KERNELS: &[&str] = &[ "grad_h_accumulate", // RL Phase E.2: element-wise grad_h_encoder += λ × grad_h_head accumulator (one head at a time, serialised by stream) "bellman_target_projection", // RL Phase E.2-DEFER: C51 categorical projection of Bellman target Z(s_{t+1}, a*) onto the discrete support, reads γ from ISV[400]; replaces host-side build_synthetic_bellman_target stand-in "bellman_scalar_target_project", // Phase 2.1 (2026-05-28): Dueling-DQN scalar Bellman target y = r + γ × V_target(s_{t+1}); projects (y − V(s_t)) advantage onto action's atom support; replaces fused_select_and_project_bellman in the trainer's Bellman build per spec §3 Sub-phase 2.1 + "mean_b_reduce", // 2026-05-29: deterministic [B]→[1] mean reducer replacing atomicAdd in dqn_distributional_q (l_q) and ppo_clipped_surrogate (l_pi, l_ent); single block + warp-shuffle per feedback_no_atomicadd "rl_lr_controller", // RL Phase E.2-DEFER: per-head learning-rate ISV emitter — bootstraps ISV[412..417] with 1e-3 (BCE/Q/π/V/aux); replaces hardcoded PHASE_E2_DEFAULT_LR "rl_rollout_steps_controller", // RL Phase E.3b: rollout-buffer-length ISV emitter — emits ISV[RL_N_ROLLOUT_STEPS_INDEX=404] from var(advantage)/|mean A| EMA; bootstraps 2048 "rl_per_alpha_controller", // RL Phase E.3b: PER priority-exponent ISV emitter — emits ISV[RL_PER_ALPHA_INDEX=405] from TD-error kurtosis EMA; bootstraps 0.6 diff --git a/crates/ml-alpha/cuda/dqn_distributional_q.cu b/crates/ml-alpha/cuda/dqn_distributional_q.cu index 9cfd373b4..11a17c529 100644 --- a/crates/ml-alpha/cuda/dqn_distributional_q.cu +++ b/crates/ml-alpha/cuda/dqn_distributional_q.cu @@ -275,14 +275,18 @@ extern "C" __global__ void dqn_distributional_q_bwd( __syncthreads(); } - // Thread 0 writes the per-batch CE and accumulates into loss_out. - // The cross-batch atomicAdd on loss_out is the Phase C deferred - // fix (see ATOMIC NOTE in the file header). grad_logits is never - // atomicAdded. + // Thread 0 writes the per-batch CE — sole writer per batch index. + // 2026-05-29: cross-batch reduction now done by `mean_b_reduce` + // kernel (deterministic block-tree-reduce on `loss_per_batch[]`) + // instead of atomicAdd into `loss_out`, per `feedback_no_atomicadd` + // and the smoke-bisect that showed atomicAdd order-dependence + // flipped finite→NaN non-deterministically. Caller MUST invoke + // `mean_b_reduce(loss_per_batch, B, loss_out)` after this kernel + // returns. `loss_out` parameter retained for ABI stability but + // unused here (zeroed by caller, written by mean_b_reduce). + (void)loss_out; if (tid == 0) { - const float ce = s_ce[0]; - loss_per_batch[batch] = ce; - atomicAdd(loss_out, ce / (float)B); + loss_per_batch[batch] = s_ce[0]; } } diff --git a/crates/ml-alpha/cuda/mean_b_reduce.cu b/crates/ml-alpha/cuda/mean_b_reduce.cu new file mode 100644 index 000000000..0eeccde3e --- /dev/null +++ b/crates/ml-alpha/cuda/mean_b_reduce.cu @@ -0,0 +1,59 @@ +// mean_b_reduce.cu — deterministic mean of [B] floats → scalar [1]. +// +// Replaces cross-block atomicAdd accumulation in ppo_clipped_surrogate +// (loss_pi, loss_entropy) and dqn_distributional_q (loss_out) which +// violated `feedback_no_atomicadd.md`. Float addition is non-associative +// so atomicAdd across blocks produced order-dependent sums, and when +// inputs were extreme (e.g. l_v=6+ early in training driving wide PPO +// ratios) the order-dependence flipped finite→NaN non-deterministically. +// +// Same SHA + same seed + same params → different NaN outcomes was the +// smoking gun (session 2026-05-29). The fix is structural: per-batch +// sole-writer outputs (already present in the violator kernels as +// `loss_per_batch[B]`, `entropy[B]`, `pi_log_prob[B]`) plus this +// deterministic reducer. No atomicAdd anywhere in the loss path. +// +// Layout: single block, single warp (32 threads), grid-stride loop + +// warp-shuffle reduce — mirrors `compute_advantage_rms.cu` exactly so +// the codebase idiom stays single-source-of-truth. +// +// Inputs: +// per_batch [B] — per-batch scalar values to average +// B int — batch size (used as divisor) +// Outputs: +// out_mean [1] — sum(per_batch) / B +// +// Launch contract (caller responsibility): +// grid_dim = (1, 1, 1) +// block_dim = (32, 1, 1) ← single warp; matches compute_advantage_rms +// smem = 0 ← warp-shuffle, no shared mem needed + +extern "C" __global__ void mean_b_reduce( + const float* __restrict__ per_batch, // [B] + int B, + float* __restrict__ out_mean // [1] OVERWRITE +) { + // Single block guard (grid is (1,1,1) by launch contract). + if (blockIdx.x != 0) return; + const int tid = threadIdx.x; + const int nt = blockDim.x; + + // Each thread accumulates partial over a strided slice of B. + float partial = 0.0f; + for (int b = tid; b < B; b += nt) { + partial += per_batch[b]; + } + + // Warp-shuffle reduce within the 32-lane warp. + // No __syncthreads needed: the shfl_xor_sync masks all 32 lanes. + for (int off = 16; off > 0; off >>= 1) { + partial += __shfl_xor_sync(0xffffffff, partial, off); + } + + // Lane 0 writes the mean. Guard against B==0 to avoid div-by-zero; + // callers always pass B>0 in practice but defensive is cheap. + if (tid == 0) { + const float denom = (B > 0) ? (float)B : 1.0f; + out_mean[0] = partial / denom; + } +} diff --git a/crates/ml-alpha/cuda/ppo_clipped_surrogate.cu b/crates/ml-alpha/cuda/ppo_clipped_surrogate.cu index 9e5766451..4f0f824ce 100644 --- a/crates/ml-alpha/cuda/ppo_clipped_surrogate.cu +++ b/crates/ml-alpha/cuda/ppo_clipped_surrogate.cu @@ -142,23 +142,30 @@ extern "C" __global__ void ppo_policy_logits_fwd( // isv [≥ 404] — reads ε at 402, coef at 403 // B — batch size // Outputs: -// pi_log_prob [B] — log π_new(a_taken|s) for diagnostics -// and Phase E KL EMA -// entropy [B] — H(π_new) per batch sample -// loss_pi [1] — Σ_b L_π (ATOMIC; see header) -// loss_entropy [1] — Σ_b L_entropy +// pi_log_prob [B] — log π_new(a_taken|s) for diagnostics + KL EMA +// entropy [B] — H(π_new) per batch sample +// l_pi_per_batch [B] — per-batch L_π (sole writer; caller reduces +// via `mean_b_reduce` per feedback_no_atomicadd) +// l_ent_per_batch [B] — per-batch L_entropy (sole writer) +// +// 2026-05-29: cross-batch reduction moved out of this kernel; the +// `loss_pi`/`loss_entropy` scalars are now written by `mean_b_reduce` +// invoked by the caller AFTER this kernel returns. Previous design used +// atomicAdd into single floats — non-deterministic across warp execution +// orders, which surfaced as flip-flopping NaN at step 4 when V was wild +// early in training. Per-batch sole-writer is the canonical fix. // ───────────────────────────────────────────────────────────────────── extern "C" __global__ void ppo_clipped_surrogate_fwd( - const float* __restrict__ logits, // [B * N_ACTIONS] - const float* __restrict__ log_pi_old, // [B] - const int* __restrict__ actions, // [B] - const float* __restrict__ advantages, // [B] - const float* __restrict__ isv, // [>= 404] + const float* __restrict__ logits, // [B * N_ACTIONS] + const float* __restrict__ log_pi_old, // [B] + const int* __restrict__ actions, // [B] + const float* __restrict__ advantages, // [B] + const float* __restrict__ isv, // [>= 404] int B, - float* __restrict__ pi_log_prob, // [B] - float* __restrict__ entropy, // [B] - float* __restrict__ loss_pi, // [1] - float* __restrict__ loss_entropy // [1] + float* __restrict__ pi_log_prob, // [B] + float* __restrict__ entropy, // [B] + float* __restrict__ l_pi_per_batch, // [B] sole writer + float* __restrict__ l_ent_per_batch // [B] sole writer ) { const int batch = blockIdx.x; const int act = threadIdx.x; @@ -264,14 +271,19 @@ extern "C" __global__ void ppo_clipped_surrogate_fwd( const float coef = isv[RL_ENTROPY_COEF_INDEX]; const float l_ent = -coef * h; - // ATOMIC NOTE: see header. Phase E warp-shuffles these out. - // /B for batch-invariant loss reporting. - const float b_inv = 1.0f / (float)B; - atomicAdd(loss_pi, l_pi * b_inv); - atomicAdd(loss_entropy, l_ent * b_inv); + // 2026-05-29: sole-writer per-batch outputs replace the + // prior atomicAdd into single-float accumulators. The caller + // invokes `mean_b_reduce(l_pi_per_batch, B, loss_pi)` and + // `mean_b_reduce(l_ent_per_batch, B, loss_entropy)` after + // this kernel returns to compute the deterministic mean. + // Per `feedback_no_atomicadd`. + l_pi_per_batch[batch] = l_pi; + l_ent_per_batch[batch] = l_ent; } else { // Invalid action; leave diagnostics at 0, skip loss accum. - pi_log_prob[batch] = 0.0f; + pi_log_prob[batch] = 0.0f; + l_pi_per_batch[batch] = 0.0f; + l_ent_per_batch[batch] = 0.0f; } } } diff --git a/crates/ml-alpha/src/rl/dqn.rs b/crates/ml-alpha/src/rl/dqn.rs index 4b36799e3..fef49d51c 100644 --- a/crates/ml-alpha/src/rl/dqn.rs +++ b/crates/ml-alpha/src/rl/dqn.rs @@ -119,6 +119,13 @@ const BELLMAN_PROJ_CUBIN: &[u8] = include_bytes!(concat!( env!("OUT_DIR"), "/bellman_target_projection.cubin" )); +/// 2026-05-29: deterministic mean reducer to replace atomicAdd into +/// `loss_out` in `dqn_distributional_q_bwd`. See header of +/// `mean_b_reduce.cu` for rationale (no-atomicadd per codebase rule). +const MEAN_B_REDUCE_CUBIN: &[u8] = include_bytes!(concat!( + env!("OUT_DIR"), + "/mean_b_reduce.cubin" +)); const BELLMAN_SCALAR_CUBIN: &[u8] = include_bytes!(concat!( env!("OUT_DIR"), "/bellman_scalar_target_project.cubin" @@ -208,6 +215,12 @@ pub struct DqnHead { /// `softmax(A_dist(s_t, a_t))` and this projected target trains A. pub bellman_scalar_target_fn: CudaFunction, _bellman_scalar_module: Arc, + /// 2026-05-29: deterministic mean reducer for cross-batch loss + /// accumulation in `backward_logits`. Replaces atomicAdd into + /// `loss_out` (which was non-deterministic across warp execution + /// orders) per `feedback_no_atomicadd`. + pub mean_b_reduce_fn: CudaFunction, + _mean_b_reduce_module: Arc, /// Phase R5: element-wise target-net soft update kernel handle. /// `target[i] = (1 - τ) · target[i] + τ · current[i]` reading τ @@ -286,6 +299,13 @@ impl DqnHead { let bellman_scalar_target_fn = bellman_scalar_module .load_function("bellman_scalar_target_project") .context("load bellman_scalar_target_project")?; + // 2026-05-29: deterministic mean reducer (no-atomicAdd fix). + let mean_b_reduce_module = ctx + .load_cubin(MEAN_B_REDUCE_CUBIN.to_vec()) + .context("load mean_b_reduce cubin (DqnHead)")?; + let mean_b_reduce_fn = mean_b_reduce_module + .load_function("mean_b_reduce") + .context("load mean_b_reduce (DqnHead)")?; // Phase R5: target soft-update kernel. let target_soft_update_module = ctx @@ -358,6 +378,8 @@ impl DqnHead { _bellman_module: bellman_module, bellman_scalar_target_fn, _bellman_scalar_module: bellman_scalar_module, + mean_b_reduce_fn, + _mean_b_reduce_module: mean_b_reduce_module, target_soft_update_fn, _target_soft_update_module: target_soft_update_module, w_d, @@ -515,22 +537,45 @@ impl DqnHead { debug_assert_eq!(grad_logits.len(), b_size * N_ACTIONS * Q_N_ATOMS); let b_i = b_size as i32; - let mut args = RawArgs::new(); - args.push_ptr(logits.raw_ptr()); - args.push_ptr(target_dist.raw_ptr()); - args.push_ptr(actions_taken.raw_ptr()); - args.push_i32(b_i); - args.push_ptr(*loss_out_dev_ptr); - args.push_ptr(loss_per_batch.raw_ptr()); - args.push_ptr(grad_logits.raw_ptr()); - let mut ptrs = args.build_arg_ptrs(); - unsafe { - raw_launch( - self.bwd_fn.cu_function(), - (b_size as u32, 1, 1), (BWD_BLOCK_SIZE, 1, 1), 0, - self.raw_stream, - &mut ptrs[..args.len()], - ).map_err(|e| anyhow::anyhow!("dqn_distributional_q_bwd: {:?}", e))?; + { + let mut args = RawArgs::new(); + args.push_ptr(logits.raw_ptr()); + args.push_ptr(target_dist.raw_ptr()); + args.push_ptr(actions_taken.raw_ptr()); + args.push_i32(b_i); + args.push_ptr(*loss_out_dev_ptr); + args.push_ptr(loss_per_batch.raw_ptr()); + args.push_ptr(grad_logits.raw_ptr()); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + self.bwd_fn.cu_function(), + (b_size as u32, 1, 1), (BWD_BLOCK_SIZE, 1, 1), 0, + self.raw_stream, + &mut ptrs[..args.len()], + ).map_err(|e| anyhow::anyhow!("dqn_distributional_q_bwd: {:?}", e))?; + } + } + + // 2026-05-29: deterministic cross-batch mean reduce on + // `loss_per_batch[B]` → `loss_out_dev_ptr` scalar. Replaces the + // prior atomicAdd in `dqn_distributional_q_bwd` (the kernel arg + // `loss_out` is still passed but now unused; caller zeros it). + // Per `feedback_no_atomicadd`. + { + let mut args = RawArgs::new(); + args.push_ptr(loss_per_batch.raw_ptr()); + args.push_i32(b_i); + args.push_ptr(*loss_out_dev_ptr); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + self.mean_b_reduce_fn.cu_function(), + (1, 1, 1), (32, 1, 1), 0, + self.raw_stream, + &mut ptrs[..args.len()], + ).map_err(|e| anyhow::anyhow!("mean_b_reduce (dqn loss_q): {:?}", e))?; + } } Ok(()) } diff --git a/crates/ml-alpha/src/rl/ppo.rs b/crates/ml-alpha/src/rl/ppo.rs index 8863b3900..5277646f1 100644 --- a/crates/ml-alpha/src/rl/ppo.rs +++ b/crates/ml-alpha/src/rl/ppo.rs @@ -54,6 +54,11 @@ const PPO_SURR_CUBIN: &[u8] = include_bytes!(concat!( )); const V_HEAD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/v_head_fwd_bwd.cubin")); +/// 2026-05-29: deterministic [B]→[1] mean reducer that replaces atomicAdd +/// in the PPO surrogate cross-batch loss accumulation. Single-warp +/// block-tree-reduce per `feedback_no_atomicadd`. +const MEAN_B_REDUCE_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/mean_b_reduce.cubin")); /// Construction config for [`PolicyHead`] and [`ValueHead`]. #[derive(Clone, Debug)] @@ -107,6 +112,13 @@ pub struct PolicyHead { /// caller's responsibility via `reduce_axis0`. pub grad_w_b_h_t_fn: CudaFunction, + /// 2026-05-29: deterministic mean reducer used in `surrogate_forward` + /// to compute `loss_pi` and `loss_entropy` scalars from per-batch + /// sole-writer scratch buffers, replacing the prior atomicAdd into + /// `loss_pi`/`loss_entropy` scalars (non-deterministic across warp + /// execution orders — see header of `ppo_clipped_surrogate.cu`). + pub mean_b_reduce_fn: CudaFunction, + /// Online weights `[N_ACTIONS, HIDDEN_DIM]`, row-major. Each row is /// one action's projection vector. pub w_d: CudaSlice, @@ -140,6 +152,15 @@ impl PolicyHead { let grad_w_b_h_t_fn = module .load_function("ppo_grad_w_b_h_t") .context("load ppo_grad_w_b_h_t")?; + // 2026-05-29: load the deterministic mean reducer cubin. + let mean_b_module = ctx + .load_cubin(MEAN_B_REDUCE_CUBIN.to_vec()) + .context("load mean_b_reduce cubin (PolicyHead)")?; + let mean_b_reduce_fn = mean_b_module + .load_function("mean_b_reduce") + .context("load mean_b_reduce (PolicyHead)")?; + // Module is dropped here; the `CudaFunction` keeps the cubin + // alive by holding a strong ref to it (cudarc pattern). // Per pearl_scoped_init_seed_for_reproducibility: install the // scoped seed guard BEFORE drawing any Xavier samples. @@ -167,6 +188,7 @@ impl PolicyHead { surrogate_bwd_fn, policy_linear_fwd_fn, grad_w_b_h_t_fn, + mean_b_reduce_fn, w_d, b_d, }) @@ -232,6 +254,7 @@ impl PolicyHead { /// longer accepts `returns` / `v_pred` / `loss_v` (greenfield removal /// 2026-05-22). #[allow(clippy::too_many_arguments)] + #[allow(clippy::too_many_arguments)] pub fn surrogate_forward( &self, logits: &CudaSlice, @@ -242,6 +265,13 @@ impl PolicyHead { b_size: usize, pi_log_prob: &mut CudaSlice, entropy: &mut CudaSlice, + // 2026-05-29: per-batch sole-writer scratch buffers replace the + // prior atomicAdd-into-scalar pattern. The scalar outputs + // (`loss_pi_dev_ptr`, `loss_entropy_dev_ptr`) are now written by + // the `mean_b_reduce` kernel invoked at the end of this method. + // Per `feedback_no_atomicadd`. + l_pi_per_batch: &mut CudaSlice, + l_ent_per_batch: &mut CudaSlice, loss_pi_dev_ptr: &u64, loss_entropy_dev_ptr: &u64, ) -> Result<()> { @@ -251,6 +281,8 @@ impl PolicyHead { debug_assert_eq!(advantages.len(), b_size); debug_assert_eq!(pi_log_prob.len(), b_size); debug_assert_eq!(entropy.len(), b_size); + debug_assert_eq!(l_pi_per_batch.len(), b_size); + debug_assert_eq!(l_ent_per_batch.len(), b_size); let b_i = b_size as i32; { @@ -263,8 +295,8 @@ impl PolicyHead { args.push_i32(b_i); args.push_ptr(pi_log_prob.raw_ptr()); args.push_ptr(entropy.raw_ptr()); - args.push_ptr(*loss_pi_dev_ptr); - args.push_ptr(*loss_entropy_dev_ptr); + args.push_ptr(l_pi_per_batch.raw_ptr()); + args.push_ptr(l_ent_per_batch.raw_ptr()); let mut ptrs = args.build_arg_ptrs(); unsafe { raw_launch( @@ -278,6 +310,31 @@ impl PolicyHead { .map_err(|e| anyhow::anyhow!("ppo_clipped_surrogate_fwd: {:?}", e))?; } } + + // 2026-05-29: deterministic cross-batch reduce. Single-block, + // single-warp grid-stride reduce in `mean_b_reduce.cu`. Two + // launches — one per per-batch loss buffer. + for (per_batch_ptr, scalar_ptr) in [ + (l_pi_per_batch.raw_ptr(), *loss_pi_dev_ptr), + (l_ent_per_batch.raw_ptr(), *loss_entropy_dev_ptr), + ] { + let mut args = RawArgs::new(); + args.push_ptr(per_batch_ptr); + args.push_i32(b_i); + args.push_ptr(scalar_ptr); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + self.mean_b_reduce_fn.cu_function(), + (1, 1, 1), + (32, 1, 1), + 0, + self.raw_stream, + &mut ptrs[..args.len()], + ) + .map_err(|e| anyhow::anyhow!("mean_b_reduce: {:?}", e))?; + } + } Ok(()) } diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index c593b43f9..90f9d7e58 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -1185,6 +1185,16 @@ pub struct IntegratedTrainer { pub ss_v_loss_per_batch_d: CudaSlice, pub ss_pi_log_prob_d: CudaSlice, pub ss_entropy_d: CudaSlice, + /// 2026-05-29: per-batch L_π scratch — sole writer is + /// `ppo_clipped_surrogate_fwd`, reader is `mean_b_reduce` which + /// computes the deterministic cross-batch mean into + /// `ss_pi_loss_dev_ptr`. Replaces atomicAdd-into-scalar per + /// `feedback_no_atomicadd` (smoke bisect 2026-05-29 showed the + /// atomicAdd flipped finite→NaN non-deterministically). + pub ss_pi_l_pi_per_batch_d: CudaSlice, + /// 2026-05-29: per-batch L_entropy scratch — same pattern as + /// `ss_pi_l_pi_per_batch_d`. + pub ss_pi_l_ent_per_batch_d: CudaSlice, // Q-head gradient buffers. backward_gemm produces grad_w / grad_b // directly reduced (no per-batch scratch needed — cuBLAS SGEMM @@ -2296,6 +2306,11 @@ impl IntegratedTrainer { .context("alloc ss_pi_log_prob_d")?; let ss_entropy_d = stream.alloc_zeros::(b_size) .context("alloc ss_entropy_d")?; + // 2026-05-29: per-batch loss scratch buffers (no-atomicAdd fix). + let ss_pi_l_pi_per_batch_d = stream.alloc_zeros::(b_size) + .context("alloc ss_pi_l_pi_per_batch_d")?; + let ss_pi_l_ent_per_batch_d = stream.alloc_zeros::(b_size) + .context("alloc ss_pi_l_ent_per_batch_d")?; // Q-head gradient buffers. backward_gemm produces grad_w / grad_b // directly reduced via cuBLAS SGEMM — no per-batch scratch needed. @@ -2742,6 +2757,8 @@ impl IntegratedTrainer { ss_v_loss_per_batch_d, ss_pi_log_prob_d, ss_entropy_d, + ss_pi_l_pi_per_batch_d, + ss_pi_l_ent_per_batch_d, ss_q_grad_logits_d, ss_q_grad_h_t_d, ss_q_grad_w_d, @@ -4352,6 +4369,8 @@ impl IntegratedTrainer { b_size, &mut self.ss_pi_log_prob_d, &mut self.ss_entropy_d, + &mut self.ss_pi_l_pi_per_batch_d, + &mut self.ss_pi_l_ent_per_batch_d, &self.ss_pi_loss_dev_ptr, &self.ss_pi_loss_entropy_dev_ptr, )