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

fix(cuda): remove atomicAdd violations in ppo/dqn loss accumulators

Browse Source 
Three atomicAdd sites violated feedback_no_atomicadd (deferred Phase C/E
fixes per file headers):
- ppo_clipped_surrogate.cu lines 270-271: atomicAdd(loss_pi/loss_entropy)
- dqn_distributional_q.cu line 285: atomicAdd(loss_out, ce/B)

Float addition is non-associative; atomicAdd across blocks produces
order-dependent sums. When inputs were extreme (l_v=6+ early in training
driving wide PPO ratios), the order-dependence flipped finite→NaN
non-deterministically — same SHA + same seed + same params produced
different NaN outcomes (smoke bisect 2026-05-29).

Fix is structural per the codebase pattern (compute_advantage_rms,
rl_q_bias_correction): per-batch sole-writer outputs + a deterministic
single-warp block-tree reducer. New kernel `mean_b_reduce.cu` performs
the [B]→[1] mean via grid-stride loop + warp-shuffle reduce.

Changes:
- cuda/mean_b_reduce.cu (NEW): single-block single-warp deterministic
  mean reducer mirroring compute_advantage_rms.cu pattern
- cuda/dqn_distributional_q.cu: remove atomicAdd; loss_per_batch[]
  already sole-writer, caller invokes mean_b_reduce after
- cuda/ppo_clipped_surrogate.cu: replace `loss_pi`/`loss_entropy` scalar
  args with `l_pi_per_batch`/`l_ent_per_batch` per-batch sole-writer
  outputs; caller invokes mean_b_reduce twice after kernel
- build.rs: register mean_b_reduce cubin
- src/rl/dqn.rs: DqnHead loads mean_b_reduce_fn; backward_logits invokes
  it on loss_per_batch → loss_out_dev_ptr
- src/rl/ppo.rs: PolicyHead loads mean_b_reduce_fn; surrogate_forward
  signature gains l_pi_per_batch + l_ent_per_batch buffer args; invokes
  mean_b_reduce on each → scalar dev ptrs
- src/trainer/integrated.rs: IntegratedTrainer gains
  ss_pi_l_pi_per_batch_d + ss_pi_l_ent_per_batch_d scratch fields;
  allocated at b_size; passed to surrogate_forward

Smoke status: l_q/l_v values now reproducible bit-for-bit across runs
(verified by comparing two runs of the same seed). NaN at step 4
still reproduces — proximate cause is elsewhere (likely in V head
forward dynamics, not the loss summation). The atomicAdd removal is
still required: it was a real correctness violation independent of the
NaN symptom, and the deterministic loss values are now a precondition
for diagnosing the remaining instability.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-29 01:24:21 +02:00
parent db4d9a16ff
commit 10d4614fb4
7 changed files with 242 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/ml-alpha/build.rs
View File

@@ -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

18
crates/ml-alpha/cuda/dqn_distributional_q.cu
View File

@@ -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];
}
}

59
crates/ml-alpha/cuda/mean_b_reduce.cu Normal file
View File

@@ -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;
}
}

52
crates/ml-alpha/cuda/ppo_clipped_surrogate.cu
View File

@@ -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;
}
}
}

77
crates/ml-alpha/src/rl/dqn.rs
View File

@@ -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<CudaModule>,
/// 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<CudaModule>,
/// 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(())
}

61
crates/ml-alpha/src/rl/ppo.rs
View File

@@ -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<f32>,
@@ -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<f32>,
@@ -242,6 +265,13 @@ impl PolicyHead {
b_size: usize,
pi_log_prob: &mut CudaSlice<f32>,
entropy: &mut CudaSlice<f32>,
// 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<f32>,
l_ent_per_batch: &mut CudaSlice<f32>,
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(())
}

19
crates/ml-alpha/src/trainer/integrated.rs
View File

@@ -1185,6 +1185,16 @@ pub struct IntegratedTrainer {
pub ss_v_loss_per_batch_d: CudaSlice<f32>,
pub ss_pi_log_prob_d: CudaSlice<f32>,
pub ss_entropy_d: CudaSlice<f32>,
/// 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<f32>,
/// 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<f32>,
// 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::<f32>(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::<f32>(b_size)
.context("alloc ss_pi_l_pi_per_batch_d")?;
let ss_pi_l_ent_per_batch_d = stream.alloc_zeros::<f32>(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,
)