From 5232a1ae31489489509e06e9b488f532d7a82e4b Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 29 Mar 2026 00:42:06 +0100 Subject: [PATCH] =?UTF-8?q?fix(bf16):=20ROOT=20CAUSE=20=E2=80=94=20float?= =?UTF-8?q?=20experience=20features=20+=20IS-weight=20overflow=20clamp?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two root causes of intermittent training NaN (1/3000 steps) identified and fixed: 1. BF16 portfolio/market feature overflow in experience_kernels.cu: - 6 portfolio features (lines 220-226) computed with bf16 divisions that overflow when equity/position values are large (ES at ~5000) - 16 multi-timeframe market features computed with bf16 subtraction of similar close prices → precision loss and overflow - Fix: ALL portfolio + market feature computation now in float (read bf16 inputs → float arithmetic → write bf16 output) - NaN states in replay buffer → NaN GemmEx output → NaN loss (eliminated) 2. PER IS-weight Inf→NaN cascade in replay_buffer_kernels.cu: - powf(tiny_prob, -beta) produces Inf when priorities are very skewed - normalize_weights_f32 divides all weights by max_weight - Inf / Inf = NaN (IEEE 754) → ENTIRE batch has NaN IS-weights - Fix: clamp IS-weight to 1e6 before normalization (well within f32, normalized to ≤1.0 by max division) - prob floor at 1e-12 and total_sum floor at 1e-8 prevent division by zero NaN guards REMOVED from loss kernels (no longer needed): - mse_loss_kernel.cu: removed fast_isfinite guard on weighted_loss - c51_loss_kernel.cu: removed fast_isfinite guard on weighted_loss/clamped_ce 895/895 unit + 9/9 smoke tests pass. Zero NaN guards in the training path. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml-dqn/src/replay_buffer_kernels.cu | 9 +- .../ml/src/cuda_pipeline/c51_loss_kernel.cu | 2 - .../src/cuda_pipeline/experience_kernels.cu | 112 ++++++++++-------- .../ml/src/cuda_pipeline/mse_loss_kernel.cu | 4 - 4 files changed, 70 insertions(+), 57 deletions(-) diff --git a/crates/ml-dqn/src/replay_buffer_kernels.cu b/crates/ml-dqn/src/replay_buffer_kernels.cu index a5c4bd465..5c5625606 100644 --- a/crates/ml-dqn/src/replay_buffer_kernels.cu +++ b/crates/ml-dqn/src/replay_buffer_kernels.cu @@ -159,8 +159,13 @@ void is_weights_f32( int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= batch_size) return; float ts = total_sum_buf[0]; - float prob = (sampled_prios[i] * (float)n_buffer) / ts; - weights[i] = powf(prob, neg_beta); + float prob = fmaxf((sampled_prios[i] * (float)n_buffer) / fmaxf(ts, 1e-8f), 1e-12f); + /* Clamp weight to prevent Inf → NaN cascade in normalize_weights_f32. + * powf(tiny_prob, -beta) can exceed f32 range when priorities are very + * skewed. Inf weight → max_weight = Inf → Inf/Inf = NaN for ALL samples. + * Cap at 1e6 (well within f32, normalized to ≤1.0 by max division). */ + float w = powf(prob, neg_beta); + weights[i] = fminf(w, 1e6f); } // Normalize weights by max (two-pass: first find max, then divide) diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index f5c3199ff..ff6579780 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -410,8 +410,6 @@ 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; - 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); diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 7e9ec8b0f..2e5ebca62 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -191,39 +191,53 @@ extern "C" __global__ void experience_state_gather( __nv_bfloat16 entry_price = ps[12]; __nv_bfloat16 trade_start_pnl = ps[13]; - __nv_bfloat16 one = bf16_one(); - __nv_bfloat16 equity = (portfolio_value > one) ? portfolio_value : one; - __nv_bfloat16 drawdown = (peak_equity > one) - ? (peak_equity - portfolio_value) / peak_equity - : bf16_zero(); - drawdown = (drawdown > bf16_zero()) ? drawdown : bf16_zero(); + /* ── Portfolio features computed in FLOAT to prevent bf16 overflow ── + * BF16 max ~65504. ES position × price_diff can exceed this (2.0 × 5000 = 10000, + * but accumulated P&L or subtraction of similar prices overflow easily). + * Float arithmetic eliminates the root cause of NaN states in the replay buffer. */ + float f_position = (float)position; + float f_cash = (float)cash; + float f_portfolio_value = (float)portfolio_value; + float f_peak_equity = (float)peak_equity; + float f_hold_time = (float)hold_time; + float f_realized_pnl = (float)realized_pnl; + float f_entry_price = (float)entry_price; + float f_trade_start_pnl = (float)trade_start_pnl; + + float f_equity = fmaxf(f_portfolio_value, 1.0f); + + float f_drawdown = (f_peak_equity > 1.0f) + ? (f_peak_equity - f_portfolio_value) / f_peak_equity + : 0.0f; + f_drawdown = fmaxf(f_drawdown, 0.0f); /* Unrealized P&L: current position mark-to-market minus entry cost */ - __nv_bfloat16 unrealized_pnl = (entry_price > bf16_zero() && position != bf16_zero()) - ? position * (market_features[(long long)bar_idx * market_dim] - entry_price) - : bf16_zero(); + float f_close = (float)market_features[(long long)bar_idx * market_dim]; + float f_unrealized_pnl = (f_entry_price > 0.0f && f_position != 0.0f) + ? f_position * (f_close - f_entry_price) + : 0.0f; /* Trade return since entry */ - __nv_bfloat16 trade_return = (trade_start_pnl != bf16_zero() || realized_pnl != bf16_zero()) - ? (realized_pnl - trade_start_pnl + unrealized_pnl) / equity - : bf16_zero(); + float f_trade_return = (f_trade_start_pnl != 0.0f || f_realized_pnl != 0.0f) + ? (f_realized_pnl - f_trade_start_pnl + f_unrealized_pnl) / f_equity + : 0.0f; /* Capital floor distance (75% of peak = capital floor) */ - __nv_bfloat16 floor_val = peak_equity * bf16(0.75f); - __nv_bfloat16 floor_dist = (equity > floor_val && equity > one) - ? (equity - floor_val) / equity - : bf16_zero(); + float f_floor_val = f_peak_equity * 0.75f; + float f_floor_dist = (f_equity > f_floor_val && f_equity > 1.0f) + ? (f_equity - f_floor_val) / f_equity + : 0.0f; int portfolio_base = market_dim; if (portfolio_base + 7 < state_dim) { - out[portfolio_base + 0] = position; /* raw position */ - out[portfolio_base + 1] = unrealized_pnl / equity; /* trade P&L signal */ - out[portfolio_base + 2] = drawdown; /* risk: how deep are we? */ - out[portfolio_base + 3] = hold_time / bf16(100.0f); /* how long in trade? */ - out[portfolio_base + 4] = realized_pnl / equity; /* session P&L */ - out[portfolio_base + 5] = floor_dist; /* distance to game over */ - out[portfolio_base + 6] = trade_return; /* this trade's return */ - out[portfolio_base + 7] = cash / equity; /* available capital */ + out[portfolio_base + 0] = bf16(f_position); /* raw position */ + out[portfolio_base + 1] = bf16(f_unrealized_pnl / f_equity); /* trade P&L signal */ + out[portfolio_base + 2] = bf16(f_drawdown); /* risk: how deep are we? */ + out[portfolio_base + 3] = bf16(f_hold_time / 100.0f); /* how long in trade? */ + out[portfolio_base + 4] = bf16(f_realized_pnl / f_equity); /* session P&L */ + out[portfolio_base + 5] = bf16(f_floor_dist); /* distance to game over */ + out[portfolio_base + 6] = bf16(f_trade_return); /* this trade's return */ + out[portfolio_base + 7] = bf16(f_cash / f_equity); /* available capital */ } /* -- Multi-timeframe features: [market_dim+8 .. market_dim+8+16) -- @@ -249,43 +263,43 @@ extern "C" __global__ void experience_state_gather( if (past_idx >= 0 && slot + 3 < state_dim) { const __nv_bfloat16* now_row = market_features + (long long)bar_idx * market_dim; const __nv_bfloat16* past_row = market_features + (long long)past_idx * market_dim; - __nv_bfloat16 close_now = now_row[0]; - __nv_bfloat16 close_past = past_row[0]; + /* Float arithmetic — bf16 close prices (~5000) subtracted produce + * tiny differences that lose all precision in bf16 (7-bit mantissa). */ + float f_close_now = (float)now_row[0]; + float f_close_past = (float)past_row[0]; /* Return over N bars */ - __nv_bfloat16 ret = (close_past > bf16_zero()) ? (close_now - close_past) / close_past : bf16_zero(); - __nv_bfloat16 scaled_ret = ret * bf16(100.0f); - out[slot + 0] = bf16_fmax(bf16(-10.0f), bf16_fmin(bf16(10.0f), scaled_ret)); /* clamp +/-10% */ + float f_ret = (f_close_past > 0.0f) ? (f_close_now - f_close_past) / f_close_past : 0.0f; + float f_scaled_ret = f_ret * 100.0f; + out[slot + 0] = bf16(fminf(10.0f, fmaxf(-10.0f, f_scaled_ret))); - /* Volatility: scan high/low over window (approx from close changes) */ - __nv_bfloat16 max_val = close_now; - __nv_bfloat16 min_val = close_now; - __nv_bfloat16 vol_sum = bf16_zero(); + /* Volatility: scan high/low over window */ + float f_max_val = f_close_now; + float f_min_val = f_close_now; + float f_vol_sum = 0.0f; int vol_count = 0; for (int j = past_idx; j <= bar_idx && j < total_bars; j++) { - __nv_bfloat16 v = market_features[(long long)j * market_dim]; - if (v > max_val) max_val = v; - if (v < min_val) min_val = v; - /* Volume proxy: use feature index 4 if it exists (commonly volume) */ + float fv = (float)market_features[(long long)j * market_dim]; + if (fv > f_max_val) f_max_val = fv; + if (fv < f_min_val) f_min_val = fv; if (market_dim > 4) { - vol_sum = vol_sum + market_features[(long long)j * market_dim + 4]; + f_vol_sum += (float)market_features[(long long)j * market_dim + 4]; vol_count++; } } - __nv_bfloat16 range = (close_now > bf16_zero()) ? (max_val - min_val) / close_now : bf16_zero(); - __nv_bfloat16 scaled_range = range * bf16(100.0f); - out[slot + 1] = bf16_fmax(bf16_zero(), bf16_fmin(bf16(10.0f), scaled_range)); /* clamp 0-10% */ + float f_range = (f_close_now > 0.0f) ? (f_max_val - f_min_val) / f_close_now : 0.0f; + out[slot + 1] = bf16(fminf(10.0f, fmaxf(0.0f, f_range * 100.0f))); - /* Volume trend: current vs average */ - __nv_bfloat16 avg_vol = (vol_count > 0) ? vol_sum / bf16((float)vol_count) : bf16_one(); - __nv_bfloat16 cur_vol = (market_dim > 4) ? now_row[4] : bf16_one(); - __nv_bfloat16 vol_ratio = (avg_vol > bf16_zero()) ? cur_vol / avg_vol : bf16_one(); - out[slot + 2] = bf16_fmax(bf16_zero(), bf16_fmin(bf16(5.0f), vol_ratio)); /* clamp 0-5x */ + /* Volume trend */ + float f_avg_vol = (vol_count > 0) ? f_vol_sum / (float)vol_count : 1.0f; + float f_cur_vol = (market_dim > 4) ? (float)now_row[4] : 1.0f; + float f_vol_ratio = (f_avg_vol > 0.0f) ? f_cur_vol / f_avg_vol : 1.0f; + out[slot + 2] = bf16(fminf(5.0f, fmaxf(0.0f, f_vol_ratio))); /* Momentum: position within range [0=bottom, 1=top] */ - __nv_bfloat16 range_size = max_val - min_val; - out[slot + 3] = (range_size > bf16_zero()) - ? (close_now - min_val) / range_size + float f_range_size = f_max_val - f_min_val; + out[slot + 3] = (f_range_size > 0.0f) + ? bf16((f_close_now - f_min_val) / f_range_size) : bf16(0.5f); } else { /* Not enough history — zero pad */ diff --git a/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu b/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu index 45140e8b0..0f82a9339 100644 --- a/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu @@ -354,10 +354,6 @@ extern "C" __global__ void mse_loss_batched( if (tid == 0) { float weighted_loss = avg_mse * is_weight; - /* Safety net: zero out any NaN from bf16 input overflow (rewards/dones/IS-weights). - * F32 logits eliminated the main NaN source; this catches edge cases. */ - 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);