feat(ml-alpha): ISV-driven per-horizon EMA + lambda (Phase 1+2)
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 fromeb51c0f9c. 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 to3a196382f. 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 <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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) {
|
||||
|
||||
74
crates/ml-alpha/cuda/horizon_lambda.cu
Normal file
74
crates/ml-alpha/cuda/horizon_lambda.cu
Normal file
@@ -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));
|
||||
}
|
||||
}
|
||||
@@ -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<f32>,
|
||||
/// 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<f32>,
|
||||
/// 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<f32>,
|
||||
/// 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<CudaModule>,
|
||||
/// Valid-label count for BCE normalisation.
|
||||
valid_d: CudaSlice<i32>,
|
||||
/// 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::<f32>(1)?,
|
||||
loss_host_d: unsafe { MappedF32Buffer::new(1) }.map_err(|e| anyhow::anyhow!("loss_host_d: {e}"))?,
|
||||
loss_per_horizon_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
|
||||
loss_ema_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
|
||||
lambda_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
|
||||
horizon_lambda_fn,
|
||||
_horizon_lambda_module: horizon_lambda_module,
|
||||
valid_d: stream.alloc_zeros::<i32>(1)?,
|
||||
grad_h_carry_d: stream.alloc_zeros::<f32>(cfg.n_batch * n_hid)?,
|
||||
grad_h_new_d: stream.alloc_zeros::<f32>(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::<f32>(),
|
||||
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::<f32>(),
|
||||
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")?;
|
||||
|
||||
@@ -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::<f32>() / 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.
|
||||
|
||||
Reference in New Issue
Block a user