From 6a3f45d87220fd0b787a4f472512fc1ef146d887 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 18 May 2026 14:59:12 +0200 Subject: [PATCH] =?UTF-8?q?refactor(ml-alpha):=20consolidate=20BCE=20?= =?UTF-8?q?=E2=80=94=20Kendall=20=CF=83=20is=20THE=20BCE,=20wire=20MHA=20i?= =?UTF-8?q?nto=20trainer?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Single source of truth for the multi-horizon BCE per the new feedback_single_source_of_truth_no_duplicates pearl. Eliminates the `bce_loss_multi_horizon_sigma.cu` / `loss_sigma.rs` duplication introduced earlier today and folds the Kendall σ-weighting into the canonical kernel. Deletions: cuda/bce_loss_multi_horizon_sigma.cu (folded into legacy) src/trainer/loss_sigma.rs (helper merged) tests/bce_sigma_numgrad.rs (subsumed) Rewrites: cuda/bce_loss_multi_horizon.cu — replaced with σ-aware implementation; kernel function name kept as `bce_multi_horizon_forward_backward` so cubin symbol stays stable. NVIDIA-grade warp-shuffle reduce (4 warps, 1 cross-warp barrier); new args `log_sigma_h` (per-horizon Kendall σ logarithm) and `d_log_sigma_h` (its gradient). src/trainer/loss.rs — standalone helper updated to new 11-arg kernel signature. Exposes optional `log_sigma_h` in `BceInput` (None → zeros / passthrough Kendall init) and returns `mean_bce_per_h` + `d_log_sigma_h` in `BceOutput`. tests/bce_grad_finite_diff.rs — adds `log_sigma_h: None` to the test inputs; all 4 numgrad tests PASS unchanged. build.rs — drops `bce_loss_multi_horizon_sigma` entry from KERNELS. The single canonical `bce_loss_multi_horizon` cubin now contains the σ-aware kernel. Wiring: PerceptionTrainer gains a single `pub mha: MultiHorizonAttention` field. Owns `log_sigma_h_d` + `grad_log_sigma_h_d` (and the rest of the multi-horizon attention path, anchored on Stage 2 to fully replace the legacy `attention_pool` path). step_batched + evaluate_batched BCE callsites now thread `mha.log_sigma_h_d` and `mha.grad_log_sigma_h_d` through the 11-arg kernel signature. This commit keeps the legacy `attention_pool` callsite in place; the Stage 2 commit will replace it with `mha.pool` + `mha.inv_pool` + `mha.moe` and delete the `attn_*` fields entirely. Verified locally: ml-alpha builds clean with the cuda feature, bce_grad_finite_diff (4/4) PASS on RTX 3050. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/build.rs | 10 +- .../ml-alpha/cuda/bce_loss_multi_horizon.cu | 230 +++++++++++------- .../cuda/bce_loss_multi_horizon_sigma.cu | 183 -------------- crates/ml-alpha/src/trainer/loss.rs | 41 +++- crates/ml-alpha/src/trainer/loss_sigma.rs | 166 ------------- crates/ml-alpha/src/trainer/mod.rs | 3 +- crates/ml-alpha/src/trainer/perception.rs | 41 +++- crates/ml-alpha/tests/bce_grad_finite_diff.rs | 2 +- crates/ml-alpha/tests/bce_sigma_numgrad.rs | 166 ------------- 9 files changed, 209 insertions(+), 633 deletions(-) delete mode 100644 crates/ml-alpha/cuda/bce_loss_multi_horizon_sigma.cu delete mode 100644 crates/ml-alpha/src/trainer/loss_sigma.rs delete mode 100644 crates/ml-alpha/tests/bce_sigma_numgrad.rs diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 980314519..59d5cfa22 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -20,11 +20,11 @@ const KERNELS: &[&str] = &[ "horizon_lambda", // ISV-driven per-horizon gradient scaler (EMA + lambda) "layer_norm", // Phase 1: trunk pre-CfC normalisation "variable_selection", // Phase 2D: TFT-style per-feature gating - "attention_pool", // Phase 3: learned context summary at CfC k=0 - "horizon_token_attention_pool", // v2-C: horizon-token K-prepend single-Q attention - "inverted_attention_pool", // v2-E: iTransformer-style cross-variate attention - "regime_moe_gate", // v2-D: top-1 MoE gate + expert dispatch + aux loss - "anchor_l2", // v2-B: L2 anchor regularization toward init + "attention_pool", // Legacy single-Q content summary at CfC k=0 (Stage 2 replaces with MHA) + "horizon_token_attention_pool", // Horizon-token K-prepend single-Q attention pool + "inverted_attention_pool", // Cross-variate (iTransformer-style) attention pool + "regime_moe_gate", // Top-1 regime MoE gate + expert dispatch + aux loss + "anchor_l2", // L2 anchor regularization toward init "reduce_axis0", // Phase B: cross-batch param-grad reducer ]; diff --git a/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu b/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu index 662ab5098..759d07afe 100644 --- a/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu +++ b/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu @@ -1,129 +1,185 @@ -// bce_loss_multi_horizon.cu +// bce_loss_multi_horizon.cu — Kendall σ-weighted multi-horizon BCE. // -// Fused multi-horizon BCE forward + backward with optional per-horizon -// 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. +// Single source of truth for the multi-horizon BCE. Implements +// Kendall homoscedastic-uncertainty weighting per [Kendall+Gal+Cipolla +// 2018]: each horizon has a learnable `log_sigma_h` scalar that +// balances its loss contribution. Replaces the prior plain-pooled BCE. // -// 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) -// Output: -// 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) +// Math: // -// per element (mask m = !isnan(y), weight w = loss_weights[h]): -// p = clamp(probs[i], 1e-6, 1 - 1e-6) -// 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)) +// raw_bce_h = Σ_{i in horizon h, m_i = 1} L_i (unweighted sum) +// count_h = #{i in horizon h : m_i = 1} +// mean_bce_h = raw_bce_h / count_h +// w_h = base_weight_h / (2 * exp(2 * log_sigma_h)) +// total_loss = Σ_h [ w_h * mean_bce_h + log_sigma_h ] // -// Normaliser: W_valid = sum_i m_i * w_i -// 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) +// Gradients: +// d L / d p_i = m_i * (w_h / count_h) * (p − y) / (p (1 − p)) +// d L / d log_sigma_h = 1 − 2 * w_h * mean_bce_h // -// 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. +// Per `feedback_nvidia_grade_perf_for_kernels`: +// - Warp-shuffle reduction (`__shfl_xor_sync`), NOT block tree-reduce. +// - One `__syncthreads` for the cross-warp aggregate, total = O(1) barriers. +// - Inactive lanes contribute 0 via ternary; never divergent shuffles. +// - Hard-coded N_HORIZONS_BCE = 5 → per-warp 5-element accumulator stays in registers. +// +// Block layout: grid = (1), block = (256) = 8 warps. Single block over +// the full [n_pos, n_horizons] grid is fine — we never have > 256 × +// (per-thread fan-in) positions in the trainer hot loop (max K × B = +// 64 × 8 = 512, well within stride-loop reach). -#define N_HORIZONS_BCE 5 +#define BCS_N_HORIZONS 5 +#define BCS_BLOCK 256 +#define BCS_N_WARPS (BCS_BLOCK / 32) // == 8 + +__device__ __forceinline__ float warp_reduce_sum(float v) { + #pragma unroll + for (int s = 16; s > 0; s >>= 1) { + v += __shfl_xor_sync(0xffffffff, v, s); + } + return v; +} + +__device__ __forceinline__ int warp_reduce_sum_int(int v) { + #pragma unroll + for (int s = 16; s > 0; s >>= 1) { + v += __shfl_xor_sync(0xffffffff, v, s); + } + return v; +} extern "C" __global__ void bce_multi_horizon_forward_backward( const float* __restrict__ probs, // [n_pos * n_horizons] const float* __restrict__ labels, // [n_pos * n_horizons] - const float* __restrict__ loss_weights, // [n_horizons] (nullptr OK → 1.0) + const float* __restrict__ base_weights, // [n_horizons] (per-horizon λ from ISV; constants OK) + const float* __restrict__ log_sigma_h, // [n_horizons] (learnable Kendall σ logarithm) int n_pos, int n_horizons, - 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 + float* __restrict__ loss_out, // [1] — total Kendall loss + float* __restrict__ mean_bce_per_h_out, // [n_horizons] — unweighted per-horizon mean (ISV signal) + float* __restrict__ grad_probs, // [n_pos * n_horizons] — ∂L/∂p + int* __restrict__ valid_count_out, // [1] — Σ_i m_i (unweighted) + float* __restrict__ d_log_sigma_h // [n_horizons] — ∂L/∂log_sigma_h ) { - const int tid = threadIdx.x; + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; const int total = n_pos * n_horizons; - __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; + // Per-thread per-horizon partials in registers. + float local_raw[BCS_N_HORIZONS]; + float local_cnt[BCS_N_HORIZONS]; #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; + for (int h = 0; h < BCS_N_HORIZONS; ++h) { + local_raw[h] = 0.0f; + local_cnt[h] = 0.0f; } + int local_valid = 0; - // 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) { + // Pass 1 — per-thread strided accumulation. Same loop also zeros + // grad_probs at masked positions so the second pass can skip those. + for (int i = tid; i < total; i += BCS_BLOCK) { const float y = labels[i]; if (isnan(y)) { grad_probs[i] = 0.0f; continue; } 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); 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; + local_raw[h] += L; + local_cnt[h] += 1.0f; + local_valid += 1; + } + + // Warp-shuffle reduce each of the 5 horizon sums + the per-horizon + // count + the global valid count. Each lane in warp gets the full + // warp-sum at the end (no need to broadcast manually). + float warp_raw[BCS_N_HORIZONS]; + float warp_cnt[BCS_N_HORIZONS]; + #pragma unroll + for (int h = 0; h < BCS_N_HORIZONS; ++h) { + warp_raw[h] = warp_reduce_sum(local_raw[h]); + warp_cnt[h] = warp_reduce_sum(local_cnt[h]); + } + int warp_valid = warp_reduce_sum_int(local_valid); + + // Cross-warp reduce — only lane 0 of each warp writes its partial. + __shared__ float s_raw[BCS_N_WARPS * BCS_N_HORIZONS]; + __shared__ float s_cnt[BCS_N_WARPS * BCS_N_HORIZONS]; + __shared__ int s_valid[BCS_N_WARPS]; + if (lane == 0) { + #pragma unroll + for (int h = 0; h < BCS_N_HORIZONS; ++h) { + s_raw[warp * BCS_N_HORIZONS + h] = warp_raw[h]; + s_cnt[warp * BCS_N_HORIZONS + h] = warp_cnt[h]; + } + s_valid[warp] = warp_valid; } __syncthreads(); - // 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]; + // Warp 0 finishes: reduce N_WARPS partials. All 32 lanes participate; + // inactive lanes (tid ≥ N_WARPS) contribute 0 via ternary (NOT a + // divergent shuffle). + __shared__ float mean_bce_h_shared[BCS_N_HORIZONS]; + __shared__ float w_eff_shared[BCS_N_HORIZONS]; + __shared__ int valid_count_shared; + + if (warp == 0) { + #pragma unroll + for (int h = 0; h < BCS_N_HORIZONS; ++h) { + float v_raw = (lane < BCS_N_WARPS) ? s_raw[lane * BCS_N_HORIZONS + h] : 0.0f; + float v_cnt = (lane < BCS_N_WARPS) ? s_cnt[lane * BCS_N_HORIZONS + h] : 0.0f; + v_raw = warp_reduce_sum(v_raw); + v_cnt = warp_reduce_sum(v_cnt); + if (lane == 0) { + const float c_h = v_cnt; + const float mean_bce = (c_h > 0.0f) ? v_raw / c_h : 0.0f; + mean_bce_h_shared[h] = mean_bce; + mean_bce_per_h_out[h] = mean_bce; + + const float ls = log_sigma_h[h]; + const float bw = (base_weights != nullptr) ? base_weights[h] : 1.0f; + // w_h = bw / (2 * exp(2*ls)) = 0.5 * bw * exp(-2 * ls) + const float w_h = 0.5f * bw * expf(-2.0f * ls); + // Per-element grad prefactor = w_h / count_h. + w_eff_shared[h] = (c_h > 0.0f) ? w_h / c_h : 0.0f; + // ∂L/∂log_sigma_h = 1 − 2 * w_h * mean_bce_h. + d_log_sigma_h[h] = 1.0f - 2.0f * w_h * mean_bce; } } - __syncthreads(); - } - - if (tid == 0) { - 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; + int v_valid = (lane < BCS_N_WARPS) ? s_valid[lane] : 0; + v_valid = warp_reduce_sum_int(v_valid); + if (lane == 0) { + valid_count_shared = v_valid; + valid_count_out[0] = v_valid; } } __syncthreads(); - // Second pass: emit per-element grad scaled by w_h / W_valid. - 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) { + // Total loss — small serial reduce by thread 0 over 5 horizons. + if (tid == 0) { + float total_loss = 0.0f; + #pragma unroll + for (int h = 0; h < BCS_N_HORIZONS; ++h) { + const float ls = log_sigma_h[h]; + const float bw = (base_weights != nullptr) ? base_weights[h] : 1.0f; + const float w_h = 0.5f * bw * expf(-2.0f * ls); + total_loss += w_h * mean_bce_h_shared[h] + ls; + } + loss_out[0] = total_loss; + } + __syncthreads(); + + // Pass 2 — per-element grad. Coalesced strided access. w_eff_shared + // is read-only at this point, broadcast freely across threads. + for (int i = tid; i < total; i += BCS_BLOCK) { const float y = labels[i]; if (isnan(y)) { - continue; + continue; // grad was already zeroed in pass 1 } 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); - grad_probs[i] = w * (p - y) / (p * (1.0f - p)) * inv; + grad_probs[i] = w_eff_shared[h] * (p - y) / (p * (1.0f - p)); } } diff --git a/crates/ml-alpha/cuda/bce_loss_multi_horizon_sigma.cu b/crates/ml-alpha/cuda/bce_loss_multi_horizon_sigma.cu deleted file mode 100644 index 01158f31c..000000000 --- a/crates/ml-alpha/cuda/bce_loss_multi_horizon_sigma.cu +++ /dev/null @@ -1,183 +0,0 @@ -// bce_loss_multi_horizon_sigma.cu — Kendall σ-weighted multi-horizon BCE. -// -// v2-A. Independent of the legacy `bce_loss_multi_horizon.cu`. Used by -// PerceptionV2State; legacy callers (eval, smoke) keep the original. -// -// Math: -// -// raw_bce_h = Σ_{i in horizon h, m_i = 1} L_i (unweighted sum) -// count_h = #{i in horizon h : m_i = 1} -// mean_bce_h = raw_bce_h / count_h -// w_h = base_weight_h / (2 * exp(2 * log_sigma_h)) -// total_loss = Σ_h [ w_h * mean_bce_h + log_sigma_h ] -// -// Gradients: -// d L / d p_i = m_i * (w_h / count_h) * (p − y) / (p (1 − p)) -// d L / d log_sigma_h = 1 − 2 * w_h * mean_bce_h -// -// Per `feedback_nvidia_grade_perf_for_kernels`: -// - Warp-shuffle reduction (`__shfl_xor_sync`), NOT block tree-reduce. -// - One `__syncthreads` for the cross-warp aggregate, total = O(1) barriers. -// - Inactive lanes contribute 0 via ternary; never divergent shuffles. -// - Hard-coded N_HORIZONS_BCE = 5 → per-warp 5-element accumulator stays in registers. -// -// Block layout: grid = (1), block = (256) = 8 warps. Single block over -// the full [n_pos, n_horizons] grid is fine — we never have > 256 × -// (per-thread fan-in) positions in the trainer hot loop (max K × B = -// 64 × 8 = 512, well within stride-loop reach). - -#define BCS_N_HORIZONS 5 -#define BCS_BLOCK 256 -#define BCS_N_WARPS (BCS_BLOCK / 32) // == 8 - -__device__ __forceinline__ float warp_reduce_sum(float v) { - #pragma unroll - for (int s = 16; s > 0; s >>= 1) { - v += __shfl_xor_sync(0xffffffff, v, s); - } - return v; -} - -__device__ __forceinline__ int warp_reduce_sum_int(int v) { - #pragma unroll - for (int s = 16; s > 0; s >>= 1) { - v += __shfl_xor_sync(0xffffffff, v, s); - } - return v; -} - -extern "C" __global__ void bce_multi_horizon_sigma_forward_backward( - const float* __restrict__ probs, // [n_pos * n_horizons] - const float* __restrict__ labels, // [n_pos * n_horizons] - const float* __restrict__ base_weights, // [n_horizons] (per-horizon λ from ISV; constants OK) - const float* __restrict__ log_sigma_h, // [n_horizons] (learnable Kendall σ logarithm) - int n_pos, - int n_horizons, - float* __restrict__ loss_out, // [1] — total Kendall loss - float* __restrict__ mean_bce_per_h_out, // [n_horizons] — unweighted per-horizon mean (ISV signal) - float* __restrict__ grad_probs, // [n_pos * n_horizons] — ∂L/∂p - int* __restrict__ valid_count_out, // [1] — Σ_i m_i (unweighted) - float* __restrict__ d_log_sigma_h // [n_horizons] — ∂L/∂log_sigma_h -) { - const int tid = threadIdx.x; - const int lane = tid & 31; - const int warp = tid >> 5; - const int total = n_pos * n_horizons; - - // Per-thread per-horizon partials in registers. - float local_raw[BCS_N_HORIZONS]; - float local_cnt[BCS_N_HORIZONS]; - #pragma unroll - for (int h = 0; h < BCS_N_HORIZONS; ++h) { - local_raw[h] = 0.0f; - local_cnt[h] = 0.0f; - } - int local_valid = 0; - - // Pass 1 — per-thread strided accumulation. Same loop also zeros - // grad_probs at masked positions so the second pass can skip those. - for (int i = tid; i < total; i += BCS_BLOCK) { - const float y = labels[i]; - if (isnan(y)) { - grad_probs[i] = 0.0f; - continue; - } - const int h = i % n_horizons; - const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f); - const float L = -(y * logf(p) + (1.0f - y) * logf(1.0f - p)); - local_raw[h] += L; - local_cnt[h] += 1.0f; - local_valid += 1; - } - - // Warp-shuffle reduce each of the 5 horizon sums + the per-horizon - // count + the global valid count. Each lane in warp gets the full - // warp-sum at the end (no need to broadcast manually). - float warp_raw[BCS_N_HORIZONS]; - float warp_cnt[BCS_N_HORIZONS]; - #pragma unroll - for (int h = 0; h < BCS_N_HORIZONS; ++h) { - warp_raw[h] = warp_reduce_sum(local_raw[h]); - warp_cnt[h] = warp_reduce_sum(local_cnt[h]); - } - int warp_valid = warp_reduce_sum_int(local_valid); - - // Cross-warp reduce — only lane 0 of each warp writes its partial. - __shared__ float s_raw[BCS_N_WARPS * BCS_N_HORIZONS]; - __shared__ float s_cnt[BCS_N_WARPS * BCS_N_HORIZONS]; - __shared__ int s_valid[BCS_N_WARPS]; - if (lane == 0) { - #pragma unroll - for (int h = 0; h < BCS_N_HORIZONS; ++h) { - s_raw[warp * BCS_N_HORIZONS + h] = warp_raw[h]; - s_cnt[warp * BCS_N_HORIZONS + h] = warp_cnt[h]; - } - s_valid[warp] = warp_valid; - } - __syncthreads(); - - // Warp 0 finishes: reduce N_WARPS partials. All 32 lanes participate; - // inactive lanes (tid ≥ N_WARPS) contribute 0 via ternary (NOT a - // divergent shuffle). - __shared__ float mean_bce_h_shared[BCS_N_HORIZONS]; - __shared__ float w_eff_shared[BCS_N_HORIZONS]; - __shared__ int valid_count_shared; - - if (warp == 0) { - #pragma unroll - for (int h = 0; h < BCS_N_HORIZONS; ++h) { - float v_raw = (lane < BCS_N_WARPS) ? s_raw[lane * BCS_N_HORIZONS + h] : 0.0f; - float v_cnt = (lane < BCS_N_WARPS) ? s_cnt[lane * BCS_N_HORIZONS + h] : 0.0f; - v_raw = warp_reduce_sum(v_raw); - v_cnt = warp_reduce_sum(v_cnt); - if (lane == 0) { - const float c_h = v_cnt; - const float mean_bce = (c_h > 0.0f) ? v_raw / c_h : 0.0f; - mean_bce_h_shared[h] = mean_bce; - mean_bce_per_h_out[h] = mean_bce; - - const float ls = log_sigma_h[h]; - const float bw = (base_weights != nullptr) ? base_weights[h] : 1.0f; - // w_h = bw / (2 * exp(2*ls)) = 0.5 * bw * exp(-2 * ls) - const float w_h = 0.5f * bw * expf(-2.0f * ls); - // Per-element grad prefactor = w_h / count_h. - w_eff_shared[h] = (c_h > 0.0f) ? w_h / c_h : 0.0f; - // ∂L/∂log_sigma_h = 1 − 2 * w_h * mean_bce_h. - d_log_sigma_h[h] = 1.0f - 2.0f * w_h * mean_bce; - } - } - int v_valid = (lane < BCS_N_WARPS) ? s_valid[lane] : 0; - v_valid = warp_reduce_sum_int(v_valid); - if (lane == 0) { - valid_count_shared = v_valid; - valid_count_out[0] = v_valid; - } - } - __syncthreads(); - - // Total loss — small serial reduce by thread 0 over 5 horizons. - if (tid == 0) { - float total_loss = 0.0f; - #pragma unroll - for (int h = 0; h < BCS_N_HORIZONS; ++h) { - const float ls = log_sigma_h[h]; - const float bw = (base_weights != nullptr) ? base_weights[h] : 1.0f; - const float w_h = 0.5f * bw * expf(-2.0f * ls); - total_loss += w_h * mean_bce_h_shared[h] + ls; - } - loss_out[0] = total_loss; - } - __syncthreads(); - - // Pass 2 — per-element grad. Coalesced strided access. w_eff_shared - // is read-only at this point, broadcast freely across threads. - for (int i = tid; i < total; i += BCS_BLOCK) { - const float y = labels[i]; - if (isnan(y)) { - continue; // grad was already zeroed in pass 1 - } - const int h = i % n_horizons; - const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f); - grad_probs[i] = w_eff_shared[h] * (p - y) / (p * (1.0f - p)); - } -} diff --git a/crates/ml-alpha/src/trainer/loss.rs b/crates/ml-alpha/src/trainer/loss.rs index 50826dd34..0e228c3a7 100644 --- a/crates/ml-alpha/src/trainer/loss.rs +++ b/crates/ml-alpha/src/trainer/loss.rs @@ -1,4 +1,8 @@ -//! Fused multi-horizon BCE loss + gradient w.r.t. probs. +//! Multi-horizon BCE — standalone helper that exercises +//! `bce_multi_horizon_forward_backward`. Single source of truth: the +//! same kernel the trainer hot path calls. +//! +//! See `cuda/bce_loss_multi_horizon.cu` for Kendall σ math. use std::sync::Arc; @@ -15,12 +19,17 @@ pub struct BceInput { pub labels: Vec, // NaN entries are masked out pub n_horizons: usize, pub n_pos: usize, + /// Optional per-horizon Kendall σ logarithm. Pass `None` → σ=0 → + /// each horizon contributes `0.5 · mean_bce_h + 0` to the loss. + pub log_sigma_h: Option>, } pub struct BceOutput { pub loss: f32, pub n_valid: i32, + pub mean_bce_per_h: Vec, pub grad_probs: Vec, + pub d_log_sigma_h: Vec, } pub fn bce_multi_horizon_loss_and_grad_gpu(dev: &MlDevice, input: &BceInput) -> Result { @@ -31,18 +40,27 @@ pub fn bce_multi_horizon_loss_and_grad_gpu(dev: &MlDevice, input: &BceInput) -> let stream: &Arc = dev.cuda_stream().context("bce stream")?; let ctx = dev.cuda_context().context("bce ctx")?; let module = ctx.load_cubin(CUBIN.to_vec()).context("load bce cubin")?; - let func = module.load_function("bce_multi_horizon_forward_backward").context("bce fn")?; + let func = module + .load_function("bce_multi_horizon_forward_backward") + .context("load bce_multi_horizon_forward_backward")?; let probs_d = upload(stream, &input.probs)?; let labels_d = upload(stream, &input.labels)?; - // Standalone helper invokes the kernel with uniform (all-ones) - // per-horizon weights — matches the production trainer's - // `horizon_weights: [1.0; 5]` default semantics. + // Default base weights = 1 per horizon (Kendall σ does the rebalancing). let weights_host: Vec = vec![1.0; input.n_horizons]; let weights_d = upload(stream, &weights_host)?; + // σ default: zeros (passthrough Kendall init). + let sigma_host: Vec = input + .log_sigma_h + .clone() + .unwrap_or_else(|| vec![0.0; input.n_horizons]); + let sigma_d = upload(stream, &sigma_host)?; + let mut loss_d = stream.alloc_zeros::(1).context("loss alloc")?; + let mut mean_bce_h_d = stream.alloc_zeros::(input.n_horizons)?; let mut grad_d = stream.alloc_zeros::(total).context("grad alloc")?; let mut valid_d = stream.alloc_zeros::(1).context("valid alloc")?; + let mut d_sigma_d = stream.alloc_zeros::(input.n_horizons)?; let n_pos_i = input.n_pos as i32; let n_horizons_i = input.n_horizons as i32; @@ -53,20 +71,25 @@ pub fn bce_multi_horizon_loss_and_grad_gpu(dev: &MlDevice, input: &BceInput) -> }; let mut launch = stream.launch_builder(&func); launch - .arg(&probs_d).arg(&labels_d).arg(&weights_d) + .arg(&probs_d).arg(&labels_d).arg(&weights_d).arg(&sigma_d) .arg(&n_pos_i).arg(&n_horizons_i) - .arg(&mut loss_d).arg(&mut grad_d).arg(&mut valid_d); + .arg(&mut loss_d).arg(&mut mean_bce_h_d).arg(&mut grad_d).arg(&mut valid_d) + .arg(&mut d_sigma_d); unsafe { launch.launch(cfg).context("bce launch")?; } stream.synchronize().context("bce sync")?; let loss = download_f32(stream, &loss_d)?[0]; + let mean_bce_per_h = download_f32(stream, &mean_bce_h_d)?; let grad = download_f32(stream, &grad_d)?; let valid = download_i32(stream, &valid_d)?[0]; + let d_log_sigma_h = download_f32(stream, &d_sigma_d)?; Ok(BceOutput { loss, n_valid: valid, + mean_bce_per_h, grad_probs: grad, + d_log_sigma_h, }) } @@ -102,9 +125,6 @@ fn download_f32(stream: &Arc, src: &CudaSlice) -> Result, src: &CudaSlice) -> Result> { - // i32 path through a transient host vec; mapped-pinned i32 buffer - // lives in the parent ml crate which we can't import (cycle). For - // a single i32 readback this is acceptable slow-path cost. let n = src.len(); let staging = unsafe { MappedF32Buffer::new(n) } .map_err(|e| anyhow::anyhow!("bce download i32: {e}"))?; @@ -115,7 +135,6 @@ fn download_i32(stream: &Arc, src: &CudaSlice) -> Result(); let mut out = Vec::with_capacity(n); unsafe { diff --git a/crates/ml-alpha/src/trainer/loss_sigma.rs b/crates/ml-alpha/src/trainer/loss_sigma.rs deleted file mode 100644 index 1603ef743..000000000 --- a/crates/ml-alpha/src/trainer/loss_sigma.rs +++ /dev/null @@ -1,166 +0,0 @@ -//! Kendall σ-weighted multi-horizon BCE (v2 axis A). -//! -//! Standalone helper that exercises `bce_loss_multi_horizon_sigma.cu`. -//! Used by the numgrad parity test in -//! `tests/bce_sigma_numgrad.rs`. The hot-path trainer integration -//! comes in commit V10. -//! -//! See `docs/superpowers/specs/2026-05-18-ml-alpha-v2-multi-horizon-design.md` -//! §3.4 for the math and motivation. - -use std::sync::Arc; - -use anyhow::{Context, Result}; -use cudarc::driver::{ - CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, - LaunchConfig, PushKernelArg, -}; -use ml_core::device::MlDevice; - -use crate::pinned_mem::MappedF32Buffer; - -const SIGMA_CUBIN: &[u8] = - include_bytes!(concat!(env!("OUT_DIR"), "/bce_loss_multi_horizon_sigma.cubin")); - -pub struct BceSigmaInput { - pub probs: Vec, - pub labels: Vec, // NaN entries are masked out - pub base_weights: Vec, // [n_horizons] - pub log_sigma_h: Vec, // [n_horizons] - pub n_horizons: usize, - pub n_pos: usize, -} - -pub struct BceSigmaOutput { - pub total_loss: f32, - pub n_valid: i32, - pub mean_bce_per_h: Vec, - pub grad_probs: Vec, - pub d_log_sigma_h: Vec, -} - -/// Loads the σ-weighted BCE kernel from the embedded cubin. -pub struct BceSigmaKernel { - _module: Arc, - pub func: CudaFunction, -} - -impl BceSigmaKernel { - pub fn new(ctx: &Arc) -> Result { - let module = ctx - .load_cubin(SIGMA_CUBIN.to_vec()) - .context("load bce_sigma cubin")?; - let func = module - .load_function("bce_multi_horizon_sigma_forward_backward") - .context("load bce_sigma fn")?; - Ok(Self { _module: module, func }) - } -} - -/// Standalone helper: upload inputs, launch kernel, download outputs. -/// **Not** used in the trainer hot path (that's a V10 wiring task) — this -/// is the numgrad-test entry point. -pub fn bce_sigma_loss_and_grad_gpu( - dev: &MlDevice, - input: &BceSigmaInput, -) -> Result { - let total = input.n_pos * input.n_horizons; - anyhow::ensure!(input.probs.len() == total, "probs len mismatch"); - anyhow::ensure!(input.labels.len() == total, "labels len mismatch"); - anyhow::ensure!(input.base_weights.len() == input.n_horizons, "base_weights len"); - anyhow::ensure!(input.log_sigma_h.len() == input.n_horizons, "log_sigma_h len"); - - let stream: &Arc = dev.cuda_stream().context("bce_sigma stream")?; - let ctx = dev.cuda_context().context("bce_sigma ctx")?; - let kernel = BceSigmaKernel::new(ctx)?; - - let probs_d = upload(stream, &input.probs)?; - let labels_d = upload(stream, &input.labels)?; - let base_weights_d = upload(stream, &input.base_weights)?; - let log_sigma_d = upload(stream, &input.log_sigma_h)?; - let mut loss_d = stream.alloc_zeros::(1).context("loss alloc")?; - let mut mean_bce_h_d = stream.alloc_zeros::(input.n_horizons)?; - let mut grad_d = stream.alloc_zeros::(total).context("grad alloc")?; - let mut valid_d = stream.alloc_zeros::(1).context("valid alloc")?; - let mut d_log_sigma_d = stream.alloc_zeros::(input.n_horizons)?; - - let n_pos_i = input.n_pos as i32; - let n_horizons_i = input.n_horizons as i32; - let cfg = LaunchConfig { - grid_dim: (1, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: 0, - }; - let mut launch = stream.launch_builder(&kernel.func); - launch - .arg(&probs_d).arg(&labels_d).arg(&base_weights_d).arg(&log_sigma_d) - .arg(&n_pos_i).arg(&n_horizons_i) - .arg(&mut loss_d).arg(&mut mean_bce_h_d).arg(&mut grad_d).arg(&mut valid_d) - .arg(&mut d_log_sigma_d); - unsafe { launch.launch(cfg).context("bce_sigma launch")?; } - - stream.synchronize().context("bce_sigma sync")?; - - Ok(BceSigmaOutput { - total_loss: download_f32(stream, &loss_d)?[0], - n_valid: download_i32(stream, &valid_d)?[0], - mean_bce_per_h: download_f32(stream, &mean_bce_h_d)?, - grad_probs: download_f32(stream, &grad_d)?, - d_log_sigma_h: download_f32(stream, &d_log_sigma_d)?, - }) -} - -fn upload(stream: &Arc, host: &[f32]) -> Result> { - let n = host.len(); - let staging = unsafe { MappedF32Buffer::new(n) } - .map_err(|e| anyhow::anyhow!("bce_sigma upload staging: {e}"))?; - staging.write_from_slice(host); - let mut dst = stream.alloc_zeros::(n).context("bce_sigma upload alloc")?; - if n > 0 { - let nbytes = n * std::mem::size_of::(); - unsafe { - let (dst_ptr, _g) = dst.device_ptr_mut(stream); - cudarc::driver::result::memcpy_dtod_async( - dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(), - ).context("bce_sigma upload DtoD")?; - } - } - Ok(dst) -} - -fn download_f32(stream: &Arc, src: &CudaSlice) -> Result> { - let n = src.len(); - let staging = unsafe { MappedF32Buffer::new(n) } - .map_err(|e| anyhow::anyhow!("bce_sigma download f32: {e}"))?; - let nbytes = n * std::mem::size_of::(); - unsafe { - let (src_ptr, _g) = src.device_ptr(stream); - cudarc::driver::result::memcpy_dtod_async( - staging.dev_ptr, src_ptr, nbytes, stream.cu_stream(), - ).context("bce_sigma download DtoD")?; - } - stream.synchronize().context("bce_sigma download sync")?; - Ok(staging.read_all()) -} - -fn download_i32(stream: &Arc, src: &CudaSlice) -> Result> { - let n = src.len(); - let staging = unsafe { MappedF32Buffer::new(n) } - .map_err(|e| anyhow::anyhow!("bce_sigma download i32: {e}"))?; - let nbytes = n * std::mem::size_of::(); - unsafe { - let (src_ptr, _g) = src.device_ptr(stream); - cudarc::driver::result::memcpy_dtod_async( - staging.dev_ptr, src_ptr, nbytes, stream.cu_stream(), - ).context("bce_sigma i32 download DtoD")?; - } - stream.synchronize().context("bce_sigma i32 download sync")?; - let host_ptr = staging.host_ptr.cast::(); - let mut out = Vec::with_capacity(n); - unsafe { - for i in 0..n { - out.push(std::ptr::read_volatile(host_ptr.add(i))); - } - } - Ok(out) -} diff --git a/crates/ml-alpha/src/trainer/mod.rs b/crates/ml-alpha/src/trainer/mod.rs index 5e55219cf..87d0f0679 100644 --- a/crates/ml-alpha/src/trainer/mod.rs +++ b/crates/ml-alpha/src/trainer/mod.rs @@ -3,7 +3,6 @@ pub mod anchor_controller; pub mod loss; -pub mod loss_sigma; +pub mod multi_horizon_attention; pub mod optim; pub mod perception; -pub mod multi_horizon_attention; diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 4594992f3..4d1546ba0 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -375,6 +375,14 @@ pub struct PerceptionTrainer { attn_bwd_fn: CudaFunction, _attn_module: Arc, + /// MultiHorizonAttention bundle owns the v2 attention path + /// (horizon-token + inverted + MoE), the Kendall σ logarithm for + /// BCE, the anchor controller, and all corresponding optimizers. + /// At present `log_sigma_h_d` and `grad_log_sigma_h_d` are + /// consumed by the BCE callsite; full forward/backward replacement + /// of the legacy `attn_*` path lands in the Stage 2 commit. + pub mha: crate::trainer::multi_horizon_attention::MultiHorizonAttention, + // ── K-loop parallelization (Phase B) ── // Per-batch grad scratch buffers for cfc_step_backward_batched. // Zeroed once per training step; the K-loop's 64 bwd calls accumulate @@ -803,13 +811,20 @@ impl PerceptionTrainer { // Phase B: attn pool per-batch grad scratch. let attn_grad_q_scratch_d = stream.alloc_zeros::(cfg.n_batch * HIDDEN_DIM)?; - // V1 (2026-05-18): per-horizon Q_h trainer state (C24/C25) removed. - // A/B sweep at commit 83546b5c3 falsified the approach (mean_auc - // -0.019 vs baseline; h6000 essentially tied). v2 design with - // horizon-token K-prepend + inverted attention + regime-MoE + - // Kendall σ-BCE + L2 anchor replaces it. See - // docs/superpowers/specs/2026-05-18-ml-alpha-v2-multi-horizon-design.md - // and the V1-V13 plan. + // MultiHorizonAttention bundle. Single source of truth for the + // multi-horizon attention path + Kendall σ on the BCE. Init + // captures snapshots used by the L2 anchor controller. + // Stage 1 (this commit) wires `mha.log_sigma_h_d` and + // `mha.grad_log_sigma_h_d` into the BCE launch. Stage 2 + // replaces the legacy `attention_pool` forward/backward + // callsites with `mha.pool` + `mha.inv_pool` + `mha.moe`. + let mha = crate::trainer::multi_horizon_attention::MultiHorizonAttention::new( + dev, + cfg.n_batch, + cfg.seq_len, + cfg.lr_cfc, + cfg.seed.wrapping_add(0xB200_C00B), + )?; let k = cfg.seq_len; Ok(Self { @@ -866,6 +881,7 @@ impl PerceptionTrainer { attn_fwd_fn, attn_bwd_fn, _attn_module: attn_module, + mha, // Phase B: cfc per-batch grad scratch + reducer. cfc_grad_w_in_scratch_d, cfc_grad_w_rec_scratch_d, @@ -1670,17 +1686,16 @@ impl PerceptionTrainer { launch .arg(&self.probs_per_k_d).arg(&self.labels_per_k_d) .arg(&self.loss_weights_d) + .arg(&self.mha.log_sigma_h_d) .arg(&n_pos_i).arg(&n_h_i) .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); + .arg(&mut self.valid_d) + .arg(&mut self.mha.grad_log_sigma_h_d); unsafe { launch.launch(bce_cfg).context("bce launch")?; } } - // ── 5a. (v2 reserved) — per-horizon Q_h backward removed at V1. - // v2 backward path will live here in commit V10. - // ── 5b. ISV-driven per-horizon EMA + lambda. Updates the EMA // of unweighted per-horizon BCE and emits a clamped // multiplier `lambda_d[h]` per @@ -2586,10 +2601,12 @@ impl PerceptionTrainer { let mut launch = self.stream.launch_builder(&self.bce_fn); launch .arg(&self.probs_per_k_d).arg(&self.labels_per_k_d).arg(&self.loss_weights_d) + .arg(&self.mha.log_sigma_h_d) .arg(&n_pos_i).arg(&n_h_i) .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); + .arg(&mut self.grad_probs_per_k_d).arg(&mut self.valid_d) + .arg(&mut self.mha.grad_log_sigma_h_d); launch.launch(bce_cfg).context("eval bce")?; } self.stream.synchronize().context("eval sync")?; diff --git a/crates/ml-alpha/tests/bce_grad_finite_diff.rs b/crates/ml-alpha/tests/bce_grad_finite_diff.rs index 9ae139f36..212537879 100644 --- a/crates/ml-alpha/tests/bce_grad_finite_diff.rs +++ b/crates/ml-alpha/tests/bce_grad_finite_diff.rs @@ -14,7 +14,7 @@ fn test_device() -> MlDevice { } fn make_input(probs: Vec, labels: Vec) -> BceInput { - BceInput { probs, labels, n_horizons: 5, n_pos: 4 } + BceInput { probs, labels, n_horizons: 5, n_pos: 4, log_sigma_h: None } } #[test] diff --git a/crates/ml-alpha/tests/bce_sigma_numgrad.rs b/crates/ml-alpha/tests/bce_sigma_numgrad.rs deleted file mode 100644 index 225c6ae4b..000000000 --- a/crates/ml-alpha/tests/bce_sigma_numgrad.rs +++ /dev/null @@ -1,166 +0,0 @@ -//! Numgrad parity test for the Kendall σ-weighted multi-horizon BCE -//! kernel (`bce_multi_horizon_sigma_forward_backward`, v2 axis A). -//! -//! Verifies: -//! 1. Analytical d_log_sigma_h matches central-difference numgrad -//! within 5e-2 rel-tol / 5e-3 abs-floor. -//! 2. Analytical grad_probs matches central-difference numgrad on -//! a random subset of positions. -//! 3. Total loss equals the closed-form Σ_h (w_h · mean_bce_h + log σ_h) -//! reconstructed from `mean_bce_per_h`. - -#![cfg(feature = "cuda")] - -use anyhow::Result; -use ml_alpha::trainer::loss_sigma::{bce_sigma_loss_and_grad_gpu, BceSigmaInput}; -use ml_core::device::MlDevice; - -const N_HORIZONS: usize = 5; -const N_POS: usize = 24; - -fn make_inputs(seed: u64) -> BceSigmaInput { - // Deterministic linear-congruential mock RNG; avoids a rand dep. - let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15); - let mut nxt = || -> f32 { - s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); - let x = ((s >> 16) & 0xFFFF) as f32 / 65536.0; - x - }; - - let total = N_POS * N_HORIZONS; - let mut probs = Vec::with_capacity(total); - let mut labels = Vec::with_capacity(total); - for _ in 0..total { - let p = 0.05_f32 + 0.9 * nxt(); // (0.05, 0.95) - let y = if nxt() < 0.5 { 0.0_f32 } else { 1.0_f32 }; - probs.push(p); - labels.push(y); - } - - // Non-uniform base weights, non-uniform log_sigma. - let base_weights: Vec = (0..N_HORIZONS) - .map(|i| 0.5 + i as f32 * 0.25) - .collect(); - let log_sigma_h: Vec = (0..N_HORIZONS) - .map(|i| -0.4 + i as f32 * 0.15) - .collect(); - - BceSigmaInput { - probs, - labels, - base_weights, - log_sigma_h, - n_horizons: N_HORIZONS, - n_pos: N_POS, - } -} - -fn forward_only_loss(dev: &MlDevice, input: &BceSigmaInput) -> Result { - Ok(bce_sigma_loss_and_grad_gpu(dev, input)?.total_loss) -} - -#[test] -#[ignore = "requires CUDA"] -fn d_log_sigma_matches_central_difference() -> Result<()> { - let dev = MlDevice::cuda(0)?; - let input = make_inputs(20260518); - - let out = bce_sigma_loss_and_grad_gpu(&dev, &input)?; - assert!(out.total_loss.is_finite(), "loss not finite: {}", out.total_loss); - - let eps = 1e-3_f32; - for h in 0..N_HORIZONS { - let mut input_plus = input.clone_for_numgrad(); - let mut input_minus = input.clone_for_numgrad(); - input_plus.log_sigma_h[h] += eps; - input_minus.log_sigma_h[h] -= eps; - let l_plus = forward_only_loss(&dev, &input_plus)?; - let l_minus = forward_only_loss(&dev, &input_minus)?; - let numgrad = (l_plus - l_minus) / (2.0 * eps); - let analytic = out.d_log_sigma_h[h]; - let abs_err = (analytic - numgrad).abs(); - let rel_err = abs_err / numgrad.abs().max(1e-3); - assert!( - abs_err < 5e-3 || rel_err < 5e-2, - "d_log_sigma[{h}] mismatch: analytic={analytic:.6} numgrad={numgrad:.6} abs={abs_err:.3e} rel={rel_err:.3e}" - ); - } - Ok(()) -} - -#[test] -#[ignore = "requires CUDA"] -fn grad_probs_matches_central_difference() -> Result<()> { - let dev = MlDevice::cuda(0)?; - let input = make_inputs(42); - - let out = bce_sigma_loss_and_grad_gpu(&dev, &input)?; - assert!(out.total_loss.is_finite(), "loss not finite: {}", out.total_loss); - - // Sample 8 random positions per the LCG. - let eps = 1e-3_f32; - for k in 0..8 { - let i = (k * 7919 + 3) % (N_POS * N_HORIZONS); - - let mut input_plus = input.clone_for_numgrad(); - let mut input_minus = input.clone_for_numgrad(); - input_plus.probs[i] += eps; - input_minus.probs[i] -= eps; - let l_plus = forward_only_loss(&dev, &input_plus)?; - let l_minus = forward_only_loss(&dev, &input_minus)?; - let numgrad = (l_plus - l_minus) / (2.0 * eps); - let analytic = out.grad_probs[i]; - let abs_err = (analytic - numgrad).abs(); - let rel_err = abs_err / numgrad.abs().max(1e-3); - assert!( - abs_err < 5e-3 || rel_err < 5e-2, - "grad_probs[{i}] mismatch: analytic={analytic:.6} numgrad={numgrad:.6} abs={abs_err:.3e} rel={rel_err:.3e}" - ); - } - Ok(()) -} - -#[test] -#[ignore = "requires CUDA"] -fn total_loss_equals_closed_form() -> Result<()> { - let dev = MlDevice::cuda(0)?; - let input = make_inputs(123); - let out = bce_sigma_loss_and_grad_gpu(&dev, &input)?; - - // Closed-form Σ_h (w_h · mean_bce_h + log σ_h) - let mut expected = 0.0_f64; - for h in 0..N_HORIZONS { - let ls = input.log_sigma_h[h] as f64; - let bw = input.base_weights[h] as f64; - let w_h = 0.5 * bw * (-2.0 * ls).exp(); - let mean = out.mean_bce_per_h[h] as f64; - expected += w_h * mean + ls; - } - let abs_err = (out.total_loss as f64 - expected).abs(); - let rel_err = abs_err / expected.abs().max(1e-3); - assert!( - abs_err < 1e-3 || rel_err < 5e-3, - "total_loss mismatch: kernel={} closed_form={expected} abs={abs_err:.3e} rel={rel_err:.3e}", - out.total_loss - ); - Ok(()) -} - -// Small helper trait — clones the input by hand because BceSigmaInput -// doesn't implement Clone (deliberately to keep the test crate's -// dependency surface clean). -trait CloneForNumgrad { - fn clone_for_numgrad(&self) -> Self; -} -impl CloneForNumgrad for BceSigmaInput { - fn clone_for_numgrad(&self) -> Self { - Self { - probs: self.probs.clone(), - labels: self.labels.clone(), - base_weights: self.base_weights.clone(), - log_sigma_h: self.log_sigma_h.clone(), - n_horizons: self.n_horizons, - n_pos: self.n_pos, - } - } -}