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

cleanup: remove BUFFER_DIAG, RECAPTURE_DIAG, GRAD_DIAG diagnostic infrastructure

Browse Source 
Root causes are fixed (cuBLASLt heuristic, entropy inv_batch, overflow,
IQL OOB). Removes: run_buffer_diagnostics, run_recapture_diagnostics,
debug_buffer_norm_f32, diag_recapture_remaining field,
gradient_stage_diagnostics field + FOXHUNT_GRAD_DIAG env var.
Eliminates ~4 stream syncs per step for first 3 steps each fold.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-14 08:10:46 +02:00
parent a5a754ec53
commit bf4e92bdb7
2 changed files with 1 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -613,7 +613,6 @@ pub struct GpuDqnTrainer {
m_buf: CudaSlice<f32>, // [TOTAL_PARAMS] Adam first moment (f32 for precision)
v_buf: CudaSlice<f32>, // [TOTAL_PARAMS] Adam second moment (f32 for precision)
pub(crate) grad_norm_buf: CudaSlice<f32>, // [1] f32 L2 norm (written by finalize)
diag_recapture_remaining: u32, // diagnostic steps remaining after graph recapture
grad_norm_partials: CudaSlice<f32>, // [grad_norm_blocks] per-block partial sums
grad_norm_blocks: usize, // number of blocks for grad_norm kernel
cql_grad_scratch: CudaSlice<f32>, // [TOTAL_PARAMS] f32 CQL gradient isolation buffer
@@ -2124,89 +2123,6 @@ impl GpuDqnTrainer {
Ok(norm_f32[0])
}
/// Per-step buffer diagnostics: check every critical buffer after forward+backward.
/// Logs L2 norms of d_value_logits, d_adv_logits, grad_buf, and h_s2 activations.
/// Costs ~4 stream syncs. Gate on step count in caller.
pub fn run_buffer_diagnostics(&mut self, step: usize) {
let b = self.config.batch_size;
let na = self.config.num_atoms;
let tba = (self.config.branch_0_size + self.config.branch_1_size
+ self.config.branch_2_size + self.config.branch_3_size) * na;
let d_val_norm = self.debug_buffer_norm_f32(self.d_value_logits_buf.raw_ptr(), b * na).unwrap_or(f32::NAN);
let d_adv_norm = self.debug_buffer_norm_f32(self.d_adv_logits_buf.raw_ptr(), b * tba).unwrap_or(f32::NAN);
let grad_norm = self.debug_buffer_norm_f32(self.ptrs.grad_buf, self.total_params).unwrap_or(f32::NAN);
tracing::warn!(
step,
d_val_norm,
d_adv_norm,
grad_norm,
c51_alpha = self.c51_alpha,
"BUFFER_DIAG: per-step buffer norms after forward+backward"
);
}
/// Extended diagnostics after graph recapture: checks MSE scratch, C51 d_logits,
/// blended d_logits, and grad_buf. Fires for first 3 steps after any graph recapture.
pub fn run_recapture_diagnostics(&mut self) {
if self.diag_recapture_remaining == 0 { return; }
self.diag_recapture_remaining -= 1;
let b = self.config.batch_size;
let na = self.config.num_atoms;
let tba = (self.config.branch_0_size + self.config.branch_1_size
+ self.config.branch_2_size + self.config.branch_3_size) * na;
// After graph replay: d_value/adv_logits contain BLENDED gradients,
// MSE scratch contains MSE-only gradients, grad_buf contains backward output.
let d_val_norm = self.debug_buffer_norm_f32(self.d_value_logits_buf.raw_ptr(), b * na).unwrap_or(f32::NAN);
let d_adv_norm = self.debug_buffer_norm_f32(self.d_adv_logits_buf.raw_ptr(), b * tba).unwrap_or(f32::NAN);
let mse_val_norm = self.debug_buffer_norm_f32(self.d_value_logits_mse.raw_ptr(), b * na).unwrap_or(f32::NAN);
let mse_adv_norm = self.debug_buffer_norm_f32(self.d_adv_logits_mse.raw_ptr(), b * tba).unwrap_or(f32::NAN);
let grad_norm = self.debug_buffer_norm_f32(self.ptrs.grad_buf, self.total_params).unwrap_or(f32::NAN);
tracing::warn!(
remaining = self.diag_recapture_remaining,
c51_alpha = self.c51_alpha,
d_val_norm,
d_adv_norm,
mse_val_norm,
mse_adv_norm,
grad_norm,
"RECAPTURE_DIAG: buffer norms after graph recapture replay"
);
}
/// Compute L2 norm of an arbitrary f32 GPU buffer (diagnostic use only).
/// Reuses the standalone grad_norm two-phase reduction. Costs 1 stream sync.
pub fn debug_buffer_norm_f32(&mut self, ptr: u64, n_elems: usize) -> Result<f32, MLError> {
let _evt_guard = EventTrackingGuard::new(self.stream.context());
let n_i32 = n_elems as i32;
let blocks = ((n_elems + 255) / 256) as u32;
let partials_ptr = self.grad_norm_partials.raw_ptr();
let buf_ptr = self.aux_norm_scratch.raw_ptr();
let nb = blocks as i32;
unsafe {
self.stream.launch_builder(&self.grad_norm_standalone)
.arg(&ptr).arg(&partials_ptr).arg(&n_i32)
.launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })
.map_err(|e| MLError::ModelError(format!("debug_norm phase1: {e}")))?;
self.stream.launch_builder(&self.grad_norm_finalize_kernel)
.arg(&partials_ptr).arg(&buf_ptr).arg(&nb)
.launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })
.map_err(|e| MLError::ModelError(format!("debug_norm phase2: {e}")))?;
cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream());
}
let mut norm_val = [0.0_f32; 1];
unsafe {
cudarc::driver::sys::cuMemcpyDtoH_v2(
norm_val.as_mut_ptr().cast(), self.aux_norm_scratch.raw_ptr(), 4,
);
}
Ok(norm_val[0])
}
/// Apply GPU multi-head feature attention to `save_h_s2` (post-graph).
///
/// Runs 4-head self-attention over the trunk output `h_s2 [B, SHARED_H2]`
@@ -3074,7 +2990,6 @@ impl GpuDqnTrainer {
m_buf,
v_buf,
grad_norm_buf,
diag_recapture_remaining: 3,
grad_norm_partials,
grad_norm_blocks,
cql_grad_scratch,
@@ -4635,7 +4550,6 @@ impl GpuDqnTrainer {
self.graph_forward = Some(SendSyncGraph(graph_fwd));
self.graph_forward_ddqn = Some(SendSyncGraph(graph_ddqn));
self.graph_adam = Some(SendSyncGraph(graph_adam));
self.diag_recapture_remaining = 3; // fire diagnostics for first 3 steps after recapture
Ok(())
}

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

@@ -286,10 +286,6 @@ pub(crate) struct FusedTrainingCtx {
pub(crate) hindsight_fraction: f64,
/// v8: Curriculum learning enabled (false by default).
pub(crate) curriculum_enabled: bool,
/// When true, disables mega-graph/graph_aux capture and logs per-stage
/// gradient norms (raw → post-primary-clip → post-IQN → post-CQL).
/// Costs ~4 stream syncs per step. For debugging H100 gradient collapse.
pub(crate) gradient_stage_diagnostics: bool,
}
impl Drop for FusedTrainingCtx {
@@ -743,7 +739,6 @@ impl FusedTrainingCtx {
prev_popart_var: 0.0,
hindsight_fraction: hyperparams.hindsight_fraction,
curriculum_enabled: hyperparams.curriculum_enabled,
gradient_stage_diagnostics: std::env::var("FOXHUNT_GRAD_DIAG").is_ok(),
})
}
@@ -1018,13 +1013,6 @@ impl FusedTrainingCtx {
.map_err(|e| anyhow::anyhow!("pre-Adam grad_buf clip: {e}"))?;
}
// ── BUFFER_DIAG: check where gradients die (first 3 steps each fold) ──
if self.steps_since_varmap_sync < 3 {
self.trainer.run_buffer_diagnostics(self.steps_since_varmap_sync);
}
// ── RECAPTURE_DIAG: check C51 gradient path after graph recapture ──
self.trainer.run_recapture_diagnostics();
// ── Step 5: Pruning + Adam ───────────────────────────────────────
self.trainer.apply_pruning_mask()
.map_err(|e| anyhow::anyhow!("Pruning mask apply: {e}"))?;
@@ -1073,12 +1061,11 @@ impl FusedTrainingCtx {
// Capture mega-graph at step 2 (all sub-trainers initialized).
// Mega-graph fuses spectral + forward + backward + aux → 1 launch.
// Skip when gradient_stage_diagnostics is on — need ungraphed path for per-stage reads.
// Capture graph_mega on step 0 — ONE capture session for determinism.
// Individual graphs at steps 0-1 caused cross-process non-determinism:
// each separate capture runs cublasLtMatmul ungraphed, contaminating
// the shared handle's internal state for subsequent captures.
if self.graph_mega.is_none() && !self.gradient_stage_diagnostics {
if self.graph_mega.is_none() {
if let Err(e) = self.capture_graph_mega(agent, gpu_batch) {
tracing::warn!(
"graph_mega capture failed (non-fatal, falling back to individual graphs): {e}"
@@ -1306,17 +1293,6 @@ impl FusedTrainingCtx {
iqn.compute_iqr()
.map_err(|e| anyhow::anyhow!("IQN IQR: {e}"))?;
// ── DIAG: post-IQN trunk SAXPY gradient norm ──────────
if self.gradient_stage_diagnostics {
let post_iqn_norm = self.trainer.read_grad_norm_standalone_sync()
.unwrap_or(f32::NAN);
tracing::warn!(
batch_size = self.batch_size,
post_iqn_norm,
stage = "post_iqn_saxpy",
"GRAD_DIAG: gradient norm AFTER IQN trunk SAXPY"
);
}
}
Err(e) => {
tracing::warn!("IQN step failed (non-fatal): {e}");
@@ -1346,18 +1322,6 @@ impl FusedTrainingCtx {
}
}
// ── DIAG: final gradient norm (after all component injections) ────
if self.gradient_stage_diagnostics {
let final_norm = self.trainer.read_grad_norm_standalone_sync()
.unwrap_or(f32::NAN);
tracing::warn!(
batch_size = self.batch_size,
final_norm,
stage = "final_pre_adam",
"GRAD_DIAG: final gradient norm before Adam (all components injected)"
);
}
// Regime-adaptive PER scaling.
self.trainer.regime_scale_td_errors()
.map_err(|e| anyhow::anyhow!("Regime PER scaling: {e}"))?;