diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index 296b12f5b..a671073c2 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -1082,8 +1082,18 @@ impl GpuReplayBuffer { /// Update PER priorities from TD errors using the prefix-sum kernel. /// /// `per_update_pa` reads f32 td_errors directly (no cast needed), - /// computes `(|td|^alpha + eps)`, and writes to both `priorities` and - /// `priorities_pa`. Prefix sum is rebuilt lazily before sampling. + /// computes `priorities[i] = |td| + eps` (raw PER priority) and + /// `priorities_pa[i] = (|td| + eps)^alpha` (sampling weight). The + /// sum-tree (per_prefix_scan + per_sample) walks `priorities_pa`; + /// IS-weight normalises by the same column. Prefix sum is rebuilt + /// lazily before sampling. + /// + /// Class B Finding #2 fix (2026-05-08): prior implementation applied + /// alpha twice — `priorities = |td|^alpha + eps` then `priorities_pa + /// = priorities^alpha = (|td|^alpha + eps)^alpha ≈ |td|^(alpha²)`. + /// Effective sampling exponent was alpha² = 0.36 instead of the + /// configured alpha = 0.6 — high-TD events were sampled only ~2.5× + /// more than low-TD instead of ~10× as PER intends. /// /// If `ext_stream` is provided, all GPU work runs on that stream instead of /// the replay buffer's own stream. This avoids cross-stream synchronization @@ -1496,4 +1506,126 @@ mod tests { ); } } + + /// Class B Finding #2 (2026-05-08): PER priority double-exponent fix. + /// + /// Standard PER (Schaul et al. 2016): sampling weight `priorities_pa[i] + /// = (|td_i| + ε)^α`. Sum-tree (per_prefix_scan + per_sample) walks + /// `priorities_pa`, so the effective sampling exponent IS α. + /// + /// Pre-fix bug: `per_update_pa` stored `priorities[idx] = |td|^α + ε` + /// then `priorities_pa[idx] = priorities^α`, applying α twice. The + /// effective sampling exponent was α² instead of α — for α=0.6, that's + /// 0.36 instead of 0.6, flattening prioritization from a ~10× ratio + /// (high-TD vs low-TD) down to ~2.5×. + /// + /// This test inserts a buffer at max-priority, then fires + /// `update_priorities_gpu` with TD errors spanning four orders of + /// magnitude (0.001 → 100.0), and verifies the resulting + /// `priorities_pa` ratio matches `(|td_max|/|td_min|)^α` within + /// tolerance — i.e. the single-α formula. Pre-fix, the ratio would + /// have been `(|td_max|/|td_min|)^(α²) ≈ 100×` (off by 20×). + #[test] + #[ignore] // Requires CUDA GPU + fn test_per_priority_single_exponent_no_double_alpha() { + let stream = make_stream(); + let cap = 256usize; + let bs = 16usize; + let alpha = 0.6_f32; + let epsilon = 1e-6_f32; + let sd = ml_core::state_layout::STATE_DIM; + let c = GpuReplayBufferConfig { + capacity: cap, alpha, beta_start: 0.4, beta_max: 1.0, + beta_annealing_steps: 1000, epsilon, + max_memory_bytes: 4 << 30, max_batch_size: bs, + }; + let mut b = GpuReplayBuffer::new(c, &stream).expect("buf"); + + // 1. Seed `bs` slots with arbitrary state — all start at + // max_priority (1.0) per `per_insert_pa` semantics. + let s = stream.alloc_zeros::(bs * sd).unwrap(); + let ns = stream.alloc_zeros::(bs * sd).unwrap(); + let a = stream.alloc_zeros::(bs).unwrap(); + let r = stream.alloc_zeros::(bs).unwrap(); + let d = stream.alloc_zeros::(bs).unwrap(); + let aux = stream.alloc_zeros::(bs).unwrap(); + b.insert_batch(&s, &ns, &a, &r, &d, &aux, bs).unwrap(); + + // 2. Fire `update_priorities_gpu` with TD errors spanning [td_min, td_max]. + // Health stays at default 1.0, so the standard `per_update_pa` path + // is taken (NOT the diverse fallback). All bs slots get updated. + let td_min = 0.001_f32; + let td_max = 100.0_f32; + let td_host: Vec = (0..bs) + .map(|i| { + let t = i as f32 / (bs as f32 - 1.0); + // Geometric interpolation in log-space — endpoints land + // exactly on td_min, td_max. + td_min * (td_max / td_min).powf(t) + }) + .collect(); + let mut td_dev = stream.alloc_zeros::(bs).unwrap(); + stream.memcpy_htod(&td_host, &mut td_dev).unwrap(); + + // Indices 0..bs (all inserted slots). + let idx_host: Vec = (0..bs as u32).collect(); + let mut idx_dev = stream.alloc_zeros::(bs).unwrap(); + stream.memcpy_htod(&idx_host, &mut idx_dev).unwrap(); + + b.update_priorities_gpu(&idx_dev, &td_dev, bs, None) + .expect("update_priorities_gpu"); + + // 3. Read back priorities and priorities_pa — verify the + // single-α contract. + unsafe { + cudarc::driver::sys::cuStreamSynchronize(stream.cu_stream()); + } + let mut prio_host = vec![0.0_f32; bs]; + let mut prio_pa_host = vec![0.0_f32; bs]; + stream.memcpy_dtoh(&b.priorities_slice().slice(..bs), &mut prio_host).unwrap(); + stream.memcpy_dtoh(&b.priorities_pa_slice().slice(..bs), &mut prio_pa_host).unwrap(); + unsafe { + cudarc::driver::sys::cuStreamSynchronize(stream.cu_stream()); + } + + // 3a. priorities[i] == |td_i| + ε (raw priority, NO alpha). + for (i, &p) in prio_host.iter().enumerate() { + let expected = td_host[i] + epsilon; + let rel = ((p - expected) / expected).abs(); + assert!( + rel < 1e-4, + "priorities[{i}] = {p}, expected raw |td|+ε = {expected} (rel err {rel:.2e}); \ + if alpha was applied here, the double-exponent bug is back" + ); + } + + // 3b. priorities_pa[i] == (|td_i| + ε)^α (single alpha). + for (i, &pa) in prio_pa_host.iter().enumerate() { + let expected = (td_host[i] + epsilon).powf(alpha); + let rel = ((pa - expected) / expected).abs(); + assert!( + rel < 1e-3, + "priorities_pa[{i}] = {pa}, expected (|td|+ε)^α = {expected} (rel err {rel:.2e}); \ + if alpha is applied twice, ratio drops by factor of (|td|+ε)^α" + ); + } + + // 3c. The behavioral assertion that proves single-α: ratio of + // priorities_pa[max_td] / priorities_pa[min_td] equals + // (td_max/td_min)^α — NOT (td_max/td_min)^(α²). + let pa_max = prio_pa_host[bs - 1]; + let pa_min = prio_pa_host[0]; + let actual_ratio = pa_max / pa_min; + let expected_single_alpha = (td_max / td_min).powf(alpha); + let expected_double_alpha = (td_max / td_min).powf(alpha * alpha); + let rel = ((actual_ratio - expected_single_alpha) / expected_single_alpha).abs(); + assert!( + rel < 0.05, + "priorities_pa ratio = {actual_ratio:.4}, expected single-α \ + ({td_max}/{td_min})^{alpha} = {expected_single_alpha:.4} \ + (rel err {rel:.4}). Pre-fix double-α would give \ + ({td_max}/{td_min})^(α²) = {expected_double_alpha:.4} — if \ + this assertion fires near the pre-fix value, the bug is back" + ); + } } diff --git a/crates/ml-dqn/src/per_kernels.cu b/crates/ml-dqn/src/per_kernels.cu index 27af3e840..689d20965 100644 --- a/crates/ml-dqn/src/per_kernels.cu +++ b/crates/ml-dqn/src/per_kernels.cu @@ -10,10 +10,25 @@ // 4. per_sample — proportional sampling via binary search // --------------------------------------------------------------------------- -// 1. per_update_pa: priority leaf update from BF16 TD errors -// Reads bf16 TD errors directly (no separate cast kernel). -// Computes p = |td|^alpha + eps, writes raw priority, computes pa = p^alpha, -// writes to priorities_pa. atomicMax for batch max. +// 1. per_update_pa: priority leaf update from f32 TD errors +// Reads f32 TD errors directly (no separate cast kernel). +// Standard PER (Schaul et al. 2016): priority p_i = |td_i| + eps, +// sampling weight P(i) ∝ p_i^alpha. Stores BOTH columns of the +// (priority, priority^alpha) row pair so the contract documented at +// `gpu_replay_buffer.rs::priorities_pa_slice` (`priorities_pa[i] = +// priorities[i]^alpha`) holds. atomicMax tracks the running raw- +// priority max for `per_insert_pa`'s max-priority broadcast. +// +// NOTE on prior bug (Class B Finding #2, fixed 2026-05-08): this +// kernel previously stored `priorities[idx] = |td|^alpha + eps` and +// then `priorities_pa[idx] = priorities^alpha`, applying alpha twice. +// Sum-tree sampler reads `priorities_pa`, so the effective sampling +// exponent was alpha² = 0.36 instead of alpha = 0.6 — high-TD events +// were sampled only ~2.5× more than low-TD instead of ~10× as PER +// intends. Fix below stores raw priority in `priorities` and the +// α-weighted sampling weight in `priorities_pa`, matching the +// Rainbow / Schaul-2016 standard. +// // Launch: grid=(ceil(batch_size/256)), block=(256). // --------------------------------------------------------------------------- extern "C" __global__ void per_update_pa( @@ -31,16 +46,23 @@ extern "C" __global__ void per_update_pa( unsigned int idx = indices[i]; if (idx >= (unsigned int)capacity) return; + /* Raw PER priority p_i = |td_i| + eps (Schaul et al. 2016). */ float td = fabsf((td_errors_bf16[i])); - float new_prio = powf(td, alpha) + epsilon; + float new_prio = td + epsilon; if (new_prio < epsilon) new_prio = epsilon; if (new_prio > 1e6f) new_prio = 1e6f; priorities[idx] = new_prio; + /* atomicMax tracks raw-priority max — matches `per_insert_pa`'s + * `priorities[idx] = effective` (also raw) and the post-flush + * `max_priority` broadcast that re-seeds new inserts. */ int ival = __float_as_int(new_prio); atomicMax((int*)batch_max, ival); + /* Sampling weight = priority^alpha. Sum-tree (per_prefix_scan + + * per_sample) walks this buffer; IS-weight (is_weights_f32) reads + * the same column for normalisation. */ float pa = powf(new_prio, alpha); priorities_pa[idx] = pa; } diff --git a/crates/ml-dqn/src/replay_buffer_kernels.cu b/crates/ml-dqn/src/replay_buffer_kernels.cu index 2f793c0f6..5fea2e1fb 100644 --- a/crates/ml-dqn/src/replay_buffer_kernels.cu +++ b/crates/ml-dqn/src/replay_buffer_kernels.cu @@ -535,18 +535,28 @@ extern "C" __global__ void actions_sum_scale_reduce_u32( /* C1/P1: Health-weighted PER priority — boosts priorities of diverse-action * experiences during collapse (health < 0.8). * - * new_prio = clamp((|td_error[i]| + epsilon)^alpha + * Standard PER (Schaul et al. 2016) with diversity multiplier: + * raw_prio = clamp((|td_error[i]| + epsilon) * × (1 + 2×(1−health) × |action[i] − mean_action|), * epsilon, 1e6) - * priorities[indices[i]] = new_prio - * priorities_pa[indices[i]] = new_prio^alpha - * atomicMax(batch_max, new_prio) (IEEE 754 int trick) + * priorities[indices[i]] = raw_prio + * priorities_pa[indices[i]] = raw_prio^alpha ← sampling weight + * atomicMax(batch_max, raw_prio) ← IEEE 754 int trick + * + * Sum-tree sampler walks `priorities_pa`, IS-weight normalises by it. + * The `priorities_pa[i] = priorities[i]^alpha` contract documented at + * `gpu_replay_buffer.rs::priorities_pa_slice` is preserved — drop-in + * replacement for per_update_pa when health < 0.8. * * `mean_action_scaled_ptr` points to a device-side i32 set by the companion * kernel `actions_sum_scale_reduce_u32` (same batch, scaled mean × 1000). * Device-pointer passthrough keeps the update graph-safe — no host readback. * - * Drop-in replacement for per_update_pa when health < 0.8. + * NOTE on prior bug (Class B Finding #2, fixed 2026-05-08): this kernel + * previously computed `base = (|td|+eps)^alpha` then `priorities_pa = + * (base × diversity)^alpha`, applying alpha twice. Effective sampling + * exponent was alpha² = 0.36 instead of alpha = 0.6 — same bug as + * per_update_pa, fixed atomically. */ extern "C" __global__ void pow_alpha_diverse_f32( float* __restrict__ priorities_pa, @@ -568,7 +578,9 @@ extern "C" __global__ void pow_alpha_diverse_f32( unsigned int idx = indices[i]; if (idx >= (unsigned int)capacity) return; - float base = powf(fabsf(td_errors[i]) + epsilon, alpha); + /* Raw PER priority (NO alpha here; alpha is applied once below for + * the sampling weight). */ + float base = fabsf(td_errors[i]) + epsilon; /* Plain load (not __ldg): the backing buffer is device-mapped pinned host * memory (cuMemHostAlloc + cuMemHostGetDevicePointer). __ldg routes through @@ -589,6 +601,8 @@ extern "C" __global__ void pow_alpha_diverse_f32( int ival = __float_as_int(new_prio); atomicMax((int*)batch_max, ival); + /* Sampling weight = priority^alpha (single application). Sum-tree + * (per_prefix_scan + per_sample) walks this column. */ float pa = powf(new_prio, alpha); priorities_pa[idx] = pa; } diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 32f7019c4..8e4950e0a 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -8385,3 +8385,100 @@ No mocks; exercises the actual `memset_zeros` + DtoH path. Runs at deploy time o 1. **`DQNTrainer::clear_replay_buffer` (async, public) is now only called externally**: it remains as a public API surface for callers who want to clear the buffer outside of a fold boundary. Internally it's no longer called from `reset_for_fold`. Not removed because it's a useful operation for ad-hoc replay-state reset (e.g., debug paths, future contract tests). 2. **`priorities_slice()` consumers**: if any future code reads `priorities[]` directly (the public accessor), they will now see all-zeros immediately post-clear (consistent with `priorities_pa`). No current call site reads `priorities_slice()` outside the kernel-update path, but the invariant is now explicit. 3. **Audit-stated leak severity vs trainer-level reality**: the audit framed buffer contamination as "fold N+1 trains on a 50/50 mix" — but the higher-level `DQNTrainer::reset_for_fold` async path was already calling `clear_replay_buffer().await?` BEFORE the agent reset (line 2085 pre-fix). So the leak was actually CLOSED at the highest-level async path, but OPEN at the lower-level sync `DQN::reset_for_fold` and `DQNAgentType::reset_for_fold` entry points. The fix tightens the contract at the lowest-level single-source-of-truth so all paths are atomically guarded. The plateau-causation hypothesis stands but the leak surface was narrower than initially audited — reporting back per hard rule #8. + +--- + +## Class B Finding #2: PER priority double-exponent fix (2026-05-08) + +### Problem statement + +Per the Class B audit, the PER priority kernels were applying alpha twice — corrupting the sampling distribution. Standard Schaul-2016 PER is: + +> Sampling weight P(i) ∝ p_i^α, where p_i = |TD_i| + ε. + +`per_update_pa` (and `pow_alpha_diverse_f32`) computed: + +```c +float new_prio = powf(td, alpha) + epsilon; // |td|^α + ε +priorities[idx] = new_prio; // raw = |td|^α + ε ← α already applied +priorities_pa[idx] = powf(new_prio, alpha); // (|td|^α + ε)^α ← α applied AGAIN +``` + +The sum-tree sampler (`per_prefix_scan` + `per_sample`) walks `priorities_pa` (verified — `gpu_replay_buffer.rs:778, 799`). IS-weight (`is_weights_f32`) reads `sample_priorities` from the same column for normalisation. So the sampler's effective exponent was `α² = 0.36` instead of `α = 0.6`, **flattening prioritization from a ~10× ratio (high-TD vs low-TD) down to ~2.5×**. Sparse trade-close events (3% of buffer) — exactly the events PER is supposed to lift — were being washed out by the squashed exponent. + +### Consumer trace verification (the bug is real, not illusory) + +| Caller | Reads `priorities` | Reads `priorities_pa` | +|--------|--------------------|------------------------| +| `per_prefix_scan` (sum-tree input) | — | YES (`gpu_replay_buffer.rs:778`) | +| `per_sample` (binary search + sample priority emit) | — | YES (`gpu_replay_buffer.rs:799`, `per_kernels.cu:351`) | +| `is_weights_f32` (IS correction) | — | YES (indirectly via `sample_priorities` from `per_sample`) | +| `apply_max_priority_scalar` / max-priority broadcast | YES (via `update_batch_max` atomicMax over `new_prio`) | — | + +`priorities_pa` IS the sampling weight buffer; the bug directly affected the training distribution. + +### Fix (Option B — preserves variable semantics) + +The contract documented at `gpu_replay_buffer.rs:1285` is `priorities_pa[i] = priorities[i]^alpha`. To preserve this semantic without applying α twice, store the raw priority `|td|+ε` in `priorities` (one source of truth) and compute α once for `priorities_pa`: + +```c +/* per_update_pa (per_kernels.cu) */ +float new_prio = td + epsilon; // raw |td|+ε (NO α here) +priorities[idx] = new_prio; +priorities_pa[idx] = powf(new_prio, alpha); // (|td|+ε)^α (α applied once) +``` + +Same shape applied to `pow_alpha_diverse_f32` (replay_buffer_kernels.cu) — the diversity multiplier slots into the raw-priority pre-multiplication, NOT inside the `^α`: + +```c +float base = fabsf(td_errors[i]) + epsilon; // raw (NO α) +float new_prio = base * diversity_mult; +priorities[idx] = new_prio; +priorities_pa[idx] = powf(new_prio, alpha); // (|td|+ε)·div)^α (α applied once) +``` + +### Kernels NOT affected by the bug + +- `per_insert_pa` (per_kernels.cu:108): already correct — `priorities = effective`, `priorities_pa = effective^α`. The `effective = max_priority × boost` line uses `max_priority` as-is (raw), so this kernel is α-once. **Unchanged.** +- `priority_update_f32` (replay_buffer_kernels.cu:233): single-buffer kernel, doesn't write `priorities_pa`. Verified UNUSED — no Rust caller. **Unchanged.** + +### Files touched + +| File | Δ | Notes | +|------|---|-------| +| `crates/ml-dqn/src/per_kernels.cu` | +30 / -7 | `per_update_pa` raw-priority storage + Class B #2 docstring | +| `crates/ml-dqn/src/replay_buffer_kernels.cu` | +25 / -8 | `pow_alpha_diverse_f32` raw-priority storage + Class B #2 docstring | +| `crates/ml-dqn/src/gpu_replay_buffer.rs` | +120 / -2 | Behavioral test (single-α ratio assertion) + `update_priorities_gpu` doc update | +| `docs/dqn-wire-up-audit.md` | +75 / -0 | This entry | + +### Behavioral test + +`gpu_replay_buffer::tests::test_per_priority_single_exponent_no_double_alpha` (GPU-gated `#[ignore]`): + +- Inserts 16 transitions at max_priority. +- Fires `update_priorities_gpu` with TD errors geometrically distributed across `[0.001, 100.0]` (4 orders of magnitude). +- Reads back both `priorities` and `priorities_pa`. +- Asserts: + - `priorities[i] == |td_i| + ε` (raw, NO alpha) — within 1e-4 relative. + - `priorities_pa[i] == (|td_i| + ε)^α` — within 1e-3 relative. + - Ratio `priorities_pa[max_td] / priorities_pa[min_td] == (100/0.001)^0.6 ≈ 1995×` — within 5%. **Pre-fix would yield ~100× (off by 20×)**, an obvious failure mode the assertion catches. + +The ratio assertion makes the bug regression-proof: any future kernel edit that re-introduces a double-α (or any exponent-arithmetic mistake) will fail this test immediately. + +### Verification + +- `SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo check -p ml-dqn -p ml --tests --all-targets` — clean. +- `cargo test -p ml-dqn --release --lib gpu_replay_buffer::tests:: -- --ignored` — 3/3 pass (new test + Class B #1 fold-clear test + large-capacity prefix scan test). +- `cargo test -p ml --test sp14_oracle_tests --release -- --ignored` — 24/24 pass. +- `cargo test -p ml --test sp15_phase1_oracle_tests --release -- --ignored` — 36/36 pass, including `per_sampler_weights_per_transition_not_uniform_across_batch` and `per_sampler_weights_recovery_transitions_higher` which exercise the sampling path. + +### Cumulative WR-plateau fix series (commit 11) + +- ... (commits 1-10 from Class B #1 entry above) +- **Class B Finding #2 (this commit): PER priority double-exponent fix**. + +### Concerns + +1. **Magnitude of the upstream effect**: with α=0.6 default and the typical TD-error span of 3-4 orders of magnitude in trade-close vs no-trade transitions, the effective sampling lift on high-TD events drops from ~10× (single α) to ~2.5× (double α). That's ~4× attenuation of PER's prioritization power — a meaningful but not catastrophic miscalibration. The plateau-causation hypothesis (alongside Class B #1's fold-boundary contamination) is plausible but not proven by audit alone; smoke run will quantify. +2. **`max_priority` semantic shift**: pre-fix `max_priority` tracked the running max of `|td|^α + ε`. Post-fix it tracks max of `|td| + ε`. The downstream `per_insert_pa` broadcasts this value as `raw_priorities` for new inserts and computes `effective^α` for `priorities_pa` — so the shift is internally consistent (raw priority in, α applied once on insert). External callers of `apply_max_priority_scalar` (currently none in production paths, only test/seg_tree paths) would need to pass raw priority values — no caller change required because no production caller exists. +3. **`priority_update_f32` left in place despite being dead**: the kernel compiles but has no Rust caller. Per `feedback_no_legacy_aliases` it should be deleted, but doing so requires removing the kernel registration in the loader (`crates/ml-dqn/src/replay_buffer_kernels.rs` or similar) and that's outside the surface of this fix. Filed as a follow-up cleanup; the kernel is harmless idle code (no execution path).