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

perf: add per-child graph CUDA event timing breakdown

Browse Source 
Records cuEventRecord between each of the 5 child graph launches.
At epoch end, cuEventElapsedTime reports per-child GPU time:
spectral, forward, ddqn, aux, adam.

This tells us exactly which child dominates the 4s/step graph replay
cost, informing Phase 2 multi-stream parallelism priorities.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-18 14:59:39 +02:00
parent 61e4d8730c
commit 4295cbde79
1 changed files with 23 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -960,9 +960,21 @@ impl FusedTrainingCtx {
// 5 launches instead of 1, ~10µs overhead total, negligible vs step time.
if self.spectral_child.is_some() {
self.spectral_child.as_ref().unwrap().launch(cu_stream)?;
if let Some(ref pe) = self.phase_events {
PhaseEvents::record(pe.spectral_end, cu_stream);
}
self.forward_child.as_ref().unwrap().launch(cu_stream)?;
if let Some(ref pe) = self.phase_events {
PhaseEvents::record(pe.forward_end, cu_stream);
}
self.ddqn_child.as_ref().unwrap().launch(cu_stream)?;
if let Some(ref pe) = self.phase_events {
PhaseEvents::record(pe.loss_end, cu_stream);
}
self.aux_child.as_ref().unwrap().launch(cu_stream)?;
if let Some(ref pe) = self.phase_events {
PhaseEvents::record(pe.backward_end, cu_stream);
}
self.adam_child.as_ref().unwrap().launch(cu_stream)?;
} else {
// Ungraphed step 0: run everything, then capture.
@@ -1162,10 +1174,21 @@ impl FusedTrainingCtx {
let per_update = elapsed(pe.per_update_start, pe.per_update_end);
let total = upload + fwd_bwd + adam + per_update;
// Per-child graph breakdown (events recorded between launches)
let spectral_ms = elapsed(pe.fwd_bwd_start, pe.spectral_end);
let forward_ms = elapsed(pe.spectral_end, pe.forward_end);
let ddqn_ms = elapsed(pe.forward_end, pe.loss_end);
let aux_ms = elapsed(pe.loss_end, pe.backward_end);
let adam_child_ms = elapsed(pe.backward_end, pe.adam_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!(
"GPU child graph breakdown: spectral={:.1}ms forward={:.1}ms ddqn={:.1}ms aux={:.1}ms adam={:.1}ms",
spectral_ms, forward_ms, ddqn_ms, aux_ms, adam_child_ms,
);
self.phase_batch_count = 0;
}