From 7deaa88d833fa14b8acd911e30a4c954b6030ad3 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 16 Apr 2026 23:52:38 +0200 Subject: [PATCH] =?UTF-8?q?feat(isv):=20target=20network=20=E2=80=94=20ISV?= =?UTF-8?q?=20weights=20online-only?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit EMA sync restricted to indices 0-67 (non-ISV params). ISV weights (68-77) exist only in online params_buf. Target forward uses online gamma_buf (drift-conditioned) and online branch_gate_buf — correct for Bellman projection. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 36 ++++++++++++++----- 1 file changed, 28 insertions(+), 8 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 86929aa8a..1de5bdb57 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -83,6 +83,9 @@ const MAMBA2_STATE_DIM: usize = 16; // SSM state dimension /// Introspective State Vector (ISV) configuration. const ISV_K: usize = 4; // Temporal ISV history length const ISV_DIM: usize = 8; // ISV signal count +/// First ISV tensor index in the flat param buffer. +/// ISV weights (68-77) are online-only — NOT synced to the target network. +const FIRST_ISV_TENSOR: usize = 68; const NEW_COMPONENT_WARMUP_STEPS: u32 = 500; /// Homeostatic regularizer: number of observable signals. @@ -501,6 +504,12 @@ pub(crate) fn compute_total_params(cfg: &GpuDqnTrainConfig) -> usize { compute_param_sizes(cfg).iter().map(|&s| align4(s)).sum() } +/// Non-ISV parameter count (indices 0..FIRST_ISV_TENSOR). +/// EMA target sync is restricted to this range — ISV weights are online-only. +pub(crate) fn compute_non_isv_params(cfg: &GpuDqnTrainConfig) -> usize { + compute_param_sizes(cfg)[..FIRST_ISV_TENSOR].iter().map(|&s| align4(s)).sum() +} + /// Cumulative byte offset to tensor `idx` in the padded flat buffer. pub(crate) fn padded_byte_offset(param_sizes: &[usize], idx: usize) -> u64 { param_sizes[..idx].iter() @@ -894,6 +903,8 @@ pub struct GpuDqnTrainer { // ── Training state ────────────────────────────────────────────── pub(crate) adam_step: i32, total_params: usize, + /// Non-ISV parameter count: EMA target sync restricted to this range. + non_isv_params: usize, pub(crate) params_initialized: bool, target_params_initialized: bool, attention_initialized: bool, @@ -3800,6 +3811,7 @@ impl GpuDqnTrainer { ) -> Result { let b = config.batch_size; let total_params = compute_total_params(&config); + let non_isv_params = compute_non_isv_params(&config); // Event tracking: kept ENABLED during normal ops (buffer allocation, DtoD). // DISABLED only during CUDA Graph capture (in capture_training_graphs). @@ -5520,6 +5532,7 @@ impl GpuDqnTrainer { q_div_ema: 0.0, adam_step: 0, total_params, + non_isv_params, params_initialized: false, target_params_initialized: false, attention_initialized: false, @@ -9109,7 +9122,9 @@ impl GpuDqnTrainer { /// matches the GOFF_* defines exactly, with `align4()` padding per tensor. /// /// Both params_buf and target_params_buf receive identical weights — the EMA - /// kernel will diverge them during training. + /// kernel will diverge them during training. ISV weights (indices 68-77) in + /// target_params_buf are inert: EMA skips them, and the target forward never + /// reads them. pub(crate) fn xavier_init_params_buf(&mut self) -> Result<(), MLError> { let cfg = &self.config; let sizes = compute_param_sizes(cfg); @@ -9401,11 +9416,13 @@ impl GpuDqnTrainer { /// GPU-native Polyak EMA: `target[i] = (1-tau)*target[i] + tau*online[i]` /// - /// Fused single-kernel update over flat parameter buffers. On first call, - /// flattens target weights into `target_params_buf` (same GOFF_* layout as - /// `params_buf`). Then launches ONE EMA kernel over the entire flat buffer - /// instead of 20 per-tensor launches. After the kernel, scatters the updated - /// flat target weights back to individual tensors via `unflatten_target_weights()`. + /// Fused single-kernel update over flat parameter buffers, restricted to + /// non-ISV params (indices 0-67). ISV weights (68-77) are online-only and + /// never synced to the target network. On first call, flattens target weights + /// into `target_params_buf` (same GOFF_* layout as `params_buf`). Then + /// launches ONE EMA kernel over the non-ISV portion instead of 20 per-tensor + /// launches. After the kernel, scatters the updated flat target weights back + /// to individual tensors via `unflatten_target_weights()`. /// /// Runs OUTSIDE the captured CUDA Graph -- device pointers are stable so /// the graph stays valid. @@ -9432,8 +9449,11 @@ impl GpuDqnTrainer { tau }; - let n = self.total_params as i32; - let blocks = ((self.total_params + 255) / 256) as u32; + // EMA only over non-ISV params (indices 0-67). ISV weights (68-77) are + // online-only — the target network uses neutral defaults (base_gamma, + // uniform gates) instead of ISV-modulated values. + let n = self.non_isv_params as i32; + let blocks = ((self.non_isv_params + 255) / 256) as u32; let launch_cfg = LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1),