From e18325aaf1cd11ba8a98f0edcfa4099f34a56aca Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 21 May 2026 23:48:15 +0200 Subject: [PATCH] Revert "fix(per-horizon-cfc): extend block-diagonal mask to heads_w1 (Smoke 2 fix)" This reverts commit 4c1ed1d953087c0559d006a03e0b1871f55120a0. --- .../cuda/bucket_transition_kernels.cu | 104 --------- crates/ml-alpha/src/trainer/perception.rs | 192 ---------------- .../tests/bucket_transition_kernels.rs | 210 ------------------ 3 files changed, 506 deletions(-) diff --git a/crates/ml-alpha/cuda/bucket_transition_kernels.cu b/crates/ml-alpha/cuda/bucket_transition_kernels.cu index 4facb661a..bb68ff602 100644 --- a/crates/ml-alpha/cuda/bucket_transition_kernels.cu +++ b/crates/ml-alpha/cuda/bucket_transition_kernels.cu @@ -5,7 +5,6 @@ // boundaries for HIDDEN_DIM=128: [0, 25, 50, 75, 100, 128]. #define HIDDEN_DIM 128 -#define HEAD_MID_DIM 64 #define N_HORIZONS 5 #define BUCKET_DIM_LAST 28 // last bucket absorbs HIDDEN_DIM - 4*25 = 28 #define MAX_BUCKET_DIM 28 @@ -533,106 +532,3 @@ extern "C" __global__ void bucket_iqr_double_widening_kernel( } } } - -// ───────────────────────────────────────────────────────────────────── -// heads_w1_mask_init_kernel — Phase 1→2 transition one-shot for -// `heads_w1` (GRN input projection W1: HIDDEN → HEAD_MID per horizon). -// -// Diagnostic smoke 2 (`lob-backtest-sweep-kw7f6` at d2d34b6d4) confirmed -// that with only `heads_w_skip` block-diagonalised, `heads_w1` remained -// 100% dense at off-bucket positions (32,768/32,768 nonzero) and the -// per-horizon logit distributions collapsed to ~uniform (mean spread -// 0.05). `heads_w1` is the SOLE remaining HIDDEN_DIM-reading path in the -// GRN forward (`multi_horizon_heads_grn_fwd_batched` Pass 1: line 497 of -// `multi_horizon_heads.cu`) that bypasses bucket routing — every other -// HIDDEN_DIM-reading weight (`heads_w_skip`) is already restricted, and -// `heads_w2`, `heads_w_gate`, `heads_w_main` only see HEAD_MID-dim -// tensors. Mirrors the `heads_w_skip` mask init: build the routing mask -// AND zero-multiply off-bucket entries of `heads_w1` in place. -// -// Layout: `heads_w1` is `[N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM]` -// row-major (`heads_w1[h, m, c]` flat-index = h*HEAD_MID_DIM*HIDDEN_DIM -// + m*HIDDEN_DIM + c). The mask depends only on (h, c) — the m-axis is -// pure replication. Position `(h, m, c)` is in-bucket iff -// `bucket_id_per_channel[c] == h`. -// -// Launch: grid = (N_HORIZONS, HEAD_MID_DIM, 1), block = (HIDDEN_DIM, 1, 1). -// One block per (horizon, mid-row) × one thread per channel. No -// reductions, no shared memory, no atomicAdd — fully data-parallel. -// ───────────────────────────────────────────────────────────────────── -extern "C" __global__ void heads_w1_mask_init_kernel( - const unsigned char* __restrict__ bucket_id_per_channel, // [HIDDEN_DIM] - float* __restrict__ heads_w1, // [N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM] - float* __restrict__ mask // [N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM] -) { - int h = blockIdx.x; - int m = blockIdx.y; - int c = threadIdx.x; - if (h >= N_HORIZONS || m >= HEAD_MID_DIM || c >= HIDDEN_DIM) return; - long long idx = ((long long)h * HEAD_MID_DIM + m) * HIDDEN_DIM + c; - float msk = (bucket_id_per_channel[c] == (unsigned char)h) ? 1.0f : 0.0f; - mask[idx] = msk; - heads_w1[idx] *= msk; // zero off-bucket entries in place -} - -// ───────────────────────────────────────────────────────────────────── -// heads_w1_grad_mask_apply_kernel — per-step Phase 2 gradient mask for -// `heads_w1`. -// -// Mirrors `heads_w_skip_grad_mask_apply_kernel`: multiply the post- -// reduction `grad_heads_w1` element-wise by the routing mask BEFORE the -// Adam step. Off-bucket positions receive zero gradient, never update, -// and (since they started at zero after the init kernel) remain at zero -// throughout training. This delivers block-diagonal `heads_w1` without -// touching the GRN forward/backward kernels. -// -// Launch: grid = (N_HORIZONS, HEAD_MID_DIM, 1), block = (HIDDEN_DIM, 1, 1). -// Same layout as `heads_w1_mask_init_kernel`. -// ───────────────────────────────────────────────────────────────────── -extern "C" __global__ void heads_w1_grad_mask_apply_kernel( - const float* __restrict__ mask, // [N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM] - float* __restrict__ grad_heads_w1 // [N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM] -) { - int h = blockIdx.x; - int m = blockIdx.y; - int c = threadIdx.x; - if (h >= N_HORIZONS || m >= HEAD_MID_DIM || c >= HIDDEN_DIM) return; - long long idx = ((long long)h * HEAD_MID_DIM + m) * HIDDEN_DIM + c; - grad_heads_w1[idx] *= mask[idx]; -} - -// ───────────────────────────────────────────────────────────────────── -// heads_w1_zero_off_bucket_kernel — block-diagonal projection on -// `[N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM]` tensors (params + Adam -// moments). -// -// Same role as `zero_off_bucket_kernel` but for the 3-D heads_w1 shape. -// Launched on: -// 1. `heads_w1_d` at transition (one-shot, zeros off-bucket weights; -// redundant with `heads_w1_mask_init_kernel` but explicit for -// traceability) -// 2. `heads_w1` Adam `m` and `v` buffers at transition (one-shot; zeros -// pre-transition momentum so post-transition Adam steps cannot -// revive off-bucket positions) -// 3. `heads_w1_d` after every Phase 2 Adam step (belt-and-suspenders -// against ε-driven drift from zero moments) -// -// Layout matches `heads_w1_mask_init_kernel`: `tensor[h, m, c]` flat- -// index = `h*HEAD_MID_DIM*HIDDEN_DIM + m*HIDDEN_DIM + c`. Off-bucket iff -// `bucket_id_per_channel[c] != h`. -// -// Launch: grid = (N_HORIZONS, HEAD_MID_DIM, 1), block = (HIDDEN_DIM, 1, 1). -// ───────────────────────────────────────────────────────────────────── -extern "C" __global__ void heads_w1_zero_off_bucket_kernel( - const unsigned char* __restrict__ bucket_id_per_channel, // [HIDDEN_DIM] - float* __restrict__ tensor // [N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM] -) { - int h = blockIdx.x; - int m = blockIdx.y; - int c = threadIdx.x; - if (h >= N_HORIZONS || m >= HEAD_MID_DIM || c >= HIDDEN_DIM) return; - long long idx = ((long long)h * HEAD_MID_DIM + m) * HIDDEN_DIM + c; - if (bucket_id_per_channel[c] != (unsigned char)h) { - tensor[idx] = 0.0f; - } -} diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index b58c35c5d..6ced8e07e 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -792,31 +792,6 @@ pub struct PerceptionTrainer { /// exactly zero even under Adam's ε-driven step from zero /// moments). zero_off_bucket_fn: CudaFunction, - /// Per-channel routing mask for `heads_w1` (GRN-wide block-diagonal - /// extension; Smoke 2 fix 2026-05-21). Built once at Phase 1→2 - /// transition by `heads_w1_mask_init_kernel`: - /// `mask[h, m, c] = 1.0` iff `bucket_id_per_channel[c] == h`, else 0.0. - /// At the same launch the kernel zero-multiplies off-bucket entries of - /// `heads_w1_d` so the GRN forward sees only block-diagonal weights on - /// the HIDDEN-reading W1 path (Pass 1 of `multi_horizon_heads_grn_fwd_ - /// batched`). Each Phase 2 backward step - /// `heads_w1_grad_mask_apply_kernel` multiplies `grad_heads_w1_d` by - /// this mask before Adam — off-bucket entries stay at 0 throughout - /// training. Shape `[N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM]`. - heads_w1_mask_d: CudaSlice, - /// Cached function handle for `heads_w1_mask_init_kernel` — one-shot - /// at Phase 1→2 transition (builds mask + zeros off-bucket entries of - /// `heads_w1_d` in place). - heads_w1_mask_init_fn: CudaFunction, - /// Cached function handle for `heads_w1_grad_mask_apply_kernel` — - /// per-step Phase 2 (multiplies `grad_heads_w1_d` by the routing mask - /// before Adam, enforcing block-diagonal updates on the W1 path). - heads_w1_grad_mask_apply_fn: CudaFunction, - /// Cached function handle for `heads_w1_zero_off_bucket_kernel` — - /// 3-D analogue of `zero_off_bucket_kernel`, used on - /// `[N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM]` tensors (heads_w1 - /// params + Adam m/v moments) at transition and post-Adam in Phase 2. - heads_w1_zero_off_bucket_fn: CudaFunction, /// Per-horizon LR multiplier vector consumed by /// `heads_lr_multiplier_scale_kernel`. Mapped-pinned per /// `feedback_no_htod_htoh_only_mapped_pinned`: host computes @@ -1097,15 +1072,6 @@ impl PerceptionTrainer { let zero_off_bucket_fn = bucket_transition_module .load_function("zero_off_bucket_kernel") .context("load zero_off_bucket_kernel")?; - let heads_w1_mask_init_fn = bucket_transition_module - .load_function("heads_w1_mask_init_kernel") - .context("load heads_w1_mask_init_kernel")?; - let heads_w1_grad_mask_apply_fn = bucket_transition_module - .load_function("heads_w1_grad_mask_apply_kernel") - .context("load heads_w1_grad_mask_apply_kernel")?; - let heads_w1_zero_off_bucket_fn = bucket_transition_module - .load_function("heads_w1_zero_off_bucket_kernel") - .context("load heads_w1_zero_off_bucket_kernel")?; // Mamba2 stacks live on the trunk from new_random; we wire optimizer // + training scratches against the trunk-owned blocks. @@ -1495,16 +1461,6 @@ impl PerceptionTrainer { let heads_w_skip_mask_d = stream .alloc_zeros::(N_HORIZONS * crate::cfc::bucket_routing::HIDDEN_DIM) .context("heads_w_skip_mask_d alloc")?; - // GRN-wide block-diagonal: heads_w1 mask (Smoke 2 fix 2026-05-21). - // Same lifecycle as `heads_w_skip_mask_d`: zero-allocated here, - // populated at Phase 1→2 transition by `heads_w1_mask_init_kernel`, - // read each Phase 2 step by `heads_w1_grad_mask_apply_kernel`. - // Shape `[N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM]` = 40,960 floats. - let heads_w1_mask_d = stream - .alloc_zeros::( - N_HORIZONS * HEAD_MID_DIM * crate::cfc::bucket_routing::HIDDEN_DIM, - ) - .context("heads_w1_mask_d alloc")?; // Per-horizon LR multiplier staging. Mapped-pinned so the host // writes from `target_jitter_per_h` + jitter EMA each Phase 2 step // and the kernel reads the device-visible pointer without DtoH. @@ -1904,10 +1860,6 @@ impl PerceptionTrainer { heads_w_skip_mask_init_fn, heads_w_skip_grad_mask_apply_fn, zero_off_bucket_fn, - heads_w1_mask_d, - heads_w1_mask_init_fn, - heads_w1_grad_mask_apply_fn, - heads_w1_zero_off_bucket_fn, lr_mult_per_horizon_staging, smoothness_jitter_ema_host_d, target_jitter_per_h: [0.0; N_HORIZONS], @@ -2433,94 +2385,6 @@ impl PerceptionTrainer { } } - // 7c. GRN-wide block-diagonal extension (Smoke 2 fix - // 2026-05-21). Diagnostic - // `lob-backtest-sweep-kw7f6` at d2d34b6d4 confirmed - // `heads_w1` is the sole remaining HIDDEN_DIM-reading - // path in the GRN forward (`multi_horizon_heads_grn_ - // fwd_batched` Pass 1) that bypasses bucket routing, - // and that it stayed 100% dense (32,768/32,768 nonzero - // at off-bucket positions) when only `heads_w_skip` - // was restricted — collapsing per-horizon logits to - // ~uniform (mean spread 0.05, mean_run_len ratio - // 1.04×). Mirror the heads_w_skip defense exactly: - // mask init (zero off-bucket params + build mask), - // Adam moment projection (m, v), grad mask per step, - // post-Adam projection per step (belt+suspenders). - { - let metadata_ref = self - .bucket_routing_metadata - .as_ref() - .expect("metadata persisted in step 6"); - let cfg_mask_init_w1 = LaunchConfig { - grid_dim: (N_HORIZONS as u32, HEAD_MID_DIM as u32, 1), - block_dim: (HIDDEN_DIM as u32, 1, 1), - shared_mem_bytes: 0, - }; - // (i) Build mask + zero off-bucket params of heads_w1 - // in place. - { - let mut launch = self - .stream - .launch_builder(&self.heads_w1_mask_init_fn); - launch - .arg(&metadata_ref.bucket_id_per_channel_d) - .arg(&mut self.trunk.heads_w1_d) - .arg(&mut self.heads_w1_mask_d); - unsafe { - launch - .launch(cfg_mask_init_w1) - .context("heads_w1_mask_init_kernel launch (Phase 1→2 block-diagonal heads_w1)")?; - } - } - // (ii) Traceability re-zero of heads_w1_d (redundant - // with the mask_init zero above; same role as the - // heads_w_skip step 7b.(a) launch). - { - let mut launch = self - .stream - .launch_builder(&self.heads_w1_zero_off_bucket_fn); - launch - .arg(&metadata_ref.bucket_id_per_channel_d) - .arg(&mut self.trunk.heads_w1_d); - unsafe { - launch - .launch(cfg_mask_init_w1) - .context("heads_w1_zero_off_bucket_kernel launch (heads_w1_d at transition)")?; - } - } - // (iii) Project Adam's first-moment buffer `m` for - // heads_w1. - { - let mut launch = self - .stream - .launch_builder(&self.heads_w1_zero_off_bucket_fn); - launch - .arg(&metadata_ref.bucket_id_per_channel_d) - .arg(self.opt_heads_w1.m_mut()); - unsafe { - launch - .launch(cfg_mask_init_w1) - .context("heads_w1_zero_off_bucket_kernel launch (heads_w1 Adam m at transition)")?; - } - } - // (iv) Project Adam's second-moment buffer `v` for - // heads_w1. - { - let mut launch = self - .stream - .launch_builder(&self.heads_w1_zero_off_bucket_fn); - launch - .arg(&metadata_ref.bucket_id_per_channel_d) - .arg(self.opt_heads_w1.v_mut()); - unsafe { - launch - .launch(cfg_mask_init_w1) - .context("heads_w1_zero_off_bucket_kernel launch (heads_w1 Adam v at transition)")?; - } - } - } - // 8. Latch trainer into Phase 2. // // ALPHA fix 2026-05-21: τ stays in ORIGINAL channel @@ -4101,63 +3965,7 @@ impl PerceptionTrainer { } } } - // ── Block-diagonal heads_w1: grad mask apply (Phase 2 only) ── - // GRN-wide block-diagonal extension (Smoke 2 fix 2026-05-21). - // Mirrors the heads_w_skip grad mask below: multiplies - // `grad_heads_w1_d` element-wise by the routing mask built at - // transition. Off-bucket entries see zero gradient → Adam never - // updates them → they stay at 0 forever (init zero-multiplied - // them at the transition). Realises block-diagonal heads_w1 on - // the GRN's Pass-1 HIDDEN-reading path without touching - // `multi_horizon_heads_grn_fwd_batched` itself. MUST run BEFORE - // `opt_heads_w1.step` so the masked grad is the one Adam reads. - if self.phase == TrainingPhase::Phase2Routed { - let cfg_grad_mask_w1 = LaunchConfig { - grid_dim: (N_HORIZONS as u32, HEAD_MID_DIM as u32, 1), - block_dim: (HIDDEN_DIM as u32, 1, 1), - shared_mem_bytes: 0, - }; - let mut launch = self - .stream - .launch_builder(&self.heads_w1_grad_mask_apply_fn); - launch - .arg(&self.heads_w1_mask_d) - .arg(&mut self.grad_heads_w1_d); - unsafe { - launch - .launch(cfg_grad_mask_w1) - .context("heads_w1_grad_mask_apply_kernel launch (block-diagonal heads_w1)")?; - } - } self.opt_heads_w1.step(&mut self.trunk.heads_w1_d, &self.grad_heads_w1_d)?; - // ── Block-diagonal heads_w1: post-Adam projection (Phase 2 only) ── - // Belt-and-suspenders re-zero of off-bucket `heads_w1_d` positions - // after every Phase 2 Adam step. Same role as the heads_w_skip - // post-Adam zero below: even with the grad mask + Adam-moment - // mask at transition, mixed-precision rounding and finite-precision - // moment buffers can drift; we project hard per - // `feedback_no_quickfixes`. Effectively realises - // `heads_w1[h, m, c] = 0 for c not in bucket h`. - if self.phase == TrainingPhase::Phase2Routed { - if let Some(metadata) = self.bucket_routing_metadata.as_ref() { - let cfg_zero_w1 = LaunchConfig { - grid_dim: (N_HORIZONS as u32, HEAD_MID_DIM as u32, 1), - block_dim: (HIDDEN_DIM as u32, 1, 1), - shared_mem_bytes: 0, - }; - let mut launch = self - .stream - .launch_builder(&self.heads_w1_zero_off_bucket_fn); - launch - .arg(&metadata.bucket_id_per_channel_d) - .arg(&mut self.trunk.heads_w1_d); - unsafe { - launch - .launch(cfg_zero_w1) - .context("heads_w1_zero_off_bucket_kernel launch (post-Adam block-diagonal projection)")?; - } - } - } 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)?; diff --git a/crates/ml-alpha/tests/bucket_transition_kernels.rs b/crates/ml-alpha/tests/bucket_transition_kernels.rs index e38107b25..19f3b4800 100644 --- a/crates/ml-alpha/tests/bucket_transition_kernels.rs +++ b/crates/ml-alpha/tests/bucket_transition_kernels.rs @@ -873,213 +873,3 @@ fn slack_factor_apply_kernel_widens_iqr_geometrically() -> Result<()> { } Ok(()) } - -// ───────────────────────────────────────────────────────────────────── -// heads_w1 block-diagonal kernels — GRN-wide extension (Smoke 2 fix -// 2026-05-21). heads_w1 has shape [N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM]. -// The mask depends only on (h, c) — replicated across the m axis. -// ───────────────────────────────────────────────────────────────────── - -const HEAD_MID_DIM: usize = 64; - -#[test] -#[ignore = "requires CUDA"] -fn heads_w1_mask_init_zeros_off_bucket() -> Result<()> { - let dev = MlDevice::cuda(0)?; - let stream = dev.cuda_stream()?.clone(); - let func = load_kernel(&stream, "heads_w1_mask_init_kernel")?; - - // Bucket assignment: channel c → bucket c/25 capped at 4 (same - // quintile layout used by the heads_w_skip mask init test above). - // heads_w1 starts all 1.0 across [N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM]; - // result is 1.0 at in-bucket positions, 0.0 elsewhere. - let bucket_id_host: Vec = - (0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect(); - let total = N_HORIZONS * HEAD_MID_DIM * HIDDEN_DIM; - let heads_w1_host: Vec = vec![1.0; total]; - - let bucket_id_d = upload(&stream, &bucket_id_host)?; - let mut heads_w1_d = upload(&stream, &heads_w1_host)?; - let mut mask_d = stream.alloc_zeros::(total)?; - - let cfg = LaunchConfig { - grid_dim: (N_HORIZONS as u32, HEAD_MID_DIM as u32, 1), - block_dim: (HIDDEN_DIM as u32, 1, 1), - shared_mem_bytes: 0, - }; - let mut launch = stream.launch_builder(&func); - launch - .arg(&bucket_id_d) - .arg(&mut heads_w1_d) - .arg(&mut mask_d); - unsafe { - launch.launch(cfg)?; - } - stream.synchronize()?; - - let mask = download(&stream, &mask_d)?; - let heads_w1_out = download(&stream, &heads_w1_d)?; - for h in 0..N_HORIZONS { - for m in 0..HEAD_MID_DIM { - for c in 0..HIDDEN_DIM { - let idx = (h * HEAD_MID_DIM + m) * HIDDEN_DIM + c; - let in_bucket = (c / 25).min(4) == h; - let expected = if in_bucket { 1.0 } else { 0.0 }; - assert!( - (mask[idx] - expected).abs() < 1e-6, - "mask[{},{},{}]={} expected {}", - h, m, c, mask[idx], expected - ); - assert!( - (heads_w1_out[idx] - expected).abs() < 1e-6, - "heads_w1[{},{},{}]={} expected {}", - h, m, c, heads_w1_out[idx], expected - ); - } - } - } - Ok(()) -} - -#[test] -#[ignore = "requires CUDA"] -fn heads_w1_grad_mask_apply_zeros_off_bucket_grad() -> Result<()> { - let dev = MlDevice::cuda(0)?; - let stream = dev.cuda_stream()?.clone(); - let func = load_kernel(&stream, "heads_w1_grad_mask_apply_kernel")?; - - // Build a mask manually for the same quintile layout; feed an all- - // ones grad tensor and expect off-bucket entries to be zeroed across - // the full [N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM] shape. - let total = N_HORIZONS * HEAD_MID_DIM * HIDDEN_DIM; - let mut mask_host: Vec = vec![0.0; total]; - for h in 0..N_HORIZONS { - for m in 0..HEAD_MID_DIM { - for c in 0..HIDDEN_DIM { - if (c / 25).min(4) == h { - mask_host[(h * HEAD_MID_DIM + m) * HIDDEN_DIM + c] = 1.0; - } - } - } - } - let grad_host: Vec = vec![1.0; total]; - - let mask_d = upload(&stream, &mask_host)?; - let mut grad_d = upload(&stream, &grad_host)?; - - let cfg = LaunchConfig { - grid_dim: (N_HORIZONS as u32, HEAD_MID_DIM as u32, 1), - block_dim: (HIDDEN_DIM as u32, 1, 1), - shared_mem_bytes: 0, - }; - let mut launch = stream.launch_builder(&func); - launch.arg(&mask_d).arg(&mut grad_d); - unsafe { - launch.launch(cfg)?; - } - stream.synchronize()?; - - let grad_out = download(&stream, &grad_d)?; - for h in 0..N_HORIZONS { - for m in 0..HEAD_MID_DIM { - for c in 0..HIDDEN_DIM { - let idx = (h * HEAD_MID_DIM + m) * HIDDEN_DIM + c; - let in_bucket = (c / 25).min(4) == h; - let expected = if in_bucket { 1.0 } else { 0.0 }; - assert!( - (grad_out[idx] - expected).abs() < 1e-6, - "grad[{},{},{}]={} expected {}", - h, m, c, grad_out[idx], expected - ); - } - } - } - Ok(()) -} - -#[test] -#[ignore = "requires CUDA"] -fn heads_w1_zero_off_bucket_keeps_off_bucket_zero_after_many_mock_adam_steps() -> Result<()> { - // ALPHA-style invariant test for the 3-D heads_w1 projection. Start - // with heads_w1 = all-ones at in-bucket positions, zero elsewhere - // (post-mask_init state). Run N "Adam steps" of the form - // heads_w1 += eps (uniform tiny push that would breach the - // block-diagonal invariant without projection) - // and verify that after each post-step zero call, off-bucket entries - // return to exactly 0.0 while in-bucket entries accumulate the push. - let dev = MlDevice::cuda(0)?; - let stream = dev.cuda_stream()?.clone(); - let func_zero = load_kernel(&stream, "heads_w1_zero_off_bucket_kernel")?; - let func_mask_init = load_kernel(&stream, "heads_w1_mask_init_kernel")?; - - let total = N_HORIZONS * HEAD_MID_DIM * HIDDEN_DIM; - let bucket_id_host: Vec = - (0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect(); - let bucket_id_d = upload(&stream, &bucket_id_host)?; - let mut heads_w1_d = upload(&stream, &vec![1.0_f32; total])?; - let mut mask_d = stream.alloc_zeros::(total)?; - - let cfg_w1 = LaunchConfig { - grid_dim: (N_HORIZONS as u32, HEAD_MID_DIM as u32, 1), - block_dim: (HIDDEN_DIM as u32, 1, 1), - shared_mem_bytes: 0, - }; - // Initialise: zero off-bucket weights + populate mask. - { - let mut launch = stream.launch_builder(&func_mask_init); - launch - .arg(&bucket_id_d) - .arg(&mut heads_w1_d) - .arg(&mut mask_d); - unsafe { - launch.launch(cfg_w1)?; - } - } - stream.synchronize()?; - - // 10 mock steps (keep runtime modest given 40,960-element tensor) of - // uniform-push + projection. - let n_steps = 10; - let eps_push: f32 = 1e-3; - for step in 0..n_steps { - let mut current = download(&stream, &heads_w1_d)?; - for i in 0..current.len() { - current[i] += eps_push; - } - stream.memcpy_htod(¤t, &mut heads_w1_d)?; - let mut launch = stream.launch_builder(&func_zero); - launch.arg(&bucket_id_d).arg(&mut heads_w1_d); - unsafe { - launch.launch(cfg_w1)?; - } - stream.synchronize()?; - let post = download(&stream, &heads_w1_d)?; - for h in 0..N_HORIZONS { - // Sample a few m rows to keep test runtime bounded. - for m in [0usize, HEAD_MID_DIM / 2, HEAD_MID_DIM - 1] { - for c in 0..HIDDEN_DIM { - let idx = (h * HEAD_MID_DIM + m) * HIDDEN_DIM + c; - let in_bucket = (c / 25).min(4) == h; - if !in_bucket { - assert_eq!( - post[idx], 0.0, - "off-bucket [{}, {}, {}] = {} after step {} — projection failed", - h, m, c, post[idx], step - ); - } else { - // In-bucket positions accumulate the push: - // starts at 1.0; after (step + 1) pushes, - // value = 1.0 + (step+1) * eps_push. - let expected = 1.0_f32 + ((step + 1) as f32) * eps_push; - assert!( - (post[idx] - expected).abs() < 1e-5, - "in-bucket [{}, {}, {}] = {} after step {} expected {}", - h, m, c, post[idx], step, expected - ); - } - } - } - } - } - Ok(()) -}