diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index a89c2fb7d..6da91486c 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -2134,7 +2134,13 @@ impl GpuDqnTrainer { + self.config.branch_2_size + self.config.branch_3_size) * na; let d_val_norm = self.debug_buffer_norm_f32(self.d_value_logits_buf.raw_ptr(), b * na).unwrap_or(f32::NAN); - let d_adv_norm = self.debug_buffer_norm_f32(self.d_adv_logits_buf.raw_ptr(), b * tba).unwrap_or(f32::NAN); + // Guard: d_adv needs (b*tba+255)/256 blocks but grad_norm_partials only + // has grad_norm_blocks slots. Overflow corrupts adjacent GPU memory. + let d_adv_norm = if ((b * tba + 255) / 256) <= self.grad_norm_blocks { + self.debug_buffer_norm_f32(self.d_adv_logits_buf.raw_ptr(), b * tba).unwrap_or(f32::NAN) + } else { + f32::NAN // skip — would overflow partials buffer + }; let grad_norm = self.debug_buffer_norm_f32(self.ptrs.grad_buf, self.total_params).unwrap_or(f32::NAN); tracing::warn!( @@ -2160,10 +2166,11 @@ impl GpuDqnTrainer { // After graph replay: d_value/adv_logits contain BLENDED gradients, // MSE scratch contains MSE-only gradients, grad_buf contains backward output. + let adv_safe = ((b * tba + 255) / 256) <= self.grad_norm_blocks; let d_val_norm = self.debug_buffer_norm_f32(self.d_value_logits_buf.raw_ptr(), b * na).unwrap_or(f32::NAN); - let d_adv_norm = self.debug_buffer_norm_f32(self.d_adv_logits_buf.raw_ptr(), b * tba).unwrap_or(f32::NAN); + let d_adv_norm = if adv_safe { self.debug_buffer_norm_f32(self.d_adv_logits_buf.raw_ptr(), b * tba).unwrap_or(f32::NAN) } else { f32::NAN }; let mse_val_norm = self.debug_buffer_norm_f32(self.d_value_logits_mse.raw_ptr(), b * na).unwrap_or(f32::NAN); - let mse_adv_norm = self.debug_buffer_norm_f32(self.d_adv_logits_mse.raw_ptr(), b * tba).unwrap_or(f32::NAN); + let mse_adv_norm = if adv_safe { self.debug_buffer_norm_f32(self.d_adv_logits_mse.raw_ptr(), b * tba).unwrap_or(f32::NAN) } else { f32::NAN }; let grad_norm = self.debug_buffer_norm_f32(self.ptrs.grad_buf, self.total_params).unwrap_or(f32::NAN); tracing::warn!(