From 842e90aeb06d40a3dfa15a90c3c9df24121197ad Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 22 May 2026 11:16:29 +0200 Subject: [PATCH] feat(aux-heads): rewrite for A+B paired (BCE + conditional-Huber) (CB3+CB4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CB1+CB2 swapped labels D→A+B; this swaps the kernels to match. aux_heads.cu — 4-output structure: - Forward: 12 outputs per snapshot = 4 per direction × N_HORIZONS (prof_long_logit, size_long_pred, prof_short_logit, size_short_pred, each [N_HORIZONS]). Linear projections; sigmoid applied in BCE kernel. - Backward: accepts 4 grad_y inputs, produces 8 grad_W + 8 grad_b + grad_h_aux. Cooperative h_aux staging in shmem once per block. - 8 weight matrices total, Xavier × 0.1 init under scoped_init_seed. aux_loss.cu — 2 kernels: - aux_bce_loss_fwd_bwd: class-weighted BCE+sigmoid fused. pos_weight in shared mem; scales positive-class gradient. NaN-mask y_true. - aux_huber_masked_fwd_bwd: Huber w/ NaN-mask. CB1's y_size=NaN at y_prof=0 provides the conditional-Huber semantics naturally — no separate mask buffer needed. aux_heads.rs: - AuxHeads + AuxHeadsWeights: 8 buffer fields (4 W + 4 b) - AuxBceLoss + AuxMaskedHuberLoss wrappers replace AuxHuberLoss - POS_WEIGHT_MIN/MAX = [1.0, 50.0] clamps per E3 - aux_heads_fwd_gpu/aux_heads_bwd_gpu/aux_bce_loss_gpu/aux_huber_masked_loss_gpu perception.rs (minimal compile-keeping signature updates only): - Renamed/added buffers: 4 prediction (prof/size × long/short), 4 label staging, 4 grad_y per-K, 8 head grad scratches, 8 head Adam optimizers, 2 pos_weight buffers (device + staging) - HOLDING PATTERN: bwd zeroes the 4 grad_y_per_K buffers each step so the head Adam updates are no-ops on grad=0 (no aux gradient signal this commit). CB5 wires the actual aux_bce + aux_huber_masked calls. - BCE direction signal + dir_acc readouts updated to use the new prof_long/prof_short prediction buffers (so existing perception_overfit aux test still passes). 7 GPU oracle tests on RTX 3050 sm_86, all pass in 2.43s: - fwd_matches_naive_reference, bwd_finite_diff_matches_bias_sample, aux_bce_loss_matches_naive_reference, aux_bce_pos_weight_scales_positive_gradient (verified: pos_weight=10 → 10× gradient ratio within 1e-4), aux_huber_masked_does_not_propagate, aux_huber_masked_covers_both_branches, aux_bce_nan_mask_does_not_propagate Cubins rebuilt: aux_heads (8736→10528 bytes, +20% for 4-head fwd/bwd), aux_loss (6944→13344 bytes, +92% for 2 kernels). Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml-alpha/build.rs | 2 +- crates/ml-alpha/cuda/aux_heads.cu | 245 +++-- crates/ml-alpha/cuda/aux_loss.cu | 260 +++-- crates/ml-alpha/src/aux_heads.rs | 499 ++++++---- crates/ml-alpha/src/trainer/perception.rs | 1058 +++++++++++++-------- crates/ml-alpha/tests/aux_heads.rs | 684 +++++++++---- 6 files changed, 1840 insertions(+), 908 deletions(-) diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index c83fd2f4a..de4633cf5 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -33,7 +33,7 @@ const KERNELS: &[&str] = &[ "aux_vec_add", // SDD-3 Layer B5: element-wise dst += src for aux→encoder gradient accumulation (lifted stop-grad) ]; -// Cache bust v18 (2026-05-22): SDD-3 Layer B5 — aux_heads + aux_trunk bwd grad scratches use += for K-loop accumulation; aux_vec_add cubin for stop-grad-lifted encoder accumulation. +// Cache bust v19 (2026-05-22): SDD-3 CB3+CB4 — aux_heads.cu rewritten for A+B paired (4 heads × 2 dirs = 8 weight matrices, 4 outputs per direction); aux_loss.cu replaced Huber-only with class-weighted BCE (sigmoid-fused, pos_weight per horizon) + NaN-masked Huber (conditional via loader's y_size = NaN @ y_prof = 0). fn main() { println!("cargo:rerun-if-changed=build.rs"); diff --git a/crates/ml-alpha/cuda/aux_heads.cu b/crates/ml-alpha/cuda/aux_heads.cu index 47320eb1a..94c0bbce6 100644 --- a/crates/ml-alpha/cuda/aux_heads.cu +++ b/crates/ml-alpha/cuda/aux_heads.cu @@ -1,17 +1,19 @@ -// aux_heads.cu — linear regression heads on the AuxTrunk output. +// aux_heads.cu — A+B paired linear heads on the AuxTrunk output (CB3). // -// Per `pearl_separate_aux_trunk_when_shared_starves` (SDD-3 Layer B3): -// the aux trunk runs alongside the main per-branch CfC trunk and feeds -// its 64-dim hidden state to a pair of LINEAR projection heads: +// Per SDD-3 CB1+CB3: each direction (long, short) emits TWO independent +// heads on the shared aux trunk hidden state: // -// y_long_hat [B, N_AUX_HORIZONS] = h_aux [B, AUX_HIDDEN] @ W_long^T + b_long -// y_short_hat[B, N_AUX_HORIZONS] = h_aux [B, AUX_HIDDEN] @ W_short^T + b_short +// prof_long_logit [B, N] = h_aux [B, H] @ W_prof_long^T + b_prof_long +// size_long_pred [B, N] = h_aux [B, H] @ W_size_long^T + b_size_long +// prof_short_logit[B, N] = h_aux [B, H] @ W_prof_short^T + b_prof_short +// size_short_pred [B, N] = h_aux [B, H] @ W_size_short^T + b_size_short // -// where (W_long, b_long, W_short, b_short) are independent per-direction -// parameters. No activation — these are regression outputs supervised by -// the D-style asymmetric-loss-aversion labels generated in -// `multi_horizon_labels::generate_outcome_labels_d` (SDD-3 Layer B1) and -// scored by the Huber kernel in `aux_loss.cu` (this same task). +// where H = AUX_HIDDEN = 64 and N = N_AUX_HORIZONS = 3. The prof heads +// emit RAW LOGITS — sigmoid is fused into the BCE loss kernel in +// `aux_loss.cu`. The size heads emit raw linear regression outputs +// supervised by NaN-masked Huber (conditional on the matching prof label; +// CB1's loader sets `y_size = NaN` whenever `y_prof = 0`, which gives +// the conditional-Huber semantics for free). // // Design constraints (per CLAUDE.md + repo memory): // * `feedback_no_atomicadd.md` — no atomicAdd. Each thread is @@ -27,19 +29,22 @@ // that affect control flow. // // Block layout: one block per batch sample, AUX_HIDDEN=64 threads per -// block. Forward issues a single matmul per direction with thread role -// = output channel k. Backward uses thread role = hidden unit c so that -// grad_h_aux (which sums over k) is the sole-writer responsibility of -// thread c; the parameter-grad writes (grad_w[k, c], grad_b[k]) then -// loop k inside each thread c. +// block. Forward issues a fused matmul that produces all four outputs +// per direction by routing the first N_OUTPUTS_PER_BLOCK threads to +// distinct (head, horizon) slots. Backward uses thread role = hidden +// unit c so that grad_h_aux (which sums over k across all four heads) +// is the sole-writer responsibility of thread c; the parameter-grad +// writes (grad_w[k, c], grad_b[k]) then loop k inside each thread c. #define AUX_HIDDEN 64 #define N_AUX_HORIZONS 3 -#define N_DIRS 2 // long, short — laid out as two contiguous blocks +#define N_DIRS 2 // long, short +#define N_HEADS_PER_DIR 2 // prof, size +#define N_OUTPUTS (N_DIRS * N_HEADS_PER_DIR * N_AUX_HORIZONS) // 12 // ───────────────────────────────────────────────────────────────────── -// aux_heads_fwd: per-direction linear projection, batched. +// aux_heads_fwd: per-direction × per-head linear projection, batched. // // Launch: // grid = (B, 1, 1) @@ -47,20 +52,29 @@ // shared_mem_bytes = AUX_HIDDEN * sizeof(float) (h_aux row stage) // // Outputs are written in one kernel launch: -// y_long_hat [batch, k] = b_long[k] + Σ_c W_long [k, c] * h_aux[batch, c] -// y_short_hat[batch, k] = b_short[k] + Σ_c W_short[k, c] * h_aux[batch, c] +// prof_long_logit [batch, k] = b_prof_long [k] + Σ_c W_prof_long [k, c] * h_aux[batch, c] +// size_long_pred [batch, k] = b_size_long [k] + Σ_c W_size_long [k, c] * h_aux[batch, c] +// prof_short_logit[batch, k] = b_prof_short[k] + Σ_c W_prof_short[k, c] * h_aux[batch, c] +// size_short_pred [batch, k] = b_size_short[k] + Σ_c W_size_short[k, c] * h_aux[batch, c] // -// Layout: W_long, W_short are row-major [N_AUX_HORIZONS, AUX_HIDDEN]. +// Layout: all four W_* are row-major [N_AUX_HORIZONS, AUX_HIDDEN]; all +// four b_* are length N_AUX_HORIZONS. // ───────────────────────────────────────────────────────────────────── extern "C" __global__ void aux_heads_fwd( - const float* __restrict__ w_long, // [N_AUX_HORIZONS × AUX_HIDDEN] - const float* __restrict__ b_long, // [N_AUX_HORIZONS] - const float* __restrict__ w_short, // [N_AUX_HORIZONS × AUX_HIDDEN] - const float* __restrict__ b_short, // [N_AUX_HORIZONS] - const float* __restrict__ h_aux, // [B × AUX_HIDDEN] + const float* __restrict__ w_prof_long, // [N_AUX_HORIZONS × AUX_HIDDEN] + const float* __restrict__ b_prof_long, // [N_AUX_HORIZONS] + const float* __restrict__ w_size_long, // [N_AUX_HORIZONS × AUX_HIDDEN] + const float* __restrict__ b_size_long, // [N_AUX_HORIZONS] + const float* __restrict__ w_prof_short, // [N_AUX_HORIZONS × AUX_HIDDEN] + const float* __restrict__ b_prof_short, // [N_AUX_HORIZONS] + const float* __restrict__ w_size_short, // [N_AUX_HORIZONS × AUX_HIDDEN] + const float* __restrict__ b_size_short, // [N_AUX_HORIZONS] + const float* __restrict__ h_aux, // [B × AUX_HIDDEN] int B, - float* __restrict__ y_long_hat, // [B × N_AUX_HORIZONS] - float* __restrict__ y_short_hat // [B × N_AUX_HORIZONS] + float* __restrict__ prof_long_logit, // [B × N_AUX_HORIZONS] + float* __restrict__ size_long_pred, // [B × N_AUX_HORIZONS] + float* __restrict__ prof_short_logit, // [B × N_AUX_HORIZONS] + float* __restrict__ size_short_pred // [B × N_AUX_HORIZONS] ) { extern __shared__ float s_h[]; // [AUX_HIDDEN] @@ -71,89 +85,123 @@ extern "C" __global__ void aux_heads_fwd( // Cooperative staging of h_aux[batch, *] into shared mem so every // thread's later dot-product reads from smem rather than re-fetching - // the same row N_DIRS × N_AUX_HORIZONS times from DRAM. + // the same row N_OUTPUTS times from DRAM. s_h[c] = h_aux[batch * AUX_HIDDEN + c]; __syncthreads(); - // Compute the 6 outputs (2 dirs × 3 horizons) via 6 threads (k=0..5). - // Remaining 58 threads sit idle this phase — block is intentionally - // sized to AUX_HIDDEN so the *backward* kernel (which is the heavier - // pass) gets full occupancy on h_aux/grad_h_aux. The forward's idle - // lanes are a constant per-block overhead, not a per-batch cost. - if (c < N_DIRS * N_AUX_HORIZONS) { - const int dir = c / N_AUX_HORIZONS; // 0 = long, 1 = short - const int k = c % N_AUX_HORIZONS; + // Compute the 12 outputs (2 dirs × 2 heads × 3 horizons) via 12 + // threads. Remaining 52 threads sit idle this phase — block is + // intentionally sized to AUX_HIDDEN so the *backward* kernel (which + // is the heavier pass) gets full occupancy on h_aux/grad_h_aux. + if (c < N_OUTPUTS) { + // Output layout: thread index c encodes (dir, head, k) as + // dir = c / (N_HEADS_PER_DIR * N_AUX_HORIZONS) ∈ {0,1} + // head = (c / N_AUX_HORIZONS) % N_HEADS_PER_DIR ∈ {0=prof,1=size} + // k = c % N_AUX_HORIZONS ∈ {0..2} + const int dir = c / (N_HEADS_PER_DIR * N_AUX_HORIZONS); + const int head = (c / N_AUX_HORIZONS) % N_HEADS_PER_DIR; + const int k = c % N_AUX_HORIZONS; - const float* w = (dir == 0) ? w_long : w_short; - const float* bv = (dir == 0) ? b_long : b_short; + // Route to the correct (W, b, out) triple. Branch is uniform + // across the warp because all four if-branches dispatch on the + // same (dir, head) — the warp uniformly executes a single chain. + const float* w; + const float* bv; + float* out; + if (dir == 0 && head == 0) { + w = w_prof_long; bv = b_prof_long; out = prof_long_logit; + } else if (dir == 0 && head == 1) { + w = w_size_long; bv = b_size_long; out = size_long_pred; + } else if (dir == 1 && head == 0) { + w = w_prof_short; bv = b_prof_short; out = prof_short_logit; + } else { + w = w_size_short; bv = b_size_short; out = size_short_pred; + } float acc = bv[k]; #pragma unroll for (int ci = 0; ci < AUX_HIDDEN; ++ci) { acc += w[k * AUX_HIDDEN + ci] * s_h[ci]; } - float* out = (dir == 0) ? y_long_hat : y_short_hat; out[batch * N_AUX_HORIZONS + k] = acc; } } // ───────────────────────────────────────────────────────────────────── -// aux_heads_bwd: backward through the two per-direction linear heads. +// aux_heads_bwd: backward through the four per-direction × per-head +// linear heads. // // Launch: // grid = (B, 1, 1) // block = (AUX_HIDDEN = 64, 1, 1) // shared_mem_bytes = AUX_HIDDEN * sizeof(float) (h_aux row stage) // -// Grad shapes (per-batch scratch; OVERWRITE semantics — caller reduces -// axis 0 via existing `reduce_axis0` infrastructure): -// grad_w_long : [B × N_AUX_HORIZONS × AUX_HIDDEN] -// grad_b_long : [B × N_AUX_HORIZONS] -// grad_w_short : [B × N_AUX_HORIZONS × AUX_HIDDEN] -// grad_b_short : [B × N_AUX_HORIZONS] -// grad_h_aux : [B × AUX_HIDDEN] — to be added by the -// trainer to the aux trunk's grad_h_new (cross-batch -// reduction NOT needed — h_aux is per-batch). +// Grad shapes (per-batch scratch; `+=` accumulate semantics — caller +// reduces axis 0 via existing `reduce_axis0` infrastructure): +// grad_w_prof_long : [B × N_AUX_HORIZONS × AUX_HIDDEN] +// grad_b_prof_long : [B × N_AUX_HORIZONS] +// grad_w_size_long : [B × N_AUX_HORIZONS × AUX_HIDDEN] +// grad_b_size_long : [B × N_AUX_HORIZONS] +// grad_w_prof_short : [B × N_AUX_HORIZONS × AUX_HIDDEN] +// grad_b_prof_short : [B × N_AUX_HORIZONS] +// grad_w_size_short : [B × N_AUX_HORIZONS × AUX_HIDDEN] +// grad_b_size_short : [B × N_AUX_HORIZONS] +// grad_h_aux : [B × AUX_HIDDEN] — to be added by the +// trainer to the aux trunk's grad_h_new (no cross- +// batch reduction needed — h_aux is per-batch). // -// Math (per batch, per direction d ∈ {long, short}): -// grad_w_d[k, c] = grad_y_d_hat[k] * h_aux[c] -// grad_b_d[k] = grad_y_d_hat[k] -// grad_h_aux[c] += Σ_k grad_y_long_hat [k] * W_long [k, c] -// + Σ_k grad_y_short_hat[k] * W_short[k, c] +// Inputs (per-element gradients from the loss kernels): +// grad_prof_long_logit : [B × N_AUX_HORIZONS] — dL_BCE/dz, sigmoid fused +// grad_size_long_pred : [B × N_AUX_HORIZONS] — dL_Huber/dy_size +// grad_prof_short_logit : [B × N_AUX_HORIZONS] +// grad_size_short_pred : [B × N_AUX_HORIZONS] +// +// Math (per batch, per head): +// grad_w[k, c] = grad_out[k] * h_aux[c] +// grad_b[k] = grad_out[k] +// grad_h_aux[c] += Σ_k Σ_{head ∈ all four} grad_out[head][k] * W[head][k, c] // // Thread role: c = threadIdx.x is the hidden-unit dim. Each thread c // is the sole writer of: -// * grad_w_{long,short}[batch, k, c] for all k (loops k internally) +// * grad_w_*_*[batch, k, c] for all k (loops k internally) // * grad_h_aux[batch, c] // And the first N_AUX_HORIZONS threads additionally write grad_b for -// each direction. +// each of the four heads. // // No atomicAdd; no cross-thread reduction needed for grad_w / grad_h_aux // because each thread owns disjoint output slots (per // `feedback_no_atomicadd.md`). // // Accumulation semantics (SDD-3 Layer B5 callers): grad_w / grad_b are -// `+=` accumulated; grad_h_aux is `=` overwrite (per-K consumer downstream). -// The trainer's per-K backward loop launches this kernel multiple times -// over the K axis (gradient at each step k re-enters the same head with -// a different h_aux), and the per-batch param-grad sum across K is the -// correct gradient. Caller is responsible for zeroing the grad_w / grad_b -// scratches before the first K-iteration (the trainer does this with the -// rest of its per-step scratch memsets). +// `+=` accumulated across the K-loop; grad_h_aux is `=` overwrite +// (per-K consumer downstream). The trainer's per-K backward loop launches +// this kernel multiple times over the K axis (gradient at each step k +// re-enters the head with a different h_aux), and the per-batch param- +// grad sum across K is the correct gradient. Caller is responsible for +// zeroing the grad_w / grad_b scratches before the first K-iteration +// (the trainer does this with the rest of its per-step scratch memsets). // ───────────────────────────────────────────────────────────────────── extern "C" __global__ void aux_heads_bwd( - const float* __restrict__ w_long, // [N_AUX_HORIZONS × AUX_HIDDEN] - const float* __restrict__ w_short, // [N_AUX_HORIZONS × AUX_HIDDEN] - const float* __restrict__ h_aux, // [B × AUX_HIDDEN] - const float* __restrict__ grad_y_long_hat, // [B × N_AUX_HORIZONS] - const float* __restrict__ grad_y_short_hat, // [B × N_AUX_HORIZONS] + const float* __restrict__ w_prof_long, // [N_AUX_HORIZONS × AUX_HIDDEN] + const float* __restrict__ w_size_long, // [N_AUX_HORIZONS × AUX_HIDDEN] + const float* __restrict__ w_prof_short, // [N_AUX_HORIZONS × AUX_HIDDEN] + const float* __restrict__ w_size_short, // [N_AUX_HORIZONS × AUX_HIDDEN] + const float* __restrict__ h_aux, // [B × AUX_HIDDEN] + const float* __restrict__ grad_prof_long_logit, // [B × N_AUX_HORIZONS] + const float* __restrict__ grad_size_long_pred, // [B × N_AUX_HORIZONS] + const float* __restrict__ grad_prof_short_logit, // [B × N_AUX_HORIZONS] + const float* __restrict__ grad_size_short_pred, // [B × N_AUX_HORIZONS] int B, - float* __restrict__ grad_w_long, // [B × N_AUX_HORIZONS × AUX_HIDDEN] - float* __restrict__ grad_b_long, // [B × N_AUX_HORIZONS] - float* __restrict__ grad_w_short, // [B × N_AUX_HORIZONS × AUX_HIDDEN] - float* __restrict__ grad_b_short, // [B × N_AUX_HORIZONS] - float* __restrict__ grad_h_aux // [B × AUX_HIDDEN] + float* __restrict__ grad_w_prof_long, // [B × N_AUX_HORIZONS × AUX_HIDDEN] + float* __restrict__ grad_b_prof_long, // [B × N_AUX_HORIZONS] + float* __restrict__ grad_w_size_long, // [B × N_AUX_HORIZONS × AUX_HIDDEN] + float* __restrict__ grad_b_size_long, // [B × N_AUX_HORIZONS] + float* __restrict__ grad_w_prof_short, // [B × N_AUX_HORIZONS × AUX_HIDDEN] + float* __restrict__ grad_b_prof_short, // [B × N_AUX_HORIZONS] + float* __restrict__ grad_w_size_short, // [B × N_AUX_HORIZONS × AUX_HIDDEN] + float* __restrict__ grad_b_size_short, // [B × N_AUX_HORIZONS] + float* __restrict__ grad_h_aux // [B × AUX_HIDDEN] ) { extern __shared__ float s_h[]; // [AUX_HIDDEN] @@ -162,25 +210,31 @@ extern "C" __global__ void aux_heads_bwd( if (batch >= B) return; - // Stage h_aux[batch, *] and the per-direction grad_y vectors. The - // grad_y vectors are small (N_AUX_HORIZONS = 3 each) so we keep them - // in shared via the first N_AUX_HORIZONS threads. + // Stage h_aux[batch, *] and the four per-head grad_out vectors. The + // grad_out vectors are small (N_AUX_HORIZONS = 3 each) so we keep + // them in shared via the first N_AUX_HORIZONS threads. s_h[c] = h_aux[batch * AUX_HIDDEN + c]; - __shared__ float s_gyl[N_AUX_HORIZONS]; - __shared__ float s_gys[N_AUX_HORIZONS]; + __shared__ float s_gpl[N_AUX_HORIZONS]; // grad_prof_long_logit + __shared__ float s_gsl[N_AUX_HORIZONS]; // grad_size_long_pred + __shared__ float s_gps[N_AUX_HORIZONS]; // grad_prof_short_logit + __shared__ float s_gss[N_AUX_HORIZONS]; // grad_size_short_pred if (c < N_AUX_HORIZONS) { - s_gyl[c] = grad_y_long_hat [batch * N_AUX_HORIZONS + c]; - s_gys[c] = grad_y_short_hat[batch * N_AUX_HORIZONS + c]; + s_gpl[c] = grad_prof_long_logit [batch * N_AUX_HORIZONS + c]; + s_gsl[c] = grad_size_long_pred [batch * N_AUX_HORIZONS + c]; + s_gps[c] = grad_prof_short_logit[batch * N_AUX_HORIZONS + c]; + s_gss[c] = grad_size_short_pred [batch * N_AUX_HORIZONS + c]; } __syncthreads(); - // grad_b_{long,short}[batch, k] += grad_y_{long,short}_hat[batch, k] + // grad_b_*[batch, k] += grad_out_*[batch, k] // Sole writer: thread c == k (for k in [0, N_AUX_HORIZONS)). // `+=` per K-loop iteration sums each step's bias gradient. if (c < N_AUX_HORIZONS) { - grad_b_long [batch * N_AUX_HORIZONS + c] += s_gyl[c]; - grad_b_short[batch * N_AUX_HORIZONS + c] += s_gys[c]; + grad_b_prof_long [batch * N_AUX_HORIZONS + c] += s_gpl[c]; + grad_b_size_long [batch * N_AUX_HORIZONS + c] += s_gsl[c]; + grad_b_prof_short[batch * N_AUX_HORIZONS + c] += s_gps[c]; + grad_b_size_short[batch * N_AUX_HORIZONS + c] += s_gss[c]; } // For thread c (the hidden-unit dim) — accumulate grad_w[batch, k, c] @@ -192,19 +246,26 @@ extern "C" __global__ void aux_heads_bwd( float gh_acc = 0.0f; #pragma unroll for (int k = 0; k < N_AUX_HORIZONS; ++k) { - const float gyl_k = s_gyl[k]; - const float gys_k = s_gys[k]; + const float gpl_k = s_gpl[k]; + const float gsl_k = s_gsl[k]; + const float gps_k = s_gps[k]; + const float gss_k = s_gss[k]; - // grad_w_d[batch, k, c] += grad_y_d_hat[batch, k] * h_aux[batch, c] const long long off = (long long)batch * N_AUX_HORIZONS * AUX_HIDDEN + (long long)k * AUX_HIDDEN + c; - grad_w_long [off] += gyl_k * h_c; - grad_w_short[off] += gys_k * h_c; - // grad_h_aux[batch, c] += grad_y_d_hat[k] * W_d[k, c] for both dirs. - gh_acc += gyl_k * w_long [k * AUX_HIDDEN + c]; - gh_acc += gys_k * w_short[k * AUX_HIDDEN + c]; + // grad_w_*[batch, k, c] += grad_out_*[batch, k] * h_aux[batch, c] + grad_w_prof_long [off] += gpl_k * h_c; + grad_w_size_long [off] += gsl_k * h_c; + grad_w_prof_short[off] += gps_k * h_c; + grad_w_size_short[off] += gss_k * h_c; + + // grad_h_aux[batch, c] += Σ over all four heads + gh_acc += gpl_k * w_prof_long [k * AUX_HIDDEN + c]; + gh_acc += gsl_k * w_size_long [k * AUX_HIDDEN + c]; + gh_acc += gps_k * w_prof_short[k * AUX_HIDDEN + c]; + gh_acc += gss_k * w_size_short[k * AUX_HIDDEN + c]; } grad_h_aux[batch * AUX_HIDDEN + c] = gh_acc; } diff --git a/crates/ml-alpha/cuda/aux_loss.cu b/crates/ml-alpha/cuda/aux_loss.cu index c51b1c294..f047eed75 100644 --- a/crates/ml-alpha/cuda/aux_loss.cu +++ b/crates/ml-alpha/cuda/aux_loss.cu @@ -1,23 +1,49 @@ -// aux_loss.cu — Huber loss kernel for the aux trade-outcome heads. +// aux_loss.cu — class-weighted BCE + NaN-masked Huber for the A+B +// paired prof-binary + size-regression aux supervision (SDD-3 CB1/CB3/CB4). // -// Per E2's design memo (referenced from `pearl_separate_aux_trunk_when_shared_starves`): -// the aux head predicts continuous-valued trade-outcome targets -// (D-style asymmetric loss-aversion labels from -// `multi_horizon_labels::generate_outcome_labels_d`). Huber is the -// canonical regression loss when the target distribution has heavy -// tails — squared near zero (smooth gradient), absolute far from zero -// (bounded gradient magnitude). δ = 1.0 by default per E2. +// Per CB1's relabel: each (direction, horizon) emits TWO labels: +// y_prof ∈ {0, 1} or NaN — did the trade hit profit before stop? +// y_size ∈ ℝ or NaN — σ-normalised realised return (only when prof=1) // -// Formulation (per element): -// r = y_hat - y_true -// L_huber = (|r| <= δ) ? 0.5 * r² : δ * (|r| - 0.5 * δ) -// dL/dy_hat = (|r| <= δ) ? r : δ * sign(r) +// Two loss kernels here: // -// NaN handling: y_true may be NaN at positions where the D-label -// generator dropped the target (edge/invalid samples). We MUST mask -// these out — the loss contribution is 0 AND the grad is 0 so the +// aux_bce_loss_fwd_bwd — sigmoid+BCE on (prof_logit, y_prof), with +// per-horizon `pos_weight` for class balance +// (`pos_weight = clamp(n_neg/n_pos, 1.0, 50.0)` +// computed once per epoch on the loader side). +// NaN in y_prof zeroes the gradient and +// skips the loss contribution. +// +// aux_huber_masked_fwd_bwd — Huber on (y_size_hat, y_size_true). The +// loader emits y_size = NaN whenever +// y_prof = 0 (or the sample is invalid), +// which gives the conditional-Huber +// semantics for free: the NaN mask both +// zeroes the gradient and skips the loss +// contribution. +// +// Formulation: +// +// BCE (sigmoid fused, class-weighted): +// p = 1 / (1 + exp(-z)) +// L_BCE = -[pos_weight * y * log(p) + (1-y) * log(1-p)] +// dL_BCE/dz = (y > 0.5) ? pos_weight * (p - 1) +// : p +// (clamp log(·) at log(1e-7) to guard against p ∈ {0, 1} blowup; +// `--use_fast_math` already produces approximate sigmoid so the +// clamp is a defence-in-depth.) +// +// Huber (δ = 1.0 default): +// r = y_hat - y_true +// L_Huber = (|r| <= δ) ? 0.5 * r² : δ * (|r| - 0.5 * δ) +// dL/dy_hat = (|r| <= δ) ? r : δ * sign(r) +// +// NaN-mask discipline: y_true may be NaN at masked positions; we MUST +// zero the gradient AND skip the loss accumulator at those slots so // downstream `aux_heads_bwd` propagates zero through the head and the -// `aux_trunk_bwd` sees no spurious gradient. +// `aux_trunk_bwd` sees no spurious gradient. The optimiser cannot eat +// a NaN scalar — a single masked label without the mask would +// poison the entire training step. // // Reduction discipline (per `feedback_no_atomicadd.md`): // * Per-thread strided accumulation into registers. @@ -26,11 +52,12 @@ // * Final warp-0 reduction to total loss / total valid count. // * No atomicAdd anywhere. // -// Per-direction split: this kernel processes ONE direction's tensor at -// a time. The trainer calls it twice (once for long, once for short) -// and sums the two scalars into the total aux loss. Keeping the kernel -// per-direction keeps the block layout simple and lets the per-direction -// telemetry (mean Huber, valid count) emerge for free. +// Per-direction split: each kernel processes ONE direction × ONE head's +// tensor at a time. The trainer calls each twice (once for long, once +// for short) and sums the four scalars into the total aux loss. Keeping +// the kernels per-(direction, head) keeps the block layout simple and +// lets the per-direction telemetry (mean loss, valid count) emerge +// for free. // // Block layout: grid = (1), block = (BLOCK = 256). Single block handles // the full [B × N_AUX_HORIZONS] grid (max ~32 × 3 = 96, easily within @@ -38,6 +65,7 @@ #define AUX_LOSS_BLOCK 256 #define AUX_LOSS_WARPS (AUX_LOSS_BLOCK / 32) // == 8 +#define N_AUX_HORIZONS 3 __device__ __forceinline__ float warp_reduce_sum_f32(float v) { @@ -58,68 +86,178 @@ __device__ __forceinline__ int warp_reduce_sum_i32(int v) { // ───────────────────────────────────────────────────────────────────── -// aux_huber_loss_fwd_bwd: -// Fused Huber-loss forward + per-element gradient writeback for ONE -// direction (long or short). Inputs/outputs are [B × N_AUX_HORIZONS] -// layout. NaN-masked positions contribute zero loss and write zero -// gradient (no contamination of downstream consumers). +// aux_bce_loss_fwd_bwd: +// Fused sigmoid + class-weighted BCE forward + per-element logit +// gradient writeback for ONE direction (long or short) × the prof +// head. Inputs/outputs are [B × N_AUX_HORIZONS] layout. NaN-masked +// positions contribute zero loss and write zero gradient (no +// contamination of downstream consumers). +// +// Inputs: +// prof_logit [B × N_AUX_HORIZONS] — RAW logit z (no sigmoid) +// y_prof_true [B × N_AUX_HORIZONS] — labels in {0, 1} or NaN +// pos_weight [N_AUX_HORIZONS] — per-horizon positive-class +// weight (uploaded host-side +// from loader epoch stats). +// n_total = B × N_AUX_HORIZONS // // Outputs: -// loss_out [1] — Σ_valid L_huber (UNREDUCED; -// caller divides by valid_count -// if a mean is desired). -// grad_y_hat [B × N_AUX_HORIZONS] — dL/dy_hat per element. -// valid_count_out [1] — #{(b, k) : isfinite(y_true)}. +// loss_out [1] — Σ_valid L_BCE (UNREDUCED; +// caller divides by valid_count +// if a mean is desired). +// grad_prof_logit [B × N_AUX_HORIZONS] — dL_BCE/dz per element. +// valid_count_out [1] — #{(b, k) : isfinite(y_prof_true)}. // -// `loss_out` is intentionally a SUM, not a mean, so the caller controls -// the reduction policy (mean over valid, mean over total elements, etc.) -// without making the kernel re-runnable for different conventions. +// `loss_out` is intentionally a SUM, not a mean — caller controls the +// reduction policy without re-running the kernel. // ───────────────────────────────────────────────────────────────────── -extern "C" __global__ void aux_huber_loss_fwd_bwd( - const float* __restrict__ y_hat, // [B × N_AUX_HORIZONS] - const float* __restrict__ y_true, // [B × N_AUX_HORIZONS] (NaN-maskable) - float delta, // Huber transition point (default 1.0) - int n_total, // = B × N_AUX_HORIZONS - float* __restrict__ loss_out, // [1] — Σ L (unreduced) - float* __restrict__ grad_y_hat, // [B × N_AUX_HORIZONS] - int* __restrict__ valid_count_out // [1] +extern "C" __global__ void aux_bce_loss_fwd_bwd( + const float* __restrict__ prof_logit, // [B × N_AUX_HORIZONS] + const float* __restrict__ y_prof_true, // [B × N_AUX_HORIZONS] (NaN-maskable) + const float* __restrict__ pos_weight, // [N_AUX_HORIZONS] + int n_total, // = B × N_AUX_HORIZONS + float* __restrict__ loss_out, // [1] — Σ L_BCE (unreduced) + float* __restrict__ grad_prof_logit, // [B × N_AUX_HORIZONS] + int* __restrict__ valid_count_out // [1] ) { const int tid = threadIdx.x; const int lane = tid & 31; const int warp = tid >> 5; - // Per-thread strided accumulation into registers. + // Stage `pos_weight[N_AUX_HORIZONS]` into shared so every strided + // iteration reads from smem rather than re-fetching from DRAM. + __shared__ float s_pw[N_AUX_HORIZONS]; + if (tid < N_AUX_HORIZONS) { + s_pw[tid] = pos_weight[tid]; + } + __syncthreads(); + float local_loss = 0.0f; int local_valid = 0; for (int i = tid; i < n_total; i += AUX_LOSS_BLOCK) { - const float y_t = y_true[i]; + const float y_t = y_prof_true[i]; if (!isfinite(y_t)) { - grad_y_hat[i] = 0.0f; + grad_prof_logit[i] = 0.0f; continue; } - const float y_p = y_hat[i]; - const float r = y_p - y_t; - const float ar = fabsf(r); + const int h = i % N_AUX_HORIZONS; + const float pw = s_pw[h]; + const float z = prof_logit[i]; + // sigmoid(z) = 1/(1+exp(-z)). With --use_fast_math nvcc emits + // the fast intrinsic; we clamp the log argument to avoid -inf + // when p saturates to 0 or 1. + const float p = 1.0f / (1.0f + expf(-z)); + const float p_c = fmaxf(p, 1e-7f); + const float one_m_p = fmaxf(1.0f - p, 1e-7f); - float L, g; - if (ar <= delta) { - L = 0.5f * r * r; - g = r; + // Class-weighted BCE: positive class gets weight `pw`, negative + // class gets weight 1. Equivalent to scaling the positive-class + // term in the standard formula. + float L; + float grad_z; + if (y_t > 0.5f) { + L = -pw * logf(p_c); + grad_z = pw * (p - 1.0f); } else { - L = delta * (ar - 0.5f * delta); - // δ * sign(r). sign uses copysign to keep g in the same sign - // bucket as r (and to give 0 when r == 0, though the |r| > δ - // branch already excludes r == 0). - g = copysignf(delta, r); + L = -logf(one_m_p); + grad_z = p; } - local_loss += L; - local_valid += 1; - grad_y_hat[i] = g; + local_loss += L; + local_valid += 1; + grad_prof_logit[i] = grad_z; } // Warp-shuffle reduce. - float warp_loss = warp_reduce_sum_f32(local_loss); + float warp_loss = warp_reduce_sum_f32(local_loss); + int warp_valid = warp_reduce_sum_i32(local_valid); + + __shared__ float s_loss[AUX_LOSS_WARPS]; + __shared__ int s_valid[AUX_LOSS_WARPS]; + if (lane == 0) { + s_loss[warp] = warp_loss; + s_valid[warp] = warp_valid; + } + __syncthreads(); + + if (warp == 0) { + float v_l = (lane < AUX_LOSS_WARPS) ? s_loss[lane] : 0.0f; + int v_v = (lane < AUX_LOSS_WARPS) ? s_valid[lane] : 0; + v_l = warp_reduce_sum_f32(v_l); + v_v = warp_reduce_sum_i32(v_v); + if (lane == 0) { + loss_out[0] = v_l; + valid_count_out[0] = v_v; + } + } +} + + +// ───────────────────────────────────────────────────────────────────── +// aux_huber_masked_fwd_bwd: +// Fused Huber-loss forward + per-element gradient writeback for ONE +// direction (long or short) × the size head. NaN-masked positions +// contribute zero loss and write zero gradient. +// +// Conditional-Huber semantics: the loader (CB1) sets `y_size = NaN` +// whenever `y_prof = 0` (no profit hit so the realised-return target is +// undefined). The NaN-mask path here therefore implements the +// conditional-Huber loss `L_Huber if y_prof=1 else 0` without needing a +// separate prof-mask buffer — the loader has already encoded the +// condition into the label tensor. +// +// Inputs: +// y_size_hat [B × N_AUX_HORIZONS] — raw linear regression output +// y_size_true [B × N_AUX_HORIZONS] — σ-normalised return or NaN +// delta — Huber transition point (default 1.0) +// n_total = B × N_AUX_HORIZONS +// +// Outputs: +// loss_out [1] — Σ_valid L_Huber (unreduced) +// grad_y_size_hat [B × N_AUX_HORIZONS] — dL/dy_size_hat per element +// valid_count_out [1] — #{(b, k) : isfinite(y_size_true)}. +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void aux_huber_masked_fwd_bwd( + const float* __restrict__ y_size_hat, // [B × N_AUX_HORIZONS] + const float* __restrict__ y_size_true, // [B × N_AUX_HORIZONS] (NaN-maskable) + float delta, // Huber transition point + int n_total, // = B × N_AUX_HORIZONS + float* __restrict__ loss_out, // [1] — Σ L_Huber (unreduced) + float* __restrict__ grad_y_size_hat, // [B × N_AUX_HORIZONS] + int* __restrict__ valid_count_out // [1] +) { + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + + float local_loss = 0.0f; + int local_valid = 0; + + for (int i = tid; i < n_total; i += AUX_LOSS_BLOCK) { + const float y_t = y_size_true[i]; + if (!isfinite(y_t)) { + grad_y_size_hat[i] = 0.0f; + continue; + } + const float y_p = y_size_hat[i]; + const float r = y_p - y_t; + const float ar = fabsf(r); + + float L, g; + if (ar <= delta) { + L = 0.5f * r * r; + g = r; + } else { + L = delta * (ar - 0.5f * delta); + g = copysignf(delta, r); + } + local_loss += L; + local_valid += 1; + grad_y_size_hat[i] = g; + } + + // Warp-shuffle reduce. + float warp_loss = warp_reduce_sum_f32(local_loss); int warp_valid = warp_reduce_sum_i32(local_valid); __shared__ float s_loss[AUX_LOSS_WARPS]; diff --git a/crates/ml-alpha/src/aux_heads.rs b/crates/ml-alpha/src/aux_heads.rs index b1a935590..7e4e53b82 100644 --- a/crates/ml-alpha/src/aux_heads.rs +++ b/crates/ml-alpha/src/aux_heads.rs @@ -1,36 +1,39 @@ -//! SDD-3 Layer B4 — Aux heads (linear regression) + Huber loss. +//! SDD-3 CB3+CB4 — Aux heads (A+B paired: prof BCE + size Huber) on the +//! aux trunk hidden state. //! -//! Sits on top of [`crate::cfc::AuxTrunk`] and predicts continuous-valued -//! trade-outcome targets emitted by [`crate::multi_horizon_labels::generate_outcome_labels_ab`] -//! (Candidate-B paired prof-binary + size-regression labels from SDD-3 CB1). -//! -//! CB2 transition state: the head signature here is unchanged (still emits -//! `y_{long,short}_hat[B, N]` regression outputs). CB5 will rewrite the loss -//! kernel to consume the prof-binary head via BCE and the size head via Huber. -//! Until then, the regression outputs are supervised against `y_prof_{dir}` -//! through the existing Huber path — a holding-pattern signal, not the final -//! A+B losses. +//! Sits on top of [`crate::cfc::AuxTrunk`] and predicts BOTH a binary +//! "did the trade hit profit?" signal AND a σ-normalised realised-return +//! magnitude — per direction (long, short) × per horizon. The label +//! schema is defined by +//! [`crate::multi_horizon_labels::generate_outcome_labels_ab`] (CB1). //! //! ## Architecture //! -//! Two INDEPENDENT per-direction linear heads: +//! Four INDEPENDENT linear heads on the shared aux trunk hidden state: //! //! ```text -//! y_long_hat [B, N] = h_aux [B, H] @ W_long^T + b_long -//! y_short_hat[B, N] = h_aux [B, H] @ W_short^T + b_short +//! prof_long_logit [B, N] = h_aux [B, H] @ W_prof_long^T + b_prof_long +//! size_long_pred [B, N] = h_aux [B, H] @ W_size_long^T + b_size_long +//! prof_short_logit[B, N] = h_aux [B, H] @ W_prof_short^T + b_prof_short +//! size_short_pred [B, N] = h_aux [B, H] @ W_size_short^T + b_size_short //! ``` //! -//! where `H = AUX_HIDDEN = 64` and `N = N_AUX_HORIZONS = 3`. No activation -//! — these are regression outputs supervised by Huber loss with NaN masking. +//! where `H = AUX_HIDDEN = 64` and `N = N_AUX_HORIZONS = 3`. The prof +//! heads emit raw logits — sigmoid is fused into the BCE loss kernel +//! ([`AuxBceLoss`]). The size heads emit raw linear regression outputs +//! supervised by NaN-masked Huber ([`AuxMaskedHuberLoss`]). CB1 sets +//! `y_size = NaN` whenever `y_prof = 0`, which provides the +//! conditional-Huber semantics without an explicit mask buffer. //! -//! ## Why linear (not GRN) +//! ## Why two heads per direction //! -//! Per E2's design memo: the asymmetric loss-aversion labels are -//! already non-linear functions of price + drawdown over the K-horizon; -//! the head's role is to predict a calibrated scalar, not to re-derive -//! the asymmetry. A linear projection keeps the head's gradient signal -//! interpretable and avoids the extra-parameter risk of the main BCE -//! head's GRN structure. +//! Per the E3→CB1 redesign: a single magnitude head over signed returns +//! conflates "direction got the sign right" with "direction sized up +//! correctly" — making the loss-aversion asymmetry impossible to learn +//! and producing the WR/PF anticorrelation observed in the D-label +//! sweeps. Splitting prof (classification, asymmetric class weights) +//! from size (regression, conditional on prof=1) lets each loss own a +//! single, well-conditioned objective. //! //! ## Initialisation //! @@ -38,8 +41,11 @@ //! `AuxHeads::new` installs a [`scoped_init_seed`] guard so the Xavier //! draws are deterministic given the seed. //! -//! * `w_long`, `w_short`: Xavier uniform `[-sqrt(6/(H+N)), +]` -//! * `b_long`, `b_short`: zeros +//! * All four `W_*` matrices: Xavier uniform × 0.1 +//! (`scale = 0.1 * sqrt(6 / (H + N))`). The 0.1 down-scale keeps the +//! initial logits + size predictions near zero so the loss gradient +//! is dominated by the targets, not by noise. +//! * All four `b_*` vectors: zeros. //! //! ## Constraints honoured //! @@ -49,6 +55,9 @@ //! * `feedback_no_htod_htoh_only_mapped_pinned.md` — uploads go //! through mapped-pinned staging. //! * `feedback_no_nvrtc.md` — pre-compiled cubin via `build.rs`. +//! * `feedback_no_partial_refactor.md` — the old `AuxHuberLoss` is +//! fully replaced by [`AuxBceLoss`] + [`AuxMaskedHuberLoss`] in +//! the same commit; no leftover aliases. use std::sync::Arc; @@ -69,25 +78,37 @@ const AUX_HEADS_CUBIN: &[u8] = const AUX_LOSS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/aux_loss.cubin")); -/// Number of aux-horizon regression outputs per direction. Matches the +/// Number of aux-horizon outputs per (direction, head). Matches the /// main BCE head's [`crate::heads::N_HORIZONS`] (= 3) by construction: -/// the D-label generator emits one outcome per (direction, horizon) -/// across the same horizon grid as the BCE supervision so the two -/// heads can share the loader pipeline. +/// the A+B label generator emits one (prof, size) pair per (direction, +/// horizon) across the same horizon grid as the BCE supervision so the +/// two heads can share the loader pipeline. pub const N_AUX_HORIZONS: usize = crate::heads::N_HORIZONS; -/// Default Huber transition point per E2's recommendation. The -/// `aux_huber_loss_gpu` wrapper exposes it as a runtime argument so -/// downstream callers can sweep without rebuilding the cubin. +/// Default Huber transition point for the size head. CB5 exposes it as +/// a runtime argument so downstream callers can sweep without rebuilding +/// the cubin. pub const DEFAULT_HUBER_DELTA: f32 = 1.0; +/// Class-weight clamp bounds for the BCE positive-class weight. The +/// loader computes `pos_weight = clamp(n_neg / n_pos, 1.0, 50.0)` per +/// horizon; the clamp prevents single-positive epochs from producing +/// runaway logit gradients while still letting >5:1 class imbalance +/// scale the positive-class loss proportionally. +pub const POS_WEIGHT_MIN: f32 = 1.0; +pub const POS_WEIGHT_MAX: f32 = 50.0; + /// Host-side weight storage for [`AuxHeads`] (checkpoint round-trip). #[derive(Clone, Debug)] pub struct AuxHeadsWeights { - pub w_long: Vec, // [N_AUX_HORIZONS × AUX_HIDDEN] - pub b_long: Vec, // [N_AUX_HORIZONS] - pub w_short: Vec, // [N_AUX_HORIZONS × AUX_HIDDEN] - pub b_short: Vec, // [N_AUX_HORIZONS] + pub w_prof_long: Vec, // [N_AUX_HORIZONS × AUX_HIDDEN] + pub b_prof_long: Vec, // [N_AUX_HORIZONS] + pub w_size_long: Vec, // [N_AUX_HORIZONS × AUX_HIDDEN] + pub b_size_long: Vec, // [N_AUX_HORIZONS] + pub w_prof_short: Vec, // [N_AUX_HORIZONS × AUX_HIDDEN] + pub b_prof_short: Vec, // [N_AUX_HORIZONS] + pub w_size_short: Vec, // [N_AUX_HORIZONS × AUX_HIDDEN] + pub b_size_short: Vec, // [N_AUX_HORIZONS] } /// Construction config for [`AuxHeads`]. @@ -98,8 +119,9 @@ pub struct AuxHeadsConfig { pub seed: u64, } -/// Aux-direction regression heads. Owns device weights + cached cubin -/// function handles for forward / backward. +/// Aux per-direction × per-head linear regression / classification +/// heads. Owns device weights + cached cubin function handles for +/// forward / backward. pub struct AuxHeads { cfg: AuxHeadsConfig, stream: Arc, @@ -108,15 +130,29 @@ pub struct AuxHeads { pub fwd_fn: CudaFunction, pub bwd_fn: CudaFunction, - pub w_long_d: CudaSlice, - pub b_long_d: CudaSlice, - pub w_short_d: CudaSlice, - pub b_short_d: CudaSlice, + pub w_prof_long_d: CudaSlice, + pub b_prof_long_d: CudaSlice, + pub w_size_long_d: CudaSlice, + pub b_size_long_d: CudaSlice, + pub w_prof_short_d: CudaSlice, + pub b_prof_short_d: CudaSlice, + pub w_size_short_d: CudaSlice, + pub b_size_short_d: CudaSlice, } -/// Wrapper around the Huber loss cubin. Cached once per trainer so the -/// per-step launch path doesn't reload the module. -pub struct AuxHuberLoss { +/// Wrapper around the class-weighted BCE loss cubin (sigmoid fused). +/// Cached once per trainer so the per-step launch path doesn't reload +/// the module. +pub struct AuxBceLoss { + stream: Arc, + _module: Arc, + pub fn_: CudaFunction, +} + +/// Wrapper around the NaN-masked Huber loss cubin (used for size head; +/// conditional-Huber semantics come from the loader marking `y_size = NaN` +/// at `y_prof = 0`). +pub struct AuxMaskedHuberLoss { stream: Arc, _module: Arc, pub fn_: CudaFunction, @@ -136,22 +172,36 @@ impl AuxHeads { .load_function("aux_heads_bwd") .context("load aux_heads_bwd")?; + // Per `pearl_scoped_init_seed_for_reproducibility`: install the + // scoped seed BEFORE drawing any Xavier samples so the GPU init + // helpers downstream see the same RNG state across runs. let _seed_guard = scoped_init_seed(cfg.seed); let mut r = ChaCha8Rng::seed_from_u64(cfg.seed); - let scale = (6.0_f32 / (AUX_HIDDEN + N_AUX_HORIZONS) as f32).sqrt(); - let w_long: Vec = (0..N_AUX_HORIZONS * AUX_HIDDEN) - .map(|_| r.gen_range(-scale..scale)) - .collect(); - let w_short: Vec = (0..N_AUX_HORIZONS * AUX_HIDDEN) - .map(|_| r.gen_range(-scale..scale)) - .collect(); - let b_long: Vec = vec![0.0; N_AUX_HORIZONS]; - let b_short: Vec = vec![0.0; N_AUX_HORIZONS]; + // Xavier × 0.1: 10× smaller than standard Xavier so the initial + // prof logits sit near 0 (sigmoid ≈ 0.5, uniform-prior-like) and + // the initial size predictions sit near 0 (Huber gradient + // dominated by the σ-normalised target). + let scale = 0.1_f32 * (6.0_f32 / (AUX_HIDDEN + N_AUX_HORIZONS) as f32).sqrt(); + let xavier = + |r: &mut ChaCha8Rng| -> Vec { + (0..N_AUX_HORIZONS * AUX_HIDDEN) + .map(|_| r.gen_range(-scale..scale)) + .collect() + }; + let w_prof_long = xavier(&mut r); + let w_size_long = xavier(&mut r); + let w_prof_short = xavier(&mut r); + let w_size_short = xavier(&mut r); + let zeros_n: Vec = vec![0.0; N_AUX_HORIZONS]; - let w_long_d = upload(&stream, &w_long)?; - let b_long_d = upload(&stream, &b_long)?; - let w_short_d = upload(&stream, &w_short)?; - let b_short_d = upload(&stream, &b_short)?; + let w_prof_long_d = upload(&stream, &w_prof_long)?; + let b_prof_long_d = upload(&stream, &zeros_n)?; + let w_size_long_d = upload(&stream, &w_size_long)?; + let b_size_long_d = upload(&stream, &zeros_n)?; + let w_prof_short_d = upload(&stream, &w_prof_short)?; + let b_prof_short_d = upload(&stream, &zeros_n)?; + let w_size_short_d = upload(&stream, &w_size_short)?; + let b_size_short_d = upload(&stream, &zeros_n)?; Ok(Self { cfg, @@ -159,10 +209,14 @@ impl AuxHeads { _module: module, fwd_fn, bwd_fn, - w_long_d, - b_long_d, - w_short_d, - b_short_d, + w_prof_long_d, + b_prof_long_d, + w_size_long_d, + b_size_long_d, + w_prof_short_d, + b_prof_short_d, + w_size_short_d, + b_size_short_d, }) } @@ -177,41 +231,61 @@ impl AuxHeads { /// Download the head weights into host vectors. Used by checkpoint /// save and unit tests. pub fn download_weights(&self) -> Result { - let mut w_long = vec![0.0_f32; self.w_long_d.len()]; - let mut b_long = vec![0.0_f32; self.b_long_d.len()]; - let mut w_short = vec![0.0_f32; self.w_short_d.len()]; - let mut b_short = vec![0.0_f32; self.b_short_d.len()]; - self.stream - .memcpy_dtoh(&self.w_long_d, w_long.as_mut_slice()) - .context("aux_heads dtoh w_long")?; - self.stream - .memcpy_dtoh(&self.b_long_d, b_long.as_mut_slice()) - .context("aux_heads dtoh b_long")?; - self.stream - .memcpy_dtoh(&self.w_short_d, w_short.as_mut_slice()) - .context("aux_heads dtoh w_short")?; - self.stream - .memcpy_dtoh(&self.b_short_d, b_short.as_mut_slice()) - .context("aux_heads dtoh b_short")?; + let dtoh = |d: &CudaSlice, name: &str| -> Result> { + let mut h = vec![0.0_f32; d.len()]; + self.stream + .memcpy_dtoh(d, h.as_mut_slice()) + .with_context(|| format!("aux_heads dtoh {name}"))?; + Ok(h) + }; Ok(AuxHeadsWeights { - w_long, - b_long, - w_short, - b_short, + w_prof_long: dtoh(&self.w_prof_long_d, "w_prof_long")?, + b_prof_long: dtoh(&self.b_prof_long_d, "b_prof_long")?, + w_size_long: dtoh(&self.w_size_long_d, "w_size_long")?, + b_size_long: dtoh(&self.b_size_long_d, "b_size_long")?, + w_prof_short: dtoh(&self.w_prof_short_d, "w_prof_short")?, + b_prof_short: dtoh(&self.b_prof_short_d, "b_prof_short")?, + w_size_short: dtoh(&self.w_size_short_d, "w_size_short")?, + b_size_short: dtoh(&self.b_size_short_d, "b_size_short")?, }) } } -impl AuxHuberLoss { +impl AuxBceLoss { pub fn new(dev: &MlDevice) -> Result { - let stream: Arc = dev.cuda_stream().context("aux_loss stream")?.clone(); - let ctx = dev.cuda_context().context("aux_loss ctx")?; + let stream: Arc = dev.cuda_stream().context("aux_bce_loss stream")?.clone(); + let ctx = dev.cuda_context().context("aux_bce_loss ctx")?; let module = ctx .load_cubin(AUX_LOSS_CUBIN.to_vec()) .context("load aux_loss cubin")?; let fn_ = module - .load_function("aux_huber_loss_fwd_bwd") - .context("load aux_huber_loss_fwd_bwd")?; + .load_function("aux_bce_loss_fwd_bwd") + .context("load aux_bce_loss_fwd_bwd")?; + Ok(Self { + stream, + _module: module, + fn_, + }) + } + + pub fn stream(&self) -> &Arc { + &self.stream + } +} + +impl AuxMaskedHuberLoss { + pub fn new(dev: &MlDevice) -> Result { + let stream: Arc = dev + .cuda_stream() + .context("aux_huber_masked_loss stream")? + .clone(); + let ctx = dev.cuda_context().context("aux_huber_masked_loss ctx")?; + let module = ctx + .load_cubin(AUX_LOSS_CUBIN.to_vec()) + .context("load aux_loss cubin (masked huber)")?; + let fn_ = module + .load_function("aux_huber_masked_fwd_bwd") + .context("load aux_huber_masked_fwd_bwd")?; Ok(Self { stream, _module: module, @@ -226,12 +300,14 @@ impl AuxHuberLoss { // ── Launch wrappers ────────────────────────────────────────────────── -/// Launch the fused aux-heads forward kernel. +/// Launch the fused aux-heads forward kernel for ALL four heads. /// /// Buffer contract: -/// * `h_aux_d` : `[B × AUX_HIDDEN]` — aux trunk hidden state -/// * `y_long_hat_d` : `[B × N_AUX_HORIZONS]` — long predictions (overwrite) -/// * `y_short_hat_d` : `[B × N_AUX_HORIZONS]` — short predictions (overwrite) +/// * `h_aux_d` : `[B × AUX_HIDDEN]` — aux trunk hidden state +/// * `prof_long_logit_d` : `[B × N_AUX_HORIZONS]` — prof_long raw logit (overwrite) +/// * `size_long_pred_d` : `[B × N_AUX_HORIZONS]` — size_long raw linear (overwrite) +/// * `prof_short_logit_d` : `[B × N_AUX_HORIZONS]` — prof_short raw logit (overwrite) +/// * `size_short_pred_d` : `[B × N_AUX_HORIZONS]` — size_short raw linear (overwrite) /// /// Shared memory layout: `AUX_HIDDEN * sizeof(float)` covering the /// cooperative-staged `h_aux` row. @@ -239,23 +315,35 @@ impl AuxHuberLoss { pub fn aux_heads_fwd_gpu( stream: &Arc, func: &CudaFunction, - w_long_d: &CudaSlice, - b_long_d: &CudaSlice, - w_short_d: &CudaSlice, - b_short_d: &CudaSlice, + w_prof_long_d: &CudaSlice, + b_prof_long_d: &CudaSlice, + w_size_long_d: &CudaSlice, + b_size_long_d: &CudaSlice, + w_prof_short_d: &CudaSlice, + b_prof_short_d: &CudaSlice, + w_size_short_d: &CudaSlice, + b_size_short_d: &CudaSlice, h_aux_d: &CudaSlice, b_sz: i32, - y_long_hat_d: &mut CudaSlice, - y_short_hat_d: &mut CudaSlice, + prof_long_logit_d: &mut CudaSlice, + size_long_pred_d: &mut CudaSlice, + prof_short_logit_d: &mut CudaSlice, + size_short_pred_d: &mut CudaSlice, ) -> Result<()> { let b_sz_u = b_sz as usize; - debug_assert_eq!(w_long_d.len(), N_AUX_HORIZONS * AUX_HIDDEN); - debug_assert_eq!(b_long_d.len(), N_AUX_HORIZONS); - debug_assert_eq!(w_short_d.len(), N_AUX_HORIZONS * AUX_HIDDEN); - debug_assert_eq!(b_short_d.len(), N_AUX_HORIZONS); - debug_assert_eq!(h_aux_d.len(), b_sz_u * AUX_HIDDEN); - debug_assert_eq!(y_long_hat_d.len(), b_sz_u * N_AUX_HORIZONS); - debug_assert_eq!(y_short_hat_d.len(), b_sz_u * N_AUX_HORIZONS); + debug_assert_eq!(w_prof_long_d.len(), N_AUX_HORIZONS * AUX_HIDDEN); + debug_assert_eq!(b_prof_long_d.len(), N_AUX_HORIZONS); + debug_assert_eq!(w_size_long_d.len(), N_AUX_HORIZONS * AUX_HIDDEN); + debug_assert_eq!(b_size_long_d.len(), N_AUX_HORIZONS); + debug_assert_eq!(w_prof_short_d.len(), N_AUX_HORIZONS * AUX_HIDDEN); + debug_assert_eq!(b_prof_short_d.len(), N_AUX_HORIZONS); + debug_assert_eq!(w_size_short_d.len(), N_AUX_HORIZONS * AUX_HIDDEN); + debug_assert_eq!(b_size_short_d.len(), N_AUX_HORIZONS); + debug_assert_eq!(h_aux_d.len(), b_sz_u * AUX_HIDDEN); + debug_assert_eq!(prof_long_logit_d.len(), b_sz_u * N_AUX_HORIZONS); + debug_assert_eq!(size_long_pred_d.len(), b_sz_u * N_AUX_HORIZONS); + debug_assert_eq!(prof_short_logit_d.len(), b_sz_u * N_AUX_HORIZONS); + debug_assert_eq!(size_short_pred_d.len(), b_sz_u * N_AUX_HORIZONS); let cfg = LaunchConfig { grid_dim: (b_sz as u32, 1, 1), @@ -264,14 +352,20 @@ pub fn aux_heads_fwd_gpu( }; let mut launch = stream.launch_builder(func); launch - .arg(w_long_d) - .arg(b_long_d) - .arg(w_short_d) - .arg(b_short_d) + .arg(w_prof_long_d) + .arg(b_prof_long_d) + .arg(w_size_long_d) + .arg(b_size_long_d) + .arg(w_prof_short_d) + .arg(b_prof_short_d) + .arg(w_size_short_d) + .arg(b_size_short_d) .arg(h_aux_d) .arg(&b_sz) - .arg(y_long_hat_d) - .arg(y_short_hat_d); + .arg(prof_long_logit_d) + .arg(size_long_pred_d) + .arg(prof_short_logit_d) + .arg(size_short_pred_d); unsafe { launch.launch(cfg).context("aux_heads_fwd launch")?; } @@ -280,11 +374,11 @@ pub fn aux_heads_fwd_gpu( /// Launch the fused aux-heads backward kernel. /// -/// Per-batch grad scratch is OVERWRITTEN by this kernel; the caller is -/// responsible for the cross-batch reduction (e.g. `reduce_axis0`). -/// `grad_h_aux_d` is per-batch and does NOT require cross-batch reduction -/// — it feeds back into the aux-trunk's `grad_h_new` as a per-sample -/// signal. +/// Per-batch grad scratch is `+=`-accumulated across the K-loop by this +/// kernel; the caller is responsible for the cross-batch reduction +/// (e.g. `reduce_axis0`). `grad_h_aux_d` is per-batch and does NOT +/// require cross-batch reduction — it feeds back into the aux-trunk's +/// `grad_h_new` as a per-sample signal. /// /// Shared memory layout: `AUX_HIDDEN * sizeof(float)` covering the /// cooperative-staged `h_aux` row. @@ -292,29 +386,45 @@ pub fn aux_heads_fwd_gpu( pub fn aux_heads_bwd_gpu( stream: &Arc, func: &CudaFunction, - w_long_d: &CudaSlice, - w_short_d: &CudaSlice, + w_prof_long_d: &CudaSlice, + w_size_long_d: &CudaSlice, + w_prof_short_d: &CudaSlice, + w_size_short_d: &CudaSlice, h_aux_d: &CudaSlice, - grad_y_long_hat_d: &CudaSlice, - grad_y_short_hat_d: &CudaSlice, + grad_prof_long_logit_d: &CudaSlice, + grad_size_long_pred_d: &CudaSlice, + grad_prof_short_logit_d: &CudaSlice, + grad_size_short_pred_d: &CudaSlice, b_sz: i32, - grad_w_long_d: &mut CudaSlice, - grad_b_long_d: &mut CudaSlice, - grad_w_short_d: &mut CudaSlice, - grad_b_short_d: &mut CudaSlice, + grad_w_prof_long_d: &mut CudaSlice, + grad_b_prof_long_d: &mut CudaSlice, + grad_w_size_long_d: &mut CudaSlice, + grad_b_size_long_d: &mut CudaSlice, + grad_w_prof_short_d: &mut CudaSlice, + grad_b_prof_short_d: &mut CudaSlice, + grad_w_size_short_d: &mut CudaSlice, + grad_b_size_short_d: &mut CudaSlice, grad_h_aux_d: &mut CudaSlice, ) -> Result<()> { let b_sz_u = b_sz as usize; - debug_assert_eq!(w_long_d.len(), N_AUX_HORIZONS * AUX_HIDDEN); - debug_assert_eq!(w_short_d.len(), N_AUX_HORIZONS * AUX_HIDDEN); - debug_assert_eq!(h_aux_d.len(), b_sz_u * AUX_HIDDEN); - debug_assert_eq!(grad_y_long_hat_d.len(), b_sz_u * N_AUX_HORIZONS); - debug_assert_eq!(grad_y_short_hat_d.len(), b_sz_u * N_AUX_HORIZONS); - debug_assert_eq!(grad_w_long_d.len(), b_sz_u * N_AUX_HORIZONS * AUX_HIDDEN); - debug_assert_eq!(grad_b_long_d.len(), b_sz_u * N_AUX_HORIZONS); - debug_assert_eq!(grad_w_short_d.len(), b_sz_u * N_AUX_HORIZONS * AUX_HIDDEN); - debug_assert_eq!(grad_b_short_d.len(), b_sz_u * N_AUX_HORIZONS); - debug_assert_eq!(grad_h_aux_d.len(), b_sz_u * AUX_HIDDEN); + debug_assert_eq!(w_prof_long_d.len(), N_AUX_HORIZONS * AUX_HIDDEN); + debug_assert_eq!(w_size_long_d.len(), N_AUX_HORIZONS * AUX_HIDDEN); + debug_assert_eq!(w_prof_short_d.len(), N_AUX_HORIZONS * AUX_HIDDEN); + debug_assert_eq!(w_size_short_d.len(), N_AUX_HORIZONS * AUX_HIDDEN); + debug_assert_eq!(h_aux_d.len(), b_sz_u * AUX_HIDDEN); + debug_assert_eq!(grad_prof_long_logit_d.len(), b_sz_u * N_AUX_HORIZONS); + debug_assert_eq!(grad_size_long_pred_d.len(), b_sz_u * N_AUX_HORIZONS); + debug_assert_eq!(grad_prof_short_logit_d.len(), b_sz_u * N_AUX_HORIZONS); + debug_assert_eq!(grad_size_short_pred_d.len(), b_sz_u * N_AUX_HORIZONS); + debug_assert_eq!(grad_w_prof_long_d.len(), b_sz_u * N_AUX_HORIZONS * AUX_HIDDEN); + debug_assert_eq!(grad_b_prof_long_d.len(), b_sz_u * N_AUX_HORIZONS); + debug_assert_eq!(grad_w_size_long_d.len(), b_sz_u * N_AUX_HORIZONS * AUX_HIDDEN); + debug_assert_eq!(grad_b_size_long_d.len(), b_sz_u * N_AUX_HORIZONS); + debug_assert_eq!(grad_w_prof_short_d.len(), b_sz_u * N_AUX_HORIZONS * AUX_HIDDEN); + debug_assert_eq!(grad_b_prof_short_d.len(), b_sz_u * N_AUX_HORIZONS); + debug_assert_eq!(grad_w_size_short_d.len(), b_sz_u * N_AUX_HORIZONS * AUX_HIDDEN); + debug_assert_eq!(grad_b_size_short_d.len(), b_sz_u * N_AUX_HORIZONS); + debug_assert_eq!(grad_h_aux_d.len(), b_sz_u * AUX_HIDDEN); let cfg = LaunchConfig { grid_dim: (b_sz as u32, 1, 1), @@ -323,16 +433,24 @@ pub fn aux_heads_bwd_gpu( }; let mut launch = stream.launch_builder(func); launch - .arg(w_long_d) - .arg(w_short_d) + .arg(w_prof_long_d) + .arg(w_size_long_d) + .arg(w_prof_short_d) + .arg(w_size_short_d) .arg(h_aux_d) - .arg(grad_y_long_hat_d) - .arg(grad_y_short_hat_d) + .arg(grad_prof_long_logit_d) + .arg(grad_size_long_pred_d) + .arg(grad_prof_short_logit_d) + .arg(grad_size_short_pred_d) .arg(&b_sz) - .arg(grad_w_long_d) - .arg(grad_b_long_d) - .arg(grad_w_short_d) - .arg(grad_b_short_d) + .arg(grad_w_prof_long_d) + .arg(grad_b_prof_long_d) + .arg(grad_w_size_long_d) + .arg(grad_b_size_long_d) + .arg(grad_w_prof_short_d) + .arg(grad_b_prof_short_d) + .arg(grad_w_size_short_d) + .arg(grad_b_size_short_d) .arg(grad_h_aux_d); unsafe { launch.launch(cfg).context("aux_heads_bwd launch")?; @@ -340,35 +458,37 @@ pub fn aux_heads_bwd_gpu( Ok(()) } -/// Launch the fused Huber loss forward + backward kernel for ONE -/// direction. The kernel writes: -/// * `loss_out_d[0]` — Σ_valid Huber loss (unreduced sum; -/// caller divides by valid_count if a mean -/// is desired). -/// * `grad_y_hat_d` — per-element dL/dy_hat (NaN-masked zero). -/// * `valid_count_out_d[0]` — #{(b, k) : isfinite(y_true)}. +/// Launch the class-weighted BCE loss kernel for ONE direction × the +/// prof head. The kernel writes: +/// * `loss_out_d[0]` — Σ_valid L_BCE (unreduced sum; caller +/// divides by valid_count if a mean is +/// desired). +/// * `grad_prof_logit_d` — per-element dL/dz (sigmoid+BCE fused; +/// NaN-masked zero). +/// * `valid_count_out_d[0]` — #{(b, k) : isfinite(y_prof_true)}. /// /// Call twice — once for long, once for short — and sum the two loss -/// scalars to get the total aux loss. The per-direction split keeps the -/// telemetry (mean Huber per direction, valid count per direction) -/// available for free without extra kernels. +/// scalars to get the prof-head total. `pos_weight_d` is a length +/// `N_AUX_HORIZONS` device buffer with the per-horizon positive-class +/// weight (clamped to `[POS_WEIGHT_MIN, POS_WEIGHT_MAX]` by the loader). #[allow(clippy::too_many_arguments)] -pub fn aux_huber_loss_gpu( +pub fn aux_bce_loss_gpu( stream: &Arc, func: &CudaFunction, - y_hat_d: &CudaSlice, - y_true_d: &CudaSlice, - delta: f32, + prof_logit_d: &CudaSlice, + y_prof_true_d: &CudaSlice, + pos_weight_d: &CudaSlice, n_total: i32, loss_out_d: &mut CudaSlice, - grad_y_hat_d: &mut CudaSlice, + grad_prof_logit_d: &mut CudaSlice, valid_count_out_d: &mut CudaSlice, ) -> Result<()> { let n_total_u = n_total as usize; - debug_assert_eq!(y_hat_d.len(), n_total_u); - debug_assert_eq!(y_true_d.len(), n_total_u); - debug_assert_eq!(grad_y_hat_d.len(), n_total_u); - debug_assert_eq!(loss_out_d.len(), 1); + debug_assert_eq!(prof_logit_d.len(), n_total_u); + debug_assert_eq!(y_prof_true_d.len(), n_total_u); + debug_assert_eq!(pos_weight_d.len(), N_AUX_HORIZONS); + debug_assert_eq!(grad_prof_logit_d.len(), n_total_u); + debug_assert_eq!(loss_out_d.len(), 1); debug_assert_eq!(valid_count_out_d.len(), 1); const AUX_LOSS_BLOCK: u32 = 256; @@ -379,15 +499,66 @@ pub fn aux_huber_loss_gpu( }; let mut launch = stream.launch_builder(func); launch - .arg(y_hat_d) - .arg(y_true_d) + .arg(prof_logit_d) + .arg(y_prof_true_d) + .arg(pos_weight_d) + .arg(&n_total) + .arg(loss_out_d) + .arg(grad_prof_logit_d) + .arg(valid_count_out_d); + unsafe { + launch.launch(cfg).context("aux_bce_loss launch")?; + } + Ok(()) +} + +/// Launch the NaN-masked Huber loss kernel for ONE direction × the +/// size head. Conditional-Huber semantics: the loader emits +/// `y_size_true = NaN` whenever `y_prof_true = 0` (no profit hit, so +/// the σ-normalised return is undefined). The NaN mask zeroes the +/// gradient AND skips the loss contribution at those slots, giving +/// `L_Huber if y_prof=1 else 0` without an explicit prof-mask buffer. +/// +/// Outputs match the BCE wrapper: +/// * `loss_out_d[0]` — Σ_valid L_Huber +/// * `grad_y_size_hat_d` — per-element dL/dy_size_hat (NaN-masked zero) +/// * `valid_count_out_d[0]` — #{(b, k) : isfinite(y_size_true)} +#[allow(clippy::too_many_arguments)] +pub fn aux_huber_masked_loss_gpu( + stream: &Arc, + func: &CudaFunction, + y_size_hat_d: &CudaSlice, + y_size_true_d: &CudaSlice, + delta: f32, + n_total: i32, + loss_out_d: &mut CudaSlice, + grad_y_size_hat_d: &mut CudaSlice, + valid_count_out_d: &mut CudaSlice, +) -> Result<()> { + let n_total_u = n_total as usize; + debug_assert_eq!(y_size_hat_d.len(), n_total_u); + debug_assert_eq!(y_size_true_d.len(), n_total_u); + debug_assert_eq!(grad_y_size_hat_d.len(), n_total_u); + debug_assert_eq!(loss_out_d.len(), 1); + debug_assert_eq!(valid_count_out_d.len(), 1); + + const AUX_LOSS_BLOCK: u32 = 256; + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (AUX_LOSS_BLOCK, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = stream.launch_builder(func); + launch + .arg(y_size_hat_d) + .arg(y_size_true_d) .arg(&delta) .arg(&n_total) .arg(loss_out_d) - .arg(grad_y_hat_d) + .arg(grad_y_size_hat_d) .arg(valid_count_out_d); unsafe { - launch.launch(cfg).context("aux_huber_loss launch")?; + launch.launch(cfg).context("aux_huber_masked_loss launch")?; } Ok(()) } diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 32c5c99cf..4792db7dd 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -46,8 +46,7 @@ use rand::{Rng, SeedableRng}; use rand_chacha::ChaCha8Rng; use crate::aux_heads::{ - aux_huber_loss_gpu, AuxHeads, AuxHeadsConfig, AuxHuberLoss, DEFAULT_HUBER_DELTA, - N_AUX_HORIZONS, + AuxBceLoss, AuxHeads, AuxHeadsConfig, AuxMaskedHuberLoss, N_AUX_HORIZONS, }; use crate::cfc::aux_trunk::AUX_HIDDEN; use crate::cfc::snap_features::{Mbp10RawInput, ES_TICK_SIZE, FEATURE_DIM, REGIME_DIM}; @@ -883,54 +882,101 @@ pub struct PerceptionTrainer { /// pattern (per `feedback_no_htod_htoh_only_mapped_pinned`). h_mag_per_bucket_staged: MappedF32Buffer, - // ── SDD-3 Layer B5: aux supervision (parallel trunk + asymmetric stop-grad) ── + // ── SDD-3 CB3+CB4: aux supervision (A+B paired prof BCE + size Huber) ── // // A SECOND, smaller CfC trunk (`aux_trunk`) consumes the SAME encoder - // output as the main BCE trunk and feeds two linear regression heads - // (`aux_heads`, one long / one short) supervised by Huber loss against - // the D-style outcome labels from the loader. Per - // `pearl_separate_aux_trunk_when_shared_starves` + `pearl_adam_normalizes_loss_weights`, - // aux gets its own optimizer group; per the E3 memo, gradient flow - // from aux back into the shared encoder is MASKED at first - // (asymmetric stop-grad — BCE drives the encoder) and conditionally - // lifted once aux supervision converges. + // output as the main BCE trunk and feeds FOUR linear heads on top: + // (prof_long, size_long, prof_short, size_short). Each prof head is + // supervised by class-weighted BCE (sigmoid-fused) against the + // CB1-generated `y_prof_{dir}` labels; each size head is supervised + // by NaN-masked Huber against `y_size_{dir}` (CB1 sets `y_size = NaN` + // when `y_prof = 0`, which gives conditional-Huber semantics). + // + // Per `pearl_separate_aux_trunk_when_shared_starves` + `pearl_adam_normalizes_loss_weights`, + // aux gets its own optimizer group; per the E3 memo + CB1 redesign, + // gradient flow from aux back into the shared encoder is MASKED at + // first (asymmetric stop-grad — BCE drives the encoder) and + // conditionally lifted once aux supervision converges. + // + // CB3+CB4 transition state: the kernels are rewritten (aux_heads_fwd + // emits 4 outputs; aux_loss exposes BCE + masked-Huber). The trainer + // here is in a HOLDING PATTERN — the new fwd kernel runs each step + // and writes 4 outputs, but the BCE + masked-Huber CALLS are CB5's + // scope. Until CB5 lands, the four grad_y buffers are zeroed each + // step so the bwd path produces zero gradients (Adam no-op). /// Aux trunk — smaller CfC (AUX_HIDDEN=64, single bucket). pub aux_trunk: AuxTrunk, - /// Aux heads — long + short linear regression on aux trunk output. + /// Aux heads — 4 linear projections (prof_long, size_long, prof_short, + /// size_short), each `[N_AUX_HORIZONS × AUX_HIDDEN]`. pub aux_heads: AuxHeads, - /// Aux Huber loss kernel handle (cached). - aux_huber: AuxHuberLoss, + /// Aux class-weighted BCE loss kernel handle (cached). Wired by CB5. + aux_bce: AuxBceLoss, + /// Aux NaN-masked Huber loss kernel handle (cached). Wired by CB5. + aux_huber_masked: AuxMaskedHuberLoss, /// Aux trunk per-K hidden state buffer (overwritten each step). /// Slot k holds `[B × AUX_HIDDEN]` aux-trunk output at position k. aux_h_per_k_d: CudaSlice, - /// Aux long predictions per K: `[K × B × N_AUX_HORIZONS]`. - aux_y_long_hat_per_k_d: CudaSlice, - /// Aux short predictions per K: `[K × B × N_AUX_HORIZONS]`. - aux_y_short_hat_per_k_d: CudaSlice, - /// Aux long outcome labels (uploaded each step): `[K × B × N_AUX_HORIZONS]`. - aux_y_long_per_k_d: CudaSlice, - /// Aux short outcome labels (uploaded each step): `[K × B × N_AUX_HORIZONS]`. - aux_y_short_per_k_d: CudaSlice, - /// Mapped-pinned staging for aux long outcomes (host → DtoD → device). - stg_aux_y_long: MappedF32Buffer, - /// Mapped-pinned staging for aux short outcomes. - stg_aux_y_short: MappedF32Buffer, - /// Aux Huber per-K grad output for long preds (consumed by aux_heads bwd). - aux_grad_y_long_hat_per_k_d: CudaSlice, - /// Aux Huber per-K grad output for short preds. - aux_grad_y_short_hat_per_k_d: CudaSlice, - /// Aux Huber scalar loss outputs (one per direction). Unreduced sums. - aux_loss_long_d: CudaSlice, - aux_loss_short_d: CudaSlice, - /// Aux valid counts (one per direction). - aux_valid_long_d: CudaSlice, - aux_valid_short_d: CudaSlice, - /// Mapped-pinned host shadows for aux per-direction loss + valid count. - /// Read post-sync each step to drive the per-horizon Huber EMA. - aux_loss_long_host_d: MappedF32Buffer, - aux_loss_short_host_d: MappedF32Buffer, - aux_valid_long_host_d: MappedI32Buffer, - aux_valid_short_host_d: MappedI32Buffer, + /// Aux prof_long raw logits per K: `[K × B × N_AUX_HORIZONS]`. + aux_prof_long_logit_per_k_d: CudaSlice, + /// Aux size_long linear predictions per K: `[K × B × N_AUX_HORIZONS]`. + aux_size_long_pred_per_k_d: CudaSlice, + /// Aux prof_short raw logits per K: `[K × B × N_AUX_HORIZONS]`. + aux_prof_short_logit_per_k_d: CudaSlice, + /// Aux size_short linear predictions per K: `[K × B × N_AUX_HORIZONS]`. + aux_size_short_pred_per_k_d: CudaSlice, + /// Aux prof_long labels (binary {0, 1} or NaN): `[K × B × N_AUX_HORIZONS]`. + aux_y_prof_long_per_k_d: CudaSlice, + /// Aux size_long σ-normalised return labels (NaN at y_prof=0). + aux_y_size_long_per_k_d: CudaSlice, + /// Aux prof_short labels. + aux_y_prof_short_per_k_d: CudaSlice, + /// Aux size_short labels. + aux_y_size_short_per_k_d: CudaSlice, + /// Mapped-pinned staging for aux prof_long labels. + stg_aux_y_prof_long: MappedF32Buffer, + /// Mapped-pinned staging for aux size_long labels. + stg_aux_y_size_long: MappedF32Buffer, + /// Mapped-pinned staging for aux prof_short labels. + stg_aux_y_prof_short: MappedF32Buffer, + /// Mapped-pinned staging for aux size_short labels. + stg_aux_y_size_short: MappedF32Buffer, + /// Per-horizon BCE positive-class weight (device-side, length + /// `N_AUX_HORIZONS`). Uploaded each step from the loader's + /// `pos_fraction` via `last_pos_fraction` → clamped to + /// `[POS_WEIGHT_MIN, POS_WEIGHT_MAX]`. CB5 wires the upload + the + /// BCE call site; CB3+CB4 just provisions the buffer. + aux_pos_weight_long_d: CudaSlice, + aux_pos_weight_short_d: CudaSlice, + /// Mapped-pinned staging for `pos_weight` uploads (length `N_AUX_HORIZONS` each). + stg_aux_pos_weight_long: MappedF32Buffer, + stg_aux_pos_weight_short: MappedF32Buffer, + /// Per-K grad outputs from the four loss kernels (consumed by aux_heads bwd). + /// All four are `[K × B × N_AUX_HORIZONS]`. In CB3+CB4 these are zeroed + /// each step (holding pattern); CB5 wires the BCE + masked-Huber calls. + aux_grad_prof_long_logit_per_k_d: CudaSlice, + aux_grad_size_long_pred_per_k_d: CudaSlice, + aux_grad_prof_short_logit_per_k_d: CudaSlice, + aux_grad_size_short_pred_per_k_d: CudaSlice, + /// Aux scalar loss outputs — one per (direction × head). Unreduced sums. + aux_loss_prof_long_d: CudaSlice, + aux_loss_size_long_d: CudaSlice, + aux_loss_prof_short_d: CudaSlice, + aux_loss_size_short_d: CudaSlice, + /// Aux valid counts — one per (direction × head). + aux_valid_prof_long_d: CudaSlice, + aux_valid_size_long_d: CudaSlice, + aux_valid_prof_short_d: CudaSlice, + aux_valid_size_short_d: CudaSlice, + /// Mapped-pinned host shadows for aux per-(direction, head) loss + valid count. + /// Read post-sync each step to drive the per-horizon EMAs. + aux_loss_prof_long_host_d: MappedF32Buffer, + aux_loss_size_long_host_d: MappedF32Buffer, + aux_loss_prof_short_host_d: MappedF32Buffer, + aux_loss_size_short_host_d: MappedF32Buffer, + aux_valid_prof_long_host_d: MappedI32Buffer, + aux_valid_size_long_host_d: MappedI32Buffer, + aux_valid_prof_short_host_d: MappedI32Buffer, + aux_valid_size_short_host_d: MappedI32Buffer, /// Per-K aux trunk param-grad scratch (reduce_axis0 collapses → final). /// Shape: `[B × AUX_HIDDEN × MAX_AUX_FEAT_DIM]` etc. — sized so the /// per-batch backward kernel writes one row per sample. @@ -959,29 +1005,43 @@ pub struct PerceptionTrainer { /// Zero-init buffer used as h_old at the first K step of aux trunk. aux_zero_h_d: CudaSlice, /// Aux head per-batch param-grad scratch (reduce across B for final grad). - aux_heads_grad_w_long_scratch_d: CudaSlice, - aux_heads_grad_b_long_scratch_d: CudaSlice, - aux_heads_grad_w_short_scratch_d: CudaSlice, - aux_heads_grad_b_short_scratch_d: CudaSlice, + /// Eight scratches: prof_{long,short} × {w, b} + size_{long,short} × {w, b}. + aux_heads_grad_w_prof_long_scratch_d: CudaSlice, + aux_heads_grad_b_prof_long_scratch_d: CudaSlice, + aux_heads_grad_w_size_long_scratch_d: CudaSlice, + aux_heads_grad_b_size_long_scratch_d: CudaSlice, + aux_heads_grad_w_prof_short_scratch_d: CudaSlice, + aux_heads_grad_b_prof_short_scratch_d: CudaSlice, + aux_heads_grad_w_size_short_scratch_d: CudaSlice, + aux_heads_grad_b_size_short_scratch_d: CudaSlice, /// Aux head grad accumulators reduced across B (consumed by Adam). - aux_heads_grad_w_long_d: CudaSlice, - aux_heads_grad_b_long_d: CudaSlice, - aux_heads_grad_w_short_d: CudaSlice, - aux_heads_grad_b_short_d: CudaSlice, + aux_heads_grad_w_prof_long_d: CudaSlice, + aux_heads_grad_b_prof_long_d: CudaSlice, + aux_heads_grad_w_size_long_d: CudaSlice, + aux_heads_grad_b_size_long_d: CudaSlice, + aux_heads_grad_w_prof_short_d: CudaSlice, + aux_heads_grad_b_prof_short_d: CudaSlice, + aux_heads_grad_w_size_short_d: CudaSlice, + aux_heads_grad_b_size_short_d: CudaSlice, /// `[B × AUX_HIDDEN]` grad_h_aux from aux_heads bwd (consumed as /// grad_h_new by aux_trunk bwd). aux_grad_h_aux_d: CudaSlice, /// Aux Adam optimizers — SEPARATE group from main trunk + heads /// per `pearl_adam_normalizes_loss_weights` (independent m/v moments - /// so per-group LR is the actual lever). + /// so per-group LR is the actual lever). Eight head optimizers: + /// prof_{long,short} × {w, b} + size_{long,short} × {w, b}. pub aux_opt_w_in: AdamW, pub aux_opt_w_rec: AdamW, pub aux_opt_b: AdamW, pub aux_opt_tau: AdamW, - pub aux_opt_w_long: AdamW, - pub aux_opt_b_long: AdamW, - pub aux_opt_w_short: AdamW, - pub aux_opt_b_short: AdamW, + pub aux_opt_w_prof_long: AdamW, + pub aux_opt_b_prof_long: AdamW, + pub aux_opt_w_size_long: AdamW, + pub aux_opt_b_size_long: AdamW, + pub aux_opt_w_prof_short: AdamW, + pub aux_opt_b_prof_short: AdamW, + pub aux_opt_w_size_short: AdamW, + pub aux_opt_b_size_short: AdamW, /// Asymmetric stop-grad flag for the aux→encoder gradient path. `true` /// at construction: aux drives only its own params (passive auditor). /// Per E3, lifted to `false` after `aux_huber_ema_per_h.all(< threshold)` @@ -1773,8 +1833,10 @@ impl PerceptionTrainer { }, ) .context("PerceptionTrainer: AuxHeads::new")?; - let aux_huber = AuxHuberLoss::new(dev) - .context("PerceptionTrainer: AuxHuberLoss::new")?; + let aux_bce = AuxBceLoss::new(dev) + .context("PerceptionTrainer: AuxBceLoss::new")?; + let aux_huber_masked = AuxMaskedHuberLoss::new(dev) + .context("PerceptionTrainer: AuxMaskedHuberLoss::new")?; // Aux backward grad scratch sizes (per-batch; reduce_axis0 collapses). // Aux trunk weight tensor shapes (mirror AuxTrunk::new_random): @@ -1815,28 +1877,43 @@ impl PerceptionTrainer { let aux_zero_h_d = stream.alloc_zeros::( cfg.n_batch * AUX_HIDDEN, )?; - // Aux head grad scratches: per-batch w_long/short [B × N_AUX_HORIZONS × AUX_HIDDEN], - // bias [B × N_AUX_HORIZONS]; reduce_axis0 collapses across B. - let aux_heads_grad_w_long_scratch_d = stream.alloc_zeros::( - cfg.n_batch * N_AUX_HORIZONS * AUX_HIDDEN, - )?; - let aux_heads_grad_b_long_scratch_d = stream.alloc_zeros::( - cfg.n_batch * N_AUX_HORIZONS, - )?; - let aux_heads_grad_w_short_scratch_d = stream.alloc_zeros::( - cfg.n_batch * N_AUX_HORIZONS * AUX_HIDDEN, - )?; - let aux_heads_grad_b_short_scratch_d = stream.alloc_zeros::( - cfg.n_batch * N_AUX_HORIZONS, - )?; - let aux_heads_grad_w_long_d = stream.alloc_zeros::( - N_AUX_HORIZONS * AUX_HIDDEN, - )?; - let aux_heads_grad_b_long_d = stream.alloc_zeros::(N_AUX_HORIZONS)?; - let aux_heads_grad_w_short_d = stream.alloc_zeros::( - N_AUX_HORIZONS * AUX_HIDDEN, - )?; - let aux_heads_grad_b_short_d = stream.alloc_zeros::(N_AUX_HORIZONS)?; + // Aux head grad scratches: per-batch w_{prof,size}_{long,short} + // [B × N_AUX_HORIZONS × AUX_HIDDEN] per matrix, + // bias [B × N_AUX_HORIZONS] per vector; reduce_axis0 collapses across B. + let alloc_w_scratch = || -> Result> { + stream + .alloc_zeros::(cfg.n_batch * N_AUX_HORIZONS * AUX_HIDDEN) + .map_err(Into::into) + }; + let alloc_b_scratch = || -> Result> { + stream + .alloc_zeros::(cfg.n_batch * N_AUX_HORIZONS) + .map_err(Into::into) + }; + let aux_heads_grad_w_prof_long_scratch_d = alloc_w_scratch()?; + let aux_heads_grad_b_prof_long_scratch_d = alloc_b_scratch()?; + let aux_heads_grad_w_size_long_scratch_d = alloc_w_scratch()?; + let aux_heads_grad_b_size_long_scratch_d = alloc_b_scratch()?; + let aux_heads_grad_w_prof_short_scratch_d = alloc_w_scratch()?; + let aux_heads_grad_b_prof_short_scratch_d = alloc_b_scratch()?; + let aux_heads_grad_w_size_short_scratch_d = alloc_w_scratch()?; + let aux_heads_grad_b_size_short_scratch_d = alloc_b_scratch()?; + let alloc_w_final = || -> Result> { + stream + .alloc_zeros::(N_AUX_HORIZONS * AUX_HIDDEN) + .map_err(Into::into) + }; + let alloc_b_final = || -> Result> { + stream.alloc_zeros::(N_AUX_HORIZONS).map_err(Into::into) + }; + let aux_heads_grad_w_prof_long_d = alloc_w_final()?; + let aux_heads_grad_b_prof_long_d = alloc_b_final()?; + let aux_heads_grad_w_size_long_d = alloc_w_final()?; + let aux_heads_grad_b_size_long_d = alloc_b_final()?; + let aux_heads_grad_w_prof_short_d = alloc_w_final()?; + let aux_heads_grad_b_prof_short_d = alloc_b_final()?; + let aux_heads_grad_w_size_short_d = alloc_w_final()?; + let aux_heads_grad_b_size_short_d = alloc_b_final()?; let aux_grad_h_aux_d = stream.alloc_zeros::( cfg.n_batch * AUX_HIDDEN, )?; @@ -1850,12 +1927,21 @@ impl PerceptionTrainer { // policy for log-uniformly initialised time constants. let mut aux_opt_tau = AdamW::new(dev, AUX_HIDDEN, cfg.lr_aux * 0.1)?; aux_opt_tau.wd = 0.0; - let aux_opt_w_long = AdamW::new(dev, N_AUX_HORIZONS * AUX_HIDDEN, cfg.lr_aux)?; - let mut aux_opt_b_long = AdamW::new(dev, N_AUX_HORIZONS, cfg.lr_aux)?; - aux_opt_b_long.wd = 0.0; - let aux_opt_w_short = AdamW::new(dev, N_AUX_HORIZONS * AUX_HIDDEN, cfg.lr_aux)?; - let mut aux_opt_b_short = AdamW::new(dev, N_AUX_HORIZONS, cfg.lr_aux)?; - aux_opt_b_short.wd = 0.0; + // Eight head optimizers — prof_{long,short} × {w, b} + + // size_{long,short} × {w, b}. Biases use `wd = 0` mirroring the + // trunk-bias policy. + let aux_opt_w_prof_long = AdamW::new(dev, N_AUX_HORIZONS * AUX_HIDDEN, cfg.lr_aux)?; + let mut aux_opt_b_prof_long = AdamW::new(dev, N_AUX_HORIZONS, cfg.lr_aux)?; + aux_opt_b_prof_long.wd = 0.0; + let aux_opt_w_size_long = AdamW::new(dev, N_AUX_HORIZONS * AUX_HIDDEN, cfg.lr_aux)?; + let mut aux_opt_b_size_long = AdamW::new(dev, N_AUX_HORIZONS, cfg.lr_aux)?; + aux_opt_b_size_long.wd = 0.0; + let aux_opt_w_prof_short = AdamW::new(dev, N_AUX_HORIZONS * AUX_HIDDEN, cfg.lr_aux)?; + let mut aux_opt_b_prof_short = AdamW::new(dev, N_AUX_HORIZONS, cfg.lr_aux)?; + aux_opt_b_prof_short.wd = 0.0; + let aux_opt_w_size_short = AdamW::new(dev, N_AUX_HORIZONS * AUX_HIDDEN, cfg.lr_aux)?; + let mut aux_opt_b_size_short = AdamW::new(dev, N_AUX_HORIZONS, cfg.lr_aux)?; + aux_opt_b_size_short.wd = 0.0; // Load the aux→encoder accumulator cubin (single kernel, // `aux_vec_add_inplace`). Loaded once at construction so the // captured-graph hot path doesn't pay a load-cubin cost. @@ -1873,38 +1959,73 @@ impl PerceptionTrainer { // borrow-after-move). All allocations use the trainer's primary // stream, matching the rest of the trainer's buffer ownership pattern. let aux_h_per_k_d = stream.alloc_zeros::(k * cfg.n_batch * AUX_HIDDEN)?; - let aux_y_long_hat_per_k_d = stream - .alloc_zeros::(k * cfg.n_batch * N_AUX_HORIZONS)?; - let aux_y_short_hat_per_k_d = stream - .alloc_zeros::(k * cfg.n_batch * N_AUX_HORIZONS)?; - let aux_y_long_per_k_d = stream - .alloc_zeros::(k * cfg.n_batch * N_AUX_HORIZONS)?; - let aux_y_short_per_k_d = stream - .alloc_zeros::(k * cfg.n_batch * N_AUX_HORIZONS)?; - let stg_aux_y_long = unsafe { - MappedF32Buffer::new(k * cfg.n_batch * N_AUX_HORIZONS) - } - .map_err(|e| anyhow::anyhow!("stg_aux_y_long: {e}"))?; - let stg_aux_y_short = unsafe { - MappedF32Buffer::new(k * cfg.n_batch * N_AUX_HORIZONS) - } - .map_err(|e| anyhow::anyhow!("stg_aux_y_short: {e}"))?; - let aux_grad_y_long_hat_per_k_d = stream - .alloc_zeros::(k * cfg.n_batch * N_AUX_HORIZONS)?; - let aux_grad_y_short_hat_per_k_d = stream - .alloc_zeros::(k * cfg.n_batch * N_AUX_HORIZONS)?; - let aux_loss_long_d = stream.alloc_zeros::(1)?; - let aux_loss_short_d = stream.alloc_zeros::(1)?; - let aux_valid_long_d = stream.alloc_zeros::(1)?; - let aux_valid_short_d = stream.alloc_zeros::(1)?; - let aux_loss_long_host_d = unsafe { MappedF32Buffer::new(1) } - .map_err(|e| anyhow::anyhow!("aux_loss_long_host: {e}"))?; - let aux_loss_short_host_d = unsafe { MappedF32Buffer::new(1) } - .map_err(|e| anyhow::anyhow!("aux_loss_short_host: {e}"))?; - let aux_valid_long_host_d = unsafe { MappedI32Buffer::new(1) } - .map_err(|e| anyhow::anyhow!("aux_valid_long_host: {e}"))?; - let aux_valid_short_host_d = unsafe { MappedI32Buffer::new(1) } - .map_err(|e| anyhow::anyhow!("aux_valid_short_host: {e}"))?; + let alloc_per_k = || -> Result> { + stream + .alloc_zeros::(k * cfg.n_batch * N_AUX_HORIZONS) + .map_err(Into::into) + }; + let stg_per_k = |name: &str| -> Result { + unsafe { MappedF32Buffer::new(k * cfg.n_batch * N_AUX_HORIZONS) } + .map_err(|e| anyhow::anyhow!("{name}: {e}")) + }; + // Four prediction buffers (prof_long_logit, size_long_pred, + // prof_short_logit, size_short_pred) — written by the new aux + // forward kernel. + let aux_prof_long_logit_per_k_d = alloc_per_k()?; + let aux_size_long_pred_per_k_d = alloc_per_k()?; + let aux_prof_short_logit_per_k_d = alloc_per_k()?; + let aux_size_short_pred_per_k_d = alloc_per_k()?; + // Four label buffers (prof × {long, short} + size × {long, short}) + // — uploaded each step from the loader's A+B label set. + let aux_y_prof_long_per_k_d = alloc_per_k()?; + let aux_y_size_long_per_k_d = alloc_per_k()?; + let aux_y_prof_short_per_k_d = alloc_per_k()?; + let aux_y_size_short_per_k_d = alloc_per_k()?; + let stg_aux_y_prof_long = stg_per_k("stg_aux_y_prof_long")?; + let stg_aux_y_size_long = stg_per_k("stg_aux_y_size_long")?; + let stg_aux_y_prof_short = stg_per_k("stg_aux_y_prof_short")?; + let stg_aux_y_size_short = stg_per_k("stg_aux_y_size_short")?; + // Per-horizon BCE positive-class weight buffers (length + // `N_AUX_HORIZONS` each). Uploaded each step from + // `last_pos_fraction` → clamp. CB5 wires the upload + call. + let aux_pos_weight_long_d = stream.alloc_zeros::(N_AUX_HORIZONS)?; + let aux_pos_weight_short_d = stream.alloc_zeros::(N_AUX_HORIZONS)?; + let stg_aux_pos_weight_long = unsafe { MappedF32Buffer::new(N_AUX_HORIZONS) } + .map_err(|e| anyhow::anyhow!("stg_aux_pos_weight_long: {e}"))?; + let stg_aux_pos_weight_short = unsafe { MappedF32Buffer::new(N_AUX_HORIZONS) } + .map_err(|e| anyhow::anyhow!("stg_aux_pos_weight_short: {e}"))?; + // Four grad_out buffers (per-K) — written by the loss kernels; + // consumed by aux_heads_bwd. CB3+CB4 holding pattern: zeroed each + // step (no loss call); CB5 wires the BCE + Huber calls. + let aux_grad_prof_long_logit_per_k_d = alloc_per_k()?; + let aux_grad_size_long_pred_per_k_d = alloc_per_k()?; + let aux_grad_prof_short_logit_per_k_d = alloc_per_k()?; + let aux_grad_size_short_pred_per_k_d = alloc_per_k()?; + // Four scalar loss outputs + four valid counts (per (direction × head)). + let aux_loss_prof_long_d = stream.alloc_zeros::(1)?; + let aux_loss_size_long_d = stream.alloc_zeros::(1)?; + let aux_loss_prof_short_d = stream.alloc_zeros::(1)?; + let aux_loss_size_short_d = stream.alloc_zeros::(1)?; + let aux_valid_prof_long_d = stream.alloc_zeros::(1)?; + let aux_valid_size_long_d = stream.alloc_zeros::(1)?; + let aux_valid_prof_short_d = stream.alloc_zeros::(1)?; + let aux_valid_size_short_d = stream.alloc_zeros::(1)?; + let aux_loss_prof_long_host_d = unsafe { MappedF32Buffer::new(1) } + .map_err(|e| anyhow::anyhow!("aux_loss_prof_long_host: {e}"))?; + let aux_loss_size_long_host_d = unsafe { MappedF32Buffer::new(1) } + .map_err(|e| anyhow::anyhow!("aux_loss_size_long_host: {e}"))?; + let aux_loss_prof_short_host_d = unsafe { MappedF32Buffer::new(1) } + .map_err(|e| anyhow::anyhow!("aux_loss_prof_short_host: {e}"))?; + let aux_loss_size_short_host_d = unsafe { MappedF32Buffer::new(1) } + .map_err(|e| anyhow::anyhow!("aux_loss_size_short_host: {e}"))?; + let aux_valid_prof_long_host_d = unsafe { MappedI32Buffer::new(1) } + .map_err(|e| anyhow::anyhow!("aux_valid_prof_long_host: {e}"))?; + let aux_valid_size_long_host_d = unsafe { MappedI32Buffer::new(1) } + .map_err(|e| anyhow::anyhow!("aux_valid_size_long_host: {e}"))?; + let aux_valid_prof_short_host_d = unsafe { MappedI32Buffer::new(1) } + .map_err(|e| anyhow::anyhow!("aux_valid_prof_short_host: {e}"))?; + let aux_valid_size_short_host_d = unsafe { MappedI32Buffer::new(1) } + .map_err(|e| anyhow::anyhow!("aux_valid_size_short_host: {e}"))?; Ok(Self { cfg: cfg.clone(), h_new_per_k_d: stream.alloc_zeros::(k * cfg.n_batch * n_hid)?, @@ -2188,27 +2309,48 @@ impl PerceptionTrainer { h_mag_per_bucket_d, h_mag_per_bucket_staged, - // SDD-3 Layer B5: aux supervision fields. + // SDD-3 CB3+CB4: aux supervision fields (A+B paired heads). aux_trunk, aux_heads, - aux_huber, + aux_bce, + aux_huber_masked, aux_h_per_k_d, - aux_y_long_hat_per_k_d, - aux_y_short_hat_per_k_d, - aux_y_long_per_k_d, - aux_y_short_per_k_d, - stg_aux_y_long, - stg_aux_y_short, - aux_grad_y_long_hat_per_k_d, - aux_grad_y_short_hat_per_k_d, - aux_loss_long_d, - aux_loss_short_d, - aux_valid_long_d, - aux_valid_short_d, - aux_loss_long_host_d, - aux_loss_short_host_d, - aux_valid_long_host_d, - aux_valid_short_host_d, + aux_prof_long_logit_per_k_d, + aux_size_long_pred_per_k_d, + aux_prof_short_logit_per_k_d, + aux_size_short_pred_per_k_d, + aux_y_prof_long_per_k_d, + aux_y_size_long_per_k_d, + aux_y_prof_short_per_k_d, + aux_y_size_short_per_k_d, + stg_aux_y_prof_long, + stg_aux_y_size_long, + stg_aux_y_prof_short, + stg_aux_y_size_short, + aux_pos_weight_long_d, + aux_pos_weight_short_d, + stg_aux_pos_weight_long, + stg_aux_pos_weight_short, + aux_grad_prof_long_logit_per_k_d, + aux_grad_size_long_pred_per_k_d, + aux_grad_prof_short_logit_per_k_d, + aux_grad_size_short_pred_per_k_d, + aux_loss_prof_long_d, + aux_loss_size_long_d, + aux_loss_prof_short_d, + aux_loss_size_short_d, + aux_valid_prof_long_d, + aux_valid_size_long_d, + aux_valid_prof_short_d, + aux_valid_size_short_d, + aux_loss_prof_long_host_d, + aux_loss_size_long_host_d, + aux_loss_prof_short_host_d, + aux_loss_size_short_host_d, + aux_valid_prof_long_host_d, + aux_valid_size_long_host_d, + aux_valid_prof_short_host_d, + aux_valid_size_short_host_d, aux_trunk_grad_w_in_scratch_d, aux_trunk_grad_w_rec_scratch_d, aux_trunk_grad_b_scratch_d, @@ -2221,23 +2363,35 @@ impl PerceptionTrainer { aux_grad_h_old_step_d, aux_grad_x_step_d, aux_zero_h_d, - aux_heads_grad_w_long_scratch_d, - aux_heads_grad_b_long_scratch_d, - aux_heads_grad_w_short_scratch_d, - aux_heads_grad_b_short_scratch_d, - aux_heads_grad_w_long_d, - aux_heads_grad_b_long_d, - aux_heads_grad_w_short_d, - aux_heads_grad_b_short_d, + aux_heads_grad_w_prof_long_scratch_d, + aux_heads_grad_b_prof_long_scratch_d, + aux_heads_grad_w_size_long_scratch_d, + aux_heads_grad_b_size_long_scratch_d, + aux_heads_grad_w_prof_short_scratch_d, + aux_heads_grad_b_prof_short_scratch_d, + aux_heads_grad_w_size_short_scratch_d, + aux_heads_grad_b_size_short_scratch_d, + aux_heads_grad_w_prof_long_d, + aux_heads_grad_b_prof_long_d, + aux_heads_grad_w_size_long_d, + aux_heads_grad_b_size_long_d, + aux_heads_grad_w_prof_short_d, + aux_heads_grad_b_prof_short_d, + aux_heads_grad_w_size_short_d, + aux_heads_grad_b_size_short_d, aux_grad_h_aux_d, aux_opt_w_in, aux_opt_w_rec, aux_opt_b, aux_opt_tau, - aux_opt_w_long, - aux_opt_b_long, - aux_opt_w_short, - aux_opt_b_short, + aux_opt_w_prof_long, + aux_opt_b_prof_long, + aux_opt_w_size_long, + aux_opt_b_size_long, + aux_opt_w_prof_short, + aux_opt_b_prof_short, + aux_opt_w_size_short, + aux_opt_b_size_short, // E3 design memo: aux is a PASSIVE auditor at Phase 1 — gradient // to encoder is masked. Lifted conditionally once aux Huber EMA // falls below the threshold AND directional accuracy clears @@ -2303,10 +2457,14 @@ impl PerceptionTrainer { self.aux_opt_w_rec.lr = lr; self.aux_opt_b.lr = lr; self.aux_opt_tau.lr = lr * 0.1; - self.aux_opt_w_long.lr = lr; - self.aux_opt_b_long.lr = lr; - self.aux_opt_w_short.lr = lr; - self.aux_opt_b_short.lr = lr; + self.aux_opt_w_prof_long.lr = lr; + self.aux_opt_b_prof_long.lr = lr; + self.aux_opt_w_size_long.lr = lr; + self.aux_opt_b_size_long.lr = lr; + self.aux_opt_w_prof_short.lr = lr; + self.aux_opt_b_prof_short.lr = lr; + self.aux_opt_w_size_short.lr = lr; + self.aux_opt_b_size_short.lr = lr; } /// One training step on a single sequence — thin wrapper around @@ -2558,23 +2716,26 @@ impl PerceptionTrainer { } } } - // SDD-3 CB2 transition state: aux outcome staging now sources the - // prof-binary arrays (long/short) from the A+B label set. Layout is - // unchanged ([K × B × N_AUX_HORIZONS] row-major, horizon innermost) - // so the existing per-K backward slot pointers in the captured graph - // still resolve. The σ-normalized size targets are routed through - // CB5 — for now the holding-pattern D-era Huber kernel runs on the - // prof-binary arrays which provides a finite, NaN-masked signal - // (no convergence guarantees until CB5 lands the proper BCE+Huber). + // SDD-3 CB3+CB4: aux label staging now routes ALL four A+B + // arrays (prof × {long, short} + size × {long, short}) into + // distinct mapped-pinned buffers. Layout is `[K × B × N_AUX_HORIZONS]` + // row-major (horizon innermost) per the existing convention. The + // per-K backward slot pointers consume these via the new BCE + + // masked-Huber kernels (call sites land in CB5; CB3+CB4 stages + // the data and provisions the buffers). { - let dst_long = self.stg_aux_y_long.host_slice_mut(); - let dst_short = self.stg_aux_y_short.host_slice_mut(); + let dst_pl = self.stg_aux_y_prof_long.host_slice_mut(); + let dst_sl = self.stg_aux_y_size_long.host_slice_mut(); + let dst_ps = self.stg_aux_y_prof_short.host_slice_mut(); + let dst_ss = self.stg_aux_y_size_short.host_slice_mut(); for k in 0..k_seq { for b_idx in 0..b_sz { for h in 0..N_AUX_HORIZONS { let idx = (k * b_sz + b_idx) * N_AUX_HORIZONS + h; - dst_long[idx] = outcome_prof_long_batch[b_idx][k][h]; - dst_short[idx] = outcome_prof_short_batch[b_idx][k][h]; + dst_pl[idx] = outcome_prof_long_batch[b_idx][k][h]; + dst_sl[idx] = outcome_size_long_batch[b_idx][k][h]; + dst_ps[idx] = outcome_prof_short_batch[b_idx][k][h]; + dst_ss[idx] = outcome_size_short_batch[b_idx][k][h]; } } } @@ -3289,117 +3450,108 @@ impl PerceptionTrainer { } } - // ── 6. SDD-3 Layer B5: aux supervision host-side ISV + lift check ── + // ── 6. SDD-3 CB3+CB4: aux supervision host-side ISV + lift check ── // - // The captured graph has already populated - // `aux_loss_long_host_d` — Σ Huber loss over long-direction valid entries - // `aux_loss_short_host_d` — Σ Huber loss over short-direction valid entries - // `aux_valid_long_host_d` — #{(k, b, h) : isfinite(y_long_true)} - // `aux_valid_short_host_d` — #{(k, b, h) : isfinite(y_short_true)} - // via DtoD shadows in the dispatch path. We compute the per-step - // average Huber and a coarse per-horizon directional-accuracy - // proxy from the staged labels + predictions, then update Wiener-α - // EMAs per `pearl_first_observation_bootstrap` (sentinel = 0, - // first observation replaces directly; subsequent steps blend). + // The captured graph has already populated (when CB5 lands the + // BCE + masked-Huber calls): + // `aux_loss_prof_long_host_d` — Σ BCE loss over prof_long valid entries + // `aux_loss_size_long_host_d` — Σ Huber loss over size_long valid entries + // `aux_loss_prof_short_host_d` — Σ BCE loss over prof_short valid entries + // `aux_loss_size_short_host_d` — Σ Huber loss over size_short valid entries + // (and matching valid counts). // - // Per-horizon dir_acc is computed from the staged predictions - // (mapped-pinned host read on the long/short hat per-K buffers — - // a small one-shot DtoH after the end-of-step sync is acceptable - // for the per-horizon proxy; the value gates the lift, not the - // hot path). + // CB3+CB4 holding pattern: the four loss/valid host buffers are + // never written by the captured graph (the loss kernels are not + // yet called); they retain their zero-init contents, so + // `n_total = 0` and the EMA path is a no-op. Once CB5 wires the + // loss kernel calls + DtoD shadows, this branch fires every step. // - // No-op (sentinel + flag stays at construction defaults) on any - // step where neither direction had valid labels (e.g. very first - // sequences before D-labels enter the window). + // Per-horizon dir_acc is computed from the staged labels + + // prof-head predictions (a small one-shot DtoH after the + // end-of-step sync is acceptable; the value gates the lift, not + // the hot path). { - let total_long_loss = unsafe { - std::ptr::read_volatile(self.aux_loss_long_host_d.host_ptr) + let total_loss_pl = unsafe { + std::ptr::read_volatile(self.aux_loss_prof_long_host_d.host_ptr) }; - let total_short_loss = unsafe { - std::ptr::read_volatile(self.aux_loss_short_host_d.host_ptr) + let total_loss_sl = unsafe { + std::ptr::read_volatile(self.aux_loss_size_long_host_d.host_ptr) }; - let n_long = unsafe { - std::ptr::read_volatile(self.aux_valid_long_host_d.host_ptr) + let total_loss_ps = unsafe { + std::ptr::read_volatile(self.aux_loss_prof_short_host_d.host_ptr) + }; + let total_loss_ss = unsafe { + std::ptr::read_volatile(self.aux_loss_size_short_host_d.host_ptr) + }; + let n_pl = unsafe { + std::ptr::read_volatile(self.aux_valid_prof_long_host_d.host_ptr) } as i64; - let n_short = unsafe { - std::ptr::read_volatile(self.aux_valid_short_host_d.host_ptr) + let n_sl = unsafe { + std::ptr::read_volatile(self.aux_valid_size_long_host_d.host_ptr) } as i64; - let n_total = n_long + n_short; + let n_ps = unsafe { + std::ptr::read_volatile(self.aux_valid_prof_short_host_d.host_ptr) + } as i64; + let n_ss = unsafe { + std::ptr::read_volatile(self.aux_valid_size_short_host_d.host_ptr) + } as i64; + let n_total = n_pl + n_sl + n_ps + n_ss; if n_total > 0 { - let mean_huber = (total_long_loss + total_short_loss) + // Joint loss mean — CB3+CB4 fields `aux_huber_ema_per_h` + // is still named after the D-era Huber but now tracks the + // joint BCE+Huber aux signal. CB5 may split the EMAs. + let mean_loss = (total_loss_pl + + total_loss_sl + + total_loss_ps + + total_loss_ss) / (n_total as f32); - // Per-horizon mean (rough proxy: split the joint sum - // proportional to per-horizon valid counts derived from the - // staged labels). Reads `stg_aux_y_long` / `stg_aux_y_short` - // host shadows (already populated above the captured - // region). Per `pearl_blend_formulas_must_have_permanent_floor`, - // the per-horizon mean uses a floor of 1 on the valid count - // to avoid divide-by-zero on horizons with no D-labels in - // the current batch. - let labels_long = self.stg_aux_y_long.read_all(); - let labels_short = self.stg_aux_y_short.read_all(); + let labels_pl = self.stg_aux_y_prof_long.read_all(); + let labels_ps = self.stg_aux_y_prof_short.read_all(); let n_per_kb = k_seq * b_sz; - let mut per_h_count = [0_i64; N_AUX_HORIZONS]; - for h in 0..N_AUX_HORIZONS { - for kb in 0..n_per_kb { - if labels_long[kb * N_AUX_HORIZONS + h].is_finite() { - per_h_count[h] += 1; - } - if labels_short[kb * N_AUX_HORIZONS + h].is_finite() { - per_h_count[h] += 1; - } - } - } // First-observation bootstrap per `pearl_first_observation_bootstrap`. // Wiener-α floor = 0.4 per `pearl_wiener_alpha_floor_for_nonstationary` // (the gradient-driven trunk is non-stationary — aux loss // trajectory drifts as encoder representations adapt). - const AUX_HUBER_ALPHA: f32 = 0.4; + const AUX_LOSS_ALPHA: f32 = 0.4; if !self.aux_huber_first_obs { for h in 0..N_AUX_HORIZONS { - self.aux_huber_ema_per_h[h] = mean_huber; + self.aux_huber_ema_per_h[h] = mean_loss; } self.aux_huber_first_obs = true; } else { for h in 0..N_AUX_HORIZONS { - // Per-horizon weight = per_h_count / total_count; - // no per-h fidelity beyond the joint mean is - // available without splitting the kernel, so we - // anchor every horizon's EMA on the joint mean and - // let the dir_acc gate provide per-horizon - // discrimination instead. - self.aux_huber_ema_per_h[h] = (1.0 - AUX_HUBER_ALPHA) + self.aux_huber_ema_per_h[h] = (1.0 - AUX_LOSS_ALPHA) * self.aux_huber_ema_per_h[h] - + AUX_HUBER_ALPHA * mean_huber; + + AUX_LOSS_ALPHA * mean_loss; } } - // Per-horizon dir_acc proxy: count fraction of (k, b) - // entries where sign(y_long_hat - y_short_hat) matches - // sign(y_long_true - y_short_true) within this batch. - // Done host-side from a one-shot DtoH read of the per-K - // prediction buffers (size = K × B × N_AUX_HORIZONS each, - // total ≤ 32 × 8 × 3 = 768 floats per side at typical - // configs — a trivial sync cost off the captured path). - let mut y_long_hat = vec![0.0_f32; k_seq * b_sz * N_AUX_HORIZONS]; - let mut y_short_hat = vec![0.0_f32; k_seq * b_sz * N_AUX_HORIZONS]; + // Per-horizon dir_acc proxy on the PROF heads: count + // fraction of (k, b) entries where + // sign(prof_long_logit - prof_short_logit) matches + // sign(y_prof_long - y_prof_short) within this batch. + // Reads the prof_*_logit per-K buffers via a one-shot DtoH + // after end-of-step sync (≤ 32 × 8 × 3 = 768 floats per + // side — trivial cost off the captured path). + let mut prof_long = vec![0.0_f32; k_seq * b_sz * N_AUX_HORIZONS]; + let mut prof_short = vec![0.0_f32; k_seq * b_sz * N_AUX_HORIZONS]; self.stream - .memcpy_dtoh(&self.aux_y_long_hat_per_k_d, y_long_hat.as_mut_slice()) - .context("aux y_long_hat dtoh for dir_acc")?; + .memcpy_dtoh(&self.aux_prof_long_logit_per_k_d, prof_long.as_mut_slice()) + .context("aux prof_long_logit dtoh for dir_acc")?; self.stream - .memcpy_dtoh(&self.aux_y_short_hat_per_k_d, y_short_hat.as_mut_slice()) - .context("aux y_short_hat dtoh for dir_acc")?; + .memcpy_dtoh(&self.aux_prof_short_logit_per_k_d, prof_short.as_mut_slice()) + .context("aux prof_short_logit dtoh for dir_acc")?; let mut per_h_match = [0_i64; N_AUX_HORIZONS]; let mut per_h_valid = [0_i64; N_AUX_HORIZONS]; for h in 0..N_AUX_HORIZONS { for kb in 0..n_per_kb { let idx = kb * N_AUX_HORIZONS + h; - let lt = labels_long[idx]; - let st = labels_short[idx]; + let lt = labels_pl[idx]; + let st = labels_ps[idx]; if !lt.is_finite() || !st.is_finite() { continue; } let true_diff = lt - st; - let pred_diff = y_long_hat[idx] - y_short_hat[idx]; + let pred_diff = prof_long[idx] - prof_short[idx]; if (true_diff > 0.0 && pred_diff > 0.0) || (true_diff < 0.0 && pred_diff < 0.0) || (true_diff == 0.0 && pred_diff == 0.0) @@ -3721,28 +3873,44 @@ impl PerceptionTrainer { dst_ptr, self.stg_labels.dev_ptr, nbytes, self.stream.cu_stream(), ).context("labels upload DtoD")?; } - // SDD-3 Layer B5: aux outcome label uploads. Same DtoD pattern as + // SDD-3 CB3+CB4: aux A+B label uploads. Same DtoD pattern as // `labels_per_k_d` (mapped-pinned staging → device buffer via - // captured async memcpy). + // captured async memcpy). Four label tensors total — prof + size + // per direction. let total_aux_labels = k_seq * b_sz * N_AUX_HORIZONS; unsafe { - let (dst_ptr, _g) = self.aux_y_long_per_k_d.device_ptr_mut(&self.stream); let nbytes = total_aux_labels * std::mem::size_of::(); - cudarc::driver::result::memcpy_dtod_async( - dst_ptr, - self.stg_aux_y_long.dev_ptr, - nbytes, - self.stream.cu_stream(), - ) - .context("aux y_long upload DtoD")?; - let (dst_ptr, _g) = self.aux_y_short_per_k_d.device_ptr_mut(&self.stream); - cudarc::driver::result::memcpy_dtod_async( - dst_ptr, - self.stg_aux_y_short.dev_ptr, - nbytes, - self.stream.cu_stream(), - ) - .context("aux y_short upload DtoD")?; + let upload_dtod = + |dst: &mut CudaSlice, src_ptr: u64, ctx: &str| -> Result<()> { + let (dst_ptr, _g) = dst.device_ptr_mut(&self.stream); + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, + src_ptr, + nbytes, + self.stream.cu_stream(), + ) + .with_context(|| ctx.to_string()) + }; + upload_dtod( + &mut self.aux_y_prof_long_per_k_d, + self.stg_aux_y_prof_long.dev_ptr, + "aux y_prof_long upload DtoD", + )?; + upload_dtod( + &mut self.aux_y_size_long_per_k_d, + self.stg_aux_y_size_long.dev_ptr, + "aux y_size_long upload DtoD", + )?; + upload_dtod( + &mut self.aux_y_prof_short_per_k_d, + self.stg_aux_y_prof_short.dev_ptr, + "aux y_prof_short upload DtoD", + )?; + upload_dtod( + &mut self.aux_y_size_short_per_k_d, + self.stg_aux_y_size_short.dev_ptr, + "aux y_size_short upload DtoD", + )?; } // CfC per-batch grad scratch (Phase B): zero ONCE per step; K-loop @@ -3797,14 +3965,34 @@ impl PerceptionTrainer { .map_err(|e| anyhow::anyhow!("zero aux_trunk_grad_b_scratch: {e}"))?; self.stream.memset_zeros(&mut self.aux_trunk_grad_tau_scratch_d) .map_err(|e| anyhow::anyhow!("zero aux_trunk_grad_tau_scratch: {e}"))?; - self.stream.memset_zeros(&mut self.aux_heads_grad_w_long_scratch_d) - .map_err(|e| anyhow::anyhow!("zero aux_heads_grad_w_long_scratch: {e}"))?; - self.stream.memset_zeros(&mut self.aux_heads_grad_b_long_scratch_d) - .map_err(|e| anyhow::anyhow!("zero aux_heads_grad_b_long_scratch: {e}"))?; - self.stream.memset_zeros(&mut self.aux_heads_grad_w_short_scratch_d) - .map_err(|e| anyhow::anyhow!("zero aux_heads_grad_w_short_scratch: {e}"))?; - self.stream.memset_zeros(&mut self.aux_heads_grad_b_short_scratch_d) - .map_err(|e| anyhow::anyhow!("zero aux_heads_grad_b_short_scratch: {e}"))?; + // CB3+CB4: 8 aux-head grad scratches (prof_{long,short} × {w,b} + + // size_{long,short} × {w,b}). Zeroed once per step; the per-K + // backward kernel `+=` accumulates into them, then reduce_axis0 + // collapses to the 8 reduced gradients consumed by Adam. + self.stream + .memset_zeros(&mut self.aux_heads_grad_w_prof_long_scratch_d) + .map_err(|e| anyhow::anyhow!("zero aux_heads_grad_w_prof_long_scratch: {e}"))?; + self.stream + .memset_zeros(&mut self.aux_heads_grad_b_prof_long_scratch_d) + .map_err(|e| anyhow::anyhow!("zero aux_heads_grad_b_prof_long_scratch: {e}"))?; + self.stream + .memset_zeros(&mut self.aux_heads_grad_w_size_long_scratch_d) + .map_err(|e| anyhow::anyhow!("zero aux_heads_grad_w_size_long_scratch: {e}"))?; + self.stream + .memset_zeros(&mut self.aux_heads_grad_b_size_long_scratch_d) + .map_err(|e| anyhow::anyhow!("zero aux_heads_grad_b_size_long_scratch: {e}"))?; + self.stream + .memset_zeros(&mut self.aux_heads_grad_w_prof_short_scratch_d) + .map_err(|e| anyhow::anyhow!("zero aux_heads_grad_w_prof_short_scratch: {e}"))?; + self.stream + .memset_zeros(&mut self.aux_heads_grad_b_prof_short_scratch_d) + .map_err(|e| anyhow::anyhow!("zero aux_heads_grad_b_prof_short_scratch: {e}"))?; + self.stream + .memset_zeros(&mut self.aux_heads_grad_w_size_short_scratch_d) + .map_err(|e| anyhow::anyhow!("zero aux_heads_grad_w_size_short_scratch: {e}"))?; + self.stream + .memset_zeros(&mut self.aux_heads_grad_b_size_short_scratch_d) + .map_err(|e| anyhow::anyhow!("zero aux_heads_grad_b_size_short_scratch: {e}"))?; self.stream.memset_zeros(&mut self.aux_grad_h_carry_d) .map_err(|e| anyhow::anyhow!("zero aux_grad_h_carry: {e}"))?; self.stream.memset_zeros(&mut self.aux_zero_h_d) @@ -4033,10 +4221,14 @@ impl PerceptionTrainer { let dt_s_aux = 1.0_f32; { let (aux_h_base, _g_aux_h) = self.aux_h_per_k_d.device_ptr_mut(&self.stream); - let (aux_ylh_base, _g_aux_ylh) = - self.aux_y_long_hat_per_k_d.device_ptr_mut(&self.stream); - let (aux_ysh_base, _g_aux_ysh) = - self.aux_y_short_hat_per_k_d.device_ptr_mut(&self.stream); + let (aux_pl_base, _g_aux_pl) = + self.aux_prof_long_logit_per_k_d.device_ptr_mut(&self.stream); + let (aux_sl_base, _g_aux_sl) = + self.aux_size_long_pred_per_k_d.device_ptr_mut(&self.stream); + let (aux_ps_base, _g_aux_ps) = + self.aux_prof_short_logit_per_k_d.device_ptr_mut(&self.stream); + let (aux_ss_base, _g_aux_ss) = + self.aux_size_short_pred_per_k_d.device_ptr_mut(&self.stream); let henr_t_base_aux = { let (p, _g) = self.h_enriched_seq_t_d.cuda_data().device_ptr(&self.stream); p @@ -4053,8 +4245,10 @@ impl PerceptionTrainer { aux_h_base + ((k - 1) * kb_aux_hid_bytes) as u64 }; let x_k_ptr = henr_t_base_aux + (k * kb_hid_bytes) as u64; - let aux_ylh_k_ptr = aux_ylh_base + (k * kb_aux_nh_bytes) as u64; - let aux_ysh_k_ptr = aux_ysh_base + (k * kb_aux_nh_bytes) as u64; + let aux_pl_k_ptr = aux_pl_base + (k * kb_aux_nh_bytes) as u64; + let aux_sl_k_ptr = aux_sl_base + (k * kb_aux_nh_bytes) as u64; + let aux_ps_k_ptr = aux_ps_base + (k * kb_aux_nh_bytes) as u64; + let aux_ss_k_ptr = aux_ss_base + (k * kb_aux_nh_bytes) as u64; // Aux trunk forward — block-per-batch, AUX_HIDDEN threads. // Shared mem = (feat_dim + AUX_HIDDEN) * sizeof(float) @@ -4085,8 +4279,10 @@ impl PerceptionTrainer { } } - // Aux heads forward — block-per-batch, AUX_HIDDEN threads. - // Writes y_long_hat / y_short_hat for slot k. + // Aux heads forward (A+B paired) — block-per-batch, + // AUX_HIDDEN threads. Writes all four outputs + // (prof_long_logit, size_long_pred, prof_short_logit, + // size_short_pred) for slot k. { let cfg_aux_heads_fwd = LaunchConfig { grid_dim: (b_sz as u32, 1, 1), @@ -4096,54 +4292,56 @@ impl PerceptionTrainer { }; let mut launch = self.stream.launch_builder(&self.aux_heads.fwd_fn); launch - .arg(&self.aux_heads.w_long_d) - .arg(&self.aux_heads.b_long_d) - .arg(&self.aux_heads.w_short_d) - .arg(&self.aux_heads.b_short_d) + .arg(&self.aux_heads.w_prof_long_d) + .arg(&self.aux_heads.b_prof_long_d) + .arg(&self.aux_heads.w_size_long_d) + .arg(&self.aux_heads.b_size_long_d) + .arg(&self.aux_heads.w_prof_short_d) + .arg(&self.aux_heads.b_prof_short_d) + .arg(&self.aux_heads.w_size_short_d) + .arg(&self.aux_heads.b_size_short_d) .arg(&aux_h_new_k_ptr) .arg(&n_batch_i) - .arg(&aux_ylh_k_ptr) - .arg(&aux_ysh_k_ptr); + .arg(&aux_pl_k_ptr) + .arg(&aux_sl_k_ptr) + .arg(&aux_ps_k_ptr) + .arg(&aux_ss_k_ptr); unsafe { launch.launch(cfg_aux_heads_fwd).context("aux heads fwd k")?; } } } - drop((_g_aux_h, _g_aux_ylh, _g_aux_ysh)); + drop((_g_aux_h, _g_aux_pl, _g_aux_sl, _g_aux_ps, _g_aux_ss)); } - // ── 4c. SDD-3 Layer B5: aux Huber loss + grad (long + short) ── + // ── 4c. SDD-3 CB3+CB4: aux BCE + Huber loss + grad ── // - // Two launches over the flattened [K × B × N_AUX_HORIZONS] arrays. - // Each emits the unreduced Huber sum, the per-element dL/dy_hat - // (NaN-masked zero), and the valid-element count for that direction. - // The summed loss + valid count are shadowed to mapped-pinned - // host buffers (post-step sync) to drive the per-horizon Huber - // EMA + lift logic in `step_batched`. + // CB3+CB4 HOLDING PATTERN: the new BCE (prof heads) + NaN-masked + // Huber (size heads) kernels exist but are not yet called from + // the captured graph — that's CB5's scope. To keep the bwd path + // safe in the meantime we zero the four `aux_grad_*_per_k_d` + // tensors here; the aux_heads bwd kernel then writes zero + // gradients through the head + trunk paths and the eight Adam + // optimizers become no-ops for the aux head weights this step. + // + // The four `aux_loss_*_d` / `aux_valid_*_d` scalars are left at + // their zero-init contents (DtoD shadows fire each step but + // copy the previous zero through the host buffers), which keeps + // the host-side EMA bootstrap branch dormant until CB5 lands the + // real loss calls. { - let n_total_aux = (k_seq * b_sz * N_AUX_HORIZONS) as i32; - aux_huber_loss_gpu( - &self.stream, - &self.aux_huber.fn_, - &self.aux_y_long_hat_per_k_d, - &self.aux_y_long_per_k_d, - DEFAULT_HUBER_DELTA, - n_total_aux, - &mut self.aux_loss_long_d, - &mut self.aux_grad_y_long_hat_per_k_d, - &mut self.aux_valid_long_d, - )?; - aux_huber_loss_gpu( - &self.stream, - &self.aux_huber.fn_, - &self.aux_y_short_hat_per_k_d, - &self.aux_y_short_per_k_d, - DEFAULT_HUBER_DELTA, - n_total_aux, - &mut self.aux_loss_short_d, - &mut self.aux_grad_y_short_hat_per_k_d, - &mut self.aux_valid_short_d, - )?; + self.stream + .memset_zeros(&mut self.aux_grad_prof_long_logit_per_k_d) + .map_err(|e| anyhow::anyhow!("zero aux_grad_prof_long_logit: {e}"))?; + self.stream + .memset_zeros(&mut self.aux_grad_size_long_pred_per_k_d) + .map_err(|e| anyhow::anyhow!("zero aux_grad_size_long_pred: {e}"))?; + self.stream + .memset_zeros(&mut self.aux_grad_prof_short_logit_per_k_d) + .map_err(|e| anyhow::anyhow!("zero aux_grad_prof_short_logit: {e}"))?; + self.stream + .memset_zeros(&mut self.aux_grad_size_short_pred_per_k_d) + .map_err(|e| anyhow::anyhow!("zero aux_grad_size_short_pred: {e}"))?; } // ── Controller D signal: per-bucket mean(|h_state|) at K-1 ── @@ -5015,11 +5213,17 @@ impl PerceptionTrainer { let aux_lift_active = !self.stop_grad_aux_to_encoder; { let (aux_h_base_bwd, _g_aux_h) = self.aux_h_per_k_d.device_ptr_mut(&self.stream); - let (aux_grad_ylh_base, _g_aux_glh) = self - .aux_grad_y_long_hat_per_k_d + let (aux_grad_pl_base, _g_aux_gpl) = self + .aux_grad_prof_long_logit_per_k_d .device_ptr_mut(&self.stream); - let (aux_grad_ysh_base, _g_aux_gsh) = self - .aux_grad_y_short_hat_per_k_d + let (aux_grad_sl_base, _g_aux_gsl) = self + .aux_grad_size_long_pred_per_k_d + .device_ptr_mut(&self.stream); + let (aux_grad_ps_base, _g_aux_gps) = self + .aux_grad_prof_short_logit_per_k_d + .device_ptr_mut(&self.stream); + let (aux_grad_ss_base, _g_aux_gss) = self + .aux_grad_size_short_pred_per_k_d .device_ptr_mut(&self.stream); let henr_t_base_aux_bwd = { let (p, _g) = self.h_enriched_seq_t_d.cuda_data().device_ptr(&self.stream); @@ -5042,17 +5246,19 @@ impl PerceptionTrainer { aux_h_base_bwd + ((k - 1) * kb_aux_hid_bytes) as u64 }; let x_k_ptr = henr_t_base_aux_bwd + (k * kb_hid_bytes) as u64; - let aux_grad_ylh_k_ptr = aux_grad_ylh_base + (k * kb_aux_nh_bytes) as u64; - let aux_grad_ysh_k_ptr = aux_grad_ysh_base + (k * kb_aux_nh_bytes) as u64; + let aux_grad_pl_k_ptr = aux_grad_pl_base + (k * kb_aux_nh_bytes) as u64; + let aux_grad_sl_k_ptr = aux_grad_sl_base + (k * kb_aux_nh_bytes) as u64; + let aux_grad_ps_k_ptr = aux_grad_ps_base + (k * kb_aux_nh_bytes) as u64; + let aux_grad_ss_k_ptr = aux_grad_ss_base + (k * kb_aux_nh_bytes) as u64; - // Aux heads backward — block-per-batch. - // Per-batch param grad scratch += accumulate via the - // per-k bwd kernel; reduce_axis0 collapses across B later. - // We write per-batch slot at offset (b * N_AUX_HORIZONS * AUX_HIDDEN). - // The per-K iterations use the FULL scratch buffer (no per-k - // stride) because the gradient sums over all K positions - // into the same [B, N_AUX_HORIZONS, AUX_HIDDEN] tensor — - // matching the main heads-bwd pattern. + // Aux heads backward — block-per-batch (A+B paired). + // Per-batch param grad scratch `+=` accumulates via the + // per-k bwd kernel for all eight head matrices; reduce_axis0 + // collapses across B later. The per-K iterations use the + // FULL scratch buffer (no per-k stride) because the + // gradient sums over all K positions into the same + // [B, N_AUX_HORIZONS, AUX_HIDDEN] tensor — matching the + // main heads-bwd pattern. { let cfg_aux_heads_bwd = LaunchConfig { grid_dim: (b_sz as u32, 1, 1), @@ -5062,16 +5268,24 @@ impl PerceptionTrainer { }; let mut launch = self.stream.launch_builder(&self.aux_heads.bwd_fn); launch - .arg(&self.aux_heads.w_long_d) - .arg(&self.aux_heads.w_short_d) + .arg(&self.aux_heads.w_prof_long_d) + .arg(&self.aux_heads.w_size_long_d) + .arg(&self.aux_heads.w_prof_short_d) + .arg(&self.aux_heads.w_size_short_d) .arg(&aux_h_new_k_ptr) - .arg(&aux_grad_ylh_k_ptr) - .arg(&aux_grad_ysh_k_ptr) + .arg(&aux_grad_pl_k_ptr) + .arg(&aux_grad_sl_k_ptr) + .arg(&aux_grad_ps_k_ptr) + .arg(&aux_grad_ss_k_ptr) .arg(&n_batch_i) - .arg(&mut self.aux_heads_grad_w_long_scratch_d) - .arg(&mut self.aux_heads_grad_b_long_scratch_d) - .arg(&mut self.aux_heads_grad_w_short_scratch_d) - .arg(&mut self.aux_heads_grad_b_short_scratch_d) + .arg(&mut self.aux_heads_grad_w_prof_long_scratch_d) + .arg(&mut self.aux_heads_grad_b_prof_long_scratch_d) + .arg(&mut self.aux_heads_grad_w_size_long_scratch_d) + .arg(&mut self.aux_heads_grad_b_size_long_scratch_d) + .arg(&mut self.aux_heads_grad_w_prof_short_scratch_d) + .arg(&mut self.aux_heads_grad_b_prof_short_scratch_d) + .arg(&mut self.aux_heads_grad_w_size_short_scratch_d) + .arg(&mut self.aux_heads_grad_b_size_short_scratch_d) .arg(&mut self.aux_grad_h_aux_d); unsafe { launch @@ -5184,7 +5398,14 @@ impl PerceptionTrainer { } } } - drop((_g_aux_h, _g_aux_glh, _g_aux_gsh, _g_aux_grad_henr)); + drop(( + _g_aux_h, + _g_aux_gpl, + _g_aux_gsl, + _g_aux_gps, + _g_aux_gss, + _g_aux_grad_henr, + )); } // Reduce per-batch aux grad scratches → final Adam-consumed grads. @@ -5234,33 +5455,65 @@ impl PerceptionTrainer { &mut self.aux_trunk_grad_tau_d, "reduce aux_trunk_grad_tau", )?; + // Eight head-grad reducers (prof_{long,short} × {w,b} + + // size_{long,short} × {w,b}). reduce_at_aux( N_AUX_HORIZONS * AUX_HIDDEN, - &self.aux_heads_grad_w_long_scratch_d, - &mut self.aux_heads_grad_w_long_d, - "reduce aux_heads_grad_w_long", + &self.aux_heads_grad_w_prof_long_scratch_d, + &mut self.aux_heads_grad_w_prof_long_d, + "reduce aux_heads_grad_w_prof_long", )?; reduce_at_aux( N_AUX_HORIZONS, - &self.aux_heads_grad_b_long_scratch_d, - &mut self.aux_heads_grad_b_long_d, - "reduce aux_heads_grad_b_long", + &self.aux_heads_grad_b_prof_long_scratch_d, + &mut self.aux_heads_grad_b_prof_long_d, + "reduce aux_heads_grad_b_prof_long", )?; reduce_at_aux( N_AUX_HORIZONS * AUX_HIDDEN, - &self.aux_heads_grad_w_short_scratch_d, - &mut self.aux_heads_grad_w_short_d, - "reduce aux_heads_grad_w_short", + &self.aux_heads_grad_w_size_long_scratch_d, + &mut self.aux_heads_grad_w_size_long_d, + "reduce aux_heads_grad_w_size_long", )?; reduce_at_aux( N_AUX_HORIZONS, - &self.aux_heads_grad_b_short_scratch_d, - &mut self.aux_heads_grad_b_short_d, - "reduce aux_heads_grad_b_short", + &self.aux_heads_grad_b_size_long_scratch_d, + &mut self.aux_heads_grad_b_size_long_d, + "reduce aux_heads_grad_b_size_long", + )?; + reduce_at_aux( + N_AUX_HORIZONS * AUX_HIDDEN, + &self.aux_heads_grad_w_prof_short_scratch_d, + &mut self.aux_heads_grad_w_prof_short_d, + "reduce aux_heads_grad_w_prof_short", + )?; + reduce_at_aux( + N_AUX_HORIZONS, + &self.aux_heads_grad_b_prof_short_scratch_d, + &mut self.aux_heads_grad_b_prof_short_d, + "reduce aux_heads_grad_b_prof_short", + )?; + reduce_at_aux( + N_AUX_HORIZONS * AUX_HIDDEN, + &self.aux_heads_grad_w_size_short_scratch_d, + &mut self.aux_heads_grad_w_size_short_d, + "reduce aux_heads_grad_w_size_short", + )?; + reduce_at_aux( + N_AUX_HORIZONS, + &self.aux_heads_grad_b_size_short_scratch_d, + &mut self.aux_heads_grad_b_size_short_d, + "reduce aux_heads_grad_b_size_short", )?; } // Aux Adam updates — separate group per `pearl_adam_normalizes_loss_weights`. + // Eight head optimizers + four trunk optimizers. In the CB3+CB4 + // holding pattern the head grads are all zero (the loss kernels + // are not yet called), so the Adam updates are no-ops on the head + // weights. The trunk path still runs since the prediction outputs + // travel through aux_heads bwd → grad_h_aux (all-zero grad + // upstream → all-zero grad downstream). self.aux_opt_w_in .step(&mut self.aux_trunk.w_in_d, &self.aux_trunk_grad_w_in_d)?; self.aux_opt_w_rec @@ -5269,60 +5522,85 @@ impl PerceptionTrainer { .step(&mut self.aux_trunk.b_d, &self.aux_trunk_grad_b_d)?; self.aux_opt_tau .step(&mut self.aux_trunk.tau_d, &self.aux_trunk_grad_tau_d)?; - self.aux_opt_w_long.step( - &mut self.aux_heads.w_long_d, - &self.aux_heads_grad_w_long_d, + self.aux_opt_w_prof_long.step( + &mut self.aux_heads.w_prof_long_d, + &self.aux_heads_grad_w_prof_long_d, )?; - self.aux_opt_b_long.step( - &mut self.aux_heads.b_long_d, - &self.aux_heads_grad_b_long_d, + self.aux_opt_b_prof_long.step( + &mut self.aux_heads.b_prof_long_d, + &self.aux_heads_grad_b_prof_long_d, )?; - self.aux_opt_w_short.step( - &mut self.aux_heads.w_short_d, - &self.aux_heads_grad_w_short_d, + self.aux_opt_w_size_long.step( + &mut self.aux_heads.w_size_long_d, + &self.aux_heads_grad_w_size_long_d, )?; - self.aux_opt_b_short.step( - &mut self.aux_heads.b_short_d, - &self.aux_heads_grad_b_short_d, + self.aux_opt_b_size_long.step( + &mut self.aux_heads.b_size_long_d, + &self.aux_heads_grad_b_size_long_d, + )?; + self.aux_opt_w_prof_short.step( + &mut self.aux_heads.w_prof_short_d, + &self.aux_heads_grad_w_prof_short_d, + )?; + self.aux_opt_b_prof_short.step( + &mut self.aux_heads.b_prof_short_d, + &self.aux_heads_grad_b_prof_short_d, + )?; + self.aux_opt_w_size_short.step( + &mut self.aux_heads.w_size_short_d, + &self.aux_heads_grad_w_size_short_d, + )?; + self.aux_opt_b_size_short.step( + &mut self.aux_heads.b_size_short_d, + &self.aux_heads_grad_b_size_short_d, )?; // Aux loss + valid-count shadows into mapped-pinned host buffers. // Captured DtoD on the same stream; readable post-sync via // `read_volatile` of `host_ptr` per the standard pattern - // (`feedback_no_htod_htoh_only_mapped_pinned`). + // (`feedback_no_htod_htoh_only_mapped_pinned`). Four shadows + // (one per (direction × head)). + // + // CB3+CB4 holding pattern: the `aux_loss_*_d` / `aux_valid_*_d` + // device buffers are never written (the loss kernels are not yet + // called); the DtoD copies move the zero-init contents through + // each step. The host EMA logic in `step_batched` short-circuits + // when `n_total == 0`. unsafe { - let (src_ptr, _g) = self.aux_loss_long_d.device_ptr(&self.stream); - cudarc::driver::result::memcpy_dtod_async( - self.aux_loss_long_host_d.dev_ptr, - src_ptr, - std::mem::size_of::(), - self.stream.cu_stream(), - ) - .context("aux_loss_long shadow")?; - let (src_ptr, _g) = self.aux_loss_short_d.device_ptr(&self.stream); - cudarc::driver::result::memcpy_dtod_async( - self.aux_loss_short_host_d.dev_ptr, - src_ptr, - std::mem::size_of::(), - self.stream.cu_stream(), - ) - .context("aux_loss_short shadow")?; - let (src_ptr, _g) = self.aux_valid_long_d.device_ptr(&self.stream); - cudarc::driver::result::memcpy_dtod_async( - self.aux_valid_long_host_d.dev_ptr, - src_ptr, - std::mem::size_of::(), - self.stream.cu_stream(), - ) - .context("aux_valid_long shadow")?; - let (src_ptr, _g) = self.aux_valid_short_d.device_ptr(&self.stream); - cudarc::driver::result::memcpy_dtod_async( - self.aux_valid_short_host_d.dev_ptr, - src_ptr, - std::mem::size_of::(), - self.stream.cu_stream(), - ) - .context("aux_valid_short shadow")?; + let dtod_f32 = |dst: u64, src: u64, ctx: &str| -> Result<()> { + cudarc::driver::result::memcpy_dtod_async( + dst, + src, + std::mem::size_of::(), + self.stream.cu_stream(), + ) + .with_context(|| ctx.to_string()) + }; + let dtod_i32 = |dst: u64, src: u64, ctx: &str| -> Result<()> { + cudarc::driver::result::memcpy_dtod_async( + dst, + src, + std::mem::size_of::(), + self.stream.cu_stream(), + ) + .with_context(|| ctx.to_string()) + }; + let (src, _g) = self.aux_loss_prof_long_d.device_ptr(&self.stream); + dtod_f32(self.aux_loss_prof_long_host_d.dev_ptr, src, "aux_loss_prof_long shadow")?; + let (src, _g) = self.aux_loss_size_long_d.device_ptr(&self.stream); + dtod_f32(self.aux_loss_size_long_host_d.dev_ptr, src, "aux_loss_size_long shadow")?; + let (src, _g) = self.aux_loss_prof_short_d.device_ptr(&self.stream); + dtod_f32(self.aux_loss_prof_short_host_d.dev_ptr, src, "aux_loss_prof_short shadow")?; + let (src, _g) = self.aux_loss_size_short_d.device_ptr(&self.stream); + dtod_f32(self.aux_loss_size_short_host_d.dev_ptr, src, "aux_loss_size_short shadow")?; + let (src, _g) = self.aux_valid_prof_long_d.device_ptr(&self.stream); + dtod_i32(self.aux_valid_prof_long_host_d.dev_ptr, src, "aux_valid_prof_long shadow")?; + let (src, _g) = self.aux_valid_size_long_d.device_ptr(&self.stream); + dtod_i32(self.aux_valid_size_long_host_d.dev_ptr, src, "aux_valid_size_long shadow")?; + let (src, _g) = self.aux_valid_prof_short_d.device_ptr(&self.stream); + dtod_i32(self.aux_valid_prof_short_host_d.dev_ptr, src, "aux_valid_prof_short shadow")?; + let (src, _g) = self.aux_valid_size_short_d.device_ptr(&self.stream); + dtod_i32(self.aux_valid_size_short_host_d.dev_ptr, src, "aux_valid_size_short shadow")?; } // Queue loss_d → loss_host_d (mapped-pinned). Captured inside diff --git a/crates/ml-alpha/tests/aux_heads.rs b/crates/ml-alpha/tests/aux_heads.rs index d98fc5563..d74d9b0a1 100644 --- a/crates/ml-alpha/tests/aux_heads.rs +++ b/crates/ml-alpha/tests/aux_heads.rs @@ -1,4 +1,4 @@ -//! GPU oracle tests for `aux_heads.cu` + `aux_loss.cu` (SDD-3 Layer B4). +//! GPU oracle tests for `aux_heads.cu` + `aux_loss.cu` (SDD-3 CB3+CB4). //! //! Run with: //! SQLX_OFFLINE=true cargo test -p ml-alpha --test aux_heads \ @@ -10,11 +10,19 @@ //! CPU naive reference (per `feedback_no_cpu_test_fallbacks.md` the //! production path stays GPU-only; the CPU implementation here exists //! ONLY as a unit-test oracle). +//! +//! Covers the A+B paired-head architecture: +//! * 4 outputs per direction (prof_long_logit, size_long_pred, +//! prof_short_logit, size_short_pred) +//! * Class-weighted BCE on the prof heads (with per-horizon +//! `pos_weight` scaling positive-class gradient) +//! * NaN-masked Huber on the size heads (conditional-Huber via the +//! loader's `y_size = NaN` at `y_prof = 0`) use anyhow::Result; use ml_alpha::aux_heads::{ - aux_heads_bwd_gpu, aux_heads_fwd_gpu, aux_huber_loss_gpu, AuxHeads, AuxHeadsConfig, - AuxHuberLoss, DEFAULT_HUBER_DELTA, N_AUX_HORIZONS, + aux_bce_loss_gpu, aux_heads_bwd_gpu, aux_heads_fwd_gpu, aux_huber_masked_loss_gpu, AuxBceLoss, + AuxHeads, AuxHeadsConfig, AuxMaskedHuberLoss, DEFAULT_HUBER_DELTA, N_AUX_HORIZONS, }; use ml_alpha::cfc::AUX_HIDDEN; use ml_core::device::MlDevice; @@ -58,10 +66,13 @@ fn download_i32( /// Forward shape correctness + value match against a naive CPU oracle. /// /// Validates: -/// * Output shapes ([B × N_AUX_HORIZONS] per direction). -/// * Per-element value within tolerance (5e-4) of the naive reference. -/// * Long and short heads use INDEPENDENT parameter sets — flipping a -/// single long weight does not perturb a short output and vice versa. +/// * Output shapes ([B × N_AUX_HORIZONS] per (direction, head)). +/// * Per-element value within tolerance (5e-4) of the naive reference +/// for ALL FOUR outputs (prof_long, size_long, prof_short, size_short). +/// * The four heads use INDEPENDENT parameter sets — flipping a single +/// prof_long weight does not perturb a size_long output nor any +/// short-direction output (verified implicitly via the closed-form +/// match). #[test] #[ignore = "requires CUDA"] fn fwd_matches_naive_reference() -> Result<()> { @@ -74,53 +85,67 @@ fn fwd_matches_naive_reference() -> Result<()> { .map(|i| ((i as f32) * 0.013).sin() * 0.6) .collect(); let h_aux_d = upload(&stream, &h_aux)?; - let mut y_long_d = stream.alloc_zeros::(b_sz * N_AUX_HORIZONS)?; - let mut y_short_d = stream.alloc_zeros::(b_sz * N_AUX_HORIZONS)?; + let mut prof_long_d = stream.alloc_zeros::(b_sz * N_AUX_HORIZONS)?; + let mut size_long_d = stream.alloc_zeros::(b_sz * N_AUX_HORIZONS)?; + let mut prof_short_d = stream.alloc_zeros::(b_sz * N_AUX_HORIZONS)?; + let mut size_short_d = stream.alloc_zeros::(b_sz * N_AUX_HORIZONS)?; aux_heads_fwd_gpu( &stream, &heads.fwd_fn, - &heads.w_long_d, - &heads.b_long_d, - &heads.w_short_d, - &heads.b_short_d, + &heads.w_prof_long_d, + &heads.b_prof_long_d, + &heads.w_size_long_d, + &heads.b_size_long_d, + &heads.w_prof_short_d, + &heads.b_prof_short_d, + &heads.w_size_short_d, + &heads.b_size_short_d, &h_aux_d, b_sz as i32, - &mut y_long_d, - &mut y_short_d, + &mut prof_long_d, + &mut size_long_d, + &mut prof_short_d, + &mut size_short_d, )?; stream.synchronize()?; - let y_long = download(&stream, &y_long_d)?; - let y_short = download(&stream, &y_short_d)?; + let prof_long = download(&stream, &prof_long_d)?; + let size_long = download(&stream, &size_long_d)?; + let prof_short = download(&stream, &prof_short_d)?; + let size_short = download(&stream, &size_short_d)?; let w = heads.download_weights()?; for batch in 0..b_sz { for k in 0..N_AUX_HORIZONS { - let mut expected_long = w.b_long[k]; - let mut expected_short = w.b_short[k]; + // Closed-form expected outputs for all four heads. + let mut e_pl = w.b_prof_long[k]; + let mut e_sl = w.b_size_long[k]; + let mut e_ps = w.b_prof_short[k]; + let mut e_ss = w.b_size_short[k]; for c in 0..AUX_HIDDEN { - expected_long += w.w_long[k * AUX_HIDDEN + c] * h_aux[batch * AUX_HIDDEN + c]; - expected_short += w.w_short[k * AUX_HIDDEN + c] * h_aux[batch * AUX_HIDDEN + c]; + let hc = h_aux[batch * AUX_HIDDEN + c]; + e_pl += w.w_prof_long [k * AUX_HIDDEN + c] * hc; + e_sl += w.w_size_long [k * AUX_HIDDEN + c] * hc; + e_ps += w.w_prof_short[k * AUX_HIDDEN + c] * hc; + e_ss += w.w_size_short[k * AUX_HIDDEN + c] * hc; + } + let g_pl = prof_long [batch * N_AUX_HORIZONS + k]; + let g_sl = size_long [batch * N_AUX_HORIZONS + k]; + let g_ps = prof_short[batch * N_AUX_HORIZONS + k]; + let g_ss = size_short[batch * N_AUX_HORIZONS + k]; + for (name, got, expected) in [ + ("prof_long", g_pl, e_pl), + ("size_long", g_sl, e_sl), + ("prof_short", g_ps, e_ps), + ("size_short", g_ss, e_ss), + ] { + let err = (got - expected).abs(); + assert!( + err < 5e-4, + "{name}[{batch},{k}] = {got} vs expected {expected} (err={err})", + ); } - let got_long = y_long[batch * N_AUX_HORIZONS + k]; - let got_short = y_short[batch * N_AUX_HORIZONS + k]; - assert!( - (got_long - expected_long).abs() < 5e-4, - "y_long[{},{}] = {got_long} vs expected {expected_long} \ - (diff {})", - batch, - k, - (got_long - expected_long).abs(), - ); - assert!( - (got_short - expected_short).abs() < 5e-4, - "y_short[{},{}] = {got_short} vs expected {expected_short} \ - (diff {})", - batch, - k, - (got_short - expected_short).abs(), - ); } } @@ -129,7 +154,7 @@ fn fwd_matches_naive_reference() -> Result<()> { /// Backward grad shapes finite, no NaN/Inf, and bias-grad finite-difference /// matches within tolerance against the implicit loss -/// L = Σ grad_y_*_hat[batch, k] * y_*_hat[batch, k] (linear in y_hat so +/// `L = Σ Σ grad_out[head] * out[head]` (linear in each output so all /// dL/dW and dL/db are closed-form and stable under FD). /// /// Also asserts the `grad_h_aux` per-batch slot is non-zero and finite, @@ -146,54 +171,84 @@ fn bwd_finite_diff_matches_bias_sample() -> Result<()> { .map(|i| ((i as f32) * 0.011).cos() * 0.4) .collect(); // Structured per-element grad (non-zero everywhere) so every grad - // path exercises a non-trivial accumulation. - let grad_y_long: Vec = (0..b_sz * N_AUX_HORIZONS) - .map(|i| ((i as f32) * 0.17).sin() * 0.5 + 0.1) + // path exercises a non-trivial accumulation. Distinct seeds across + // the four heads catch any cross-head wiring mistake. + let g_pl: Vec = (0..b_sz * N_AUX_HORIZONS) + .map(|i| ((i as f32) * 0.17).sin() * 0.5 + 0.10) .collect(); - let grad_y_short: Vec = (0..b_sz * N_AUX_HORIZONS) + let g_sl: Vec = (0..b_sz * N_AUX_HORIZONS) + .map(|i| ((i as f32) * 0.21).cos() * 0.4 + 0.05) + .collect(); + let g_ps: Vec = (0..b_sz * N_AUX_HORIZONS) .map(|i| ((i as f32) * 0.23).cos() * 0.4 - 0.05) .collect(); + let g_ss: Vec = (0..b_sz * N_AUX_HORIZONS) + .map(|i| ((i as f32) * 0.29).sin() * 0.3 - 0.02) + .collect(); let h_aux_d = upload(&stream, &h_aux)?; - let grad_y_long_d = upload(&stream, &grad_y_long)?; - let grad_y_short_d = upload(&stream, &grad_y_short)?; + let g_pl_d = upload(&stream, &g_pl)?; + let g_sl_d = upload(&stream, &g_sl)?; + let g_ps_d = upload(&stream, &g_ps)?; + let g_ss_d = upload(&stream, &g_ss)?; - let mut grad_w_long_d = - stream.alloc_zeros::(b_sz * N_AUX_HORIZONS * AUX_HIDDEN)?; - let mut grad_b_long_d = stream.alloc_zeros::(b_sz * N_AUX_HORIZONS)?; - let mut grad_w_short_d = - stream.alloc_zeros::(b_sz * N_AUX_HORIZONS * AUX_HIDDEN)?; - let mut grad_b_short_d = stream.alloc_zeros::(b_sz * N_AUX_HORIZONS)?; + let alloc_w = || stream.alloc_zeros::(b_sz * N_AUX_HORIZONS * AUX_HIDDEN); + let alloc_b = || stream.alloc_zeros::(b_sz * N_AUX_HORIZONS); + let mut gw_pl = alloc_w()?; + let mut gb_pl = alloc_b()?; + let mut gw_sl = alloc_w()?; + let mut gb_sl = alloc_b()?; + let mut gw_ps = alloc_w()?; + let mut gb_ps = alloc_b()?; + let mut gw_ss = alloc_w()?; + let mut gb_ss = alloc_b()?; let mut grad_h_aux_d = stream.alloc_zeros::(b_sz * AUX_HIDDEN)?; aux_heads_bwd_gpu( &stream, &heads.bwd_fn, - &heads.w_long_d, - &heads.w_short_d, + &heads.w_prof_long_d, + &heads.w_size_long_d, + &heads.w_prof_short_d, + &heads.w_size_short_d, &h_aux_d, - &grad_y_long_d, - &grad_y_short_d, + &g_pl_d, + &g_sl_d, + &g_ps_d, + &g_ss_d, b_sz as i32, - &mut grad_w_long_d, - &mut grad_b_long_d, - &mut grad_w_short_d, - &mut grad_b_short_d, + &mut gw_pl, + &mut gb_pl, + &mut gw_sl, + &mut gb_sl, + &mut gw_ps, + &mut gb_ps, + &mut gw_ss, + &mut gb_ss, &mut grad_h_aux_d, )?; stream.synchronize()?; - let gwl = download(&stream, &grad_w_long_d)?; - let gbl = download(&stream, &grad_b_long_d)?; - let gws = download(&stream, &grad_w_short_d)?; - let gbs = download(&stream, &grad_b_short_d)?; + let gwl_pl = download(&stream, &gw_pl)?; + let gbl_pl = download(&stream, &gb_pl)?; + let gwl_sl = download(&stream, &gw_sl)?; + let gbl_sl = download(&stream, &gb_sl)?; + let gwl_ps = download(&stream, &gw_ps)?; + let gbl_ps = download(&stream, &gb_ps)?; + let gwl_ss = download(&stream, &gw_ss)?; + let gbl_ss = download(&stream, &gb_ss)?; let gha = download(&stream, &grad_h_aux_d)?; + for (name, v) in [ - ("grad_w_long", &gwl), - ("grad_b_long", &gbl), - ("grad_w_short", &gws), - ("grad_b_short", &gbs), - ("grad_h_aux", &gha), + ("grad_w_prof_long", &gwl_pl), + ("grad_b_prof_long", &gbl_pl), + ("grad_w_size_long", &gwl_sl), + ("grad_b_size_long", &gbl_sl), + ("grad_w_prof_short", &gwl_ps), + ("grad_b_prof_short", &gbl_ps), + ("grad_w_size_short", &gwl_ss), + ("grad_b_size_short", &gbl_ss), + ("grad_h_aux", &gha), ] { assert!( v.iter().all(|x| x.is_finite()), @@ -205,87 +260,88 @@ fn bwd_finite_diff_matches_bias_sample() -> Result<()> { "grad_h_aux is all-zero — kernel didn't write the trunk feedback?", ); - // Closed-form expected grads from L = Σ grad_y * y_hat: - // dL/db_long[k] = Σ_batch grad_y_long[batch, k] - // dL/db_short[k] = Σ_batch grad_y_short[batch, k] - // dL/dW_long[k, c] = Σ_batch grad_y_long[batch, k] * h_aux[batch, c] - // dL/dW_short[k, c] = Σ_batch grad_y_short[batch, k] * h_aux[batch, c] - // dL/dh_aux[batch, c] = Σ_k grad_y_long[batch, k]*W_long[k,c] - // + Σ_k grad_y_short[batch, k]*W_short[k,c] - // - // The kernel writes PER-BATCH grad scratch (caller reduces axis 0) - // for the parameter grads — the assertion sums the per-batch entries - // and compares against the closed-form aggregate. + // Closed-form bias-grad: dL/db[head][k] = Σ_batch grad_out[head][batch, k] let w = heads.download_weights()?; for k in 0..N_AUX_HORIZONS { - let mut got_bl = 0.0_f32; - let mut got_bs = 0.0_f32; - let mut expected_bl = 0.0_f32; - let mut expected_bs = 0.0_f32; + // Per-head closed-form bias sum across batch. + let mut got = [0.0_f32; 4]; + let mut expected = [0.0_f32; 4]; for batch in 0..b_sz { - got_bl += gbl[batch * N_AUX_HORIZONS + k]; - got_bs += gbs[batch * N_AUX_HORIZONS + k]; - expected_bl += grad_y_long[batch * N_AUX_HORIZONS + k]; - expected_bs += grad_y_short[batch * N_AUX_HORIZONS + k]; + let idx = batch * N_AUX_HORIZONS + k; + got[0] += gbl_pl[idx]; + got[1] += gbl_sl[idx]; + got[2] += gbl_ps[idx]; + got[3] += gbl_ss[idx]; + expected[0] += g_pl[idx]; + expected[1] += g_sl[idx]; + expected[2] += g_ps[idx]; + expected[3] += g_ss[idx]; + } + for (name, gv, ev) in [ + ("grad_b_prof_long", got[0], expected[0]), + ("grad_b_size_long", got[1], expected[1]), + ("grad_b_prof_short", got[2], expected[2]), + ("grad_b_size_short", got[3], expected[3]), + ] { + let err = (gv - ev).abs(); + assert!( + err < 1e-5, + "{name}[{k}] mismatch: got={gv} expected={ev}", + ); } - let err_bl = (got_bl - expected_bl).abs(); - let err_bs = (got_bs - expected_bs).abs(); - assert!( - err_bl < 1e-5, - "grad_b_long[{k}] mismatch: got={got_bl} expected={expected_bl}", - ); - assert!( - err_bs < 1e-5, - "grad_b_short[{k}] mismatch: got={got_bs} expected={expected_bs}", - ); } - // Spot-check four weight entries against the closed form (full - // sweep is O(N_AUX_HORIZONS × AUX_HIDDEN × b_sz) — fine here, but - // sampling is sufficient to catch wiring/permutation breakage). + // Spot-check weight entries against the closed form. Per head: + // dL/dW[head][k, c] = Σ_batch grad_out[head][batch, k] * h_aux[batch, c] for &(k, c) in &[(0_usize, 0_usize), (1, 7), (2, 33), (0, 63)] { - let mut got_wl = 0.0_f32; - let mut got_ws = 0.0_f32; - let mut expected_wl = 0.0_f32; - let mut expected_ws = 0.0_f32; + let mut got = [0.0_f32; 4]; + let mut expected = [0.0_f32; 4]; for batch in 0..b_sz { let off = batch * N_AUX_HORIZONS * AUX_HIDDEN + k * AUX_HIDDEN + c; - got_wl += gwl[off]; - got_ws += gws[off]; - expected_wl += - grad_y_long[batch * N_AUX_HORIZONS + k] * h_aux[batch * AUX_HIDDEN + c]; - expected_ws += - grad_y_short[batch * N_AUX_HORIZONS + k] * h_aux[batch * AUX_HIDDEN + c]; + got[0] += gwl_pl[off]; + got[1] += gwl_sl[off]; + got[2] += gwl_ps[off]; + got[3] += gwl_ss[off]; + let g_idx = batch * N_AUX_HORIZONS + k; + let h_c = h_aux[batch * AUX_HIDDEN + c]; + expected[0] += g_pl[g_idx] * h_c; + expected[1] += g_sl[g_idx] * h_c; + expected[2] += g_ps[g_idx] * h_c; + expected[3] += g_ss[g_idx] * h_c; + } + for (name, gv, ev) in [ + ("grad_w_prof_long", got[0], expected[0]), + ("grad_w_size_long", got[1], expected[1]), + ("grad_w_prof_short", got[2], expected[2]), + ("grad_w_size_short", got[3], expected[3]), + ] { + let err = (gv - ev).abs(); + let tol = (5e-4_f32 * ev.abs().max(1e-3)).max(5e-5); + assert!( + err < tol, + "{name}[{k},{c}] mismatch: got={gv} expected={ev} (err={err})", + ); } - let err_wl = (got_wl - expected_wl).abs(); - let err_ws = (got_ws - expected_ws).abs(); - let tol = 5e-4_f32 * expected_wl.abs().max(1e-3); - let tol_s = 5e-4_f32 * expected_ws.abs().max(1e-3); - assert!( - err_wl < tol.max(5e-5), - "grad_w_long[{k},{c}] mismatch: got={got_wl} expected={expected_wl} (err={err_wl})", - ); - assert!( - err_ws < tol_s.max(5e-5), - "grad_w_short[{k},{c}] mismatch: got={got_ws} expected={expected_ws} (err={err_ws})", - ); } - // grad_h_aux is per-batch — closed-form per (batch, c). + // grad_h_aux is per-batch — closed-form per (batch, c). The + // expected value sums across all FOUR heads + N_AUX_HORIZONS. for &batch in &[0_usize, b_sz - 1] { for &c in &[0_usize, 17, AUX_HIDDEN - 1] { let mut expected = 0.0_f32; for k in 0..N_AUX_HORIZONS { - expected += grad_y_long[batch * N_AUX_HORIZONS + k] - * w.w_long[k * AUX_HIDDEN + c]; - expected += grad_y_short[batch * N_AUX_HORIZONS + k] - * w.w_short[k * AUX_HIDDEN + c]; + let g_idx = batch * N_AUX_HORIZONS + k; + expected += g_pl[g_idx] * w.w_prof_long [k * AUX_HIDDEN + c]; + expected += g_sl[g_idx] * w.w_size_long [k * AUX_HIDDEN + c]; + expected += g_ps[g_idx] * w.w_prof_short[k * AUX_HIDDEN + c]; + expected += g_ss[g_idx] * w.w_size_short[k * AUX_HIDDEN + c]; } let got = gha[batch * AUX_HIDDEN + c]; let err = (got - expected).abs(); + let tol = (5e-4_f32 * expected.abs().max(1e-3)).max(5e-4); assert!( - err < 5e-4_f32.max(5e-4 * expected.abs()), - "grad_h_aux[{batch},{c}] mismatch: got={got} expected={expected}", + err < tol, + "grad_h_aux[{batch},{c}] mismatch: got={got} expected={expected} (err={err})", ); } } @@ -293,19 +349,270 @@ fn bwd_finite_diff_matches_bias_sample() -> Result<()> { Ok(()) } -/// Huber loss forward + backward correctness on a structured target -/// tensor that exercises BOTH the quadratic branch (|r| ≤ δ) and the -/// linear branch (|r| > δ). +/// Class-weighted BCE forward + backward correctness on a structured +/// target tensor that exercises BOTH y=0 and y=1 paths and verifies the +/// `pos_weight` parameter scales the positive-class gradient correctly. +/// +/// Validates: +/// * Sigmoid+BCE math matches naive CPU oracle within 5e-4. +/// * `dL/dz` matches `pos_weight*(p-1)` for y=1 and `p` for y=0. +/// * `pos_weight=10` produces positive-class gradient scaled by ~10× +/// relative to `pos_weight=1` for the same logit/label. #[test] #[ignore = "requires CUDA"] -fn huber_loss_matches_naive_reference() -> Result<()> { +fn aux_bce_loss_matches_naive_reference() -> Result<()> { let dev = MlDevice::cuda(0)?; let stream = dev.cuda_stream()?.clone(); - let loss = AuxHuberLoss::new(&dev)?; + let bce = AuxBceLoss::new(&dev)?; + + let b_sz: usize = 8; + let n_total = b_sz * N_AUX_HORIZONS; + // Mix of positive (y=1) and negative (y=0) labels across horizons + // so the per-horizon `pos_weight` routing exercises every slot. + let y_prof: Vec = (0..n_total) + .map(|i| if (i % 3) != 0 { 1.0 } else { 0.0 }) + .collect(); + // Spread logits across the saturation regions so we hit the clamp + // (p → 0 or 1) and the linear regime. + let logits: Vec = (0..n_total) + .map(|i| (i as f32 - (n_total as f32 / 2.0)) * 0.3) + .collect(); + // Per-horizon positive-class weights — distinct values catch a + // mis-indexing bug where the kernel reads the wrong horizon's weight. + let pos_weight: [f32; N_AUX_HORIZONS] = [1.0, 5.0, 12.0]; + + let logits_d = upload(&stream, &logits)?; + let y_prof_d = upload(&stream, &y_prof)?; + let pw_d = upload(&stream, &pos_weight)?; + let mut loss_d = stream.alloc_zeros::(1)?; + let mut grad_d = stream.alloc_zeros::(n_total)?; + let mut valid_d = upload_i32(&stream, &[0_i32])?; + + aux_bce_loss_gpu( + &stream, + &bce.fn_, + &logits_d, + &y_prof_d, + &pw_d, + n_total as i32, + &mut loss_d, + &mut grad_d, + &mut valid_d, + )?; + stream.synchronize()?; + + let loss_v = download(&stream, &loss_d)?[0]; + let grad_v = download(&stream, &grad_d)?; + let valid_v = download_i32(&stream, &valid_d)?[0]; + assert_eq!(valid_v as usize, n_total, "valid count mismatch"); + + let mut expected_loss = 0.0_f32; + let mut saw_pos = false; + let mut saw_neg = false; + for i in 0..n_total { + let z = logits[i]; + let y = y_prof[i]; + let h = i % N_AUX_HORIZONS; + let pw = pos_weight[h]; + let p = 1.0_f32 / (1.0 + (-z).exp()); + let p_c = p.max(1e-7); + let one_m_p = (1.0 - p).max(1e-7); + let (l_e, g_e) = if y > 0.5 { + saw_pos = true; + (-pw * p_c.ln(), pw * (p - 1.0)) + } else { + saw_neg = true; + (-one_m_p.ln(), p) + }; + expected_loss += l_e; + let err = (grad_v[i] - g_e).abs(); + let tol = (5e-5_f32 * g_e.abs().max(1.0)).max(5e-5); + assert!( + err < tol, + "grad_prof_logit[{i}] = {} vs expected {} (z={}, y={}, pw={})", + grad_v[i], g_e, z, y, pw, + ); + } + assert!(saw_pos, "test fixture missed any y=1 sample"); + assert!(saw_neg, "test fixture missed any y=0 sample"); + let err = (loss_v - expected_loss).abs(); + assert!( + err < 5e-4, + "total BCE loss mismatch: got={loss_v} expected={expected_loss}", + ); + + Ok(()) +} + +/// `pos_weight = 10` scales the positive-class gradient by exactly 10× +/// relative to `pos_weight = 1` for the same logit/label. This is the +/// load-bearing class-balance signal — if the kernel ignores the +/// pos_weight buffer, the test catches it immediately. +#[test] +#[ignore = "requires CUDA"] +fn aux_bce_pos_weight_scales_positive_gradient() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream()?.clone(); + let bce = AuxBceLoss::new(&dev)?; + + // Single batch, all-positive labels at the same logit so the gradient + // is purely the positive-class term `pw * (p - 1)`. Three horizons + // with pos_weights [1, 10, 1] — only horizon 1's gradient should be + // 10× the horizon 0 gradient. + let b_sz: usize = 1; + let n_total = b_sz * N_AUX_HORIZONS; + let logits: Vec = vec![0.5; n_total]; // identical logit per horizon + let y_prof: Vec = vec![1.0; n_total]; // all positives + let pos_weight: [f32; N_AUX_HORIZONS] = [1.0, 10.0, 1.0]; + + let logits_d = upload(&stream, &logits)?; + let y_prof_d = upload(&stream, &y_prof)?; + let pw_d = upload(&stream, &pos_weight)?; + let mut loss_d = stream.alloc_zeros::(1)?; + let mut grad_d = stream.alloc_zeros::(n_total)?; + let mut valid_d = upload_i32(&stream, &[0_i32])?; + + aux_bce_loss_gpu( + &stream, + &bce.fn_, + &logits_d, + &y_prof_d, + &pw_d, + n_total as i32, + &mut loss_d, + &mut grad_d, + &mut valid_d, + )?; + stream.synchronize()?; + + let g = download(&stream, &grad_d)?; + // grad_h0 with pos_weight=1, grad_h1 with pos_weight=10, grad_h2 with pos_weight=1. + // Expected ratio grad_h1 / grad_h0 == 10. + let ratio = g[1] / g[0]; + assert!( + (ratio - 10.0).abs() < 1e-4, + "pos_weight=10 expected to scale gradient by 10× — got ratio={ratio} \ + (g[0]={}, g[1]={}, g[2]={})", + g[0], g[1], g[2], + ); + // Sanity: h0 and h2 share pos_weight=1 → identical gradient at + // identical (logit, label). + assert!( + (g[0] - g[2]).abs() < 1e-6, + "same pos_weight should give same gradient: g[0]={}, g[2]={}", + g[0], g[2], + ); + Ok(()) +} + +/// Conditional-Huber: `y_size = NaN` at masked slots must produce ZERO +/// gradient and contribute nothing to the loss accumulator. This is the +/// load-bearing semantic — without the mask a NaN label would poison +/// the trunk gradient and the optimiser update would explode to NaN. +#[test] +#[ignore = "requires CUDA"] +fn aux_huber_masked_does_not_propagate() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream()?.clone(); + let huber = AuxMaskedHuberLoss::new(&dev)?; + + let n: usize = 12; + let y_hat: Vec = (0..n).map(|i| (i as f32) * 0.1 - 0.5).collect(); + // Half the labels are NaN-masked (simulating `y_size = NaN @ + // y_prof = 0`); the other half are 0.0. The kernel must skip the + // masked slots and write a finite loss covering only the unmasked. + let y_size_true: Vec = (0..n) + .map(|i| if i % 2 == 0 { f32::NAN } else { 0.0 }) + .collect(); + + let y_hat_d = upload(&stream, &y_hat)?; + let y_true_d = upload(&stream, &y_size_true)?; + let mut loss_d = stream.alloc_zeros::(1)?; + let mut grad_d = stream.alloc_zeros::(n)?; + // Pre-seed with sentinels so we can verify the kernel actually wrote. + let mut valid_d = upload_i32(&stream, &[-1_i32])?; + + aux_huber_masked_loss_gpu( + &stream, + &huber.fn_, + &y_hat_d, + &y_true_d, + DEFAULT_HUBER_DELTA, + n as i32, + &mut loss_d, + &mut grad_d, + &mut valid_d, + )?; + stream.synchronize()?; + + let loss_v = download(&stream, &loss_d)?[0]; + let grad_v = download(&stream, &grad_d)?; + let valid_v = download_i32(&stream, &valid_d)?[0]; + + assert_eq!(valid_v, (n / 2) as i32, "NaN-masked valid count wrong"); + assert!(loss_v.is_finite(), "loss is NaN/Inf"); + + let mut expected_loss = 0.0_f32; + let mut hit_quadratic = false; + let mut hit_linear = false; + for i in 0..n { + if !y_size_true[i].is_finite() { + // NaN label: grad must be exactly 0, no loss contribution. + assert_eq!( + grad_v[i], 0.0, + "NaN-masked grad_y_size_hat[{i}] = {} (must be 0)", + grad_v[i], + ); + } else { + let r = y_hat[i] - y_size_true[i]; + let ar = r.abs(); + let (l, g) = if ar <= DEFAULT_HUBER_DELTA { + hit_quadratic = true; + (0.5 * r * r, r) + } else { + hit_linear = true; + ( + DEFAULT_HUBER_DELTA * (ar - 0.5 * DEFAULT_HUBER_DELTA), + DEFAULT_HUBER_DELTA * r.signum(), + ) + }; + expected_loss += l; + let err = (grad_v[i] - g).abs(); + assert!( + err < 5e-5, + "grad_y_size_hat[{i}] = {} vs expected {} (r={}, |r|={})", + grad_v[i], g, r, ar, + ); + } + } + // Confirm the test fixture exercised both branches at least once. + assert!(hit_quadratic, "test fixture missed the quadratic branch"); + // Linear branch may not fire on this small fixture; widen y_hat + // range to also test it. + let _ = hit_linear; // signal-only; explicit linear-branch test below. + + let err = (loss_v - expected_loss).abs(); + assert!( + err < 5e-4, + "NaN-masked Huber loss mismatch: got={loss_v} expected={expected_loss}", + ); + + Ok(()) +} + +/// Huber math correctness across BOTH branches (|r| ≤ δ quadratic and +/// |r| > δ linear). Mirrors the original B4 Huber test but for the new +/// `aux_huber_masked_fwd_bwd` kernel signature. +#[test] +#[ignore = "requires CUDA"] +fn aux_huber_masked_covers_both_branches() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream()?.clone(); + let huber = AuxMaskedHuberLoss::new(&dev)?; // Construct y_hat and y_true so that some residuals land in each // branch. With δ = 1.0: residuals in [-0.5, 0.5] are quadratic, - // residuals in [-2.0, -1.0] ∪ [1.0, 2.0] are linear. + // residuals in [-2.4, -1.0] ∪ [1.0, 2.2] are linear. let n: usize = 24; let y_hat: Vec = (0..n) .map(|i| (i as f32 - 12.0) * 0.2) // -2.4 .. 2.2 step 0.2 @@ -319,9 +626,9 @@ fn huber_loss_matches_naive_reference() -> Result<()> { let mut grad_d = stream.alloc_zeros::(n)?; let mut valid_d = upload_i32(&stream, &[0_i32])?; - aux_huber_loss_gpu( + aux_huber_masked_loss_gpu( &stream, - &loss.fn_, + &huber.fn_, &y_hat_d, &y_true_d, delta, @@ -337,7 +644,6 @@ fn huber_loss_matches_naive_reference() -> Result<()> { let valid_v = download_i32(&stream, &valid_d)?[0]; assert_eq!(valid_v as usize, n, "valid count mismatch"); - // CPU oracle — sum of per-element Huber. let mut expected_loss = 0.0_f32; let mut hit_quadratic = false; let mut hit_linear = false; @@ -355,11 +661,8 @@ fn huber_loss_matches_naive_reference() -> Result<()> { let err = (grad_v[i] - g).abs(); assert!( err < 5e-5, - "grad_y_hat[{i}] = {} vs expected {} (r={}, |r|={})", - grad_v[i], - g, - r, - ar, + "grad_y_size_hat[{i}] = {} vs expected {} (r={}, |r|={})", + grad_v[i], g, r, ar, ); } assert!(hit_quadratic, "test fixture missed the quadratic branch"); @@ -367,47 +670,52 @@ fn huber_loss_matches_naive_reference() -> Result<()> { let err = (loss_v - expected_loss).abs(); assert!( err < 5e-4, - "total loss mismatch: got={loss_v} expected={expected_loss}", + "Huber total loss mismatch: got={loss_v} expected={expected_loss}", ); Ok(()) } -/// NaN-masked labels must produce zero gradient at the masked slot and -/// must NOT contribute to the loss accumulator or the valid count. -/// Without masking, a single NaN label would poison the whole batch: -/// * grad_y_hat[i] = NaN → propagates through aux_heads_bwd into -/// grad_h_aux, then into aux_trunk_bwd's grad_h_new → kills the -/// trunk for that batch sample. -/// * Σ Huber would become NaN → optimiser update goes NaN globally. +/// NaN-masked BCE: `y_prof = NaN` must produce ZERO gradient and skip +/// the loss accumulator — same defence as the Huber mask path. #[test] #[ignore = "requires CUDA"] -fn huber_loss_nan_mask_does_not_propagate() -> Result<()> { +fn aux_bce_nan_mask_does_not_propagate() -> Result<()> { let dev = MlDevice::cuda(0)?; let stream = dev.cuda_stream()?.clone(); - let loss = AuxHuberLoss::new(&dev)?; + let bce = AuxBceLoss::new(&dev)?; - let n: usize = 12; - let y_hat: Vec = (0..n).map(|i| (i as f32) * 0.1 - 0.5).collect(); - // Half the labels are NaN-masked; the other half are 0.0. - let y_true: Vec = (0..n) - .map(|i| if i % 2 == 0 { f32::NAN } else { 0.0 }) + let b_sz: usize = 4; + let n_total = b_sz * N_AUX_HORIZONS; + let logits: Vec = (0..n_total).map(|i| (i as f32) * 0.05 - 0.2).collect(); + // Mark even-indexed slots as NaN, odd as 0/1 alternating. + let y_prof: Vec = (0..n_total) + .map(|i| { + if i % 2 == 0 { + f32::NAN + } else if (i / 2) % 2 == 0 { + 0.0 + } else { + 1.0 + } + }) .collect(); + let pos_weight: [f32; N_AUX_HORIZONS] = [3.0, 3.0, 3.0]; - let y_hat_d = upload(&stream, &y_hat)?; - let y_true_d = upload(&stream, &y_true)?; + let logits_d = upload(&stream, &logits)?; + let y_prof_d = upload(&stream, &y_prof)?; + let pw_d = upload(&stream, &pos_weight)?; let mut loss_d = stream.alloc_zeros::(1)?; - let mut grad_d = stream.alloc_zeros::(n)?; - // Pre-seed with sentinels so we can verify the kernel actually wrote. + let mut grad_d = stream.alloc_zeros::(n_total)?; let mut valid_d = upload_i32(&stream, &[-1_i32])?; - aux_huber_loss_gpu( + aux_bce_loss_gpu( &stream, - &loss.fn_, - &y_hat_d, - &y_true_d, - DEFAULT_HUBER_DELTA, - n as i32, + &bce.fn_, + &logits_d, + &y_prof_d, + &pw_d, + n_total as i32, &mut loss_d, &mut grad_d, &mut valid_d, @@ -418,43 +726,19 @@ fn huber_loss_nan_mask_does_not_propagate() -> Result<()> { let grad_v = download(&stream, &grad_d)?; let valid_v = download_i32(&stream, &valid_d)?[0]; - // Valid count is the non-NaN half. - assert_eq!(valid_v, (n / 2) as i32, "NaN-masked valid count wrong"); - // Loss is finite and equals the closed-form sum over the non-NaN - // entries — proves the NaN contributions were skipped. - assert!(loss_v.is_finite(), "loss is NaN/Inf"); - let mut expected_loss = 0.0_f32; - for i in 0..n { - if !y_true[i].is_finite() { - // NaN label: grad must be exactly 0, no loss contribution. + let expected_valid = y_prof.iter().filter(|y| y.is_finite()).count() as i32; + assert_eq!(valid_v, expected_valid, "BCE NaN-masked valid count wrong"); + assert!(loss_v.is_finite(), "BCE loss is NaN/Inf"); + for i in 0..n_total { + if !y_prof[i].is_finite() { assert_eq!( grad_v[i], 0.0, - "NaN-masked grad_y_hat[{i}] = {} (must be 0)", + "NaN-masked grad_prof_logit[{i}] = {} (must be 0)", grad_v[i], ); } else { - // Standard Huber contribution. - let r = y_hat[i] - y_true[i]; - let ar = r.abs(); - expected_loss += if ar <= DEFAULT_HUBER_DELTA { - 0.5 * r * r - } else { - DEFAULT_HUBER_DELTA * (ar - 0.5 * DEFAULT_HUBER_DELTA) - }; - // Live grad must be finite (already true if loss is finite, - // but the per-slot check pins the wiring down). - assert!( - grad_v[i].is_finite(), - "non-masked grad_y_hat[{i}] = {} (NaN/Inf)", - grad_v[i], - ); + assert!(grad_v[i].is_finite(), "non-masked grad NaN at i={i}"); } } - let err = (loss_v - expected_loss).abs(); - assert!( - err < 5e-4, - "NaN-masked loss mismatch: got={loss_v} expected={expected_loss}", - ); - Ok(()) }