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

debug: split adam into adam_grad + adam_update to isolate 3.1s bottleneck

Browse Source 
adam_child = 3113ms even after CUfunction fix. Split into:
- adam_grad: mamba2_backward + step_mamba2_adam + pruning + grad_norm
- adam_update: Adam kernel + unflatten + ISV

If adam_grad is fast and adam_update is slow, the bottleneck is Adam/unflatten.
If adam_grad is slow, grad_norm_kernel has an issue.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-18 19:13:32 +02:00
parent b4cf279b73
commit 7b0280a187
1 changed files with 20 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -335,6 +335,7 @@ pub(crate) struct FusedTrainingCtx {
forward_child: Option<ChildGraph>,
ddqn_child: Option<ChildGraph>,
aux_child: Option<ChildGraph>,
adam_grad_child: Option<ChildGraph>,
adam_child: Option<ChildGraph>,
/// Composed parent -- single launch replays all children sequentially.
parent_graph: Option<ParentGraph>,
@@ -839,6 +840,7 @@ impl FusedTrainingCtx {
forward_child: None,
ddqn_child: None,
aux_child: None,
adam_grad_child: None,
adam_child: None,
parent_graph: None,
eval_forward_exec: None,
@@ -872,6 +874,7 @@ impl FusedTrainingCtx {
self.forward_child = None;
self.ddqn_child = None;
self.aux_child = None;
self.adam_grad_child = None;
self.adam_child = None;
// Phase 2 aux streams, workspaces, and events persist across folds.
// They are reusable GPU resources — no reset or reallocation needed.
@@ -1027,6 +1030,15 @@ impl FusedTrainingCtx {
if let Some(ref pe) = self.phase_events {
PhaseEvents::record(pe.backward_end, cu_stream);
}
// DEBUG: adam split into grad + update to isolate 3.1s bottleneck
if let Some(ref g) = self.adam_grad_child {
g.launch(cu_stream)?;
}
if let Some(ref pe) = self.phase_events {
// Reuse loss_end event (not used after ddqn since we have backward_end for aux)
// Actually add a new usage: backward_end = after aux, we need a new event for adam_grad
// For now, measure adam_update only by using adam_child_end after adam launch
}
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);
@@ -1692,14 +1704,18 @@ impl FusedTrainingCtx {
s.submit_aux_ops(agent, gpu_batch)
})?;
// Capture adam child (mamba2 bwd + pruning + grad_norm + adam + ISV).
let adam = self.capture_child_graph("adam", |s| {
// DEBUG: Split adam into two children to isolate the 3.1s bottleneck.
// adam_grad: mamba2 bwd + pruning + grad_norm
// adam_update: Adam + unflatten + ISV
let adam_grad = self.capture_child_graph("adam_grad", |s| {
s.mamba2_backward(s.batch_size)?;
s.step_mamba2_adam()?;
s.trainer.apply_pruning_mask()
.map_err(|e| anyhow::anyhow!("{e}"))?;
s.trainer.compute_grad_norm_for_adam()
.map_err(|e| anyhow::anyhow!("{e}"))?;
.map_err(|e| anyhow::anyhow!("{e}"))
})?;
let adam = self.capture_child_graph("adam_update", |s| {
s.trainer.submit_adam_ops(&s.online_dueling, &s.online_branching)
.map_err(|e| anyhow::anyhow!("{e}"))?;
s.trainer.update_isv_signals()
@@ -1711,6 +1727,7 @@ impl FusedTrainingCtx {
self.forward_child = Some(forward);
self.ddqn_child = Some(ddqn);
self.aux_child = Some(aux);
self.adam_grad_child = Some(adam_grad);
self.adam_child = Some(adam);
// Children launched individually -- parent composition hangs on Hopper.