From 021bb0ef736ab380e02fa8a1646f8cb84ec20593 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 27 Apr 2026 00:27:46 +0200 Subject: [PATCH] spec(dqn): pivot Phase 0+2 tests to GPU-direct (no CPU mirror) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit User correctly identified that CPU mirror function tests don't test the production GPU code path. A bug shared between mirror and kernel (translated identically wrong) would slip through. Mirror tests + a single GPU bridge test were a weak compromise. GPU-direct testing strategy: - All Phase 0 kernel-correctness tests (0.A, 0.B, 0.C, 0.D, 0.F): launch tiny test-only kernels with the SAME math the Phase 2 production kernel will use; assert properties of the output. - Test 0.E (synthetic edge discovery): stays CPU. It tests an ALGORITHMIC PROPERTY of Thompson exploration (does it discover edge if edge exists?), not a kernel correctness property. - All Phase 2 unit tests (2.A-2.D): GPU-direct against the modified production kernel. - Phase 0.F (real checkpoint extraction): unchanged — already GPU. Local development uses RTX 3050 GPU (per memory user_dev_environment.md). CI runs --ignored flag to skip GPU tests on CPU-only runners. Time budget: Phase 0 was 1-2 days (CPU mirror); now 2-3 days (GPU-direct, includes kernel wrapper setup half-day). Other delta: - Phase 0 deliverable file renamed: distributional_q.rs -> distributional_q_tests.rs (no mirror functions, just tests + kernel wrappers). - Phase 2 unit tests rephrased to launch production kernel rather than compare against CPU mirror. - L1 verification gate runtime: seconds -> minutes (GPU launch overhead per test). The user's intuition was right: testing production directly is the honest approach. Mirror was an optimization that traded correctness for speed; with local GPU available the optimization isn't needed. Co-Authored-By: Claude Opus 4.7 (1M context) --- ...26-distributional-rl-aggregation-design.md | 58 +++++++++---------- 1 file changed, 28 insertions(+), 30 deletions(-) diff --git a/docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md b/docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md index c693ea026..35116768c 100644 --- a/docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md +++ b/docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md @@ -286,45 +286,43 @@ The decision to defer v2 keeps Phase 2 verification clean: we measure pure Thomp ## Phased Implementation -### Phase 0 — TDD Hypothesis Verification + Synthetic Edge Test +### Phase 0 — GPU-Direct Hypothesis Verification + Synthetic Edge Test -**Goal**: Prove (or disprove) the bias hypothesis with controlled inputs BEFORE changing any kernel behaviour. Verify Thompson sampling reverses the bias AND can discover synthetic edge. +**Goal**: Prove (or disprove) the bias hypothesis by exercising the actual production kernels (or thin wrappers around them) with controlled inputs. No CPU mirror functions — tests target the real GPU code path so kernel-only bugs (CUDA index errors, race conditions, type cast mistakes) are caught at the unit-test level rather than waiting for production training. **Deliverables**: -CPU mirror functions in `crates/ml/src/trainers/dqn/distributional_q.rs`: -- `expected_value_c51(atoms: &[f32], probs: &[f32]) -> f32` -- `expected_value_iqn(quantiles: &[f32]) -> f32` -- `sample_c51(atoms: &[f32], probs: &[f32], rng: &mut impl Rng) -> f32` -- `sample_iqn(quantiles: &[f32], rng: &mut impl Rng) -> f32` -- `thompson_pick_direction(c51_per_dir: &[(&[f32], &[f32])], iqn_per_dir: &[&[f32]], rng: &mut impl Rng) -> usize` -- `argmax_e_q(c51_per_dir: &[(&[f32], &[f32])], iqn_per_dir: &[&[f32]]) -> usize` +A test module `crates/ml/src/trainers/dqn/distributional_q_tests.rs` containing: +- A small kernel wrapper `launch_thompson_direction(c51_probs, atom_values, iqn_quantiles, rng_seed) -> [u8; b0_size]` that launches a tiny test-only CUDA kernel with the SAME math the Phase 2 kernel will use. This becomes the actual production code in Phase 2; Phase 0 tests it standalone first. +- A small kernel wrapper `launch_argmax_eq(c51_probs, atom_values, iqn_quantiles) -> u8` for eval-mode argmax of E[Q_C51]+E[Q_IQN]. +- Helper `compute_sigma_from_c51_atoms(probs, atoms) -> f32` and `compute_sigma_from_iqn_quantiles(qs) -> f32` — these run on the GPU via thin kernel wrappers and exist as observation utilities for Test 0.F. -**Six unit tests**: +Each test launches the kernel(s), reads outputs back, asserts properties. All tests are `#[ignore]` (GPU-required); explicitly run via `cargo test -p ml --lib distributional_q_tests -- --ignored --nocapture`. -- **0.A — Bias mechanism reproduces under both `argmax(E[Q])` AND `Boltzmann(E[Q])` (current eval semantics)**: vanilla argmax picks Flat ~100%; current Boltzmann at typical τ ranges produces sharp picks toward Flat (>60% Flat); Thompson sampling produces non-degenerate distribution where Long/Short combined ≥ 30%. Three distributions to compare; assertion is on relative ordering. +**Six tests** (5 GPU-direct kernel tests + 1 CPU algorithmic-property test): -- **0.B — `sample_c51` distribution matches input probabilities**: 100k samples; histogram matches `probs` within 1% per atom. +- **0.A — Bias mechanism reproduces (GPU)**: build a synthetic `c51_probs[4 dirs × n_atoms]` and `iqn_quantiles[4 dirs × 5]` representing the failure mode (Flat = δ(0); Long/Short = wide distribution centered at −0.001). Launch `launch_argmax_eq` AND launch `launch_thompson_direction` 10000× with different seeds. Assert: argmax picks Flat 100% of trials; Thompson picks (Long+Short) ≥ 30% of trials. -- **0.C — `sample_iqn` empirical CDF matches quantile inputs**: 100k samples; quantile estimates match input within 5% on standard test distribution. +- **0.B — Thompson C51 sample distribution matches input (GPU)**: launch `launch_thompson_direction` 100k× with single direction having known atom probabilities `p = [0.1, 0.7, 0.2]` over atoms `[−1, 0, +1]`. Histogram returned samples (binned to atom indices). Assert: each bin within 1% of input probability. -- **0.D — Thompson reverses bias on synthetic failure mode**: construct `(E_flat=0, σ_flat=0, E_long=−0.001, σ_long=0.05)`. The σ_long=0.05 value is chosen to match the order of magnitude expected when distributional Q is trained on per-bar returns where typical |return| ~ 50 bps; Phase 0.F validates this against a real checkpoint. Confirm: `P(Thompson picks Long) > 0.20` over 10k trials, while argmax(E[Q]) picks Long with probability 0 (deterministic). +- **0.C — Thompson IQN sample distribution matches quantile inputs (GPU)**: similar to 0.B but with synthetic quantile values matching a known Gaussian. Empirical quantiles of 100k samples within 5% of analytic quantiles (Kolmogorov-Smirnov–style check). -- **0.E — SYNTHETIC EDGE DISCOVERY (NEW, critical)**: simulate 100 iterations of Q-learning on a tiny synthetic bandit (1-state, 2-action: Long with KNOWN +0.005 edge after tx_cost, Flat at 0). ~50 lines of Rust. Confirm: - - With argmax-only training: model NEVER picks Long enough to discover edge; Q[Long] stays below Q[Flat] = 0 - - With Thompson training: model picks Long with probability proportional to σ; Q[Long] eventually exceeds Q[Flat] (concretely: Q[Long] > 0 by iteration 50) - - **If 0.E fails → fundamental issue beyond aggregation; halt and reassess reward shaping** +- **0.D — Thompson reverses bias on synthetic failure mode (GPU)**: same setup as 0.A but assertion focused on directional preference. With `(E_flat=0, σ_flat=0, E_long=−0.001, σ_long=0.05)`: argmax picks Long with probability 0 (deterministic Flat); Thompson picks Long with probability ≥ 0.20 over 10k seeds. Quantitatively confirms σ-driven exploration. -- **0.F — Real checkpoint extraction (`#[ignore]`, GPU)**: load a "converged" checkpoint. **"Converged" definition**: a checkpoint from a training run of ≥60 epochs (matching prior validation runs in project history per memory `dqn-session-2026-03-25.md`) where val_sharpe has stabilised (no >10% change over the last 10 epochs). Concrete assertions: +- **0.E — Synthetic edge discovery (CPU, Q-learning simulation)**: a 100-iteration Q-learning loop on a 1-state 2-action bandit with KNOWN edge. Tests an ALGORITHMIC PROPERTY (does Thompson exploration converge to right answer when edge exists?), not a kernel correctness property. CPU is appropriate because no production kernel is being tested — only the math of "samples from a tightening distribution converge to the underlying Q if edge exists." Confirm: with argmax-only updates, Q[Long] stays below Q[Flat]; with Thompson updates, Q[Long] > Q[Flat] by iteration 50. + +- **0.F — Real checkpoint extraction (GPU, integration)**: load a "converged" checkpoint. **"Converged" definition**: a checkpoint from a training run of ≥60 epochs (matching prior validation runs in project history per memory `dqn-session-2026-03-25.md`) where val_sharpe has stabilised (no >10% change over the last 10 epochs). Run the existing forward pass to extract C51 atoms + IQN quantiles per direction on a representative val state batch. Use `compute_sigma_from_c51_atoms` and `compute_sigma_from_iqn_quantiles` to check structural assertions: - σ_C51[FLAT] < 0.01 × σ_C51[LONG] (Flat dist is much tighter than Long; structural) - σ_IQN[FLAT] < 0.01 × σ_IQN[LONG] (same for IQN) - E[Q_FLAT] > E[Q_LONG] (the bias actually exists in this checkpoint) - - argmax(E_C51 + E_IQN) picks FLAT - - Thompson sampling over 1000 trials produces P(LONG) + P(SHORT) ≥ 0.20 + - `launch_argmax_eq` picks FLAT + - `launch_thompson_direction` over 1000 trials produces P(LONG) + P(SHORT) ≥ 0.20 **Exit gate**: All 6 tests pass. Phase 0 first ENUMERATES existing checkpoints (project history has multiple long training runs from `train-7rgqd` 60-epoch validation per task #80, `train-bv2n5`, etc.); if any qualifies as "converged" per the above definition, use it. Only if none qualifies does Phase 0 train a new 60-epoch checkpoint (~3 hours on L40S, not "1 day"). -**Time budget**: 1-2 days. Tests 0.A-0.E ~1 hour each (5 hours); checkpoint enumeration ~1 hour; Test 0.F implementation + run ~2 hours; new checkpoint training only if needed (~3 hours runtime, can run while implementer works on other tests). +**Time budget**: 2-3 days. GPU test setup (kernel wrappers, fixture loading) ~half day; tests 0.A-0.D ~2 hours each (8 hours); checkpoint enumeration ~1 hour; Test 0.F ~3 hours; Test 0.E (CPU simulation) ~2 hours; new checkpoint training only if needed (~3 hours runtime, can run while implementer works on other tests). + +**Why GPU-direct rather than CPU mirror**: tests now exercise the actual production kernel math, catching CUDA-specific bugs (index errors, race conditions, type cast mistakes) that CPU mirror tests would miss. Local development uses RTX 3050 GPU per `memory/user_dev_environment.md`. CI runs `--ignored` flag to skip GPU tests on CPU-only runners; GPU runners (Argo training pods, local) run them as part of the L1 verification gate. ### Phase 1 — Audit & Repair Existing Reward Levers (TDD, conditional) @@ -416,12 +414,12 @@ The existing conviction computation (`experience_action_select` after action sel Even though Thompson selects the action at training, the policy's CONFIDENCE about the choice is still E[Q]-based — Thompson samples are stochastic and would jitter the Kelly cap warmup floor. -#### Component 4.4: Phase 2 unit tests +#### Component 4.4: Phase 2 unit tests (all GPU-direct against the production kernel) -- **2.A — End-to-end Thompson + argmax modes** (CPU mirror): given synthetic C51+IQN distributions, confirm training mode produces stochastic picks matching distribution shapes; eval mode produces deterministic argmax picks. -- **2.B — GPU Thompson sampler matches CPU mirror** (GPU integration, `#[ignore]`): run kernel with controlled inputs and seeded RNG; confirm output matches CPU mirror. -- **2.C — Magnitude/order/urgency unaffected**: numerical equality after seeded RNG. -- **2.D — Real-batch e2e on converged checkpoint**: training mode produces `active_frac` (Long+Short combined) > 40% over the batch; eval mode reproduces previously-observed argmax picks (i.e., picks the same direction as `argmax(E_C51 + E_IQN)` on the loaded checkpoint). +- **2.A — End-to-end Thompson + argmax modes (GPU)**: launch the modified `experience_action_select` kernel with synthetic C51+IQN distributions for `eval_mode=false` (training) and `eval_mode=true` (eval). Assert: training mode produces stochastic picks matching distribution shapes (Long+Short ≥ 30% over many seeds when distribution favours them); eval mode produces deterministic argmax picks (same direction across all seeds for given inputs). +- **2.B — Phase 2 kernel matches Phase 0 standalone wrapper (GPU)**: run the Phase 2 production kernel and the Phase 0 `launch_thompson_direction` wrapper with identical inputs and seeded RNG. Confirm: outputs are bit-identical. This catches drift between standalone test wrapper and integrated production kernel. +- **2.C — Magnitude/order/urgency unaffected (GPU)**: launch the kernel with eval_mode=true and given inputs. Confirm: mag/ord/urg picks are bit-identical to a baseline run BEFORE Phase 2 (use a stored snapshot of pre-Phase-2 kernel output for the same inputs). +- **2.D — Real-batch e2e on converged checkpoint (GPU)**: load the converged checkpoint from Phase 0.F; run a full action-select pass over a representative val batch. Training mode: `active_frac` (Long+Short combined) > 40% over the batch. Eval mode: reproduces previously-observed argmax picks (picks the same direction as `launch_argmax_eq` on the same checkpoint). #### Component 4.5: Documentation + audit @@ -469,8 +467,8 @@ The single 5-epoch L40S smoke is INSUFFICIENT to verify Thompson — long-term e | Layer | Scope | Pass criterion | Runtime | |---|---|---|---| -| **L1**: CPU unit tests | Phase 0+1+2 mirror functions | All asserts pass | seconds | -| **L2**: GPU integration test | Tests 0.F + 2.B against converged checkpoint | sampling matches CPU mirror | seconds | +| **L1**: GPU-direct unit tests | Phase 0 (5 GPU + 1 CPU) + Phase 1 (6) + Phase 2 (4 GPU) | All asserts pass | minutes (GPU launch overhead) | +| **L2**: GPU integration test | Tests 0.F + 2.B + 2.D against converged checkpoint | kernel output matches structural assertions | seconds | | **L3**: Short smoke (3-fold × 5-epoch L40S) | Immediate Thompson effect | **training-time `train_active_frac` > 40%** across all epochs (NEW HEALTH_DIAG metric — see Phase 2 deliverable below) | ~12 min | | **L4**: Long smoke (1-fold × 30-epoch L40S) | Long-term edge discovery | val_sharpe trends positive; val_active_frac increases as edge discovered | ~1 hour | | **L5**: Full validation matrix (per Plan 5 Task 5 spec) | Tier 1/2/3 exit criteria | per Plan 5 spec | ~6 hours | @@ -574,7 +572,7 @@ The single 5-epoch L40S smoke is INSUFFICIENT to verify Thompson — long-term e | Plan | Description | Time | |---|---|---| -| Plan A | Phase 0: TDD hypothesis + synthetic edge verification | 1-2 days | +| Plan A | Phase 0: GPU-direct hypothesis + synthetic edge verification | 2-3 days | | Plan B | Phase 1: existing-lever audit | 1 day (longer if bug) | | Plan C | Phase 2: Thompson sampling integration | 2-3 days | | Plan D | Phase 3: long verification + direction-branch dead-code cleanup | per Plan 5 + 0.5 day |