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

chore: remove runtime NaN tracing stream syncs

Browse Source 
Removes run_nan_tracing (3 stream syncs per step on first 20 steps)
and unused grad_buf_ptr/params_buf_ptr accessors. Root causes are
fixed — runtime detection no longer needed. isfinite guards in CUDA
kernels remain as silent defense-in-depth.

17/17 smoke tests pass.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-05 15:33:36 +02:00
parent a5eeeaa477
commit b7ea7c0760
2 changed files with 0 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1000,15 +1000,6 @@ impl GpuDqnTrainer {
/// Raw device pointer to `grad_buf` for auxiliary gradient injection between graph phases.
///
/// Between `graph_forward` replay (which populates `grad_buf` with C51 gradients)
/// and `graph_adam` replay (which reads `grad_buf` for Adam), external code can ADD
/// auxiliary gradients (IQN trunk, ensemble heads) to `grad_buf` via SAXPY.
///
/// Caller must use `EventTrackingGuard` and operate on the same stream.
pub fn grad_buf_ptr(&self) -> u64 {
self.grad_buf.raw_ptr()
}
/// Mutable reference to grad_buf for test injection.
#[cfg(test)]
pub fn grad_buf_mut(&mut self) -> &mut CudaSlice<f32> {
@@ -1021,11 +1012,6 @@ 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()
}
/// Total number of trainable parameters.
pub fn total_params(&self) -> usize {
self.total_params

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

@@ -655,9 +655,6 @@ impl FusedTrainingCtx {
let gpu_batch = batch.gpu_batch.as_ref()
.ok_or_else(|| anyhow::anyhow!("Fused training requires gpu_batch (GPU PER)"))?;
// ── NaN detection: GPU-side checks on first 20 steps ────────────
let nan_diag = self.steps_since_varmap_sync < 20;
// ── 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}"))?;
@@ -695,11 +692,6 @@ 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() {
self.pre_replay_state_update(agent);
@@ -709,11 +701,6 @@ impl FusedTrainingCtx {
self.submit_aux_ops(agent, gpu_batch)?;
}
// ── 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) ──────────────────────
// Ensemble diversity (variable topology, complex forward passes).
@@ -750,11 +737,6 @@ 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())
@@ -937,33 +919,6 @@ 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).