spec(dqn): pivot Phase 0+2 tests to GPU-direct (no CPU mirror)
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) <noreply@anthropic.com>
This commit is contained in:
@@ -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 |
|
||||
|
||||
Reference in New Issue
Block a user