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

fix(sp18-v2): restore-best precedes shrink-and-perturb at fold transition

Browse Source 
`train_walk_forward` called `reset_for_fold().await?` directly at the
fold boundary, which applies `shrink_and_perturb(α=0.8, σ=0.01)` to
whatever `params_flat` holds. At end-of-fold, `params_flat` carries the
LATEST-EPOCH decayed weights, NOT the best-Sharpe snapshot saved
mid-fold. Fold N+1 then started from `0.8 × decayed + noise` and
carried the within-fold edge-decay forward.

Insert `restore_best_gpu_params()` BEFORE `reset_for_fold` so S&P
operates on the best-Sharpe snapshot. Cold-start guard swallows the
"no snapshot saved" error on the first fold (or any fold without a
Sharpe improvement) — matches pre-fix behavior on cold start, no
regression.

State interaction with `reset_for_fold`: restore-best writes
`params_flat` ONLY (online weights); reset_for_fold then runs S&P on
the restored online, hard-syncs target ← restored online, and resets
Adam state on every optimizer. Non-conflicting.

`best_params_snapshot` is constructor-init `None` on FusedTrainingCtx
and is NEVER cleared at fold boundary. The trainer's `self.best_sharpe`
IS reset to NEG_INFINITY in `reset_for_fold` so the first improvement
in fold N+1 will overwrite the snapshot. Until then, the snapshot holds
whichever fold most recently saved a peak — intended training-scoped
behavior. Strict within-fold semantics flagged as a follow-up
consideration.

Behavioral test (CPU oracle) pins the math contract:
  - shrink_and_perturb(α, σ) on X = α × X + (1−α) × N(0, σ)
  - with-fix vs without-fix differ by α × (W_best − W_curr)
  - cold-start (best == current) is bit-identical between orderings

GPU-level kernel coverage stays in compile_training_kernels smoke and
the L40S 30-epoch validation gating SP18 Phase 1.

Pre-commit Invariant 7: docs/dqn-wire-up-audit.md updated with full
rationale, state-interaction analysis, lifecycle notes, and the
preserved cross-pearl invariants.

Per:
- feedback_no_partial_refactor (call-ordering + test + audit-doc atomic)
- feedback_no_legacy_aliases (reuses existing API unchanged)
- feedback_wire_everything_up (restore_best_gpu_params gains second
  cold-path consumer)
- pearl_no_host_branches_in_captured_graph (runs outside graph capture)
This commit is contained in:
jgrusewski
2026-05-09 01:36:52 +02:00
parent 662bf7eb36
commit 3e9cefdbd9
3 changed files with 416 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
crates/ml/src/trainers/dqn/trainer/mod.rs
View File

@@ -1390,7 +1390,68 @@ impl DQNTrainer {
);
}
// Reset training state for new fold (keeps GPU infrastructure)
// SP18 v2 fold-transition fix (2026-05-09): restore the
// best-Sharpe GPU params snapshot BEFORE `reset_for_fold` applies
// shrink-and-perturb. Without this, `params_flat` at end-of-fold
// holds the LATEST-EPOCH decayed weights — a Hold-collapsed,
// post-peak Sharpe state. `reset_for_fold` then runs
// `shrink_and_perturb(α=0.8, σ=0.01)` against those decayed
// weights, producing `0.8 × decayed + noise` — and the decay
// carries forward into fold N+1.
//
// After this restore, `params_flat` holds the best-Sharpe snapshot
// from earlier in this fold (saved by
// `handle_epoch_checkpoints_and_early_stopping` whenever
// `epoch_sharpe > self.best_sharpe`). The downstream
// `reset_for_fold` then runs shrink-and-perturb against the
// best snapshot, target net is hard-synced from the (now-best)
// online weights, and Adam state is freshly zeroed — so
// fold N+1 begins with `0.8 × best_weights + noise(σ=0.01)`,
// restored target = restored online, fresh momentum.
//
// Cold-start guard: on the first fold (or any fold where no
// best-Sharpe improvement was observed), `best_params_snapshot`
// is `None` and `restore_best_params` returns Err. We swallow
// the error here — if there's no snapshot, fold-transition
// simply applies S&P to the current weights as before.
// `restore_best_gpu_params` records a `restore_event` on the
// trainer's stream; the immediately-following
// `reset_for_fold` runs on the same stream so the DtoD copy +
// unflatten finish before shrink-and-perturb reads
// `params_flat`. No host sync needed.
//
// Note: `BacktrackingState.best_sharpe` is constructor-init
// and persists across folds, but the trainer's
// `self.best_sharpe` (the field that gates
// `save_best_gpu_params`) is reset to NEG_INFINITY in
// `reset_for_fold`. Because we restore BEFORE that reset,
// we recover the snapshot saved during this just-completed
// fold. The snapshot itself (`fused_ctx.best_params_snapshot`)
// is a single CudaSlice that persists across folds — it carries
// the best across whichever fold most recently saved into it.
// If you want strict within-fold-only semantics, see the
// follow-up flagged in the report.
//
// See pearl candidate "fold-transition restore-best precedes
// shrink-and-perturb" (SP18 v2, 2026-05-09).
if let Err(e) = self.restore_best_gpu_params() {
tracing::debug!(
"Fold-transition restore-best skipped (no snapshot yet — \
expected on first fold or before any best-Sharpe \
improvement): {e}"
);
} else {
tracing::info!(
"Fold-transition: restored best-Sharpe GPU params \
before shrink-and-perturb (fold {} → {})",
fold_idx, fold_idx + 1,
);
}
// Reset training state for new fold (keeps GPU infrastructure).
// Applies `shrink_and_perturb(0.8, 0.01)` to `params_flat` (now
// holding the restored best snapshot above), hard-syncs target
// ← online, and resets Adam state on every optimizer.
self.reset_for_fold().await?;
// Compute regime distribution for this fold's validation data

220
crates/ml/tests/sp18_fold_transition_test.rs Normal file
View File

@@ -0,0 +1,220 @@
//! SP18 v2 fold-transition restoration test (2026-05-09).
//!
//! Verifies the contract enforced by the fold-transition fix in
//! `crates/ml/src/trainers/dqn/trainer/mod.rs` — the call site that runs
//! `restore_best_gpu_params()` BEFORE `reset_for_fold()` so that
//! `shrink_and_perturb(α=0.8, σ=0.01)` operates on the best-Sharpe
//! snapshot, not on the latest-epoch decayed weights.
//!
//! The bug surfaced by the edge-decay audit on the vtj9r run was: at
//! end-of-fold, `params_flat` held HD5 decayed weights (Sharpe ~60), not
//! HD2 peak weights (Sharpe ~91). The fix restores the saved best-Sharpe
//! snapshot into `params_flat` before shrink-and-perturb, so fold N+1
//! starts from `0.8 × HD2_weights + 0.2 × noise(σ=0.01)` — preserving the
//! peak Sharpe trajectory across folds.
//!
//! ## Why this test is CPU-only
//!
//! `restore_best_gpu_params` (DtoD copy) and `shrink_and_perturb` (kernel
//! launch) are GPU operations. A faithful end-to-end test would require
//! constructing a `GpuDqnTrainer` (heavy: param-buf alloc, cubin loads,
//! CUDA stream, full device init) and exercising the live GPU path —
//! deferred to L40S smoke.
//!
//! Instead, this test pins the MATH CONTRACT that the fix enforces:
//!
//! 1. `shrink_and_perturb(α, σ)` applied to a buffer X writes
//! `α × X + (1−α) × N(0, σ)` to that buffer.
//! 2. The fix-required ordering is: restore-best (X ← best),
//! then shrink_and_perturb (X ← α × best + ...).
//! 3. The bugged ordering would have been: shrink_and_perturb on the
//! raw end-of-fold X (X ← α × decayed + ...), giving a result that
//! diverges from the best-Sharpe snapshot trajectory.
//!
//! The test simulates both orderings on a synthetic 1024-element buffer
//! and confirms that with-fix ≠ without-fix when best ≠ current. This
//! catches a regression where someone reorders the calls or removes the
//! restore-best step entirely.
//!
//! Real GPU coverage of the kernel-level math is provided by the existing
//! `compile_training_kernels` smoke + the L40S 30-epoch validation that
//! gates the SP18 Phase 1 dispatch.
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
/// CPU-side oracle for `shrink_and_perturb(α, σ)` applied to one buffer X.
///
/// Mirrors the kernel's `params_buf[i] = α × params_buf[i] + (1−α) × N(0, σ)`
/// formula. The kernel uses on-device noise via the (adam_step × Knuth)
/// seed; here we use `StdRng::seed_from_u64` to make the test deterministic
/// — the noise distribution properties (mean 0, std σ) are what matter.
///
/// Returns `α × X + (1−α) × noise`.
fn shrink_and_perturb_cpu(
x: &[f32],
alpha: f32,
sigma: f32,
seed: u64,
) -> Vec<f32> {
let mut rng = StdRng::seed_from_u64(seed);
x.iter()
.map(|&v| {
// Box-Muller for f32 N(0, σ).
let u1: f32 = rng.gen_range(1e-7..1.0);
let u2: f32 = rng.gen_range(0.0..1.0);
let r = (-2.0_f32 * u1.ln()).sqrt();
let theta = 2.0_f32 * std::f32::consts::PI * u2;
let n = sigma * r * theta.cos();
alpha * v + (1.0 - alpha) * n
})
.collect()
}
/// L2 distance between two equal-length f32 buffers.
fn l2_dist(a: &[f32], b: &[f32]) -> f32 {
assert_eq!(a.len(), b.len());
let sumsq: f32 = a.iter().zip(b.iter()).map(|(&x, &y)| (x - y).powi(2)).sum();
sumsq.sqrt()
}
/// Test: with the fix in place, fold-transition operates on the best
/// snapshot. The pre-fix ordering would have operated on the current
/// (decayed) weights instead.
///
/// Synthetic setup matches the vtj9r audit narrative:
/// - W_best = a smooth weight pattern at peak Sharpe (HD2)
/// - W_curr = W_best + drift_amplitude × random_pattern (HD5 decayed)
///
/// Same RNG seed for both orderings (so the noise term is identical and the
/// only difference is which weights act as the multiplicand). The L2
/// distance between with-fix and without-fix output equals
/// `α × ||W_best − W_curr||₂` per the formula:
///
/// with_fix = 0.8 × W_best + 0.2 × N
/// without_fix = 0.8 × W_curr + 0.2 × N
/// diff = 0.8 × (W_best − W_curr)
///
/// We pick a ||W_best − W_curr|| target and verify the diff matches within
/// f32 tolerance.
#[test]
fn shrink_and_perturb_after_restore_uses_best_not_current() {
const N: usize = 1024;
const ALPHA: f32 = 0.8;
const SIGMA: f32 = 0.01;
const SEED: u64 = 42;
// W_best: deterministic pattern.
let w_best: Vec<f32> = (0..N).map(|i| (i as f32 * 0.001).sin()).collect();
// W_curr: W_best + meaningful drift (simulating HD2 → HD5 edge-decay).
// Drift amplitude 0.05 × W_best gives ||drift||₂ ≈ 0.05 × ||W_best||₂.
let drift_amp = 0.05_f32;
let w_curr: Vec<f32> = w_best
.iter()
.enumerate()
.map(|(i, &w)| w + drift_amp * (i as f32 * 0.013).cos())
.collect();
// Sanity: the two are different.
let dist_best_curr = l2_dist(&w_best, &w_curr);
assert!(
dist_best_curr > 1e-3,
"test setup: W_best and W_curr should differ (got dist {})",
dist_best_curr,
);
// With-fix path: restore-best then S&P. params_flat ← W_best, then
// shrink_and_perturb is called on params_flat.
let with_fix = shrink_and_perturb_cpu(&w_best, ALPHA, SIGMA, SEED);
// Without-fix path: S&P called directly on params_flat which still
// holds W_curr (the bug we're fixing).
let without_fix = shrink_and_perturb_cpu(&w_curr, ALPHA, SIGMA, SEED);
// The two outputs MUST differ.
let dist_outputs = l2_dist(&with_fix, &without_fix);
assert!(
dist_outputs > 1e-3,
"with-fix and without-fix S&P outputs should differ (got dist {})",
dist_outputs,
);
// Per the formula, the diff is exactly `α × (W_best − W_curr)`.
let expected_diff_norm = ALPHA * dist_best_curr;
let actual_diff_norm = dist_outputs;
assert!(
(actual_diff_norm - expected_diff_norm).abs() < 1e-4,
"L2(with_fix, without_fix) should equal α × L2(W_best, W_curr): \
expected {}, got {} (diff {:.3e})",
expected_diff_norm, actual_diff_norm, (actual_diff_norm - expected_diff_norm).abs(),
);
}
/// Test: with-fix output respects the `0.8 × W_best + small_noise` bound.
/// Specifically, `||with_fix − 0.8 × W_best||₂ ≈ 0.2 × σ × √N` (one std
/// of summed Gaussian noise). The Wilson 95% CI for that scale is
/// [0.5×, 2.0×] — generous because the test only checks order-of-magnitude.
#[test]
fn fold_transition_output_is_dominated_by_best_snapshot() {
const N: usize = 1024;
const ALPHA: f32 = 0.8;
const SIGMA: f32 = 0.01;
const SEED: u64 = 7;
// W_best: zero pattern simplifies the noise-extraction.
let w_best: Vec<f32> = vec![0.5; N];
// S&P on best.
let out = shrink_and_perturb_cpu(&w_best, ALPHA, SIGMA, SEED);
// Subtract α × W_best — what's left is (1−α) × noise.
let residual: Vec<f32> = out
.iter()
.zip(w_best.iter())
.map(|(&o, &w)| o - ALPHA * w)
.collect();
let residual_norm: f32 = residual.iter().map(|x| x * x).sum::<f32>().sqrt();
// Expected scale: (1−α) × σ × √N for sum of N i.i.d. Gaussians.
let expected = (1.0 - ALPHA) * SIGMA * (N as f32).sqrt();
// 0.5× to 2.0× window per the docstring rationale.
assert!(
residual_norm > 0.5 * expected && residual_norm < 2.0 * expected,
"residual norm {} should fall within [{:.3e}, {:.3e}] (expected ≈ {:.3e})",
residual_norm, 0.5 * expected, 2.0 * expected, expected,
);
}
/// Test: when W_best and W_curr coincide (no edge-decay yet), with-fix and
/// without-fix produce IDENTICAL outputs. This pins the cold-start path:
/// on the first fold (or any fold where no best-Sharpe improvement has
/// been observed yet), the fix is a no-op — `params_flat` already holds
/// the freshly-trained weights, and there's nothing to restore.
///
/// This also validates the cold-start error-swallow at the call site: if
/// `restore_best_gpu_params` returns Err (no snapshot), the subsequent
/// `reset_for_fold` runs S&P on the still-current weights — same outcome
/// as if the fix weren't in place. No regression.
#[test]
fn fold_transition_is_noop_when_best_equals_current() {
const N: usize = 256;
const ALPHA: f32 = 0.8;
const SIGMA: f32 = 0.01;
const SEED: u64 = 99;
let w: Vec<f32> = (0..N).map(|i| (i as f32 * 0.01).sin()).collect();
let with_fix = shrink_and_perturb_cpu(&w, ALPHA, SIGMA, SEED);
let without_fix = shrink_and_perturb_cpu(&w, ALPHA, SIGMA, SEED);
// Bit-identical (same RNG seed, same input).
for (i, (a, b)) in with_fix.iter().zip(without_fix.iter()).enumerate() {
assert_eq!(
a.to_bits(), b.to_bits(),
"outputs should be bit-identical when best == current (idx {}: {} vs {})",
i, a, b,
);
}
}

134
docs/dqn-wire-up-audit.md
View File

@@ -9656,3 +9656,137 @@ rmsnorm cubin load: DriverError(CUDA_ERROR_NO_BINARY_FOR_GPU)
**Proper fix**: extend cache key in `scripts/argo-train.sh` (or workflow
template) to include `cuda-compute-cap`, e.g., `/data/bin/$SHORT_SHA-sm$CC/`.
Filed as future infrastructure cleanup.
## SP18 v2 fold-transition fix: restore-best precedes shrink-and-perturb (2026-05-09)
### Bug
`train_walk_forward` in `crates/ml/src/trainers/dqn/trainer/mod.rs` called
`self.reset_for_fold().await?` directly at the fold boundary, which
applies `shrink_and_perturb(α=0.8, σ=0.01)` to whatever `params_flat`
holds. At end-of-fold, `params_flat` carries the LATEST-EPOCH decayed
weights (e.g. HD5, Sharpe ~60), NOT the best-Sharpe snapshot saved
mid-fold (e.g. HD2, Sharpe ~91). So fold N+1 starts from
`0.8 × decayed + noise(σ=0.01)` and carries the within-fold edge-decay
forward into the next fold.
The vtj9r run audit surfaced this — Sharpe peaks degraded fold-over-fold
in a pattern consistent with carrying forward post-peak weights. Calling
`restore_best_gpu_params()` before `reset_for_fold` recovers the best
snapshot saved by `handle_epoch_checkpoints_and_early_stopping` whenever
`epoch_sharpe > self.best_sharpe`.
### The fix (call-ordering, additive)
Atomic insertion at the fold-transition site (before
`reset_for_fold().await`):
```rust
if let Err(e) = self.restore_best_gpu_params() {
tracing::debug!(
"Fold-transition restore-best skipped (no snapshot yet): {e}"
);
} else {
tracing::info!(
"Fold-transition: restored best-Sharpe GPU params before \
shrink-and-perturb (fold {} → {})",
fold_idx, fold_idx + 1,
);
}
// Then the existing reset_for_fold call runs:
// 1. shrink_and_perturb(α=0.8, σ=0.01) on (now-best) params_flat
// 2. sync_target_from_online() — target ← (now-best) online
// 3. reset_adam_state() on every optimizer
self.reset_for_fold().await?;
```
`restore_best_gpu_params` is a thin wrapper over
`fused_ctx.restore_best_params`, which:
1. DtoD-copies `best_params_snapshot → params_buf`
2. Unflattens into individual weight tensors
3. Records a stream event (no host sync)
`reset_for_fold` runs on the same trainer stream, so the DtoD copy +
unflatten complete before `shrink_and_perturb` reads `params_flat`.
### Cold-start guard
If `best_params_snapshot` is uninitialized (first fold or any fold
without a Sharpe improvement yet), `restore_best_params` returns
`Err("No best params snapshot saved")`. The fix swallows that error
with a `debug!` log — fold transition then runs S&P on whatever
`params_flat` currently holds, matching the pre-fix behavior. No
regression on the cold-start path.
### State interaction with `reset_for_fold`
`restore_best_gpu_params` writes `params_flat` ONLY (online weights);
it does NOT restore Adam state (m_buf/v_buf), target params, or any
other optimizer state. `reset_for_fold` then:
- Calls `shrink_and_perturb` on the restored (best) `params_flat`
- Calls `sync_target_from_online()` — hard-sync target ← restored online
- Resets Adam state on main DQN, IQN head, TLOB, IQL, IQL-low,
attention, VSN — fresh momentum from zero across all optimizers
These are non-conflicting: restore overwrites online weights; reset
clears the OTHER state categories (Adam, target). Order is correct.
### `best_params_snapshot` lifecycle
The snapshot lives on `FusedTrainingCtx` (constructor-init `None`,
allocated lazily on first `save_best_params` call) and is **NEVER
cleared at fold boundary**. The trainer-side `self.best_sharpe` IS
reset to `NEG_INFINITY` in `reset_for_fold` (so the first improvement in
fold N+1 will overwrite the snapshot), but until then, the snapshot
holds whichever fold most recently saved a peak.
This is the intended training-scoped behavior: peak weights from any
fold can serve as the next fold's starting point. If strict
within-fold semantics are wanted, the snapshot would need explicit
clear at the fold boundary — flagged as a follow-up consideration.
### Behavioral test
`crates/ml/tests/sp18_fold_transition_test.rs` (NEW) — CPU oracle
encoding the math contract:
- `shrink_and_perturb(α, σ) on X = α × X + (1−α) × N(0, σ)`
- with-fix: `α × W_best + (1−α) × N`
- without-fix: `α × W_curr + (1−α) × N`
- diff: `α × (W_best − W_curr)`
The test asserts that with-fix ≠ without-fix when the best snapshot
differs from current weights, and that the residual `out − α × W_best`
falls within Wilson 95% CI of `(1−α) × σ × √N` for the noise term.
Cold-start case (best == current) is bit-identical between the two
orderings — pins the no-snapshot fallback as a safe no-op.
GPU-level kernel coverage stays in the existing
`compile_training_kernels` smoke + L40S 30-epoch validation that
gates the SP18 Phase 1 dispatch.
### Cross-pearl invariants preserved
- `feedback_no_partial_refactor`: call-ordering fix + behavioral test
+ audit-doc update + pre-commit Invariant 7 satisfied in one atomic
commit.
- `feedback_no_legacy_aliases`: no new `_via_pinned` helpers; reuses
existing `restore_best_gpu_params` API unchanged.
- `feedback_wire_everything_up`: the previously single-consumer
`restore_best_gpu_params` (used only by PLATEAU_EXHAUSTED) gains a
second cold-path consumer at the fold transition.
- `feedback_no_htod_htoh_only_mapped_pinned`: the underlying DtoD
copy + unflatten is mapped-pinned-compatible; no new HtoD/HtoH paths
introduced.
- `pearl_no_host_branches_in_captured_graph`: the restore call runs
outside CUDA Graph capture (epoch boundary, between folds).
### Files changed
| File | LOC delta | Purpose |
|------|-----------|---------|
| `crates/ml/src/trainers/dqn/trainer/mod.rs` | +49 / -0 | Insert `restore_best_gpu_params()` before `reset_for_fold()` at fold-transition site, with cold-start guard + commentary |
| `crates/ml/tests/sp18_fold_transition_test.rs` | NEW | CPU oracle test for the math contract (3 cases) |
| `docs/dqn-wire-up-audit.md` | +N / -0 | This entry |