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

fix: IQN/ensemble trunk finalize no longer overwrites main grad_norm_buf

Browse Source 
The grad_norm_finalize kernel writes both f32 (sum-of-squares) and bf16
(L2 norm). The IQN and ensemble trunk gradient computations were passing
the MAIN grad_norm_buf for the bf16 output, overwriting it mid-pipeline.

While compute_grad_norm_outside_graph() overwrites it again before Adam,
this was still a correctness hazard. Now uses a dedicated 1-element bf16
scratch buffer (aux_norm_bf16_scratch) for IQN/ensemble finalize calls.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-09 20:48:30 +02:00
parent 0b3cb35ef4
commit 3964e111e7
1 changed files with 9 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -542,6 +542,8 @@ pub struct GpuDqnTrainer {
iqn_trunk_adam_step: i32,
/// IQN trunk grad norm scratch [1] (f32 accumulator).
iqn_trunk_grad_norm: CudaSlice<f32>,
/// Scratch bf16 [1] for IQN/ensemble finalize — prevents overwriting main grad_norm_buf.
aux_norm_bf16_scratch: CudaSlice<half::bf16>,
/// IQN trunk Adam step counter on device [1].
iqn_trunk_t_buf: CudaSlice<i32>,
/// Number of trunk parameters (w_s1 + b_s1 + w_s2 + b_s2).
@@ -1217,7 +1219,8 @@ impl GpuDqnTrainer {
let scratch_ptr = self.ptrs.iqn_trunk_m;
let partials_ptr = self.grad_norm_partials.raw_ptr();
let norm_ptr = self.ptrs.iqn_trunk_grad_norm;
let bf16_ptr = self.grad_norm_buf.raw_ptr();
// Use dedicated scratch for bf16 output — NOT main grad_norm_buf
let bf16_ptr = self.aux_norm_bf16_scratch.raw_ptr();
let n_i32 = trunk_grad_total as i32;
let blocks = ((trunk_grad_total + 255) / 256) as u32;
let nb = blocks as i32;
@@ -1235,7 +1238,7 @@ impl GpuDqnTrainer {
})
.map_err(|e| MLError::ModelError(format!("IQN trunk grad_norm phase1: {e}")))?;
}
// Phase 2: reduce → iqn_trunk_grad_norm
// Phase 2: reduce → iqn_trunk_grad_norm (f32) + scratch (bf16)
unsafe {
self.stream
.launch_builder(&self.grad_norm_finalize_kernel)
@@ -1477,7 +1480,7 @@ impl GpuDqnTrainer {
let scratch_ptr = self.ptrs.iqn_trunk_m;
let partials_ptr = self.grad_norm_partials.raw_ptr();
let norm_ptr = self.ptrs.iqn_trunk_grad_norm;
let bf16_ptr = self.grad_norm_buf.raw_ptr();
let bf16_ptr = self.aux_norm_bf16_scratch.raw_ptr();
let n_i32 = trunk_grad_total as i32;
let blocks = ((trunk_grad_total + 255) / 256) as u32;
let nb = blocks as i32;
@@ -2507,6 +2510,8 @@ impl GpuDqnTrainer {
let iqn_trunk_v = alloc_bf16(&stream, trunk_params, "iqn_trunk_v")?;
let iqn_trunk_grad_norm = stream.alloc_zeros::<f32>(1)
.map_err(|e| MLError::ModelError(format!("alloc iqn_trunk_grad_norm f32: {e}")))?;
// Scratch bf16 for IQN/ensemble trunk finalize — prevents overwriting main grad_norm_buf.
let aux_norm_bf16_scratch = alloc_bf16(&stream, 1, "aux_norm_bf16_scratch")?;
let iqn_trunk_t_buf = alloc_i32(&stream, 1, "iqn_trunk_t_buf")?;
let mut spec_u_s1 = alloc_bf16(&stream, config.shared_h1, "spec_u_s1")?;
let mut spec_v_s1 = alloc_bf16(&stream, config.state_dim, "spec_v_s1")?;
@@ -2915,6 +2920,7 @@ impl GpuDqnTrainer {
iqn_trunk_v,
iqn_trunk_adam_step: 0,
iqn_trunk_grad_norm,
aux_norm_bf16_scratch,
iqn_trunk_t_buf,
trunk_param_count: trunk_params,
states_buf,