gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(ml-alpha): move GRN heads from PerceptionTrainer to CfcTrunk (X9)

Browse Source 
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.
This commit is contained in:
jgrusewski
2026-05-19 08:24:48 +02:00
parent ef29e13c74
commit 5bd8897880
2 changed files with 52 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
crates/ml-alpha/src/cfc/trunk.rs
View File

@@ -120,6 +120,13 @@ pub struct CfcTrunk {
pub ln_b_gain_d: CudaSlice<f32>,
pub ln_b_bias_d: CudaSlice<f32>,
pub attn_q_d: CudaSlice<f32>,
// 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<f32>,
pub heads_b1_d: CudaSlice<f32>,
pub heads_w2_d: CudaSlice<f32>,
pub heads_b2_d: CudaSlice<f32>,
pub heads_w_gate_d: CudaSlice<f32>,
pub heads_b_gate_d: CudaSlice<f32>,
pub heads_w_main_d: CudaSlice<f32>,
@@ -254,6 +261,14 @@ impl CfcTrunk {
.context("v2 ln_b_bias alloc")?,
attn_q_d: stream.alloc_zeros::<f32>(cfg.n_hid)
.context("v2 attn_q alloc")?,
heads_w1_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM * cfg.n_hid)
.context("v2 heads_w1 alloc")?,
heads_b1_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM)
.context("v2 heads_b1 alloc")?,
heads_w2_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM * HEAD_MID_DIM)
.context("v2 heads_w2 alloc")?,
heads_b2_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM)
.context("v2 heads_b2 alloc")?,
heads_w_gate_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM)
.context("v2 heads_w_gate alloc")?,
heads_b_gate_d: stream.alloc_zeros::<f32>(N_HORIZONS)

90
crates/ml-alpha/src/trainer/perception.rs
View File

@@ -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<f32>, // [N_HORIZONS, HEAD_MID, HIDDEN]
pub heads_b1_d: CudaSlice<f32>, // [N_HORIZONS, HEAD_MID]
pub heads_w2_d: CudaSlice<f32>, // [N_HORIZONS, HEAD_MID, HEAD_MID]
pub heads_b2_d: CudaSlice<f32>, // [N_HORIZONS, HEAD_MID]
pub heads_w_gate_d: CudaSlice<f32>, // [N_HORIZONS, HEAD_MID]
pub heads_b_gate_d: CudaSlice<f32>, // [N_HORIZONS]
pub heads_w_main_d: CudaSlice<f32>, // [N_HORIZONS, HEAD_MID]
pub heads_b_main_d: CudaSlice<f32>, // [N_HORIZONS]
pub heads_w_skip_d: CudaSlice<f32>, // [N_HORIZONS, HIDDEN]
pub heads_b_skip_d: CudaSlice<f32>, // [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::<i64>(cfg.n_batch * k)?,
prev_ts_ns_all_d: stream.alloc_zeros::<i64>(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)