From fc3f6f9df7e9503715f76c6d7ef5699e3a3161d7 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 5 Apr 2026 12:27:31 +0200 Subject: [PATCH] fix: 7 NaN root causes in GPU training pipeline + regression test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../attention_backward_kernel.cu | 6 + .../ml/src/cuda_pipeline/batched_backward.rs | 56 ++++- .../cuda_pipeline/common_device_functions.cuh | 11 +- .../curiosity_training_kernel.cu | 8 + .../src/cuda_pipeline/dqn_utility_kernels.cu | 11 +- .../src/cuda_pipeline/experience_kernels.cu | 14 +- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 53 +++-- .../src/cuda_pipeline/iqn_dual_head_kernel.cu | 27 ++- crates/ml/src/trainers/dqn/fused_training.rs | 194 +++++++++++++++--- .../trainers/dqn/smoke_tests/regression.rs | 87 ++++++++ 10 files changed, 394 insertions(+), 73 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu b/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu index ec1e5d364..bafcb0744 100644 --- a/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu +++ b/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu @@ -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()) diff --git a/crates/ml/src/cuda_pipeline/batched_backward.rs b/crates/ml/src/cuda_pipeline/batched_backward.rs index ea995f6e1..6d1d12080 100644 --- a/crates/ml/src/cuda_pipeline/batched_backward.rs +++ b/crates/ml/src/cuda_pipeline/batched_backward.rs @@ -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, + 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, )?; diff --git a/crates/ml/src/cuda_pipeline/common_device_functions.cuh b/crates/ml/src/cuda_pipeline/common_device_functions.cuh index 4c8c0bc21..dce118b9f 100644 --- a/crates/ml/src/cuda_pipeline/common_device_functions.cuh +++ b/crates/ml/src/cuda_pipeline/common_device_functions.cuh @@ -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); + } } /** diff --git a/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu b/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu index abc4bc2ff..4a68e4f90 100644 --- a/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu +++ b/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu @@ -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)); diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index f735ab27f..6ebfb7d70 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -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(); } /* ══════════════════════════════════════════════════════════════════════ diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index c19bdcdde..2d174a538 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -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; } diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index ecb1d6e52..a794dd505 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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)" ); diff --git a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu index b8ed59999..5a24fda43 100644 --- a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu @@ -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) { diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index ca3b47a32..15a3cdee2 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -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::()); + } + 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::()); + } + 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::()); + } + 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). diff --git a/crates/ml/src/trainers/dqn/smoke_tests/regression.rs b/crates/ml/src/trainers/dqn/smoke_tests/regression.rs index b356ae4d0..cd9d2a21d 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/regression.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/regression.rs @@ -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