diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 056252ac4..67db881e8 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -892,6 +892,33 @@ pub struct GpuDqnTrainer { /// v8: GPU-resident aux weight scalar [1] — enables graph capture of aux GEMM. pub(crate) exposure_aux_weight_gpu: CudaSlice, + // ── Separate Adam optimizer for exposure aux loss (breaks i%3 degeneracy) ── + // The aux loss gets its OWN Adam state and updates exposure output weights + // INDEPENDENTLY from the main C51 optimizer. No shared gradient buffer, + // no gradient clipping competition. + /// First moment [aux_param_count] — w_b0out + b_b0out. + aux_adam_m: CudaSlice, + /// Second moment [aux_param_count]. + aux_adam_v: CudaSlice, + /// Step counter [1]. + aux_adam_t: CudaSlice, + /// Separate gradient accumulator [aux_param_count]. + aux_grad_buf: CudaSlice, + /// Grad norm f32 accumulator [1]. + aux_grad_norm_f32: CudaSlice, + /// Grad norm bf16 output [1]. + aux_grad_norm_bf16: CudaSlice, + /// Per-block partial sums for grad norm reduction. + aux_grad_norm_partials: CudaSlice, + /// Number of exposure output parameters (w_b0out + b_b0out). + aux_param_count: usize, + /// Host-side step counter for aux Adam. + aux_adam_step: i32, + /// Standalone Adam kernel for aux (not captured in any CUDA graph). + aux_adam_update_kernel: CudaFunction, + /// Standalone grad_norm_finalize for aux (not captured in any CUDA graph). + aux_grad_norm_finalize_kernel: CudaFunction, + /// v8: PopArt reward normalization kernel (running mean/variance). popart_normalize_kernel: CudaFunction, /// v8: PopArt running mean [1]. @@ -926,6 +953,15 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("reset iqn t: {e}")))?; self.iqn_trunk_adam_step = 0; + // Reset exposure aux Adam state (separate optimizer for exposure output weights). + self.stream.memset_zeros(&mut self.aux_adam_m) + .map_err(|e| MLError::ModelError(format!("reset aux_adam_m: {e}")))?; + self.stream.memset_zeros(&mut self.aux_adam_v) + .map_err(|e| MLError::ModelError(format!("reset aux_adam_v: {e}")))?; + self.stream.memset_zeros(&mut self.aux_adam_t) + .map_err(|e| MLError::ModelError(format!("reset aux_adam_t: {e}")))?; + self.aux_adam_step = 0; + tracing::info!("Adam optimizer state reset for new fold"); Ok(()) } @@ -2001,12 +2037,13 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("f32_to_bf16 exposure_aux: {e}")))?; } - // Phase 3: Backward FC layer -- accumulate weight + bias gradients - let sizes = compute_param_sizes(&self.config); - let goff_w_b0out: usize = sizes[..10].iter().sum::(); - let goff_b_b0out: usize = sizes[..11].iter().sum::(); + // Phase 3: Zero aux_grad_buf before backward (backward_fc_layer beta=1.0 accumulates) + self.stream.memset_zeros(&mut self.aux_grad_buf) + .map_err(|e| MLError::ModelError(format!("zero aux_grad_buf: {e}")))?; - let grad_base = self.ptrs.grad_buf; + // Phase 4: Backward FC layer -> aux_grad_buf (NOT shared grad_buf) + let sizes = compute_param_sizes(&self.config); + let aux_grad_ptr = self.aux_grad_buf.raw_ptr(); let h_b0_ptr = self.save_h_b0.raw_ptr(); self.cublas_backward.backward_fc_layer( @@ -2014,14 +2051,117 @@ impl GpuDqnTrainer { staging_ptr, h_b0_ptr, 0u64, - grad_base + (goff_w_b0out as u64) * 4, - grad_base + (goff_b_b0out as u64) * 4, + aux_grad_ptr, // dW at offset 0 + aux_grad_ptr + (sizes[10] as u64) * 4, // db at offset w_b0out_size 0u64, b0 * na, ah, b, )?; + // Phase 5: Compute grad norm of aux_grad_buf + let aux_pc = self.aux_param_count; + let aux_blocks = ((aux_pc + 255) / 256) as u32; + let aux_partials_ptr = self.aux_grad_norm_partials.raw_ptr(); + let aux_norm_f32_ptr = self.aux_grad_norm_f32.raw_ptr(); + let aux_norm_bf16_ptr = self.aux_grad_norm_bf16.raw_ptr(); + let aux_n_i32 = aux_pc as i32; + let aux_nb = aux_blocks as i32; + // Phase 1: per-block partials (reuse grad_norm_standalone — not captured in any graph) + unsafe { + self.stream + .launch_builder(&self.grad_norm_standalone) + .arg(&aux_grad_ptr) + .arg(&aux_partials_ptr) + .arg(&aux_n_i32) + .launch(LaunchConfig { + grid_dim: (aux_blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("aux grad_norm phase1: {e}")))?; + } + // Phase 2: reduce partials → f32 sum-of-squares + bf16 L2 norm + unsafe { + self.stream + .launch_builder(&self.aux_grad_norm_finalize_kernel) + .arg(&aux_partials_ptr) + .arg(&aux_norm_f32_ptr) + .arg(&aux_norm_bf16_ptr) + .arg(&aux_nb) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("aux grad_norm phase2: {e}")))?; + } + + // Phase 6: Separate Adam update on exposure output weights only + self.aux_adam_step += 1; + unsafe { + cudarc::driver::sys::cuMemcpyHtoDAsync_v2( + self.aux_adam_t.raw_ptr(), + (&self.aux_adam_step as *const i32).cast(), + std::mem::size_of::(), + self.stream.cu_stream(), + ); + } + + let goff_w_b0out: usize = sizes[..10].iter().sum::(); + let aux_params_ptr = self.ptrs.params_f32_ptr + (goff_w_b0out as u64) * 4; + let aux_lr = 0.001_f32; // higher LR for aux (independent from main) + let aux_beta1 = 0.9_f32; + let aux_beta2 = 0.999_f32; + let aux_eps = 1e-8_f32; + let aux_wd = 0.0_f32; // no weight decay for aux + let aux_max_gn = 1.0_f32; // already small gradients + let aux_m_ptr = self.aux_adam_m.raw_ptr(); + let aux_v_ptr = self.aux_adam_v.raw_ptr(); + let aux_t_ptr = self.aux_adam_t.raw_ptr(); + + unsafe { + self.stream + .launch_builder(&self.aux_adam_update_kernel) + .arg(&aux_params_ptr) + .arg(&aux_grad_ptr) + .arg(&aux_m_ptr) + .arg(&aux_v_ptr) + .arg(&aux_norm_f32_ptr) + .arg(&aux_lr) + .arg(&aux_beta1) + .arg(&aux_beta2) + .arg(&aux_eps) + .arg(&aux_wd) + .arg(&aux_max_gn) + .arg(&aux_t_ptr) + .arg(&aux_n_i32) + .launch(LaunchConfig { + grid_dim: (aux_blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("aux Adam update: {e}")))?; + } + + // Phase 7: Cast updated f32 weights → bf16 shadow + let bf16_w_ptr = self.ptrs.params_buf + (goff_w_b0out as u64) * 2; + let cast_count = aux_pc as i32; + let cast_blk = aux_blocks; + unsafe { + self.stream + .launch_builder(&self.f32_to_bf16_kernel) + .arg(&aux_params_ptr) + .arg(&bf16_w_ptr) + .arg(&cast_count) + .launch(LaunchConfig { + grid_dim: (cast_blk, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("aux f32→bf16 cast: {e}")))?; + } + Ok(()) } @@ -2966,6 +3106,38 @@ impl GpuDqnTrainer { let exposure_aux_weight_gpu = stream.alloc_zeros::(1) .map_err(|e| MLError::ModelError(format!("alloc exposure_aux_weight_gpu: {e}")))?; + // Separate Adam optimizer for exposure aux loss (breaks i%3 degeneracy) + let sizes_for_aux = compute_param_sizes(&config); + let aux_param_count = sizes_for_aux[10] + sizes_for_aux[11]; // w_b0out + b_b0out + let aux_adam_m = stream.alloc_zeros::(aux_param_count) + .map_err(|e| MLError::ModelError(format!("alloc aux_adam_m: {e}")))?; + let aux_adam_v = stream.alloc_zeros::(aux_param_count) + .map_err(|e| MLError::ModelError(format!("alloc aux_adam_v: {e}")))?; + let aux_adam_t = stream.alloc_zeros::(1) + .map_err(|e| MLError::ModelError(format!("alloc aux_adam_t: {e}")))?; + let aux_grad_buf = stream.alloc_zeros::(aux_param_count) + .map_err(|e| MLError::ModelError(format!("alloc aux_grad_buf: {e}")))?; + let aux_grad_norm_f32 = stream.alloc_zeros::(1) + .map_err(|e| MLError::ModelError(format!("alloc aux_grad_norm_f32: {e}")))?; + let aux_grad_norm_bf16 = stream.alloc_zeros::(1) + .map_err(|e| MLError::ModelError(format!("alloc aux_grad_norm_bf16: {e}")))?; + let aux_gn_blocks = (aux_param_count + 255) / 256; + let aux_grad_norm_partials = stream.alloc_zeros::(aux_gn_blocks) + .map_err(|e| MLError::ModelError(format!("alloc aux_grad_norm_partials: {e}")))?; + // Standalone kernel handles for aux Adam (not captured in any CUDA graph) + let aux_adam_update_kernel = { + let module = stream.context().load_cubin(DQN_UTILITY_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("aux_adam cubin: {e}")))?; + module.load_function("dqn_adam_update_kernel") + .map_err(|e| MLError::ModelError(format!("aux_adam load: {e}")))? + }; + let aux_grad_norm_finalize_kernel = { + let module = stream.context().load_cubin(DQN_UTILITY_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("aux_gnf cubin: {e}")))?; + module.load_function("dqn_grad_norm_finalize") + .map_err(|e| MLError::ModelError(format!("aux_gnf load: {e}")))? + }; + // v8: PopArt running statistics buffers let popart_mean = stream.alloc_zeros::(1) .map_err(|e| MLError::ModelError(format!("alloc popart_mean: {e}")))?; @@ -3214,6 +3386,17 @@ impl GpuDqnTrainer { exposure_aux_grad_kernel, exposure_aux_scratch, exposure_aux_weight_gpu, + aux_adam_m, + aux_adam_v, + aux_adam_t, + aux_grad_buf, + aux_grad_norm_f32, + aux_grad_norm_bf16, + aux_grad_norm_partials, + aux_param_count, + aux_adam_step: 0, + aux_adam_update_kernel, + aux_grad_norm_finalize_kernel, popart_normalize_kernel, popart_mean, popart_var, diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 63bbf618a..fff326901 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1052,8 +1052,8 @@ impl FusedTrainingCtx { .map_err(|e| anyhow::anyhow!("Regime PER scaling: {e}"))?; // v7.1: Exposure auxiliary gradient (breaks i%3 degeneracy). - // Moved here from run_full_step so it's captured inside graph_aux / graph_mega. - // The kernel reads aux_weight from GPU-resident scalar (set_exposure_aux_weight). + // Now uses a SEPARATE Adam optimizer — aux gradients go into aux_grad_buf, + // not the shared grad_buf. No gradient clipping competition with C51. if self.exposure_aux_weight > 1e-6 { if let Err(e) = self.trainer.launch_exposure_aux_grad( &self.exposure_targets, @@ -1061,10 +1061,8 @@ impl FusedTrainingCtx { ) { tracing::warn!("Exposure aux grad failed (non-fatal): {e}"); } - // Clip combined gradient to prevent aux GEMM backward from dominating. - // Same pattern as CQL gradient budget — keeps total norm within max_grad_norm. - self.trainer.clip_grad_buf_inplace(self.trainer.config().max_grad_norm) - .map_err(|e| anyhow::anyhow!("Aux gradient clip: {e}"))?; + // No clip_grad_buf_inplace needed — aux gradients are in a separate + // buffer with their own Adam optimizer. They never enter grad_buf. } Ok(())