From dd075ce02a2c0c16eadbe9244db7b2d8a5a94aa8 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 9 Mar 2026 22:43:34 +0100 Subject: [PATCH] fix(gpu): cap auto-scaled episode count at GPU buffer limit (4096) optimal_n_episodes() returns 8192 on H100 (132 SMs) but the GPU experience collector pre-allocates buffers for MAX_EPISODES_LIMIT=4096. Exceeding that caused "Episode count exceeds allocated buffer size" and silent fallback to CPU. Now all 3 auto-scaling sites clamp at 4096. Co-Authored-By: Claude Opus 4.6 --- crates/ml/src/trainers/dqn/trainer.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index 5ccc5c972..e7093d4b0 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/crates/ml/src/trainers/dqn/trainer.rs @@ -1850,10 +1850,11 @@ impl DQNTrainer { let configured = self.hyperparams.gpu_n_episodes; if configured >= 128 { match detect_gpu_hardware() { + // Cap at GPU experience collector's buffer limit (4096) Ok(hw) => configured.max(hw.optimal_n_episodes( state_dim, self.hyperparams.gpu_timesteps_per_episode, - )), + )).min(4096), Err(_) => configured, } } else { @@ -1892,10 +1893,11 @@ impl DQNTrainer { let configured = self.hyperparams.gpu_n_episodes; if configured >= 128 { match detect_gpu_hardware() { + // Cap at GPU experience collector's buffer limit (4096) Ok(hw) => configured.max(hw.optimal_n_episodes( state_dim, self.hyperparams.gpu_timesteps_per_episode, - )), + )).min(4096), Err(_) => configured, } } else { @@ -2001,7 +2003,8 @@ impl DQNTrainer { aligned_sd, self.hyperparams.gpu_timesteps_per_episode, ); - let chosen = configured.max(optimal); + // Cap at GPU experience collector's buffer limit (4096) + let chosen = configured.max(optimal).min(4096); if chosen != configured { info!( "GPU auto-scaled n_episodes: {} → {} (SMs={}, VRAM={:.0}MB)",