fix: replace aliased SAXPY with in-place scale kernel (fixes __restrict__ NaN bug)
The C51+MSE gradient blending passed the same pointer as both y (output) and x (input) to dqn_saxpy_f32_kernel. That kernel has __restrict__ on both parameters, promising the compiler they don't alias. When they DO alias, NVCC's no-alias optimizations produce undefined behavior (NaN from register corruption). Added dqn_scale_f32_kernel (y[i] *= alpha) — single-pointer variant with no __restrict__, no aliasing possible. The two self-aliased SAXPY calls are replaced with scale+SAXPY pairs using distinct pointers. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -165,6 +165,28 @@ extern "C" __global__ void dqn_saxpy_f32_kernel(
|
||||
if (i < n) y[i] = y[i] + alpha * x[i];
|
||||
}
|
||||
|
||||
/* ══════════════════════════════════════════════════════════════════════
|
||||
* F32 IN-PLACE SCALE KERNEL
|
||||
*
|
||||
* y[i] *= alpha for i = 0..n-1
|
||||
*
|
||||
* Single-pointer variant — no aliasing possible, no __restrict__ needed.
|
||||
* Used for gradient scaling where SAXPY y+=α*x with x==y violates
|
||||
* __restrict__ and causes undefined behavior (NaN from register
|
||||
* corruption under NVCC's no-alias optimizations).
|
||||
*
|
||||
* Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
|
||||
* ══════════════════════════════════════════════════════════════════════ */
|
||||
|
||||
extern "C" __global__ void dqn_scale_f32_kernel(
|
||||
float* y,
|
||||
float alpha,
|
||||
int n
|
||||
) {
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i < n) y[i] *= alpha;
|
||||
}
|
||||
|
||||
/* ══════════════════════════════════════════════════════════════════════
|
||||
* CLIPPED SAXPY KERNEL
|
||||
*
|
||||
|
||||
@@ -497,6 +497,7 @@ pub struct GpuDqnTrainer {
|
||||
bf16_to_f32_kernel: CudaFunction,
|
||||
saxpy_kernel: CudaFunction,
|
||||
saxpy_f32_kernel: CudaFunction,
|
||||
scale_f32_kernel: CudaFunction,
|
||||
zero_kernel: CudaFunction,
|
||||
regime_scale_kernel: CudaFunction,
|
||||
shrink_perturb_kernel: CudaFunction,
|
||||
@@ -2166,7 +2167,7 @@ impl GpuDqnTrainer {
|
||||
// per array. Stack is set once in DQNTrainer::new() (64KB for all kernels).
|
||||
|
||||
// ── Compile 4 utility kernels (grad_norm, adam_update, BF16 converters) ─
|
||||
let (grad_norm_kernel, grad_norm_finalize_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clipped_saxpy_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, causal_intervene_kernel_fn, causal_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel) =
|
||||
let (grad_norm_kernel, grad_norm_finalize_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clipped_saxpy_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, causal_intervene_kernel_fn, causal_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel) =
|
||||
compile_training_kernels(&stream, &config)?;
|
||||
|
||||
// Separate grad_norm instance for non-graph launches (clip_grad_buf_inplace).
|
||||
@@ -2666,6 +2667,7 @@ impl GpuDqnTrainer {
|
||||
bf16_to_f32_kernel,
|
||||
saxpy_kernel,
|
||||
saxpy_f32_kernel,
|
||||
scale_f32_kernel,
|
||||
zero_kernel,
|
||||
regime_scale_kernel,
|
||||
shrink_perturb_kernel: shrink_perturb,
|
||||
@@ -4093,7 +4095,6 @@ impl GpuDqnTrainer {
|
||||
let b2 = self.config.branch_2_size;
|
||||
let n_val = (b * na) as i32;
|
||||
let n_adv = (b * (b0 + b1 + b2) * na) as i32;
|
||||
let scale_c51 = alpha - 1.0; // SAXPY: y += (α-1)*y → y *= α
|
||||
let scale_mse = 1.0 - alpha;
|
||||
let val_blocks = ((b * na + 255) / 256) as u32;
|
||||
let adv_blocks = ((b * (b0 + b1 + b2) * na + 255) / 256) as u32;
|
||||
@@ -4107,22 +4108,22 @@ impl GpuDqnTrainer {
|
||||
let adv_mse_ptr = self.d_adv_logits_mse.raw_ptr();
|
||||
|
||||
unsafe {
|
||||
// d_value += (α-1) * d_value → d_value *= α (f32 SAXPY)
|
||||
self.stream.launch_builder(&self.saxpy_f32_kernel)
|
||||
.arg(&val_ptr).arg(&val_ptr)
|
||||
.arg(&scale_c51).arg(&n_val)
|
||||
.launch(cfg_val).map_err(|e| MLError::ModelError(format!("blend c51 val: {e}")))?;
|
||||
// d_value += (1-α) * mse_scratch
|
||||
// d_value *= α (in-place scale — single pointer, no aliasing)
|
||||
self.stream.launch_builder(&self.scale_f32_kernel)
|
||||
.arg(&val_ptr)
|
||||
.arg(&alpha).arg(&n_val)
|
||||
.launch(cfg_val).map_err(|e| MLError::ModelError(format!("scale c51 val: {e}")))?;
|
||||
// d_value += (1-α) * mse_scratch (SAXPY with distinct pointers — safe)
|
||||
self.stream.launch_builder(&self.saxpy_f32_kernel)
|
||||
.arg(&val_ptr).arg(&val_mse_ptr)
|
||||
.arg(&scale_mse).arg(&n_val)
|
||||
.launch(cfg_val).map_err(|e| MLError::ModelError(format!("blend mse val: {e}")))?;
|
||||
// d_adv += (α-1) * d_adv → d_adv *= α
|
||||
self.stream.launch_builder(&self.saxpy_f32_kernel)
|
||||
.arg(&adv_ptr).arg(&adv_ptr)
|
||||
.arg(&scale_c51).arg(&n_adv)
|
||||
.launch(cfg_adv).map_err(|e| MLError::ModelError(format!("blend c51 adv: {e}")))?;
|
||||
// d_adv += (1-α) * mse_scratch
|
||||
// d_adv *= α (in-place scale — single pointer, no aliasing)
|
||||
self.stream.launch_builder(&self.scale_f32_kernel)
|
||||
.arg(&adv_ptr)
|
||||
.arg(&alpha).arg(&n_adv)
|
||||
.launch(cfg_adv).map_err(|e| MLError::ModelError(format!("scale c51 adv: {e}")))?;
|
||||
// d_adv += (1-α) * mse_scratch (SAXPY with distinct pointers — safe)
|
||||
self.stream.launch_builder(&self.saxpy_f32_kernel)
|
||||
.arg(&adv_ptr).arg(&adv_mse_ptr)
|
||||
.arg(&scale_mse).arg(&n_adv)
|
||||
@@ -5567,7 +5568,7 @@ impl GpuDqnTrainer {
|
||||
fn compile_training_kernels(
|
||||
stream: &Arc<CudaStream>,
|
||||
config: &GpuDqnTrainConfig,
|
||||
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
|
||||
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
|
||||
info!(
|
||||
state_dim = config.state_dim,
|
||||
total_params = compute_total_params(config),
|
||||
@@ -5593,6 +5594,8 @@ fn compile_training_kernels(
|
||||
.map_err(|e| MLError::ModelError(format!("dqn_saxpy_kernel load: {e}")))?;
|
||||
let saxpy_f32 = module.load_function("dqn_saxpy_f32_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("dqn_saxpy_f32_kernel load: {e}")))?;
|
||||
let scale_f32 = module.load_function("dqn_scale_f32_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("dqn_scale_f32_kernel load: {e}")))?;
|
||||
let zero = module.load_function("dqn_zero_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("dqn_zero_kernel load: {e}")))?;
|
||||
let regime_scale = module.load_function("dqn_regime_scale_kernel")
|
||||
@@ -5635,8 +5638,8 @@ fn compile_training_kernels(
|
||||
let her_inplace = module.load_function("her_inplace_relabel")
|
||||
.map_err(|e| MLError::ModelError(format!("her_inplace_relabel load: {e}")))?;
|
||||
|
||||
info!("GpuDqnTrainer: 27 utility kernels loaded from precompiled cubin");
|
||||
Ok((grad_norm, grad_norm_finalize, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clipped_saxpy, clip_grad, pad_states, saxpy_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, causal_intervene, causal_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace))
|
||||
info!("GpuDqnTrainer: 28 utility kernels loaded from precompiled cubin");
|
||||
Ok((grad_norm, grad_norm_finalize, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clipped_saxpy, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, causal_intervene, causal_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace))
|
||||
}
|
||||
|
||||
/// Load the standalone Polyak EMA kernel from precompiled cubin.
|
||||
|
||||
Reference in New Issue
Block a user