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

debug: add eprintln diagnostics to find H100 hang point

Browse Source 
Temporary diagnostics between each fused training step to identify
which operation hangs after replay_forward on H100.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-06 19:55:54 +02:00
parent 29d58efca2
commit 927460b9e8
2 changed files with 10 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -701,6 +701,7 @@ impl FusedTrainingCtx {
).map_err(|e| anyhow::anyhow!("Fused train_step_gpu (forward only): {e}"))?;
self.trainer.run_nan_checks_post_forward(self.batch_size)?;
eprintln!("FUSED_DIAG: step {} — post-forward NaN checks done", self.steps_since_varmap_sync);
// ── Step 2b: HER donor computation (outside graph_aux) ───────────
// Donor indices vary per step (random/future/final). The computation
@@ -726,6 +727,7 @@ impl FusedTrainingCtx {
}
}
eprintln!("FUSED_DIAG: step {} — HER done, entering aux ops", self.steps_since_varmap_sync);
// ── Step 3: Auxiliary ops (graph_aux captured or replayed) ────────
if self.graph_aux.is_some() {
self.pre_replay_state_update(agent);
@@ -735,6 +737,7 @@ impl FusedTrainingCtx {
self.submit_aux_ops(agent, gpu_batch)?;
}
eprintln!("FUSED_DIAG: step {} — aux ops done", self.steps_since_varmap_sync);
// ── Step 4: Conditional ops (outside graph) ──────────────────────
// Ensemble diversity (variable topology, complex forward passes).
@@ -764,12 +767,14 @@ impl FusedTrainingCtx {
}
self.pending_vaccine_batch = None;
eprintln!("FUSED_DIAG: step {} — conditional ops done, entering Adam", self.steps_since_varmap_sync);
// ── Step 5: Pruning + Adam ───────────────────────────────────────
self.trainer.apply_pruning_mask()
.map_err(|e| anyhow::anyhow!("Pruning mask apply: {e}"))?;
let fused_result = self.trainer.replay_adam_and_readback()
.map_err(|e| anyhow::anyhow!("graph_adam replay: {e}"))?;
eprintln!("FUSED_DIAG: step {} — Adam done, entering PER update", self.steps_since_varmap_sync);
// ── Step 6: PER priority update via seg_tree_update ─────────────
// Uses the replay buffer's segment tree kernel: writes priorities,
@@ -780,6 +785,7 @@ impl FusedTrainingCtx {
&self.stream,
).map_err(|e| anyhow::anyhow!("PER seg_tree_update: {e}"))?;
eprintln!("FUSED_DIAG: step {} — PER update done", self.steps_since_varmap_sync);
// Bookkeeping.
agent.fused_post_step_gpu_bookkeeping()
.map_err(|e| anyhow::anyhow!("Fused GPU bookkeeping: {e}"))?;

5
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -1175,6 +1175,7 @@ impl DQNTrainer {
}
for _step in 0..num_training_steps {
eprintln!("TRAIN_DIAG: step {_step}/{num_training_steps} — sampling");
let sample_start = std::time::Instant::now();
let (batch, vaccine_batch) = {
let agent = self.agent.read().await;
@@ -1187,7 +1188,9 @@ impl DQNTrainer {
let vb = buffer.sample(self.current_batch_size).ok();
(b, vb)
};
sample_total_us += sample_start.elapsed().as_micros() as u64;
let sample_us = sample_start.elapsed().as_micros();
eprintln!("TRAIN_DIAG: step {_step} — sample done in {sample_us}µs");
sample_total_us += sample_us as u64;
{
let mut agent = self.agent.write().await;