diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 1488ce6dc..096ce72be 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -340,8 +340,12 @@ pub struct PerceptionTrainer { // gates = softmax(W_vsn @ x + b_vsn); y = x * gates // Lets the model learn which snap_features matter per regime. /// `[FEATURE_DIM, FEATURE_DIM]` linear layer for gate logits. - pub vsn_w_d: CudaSlice, - pub vsn_b_d: CudaSlice, + // X2 migration: VSN weights now live on `self.trunk` (CfcTrunk's + // pub vsn_w_d / vsn_b_d). PerceptionTrainer wraps a trunk to host + // the v2 inference graph; the trunk's V1 fields stay unused for now + // (X8 consolidates against PerceptionTrainer's CfC weights). Gradient + // buffers + AdamW state for VSN remain at trainer level. + pub trunk: crate::cfc::trunk::CfcTrunk, /// VSN forward output `[B, K, FEATURE_DIM]` — replaces window_tensor_d /// as Mamba2's input. window_tensor_d still holds the RAW features /// (consumed by VSN bwd as `x`). @@ -511,6 +515,23 @@ impl PerceptionTrainer { // dqn / ppo crates that also use OwnedGpuLinear::xavier). let _seed_guard = ml_core::cuda_autograd::init::scoped_init_seed(cfg.seed); + // X2: construct the wrapped CfcTrunk that owns the v2 inference + // graph weights. Trunk's V1 fields (CfC + V1 heads + projection) + // stay unused at this commit; subsequent migrations (X3..X9) move + // PerceptionTrainer's per-layer weights into the trunk one group + // at a time, gated by the perception_forward_golden bit-equivalence + // test. + let mut trunk = crate::cfc::trunk::CfcTrunk::new_random( + dev, + &crate::cfc::trunk::CfcConfig { + n_in: FEATURE_DIM, + n_hid: HIDDEN_DIM, + mamba2_state_dim: cfg.mamba2_state_dim, + }, + cfg.seed, + ) + .context("PerceptionTrainer: construct wrapped CfcTrunk")?; + let stream = dev.cuda_stream().context("trainer stream")?.clone(); let ctx = dev.cuda_context().context("trainer ctx")?; // Disable cudarc's automatic per-allocation read/write event @@ -800,12 +821,28 @@ impl PerceptionTrainer { // W_vsn ∈ [FEATURE_DIM, FEATURE_DIM] init near zero so initial // gates ≈ softmax(0) = uniform 1/FEATURE_DIM. The model can then // learn departures from uniform. Scale = sqrt(1/FEATURE_DIM). + // + // X2 migration: VSN weights now live on `self.trunk` (the wrapped + // CfcTrunk). We still draw values from the trainer's PRNG chain + // here to keep the chain state — and therefore every downstream + // weight (attn_q, etc.) — bit-identical to the pre-X2 layout. The + // draws are uploaded directly into the trunk's zero-initialised + // VSN slots via memcpy_htod, replacing the prior trainer-owned + // vsn_w_d / vsn_b_d allocations. let vsn_scale = (1.0_f32 / FEATURE_DIM as f32).sqrt(); let vsn_w_init: Vec = (0..FEATURE_DIM * FEATURE_DIM) .map(|_| r.gen_range(-vsn_scale..vsn_scale)).collect(); let vsn_b_init: Vec = vec![0.0; FEATURE_DIM]; - let vsn_w_d = upload(&stream, &vsn_w_init)?; - let vsn_b_d = upload(&stream, &vsn_b_init)?; + // Upload the trainer-driven VSN draws into the trunk's + // zero-initialised VSN slots from X1. From here on, + // `self.trunk.vsn_w_d` / `self.trunk.vsn_b_d` are the canonical + // VSN weights for both forward and backward paths. + stream + .memcpy_htod(&vsn_w_init, &mut trunk.vsn_w_d) + .context("trunk.vsn_w_d upload")?; + stream + .memcpy_htod(&vsn_b_init, &mut trunk.vsn_b_d) + .context("trunk.vsn_b_d upload")?; let opt_vsn_w = AdamW::new(dev, FEATURE_DIM * FEATURE_DIM, cfg.lr_cfc)?; let mut opt_vsn_b = AdamW::new(dev, FEATURE_DIM, cfg.lr_cfc)?; opt_vsn_b.wd = 0.0; @@ -881,8 +918,7 @@ impl PerceptionTrainer { opt_ln_gain, opt_ln_bias, // VSN (Phase 2D) — params + per-K gates + scratch. - vsn_w_d, - vsn_b_d, + trunk, vsn_out_d: GpuTensor::zeros(&[cfg.n_batch, k, FEATURE_DIM], &stream) .map_err(|e| anyhow::anyhow!("vsn_out_d alloc: {e}"))?, vsn_gates_d: stream.alloc_zeros::(cfg.n_batch * k * FEATURE_DIM)?, @@ -1385,8 +1421,8 @@ impl PerceptionTrainer { }; let mut launch = self.stream.launch_builder(&self.vsn_fwd_fn); launch - .arg(&self.vsn_w_d) - .arg(&self.vsn_b_d) + .arg(&self.trunk.vsn_w_d) + .arg(&self.trunk.vsn_b_d) .arg(self.window_tensor_d.cuda_data()) .arg(&n_rows_vsn) .arg(self.vsn_out_d.data_mut()) @@ -2113,7 +2149,7 @@ impl PerceptionTrainer { }; let mut launch = self.stream.launch_builder(&self.vsn_bwd_fn); launch - .arg(&self.vsn_w_d) + .arg(&self.trunk.vsn_w_d) .arg(self.window_tensor_d.cuda_data()) .arg(&self.vsn_gates_d) .arg(self.mamba2_grads_buffers.d_x_from_in.cuda_data()) @@ -2243,8 +2279,8 @@ impl PerceptionTrainer { self.opt_ln_bias.step(&mut self.ln_bias_d, &self.grad_ln_bias_d)?; self.opt_ln_a_gain.step(&mut self.ln_a_gain_d, &self.grad_ln_a_gain_d)?; self.opt_ln_a_bias.step(&mut self.ln_a_bias_d, &self.grad_ln_a_bias_d)?; - self.opt_vsn_w.step(&mut self.vsn_w_d, &self.grad_vsn_w_d)?; - self.opt_vsn_b.step(&mut self.vsn_b_d, &self.grad_vsn_b_d)?; + self.opt_vsn_w.step(&mut self.trunk.vsn_w_d, &self.grad_vsn_w_d)?; + self.opt_vsn_b.step(&mut self.trunk.vsn_b_d, &self.grad_vsn_b_d)?; self.opt_attn_q.step(&mut self.attn_q_d, &self.grad_attn_q_d)?; // Kendall σ is ISV-driven (closed form from loss_ema in // horizon_ema_and_lambda) — no Adam step. grad_log_sigma_h_d @@ -2424,8 +2460,8 @@ impl PerceptionTrainer { }; let mut launch = self.stream.launch_builder(&self.vsn_fwd_fn); launch - .arg(&self.vsn_w_d) - .arg(&self.vsn_b_d) + .arg(&self.trunk.vsn_w_d) + .arg(&self.trunk.vsn_b_d) .arg(self.window_tensor_d.cuda_data()) .arg(&n_rows_vsn) .arg(self.vsn_out_d.data_mut())