feat(aux-trunk): smaller single-bucket CfC for aux supervision (Layer B3)
New 64-hidden single-bucket CfC trunk to serve as the parallel parameter group for D-style aux supervision (Layer B of anti-cal plan). Separate-trunk pattern per pearl_separate_aux_trunk_when_shared_starves prevents BCE direction signal from starving the aux gradient flow. Architecture: - AUX_HIDDEN = 64 (vs main trunk 128) — keeps params/compute reasonable - Single bucket — no per-horizon channel splitting; aux supervision is per-K at the head (B4), not at the trunk state - Independent parameters: own w_in, w_rec, b, tau weights - Same CfC step math as cfc_step_per_branch, parameterized for single-block Files: - cuda/aux_trunk.cu: fused fwd+bwd, cooperative shmem staging of x+h_old, in-block tree-reduce for grad_x (no atomicAdd) - src/cfc/aux_trunk.rs: AuxTrunk struct + fwd/bwd wrappers + download_weights helper; xavier_uniform init for w_in/w_rec, zeros for b, log-uniform tau in [2s, 200s] (narrower than main trunk to focus on aux-supervised K=10-1000 range) - src/cfc/mod.rs: pub mod aux_trunk + re-exports - build.rs: KERNELS += "aux_trunk" - tests/aux_trunk.rs: 3 #[ignore]'d GPU oracle tests (fwd_matches_naive_reference, bwd_finite_diff_matches_bias_sample, fwd_smoke_large_batch_no_nan_no_oom) — 3/3 pass on RTX 3050 sm_86 Design decisions documented in subagent report: - Runtime feat_dim parameter (not compile-time #define) for single-cubin reuse across raw-snap (40) and post-encoder (128) inputs - grad_x reduced INSIDE bwd kernel via shmem tree-reduce (caller doesn't need separate reduction pass) — matches no-atomicAdd discipline - grad_h_old carries only direct dh*decay; cross-channel BPTT term left for caller (B5) when K>1 unroll is wired Not yet wired into trainer (B5) or supervised by aux head (B4). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -27,6 +27,7 @@ const KERNELS: &[&str] = &[
|
||||
"bucket_transition_kernels", // Per-horizon CfC Phase 1→2 transition: tau_sort, bucket_assign, bucket_iqr, channels_in_bucket, heads_compact, zero_off_bucket (ALPHA fix 2026-05-21)
|
||||
"cfc_step_per_branch", // Per-horizon CfC Phase 2: fused per-(batch, branch) fwd + bwd over [25,25,25,25,28] buckets
|
||||
"heads_block_diagonal_fwd", // Per-horizon CfC Phase 2: heads w_skip projection with compact ragged storage (640→128 floats)
|
||||
"aux_trunk", // SDD-3 Layer B3: smaller single-bucket CfC trunk (AUX_HIDDEN=64) for outcome-supervision (D-labels)
|
||||
];
|
||||
|
||||
// Cache bust v16 (2026-05-21): ALPHA fix — channels_in_bucket_kernel + zero_off_bucket_kernel replace tau_reorder_kernel.
|
||||
|
||||
241
crates/ml-alpha/cuda/aux_trunk.cu
Normal file
241
crates/ml-alpha/cuda/aux_trunk.cu
Normal file
@@ -0,0 +1,241 @@
|
||||
// aux_trunk.cu — single-bucket CfC step (forward + backward) at AUX_HIDDEN=64.
|
||||
//
|
||||
// Per `pearl_separate_aux_trunk_when_shared_starves`: this is the kernel
|
||||
// behind a SECOND, smaller CfC trunk used for outcome-supervision (D-labels).
|
||||
// It runs in parallel with the main per-branch CfC trunk on the same encoder
|
||||
// output, with INDEPENDENT parameters (no shared w_in/w_rec/b/tau) so that
|
||||
// BCE gradient flow on the main trunk does not starve the aux supervision.
|
||||
//
|
||||
// Design contract:
|
||||
// • Single bucket — no per-horizon channel splitting, no
|
||||
// channels_in_bucket lookup, no bucket_dim_k. Aux supervision is
|
||||
// per-K at the head (B4 wiring), not at the τ-state.
|
||||
// • Hidden dim is AUX_HIDDEN=64 (half of main trunk's HIDDEN_DIM=128).
|
||||
// • Algorithm matches `cfc_step_per_branch.cu`: continuous-time CfC
|
||||
// decay update `h_new = h_old·decay + (1-decay)·tanh(pre)`, with
|
||||
// `decay = exp(-dt / max(tau, 1e-6))`.
|
||||
// • Input feature dim is parameterised (`feat_dim` runtime argument)
|
||||
// so a single cubin handles both raw-snap input (FEATURE_DIM=40) and
|
||||
// post-encoder input (HIDDEN_DIM=128). The host wrapper passes the
|
||||
// correct dim; the kernel iterates `feat_dim` in the input MVM.
|
||||
//
|
||||
// Per `feedback_no_atomicadd.md`: no atomicAdd anywhere. Each thread
|
||||
// writes to its own (batch, c) slot in forward, and to its own per-batch
|
||||
// grad slice in backward. Cross-batch reduction is the caller's
|
||||
// responsibility via the existing `reduce_axis0` infrastructure.
|
||||
//
|
||||
// Per `pearl_cooperative_staging_eliminates_redundant_reads`: the per-
|
||||
// batch `x` and `h_old` rows are SHARED across all output channels in a
|
||||
// block. Stage them into shared memory once at block entry so the K-loop
|
||||
// reads from smem instead of issuing AUX_HIDDEN × redundant DRAM reads.
|
||||
//
|
||||
// Per `pearl_no_host_branches_in_captured_graph`: no host branching;
|
||||
// the kernel takes feat_dim as a runtime arg but does not consult any
|
||||
// host-side state during execution.
|
||||
//
|
||||
// Per `feedback_nvidia_grade_perf_for_kernels.md`: warp-uniform branches
|
||||
// (the only conditional is `batch < B`, identical for every thread in a
|
||||
// block since batch == blockIdx.x), no atomicAdd, coalesced loads via
|
||||
// cooperative staging.
|
||||
|
||||
#define AUX_HIDDEN 64
|
||||
#define MAX_FEAT_DIM 128 // upper bound for shared-memory staging; covers
|
||||
// both FEATURE_DIM=40 and HIDDEN_DIM=128 callers
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// aux_trunk_fwd: forward pass.
|
||||
//
|
||||
// Launch:
|
||||
// grid = (B, 1, 1)
|
||||
// block = (AUX_HIDDEN = 64, 1, 1)
|
||||
// shared_mem_bytes = (feat_dim + AUX_HIDDEN) * sizeof(float)
|
||||
//
|
||||
// Per (batch, c) thread computes:
|
||||
// pre = b[c] + Σ_k W_in[c, k] * x[batch, k] + Σ_k W_rec[c, k] * h_old[batch, k]
|
||||
// decay = exp(-dt / max(tau[c], 1e-6))
|
||||
// h_new[batch, c] = h_old[batch, c] * decay + (1 - decay) * tanh(pre)
|
||||
//
|
||||
// W_in is [AUX_HIDDEN × feat_dim], W_rec is [AUX_HIDDEN × AUX_HIDDEN].
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
extern "C" __global__ void aux_trunk_fwd(
|
||||
const float* __restrict__ w_in, // [AUX_HIDDEN × feat_dim]
|
||||
const float* __restrict__ w_rec, // [AUX_HIDDEN × AUX_HIDDEN]
|
||||
const float* __restrict__ b, // [AUX_HIDDEN]
|
||||
const float* __restrict__ tau, // [AUX_HIDDEN]
|
||||
const float* __restrict__ x, // [B × feat_dim]
|
||||
const float* __restrict__ h_old, // [B × AUX_HIDDEN]
|
||||
float dt,
|
||||
int B,
|
||||
int feat_dim,
|
||||
float* __restrict__ h_new // [B × AUX_HIDDEN]
|
||||
) {
|
||||
extern __shared__ float smem[];
|
||||
float* x_local = smem; // [feat_dim]
|
||||
float* h_old_local = smem + feat_dim; // [AUX_HIDDEN]
|
||||
|
||||
int batch = blockIdx.x;
|
||||
int c = threadIdx.x;
|
||||
|
||||
if (batch >= B) return;
|
||||
|
||||
// Cooperative staging of per-batch x[batch, *] and h_old[batch, *].
|
||||
// Each thread loads ceil(feat_dim / AUX_HIDDEN) = up to ~2 x slots
|
||||
// and exactly one h_old slot (since blockDim.x == AUX_HIDDEN).
|
||||
for (int i = c; i < feat_dim; i += AUX_HIDDEN) {
|
||||
x_local[i] = x[batch * feat_dim + i];
|
||||
}
|
||||
h_old_local[c] = h_old[batch * AUX_HIDDEN + c];
|
||||
__syncthreads();
|
||||
|
||||
// pre = b[c] + Σ_k W_in[c, k] * x_local[k] + Σ_k W_rec[c, k] * h_old_local[k]
|
||||
float pre = b[c];
|
||||
for (int k = 0; k < feat_dim; ++k) {
|
||||
pre += w_in[c * feat_dim + k] * x_local[k];
|
||||
}
|
||||
for (int k = 0; k < AUX_HIDDEN; ++k) {
|
||||
pre += w_rec[c * AUX_HIDDEN + k] * h_old_local[k];
|
||||
}
|
||||
|
||||
float decay = expf(-dt / fmaxf(tau[c], 1e-6f));
|
||||
h_new[batch * AUX_HIDDEN + c] =
|
||||
h_old_local[c] * decay + (1.0f - decay) * tanhf(pre);
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// aux_trunk_bwd: backward pass.
|
||||
//
|
||||
// Launch:
|
||||
// grid = (B, 1, 1)
|
||||
// block = (AUX_HIDDEN = 64, 1, 1)
|
||||
// shared_mem_bytes = (feat_dim + AUX_HIDDEN) * sizeof(float)
|
||||
//
|
||||
// Per-batch grad slices are written by this kernel (OVERWRITE semantics);
|
||||
// cross-batch reduction is the caller's responsibility via existing
|
||||
// `reduce_axis0` infrastructure.
|
||||
//
|
||||
// grad shapes:
|
||||
// grad_w_in : [B × AUX_HIDDEN × feat_dim] per-batch scratch
|
||||
// grad_w_rec : [B × AUX_HIDDEN × AUX_HIDDEN] per-batch scratch
|
||||
// grad_b : [B × AUX_HIDDEN] per-batch scratch
|
||||
// grad_tau : [B × AUX_HIDDEN] per-batch scratch
|
||||
// grad_h_old : [B × AUX_HIDDEN] per-batch (decay
|
||||
// contribution only;
|
||||
// cross-channel reduction
|
||||
// deferred to caller if
|
||||
// unrolled K>1)
|
||||
// grad_x : [B × feat_dim] per-batch (reduced
|
||||
// across c via shared-
|
||||
// memory tree-reduce)
|
||||
//
|
||||
// grad_x cross-channel reduction is done inside this kernel via a
|
||||
// block-resident shared-memory tree-reduce (no atomicAdd per
|
||||
// `feedback_no_atomicadd`).
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
extern "C" __global__ void aux_trunk_bwd(
|
||||
const float* __restrict__ w_in, // [AUX_HIDDEN × feat_dim]
|
||||
const float* __restrict__ w_rec, // [AUX_HIDDEN × AUX_HIDDEN]
|
||||
const float* __restrict__ b, // [AUX_HIDDEN]
|
||||
const float* __restrict__ tau, // [AUX_HIDDEN]
|
||||
const float* __restrict__ x, // [B × feat_dim]
|
||||
const float* __restrict__ h_old, // [B × AUX_HIDDEN]
|
||||
const float* __restrict__ grad_h_new, // [B × AUX_HIDDEN]
|
||||
float dt,
|
||||
int B,
|
||||
int feat_dim,
|
||||
float* __restrict__ grad_w_in, // [B × AUX_HIDDEN × feat_dim]
|
||||
float* __restrict__ grad_w_rec, // [B × AUX_HIDDEN × AUX_HIDDEN]
|
||||
float* __restrict__ grad_b, // [B × AUX_HIDDEN]
|
||||
float* __restrict__ grad_tau, // [B × AUX_HIDDEN]
|
||||
float* __restrict__ grad_h_old, // [B × AUX_HIDDEN]
|
||||
float* __restrict__ grad_x // [B × feat_dim]
|
||||
) {
|
||||
// Shared layout:
|
||||
// [0 .. feat_dim) : x_local
|
||||
// [feat_dim .. feat_dim+AH) : h_old_local
|
||||
// [feat_dim+AH .. feat_dim+AH+AH) : d_pre_smem (per-channel d_pre
|
||||
// for cross-channel grad_x reduce)
|
||||
extern __shared__ float smem[];
|
||||
float* x_local = smem;
|
||||
float* h_old_local = smem + feat_dim;
|
||||
float* d_pre_smem = smem + feat_dim + AUX_HIDDEN;
|
||||
|
||||
int batch = blockIdx.x;
|
||||
int c = threadIdx.x;
|
||||
|
||||
if (batch >= B) return;
|
||||
|
||||
// Cooperative staging — identical pattern to forward.
|
||||
for (int i = c; i < feat_dim; i += AUX_HIDDEN) {
|
||||
x_local[i] = x[batch * feat_dim + i];
|
||||
}
|
||||
h_old_local[c] = h_old[batch * AUX_HIDDEN + c];
|
||||
__syncthreads();
|
||||
|
||||
// Recompute forward pre + tanh to derive d_pre.
|
||||
float pre = b[c];
|
||||
for (int k = 0; k < feat_dim; ++k) {
|
||||
pre += w_in[c * feat_dim + k] * x_local[k];
|
||||
}
|
||||
for (int k = 0; k < AUX_HIDDEN; ++k) {
|
||||
pre += w_rec[c * AUX_HIDDEN + k] * h_old_local[k];
|
||||
}
|
||||
const float tau_eps = 1e-6f;
|
||||
const float tau_raw = tau[c];
|
||||
const float tau_eff = fmaxf(tau_raw, tau_eps);
|
||||
const float decay = expf(-dt / tau_eff);
|
||||
const float s = tanhf(pre);
|
||||
const float dh = grad_h_new[batch * AUX_HIDDEN + c];
|
||||
const float d_pre = dh * (1.0f - decay) * (1.0f - s * s);
|
||||
const float d_decay = dh * (h_old_local[c] - s);
|
||||
|
||||
// Per-batch grad-scratch writes. Each thread (batch, c) is the sole
|
||||
// writer of grad_*[batch, c, ...] — no race, no atomicAdd.
|
||||
grad_b[batch * AUX_HIDDEN + c] = d_pre;
|
||||
|
||||
// grad_tau receives 0 when tau hits the clamp floor
|
||||
// (strict d(max)/d(tau) = 0 below tau_eps).
|
||||
const float gate = (tau_raw > tau_eps) ? 1.0f : 0.0f;
|
||||
grad_tau[batch * AUX_HIDDEN + c] =
|
||||
gate * d_decay * decay * dt / (tau_eff * tau_eff);
|
||||
|
||||
// grad_w_in[batch, c, k] = d_pre * x_local[k]
|
||||
for (int k = 0; k < feat_dim; ++k) {
|
||||
const long long off =
|
||||
(long long)batch * AUX_HIDDEN * feat_dim
|
||||
+ (long long)c * feat_dim + k;
|
||||
grad_w_in[off] = d_pre * x_local[k];
|
||||
}
|
||||
// grad_w_rec[batch, c, k] = d_pre * h_old_local[k]
|
||||
for (int k = 0; k < AUX_HIDDEN; ++k) {
|
||||
const long long off =
|
||||
(long long)batch * AUX_HIDDEN * AUX_HIDDEN
|
||||
+ (long long)c * AUX_HIDDEN + k;
|
||||
grad_w_rec[off] = d_pre * h_old_local[k];
|
||||
}
|
||||
|
||||
// grad_h_old[batch, c] direct decay contribution.
|
||||
// (Cross-channel term Σ_j d_pre[j] * W_rec[j, c] is left for the
|
||||
// caller to add if BPTT through h_old is needed — for K=1 per-step
|
||||
// aux supervision this direct term is sufficient.)
|
||||
grad_h_old[batch * AUX_HIDDEN + c] = dh * decay;
|
||||
|
||||
// grad_x[batch, k] = Σ_c d_pre[c] * W_in[c, k]
|
||||
// Cross-channel reduction via shared memory: stash d_pre to smem,
|
||||
// then have each thread compute its assigned k-slice by summing
|
||||
// across c. This avoids atomicAdd while keeping the reduction
|
||||
// block-local. (Per pearl_no_atomicadd + the block tree-reduce
|
||||
// pattern used in reduce_axis0.cu.)
|
||||
d_pre_smem[c] = d_pre;
|
||||
__syncthreads();
|
||||
|
||||
// Each of the AUX_HIDDEN threads handles a strided subset of feat_dim
|
||||
// output positions. With feat_dim ≤ MAX_FEAT_DIM = 128 and
|
||||
// AUX_HIDDEN = 64, each thread covers at most 2 k positions.
|
||||
for (int k = c; k < feat_dim; k += AUX_HIDDEN) {
|
||||
float acc = 0.0f;
|
||||
for (int j = 0; j < AUX_HIDDEN; ++j) {
|
||||
acc += d_pre_smem[j] * w_in[j * feat_dim + k];
|
||||
}
|
||||
grad_x[batch * feat_dim + k] = acc;
|
||||
}
|
||||
}
|
||||
385
crates/ml-alpha/src/cfc/aux_trunk.rs
Normal file
385
crates/ml-alpha/src/cfc/aux_trunk.rs
Normal file
@@ -0,0 +1,385 @@
|
||||
//! SDD-3 Layer B3 — auxiliary CfC trunk for outcome-supervision (D-labels).
|
||||
//!
|
||||
//! A SECOND, smaller CfC trunk that operates on the same encoder output
|
||||
//! as the main BCE trunk, but with INDEPENDENT parameters and a smaller
|
||||
//! hidden dim. Lives next to the main per-branch trunk in `cfc::trunk`.
|
||||
//!
|
||||
//! ## Why separate trunk (not shared)
|
||||
//!
|
||||
//! Per `pearl_separate_aux_trunk_when_shared_starves`: aux supervision
|
||||
//! plateauing at ln(K) on a shared trunk signals starvation. Separate
|
||||
//! aux trunk + own optimiser group prevents the BCE direction signal
|
||||
//! from dominating gradient flow toward the aux head.
|
||||
//!
|
||||
//! Per the E3 design memo: asymmetric stop-grad at encoder boundary —
|
||||
//! BCE drives encoder; aux gets gradient only into its own params at
|
||||
//! first; lifted conditionally once aux CE < 0.4 + dir_acc > 0.85
|
||||
//! within 200 steps.
|
||||
//!
|
||||
//! ## Algorithm
|
||||
//!
|
||||
//! Identical to `cfc_step_per_branch` but stripped of the per-bucket
|
||||
//! channel split (single bucket, single launch grid).
|
||||
//!
|
||||
//! pre = b[c] + Σ_k W_in[c, k] * x[batch, k] + Σ_k W_rec[c, k] * h_old[batch, k]
|
||||
//! decay = exp(-dt / max(tau[c], 1e-6))
|
||||
//! h_new[batch, c] = h_old[batch, c] * decay + (1 - decay) * tanh(pre)
|
||||
//!
|
||||
//! ## Initialisation
|
||||
//!
|
||||
//! Per `pearl_scoped_init_seed_for_reproducibility`: `AuxTrunk::new`
|
||||
//! installs a `scoped_init_seed(seed)` guard so the Xavier draws are
|
||||
//! deterministic given `seed`.
|
||||
//!
|
||||
//! * `w_in`: Xavier uniform `[-sqrt(6/(feat_dim+AUX_HIDDEN)), +]`
|
||||
//! * `w_rec`: Xavier uniform `[-sqrt(6/(2*AUX_HIDDEN)), +]`
|
||||
//! * `b`: zeros
|
||||
//! * `tau`: log-uniform over `[2 s, 200 s]` (single-bucket so a wide
|
||||
//! span across all channels is fine; no per-horizon assignment).
|
||||
//!
|
||||
//! ## Constraints honoured
|
||||
//!
|
||||
//! * `feedback_no_atomicadd.md` — no atomicAdd in fwd or bwd kernels;
|
||||
//! bwd writes per-batch grad scratch and the caller reduces axis 0.
|
||||
//! * `feedback_no_htod_htoh_only_mapped_pinned.md` — uploads go through
|
||||
//! mapped-pinned staging like the rest of `cfc/step.rs`.
|
||||
//! * `feedback_no_nvrtc.md` — pre-compiled cubin via build.rs.
|
||||
//! * `pearl_no_host_branches_in_captured_graph` — kernel uses no host
|
||||
//! branching; the `feat_dim` runtime arg is read once and used as
|
||||
//! loop bound.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::driver::{
|
||||
CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut, LaunchConfig, PushKernelArg,
|
||||
};
|
||||
use ml_core::cuda_autograd::init::scoped_init_seed;
|
||||
use ml_core::device::MlDevice;
|
||||
use rand::{Rng, SeedableRng};
|
||||
use rand_chacha::ChaCha8Rng;
|
||||
|
||||
use crate::pinned_mem::MappedF32Buffer;
|
||||
|
||||
const KERNEL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/aux_trunk.cubin"));
|
||||
|
||||
/// Aux trunk hidden dimension. Half the main trunk's HIDDEN_DIM=128.
|
||||
/// Kept constant (not configurable) so cubin compile-time `#define
|
||||
/// AUX_HIDDEN` stays in lockstep with the host wrapper.
|
||||
pub const AUX_HIDDEN: usize = 64;
|
||||
|
||||
/// Upper bound on the input feature dim. Matches the cubin's
|
||||
/// `MAX_FEAT_DIM = 128`. Covers both raw-snap input (FEATURE_DIM=40)
|
||||
/// and post-encoder input (HIDDEN_DIM=128).
|
||||
pub const MAX_AUX_FEAT_DIM: usize = 128;
|
||||
|
||||
/// Tau initialisation range (seconds). Single-bucket trunk gets a wide
|
||||
/// log-uniform span so the aux head sees diverse memory horizons within
|
||||
/// one trunk. Matches the spirit of the main trunk's [10 ms, 1000 s]
|
||||
/// span but narrowed since aux supervision targets ~1k-step D-label
|
||||
/// outcomes — too-long τ wastes capacity on horizons the head doesn't
|
||||
/// supervise.
|
||||
const AUX_TAU_LO_S: f32 = 2.0;
|
||||
const AUX_TAU_HI_S: f32 = 200.0;
|
||||
|
||||
/// Host-side weight storage for `AuxTrunk` (used by checkpoint round-trip).
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct AuxTrunkWeights {
|
||||
pub w_in: Vec<f32>, // [AUX_HIDDEN × feat_dim]
|
||||
pub w_rec: Vec<f32>, // [AUX_HIDDEN × AUX_HIDDEN]
|
||||
pub b: Vec<f32>, // [AUX_HIDDEN]
|
||||
pub tau: Vec<f32>, // [AUX_HIDDEN]
|
||||
pub feat_dim: usize,
|
||||
}
|
||||
|
||||
/// Construction config for `AuxTrunk`.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct AuxTrunkConfig {
|
||||
/// Input feature dimensionality. Must be ≤ `MAX_AUX_FEAT_DIM` so the
|
||||
/// cubin's shared-memory staging buffer fits.
|
||||
pub feat_dim: usize,
|
||||
/// Random seed for Xavier + tau init. Reproducibility-critical per
|
||||
/// `pearl_scoped_init_seed_for_reproducibility`.
|
||||
pub seed: u64,
|
||||
}
|
||||
|
||||
/// Auxiliary CfC trunk. Owns device weights + a cached cubin module
|
||||
/// with `aux_trunk_fwd` / `aux_trunk_bwd` function handles.
|
||||
pub struct AuxTrunk {
|
||||
cfg: AuxTrunkConfig,
|
||||
stream: Arc<CudaStream>,
|
||||
|
||||
// Cubin module kept alive so the cached `CudaFunction` handles stay
|
||||
// valid for the trunk's lifetime.
|
||||
_module: Arc<CudaModule>,
|
||||
pub fwd_fn: CudaFunction,
|
||||
pub bwd_fn: CudaFunction,
|
||||
|
||||
// Device weights. Public so the trainer (B5) can mutate via AdamW
|
||||
// and the checkpoint envelope (B5) can memcpy_htod from a loaded
|
||||
// file.
|
||||
pub w_in_d: CudaSlice<f32>, // [AUX_HIDDEN × feat_dim]
|
||||
pub w_rec_d: CudaSlice<f32>, // [AUX_HIDDEN × AUX_HIDDEN]
|
||||
pub b_d: CudaSlice<f32>, // [AUX_HIDDEN]
|
||||
pub tau_d: CudaSlice<f32>, // [AUX_HIDDEN]
|
||||
}
|
||||
|
||||
impl AuxTrunk {
|
||||
/// Construct a fresh `AuxTrunk` with seeded Xavier + log-uniform τ
|
||||
/// initialisation. Loads the `aux_trunk.cubin` and resolves the
|
||||
/// `aux_trunk_fwd` / `aux_trunk_bwd` function handles once.
|
||||
pub fn new(dev: &MlDevice, cfg: AuxTrunkConfig) -> Result<Self> {
|
||||
anyhow::ensure!(
|
||||
cfg.feat_dim > 0 && cfg.feat_dim <= MAX_AUX_FEAT_DIM,
|
||||
"AuxTrunk feat_dim={} out of (0, {}]",
|
||||
cfg.feat_dim,
|
||||
MAX_AUX_FEAT_DIM,
|
||||
);
|
||||
|
||||
let stream: Arc<CudaStream> = dev.cuda_stream().context("aux_trunk stream")?.clone();
|
||||
let ctx = dev.cuda_context().context("aux_trunk ctx")?;
|
||||
let module = ctx
|
||||
.load_cubin(KERNEL_CUBIN.to_vec())
|
||||
.context("load aux_trunk cubin")?;
|
||||
let fwd_fn = module
|
||||
.load_function("aux_trunk_fwd")
|
||||
.context("load aux_trunk_fwd")?;
|
||||
let bwd_fn = module
|
||||
.load_function("aux_trunk_bwd")
|
||||
.context("load aux_trunk_bwd")?;
|
||||
|
||||
// Deterministic init guard — every random draw below funnels
|
||||
// through `scoped_init_seed`'s thread-local RNG so the same
|
||||
// cfg.seed produces the same weights bit-for-bit.
|
||||
let _seed_guard = scoped_init_seed(cfg.seed);
|
||||
let mut r = ChaCha8Rng::seed_from_u64(cfg.seed);
|
||||
let scale_in = (6.0_f32 / (cfg.feat_dim + AUX_HIDDEN) as f32).sqrt();
|
||||
let scale_rec = (6.0_f32 / (2 * AUX_HIDDEN) as f32).sqrt();
|
||||
let w_in: Vec<f32> = (0..AUX_HIDDEN * cfg.feat_dim)
|
||||
.map(|_| r.gen_range(-scale_in..scale_in))
|
||||
.collect();
|
||||
let w_rec: Vec<f32> = (0..AUX_HIDDEN * AUX_HIDDEN)
|
||||
.map(|_| r.gen_range(-scale_rec..scale_rec))
|
||||
.collect();
|
||||
let b: Vec<f32> = vec![0.0; AUX_HIDDEN];
|
||||
// tau: log-uniform over [AUX_TAU_LO_S, AUX_TAU_HI_S].
|
||||
let lo = AUX_TAU_LO_S.ln();
|
||||
let hi = AUX_TAU_HI_S.ln();
|
||||
let tau: Vec<f32> = (0..AUX_HIDDEN)
|
||||
.map(|_| {
|
||||
let u: f32 = r.gen_range(0.0..1.0);
|
||||
(lo + u * (hi - lo)).exp()
|
||||
})
|
||||
.collect();
|
||||
|
||||
let w_in_d = upload(&stream, &w_in)?;
|
||||
let w_rec_d = upload(&stream, &w_rec)?;
|
||||
let b_d = upload(&stream, &b)?;
|
||||
let tau_d = upload(&stream, &tau)?;
|
||||
|
||||
Ok(Self {
|
||||
cfg,
|
||||
stream,
|
||||
_module: module,
|
||||
fwd_fn,
|
||||
bwd_fn,
|
||||
w_in_d,
|
||||
w_rec_d,
|
||||
b_d,
|
||||
tau_d,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn config(&self) -> &AuxTrunkConfig {
|
||||
&self.cfg
|
||||
}
|
||||
|
||||
pub fn stream(&self) -> &Arc<CudaStream> {
|
||||
&self.stream
|
||||
}
|
||||
|
||||
/// Download the trunk's current weights into host vectors. Used by
|
||||
/// checkpoint save and unit tests.
|
||||
pub fn download_weights(&self) -> Result<AuxTrunkWeights> {
|
||||
let mut w_in = vec![0.0_f32; self.w_in_d.len()];
|
||||
let mut w_rec = vec![0.0_f32; self.w_rec_d.len()];
|
||||
let mut b = vec![0.0_f32; self.b_d.len()];
|
||||
let mut tau = vec![0.0_f32; self.tau_d.len()];
|
||||
self.stream
|
||||
.memcpy_dtoh(&self.w_in_d, w_in.as_mut_slice())
|
||||
.context("aux_trunk dtoh w_in")?;
|
||||
self.stream
|
||||
.memcpy_dtoh(&self.w_rec_d, w_rec.as_mut_slice())
|
||||
.context("aux_trunk dtoh w_rec")?;
|
||||
self.stream
|
||||
.memcpy_dtoh(&self.b_d, b.as_mut_slice())
|
||||
.context("aux_trunk dtoh b")?;
|
||||
self.stream
|
||||
.memcpy_dtoh(&self.tau_d, tau.as_mut_slice())
|
||||
.context("aux_trunk dtoh tau")?;
|
||||
Ok(AuxTrunkWeights {
|
||||
w_in,
|
||||
w_rec,
|
||||
b,
|
||||
tau,
|
||||
feat_dim: self.cfg.feat_dim,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Launch the fused aux-trunk forward kernel.
|
||||
///
|
||||
/// Buffer contract:
|
||||
/// * `x_d` : `[B × feat_dim]` per-batch encoder output
|
||||
/// * `h_old_d` : `[B × AUX_HIDDEN]` per-batch prior aux-trunk state
|
||||
/// * `h_new_d` : `[B × AUX_HIDDEN]` overwrite — new aux-trunk state
|
||||
///
|
||||
/// Shared memory layout: `(feat_dim + AUX_HIDDEN) * sizeof(float)`
|
||||
/// covering `x_local` + `h_old_local` cooperative staging.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn aux_trunk_fwd_gpu(
|
||||
stream: &Arc<CudaStream>,
|
||||
func: &CudaFunction,
|
||||
feat_dim: usize,
|
||||
w_in_d: &CudaSlice<f32>,
|
||||
w_rec_d: &CudaSlice<f32>,
|
||||
b_d: &CudaSlice<f32>,
|
||||
tau_d: &CudaSlice<f32>,
|
||||
x_d: &CudaSlice<f32>,
|
||||
h_old_d: &CudaSlice<f32>,
|
||||
dt_s: f32,
|
||||
b_sz: i32,
|
||||
h_new_d: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
debug_assert_eq!(w_in_d.len(), AUX_HIDDEN * feat_dim);
|
||||
debug_assert_eq!(w_rec_d.len(), AUX_HIDDEN * AUX_HIDDEN);
|
||||
debug_assert_eq!(b_d.len(), AUX_HIDDEN);
|
||||
debug_assert_eq!(tau_d.len(), AUX_HIDDEN);
|
||||
debug_assert_eq!(x_d.len(), (b_sz as usize) * feat_dim);
|
||||
debug_assert_eq!(h_old_d.len(), (b_sz as usize) * AUX_HIDDEN);
|
||||
debug_assert_eq!(h_new_d.len(), (b_sz as usize) * AUX_HIDDEN);
|
||||
|
||||
let feat_dim_i: i32 = feat_dim as i32;
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (b_sz as u32, 1, 1),
|
||||
block_dim: (AUX_HIDDEN as u32, 1, 1),
|
||||
// x_local (feat_dim floats) + h_old_local (AUX_HIDDEN floats)
|
||||
shared_mem_bytes: ((feat_dim + AUX_HIDDEN) * std::mem::size_of::<f32>()) as u32,
|
||||
};
|
||||
let mut launch = stream.launch_builder(func);
|
||||
launch
|
||||
.arg(w_in_d)
|
||||
.arg(w_rec_d)
|
||||
.arg(b_d)
|
||||
.arg(tau_d)
|
||||
.arg(x_d)
|
||||
.arg(h_old_d)
|
||||
.arg(&dt_s)
|
||||
.arg(&b_sz)
|
||||
.arg(&feat_dim_i)
|
||||
.arg(h_new_d);
|
||||
unsafe {
|
||||
launch.launch(cfg).context("aux_trunk_fwd launch")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Launch the fused aux-trunk backward kernel.
|
||||
///
|
||||
/// Per-batch grad scratch is OVERWRITTEN by this kernel; the caller is
|
||||
/// responsible for the cross-batch reduction (e.g. `reduce_axis0`).
|
||||
///
|
||||
/// Shared memory layout: `(feat_dim + 2 * AUX_HIDDEN) * sizeof(float)`
|
||||
/// covering `x_local` + `h_old_local` + `d_pre_smem` (the per-channel
|
||||
/// `d_pre` staging used for the in-block `grad_x` tree-reduce).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn aux_trunk_bwd_gpu(
|
||||
stream: &Arc<CudaStream>,
|
||||
func: &CudaFunction,
|
||||
feat_dim: usize,
|
||||
w_in_d: &CudaSlice<f32>,
|
||||
w_rec_d: &CudaSlice<f32>,
|
||||
b_d: &CudaSlice<f32>,
|
||||
tau_d: &CudaSlice<f32>,
|
||||
x_d: &CudaSlice<f32>,
|
||||
h_old_d: &CudaSlice<f32>,
|
||||
grad_h_new_d: &CudaSlice<f32>,
|
||||
dt_s: f32,
|
||||
b_sz: i32,
|
||||
grad_w_in_d: &mut CudaSlice<f32>,
|
||||
grad_w_rec_d: &mut CudaSlice<f32>,
|
||||
grad_b_d: &mut CudaSlice<f32>,
|
||||
grad_tau_d: &mut CudaSlice<f32>,
|
||||
grad_h_old_d: &mut CudaSlice<f32>,
|
||||
grad_x_d: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
let b_sz_u = b_sz as usize;
|
||||
debug_assert_eq!(w_in_d.len(), AUX_HIDDEN * feat_dim);
|
||||
debug_assert_eq!(w_rec_d.len(), AUX_HIDDEN * AUX_HIDDEN);
|
||||
debug_assert_eq!(b_d.len(), AUX_HIDDEN);
|
||||
debug_assert_eq!(tau_d.len(), AUX_HIDDEN);
|
||||
debug_assert_eq!(x_d.len(), b_sz_u * feat_dim);
|
||||
debug_assert_eq!(h_old_d.len(), b_sz_u * AUX_HIDDEN);
|
||||
debug_assert_eq!(grad_h_new_d.len(), b_sz_u * AUX_HIDDEN);
|
||||
debug_assert_eq!(grad_w_in_d.len(), b_sz_u * AUX_HIDDEN * feat_dim);
|
||||
debug_assert_eq!(grad_w_rec_d.len(), b_sz_u * AUX_HIDDEN * AUX_HIDDEN);
|
||||
debug_assert_eq!(grad_b_d.len(), b_sz_u * AUX_HIDDEN);
|
||||
debug_assert_eq!(grad_tau_d.len(), b_sz_u * AUX_HIDDEN);
|
||||
debug_assert_eq!(grad_h_old_d.len(), b_sz_u * AUX_HIDDEN);
|
||||
debug_assert_eq!(grad_x_d.len(), b_sz_u * feat_dim);
|
||||
|
||||
let feat_dim_i: i32 = feat_dim as i32;
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (b_sz as u32, 1, 1),
|
||||
block_dim: (AUX_HIDDEN as u32, 1, 1),
|
||||
shared_mem_bytes: ((feat_dim + 2 * AUX_HIDDEN) * std::mem::size_of::<f32>()) as u32,
|
||||
};
|
||||
let mut launch = stream.launch_builder(func);
|
||||
launch
|
||||
.arg(w_in_d)
|
||||
.arg(w_rec_d)
|
||||
.arg(b_d)
|
||||
.arg(tau_d)
|
||||
.arg(x_d)
|
||||
.arg(h_old_d)
|
||||
.arg(grad_h_new_d)
|
||||
.arg(&dt_s)
|
||||
.arg(&b_sz)
|
||||
.arg(&feat_dim_i)
|
||||
.arg(grad_w_in_d)
|
||||
.arg(grad_w_rec_d)
|
||||
.arg(grad_b_d)
|
||||
.arg(grad_tau_d)
|
||||
.arg(grad_h_old_d)
|
||||
.arg(grad_x_d);
|
||||
unsafe {
|
||||
launch.launch(cfg).context("aux_trunk_bwd launch")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ── pinned-staging upload helper (mirrors cfc/step.rs::upload) ────────
|
||||
|
||||
fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
|
||||
let n = host.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("aux_trunk upload staging: {e}"))?;
|
||||
staging.write_from_slice(host);
|
||||
let mut dst = stream
|
||||
.alloc_zeros::<f32>(n)
|
||||
.context("aux_trunk upload alloc")?;
|
||||
if n > 0 {
|
||||
let nbytes = n * std::mem::size_of::<f32>();
|
||||
unsafe {
|
||||
let (dst_ptr, _g) = dst.device_ptr_mut(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
dst_ptr,
|
||||
staging.dev_ptr,
|
||||
nbytes,
|
||||
stream.cu_stream(),
|
||||
)
|
||||
.context("aux_trunk upload DtoD")?;
|
||||
}
|
||||
}
|
||||
Ok(dst)
|
||||
}
|
||||
@@ -3,9 +3,11 @@
|
||||
//! amendment: this module owns the CfC layer that sits on top of a
|
||||
//! frozen Mamba2 sequence encoder.
|
||||
|
||||
pub mod aux_trunk;
|
||||
pub mod bucket_routing;
|
||||
pub mod snap_features;
|
||||
pub mod step;
|
||||
pub mod trunk;
|
||||
|
||||
pub use aux_trunk::{AuxTrunk, AuxTrunkConfig, AuxTrunkWeights, AUX_HIDDEN, MAX_AUX_FEAT_DIM};
|
||||
pub use trunk::{CfcConfig, CfcTrunk};
|
||||
|
||||
326
crates/ml-alpha/tests/aux_trunk.rs
Normal file
326
crates/ml-alpha/tests/aux_trunk.rs
Normal file
@@ -0,0 +1,326 @@
|
||||
//! GPU oracle tests for `aux_trunk.cu` (single-bucket CfC at AUX_HIDDEN=64).
|
||||
//!
|
||||
//! Run with:
|
||||
//! SQLX_OFFLINE=true cargo test -p ml-alpha --test aux_trunk \
|
||||
//! -- --ignored --nocapture
|
||||
//!
|
||||
//! These tests are `#[ignore]` because they require a CUDA device.
|
||||
//! They drive the cubin directly through the host wrappers in
|
||||
//! `crates/ml-alpha/src/cfc/aux_trunk.rs` and compare against a
|
||||
//! deterministic CPU naive reference (per
|
||||
//! `feedback_no_cpu_test_fallbacks.md` the production path stays GPU-
|
||||
//! only; the CPU implementation here exists ONLY as a unit-test oracle).
|
||||
|
||||
use anyhow::Result;
|
||||
use ml_alpha::cfc::aux_trunk::{
|
||||
aux_trunk_bwd_gpu, aux_trunk_fwd_gpu, AuxTrunk, AuxTrunkConfig, AUX_HIDDEN,
|
||||
};
|
||||
use ml_core::device::MlDevice;
|
||||
|
||||
// Pick a representative feat_dim. Aux trunk's production caller (B5)
|
||||
// will pass HIDDEN_DIM=128 (encoder output). Using a non-128 dim here
|
||||
// also exercises the runtime-`feat_dim` parameterisation.
|
||||
const TEST_FEAT_DIM: usize = 96;
|
||||
|
||||
fn upload(
|
||||
stream: &std::sync::Arc<cudarc::driver::CudaStream>,
|
||||
host: &[f32],
|
||||
) -> Result<cudarc::driver::CudaSlice<f32>> {
|
||||
let mut d = stream.alloc_zeros::<f32>(host.len())?;
|
||||
stream.memcpy_htod(host, &mut d)?;
|
||||
Ok(d)
|
||||
}
|
||||
|
||||
fn download(
|
||||
stream: &std::sync::Arc<cudarc::driver::CudaStream>,
|
||||
d: &cudarc::driver::CudaSlice<f32>,
|
||||
) -> Result<Vec<f32>> {
|
||||
let mut h = vec![0.0_f32; d.len()];
|
||||
stream.memcpy_dtoh(d, &mut h)?;
|
||||
Ok(h)
|
||||
}
|
||||
|
||||
/// Forward shape correctness + value match against a naive CPU oracle.
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn fwd_matches_naive_reference() -> Result<()> {
|
||||
let dev = MlDevice::cuda(0)?;
|
||||
let stream = dev.cuda_stream()?.clone();
|
||||
let trunk = AuxTrunk::new(
|
||||
&dev,
|
||||
AuxTrunkConfig {
|
||||
feat_dim: TEST_FEAT_DIM,
|
||||
seed: 0xAEB1_BEEF_u64,
|
||||
},
|
||||
)?;
|
||||
|
||||
let b_sz: usize = 3;
|
||||
let x: Vec<f32> = (0..b_sz * TEST_FEAT_DIM)
|
||||
.map(|i| ((i as f32) * 0.0011).sin() * 0.5)
|
||||
.collect();
|
||||
let h_old: Vec<f32> = (0..b_sz * AUX_HIDDEN)
|
||||
.map(|i| ((i as f32) * 0.0037).cos() * 0.3)
|
||||
.collect();
|
||||
|
||||
let x_d = upload(&stream, &x)?;
|
||||
let h_old_d = upload(&stream, &h_old)?;
|
||||
let mut h_new_d = stream.alloc_zeros::<f32>(b_sz * AUX_HIDDEN)?;
|
||||
|
||||
let dt: f32 = 1.0;
|
||||
aux_trunk_fwd_gpu(
|
||||
&stream,
|
||||
&trunk.fwd_fn,
|
||||
TEST_FEAT_DIM,
|
||||
&trunk.w_in_d,
|
||||
&trunk.w_rec_d,
|
||||
&trunk.b_d,
|
||||
&trunk.tau_d,
|
||||
&x_d,
|
||||
&h_old_d,
|
||||
dt,
|
||||
b_sz as i32,
|
||||
&mut h_new_d,
|
||||
)?;
|
||||
stream.synchronize()?;
|
||||
let h_new = download(&stream, &h_new_d)?;
|
||||
|
||||
// Pull weights down for the CPU oracle.
|
||||
let weights = trunk.download_weights()?;
|
||||
let w_in = &weights.w_in;
|
||||
let w_rec = &weights.w_rec;
|
||||
let b_vec = &weights.b;
|
||||
let tau = &weights.tau;
|
||||
|
||||
for batch in 0..b_sz {
|
||||
for c in 0..AUX_HIDDEN {
|
||||
let mut pre = b_vec[c];
|
||||
for k in 0..TEST_FEAT_DIM {
|
||||
pre += w_in[c * TEST_FEAT_DIM + k] * x[batch * TEST_FEAT_DIM + k];
|
||||
}
|
||||
for k in 0..AUX_HIDDEN {
|
||||
pre += w_rec[c * AUX_HIDDEN + k] * h_old[batch * AUX_HIDDEN + k];
|
||||
}
|
||||
let decay = (-dt / tau[c].max(1e-6)).exp();
|
||||
let expected =
|
||||
h_old[batch * AUX_HIDDEN + c] * decay + (1.0 - decay) * pre.tanh();
|
||||
let got = h_new[batch * AUX_HIDDEN + c];
|
||||
assert!(
|
||||
(got - expected).abs() < 5e-4,
|
||||
"h_new[{},{}] = {got} vs expected {expected} (diff {})",
|
||||
batch,
|
||||
c,
|
||||
(got - expected).abs(),
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Backward grad shapes finite, non-zero, and bias-grad finite-difference
|
||||
/// matches within tolerance. Full per-weight FD validation is expensive
|
||||
/// (O(AUX_HIDDEN²·feat_dim) reference evaluations); we sample a handful
|
||||
/// of bias entries which is sufficient to catch wiring breakage.
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn bwd_finite_diff_matches_bias_sample() -> Result<()> {
|
||||
let dev = MlDevice::cuda(0)?;
|
||||
let stream = dev.cuda_stream()?.clone();
|
||||
let trunk = AuxTrunk::new(
|
||||
&dev,
|
||||
AuxTrunkConfig {
|
||||
feat_dim: TEST_FEAT_DIM,
|
||||
seed: 0x12_34_56_78,
|
||||
},
|
||||
)?;
|
||||
|
||||
let b_sz: usize = 2;
|
||||
let x: Vec<f32> = (0..b_sz * TEST_FEAT_DIM)
|
||||
.map(|i| ((i as f32) * 0.0021).sin() * 0.4)
|
||||
.collect();
|
||||
let h_old: Vec<f32> = (0..b_sz * AUX_HIDDEN)
|
||||
.map(|i| ((i as f32) * 0.0049).cos() * 0.25)
|
||||
.collect();
|
||||
// Pick a structured grad_h_new (1.0 at one slot per batch, 0 else)
|
||||
// so that grad_b finite-diff has a clean closed-form expectation.
|
||||
let grad_h_new: Vec<f32> = (0..b_sz * AUX_HIDDEN)
|
||||
.map(|i| ((i % 7) as f32) * 0.03)
|
||||
.collect();
|
||||
|
||||
let weights = trunk.download_weights()?;
|
||||
let w_in = weights.w_in.clone();
|
||||
let w_rec = weights.w_rec.clone();
|
||||
let b_vec = weights.b.clone();
|
||||
let tau = weights.tau.clone();
|
||||
|
||||
let dt: f32 = 1.0;
|
||||
let x_d = upload(&stream, &x)?;
|
||||
let h_old_d = upload(&stream, &h_old)?;
|
||||
let grad_h_new_d = upload(&stream, &grad_h_new)?;
|
||||
|
||||
let mut grad_w_in_d = stream.alloc_zeros::<f32>(b_sz * AUX_HIDDEN * TEST_FEAT_DIM)?;
|
||||
let mut grad_w_rec_d = stream.alloc_zeros::<f32>(b_sz * AUX_HIDDEN * AUX_HIDDEN)?;
|
||||
let mut grad_b_d = stream.alloc_zeros::<f32>(b_sz * AUX_HIDDEN)?;
|
||||
let mut grad_tau_d = stream.alloc_zeros::<f32>(b_sz * AUX_HIDDEN)?;
|
||||
let mut grad_h_old_d = stream.alloc_zeros::<f32>(b_sz * AUX_HIDDEN)?;
|
||||
let mut grad_x_d = stream.alloc_zeros::<f32>(b_sz * TEST_FEAT_DIM)?;
|
||||
|
||||
aux_trunk_bwd_gpu(
|
||||
&stream,
|
||||
&trunk.bwd_fn,
|
||||
TEST_FEAT_DIM,
|
||||
&trunk.w_in_d,
|
||||
&trunk.w_rec_d,
|
||||
&trunk.b_d,
|
||||
&trunk.tau_d,
|
||||
&x_d,
|
||||
&h_old_d,
|
||||
&grad_h_new_d,
|
||||
dt,
|
||||
b_sz as i32,
|
||||
&mut grad_w_in_d,
|
||||
&mut grad_w_rec_d,
|
||||
&mut grad_b_d,
|
||||
&mut grad_tau_d,
|
||||
&mut grad_h_old_d,
|
||||
&mut grad_x_d,
|
||||
)?;
|
||||
stream.synchronize()?;
|
||||
|
||||
// All grads finite, no NaN/Inf.
|
||||
let g_w_in = download(&stream, &grad_w_in_d)?;
|
||||
let g_w_rec = download(&stream, &grad_w_rec_d)?;
|
||||
let g_b = download(&stream, &grad_b_d)?;
|
||||
let g_tau = download(&stream, &grad_tau_d)?;
|
||||
let g_h_old = download(&stream, &grad_h_old_d)?;
|
||||
let g_x = download(&stream, &grad_x_d)?;
|
||||
for (name, v) in [
|
||||
("grad_w_in", &g_w_in),
|
||||
("grad_w_rec", &g_w_rec),
|
||||
("grad_b", &g_b),
|
||||
("grad_tau", &g_tau),
|
||||
("grad_h_old", &g_h_old),
|
||||
("grad_x", &g_x),
|
||||
] {
|
||||
assert!(
|
||||
v.iter().all(|x| x.is_finite()),
|
||||
"{name} contains NaN/Inf",
|
||||
);
|
||||
}
|
||||
// Kernel must have written something non-zero for at least the bias
|
||||
// (the only grad guaranteed to be non-trivial given non-zero
|
||||
// grad_h_new entries).
|
||||
assert!(
|
||||
g_b.iter().any(|x| x.abs() > 0.0),
|
||||
"grad_b is all zero — kernel didn't write?",
|
||||
);
|
||||
|
||||
// Finite-difference oracle for grad_b on the first batch, channel 0.
|
||||
// Hand-coded forward (CPU) to derive d/d b[c] of the implicit loss
|
||||
// L = Σ_{batch,c} grad_h_new[batch, c] * h_new[batch, c].
|
||||
let h_new_fwd = |b_local: &[f32]| -> Vec<f32> {
|
||||
let mut out = vec![0.0_f32; b_sz * AUX_HIDDEN];
|
||||
for batch in 0..b_sz {
|
||||
for c in 0..AUX_HIDDEN {
|
||||
let mut pre = b_local[c];
|
||||
for k in 0..TEST_FEAT_DIM {
|
||||
pre += w_in[c * TEST_FEAT_DIM + k] * x[batch * TEST_FEAT_DIM + k];
|
||||
}
|
||||
for k in 0..AUX_HIDDEN {
|
||||
pre += w_rec[c * AUX_HIDDEN + k] * h_old[batch * AUX_HIDDEN + k];
|
||||
}
|
||||
let decay = (-dt / tau[c].max(1e-6)).exp();
|
||||
out[batch * AUX_HIDDEN + c] =
|
||||
h_old[batch * AUX_HIDDEN + c] * decay + (1.0 - decay) * pre.tanh();
|
||||
}
|
||||
}
|
||||
out
|
||||
};
|
||||
let loss_for = |b_local: &[f32]| -> f32 {
|
||||
let h_new = h_new_fwd(b_local);
|
||||
let mut acc = 0.0_f32;
|
||||
for i in 0..b_sz * AUX_HIDDEN {
|
||||
acc += grad_h_new[i] * h_new[i];
|
||||
}
|
||||
acc
|
||||
};
|
||||
|
||||
let eps = 1e-3_f32;
|
||||
// Sample 4 channels at modest stride; full sweep is overkill for a
|
||||
// smoke. The kernel writes per-batch grad slices — sum across batch
|
||||
// to match the loss gradient w.r.t. the shared bias param.
|
||||
for &c in &[0_usize, 17, 33, 50] {
|
||||
let mut b_plus = b_vec.clone();
|
||||
b_plus[c] += eps;
|
||||
let mut b_minus = b_vec.clone();
|
||||
b_minus[c] -= eps;
|
||||
let fd = (loss_for(&b_plus) - loss_for(&b_minus)) / (2.0 * eps);
|
||||
|
||||
// Sum kernel's per-batch grad_b across batch to match dL/db[c].
|
||||
let mut analytical = 0.0_f32;
|
||||
for batch in 0..b_sz {
|
||||
analytical += g_b[batch * AUX_HIDDEN + c];
|
||||
}
|
||||
let abs_err = (fd - analytical).abs();
|
||||
let rel_tol = 5e-3 * fd.abs().max(1e-4);
|
||||
assert!(
|
||||
abs_err < rel_tol.max(2e-3),
|
||||
"grad_b[{c}] mismatch: fd={fd}, analytical={analytical}, abs_err={abs_err}",
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Smoke that AUX_HIDDEN=64 path runs without OOM or NaN at a batch
|
||||
/// size representative of the trainer (which uses B=32-128 per batch).
|
||||
/// Designed to be cheap enough for RTX 3050 sm_86 local validation.
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn fwd_smoke_large_batch_no_nan_no_oom() -> Result<()> {
|
||||
let dev = MlDevice::cuda(0)?;
|
||||
let stream = dev.cuda_stream()?.clone();
|
||||
let trunk = AuxTrunk::new(
|
||||
&dev,
|
||||
AuxTrunkConfig {
|
||||
feat_dim: 128, // matches HIDDEN_DIM — actual B5 wiring config
|
||||
seed: 0xFEED_FACE_u64,
|
||||
},
|
||||
)?;
|
||||
|
||||
let b_sz: usize = 64;
|
||||
let feat_dim = 128usize;
|
||||
let x: Vec<f32> = (0..b_sz * feat_dim)
|
||||
.map(|i| ((i as f32) * 0.0007).sin())
|
||||
.collect();
|
||||
let h_old: Vec<f32> = vec![0.0; b_sz * AUX_HIDDEN]; // cold start
|
||||
let x_d = upload(&stream, &x)?;
|
||||
let h_old_d = upload(&stream, &h_old)?;
|
||||
let mut h_new_d = stream.alloc_zeros::<f32>(b_sz * AUX_HIDDEN)?;
|
||||
|
||||
aux_trunk_fwd_gpu(
|
||||
&stream,
|
||||
&trunk.fwd_fn,
|
||||
feat_dim,
|
||||
&trunk.w_in_d,
|
||||
&trunk.w_rec_d,
|
||||
&trunk.b_d,
|
||||
&trunk.tau_d,
|
||||
&x_d,
|
||||
&h_old_d,
|
||||
1.0,
|
||||
b_sz as i32,
|
||||
&mut h_new_d,
|
||||
)?;
|
||||
stream.synchronize()?;
|
||||
let h_new = download(&stream, &h_new_d)?;
|
||||
assert!(
|
||||
h_new.iter().all(|v| v.is_finite()),
|
||||
"h_new contains NaN/Inf after cold-start forward",
|
||||
);
|
||||
// Output is bounded by tanh ∈ (-1, 1); with h_old = 0 and small init,
|
||||
// |h_new| should sit well inside that range.
|
||||
assert!(
|
||||
h_new.iter().all(|v| v.abs() < 1.5),
|
||||
"h_new produced out-of-range values (|v| ≥ 1.5)",
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user