d472b2ac3a
bce_loss_multi_horizon: fused forward+backward, block tree-reduce (no
atomicAdd), NaN labels masked (drop). Loss = mean over valid; grad =
(p-y)/(p(1-p)) scaled by 1/N_valid.
adamw_step: element-wise AdamW with weight decay; one thread per param.
Tests pass on sm_86:
BCE (4/4): positive+finite loss, near-zero loss when probs match
labels, analytic grad matches GPU-computed finite-difference at
eps=1e-3 / max_relative=5e-2 across 5 perturbation points, NaN
labels mask grad and contribute zero to loss/N_valid.
AdamW (4/4): zero-grad moves param only by weight-decay, positive
grad decreases param, step counter increments, repeated descent
on grad=theta drives ‖θ‖ from 40 to <1 in 200 steps.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
86 lines
3.3 KiB
Rust
86 lines
3.3 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 }
|
|
}
|
|
|
|
#[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]);
|
|
}
|
|
}
|