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

fix: guard diagnostic buffer overflow — d_adv needed 2304 blocks, partials had 1137

Browse Source 
debug_buffer_norm_f32 reuses grad_norm_partials (sized for total_params)
as scratch. d_adv_logits with batch_size*tba=589K elements needs 2304
blocks but partials only has 1137 — GPU memory overflow corrupting
adjacent allocations. On H100 this likely trashed grad_norm_buf,
explaining the persistent grad_norm=0.000000. Guard added to skip
oversized norms.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-13 23:23:03 +02:00
parent 4159d201a0
commit fee5a7cb3a
1 changed files with 10 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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