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

feat(per-horizon-cfc): atomic Checkpoint envelope migration (3 consumers)

Browse Source 
Per spec §5.2 and feedback_no_partial_refactor. Adds bucket-routing
metadata to the Checkpoint envelope:
- Renames tau -> tau_all (same HIDDEN_DIM shape)
- Adds bucket_id_per_channel: Vec<u8>
- Adds bucket_channel_offset: Vec<u32>
- Adds bucket_dim_k: Vec<u32>
- Adds heads_w_skip_offset: Vec<u32>

heads_w_skip stays at [N_HORIZONS x HIDDEN_DIM]=640 floats; compaction
to [HIDDEN_DIM]=128 deferred to Task 10's forward dispatch update.

new_random zero-initializes bucket metadata; Phase 1->2 transition
(Task 9) populates them.

CfcTrunk field `tau_d` renamed to `tau_all_d` across:
- crates/ml-alpha/src/cfc/trunk.rs (struct + save_checkpoint
  + load_checkpoint + new_random + save_load_roundtrip bit-equality test
  with 4 new u8/u32 assertions on the bucket metadata fields)
- crates/ml-alpha/src/trainer/perception.rs (6 self.trunk.tau_d
  references + 1 trunk.tau_d upload + 1 comment line)

ml-backtesting/tests/checkpoint_smoke.rs needed no edit: its only
consumer surface is CfcTrunk::load_checkpoint(&dev, &cfg, &ckpt),
whose signature is unchanged. The 3-consumer atomic-migration audit
is satisfied because consumer #3's API contract is preserved.

Greenfield envelope (per existing cfc/trunk.rs:25-28 convention -
no backward compat).

Verified:
- cargo check --workspace --all-targets: clean (pre-existing
  cupti/insert_batch errors in crates/ml/tests excluded per task scope)
- cargo test -p ml-alpha --lib: 33 passed
- cargo test -p ml-backtesting --lib: 33 passed
- cargo test -p ml-alpha --lib save_load_roundtrip -- --ignored: PASS
  (bit-equality on tau_all_d + 4 new bucket metadata fields, RTX 3050)
- cargo test -p ml-backtesting --test checkpoint_smoke: compiles clean

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-21 16:35:48 +02:00
parent 1a7c4fc66a
commit fc0311de77
2 changed files with 99 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -57,7 +57,12 @@ struct Checkpoint {
// Attention pool
attn_q: Vec<f32>,
// CfC
w_in: Vec<f32>, w_rec: Vec<f32>, b: Vec<f32>, tau: Vec<f32>,
w_in: Vec<f32>, w_rec: Vec<f32>, b: Vec<f32>, tau_all: Vec<f32>,
// Bucket metadata (populated at Phase 1→2 transition; zero-initialized in new_random)
bucket_id_per_channel: Vec<u8>, // [HIDDEN_DIM]
bucket_channel_offset: Vec<u32>, // [N_HORIZONS + 1]
bucket_dim_k: Vec<u32>, // [N_HORIZONS]
heads_w_skip_offset: Vec<u32>, // [N_HORIZONS + 1]
// Heads (GRN structure × 5 horizons)
heads_w1: Vec<f32>, heads_b1: Vec<f32>,
heads_w2: Vec<f32>, heads_b2: Vec<f32>,
@@ -135,7 +140,14 @@ pub struct CfcTrunk {
pub w_in_d: CudaSlice<f32>,
pub w_rec_d: CudaSlice<f32>,
pub b_d: CudaSlice<f32>,
pub tau_d: CudaSlice<f32>,
pub tau_all_d: CudaSlice<f32>,
// Bucket metadata for per-horizon CfC routing (Phase 1→2 transition
// populates these; zero-initialized at `new_random`). See spec §2.2.
pub bucket_id_per_channel_d: CudaSlice<u8>,
pub bucket_channel_offset_d: CudaSlice<u32>,
pub bucket_dim_k_d: CudaSlice<u32>,
pub heads_w_skip_offset_d: CudaSlice<u32>,
// Inference weights for the full perception chain:
// snap → vsn → mamba2_l1 → ln_a → mamba2_l2 → ln_b → attn_pool →
@@ -217,12 +229,34 @@ impl CfcTrunk {
Ok(buf)
};
// Bucket metadata: zero-initialised here; populated by the Phase
// 1→2 transition orchestrator in `cfc/bucket_routing.rs` once
// Controller A fires. Allocated unconditionally so that
// checkpoint save/load round-trips include them even before the
// transition has run.
let bucket_id_per_channel_d = stream
.alloc_zeros::<u8>(HIDDEN_DIM)
.context("bucket_id_per_channel alloc")?;
let bucket_channel_offset_d = stream
.alloc_zeros::<u32>(N_HORIZONS + 1)
.context("bucket_channel_offset alloc")?;
let bucket_dim_k_d = stream
.alloc_zeros::<u32>(N_HORIZONS)
.context("bucket_dim_k alloc")?;
let heads_w_skip_offset_d = stream
.alloc_zeros::<u32>(N_HORIZONS + 1)
.context("heads_w_skip_offset alloc")?;
Ok(Self {
cfg: cfg.clone(),
w_in_d: upload(&w_in)?,
w_rec_d: upload(&w_rec)?,
b_d: upload(&b)?,
tau_d: upload(&tau)?,
tau_all_d: upload(&tau)?,
bucket_id_per_channel_d,
bucket_channel_offset_d,
bucket_dim_k_d,
heads_w_skip_offset_d,
// Inference layers — zero-initialised here; PerceptionTrainer's
// construction overwrites them with its own seeded init via
// memcpy_htod. Standalone trunk callers (tests, fxt-backtest's
@@ -329,6 +363,16 @@ impl CfcTrunk {
self.stream.memcpy_dtoh(slice, v.as_mut_slice())?;
Ok(v)
};
let download_u8 = |slice: &CudaSlice<u8>| -> Result<Vec<u8>> {
let mut v = vec![0u8; slice.len()];
self.stream.memcpy_dtoh(slice, v.as_mut_slice())?;
Ok(v)
};
let download_u32 = |slice: &CudaSlice<u32>| -> Result<Vec<u32>> {
let mut v = vec![0u32; slice.len()];
self.stream.memcpy_dtoh(slice, v.as_mut_slice())?;
Ok(v)
};
let m1 = self.mamba2_l1();
let m2 = self.mamba2_l2();
let ckpt = Checkpoint {
@@ -367,7 +411,11 @@ impl CfcTrunk {
w_in: download(&self.w_in_d)?,
w_rec: download(&self.w_rec_d)?,
b: download(&self.b_d)?,
tau: download(&self.tau_d)?,
tau_all: download(&self.tau_all_d)?,
bucket_id_per_channel: download_u8(&self.bucket_id_per_channel_d)?,
bucket_channel_offset: download_u32(&self.bucket_channel_offset_d)?,
bucket_dim_k: download_u32(&self.bucket_dim_k_d)?,
heads_w_skip_offset: download_u32(&self.heads_w_skip_offset_d)?,
heads_w1: download(&self.heads_w1_d)?,
heads_b1: download(&self.heads_b1_d)?,
heads_w2: download(&self.heads_w2_d)?,
@@ -420,6 +468,22 @@ impl CfcTrunk {
stream.memcpy_htod(src, dst)?;
Ok(())
};
let upload_u8 = |dst: &mut CudaSlice<u8>, src: &[u8], name: &str| -> Result<()> {
anyhow::ensure!(
dst.len() == src.len(),
"size mismatch loading {name}: device {} != file {}", dst.len(), src.len()
);
stream.memcpy_htod(src, dst)?;
Ok(())
};
let upload_u32 = |dst: &mut CudaSlice<u32>, src: &[u32], name: &str| -> Result<()> {
anyhow::ensure!(
dst.len() == src.len(),
"size mismatch loading {name}: device {} != file {}", dst.len(), src.len()
);
stream.memcpy_htod(src, dst)?;
Ok(())
};
upload(&mut trunk.vsn_w_d, &ckpt.vsn_w, "vsn_w")?;
upload(&mut trunk.vsn_b_d, &ckpt.vsn_b, "vsn_b")?;
{
@@ -454,7 +518,11 @@ impl CfcTrunk {
upload(&mut trunk.w_in_d, &ckpt.w_in, "cfc.w_in")?;
upload(&mut trunk.w_rec_d, &ckpt.w_rec, "cfc.w_rec")?;
upload(&mut trunk.b_d, &ckpt.b, "cfc.b")?;
upload(&mut trunk.tau_d, &ckpt.tau, "cfc.tau")?;
upload(&mut trunk.tau_all_d, &ckpt.tau_all, "cfc.tau_all")?;
upload_u8(&mut trunk.bucket_id_per_channel_d, &ckpt.bucket_id_per_channel, "bucket_id_per_channel")?;
upload_u32(&mut trunk.bucket_channel_offset_d, &ckpt.bucket_channel_offset, "bucket_channel_offset")?;
upload_u32(&mut trunk.bucket_dim_k_d, &ckpt.bucket_dim_k, "bucket_dim_k")?;
upload_u32(&mut trunk.heads_w_skip_offset_d, &ckpt.heads_w_skip_offset, "heads_w_skip_offset")?;
upload(&mut trunk.heads_w1_d, &ckpt.heads_w1, "heads_w1")?;
upload(&mut trunk.heads_b1_d, &ckpt.heads_b1, "heads_b1")?;
upload(&mut trunk.heads_w2_d, &ckpt.heads_w2, "heads_w2")?;
@@ -509,10 +577,27 @@ mod checkpoint_tests {
trunk.stream.memcpy_dtoh(slice, v.as_mut_slice()).unwrap();
v
};
let read_u8 = |trunk: &CfcTrunk, slice: &CudaSlice<u8>| -> Vec<u8> {
let mut v = vec![0u8; slice.len()];
trunk.stream.memcpy_dtoh(slice, v.as_mut_slice()).unwrap();
v
};
let read_u32 = |trunk: &CfcTrunk, slice: &CudaSlice<u32>| -> Vec<u32> {
let mut v = vec![0u32; slice.len()];
trunk.stream.memcpy_dtoh(slice, v.as_mut_slice()).unwrap();
v
};
assert_eq!(read(&trunk_a, &trunk_a.w_in_d), read(&trunk_b, &trunk_b.w_in_d));
assert_eq!(read(&trunk_a, &trunk_a.w_rec_d), read(&trunk_b, &trunk_b.w_rec_d));
assert_eq!(read(&trunk_a, &trunk_a.b_d), read(&trunk_b, &trunk_b.b_d));
assert_eq!(read(&trunk_a, &trunk_a.tau_d), read(&trunk_b, &trunk_b.tau_d));
assert_eq!(read(&trunk_a, &trunk_a.tau_all_d), read(&trunk_b, &trunk_b.tau_all_d));
// Bucket metadata: zero-initialised in new_random; round-trips must
// preserve them bit-identically (trivially zero == zero pre-transition,
// but the producer/consumer contract is verified here).
assert_eq!(read_u8(&trunk_a, &trunk_a.bucket_id_per_channel_d), read_u8(&trunk_b, &trunk_b.bucket_id_per_channel_d));
assert_eq!(read_u32(&trunk_a, &trunk_a.bucket_channel_offset_d), read_u32(&trunk_b, &trunk_b.bucket_channel_offset_d));
assert_eq!(read_u32(&trunk_a, &trunk_a.bucket_dim_k_d), read_u32(&trunk_b, &trunk_b.bucket_dim_k_d));
assert_eq!(read_u32(&trunk_a, &trunk_a.heads_w_skip_offset_d), read_u32(&trunk_b, &trunk_b.heads_w_skip_offset_d));
assert_eq!(read(&trunk_a, &trunk_a.vsn_w_d), read(&trunk_b, &trunk_b.vsn_w_d));
assert_eq!(read(&trunk_a, &trunk_a.ln_a_gain_d), read(&trunk_b, &trunk_b.ln_a_gain_d));
assert_eq!(read(&trunk_a, &trunk_a.attn_q_d), read(&trunk_b, &trunk_b.attn_q_d));

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

@@ -213,7 +213,7 @@ pub struct PerceptionTrainer {
// CfC + heads weights + their AdamWs (6 groups — tau is trained now).
// 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
// `self.trunk.w_in_d` / `w_rec_d` / `b_d` / `tau_all_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
@@ -840,7 +840,7 @@ impl PerceptionTrainer {
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")?;
stream.memcpy_htod(&tau, &mut trunk.tau_all_d).context("trunk.tau_all_d upload")?;
// 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")?;
@@ -2095,7 +2095,7 @@ impl PerceptionTrainer {
unsafe {
let mut launch = self.stream.launch_builder(&self.trunk.step_batched_fn);
launch
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_d)
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_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);
@@ -2345,7 +2345,7 @@ impl PerceptionTrainer {
unsafe {
let mut launch = self.stream.launch_builder(&self.step_bwd_batched_fn);
launch
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_d)
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_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)
@@ -2736,7 +2736,7 @@ impl PerceptionTrainer {
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_tau.step(&mut self.trunk.tau_all_d, &self.grad_tau_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)?;
@@ -3205,7 +3205,7 @@ impl PerceptionTrainer {
unsafe {
let mut launch = self.stream.launch_builder(&self.trunk.step_batched_fn);
launch
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_d)
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_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);
@@ -3480,7 +3480,7 @@ impl PerceptionTrainer {
unsafe {
let mut launch = self.stream.launch_builder(&self.trunk.step_batched_fn);
launch
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_d)
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_d)
.arg(&self.ln_b_step_out_d).arg(&self.cfc_h_state_step_d)
.arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i)
.arg(&mut self.cfc_h_new_step_d);
@@ -4118,7 +4118,7 @@ impl PerceptionTrainer {
unsafe {
let mut launch = self.stream.launch_builder(&self.trunk.step_batched_fn);
launch
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_d)
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_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);