feat(aux-heads): rewrite for A+B paired (BCE + conditional-Huber) (CB3+CB4)
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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");
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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];
|
||||
|
||||
@@ -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<f32>, // [N_AUX_HORIZONS × AUX_HIDDEN]
|
||||
pub b_long: Vec<f32>, // [N_AUX_HORIZONS]
|
||||
pub w_short: Vec<f32>, // [N_AUX_HORIZONS × AUX_HIDDEN]
|
||||
pub b_short: Vec<f32>, // [N_AUX_HORIZONS]
|
||||
pub w_prof_long: Vec<f32>, // [N_AUX_HORIZONS × AUX_HIDDEN]
|
||||
pub b_prof_long: Vec<f32>, // [N_AUX_HORIZONS]
|
||||
pub w_size_long: Vec<f32>, // [N_AUX_HORIZONS × AUX_HIDDEN]
|
||||
pub b_size_long: Vec<f32>, // [N_AUX_HORIZONS]
|
||||
pub w_prof_short: Vec<f32>, // [N_AUX_HORIZONS × AUX_HIDDEN]
|
||||
pub b_prof_short: Vec<f32>, // [N_AUX_HORIZONS]
|
||||
pub w_size_short: Vec<f32>, // [N_AUX_HORIZONS × AUX_HIDDEN]
|
||||
pub b_size_short: Vec<f32>, // [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<CudaStream>,
|
||||
@@ -108,15 +130,29 @@ pub struct AuxHeads {
|
||||
pub fwd_fn: CudaFunction,
|
||||
pub bwd_fn: CudaFunction,
|
||||
|
||||
pub w_long_d: CudaSlice<f32>,
|
||||
pub b_long_d: CudaSlice<f32>,
|
||||
pub w_short_d: CudaSlice<f32>,
|
||||
pub b_short_d: CudaSlice<f32>,
|
||||
pub w_prof_long_d: CudaSlice<f32>,
|
||||
pub b_prof_long_d: CudaSlice<f32>,
|
||||
pub w_size_long_d: CudaSlice<f32>,
|
||||
pub b_size_long_d: CudaSlice<f32>,
|
||||
pub w_prof_short_d: CudaSlice<f32>,
|
||||
pub b_prof_short_d: CudaSlice<f32>,
|
||||
pub w_size_short_d: CudaSlice<f32>,
|
||||
pub b_size_short_d: CudaSlice<f32>,
|
||||
}
|
||||
|
||||
/// 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<CudaStream>,
|
||||
_module: Arc<CudaModule>,
|
||||
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<CudaStream>,
|
||||
_module: Arc<CudaModule>,
|
||||
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<f32> = (0..N_AUX_HORIZONS * AUX_HIDDEN)
|
||||
.map(|_| r.gen_range(-scale..scale))
|
||||
.collect();
|
||||
let w_short: Vec<f32> = (0..N_AUX_HORIZONS * AUX_HIDDEN)
|
||||
.map(|_| r.gen_range(-scale..scale))
|
||||
.collect();
|
||||
let b_long: Vec<f32> = vec![0.0; N_AUX_HORIZONS];
|
||||
let b_short: Vec<f32> = 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<f32> {
|
||||
(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<f32> = 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<AuxHeadsWeights> {
|
||||
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<f32>, name: &str| -> Result<Vec<f32>> {
|
||||
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<Self> {
|
||||
let stream: Arc<CudaStream> = dev.cuda_stream().context("aux_loss stream")?.clone();
|
||||
let ctx = dev.cuda_context().context("aux_loss ctx")?;
|
||||
let stream: Arc<CudaStream> = 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<CudaStream> {
|
||||
&self.stream
|
||||
}
|
||||
}
|
||||
|
||||
impl AuxMaskedHuberLoss {
|
||||
pub fn new(dev: &MlDevice) -> Result<Self> {
|
||||
let stream: Arc<CudaStream> = 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<CudaStream>,
|
||||
func: &CudaFunction,
|
||||
w_long_d: &CudaSlice<f32>,
|
||||
b_long_d: &CudaSlice<f32>,
|
||||
w_short_d: &CudaSlice<f32>,
|
||||
b_short_d: &CudaSlice<f32>,
|
||||
w_prof_long_d: &CudaSlice<f32>,
|
||||
b_prof_long_d: &CudaSlice<f32>,
|
||||
w_size_long_d: &CudaSlice<f32>,
|
||||
b_size_long_d: &CudaSlice<f32>,
|
||||
w_prof_short_d: &CudaSlice<f32>,
|
||||
b_prof_short_d: &CudaSlice<f32>,
|
||||
w_size_short_d: &CudaSlice<f32>,
|
||||
b_size_short_d: &CudaSlice<f32>,
|
||||
h_aux_d: &CudaSlice<f32>,
|
||||
b_sz: i32,
|
||||
y_long_hat_d: &mut CudaSlice<f32>,
|
||||
y_short_hat_d: &mut CudaSlice<f32>,
|
||||
prof_long_logit_d: &mut CudaSlice<f32>,
|
||||
size_long_pred_d: &mut CudaSlice<f32>,
|
||||
prof_short_logit_d: &mut CudaSlice<f32>,
|
||||
size_short_pred_d: &mut CudaSlice<f32>,
|
||||
) -> 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<CudaStream>,
|
||||
func: &CudaFunction,
|
||||
w_long_d: &CudaSlice<f32>,
|
||||
w_short_d: &CudaSlice<f32>,
|
||||
w_prof_long_d: &CudaSlice<f32>,
|
||||
w_size_long_d: &CudaSlice<f32>,
|
||||
w_prof_short_d: &CudaSlice<f32>,
|
||||
w_size_short_d: &CudaSlice<f32>,
|
||||
h_aux_d: &CudaSlice<f32>,
|
||||
grad_y_long_hat_d: &CudaSlice<f32>,
|
||||
grad_y_short_hat_d: &CudaSlice<f32>,
|
||||
grad_prof_long_logit_d: &CudaSlice<f32>,
|
||||
grad_size_long_pred_d: &CudaSlice<f32>,
|
||||
grad_prof_short_logit_d: &CudaSlice<f32>,
|
||||
grad_size_short_pred_d: &CudaSlice<f32>,
|
||||
b_sz: i32,
|
||||
grad_w_long_d: &mut CudaSlice<f32>,
|
||||
grad_b_long_d: &mut CudaSlice<f32>,
|
||||
grad_w_short_d: &mut CudaSlice<f32>,
|
||||
grad_b_short_d: &mut CudaSlice<f32>,
|
||||
grad_w_prof_long_d: &mut CudaSlice<f32>,
|
||||
grad_b_prof_long_d: &mut CudaSlice<f32>,
|
||||
grad_w_size_long_d: &mut CudaSlice<f32>,
|
||||
grad_b_size_long_d: &mut CudaSlice<f32>,
|
||||
grad_w_prof_short_d: &mut CudaSlice<f32>,
|
||||
grad_b_prof_short_d: &mut CudaSlice<f32>,
|
||||
grad_w_size_short_d: &mut CudaSlice<f32>,
|
||||
grad_b_size_short_d: &mut CudaSlice<f32>,
|
||||
grad_h_aux_d: &mut CudaSlice<f32>,
|
||||
) -> 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<CudaStream>,
|
||||
func: &CudaFunction,
|
||||
y_hat_d: &CudaSlice<f32>,
|
||||
y_true_d: &CudaSlice<f32>,
|
||||
delta: f32,
|
||||
prof_logit_d: &CudaSlice<f32>,
|
||||
y_prof_true_d: &CudaSlice<f32>,
|
||||
pos_weight_d: &CudaSlice<f32>,
|
||||
n_total: i32,
|
||||
loss_out_d: &mut CudaSlice<f32>,
|
||||
grad_y_hat_d: &mut CudaSlice<f32>,
|
||||
grad_prof_logit_d: &mut CudaSlice<f32>,
|
||||
valid_count_out_d: &mut CudaSlice<i32>,
|
||||
) -> 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<CudaStream>,
|
||||
func: &CudaFunction,
|
||||
y_size_hat_d: &CudaSlice<f32>,
|
||||
y_size_true_d: &CudaSlice<f32>,
|
||||
delta: f32,
|
||||
n_total: i32,
|
||||
loss_out_d: &mut CudaSlice<f32>,
|
||||
grad_y_size_hat_d: &mut CudaSlice<f32>,
|
||||
valid_count_out_d: &mut CudaSlice<i32>,
|
||||
) -> 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(())
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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::<f32>(b_sz * N_AUX_HORIZONS)?;
|
||||
let mut y_short_d = stream.alloc_zeros::<f32>(b_sz * N_AUX_HORIZONS)?;
|
||||
let mut prof_long_d = stream.alloc_zeros::<f32>(b_sz * N_AUX_HORIZONS)?;
|
||||
let mut size_long_d = stream.alloc_zeros::<f32>(b_sz * N_AUX_HORIZONS)?;
|
||||
let mut prof_short_d = stream.alloc_zeros::<f32>(b_sz * N_AUX_HORIZONS)?;
|
||||
let mut size_short_d = stream.alloc_zeros::<f32>(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<f32> = (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<f32> = (0..b_sz * N_AUX_HORIZONS)
|
||||
.map(|i| ((i as f32) * 0.17).sin() * 0.5 + 0.10)
|
||||
.collect();
|
||||
let grad_y_short: Vec<f32> = (0..b_sz * N_AUX_HORIZONS)
|
||||
let g_sl: Vec<f32> = (0..b_sz * N_AUX_HORIZONS)
|
||||
.map(|i| ((i as f32) * 0.21).cos() * 0.4 + 0.05)
|
||||
.collect();
|
||||
let g_ps: Vec<f32> = (0..b_sz * N_AUX_HORIZONS)
|
||||
.map(|i| ((i as f32) * 0.23).cos() * 0.4 - 0.05)
|
||||
.collect();
|
||||
let g_ss: Vec<f32> = (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::<f32>(b_sz * N_AUX_HORIZONS * AUX_HIDDEN)?;
|
||||
let mut grad_b_long_d = stream.alloc_zeros::<f32>(b_sz * N_AUX_HORIZONS)?;
|
||||
let mut grad_w_short_d =
|
||||
stream.alloc_zeros::<f32>(b_sz * N_AUX_HORIZONS * AUX_HIDDEN)?;
|
||||
let mut grad_b_short_d = stream.alloc_zeros::<f32>(b_sz * N_AUX_HORIZONS)?;
|
||||
let alloc_w = || stream.alloc_zeros::<f32>(b_sz * N_AUX_HORIZONS * AUX_HIDDEN);
|
||||
let alloc_b = || stream.alloc_zeros::<f32>(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::<f32>(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<f32> = (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<f32> = (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::<f32>(1)?;
|
||||
let mut grad_d = stream.alloc_zeros::<f32>(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<f32> = vec![0.5; n_total]; // identical logit per horizon
|
||||
let y_prof: Vec<f32> = 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::<f32>(1)?;
|
||||
let mut grad_d = stream.alloc_zeros::<f32>(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<f32> = (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<f32> = (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::<f32>(1)?;
|
||||
let mut grad_d = stream.alloc_zeros::<f32>(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<f32> = (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::<f32>(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<f32> = (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<f32> = (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<f32> = (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<f32> = (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::<f32>(1)?;
|
||||
let mut grad_d = stream.alloc_zeros::<f32>(n)?;
|
||||
// Pre-seed with sentinels so we can verify the kernel actually wrote.
|
||||
let mut grad_d = stream.alloc_zeros::<f32>(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(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user