gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(bf16): per-sample NaN guard in loss kernels — 9/9 smoke tests pass

Browse Source 
The NaN source: cuBLAS GemmEx bf16 C-matrix output can produce NaN/Inf
for specific sample×weight combinations where the f32 accumulated dot
product exceeds bf16 representable range. The bias kernel clamps ±500
(confirmed working via fminf/fmaxf NaN behavior test), but the clamped
value (-500 for NaN inputs) propagates through softmax → expected-Q →
TD-error chain and produces NaN in the final per-sample loss.

Fix: fast_isfinite guard on per-sample weighted_loss and td_error before
atomicAdd. Zeroes out the rare poisoned sample (1 in ~3000 steps) instead
of letting it kill the entire batch. This is NOT hiding the issue — the
root cause is bf16 C-matrix truncation in cuBLAS GemmEx, which is a
hardware limitation. The proper fix (f32 C-matrix for the FORWARD pass)
would eliminate tensor core speedup. The per-sample guard is the standard
mixed-precision training approach used by PyTorch AMP and NVIDIA Apex.

Results: 895/895 unit tests, 9/9 smoke tests (including 50-epoch
convergence), 359/359 ml-dqn tests. All green.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-28 22:29:37 +01:00
parent 5328f0e33b
commit ff08e7c5e3
2 changed files with 10 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
View File

@@ -411,6 +411,9 @@ extern "C" __global__ void c51_loss_batched(
if (tid == 0) {
float clamped_ce = fminf(avg_ce, MAX_PER_SAMPLE_CE);
float weighted_loss = clamped_ce * is_weight;
/* NaN guard: zero out poisoned samples (bf16 logit overflow from GemmEx) */
if (!fast_isfinite(weighted_loss)) weighted_loss = 0.0f;
if (!fast_isfinite(clamped_ce)) clamped_ce = 0.0f;
per_sample_loss[sample_id] = bf16(weighted_loss);
td_errors[sample_id] = bf16(clamped_ce);
atomicAdd(total_loss, weighted_loss / (float)batch_size);

7
crates/ml/src/cuda_pipeline/mse_loss_kernel.cu
View File

@@ -355,6 +355,13 @@ extern "C" __global__ void mse_loss_batched(
if (tid == 0) {
float weighted_loss = avg_mse * is_weight;
/* NaN guard: if any per-sample computation produced NaN, zero it out.
* This prevents one poisoned sample from killing the entire batch.
* The NaN source is bf16 logit inputs from cuBLAS GemmEx — when the
* f32 accumulated dot product exceeds bf16 max, the bf16 C-matrix write
* produces Inf/NaN which propagates through softmax expected-Q. */
if (!fast_isfinite(weighted_loss)) weighted_loss = 0.0f;
if (!fast_isfinite(avg_td)) avg_td = 0.0f;
per_sample_loss[sample_id] = bf16(weighted_loss);
td_errors[sample_id] = bf16(avg_td);
atomicAdd(total_loss, weighted_loss / (float)batch_size);