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() {