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

fix: async adam_step HtoD + batch_size=0 auto — root cause of >128 hang

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-02 08:19:56 +02:00
parent c4a7b6ee3c
commit 68e83bef0e
2 changed files with 28 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
config/training/dqn-production.toml
View File

@@ -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

36
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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::<i32>(),
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::<i32>(),
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<FusedTrainScalars, MLError> {
// ── 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::<i32>(),
self.stream.cu_stream(),
);
}
// ── Execute training step (split graph: forward then adam) ──
if self.graph_forward.is_none() {