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:
@@ -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) ──
|
||||
|
||||
Reference in New Issue
Block a user