plan(ml-alpha): v2 multi-horizon implementation plan (V1-V13)
Concrete TDD-driven commit map for the v2 spec (2026-05-18-ml-alpha-v2-multi-horizon-design.md). Thirteen commits ordered by dependency: delete falsified path, build new kernels with numgrad parity (V2-V8), trainer state bundle (V9), wire into PerceptionTrainer (V10), local smoke (V11), cluster smoke (V12), 3-fold A/B (V13). Per-commit verification gates and explicit stop conditions. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
535
docs/superpowers/plans/2026-05-18-ml-alpha-v2-multi-horizon.md
Normal file
535
docs/superpowers/plans/2026-05-18-ml-alpha-v2-multi-horizon.md
Normal file
@@ -0,0 +1,535 @@
|
||||
# ml-alpha v2 Multi-Horizon Implementation Plan
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||
|
||||
**Goal:** Execute the v2 redesign in `docs/superpowers/specs/2026-05-18-ml-alpha-v2-multi-horizon-design.md` (A+B+C+D+E) in a sequence of validated commits ending with a cluster A/B vs the task #200 baseline.
|
||||
|
||||
**Architecture:** Replace per-horizon Q_h pool (C21/C22/C25) with horizon-token K-prepend + inverted attention + regime-MoE + Kendall σ-BCE + L2 anchor. All new kernels are GPU-only, capture-safe, warp-shuffle reduce where applicable.
|
||||
|
||||
**Tech Stack:** CUDA 12.4, cudarc 0.19, sm_86 (local) / sm_89 (L40S cluster), Rust 1.85.
|
||||
|
||||
---
|
||||
|
||||
## Commit map
|
||||
|
||||
| Commit | Subject | Files |
|
||||
|---|---|---|
|
||||
| V1 | Delete old per-horizon Q_h path | `crates/ml-alpha/cuda/per_horizon_*.cu`, `src/per_horizon_*.rs`, `src/trainer/per_horizon_state.rs`, `build.rs`, `trainer/mod.rs`, perception.rs callsites |
|
||||
| V2 | (C) horizon_token_attention_pool kernel + binding + numgrad | `cuda/horizon_token_attention_pool.cu`, `src/horizon_token_attention_pool.rs`, `tests/horizon_token_attention_pool_numgrad.rs`, `build.rs` |
|
||||
| V3 | (E) inverted_attention_pool kernel + binding + numgrad | `cuda/inverted_attention_pool.cu`, `src/inverted_attention_pool.rs`, `tests/inverted_attention_pool_numgrad.rs`, `build.rs` |
|
||||
| V4 | Fused projection kernel + binding (combines outputs from V2 + V3) | `cuda/fused_ctx_proj.cu`, `src/fused_ctx_proj.rs`, `tests/fused_ctx_proj_numgrad.rs`, `build.rs` |
|
||||
| V5 | (D) regime feature extraction wired into snap_feature_assemble | `cuda/snap_feature_assemble.cu`, `src/data/loader.rs`, `tests/regime_features_smoke.rs` |
|
||||
| V6 | (D) regime_moe_gate kernel + binding + numgrad | `cuda/regime_moe_gate.cu`, `src/regime_moe_gate.rs`, `tests/regime_moe_gate_numgrad.rs`, `build.rs` |
|
||||
| V7 | (A) modify bce_loss_multi_horizon → sigma-weighted | `cuda/bce_loss_multi_horizon.cu`, `tests/bce_sigma_numgrad.rs` |
|
||||
| V8 | (B) anchor_l2 kernel + Wiener-α lambda controller + binding | `cuda/anchor_l2.cu`, `src/anchor_l2.rs`, `src/trainer/anchor_controller.rs`, `tests/anchor_l2_numgrad.rs`, `build.rs` |
|
||||
| V9 | TrainerState v2 — bundle horizon tokens, MoE experts, log_sigma, anchor buffers, all AdamWs | `src/trainer/perception_v2_state.rs` |
|
||||
| V10 | PerceptionTrainer wiring — forward + backward + AdamW for v2 path | `src/trainer/perception.rs` (deep modification) |
|
||||
| V11 | Local end-to-end smoke + bit-identity invariant at init | `tests/v2_pipeline_smoke.rs` |
|
||||
| V12 | Cluster smoke @ 1 epoch on the new commit | `(submit via argo-alpha-perception.sh)` |
|
||||
| V13 | Cluster A/B 3-fold @ 30 epochs | `(submit 3 workflows)` |
|
||||
|
||||
---
|
||||
|
||||
## Task V1: Delete old per-horizon Q_h path
|
||||
|
||||
**Files:**
|
||||
- Delete: `crates/ml-alpha/cuda/per_horizon_attention_pool.cu`
|
||||
- Delete: `crates/ml-alpha/cuda/per_horizon_residual_head.cu`
|
||||
- Delete: `crates/ml-alpha/cuda/per_horizon_prob_blend.cu`
|
||||
- Delete: `crates/ml-alpha/src/per_horizon_attention_pool.rs`
|
||||
- Delete: `crates/ml-alpha/src/per_horizon_residual_head.rs`
|
||||
- Delete: `crates/ml-alpha/src/trainer/per_horizon_state.rs`
|
||||
- Delete: `crates/ml-alpha/tests/per_horizon_attention_pool_numgrad.rs`
|
||||
- Delete: `crates/ml-alpha/tests/per_horizon_residual_head_numgrad.rs`
|
||||
- Delete: `crates/ml-alpha/tests/per_horizon_full_pipeline_smoke.rs`
|
||||
- Modify: `crates/ml-alpha/build.rs` — drop 3 entries from KERNELS
|
||||
- Modify: `crates/ml-alpha/src/lib.rs` — drop `pub mod per_horizon_*`
|
||||
- Modify: `crates/ml-alpha/src/trainer/mod.rs` — drop `pub mod per_horizon_state`
|
||||
- Modify: `crates/ml-alpha/src/trainer/perception.rs` — drop all per_horizon. references in step_batched fwd + bwd + AdamW sections; revert to pre-C24 baseline behavior
|
||||
|
||||
- [ ] **Step 1: Identify all callsites in perception.rs**
|
||||
|
||||
Run: `grep -n "per_horizon" crates/ml-alpha/src/trainer/perception.rs`
|
||||
Expected: 3 sections (section 4.5 fwd, section 5a bwd, section 9 AdamW)
|
||||
|
||||
- [ ] **Step 2: Remove the 3 sections** keeping baseline BCE path intact
|
||||
|
||||
- [ ] **Step 3: Remove module imports + struct field**
|
||||
|
||||
Edit `perception.rs`:
|
||||
- Remove `pub per_horizon: PerHorizonTrainState` from struct
|
||||
- Remove its initialization in `new()`
|
||||
- Remove `use crate::trainer::per_horizon_state::PerHorizonTrainState;`
|
||||
|
||||
- [ ] **Step 4: Delete files + update build.rs/mod.rs**
|
||||
|
||||
```bash
|
||||
rm crates/ml-alpha/cuda/per_horizon_attention_pool.cu \
|
||||
crates/ml-alpha/cuda/per_horizon_residual_head.cu \
|
||||
crates/ml-alpha/cuda/per_horizon_prob_blend.cu \
|
||||
crates/ml-alpha/src/per_horizon_attention_pool.rs \
|
||||
crates/ml-alpha/src/per_horizon_residual_head.rs \
|
||||
crates/ml-alpha/src/trainer/per_horizon_state.rs \
|
||||
crates/ml-alpha/tests/per_horizon_attention_pool_numgrad.rs \
|
||||
crates/ml-alpha/tests/per_horizon_residual_head_numgrad.rs \
|
||||
crates/ml-alpha/tests/per_horizon_full_pipeline_smoke.rs
|
||||
```
|
||||
|
||||
Edit `build.rs` KERNELS — remove `per_horizon_attention_pool`, `per_horizon_residual_head`, `per_horizon_prob_blend`.
|
||||
|
||||
- [ ] **Step 5: Verify clean workspace check**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml-alpha`
|
||||
Expected: clean compile, no warnings about removed modules.
|
||||
|
||||
- [ ] **Step 6: Run baseline smoke to confirm regression to v0 behavior**
|
||||
|
||||
Run: `CUDA_COMPUTE_CAP=86 SQLX_OFFLINE=true cargo test -p ml-alpha --features cuda --test mamba2_cfc_smoke -- --ignored`
|
||||
Expected: passes (this is the pre-v1 smoke test that doesn't touch per-horizon path).
|
||||
|
||||
- [ ] **Step 7: Commit V1**
|
||||
|
||||
```bash
|
||||
git add -A && git commit -m "refactor(ml-alpha): remove per-horizon Q_h path (C21-C25 falsified)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task V2: (C) horizon_token_attention_pool — kernel + binding + numgrad
|
||||
|
||||
**Files:**
|
||||
- Create: `crates/ml-alpha/cuda/horizon_token_attention_pool.cu`
|
||||
- Create: `crates/ml-alpha/src/horizon_token_attention_pool.rs`
|
||||
- Create: `crates/ml-alpha/tests/horizon_token_attention_pool_numgrad.rs`
|
||||
- Modify: `crates/ml-alpha/build.rs`, `crates/ml-alpha/src/lib.rs`
|
||||
|
||||
**Math (forward):**
|
||||
```
|
||||
extended[b, i, d] = (i < N_H) ? horizon_tokens[i, d] : LN_b[b, i - N_H, d] # [B, N_H + K, H]
|
||||
scores[b, i] = Σ_d Q[d] · extended[b, i, d] # [B, N_H + K]
|
||||
attn[b, i] = softmax_i(scores[b, i]) # [B, N_H + K]
|
||||
ctx_h[b, h, d] = attn[b, h] · horizon_tokens[h, d]
|
||||
+ Σ_k attn[b, N_H + k] · LN_b[b, k, d] # [B, N_H, H]
|
||||
```
|
||||
|
||||
The kernel logic is identical to a standard single-Q attention pool with `K' = N_HORIZONS + K`. The per-horizon ctx is recovered by splicing each horizon token's contribution back into the per-horizon output slot.
|
||||
|
||||
- [ ] **Step 1: Write failing numgrad test**
|
||||
|
||||
```rust
|
||||
// tests/horizon_token_attention_pool_numgrad.rs
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn forward_then_backward_matches_central_difference() {
|
||||
// B=2, K=8, H=128, N_H=5. Init horizon_tokens + Q + ln_b with seeded random.
|
||||
// Run forward → ctx_h, save attn weights. Compute analytic gradients via
|
||||
// backward. Compare per-parameter to 2e-3 central-difference numgrad
|
||||
// within 5e-2 rel-tol / 5e-3 abs-floor.
|
||||
}
|
||||
```
|
||||
|
||||
Run: `cargo test -p ml-alpha --features cuda --test horizon_token_attention_pool_numgrad -- --ignored`
|
||||
Expected: FAIL (kernel doesn't exist).
|
||||
|
||||
- [ ] **Step 2: Write the .cu kernel**
|
||||
|
||||
```cuda
|
||||
// horizon_token_attention_pool.cu
|
||||
#define HTAP_HIDDEN_DIM 128
|
||||
#define HTAP_BLOCK 128
|
||||
#define HTAP_N_HORIZONS 5
|
||||
#define HTAP_N_WARPS 4
|
||||
|
||||
extern "C" __global__ void horizon_token_attention_pool_fwd(
|
||||
const float* __restrict__ horizon_tokens, // [N_H, H]
|
||||
const float* __restrict__ Q, // [H]
|
||||
const float* __restrict__ ln_out, // [B, K, H]
|
||||
int n_batch, int k_seq,
|
||||
float* __restrict__ ctx_h, // [B, N_H, H] OUT
|
||||
float* __restrict__ attn_weights // [B, N_H + K] OUT
|
||||
);
|
||||
|
||||
// Identical algorithm to attention_pool.cu but with K' = N_H + K and the
|
||||
// per-horizon output computed by indexing attn at positions [0, N_H) for
|
||||
// the token contribution and [N_H, N_H + K) for the time contribution.
|
||||
// Use the same warp-shuffle block_reduce_sum helper.
|
||||
```
|
||||
|
||||
Implement bwd similarly:
|
||||
```cuda
|
||||
extern "C" __global__ void horizon_token_attention_pool_bwd(
|
||||
const float* __restrict__ horizon_tokens,
|
||||
const float* __restrict__ Q,
|
||||
const float* __restrict__ ln_out,
|
||||
const float* __restrict__ attn_weights,
|
||||
const float* __restrict__ grad_ctx_h,
|
||||
int n_batch, int k_seq,
|
||||
float* __restrict__ grad_horizon_tokens_scratch, // [B, N_H, H] (+=, reduce_axis0)
|
||||
float* __restrict__ grad_Q_scratch, // [B, H] (+=, reduce_axis0)
|
||||
float* __restrict__ grad_ln_out // [B, K, H] (+=)
|
||||
);
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Add to build.rs KERNELS array**
|
||||
|
||||
```rust
|
||||
"horizon_token_attention_pool",
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Write Rust binding**
|
||||
|
||||
```rust
|
||||
// src/horizon_token_attention_pool.rs
|
||||
pub struct HorizonTokenAttentionPool { /* fwd_fn, bwd_fn, stream, module */ }
|
||||
impl HorizonTokenAttentionPool {
|
||||
pub fn new(ctx, stream) -> Result<Self> { /* ... */ }
|
||||
pub fn forward(&self, horizon_tokens, q, ln_out, n_b, k, ctx_h, attn) -> Result<()> { /* no synchronize */ }
|
||||
pub fn backward(...) -> Result<()> { /* no synchronize */ }
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Run numgrad — must pass**
|
||||
|
||||
Run: `cargo test -p ml-alpha --features cuda --test horizon_token_attention_pool_numgrad -- --ignored`
|
||||
Expected: PASS within 5e-2 rel / 5e-3 abs.
|
||||
|
||||
- [ ] **Step 6: Commit V2**
|
||||
|
||||
```bash
|
||||
git add crates/ml-alpha/cuda/horizon_token_attention_pool.cu crates/ml-alpha/src/horizon_token_attention_pool.rs crates/ml-alpha/tests/horizon_token_attention_pool_numgrad.rs crates/ml-alpha/build.rs crates/ml-alpha/src/lib.rs
|
||||
git commit -m "feat(ml-alpha): horizon_token_attention_pool kernel + numgrad (v2 C)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task V3: (E) inverted_attention_pool — kernel + binding + numgrad
|
||||
|
||||
**Files:** symmetric to V2.
|
||||
|
||||
**Math:**
|
||||
```
|
||||
X_inv[b, h, k] = ln_out[b, k, h] # [B, H, K]
|
||||
scores[b, h, j] = Σ_k X_inv[b, h, k] · Q_inv[j, k] # [B, H, H]
|
||||
attn[b, h, j] = softmax_j(scores[b, h, j]) # [B, H, H]
|
||||
feat_inv[b, h, k] = Σ_j attn[b, h, j] · X_inv[b, j, k] # [B, H, K]
|
||||
pooled_inv[b, h] = mean_k(feat_inv[b, h, k]) # [B, H]
|
||||
broadcast over N_HORIZONS → inv_ctx[b, n, h] = pooled_inv[b, h] # [B, N_H, H]
|
||||
```
|
||||
|
||||
Params: `Q_inv [HIDDEN_DIM, K_MAX]` (where K_MAX is a compile-time upper bound = 64 for our K range). Stored as `[HIDDEN_DIM, K_MAX]` but only the first k_seq columns are used.
|
||||
|
||||
- [ ] **Step 1-6: Same TDD shape as V2 (failing test → kernel → binding → passing test → commit)**
|
||||
|
||||
`git commit -m "feat(ml-alpha): inverted_attention_pool kernel + numgrad (v2 E)"`
|
||||
|
||||
---
|
||||
|
||||
## Task V4: Fused 1×1 projection — kernel + binding + numgrad
|
||||
|
||||
Concatenates `ctx_h [B, N_H, H]` (from V2) and `inv_ctx [B, N_H, H]` (from V3) along the channel axis to get `[B, N_H, 2H]`, then applies a 1×1 projection `W_fuse [2H, H]` + `b_fuse [H]` to recover `fused_ctx [B, N_H, H]`.
|
||||
|
||||
- [ ] **Step 1: failing numgrad test**
|
||||
- [ ] **Step 2: kernel `fused_ctx_proj.cu` — block-per-(batch, horizon), thread per output channel**
|
||||
- [ ] **Step 3: build.rs + binding**
|
||||
- [ ] **Step 4: numgrad passes**
|
||||
- [ ] **Step 5: commit**
|
||||
|
||||
`git commit -m "feat(ml-alpha): fused_ctx_proj kernel + numgrad (v2 C+E fuse)"`
|
||||
|
||||
---
|
||||
|
||||
## Task V5: (D) Regime feature extraction in snap_feature_assemble
|
||||
|
||||
**Goal:** Add a 12-dim regime feature vector per batch to the trainer's input.
|
||||
|
||||
- [ ] **Step 1: Define regime layout**
|
||||
|
||||
```rust
|
||||
// 12 dims: [spread_q1, spread_q2, ..., spread_q5, // 5
|
||||
// vol_low, vol_med, vol_high, // 3
|
||||
// tod_open, tod_mid_am, tod_mid_pm, tod_close] // 4
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Modify `snap_feature_assemble.cu` to emit regime_features [B, 12]**
|
||||
|
||||
Compute spread quintiles from rolling MBP-1 spread (1st level only). Vol regimes from rolling std of mid-price. Time-of-day from snapshot timestamp.
|
||||
|
||||
- [ ] **Step 3: Plumb through loader → trainer**
|
||||
|
||||
Add `regime_features_d: CudaSlice<f32>` to PerceptionTrainer state.
|
||||
|
||||
- [ ] **Step 4: Test smoke**
|
||||
|
||||
Synthetic data smoke that verifies regime values are 0/1 one-hot.
|
||||
|
||||
- [ ] **Step 5: commit**
|
||||
|
||||
`git commit -m "feat(ml-alpha): regime feature extraction (12-dim one-hot) (v2 D-1)"`
|
||||
|
||||
---
|
||||
|
||||
## Task V6: (D) regime_moe_gate kernel + numgrad
|
||||
|
||||
**Math:**
|
||||
```
|
||||
gate_logits[b, e] = Σ_r W_gate[e, r] · R[b, r] # [B, N_EXPERTS]
|
||||
gate_probs[b, e] = softmax_e(gate_logits[b, e]) # [B, N_EXPERTS]
|
||||
top_e[b] = argmax_e(gate_probs[b, e]) # [B]
|
||||
expert_out[e][b, n, d_out] = Σ_d_in W_e[e, d_out, d_in] · fused_ctx[b, n, d_in] + b_e[e, d_out]
|
||||
routed_ctx[b, n, d] = expert_out[top_e[b]][b, n, d]
|
||||
```
|
||||
|
||||
Backward: gradient routes only through `top_e[b]`'s expert. Top-1 is non-differentiable; we use straight-through estimator (`d_gate_logits[b, e] = gate_probs[b, e] · (gate_probs[b, e] − {e == top_e[b]}) · feedback_signal`, standard Switch Transformer pattern).
|
||||
|
||||
**Load balancing:**
|
||||
```
|
||||
frac[e] = mean_b({top_e[b] == e}) # fraction of batches routed to expert e
|
||||
prob_mean[e] = mean_b(gate_probs[b, e])
|
||||
aux_loss = N_EXPERTS · Σ_e (frac[e] · prob_mean[e])
|
||||
```
|
||||
Aux loss is added to total with weight `λ_aux = 0.01` (Switch Transformer default).
|
||||
|
||||
- [ ] **Step 1: failing numgrad test (verifies straight-through grad through the gate)**
|
||||
- [ ] **Step 2: kernel + binding (4 experts, N_EXPERTS=4)**
|
||||
- [ ] **Step 3: aux_loss + frac computation kernel**
|
||||
- [ ] **Step 4: numgrad passes**
|
||||
- [ ] **Step 5: commit**
|
||||
|
||||
`git commit -m "feat(ml-alpha): regime_moe_gate kernel + STE grad + numgrad (v2 D-2)"`
|
||||
|
||||
---
|
||||
|
||||
## Task V7: (A) Kendall σ-weighted BCE
|
||||
|
||||
**Modify `bce_loss_multi_horizon.cu`:** add a new arg `log_sigma_h [N_HORIZONS]`. Per-horizon loss becomes:
|
||||
|
||||
```
|
||||
sigma_h = expf(log_sigma_h[h])
|
||||
weight_h = 1.0f / (2.0f * sigma_h * sigma_h)
|
||||
weighted_bce[h] = weight_h * raw_bce[h] + log_sigma_h[h]
|
||||
```
|
||||
|
||||
Backward writes a new grad buffer `d_log_sigma_h [N_HORIZONS]`:
|
||||
```
|
||||
d_log_sigma_h[h] = (1 - weight_h * raw_bce[h])
|
||||
+ d_weighted_bce[h] * (-2 * weight_h * raw_bce[h] + 0)
|
||||
```
|
||||
(After working through `d/d(log_sigma) [(1/(2*exp(2*log_sigma))) * bce + log_sigma]`.)
|
||||
|
||||
- [ ] **Step 1: failing numgrad on the σ-weighted variant**
|
||||
- [ ] **Step 2: modify .cu, update binding signature**
|
||||
- [ ] **Step 3: trainer adds `log_sigma_h_d [N_HORIZONS]` learnable param + AdamW (LR scaled by 0.25)**
|
||||
- [ ] **Step 4: numgrad passes**
|
||||
- [ ] **Step 5: commit**
|
||||
|
||||
`git commit -m "feat(ml-alpha): Kendall sigma-weighted BCE + log_sigma param group (v2 A)"`
|
||||
|
||||
---
|
||||
|
||||
## Task V8: (B) anchor_l2 + Wiener-α lambda controller
|
||||
|
||||
**Math:**
|
||||
```
|
||||
anchor_loss[group] = λ_anchor(t) · Σ_i (p[i] − p_init[i])²
|
||||
λ_anchor(t) = max(λ_floor, Wiener_α_smoothed(loss_improvement_rate))
|
||||
λ_floor = ||p_init||₂ / (100 · sqrt(numel(p))) # ISV-driven floor per pearl
|
||||
```
|
||||
|
||||
The anchor coefficient is driven by the EMA-tracked improvement rate of val_loss (mirroring the existing horizon_lambda controller pattern in the trainer). Wiener-α with a floor of 0.4 per `pearl_wiener_alpha_floor_for_nonstationary`.
|
||||
|
||||
Parameter groups anchored: `horizon_tokens [N_H, H]`, `Q [H]`. Their initialisation snapshots are captured at trainer construction and stored in non-trainable buffers.
|
||||
|
||||
- [ ] **Step 1: failing numgrad on anchor_l2 kernel**
|
||||
- [ ] **Step 2: kernel `anchor_l2.cu` — block-per-param-group, warp-shuffle reduce, emits scalar anchor_loss + grad += `2λ(p - p_init)` on the parameter buffer**
|
||||
- [ ] **Step 3: AnchorController in `trainer/anchor_controller.rs` — Wiener-α + floor**
|
||||
- [ ] **Step 4: numgrad passes**
|
||||
- [ ] **Step 5: commit**
|
||||
|
||||
`git commit -m "feat(ml-alpha): anchor_l2 + Wiener-alpha controller (v2 B)"`
|
||||
|
||||
---
|
||||
|
||||
## Task V9: Trainer state v2 (Perceptionv2State)
|
||||
|
||||
Bundle all the new v2 components into one cohesive trainer field:
|
||||
|
||||
```rust
|
||||
pub struct PerceptionV2State {
|
||||
// (C) horizon tokens
|
||||
pub horizon_tokens_d: CudaSlice<f32>,
|
||||
pub horizon_tokens_init_d: CudaSlice<f32>,
|
||||
pub grad_horizon_tokens_d: CudaSlice<f32>,
|
||||
pub opt_horizon_tokens: AdamW,
|
||||
pub attn_pool: HorizonTokenAttentionPool,
|
||||
|
||||
// (E) inverted attention
|
||||
pub q_inv_d: CudaSlice<f32>,
|
||||
pub grad_q_inv_d: CudaSlice<f32>,
|
||||
pub opt_q_inv: AdamW,
|
||||
pub inv_pool: InvertedAttentionPool,
|
||||
|
||||
// (C+E fuse) projection
|
||||
pub w_fuse_d: CudaSlice<f32>,
|
||||
pub b_fuse_d: CudaSlice<f32>,
|
||||
pub grad_w_fuse_d: CudaSlice<f32>,
|
||||
pub grad_b_fuse_d: CudaSlice<f32>,
|
||||
pub opt_fuse: AdamW,
|
||||
pub fuse_proj: FusedCtxProj,
|
||||
|
||||
// (D) regime-MoE
|
||||
pub w_gate_d: CudaSlice<f32>,
|
||||
pub experts_d: CudaSlice<f32>, // [N_EXPERTS, H, H]
|
||||
pub experts_bias_d: CudaSlice<f32>, // [N_EXPERTS, H]
|
||||
pub grad_w_gate_d, grad_experts_d, grad_experts_bias_d
|
||||
pub opt_gate, opt_experts: AdamW
|
||||
pub moe: RegimeMoeGate,
|
||||
|
||||
// (A) Kendall σ
|
||||
pub log_sigma_h_d: CudaSlice<f32>, // [N_HORIZONS]
|
||||
pub grad_log_sigma_h_d: CudaSlice<f32>,
|
||||
pub opt_log_sigma: AdamW,
|
||||
|
||||
// (B) anchor controller
|
||||
pub q_init_d: CudaSlice<f32>, // non-trainable snapshot
|
||||
pub anchor: AnchorController, // emits λ_anchor each step
|
||||
pub anchor_l2: AnchorL2,
|
||||
|
||||
// shared bookkeeping
|
||||
n_batch: usize, k_seq: usize, stream: Arc<CudaStream>,
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 1: New module file**
|
||||
- [ ] **Step 2: `new(...)` with all init: horizon_tokens at 1/√H scale, Q_inv 1/√K, experts Xavier, gate small init, log_sigma at 0, anchor snapshots captured**
|
||||
- [ ] **Step 3: `zero_grads()` using stream.memset_zeros (capture-safe)**
|
||||
- [ ] **Step 4: `forward(ln_b_out) → fused_routed_ctx [B, N_H, H]` (all kernel calls — no synchronize)**
|
||||
- [ ] **Step 5: `backward(grad_fused_routed_ctx, ln_b_out, grad_ln_b_out)`**
|
||||
- [ ] **Step 6: `adamw_step()` — all 6 optimizer groups**
|
||||
- [ ] **Step 7: commit**
|
||||
|
||||
`git commit -m "feat(ml-alpha): PerceptionV2State (all v2 components bundled)"`
|
||||
|
||||
---
|
||||
|
||||
## Task V10: Wire v2 into PerceptionTrainer
|
||||
|
||||
Replace the existing single-Q `attention_pool` invocation + multi-horizon heads input with the v2 path:
|
||||
|
||||
```
|
||||
Old: ctx[B, H] = attention_pool(Q, LN_b_out)
|
||||
heads_in[k, b, h, :] = [ctx[b, :]; h_carry[b, :]] // shared ctx, broadcast across horizons
|
||||
|
||||
New: fused_routed_ctx[B, N_H, H] = PerceptionV2State.forward(LN_b_out)
|
||||
heads_in[k, b, h, :] = [fused_routed_ctx[b, h, :]; h_carry[b, :]] // per-horizon ctx
|
||||
```
|
||||
|
||||
GRN heads then consume per-horizon contexts naturally.
|
||||
|
||||
- [ ] **Step 1: Replace attention_pool fwd in `dispatch_train_step`**
|
||||
- [ ] **Step 2: Replace attention_pool bwd**
|
||||
- [ ] **Step 3: Replace BCE call with σ-weighted variant; thread log_sigma_h_d as 7th input**
|
||||
- [ ] **Step 4: Add anchor_l2 launch + Wiener-α λ update at the end of step_batched**
|
||||
- [ ] **Step 5: Add aux_loss accumulation from MoE; add to total loss**
|
||||
- [ ] **Step 6: Add v2.adamw_step() to section 9**
|
||||
- [ ] **Step 7: Verify graph-capture safety (no synchronize, no host allocs, no scalar arg pointer changes)**
|
||||
- [ ] **Step 8: Workspace check passes**
|
||||
- [ ] **Step 9: commit**
|
||||
|
||||
`git commit -m "feat(ml-alpha): wire PerceptionV2State into PerceptionTrainer step_batched"`
|
||||
|
||||
---
|
||||
|
||||
## Task V11: Local end-to-end smoke + init-identity invariant
|
||||
|
||||
The original `per_horizon_full_pipeline_smoke::alpha_zero_init_is_identity_to_baseline` test is gone. New invariant: at `log_sigma_h = 0` AND `λ_anchor = 0`, v2 with the horizon tokens initialized to the same Q (i.e., horizon_tokens[h, :] = Q at init) should produce contexts that match the OLD single-Q attention-pool output bit-for-bit modulo the inverted/MoE/fuse paths.
|
||||
|
||||
Since v2 also includes inverted attention + MoE + fuse, full bit-identity to v0 is impossible. The relaxed invariant: **forward runs without NaN, backward gradients are finite, AdamW does not produce NaN params**. This is the v2 smoke contract.
|
||||
|
||||
- [ ] **Step 1: Write `v2_pipeline_smoke.rs`**
|
||||
|
||||
```rust
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn v2_forward_backward_step_produces_finite_grads() {
|
||||
// Construct PerceptionV2State with B=2, K=8, H=128, N_H=5
|
||||
// Random ln_b_out, random labels, random regime features
|
||||
// Run forward → fused_routed_ctx; bce_loss; backward → all grads
|
||||
// Assert: all params finite, loss finite, no NaN
|
||||
// Run 3 AdamW steps; assert params bounded
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Run + iterate until passes**
|
||||
- [ ] **Step 3: commit**
|
||||
|
||||
`git commit -m "test(ml-alpha): v2 pipeline smoke (no-NaN invariant)"`
|
||||
|
||||
---
|
||||
|
||||
## Task V12: Cluster smoke (1 epoch)
|
||||
|
||||
- [ ] **Step 1: Push branch**
|
||||
|
||||
```bash
|
||||
git push origin ml-alpha-phase-a
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Submit smoke**
|
||||
|
||||
```bash
|
||||
./scripts/argo-alpha-perception.sh --epochs 1 --n-train-seqs 1000 --n-val-seqs 200 --cv-n-folds 1
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Poll to terminal state with robust watcher (see /tmp/ab_sweep_watcher.sh pattern)**
|
||||
|
||||
- [ ] **Step 4: Inspect logs — verify:**
|
||||
- No CUDA_ERROR
|
||||
- Epoch time ≤ 30 s
|
||||
- val_loss + per-horizon AUC printed
|
||||
- log_sigma_h evolves
|
||||
- λ_anchor evolves
|
||||
- aux_loss reported
|
||||
|
||||
- [ ] **Step 5: If any failure → diagnose root cause + fix (no quickfixes per `feedback_no_quickfixes`)**
|
||||
|
||||
- [ ] **Step 6: When smoke passes, decision point**: if any kill criterion already fails at 1 epoch (NaN, time > 60s, etc.) — STOP, fix, repeat. Otherwise proceed to V13.
|
||||
|
||||
---
|
||||
|
||||
## Task V13: Cluster A/B 3-fold @ 30 epochs
|
||||
|
||||
- [ ] **Step 1: Submit 3 folds at HEAD**
|
||||
|
||||
```bash
|
||||
./scripts/argo-alpha-perception.sh --epochs 30 --cv-fold 0 --cv-n-folds 3
|
||||
./scripts/argo-alpha-perception.sh --epochs 30 --cv-fold 1 --cv-n-folds 3
|
||||
./scripts/argo-alpha-perception.sh --epochs 30 --cv-fold 2 --cv-n-folds 3
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Watch with the robust watcher**
|
||||
|
||||
- [ ] **Step 3: Aggregate AUC mean ± std across folds; compute per-horizon variance**
|
||||
|
||||
- [ ] **Step 4: Compare against task #200 baseline (mean_auc 0.7749 ± 0.024 / h6000 0.7591 ± 0.018) and spec §5 falsifiable claims (A-E)**
|
||||
|
||||
- [ ] **Step 5: Write a project memory file `project_ml_alpha_v2_verdict.md` capturing per-axis kill-criterion pass/fail**
|
||||
|
||||
- [ ] **Step 6: commit + push**
|
||||
|
||||
---
|
||||
|
||||
## Per-commit verification gates
|
||||
|
||||
After every commit V1-V13:
|
||||
1. `SQLX_OFFLINE=true cargo check -p ml-alpha` — clean (warnings count == 0 for new code)
|
||||
2. Affected numgrad tests (where applicable): `cargo test ... -- --ignored`
|
||||
3. `pre-commit` hook passes (already enforced)
|
||||
|
||||
## Stop conditions
|
||||
|
||||
STOP and report if:
|
||||
- Any numgrad test fails after 3 attempts to fix
|
||||
- Cluster smoke (V12) crashes
|
||||
- Cluster A/B (V13) shows mean_auc regression > -0.02 vs baseline
|
||||
- Wall-time > 60 s/epoch
|
||||
- ≥ 3 of 5 axis kill criteria from spec §5 fail
|
||||
Reference in New Issue
Block a user