refactor(ml-alpha): move VSN weights from PerceptionTrainer to CfcTrunk
PerceptionTrainer gains a `trunk: CfcTrunk` field constructed at the top of `new()` (after the determinism seed guard). VSN's weight tensors (`vsn_w_d`, `vsn_b_d`) now live on `self.trunk`; the trainer-owned copies are removed. Forward + backward + AdamW access sites redirect to `self.trunk.vsn_w_d` / `self.trunk.vsn_b_d`. PRNG-state preservation: the trainer's ChaCha8Rng chain still draws VSN values at the same call position as before, then memcpy_htod's them into the trunk's zero-initialised VSN slots from X1. This keeps every downstream weight (attn_q, ...) bit-identical to the pre-X2 layout — perception_forward_golden continues to pass with max_diff = 0.000000. Gradient buffers (`grad_vsn_w_d`, `grad_vsn_b_d`) and AdamW state (`opt_vsn_w`, `opt_vsn_b`) remain at trainer level — they're training-only and shouldn't move to the inference trunk. Verification: - perception_forward_golden: PASS (max_diff = 0.000000) - ml-alpha lib tests: 34 pass - alpha_train example builds clean Per spec 2026-05-19-ml-alpha-v2-trunk-grows-and-deployability-design.md §1.1, §2.2 (X2). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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<f32>,
|
||||
pub vsn_b_d: CudaSlice<f32>,
|
||||
// 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<f32> = (0..FEATURE_DIM * FEATURE_DIM)
|
||||
.map(|_| r.gen_range(-vsn_scale..vsn_scale)).collect();
|
||||
let vsn_b_init: Vec<f32> = 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::<f32>(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())
|
||||
|
||||
Reference in New Issue
Block a user