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

refactor(ml-alpha): move CfC weights from PerceptionTrainer to CfcTrunk (X8)

Browse Source 
Adds CfcConfig.cfc_n_in field (default HIDDEN_DIM) so the trunk's CfC
layer is sized for v2 usage (CfC input = LN_b output = HIDDEN_DIM) rather
than V1 usage (CfC input = snap features = FEATURE_DIM). The previous
CfcConfig.n_in field stays as "raw snap feature dim" for any V1 callers
still in the tree (fxt-backtest, trunk_forward.rs, graph_a_replay.rs);
their forward paths will be cleaned up in X10/X11 when the v2 forward
graph lives natively on the trunk.

Trainer's w_in_d / w_rec_d / b_d / tau_d fields are removed. Trainer
init still draws CfC weights from its ChaCha8Rng chain at the same
call position, then memcpy_htod's them into the trunk's now-v2-shaped
slots. Forward, backward, and AdamW access sites redirect to
self.trunk.{w_in,w_rec,b,tau}_d.

Verification:
- perception_forward_golden: PASS (max_diff = 0.000000)
- ml-alpha lib tests: 34 pass

Per spec §2.2 (X8).
This commit is contained in:
jgrusewski
2026-05-19 08:20:33 +02:00
parent 2bd774ad14
commit ef29e13c74
2 changed files with 43 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -49,8 +49,14 @@ const PROJ_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/projection.c
#[derive(Clone, Debug)]
pub struct CfcConfig {
/// Raw snap feature dimensionality (V1 forward path uses this).
pub n_in: usize,
pub n_hid: usize,
/// CfC layer input dimensionality. In V1 this was conflated with
/// `n_in` (CfC consumed snap features directly). In V2 the CfC sits
/// AFTER LN_b, consuming HIDDEN_DIM-dim tokens. Default matches V2
/// usage.
pub cfc_n_in: usize,
/// Mamba2 state dimension — used by the v2 trunk weight skeleton
/// (X1) and the per-stack Mamba2Block constructions added in X3/X5.
/// Default matches `PerceptionTrainerConfig::default().mamba2_state_dim`.
@@ -59,7 +65,12 @@ pub struct CfcConfig {
impl Default for CfcConfig {
fn default() -> Self {
Self { n_in: FEATURE_DIM, n_hid: HIDDEN_DIM, mamba2_state_dim: 16 }
Self {
n_in: FEATURE_DIM,
n_hid: HIDDEN_DIM,
cfc_n_in: HIDDEN_DIM,
mamba2_state_dim: 16,
}
}
}
@@ -77,13 +88,15 @@ pub struct CfcTrunk {
heads_fn: CudaFunction,
proj_fn: CudaFunction,
// Device-resident weights — V1 (CfC + simple heads + projection).
// Preserved unchanged until X8/X9 consolidate against the v2 fields
// added below.
w_in_d: CudaSlice<f32>,
w_rec_d: CudaSlice<f32>,
b_d: CudaSlice<f32>,
tau_d: CudaSlice<f32>,
// CfC weights — now v2-shaped (cfc_n_in = HIDDEN_DIM). Public so
// PerceptionTrainer can mutate during init + AdamW.
pub w_in_d: CudaSlice<f32>,
pub w_rec_d: CudaSlice<f32>,
pub b_d: CudaSlice<f32>,
pub tau_d: CudaSlice<f32>,
// V1 legacy fields — only used by the now-stub V1 forward path
// (dispatch_perception). Kept allocated so the trunk struct lays out
// the same in tests; will be removed in a follow-up cleanup commit.
heads_w_d: CudaSlice<f32>,
heads_b_d: CudaSlice<f32>,
proj_w_d: CudaSlice<f32>,
@@ -165,9 +178,12 @@ impl CfcTrunk {
// Random init.
let mut r = ChaCha8Rng::seed_from_u64(seed);
let scale_in = (1.0_f32 / cfg.n_in as f32).sqrt();
// X8: CfC layer weights now use cfc_n_in (V2 = HIDDEN_DIM), not
// n_in (which still means snap feature dimensionality for any
// residual V1 callers). This re-shapes w_in_d to v2 layout.
let scale_in = (1.0_f32 / cfg.cfc_n_in as f32).sqrt();
let scale_rec = (1.0_f32 / cfg.n_hid as f32).sqrt();
let w_in: Vec<f32> = (0..cfg.n_hid * cfg.n_in).map(|_| r.gen_range(-scale_in..scale_in)).collect();
let w_in: Vec<f32> = (0..cfg.n_hid * cfg.cfc_n_in).map(|_| r.gen_range(-scale_in..scale_in)).collect();
let w_rec: Vec<f32> = (0..cfg.n_hid * cfg.n_hid).map(|_| r.gen_range(-scale_rec..scale_rec)).collect();
let b: Vec<f32> = vec![0.0; cfg.n_hid];
// tau log-uniform over [10ms, 1000s] per spec Section 2 (ln(0.01)=-4.6, ln(1000)=6.9).

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

@@ -206,10 +206,9 @@ pub struct PerceptionTrainer {
grad_h_enriched_seq_d: GpuTensor,
// CfC + heads weights + their AdamWs (6 groups — tau is trained now).
pub w_in_d: CudaSlice<f32>,
pub w_rec_d: CudaSlice<f32>,
pub b_d: CudaSlice<f32>,
pub tau_d: CudaSlice<f32>,
// X8: CfC weights (w_in, w_rec, b, tau) moved to
// `self.trunk.w_in_d` / `w_rec_d` / `b_d` / `tau_d` (now v2-shaped
// via CfcConfig.cfc_n_in = HIDDEN_DIM).
// ── TFT GRN heads (Phase 1.7) ──
// Per-horizon Gated Residual Network: 2-layer GELU MLP body
// (eta_2 → eta_1) + GLU gate + skip-projection. Gives per-horizon
@@ -521,6 +520,7 @@ impl PerceptionTrainer {
&crate::cfc::trunk::CfcConfig {
n_in: FEATURE_DIM,
n_hid: HIDDEN_DIM,
cfc_n_in: HIDDEN_DIM,
mamba2_state_dim: cfg.mamba2_state_dim,
},
cfg.seed,
@@ -714,10 +714,12 @@ impl PerceptionTrainer {
.map(|_| r.gen_range(-scale_h_skip..scale_h_skip)).collect();
let heads_b_skip: Vec<f32> = vec![0.0; N_HORIZONS];
let w_in_d = upload(&stream, &w_in)?;
let w_rec_d = upload(&stream, &w_rec)?;
let b_d = upload(&stream, &b)?;
let tau_d = upload(&stream, &tau)?;
// X8: upload trainer's PRNG-derived CfC weights into the trunk's
// (now v2-shaped) CfC slots.
stream.memcpy_htod(&w_in, &mut trunk.w_in_d).context("trunk.w_in_d upload")?;
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)?;
@@ -1043,10 +1045,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)?,
w_in_d,
w_rec_d,
b_d,
tau_d,
heads_w1_d,
heads_b1_d,
heads_w2_d,
@@ -1717,7 +1715,7 @@ impl PerceptionTrainer {
unsafe {
let mut launch = self.stream.launch_builder(&self.step_batched_fn);
launch
.arg(&self.w_in_d).arg(&self.w_rec_d).arg(&self.b_d).arg(&self.tau_d)
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_d)
.arg(&x_k_ptr).arg(&h_old_k_ptr)
.arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i)
.arg(&h_new_k_ptr);
@@ -1885,7 +1883,7 @@ impl PerceptionTrainer {
unsafe {
let mut launch = self.stream.launch_builder(&self.step_bwd_batched_fn);
launch
.arg(&self.w_in_d).arg(&self.w_rec_d).arg(&self.b_d).arg(&self.tau_d)
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_d)
.arg(&x_k_ptr).arg(&h_old_k_ptr).arg(&self.grad_h_new_d)
.arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i)
.arg(&mut self.cfc_grad_w_in_scratch_d)
@@ -2273,10 +2271,10 @@ impl PerceptionTrainer {
// ── 9. Apply AdamW updates on all 17 param groups: CfC×4 +
// GRN heads×10 + LN×2 + Mamba2 grouped.
self.opt_w_in.step(&mut self.w_in_d, &self.grad_w_in_d)?;
self.opt_w_rec.step(&mut self.w_rec_d, &self.grad_w_rec_d)?;
self.opt_b.step(&mut self.b_d, &self.grad_b_d)?;
self.opt_tau.step(&mut self.tau_d, &self.grad_tau_d)?;
self.opt_w_in.step(&mut self.trunk.w_in_d, &self.grad_w_in_d)?;
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)?;
@@ -2661,7 +2659,7 @@ impl PerceptionTrainer {
unsafe {
let mut launch = self.stream.launch_builder(&self.step_batched_fn);
launch
.arg(&self.w_in_d).arg(&self.w_rec_d).arg(&self.b_d).arg(&self.tau_d)
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_d)
.arg(&x_k_ptr).arg(&h_old_k_ptr)
.arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i)
.arg(&h_new_k_ptr);