diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index d082d4494..5b4aae5be 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -3159,11 +3159,10 @@ impl GpuDqnTrainer { self.stream.cu_stream(), ); } - if self.graph_forward.is_none() { - self.capture_training_graphs(online_dueling, online_branching)?; - } + // DIAG: skip ALL graph capture — run forward+backward ungraphed every step. + // If this fixes H100 collapse, the graph recapture path is the root cause. self.update_stochastic_depth_mask()?; - self.replay_forward()?; + self.submit_forward_ops_main()?; Ok(FusedTrainScalars { total_loss: 0.0, grad_norm: 0.0 }) } @@ -3813,6 +3812,7 @@ impl GpuDqnTrainer { /// Replay graph_adam only (grad_norm → Adam → unflatten). /// Must be called AFTER replay_forward() and any auxiliary gradient injection. pub fn replay_adam(&self) -> Result<(), MLError> { + // Keep graph_adam replay — it's captured once and never invalidated if let Some(ref graph) = self.graph_adam { graph.0.launch().map_err(|e| { MLError::ModelError(format!("CUDA graph_adam replay: {e}"))