From 408188e045e4d64c50c6792ecd48b8f92a7be2c4 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 2 Apr 2026 13:49:26 +0200 Subject: [PATCH] perf: graph_aux CUDA graph capture + fix 3 kernel ABI mismatches - Fix EMA kernel: add tau_buf device field, async HtoD via stable host address, pass pointer not scalar (was ILLEGAL_ADDRESS every step) - Fix HER relabel kernel: revert indirect ptr_buf to direct bf16 pointer - Fix PER update kernel: revert indirect ptr_buf to direct u32/bf16 pointers - Remove IQL per-step DtoH readback (cuStreamSynchronize blocks graph capture) - Permanently disable cudarc event tracking (SyncOnDrop safe for capture) - EventTrackingGuard no longer re-enables tracking on drop - Pre-allocate pass1_event/pass3_event (no cuEventCreate per step) - RawCudaGraph: raw CUDA driver API bypassing cudarc bind_to_thread - graph_aux captures ~30 aux kernel launches (HER+clip+EMA+attn+IQL+IQN+CQL) into single CUDA graph, replayed from step 3+ for zero launch overhead - IQN/GpuDqnTrainer: tau_host stable field for graph-captured HtoD - Remove 3 dead indirect pointer kernels from dqn_utility_kernels.cu - Local RTX 3050: 7.5ms/step steady state (batch=64, 200 steps/epoch) Co-Authored-By: Claude Opus 4.6 (1M context) --- .../src/cuda_pipeline/dqn_utility_kernels.cu | 58 +- crates/ml/src/cuda_pipeline/gpu_attention.rs | 8 +- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 700 +++++++----------- .../ml/src/cuda_pipeline/gpu_iql_trainer.rs | 14 +- crates/ml/src/cuda_pipeline/gpu_iqn_head.rs | 18 +- .../ml/src/cuda_pipeline/per_update_kernel.cu | 7 +- crates/ml/src/trainers/dqn/fused_training.rs | 683 +++++++++-------- ...04-02-mega-graph-training-loop-refactor.md | 479 ++++++++++++ 8 files changed, 1119 insertions(+), 848 deletions(-) create mode 100644 docs/superpowers/plans/2026-04-02-mega-graph-training-loop-refactor.md diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index e2c1764ea..3ba373c53 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -455,61 +455,6 @@ extern "C" __global__ void pad_states_kernel( dst[idx] = (j < sd) ? src[b * sd + j] : bf16_zero(); } -/* ══════════════════════════════════════════════════════════════════════ - * INDIRECT PAD STATES KERNEL (graph-capture compatible) - * - * Same as pad_states_kernel but reads the source pointer from a device - * buffer. This allows the CUDA graph to be replayed with different batch - * data each step — only the pointer in src_ptr_buf changes (via async HtoD). - * - * Launch config: grid=(ceil(batch * padded_sd / 256), 1, 1), block=(256, 1, 1). - * ══════════════════════════════════════════════════════════════════════ */ -extern "C" __global__ void pad_states_indirect_kernel( - __nv_bfloat16* __restrict__ dst, - const unsigned long long* __restrict__ src_ptr_buf, /* [1] device ptr to source */ - int batch, - int sd, - int padded_sd -) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= batch * padded_sd) return; - const __nv_bfloat16* src = (const __nv_bfloat16*)src_ptr_buf[0]; - int b = idx / padded_sd; - int j = idx % padded_sd; - dst[idx] = (j < sd) ? src[b * sd + j] : bf16_zero(); -} - -/* ══════════════════════════════════════════════════════════════════════ - * INDIRECT MEMCPY KERNEL (graph-capture compatible) - * - * Copies N elements from a source address stored in a device buffer - * to a fixed destination. Replaces async DtoD copies that can't be - * graphed because the source pointer changes per step. - * - * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1). - * ══════════════════════════════════════════════════════════════════════ */ -extern "C" __global__ void indirect_copy_f32_kernel( - float* __restrict__ dst, - const unsigned long long* __restrict__ src_ptr_buf, /* [1] device ptr */ - int n -) { - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= n) return; - const float* src = (const float*)src_ptr_buf[0]; - dst[i] = src[i]; -} - -extern "C" __global__ void indirect_copy_i32_kernel( - int* __restrict__ dst, - const unsigned long long* __restrict__ src_ptr_buf, /* [1] device ptr */ - int n -) { - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= n) return; - const int* src = (const int*)src_ptr_buf[0]; - dst[i] = src[i]; -} - /* ══════════════════════════════════════════════════════════════════════ * BOTTLENECK TANH + CONCAT KERNEL (#31) * @@ -1157,7 +1102,7 @@ extern "C" __global__ void her_inplace_relabel( __nv_bfloat16* __restrict__ dst_states, /* [batch_size, dst_stride] padded */ __nv_bfloat16* __restrict__ dst_next_states, /* [batch_size, dst_stride] padded */ float* __restrict__ dst_rewards, /* [batch_size] f32 */ - const unsigned long long* __restrict__ src_next_ptr_buf, /* [1] ptr to src next_states */ + const __nv_bfloat16* __restrict__ src_next_states, /* [batch_size, src_stride] direct ptr */ const int* __restrict__ donor_indices, /* [her_batch_size] i32 */ int offset, /* normal_count: first HER row in dst */ int goal_dim, /* number of goal columns to replace */ @@ -1169,7 +1114,6 @@ extern "C" __global__ void her_inplace_relabel( int lane = threadIdx.x; /* warp lane 0-31 */ int donor = donor_indices[i]; int dst_row = offset + i; - const __nv_bfloat16* src_next_states = (const __nv_bfloat16*)src_next_ptr_buf[0]; /* Replace goal columns (first goal_dim elements) with donor's achieved goal */ for (int d = lane; d < goal_dim; d += 32) { diff --git a/crates/ml/src/cuda_pipeline/gpu_attention.rs b/crates/ml/src/cuda_pipeline/gpu_attention.rs index be2af4a1d..df065d316 100644 --- a/crates/ml/src/cuda_pipeline/gpu_attention.rs +++ b/crates/ml/src/cuda_pipeline/gpu_attention.rs @@ -16,7 +16,7 @@ //! Adam optimizer (separate from the DQN trunk Adam). use std::sync::Arc; -use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; +use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; use tracing::info; use crate::MLError; @@ -317,6 +317,12 @@ impl GpuAttention { /// Run gradient norm + Adam update on attention parameters. /// + /// Increment the host-side adam step counter without launching kernels. + /// Used before graph_aux replay so the captured HtoD copies the updated value. + pub fn increment_adam_step(&mut self) { + self.attn_adam_step += 1; + } + /// Must be called after `backward()`. Uses the same grad_norm + Adam /// kernel pattern as IQN (separate from DQN trunk Adam). pub fn adam_step(&mut self, lr: f32, max_grad_norm: f32) -> Result<(), MLError> { diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 7a5112fc1..e2419a139 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -75,8 +75,9 @@ static CQL_GRAD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cql_gra // ── Event tracking RAII guard ──────────────────────────────────────────────── /// RAII guard that disables cudarc event tracking on creation and -/// re-enables it on drop. Prevents early-return bugs where tracking -/// is left permanently disabled. +/// Event tracking is permanently disabled for CUDA graph compatibility +/// (see FusedTrainingCtx::new). This guard now only drains stale errors +/// on drop — it does NOT re-enable tracking. pub(crate) struct EventTrackingGuard<'a> { ctx: &'a cudarc::driver::CudaContext, } @@ -90,7 +91,8 @@ impl<'a> EventTrackingGuard<'a> { impl Drop for EventTrackingGuard<'_> { fn drop(&mut self) { - unsafe { self.ctx.enable_event_tracking(); } + // Do NOT re-enable event tracking — permanently disabled for + // graph capture safety. Only drain stale errors. let _ = self.ctx.check_err(); } } @@ -101,7 +103,7 @@ impl Drop for EventTrackingGuard<'_> { /// /// The graph is only launched on the same stream/context that created it. /// The trainer is not shared across threads in practice. -pub(crate) struct SendSyncGraph(pub(crate) CudaGraph); +pub(crate) struct SendSyncGraph(CudaGraph); // Safety: CudaGraph is bound to a specific CUDA context. The trainer // that owns it is always used from the thread that created the context. @@ -180,9 +182,6 @@ pub struct GpuDqnTrainConfig { pub c51_warmup_epochs: usize, /// CQL regularization strength (0.0 = disabled, 0.1 = mild, 1.0 = full offline-RL). pub cql_alpha: f32, - pub per_alpha: f32, - pub per_epsilon: f32, - pub per_capacity: usize, /// Curiosity Q-penalty lambda: scales gamma by 1/(1 + lambda * curiosity_error). /// 0.0 = disabled. Typical range 0.5-5.0. Higher = more aggressive penalization /// of Q-values in novel states (prevents overconfident extrapolation OOS). @@ -258,9 +257,6 @@ impl Default for GpuDqnTrainConfig { entropy_coefficient: 0.001, // Must be ≤ reward magnitude (~0.001). Old 0.01 was 10x reward → entropy dominated learning. c51_warmup_epochs: 5, cql_alpha: 0.1, - per_alpha: 0.6, - per_epsilon: 1e-6, - per_capacity: 500_000, curiosity_q_penalty_lambda: 0.0, asymmetric_dd_weight: 0.0, ensemble_disagreement_penalty: 0.0, @@ -581,14 +577,17 @@ pub struct GpuDqnTrainer { // ── Adam step counter on device (CUDA Graph cannot bake scalars) ─ t_buf: CudaSlice, // [1] current Adam step + /// EMA tau on device [1] — kernel reads `tau_buf[0]` (pointer, not scalar). + tau_buf: CudaSlice, + + /// Host-side tau for graph replay. Graph captures HtoD from this address. + pub(crate) tau_host: f32, // ── Training state ────────────────────────────────────────────── - adam_step: i32, + pub(crate) adam_step: i32, total_params: usize, - params_initialized: bool, + pub(crate) params_initialized: bool, target_params_initialized: bool, - /// GPU-resident tau for EMA (async HtoD, graph-capture compatible). - tau_buf: CudaSlice, attention_initialized: bool, // ── CUDA Graphs (split: forward+backward, then optimizer) ────── @@ -596,8 +595,8 @@ pub struct GpuDqnTrainer { // Graph B: grad_norm → Adam → unflatten // Between A and B: external code can ADD auxiliary gradients to grad_buf // (IQN trunk, ensemble heads) — single Adam sees combined gradient. - graph_forward: Option, - graph_adam: Option, + pub(crate) graph_forward: Option, + pub(crate) graph_adam: Option, // ── Consolidated transfer buffers ───────────────────────────── /// Single staging buffer for batch upload consolidation (bf16 data only). @@ -656,10 +655,6 @@ pub struct GpuDqnTrainer { pruning_compute_kernel: CudaFunction, /// HER in-place goal relabel kernel (writes directly into padded staging buffers). pub(crate) her_inplace_kernel: CudaFunction, - pad_states_indirect_kernel: CudaFunction, - indirect_copy_f32_kernel: CudaFunction, - indirect_copy_i32_kernel: CudaFunction, - batch_ptr_buf: CudaSlice, /// #20 Pruning epoch (epoch at which to compute the mask). pruning_epoch: usize, /// #20 Pruning fraction (0.7 = prune 70% of smallest weights). @@ -755,6 +750,13 @@ pub struct GpuDqnTrainer { /// output buffers. double_dqn_stream: Arc, + /// Pre-allocated event: main stream records after Pass 1 + Pass 2 complete. + /// double_dqn_stream waits on this before starting Pass 3. + pass1_event: CudaEvent, + /// Pre-allocated event: double_dqn_stream records after Pass 3 complete. + /// Main stream waits on this before starting loss kernels. + pass3_event: CudaEvent, + // Scratch buffers for cuBLAS target network forward pass. // Online activations reuse the existing save_h_* buffers. // Target only needs h_s2 (for DDQN) + scratch for intermediate layers. @@ -913,7 +915,7 @@ impl GpuDqnTrainer { /// Reference to the states buffer on GPU. /// - /// Shape: `[B, STATE_DIM]` — contains the batch's states after batch upload. + /// Shape: `[B, STATE_DIM]` — contains the batch's states after `upload_batch_gpu()`. /// Used by IQL value network to read states without Candle tensor intermediaries. pub fn states_buf(&self) -> &CudaSlice { &self.states_buf @@ -939,8 +941,6 @@ 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 } @@ -1007,6 +1007,7 @@ impl GpuDqnTrainer { _online_dueling: &mut DuelingWeightSet, ) -> Result<(), MLError> { let b = self.config.batch_size; + let _eg = EventTrackingGuard::new(self.stream.context()); let sd = self.config.state_dim; let sh1 = self.config.shared_h1; let sh2 = self.config.shared_h2; @@ -1217,6 +1218,7 @@ impl GpuDqnTrainer { scale: f32, ) -> Result<(), MLError> { let b = self.config.batch_size; + let _eg = EventTrackingGuard::new(self.stream.context()); let sd = self.config.state_dim; let sh1 = self.config.shared_h1; let sh2 = self.config.shared_h2; @@ -1453,6 +1455,7 @@ impl GpuDqnTrainer { let bs = self.config.batch_size as i32; let blocks = ((self.config.batch_size + 255) / 256) as u32; + let _evt_guard = EventTrackingGuard::new(self.stream.context()); let td_ptr = self.td_errors_buf.raw_ptr(); let states_ptr = self.states_buf.raw_ptr(); @@ -1488,7 +1491,7 @@ impl GpuDqnTrainer { /// `state_dim`: unpadded state dimension pub fn launch_her_inplace_relabel( &self, - _src_next_states_ptr: u64, // unused — kernel reads from batch_ptr_buf[1] + src_next_states_ptr: u64, donor_indices_ptr: u64, normal_count: usize, her_batch_size: usize, @@ -1518,7 +1521,7 @@ impl GpuDqnTrainer { .arg(&states_ptr) .arg(&next_states_ptr) .arg(&rewards_ptr) - .arg(&(self.batch_ptr_buf.raw_ptr() + 8)) + .arg(&src_next_states_ptr) .arg(&donor_indices_ptr) .arg(&offset_i32) .arg(&goal_dim_i32) @@ -1575,6 +1578,7 @@ impl GpuDqnTrainer { pub fn apply_cql_gradient( &mut self, ) -> Result { + let _eg = EventTrackingGuard::new(self.stream.context()); let cql_kernel = match &self.cql_logit_grad_kernel { Some(k) => k.clone(), None => return Ok(false), @@ -1591,6 +1595,7 @@ impl GpuDqnTrainer { let total_actions = b0 + b1 + b2; let _total_branch_atoms = total_actions * na; + let _evt_guard = EventTrackingGuard::new(self.stream.context()); // Step 1: Compute CQL logit gradients via kernel let v_logits_ptr = self.on_v_logits_buf.raw_ptr(); @@ -1740,6 +1745,7 @@ impl GpuDqnTrainer { /// Computes norm of scratch, clips to `cql_budget`, then SAXPYs into grad_buf. pub fn apply_cql_clipped_saxpy(&mut self, cql_budget: f32) -> Result<(), MLError> { // Compute CQL gradient norm (float accumulator) + let _eg = EventTrackingGuard::new(self.stream.context()); self.stream.memset_zeros(&mut self.grad_norm_f32_buf) .map_err(|e| MLError::ModelError(format!("zero cql_grad_norm_f32: {e}")))?; @@ -1798,16 +1804,18 @@ impl GpuDqnTrainer { online_branching: &mut BranchingWeightSet, ) -> Result<(), MLError> { let sh1 = self.config.shared_h1 as i32; + let _eg = EventTrackingGuard::new(self.stream.context()); let sh2 = self.config.shared_h2 as i32; let sd = self.config.state_dim as i32; let sigma_max = self.config.spectral_norm_sigma_max; + let _evt_guard = EventTrackingGuard::new(self.stream.context()); // Spectral norm on W_s1 [shared_h1, state_dim] { - let w_ptr = online_dueling.w_s1.raw_ptr(); - let u_ptr = self.spec_u_s1.raw_ptr(); - let v_ptr = self.spec_v_s1.raw_ptr(); + let (w_ptr, _g) = online_dueling.w_s1.device_ptr(&self.stream); + let (u_ptr, _g2) = self.spec_u_s1.device_ptr(&self.stream); + let (v_ptr, _g3) = self.spec_v_s1.device_ptr(&self.stream); unsafe { self.stream .launch_builder(&self.spectral_norm_kernel) @@ -1824,9 +1832,9 @@ impl GpuDqnTrainer { // Spectral norm on W_s2 [shared_h2, shared_h1] { - let w_ptr = online_dueling.w_s2.raw_ptr(); - let u_ptr = self.spec_u_s2.raw_ptr(); - let v_ptr = self.spec_v_s2.raw_ptr(); + let (w_ptr, _g) = online_dueling.w_s2.device_ptr(&self.stream); + let (u_ptr, _g2) = self.spec_u_s2.device_ptr(&self.stream); + let (v_ptr, _g3) = self.spec_v_s2.device_ptr(&self.stream); unsafe { self.stream .launch_builder(&self.spectral_norm_kernel) @@ -1851,9 +1859,9 @@ impl GpuDqnTrainer { macro_rules! spec_norm { ($w_slice:expr, $u_slice:expr, $v_slice:expr, $out_dim:expr, $in_dim:expr, $label:literal) => {{ - let w_ptr = $w_slice.raw_ptr(); - let u_ptr = $u_slice.raw_ptr(); - let v_ptr = $v_slice.raw_ptr(); + let (w_ptr, _gw) = $w_slice.device_ptr(&self.stream); + let (u_ptr, _gu) = $u_slice.device_ptr(&self.stream); + let (v_ptr, _gv) = $v_slice.device_ptr(&self.stream); unsafe { self.stream .launch_builder(&self.spectral_norm_kernel) @@ -1969,6 +1977,7 @@ impl GpuDqnTrainer { pub fn clip_grad_buf_inplace(&mut self, max_norm: f32) -> Result<(), MLError> { // Disable cudarc event tracking — graph replay leaves stale events // that cause INVALID_VALUE on post-graph kernel launches. + let _evt_guard = EventTrackingGuard::new(self.stream.context()); // Zero the float grad_norm accumulator self.stream .memset_zeros(&mut self.grad_norm_f32_buf) @@ -2006,6 +2015,7 @@ impl GpuDqnTrainer { /// Used for per-component gradient diagnostics. Costs one stream sync + /// 4-byte DtoH transfer. Call sparingly (e.g. once per epoch or every N steps). pub fn read_grad_norm_sync(&mut self) -> Result { + let _evt_guard = EventTrackingGuard::new(self.stream.context()); self.stream .memset_zeros(&mut self.grad_norm_f32_buf) @@ -2052,9 +2062,9 @@ impl GpuDqnTrainer { // Disable cudarc event tracking — no sync needed, stream ordering // guarantees graph replay completed before these kernel launches. + let _evt_guard = EventTrackingGuard::new(self.stream.context()); // Run attention forward: save_h_s2 → attention output buffer - tracing::error!("trainer stream={:?} attn stream={:?}", self.stream.cu_stream() as u64, attention.stream_handle()); let attn_output = attention.forward(&self.save_h_s2, batch_size)?; // DtoD copy: attention output → save_h_s2 (in-place update for next graph replay) @@ -2091,6 +2101,7 @@ impl GpuDqnTrainer { // GPU-native: single kernel generates noise + blends in one pass on bf16 shadow. // params_bf16[i] = alpha * params_bf16[i] + (1-alpha) * N(0, sigma) // No CPU noise generation, no HtoD upload. + let _evt_guard = EventTrackingGuard::new(self.stream.context()); let params_ptr = self.params_bf16.raw_ptr(); unsafe { self.stream @@ -2124,6 +2135,7 @@ impl GpuDqnTrainer { }) .map_err(|e| MLError::ModelError(format!("bf16_to_f32 shrink_perturb: {e}")))?; } + drop(_evt_guard); // Invalidate both CUDA Graphs — weights changed outside captured ops self.graph_forward = None; @@ -2154,7 +2166,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, pad_states_indirect_kernel, indirect_copy_f32_kernel, indirect_copy_i32_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, 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). @@ -2248,6 +2260,8 @@ impl GpuDqnTrainer { let cql_grad_scratch = stream.alloc_zeros::(total_params) .map_err(|e| MLError::ModelError(format!("alloc cql_grad_scratch f32: {e}")))?; let t_buf = alloc_i32(&stream, 1, "adam_t")?; + let tau_buf = stream.alloc_zeros::(1) + .map_err(|e| MLError::ModelError(format!("alloc tau_buf: {e}")))?; // ── Allocate consolidated transfer buffers ───────────────── // Upload staging: states + next_states (bf16 only) @@ -2462,6 +2476,12 @@ impl GpuDqnTrainer { let double_dqn_stream = stream.fork() .map_err(|e| MLError::DeviceError(format!("fork double_dqn_stream: {e}")))?; + // Pre-allocate multi-stream sync events (no cuEventCreate during capture). + let pass1_event = stream.context().new_event(None) + .map_err(|e| MLError::DeviceError(format!("alloc pass1_event: {e}")))?; + let pass3_event = stream.context().new_event(None) + .map_err(|e| MLError::DeviceError(format!("alloc pass3_event: {e}")))?; + // ── Initialize cuBLAS backward context (required) ────────── let cublas_backward = CublasBackward::new(&stream, &config)?; info!("GpuDqnTrainer: cuBLAS batched backward initialized"); @@ -2532,7 +2552,6 @@ impl GpuDqnTrainer { }; let ptrs = { - let _evt_guard = EventTrackingGuard::new(stream.context()); let _evt_guard = EventTrackingGuard::new(stream.context()); CachedPtrs { // bf16 shadows — GemmEx reads these for forward/backward @@ -2629,10 +2648,6 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("alloc bn_d_concat: {e}")))?; let bn_d_hidden_buf = stream.alloc_zeros::(b * bn_alloc_dim) .map_err(|e| MLError::ModelError(format!("alloc bn_d_hidden: {e}")))?; - let tau_buf = stream.alloc_zeros::(1) - .map_err(|e| MLError::ModelError(format!("tau_buf alloc: {e}")))?; - let batch_ptr_buf = stream.alloc_zeros::(8) - .map_err(|e| MLError::ModelError(format!("batch_ptr_buf alloc: {e}")))?; Ok(Self { config, stream, @@ -2702,11 +2717,12 @@ impl GpuDqnTrainer { grad_norm_f32_buf, cql_grad_scratch, t_buf, + tau_buf, + tau_host: 0.0, adam_step: 0, total_params, params_initialized: false, target_params_initialized: false, - tau_buf, attention_initialized: false, graph_forward: None, graph_adam: None, @@ -2726,6 +2742,8 @@ impl GpuDqnTrainer { bf16_next_states_buf, cublas_forward, double_dqn_stream, + pass1_event, + pass3_event, tg_h_s1_scratch, tg_h_s2_buf, tg_h_v_scratch, @@ -2780,10 +2798,6 @@ impl GpuDqnTrainer { pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel, - pad_states_indirect_kernel, - indirect_copy_f32_kernel, - indirect_copy_i32_kernel, - batch_ptr_buf, pruning_epoch: prune_ep, pruning_fraction: prune_frac, causal_intervene_kernel: causal_intervene_kernel_fn, @@ -2827,17 +2841,6 @@ impl GpuDqnTrainer { &self.stream } - /// Synchronize both the main stream and double_dqn_stream. - /// Required before CUDA Graph capture to ensure all cross-stream - /// dependencies are resolved. - pub fn sync_all_streams(&self) -> Result<(), MLError> { - self.stream.synchronize() - .map_err(|e| MLError::ModelError(format!("main stream sync: {e}")))?; - self.double_dqn_stream.synchronize() - .map_err(|e| MLError::ModelError(format!("double_dqn_stream sync: {e}")))?; - Ok(()) - } - // ═══════════════════════════════════════════════════════════════════ /// GPU-direct training step — no CPU roundtrip. @@ -2851,8 +2854,8 @@ impl GpuDqnTrainer { pub fn train_step_gpu( &mut self, gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch, - online_dueling: &mut DuelingWeightSet, - online_branching: &mut BranchingWeightSet, + online_dueling: &DuelingWeightSet, + online_branching: &BranchingWeightSet, _target_dueling: &DuelingWeightSet, _target_branching: &BranchingWeightSet, ) -> Result { @@ -2864,9 +2867,8 @@ impl GpuDqnTrainer { self.params_initialized = true; } - // ── Upload batch pointers to GPU indirection buffer (async HtoD) ─ - // The indirect upload kernels in graph_forward read from batch_ptr_buf. - self.upload_batch_ptrs_6(gpu_batch); + // ── GPU-direct upload (DtoD, no CPU staging) ───────────────── + self.upload_batch_gpu(gpu_batch)?; // Upload + capture + replay forward ONLY. // Adam is deferred to replay_adam_and_readback() so the caller can @@ -2905,6 +2907,7 @@ impl GpuDqnTrainer { /// Reads grad_buf (which has been modified by all gradient injections), /// writes to grad_norm_f32_buf and grad_norm_buf. pub fn compute_grad_norm_outside_graph(&mut self) -> Result<(), MLError> { + let _eg = EventTrackingGuard::new(self.stream.context()); self.stream.memset_zeros(&mut self.grad_norm_f32_buf) .map_err(|e| MLError::ModelError(format!("zero grad_norm_f32 pre-adam: {e}")))?; self.launch_grad_norm()?; @@ -3044,10 +3047,9 @@ 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: &mut DuelingWeightSet, - online_branching: &mut BranchingWeightSet, + online_dueling: &DuelingWeightSet, + online_branching: &BranchingWeightSet, ) -> Result { let b = self.config.batch_size; @@ -3295,9 +3297,15 @@ impl GpuDqnTrainer { } } - // Sensitivity stays on GPU — no readback needed. - // The caller discards the return value (only error handling). - Ok(0.0) + // Single DtoH readback at the end — read mean sensitivity + let n_features = market_dim.min(14); + let mut host_sens = vec![0.0_f32; market_dim]; + self.stream.synchronize() + .map_err(|e| MLError::ModelError(format!("causal final sync: {e}")))?; + self.stream.memcpy_dtoh(&self.causal_sensitivity_buf, &mut host_sens[..market_dim]) + .map_err(|e| MLError::ModelError(format!("causal readback: {e}")))?; + let total: f32 = host_sens[..n_features].iter().sum(); + Ok(total / n_features as f32) } /// #32 Gradient Vaccine: project out overfitting gradient components. @@ -3332,9 +3340,8 @@ impl GpuDqnTrainer { self.stream.memset_zeros(&mut self.d_adv_logits_buf) .map_err(|e| MLError::ModelError(format!("vaccine zero d_adv: {e}")))?; - // Step 3: Upload vaccine batch via indirect pointers + forward+backward - self.upload_batch_ptrs_6(vaccine_batch); - self.submit_indirect_upload_ops()?; + // Step 3: Upload vaccine batch + forward+backward (NON-graph path) + self.upload_batch_gpu(vaccine_batch)?; // Forward pass (cuBLAS, outside graph) self.launch_cublas_forward()?; self.launch_curiosity_inference()?; @@ -3440,8 +3447,8 @@ impl GpuDqnTrainer { /// by `update_priorities_cuda()` which reads them directly via device pointer. fn execute_train_scalars_only( &mut self, - online_dueling: &mut DuelingWeightSet, - online_branching: &mut BranchingWeightSet, + online_dueling: &DuelingWeightSet, + online_branching: &BranchingWeightSet, ) -> Result { // ── Update Adam step counter on device (OUTSIDE graph) ─────── self.adam_step += 1; @@ -3502,13 +3509,13 @@ impl GpuDqnTrainer { /// * `epsilon` — PER epsilon floor pub fn update_priorities_cuda( &mut self, - _indices: &CudaSlice, // unused — kernel reads from batch_ptr_buf[6] - _priorities_data: &CudaSlice, // unused — kernel reads from batch_ptr_buf[7] + indices: &CudaSlice, + priorities: &CudaSlice, alpha: f32, epsilon: f32, ) -> Result<(), MLError> { let b = self.config.batch_size as i32; - let capacity = _priorities_data.len() as i32; + let capacity = priorities.len() as i32; let blocks = ((self.config.batch_size + 255) / 256) as u32; let launch_cfg = LaunchConfig { grid_dim: (blocks, 1, 1), @@ -3520,8 +3527,8 @@ impl GpuDqnTrainer { self.stream .launch_builder(&self.per_update_kernel) .arg(&self.td_errors_buf) - .arg(&(self.batch_ptr_buf.raw_ptr() + 6 * 8)) - .arg(&(self.batch_ptr_buf.raw_ptr() + 7 * 8)) + .arg(indices) + .arg(priorities) .arg(&self.batch_max_buf) .arg(&alpha) .arg(&epsilon) @@ -3628,6 +3635,7 @@ impl GpuDqnTrainer { states: &CudaSlice, batch_size: usize, ) -> Result<&CudaSlice, MLError> { + let _eg = EventTrackingGuard::new(self.stream.context()); if batch_size > self.config.batch_size { return Err(MLError::ModelError(format!( "compute_q_values: batch_size {batch_size} exceeds trainer batch_size {}", @@ -3707,6 +3715,7 @@ impl GpuDqnTrainer { states: &CudaSlice, batch_size: usize, ) -> Result { + let _eg = EventTrackingGuard::new(self.stream.context()); // Pad input rows to pad128(state_dim) for CUTLASS K-tile alignment, // and pad batch dimension to config.batch_size for N-tile alignment. @@ -3772,6 +3781,7 @@ impl GpuDqnTrainer { /// returns zeros. Call `flush_q_stats_readback()` at epoch end to drain /// the last in-flight readback. pub fn reduce_current_q_stats(&mut self) -> Result { + let _eg = EventTrackingGuard::new(self.stream.context()); // 1. Collect previous reduction's results (if pending) let prev = if self.q_stats_ready { @@ -3876,30 +3886,35 @@ impl GpuDqnTrainer { // CUDA Graph capture and invalidation // ═══════════════════════════════════════════════════════════════════ - /// Capture CUDA Graphs for the training step. + /// Capture two CUDA Graphs: forward (zero→backward) and adam (grad_norm→unflatten). /// - /// Two graphs are captured: - /// - `graph_forward`: forward pass, loss, gradient, backward - /// - `graph_adam`: CQL gradient, C51 clip, pruning mask, grad_norm, Adam, unflatten + /// Called on the first `train_step()` or after `invalidate_training_graph()`. + /// The split allows external code to inject auxiliary gradients (IQN, attention, + /// ensemble) into `grad_buf` between the two graph replays. /// - /// The split between forward and adam allows auxiliary gradient injection - /// (IQN, attention, ensemble) into `grad_buf` between the two replays. - /// CQL, clip, pruning, and grad_norm are captured in graph_adam because - /// they are pure element-wise ops with fixed control flow. + /// Graph A (`graph_forward`): zero → cuBLAS forward → C51 loss → C51 grad → cuBLAS backward + /// Graph B (`graph_adam`): grad_norm → Adam → unflatten (20 d2d copies) fn capture_training_graphs( &mut self, - online_d: &mut DuelingWeightSet, - online_b: &mut BranchingWeightSet, + online_d: &DuelingWeightSet, + online_b: &BranchingWeightSet, ) -> Result<(), MLError> { + // Synchronize the stream before capture to ensure all pending work + // (BF16 mirror sync, batch upload, adam_step memcpy) is complete. self.stream.synchronize() .map_err(|e| MLError::ModelError(format!("stream sync before capture: {e}")))?; + // Disable event tracking during capture — cudarc's device_ptr records + // CudaEvents which are DISALLOWED inside CUDA Graph capture. + unsafe { self.stream.context().disable_event_tracking(); } // ── Capture graph_forward ────────────────────────────────────── let begin_result = self.stream.begin_capture( cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL, ); if let Err(e) = begin_result { + unsafe { self.stream.context().enable_event_tracking(); } + let _ = self.stream.context().check_err(); return Err(MLError::ModelError(format!("CUDA graph_forward begin_capture: {e}"))); } @@ -3909,44 +3924,49 @@ impl GpuDqnTrainer { cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, ); + // Check forward submission if let Err(e) = submit_fwd_result { + unsafe { self.stream.context().enable_event_tracking(); } + let _ = self.stream.context().check_err(); return Err(e); } let graph_fwd = graph_fwd_result .map_err(|e| { + unsafe { self.stream.context().enable_event_tracking(); } + let _ = self.stream.context().check_err(); MLError::ModelError(format!("CUDA graph_forward end_capture: {e}")) })? .ok_or_else(|| { + unsafe { self.stream.context().enable_event_tracking(); } + let _ = self.stream.context().check_err(); MLError::ModelError( "CUDA graph_forward capture returned None — stream may not support capture".into() ) })?; - // ── Capture graph_adam (includes CQL + clip + pruning + grad_norm) ── + // ── Capture graph_adam ────────────────────────────────────────── let begin_result = self.stream.begin_capture( cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL, ); if let Err(e) = begin_result { + unsafe { self.stream.context().enable_event_tracking(); } + let _ = self.stream.context().check_err(); return Err(MLError::ModelError(format!("CUDA graph_adam begin_capture: {e}"))); } - // CQL gradient + backward + clipped SAXPY into grad_buf - self.submit_cql_ops()?; - - // C51 gradient budget clip (60%) - self.submit_c51_clip_ops()?; - - // Pruning mask: grad_buf *= mask - self.submit_pruning_mask_ops()?; - - // Adam optimizer: grad_norm + Adam update + unflatten - self.submit_adam_ops(online_d, online_b)?; + let submit_adam_result = self.submit_adam_ops(online_d, online_b); let graph_adam_result = self.stream.end_capture( cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, ); + // Re-enable event tracking after both captures and drain stale errors. + unsafe { self.stream.context().enable_event_tracking(); } + let _ = self.stream.context().check_err(); + + // Propagate adam submission error + submit_adam_result?; let graph_adam = graph_adam_result .map_err(|e| MLError::ModelError(format!("CUDA graph_adam end_capture: {e}")))? @@ -3954,6 +3974,9 @@ impl GpuDqnTrainer { "CUDA graph_adam capture returned None — stream may not support capture".into() ))?; + // Launch both graphs on first capture (warm-up). On subsequent steps, + // only graph_forward is replayed by train_step_gpu(). The caller then + // injects auxiliary gradients and calls replay_adam_and_readback(). graph_fwd.launch().map_err(|e| { MLError::ModelError(format!("CUDA graph_forward first launch: {e}")) })?; @@ -3964,7 +3987,7 @@ impl GpuDqnTrainer { info!( "GpuDqnTrainer: 2 CUDA graphs captured and launched \ (graph_forward: 5 memsets + forward + loss + grad + backward; \ - graph_adam: CQL + C51 clip + pruning + grad_norm + adam + 20 d2d unflatten)" + graph_adam: grad_norm + adam + 20 d2d unflatten)" ); self.graph_forward = Some(SendSyncGraph(graph_fwd)); self.graph_adam = Some(SendSyncGraph(graph_adam)); @@ -3979,6 +4002,7 @@ impl GpuDqnTrainer { /// Debug: run forward WITHOUT graph capture, with intermediate readbacks. #[allow(dead_code)] fn debug_forward_no_graph(&mut self) -> Result<(), MLError> { + let _eg = EventTrackingGuard::new(self.stream.context()); // Zero accumulators self.stream.memset_zeros(&mut self.total_loss_buf).ok(); @@ -4002,170 +4026,7 @@ 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. - /// Submit batch upload ops using indirect kernels (graph-capture compatible). - fn submit_indirect_upload_ops(&mut self) -> Result<(), MLError> { - let b = self.config.batch_size; - let sd = self.config.state_dim; - let total = b * sd; - let blocks = ((total + 255) / 256) as u32; - let cfg = LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }; - let bp = self.batch_ptr_buf.raw_ptr(); - let ptr_sz = std::mem::size_of::() as u64; - let b_i32 = b as i32; - let sd_i32 = sd as i32; - // States - unsafe { - self.stream.launch_builder(&self.pad_states_indirect_kernel) - .arg(&self.states_buf.raw_ptr()).arg(&bp) - .arg(&b_i32).arg(&sd_i32).arg(&sd_i32) - .launch(cfg).map_err(|e| MLError::ModelError(format!("indirect pad states: {e}")))?; - } - // Next_states - let bp1 = bp + ptr_sz; - unsafe { - self.stream.launch_builder(&self.pad_states_indirect_kernel) - .arg(&self.next_states_buf.raw_ptr()).arg(&bp1) - .arg(&b_i32).arg(&sd_i32).arg(&sd_i32) - .launch(cfg).map_err(|e| MLError::ModelError(format!("indirect pad next: {e}")))?; - } - let blocks_b = ((b + 255) / 256) as u32; - let cfg_b = LaunchConfig { grid_dim: (blocks_b, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }; - // Rewards - unsafe { - self.stream.launch_builder(&self.indirect_copy_f32_kernel) - .arg(&self.rewards_buf.raw_ptr()).arg(&(bp + 2*ptr_sz)).arg(&b_i32) - .launch(cfg_b).map_err(|e| MLError::ModelError(format!("indirect rewards: {e}")))?; - } - // Dones - unsafe { - self.stream.launch_builder(&self.indirect_copy_f32_kernel) - .arg(&self.dones_buf.raw_ptr()).arg(&(bp + 3*ptr_sz)).arg(&b_i32) - .launch(cfg_b).map_err(|e| MLError::ModelError(format!("indirect dones: {e}")))?; - } - // IS-weights - unsafe { - self.stream.launch_builder(&self.indirect_copy_f32_kernel) - .arg(&self.is_weights_buf.raw_ptr()).arg(&(bp + 4*ptr_sz)).arg(&b_i32) - .launch(cfg_b).map_err(|e| MLError::ModelError(format!("indirect weights: {e}")))?; - } - // Actions - unsafe { - self.stream.launch_builder(&self.indirect_copy_i32_kernel) - .arg(&self.actions_buf.raw_ptr()).arg(&(bp + 5*ptr_sz)).arg(&b_i32) - .launch(cfg_b).map_err(|e| MLError::ModelError(format!("indirect actions: {e}")))?; - } - Ok(()) - } - - /// Upload batch source pointers to GPU indirection buffer (async HtoD). - /// Upload all batch + PER pointers to GPU indirection buffer (async HtoD). - /// Layout: [0]=states, [1]=next_states, [2]=rewards, [3]=dones, - /// [4]=weights, [5]=actions, [6]=indices, [7]=priorities - /// Upload 6 batch source pointers (states through actions). - pub fn upload_batch_ptrs_6(&self, gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch) { - let ptrs: [u64; 6] = [ - gpu_batch.states.data().raw_ptr(), - gpu_batch.next_states.data().raw_ptr(), - gpu_batch.rewards.raw_ptr(), - gpu_batch.dones.raw_ptr(), - gpu_batch.weights.raw_ptr(), - gpu_batch.actions.raw_ptr(), - ]; - unsafe { - cudarc::driver::sys::cuMemcpyHtoDAsync_v2( - self.batch_ptr_buf.raw_ptr(), - ptrs.as_ptr().cast(), - 6 * std::mem::size_of::(), - self.stream.cu_stream(), - ); - } - } - - /// Upload PER pointers (indices + priorities) to slots [6] and [7]. - pub fn upload_per_ptrs(&self, indices_ptr: u64, priorities_ptr: u64) { - let ptrs: [u64; 2] = [indices_ptr, priorities_ptr]; - unsafe { - cudarc::driver::sys::cuMemcpyHtoDAsync_v2( - self.batch_ptr_buf.raw_ptr() + 6 * std::mem::size_of::() as u64, - ptrs.as_ptr().cast(), - 2 * std::mem::size_of::(), - self.stream.cu_stream(), - ); - } - } - - pub fn upload_batch_ptrs_8( - &self, - gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch, - indices_ptr: u64, - priorities_ptr: u64, - ) { - let ptrs: [u64; 8] = [ - gpu_batch.states.data().raw_ptr(), - gpu_batch.next_states.data().raw_ptr(), - gpu_batch.rewards.raw_ptr(), - gpu_batch.dones.raw_ptr(), - gpu_batch.weights.raw_ptr(), - gpu_batch.actions.raw_ptr(), - indices_ptr, - priorities_ptr, - ]; - unsafe { - cudarc::driver::sys::cuMemcpyHtoDAsync_v2( - self.batch_ptr_buf.raw_ptr(), - ptrs.as_ptr().cast(), - 8 * std::mem::size_of::(), - self.stream.cu_stream(), - ); - } - } - - 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> { + pub(crate) fn submit_forward_ops(&mut self) -> Result<(), MLError> { // ── Zero accumulators (capturable: memset_zeros uses cuMemsetD32Async) ─ self.stream .memset_zeros(&mut self.total_loss_buf) @@ -4272,176 +4133,11 @@ impl GpuDqnTrainer { Ok(()) } - /// Submit CQL gradient ops for graph capture. - /// Identical to apply_cql_gradient + apply_cql_clipped_saxpy but without - /// EventTrackingGuard (tracking already disabled during capture). - fn submit_cql_ops(&mut self) -> Result<(), MLError> { - let cql_kernel = match &self.cql_logit_grad_kernel { - Some(k) => k.clone(), - None => return Ok(()), - }; - if self.config.cql_alpha <= 0.0 { return Ok(()); } - - let b = self.config.batch_size; - let na = self.config.num_atoms; - let b0 = self.config.branch_0_size; - let b1 = self.config.branch_1_size; - let b2 = self.config.branch_2_size; - - // Zero CQL staging buffers - self.stream.memset_zeros(&mut self.cql_d_value_logits) - .map_err(|e| MLError::ModelError(format!("zero cql_d_val: {e}")))?; - self.stream.memset_zeros(&mut self.cql_d_adv_logits) - .map_err(|e| MLError::ModelError(format!("zero cql_d_adv: {e}")))?; - - // CQL logit gradient kernel - let blocks = ((b + 255) / 256) as u32; - unsafe { - self.stream.launch_builder(&cql_kernel) - .arg(&self.on_v_logits_buf.raw_ptr()) - .arg(&self.on_b_logits_buf.raw_ptr()) - .arg(&self.actions_buf.raw_ptr()) - .arg(&self.cql_d_value_logits.raw_ptr()) - .arg(&self.cql_d_adv_logits.raw_ptr()) - .arg(&self.config.cql_alpha) - .arg(&(b as i32)) - .arg(&(na as i32)) - .arg(&(b0 as i32)) - .arg(&(b1 as i32)) - .arg(&(b2 as i32)) - .arg(&self.config.v_min) - .arg(&self.config.v_max) - .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) - .map_err(|e| MLError::ModelError(format!("cql_logit_grad capture: {e}")))?; - } - - // Cast CQL d_logits f32 → bf16 for cuBLAS backward - { - let total_actions = b0 + b1 + b2; - let n_val = (b * na) as i32; - let n_adv = (b * total_actions * na) as i32; - let cfg = |n: i32| LaunchConfig { - grid_dim: (((n as u32) + 255) / 256, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0, - }; - unsafe { - self.stream.launch_builder(&self.f32_to_bf16_kernel) - .arg(&self.cql_d_value_logits.raw_ptr()) - .arg(&self.d_value_logits_bf16.raw_ptr()) - .arg(&n_val) - .launch(cfg(n_val)) - .map_err(|e| MLError::ModelError(format!("f32_to_bf16 cql_d_val capture: {e}")))?; - self.stream.launch_builder(&self.f32_to_bf16_kernel) - .arg(&self.cql_d_adv_logits.raw_ptr()) - .arg(&self.d_adv_logits_bf16.raw_ptr()) - .arg(&n_adv) - .launch(cfg(n_adv)) - .map_err(|e| MLError::ModelError(format!("f32_to_bf16 cql_d_adv capture: {e}")))?; - } - } - - // Zero cql_grad_scratch and run cuBLAS backward into it - self.stream.memset_zeros(&mut self.cql_grad_scratch) - .map_err(|e| MLError::ModelError(format!("zero cql_grad_scratch capture: {e}")))?; - { - let param_sizes = compute_param_sizes(&self.config); - let w_ptrs = bf16_weight_ptrs_from_base(self.ptrs.params_buf, ¶m_sizes); - let bf16_size = std::mem::size_of::(); - let d_val_bf16 = self.d_value_logits_bf16.raw_ptr(); - let d_adv_bf16_base = self.d_adv_logits_bf16.raw_ptr(); - let d_adv_ptrs = [ - d_adv_bf16_base, - d_adv_bf16_base + (b0 * na * bf16_size) as u64, - d_adv_bf16_base + ((b0 + b1) * na * bf16_size) as u64, - ]; - self.cublas_backward.backward_full( - &self.stream, d_val_bf16, &d_adv_ptrs, - self.states_buf.raw_ptr(), - self.save_h_s1.raw_ptr(), self.save_h_s2.raw_ptr(), self.save_h_v.raw_ptr(), - &[self.save_h_b0.raw_ptr(), self.save_h_b1.raw_ptr(), self.save_h_b2.raw_ptr()], - &w_ptrs, self.cql_grad_scratch.raw_ptr(), - self.bw_d_h_s2.raw_ptr(), self.bw_d_h_s1.raw_ptr(), self.bw_d_h_v.raw_ptr(), - &[self.bw_d_h_b0.raw_ptr(), self.bw_d_h_b1.raw_ptr(), self.bw_d_h_b2.raw_ptr()], - self.bw_dy_bf16_staging.raw_ptr(), 0, - ).map_err(|e| MLError::ModelError(format!("CQL backward_full capture: {e}")))?; - } - - // Clipped SAXPY: grad_buf += clip(cql_scratch, budget) - let cql_budget = self.config.max_grad_norm * 0.25; - self.stream.memset_zeros(&mut self.grad_norm_f32_buf) - .map_err(|e| MLError::ModelError(format!("zero cql_norm capture: {e}")))?; - { - let total = self.total_params as i32; - let blocks = ((self.total_params + 255) / 256) as u32; - let cfg = LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 256 }; - unsafe { - self.stream.launch_builder(&self.grad_norm_kernel) - .arg(&self.ptrs.cql_grad_scratch) - .arg(&self.ptrs.grad_norm_f32_buf) - .arg(&total) - .launch(cfg) - .map_err(|e| MLError::ModelError(format!("cql grad_norm capture: {e}")))?; - } - let cfg2 = LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }; - let alpha = 1.0_f32; - unsafe { - self.stream.launch_builder(&self.clipped_saxpy_kernel) - .arg(&self.ptrs.grad_buf) - .arg(&self.ptrs.cql_grad_scratch) - .arg(&alpha) - .arg(&cql_budget) - .arg(&self.ptrs.grad_norm_f32_buf) - .arg(&total) - .launch(cfg2) - .map_err(|e| MLError::ModelError(format!("cql clipped_saxpy capture: {e}")))?; - } - } - Ok(()) - } - - /// Submit C51 gradient clip ops for graph capture. - fn submit_c51_clip_ops(&mut self) -> Result<(), MLError> { - let c51_budget = self.config.max_grad_norm * 0.60; - self.stream.memset_zeros(&mut self.grad_norm_f32_buf) - .map_err(|e| MLError::ModelError(format!("zero c51_clip_norm: {e}")))?; - self.launch_grad_norm()?; - self.launch_grad_norm_finalize()?; - let total = self.total_params as i32; - let blocks = ((self.total_params + 255) / 256) as u32; - unsafe { - self.stream.launch_builder(&self.clip_grad_kernel) - .arg(&self.ptrs.grad_buf) - .arg(&self.ptrs.grad_norm_buf) - .arg(&c51_budget) - .arg(&total) - .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) - .map_err(|e| MLError::ModelError(format!("c51 clip capture: {e}")))?; - } - Ok(()) - } - - /// Submit pruning mask ops for graph capture. - fn submit_pruning_mask_ops(&mut self) -> Result<(), MLError> { - if let Some(ref mask) = self.pruning_mask { - let tp = self.total_params as i32; - let blocks = ((tp as u32 + 255) / 256) as u32; - unsafe { - self.stream.launch_builder(&self.pruning_mask_kernel) - .arg(&self.ptrs.params_f32_ptr) - .arg(&self.ptrs.params_buf) - .arg(&mask.raw_ptr()) - .arg(&tp) - .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) - .map_err(|e| MLError::ModelError(format!("pruning mask capture: {e}")))?; - } - } - Ok(()) - } - /// Submit the optimizer phase ops to the stream (captured into graph_adam). /// /// Steps: zero grad_norm → grad_norm → Adam → unflatten. /// Reads `grad_buf` which may contain combined C51 + IQN + ensemble gradients. - fn submit_adam_ops( + pub(crate) fn submit_adam_ops( &mut self, online_d: &DuelingWeightSet, online_b: &BranchingWeightSet, @@ -4459,9 +4155,6 @@ impl GpuDqnTrainer { // ── 7. Unflatten: params_bf16 → individual bf16 weight tensors ─ self.unflatten_online_weights(online_d, online_b)?; - // ── 8. Regime-adaptive PER scaling (element-wise, all pre-allocated) ─ - self.regime_scale_td_errors()?; - Ok(()) } @@ -4524,6 +4217,78 @@ impl GpuDqnTrainer { /// Upload batch data from GPU tensors — pure CudaSlice, zero Candle temporaries. + /// + /// GpuBatch layout from GpuReplayBuffer: + /// - states/next_states: BF16 (stored in BF16 ring buffer) + /// - rewards/dones: BF16 GpuTensor + /// - weights: CudaSlice (IS-weights kept as f32 to avoid bf16 overflow → Inf → NaN) + /// - actions: U32 (stored in U32 ring buffer) + /// + /// No Candle `to_dtype()` temporaries — eliminates Drop/event conflicts with forked stream. + fn upload_batch_gpu( + &mut self, + gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch, + ) -> Result<(), MLError> { + let b = self.config.batch_size; + + let f32_size = std::mem::size_of::(); + + // States + next_states: contiguous [B, SD] in GpuBatch → padded [B, pad128(SD)] + self.launch_pad_states( + self.states_buf.raw_ptr(), + gpu_batch.states.data().raw_ptr(), + b, + )?; + self.launch_pad_states( + self.next_states_buf.raw_ptr(), + gpu_batch.next_states.data().raw_ptr(), + b, + )?; + + // Rewards, dones: f32 DtoD (no bf16 NaN risk) + dtod_copy( + self.rewards_buf.raw_ptr(), + gpu_batch.rewards.raw_ptr(), + b * f32_size, &self.stream, 2, "rewards", + )?; + dtod_copy( + self.dones_buf.raw_ptr(), + gpu_batch.dones.raw_ptr(), + b * f32_size, &self.stream, 3, "dones", + )?; + + // IS-weights: f32 DtoD (bf16 overflows to Inf for PER weights) + dtod_copy( + self.is_weights_buf.raw_ptr(), + gpu_batch.weights.raw_ptr(), + b * f32_size, &self.stream, 4, "is_weights", + )?; + + // Actions: CudaSlice → i32 buf (async DtoD, zero CPU sync) + dtod_copy( + self.actions_buf.raw_ptr(), gpu_batch.actions.raw_ptr(), + b * std::mem::size_of::(), &self.stream, 4, "actions", + )?; + + Ok(()) + } + + // ═══════════════════════════════════════════════════════════════════ + // Kernel launch methods + // ═══════════════════════════════════════════════════════════════════ + + /// Launch the cuBLAS batched forward pass (online + target networks). + /// + /// Replaces `launch_forward_loss` for the forward phase only. + /// The C51 distributional loss is still computed by the existing NVRTC kernel + /// which reads the cuBLAS output buffers (`on_v_logits_buf`, `on_b_logits_buf`, + /// `tg_v_logits_buf`, `tg_b_logits_buf`) together with the saved activations. + /// + /// This raises GPU occupancy from ~1.56% to >60% on H100 by replacing the + /// 1-warp-per-sample shared-memory-bound kernel with cuBLAS DGEMM which + /// uses tensor cores and efficiently tiles across the full SM. + /// + /// Returns `Ok(false)` if cuBLAS is not initialized (caller falls back to BF16 kernel). fn launch_cublas_forward(&self) -> Result<(), MLError> { let cublas = &self.cublas_forward; @@ -4644,11 +4409,11 @@ impl GpuDqnTrainer { // Pass 2 (target fwd) and Pass 3 (online fwd on next_states for Double // DQN) are independent — they read the same next_states_buf but write to // completely separate output buffers. Run them concurrently on 2 streams. - let pass1_done = self.stream.record_event(None) - .map_err(|e| MLError::ModelError(format!("pass1 event: {e}")))?; + self.pass1_event.record(&self.stream) + .map_err(|e| MLError::ModelError(format!("pass1 event record: {e}")))?; // double_dqn_stream waits for Pass 1 + stochastic depth to finish. - self.double_dqn_stream.wait(&pass1_done) + self.double_dqn_stream.wait(&self.pass1_event) .map_err(|e| MLError::ModelError(format!("double_dqn wait pass1: {e}")))?; // ── Pass 2: Target forward on NEXT_STATES — main stream (graph-safe) @@ -4676,9 +4441,9 @@ impl GpuDqnTrainer { cublas.set_stream(&self.stream)?; // ── Join: main stream waits for Pass 3 to complete before loss kernels ── - let pass3_done = self.double_dqn_stream.record_event(None) - .map_err(|e| MLError::ModelError(format!("pass3 done event: {e}")))?; - self.stream.wait(&pass3_done) + self.pass3_event.record(&self.double_dqn_stream) + .map_err(|e| MLError::ModelError(format!("pass3 event record: {e}")))?; + self.stream.wait(&self.pass3_event) .map_err(|e| MLError::ModelError(format!("join pass3: {e}")))?; Ok(()) @@ -5437,7 +5202,7 @@ impl GpuDqnTrainer { /// /// After copying to the bf16 shadow, launches a bf16->f32 cast kernel to /// initialize the f32 master weights (used only on first call / graph invalidation). - fn flatten_online_weights( + pub(crate) fn flatten_online_weights( &self, online_d: &DuelingWeightSet, online_b: &BranchingWeightSet, @@ -5680,6 +5445,7 @@ impl GpuDqnTrainer { // No cuStreamSynchronize — stream ordering is sufficient. // The previous hang was caused by the sync memcpy_htod for adam_step // (now async), not by missing sync here. + let _eg = EventTrackingGuard::new(self.stream.context()); // First call: sync + flatten target weights into flat buffer. if !self.target_params_initialized { @@ -5700,17 +5466,19 @@ impl GpuDqnTrainer { }; // EMA on f32 master weights: target = (1-tau)*target + tau*online + // Store tau in stable host field (graph captures HtoD from this address). + self.tau_host = tau; unsafe { cudarc::driver::sys::cuMemcpyHtoDAsync_v2( self.tau_buf.raw_ptr(), - (&tau as *const f32).cast(), + (&self.tau_host as *const f32).cast(), std::mem::size_of::(), self.stream.cu_stream(), ); } - let tau_ptr = self.tau_buf.raw_ptr(); 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) @@ -5732,6 +5500,40 @@ impl GpuDqnTrainer { Ok(()) } + + /// Synchronize both streams (main + double_dqn). Used before graph capture + /// to ensure all prior work is complete. + pub(crate) fn sync_all_streams(&self) -> Result<(), MLError> { + self.stream.synchronize() + .map_err(|e| MLError::ModelError(format!("main stream sync: {e}")))?; + self.double_dqn_stream.synchronize() + .map_err(|e| MLError::ModelError(format!("double_dqn stream sync: {e}")))?; + Ok(()) + } + + /// Increment Adam step counter and async-upload to device buffer. + /// Must be called OUTSIDE the captured graph (before replay). + pub(crate) fn adam_step_async(&mut self) { + self.adam_step += 1; + unsafe { + cudarc::driver::sys::cuMemcpyHtoDAsync_v2( + self.t_buf.raw_ptr(), + (&self.adam_step as *const i32).cast(), + std::mem::size_of::(), + self.stream.cu_stream(), + ); + } + } + + /// Raw pointer to the double DQN stream (for raw CUDA API calls during capture). + pub(crate) fn double_dqn_cu_stream(&self) -> cudarc::driver::sys::CUstream { + self.double_dqn_stream.cu_stream() + } + + /// Raw pointer to the main stream. + pub(crate) fn main_cu_stream(&self) -> cudarc::driver::sys::CUstream { + self.stream.cu_stream() + } } // ── Compilation ───────────────────────────────────────────────────────────── @@ -5747,7 +5549,7 @@ impl GpuDqnTrainer { fn compile_training_kernels( stream: &Arc, 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, 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), MLError> { info!( state_dim = config.state_dim, total_params = compute_total_params(config), @@ -5814,15 +5616,9 @@ fn compile_training_kernels( let her_inplace = module.load_function("her_inplace_relabel") .map_err(|e| MLError::ModelError(format!("her_inplace_relabel load: {e}")))?; - let pad_states_indirect = module.load_function("pad_states_indirect_kernel") - .map_err(|e| MLError::ModelError(format!("pad_states_indirect load: {e}")))?; - let indirect_copy_f32 = module.load_function("indirect_copy_f32_kernel") - .map_err(|e| MLError::ModelError(format!("indirect_copy_f32 load: {e}")))?; - let indirect_copy_i32 = module.load_function("indirect_copy_i32_kernel") - .map_err(|e| MLError::ModelError(format!("indirect_copy_i32 load: {e}")))?; - info!("GpuDqnTrainer: 30 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, pad_states_indirect, indirect_copy_f32, indirect_copy_i32)) + 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)) } /// Load the standalone Polyak EMA kernel from precompiled cubin. diff --git a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs index fe1a58386..0ff5dfd09 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs @@ -369,17 +369,21 @@ impl GpuIqlTrainer { .map_err(|e| MLError::ModelError(format!("IQL adam kernel: {e}")))?; } - // Read back total loss - let mut loss_host = [0.0_f32]; - super::dtoh_bf16_to_f32(&self.stream, &self.total_loss_buf, &mut loss_host)?; - - Ok(loss_host[0]) + // Loss stays on GPU — no per-step DtoH. Caller discards the value. + // Removing the synchronous readback unblocks CUDA Graph capture. + Ok(0.0) } /// Current Adam step count. pub fn adam_step(&self) -> i32 { self.adam_step } + + /// Increment the host-side adam step counter without launching kernels. + /// Used before graph_aux replay so the captured HtoD copies the updated value. + pub fn increment_adam_step(&mut self) { + self.adam_step += 1; + } } // --------------------------------------------------------------------------- diff --git a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs index b6ba0a0c1..55f324c08 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs @@ -190,6 +190,8 @@ pub struct GpuIqnHead { adam_step: i32, t_buf: cudarc::driver::CudaSlice, tau_buf: cudarc::driver::CudaSlice, + /// Host-side tau for graph replay. Graph captures HtoD from this stable address. + tau_host: f32, /// Monotonic step counter for Philox PRNG seeding (τ sampling). rng_step: u32, total_params: usize, @@ -316,6 +318,7 @@ impl GpuIqnHead { adam_step: 0, t_buf, tau_buf, + tau_host: 0.0, rng_step: 0, total_params, }) @@ -595,10 +598,12 @@ impl GpuIqnHead { /// `target[i] = (1 - tau) * target[i] + tau * online[i]` pub fn target_ema_update(&mut self, tau: f32) -> Result<(), MLError> { let n = self.total_params; + // Store tau in stable host field (graph captures HtoD from this address). + self.tau_host = tau; unsafe { cudarc::driver::sys::cuMemcpyHtoDAsync_v2( self.tau_buf.raw_ptr(), - (&tau as *const f32).cast(), + (&self.tau_host as *const f32).cast(), std::mem::size_of::(), self.stream.cu_stream(), ); @@ -639,6 +644,17 @@ impl GpuIqnHead { self.adam_step } + /// Increment the host-side adam step counter without launching kernels. + /// Used before graph_aux replay so the captured HtoD copies the updated value. + pub fn increment_adam_step(&mut self) { + self.adam_step += 1; + } + + /// Set host-side tau for graph replay. Graph captures HtoD from `&self.tau_host`. + pub fn set_tau_host(&mut self, tau: f32) { + self.tau_host = tau; + } + /// Per-sample IQN loss buffer reference (for PER weighting). pub fn per_sample_loss(&self) -> &CudaSlice { &self.per_sample_loss diff --git a/crates/ml/src/cuda_pipeline/per_update_kernel.cu b/crates/ml/src/cuda_pipeline/per_update_kernel.cu index 6e07517f0..6f7900d6c 100644 --- a/crates/ml/src/cuda_pipeline/per_update_kernel.cu +++ b/crates/ml/src/cuda_pipeline/per_update_kernel.cu @@ -10,8 +10,8 @@ extern "C" __global__ void per_update_priorities_kernel( const __nv_bfloat16* __restrict__ td_errors, - const unsigned long long* __restrict__ indices_ptr_buf, /* [1] ptr to indices */ - const unsigned long long* __restrict__ priorities_ptr_buf, /* [1] ptr to priorities */ + const unsigned int* __restrict__ indices, + __nv_bfloat16* __restrict__ priorities, __nv_bfloat16* __restrict__ batch_max, float alpha, float epsilon, @@ -21,9 +21,6 @@ void per_update_priorities_kernel( int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= batch_size) return; - const unsigned int* indices = (const unsigned int*)indices_ptr_buf[0]; - __nv_bfloat16* priorities = (__nv_bfloat16*)priorities_ptr_buf[0]; - __nv_bfloat16 td_raw = td_errors[i]; __nv_bfloat16 td_abs = bf16_fabs(td_raw); __nv_bfloat16 alpha_bf = bf16(alpha); diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 22c8f0a62..017d633b6 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -30,6 +30,8 @@ use ml_core::device::MlDevice; use tracing::info; use cudarc::driver::DevicePtr; +use cudarc::driver::sys as cuda_sys; + use crate::cuda_pipeline::gpu_attention::{GpuAttention, GpuAttentionConfig}; use crate::cuda_pipeline::gpu_dqn_trainer::{GpuDqnTrainConfig, GpuDqnTrainer}; use crate::cuda_pipeline::gpu_her::{GpuHer, GpuHerConfig, parse_her_strategy}; @@ -43,6 +45,42 @@ use crate::dqn::replay_buffer_type::BatchSample; use super::config::DQNAgentType; use super::DQNHyperparameters; +/// Raw CUDA graph handle that bypasses cudarc's bind_to_thread/check_err. +/// +/// Uses cuGraphLaunch directly for zero-overhead graph replay. +/// cudarc's CudaGraph wrapper calls bind_to_thread → check_err which can fail +/// on stale errors left by SyncOnDrop, making it unsafe during graph capture. +struct RawCudaGraph { + exec: cuda_sys::CUgraphExec, + #[allow(dead_code)] + graph: cuda_sys::CUgraph, +} + +impl RawCudaGraph { + /// Launch the instantiated graph on the given stream. + fn launch(&self, stream: cuda_sys::CUstream) -> Result<()> { + let result = unsafe { cuda_sys::cuGraphLaunch(self.exec, stream) }; + if result != cuda_sys::cudaError_enum::CUDA_SUCCESS { + anyhow::bail!("cuGraphLaunch failed: {result:?}"); + } + Ok(()) + } +} + +impl Drop for RawCudaGraph { + fn drop(&mut self) { + unsafe { + cuda_sys::cuGraphExecDestroy(self.exec); + cuda_sys::cuGraphDestroy(self.graph); + } + } +} + +// Safety: The graph is bound to a CUDA context. FusedTrainingCtx is always used +// from the thread that created the context. No concurrent access occurs. +unsafe impl Send for RawCudaGraph {} +unsafe impl Sync for RawCudaGraph {} + /// Per-component gradient norm budget fractions for auxiliary objectives. /// C51 gets whatever remains: `1.0 - sum(active_auxiliary_budgets)`. /// When all auxiliaries are active: C51=60%, CQL=25%, IQN=10%, Ens=5%. @@ -101,17 +139,8 @@ pub(crate) struct FusedTrainingCtx { /// CVaR scales buffer (kept alive so device pointer remains valid). cvar_scales_buf: Option>, /// GPU multi-head feature attention over h_s2 (post-graph, 1-step lag). + /// When Some, applied after EMA update, before IQN, each training step. pub(crate) gpu_attention: Option, - /// CUDA Graph for attention (fwd + bwd + adam). Captured lazily on first step. - graph_attention: Option, - /// CUDA Graph for IQL value training. Captured lazily on first step. - graph_iql: Option, - /// 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, - graph_her: 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. @@ -147,6 +176,12 @@ pub(crate) struct FusedTrainingCtx { /// #32 Gradient Vaccine: pending validation batch for next training step. /// Set by the training loop, consumed by `run_full_step()`. pub(crate) pending_vaccine_batch: Option, + /// Captured auxiliary graph: HER + clip + EMA + attention + IQL + IQN + CQL. + /// Captured after the first step succeeds (ensures all conditional components + /// are initialized). Replayed on subsequent steps for zero launch overhead. + graph_aux: Option, + /// Host-side tau for graph replay. Graph captures HtoD from this address. + tau_host: f32, } impl Drop for FusedTrainingCtx { @@ -184,13 +219,6 @@ impl FusedTrainingCtx { anyhow::anyhow!("Fused CUDA training requires branching target network") })?; - // Drain any stale cudarc errors, then disable event tracking permanently. - // This ensures: (1) no stale errors poison future bind_to_thread/check_err, - // (2) all CudaSlice allocations have read=None, write=None — preventing - // SyncOnDrop guards from recording events during CUDA Graph capture. - let _ = stream.context().check_err(); - unsafe { stream.context().disable_event_tracking(); } - // Use the caller-provided forked CudaStream. All GPU components (experience // collector, fused trainer, portfolio sim, monitoring) share this single // stream via Arc::clone — no more dual-stream event tracking conflicts. @@ -232,9 +260,6 @@ impl FusedTrainingCtx { entropy_coefficient: dqn.config.entropy_coefficient as f32, c51_warmup_epochs: hyperparams.c51_warmup_epochs, cql_alpha: hyperparams.cql_alpha as f32, - per_alpha: hyperparams.per_alpha as f32, - per_epsilon: 1e-6, - per_capacity: hyperparams.buffer_size, curiosity_q_penalty_lambda: hyperparams.curiosity_q_penalty_lambda as f32, asymmetric_dd_weight: hyperparams.asymmetric_dd_weight as f32, ensemble_disagreement_penalty: hyperparams.ensemble_disagreement_penalty as f32, @@ -266,6 +291,12 @@ impl FusedTrainingCtx { gpu_weights::extract_branching_weights(target_vars, &stream) .map_err(|e| anyhow::anyhow!("Extract target branching weights: {e}"))?; + // Permanently disable cudarc event tracking BEFORE any CudaSlice allocation. + // This ensures all buffers have read=None, write=None, preventing SyncOnDrop + // from calling cuEventRecord which poisons error state during graph capture. + let _ = stream.context().check_err(); + unsafe { stream.context().disable_event_tracking(); } + // Create the fused trainer (compiles kernels, allocates buffers) let trainer = GpuDqnTrainer::new(stream.clone(), config) .map_err(|e| anyhow::anyhow!("GpuDqnTrainer init: {e}"))?; @@ -520,11 +551,6 @@ impl FusedTrainingCtx { gpu_iqn, cvar_scales_buf: None, gpu_attention, - graph_attention: None, - graph_iql: None, - graph_iqn: None, - graph_ema: None, - graph_her: None, ensemble_extra_heads, ensemble_diversity_weight: hyperparams.ensemble_diversity_weight as f32, ensemble_aggregate_kernel, @@ -536,6 +562,8 @@ impl FusedTrainingCtx { ensemble_kl_grad_kernel, ensemble_d_logits_buf, pending_vaccine_batch: None, + graph_aux: None, + tau_host: 0.0, }) } @@ -591,331 +619,80 @@ 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 _step_scalars = self.trainer.train_step_gpu( + // ── Step 2: Upload batch + replay graph_forward ────────────────── + let _fused_placeholder = self.trainer.train_step_gpu( gpu_batch, - &mut self.online_dueling, &mut self.online_branching, + &self.online_dueling, &self.online_branching, &self.target_dueling, &self.target_branching, ).map_err(|e| anyhow::anyhow!("Fused train_step_gpu (forward only): {e}"))?; - // HER relabeling via CUDA graph (Random) or ungraphed (Future/Final). + // ── Step 2b: HER donor computation (outside graph_aux) ─────────── + // Donor indices vary per step (random/future/final). The computation + // fills her.donor_indices GPU buffer. graph_aux captures the relabel + // kernel that READS from this stable buffer address. if let Some(ref mut her) = self.gpu_her { use crate::cuda_pipeline::gpu_her::HerGpuStrategy; - let her_batch_size = her.config.her_batch_size(); let batch_size = self.batch_size; - let state_dim = self.trainer.config().state_dim; - let goal_dim = her.config.goal_dim.min(state_dim); - let normal_count = batch_size.saturating_sub(her_batch_size); - match her.config.strategy { HerGpuStrategy::Random => { - if self.graph_her.is_none() { - // First step: run ungraphed - her.generate_random_donors_gpu(batch_size) - .map_err(|e| anyhow::anyhow!("HER random (first): {e}"))?; - let donor_ptr = her.donor_indices.raw_ptr(); - self.trainer.launch_her_inplace_relabel( - 0, donor_ptr, normal_count, her_batch_size, goal_dim, state_dim, - ).map_err(|e| anyhow::anyhow!("HER relabel (first): {e}"))?; - // Capture - self.trainer.sync_all_streams() - .map_err(|e| anyhow::anyhow!("sync before her capture: {e}"))?; - if let Ok(()) = self.stream.begin_capture( - cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL, - ) { - let _ = her.generate_random_donors_gpu(batch_size); - let donor_ptr = her.donor_indices.raw_ptr(); - let _ = self.trainer.launch_her_inplace_relabel( - 0, donor_ptr, normal_count, her_batch_size, goal_dim, state_dim, - ); - 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_her = Some(SendSyncGraph(graph)); - tracing::info!("Captured HER CUDA graph (random donors + relabel)"); - } - } - } else if let Some(ref graph) = self.graph_her { - graph.0.launch().map_err(|e| anyhow::anyhow!("HER graph replay: {e}"))?; - } + her.generate_random_donors_gpu(batch_size) + .map_err(|e| anyhow::anyhow!("HER random donors GPU: {e}"))?; } HerGpuStrategy::Future | HerGpuStrategy::Final => { - // Future/Final use episode_ids from gpu_batch (changes per step) let episode_ids = gpu_batch.episode_ids.as_ref().ok_or_else(|| { anyhow::anyhow!("HER {:?} requires episode_ids", her.config.strategy) })?; let source_indices = self.her_source_indices_gpu.as_ref() .ok_or_else(|| anyhow::anyhow!("HER source indices not pre-computed"))?; + let her_batch_size = her.config.her_batch_size(); her.relabel_batch_with_strategy( episode_ids, source_indices, 1, episode_ids.len(), her_batch_size, ).map_err(|e| anyhow::anyhow!("HER donor selection: {e}"))?; - let donor_ptr = her.donor_indices.raw_ptr(); - self.trainer.launch_her_inplace_relabel( - 0, donor_ptr, normal_count, her_batch_size, goal_dim, state_dim, - ).map_err(|e| anyhow::anyhow!("HER relabel: {e}"))?; } } } - // 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, - ); - // 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(), - ); - } - - 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.trainer.sync_all_streams() - .map_err(|e| anyhow::anyhow!("sync before ema capture: {e}"))?; - 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)"); - } - } - } else if let Some(ref graph) = self.graph_ema { - graph.0.launch().map_err(|e| anyhow::anyhow!("EMA graph replay: {e}"))?; - } + // ── Step 3: Auxiliary ops (graph_aux captured or replayed) ──────── + if self.graph_aux.is_some() { + // Fast path: replay captured graph. + // Pre-update host-side state (adam_steps, tau) so graphed HtoD + // copies the correct values to device buffers. + self.pre_replay_state_update(agent); + let graph_aux = self.graph_aux.as_ref().unwrap(); + graph_aux.launch(self.stream.cu_stream())?; + } else if self.steps_since_varmap_sync == 0 { + // First step: run ungraphed to initialize all sub-trainers. + self.submit_aux_ops(agent, gpu_batch)?; + } else { + // Second step: capture graph_aux. + self.submit_aux_ops(agent, gpu_batch)?; + // Capture on next call instead — we need the ops to execute first. + // Do the capture after this step completes so all buffers are initialized. + // (capture_graph_aux is called at the end of this step.) } - // ── Step 3b-5b: Auxiliary heads (attention + IQL + IQN) ───────── - // First step: run ungraphed to populate buffers, then capture graphs. - // Subsequent steps: replay captured graphs (1 launch each instead of ~21 kernels). - // IQN trunk gradient + EMA use tau that changes per step — these stay ungraphed - // but are only 2 kernel launches total. - if self.graph_attention.is_none() && self.gpu_attention.is_some() { - // First step: run attention ungraphed, then capture - if let Some(ref mut attn) = self.gpu_attention { - self.trainer.apply_attention_forward(attn, self.batch_size) - .map_err(|e| anyhow::anyhow!("Attention forward (first): {e}"))?; - let d_h_s2_bf16 = self.trainer.bw_d_h_s2_as_bf16() - .map_err(|e| anyhow::anyhow!("bw_d_h_s2 bf16 cast (first): {e}"))?; - attn.backward(d_h_s2_bf16, self.batch_size) - .map_err(|e| anyhow::anyhow!("Attention backward (first): {e}"))?; - let lr = self.trainer.config().lr; - let mgn = self.trainer.config().max_grad_norm; - attn.adam_step(lr, mgn) - .map_err(|e| anyhow::anyhow!("Attention Adam (first): {e}"))?; - } - // Capture attention graph - eprintln!("PRE-ATTN check_err: {:?}", self.stream.context().check_err()); - self.trainer.sync_all_streams() - .map_err(|e| { eprintln!("ATTN SYNC ERR: {e}"); anyhow::anyhow!("sync before attn capture: {e}") })?; - if let Ok(()) = self.stream.begin_capture( - cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL, - ) { - if let Some(ref mut attn) = self.gpu_attention { - if let Err(e) = self.trainer.apply_attention_forward(attn, self.batch_size) { - tracing::error!("CAPTURE: attention forward FAILED: {e}"); - } - match self.trainer.bw_d_h_s2_as_bf16() { - Ok(d) => { - if let Err(e) = attn.backward(d, self.batch_size) { - tracing::error!("CAPTURE: attention backward FAILED: {e}"); - } - } - Err(e) => tracing::error!("CAPTURE: bw_d_h_s2_as_bf16 FAILED: {e}"), - } - let lr = self.trainer.config().lr; - let mgn = self.trainer.config().max_grad_norm; - if let Err(e) = attn.adam_step(lr, mgn) { - tracing::error!("CAPTURE: attention adam FAILED: {e}"); - } - } - match self.stream.end_capture( - cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, - ) { - Ok(Some(graph)) => { - use crate::cuda_pipeline::gpu_dqn_trainer::SendSyncGraph; - self.graph_attention = Some(SendSyncGraph(graph)); - tracing::info!("Captured attention CUDA graph (fwd + bwd + adam)"); - } - Ok(None) => { - tracing::error!("Attention end_capture returned None"); - } - Err(e) => { - tracing::error!("Attention end_capture FAILED: {e} — stream stuck in capture mode!"); - // Force end capture via raw API to prevent cascade - unsafe { - let mut g: cudarc::driver::sys::CUgraph = std::ptr::null_mut(); - cudarc::driver::sys::cuStreamEndCapture(self.stream.cu_stream(), &mut g); - if !g.is_null() { cudarc::driver::sys::cuGraphDestroy(g); } - } - } - } - } - } else if let Some(ref graph) = self.graph_attention { - // Subsequent steps: replay attention graph - graph.0.launch().map_err(|e| anyhow::anyhow!("Attention graph replay: {e}"))?; - } - - // IQL graph - if self.graph_iql.is_none() && self.gpu_iql.is_some() { - if let Some(ref mut iql) = self.gpu_iql { - let _ = iql.train_value_step(self.trainer.states_buf(), self.trainer.rewards_buf()); - } - self.trainer.sync_all_streams() - .map_err(|e| { eprintln!("IQL SYNC ERR: {e}"); anyhow::anyhow!("sync before iql capture: {e}") })?; - if let Ok(()) = self.stream.begin_capture( - cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL, - ) { - if let Some(ref mut iql) = self.gpu_iql { - let _ = iql.train_value_step(self.trainer.states_buf(), self.trainer.rewards_buf()); - } - 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_iql = Some(SendSyncGraph(graph)); - tracing::info!("Captured IQL CUDA graph (value step)"); - } - } - } else if let Some(ref graph) = self.graph_iql { - graph.0.launch().map_err(|e| anyhow::anyhow!("IQL graph replay: {e}"))?; - } - - // IQN: most of the pipeline is graphed, but trunk gradient + target EMA - // use tau (changes per step) and write to online_dueling (shared state). - // Graph the main IQN forward+loss+backward+adam, keep trunk grad + EMA ungraphed. - // IQN: full pipeline + trunk gradient + EMA + PER DtoD in one graph. - // tau uploaded async before replay (GPU-resident tau_buf). - if self.graph_iqn.is_none() && self.gpu_iqn.is_some() { - // First step: run everything ungraphed - if let Some(ref mut iqn) = self.gpu_iqn { - let _ = iqn.train_iqn_step_gpu( - self.trainer.save_h_s2(), self.trainer.next_states_buf(), - &self.target_dueling, self.trainer.actions_buf(), - self.trainer.rewards_buf(), self.trainer.dones_buf(), - ); - let d_ptr = iqn.d_h_s2_raw_ptr(); - let _ = self.trainer.apply_iqn_trunk_gradient(d_ptr, &mut self.online_dueling); - let _ = iqn.target_ema_update(0.005); // initial tau - // IQN→PER DtoD - let bs = self.trainer.batch_size(); - let n_bytes = bs * std::mem::size_of::(); - unsafe { - let _ = cudarc::driver::result::memcpy_dtod_async( - self.trainer.td_errors_buf().raw_ptr(), - iqn.per_sample_loss().raw_ptr(), - n_bytes, self.stream.cu_stream(), - ); - } - } - // Capture the full IQN pipeline as one graph - self.trainer.sync_all_streams() - .map_err(|e| anyhow::anyhow!("sync before iqn capture: {e}"))?; - if let Ok(()) = self.stream.begin_capture( - cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL, - ) { - if let Some(ref mut iqn) = self.gpu_iqn { - let _ = iqn.train_iqn_step_gpu( - self.trainer.save_h_s2(), self.trainer.next_states_buf(), - &self.target_dueling, self.trainer.actions_buf(), - self.trainer.rewards_buf(), self.trainer.dones_buf(), - ); - let d_ptr = iqn.d_h_s2_raw_ptr(); - let _ = self.trainer.apply_iqn_trunk_gradient(d_ptr, &mut self.online_dueling); - let _ = iqn.target_ema_update(0.005); - let bs = self.trainer.batch_size(); - let n_bytes = bs * std::mem::size_of::(); - unsafe { - let _ = cudarc::driver::result::memcpy_dtod_async( - self.trainer.td_errors_buf().raw_ptr(), - iqn.per_sample_loss().raw_ptr(), - n_bytes, self.stream.cu_stream(), - ); - } - } - 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_iqn = Some(SendSyncGraph(graph)); - tracing::info!("Captured IQN CUDA graph (train + trunk grad + EMA + PER DtoD)"); - } - } - } else if let Some(ref graph) = self.graph_iqn { - // Upload tau before replay (IQN EMA reads from tau_buf) - if let Some(ref iqn) = self.gpu_iqn { - 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, - ); - unsafe { - cudarc::driver::sys::cuMemcpyHtoDAsync_v2( - iqn.tau_buf_ptr(), - (&(tau as f32) as *const f32).cast(), - std::mem::size_of::(), - self.stream.cu_stream(), - ); - } - } - graph.0.launch().map_err(|e| anyhow::anyhow!("IQN graph replay: {e}"))?; - } - - // ── Step 5b2: Ensemble multi-head diversity loss ────────────── - // Runs outside CUDA Graph. Head 0 is inside the graph; heads 1..K-1 - // run standalone cuBLAS value-head forward on save_h_s2. - // Diversity loss = λ × mean KL(head_i || head_j) across all pairs. - // Zero CPU in hot path — all ops on pre-allocated CudaSlice buffers. + // ── Step 4: Conditional ops (outside graph) ────────────────────── + // Ensemble diversity (variable topology, complex forward passes). if !self.ensemble_extra_heads.is_empty() { if let Err(e) = self.run_ensemble_step() { tracing::warn!("Ensemble diversity step failed (non-fatal): {e}"); } } - // Spectral norm moved to Step 1b (before graph_forward) — correct placement. + // Causal intervention. + { + let step = self.steps_since_varmap_sync; + if let Err(e) = self.trainer.run_causal_intervention(step) { + tracing::warn!("Causal intervention failed (non-fatal): {e}"); + } + } - // CQL gradient now captured in graph_adam — no per-step call needed. - - // ── Step 5d2: Causal Intervention (#34) — per-feature sensitivity ── - // #34 Causal intervention: always active (one production path) - // Causal intervention removed — 14 cuBLAS forwards every 100 steps - // with no readback or training decision based on the result. Pure waste. - - // ── Step 5e: Gradient Vaccine (#32) — project out overfitting gradients ── - // If a vaccine batch is provided, compute its gradient and project out - // the component of grad_buf that contradicts it. - // Vaccine runs after graphs stabilize (non-fatal on failure) - // Gradient vaccine runs every 10 steps — full ungraphed fwd+bwd per-step - // was the dominant bottleneck (~100-150ms/step). 10-step amortization cuts - // epoch wall-time by ~45% with negligible effect on gradient quality. + // Gradient vaccine (1/10 steps). if self.steps_since_varmap_sync % 10 == 0 { if let Some(ref vaccine_batch) = self.pending_vaccine_batch { if let Err(e) = self.trainer.apply_gradient_vaccine(vaccine_batch) { @@ -927,36 +704,45 @@ impl FusedTrainingCtx { } self.pending_vaccine_batch = None; - // Replay graph_adam (CQL + clip + pruning + Adam + regime_scale) - let fused_result = self.trainer.replay_adam_and_readback() - .map_err(|e| { eprintln!("!!! ADAM REPLAY FAILED: {e}"); anyhow::anyhow!("graph_adam replay: {e}") })?; + // ── Step 5: Pruning + Adam ─────────────────────────────────────── + self.trainer.apply_pruning_mask() + .map_err(|e| anyhow::anyhow!("Pruning mask apply: {e}"))?; - // PER priority update (outside graph — pointers change per step) + let fused_result = self.trainer.replay_adam_and_readback() + .map_err(|e| anyhow::anyhow!("graph_adam replay: {e}"))?; + + // ── Step 6: PER priority update (outside graph) ────────────────── if let (Some(priorities_tensor), Some((alpha, epsilon))) = (agent.priorities_tensor()?, agent.per_alpha_epsilon()) { self.trainer.update_priorities_cuda( - &gpu_batch.indices, priorities_tensor.data(), alpha, epsilon, + &gpu_batch.indices, + priorities_tensor.data(), + alpha, + epsilon, ).map_err(|e| anyhow::anyhow!("GPU PER priority update: {e}"))?; } - // training_steps++ and PER beta annealing (no CPU priority update) + // Bookkeeping. agent.fused_post_step_gpu_bookkeeping() .map_err(|e| anyhow::anyhow!("Fused GPU bookkeeping: {e}"))?; self.steps_since_varmap_sync += 1; self.last_combined_norm = fused_result.grad_norm; - // Single debug log per step — zero cost in release (compiled out) + // Capture graph_aux after the second step (all sub-trainers initialized). + if self.graph_aux.is_none() && self.steps_since_varmap_sync == 2 { + if let Err(e) = self.capture_graph_aux(agent, gpu_batch) { + tracing::warn!("graph_aux capture failed (non-fatal, continuing ungraphed): {e}"); + } + } + tracing::debug!( step = self.steps_since_varmap_sync, - iqn = self.gpu_iqn.is_some(), - cql = self.trainer.has_cql(), - ensemble_heads = self.ensemble_extra_heads.len(), + graphed = self.graph_aux.is_some(), "fused step complete" ); - // ── Step 7: Wrap raw scalars into GpuTrainResult ───────────────── GpuTrainResult::from_fused_scalars( fused_result.total_loss, fused_result.grad_norm, @@ -964,6 +750,249 @@ impl FusedTrainingCtx { ).map_err(|e| anyhow::anyhow!("Fused scalars->GpuTrainResult: {e}")) } + /// Submit all auxiliary ops between graph_forward and graph_adam. + /// + /// These ops are deterministic per-step (same kernel topology) and suitable + /// for CUDA graph capture. Conditional ops (vaccine, causal, ensemble) are + /// NOT included — they run outside the graph. + /// + /// Ops submitted: HER relabel, C51 clip, EMA target, attention, IQL, IQN, CQL, regime_scale. + fn submit_aux_ops( + &mut self, + agent: &mut DQNAgentType, + gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch, + ) -> Result<()> { + // HER in-place relabeling kernel (donor indices already pre-computed). + if let Some(ref her) = self.gpu_her { + let her_batch_size = her.config.her_batch_size(); + let state_dim = self.trainer.config().state_dim; + let goal_dim = her.config.goal_dim.min(state_dim); + let normal_count = self.batch_size.saturating_sub(her_batch_size); + + let src_next_ptr = gpu_batch.next_states.data().raw_ptr(); + let donor_ptr = her.donor_indices.raw_ptr(); + self.trainer.launch_her_inplace_relabel( + src_next_ptr, donor_ptr, + normal_count, her_batch_size, goal_dim, state_dim, + ).map_err(|e| anyhow::anyhow!("HER in-place relabel kernel: {e}"))?; + } + + // C51 gradient budget clip. + { + let cql_frac = if self.trainer.has_cql() { CQL_GRAD_BUDGET } else { 0.0 }; + let iqn_frac = if self.gpu_iqn.is_some() { IQN_GRAD_BUDGET } else { 0.0 }; + let ens_frac = if !self.ensemble_extra_heads.is_empty() { ENS_GRAD_BUDGET } else { 0.0 }; + let c51_frac = 1.0 - cql_frac - iqn_frac - ens_frac; + let c51_budget = self.trainer.config().max_grad_norm * c51_frac; + self.trainer.clip_grad_buf_inplace(c51_budget) + .map_err(|e| anyhow::anyhow!("C51 gradient budget clip: {e}"))?; + } + + // EMA target update. + { + 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, + ); + 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}"))?; + } + + // Attention forward + backward + Adam. + if let Some(ref mut attn) = self.gpu_attention { + self.trainer.apply_attention_forward(attn, self.batch_size) + .map_err(|e| anyhow::anyhow!("Attention forward: {e}"))?; + let d_h_s2_bf16 = self.trainer.bw_d_h_s2_as_bf16() + .map_err(|e| anyhow::anyhow!("bw_d_h_s2 bf16 cast: {e}"))?; + attn.backward(d_h_s2_bf16, self.batch_size) + .map_err(|e| anyhow::anyhow!("Attention backward: {e}"))?; + let lr = self.trainer.config().lr; + let mgn = self.trainer.config().max_grad_norm; + attn.adam_step(lr, mgn) + .map_err(|e| anyhow::anyhow!("Attention Adam: {e}"))?; + } + + // IQL value step. + if let Some(ref mut iql) = self.gpu_iql { + let states_f32 = self.trainer.states_buf(); + let rewards_f32 = self.trainer.rewards_buf(); + let _ = iql.train_value_step(states_f32, rewards_f32); + } + + // IQN full pipeline + trunk gradient + EMA. + if let Some(ref mut iqn) = self.gpu_iqn { + let online_h_s2 = self.trainer.save_h_s2(); + let next_states_buf = self.trainer.next_states_buf(); + let dqn_actions = self.trainer.actions_buf(); + let dqn_rewards = self.trainer.rewards_buf(); + let dqn_dones = self.trainer.dones_buf(); + + match iqn.train_iqn_step_gpu( + online_h_s2, next_states_buf, &self.target_dueling, + dqn_actions, dqn_rewards, dqn_dones, + ) { + Ok(_) => { + let d_h_s2_ptr = iqn.d_h_s2_raw_ptr(); + self.trainer.apply_iqn_trunk_gradient( + d_h_s2_ptr, &mut self.online_dueling, + ).map_err(|e| anyhow::anyhow!("IQN trunk gradient: {e}"))?; + + 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, + ); + iqn.target_ema_update(tau as f32) + .map_err(|e| anyhow::anyhow!("IQN EMA update: {e}"))?; + } + Err(e) => { + tracing::warn!("IQN step failed (non-fatal): {e}"); + } + } + } + + // IQN→PER loss DtoD. + if let Some(ref iqn) = self.gpu_iqn { + let bs = self.trainer.batch_size(); + let n_bytes = bs * std::mem::size_of::(); + let src_ptr = iqn.per_sample_loss().raw_ptr(); + let dst_ptr = self.trainer.td_errors_buf().raw_ptr(); + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, src_ptr, n_bytes, self.stream.cu_stream() + ).map_err(|e| anyhow::anyhow!("IQN→PER loss DtoD: {e}"))?; + } + } + + // CQL gradient + clip + SAXPY. + if self.trainer.has_cql() { + match self.trainer.apply_cql_gradient() { + Ok(true) => { + let cql_budget = self.trainer.config().max_grad_norm * CQL_GRAD_BUDGET; + self.trainer.apply_cql_clipped_saxpy(cql_budget) + .map_err(|e| anyhow::anyhow!("CQL clipped SAXPY: {e}"))?; + } + Ok(false) => {} + Err(e) => { + tracing::warn!("CQL gradient failed (non-fatal): {e}"); + } + } + } + + // Regime-adaptive PER scaling. + self.trainer.regime_scale_td_errors() + .map_err(|e| anyhow::anyhow!("Regime PER scaling: {e}"))?; + + Ok(()) + } + + /// Capture the auxiliary ops into a CUDA graph for zero-overhead replay. + /// + /// Called once after the first step succeeds (all sub-trainers initialized). + /// Uses raw CUDA driver API to bypass cudarc's bind_to_thread/check_err. + fn capture_graph_aux( + &mut self, + agent: &mut DQNAgentType, + gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch, + ) -> Result<()> { + // Sync both streams and drain stale errors before capture. + self.trainer.sync_all_streams() + .map_err(|e| anyhow::anyhow!("pre-capture sync: {e}"))?; + let _ = self.stream.context().check_err(); + + let cu_stream = self.stream.cu_stream(); + let mut graph: cuda_sys::CUgraph = std::ptr::null_mut(); + let mut exec: cuda_sys::CUgraphExec = std::ptr::null_mut(); + + // Begin capture. + let begin_result = unsafe { + cuda_sys::cuStreamBeginCapture_v2( + cu_stream, + cuda_sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL, + ) + }; + if begin_result != cuda_sys::cudaError_enum::CUDA_SUCCESS { + anyhow::bail!("cuStreamBeginCapture_v2 failed: {begin_result:?}"); + } + + // Submit all aux ops within capture scope. + let submit_result = self.submit_aux_ops(agent, gpu_batch); + + // End capture (must always be called, even if submit failed). + let end_result = unsafe { + cuda_sys::cuStreamEndCapture(cu_stream, &mut graph) + }; + + // Check submit result first. + submit_result?; + + if end_result != cuda_sys::cudaError_enum::CUDA_SUCCESS { + anyhow::bail!("cuStreamEndCapture failed: {end_result:?}"); + } + if graph.is_null() { + anyhow::bail!("cuStreamEndCapture returned null graph"); + } + + // Instantiate executable graph. + let inst_result = unsafe { + cuda_sys::cuGraphInstantiateWithFlags(&mut exec, graph, 0) + }; + if inst_result != cuda_sys::cudaError_enum::CUDA_SUCCESS { + unsafe { cuda_sys::cuGraphDestroy(graph); } + anyhow::bail!("cuGraphInstantiateWithFlags failed: {inst_result:?}"); + } + + info!("graph_aux captured: auxiliary ops (HER+clip+EMA+attn+IQL+IQN+CQL+regime)"); + self.graph_aux = Some(RawCudaGraph { exec, graph }); + Ok(()) + } + + /// Update sub-trainer host-side state before graph_aux replay. + /// + /// Graph captures HtoD from stable struct field addresses (adam_step, tau_host). + /// We update these host-side values before replay so the graphed HtoD copies + /// the correct values to device buffers. + fn pre_replay_state_update(&mut self, agent: &mut DQNAgentType) { + // Increment sub-trainer adam_steps so graphed HtoD copies updated values. + if let Some(ref mut attn) = self.gpu_attention { + attn.increment_adam_step(); + } + if let Some(ref mut iql) = self.gpu_iql { + iql.increment_adam_step(); + } + if let Some(ref mut iqn) = self.gpu_iqn { + iqn.increment_adam_step(); + // Update IQN tau via stable host address + 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, + ); + iqn.set_tau_host(tau as f32); + } + // Update DQN trainer tau via stable host address + { + 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, + ); + self.trainer.tau_host = tau as f32; + } + } + /// Run one ensemble diversity step (outside CUDA Graph). /// /// Heads 1..K-1 run their value/advantage head forward on the trunk activations diff --git a/docs/superpowers/plans/2026-04-02-mega-graph-training-loop-refactor.md b/docs/superpowers/plans/2026-04-02-mega-graph-training-loop-refactor.md new file mode 100644 index 000000000..3b3ef59ee --- /dev/null +++ b/docs/superpowers/plans/2026-04-02-mega-graph-training-loop-refactor.md @@ -0,0 +1,479 @@ +# Mega-Graph Training Loop Refactor Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Capture the ENTIRE DQN training step in a single CUDA graph, achieving ~1 graph launch per step with zero ungraphed kernel launches (except PER priority update and vaccine). + +**Architecture:** Move all graph capture logic from GpuDqnTrainer into FusedTrainingCtx which owns all state. Use raw CUDA driver API (cuStreamBeginCapture_v2 / cuStreamEndCapture / cuGraphLaunch) to bypass cudarc's bind_to_thread/check_err that conflicts with graph capture. Pre-allocate ALL CudaEvents before capture to avoid illegal cuEventCreate during capture. Disable cudarc event tracking permanently to prevent SyncOnDrop poisoning. + +**Tech Stack:** Rust 1.85, cudarc 0.17.3 (vendored), CUDA 13.0/12.x driver API, cuBLAS + +--- + +## Critical Bugs in Current Committed Code + +Before any refactoring, the committed code has bugs that must be fixed: + +1. **EMA kernel arg mismatch:** `ema_kernel.cu` expects `const float* tau_buf` but Rust passes `f32` scalar via `.arg(&tau)` → `CUDA_ERROR_ILLEGAL_ADDRESS` +2. **IQN EMA kernel arg mismatch:** Same issue — `iqn_ema_kernel` expects `const float* tau_buf` but may still receive scalar +3. **IQN/IQL/Attention Adam kernel arg mismatches:** Changed to `const int* t_buf` but Rust may still pass scalar `i32` +4. **PER update kernel signature changed:** `per_update_priorities_kernel` expects indirect pointers but callers pass direct pointers + +## File Structure + +### Modified Files + +| File | Changes | +|------|---------| +| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Fix kernel arg mismatches, pre-allocate events, make submit methods pub(crate), remove internal graph capture, add sync_all_streams | +| `crates/ml/src/cuda_pipeline/gpu_attention.rs` | Fix remaining device_ptr → raw_ptr | +| `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs` | Verify kernel arg compatibility | +| `crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs` | Verify kernel arg compatibility | +| `crates/ml/src/trainers/dqn/fused_training.rs` | Add RawCudaGraph, mega-graph capture in FusedTrainingCtx, rewrite run_full_step | +| `crates/ml/src/cuda_pipeline/ema_kernel.cu` | Revert to scalar tau OR fix Rust caller | +| `crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu` | Verify ema/adam kernel signatures | +| `crates/ml/src/cuda_pipeline/iql_value_kernel.cu` | Verify adam kernel signature | +| `crates/ml/src/cuda_pipeline/attention_backward_kernel.cu` | Verify adam kernel signature | +| `crates/ml/src/cuda_pipeline/per_update_kernel.cu` | Revert to direct pointers OR fix Rust caller | + +--- + +### Task 1: Fix EMA kernel arg mismatch (critical — all tests fail without this) + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:5443-5454` (target_ema_update) + +The EMA CUDA kernel was changed to read `tau` from a device buffer (`const float* tau_buf`) but the Rust code still passes a scalar `f32` via `.arg(&tau)`. The kernel interprets the scalar value as a device pointer → ILLEGAL_ADDRESS crash. + +**Decision: revert the kernel to accept scalar tau, OR fix the Rust caller.** + +For the mega-graph, we need the device buffer approach (scalar gets baked at capture time). So: fix the Rust caller to use `tau_buf`. + +- [ ] **Step 1: Check the committed EMA kernel signature** + +```bash +grep "tau_buf\|float tau\|float\* tau" crates/ml/src/cuda_pipeline/ema_kernel.cu +``` + +Expected: `const float* __restrict__ tau_buf` — kernel reads from device buffer. + +- [ ] **Step 2: Check if tau_buf field exists in GpuDqnTrainer** + +```bash +grep -n "tau_buf" crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | head -5 +``` + +Expected: Field `tau_buf: CudaSlice` exists (added in earlier commit). + +- [ ] **Step 3: Fix the EMA kernel launch to pass tau_buf pointer** + +In `gpu_dqn_trainer.rs`, find `target_ema_update` function. Replace: +```rust +.arg(&tau) +``` +With: +```rust +.arg(&tau_ptr) +``` +Where `tau_ptr` is obtained from the async HtoD upload that writes tau to `self.tau_buf`. + +Add before the launch: +```rust +// Async HtoD for tau (graph-capture compatible) +unsafe { + cudarc::driver::sys::cuMemcpyHtoDAsync_v2( + self.tau_buf.raw_ptr(), + (&tau as *const f32).cast(), + std::mem::size_of::(), + self.stream.cu_stream(), + ); +} +let tau_ptr = self.tau_buf.raw_ptr(); +``` + +And change `.arg(&tau)` to `.arg(&tau_ptr)`. + +- [ ] **Step 4: Compile and test** + +```bash +SQLX_OFFLINE=true cargo check -p ml +SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests::training_stability::test_gpu_collector_auto_initializes --ignored +``` + +Expected: PASS (if no other kernel mismatches) + +- [ ] **Step 5: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +git commit -m "fix: EMA kernel tau passed as device pointer, not scalar — fixes ILLEGAL_ADDRESS" +``` + +--- + +### Task 2: Fix ALL remaining kernel arg mismatches + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs` (IQN adam + ema) +- Modify: `crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs` (IQL adam) +- Modify: `crates/ml/src/cuda_pipeline/gpu_attention.rs` (attention adam) +- Modify: `crates/ml/src/cuda_pipeline/per_update_kernel.cu` (PER update) +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (PER update caller) + +For each modified CUDA kernel, verify the Rust caller passes the correct type (device pointer vs scalar). Fix any mismatches. + +- [ ] **Step 1: Audit IQN adam kernel (iqn_dual_head_kernel.cu)** + +Check if `iqn_adam_kernel` expects `const int* adam_t_buf` or `int adam_t`. Cross-reference with the Rust `.arg()` call in `gpu_iqn_head.rs`. + +- [ ] **Step 2: Audit IQN ema kernel (iqn_dual_head_kernel.cu)** + +Check if `iqn_ema_kernel` expects `const float* tau_buf` or `float tau`. Cross-reference with Rust caller. + +- [ ] **Step 3: Audit IQL adam kernel (iql_value_kernel.cu)** + +Check if `iql_adam_kernel` expects `const int* t_buf` or `int t`. Cross-reference with Rust caller. + +- [ ] **Step 4: Audit attention adam kernel (attention_backward_kernel.cu)** + +Check if `attn_adam_kernel` expects `const int* adam_t_buf` or `int adam_t`. Cross-reference with Rust caller. + +- [ ] **Step 5: Audit PER update kernel (per_update_kernel.cu)** + +Check if `per_update_priorities_kernel` expects indirect pointers (`const unsigned long long*`) or direct pointers (`const unsigned int*`, `__nv_bfloat16*`). If indirect, the Rust caller must pass `batch_ptr_buf` offsets. If direct, pass CudaSlice references. + +- [ ] **Step 6: Fix ALL mismatches found** + +For each mismatch, either: +a) Revert the kernel to the original signature (if we don't need graph-capture compatibility yet) +b) Fix the Rust caller to pass the correct type + +- [ ] **Step 7: Compile and run full smoke test suite** + +```bash +SQLX_OFFLINE=true cargo check --workspace +SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | tail -5 +``` + +Expected: All smoke tests pass. + +- [ ] **Step 8: Commit** + +```bash +git add -A +git commit -m "fix: all CUDA kernel arg mismatches — device ptr vs scalar compatibility" +``` + +--- + +### Task 3: Verify baseline performance with working code + +**Files:** None (testing only) + +After fixing all kernel mismatches, run the smoke test and verify per-step timing. This establishes the baseline before the mega-graph refactor. + +- [ ] **Step 1: Run full smoke test** + +```bash +SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep "PASSED\|FAILED\|per-step\|step breakdown" +``` + +Expected: Tests pass. Per-step timing logged. + +- [ ] **Step 2: Record baseline per-step time** + +Note the per-step fused time from the test output. This is the number we're trying to beat with the mega-graph. + +--- + +### Task 4: Disable cudarc event tracking permanently + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` (FusedTrainingCtx::new) + +Disable cudarc event tracking before ANY CudaSlice allocation. This ensures all buffers have `read=None, write=None`, preventing SyncOnDrop from calling cuEventRecord during graph capture. + +- [ ] **Step 1: Add event tracking disable at init** + +In `FusedTrainingCtx::new`, BEFORE `GpuDqnTrainer::new`: +```rust +let _ = stream.context().check_err(); +unsafe { stream.context().disable_event_tracking(); } +``` + +- [ ] **Step 2: Compile and test** + +```bash +SQLX_OFFLINE=true cargo check -p ml +SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests::training_stability::test_gpu_collector_auto_initializes --ignored +``` + +Expected: PASS (existing 2-graph capture still works with tracking disabled) + +- [ ] **Step 3: Commit** + +```bash +git commit -m "perf: disable cudarc event tracking permanently — SyncOnDrop-safe for graph capture" +``` + +--- + +### Task 5: Pre-allocate multi-stream sync events + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (struct + constructor + submit_forward_ops) + +Replace `self.stream.record_event(None)` calls in `submit_forward_ops` with pre-allocated events. `cuEventCreate` during capture is legal but fragile; pre-allocation is cleaner. + +- [ ] **Step 1: Add event fields to GpuDqnTrainer struct** + +```rust +pass1_event: CudaEvent, +pass3_event: CudaEvent, +``` + +- [ ] **Step 2: Allocate in constructor (BEFORE graph capture)** + +```rust +let pass1_event = stream.record_event(Some(sys::CUevent_flags::CU_EVENT_DISABLE_TIMING))?; +let pass3_event = double_dqn_stream.record_event(Some(sys::CUevent_flags::CU_EVENT_DISABLE_TIMING))?; +``` + +- [ ] **Step 3: Replace record_event calls in submit_forward_ops** + +Replace: +```rust +let pass1_done = self.stream.record_event(None)?; +``` +With: +```rust +self.pass1_event.record(&self.stream)?; +``` +And use `&self.pass1_event` instead of `&pass1_done`. + +Same for `pass3_done` → `self.pass3_event`. + +- [ ] **Step 4: Compile and test** + +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git commit -m "perf: pre-allocate multi-stream sync events — no cuEventCreate during capture" +``` + +--- + +### Task 6: Make GpuDqnTrainer submit methods pub(crate) + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +Make all submit methods accessible from FusedTrainingCtx for the mega-graph capture. + +- [ ] **Step 1: Change visibility** + +```rust +pub(crate) fn submit_forward_ops(...) +pub(crate) fn submit_adam_ops(...) +pub(crate) fn submit_indirect_upload_ops(...) +pub(crate) fn submit_cql_ops(...) +pub(crate) fn submit_c51_clip_ops(...) +pub(crate) fn submit_pruning_mask_ops(...) +pub(crate) fn submit_ema_ops(...) +pub(crate) fn flatten_online_weights(...) +pub(crate) params_initialized: bool +pub(crate) graph_forward: Option +``` + +Also add: +```rust +pub(crate) fn sync_all_streams(&self) -> Result<(), MLError> { + self.stream.synchronize()?; + self.double_dqn_stream.synchronize()?; + Ok(()) +} + +pub(crate) fn adam_step_async(&mut self) { ... } +``` + +- [ ] **Step 2: Compile and test** + +Expected: PASS (no behavior change) + +- [ ] **Step 3: Commit** + +--- + +### Task 7: Implement RawCudaGraph and mega-graph capture + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` + +This is the core refactor. Add `RawCudaGraph` struct that bypasses cudarc, implement `capture_mega_graph` on FusedTrainingCtx, and rewrite `run_full_step`. + +- [ ] **Step 1: Add RawCudaGraph struct** + +Raw CUDA graph handle that uses `cuGraphLaunch` directly (no cudarc bind_to_thread/check_err): + +```rust +struct RawCudaGraph { + exec: cudarc::driver::sys::CUgraphExec, + graph: cudarc::driver::sys::CUgraph, +} + +impl RawCudaGraph { + fn launch(&self, stream: cudarc::driver::sys::CUstream) -> Result<()> { ... } +} + +impl Drop for RawCudaGraph { + fn drop(&mut self) { cuGraphExecDestroy + cuGraphDestroy } +} + +unsafe impl Send for RawCudaGraph {} +unsafe impl Sync for RawCudaGraph {} +``` + +- [ ] **Step 2: Add mega_graph field to FusedTrainingCtx** + +Replace individual graph fields (graph_attention, graph_iql, graph_iqn, graph_ema, graph_her) with: +```rust +mega_graph: Option, +``` + +- [ ] **Step 3: Implement capture_mega_graph method** + +Uses raw `cuStreamBeginCapture_v2` / `cuStreamEndCapture` (bypasses cudarc): + +```rust +fn capture_mega_graph(&mut self) -> Result<()> { + self.trainer.sync_all_streams()?; + let _ = self.stream.context().check_err(); // drain stale errors + + // Raw begin_capture + unsafe { sys::cuStreamBeginCapture_v2(cu_stream, MODE) }; + + // Submit ALL ops in sequence: + self.trainer.submit_forward_ops()?; // includes double_dqn_stream multi-stream + // HER relabel + // Attention fwd + bwd + adam + // IQL value step + // IQN full pipeline + self.trainer.submit_cql_ops()?; + self.trainer.submit_c51_clip_ops()?; + self.trainer.submit_pruning_mask_ops()?; + self.trainer.submit_adam_ops(...)?; + self.trainer.submit_ema_ops(...)?; + self.trainer.regime_scale_td_errors()?; + + // Raw end_capture + instantiate + unsafe { sys::cuStreamEndCapture(cu_stream, &mut graph) }; + unsafe { sys::cuGraphInstantiateWithFlags(&mut exec, graph, 0) }; + + self.mega_graph = Some(RawCudaGraph { exec, graph }); +} +``` + +- [ ] **Step 4: Rewrite run_full_step** + +New flow: +1. Init weights if needed +2. Upload batch pointers + params (async HtoD) +3. If mega_graph is None: run all ops ungraphed, then capture_mega_graph +4. Else: single mega_graph.launch() +5. PER priority update (ungraphed) +6. Vaccine (ungraphed, 1/10 steps) +7. Bookkeeping + +- [ ] **Step 5: Remove individual graph fields and capture blocks** + +Delete: graph_attention, graph_iql, graph_iqn, graph_ema, graph_her fields and all their capture/replay code in run_full_step. + +- [ ] **Step 6: Compile and test** + +```bash +SQLX_OFFLINE=true cargo check --workspace +SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | tail -10 +``` + +Expected: All smoke tests pass with mega-graph. + +- [ ] **Step 7: Commit** + +```bash +git commit -m "perf: mega-graph — entire training step in single CUDA graph, 1 launch per step" +``` + +--- + +### Task 8: Remove dead code and cleanup + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (remove internal capture) +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` (remove debug logging) + +- [ ] **Step 1: Remove capture_training_graphs from GpuDqnTrainer** + +The mega-graph capture is now in FusedTrainingCtx. Remove the old method and update any remaining callers (e.g., in train_step_gpu) to not call it. + +- [ ] **Step 2: Remove EventTrackingGuard usage** + +With event tracking permanently disabled, EventTrackingGuard is a no-op. Remove all instances except the struct definition (may be used elsewhere). + +- [ ] **Step 3: Remove debug eprintln! statements** + +Remove all `eprintln!("MEGA:...")`, `eprintln!("FWD:...")`, `eprintln!("ATTN FWD:...")` debugging lines. + +- [ ] **Step 4: Compile and run full test suite** + +```bash +SQLX_OFFLINE=true cargo check --workspace +SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored +``` + +Expected: All pass, zero warnings. + +- [ ] **Step 5: Commit and push** + +```bash +git commit -m "refactor: remove dead graph capture code and debug logging" +git push origin main +``` + +--- + +### Task 9: Deploy to H100 and measure performance + +**Files:** None (deployment only) + +- [ ] **Step 1: Launch H100 training** + +```bash +argo submit -n foxhunt --from workflowtemplate/compile-and-train \ + -p commit-sha=$(git rev-parse --short HEAD) \ + -p model=dqn -p gpu-pool=ci-training-h100 \ + -p hyperopt-trials=0 -p train-epochs=5 +``` + +- [ ] **Step 2: Verify per-step timing** + +Check training logs for `Training step breakdown` line: +- Before: `per-step: fused=129.0ms` (2-graph approach with syncs) +- Target: `per-step: fused=<10ms` (mega-graph, zero ungraphed) + +- [ ] **Step 3: Verify training correctness** + +Check Sharpe, PF, Q-values, action diversity match expected ranges. + +--- + +## Risk Register + +| Risk | Impact | Mitigation | +|------|--------|-----------| +| cudarc launch_builder incompatible with capture | High | Raw cuLaunchKernel as fallback; but research shows launch_builder IS compatible with disabled event tracking | +| double_dqn_stream capture isolation | High | Pre-allocated events + cuStreamWaitEvent automatic join (proven working in current code) | +| Kernel arg type mismatches crash GPU | Critical | Task 1-2 fix ALL mismatches before any refactoring | +| PER update can't be graphed (pointer changes) | Low | Stays outside graph — 1 ungraphed kernel per step (acceptable) | +| cudarc check_err poisons during mega-graph capture | Medium | Drain errors before capture + event tracking disabled = no new errors recorded |