From fb056319979a0cbcf05079cfa388e3e37e581c34 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 2 Apr 2026 09:58:47 +0200 Subject: [PATCH] perf: capture spectral norm in graph_forward + EMA in graph_ema MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Spectral norm (10 weight matrices): - Moved into graph_forward capture (submit_spectral_norm_ops) - Runs before cuBLAS forward — graph ensures normalized weights - All device_ptr() → raw_ptr() (no cudarc events) - ~10 kernel launches eliminated from per-step ungraphed ops EMA target update: - Captured as graph_ema (EMA kernel + f32→bf16 cast + 20 unflatten DtoD) - tau read from GPU-resident tau_buf (async HtoD before replay) - ~22 kernel launches (EMA + cast + 20 DtoD) → 1 graph replay Per-step kernel launches now: 7 graph replays + ~4 ungraphed (was: 5 graph replays + ~14 ungraphed) Remaining ungraphed: - upload_batch_gpu (8 ops — batch ptrs change per step) - HER relabel (2 ops — donor indices change) - IQN trunk gradient (2 ops — tau changes) - IQN→PER DtoD (1 op) - causal/vaccine (amortized, rare) - regime_scale + PER update (2 ops) Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 65 ++++++++++++++++--- crates/ml/src/trainers/dqn/fused_training.rs | 61 +++++++++++++---- 2 files changed, 104 insertions(+), 22 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index fd6caa313..171e03576 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -929,6 +929,8 @@ impl GpuDqnTrainer { /// /// Shape: `[B]` f32 — rewards from the batch. Valid after `train_step()` or /// `train_step_gpu()`. Used by IQN for Bellman target computation on GPU. + pub fn tau_buf_ptr(&self) -> u64 { self.tau_buf.raw_ptr() } + pub fn rewards_buf(&self) -> &CudaSlice { &self.rewards_buf } @@ -2833,8 +2835,8 @@ impl GpuDqnTrainer { pub fn train_step_gpu( &mut self, gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch, - online_dueling: &DuelingWeightSet, - online_branching: &BranchingWeightSet, + online_dueling: &mut DuelingWeightSet, + online_branching: &mut BranchingWeightSet, _target_dueling: &DuelingWeightSet, _target_branching: &BranchingWeightSet, ) -> Result { @@ -3026,9 +3028,10 @@ impl GpuDqnTrainer { /// - (injection point: external code adds IQN/attention/ensemble gradients to grad_buf) /// - `graph_adam`: grad_norm → Adam → unflatten fn execute_train_and_readback( + &mut self, - online_dueling: &DuelingWeightSet, - online_branching: &BranchingWeightSet, + online_dueling: &mut DuelingWeightSet, + online_branching: &mut BranchingWeightSet, ) -> Result { let b = self.config.batch_size; @@ -3426,8 +3429,8 @@ impl GpuDqnTrainer { /// by `update_priorities_cuda()` which reads them directly via device pointer. fn execute_train_scalars_only( &mut self, - online_dueling: &DuelingWeightSet, - online_branching: &BranchingWeightSet, + online_dueling: &mut DuelingWeightSet, + online_branching: &mut BranchingWeightSet, ) -> Result { // ── Update Adam step counter on device (OUTSIDE graph) ─────── self.adam_step += 1; @@ -3877,8 +3880,8 @@ impl GpuDqnTrainer { /// they are pure element-wise ops with fixed control flow. fn capture_training_graphs( &mut self, - online_d: &DuelingWeightSet, - online_b: &BranchingWeightSet, + online_d: &mut DuelingWeightSet, + online_b: &mut BranchingWeightSet, ) -> Result<(), MLError> { self.stream.synchronize() .map_err(|e| MLError::ModelError(format!("stream sync before capture: {e}")))?; @@ -3895,6 +3898,7 @@ impl GpuDqnTrainer { return Err(MLError::ModelError(format!("CUDA graph_forward begin_capture: {e}"))); } + self.submit_spectral_norm_ops(online_d, online_b)?; let submit_fwd_result = self.submit_forward_ops(); let graph_fwd_result = self.stream.end_capture( @@ -4005,6 +4009,51 @@ impl GpuDqnTrainer { Ok(()) } + /// Submit spectral norm ops for graph capture. + /// Normalizes all weight matrices and syncs back to params_buf. + /// Must run BEFORE submit_forward_ops so cuBLAS reads normalized weights. + fn submit_spectral_norm_ops( + &mut self, + online_dueling: &mut DuelingWeightSet, + online_branching: &mut BranchingWeightSet, + ) -> Result<(), MLError> { + self.apply_spectral_norm(online_dueling, online_branching) + } + + /// Submit EMA + bf16 cast + unflatten ops for graph capture. + /// tau comes from GPU-resident tau_buf (async HtoD before replay). + pub fn submit_ema_ops( + &mut self, + target_d: &mut DuelingWeightSet, + target_b: &mut BranchingWeightSet, + ) -> Result<(), MLError> { + // EMA kernel: target_params_buf = (1-tau)*target + tau*online + let n = self.total_params as i32; + let blocks = ((self.total_params + 255) / 256) as u32; + let t_ptr = self.target_params_buf.raw_ptr(); + let o_ptr = self.params_buf.raw_ptr(); + let tau_ptr = self.tau_buf.raw_ptr(); + unsafe { + self.stream + .launch_builder(&self.ema_kernel) + .arg(&t_ptr) + .arg(&o_ptr) + .arg(&tau_ptr) + .arg(&n) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("ema capture: {e}")))?; + } + // f32 → bf16 cast + self.launch_f32_to_bf16_target_cast()?; + // Unflatten to individual weight tensors + self.unflatten_target_weights(target_d, target_b)?; + Ok(()) + } + fn submit_forward_ops(&mut self) -> Result<(), MLError> { // ── Zero accumulators (capturable: memset_zeros uses cuMemsetD32Async) ─ self.stream diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index d91482be9..d45e07c73 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -109,6 +109,8 @@ pub(crate) struct FusedTrainingCtx { /// CUDA Graph for IQN training (decode + fwd/loss + bwd + adam + trunk grad). /// Captured lazily on first step. graph_iqn: Option, + /// CUDA Graph for EMA target update + bf16 cast + unflatten. + graph_ema: Option, /// Ensemble extra heads (heads 1..K-1). Head 0 lives inside the CUDA Graph. /// Each element is an independent (DuelingWeightSet, BranchingWeightSet) pair /// that shares the same trunk but has perturbed value/advantage weights. @@ -510,6 +512,7 @@ impl FusedTrainingCtx { graph_attention: None, graph_iql: None, graph_iqn: None, + graph_ema: None, ensemble_extra_heads, ensemble_diversity_weight: hyperparams.ensemble_diversity_weight as f32, ensemble_aggregate_kernel, @@ -576,17 +579,14 @@ impl FusedTrainingCtx { let gpu_batch = batch.gpu_batch.as_ref() .ok_or_else(|| anyhow::anyhow!("Fused training requires gpu_batch (GPU PER)"))?; - // ── Step 1: Spectral normalization BEFORE forward pass ───────── - self.trainer.apply_spectral_norm(&mut self.online_dueling, &mut self.online_branching) - .map_err(|e| anyhow::anyhow!("Spectral norm (pre-forward): {e}"))?; // ── Step 2: Upload batch + replay graph_forward ONLY ───────────── // Upload the RAW PER batch into trainer staging buffers (DtoD). // HER relabeling happens AFTER upload, modifying the staging buffers in-place. // This eliminates the intermediate GpuBatch (was 17+ GPU allocs per step). - let _fused_placeholder = self.trainer.train_step_gpu( + let _step_scalars = self.trainer.train_step_gpu( gpu_batch, - &self.online_dueling, &self.online_branching, + &mut self.online_dueling, &mut self.online_branching, &self.target_dueling, &self.target_branching, ).map_err(|e| anyhow::anyhow!("Fused train_step_gpu (forward only): {e}"))?; @@ -635,21 +635,54 @@ impl FusedTrainingCtx { // C51 gradient clip now captured in graph_adam — no per-step call needed. // ── Step 3: GPU-native Polyak EMA target update ────────────────── + // tau uploaded async to GPU buffer before graph replay. { let dqn = agent.primary_dqn_mut(); - let tau = compute_cosine_annealed_tau( dqn.get_training_steps(), - dqn.config.tau, - dqn.config.tau_final, - dqn.config.tau_anneal_steps, + dqn.config.tau, dqn.config.tau_final, dqn.config.tau_anneal_steps, ); + // Async HtoD for tau (graph reads from tau_buf device address) + unsafe { + cudarc::driver::sys::cuMemcpyHtoDAsync_v2( + self.trainer.tau_buf_ptr(), + (&(tau as f32) as *const f32).cast(), + std::mem::size_of::(), + self.stream.cu_stream(), + ); + } - self.trainer.target_ema_update( - &self.online_dueling, &self.online_branching, - &mut self.target_dueling, &mut self.target_branching, - tau as f32, - ).map_err(|e| anyhow::anyhow!("GPU EMA target update: {e}"))?; + if self.graph_ema.is_none() { + // First step: run EMA ungraphed to initialize + self.trainer.target_ema_update( + &self.online_dueling, &self.online_branching, + &mut self.target_dueling, &mut self.target_branching, + tau as f32, + ).map_err(|e| anyhow::anyhow!("EMA first step: {e}"))?; + + // Capture EMA graph + self.stream.synchronize() + .map_err(|e| anyhow::anyhow!("sync before ema capture: {e}"))?; + unsafe { self.stream.context().disable_event_tracking(); } + if let Ok(()) = self.stream.begin_capture( + cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL, + ) { + let _ = self.trainer.submit_ema_ops( + &mut self.target_dueling, &mut self.target_branching, + ); + if let Ok(Some(graph)) = self.stream.end_capture( + cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ) { + use crate::cuda_pipeline::gpu_dqn_trainer::SendSyncGraph; + self.graph_ema = Some(SendSyncGraph(graph)); + tracing::info!("Captured EMA CUDA graph (ema + bf16 cast + unflatten)"); + } + } + unsafe { self.stream.context().enable_event_tracking(); } + let _ = self.stream.context().check_err(); + } else if let Some(ref graph) = self.graph_ema { + graph.0.launch().map_err(|e| anyhow::anyhow!("EMA graph replay: {e}"))?; + } } // ── Step 3b-5b: Auxiliary heads (attention + IQL + IQN) ─────────