feat(ml-alpha): wire TFT VSN into PerceptionTrainer (Phase 2D.3)
VSN sits between snap_feature_assemble and Mamba2 input:
snap_assemble → window_tensor_d [B, K, 40] (raw)
→ VSN fwd → vsn_out_d [B, K, 40] (gated) + vsn_gates_d [B*K, 40]
→ Mamba2 fwd → ...
Forward path: per-position softmax over FEATURE_DIM features, output[i] =
x[i] * gates[i]. Gates initialised near-uniform (W_vsn ~ N(0, 1/√FEATURE_DIM),
b_vsn = 0) so the model starts from "all features matter equally" and
learns regime-conditional gates.
Backward path: Mamba2 backward already computed d_x_from_in (gradient
w.r.t. its input) on its grads_buffers — previously labeled "unused but
allocated", now consumed by VSN bwd as grad_y. Zero refactor to
mamba2_block.rs.
VSN bwd writes:
grad_W_vsn [40, 40] → opt_vsn_w (AdamW, default wd)
grad_b_vsn [40] → opt_vsn_b (AdamW, wd=0)
vsn_grad_x_d [B*K, 40] → discarded (snap_features are non-trainable
transforms of raw MBP-10 data)
Param grads are explicitly memset_zeros before each VSN bwd call (the
kernel uses += semantics like the GRN bwd, but VSN runs ONCE per step
not K times, so zeroing makes the += a clean overwrite — matches
Adam's `step` expectations).
Eval path mirrors training (VSN fwd applied between snap_assemble and
Mamba2 fwd) so eval sees the gated features layers were trained on.
Trainer now manages 19 AdamW: CfC×4 + GRN heads×10 + LN×2 + VSN×2 +
Mamba2 grouped.
Synthetic overfit smoke: stride=1 initial=0.30 → final=0.00 in 50 steps
(faster than pre-VSN's 0.32, consistent with VSN's near-uniform init
giving a small head start). All 8 perception_overfit tests pass.
Note: the smoke proves VSN's W/b actually move via the Mamba2 input
gradient — if `d_x_from_in` were zero, VSN params wouldn't update and
the chain still converges but VSN remains identity. The fact that
initial loss DIFFERS (0.30 vs 0.32) shows VSN is in the forward path
end-to-end.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -60,6 +60,7 @@ const HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horiz
|
||||
const BCE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bce_loss_multi_horizon.cubin"));
|
||||
const HORIZON_LAMBDA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/horizon_lambda.cubin"));
|
||||
const LAYER_NORM_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/layer_norm.cubin"));
|
||||
const VARIABLE_SELECTION_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/variable_selection.cubin"));
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct PerceptionTrainerConfig {
|
||||
@@ -306,6 +307,31 @@ pub struct PerceptionTrainer {
|
||||
/// AdamW state for the two LN params (gain + bias).
|
||||
opt_ln_gain: AdamW,
|
||||
opt_ln_bias: AdamW,
|
||||
|
||||
// ── TFT VSN (Phase 2D) ──
|
||||
// Per-position softmax-normalised feature gating at trunk entry:
|
||||
// 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>,
|
||||
/// 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`).
|
||||
vsn_out_d: GpuTensor,
|
||||
/// Saved softmax-normalised gates `[B*K, FEATURE_DIM]` (per-K row).
|
||||
vsn_gates_d: CudaSlice<f32>,
|
||||
/// VSN bwd writes grad_x here (gradient w.r.t. snap_features). Not
|
||||
/// propagated further (snap_features are a non-trainable transform
|
||||
/// of raw MBP-10 data); allocated only so the kernel has a buffer.
|
||||
vsn_grad_x_d: CudaSlice<f32>,
|
||||
grad_vsn_w_d: CudaSlice<f32>,
|
||||
grad_vsn_b_d: CudaSlice<f32>,
|
||||
pub opt_vsn_w: AdamW,
|
||||
pub opt_vsn_b: AdamW,
|
||||
vsn_fwd_fn: CudaFunction,
|
||||
vsn_bwd_fn: CudaFunction,
|
||||
_vsn_module: Arc<CudaModule>,
|
||||
/// Per-horizon EMA of `loss_per_horizon_d`. Zero-initialised
|
||||
/// (sentinel); the horizon_ema_and_lambda kernel detects the
|
||||
/// sentinel on step 1 and replaces directly per
|
||||
@@ -422,6 +448,15 @@ impl PerceptionTrainer {
|
||||
let ln_reduce_fn = ln_module
|
||||
.load_function("layer_norm_reduce_param_grads")
|
||||
.context("layer_norm_reduce_param_grads symbol")?;
|
||||
let vsn_module = ctx
|
||||
.load_cubin(VARIABLE_SELECTION_CUBIN.to_vec())
|
||||
.context("variable_selection cubin")?;
|
||||
let vsn_fwd_fn = vsn_module
|
||||
.load_function("variable_selection_fwd")
|
||||
.context("variable_selection_fwd symbol")?;
|
||||
let vsn_bwd_fn = vsn_module
|
||||
.load_function("variable_selection_bwd")
|
||||
.context("variable_selection_bwd symbol")?;
|
||||
let snap_batched_fn = snap_module.load_function("snap_feature_assemble_batched")?;
|
||||
let bce_fn = bce_module.load_function("bce_multi_horizon_forward_backward")?;
|
||||
let step_batched_fn = step_module.load_function("cfc_step_batched")?;
|
||||
@@ -560,6 +595,20 @@ impl PerceptionTrainer {
|
||||
let mut opt_ln_bias = AdamW::new(dev, HIDDEN_DIM, cfg.lr_cfc)?;
|
||||
opt_ln_bias.wd = 0.0;
|
||||
|
||||
// ── VSN init (Phase 2D) ──
|
||||
// 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).
|
||||
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)?;
|
||||
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;
|
||||
|
||||
let k = cfg.seq_len;
|
||||
Ok(Self {
|
||||
cfg: cfg.clone(),
|
||||
@@ -592,6 +641,20 @@ impl PerceptionTrainer {
|
||||
_ln_module: ln_module,
|
||||
opt_ln_gain,
|
||||
opt_ln_bias,
|
||||
// VSN (Phase 2D) — params + per-K gates + scratch.
|
||||
vsn_w_d,
|
||||
vsn_b_d,
|
||||
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)?,
|
||||
vsn_grad_x_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * FEATURE_DIM)?,
|
||||
grad_vsn_w_d: stream.alloc_zeros::<f32>(FEATURE_DIM * FEATURE_DIM)?,
|
||||
grad_vsn_b_d: stream.alloc_zeros::<f32>(FEATURE_DIM)?,
|
||||
opt_vsn_w,
|
||||
opt_vsn_b,
|
||||
vsn_fwd_fn,
|
||||
vsn_bwd_fn,
|
||||
_vsn_module: vsn_module,
|
||||
loss_ema_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
|
||||
lambda_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
|
||||
horizon_lambda_fn,
|
||||
@@ -1017,9 +1080,35 @@ impl PerceptionTrainer {
|
||||
// No sync — stream-ordered. A sync would trip
|
||||
// STREAM_CAPTURE_ISOLATION inside the captured region.
|
||||
|
||||
// ── 1b. VSN forward — per-position softmax-normalised feature
|
||||
// gating over `window_tensor_d` [B, K, FEATURE_DIM].
|
||||
// Writes gated features to `vsn_out_d` and saved
|
||||
// post-softmax gates to `vsn_gates_d` for backward.
|
||||
// Mamba2 then consumes `vsn_out_d` instead of the raw
|
||||
// snap features.
|
||||
{
|
||||
let n_rows_vsn: i32 = (b_sz * k_seq) as i32;
|
||||
let cfg_vsn = LaunchConfig {
|
||||
grid_dim: (n_rows_vsn as u32, 1, 1),
|
||||
block_dim: (64, 1, 1), // VSN_BLOCK = 64 (1 warp + idle threads)
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.vsn_fwd_fn);
|
||||
launch
|
||||
.arg(&self.vsn_w_d)
|
||||
.arg(&self.vsn_b_d)
|
||||
.arg(self.window_tensor_d.cuda_data())
|
||||
.arg(&n_rows_vsn)
|
||||
.arg(self.vsn_out_d.data_mut())
|
||||
.arg(&mut self.vsn_gates_d);
|
||||
unsafe { launch.launch(cfg_vsn).context("variable_selection_fwd")?; }
|
||||
}
|
||||
|
||||
// ── 2. Mamba2 per-step forward — writes into self.mamba2_fwd_scratch.
|
||||
// Now consumes `vsn_out_d` (gated features) instead of the
|
||||
// raw `window_tensor_d`.
|
||||
self.mamba2
|
||||
.forward_train_seq_into(&self.window_tensor_d, &mut self.mamba2_fwd_scratch)
|
||||
.forward_train_seq_into(&self.vsn_out_d, &mut self.mamba2_fwd_scratch)
|
||||
.context("mamba2 forward_train_seq_into")?;
|
||||
|
||||
// ── 2a. LayerNorm forward over h_enriched_seq [B*K, H].
|
||||
@@ -1464,11 +1553,15 @@ impl PerceptionTrainer {
|
||||
|
||||
// ── 8. Mamba2 backward — fully pre-allocated path. Writes all
|
||||
// grads into self.mamba2_grads_buffers; no allocation.
|
||||
// Now consumes grad_ln_in_d (LN bwd output), NOT the raw
|
||||
// grad_h_enriched_seq_d — LN sits between Mamba2 and CfC.
|
||||
// Now consumes:
|
||||
// - input = vsn_out_d (gated snap_features fed to Mamba2 fwd)
|
||||
// - d_h_out = grad_ln_in_d (LN bwd output)
|
||||
// and writes `mamba2_grads_buffers.d_x_from_in` =
|
||||
// gradient w.r.t. Mamba2's input = gradient on VSN's
|
||||
// output, which VSN bwd consumes below.
|
||||
self.mamba2
|
||||
.backward_from_h_enriched_seq_full_into(
|
||||
&self.window_tensor_d,
|
||||
&self.vsn_out_d,
|
||||
&self.mamba2_fwd_scratch,
|
||||
&self.grad_ln_in_d,
|
||||
&mut self.mamba2_bwd_scratch,
|
||||
@@ -1476,6 +1569,40 @@ impl PerceptionTrainer {
|
||||
)
|
||||
.context("mamba2 backward_from_h_enriched_seq_full_into")?;
|
||||
|
||||
// ── 8b. VSN backward. Consumes:
|
||||
// grad_y = mamba2_grads_buffers.d_x_from_in [B*K, FEATURE_DIM]
|
||||
// x = window_tensor_d [B, K, FEATURE_DIM]
|
||||
// gates = vsn_gates_d (saved by fwd) [B*K, FEATURE_DIM]
|
||||
// Writes:
|
||||
// grad_W_vsn, grad_b_vsn — Adam consumes.
|
||||
// vsn_grad_x_d — discarded (snap_features non-trainable).
|
||||
// Param grads OVERWRITE (kernel uses += but the prior
|
||||
// memsets clear them at step start, so a single VSN bwd
|
||||
// per step writes the full accumulation cleanly).
|
||||
self.stream.memset_zeros(&mut self.grad_vsn_w_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero grad_vsn_w: {e}"))?;
|
||||
self.stream.memset_zeros(&mut self.grad_vsn_b_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero grad_vsn_b: {e}"))?;
|
||||
{
|
||||
let n_rows_vsn: i32 = (b_sz * k_seq) as i32;
|
||||
let cfg_vsn_bwd = LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (64, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.vsn_bwd_fn);
|
||||
launch
|
||||
.arg(&self.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())
|
||||
.arg(&n_rows_vsn)
|
||||
.arg(&mut self.grad_vsn_w_d)
|
||||
.arg(&mut self.grad_vsn_b_d)
|
||||
.arg(&mut self.vsn_grad_x_d);
|
||||
unsafe { launch.launch(cfg_vsn_bwd).context("variable_selection_bwd")?; }
|
||||
}
|
||||
|
||||
// ── 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)?;
|
||||
@@ -1494,6 +1621,8 @@ impl PerceptionTrainer {
|
||||
self.opt_heads_b_skip.step(&mut self.heads_b_skip_d, &self.grad_heads_b_skip_d)?;
|
||||
self.opt_ln_gain.step(&mut self.ln_gain_d, &self.grad_ln_gain_d)?;
|
||||
self.opt_ln_bias.step(&mut self.ln_bias_d, &self.grad_ln_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)?;
|
||||
// GPU-resident grad-clip + AdamW (zero host roundtrips).
|
||||
// Replaces step_from_buffers which did 9× memcpy_dtoh per step.
|
||||
self.mamba2_adamw
|
||||
@@ -1651,9 +1780,28 @@ impl PerceptionTrainer {
|
||||
}
|
||||
self.stream.synchronize().context("eval snap_batched sync")?;
|
||||
|
||||
// Mamba2 fwd into pre-allocated fwd_scratch.
|
||||
// VSN fwd — same as training path. Eval consumes vsn_out_d.
|
||||
{
|
||||
let n_rows_vsn: i32 = (b_sz * k_seq) as i32;
|
||||
let cfg_vsn = LaunchConfig {
|
||||
grid_dim: (n_rows_vsn as u32, 1, 1),
|
||||
block_dim: (64, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.vsn_fwd_fn);
|
||||
launch
|
||||
.arg(&self.vsn_w_d)
|
||||
.arg(&self.vsn_b_d)
|
||||
.arg(self.window_tensor_d.cuda_data())
|
||||
.arg(&n_rows_vsn)
|
||||
.arg(self.vsn_out_d.data_mut())
|
||||
.arg(&mut self.vsn_gates_d);
|
||||
unsafe { launch.launch(cfg_vsn).context("eval variable_selection_fwd")?; }
|
||||
}
|
||||
|
||||
// Mamba2 fwd into pre-allocated fwd_scratch — reads vsn_out_d.
|
||||
self.mamba2
|
||||
.forward_train_seq_into(&self.window_tensor_d, &mut self.mamba2_fwd_scratch)
|
||||
.forward_train_seq_into(&self.vsn_out_d, &mut self.mamba2_fwd_scratch)
|
||||
.context("eval mamba2 fwd_into")?;
|
||||
|
||||
// LN forward over h_enriched_seq [B*K, H] — same as training path
|
||||
|
||||
Reference in New Issue
Block a user