From 68e83bef0ecb6f9c3d2280cacb7f1ca0387bee4f Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 2 Apr 2026 08:19:56 +0200 Subject: [PATCH] =?UTF-8?q?fix:=20async=20adam=5Fstep=20HtoD=20+=20batch?= =?UTF-8?q?=5Fsize=3D0=20auto=20=E2=80=94=20root=20cause=20of=20>128=20han?= =?UTF-8?q?g?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The batch_size >128 hang was caused by cudarc's synchronous memcpy_htod for the adam_step counter. At batch_size=128 the sync completes fast enough, but at 509+ the pipeline drain from the sync interacts with CUDA Graph replay timing and deadlocks the stream. Replaced all 3 memcpy_htod(&[self.adam_step]) calls with raw cuMemcpyHtoDAsync_v2 — zero pipeline drain, zero CPU sync. Production TOML set to batch_size=0 (AutoBatchSizer drives it). AutoBatchSizer caps at 8192 for RL training. Co-Authored-By: Claude Opus 4.6 (1M context) --- config/training/dqn-production.toml | 2 +- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 36 ++++++++++++++----- 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/config/training/dqn-production.toml b/config/training/dqn-production.toml index 36ccf4d75..23ea42ce9 100644 --- a/config/training/dqn-production.toml +++ b/config/training/dqn-production.toml @@ -8,7 +8,7 @@ [training] epochs = 200 -batch_size = 128 +batch_size = 0 # 0 = auto from VRAM via AutoBatchSizer learning_rate = 0.0001 gamma = 0.99 weight_decay = 0.0001 diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 1ad7222a5..346e0312a 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -2866,9 +2866,17 @@ impl GpuDqnTrainer { // Adam is deferred to replay_adam_and_readback() so the caller can // inject IQN/attention/ensemble gradients into grad_buf first. self.adam_step += 1; - self.stream - .memcpy_htod(&[self.adam_step], &mut self.t_buf) - .map_err(|e| MLError::ModelError(format!("HtoD adam_step: {e}")))?; + // Async HtoD for adam_step counter — no pipeline drain. + // t_buf address is captured by the graph; updating its value before + // replay is correct (graph captures addresses, not values). + unsafe { + cudarc::driver::sys::cuMemcpyHtoDAsync_v2( + self.t_buf.raw_ptr(), + (&self.adam_step as *const i32).cast(), + std::mem::size_of::(), + self.stream.cu_stream(), + ); + } if self.graph_forward.is_none() { self.capture_training_graphs(online_dueling, online_branching)?; } @@ -3039,9 +3047,14 @@ impl GpuDqnTrainer { // ── Update Adam step counter on device (OUTSIDE graph) ─────── self.adam_step += 1; - self.stream - .memcpy_htod(&[self.adam_step], &mut self.t_buf) - .map_err(|e| MLError::ModelError(format!("HtoD adam_step: {e}")))?; + unsafe { + cudarc::driver::sys::cuMemcpyHtoDAsync_v2( + self.t_buf.raw_ptr(), + (&self.adam_step as *const i32).cast(), + std::mem::size_of::(), + self.stream.cu_stream(), + ); + } // ── Execute training step (split graph: forward then adam) ── if self.graph_forward.is_none() { @@ -3431,9 +3444,14 @@ impl GpuDqnTrainer { ) -> Result { // ── Update Adam step counter on device (OUTSIDE graph) ─────── self.adam_step += 1; - self.stream - .memcpy_htod(&[self.adam_step], &mut self.t_buf) - .map_err(|e| MLError::ModelError(format!("HtoD adam_step: {e}")))?; + unsafe { + cudarc::driver::sys::cuMemcpyHtoDAsync_v2( + self.t_buf.raw_ptr(), + (&self.adam_step as *const i32).cast(), + std::mem::size_of::(), + self.stream.cu_stream(), + ); + } // ── Execute training step (split graph: forward then adam) ── if self.graph_forward.is_none() {