From 8cbabcef5eee52694e86c3e992a218958d2a458a Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 9 Apr 2026 19:59:26 +0200 Subject: [PATCH] =?UTF-8?q?fix(critical):=20IQN=20d=5Fh=5Fs2=20is=20f32=20?= =?UTF-8?q?not=20bf16=20=E2=80=94=20remove=20bogus=20bf16=5Fto=5Ff32=20cas?= =?UTF-8?q?t?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The IQN backward kernel writes f32 via atomicAdd into d_h_s2_buf (allocated as CudaSlice). 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) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 24 ++++--------------- 1 file changed, 5 insertions(+), 19 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 932c6684d..ffc979061 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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::(); + 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) ──