From 5bd8897880084005097ac47b2ed4722f6dc7d4a0 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 19 May 2026 08:24:48 +0200 Subject: [PATCH] refactor(ml-alpha): move GRN heads from PerceptionTrainer to CfcTrunk (X9) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the 4 missing GRN head fields (heads_w1, heads_b1, heads_w2, heads_b2) that X1's skeleton under-modeled. Trunk now owns all 10 GRN head tensors: input projection (HIDDEN→HEAD_MID), mid→mid layer, gate/main/skip outputs. Trainer's 10 head fields removed. Trainer init still draws heads from its ChaCha8Rng chain at the same call position, then memcpy_htod's them into the trunk's now-allocated slots. Forward, backward, and AdamW access sites redirect to self.trunk.heads_*. Gradient buffers + AdamW state stay at trainer level. Verification: - perception_forward_golden: PASS (max_diff = 0.000000) - ml-alpha lib tests: 34 pass Per spec §2.2 (X9). After this commit, the trunk owns ALL v2 inference weights — the source of truth for downstream checkpoint serialization. --- crates/ml-alpha/src/cfc/trunk.rs | 15 ++++ crates/ml-alpha/src/trainer/perception.rs | 90 ++++++++++------------- 2 files changed, 52 insertions(+), 53 deletions(-) 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)