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

fix: add adam_child_end event — previous timing used wrong event (adam_end from old phase system)

Browse Source 
The child graph breakdown showed adam=-1.0ms because it was measuring
elapsed(backward_end, adam_end) where adam_end is from the OLD phase
timing system, not the adam child graph. Added proper adam_child_end
event recorded after the adam_child launch.

This will reveal the actual adam_child GPU time — currently suspected
to be 3117ms (73% of step time) based on total - measured children.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-18 16:00:06 +02:00
parent 4295cbde79
commit 8ebb65a942
1 changed files with 10 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -175,11 +175,12 @@ struct PhaseEvents {
adam_end: cuda_sys::CUevent,
per_update_start: cuda_sys::CUevent,
per_update_end: cuda_sys::CUevent,
// Sub-graph events (inside mega-graph, replayed on every step)
// Per-child-graph timing events (recorded between child launches)
spectral_end: cuda_sys::CUevent,
forward_end: cuda_sys::CUevent,
loss_end: cuda_sys::CUevent,
backward_end: cuda_sys::CUevent,
loss_end: cuda_sys::CUevent, // after ddqn_child
backward_end: cuda_sys::CUevent, // after aux_child
adam_child_end: cuda_sys::CUevent,
}
impl PhaseEvents {
@@ -208,6 +209,7 @@ impl Drop for PhaseEvents {
let _ = cuda_sys::cuEventDestroy_v2(self.forward_end);
let _ = cuda_sys::cuEventDestroy_v2(self.loss_end);
let _ = cuda_sys::cuEventDestroy_v2(self.backward_end);
let _ = cuda_sys::cuEventDestroy_v2(self.adam_child_end);
}
}
}
@@ -775,6 +777,7 @@ impl FusedTrainingCtx {
forward_end: create_event()?,
loss_end: create_event()?,
backward_end: create_event()?,
adam_child_end: create_event()?,
})
})() {
Ok(pe) => Some(pe),
@@ -976,6 +979,9 @@ impl FusedTrainingCtx {
PhaseEvents::record(pe.backward_end, cu_stream);
}
self.adam_child.as_ref().unwrap().launch(cu_stream)?;
if let Some(ref pe) = self.phase_events {
PhaseEvents::record(pe.adam_child_end, cu_stream);
}
} else {
// Ungraphed step 0: run everything, then capture.
// Ensure params are initialized for the first step.
@@ -1179,7 +1185,7 @@ impl FusedTrainingCtx {
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);
let adam_child_ms = elapsed(pe.backward_end, pe.adam_child_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",