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

fix: 7 NaN root causes in GPU training pipeline + regression test

Browse Source 
Root causes found and fixed:
- f32/bf16 type mismatch in pad_states_kernel (PER stores f32, kernel read bf16)
- States buffer stride mismatch in backward_full (lda=48 on pad128=128 buffer)
- States buffer stride mismatch in apply_iqn_trunk_gradient (same bug)
- IQN trunk forward stride mismatch in custom matvec kernel
- IQN grad_norm sqrt-per-block bug (sqrtf of partials → wrong norm + NaN)
- Uninitialized state vector slots in experience_state_gather (stale GPU memory)
- CUDA graph warmup launches producing extreme gradients from Xavier weights

Defense-in-depth guards:
- f32_to_bf16_kernel: NaN→0, clamp ±500
- Adam isfinite guards in DQN, IQN, attention, curiosity kernels
- IQN backward isfinite on d_h_s2 atomicAdd
- clipped_saxpy isfinite on auxiliary gradient input

Adds test_no_nan_after_graph_capture regression test (2s, fxcache-based).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-05 12:27:31 +02:00
parent 57844e1cc5
commit fc3f6f9df7
10 changed files with 394 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
crates/ml/src/cuda_pipeline/attention_backward_kernel.cu
View File

@@ -491,6 +491,12 @@ extern "C" __global__ void attn_adam_kernel(
__nv_bfloat16 bf_wd = bf16(weight_decay);
__nv_bfloat16 bf_max_gn = bf16(max_grad_norm);
/* Skip NaN/Inf gradients — leave param, momentum, variance unchanged. */
{
float g_check = (float)grads[tid];
if (!isfinite(g_check)) { grads[tid] = bf16_zero(); return; }
}
/* Gradient clipping */
__nv_bfloat16 grad_norm = norm[0];
__nv_bfloat16 clip_scale = (grad_norm > bf_max_gn && grad_norm > bf16_zero())

56
crates/ml/src/cuda_pipeline/batched_backward.rs
View File

@@ -132,6 +132,7 @@ pub struct CublasBackward {
// ── Network dimensions (baked at construction) ──
batch_size: usize,
state_dim: usize,
state_dim_padded: usize,
shared_h1: usize,
shared_h2: usize,
value_h: usize,
@@ -175,6 +176,7 @@ impl CublasBackward {
f32_to_bf16_cast_kernel,
batch_size: config.batch_size,
state_dim: config.state_dim,
state_dim_padded: (config.state_dim + 127) & !127,
shared_h1: config.shared_h1,
shared_h2: config.shared_h2,
value_h: config.value_h,
@@ -362,6 +364,48 @@ impl CublasBackward {
Ok(())
}
/// Like `backward_fc_layer` but with a custom leading dimension for the X
/// matrix (input activations). Used when X has a padded row stride
/// (e.g. states buffer padded to pad128(state_dim) for CUTLASS K-tile).
///
/// `x_lda` is the row stride of X in elements (>= in_dim).
/// `in_dim` is the logical column count (weight matrix width).
pub fn backward_fc_layer_lda(
&self,
stream: &Arc<CudaStream>,
dy: u64, x: u64, w: u64, dw: u64, db: u64, dx: u64,
out_dim: usize, in_dim: usize, x_lda: usize, batch: usize,
) -> Result<(), MLError> {
// Weight gradient: dW[out, in] += dY^T @ X
// X has lda=x_lda (padded stride), dW has ldc=in_dim (actual weight width).
self.gemmex_bf16_acc_f32(
cublas_sys::cublasOperation_t::CUBLAS_OP_N,
cublas_sys::cublasOperation_t::CUBLAS_OP_T,
in_dim as i32, out_dim as i32, batch as i32,
1.0, x, x_lda as i32,
dy, out_dim as i32,
1.0, dw, in_dim as i32,
"fc_dW_lda",
)?;
self.launch_bias_grad(stream, dy, db, out_dim, batch)?;
if dx != 0 {
// dX has stride x_lda (same as input for dimensional consistency).
self.gemmex_bf16_acc_f32(
cublas_sys::cublasOperation_t::CUBLAS_OP_N,
cublas_sys::cublasOperation_t::CUBLAS_OP_N,
in_dim as i32, batch as i32, out_dim as i32,
1.0, w, in_dim as i32,
dy, out_dim as i32,
0.0, dx, x_lda as i32,
"fc_dX_lda",
)?;
}
Ok(())
}
/// Apply the ReLU derivative mask in-place: `dx[i] *= (activation[i] > 0)`.
///
/// Called after `backward_fc_layer()` for layers that used ReLU activation.
@@ -719,19 +763,23 @@ impl CublasBackward {
// Cast f32 d_h_s1 → bf16 staging for shared layer 1 backward
self.cast_dx_to_staging(stream, scratch_d_h_s1, staging_bf16, b * self.shared_h1)?;
// dW for shared layer 1: input is states (or bn_concat when bottleneck active)
// dW for shared layer 1: input is states (or bn_concat when bottleneck active).
// States buffer has padded stride pad128(state_dim) for CUTLASS K-tile
// alignment. Use backward_fc_layer_lda with the padded stride so the
// backward GEMM reads the correct row offsets.
// s1_dx_output: 0 = no dX (original, states not trainable),
// non-zero = compute dX for bottleneck backward chain rule.
self.backward_fc_layer(
self.backward_fc_layer_lda(
stream,
staging_bf16, // dY [B, SH1] — bf16 staging
states, // X [B, s1_input_dim] — states or bn_concat
states, // X [B, s1_input_dim] — padded stride
w_ptrs[0], // W_s1 [SH1, s1_input_dim]
grad_buf_base + goff_w_s1,
grad_buf_base + goff_b_s1,
s1_dx_output, // dX: 0=skip, or f32 ptr for bottleneck dX
self.shared_h1,
self.state_dim, // s1_input_dim (matches w_s1 width from param_sizes)
self.state_dim, // logical in_dim (weight width)
self.state_dim_padded, // x_lda (padded row stride)
b,
)?;

11
crates/ml/src/cuda_pipeline/common_device_functions.cuh
View File

@@ -526,8 +526,15 @@ extern "C" __global__ void f32_to_bf16_kernel(
int n
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n)
dst[idx] = __float2bfloat16(src[idx]);
if (idx < n) {
float v = src[idx];
// Saturate: NaN→0, clamp to ±STAGING_SAFE_MAX.
// Prevents extreme f32 gradients (warmup / early training)
// from producing NaN/Inf in bf16 staging for cuBLAS backward.
if (!isfinite(v)) v = 0.0f;
else v = fminf(fmaxf(v, -500.0f), 500.0f);
dst[idx] = __float2bfloat16(v);
}
}
/**

8
crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu
View File

@@ -209,6 +209,10 @@ extern "C" __global__ void curiosity_adam_step(
__nv_bfloat16 one_bf = bf16_one();
__nv_bfloat16 g = grads[i] / bf16((float)batch_size);
/* Skip NaN/Inf gradients */
if (!isfinite((float)g)) return;
m[i] = beta1_bf * m[i] + (one_bf - beta1_bf) * g;
v[i] = beta2_bf * v[i] + (one_bf - beta2_bf) * g * g;
__nv_bfloat16 m_hat = m[i] / bf16(1.0f - powf(beta1, (float)step));
@@ -258,6 +262,10 @@ extern "C" __global__ void curiosity_adam_step_fused(
__nv_bfloat16 one_bf = bf16_one();
__nv_bfloat16 grad = g[local_i] / bf16((float)batch_size);
/* Skip NaN/Inf gradients */
if (!isfinite((float)grad)) return;
mm[local_i] = beta1_bf * mm[local_i] + (one_bf - beta1_bf) * grad;
vv[local_i] = beta2_bf * vv[local_i] + (one_bf - beta2_bf) * grad * grad;
__nv_bfloat16 m_hat = mm[local_i] / bf16(1.0f - powf(beta1, (float)step));

11
crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu
View File

@@ -98,6 +98,11 @@ extern "C" __global__ void dqn_adam_update_kernel(
int t = *t_ptr;
float g_f = grads[idx];
/* Skip NaN/Inf gradients — leave param, momentum, variance unchanged.
* Non-deterministic NaN can arise from IQN/CQL auxiliary gradient
* injection during early training with extreme quantile estimates. */
if (!isfinite(g_f)) return;
/* Clip using the COMPLETED sum-of-squares (native float buffer).
* dqn_grad_norm_kernel accumulates float partial sums via atomicAdd. */
float norm_sq_f = *grad_norm_f32;
@@ -508,8 +513,8 @@ extern "C" __global__ void spectral_norm_kernel(
* ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void pad_states_kernel(
__nv_bfloat16* __restrict__ dst, /* [batch, padded_sd] — padded output */
const __nv_bfloat16* __restrict__ src, /* [batch, sd] — contiguous input */
__nv_bfloat16* __restrict__ dst, /* [batch, padded_sd] — padded bf16 output */
const float* __restrict__ src, /* [batch, sd] — contiguous f32 input (PER stores f32) */
int batch,
int sd,
int padded_sd
@@ -518,7 +523,7 @@ extern "C" __global__ void pad_states_kernel(
if (idx >= batch * padded_sd) return;
int b = idx / padded_sd;
int j = idx % padded_sd;
dst[idx] = (j < sd) ? src[b * sd + j] : bf16_zero();
dst[idx] = (j < sd) ? __float2bfloat16(src[b * sd + j]) : bf16_zero();
}
/* ══════════════════════════════════════════════════════════════════════

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

@@ -409,10 +409,16 @@ extern "C" __global__ void experience_state_gather(
float* out = batch_states + (long long)i * state_dim;
/* Out-of-data: write zero state; cuBLAS will produce Q-values of ~0. */
if (bar_idx >= total_bars) {
for (int k = 0; k < state_dim; k++)
out[k] = 0.0f;
/* Zero the entire state vector first — any feature slots not explicitly
* written below stay zero instead of containing stale GPU memory.
* Without this, features beyond market_dim (portfolio, MTF, OFI) are
* uninitialized when state_dim > market_dim but those slots aren't
* reached by the feature-writing code below. */
for (int k = 0; k < state_dim; k++)
out[k] = 0.0f;
/* Out-of-data or negative index: state is already zeroed above. */
if (bar_idx < 0 || bar_idx >= total_bars) {
return;
}

53
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -1021,6 +1021,31 @@ impl GpuDqnTrainer {
&mut self.cql_grad_scratch
}
/// Raw pointer to f32 master params buffer (for NaN diagnostics).
pub fn params_buf_ptr(&self) -> u64 {
self.params_buf.raw_ptr()
}
/// Raw pointer to CQL gradient scratch buffer (for NaN diagnostics).
pub fn cql_grad_scratch_ptr(&self) -> u64 {
self.cql_grad_scratch.raw_ptr()
}
/// Raw pointer to IQN trunk gradient scratch (for NaN diagnostics).
pub fn iqn_trunk_m_ptr(&self) -> u64 {
self.ptrs.iqn_trunk_m
}
/// Raw pointer to saved h_s1 activations (for NaN diagnostics).
pub fn save_h_s1_ptr(&self) -> u64 {
self.save_h_s1.raw_ptr()
}
/// Raw CUDA stream handle (for synchronous diagnostics).
pub fn stream_cu(&self) -> cudarc::driver::sys::CUstream {
self.stream.cu_stream()
}
/// Total number of trainable parameters.
pub fn total_params(&self) -> usize {
self.total_params
@@ -1180,8 +1205,10 @@ impl GpuDqnTrainer {
let dw = scratch_base; // goff_w_s1 in scratch
let db = scratch_base + w_s1_n as u64 * f32_u; // goff_b_s1 in scratch
self.cublas_backward.backward_fc_layer(
&self.stream, staging, x, w, dw, db, 0, sh1, sd, b,
// states_buf has padded stride pad128(sd) — use _lda variant
let sd_padded = pad128(sd);
self.cublas_backward.backward_fc_layer_lda(
&self.stream, staging, x, w, dw, db, 0, sh1, sd, sd_padded, b,
)?;
}
@@ -4199,25 +4226,13 @@ impl GpuDqnTrainer {
"CUDA graph_adam capture returned None — stream may not support capture".into()
))?;
// Launch both graphs on first capture (warm-up). On subsequent steps,
// only graph_forward is replayed by train_step_gpu(). The caller then
// injects auxiliary gradients and calls replay_adam_and_readback().
// Launch forward graphs for warmup ONLY — no Adam.
// With random Xavier weights, the first forward+backward produces extreme
// gradients. Running Adam on these corrupts the initial weights with NaN.
// The first real training step runs Adam with proper gradient clipping.
graph_mse.launch().map_err(|e| {
MLError::ModelError(format!("CUDA graph_forward_mse first launch: {e}"))
})?;
graph_fwd.launch().map_err(|e| {
MLError::ModelError(format!("CUDA graph_forward first launch: {e}"))
})?;
// Zero grad_buf after C51 warmup to prevent NaN from leaking into step 0
self.stream.memset_zeros(&mut self.grad_buf)
.map_err(|e| MLError::ModelError(format!("zero grad_buf post-capture: {e}")))?;
// No warmup launches — CUDA graphs work on first launch after
// cuGraphInstantiate. Warmup with random Xavier weights produced extreme
// gradients that overflowed the bf16 d_logits cast and corrupted
// downstream buffers, causing NaN at step 0.
info!(
"GpuDqnTrainer: 3 CUDA graphs captured and launched \
"GpuDqnTrainer: 3 CUDA graphs captured \
(graph_forward: MSE+C51; graph_forward_mse: MSE-only; \
graph_adam: grad_norm + adam + unflatten)"
);

27
crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu
View File

@@ -581,8 +581,9 @@ void iqn_backward_kernel(
dL_dembed[h / 32] = dL_dcomb[h / 32] * h_dist[h / 32];
/* Accumulate dL/d(h_s2) across all quantiles.
* This flows IQN's bounded Huber gradient to the shared trunk. */
atomicAdd(&d_h_s2_out[sample * hidden_dim + h],
dL_dcomb[h / 32] * embed_dist[h / 32]);
float d_trunk = dL_dcomb[h / 32] * embed_dist[h / 32];
if (isfinite(d_trunk))
atomicAdd(&d_h_s2_out[sample * hidden_dim + h], d_trunk);
}
/* ── Gradient through ReLU ── */
@@ -599,6 +600,7 @@ void iqn_backward_kernel(
float tau = my_taus[ti];
for (int h = lane; h < hidden_dim; h += 32) {
float dL_dpre = dL_dembed[h / 32];
if (!isfinite(dL_dpre)) continue;
/* Bias gradient */
atomicAdd(&grad_buf[off[1] + h], dL_dpre);
/* Weight gradient: outer product with cosine features */
@@ -644,7 +646,7 @@ void iqn_grad_norm_kernel(
}
if (threadIdx.x == 0)
atomicAdd(norm_out, sqrtf(shared[0]));
atomicAdd(norm_out, shared[0]);
}
/* ── Adam Update Kernel ─────────────────────────────────────────────── */
@@ -654,7 +656,7 @@ void iqn_adam_kernel(
float* __restrict__ grads,
float* __restrict__ m, /* first moment */
float* __restrict__ v, /* second moment */
const float* __restrict__ norm, /* [1] gradient L2 norm */
const float* __restrict__ norm_sq, /* [1] gradient sum-of-squares (NOT sqrt'd) */
float lr, float beta1, float beta2, float eps,
float weight_decay, float max_grad_norm,
const int* __restrict__ adam_t_buf, /* [1] step counter on device */
@@ -667,11 +669,16 @@ void iqn_adam_kernel(
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid >= total_params) return;
/* Gradient clipping */
float grad_norm_val = norm[0];
float g = grads[tid];
/* Skip NaN/Inf gradients — leave param, momentum, variance unchanged. */
if (!isfinite(g)) { grads[tid] = 0.0f; return; }
/* Gradient clipping: norm_sq[0] is raw sum-of-squares from grad_norm kernel */
float grad_norm_val = sqrtf(fmaxf(norm_sq[0], 0.0f));
float clip_scale = (grad_norm_val > max_grad_norm && grad_norm_val > 0.0f)
? max_grad_norm / grad_norm_val : 1.0f;
float g = grads[tid] * clip_scale;
g *= clip_scale;
/* AdamW update */
float mi = beta1 * m[tid] + (1.0f - beta1) * g;
@@ -840,7 +847,11 @@ void iqn_trunk_forward_kernel(
extern __shared__ float shmem[];
float* h1 = shmem; /* [shared_h1] */
const __nv_bfloat16* my_state = states + sample * state_dim;
/* States buffer is padded to pad128(state_dim) per row for CUTLASS
* K-tile alignment. Use padded stride for row indexing, but only
* read state_dim valid columns (padding is zero). */
int padded_sd = (state_dim + 127) & ~127;
const __nv_bfloat16* my_state = states + sample * padded_sd;
/* Layer 1: h1 = leaky_relu(W_s1 @ state + b_s1) — bf16 inputs, f32 compute */
for (int h = lane; h < shared_h1; h += 32) {

194
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -655,34 +655,13 @@ impl FusedTrainingCtx {
let gpu_batch = batch.gpu_batch.as_ref()
.ok_or_else(|| anyhow::anyhow!("Fused training requires gpu_batch (GPU PER)"))?;
// ── NaN detection: early steps with per-stage checks ────────────
// ── NaN detection: GPU-side checks on first 20 steps ────────────
let nan_diag = self.steps_since_varmap_sync < 20;
if nan_diag {
self.trainer.run_nan_checks_pre_forward()?;
let flags = self.trainer.read_nan_flags()?;
if flags[4] != 0 || flags[5] != 0 {
tracing::error!(
"NaN_DIAG step {} PRE_FORWARD: f32_params={} bf16_params={}",
self.steps_since_varmap_sync, flags[4], flags[5]
);
}
}
// ── Step 1: Spectral normalization BEFORE forward pass ─────────
self.trainer.apply_spectral_norm(&mut self.online_dueling, &mut self.online_branching)
.map_err(|e| anyhow::anyhow!("Spectral norm (pre-forward): {e}"))?;
if nan_diag {
self.trainer.run_nan_checks_pre_forward()?;
let flags = self.trainer.read_nan_flags()?;
if flags[4] != 0 || flags[5] != 0 {
tracing::error!(
"NaN_DIAG step {} POST_SPECTRAL: f32_params={} bf16_params={}",
self.steps_since_varmap_sync, flags[4], flags[5]
);
}
}
// ── Step 2: Upload batch + replay graph_forward ──────────────────
let _fused_placeholder = self.trainer.train_step_gpu(
gpu_batch,
@@ -690,7 +669,6 @@ impl FusedTrainingCtx {
&self.target_dueling, &self.target_branching,
).map_err(|e| anyhow::anyhow!("Fused train_step_gpu (forward only): {e}"))?;
// ── NaN detection: check key buffers after graph_forward ─────────
self.trainer.run_nan_checks_post_forward(self.batch_size)?;
// ── Step 2b: HER donor computation (outside graph_aux) ───────────
@@ -717,26 +695,95 @@ impl FusedTrainingCtx {
}
}
// ── NaN check: grad_buf after forward (before aux ops touch it) ──
if nan_diag {
self.run_nan_tracing("post_forward")?;
}
// ── Step 3: Auxiliary ops (graph_aux captured or replayed) ────────
if self.graph_aux.is_some() {
// Fast path: replay captured graph.
// Pre-update host-side state (adam_steps, tau) so graphed HtoD
// copies the correct values to device buffers.
self.pre_replay_state_update(agent);
let graph_aux = self.graph_aux.as_ref().unwrap();
graph_aux.launch(self.stream.cu_stream())?;
} else if self.steps_since_varmap_sync == 0 {
// First step: run ungraphed to initialize all sub-trainers.
self.submit_aux_ops(agent, gpu_batch)?;
} else {
// Second step: capture graph_aux.
self.submit_aux_ops(agent, gpu_batch)?;
// Capture on next call instead — we need the ops to execute first.
// Do the capture after this step completes so all buffers are initialized.
// (capture_graph_aux is called at the end of this step.)
}
// ── NaN check: grad_buf after aux ops (before conditional ops) ──
if nan_diag {
self.run_nan_tracing("post_aux")?;
}
// ── Step 4: Conditional ops (outside graph) ──────────────────────
// NaN trace: grad_buf after aux ops (graphed or ungraphed)
if nan_diag {
self.trainer.reset_nan_flags()?;
self.trainer.check_nan_f32(self.trainer.grad_buf_ptr(), self.trainer.total_params(), 0)?;
let f = self.trainer.read_nan_flags()?;
if f[0] != 0 {
// Dump IQN trunk scratch to find extreme values
let cfg = self.trainer.config();
let trunk_n = cfg.shared_h1 * cfg.state_dim + cfg.shared_h1
+ cfg.shared_h2 * cfg.shared_h1 + cfg.shared_h2;
let scratch_ptr = self.trainer.iqn_trunk_m_ptr();
let mut scratch = vec![0.0_f32; trunk_n];
unsafe {
cudarc::driver::sys::cuStreamSynchronize(self.trainer.stream_cu());
cudarc::driver::sys::cuMemcpyDtoH_v2(
scratch.as_mut_ptr().cast(), scratch_ptr,
trunk_n * std::mem::size_of::<f32>());
}
let nan_c = scratch.iter().filter(|v| v.is_nan()).count();
let inf_c = scratch.iter().filter(|v| v.is_infinite()).count();
let max_abs = scratch.iter().map(|v| v.abs())
.filter(|v| v.is_finite()).fold(0.0_f32, f32::max);
let max_idx = scratch.iter().position(|v| v.abs() == max_abs).unwrap_or(0);
// Also dump grad_buf
let tp = self.trainer.total_params();
let mut grad = vec![0.0_f32; tp];
unsafe {
cudarc::driver::sys::cuMemcpyDtoH_v2(
grad.as_mut_ptr().cast(), self.trainer.grad_buf_ptr(),
tp * std::mem::size_of::<f32>());
}
let g_nan = grad.iter().filter(|v| v.is_nan()).count();
let g_inf = grad.iter().filter(|v| v.is_infinite()).count();
let g_max = grad.iter().map(|v| v.abs())
.filter(|v| v.is_finite()).fold(0.0_f32, f32::max);
let g_nan_idx = grad.iter().position(|v| v.is_nan());
// Dump states and h_s1 activations
let sd_pad = (cfg.state_dim + 127) & !127;
let b_sz = cfg.batch_size;
let mut states = vec![half::bf16::ZERO; b_sz * sd_pad];
let mut h_s1 = vec![half::bf16::ZERO; b_sz * cfg.shared_h1];
unsafe {
cudarc::driver::sys::cuMemcpyDtoH_v2(
states.as_mut_ptr().cast(), self.trainer.states_buf_ptr(),
states.len() * 2);
cudarc::driver::sys::cuMemcpyDtoH_v2(
h_s1.as_mut_ptr().cast(), self.trainer.save_h_s1_ptr(),
h_s1.len() * 2);
}
let s_max = states.iter().map(|v| v.to_f32().abs())
.filter(|v| v.is_finite()).fold(0.0_f32, f32::max);
let h1_max = h_s1.iter().map(|v| v.to_f32().abs())
.filter(|v| v.is_finite()).fold(0.0_f32, f32::max);
let s_inf = states.iter().filter(|v| !v.to_f32().is_finite()).count();
let h1_inf = h_s1.iter().filter(|v| !v.to_f32().is_finite()).count();
tracing::error!(
step = self.steps_since_varmap_sync,
"NaN in grad_buf after aux: \
scratch[nan={nan_c},inf={inf_c},max={max_abs:.2e}@{max_idx}] \
grad[nan={g_nan},inf={g_inf},max={g_max:.2e},first_nan={g_nan_idx:?}] \
states[max={s_max:.2},inf={s_inf}] h_s1[max={h1_max:.2},inf={h1_inf}]"
);
}
}
// Ensemble diversity (variable topology, complex forward passes).
if !self.ensemble_extra_heads.is_empty() {
if let Err(e) = self.run_ensemble_step() {
@@ -771,6 +818,11 @@ impl FusedTrainingCtx {
let fused_result = self.trainer.replay_adam_and_readback()
.map_err(|e| anyhow::anyhow!("graph_adam replay: {e}"))?;
// ── NaN check: after Adam update ────────────────────────────────
if nan_diag {
self.run_nan_tracing("post_adam")?;
}
// ── Step 6: PER priority update (outside graph) ──────────────────
if let (Some(priorities_tensor), Some((alpha, epsilon))) =
(agent.priorities_tensor()?, agent.per_alpha_epsilon())
@@ -868,6 +920,12 @@ impl FusedTrainingCtx {
).map_err(|e| anyhow::anyhow!("GPU EMA target update: {e}"))?;
}
// ── Granular NaN tracing: check grad_buf after each grad_buf-modifying op
let nan_trace_aux = self.steps_since_varmap_sync < 20;
if nan_trace_aux {
self.run_nan_tracing("after_grad_clip")?;
}
// Attention forward + backward + Adam.
if let Some(ref mut attn) = self.gpu_attention {
self.trainer.apply_attention_forward(attn, self.batch_size)
@@ -923,6 +981,45 @@ impl FusedTrainingCtx {
}
}
if nan_trace_aux {
self.run_nan_tracing("after_iqn_trunk")?;
// If NaN appeared, dump IQN trunk scratch + grad_buf extremes
{
let tp = self.trainer.total_params();
self.trainer.reset_nan_flags()?;
self.trainer.check_nan_f32(self.trainer.grad_buf_ptr(), tp, 0)?;
let f = self.trainer.read_nan_flags()?;
if f[0] != 0 {
// Dump IQN scratch buffer
let trunk_n = self.trainer.config().shared_h1 * self.trainer.config().state_dim
+ self.trainer.config().shared_h1
+ self.trainer.config().shared_h2 * self.trainer.config().shared_h1
+ self.trainer.config().shared_h2;
let scratch_ptr = self.trainer.iqn_trunk_m_ptr();
let mut scratch_host = vec![0.0_f32; trunk_n];
unsafe {
cudarc::driver::sys::cuStreamSynchronize(
self.trainer.stream_cu());
cudarc::driver::sys::cuMemcpyDtoH_v2(
scratch_host.as_mut_ptr().cast(), scratch_ptr,
trunk_n * std::mem::size_of::<f32>());
}
let nan_c = scratch_host.iter().filter(|v| v.is_nan()).count();
let inf_c = scratch_host.iter().filter(|v| v.is_infinite()).count();
let max_abs = scratch_host.iter().map(|v| v.abs())
.filter(|v| v.is_finite())
.fold(0.0_f32, f32::max);
tracing::error!(
step = self.steps_since_varmap_sync,
trunk_n,
nan_c, inf_c,
max_abs = format!("{max_abs:.2e}"),
"IQN trunk scratch extremes when grad_buf has NaN"
);
}
}
}
// IQN→PER loss: cast f32 per_sample_loss → bf16 td_errors_buf.
// IQN is now f32 native; PER system expects bf16 td_errors.
if let Some(ref iqn) = self.gpu_iqn {
@@ -946,6 +1043,10 @@ impl FusedTrainingCtx {
}
}
if nan_trace_aux {
self.run_nan_tracing("after_cql_saxpy")?;
}
// Regime-adaptive PER scaling.
self.trainer.regime_scale_td_errors()
.map_err(|e| anyhow::anyhow!("Regime PER scaling: {e}"))?;
@@ -953,6 +1054,33 @@ impl FusedTrainingCtx {
Ok(())
}
/// GPU-side NaN tracing across the full training step.
///
/// Checks grad_buf before Adam (to catch aux pipeline NaN) and f32 params
/// after Adam (to catch optimizer NaN). One stream sync per call (~5µs).
fn run_nan_tracing(&mut self, phase: &str) -> Result<()> {
let step = self.steps_since_varmap_sync;
let tp = self.trainer.total_params();
self.trainer.reset_nan_flags()?;
self.trainer.check_nan_f32(self.trainer.params_buf_ptr(), tp, 0)?;
self.trainer.check_nan_f32(self.trainer.grad_buf_ptr(), tp, 1)?;
let flags = self.trainer.read_nan_flags()?;
let params_nan = flags[0] != 0;
let grad_nan = flags[1] != 0;
if params_nan || grad_nan {
tracing::warn!(
step,
phase,
f32_params_nan = params_nan,
grad_buf_nan = grad_nan,
"NaN detected in training pipeline"
);
}
Ok(())
}
/// Capture the auxiliary ops into a CUDA graph for zero-overhead replay.
///
/// Called once after the first step succeeds (all sub-trainers initialized).

87
crates/ml/src/trainers/dqn/smoke_tests/regression.rs
View File

@@ -5,6 +5,93 @@
use super::helpers::*;
/// Regression: NaN must not appear in params after CUDA graph capture.
///
/// Root cause was warmup graph launches in capture_training_graphs() producing
/// extreme gradients from random Xavier weights that overflowed the bf16
/// d_logits cast and corrupted downstream buffers, causing NaN at step 0.
/// Fixed by removing warmup launches (CUDA graphs work without them after
/// cuGraphInstantiate) and adding saturation to f32_to_bf16_kernel.
///
/// Uses fxcache directly — NO DBN parsing. Caps at 500 bars for speed (~2s).
#[test]
#[ignore] // GPU required
fn test_no_nan_after_graph_capture() -> anyhow::Result<()> {
use crate::fxcache;
let cache_dir = feature_cache_dir();
if !cache_dir.exists() {
eprintln!("SKIP: no feature-cache at {:?}", cache_dir);
return Ok(());
}
let entries: Vec<_> = std::fs::read_dir(&cache_dir)?
.filter_map(|e| e.ok())
.filter(|e| e.path().extension().and_then(|s| s.to_str()) == Some("fxcache"))
.collect();
if entries.is_empty() {
eprintln!("SKIP: no .fxcache files in {:?}", cache_dir);
return Ok(());
}
let fxcache_data = fxcache::load_fxcache(&entries[0].path())?;
let n = fxcache_data.bar_count.min(500);
let train_end = (n * 80) / 100;
let mut p = smoke_params();
p.epochs = 1;
p.early_stopping_enabled = false;
p.min_epochs_before_stopping = 1;
let mut trainer = smoke_trainer_with(p)?;
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()?;
let features = &fxcache_data.features[..n];
let targets = &fxcache_data.targets[..n];
let ofi = &fxcache_data.ofi[..n.min(fxcache_data.ofi.len())];
rt.block_on(trainer.init_from_fxcache(features, targets, ofi))?;
trainer.set_training_range(0, train_end, train_end, n);
trainer.set_val_data_from_slices(
&fxcache_data.features[train_end..n],
&fxcache_data.targets[train_end..n],
train_end,
);
rt.block_on(trainer.reset_for_fold())?;
let result = rt.block_on(trainer.train_fold_from_slices(
&features[..train_end], &targets[..train_end],
|_epoch, _bytes, _best| Ok("skip".to_owned()),
));
// Training must complete — NaN at step 0 would cause an error here
let metrics = result.map_err(|e| {
anyhow::anyhow!(
"NaN REGRESSION: training failed (likely NaN at step 0 from graph capture warmup): {e}"
)
})?;
assert_finite(metrics.loss, "loss");
assert!(
metrics.epochs_trained >= 1,
"NaN REGRESSION: only {} epochs completed (expected 1)",
metrics.epochs_trained
);
let grad_norm = metrics.additional_metrics.get("avg_gradient_norm")
.copied().unwrap_or(f64::NAN);
assert_finite(grad_norm, "avg_gradient_norm");
tracing::info!(
"NaN regression PASS: loss={:.6}, grad_norm={:.2}, epochs={}",
metrics.loss, grad_norm, metrics.epochs_trained
);
drop(trainer);
drop(rt);
Ok(())
}
/// Regression: GPU training must not hang on first epoch.
///
/// Root cause was VRAM oversubscription from detect_gpu_hardware auto-scaling