6a3f45d872
Single source of truth for the multi-horizon BCE per the new
feedback_single_source_of_truth_no_duplicates pearl. Eliminates the
`bce_loss_multi_horizon_sigma.cu` / `loss_sigma.rs` duplication
introduced earlier today and folds the Kendall σ-weighting into the
canonical kernel.
Deletions:
cuda/bce_loss_multi_horizon_sigma.cu (folded into legacy)
src/trainer/loss_sigma.rs (helper merged)
tests/bce_sigma_numgrad.rs (subsumed)
Rewrites:
cuda/bce_loss_multi_horizon.cu — replaced with σ-aware
implementation; kernel function name kept as
`bce_multi_horizon_forward_backward` so cubin symbol stays stable.
NVIDIA-grade warp-shuffle reduce (4 warps, 1 cross-warp barrier);
new args `log_sigma_h` (per-horizon Kendall σ logarithm) and
`d_log_sigma_h` (its gradient).
src/trainer/loss.rs — standalone helper updated to new 11-arg
kernel signature. Exposes optional `log_sigma_h` in `BceInput`
(None → zeros / passthrough Kendall init) and returns
`mean_bce_per_h` + `d_log_sigma_h` in `BceOutput`.
tests/bce_grad_finite_diff.rs — adds `log_sigma_h: None` to the
test inputs; all 4 numgrad tests PASS unchanged.
build.rs — drops `bce_loss_multi_horizon_sigma` entry from
KERNELS. The single canonical `bce_loss_multi_horizon` cubin
now contains the σ-aware kernel.
Wiring:
PerceptionTrainer gains a single `pub mha: MultiHorizonAttention`
field. Owns `log_sigma_h_d` + `grad_log_sigma_h_d` (and the rest
of the multi-horizon attention path, anchored on Stage 2 to fully
replace the legacy `attention_pool` path).
step_batched + evaluate_batched BCE callsites now thread
`mha.log_sigma_h_d` and `mha.grad_log_sigma_h_d` through the
11-arg kernel signature.
This commit keeps the legacy `attention_pool` callsite in place; the
Stage 2 commit will replace it with `mha.pool` + `mha.inv_pool` +
`mha.moe` and delete the `attn_*` fields entirely.
Verified locally: ml-alpha builds clean with the cuda feature,
bce_grad_finite_diff (4/4) PASS on RTX 3050.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
86 lines
3.4 KiB
Rust
86 lines
3.4 KiB
Rust
//! BCE forward + backward validation.
|
|
//!
|
|
//! Per `feedback_no_cpu_test_fallbacks.md`: no CPU oracle. The gradient
|
|
//! is validated against finite-difference computed via GPU forward calls
|
|
//! at perturbed inputs (the kernel itself is the reference; we only
|
|
//! verify that the analytic grad matches its own numerical derivative).
|
|
|
|
use approx::assert_relative_eq;
|
|
use ml_alpha::trainer::loss::{bce_multi_horizon_loss_and_grad_gpu, BceInput};
|
|
use ml_core::device::MlDevice;
|
|
|
|
fn test_device() -> MlDevice {
|
|
MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests")
|
|
}
|
|
|
|
fn make_input(probs: Vec<f32>, labels: Vec<f32>) -> BceInput {
|
|
BceInput { probs, labels, n_horizons: 5, n_pos: 4, log_sigma_h: None }
|
|
}
|
|
|
|
#[test]
|
|
fn bce_loss_is_positive_and_finite() {
|
|
let dev = test_device();
|
|
let probs: Vec<f32> = (0..20).map(|i| (0.1 + i as f32 * 0.04).min(0.95)).collect();
|
|
let labels: Vec<f32> = (0..20).map(|i| (i % 2) as f32).collect();
|
|
let out = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(probs, labels)).unwrap();
|
|
assert!(out.loss.is_finite());
|
|
assert!(out.loss > 0.0);
|
|
assert_eq!(out.n_valid, 20);
|
|
}
|
|
|
|
#[test]
|
|
fn bce_grad_zero_when_probs_match_labels() {
|
|
let dev = test_device();
|
|
// probs ≈ labels (within clamp window): grad ≈ 0 wherever |p-y| ≈ 0
|
|
let mut probs = vec![0.5_f32; 20];
|
|
let labels: Vec<f32> = (0..20).map(|i| (i % 2) as f32).collect();
|
|
// Move probs very close to labels (but within (1e-6, 1-1e-6) clamp window).
|
|
for i in 0..20 {
|
|
probs[i] = if labels[i] > 0.5 { 1.0 - 1e-5 } else { 1e-5 };
|
|
}
|
|
let out = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(probs, labels)).unwrap();
|
|
// Near-perfect predictions → near-zero loss
|
|
assert!(out.loss < 1e-3, "loss should be small for matching probs/labels, got {}", out.loss);
|
|
}
|
|
|
|
#[test]
|
|
fn bce_grad_matches_finite_difference() {
|
|
let dev = test_device();
|
|
let probs: Vec<f32> = (0..20).map(|i| (0.1 + i as f32 * 0.04).min(0.95)).collect();
|
|
let labels: Vec<f32> = (0..20).map(|i| (i % 2) as f32).collect();
|
|
let base = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(probs.clone(), labels.clone())).unwrap();
|
|
|
|
let eps = 1e-3_f32;
|
|
for i in [0usize, 5, 10, 15, 19] {
|
|
let mut up = probs.clone();
|
|
up[i] += eps;
|
|
let mut dn = probs.clone();
|
|
dn[i] -= eps;
|
|
let lu = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(up, labels.clone())).unwrap().loss;
|
|
let ld = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(dn, labels.clone())).unwrap().loss;
|
|
let fd = (lu - ld) / (2.0 * eps);
|
|
assert_relative_eq!(
|
|
base.grad_probs[i],
|
|
fd,
|
|
epsilon = 5e-3,
|
|
max_relative = 5e-2
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn bce_nan_labels_mask_grad_and_loss() {
|
|
let dev = test_device();
|
|
let mut labels: Vec<f32> = (0..20).map(|i| (i % 2) as f32).collect();
|
|
// Mask out indices 5, 10, 15
|
|
labels[5] = f32::NAN;
|
|
labels[10] = f32::NAN;
|
|
labels[15] = f32::NAN;
|
|
let probs: Vec<f32> = (0..20).map(|i| (0.1 + i as f32 * 0.04).min(0.95)).collect();
|
|
let out = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(probs, labels)).unwrap();
|
|
assert_eq!(out.n_valid, 17, "should drop 3 NaN-masked entries");
|
|
for &i in &[5usize, 10, 15] {
|
|
assert_eq!(out.grad_probs[i], 0.0, "masked grad should be 0 at index {i}, got {}", out.grad_probs[i]);
|
|
}
|
|
}
|