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

fix: remove in-graph sub-phase events (CUDA_ERROR_INVALID_VALUE)

Browse Source 
Events recorded during CUDA graph capture cannot be queried with
cuEventElapsedTime after graph replay — the graph creates internal
copies. Reverted to original mega-graph capture structure.

Sub-graph timing fields kept in PhaseEvents for future use with
split-graph diagnostic mode or nsys profiling.

Also kept submit_loss_and_grad_ops() extraction and pub(crate) visibility
on launch_cublas_forward/backward for future per-phase graph splitting.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-11 02:41:41 +02:00
parent ef5be1b7a3
commit 3b2bf8d993
1 changed files with 4 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1039,21 +1039,10 @@ impl FusedTrainingCtx {
let per_update = elapsed(pe.per_update_start, pe.per_update_end);
let total = upload + fwd_bwd + adam + per_update;
// Sub-graph breakdown (inside mega-graph replay)
let spectral = elapsed(pe.fwd_bwd_start, pe.spectral_end);
let forward = elapsed(pe.spectral_end, pe.forward_end);
let loss = elapsed(pe.forward_end, pe.loss_end);
let backward = elapsed(pe.loss_end, pe.backward_end);
let aux = elapsed(pe.backward_end, pe.fwd_bwd_end);
tracing::info!(
"GPU phase timing ({} batches, last batch): upload={:.1}ms fwd_bwd={:.1}ms adam={:.1}ms per_update={:.1}ms total={:.1}ms",
self.phase_batch_count, upload, fwd_bwd, adam, per_update, total,
);
tracing::info!(
" fwd_bwd breakdown: spectral={:.1}ms forward={:.1}ms loss={:.1}ms backward={:.1}ms aux={:.1}ms",
spectral, forward, loss, backward, aux,
);
self.phase_batch_count = 0;
}
@@ -1410,39 +1399,15 @@ impl FusedTrainingCtx {
anyhow::bail!("cuStreamBeginCapture_v2 (mega) failed: {begin_result:?}");
}
// Record sub-phase events inside the graph capture — they replay every step.
let cu_stream = self.stream.cu_stream();
// Phase 1: Spectral norm
let spectral_result = self.trainer.apply_spectral_norm(
&self.online_dueling, &self.online_branching,
).map_err(|e| anyhow::anyhow!("mega capture spectral: {e}"));
if let Some(ref pe) = self.phase_events {
PhaseEvents::record(pe.spectral_end, cu_stream);
}
// Phase 2: Forward + loss + backward
let forward_result: Result<()> = if spectral_result.is_ok() {
// Forward GEMMs
self.trainer.launch_cublas_forward()
.map_err(|e| anyhow::anyhow!("mega capture fwd: {e}"))?;
if let Some(ref pe) = self.phase_events {
PhaseEvents::record(pe.forward_end, cu_stream);
}
// Loss + gradient computation
self.trainer.submit_loss_and_grad_ops()
.map_err(|e| anyhow::anyhow!("mega capture loss: {e}"))?;
if let Some(ref pe) = self.phase_events {
PhaseEvents::record(pe.loss_end, cu_stream);
}
// Backward GEMMs
self.trainer.launch_cublas_backward()
.map_err(|e| anyhow::anyhow!("mega capture bwd: {e}"))?;
if let Some(ref pe) = self.phase_events {
PhaseEvents::record(pe.backward_end, cu_stream);
}
Ok(())
// Phase 2: Forward + backward (train_step_gpu submits forward ops)
let forward_result = if spectral_result.is_ok() {
self.trainer.submit_forward_ops_main()
.map_err(|e| anyhow::anyhow!("mega capture forward: {e}"))
} else {
Ok(())
};