diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 71e43d12f..d082d4494 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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, + /// Scratch bf16 [1] for IQN/ensemble finalize — prevents overwriting main grad_norm_buf. + aux_norm_bf16_scratch: CudaSlice, /// IQN trunk Adam step counter on device [1]. iqn_trunk_t_buf: CudaSlice, /// 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::(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,