From 37c3a8f4d72ed1c58fe6e2972e98f902cf1c1bc7 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 17 May 2026 16:48:37 +0200 Subject: [PATCH] feat(ml-alpha): ISV-driven per-horizon EMA + lambda (Phase 1+2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Foundation for replacing the static `--auto-horizon-weights` formula (`min(1, K/h)`) with a signal-driven per-horizon gradient scaler. Per `feedback_isv_for_adaptive_bounds.md`: adaptive bounds live in ISV, not hardcoded constants. Per `pearl_adam_normalizes_loss_weights.md`: Adam normalizes per-loss weight lifts (SP13 saw 13× aux_w produce only 0.6%/epoch divergence), so the effective lever is scaling the GRADIENT into the shared trunk, not the BCE coefficient. This commit sets up the EMA + lambda infrastructure; Phase 3 (wiring lambda into heads_bwd to actually scale the trunk gradient) is gated on the 3-fold CV results from eb51c0f9c. Phase 1 — BCE kernel emits per-horizon UNWEIGHTED mean BCE: cuda/bce_loss_multi_horizon.cu: - New output buffer `loss_per_horizon[N_HORIZONS=5]`. - Per-horizon shared-mem accumulators (sloss_h, svalid_h) with block tree-reduce — no atomicAdd, per `feedback_no_atomicadd.md`. - Hardcoded N_HORIZONS_BCE=5; total shared-mem usage ~13 KiB (comfortable under any SM smem limit). - The aggregate `loss_out` is still the externally-weighted mean callers use for reporting; the new buffer is the UNWEIGHTED signal an EMA layer needs. Phase 2 — EMA + lambda kernel: cuda/horizon_lambda.cu (new): - Single-thread kernel (5 horizons, fixed-size loop — trivial). - First-observation bootstrap via sentinel = 0 per `pearl_first_observation_bootstrap.md`; replaces directly when `loss_ema_h <= 0` (safer than `== 0` under --use_fast_math). - Fixed α = 0.1 EMA for now; Wiener-optimal α follow-up flagged (`pearl_wiener_optimal_adaptive_alpha.md`). - lambda_h = clamp(loss_ema_h / mean(loss_ema), 0.5, 2.0). - Ratio gives natural "under-trained → boost" signal. - Clamp prevents winner-take-all per `pearl_controller_amplifies_dominant_magnitude_trap.md` and bounded-modifier safety per `pearl_audit_unboundedness_for_implicit_asymmetry.md`. Trainer wiring (trainer/perception.rs): - 3 new fields: `loss_per_horizon_d`, `loss_ema_d`, `lambda_d` (all 5-element f32 CudaSlices; pre-allocated, zero-initialised). - Cached `horizon_lambda_fn` + module handle per the `BiasKernels`-style pattern (no per-call cuModuleLoadData). - BCE callsite (train + eval paths) updated to pass `loss_per_horizon_d`. - `dispatch_train_step` launches `horizon_ema_and_lambda` right after BCE, BEFORE the K-loop backward. Inside the captured graph; per-step launch overhead is ~1 µs. - `loss_ema_snapshot()` + `lambda_snapshot()` test-only accessors (mapped-pinned readback, not for hot path) for diagnostics. Smoke test — `horizon_ema_and_lambda_track_after_training`: - Verifies pre-step EMA + lambda are zero (sentinel). - After 5 training steps: loss_ema = [0.59, 0.66, 0.45, 0.62, 0.50] — finite + positive. lambda = [1.05, 1.16, 0.80, 1.10, 0.89] — mean ≈ 1.0, all inside the [0.5, 2.0] clamp envelope. - Lower per-horizon BCE → lower lambda (de-emphasize); higher BCE → higher lambda (boost). Exactly the ISV semantics we want. Validation: 6 perception_overfit tests pass, synthetic overfit still shrinks (0.33 → 0.0006), 26 ml-alpha lib + 23 integration tests green. lambda_d is computed every step but NOT YET CONSUMED by heads_bwd; training behavior is bit-identical to 3a196382f. Phase 3 (consume lambda_d in heads_bwd_batched to scale the per-horizon gradient into the trunk) follows once CV confirms the foundation is stable. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/build.rs | 1 + .../ml-alpha/cuda/bce_loss_multi_horizon.cu | 69 ++++++++---- crates/ml-alpha/cuda/horizon_lambda.cu | 74 +++++++++++++ crates/ml-alpha/src/trainer/perception.rs | 103 +++++++++++++++++- crates/ml-alpha/tests/perception_overfit.rs | 54 +++++++++ 5 files changed, 279 insertions(+), 22 deletions(-) create mode 100644 crates/ml-alpha/cuda/horizon_lambda.cu diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index b58fe184c..2de5f52f7 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -16,6 +16,7 @@ const KERNELS: &[&str] = &[ "bce_loss_multi_horizon", "adamw_step", "grad_norm", + "horizon_lambda", // ISV-driven per-horizon gradient scaler (EMA + lambda) ]; // Cache bust v3 (2026-05-17): batched cfc/heads kernels + transpose_3d_swap_01 diff --git a/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu b/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu index 6be6f8958..662ab5098 100644 --- a/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu +++ b/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu @@ -1,33 +1,38 @@ // bce_loss_multi_horizon.cu // // Fused multi-horizon BCE forward + backward with optional per-horizon -// loss weighting. +// loss weighting. Also emits per-horizon UNWEIGHTED mean BCE — that +// is the signal an ISV-driven horizon weighting layer EMA-tracks to +// produce dynamic per-horizon weights (high per-horizon BCE → horizon +// is under-trained → boost its weight). The total scalar loss_out +// remains the externally-weighted aggregate used by callers for +// reporting and early-stopping. // // Input layout (row-major): // probs [n_pos, n_horizons] — model outputs in (0, 1) // labels [n_pos, n_horizons] — binary {0.0, 1.0}; NaN = mask (drop) // loss_weights [n_horizons] — per-horizon weight (nullptr = all 1.0) -// Down-weights horizons whose forward -// labels are correlated across K positions -// (e.g. h ≫ K → w_h ≪ 1 to avoid -// inflating gradient pressure on -// near-identical labels). // Output: -// loss_out[1] — weighted mean BCE over valid entries -// grad_probs[n_pos, n_horizons]— d loss / d prob (per-element scaled) -// valid_count_out[1] — number of non-NaN labels +// loss_out[1] — weighted mean BCE over valid entries +// loss_per_horizon[n_horizons] — UNWEIGHTED mean BCE per horizon (ISV signal) +// grad_probs[n_pos, n_horizons] — d loss / d prob (per-element scaled) +// valid_count_out[1] — number of non-NaN labels (unweighted) // // per element (mask m = !isnan(y), weight w = loss_weights[h]): // p = clamp(probs[i], 1e-6, 1 - 1e-6) -// L_i = m * w * (-[y log p + (1-y) log(1-p)]) +// L_i = -[y log p + (1-y) log(1-p)] (unweighted; weighted form scales L_i by w) // dL/dp = m * w * (p - y) / (p (1-p)) // // Normaliser: W_valid = sum_i m_i * w_i -// total_loss = (1/W_valid) * sum_i L_i -// grad_probs[i] = L_i'_unnormalized / W_valid +// total_loss = (1/W_valid) * sum_i (m_i * w_i * L_i) +// grad_probs[i] = m_i * w * (p - y) / (p (1-p)) / W_valid +// loss_per_horizon[h] = (sum_{i: m_i, h_i=h} L_i) / (sum_{i: m_i, h_i=h} 1) // -// With all weights = 1.0 the kernel reduces to the original -// uniform-mean BCE (weighted sum / valid count = unweighted mean). +// Caller contract: n_horizons MUST equal N_HORIZONS_BCE (5). Hardcoded +// rather than dynamic-shared-mem because the per-horizon accumulator +// matrix has a fixed shape (5 × 256) at our trainer's scale. + +#define N_HORIZONS_BCE 5 extern "C" __global__ void bce_multi_horizon_forward_backward( const float* __restrict__ probs, // [n_pos * n_horizons] @@ -35,7 +40,8 @@ extern "C" __global__ void bce_multi_horizon_forward_backward( const float* __restrict__ loss_weights, // [n_horizons] (nullptr OK → 1.0) int n_pos, int n_horizons, - float* __restrict__ loss_out, // [1] + float* __restrict__ loss_out, // [1] — externally-weighted mean + float* __restrict__ loss_per_horizon, // [n_horizons] — UNWEIGHTED per-horizon mean float* __restrict__ grad_probs, // [n_pos * n_horizons] int* __restrict__ valid_count_out // [1] — number of non-NaN labels ) { @@ -45,11 +51,23 @@ extern "C" __global__ void bce_multi_horizon_forward_backward( __shared__ float sloss[256]; // weighted-loss numerator accumulator __shared__ float sw_valid[256]; // weighted-mask normaliser accumulator __shared__ int svalid[256]; // unweighted valid-label counter (reported) + // Per-horizon accumulators (unweighted): sloss_h[h*256 + tid], + // svalid_h[h*256 + tid]. Hardcoded width N_HORIZONS_BCE=5 → 13 KiB + // total shared usage including the three scalars above; comfortably + // under any SM smem limit (48-100 KiB depending on arch). + __shared__ float sloss_h[N_HORIZONS_BCE * 256]; + __shared__ float svalid_h[N_HORIZONS_BCE * 256]; sloss[tid] = 0.0f; sw_valid[tid] = 0.0f; svalid[tid] = 0; + #pragma unroll + for (int h = 0; h < N_HORIZONS_BCE; ++h) { + sloss_h[h * 256 + tid] = 0.0f; + svalid_h[h * 256 + tid] = 0.0f; + } - // First pass: count valid entries and accumulate weighted loss + weight sum. + // First pass: count valid entries and accumulate weighted aggregate + // loss + unweighted per-horizon loss + weight sums. for (int i = tid; i < total; i += blockDim.x) { const float y = labels[i]; if (isnan(y)) { @@ -59,18 +77,26 @@ extern "C" __global__ void bce_multi_horizon_forward_backward( const int h = i % n_horizons; const float w = (loss_weights != nullptr) ? loss_weights[h] : 1.0f; const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f); - sloss[tid] -= w * (y * logf(p) + (1.0f - y) * logf(1.0f - p)); + const float L = -(y * logf(p) + (1.0f - y) * logf(1.0f - p)); + sloss[tid] += w * L; + sloss_h[h * 256 + tid] += L; // UNWEIGHTED — the ISV signal + svalid_h[h * 256 + tid] += 1.0f; sw_valid[tid] += w; svalid[tid] += 1; } __syncthreads(); - // Block tree-reduce on the three accumulators. + // Block tree-reduce on all accumulators (incl. per-horizon arrays). for (int s = blockDim.x / 2; s > 0; s >>= 1) { if (tid < s) { sloss[tid] += sloss[tid + s]; sw_valid[tid] += sw_valid[tid + s]; svalid[tid] += svalid[tid + s]; + #pragma unroll + for (int h = 0; h < N_HORIZONS_BCE; ++h) { + sloss_h[h * 256 + tid] += sloss_h[h * 256 + tid + s]; + svalid_h[h * 256 + tid] += svalid_h[h * 256 + tid + s]; + } } __syncthreads(); } @@ -79,12 +105,15 @@ extern "C" __global__ void bce_multi_horizon_forward_backward( const float w_valid = sw_valid[0]; loss_out[0] = (w_valid > 0.0f) ? sloss[0] / w_valid : 0.0f; valid_count_out[0] = svalid[0]; + #pragma unroll + for (int h = 0; h < N_HORIZONS_BCE; ++h) { + const float c_h = svalid_h[h * 256]; + loss_per_horizon[h] = (c_h > 0.0f) ? sloss_h[h * 256] / c_h : 0.0f; + } } __syncthreads(); // Second pass: emit per-element grad scaled by w_h / W_valid. - // (Re-derive W_valid via shared mem instead of an extra load — we - // already have it in sw_valid[0] from the reduction above.) const float w_valid = sw_valid[0]; const float inv = (w_valid > 0.0f) ? (1.0f / w_valid) : 0.0f; for (int i = tid; i < total; i += blockDim.x) { diff --git a/crates/ml-alpha/cuda/horizon_lambda.cu b/crates/ml-alpha/cuda/horizon_lambda.cu new file mode 100644 index 000000000..cb17837ed --- /dev/null +++ b/crates/ml-alpha/cuda/horizon_lambda.cu @@ -0,0 +1,74 @@ +// horizon_lambda.cu — ISV-driven per-horizon gradient scaler. +// +// Step 1: maintain an EMA of the UNWEIGHTED per-horizon BCE loss that +// the BCE kernel emits each training step into +// `loss_per_horizon[N_HORIZONS]`. +// Step 2: convert the EMA into a per-horizon multiplicative lambda +// used by the backward path to scale how strongly each +// horizon influences the shared trunk gradient. +// +// Why ISV: the current static `auto-horizon-weights` formula +// (`min(1, K/h)`) is a closed-form heuristic that ignores actual +// per-horizon learning difficulty. Empirically mhzs7 spent most of +// training over-weighting short horizons (whose label correlation +// within the K-snapshot window dominates the gradient signal) while +// h6000 — the deployment-relevant multi-minute horizon — stayed at +// AUC≈0.69. Tracking per-horizon BCE directly lets the lambda boost +// the horizons that the model is currently failing to learn, without +// hand-tuned constants. +// +// Why not just lift BCE coefficients: per +// `pearl_adam_normalizes_loss_weights.md`, Adam's m/sqrt(v) cancels +// per-loss weight lifts (SP13: 13× aux_w produced only 0.6%/epoch +// divergence). The effective lever is to scale the GRADIENT into the +// shared trunk, not the loss aggregate. heads_bwd will multiply the +// per-horizon `d_z` contribution by lambda[h] before accumulating +// into `grad_h`, bypassing Adam normalization. +// +// First-observation bootstrap: loss_ema is zero-initialised; the +// kernel detects `prev <= 0` and replaces (rather than blends) on +// the first step. After that it uses a fixed α — Wiener-optimal α is +// a Phase 3 follow-up; for now a conservative 0.1 keeps the EMA +// stable across training noise. +// +// Lambda safety: clamped to `[LAMBDA_FLOOR, LAMBDA_CEILING]` per +// `pearl_audit_unboundedness_for_implicit_asymmetry.md` — the ratio +// loss_ema_h / mean_loss_ema is unbounded above when one horizon +// stalls. Capping at 2× prevents winner-take-all amplification per +// `pearl_controller_amplifies_dominant_magnitude_trap.md`. + +#define N_HORIZONS_LAMBDA 5 +#define ALPHA_FIXED 0.1f +#define LAMBDA_FLOOR 0.5f +#define LAMBDA_CEILING 2.0f + +extern "C" __global__ void horizon_ema_and_lambda( + const float* __restrict__ loss_per_horizon, // [5] — current step UNWEIGHTED BCE + float* __restrict__ loss_ema, // [5] — EMA state (read + write) + float* __restrict__ lambda // [5] — output multiplier +) { + if (threadIdx.x != 0 || blockIdx.x != 0) return; + + float new_ema[N_HORIZONS_LAMBDA]; + float sum = 0.0f; + #pragma unroll + for (int h = 0; h < N_HORIZONS_LAMBDA; ++h) { + const float cur = loss_per_horizon[h]; + const float prev = loss_ema[h]; + // Sentinel = 0 ⇒ first observation replaces directly. + // `prev <= 0` is safer than `== 0` under --use_fast_math. + new_ema[h] = (prev <= 0.0f) ? cur : (prev + ALPHA_FIXED * (cur - prev)); + loss_ema[h] = new_ema[h]; + sum += new_ema[h]; + } + + const float mean = sum / (float)N_HORIZONS_LAMBDA; + // mean > 1e-12 is essentially always true once the first step + // has run (BCE for a random init is ~0.69), but guard anyway. + const float inv_mean = (mean > 1e-12f) ? (1.0f / mean) : 0.0f; + #pragma unroll + for (int h = 0; h < N_HORIZONS_LAMBDA; ++h) { + const float ratio = (inv_mean > 0.0f) ? new_ema[h] * inv_mean : 1.0f; + lambda[h] = fmaxf(LAMBDA_FLOOR, fminf(LAMBDA_CEILING, ratio)); + } +} diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 53bfea01d..d8a123607 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -58,6 +58,7 @@ const SNAP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/snap_feature const STEP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cfc_step.cubin")); const HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horizon_heads.cubin")); const BCE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bce_loss_multi_horizon.cubin")); +const HORIZON_LAMBDA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/horizon_lambda.cubin")); #[derive(Clone, Debug)] pub struct PerceptionTrainerConfig { @@ -225,6 +226,26 @@ pub struct PerceptionTrainer { /// Mapped-pinned host shadow of `loss_d`. After the post-step sync, /// host_ptr holds the freshly-computed loss — no extra dtoh sync. loss_host_d: MappedF32Buffer, + /// Per-horizon UNWEIGHTED mean BCE — the ISV signal an EMA layer + /// tracks to compute dynamic per-horizon weights. Refreshed by + /// the BCE kernel on every training step. Shape: [N_HORIZONS]. + loss_per_horizon_d: CudaSlice, + /// Per-horizon EMA of `loss_per_horizon_d`. Zero-initialised + /// (sentinel); the horizon_ema_and_lambda kernel detects the + /// sentinel on step 1 and replaces directly per + /// `pearl_first_observation_bootstrap.md`. + loss_ema_d: CudaSlice, + /// Per-horizon multiplicative gradient scaler. Updated each step + /// by horizon_ema_and_lambda. Clamped to [0.5, 2.0]. Wired into + /// heads_bwd in a follow-up commit; until then this buffer is + /// computed but unused — keeps the EMA infrastructure runnable + /// and validatable without changing training math. + lambda_d: CudaSlice, + /// Cached function handle for the horizon_ema_and_lambda kernel. + horizon_lambda_fn: CudaFunction, + /// Module that owns `horizon_lambda_fn`; kept alive so the function + /// handle stays valid for the trainer's lifetime. + _horizon_lambda_module: Arc, /// Valid-label count for BCE normalisation. valid_d: CudaSlice, /// Per-horizon BCE loss weights, device-resident. Filled from @@ -289,6 +310,12 @@ impl PerceptionTrainer { let step_module = ctx.load_cubin(STEP_CUBIN.to_vec()).context("step cubin")?; let heads_module = ctx.load_cubin(HEADS_CUBIN.to_vec()).context("heads cubin")?; let bce_module = ctx.load_cubin(BCE_CUBIN.to_vec()).context("bce cubin")?; + let horizon_lambda_module = ctx + .load_cubin(HORIZON_LAMBDA_CUBIN.to_vec()) + .context("horizon_lambda cubin")?; + let horizon_lambda_fn = horizon_lambda_module + .load_function("horizon_ema_and_lambda") + .context("horizon_ema_and_lambda symbol")?; let snap_batched_fn = snap_module.load_function("snap_feature_assemble_batched")?; let bce_fn = bce_module.load_function("bce_multi_horizon_forward_backward")?; let step_batched_fn = step_module.load_function("cfc_step_batched")?; @@ -383,6 +410,11 @@ impl PerceptionTrainer { loss_weights_d: upload(&stream, &cfg.horizon_weights)?, loss_d: stream.alloc_zeros::(1)?, loss_host_d: unsafe { MappedF32Buffer::new(1) }.map_err(|e| anyhow::anyhow!("loss_host_d: {e}"))?, + loss_per_horizon_d: stream.alloc_zeros::(N_HORIZONS)?, + loss_ema_d: stream.alloc_zeros::(N_HORIZONS)?, + lambda_d: stream.alloc_zeros::(N_HORIZONS)?, + horizon_lambda_fn, + _horizon_lambda_module: horizon_lambda_module, valid_d: stream.alloc_zeros::(1)?, grad_h_carry_d: stream.alloc_zeros::(cfg.n_batch * n_hid)?, grad_h_new_d: stream.alloc_zeros::(cfg.n_batch * n_hid)?, @@ -513,6 +545,48 @@ impl PerceptionTrainer { /// /// Returns mean BCE over the valid (b, position, horizon) triples /// in the batch, weighted by `cfg.horizon_weights`. + /// Test/diagnostic accessor for the per-horizon BCE EMA buffer. + /// Forces a stream sync + mapped-pinned readback — NOT for the + /// hot training path. Used by smoke tests to verify the + /// `horizon_ema_and_lambda` kernel is tracking correctly. + pub fn loss_ema_snapshot(&self) -> Result<[f32; N_HORIZONS]> { + let staging = unsafe { MappedF32Buffer::new(N_HORIZONS) } + .map_err(|e| anyhow::anyhow!("loss_ema staging: {e}"))?; + unsafe { + let (src_ptr, _g) = self.loss_ema_d.device_ptr(&self.stream); + cudarc::driver::result::memcpy_dtod_async( + staging.dev_ptr, src_ptr, + N_HORIZONS * std::mem::size_of::(), + self.stream.cu_stream(), + ).context("loss_ema dtod")?; + } + self.stream.synchronize().context("loss_ema sync")?; + let host = staging.read_all(); + let mut out = [0.0_f32; N_HORIZONS]; + out.copy_from_slice(&host[..N_HORIZONS]); + Ok(out) + } + + /// Test/diagnostic accessor for the per-horizon lambda buffer. + /// Same constraints as [`Self::loss_ema_snapshot`]. + pub fn lambda_snapshot(&self) -> Result<[f32; N_HORIZONS]> { + let staging = unsafe { MappedF32Buffer::new(N_HORIZONS) } + .map_err(|e| anyhow::anyhow!("lambda staging: {e}"))?; + unsafe { + let (src_ptr, _g) = self.lambda_d.device_ptr(&self.stream); + cudarc::driver::result::memcpy_dtod_async( + staging.dev_ptr, src_ptr, + N_HORIZONS * std::mem::size_of::(), + self.stream.cu_stream(), + ).context("lambda dtod")?; + } + self.stream.synchronize().context("lambda sync")?; + let host = staging.read_all(); + let mut out = [0.0_f32; N_HORIZONS]; + out.copy_from_slice(&host[..N_HORIZONS]); + Ok(out) + } + pub fn step_batched( &mut self, snapshots_batch: &[&[Mbp10RawInput]], @@ -883,11 +957,34 @@ impl PerceptionTrainer { .arg(&self.probs_per_k_d).arg(&self.labels_per_k_d) .arg(&self.loss_weights_d) .arg(&n_pos_i).arg(&n_h_i) - .arg(&mut self.loss_d).arg(&mut self.grad_probs_per_k_d) + .arg(&mut self.loss_d) + .arg(&mut self.loss_per_horizon_d) + .arg(&mut self.grad_probs_per_k_d) .arg(&mut self.valid_d); unsafe { launch.launch(bce_cfg).context("bce launch")?; } } + // ── 5b. ISV-driven per-horizon EMA + lambda. Updates the EMA + // of unweighted per-horizon BCE and emits a clamped + // multiplier `lambda_d[h]` per + // `pearl_adam_normalizes_loss_weights.md` strategy of + // scaling effective gradient instead of loss weight. + // Currently captured into the graph; lambda_d will be + // consumed by heads_bwd in the next commit. + { + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.horizon_lambda_fn); + launch + .arg(&self.loss_per_horizon_d) + .arg(&mut self.loss_ema_d) + .arg(&mut self.lambda_d); + unsafe { launch.launch(cfg).context("horizon_ema_and_lambda launch")?; } + } + // ── 6. Reverse-order backward K loop using pre-allocated // grad_h_enriched_seq_t_d as the per-K slot output. self.stream.memset_zeros(&mut self.grad_h_carry_d) @@ -1264,7 +1361,9 @@ impl PerceptionTrainer { launch .arg(&self.probs_per_k_d).arg(&self.labels_per_k_d).arg(&self.loss_weights_d) .arg(&n_pos_i).arg(&n_h_i) - .arg(&mut self.loss_d).arg(&mut self.grad_probs_per_k_d).arg(&mut self.valid_d); + .arg(&mut self.loss_d) + .arg(&mut self.loss_per_horizon_d) + .arg(&mut self.grad_probs_per_k_d).arg(&mut self.valid_d); launch.launch(bce_cfg).context("eval bce")?; } self.stream.synchronize().context("eval sync")?; diff --git a/crates/ml-alpha/tests/perception_overfit.rs b/crates/ml-alpha/tests/perception_overfit.rs index ef56f708e..105d2e459 100644 --- a/crates/ml-alpha/tests/perception_overfit.rs +++ b/crates/ml-alpha/tests/perception_overfit.rs @@ -218,6 +218,60 @@ fn evaluate_works_after_capture_no_replay() { assert!(loss.is_finite(), "eval loss must be finite, got {loss}"); } +/// ISV-driven per-horizon EMA + lambda — after a few training steps, +/// `loss_ema` should contain positive BCE values for every horizon +/// (proves the BCE kernel writes the per-horizon channel correctly), +/// and `lambda` should land inside [0.5, 2.0] (clamp invariant). +#[test] +fn horizon_ema_and_lambda_track_after_training() { + let dev = test_device(); + let cfg = PerceptionTrainerConfig { + seq_len: 16, + mamba2_state_dim: 8, + lr_cfc: 3e-3, + lr_mamba2: 1e-3, + seed: 0x9292, + horizon_weights: [1.0; 5], + n_batch: 1, + }; + let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); + + // EMA / lambda before any training step. + let ema0 = trainer.loss_ema_snapshot().expect("ema initial"); + let lam0 = trainer.lambda_snapshot().expect("lambda initial"); + assert!(ema0.iter().all(|v| *v == 0.0), "ema should be zero-init sentinel, got {ema0:?}"); + assert!(lam0.iter().all(|v| *v == 0.0), "lambda should be zero before first kernel run, got {lam0:?}"); + + let mut ts = 1_000_000u64; + let mut prev_mid = 5500.0_f32; + for _ in 0..5 { + let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts); + trainer.step(&seq, labels.as_slice()).expect("train step"); + prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]); + ts = seq.last().unwrap().ts_ns; + } + + let ema = trainer.loss_ema_snapshot().expect("ema after train"); + let lam = trainer.lambda_snapshot().expect("lambda after train"); + eprintln!("loss_ema = {ema:?}"); + eprintln!("lambda = {lam:?}"); + // Each horizon must have a positive finite EMA (BCE > 0 unless model is perfect). + for (h, &v) in ema.iter().enumerate() { + assert!(v.is_finite(), "ema[{h}] not finite: {v}"); + assert!(v > 0.0, "ema[{h}] should be positive, got {v}"); + } + // Lambda clamp: every entry in [0.5, 2.0]. + for (h, &v) in lam.iter().enumerate() { + assert!(v.is_finite(), "lambda[{h}] not finite: {v}"); + assert!(v >= 0.5 - 1e-6 && v <= 2.0 + 1e-6, + "lambda[{h}] = {v} outside clamp [0.5, 2.0]"); + } + // Mean of lambda should be ≈ 1.0 (since lambda is ratio to mean, clamped). + let lam_mean: f32 = lam.iter().sum::() / lam.len() as f32; + assert!((lam_mean - 1.0).abs() < 0.5, + "lambda mean {lam_mean} should be close to 1.0 (clamp envelope, ratio-of-mean)"); +} + /// Does the bug surface after just ONE step (warmup only, no capture)? /// If yes, the warmup dispatch path itself breaks subsequent eval. /// If no, the captured graph instantiation / launch is what breaks eval.