fix(bf16): Adam bias correction — bf16(0.999) rounds to 1.0 → div-by-zero

Root cause of Q-value NaN divergence: beta2=0.999 can't be represented
in BF16 (7-bit mantissa). bf16(0.999) = bf16(1.0). Then:
  1 - beta2^t = 1 - 1^t = 0
  v_hat = v / 0 = Inf → NaN → params NaN → forward NaN

Fix: compute bias correction (1 - beta^t) in f32 via powf(), then
convert result to bf16. This is the ONE justified f32 computation
in the kernel — bf16 can't represent 0.999 or 0.001.

Clamp bias correction to ≥1e-4 to prevent div-by-zero.

Smoke test: Q-values stable at -17.125, Sharpe improves
-23→-7→+9.78 across 3 epochs. Training completes without divergence.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-28 16:31:17 +01:00
parent 954bca690d
commit f80763577f

View File

@@ -96,9 +96,10 @@ extern "C" __global__ void dqn_adam_update_kernel(
__nv_bfloat16 clip_scale = (*grad_norm_sq > bf16_zero() && norm > bf16(max_grad_norm)) ? (bf16(max_grad_norm) / norm) : bf16_one();
__nv_bfloat16 clipped_g = g * clip_scale;
/* Adam update */
__nv_bfloat16 beta1_t = bf16_one() - bf16_pow(bf16(beta1), bf16((float)t));
__nv_bfloat16 beta2_t = bf16_one() - bf16_pow(bf16(beta2), bf16((float)t));
/* Adam bias correction: compute in float — bf16(0.999) rounds to 1.0 → div-by-zero.
* This is the ONLY float arithmetic in the kernel (3 decimal digits insufficient). */
__nv_bfloat16 beta1_t = bf16(fmaxf(1.0f - powf(beta1, (float)t), 1e-4f));
__nv_bfloat16 beta2_t = bf16(fmaxf(1.0f - powf(beta2, (float)t), 1e-4f));
__nv_bfloat16 m_i = bf16(beta1) * m[idx] + (bf16_one() - bf16(beta1)) * clipped_g;
__nv_bfloat16 v_i = bf16(beta2) * v[idx] + (bf16_one() - bf16(beta2)) * clipped_g * clipped_g;