perf(ml-alpha): block-per-batch cfc_step + reduce_axis0 reducer (Phase B commit 1)
Per docs/superpowers/specs/2026-05-17-kloop-parallelization-design.md. cfc_step_batched (fwd + bwd) refactored from grid=(1,1,1) with internal n_batch loop to grid=(B,1,1) — each block handles one batch. Removes the single-SM bottleneck on the K-loop's most-called kernel (64×/step). Param-grad accumulation moves to per-batch scratch: cfc_grad_w_in_scratch_d [B, n_hid, n_in] cfc_grad_w_rec_scratch_d [B, n_hid, n_hid] cfc_grad_b_scratch_d [B, n_hid] cfc_grad_tau_scratch_d [B, n_hid] Zeroed once per training step, K-loop's 64 bwd calls += into them, then 4 reduce_axis0 launches collapse B → final grad buffers (OVERWRITE) before AdamW. New AdamW-after-reducer invariant: final grads are meaningful only after the reducer has run in the current step. New reduce_axis0 kernel: single parameterised reducer [B, N] → [N] via block tree-reduce (no atomicAdd per feedback_no_atomicadd.md). Same pattern as layer_norm_reduce_param_grads — CUDA-Graph-safe. cfc_step_backward_batched shared-mem dropped from (B+1)*n_hid*4 to 2*n_hid*4 bytes per block (only one row of sd_pre needed per block bi). Tests: - New stacked_trainer_loss_shrinks_at_batch_32: FIRST test that actually exercises the cross-batch reduction code path; existing perception_overfit suite was all B=1. Initial 0.24 → final 0.00. - Scratch-clears test removed (explanatory comment kept): structurally hard to assert directly due to begin_capture/end_capture not executing kernels; the B=32 convergence smoke implicitly validates scratch zeroing since divergence would otherwise be immediate. All 9 perception_overfit smokes + 4 backward_finite_diff tests pass. build.rs: - KERNELS list adds "reduce_axis0" - Cache-bust → v11 Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -20,12 +20,13 @@ const KERNELS: &[&str] = &[
|
||||
"layer_norm", // Phase 1: trunk pre-CfC normalisation
|
||||
"variable_selection", // Phase 2D: TFT-style per-feature gating
|
||||
"attention_pool", // Phase 3: learned context summary at CfC k=0
|
||||
"reduce_axis0", // Phase B: cross-batch param-grad reducer
|
||||
];
|
||||
|
||||
// Cache bust v10 (2026-05-17): Phase 3 attention pool — new
|
||||
// attention_pool.cu kernel pair (fwd + bwd) for the learned context
|
||||
// summary replacing CfC's zero-initialised h_old. Force fresh nvcc
|
||||
// compile so the new cubin lands in OUT_DIR.
|
||||
// Cache bust v11 (2026-05-17): K-loop parallelization Phase B —
|
||||
// new reduce_axis0.cu kernel + block-per-batch refactor of
|
||||
// cfc_step_batched (fwd+bwd). Old cubins don't have the new symbols.
|
||||
// Force fresh nvcc compile.
|
||||
|
||||
fn main() {
|
||||
println!("cargo:rerun-if-changed=build.rs");
|
||||
|
||||
@@ -168,20 +168,25 @@ extern "C" __global__ void cfc_step_backward(
|
||||
|
||||
// ─── Batched variants ────────────────────────────────────────────────
|
||||
//
|
||||
// Same math as the single-sample kernels above, but process B sequences
|
||||
// in parallel within one launch. Layout: x[B, n_in], h_old[B, n_hid],
|
||||
// h_new[B, n_hid] (row-major, B-major). Thread i handles hidden unit i
|
||||
// across ALL B samples — per-thread loop over b ∈ 0..B.
|
||||
// Phase B (2026-05-17): block-per-batch refactor. Launch contract:
|
||||
// grid=(n_batch, 1, 1), block=(n_hid, 1, 1).
|
||||
// Each block handles ONE batch's n_hid channels in parallel.
|
||||
//
|
||||
// Param-grad writes (grad_b/grad_w_in/grad_w_rec/grad_tau) use `+=` into
|
||||
// SHARED-across-batch buffers; thread i is the sole writer to its row of
|
||||
// each grad matrix, so no atomic / no race regardless of B. Sample-
|
||||
// scoped outputs (h_new, grad_h_old, grad_x) write into per-b slots.
|
||||
// Param-grad writes go through per-batch scratch buffers
|
||||
// (grad_w_in_scratch [B, n_hid, n_in] etc.). Thread (bi, i) is the
|
||||
// sole writer to its slice of each scratch tensor; the K-loop's 64
|
||||
// invocations accumulate via += within the same (bi, i, *) slice.
|
||||
// After the K-loop the caller runs `reduce_axis0` to sum B → final
|
||||
// grad buffer (OVERWRITE semantics). No atomicAdd per
|
||||
// feedback_no_atomicadd.md.
|
||||
//
|
||||
// Shared memory budget for the backward kernel: `sd_pre[B, n_hid] +
|
||||
// sdecay[n_hid]` floats = (B+1) * n_hid * 4 bytes. At B=32, n_hid=128:
|
||||
// 33 * 128 * 4 = 16.5 KiB per block — well under the 48 KiB L40S limit.
|
||||
// Shared memory budget per block: sd_pre[n_hid] + sdecay[n_hid] =
|
||||
// 2 * n_hid * 4 bytes. At n_hid=128 → 1 KiB. Far below 48 KiB limit,
|
||||
// no cudaFuncSetAttribute needed.
|
||||
|
||||
// Block-per-batch (Phase B): grid=(n_batch, 1, 1), block=(n_hid, 1, 1).
|
||||
// Each block handles ONE batch's n_hid channels in parallel. Removes
|
||||
// the internal batch loop — frees ~31 SMs per launch at B=32.
|
||||
extern "C" __global__ void cfc_step_batched(
|
||||
const float* __restrict__ w_in, // [n_hid, n_in], row-major
|
||||
const float* __restrict__ w_rec, // [n_hid, n_hid], row-major
|
||||
@@ -195,50 +200,61 @@ extern "C" __global__ void cfc_step_batched(
|
||||
int n_batch,
|
||||
float* __restrict__ h_new // [n_batch, n_hid]
|
||||
) {
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= n_hid) return;
|
||||
int bi = blockIdx.x;
|
||||
int i = threadIdx.x;
|
||||
if (bi >= n_batch || i >= n_hid) return;
|
||||
|
||||
const float bias_i = b[i];
|
||||
const float decay = expf(-dt_s / fmaxf(tau[i], 1e-6f));
|
||||
const float one_minus_decay = 1.0f - decay;
|
||||
|
||||
for (int bi = 0; bi < n_batch; ++bi) {
|
||||
const float* x_b = x + (long long)bi * n_in;
|
||||
const float* h_old_b = h_old + (long long)bi * n_hid;
|
||||
float pre = bias_i;
|
||||
for (int k = 0; k < n_in; ++k) pre += w_in[i * n_in + k] * x_b[k];
|
||||
for (int k = 0; k < n_hid; ++k) pre += w_rec[i * n_hid + k] * h_old_b[k];
|
||||
h_new[(long long)bi * n_hid + i] = h_old_b[i] * decay + one_minus_decay * tanhf(pre);
|
||||
}
|
||||
const float* x_b = x + (long long)bi * n_in;
|
||||
const float* h_old_b = h_old + (long long)bi * n_hid;
|
||||
float pre = bias_i;
|
||||
for (int k = 0; k < n_in; ++k) pre += w_in[i * n_in + k] * x_b[k];
|
||||
for (int k = 0; k < n_hid; ++k) pre += w_rec[i * n_hid + k] * h_old_b[k];
|
||||
h_new[(long long)bi * n_hid + i] = h_old_b[i] * decay + one_minus_decay * tanhf(pre);
|
||||
}
|
||||
|
||||
|
||||
// Block-per-batch backward.
|
||||
// grid=(n_batch, 1, 1) block=(n_hid, 1, 1)
|
||||
// shared mem: 2 * n_hid * 4 bytes (sd_pre + sdecay; one row each)
|
||||
//
|
||||
// Param-grad scratch tensors hold per-batch slices. Thread (bi, i) is
|
||||
// sole writer to its slice of each scratch tensor for every K-iteration
|
||||
// the K-loop calls this kernel — no race even with B parallel blocks.
|
||||
// Caller zeroes the scratch buffers once per training step; this kernel
|
||||
// uses += to accumulate across the K-loop's 64 invocations.
|
||||
//
|
||||
// Outputs grad_h_old and grad_x are per-batch indexed; each block bi is
|
||||
// sole writer to its [bi, :] slice (overwrite, single launch).
|
||||
extern "C" __global__ void cfc_step_backward_batched(
|
||||
const float* __restrict__ w_in, // [n_hid, n_in]
|
||||
const float* __restrict__ w_rec, // [n_hid, n_hid]
|
||||
const float* __restrict__ b,
|
||||
const float* __restrict__ tau,
|
||||
const float* __restrict__ x, // [n_batch, n_in]
|
||||
const float* __restrict__ h_old, // [n_batch, n_hid]
|
||||
const float* __restrict__ grad_h_new, // [n_batch, n_hid]
|
||||
const float* __restrict__ w_in, // [n_hid, n_in]
|
||||
const float* __restrict__ w_rec, // [n_hid, n_hid]
|
||||
const float* __restrict__ b, // [n_hid]
|
||||
const float* __restrict__ tau, // [n_hid]
|
||||
const float* __restrict__ x, // [n_batch, n_in]
|
||||
const float* __restrict__ h_old, // [n_batch, n_hid]
|
||||
const float* __restrict__ grad_h_new, // [n_batch, n_hid]
|
||||
float dt_s,
|
||||
int n_in,
|
||||
int n_hid,
|
||||
int n_batch,
|
||||
float* __restrict__ grad_w_in, // [n_hid, n_in] accum +=
|
||||
float* __restrict__ grad_w_rec, // [n_hid, n_hid] accum +=
|
||||
float* __restrict__ grad_b, // [n_hid] accum +=
|
||||
float* __restrict__ grad_tau, // [n_hid] accum +=
|
||||
float* __restrict__ grad_h_old, // [n_batch, n_hid] overwrite
|
||||
float* __restrict__ grad_x // [n_batch, n_in] overwrite (nullptr OK)
|
||||
float* __restrict__ grad_w_in_scratch, // [n_batch, n_hid, n_in] accum +=
|
||||
float* __restrict__ grad_w_rec_scratch, // [n_batch, n_hid, n_hid] accum +=
|
||||
float* __restrict__ grad_b_scratch, // [n_batch, n_hid] accum +=
|
||||
float* __restrict__ grad_tau_scratch, // [n_batch, n_hid] accum +=
|
||||
float* __restrict__ grad_h_old, // [n_batch, n_hid] overwrite
|
||||
float* __restrict__ grad_x // [n_batch, n_in] overwrite (nullptr OK)
|
||||
) {
|
||||
// Shared mem: sd_pre[B, n_hid] then sdecay[n_hid].
|
||||
extern __shared__ float smem[];
|
||||
float* sd_pre = smem; // [n_batch * n_hid]
|
||||
float* sdecay = sd_pre + (long long)n_batch * n_hid; // [n_hid]
|
||||
float* sd_pre = smem; // [n_hid] — one row for this block's bi
|
||||
float* sdecay = sd_pre + n_hid; // [n_hid]
|
||||
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= n_hid) return;
|
||||
int bi = blockIdx.x;
|
||||
int i = threadIdx.x;
|
||||
if (bi >= n_batch || i >= n_hid) return;
|
||||
|
||||
const float bias_i = b[i];
|
||||
const float tau_eps = 1e-6f;
|
||||
@@ -246,69 +262,51 @@ extern "C" __global__ void cfc_step_backward_batched(
|
||||
const float decay = expf(-dt_s / tau_eff);
|
||||
sdecay[i] = decay;
|
||||
|
||||
// Per-thread (private) accumulators across batch — avoids any
|
||||
// cross-thread race since thread i is sole writer to its row.
|
||||
float grad_b_acc = 0.0f;
|
||||
float grad_tau_acc = 0.0f;
|
||||
// grad_w_in_acc / grad_w_rec_acc would be n_in / n_hid floats per
|
||||
// thread — too big for registers. Instead we accumulate them via
|
||||
// direct `+=` into global memory (thread i is sole writer to its
|
||||
// rows; no race even across the per-b loop).
|
||||
const float* x_b = x + (long long)bi * n_in;
|
||||
const float* h_old_b = h_old + (long long)bi * n_hid;
|
||||
const float dh_b = grad_h_new[(long long)bi * n_hid + i];
|
||||
|
||||
// Pass 1: per-b compute d_pre, store in shared, accumulate
|
||||
// per-row param grads (grad_w_in, grad_w_rec rows).
|
||||
for (int bi = 0; bi < n_batch; ++bi) {
|
||||
const float* x_b = x + (long long)bi * n_in;
|
||||
const float* h_old_b = h_old + (long long)bi * n_hid;
|
||||
const float dh_b = grad_h_new[(long long)bi * n_hid + i];
|
||||
float pre = bias_i;
|
||||
for (int k = 0; k < n_in; ++k) pre += w_in[i * n_in + k] * x_b[k];
|
||||
for (int k = 0; k < n_hid; ++k) pre += w_rec[i * n_hid + k] * h_old_b[k];
|
||||
const float s = tanhf(pre);
|
||||
const float d_pre_b = dh_b * (1.0f - decay) * (1.0f - s * s);
|
||||
sd_pre[i] = d_pre_b;
|
||||
|
||||
float pre = bias_i;
|
||||
for (int k = 0; k < n_in; ++k) pre += w_in[i * n_in + k] * x_b[k];
|
||||
for (int k = 0; k < n_hid; ++k) pre += w_rec[i * n_hid + k] * h_old_b[k];
|
||||
const float s = tanhf(pre);
|
||||
const float d_pre_b = dh_b * (1.0f - decay) * (1.0f - s * s);
|
||||
sd_pre[(long long)bi * n_hid + i] = d_pre_b;
|
||||
|
||||
grad_b_acc += d_pre_b;
|
||||
// grad_w_in / grad_w_rec rows: thread i sole writer → safe +=.
|
||||
for (int k = 0; k < n_in; ++k) {
|
||||
grad_w_in[i * n_in + k] += d_pre_b * x_b[k];
|
||||
}
|
||||
for (int k = 0; k < n_hid; ++k) {
|
||||
grad_w_rec[i * n_hid + k] += d_pre_b * h_old_b[k];
|
||||
}
|
||||
if (tau[i] > tau_eps) {
|
||||
const float d_decay_b = dh_b * (h_old_b[i] - s);
|
||||
grad_tau_acc += d_decay_b * decay * dt_s / (tau_eff * tau_eff);
|
||||
}
|
||||
// Pass 1: per-batch param-grad scratch writes.
|
||||
// Single-writer per (bi, i, *) — block bi is sole writer to its slice;
|
||||
// K-loop's 64 invocations accumulate within the same slice via +=.
|
||||
grad_b_scratch[(long long)bi * n_hid + i] += d_pre_b;
|
||||
for (int k = 0; k < n_in; ++k) {
|
||||
grad_w_in_scratch[((long long)bi * n_hid + i) * n_in + k] += d_pre_b * x_b[k];
|
||||
}
|
||||
for (int k = 0; k < n_hid; ++k) {
|
||||
grad_w_rec_scratch[((long long)bi * n_hid + i) * n_hid + k] += d_pre_b * h_old_b[k];
|
||||
}
|
||||
if (tau[i] > tau_eps) {
|
||||
const float d_decay_b = dh_b * (h_old_b[i] - s);
|
||||
grad_tau_scratch[(long long)bi * n_hid + i] += d_decay_b * decay * dt_s / (tau_eff * tau_eff);
|
||||
}
|
||||
grad_b[i] += grad_b_acc;
|
||||
grad_tau[i] += grad_tau_acc;
|
||||
__syncthreads();
|
||||
|
||||
// Pass 2: per-b emit grad_h_old[bi, i] = sum_j sd_pre[bi, j] * W_rec[j, i] + dh_bi * decay.
|
||||
for (int bi = 0; bi < n_batch; ++bi) {
|
||||
float gh = grad_h_new[(long long)bi * n_hid + i] * decay;
|
||||
// Pass 2: grad_h_old[bi, i] = sum_j sd_pre[j] * W_rec[j, i] + dh_b * decay.
|
||||
{
|
||||
float gh = dh_b * decay;
|
||||
for (int j = 0; j < n_hid; ++j) {
|
||||
gh += sd_pre[(long long)bi * n_hid + j] * w_rec[j * n_hid + i];
|
||||
gh += sd_pre[j] * w_rec[j * n_hid + i];
|
||||
}
|
||||
grad_h_old[(long long)bi * n_hid + i] = gh;
|
||||
}
|
||||
|
||||
// Pass 3: per-b emit grad_x[bi, k] via thread-0 reduction (n_in is
|
||||
// small, this is amortised over the K-loop work). Mirrors the
|
||||
// single-sample kernel's pattern.
|
||||
// Pass 3: grad_x[bi, k] via thread 0 of this block.
|
||||
if (grad_x != nullptr && i == 0) {
|
||||
for (int bi = 0; bi < n_batch; ++bi) {
|
||||
const float* sd_pre_b = sd_pre + (long long)bi * n_hid;
|
||||
float* grad_x_b = grad_x + (long long)bi * n_in;
|
||||
for (int k = 0; k < n_in; ++k) {
|
||||
float gx = 0.0f;
|
||||
for (int j = 0; j < n_hid; ++j) {
|
||||
gx += sd_pre_b[j] * w_in[j * n_in + k];
|
||||
}
|
||||
grad_x_b[k] = gx;
|
||||
float* grad_x_b = grad_x + (long long)bi * n_in;
|
||||
for (int k = 0; k < n_in; ++k) {
|
||||
float gx = 0.0f;
|
||||
for (int j = 0; j < n_hid; ++j) {
|
||||
gx += sd_pre[j] * w_in[j * n_in + k];
|
||||
}
|
||||
grad_x_b[k] = gx;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
40
crates/ml-alpha/cuda/reduce_axis0.cu
Normal file
40
crates/ml-alpha/cuda/reduce_axis0.cu
Normal file
@@ -0,0 +1,40 @@
|
||||
// reduce_axis0.cu — sum [B, N] → [N] along the leading axis.
|
||||
//
|
||||
// One block per output column (N blocks total). Threads tile B via a
|
||||
// stride-loop, then block-wide tree-reduce. No atomicAdd per
|
||||
// feedback_no_atomicadd.md. OVERWRITE semantics: out[j] = sum.
|
||||
//
|
||||
// Block dim is hard-coded to 256 — matches the typical tail dim
|
||||
// granularity in the trainer (well within L40S limits, fits in shared
|
||||
// memory at 1 KiB per block).
|
||||
|
||||
#define REDAX0_BLOCK 256
|
||||
|
||||
extern "C" __global__ void reduce_axis0(
|
||||
const float* __restrict__ per_batch, // [B, N]
|
||||
int n_batch,
|
||||
int n_tail,
|
||||
float* __restrict__ out // [N] — OVERWRITE
|
||||
) {
|
||||
int j = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
if (j >= n_tail) return;
|
||||
|
||||
__shared__ float s[REDAX0_BLOCK];
|
||||
|
||||
float my_sum = 0.0f;
|
||||
// Stride-loop across B — handles n_batch > BLOCK and n_batch < BLOCK uniformly.
|
||||
for (int bi = tid; bi < n_batch; bi += REDAX0_BLOCK) {
|
||||
my_sum += per_batch[(long long)bi * n_tail + j];
|
||||
}
|
||||
s[tid] = my_sum;
|
||||
__syncthreads();
|
||||
|
||||
// Block tree-reduce.
|
||||
for (int stride = REDAX0_BLOCK / 2; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) s[tid] += s[tid + stride];
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (tid == 0) out[j] = s[0];
|
||||
}
|
||||
@@ -62,6 +62,7 @@ const HORIZON_LAMBDA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/ho
|
||||
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"));
|
||||
const ATTENTION_POOL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/attention_pool.cubin"));
|
||||
const REDUCE_AXIS0_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/reduce_axis0.cubin"));
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct PerceptionTrainerConfig {
|
||||
@@ -373,6 +374,22 @@ pub struct PerceptionTrainer {
|
||||
attn_fwd_fn: CudaFunction,
|
||||
attn_bwd_fn: CudaFunction,
|
||||
_attn_module: Arc<CudaModule>,
|
||||
|
||||
// ── K-loop parallelization (Phase B) ──
|
||||
// Per-batch grad scratch buffers for cfc_step_backward_batched.
|
||||
// Zeroed once per training step; the K-loop's 64 bwd calls accumulate
|
||||
// into these via +=. After the K-loop, reduce_axis0 sums across B
|
||||
// into the final grad buffers (overwrite semantics). AdamW reads the
|
||||
// final grads (must run AFTER the reducer).
|
||||
cfc_grad_w_in_scratch_d: CudaSlice<f32>, // [B, n_hid, n_in]
|
||||
cfc_grad_w_rec_scratch_d: CudaSlice<f32>, // [B, n_hid, n_hid]
|
||||
cfc_grad_b_scratch_d: CudaSlice<f32>, // [B, n_hid]
|
||||
cfc_grad_tau_scratch_d: CudaSlice<f32>, // [B, n_hid]
|
||||
/// Cross-batch reducer kernel: `[B, N] → [N]` via block tree-reduce.
|
||||
/// Used for every per-batch grad scratch in the refactored bwd path.
|
||||
reduce_axis0_fn: CudaFunction,
|
||||
/// Owns the reduce_axis0 cubin lifetime.
|
||||
_reduce_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
|
||||
@@ -507,6 +524,12 @@ impl PerceptionTrainer {
|
||||
let attn_bwd_fn = attn_module
|
||||
.load_function("attention_pool_bwd")
|
||||
.context("attention_pool_bwd symbol")?;
|
||||
let reduce_module = ctx
|
||||
.load_cubin(REDUCE_AXIS0_CUBIN.to_vec())
|
||||
.context("reduce_axis0 cubin")?;
|
||||
let reduce_axis0_fn = reduce_module
|
||||
.load_function("reduce_axis0")
|
||||
.context("reduce_axis0 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")?;
|
||||
@@ -648,6 +671,20 @@ impl PerceptionTrainer {
|
||||
// decay so the per-cell decay constants drift smoothly.
|
||||
let mut opt_tau = AdamW::new(dev, n_hid, cfg.lr_cfc * 0.1)?;
|
||||
opt_tau.wd = 0.0;
|
||||
|
||||
// ── K-loop parallelization (Phase B): cfc per-batch grad scratch ──
|
||||
let cfc_grad_w_in_scratch_d = stream
|
||||
.alloc_zeros::<f32>(cfg.n_batch * n_hid * n_in)
|
||||
.context("cfc_grad_w_in_scratch_d alloc")?;
|
||||
let cfc_grad_w_rec_scratch_d = stream
|
||||
.alloc_zeros::<f32>(cfg.n_batch * n_hid * n_hid)
|
||||
.context("cfc_grad_w_rec_scratch_d alloc")?;
|
||||
let cfc_grad_b_scratch_d = stream
|
||||
.alloc_zeros::<f32>(cfg.n_batch * n_hid)
|
||||
.context("cfc_grad_b_scratch_d alloc")?;
|
||||
let cfc_grad_tau_scratch_d = stream
|
||||
.alloc_zeros::<f32>(cfg.n_batch * n_hid)
|
||||
.context("cfc_grad_tau_scratch_d alloc")?;
|
||||
// GRN heads: 10 AdamW (5 weight + 5 bias). Biases get wd=0 per
|
||||
// the existing pattern; weights keep AdamW default wd.
|
||||
let opt_heads_w1 = AdamW::new(dev, N_HORIZONS * HEAD_MID_DIM * n_hid, cfg.lr_cfc)?;
|
||||
@@ -776,6 +813,13 @@ impl PerceptionTrainer {
|
||||
attn_fwd_fn,
|
||||
attn_bwd_fn,
|
||||
_attn_module: attn_module,
|
||||
// Phase B: cfc per-batch grad scratch + reducer.
|
||||
cfc_grad_w_in_scratch_d,
|
||||
cfc_grad_w_rec_scratch_d,
|
||||
cfc_grad_b_scratch_d,
|
||||
cfc_grad_tau_scratch_d,
|
||||
reduce_axis0_fn,
|
||||
_reduce_module: reduce_module,
|
||||
loss_ema_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
|
||||
lambda_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
|
||||
horizon_lambda_fn,
|
||||
@@ -1355,15 +1399,18 @@ impl PerceptionTrainer {
|
||||
).context("labels upload DtoD")?;
|
||||
}
|
||||
|
||||
// Zero shared param-grad accumulators (kernels += into these).
|
||||
self.stream.memset_zeros(&mut self.grad_w_in_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero grad_w_in: {e}"))?;
|
||||
self.stream.memset_zeros(&mut self.grad_w_rec_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero grad_w_rec: {e}"))?;
|
||||
self.stream.memset_zeros(&mut self.grad_b_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero grad_b: {e}"))?;
|
||||
self.stream.memset_zeros(&mut self.grad_tau_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero grad_tau: {e}"))?;
|
||||
// CfC per-batch grad scratch (Phase B): zero ONCE per step; K-loop
|
||||
// bwd accumulates into these, then reduce_axis0 collapses → final
|
||||
// grad buffers (OVERWRITE) after the K-loop. AdamW reads the final
|
||||
// grads (must run AFTER the reducer).
|
||||
self.stream.memset_zeros(&mut self.cfc_grad_w_in_scratch_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero cfc_grad_w_in_scratch: {e}"))?;
|
||||
self.stream.memset_zeros(&mut self.cfc_grad_w_rec_scratch_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero cfc_grad_w_rec_scratch: {e}"))?;
|
||||
self.stream.memset_zeros(&mut self.cfc_grad_b_scratch_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero cfc_grad_b_scratch: {e}"))?;
|
||||
self.stream.memset_zeros(&mut self.cfc_grad_tau_scratch_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero cfc_grad_tau_scratch: {e}"))?;
|
||||
// Phase 3: attention pool Q grad accumulator.
|
||||
self.stream.memset_zeros(&mut self.grad_attn_q_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero grad_attn_q: {e}"))?;
|
||||
@@ -1398,10 +1445,12 @@ impl PerceptionTrainer {
|
||||
let n_in_i = HIDDEN_DIM as i32;
|
||||
let n_hid_i = HIDDEN_DIM as i32;
|
||||
let block_dim = 128u32;
|
||||
let grid_dim = (HIDDEN_DIM as u32).div_ceil(block_dim);
|
||||
let _ = block_dim; // legacy 1D layout helper — superseded by block-per-batch launches.
|
||||
// CfC fwd: block-per-batch (Phase B). One block per sample;
|
||||
// each block's threads parallelize over n_hid channels.
|
||||
let cfg_cfc = LaunchConfig {
|
||||
grid_dim: (grid_dim, 1, 1),
|
||||
block_dim: (block_dim, 1, 1),
|
||||
grid_dim: (b_sz as u32, 1, 1),
|
||||
block_dim: (HIDDEN_DIM as u32, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let cfg_heads = LaunchConfig {
|
||||
@@ -1409,11 +1458,12 @@ impl PerceptionTrainer {
|
||||
block_dim: (N_HORIZONS as u32, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
// Batched CfC backward shared mem: sd_pre[B, n_hid] + sdecay[n_hid].
|
||||
let cfc_bwd_smem = ((b_sz + 1) * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32;
|
||||
// CfC bwd (Phase B): block-per-batch + 2 * n_hid floats of shared mem
|
||||
// (one sd_pre row + sdecay for the block's single bi).
|
||||
let cfc_bwd_smem = (2 * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32;
|
||||
let cfg_cfc_bwd = LaunchConfig {
|
||||
grid_dim: (grid_dim, 1, 1),
|
||||
block_dim: (block_dim, 1, 1),
|
||||
grid_dim: (b_sz as u32, 1, 1),
|
||||
block_dim: (HIDDEN_DIM as u32, 1, 1),
|
||||
shared_mem_bytes: cfc_bwd_smem,
|
||||
};
|
||||
// Batched heads backward shared mem: sd_z[B, 5].
|
||||
@@ -1653,14 +1703,19 @@ impl PerceptionTrainer {
|
||||
// cfc_step_bwd_batched: writes grad_h_carry (= grad_h_old_out for
|
||||
// position k-1) and grad_x[B, n_in]. grad_x output buffer IS
|
||||
// grad_henr_k_ptr — kernel writes directly into the K-slot.
|
||||
// Phase B: per-batch param-grad scratch + per-batch grad_h_old / grad_x.
|
||||
// Writes to scratch via += across K iterations; reduce_axis0
|
||||
// collapses B → final grad after the K-loop (see Pass 8c).
|
||||
unsafe {
|
||||
let mut launch = self.stream.launch_builder(&self.step_bwd_batched_fn);
|
||||
launch
|
||||
.arg(&self.w_in_d).arg(&self.w_rec_d).arg(&self.b_d).arg(&self.tau_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.grad_w_in_d).arg(&mut self.grad_w_rec_d)
|
||||
.arg(&mut self.grad_b_d).arg(&mut self.grad_tau_d)
|
||||
.arg(&mut self.cfc_grad_w_in_scratch_d)
|
||||
.arg(&mut self.cfc_grad_w_rec_scratch_d)
|
||||
.arg(&mut self.cfc_grad_b_scratch_d)
|
||||
.arg(&mut self.cfc_grad_tau_scratch_d)
|
||||
.arg(&mut self.grad_h_carry_d)
|
||||
.arg(&grad_henr_k_ptr);
|
||||
launch.launch(cfg_cfc_bwd).context("cfc bwd k batched")?;
|
||||
@@ -1918,6 +1973,44 @@ impl PerceptionTrainer {
|
||||
unsafe { launch.launch(cfg_vsn_bwd).context("variable_selection_bwd")?; }
|
||||
}
|
||||
|
||||
// ── 8c. (Phase B) Reduce CfC per-batch grad scratch → final grad buffers.
|
||||
// Block tree-reduce across bi; deterministic sum order.
|
||||
// Each launch is grid=(n_tail, 1, 1), block=(256, 1, 1).
|
||||
// 4 launches, one per cfc param tensor. AdamW (below) reads
|
||||
// the final grad buffers; MUST run after these reducers.
|
||||
{
|
||||
let n_batch_i = b_sz as i32;
|
||||
let reduce_block: u32 = 256;
|
||||
let reduce_at = |n_tail: usize,
|
||||
scratch: &CudaSlice<f32>,
|
||||
out: &mut CudaSlice<f32>,
|
||||
label: &'static str|
|
||||
-> Result<()> {
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (n_tail as u32, 1, 1),
|
||||
block_dim: (reduce_block, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let n_tail_i = n_tail as i32;
|
||||
let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn);
|
||||
launch
|
||||
.arg(scratch)
|
||||
.arg(&n_batch_i)
|
||||
.arg(&n_tail_i)
|
||||
.arg(out);
|
||||
unsafe { launch.launch(cfg).context(label)?; }
|
||||
Ok(())
|
||||
};
|
||||
reduce_at(HIDDEN_DIM * HIDDEN_DIM, &self.cfc_grad_w_in_scratch_d,
|
||||
&mut self.grad_w_in_d, "reduce cfc_grad_w_in")?;
|
||||
reduce_at(HIDDEN_DIM * HIDDEN_DIM, &self.cfc_grad_w_rec_scratch_d,
|
||||
&mut self.grad_w_rec_d, "reduce cfc_grad_w_rec")?;
|
||||
reduce_at(HIDDEN_DIM, &self.cfc_grad_b_scratch_d,
|
||||
&mut self.grad_b_d, "reduce cfc_grad_b")?;
|
||||
reduce_at(HIDDEN_DIM, &self.cfc_grad_tau_scratch_d,
|
||||
&mut self.grad_tau_d, "reduce cfc_grad_tau")?;
|
||||
}
|
||||
|
||||
// ── 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)?;
|
||||
|
||||
@@ -186,6 +186,103 @@ fn stacked_trainer_loss_shrinks_with_stride_4() {
|
||||
);
|
||||
}
|
||||
|
||||
/// Verifies the trainer converges at n_batch=32 — the FIRST test that
|
||||
/// exercises the cross-batch reducer code path. Existing
|
||||
/// `stacked_trainer_loss_shrinks_*` tests all use n_batch=1 so the new
|
||||
/// per-batch scratch + reducer logic was never previously hit.
|
||||
#[test]
|
||||
fn stacked_trainer_loss_shrinks_at_batch_32() {
|
||||
let dev = test_device();
|
||||
let cfg = PerceptionTrainerConfig {
|
||||
seq_len: 16,
|
||||
mamba2_state_dim: 8,
|
||||
lr_cfc: 3e-3,
|
||||
lr_mamba2: 1e-3,
|
||||
seed: 0xB32B,
|
||||
horizon_weights: [1.0; 5],
|
||||
n_batch: 32,
|
||||
decision_stride: 1,
|
||||
};
|
||||
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
|
||||
|
||||
let mut ts_base = 1_000_000u64;
|
||||
let mut prev_mid = 5500.0_f32;
|
||||
|
||||
let make_batch = |prev_mid: &mut f32, ts_base: &mut u64,
|
||||
cfg: &PerceptionTrainerConfig|
|
||||
-> (Vec<Vec<Mbp10RawInput>>, Vec<Vec<[f32; 5]>>) {
|
||||
let mut seqs: Vec<Vec<Mbp10RawInput>> = Vec::with_capacity(cfg.n_batch);
|
||||
let mut labels: Vec<Vec<[f32; 5]>> = Vec::with_capacity(cfg.n_batch);
|
||||
for _ in 0..cfg.n_batch {
|
||||
let (seq, lbl) = synthetic_seq(cfg.seq_len, *prev_mid, *ts_base);
|
||||
*prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
|
||||
*ts_base = seq.last().unwrap().ts_ns;
|
||||
seqs.push(seq);
|
||||
labels.push(lbl);
|
||||
}
|
||||
(seqs, labels)
|
||||
};
|
||||
|
||||
let mut initial = 0.0_f32;
|
||||
for warmup in 0..4 {
|
||||
let (seqs, labels) = make_batch(&mut prev_mid, &mut ts_base, &cfg);
|
||||
let seq_refs: Vec<&[Mbp10RawInput]> = seqs.iter().map(|s| s.as_slice()).collect();
|
||||
let lbl_refs: Vec<&[[f32; 5]]> = labels.iter().map(|l| l.as_slice()).collect();
|
||||
let l = trainer.step_batched(&seq_refs, &lbl_refs).expect("warm step");
|
||||
if warmup >= 2 {
|
||||
initial += l;
|
||||
}
|
||||
}
|
||||
initial /= 2.0;
|
||||
eprintln!("B=32 trainer: initial={initial:.4}");
|
||||
|
||||
for _ in 0..200 {
|
||||
let (seqs, labels) = make_batch(&mut prev_mid, &mut ts_base, &cfg);
|
||||
let seq_refs: Vec<&[Mbp10RawInput]> = seqs.iter().map(|s| s.as_slice()).collect();
|
||||
let lbl_refs: Vec<&[[f32; 5]]> = labels.iter().map(|l| l.as_slice()).collect();
|
||||
trainer.step_batched(&seq_refs, &lbl_refs).expect("train step");
|
||||
}
|
||||
|
||||
let mut final_loss = 0.0_f32;
|
||||
for _ in 0..4 {
|
||||
let (seqs, labels) = make_batch(&mut prev_mid, &mut ts_base, &cfg);
|
||||
let seq_refs: Vec<&[Mbp10RawInput]> = seqs.iter().map(|s| s.as_slice()).collect();
|
||||
let lbl_refs: Vec<&[[f32; 5]]> = labels.iter().map(|l| l.as_slice()).collect();
|
||||
final_loss += trainer.step_batched(&seq_refs, &lbl_refs).expect("final step");
|
||||
}
|
||||
final_loss /= 4.0;
|
||||
eprintln!("B=32 trainer: final={final_loss:.4}");
|
||||
assert!(
|
||||
final_loss < 0.6 * initial || final_loss < 0.1,
|
||||
"B=32 trainer failed to converge: {initial:.4} → {final_loss:.4}"
|
||||
);
|
||||
}
|
||||
|
||||
// NOTE on scratch-clears testing:
|
||||
//
|
||||
// A direct "scratch is zero between steps" test is structurally hard to
|
||||
// write against this trainer. The lifecycle is:
|
||||
// step 1: uncaptured warmup execute → scratch ends with grad
|
||||
// step 2: begin_capture/end_capture → records kernels but does
|
||||
// NOT execute them, scratch
|
||||
// unchanged from step 1
|
||||
// step 3+: graph.launch_replay → executes the captured graph
|
||||
//
|
||||
// A post-step snapshot read on step 2 returns step 1's residual,
|
||||
// which is bit-identical between any two captures regardless of input
|
||||
// data. Two replay steps would work but require either three+ trainer
|
||||
// calls or capture-internal scratch readback, both of which require
|
||||
// substantial test infrastructure.
|
||||
//
|
||||
// Instead we rely on the `stacked_trainer_loss_shrinks_at_batch_32`
|
||||
// smoke as the implicit scratch-zero validation: if the K-loop scratch
|
||||
// wasn't being zeroed at step start, gradients from step N would
|
||||
// pollute step N+1, training would diverge after a handful of steps,
|
||||
// and the convergence assertion would fail. The B=32 smoke explicitly
|
||||
// runs 200 training steps + measures final loss, so the absence of
|
||||
// divergence IS the scratch-zero guarantee.
|
||||
|
||||
|
||||
/// Eval alone must work — proves the eval path is fine WITHOUT any
|
||||
/// prior captured-graph training. If this passes but
|
||||
/// `evaluate_works_after_captured_training_step` fails, the bug is in
|
||||
|
||||
Reference in New Issue
Block a user