From ad5c2b7ded5871ff6c48196e821e4b5b389c5db1 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 11 Apr 2026 23:10:25 +0200 Subject: [PATCH] =?UTF-8?q?refactor:=20exp=20collector=20zero-copy=20?= =?UTF-8?q?=E2=80=94=20direct=20pointer=20into=20trainer=20params=5Fbuf?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Dead code sweep (167 lines deleted): - online_params_flat (synced but never read by forward) - online_params_f32 (never synced, always zeros — ROOT CAUSE of garbage Q-values) - DuelingWeightSet/BranchingWeightSet/CuriosityWeightSet/RmsNormWeightSet fields - sync_weights_flat, sync_weights_f32, flatten_online_weights methods - sync_gpu_weights epoch-boundary DtoD copy Zero-copy replacement: - trainer_params_ptr: u64 points directly into trainer's params_buf - param_sizes computed from real config (bottleneck_dim=16, market_dim=42) - CublasForward with correct s1_input_dim=54 for bottleneck - Bottleneck forward activated (bn_hidden + tanh + concat) - Fused context init moved before Phase 2 (experience collection) Co-Authored-By: Claude Opus 4.6 (1M context) --- .../cuda_pipeline/gpu_experience_collector.rs | 271 ++++-------------- .../src/trainers/dqn/trainer/training_loop.rs | 70 +++-- crates/ml/tests/smoke_test_real_data.rs | 48 +--- .../2026-04-11-exp-collector-zero-copy.md | 251 ++++++++++++++++ 4 files changed, 343 insertions(+), 297 deletions(-) create mode 100644 docs/superpowers/plans/2026-04-11-exp-collector-zero-copy.md diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 05699a168..6f1c5d051 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -26,12 +26,9 @@ use crate::MLError; use super::batched_forward::CublasForward; use super::gpu_curiosity_trainer::GpuCuriosityTrainer; use super::gpu_dqn_trainer::{ - GpuDqnTrainConfig, compute_param_sizes, compute_total_params, dtod_copy, -}; -use super::gpu_weights::{ - BranchingWeightSet, CuriosityWeightSet, DuelingWeightSet, - RmsNormWeightSet, + GpuDqnTrainConfig, compute_param_sizes, compute_total_params, }; +use super::gpu_weights::CuriosityWeightSet; // --------------------------------------------------------------------------- // Constants @@ -473,10 +470,13 @@ pub struct GpuExperienceCollector { // ── cuBLAS forward pass context ───────────────────────────────── cublas_forward: CublasForward, - // ── Flat F32 online weight buffer for cuBLAS ──────────────────── - /// Flattened online network parameters [total_params] — single DtoD copy for sync. - online_params_flat: CudaSlice, - /// Total number of F32 parameters in the flat buffer. + // ── Zero-copy trainer params pointer ───────────────────────────── + /// Raw device pointer into the trainer's `params_buf`. + /// The collector reads weights directly — no DtoD copy needed. + /// Set to 0 initially; updated via `set_trainer_params_ptr()` once + /// the fused training context is created. + trainer_params_ptr: u64, + /// Total number of F32 parameters (for info logging only). total_params: usize, /// Per-tensor sizes for computing weight pointers. param_sizes: [usize; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS], @@ -535,14 +535,8 @@ pub struct GpuExperienceCollector { exp_v_logits: CudaSlice, // [N, num_atoms] f32 (output layer — no f32 truncation) exp_b_logits: CudaSlice, // [N, total_branch_atoms] f32 - // ── Weight sets (pointer views + backing storage) ── - online_weights: DuelingWeightSet, - target_weights: DuelingWeightSet, + // ── Weight sets (only curiosity — dueling/branching/rmsnorm deleted, zero-copy from trainer) ── curiosity_weights: CuriosityWeightSet, - online_rmsnorm: RmsNormWeightSet, - target_rmsnorm: RmsNormWeightSet, - online_branching: BranchingWeightSet, - target_branching: BranchingWeightSet, /// GPU buffer for trade_stats_reduce output: [TRADE_STATS_FLOATS] = 6 floats. trade_stats_buf: CudaSlice, @@ -592,9 +586,6 @@ pub struct GpuExperienceCollector { exp_bn_hidden: CudaSlice, /// #31 Bottleneck concat buffer [alloc_episodes, bn_dim + portfolio_dim] f32. exp_bn_concat: CudaSlice, - /// #30 F32 master weight buffer — full-precision Q-network for experience collection. - /// Synced from trainer's f32 master params (no f32 quantization noise). - online_params_f32: CudaSlice, /// #30 F32 activation buffers [alloc_episodes, dim]. exp_h_s1_f32: CudaSlice, exp_h_s2_f32: CudaSlice, @@ -676,14 +667,16 @@ impl GpuExperienceCollector { /// Create a new GPU experience collector with cuBLAS-based Q-network forward pass. /// /// Compiles 3 focused NVRTC kernels, creates a `CublasForward` context for - /// the Q-network, extracts initial weights, and allocates all GPU buffers. + /// the Q-network, and allocates all GPU buffers. + /// + /// The collector reads Q-network weights directly from the trainer's flat + /// `params_buf` via `trainer_params_ptr` (zero-copy). Call + /// `set_trainer_params_ptr()` after the fused training context is created. /// /// # Arguments /// * `stream` - CUDA stream for all GPU operations - /// * `online_dueling` - Dueling weight set for the online Q-network (pointer views) - /// * `target_dueling` - Dueling weight set for the target Q-network (pointer views) - /// * `online_branching` - Branching weight set for the online Q-network (pointer views) - /// * `target_branching` - Branching weight set for the target Q-network (pointer views) + /// * `bottleneck_dim` - Temporal causal bottleneck dimension (16 in production, 0 = disabled) + /// * `market_dim_cfg` - Market feature dimension for bottleneck separation (42) /// * `initial_capital` - Starting capital for each episode /// * `_avg_spread` - Average bid-ask spread /// * `_cash_reserve_pct` - Cash reserve percentage @@ -694,10 +687,8 @@ impl GpuExperienceCollector { #[allow(clippy::too_many_arguments)] pub fn new( stream: Arc, - online_dueling: &DuelingWeightSet, - target_dueling: &DuelingWeightSet, - online_branching: &BranchingWeightSet, - target_branching: &BranchingWeightSet, + bottleneck_dim: usize, + market_dim_cfg: usize, initial_capital: f32, _avg_spread: f32, _cash_reserve_pct: f32, @@ -741,6 +732,14 @@ impl GpuExperienceCollector { // are plain Q-values. For C51 (num_atoms_max>1), we compute // expected Q via compute_expected_q kernel. let num_atoms = num_atoms_max.max(1); + // s1_input_dim must match the trainer's layout: + // with bottleneck: compressed market → bn_dim, then concat with portfolio → bn_dim + portfolio_dim + // without bottleneck: raw state_dim + let s1_input_dim = if bottleneck_dim > 0 { + bottleneck_dim + state_dim.saturating_sub(market_dim_cfg) + } else { + state_dim + }; let cublas_forward = CublasForward::new( &stream, alloc_episodes, // batch_size = number of episodes @@ -754,7 +753,7 @@ impl GpuExperienceCollector { branch_sizes[1], branch_sizes[2], branch_sizes[3], - state_dim, // s1_input_dim: experience collector uses full state_dim (no bottleneck) + s1_input_dim, )?; info!( alloc_episodes, @@ -763,10 +762,8 @@ impl GpuExperienceCollector { ); // ── Step 3: Compute flat parameter layout ─────────────────────── - // bottleneck_dim and market_dim MUST be 0 here — the experience collector - // doesn't use a bottleneck (it feeds raw states to cuBLAS). - // Using Default::default() gave bottleneck_dim=2, market_dim=42, which made - // s1_input_dim=40 instead of state_dim=80 → cuBLAS OOB reads on W1. + // Must match the trainer's layout exactly (same bottleneck_dim, market_dim) + // so that f32_weight_ptrs_from_base() returns correct offsets. let train_cfg = GpuDqnTrainConfig { state_dim, shared_h1, @@ -778,102 +775,25 @@ impl GpuExperienceCollector { branch_1_size: branch_sizes[1], branch_2_size: branch_sizes[2], branch_3_size: 3, // urgency — fixed at 3 (4-branch refactor) - bottleneck_dim: 0, - market_dim: 0, + bottleneck_dim, + market_dim: market_dim_cfg, ..GpuDqnTrainConfig::default() }; let param_sizes = compute_param_sizes(&train_cfg); let total_params = compute_total_params(&train_cfg); - // #31 Bottleneck dimension (derived from param_sizes[24] = w_bn) - let bn_dim_from_params = if param_sizes[24] > 0 && market_dim > 0 { - param_sizes[24] / market_dim - } else { - 0 - }; - - // Allocate flat online parameter buffer - let online_params_flat = stream - .alloc_zeros::(total_params) - .map_err(|e| MLError::ModelError(format!("alloc online_params_flat: {e}")))?; - info!( total_params, total_bytes = total_params * 4, - "GPU experience collector: flat parameter buffer allocated" + bottleneck_dim, + market_dim_cfg, + "GPU experience collector: parameter layout computed (zero-copy from trainer)" ); - // ── Step 4: Copy weight set pointer views from caller ───────── - // Weight sets are raw u64 device-pointer + usize element-count pairs. - // They point into the caller's params_buf (zero-copy). We just copy - // the struct fields — no GPU allocation needed. - let online_weights = DuelingWeightSet { - w_s1: online_dueling.w_s1, w_s1_len: online_dueling.w_s1_len, - b_s1: online_dueling.b_s1, b_s1_len: online_dueling.b_s1_len, - w_s2: online_dueling.w_s2, w_s2_len: online_dueling.w_s2_len, - b_s2: online_dueling.b_s2, b_s2_len: online_dueling.b_s2_len, - w_v1: online_dueling.w_v1, w_v1_len: online_dueling.w_v1_len, - b_v1: online_dueling.b_v1, b_v1_len: online_dueling.b_v1_len, - w_v2: online_dueling.w_v2, w_v2_len: online_dueling.w_v2_len, - b_v2: online_dueling.b_v2, b_v2_len: online_dueling.b_v2_len, - w_a1: online_dueling.w_a1, w_a1_len: online_dueling.w_a1_len, - b_a1: online_dueling.b_a1, b_a1_len: online_dueling.b_a1_len, - w_a2: online_dueling.w_a2, w_a2_len: online_dueling.w_a2_len, - b_a2: online_dueling.b_a2, b_a2_len: online_dueling.b_a2_len, - }; - let target_weights = DuelingWeightSet { - w_s1: target_dueling.w_s1, w_s1_len: target_dueling.w_s1_len, - b_s1: target_dueling.b_s1, b_s1_len: target_dueling.b_s1_len, - w_s2: target_dueling.w_s2, w_s2_len: target_dueling.w_s2_len, - b_s2: target_dueling.b_s2, b_s2_len: target_dueling.b_s2_len, - w_v1: target_dueling.w_v1, w_v1_len: target_dueling.w_v1_len, - b_v1: target_dueling.b_v1, b_v1_len: target_dueling.b_v1_len, - w_v2: target_dueling.w_v2, w_v2_len: target_dueling.w_v2_len, - b_v2: target_dueling.b_v2, b_v2_len: target_dueling.b_v2_len, - w_a1: target_dueling.w_a1, w_a1_len: target_dueling.w_a1_len, - b_a1: target_dueling.b_a1, b_a1_len: target_dueling.b_a1_len, - w_a2: target_dueling.w_a2, w_a2_len: target_dueling.w_a2_len, - b_a2: target_dueling.b_a2, b_a2_len: target_dueling.b_a2_len, - }; - // Curiosity always zero-initialized (hot-path curiosity uses fused CUDA). info!("GPU experience collector: curiosity disabled, using zero weights"); let curiosity_weights = CuriosityWeightSet::zeros(&stream)?; - // RMSNorm: all-ones (no-op scaling). Real RMSNorm weights are in params_buf. - let online_rmsnorm = RmsNormWeightSet::ones(&stream, network_dims)?; - let target_rmsnorm = RmsNormWeightSet::ones(&stream, network_dims)?; - - info!("GPU experience collector: copying branching weight set pointer views"); - let online_branching_ws = BranchingWeightSet { - w_bo1: online_branching.w_bo1, w_bo1_len: online_branching.w_bo1_len, - b_bo1: online_branching.b_bo1, b_bo1_len: online_branching.b_bo1_len, - w_bo2: online_branching.w_bo2, w_bo2_len: online_branching.w_bo2_len, - b_bo2: online_branching.b_bo2, b_bo2_len: online_branching.b_bo2_len, - w_bu1: online_branching.w_bu1, w_bu1_len: online_branching.w_bu1_len, - b_bu1: online_branching.b_bu1, b_bu1_len: online_branching.b_bu1_len, - w_bu2: online_branching.w_bu2, w_bu2_len: online_branching.w_bu2_len, - b_bu2: online_branching.b_bu2, b_bu2_len: online_branching.b_bu2_len, - w_bg1: online_branching.w_bg1, w_bg1_len: online_branching.w_bg1_len, - b_bg1: online_branching.b_bg1, b_bg1_len: online_branching.b_bg1_len, - w_bg2: online_branching.w_bg2, w_bg2_len: online_branching.w_bg2_len, - b_bg2: online_branching.b_bg2, b_bg2_len: online_branching.b_bg2_len, - }; - let target_branching_ws = BranchingWeightSet { - w_bo1: target_branching.w_bo1, w_bo1_len: target_branching.w_bo1_len, - b_bo1: target_branching.b_bo1, b_bo1_len: target_branching.b_bo1_len, - w_bo2: target_branching.w_bo2, w_bo2_len: target_branching.w_bo2_len, - b_bo2: target_branching.b_bo2, b_bo2_len: target_branching.b_bo2_len, - w_bu1: target_branching.w_bu1, w_bu1_len: target_branching.w_bu1_len, - b_bu1: target_branching.b_bu1, b_bu1_len: target_branching.b_bu1_len, - w_bu2: target_branching.w_bu2, w_bu2_len: target_branching.w_bu2_len, - b_bu2: target_branching.b_bu2, b_bu2_len: target_branching.b_bu2_len, - w_bg1: target_branching.w_bg1, w_bg1_len: target_branching.w_bg1_len, - b_bg1: target_branching.b_bg1, b_bg1_len: target_branching.b_bg1_len, - w_bg2: target_branching.w_bg2, w_bg2_len: target_branching.w_bg2_len, - b_bg2: target_branching.b_bg2, b_bg2_len: target_branching.b_bg2_len, - }; - // ── Step 5: L2 cache persistence hints (Hopper GPUs) ──────────── { let caps = ml_core::gpu::capabilities::cached_capabilities(); @@ -1051,8 +971,6 @@ impl GpuExperienceCollector { // #30 F32 forward pass buffers for full-precision experience collection let pad_sd = (state_dim + 127) & !127; - let online_params_f32 = stream.alloc_zeros::(total_params) - .map_err(|e| MLError::ModelError(format!("alloc online_params_f32: {e}")))?; let exp_states_f32 = stream.alloc_zeros::(alloc_episodes * pad_sd) .map_err(|e| MLError::ModelError(format!("alloc exp_states_f32: {e}")))?; let exp_h_s1_f32 = stream.alloc_zeros::(alloc_episodes * network_dims.0) @@ -1111,11 +1029,11 @@ impl GpuExperienceCollector { let bn_tanh_concat_fn = util_module.load_function("bn_tanh_concat_f32_kernel") .map_err(|e| MLError::ModelError(format!("bn_tanh_concat_f32_kernel load: {e}")))?; - let bn_alloc = bn_dim_from_params.max(1); - let portfolio_dim_bn = state_dim - market_dim; - let exp_bn_hidden = stream.alloc_zeros::(alloc_episodes * bn_alloc) + let bn_alloc = bottleneck_dim.max(1); + let portfolio_dim_bn = state_dim.saturating_sub(market_dim_cfg); + let exp_bn_hidden = stream.alloc_zeros::(alloc_episodes * bn_alloc + 128) .map_err(|e| MLError::ModelError(format!("alloc exp_bn_hidden: {e}")))?; - let exp_bn_concat = stream.alloc_zeros::(alloc_episodes * (bn_alloc + portfolio_dim_bn)) + let exp_bn_concat = stream.alloc_zeros::(alloc_episodes * (bn_alloc + portfolio_dim_bn) + 128) .map_err(|e| MLError::ModelError(format!("alloc exp_bn_concat: {e}")))?; // Task 8: GPU-resident step counter for experience collection loop @@ -1135,7 +1053,7 @@ impl GpuExperienceCollector { alloc_episodes, alloc_timesteps, cublas_forward, - online_params_flat, + trainer_params_ptr: 0, // Set via set_trainer_params_ptr() after fused ctx init total_params, param_sizes, cvar_scales_ptr: 0, // NULL = no CVaR scaling initially @@ -1165,13 +1083,7 @@ impl GpuExperienceCollector { exp_h_v, exp_v_logits, exp_b_logits, - online_weights, - target_weights, curiosity_weights, - online_rmsnorm, - target_rmsnorm, - online_branching: online_branching_ws, - target_branching: target_branching_ws, portfolio_states, rng_states, episode_starts_buf, @@ -1189,7 +1101,6 @@ impl GpuExperienceCollector { curiosity_trainer, feature_mask_buf: None, position_histogram, - online_params_f32, exp_h_s1_f32, exp_h_s2_f32, exp_h_v_f32, @@ -1211,8 +1122,8 @@ impl GpuExperienceCollector { saboteur_perturbation_scale: 0.3, exp_bn_hidden, exp_bn_concat, - bottleneck_dim: bn_dim_from_params, - market_dim_bn: market_dim, + bottleneck_dim, + market_dim_bn: market_dim_cfg, bn_tanh_concat_fn, step_counter_gpu, step_counter_kernel, @@ -1924,7 +1835,7 @@ impl GpuExperienceCollector { { use crate::cuda_pipeline::batched_forward::f32_weight_ptrs_from_base; let f32_w = f32_weight_ptrs_from_base( - self.online_params_f32.raw_ptr(), + self.trainer_params_ptr, &self.param_sizes, ); self.cublas_forward.sgemm_f32_ldb( @@ -1980,7 +1891,7 @@ impl GpuExperienceCollector { { use crate::cuda_pipeline::batched_forward::f32_weight_ptrs_from_base; let f32_w_ptrs = f32_weight_ptrs_from_base( - self.online_params_f32.raw_ptr(), + self.trainer_params_ptr, &self.param_sizes, ); self.cublas_forward.forward_online_f32( @@ -2278,86 +2189,17 @@ impl GpuExperienceCollector { } // ═══════════════════════════════════════════════════════════════════════ - // Weight synchronization + // Weight pointer management (zero-copy from trainer) // ═══════════════════════════════════════════════════════════════════════ - /// Sync from a flat parameter buffer (single DtoD copy). + /// Set the raw device pointer to the trainer's flat params buffer. /// - /// Use this when the caller already has a contiguous F32 parameter buffer - /// (e.g. from `GpuDqnTrainer`). Avoids the per-tensor extraction overhead. - pub fn sync_weights_flat(&mut self, params_buf: &CudaSlice) -> Result<(), MLError> { - let num_bytes = self.total_params * std::mem::size_of::(); - let src = params_buf.raw_ptr(); - let dst = self.online_params_flat.raw_ptr(); - dtod_copy(dst, src, num_bytes, &self.stream, 0, "weight_sync_flat") - } - - /// #30 Sync f32 master weights directly — no f32 quantization. - /// Called by fused_training to copy f32 master params to the experience - /// collector's f32 weight buffer for full-precision Q-value computation. - pub fn sync_weights_f32(&mut self, f32_params: u64, total_params: usize) -> Result<(), MLError> { - let num_bytes = total_params * std::mem::size_of::(); - let dst = self.online_params_f32.raw_ptr(); - dtod_copy(dst, f32_params, num_bytes, &self.stream, 0, "weight_sync_f32") - } - - /// Flatten the legacy per-tensor weight sets into the contiguous `online_params_flat`. - /// - /// Copies each of the 24 weight tensors from the DuelingWeightSet and - /// BranchingWeightSet into their respective positions in the flat buffer. - /// Indices 24-25 (bottleneck) are not backed by individual CudaSlice fields. - #[allow(unused_assignments)] - fn flatten_online_weights(&mut self) -> Result<(), MLError> { - let mut byte_offset: u64 = 0; - let dst_base = self.online_params_flat.raw_ptr(); - - // Helper: copy a single weight tensor from its device pointer into the flat buffer. - macro_rules! copy_weight { - ($ptr:expr, $idx:expr) => {{ - let n_elems = self.param_sizes[$idx]; - let n_bytes = n_elems * std::mem::size_of::(); - dtod_copy(dst_base + byte_offset, $ptr, n_bytes, &self.stream, $idx, "flatten")?; - byte_offset += n_bytes as u64; - }}; - } - - // Shared trunk: w_s1, b_s1, w_s2, b_s2 - copy_weight!(self.online_weights.w_s1, 0); - copy_weight!(self.online_weights.b_s1, 1); - copy_weight!(self.online_weights.w_s2, 2); - copy_weight!(self.online_weights.b_s2, 3); - - // Value head: w_v1, b_v1, w_v2, b_v2 - copy_weight!(self.online_weights.w_v1, 4); - copy_weight!(self.online_weights.b_v1, 5); - copy_weight!(self.online_weights.w_v2, 6); - copy_weight!(self.online_weights.b_v2, 7); - - // Branch 0 (direction — from DuelingWeightSet): w_b0fc, b_b0fc, w_b0out, b_b0out - copy_weight!(self.online_weights.w_a1, 8); - copy_weight!(self.online_weights.b_a1, 9); - copy_weight!(self.online_weights.w_a2, 10); - copy_weight!(self.online_weights.b_a2, 11); - - // Branch 1 (order — from BranchingWeightSet): w_b1fc, b_b1fc, w_b1out, b_b1out - copy_weight!(self.online_branching.w_bo1, 12); - copy_weight!(self.online_branching.b_bo1, 13); - copy_weight!(self.online_branching.w_bo2, 14); - copy_weight!(self.online_branching.b_bo2, 15); - - // Branch 2 (urgency — from BranchingWeightSet): w_b2fc, b_b2fc, w_b2out, b_b2out - copy_weight!(self.online_branching.w_bu1, 16); - copy_weight!(self.online_branching.b_bu1, 17); - copy_weight!(self.online_branching.w_bu2, 18); - copy_weight!(self.online_branching.b_bu2, 19); - - // Branch 3 (magnitude — from BranchingWeightSet): w_b3fc, b_b3fc, w_b3out, b_b3out - copy_weight!(self.online_branching.w_bg1, 20); - copy_weight!(self.online_branching.b_bg1, 21); - copy_weight!(self.online_branching.w_bg2, 22); - copy_weight!(self.online_branching.b_bg2, 23); - - Ok(()) + /// Must be called after the fused training context is created (the pointer + /// is stable for the lifetime of the fused context). The experience + /// collector reads weights directly from this address — no DtoD copy. + pub fn set_trainer_params_ptr(&mut self, ptr: u64) { + self.trainer_params_ptr = ptr; + info!(trainer_params_ptr = ptr, "Experience collector: zero-copy trainer params pointer set"); } // ═══════════════════════════════════════════════════════════════════════ @@ -2596,15 +2438,6 @@ impl GpuExperienceCollector { Ok(()) } - /// Raw device pointer to the flat f32 online params buffer (destination for DtoD sync). - pub fn online_params_flat_ptr(&self) -> u64 { - self.online_params_flat.raw_ptr() - } - - /// Total size in bytes of the online params flat buffer. - pub fn total_param_bytes(&self) -> usize { - self.online_params_flat.len() * std::mem::size_of::() - } } // --------------------------------------------------------------------------- diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 728a76691..e15146e21 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -9,7 +9,7 @@ //! - `collect_gpu_experiences_slices`: run the GPU experience collection kernel (Phase 3) //! - `run_training_steps_slices`: batched training from replay buffer with guard kernel //! - `process_epoch_boundary`: single readback + safety checks at epoch end -//! - `sync_gpu_weights`: push updated weights to GPU collector +//! - (sync_gpu_weights deleted — zero-copy: collector reads trainer's params_buf directly) //! - `refresh_stale_per_priorities`: M2 PER staleness refresh //! - `log_epoch_metrics_and_financials`: logging, Prometheus, QuestDB, VaR/CVaR //! - `handle_epoch_checkpoints_and_early_stopping`: best-checkpoint + early stopping @@ -336,6 +336,26 @@ impl DQNTrainer { 0.0 }; + // ── Ensure fused CUDA context exists before experience collection ── + // The collector reads weights directly from trainer's params_buf via raw pointer. + // Must be initialized before Phase 2 (experience collection), not lazily in Phase 3. + if self.fused_ctx.is_none() && self.device.is_cuda() { + if let Some(ref stream) = self.cuda_stream { + let agent = self.agent.read().await; + match super::super::fused_training::FusedTrainingCtx::new( + &self.device, &*agent, &self.hyperparams, self.current_batch_size, std::sync::Arc::clone(stream), + ) { + Ok(ctx) => { + if let Some(ref mut collector) = self.gpu_experience_collector { + collector.set_trainer_params_ptr(ctx.params_flat_ptr()); + } + self.fused_ctx = Some(ctx); + } + Err(e) => { tracing::error!("Fused CUDA context init failed: {e}"); } + } + } + } + // ── Phase 2: GPU experience collection ── let phase2_start = std::time::Instant::now(); let gpu_experiences_collected = self.collect_gpu_experiences_slices( @@ -384,7 +404,7 @@ impl DQNTrainer { None }; - self.sync_gpu_weights().await?; + // Zero-copy: collector reads directly from trainer's params_buf — no sync needed. // Log GPU per-phase CUDA event timing after epoch boundary (DtoH readbacks sync stream). if let Some(ref mut fused) = self.fused_ctx { @@ -706,7 +726,7 @@ impl DQNTrainer { let kernel_dims = (state_dim, market_dim, num_atoms_max); // Priority: branching > plain dueling > hybrid (distributional+dueling) - if let (Some(online_br), Some(target_br)) = ( + if let (Some(online_br), Some(_target_br)) = ( dqn.branching_q_network.as_ref(), dqn.branching_target_network.as_ref(), ) { @@ -718,16 +738,10 @@ impl DQNTrainer { cfg.branch_hidden_dim, ); let alloc_episodes = self.compute_alloc_episodes(); - let (online_dueling, online_branching) = - crate::cuda_pipeline::gpu_weights::weight_sets_from_branching(online_br); - let (target_dueling, target_branching) = - crate::cuda_pipeline::gpu_weights::weight_sets_from_branching(target_br); Some(GpuExperienceCollector::new( stream, - &online_dueling, - &target_dueling, - &online_branching, - &target_branching, + self.hyperparams.bottleneck_dim, + 42, // market_dim_cfg self.hyperparams.initial_capital as f32, self.hyperparams.avg_spread as f32, self.hyperparams.cash_reserve_percent as f32, @@ -1199,7 +1213,13 @@ impl DQNTrainer { match super::super::fused_training::FusedTrainingCtx::new( &self.device, &*agent, &self.hyperparams, self.current_batch_size, std::sync::Arc::clone(stream), ) { - Ok(ctx) => { self.fused_ctx = Some(ctx); } + Ok(ctx) => { + // Zero-copy: wire collector to read weights directly from trainer's params_buf. + if let Some(ref mut collector) = self.gpu_experience_collector { + collector.set_trainer_params_ptr(ctx.params_flat_ptr()); + } + self.fused_ctx = Some(ctx); + } Err(e) => { tracing::error!("Fused CUDA context init failed: {e}"); } } } @@ -1468,31 +1488,7 @@ impl DQNTrainer { }) } - // ═══════════════════════════════════════════════════════════════════════ - // Helper: Sync GPU weight copies after training updates - // ═══════════════════════════════════════════════════════════════════════ - - pub(crate) async fn sync_gpu_weights(&mut self) -> Result<()> { - let Some(ref mut collector) = self.gpu_experience_collector else { - return Ok(()); - }; - let Some(ref fused) = self.fused_ctx else { - return Ok(()); - }; - let stream = self.cuda_stream.as_ref() - .ok_or_else(|| anyhow::anyhow!("CUDA stream required for weight sync"))?; - - // v8 perf: single DtoD copy replaces 32+ Candle hash map lookups + individual copies + flatten. - // fused trainer's params_flat has the SAME flat layout as collector's online_params_flat. - let src = fused.params_flat_ptr(); - let dst = collector.online_params_flat_ptr(); - let bytes = collector.total_param_bytes(); - unsafe { - cudarc::driver::sys::cuMemcpyDtoDAsync_v2(dst, src, bytes as usize, stream.cu_stream()); - } - - Ok(()) - } + // sync_gpu_weights deleted — zero-copy: collector reads directly from trainer's params_buf. // ═══════════════════════════════════════════════════════════════════════ // Helper: M2 — Refresh stale PER priorities diff --git a/crates/ml/tests/smoke_test_real_data.rs b/crates/ml/tests/smoke_test_real_data.rs index c574a6d9d..bfcb231b0 100644 --- a/crates/ml/tests/smoke_test_real_data.rs +++ b/crates/ml/tests/smoke_test_real_data.rs @@ -392,8 +392,6 @@ mod gpu_smoke { use ml::cuda_pipeline::gpu_experience_collector::{ ExperienceCollectorConfig, GpuExperienceCollector, }; - use ml::cuda_pipeline::gpu_weights::weight_sets_from_branching; - use ml::dqn::branching::{BranchingConfig, BranchingDuelingQNetwork}; type CudaStream = cudarc::driver::CudaStream; type CudaSlice = cudarc::driver::CudaSlice; @@ -526,21 +524,11 @@ mod gpu_smoke { info!(num_bars, market_dim = 51, "GPU buffer allocated"); - // Create dueling networks (state_dim=54: 51 market + 3 portfolio) - let config = BranchingConfig::trading_default( - 54, vec![64, 64], 32, vec![9, 3, 3], - ); - let online = BranchingDuelingQNetwork::new(config.clone(), stream.clone()).unwrap(); - let target_net = BranchingDuelingQNetwork::new(config, stream.clone()).unwrap(); - let (online_dueling, online_branching) = weight_sets_from_branching(&online); - let (target_dueling, target_branching) = weight_sets_from_branching(&target_net); - + // Create collector (zero-copy: no weight sets needed at construction) let mut collector = GpuExperienceCollector::new( stream.clone(), - &online_dueling, - &target_dueling, - &online_branching, - &target_branching, + 0, // bottleneck_dim (disabled for test) + 51, // market_dim_cfg 100_000.0, 0.01, 0.05, (64, 64, 32, 32), @@ -611,20 +599,10 @@ mod gpu_smoke { info!(num_bars, market_dim = 51, "GPU buffer allocated (with OFI)"); - let config = BranchingConfig::trading_default( - 54, vec![64, 64], 32, vec![9, 3, 3], - ); - let online = BranchingDuelingQNetwork::new(config.clone(), stream.clone()).unwrap(); - let target_net = BranchingDuelingQNetwork::new(config, stream.clone()).unwrap(); - let (online_dueling, online_branching) = weight_sets_from_branching(&online); - let (target_dueling, target_branching) = weight_sets_from_branching(&target_net); - let mut collector = GpuExperienceCollector::new( stream.clone(), - &online_dueling, - &target_dueling, - &online_branching, - &target_branching, + 0, // bottleneck_dim (disabled for test) + 51, // market_dim_cfg 100_000.0, 0.01, 0.05, (64, 64, 32, 32), @@ -675,28 +653,16 @@ mod gpu_smoke { #[test] fn smoke_gpu_real_data_noisy_distributional() -> Result<(), anyhow::Error> { - use ml::dqn::branching::{BranchingConfig, BranchingDuelingQNetwork}; - let Some(ohlcv_dir) = try_data("ohlcv") else { return Ok(()) }; let Some((_device, stream)) = try_cuda() else { return Ok(()) }; let (market_buf, target_buf, num_bars) = real_market_data(&ohlcv_dir, None, &stream)?; - let config = BranchingConfig::trading_default( - 54, vec![64, 64], 32, vec![9, 3, 3], - ); - let online = BranchingDuelingQNetwork::new(config.clone(), stream.clone()).unwrap(); - let target_net = BranchingDuelingQNetwork::new(config, stream.clone()).unwrap(); - let (online_dueling, online_branching) = weight_sets_from_branching(&online); - let (target_dueling, target_branching) = weight_sets_from_branching(&target_net); - let mut collector = GpuExperienceCollector::new( stream.clone(), - &online_dueling, - &target_dueling, - &online_branching, - &target_branching, + 0, // bottleneck_dim (disabled for test) + 51, // market_dim_cfg 100_000.0, 0.01, 0.05, (64, 64, 32, 32), diff --git a/docs/superpowers/plans/2026-04-11-exp-collector-zero-copy.md b/docs/superpowers/plans/2026-04-11-exp-collector-zero-copy.md new file mode 100644 index 000000000..d07c25dbf --- /dev/null +++ b/docs/superpowers/plans/2026-04-11-exp-collector-zero-copy.md @@ -0,0 +1,251 @@ +# Experience Collector Zero-Copy Refactor + +> **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:** Eliminate all DtoD weight copies between trainer and experience collector. Collector reads trainer's `params_buf` via direct u64 pointer. Enable bottleneck forward in collector. Delete all dead weight buffers and sync methods. + +**Architecture:** Replace `online_params_f32` (never synced, always zeros) and `online_params_flat` (synced but never read) with a single `trainer_params_ptr: u64` that points into the trainer's `params_buf`. Recompute `param_sizes` from the trainer's real config (bottleneck_dim=16). Collector gets its own `CublasForward` with correct `s1_input_dim=54`. Bottleneck forward (GEMM + tanh + concat) activates using the existing `bn_tanh_concat_fn` kernel. + +**Tech Stack:** Rust 1.85, cudarc, CUDA 12.4 + +--- + +## File Map + +| File | Action | Changes | +|------|--------|---------| +| `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` | Modify | Delete 5 dead fields + 3 dead methods, add `trainer_params_ptr`, fix `param_sizes` + `CublasForward` + bottleneck forward, resize bn buffers | +| `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | Modify | Delete `sync_gpu_weights`, pass `trainer_params_ptr` + `s1_input_dim` + `bottleneck_dim` + `market_dim` to collector constructor | +| `crates/ml/src/trainers/dqn/fused_training.rs` | Modify | Expose `params_buf_ptr()` + `param_sizes` for collector construction | + +--- + +### Task 1: Delete dead weight fields and methods from collector + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` + +- [ ] **Step 1: Delete `online_params_flat` field and its allocation** + +Remove field `online_params_flat: CudaSlice` (line 478), its allocation at lines 796-798, and its entry in the `Self { ... }` constructor return at line 1138. + +- [ ] **Step 2: Delete `online_params_f32` field and its allocation** + +Remove field `online_params_f32: CudaSlice` (line 597), its allocation at line 1054, and its entry in the constructor return at line 1192. + +- [ ] **Step 3: Delete `online_weights` and `target_weights` (DuelingWeightSet)** + +Remove fields `online_weights` and `target_weights` (used only by dead `flatten_online_weights`). Remove the construction at lines 810-837 (the 50-line copy block). Remove constructor params `online_dueling` and `target_dueling` from `fn new()` signature (line 697-698). Remove from `Self { ... }` at lines 1168-1169. + +- [ ] **Step 4: Delete `online_branching` and `target_branching` (BranchingWeightSet)** + +Remove fields, construction at lines 848-861, constructor params `online_branching` and `target_branching` from `fn new()` (lines 700-701), and from `Self { ... }` at lines 1173-1174. + +- [ ] **Step 5: Delete `curiosity_weights`, `online_rmsnorm`, `target_rmsnorm`** + +These are all-zero / all-ones stubs never read by the cuBLAS forward. Remove fields, construction at lines 841-845, and from `Self { ... }` at lines 1170-1172. + +- [ ] **Step 6: Delete dead sync methods** + +Remove `sync_weights_flat()` (lines 2288-2293), `sync_weights_f32()` (lines 2298-2302), `flatten_online_weights()` (lines 2310-2361). Also remove `online_params_flat_ptr()` accessor and `total_param_bytes()` accessor (search for them). + +- [ ] **Step 7: Verify compilation** + +```bash +SQLX_OFFLINE=true cargo check -p ml --lib 2>&1 | tail -5 +``` + +Expect: compilation errors in `training_loop.rs` (calls deleted methods) and anywhere else that passes the removed constructor params. Those are fixed in Task 2. + +- [ ] **Step 8: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +git commit -m "refactor: delete 5 dead weight fields + 3 dead sync methods from exp collector" +``` + +--- + +### Task 2: Add `trainer_params_ptr` and fix constructor signature + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` +- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs` +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` + +- [ ] **Step 1: Add new constructor params and field** + +Replace the deleted `online_dueling`, `target_dueling`, `online_branching`, `target_branching` params in `fn new()` with: + +```rust +pub fn new( + stream: Arc, + trainer_params_ptr: u64, // raw pointer into trainer's params_buf (stable) + bottleneck_dim: usize, // trainer's bottleneck_dim (16 in production) + market_dim_cfg: usize, // trainer's market_dim (42) + initial_capital: f32, + _avg_spread: f32, + _cash_reserve_pct: f32, + network_dims: (usize, usize, usize, usize), + kernel_dims: (usize, usize, usize), + n_episodes: usize, + timesteps_per_episode: usize, +) -> Result { +``` + +Add field `trainer_params_ptr: u64` to the struct (replacing `online_params_flat`). + +- [ ] **Step 2: Fix `param_sizes` computation to use real bottleneck config** + +Replace the synthetic `GpuDqnTrainConfig` at lines 775-786: + +```rust +let train_cfg = GpuDqnTrainConfig { + state_dim, + shared_h1, + shared_h2, + value_h, + adv_h, + num_atoms, + branch_0_size: branch_sizes[0], + branch_1_size: branch_sizes[1], + branch_2_size: branch_sizes[2], + branch_3_size: 3, + bottleneck_dim: bottleneck_dim, // was: 0 (WRONG — different layout) + market_dim: market_dim_cfg, // was: 0 + ..GpuDqnTrainConfig::default() +}; +let param_sizes = compute_param_sizes(&train_cfg); +let total_params = compute_total_params(&train_cfg); +``` + +- [ ] **Step 3: Fix `CublasForward` construction with correct `s1_input_dim`** + +The collector's `CublasForward::new` call (around line 744) needs the correct `s1_input_dim`: + +```rust +let s1_input_dim = if bottleneck_dim > 0 { + bottleneck_dim + state_dim.saturating_sub(market_dim_cfg) +} else { + state_dim +}; +let cublas_forward = CublasForward::new( + &stream, + alloc_episodes, + state_dim, + shared_h1, shared_h2, value_h, adv_h, + num_atoms, + branch_sizes[0], branch_sizes[1], branch_sizes[2], branch_sizes[3], + s1_input_dim, +)?; +``` + +- [ ] **Step 4: Fix bottleneck buffer allocation** + +Replace the `bn_alloc = bn_dim_from_params.max(1)` at line 1114 with: + +```rust +let bn_alloc = bottleneck_dim.max(1); +let portfolio_dim_bn = state_dim.saturating_sub(market_dim_cfg); +let exp_bn_hidden = stream.alloc_zeros::(alloc_episodes * bn_alloc) + .map_err(|e| MLError::ModelError(format!("alloc exp_bn_hidden: {e}")))?; +let exp_bn_concat = stream.alloc_zeros::(alloc_episodes * (bn_alloc + portfolio_dim_bn) + 128) + .map_err(|e| MLError::ModelError(format!("alloc exp_bn_concat: {e}")))?; +``` + +Set `bottleneck_dim: bottleneck_dim` (not `bn_dim_from_params`) in the `Self { ... }` constructor return. And `market_dim_bn: market_dim_cfg`. + +- [ ] **Step 5: Replace `online_params_f32` reads with `trainer_params_ptr`** + +In the forward pass (lines 1926-1984), replace: +```rust +self.online_params_f32.raw_ptr() +``` +with: +```rust +self.trainer_params_ptr +``` + +There are 2 occurrences: line 1927 (bottleneck SGEMM) and line 1983 (main forward). + +Store `trainer_params_ptr` in `Self { ... }` replacing `online_params_flat`. + +- [ ] **Step 6: Delete `sync_gpu_weights` from training_loop.rs** + +Remove the method at lines 1475-1495 and the call at line 387. + +- [ ] **Step 7: Fix collector construction call site in training_loop.rs** + +Find where `GpuExperienceCollector::new(...)` is called. Replace the old params (weight sets) with the new params: + +```rust +let collector = GpuExperienceCollector::new( + stream.clone(), + fused.params_flat_ptr(), // trainer_params_ptr + hyperparams.bottleneck_dim, // bottleneck_dim + 42, // market_dim (always 42) + initial_capital, + avg_spread, + cash_reserve_pct, + network_dims, + kernel_dims, + n_episodes, + timesteps_per_episode, +)?; +``` + +- [ ] **Step 8: Verify compilation and tests** + +```bash +SQLX_OFFLINE=true cargo check -p ml --lib +SQLX_OFFLINE=true cargo test -p ml-dqn --lib +SQLX_OFFLINE=true cargo test -p ml --lib -- gradient_budget config monitoring +``` + +- [ ] **Step 9: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/gpu_experience_collector.rs crates/ml/src/trainers/dqn/trainer/training_loop.rs crates/ml/src/trainers/dqn/fused_training.rs +git commit -m "refactor: exp collector reads trainer params_buf directly — zero DtoD copies" +``` + +--- + +### Task 3: Smoke test + compute-sanitizer verification + +**Files:** (no changes — verification only) + +- [ ] **Step 1: Run smoke test and check action diversity** + +```bash +FOXHUNT_TEST_DATA=test_data/futures-baseline SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests::feature_coverage --ignored --test-threads=1 --nocapture 2>&1 | grep -iE "Action diversity|LOW ORDER|test result" +``` + +Expected: order diversity should improve from 1/3 (zero weights → no Q-gap) to 2/3 or 3/3 (real Xavier-init weights → Boltzmann distributes across order types). + +- [ ] **Step 2: Run compute-sanitizer** + +```bash +SQLX_OFFLINE=true cargo test -p ml --lib --no-run 2>&1 | tail -1 +FOXHUNT_TEST_DATA=test_data/futures-baseline compute-sanitizer --tool memcheck target/debug/deps/ml-*.test "smoke_tests::feature_coverage" --ignored --test-threads=1 +``` + +Expected: `ERROR SUMMARY: 0 errors` + +- [ ] **Step 3: Run full test suite** + +```bash +SQLX_OFFLINE=true cargo test -p ml-dqn --lib +SQLX_OFFLINE=true cargo test -p ml --lib +FOXHUNT_TEST_DATA=test_data/futures-baseline SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests --ignored --test-threads=1 +``` + +Expected: All pass. + +- [ ] **Step 4: Commit and push** + +```bash +git add -A +git commit -m "fix: exp collector zero-copy — direct pointer into trainer params_buf, bottleneck active" +git push origin main +```