fix(critical): IQN d_h_s2 is f32 not bf16 — remove bogus bf16_to_f32 cast

The IQN backward kernel writes f32 via atomicAdd into d_h_s2_buf
(allocated as CudaSlice<f32>). But apply_iqn_trunk_gradient called
bf16_to_f32_kernel on this buffer, reinterpreting f32 bits as bf16.
This produced garbage/NaN that poisoned the trunk gradient on H100.

Replace with direct f32 DtoD copy — no cast needed.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-09 19:59:26 +02:00
parent 5b4561fc27
commit 8cbabcef5e

View File

@@ -1111,26 +1111,12 @@ impl GpuDqnTrainer {
}
}
// ── 2. Cast IQN d_h_s2 (bf16) → bw_d_h_s2 (f32) ─────────────────
// IQN head produces bf16 gradient. Cast to f32 for the backward scratch.
// ── 2. DtoD copy IQN d_h_s2 (f32) → bw_d_h_s2 (f32) ────────────
// IQN d_h_s2_buf is f32 (atomicAdd in backward kernel). No cast needed.
{
let src = iqn_d_h_s2_ptr;
let dst = self.ptrs.bw_d_h_s2;
let n_elems = (b * sh2) as i32;
let blocks = ((b * sh2 + 255) / 256) as u32;
unsafe {
self.stream
.launch_builder(&self.bf16_to_f32_kernel)
.arg(&src)
.arg(&dst)
.arg(&n_elems)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("IQN d_h_s2 bf16→f32: {e}")))?;
}
let n_bytes = b * sh2 * std::mem::size_of::<f32>();
dtod_copy(self.ptrs.bw_d_h_s2, iqn_d_h_s2_ptr, n_bytes, &self.stream, 0, "iqn_d_h_s2_copy")
.map_err(|e| MLError::ModelError(format!("IQN d_h_s2 f32 DtoD: {e}")))?;
}
// ── 3. ReLU mask: bw_d_h_s2 *= (save_h_s2 > 0) (f32 dx, bf16 act) ──