diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 1d6419527..2c385babf 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -335,6 +335,7 @@ pub(crate) struct FusedTrainingCtx { forward_child: Option, ddqn_child: Option, aux_child: Option, + adam_grad_child: Option, adam_child: Option, /// Composed parent -- single launch replays all children sequentially. parent_graph: Option, @@ -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.