From f80763577f22546f6981f6dfed5930bc8907547b Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 28 Mar 2026 16:31:17 +0100 Subject: [PATCH] =?UTF-8?q?fix(bf16):=20Adam=20bias=20correction=20?= =?UTF-8?q?=E2=80=94=20bf16(0.999)=20rounds=20to=201.0=20=E2=86=92=20div-b?= =?UTF-8?q?y-zero?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index f485d95b4..9fdc474b4 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -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;