diff --git a/crates/ml-alpha/src/cfc/trunk.rs b/crates/ml-alpha/src/cfc/trunk.rs index 8a31c4140..3292580fe 100644 --- a/crates/ml-alpha/src/cfc/trunk.rs +++ b/crates/ml-alpha/src/cfc/trunk.rs @@ -120,6 +120,13 @@ pub struct CfcTrunk { pub ln_b_gain_d: CudaSlice, pub ln_b_bias_d: CudaSlice, pub attn_q_d: CudaSlice, + // X9: full GRN head structure. heads_w1/b1 is the input projection + // (HIDDEN → HEAD_MID), heads_w2/b2 is the mid→mid layer, then the + // gate/main/skip outputs feed multi-horizon probabilities. + pub heads_w1_d: CudaSlice, + pub heads_b1_d: CudaSlice, + pub heads_w2_d: CudaSlice, + pub heads_b2_d: CudaSlice, pub heads_w_gate_d: CudaSlice, pub heads_b_gate_d: CudaSlice, pub heads_w_main_d: CudaSlice, @@ -254,6 +261,14 @@ impl CfcTrunk { .context("v2 ln_b_bias alloc")?, attn_q_d: stream.alloc_zeros::(cfg.n_hid) .context("v2 attn_q alloc")?, + heads_w1_d: stream.alloc_zeros::(N_HORIZONS * HEAD_MID_DIM * cfg.n_hid) + .context("v2 heads_w1 alloc")?, + heads_b1_d: stream.alloc_zeros::(N_HORIZONS * HEAD_MID_DIM) + .context("v2 heads_b1 alloc")?, + heads_w2_d: stream.alloc_zeros::(N_HORIZONS * HEAD_MID_DIM * HEAD_MID_DIM) + .context("v2 heads_w2 alloc")?, + heads_b2_d: stream.alloc_zeros::(N_HORIZONS * HEAD_MID_DIM) + .context("v2 heads_b2 alloc")?, heads_w_gate_d: stream.alloc_zeros::(N_HORIZONS * HEAD_MID_DIM) .context("v2 heads_w_gate alloc")?, heads_b_gate_d: stream.alloc_zeros::(N_HORIZONS) diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index adcf7ced9..108c274a7 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -214,16 +214,9 @@ pub struct PerceptionTrainer { // (eta_2 → eta_1) + GLU gate + skip-projection. Gives per-horizon // "linear vs deeper-transform" specialisation, matching the // regime-conditional alpha pattern (pearl_snapshot_alpha_is_regime_conditional). - pub heads_w1_d: CudaSlice, // [N_HORIZONS, HEAD_MID, HIDDEN] - pub heads_b1_d: CudaSlice, // [N_HORIZONS, HEAD_MID] - pub heads_w2_d: CudaSlice, // [N_HORIZONS, HEAD_MID, HEAD_MID] - pub heads_b2_d: CudaSlice, // [N_HORIZONS, HEAD_MID] - pub heads_w_gate_d: CudaSlice, // [N_HORIZONS, HEAD_MID] - pub heads_b_gate_d: CudaSlice, // [N_HORIZONS] - pub heads_w_main_d: CudaSlice, // [N_HORIZONS, HEAD_MID] - pub heads_b_main_d: CudaSlice, // [N_HORIZONS] - pub heads_w_skip_d: CudaSlice, // [N_HORIZONS, HIDDEN] - pub heads_b_skip_d: CudaSlice, // [N_HORIZONS] + // X9: GRN head weights moved to `self.trunk.heads_w1_d` / `heads_b1_d` + // / ... / `heads_b_skip_d`. Gradient buffers + AdamW state stay at + // trainer level. pub opt_w_in: AdamW, pub opt_w_rec: AdamW, pub opt_b: AdamW, @@ -720,16 +713,17 @@ impl PerceptionTrainer { stream.memcpy_htod(&w_rec, &mut trunk.w_rec_d).context("trunk.w_rec_d upload")?; stream.memcpy_htod(&b, &mut trunk.b_d).context("trunk.b_d upload")?; stream.memcpy_htod(&tau, &mut trunk.tau_d).context("trunk.tau_d upload")?; - let heads_w1_d = upload(&stream, &heads_w1)?; - let heads_b1_d = upload(&stream, &heads_b1)?; - let heads_w2_d = upload(&stream, &heads_w2)?; - let heads_b2_d = upload(&stream, &heads_b2)?; - let heads_w_gate_d = upload(&stream, &heads_w_gate)?; - let heads_b_gate_d = upload(&stream, &heads_b_gate)?; - let heads_w_main_d = upload(&stream, &heads_w_main)?; - let heads_b_main_d = upload(&stream, &heads_b_main)?; - let heads_w_skip_d = upload(&stream, &heads_w_skip)?; - let heads_b_skip_d = upload(&stream, &heads_b_skip)?; + // X9: upload trainer's GRN heads into the trunk's slots. + stream.memcpy_htod(&heads_w1, &mut trunk.heads_w1_d).context("trunk.heads_w1 upload")?; + stream.memcpy_htod(&heads_b1, &mut trunk.heads_b1_d).context("trunk.heads_b1 upload")?; + stream.memcpy_htod(&heads_w2, &mut trunk.heads_w2_d).context("trunk.heads_w2 upload")?; + stream.memcpy_htod(&heads_b2, &mut trunk.heads_b2_d).context("trunk.heads_b2 upload")?; + stream.memcpy_htod(&heads_w_gate, &mut trunk.heads_w_gate_d).context("trunk.heads_w_gate upload")?; + stream.memcpy_htod(&heads_b_gate, &mut trunk.heads_b_gate_d).context("trunk.heads_b_gate upload")?; + stream.memcpy_htod(&heads_w_main, &mut trunk.heads_w_main_d).context("trunk.heads_w_main upload")?; + stream.memcpy_htod(&heads_b_main, &mut trunk.heads_b_main_d).context("trunk.heads_b_main upload")?; + stream.memcpy_htod(&heads_w_skip, &mut trunk.heads_w_skip_d).context("trunk.heads_w_skip upload")?; + stream.memcpy_htod(&heads_b_skip, &mut trunk.heads_b_skip_d).context("trunk.heads_b_skip upload")?; let opt_w_in = AdamW::new(dev, n_hid * n_in, cfg.lr_cfc)?; let opt_w_rec = AdamW::new(dev, n_hid * n_hid, cfg.lr_cfc)?; @@ -1045,16 +1039,6 @@ impl PerceptionTrainer { ts_ns_all_d: stream.alloc_zeros::(cfg.n_batch * k)?, prev_ts_ns_all_d: stream.alloc_zeros::(cfg.n_batch * k)?, - heads_w1_d, - heads_b1_d, - heads_w2_d, - heads_b2_d, - heads_w_gate_d, - heads_b_gate_d, - heads_w_main_d, - heads_b_main_d, - heads_w_skip_d, - heads_b_skip_d, stream, _snap_module: snap_module, _step_module: step_module, @@ -1727,11 +1711,11 @@ impl PerceptionTrainer { unsafe { let mut launch = self.stream.launch_builder(&self.heads_grn_fwd_fn); launch - .arg(&self.heads_w1_d).arg(&self.heads_b1_d) - .arg(&self.heads_w2_d).arg(&self.heads_b2_d) - .arg(&self.heads_w_gate_d).arg(&self.heads_b_gate_d) - .arg(&self.heads_w_main_d).arg(&self.heads_b_main_d) - .arg(&self.heads_w_skip_d).arg(&self.heads_b_skip_d) + .arg(&self.trunk.heads_w1_d).arg(&self.trunk.heads_b1_d) + .arg(&self.trunk.heads_w2_d).arg(&self.trunk.heads_b2_d) + .arg(&self.trunk.heads_w_gate_d).arg(&self.trunk.heads_b_gate_d) + .arg(&self.trunk.heads_w_main_d).arg(&self.trunk.heads_b_main_d) + .arg(&self.trunk.heads_w_skip_d).arg(&self.trunk.heads_b_skip_d) .arg(&h_new_k_ptr).arg(&n_batch_i) .arg(&probs_k_ptr) .arg(&z1_k_ptr).arg(&a1_k_ptr).arg(&z2_k_ptr) @@ -1855,9 +1839,9 @@ impl PerceptionTrainer { unsafe { let mut launch = self.stream.launch_builder(&self.heads_grn_bwd_fn); launch - .arg(&self.heads_w1_d).arg(&self.heads_w2_d) - .arg(&self.heads_w_gate_d).arg(&self.heads_w_main_d) - .arg(&self.heads_w_skip_d) + .arg(&self.trunk.heads_w1_d).arg(&self.trunk.heads_w2_d) + .arg(&self.trunk.heads_w_gate_d).arg(&self.trunk.heads_w_main_d) + .arg(&self.trunk.heads_w_skip_d) .arg(&probs_k_ptr).arg(&gprobs_k_ptr) .arg(&z1_k_ptr).arg(&a1_k_ptr).arg(&z2_k_ptr) .arg(&gate_k_ptr).arg(&main_k_ptr) @@ -2275,16 +2259,16 @@ impl PerceptionTrainer { self.opt_w_rec.step(&mut self.trunk.w_rec_d, &self.grad_w_rec_d)?; self.opt_b.step(&mut self.trunk.b_d, &self.grad_b_d)?; self.opt_tau.step(&mut self.trunk.tau_d, &self.grad_tau_d)?; - self.opt_heads_w1.step(&mut self.heads_w1_d, &self.grad_heads_w1_d)?; - self.opt_heads_b1.step(&mut self.heads_b1_d, &self.grad_heads_b1_d)?; - self.opt_heads_w2.step(&mut self.heads_w2_d, &self.grad_heads_w2_d)?; - self.opt_heads_b2.step(&mut self.heads_b2_d, &self.grad_heads_b2_d)?; - self.opt_heads_w_gate.step(&mut self.heads_w_gate_d, &self.grad_heads_w_gate_d)?; - self.opt_heads_b_gate.step(&mut self.heads_b_gate_d, &self.grad_heads_b_gate_d)?; - self.opt_heads_w_main.step(&mut self.heads_w_main_d, &self.grad_heads_w_main_d)?; - self.opt_heads_b_main.step(&mut self.heads_b_main_d, &self.grad_heads_b_main_d)?; - self.opt_heads_w_skip.step(&mut self.heads_w_skip_d, &self.grad_heads_w_skip_d)?; - self.opt_heads_b_skip.step(&mut self.heads_b_skip_d, &self.grad_heads_b_skip_d)?; + self.opt_heads_w1.step(&mut self.trunk.heads_w1_d, &self.grad_heads_w1_d)?; + self.opt_heads_b1.step(&mut self.trunk.heads_b1_d, &self.grad_heads_b1_d)?; + self.opt_heads_w2.step(&mut self.trunk.heads_w2_d, &self.grad_heads_w2_d)?; + self.opt_heads_b2.step(&mut self.trunk.heads_b2_d, &self.grad_heads_b2_d)?; + self.opt_heads_w_gate.step(&mut self.trunk.heads_w_gate_d, &self.grad_heads_w_gate_d)?; + self.opt_heads_b_gate.step(&mut self.trunk.heads_b_gate_d, &self.grad_heads_b_gate_d)?; + self.opt_heads_w_main.step(&mut self.trunk.heads_w_main_d, &self.grad_heads_w_main_d)?; + self.opt_heads_b_main.step(&mut self.trunk.heads_b_main_d, &self.grad_heads_b_main_d)?; + self.opt_heads_w_skip.step(&mut self.trunk.heads_w_skip_d, &self.grad_heads_w_skip_d)?; + self.opt_heads_b_skip.step(&mut self.trunk.heads_b_skip_d, &self.grad_heads_b_skip_d)?; self.opt_ln_gain.step(&mut self.trunk.ln_b_gain_d, &self.grad_ln_gain_d)?; self.opt_ln_bias.step(&mut self.trunk.ln_b_bias_d, &self.grad_ln_bias_d)?; self.opt_ln_a_gain.step(&mut self.trunk.ln_a_gain_d, &self.grad_ln_a_gain_d)?; @@ -2668,11 +2652,11 @@ impl PerceptionTrainer { unsafe { let mut launch = self.stream.launch_builder(&self.heads_grn_fwd_fn); launch - .arg(&self.heads_w1_d).arg(&self.heads_b1_d) - .arg(&self.heads_w2_d).arg(&self.heads_b2_d) - .arg(&self.heads_w_gate_d).arg(&self.heads_b_gate_d) - .arg(&self.heads_w_main_d).arg(&self.heads_b_main_d) - .arg(&self.heads_w_skip_d).arg(&self.heads_b_skip_d) + .arg(&self.trunk.heads_w1_d).arg(&self.trunk.heads_b1_d) + .arg(&self.trunk.heads_w2_d).arg(&self.trunk.heads_b2_d) + .arg(&self.trunk.heads_w_gate_d).arg(&self.trunk.heads_b_gate_d) + .arg(&self.trunk.heads_w_main_d).arg(&self.trunk.heads_b_main_d) + .arg(&self.trunk.heads_w_skip_d).arg(&self.trunk.heads_b_skip_d) .arg(&h_new_k_ptr).arg(&n_batch_i) .arg(&probs_k_ptr) .arg(&z1_k_ptr).arg(&a1_k_ptr).arg(&z2_k_ptr)