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

fix(bf16): ROOT CAUSE — float experience features + IS-weight overflow clamp

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-29 00:42:06 +01:00
parent 51dd200e39
commit 5232a1ae31
4 changed files with 70 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
crates/ml-dqn/src/replay_buffer_kernels.cu
View File

@@ -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)

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

@@ -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);

112
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -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 */

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

@@ -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);