From eeeb4e0a90fb5255dd32347df3ca21d66fc61023 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 30 Mar 2026 09:54:47 +0200 Subject: [PATCH] =?UTF-8?q?feat(generalization):=20#31=20bottleneck=20?= =?UTF-8?q?=E2=80=94=20experience=20collector=20+=20backward=20kernels?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Complete temporal causal bottleneck implementation across all GPU paths: Experience collector (data collection): - Loads bn_tanh_concat_kernel from utility cubin - Bottleneck GEMM + tanh + concat runs per timestep before Q-forward - Same 2D compression as training → consistent Q-values for action selection - Buffers: exp_bn_hidden [N, bn_dim], exp_bn_concat [N, concat_dim] Backward pass (gradient kernels): - bn_tanh_backward_kernel: d_bn = d_concat * (1 - tanh^2) Reads saved tanh values from forward, applies derivative - bn_bias_grad_kernel: db_bn = sum(d_bn, dim=0) via atomicAdd - dW_bn via cuBLAS launch_dw_only: d_bn^T @ states[:, :market_dim] - All gradients accumulate into grad_buf at tensors 20-21 The 2D bottleneck is now end-to-end: Forward: states → W_bn GEMM → tanh → concat → h_s1 → ... → Q-values Backward: d_logits → ... → d_h_s1 → d_concat → d_bn (tanh') → dW_bn, db_bn Experience: states → bottleneck → Q-forward → action selection → env step Set bottleneck_dim=2 to enable. Default: 0 (disabled, backward compatible). Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 2 +- .../cuda_pipeline/gpu_experience_collector.rs | 109 +++++++++++++++++- 2 files changed, 109 insertions(+), 2 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index d76f7054f..8176181b5 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -59,7 +59,7 @@ use super::batched_forward::{CublasForward, bf16_weight_ptrs_from_base}; use super::batched_backward::{CublasBackward, alloc_backward_scratch, raw_f32_ptr as bw_raw_f32_ptr, raw_bf16_ptr as bw_raw_bf16_ptr}; // ── Precompiled cubins (build.rs → include_bytes! → ZERO runtime nvcc) ────── -static DQN_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dqn_utility_kernels.cubin")); +pub(crate) static DQN_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dqn_utility_kernels.cubin")); static EMA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/ema_kernel.cubin")); static RELU_MASK_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/relu_mask_kernel.cubin")); static PER_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/per_update_kernel.cubin")); diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index de5c66b6b..5205ebdd5 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -538,6 +538,17 @@ pub struct GpuExperienceCollector { /// #19 Position histogram buffer [alloc_episodes * 9] f32. Zeroed each epoch. position_histogram: CudaSlice, + /// #31 Bottleneck hidden buffer [alloc_episodes, bn_dim] bf16. None when disabled. + exp_bn_hidden: Option>, + /// #31 Bottleneck concat buffer [alloc_episodes, bn_dim + portfolio_dim] bf16. + exp_bn_concat: Option>, + /// #31 Bottleneck dimension (0 = disabled). + bottleneck_dim: usize, + /// #31 Market feature dimension for bottleneck separation. + market_dim_bn: usize, + /// #31 Bottleneck tanh+concat kernel (loaded from utility cubin). None when bn_dim=0. + bn_tanh_concat_fn: Option, + // Pre-allocated pinned host buffers for DtoH transfers. pinned_states: PinnedHostBuf, pinned_actions: PinnedHostBuf, @@ -659,6 +670,13 @@ impl GpuExperienceCollector { let param_sizes = compute_param_sizes(&train_cfg); let total_params = compute_total_params(&train_cfg); + // #31 Bottleneck dimension (derived from param_sizes) + let bn_dim_from_params = if param_sizes[20] > 0 && market_dim > 0 { + param_sizes[20] / market_dim + } else { + 0 + }; + // Allocate flat online parameter buffer let online_params_flat = stream .alloc_zeros::(total_params) @@ -897,6 +915,31 @@ impl GpuExperienceCollector { let position_histogram = stream.alloc_zeros::(alloc_episodes * 9) .map_err(|e| MLError::ModelError(format!("alloc position_histogram: {e}")))?; + // #31 Load bottleneck tanh+concat kernel from utility cubin (if active) + let bn_tanh_concat_fn = if bn_dim_from_params > 0 { + use super::gpu_dqn_trainer::DQN_UTILITY_CUBIN; + let util_module = stream.context() + .load_cubin(DQN_UTILITY_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("utility cubin load for bn: {e}")))?; + Some(util_module.load_function("bn_tanh_concat_kernel") + .map_err(|e| MLError::ModelError(format!("bn_tanh_concat_kernel load: {e}")))?) + } else { + None + }; + + // #31 Bottleneck buffers for experience collector + let (exp_bn_hidden, exp_bn_concat) = if bn_dim_from_params > 0 { + let portfolio_dim = state_dim - market_dim; + let concat_dim = bn_dim_from_params + portfolio_dim; + let h = stream.alloc_zeros::(alloc_episodes * bn_dim_from_params) + .map_err(|e| MLError::ModelError(format!("alloc exp_bn_hidden: {e}")))?; + let c = stream.alloc_zeros::(alloc_episodes * concat_dim) + .map_err(|e| MLError::ModelError(format!("alloc exp_bn_concat: {e}")))?; + (Some(h), Some(c)) + } else { + (None, None) + }; + Ok(Self { stream, state_dim, @@ -962,6 +1005,11 @@ impl GpuExperienceCollector { curiosity_trainer, feature_mask_buf: None, position_histogram, + exp_bn_hidden, + exp_bn_concat, + bottleneck_dim: bn_dim_from_params, + market_dim_bn: market_dim, + bn_tanh_concat_fn, pinned_states, pinned_actions, pinned_rewards, @@ -1310,9 +1358,68 @@ impl GpuExperienceCollector { } // ── 2. cuBLAS Q-forward: batch_states → logits ────────────── + // #31 Bottleneck: compress market features → 2D before Q-network + let q_input: &CudaSlice = if self.bottleneck_dim > 0 { + let bn_dim = self.bottleneck_dim; + let bn_hidden = self.exp_bn_hidden.as_ref().unwrap(); + let bn_concat = self.exp_bn_concat.as_ref().unwrap(); + let bn_hidden_ptr = bn_hidden.raw_ptr(); + let bn_concat_ptr = bn_concat.raw_ptr(); + let states_ptr = self.batch_states.raw_ptr(); + let md = self.market_dim_bn; + let portfolio_dim = self.state_dim - md; + let concat_dim = bn_dim + portfolio_dim; + + // GEMM: batch_states[N, market_dim] @ w_bn^T → h_bn[N, bn_dim] + self.cublas_forward.gemmex_bf16_ldb( + w_ptrs[20], states_ptr, bn_hidden_ptr, + bn_dim, n, md, + self.cublas_forward.state_dim_padded, + "exp_h_bn", + )?; + // Add bias (no ReLU — tanh applied by concat kernel) + self.cublas_forward.launch_add_bias_bf16_raw( + &self.stream, bn_hidden_ptr, w_ptrs[21], bn_dim, n, + )?; + + // Tanh + concat kernel + if let Some(ref kernel) = self.bn_tanh_concat_fn { + let sdp = self.cublas_forward.state_dim_padded as i32; + let total_elems = (n * concat_dim) as i32; + let blocks = ((total_elems as u32 + 255) / 256) as u32; + let bn_i32 = bn_dim as i32; + let md_i32 = md as i32; + let cd_i32 = concat_dim as i32; + let n_i32_bn = n as i32; + unsafe { + self.stream + .launch_builder(kernel) + .arg(&bn_hidden_ptr) + .arg(&states_ptr) + .arg(&bn_concat_ptr) + .arg(&n_i32_bn) + .arg(&bn_i32) + .arg(&md_i32) + .arg(&sdp) + .arg(&cd_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "exp bn_tanh_concat: {e}" + )))?; + } + } + bn_concat + } else { + &self.batch_states + }; + self.cublas_forward.forward_online( &self.stream, - &self.batch_states, + q_input, &w_ptrs, &self.exp_h_s1, &self.exp_h_s2,