feat(sp14-c.5b): atomic contract migration — h_s2 → h_s2_aux + revert zero-fills
Atomic flip of the aux heads' input from Q's GRN trunk output `save_h_s2` to the SEPARATE aux trunk's output `h_s2_aux`. The aux trunk now trains its own w1/w2/w3/b1/b2/b3 from CE loss (next-bar + regime); Q's encoder is structurally protected by `aux_trunk_backward`'s missing `dx_in` output param (encoder boundary stop-grad enforced at the kernel-set level). Reverts the C.0 stop-grad band-aid commits (`872bd7392`, `411a30473`): the zero-fills in `aux_next_bar_backward` + `aux_regime_backward` Step 3 are replaced with the genuine SAXPY-back-to-input gradient (`dh_s2_aux[b,j] = sum_k sh_dh_pre[k] * w1[k,j]`). The leak that motivated stop-grad is now blocked structurally rather than by data zero-fill — aux gradient flows through the aux trunk's own params, never into Q's encoder. Wired in this commit (atomic, ~330 LOC): - 4 kernel signatures renamed `h_s2 → h_s2_aux` / `dh_s2_out → dh_s2_aux_out` (`aux_next_bar_forward`, `aux_regime_forward`, `aux_next_bar_backward`, `aux_regime_backward`); Rust wrappers in `gpu_aux_heads.rs` follow - Trainer fwd: insert `aux_trunk_forward_ops.launch(...)` in `aux_heads_forward` Step 0, populating `h_s2_aux` from `save_h_s1` (encoder layer-1 output, dim=shared_h1=256). Both head fwds redirect input pointer from `save_h_s2` to `h_s2_aux` - Trainer bwd: SAXPY both `aux_dh_s2_*_buf` into `dh_s2_aux_accum` (pre-zeroed each step via graph-safe `cuMemsetD32Async`); then `aux_trunk_backward_ops.launch(...)` propagates through w3/w2/w1 + b3/b2/b1; then `launch_aux_trunk_adam_update` applies global L2-norm clip + per-tensor Adam updates over 6 grad tensors - Collector fwd: insert `exp_aux_trunk_forward_ops.launch(...)` after `forward_online_f32`, reading `exp_h_s1_f32` and writing `exp_h_s2_aux`; redirect `exp_aux_heads_fwd.forward_next_bar` input from `exp_h_s2_f32` to `exp_h_s2_aux` - Pre-capture host-write of ISV-driven LR + grad-clip + step counter into mapped-pinned buffers in `launch_cublas_backward_to` (BEFORE `aux_heads_backward`); same `&mut self` pattern as `step_ofi_embed_adam` Verification: - `cargo check -p ml --tests` clean (1m02s, only pre-existing warnings) - `aux_trunk_oracle_tests` + `sp14_oracle_tests` 12/12 pass: - aux_trunk gradient check: max_rel_err=1.33e-2 (tol=2e-2) — matches C.4 baseline - aux_trunk_backward_does_not_write_dx: kernel source clean of dx_in/dx_in_out - aux_sign_label_lookahead_mask: 60/100 masked, 40/100 valid - 9 other oracle tests pass bit-identically Plan: docs/superpowers/plans/2026-05-07-sp14-layer-c-separate-aux-trunk.md §C.5b Audit: docs/dqn-wire-up-audit.md "SP14 Layer C Phase C.5b" section
This commit is contained in:
jgrusewski
@@ -106,6 +106,16 @@ __device__ __forceinline__ float aux_elu_bwd_from_post(float y) {
|
||||
* - aux_dir_acc_reduce_kernel (argmax-vs-label compare)
|
||||
* - aux_pred_to_isv_tanh_kernel (mean(softmax[1] - softmax[0]) → ISV[375])
|
||||
*
|
||||
* SP14 Layer C Phase C.5b (2026-05-08): the input parameter was renamed
|
||||
* `h_s2` → `h_s2_aux`. The aux heads now read the SEPARATE aux trunk's
|
||||
* output rather than Q's GRN trunk output. The Rust caller passes the
|
||||
* trainer's `h_s2_aux` buffer (populated by `aux_trunk_forward` per
|
||||
* `gpu_aux_trunk.rs`). Q's `save_h_s2` is no longer wired here. Symmetric
|
||||
* rename in `aux_next_bar_backward` (`dh_s2_out` → `dh_s2_aux_out`) so
|
||||
* the gradient flows back through the aux trunk to its own params, NOT
|
||||
* into Q's encoder trunk. The encoder boundary stop-gradient is enforced
|
||||
* structurally by `aux_trunk_backward`'s missing `dx_in` output param.
|
||||
*
|
||||
* Mirroring `aux_regime_forward`'s K-fanout structure — thread 0 computes
|
||||
* all K logits serially (K=2 / K=5; the matmul is small enough that a
|
||||
* tree reduce wastes more cycles in shfl bookkeeping than the serial
|
||||
@@ -115,7 +125,7 @@ __device__ __forceinline__ float aux_elu_bwd_from_post(float y) {
|
||||
* (h cache for Linear_2 + temporary K-vector for stable softmax).
|
||||
* ===================================================================== */
|
||||
extern "C" __global__ void aux_next_bar_forward(
|
||||
const float* __restrict__ h_s2, /* [B, SH2] row-major */
|
||||
const float* __restrict__ h_s2_aux, /* [B, SH2] row-major (SP14-C.5b: aux trunk output, NOT Q-trunk h_s2) */
|
||||
const float* __restrict__ w1, /* [H, SH2] row-major */
|
||||
const float* __restrict__ b1, /* [H] */
|
||||
const float* __restrict__ w2, /* [K, H] row-major */
|
||||
@@ -143,7 +153,7 @@ extern "C" __global__ void aux_next_bar_forward(
|
||||
for (int k = tid; k < H; k += blockDim.x) {
|
||||
float acc = b1[k];
|
||||
const float* w_row = w1 + (size_t)k * SH2;
|
||||
const float* h_row = h_s2 + (size_t)b * SH2;
|
||||
const float* h_row = h_s2_aux + (size_t)b * SH2;
|
||||
for (int j = 0; j < SH2; ++j) {
|
||||
acc += w_row[j] * h_row[j];
|
||||
}
|
||||
@@ -200,7 +210,7 @@ extern "C" __global__ void aux_next_bar_forward(
|
||||
* Block: AUX_BLOCK threads. Grid: (B, 1, 1). Shared memory: H floats.
|
||||
* ===================================================================== */
|
||||
extern "C" __global__ void aux_regime_forward(
|
||||
const float* __restrict__ h_s2, /* [B, SH2] row-major */
|
||||
const float* __restrict__ h_s2_aux, /* [B, SH2] row-major (SP14-C.5b: aux trunk output, NOT Q-trunk h_s2) */
|
||||
const float* __restrict__ w1, /* [H, SH2] row-major */
|
||||
const float* __restrict__ b1, /* [H] */
|
||||
const float* __restrict__ w2, /* [K, H] row-major (K = 5) */
|
||||
@@ -221,7 +231,7 @@ extern "C" __global__ void aux_regime_forward(
|
||||
for (int k = tid; k < H; k += blockDim.x) {
|
||||
float acc = b1[k];
|
||||
const float* w_row = w1 + (size_t)k * SH2;
|
||||
const float* h_row = h_s2 + (size_t)b * SH2;
|
||||
const float* h_row = h_s2_aux + (size_t)b * SH2;
|
||||
for (int j = 0; j < SH2; ++j) {
|
||||
acc += w_row[j] * h_row[j];
|
||||
}
|
||||
@@ -480,7 +490,7 @@ extern "C" __global__ void aux_regime_label_from_states(
|
||||
* ===================================================================== */
|
||||
extern "C" __global__ void aux_next_bar_backward(
|
||||
/* Forward inputs (re-read for dW). */
|
||||
const float* __restrict__ h_s2, /* [B, SH2] */
|
||||
const float* __restrict__ h_s2_aux, /* [B, SH2] (SP14-C.5b: aux trunk output) */
|
||||
const float* __restrict__ w1, /* [H, SH2] */
|
||||
const float* __restrict__ w2, /* [K, H] row-major */
|
||||
const float* __restrict__ hidden_post, /* [B, H] saved h_post */
|
||||
@@ -496,9 +506,12 @@ extern "C" __global__ void aux_next_bar_backward(
|
||||
float* __restrict__ db1_partial, /* [B, H] */
|
||||
float* __restrict__ dW2_partial, /* [B, K, H] flat */
|
||||
float* __restrict__ db2_partial, /* [B, K] */
|
||||
/* Trunk gradient (contiguous in B, no batch-reduce needed — caller
|
||||
* accumulates into trunk's main-path d_h_s2 via SAXPY). */
|
||||
float* __restrict__ dh_s2_out /* [B, SH2] */
|
||||
/* Aux trunk-output gradient: SP14-C.5b accumulator
|
||||
* `dh_s2_aux_accum` (NOT Q's `bw_d_h_s2`). Caller SAXPYs both heads'
|
||||
* dh_s2_aux into this buffer, then aux_trunk_backward reads it as the
|
||||
* upstream gradient feeding `dh_s2_aux_in_ptr`. Encoder stop-grad is
|
||||
* enforced structurally by aux_trunk_backward (no dx_in output). */
|
||||
float* __restrict__ dh_s2_aux_out /* [B, SH2] */
|
||||
) {
|
||||
const int b = blockIdx.x;
|
||||
if (b >= B) return;
|
||||
@@ -586,7 +599,7 @@ extern "C" __global__ void aux_next_bar_backward(
|
||||
|
||||
/* dW1 partial — per-(k, j) write. Stride loop over (H * SH2). */
|
||||
{
|
||||
const float* h_row = h_s2 + (size_t)b * SH2;
|
||||
const float* h_row = h_s2_aux + (size_t)b * SH2;
|
||||
float* dW1_b = dW1_partial + (size_t)b * H * SH2;
|
||||
const int total = H * SH2;
|
||||
for (int idx = tid; idx < total; idx += blockDim.x) {
|
||||
@@ -596,44 +609,31 @@ extern "C" __global__ void aux_next_bar_backward(
|
||||
}
|
||||
}
|
||||
|
||||
/* Step 3: SP14 stop-gradient (2026-05-07).
|
||||
/* Step 3: dh_s2_aux[b, j] = sum_k sh_dh_pre[k] * w1[k, j].
|
||||
*
|
||||
* Diagnostic from L40S `train-v8ztm` 9-epoch HEALTH_DIAG aux next_bar_mse:
|
||||
* Ep 0: 0.352 (learnable signal — below random baseline ln(2)≈0.693)
|
||||
* Ep 9: 0.717 (above random baseline — aux is now WORSE than random)
|
||||
* aux_dir_acc_long stuck at 0.19 (anti-correlated with truth)
|
||||
* SP14 Layer C Phase C.5b (2026-05-08): the prior commit `872bd7392`
|
||||
* zero-filled this buffer as a stop-grad band-aid because the aux
|
||||
* head's gradient was leaking into Q's shared `h_s2`. Phase C.5b
|
||||
* REVERTS that band-aid: aux heads now read the SEPARATE aux trunk's
|
||||
* output `h_s2_aux`, and `dh_s2_aux_out` accumulates into the aux
|
||||
* trunk's own input gradient (`dh_s2_aux_accum`), NOT into Q's
|
||||
* `bw_d_h_s2`. The aux trunk's backward kernel terminates at the
|
||||
* encoder boundary (no `dx_in` output param), so Q's encoder is
|
||||
* structurally protected — no leak possible. With that structural
|
||||
* guarantee the genuine SAXPY-back-to-input gradient is restored.
|
||||
*
|
||||
* Root cause: the aux head's backward gradient was flowing back into the
|
||||
* shared trunk activation `h_s2` via the SAXPY of `dh_s2_out` into the
|
||||
* trunk's main-path d_h_s2. Q-loss gradient on h_s2 dominates (larger
|
||||
* magnitude, structurally different objective: cumulative discounted
|
||||
* reward vs next-bar direction). h_s2 evolves to support Q's task; aux's
|
||||
* CE loss climbs as h_s2 features become anti-aligned with direction
|
||||
* prediction.
|
||||
*
|
||||
* Fix: stop-gradient on aux's input. We zero `dh_s2_out` rather than
|
||||
* computing the sum-of-w1·dh_pre. The aux head still trains its own
|
||||
* params (w1/b1/w2/b2) from the CE loss via the dW1/db1/dW2/db2
|
||||
* partials written above — those are unaffected. But h_s2 is no longer
|
||||
* pulled by aux's gradient. Q-loss is the sole shaping force on h_s2;
|
||||
* aux must adapt to whatever h_s2 happens to be.
|
||||
*
|
||||
* If aux still cannot learn after this fix, the deeper diagnosis is that
|
||||
* h_s2 carries no direction signal at all — separate-trunk Option 2
|
||||
* (deferred). The trunk SAXPY adds zero either way (same arithmetic the
|
||||
* masked-row branch already produced for skipped samples), so no other
|
||||
* consumer is affected — `dh_s2_out` is read only by the trunk
|
||||
* accumulation step.
|
||||
*
|
||||
* NOTE: `aux_regime_backward` (the K=5 regime CE head) has the same
|
||||
* architecture and received the SAME fix in the immediately-following
|
||||
* commit (mirrors this kernel's Step 3 zero-write). Both auxiliary
|
||||
* heads now train their own params from CE loss without pulling on
|
||||
* shared h_s2 — Q-loss is the sole trunk-shaping force. */
|
||||
* Per-sample, per-feature serial reduction over H lanes — H is small
|
||||
* (32) and SH2 is moderate; stride loop over j. Masked rows propagate
|
||||
* zero d_h_pre through this sum, so dh_s2_aux row is zero for skipped
|
||||
* samples (the trunk SAXPY adds zero — no spurious gradient). */
|
||||
{
|
||||
float* dh_row = dh_s2_out + (size_t)b * SH2;
|
||||
float* dh_row = dh_s2_aux_out + (size_t)b * SH2;
|
||||
for (int j = tid; j < SH2; j += blockDim.x) {
|
||||
dh_row[j] = 0.0f;
|
||||
float acc = 0.0f;
|
||||
for (int k = 0; k < H; ++k) {
|
||||
acc += sh_dh_pre[k] * w1[(size_t)k * SH2 + j];
|
||||
}
|
||||
dh_row[j] = acc;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -658,7 +658,7 @@ extern "C" __global__ void aux_next_bar_backward(
|
||||
* ===================================================================== */
|
||||
extern "C" __global__ void aux_regime_backward(
|
||||
/* Forward inputs. */
|
||||
const float* __restrict__ h_s2, /* [B, SH2] */
|
||||
const float* __restrict__ h_s2_aux, /* [B, SH2] (SP14-C.5b: aux trunk output) */
|
||||
const float* __restrict__ w1, /* [H, SH2] */
|
||||
const float* __restrict__ w2, /* [K, H] row-major */
|
||||
const float* __restrict__ hidden_post, /* [B, H] saved h_post */
|
||||
@@ -673,8 +673,9 @@ extern "C" __global__ void aux_regime_backward(
|
||||
float* __restrict__ db1_partial, /* [B, H] */
|
||||
float* __restrict__ dW2_partial, /* [B, K, H] */
|
||||
float* __restrict__ db2_partial, /* [B, K] */
|
||||
/* Trunk gradient. */
|
||||
float* __restrict__ dh_s2_out /* [B, SH2] */
|
||||
/* Aux trunk-output gradient (SAXPYed into `dh_s2_aux_accum`; encoder
|
||||
* stop-grad enforced structurally by `aux_trunk_backward`). */
|
||||
float* __restrict__ dh_s2_aux_out /* [B, SH2] */
|
||||
) {
|
||||
const int b = blockIdx.x;
|
||||
if (b >= B) return;
|
||||
@@ -742,9 +743,9 @@ extern "C" __global__ void aux_regime_backward(
|
||||
db1_partial[(size_t)b * H + k] = sh_dh_pre[k];
|
||||
}
|
||||
|
||||
/* dW1_partial[b, k, j] = sh_dh_pre[k] * h_s2[b, j]. */
|
||||
/* dW1_partial[b, k, j] = sh_dh_pre[k] * h_s2_aux[b, j]. */
|
||||
{
|
||||
const float* h_row = h_s2 + (size_t)b * SH2;
|
||||
const float* h_row = h_s2_aux + (size_t)b * SH2;
|
||||
float* dW1_b = dW1_partial + (size_t)b * H * SH2;
|
||||
const int total = H * SH2;
|
||||
for (int idx = tid; idx < total; idx += blockDim.x) {
|
||||
@@ -754,38 +755,24 @@ extern "C" __global__ void aux_regime_backward(
|
||||
}
|
||||
}
|
||||
|
||||
/* Step 3: SP14 stop-gradient (2026-05-07, mirrors next-bar fix at commit
|
||||
* 872bd7392).
|
||||
/* Step 3: dh_s2_aux[b, j] = sum_k sh_dh_pre[k] * w1[k, j].
|
||||
*
|
||||
* The 5-class regime CE head's backward gradient was flowing back into
|
||||
* the shared trunk activation `h_s2` via the SAXPY of `dh_s2_out` into
|
||||
* the trunk's main-path d_h_s2 — IDENTICAL architectural conflict to
|
||||
* the next-bar head closed earlier the same day. Q-loss gradient on
|
||||
* h_s2 dominates (larger magnitude, structurally different objective:
|
||||
* cumulative discounted reward vs 5-class regime classification).
|
||||
* h_s2 evolves to support Q's task; regime's CE loss can't shape
|
||||
* h_s2 without conflicting with Q's directional pressure.
|
||||
*
|
||||
* Fix: stop-gradient. Regime reads h_s2 via forward, trains its own
|
||||
* w1/b1/w2/b2 from CE loss via the dW1_partial / db1_partial /
|
||||
* dW2_partial / db2_partial writes above (those are unaffected — they
|
||||
* read sh_dh_pre, sh_dlogits, sh_h_post, h_row, all computed
|
||||
* pre-write); only the trunk-bound gradient is severed. Q-loss is the
|
||||
* sole shaping force on h_s2; regime adapts to whatever h_s2 happens
|
||||
* to be.
|
||||
*
|
||||
* If regime can't learn even with this fix, the deeper bug is that
|
||||
* h_s2 has no regime-relevant signal — would need separate aux trunk
|
||||
* (Option 2, deferred). No other consumer reads `dh_s2_out` — only
|
||||
* the trunk accumulation step.
|
||||
*
|
||||
* Today's fix has been validated by passing oracle tests; we're not
|
||||
* waiting for L40S validation to apply the same fix to regime since
|
||||
* the fix is mechanical and identical. */
|
||||
* SP14 Layer C Phase C.5b (2026-05-08): the prior commit `411a30473`
|
||||
* zero-filled this buffer as a stop-grad band-aid because the regime
|
||||
* head's gradient was leaking into Q's shared `h_s2`. Phase C.5b
|
||||
* REVERTS that band-aid: regime now reads `h_s2_aux` (aux trunk
|
||||
* output), and `dh_s2_aux_out` SAXPYs into the aux trunk's own input
|
||||
* gradient `dh_s2_aux_accum`. The aux trunk's backward kernel has
|
||||
* NO `dx_in` output param so Q's encoder is structurally protected.
|
||||
* The genuine SAXPY-back-to-input gradient is restored. */
|
||||
{
|
||||
float* dh_row = dh_s2_out + (size_t)b * SH2;
|
||||
float* dh_row = dh_s2_aux_out + (size_t)b * SH2;
|
||||
for (int j = tid; j < SH2; j += blockDim.x) {
|
||||
dh_row[j] = 0.0f;
|
||||
float acc = 0.0f;
|
||||
for (int k = 0; k < H; ++k) {
|
||||
acc += sh_dh_pre[k] * w1[(size_t)k * SH2 + j];
|
||||
}
|
||||
dh_row[j] = acc;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -142,7 +142,8 @@ impl AuxHeadsForwardOps {
|
||||
///
|
||||
/// Caller-owned buffers (raw `u64` device pointers for graph-capture
|
||||
/// safety; mirrors `gpu_grn.rs::forward_raw`):
|
||||
/// * `h_s2_ptr` — `[B, SH2]` row-major.
|
||||
/// * `h_s2_aux_ptr` — `[B, SH2]` row-major (SP14-C.5b: aux trunk
|
||||
/// output, NOT Q's GRN trunk `save_h_s2`).
|
||||
/// * `w1_ptr`/`b1_ptr` — `[H, SH2]` / `[H]`. Slice from `params_buf[119]`/`[120]`.
|
||||
/// * `w2_ptr`/`b2_ptr` — `[K, H]` / `[K]`. Slice from `params_buf[121]`/`[122]`.
|
||||
/// * `hidden_out_ptr` — `[B, H]` SAVED post-ELU buffer (consumed by backward).
|
||||
@@ -161,7 +162,7 @@ impl AuxHeadsForwardOps {
|
||||
pub(crate) fn forward_next_bar(
|
||||
&self,
|
||||
stream: &Arc<CudaStream>,
|
||||
h_s2_ptr: u64,
|
||||
h_s2_aux_ptr: u64,
|
||||
w1_ptr: u64,
|
||||
b1_ptr: u64,
|
||||
w2_ptr: u64,
|
||||
@@ -180,7 +181,7 @@ impl AuxHeadsForwardOps {
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(&self.next_bar_forward_kernel)
|
||||
.arg(&h_s2_ptr)
|
||||
.arg(&h_s2_aux_ptr)
|
||||
.arg(&w1_ptr)
|
||||
.arg(&b1_ptr)
|
||||
.arg(&w2_ptr)
|
||||
@@ -209,7 +210,7 @@ impl AuxHeadsForwardOps {
|
||||
pub(crate) fn forward_regime(
|
||||
&self,
|
||||
stream: &Arc<CudaStream>,
|
||||
h_s2_ptr: u64,
|
||||
h_s2_aux_ptr: u64,
|
||||
w1_ptr: u64,
|
||||
b1_ptr: u64,
|
||||
w2_ptr: u64,
|
||||
@@ -227,7 +228,7 @@ impl AuxHeadsForwardOps {
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(&self.regime_forward_kernel)
|
||||
.arg(&h_s2_ptr)
|
||||
.arg(&h_s2_aux_ptr)
|
||||
.arg(&w1_ptr)
|
||||
.arg(&b1_ptr)
|
||||
.arg(&w2_ptr)
|
||||
@@ -469,7 +470,7 @@ impl AuxHeadsBackwardOps {
|
||||
pub(crate) fn backward_next_bar(
|
||||
&self,
|
||||
stream: &Arc<CudaStream>,
|
||||
h_s2_ptr: u64,
|
||||
h_s2_aux_ptr: u64,
|
||||
w1_ptr: u64,
|
||||
w2_ptr: u64,
|
||||
hidden_post_ptr: u64,
|
||||
@@ -483,7 +484,7 @@ impl AuxHeadsBackwardOps {
|
||||
db1_partial_ptr: u64,
|
||||
dw2_partial_ptr: u64,
|
||||
db2_partial_ptr: u64,
|
||||
dh_s2_out_ptr: u64,
|
||||
dh_s2_aux_out_ptr: u64,
|
||||
) -> Result<(), MLError> {
|
||||
let b_i32 = b as i32;
|
||||
let sh2_i32 = sh2 as i32;
|
||||
@@ -494,7 +495,7 @@ impl AuxHeadsBackwardOps {
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(&self.next_bar_backward_kernel)
|
||||
.arg(&h_s2_ptr)
|
||||
.arg(&h_s2_aux_ptr)
|
||||
.arg(&w1_ptr)
|
||||
.arg(&w2_ptr)
|
||||
.arg(&hidden_post_ptr)
|
||||
@@ -508,7 +509,7 @@ impl AuxHeadsBackwardOps {
|
||||
.arg(&db1_partial_ptr)
|
||||
.arg(&dw2_partial_ptr)
|
||||
.arg(&db2_partial_ptr)
|
||||
.arg(&dh_s2_out_ptr)
|
||||
.arg(&dh_s2_aux_out_ptr)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (b as u32, 1, 1),
|
||||
block_dim: (AUX_BLOCK, 1, 1),
|
||||
@@ -526,7 +527,7 @@ impl AuxHeadsBackwardOps {
|
||||
pub(crate) fn backward_regime(
|
||||
&self,
|
||||
stream: &Arc<CudaStream>,
|
||||
h_s2_ptr: u64,
|
||||
h_s2_aux_ptr: u64,
|
||||
w1_ptr: u64,
|
||||
w2_ptr: u64,
|
||||
hidden_post_ptr: u64,
|
||||
@@ -539,7 +540,7 @@ impl AuxHeadsBackwardOps {
|
||||
db1_partial_ptr: u64,
|
||||
dw2_partial_ptr: u64,
|
||||
db2_partial_ptr: u64,
|
||||
dh_s2_out_ptr: u64,
|
||||
dh_s2_aux_out_ptr: u64,
|
||||
) -> Result<(), MLError> {
|
||||
let b_i32 = b as i32;
|
||||
let sh2_i32 = sh2 as i32;
|
||||
@@ -549,7 +550,7 @@ impl AuxHeadsBackwardOps {
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(&self.regime_backward_kernel)
|
||||
.arg(&h_s2_ptr)
|
||||
.arg(&h_s2_aux_ptr)
|
||||
.arg(&w1_ptr)
|
||||
.arg(&w2_ptr)
|
||||
.arg(&hidden_post_ptr)
|
||||
@@ -562,7 +563,7 @@ impl AuxHeadsBackwardOps {
|
||||
.arg(&db1_partial_ptr)
|
||||
.arg(&dw2_partial_ptr)
|
||||
.arg(&db2_partial_ptr)
|
||||
.arg(&dh_s2_out_ptr)
|
||||
.arg(&dh_s2_aux_out_ptr)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (b as u32, 1, 1),
|
||||
block_dim: (AUX_BLOCK, 1, 1),
|
||||
|
||||
@@ -17593,8 +17593,18 @@ impl GpuDqnTrainer {
|
||||
pub(crate) fn aux_heads_forward(&self) -> Result<(), MLError> {
|
||||
let b = self.config.batch_size;
|
||||
let sh2 = self.config.shared_h2;
|
||||
let h_s2_ptr = self.ptrs.save_h_s2;
|
||||
// SP14 Layer C Phase C.5b (2026-05-08): aux heads now read the
|
||||
// SEPARATE aux trunk's output `h_s2_aux` rather than Q's GRN trunk
|
||||
// output `save_h_s2`. The aux trunk forward (below) consumes
|
||||
// `save_h_s1` (encoder output) and writes `h_s2_aux`. Symmetric
|
||||
// contract change in `aux_heads_backward` + the kernel signature
|
||||
// rename `h_s2 → h_s2_aux` lock-in step.
|
||||
let h_s2_aux_ptr = self.h_s2_aux.raw_ptr();
|
||||
let states_ptr = self.ptrs.states_buf;
|
||||
let encoder_out_ptr = self.ptrs.save_h_s1;
|
||||
let encoder_out_dim = self.config.shared_h1;
|
||||
let aux_h1 = super::gpu_aux_trunk::AUX_TRUNK_H1;
|
||||
let aux_h2 = super::gpu_aux_trunk::AUX_TRUNK_H2;
|
||||
|
||||
let param_sizes = compute_param_sizes(&self.config);
|
||||
let on_w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes);
|
||||
@@ -17608,6 +17618,32 @@ impl GpuDqnTrainer {
|
||||
let rg_w2 = on_w_ptrs[125];
|
||||
let rg_b2 = on_w_ptrs[126];
|
||||
|
||||
// Step 0: aux trunk forward — Linear→ELU→Linear→ELU→Linear MLP that
|
||||
// produces `h_s2_aux [B, SH2]` from the encoder's layer-1 output
|
||||
// `save_h_s1 [B, SH1]`. Runs BEFORE the regime label builder + per-
|
||||
// head forwards so both heads see the freshly-written `h_s2_aux`.
|
||||
// Aux trunk weights live in dedicated `aux_trunk_w{1,2,3}` /
|
||||
// `aux_trunk_b{1,2,3}` slabs (NOT in the flat `params_buf`); they
|
||||
// are trained by `launch_aux_trunk_adam_update` in the backward.
|
||||
self.aux_trunk_forward_ops.launch(
|
||||
&self.stream,
|
||||
encoder_out_ptr,
|
||||
self.aux_trunk_w1.raw_ptr(),
|
||||
self.aux_trunk_b1.raw_ptr(),
|
||||
self.aux_trunk_w2.raw_ptr(),
|
||||
self.aux_trunk_b2.raw_ptr(),
|
||||
self.aux_trunk_w3.raw_ptr(),
|
||||
self.aux_trunk_b3.raw_ptr(),
|
||||
self.h_aux1.raw_ptr(),
|
||||
self.h_aux2.raw_ptr(),
|
||||
h_s2_aux_ptr,
|
||||
b,
|
||||
encoder_out_dim,
|
||||
aux_h1,
|
||||
aux_h2,
|
||||
sh2,
|
||||
)?;
|
||||
|
||||
let state_dim_padded = ml_core::state_layout::STATE_DIM_PADDED;
|
||||
let regime_score_idx = ml_core::state_layout::OFI_START + 29;
|
||||
|
||||
@@ -17639,9 +17675,10 @@ impl GpuDqnTrainer {
|
||||
|
||||
// Step 3: next-bar direction-classification forward (SP13 B1.1a:
|
||||
// K=2 softmax). Writes hidden tile + logits tile + softmax tile.
|
||||
// SP14-C.5b: input is `h_s2_aux` (aux trunk output), NOT Q's `h_s2`.
|
||||
self.aux_heads_fwd.forward_next_bar(
|
||||
&self.stream,
|
||||
h_s2_ptr,
|
||||
h_s2_aux_ptr,
|
||||
nb_w1, nb_b1, nb_w2, nb_b2,
|
||||
b, sh2, AUX_NEXT_BAR_K,
|
||||
self.aux_nb_hidden_buf.raw_ptr(),
|
||||
@@ -17650,9 +17687,10 @@ impl GpuDqnTrainer {
|
||||
)?;
|
||||
|
||||
// Step 4: regime classification forward.
|
||||
// SP14-C.5b: input is `h_s2_aux` (aux trunk output), NOT Q's `h_s2`.
|
||||
self.aux_heads_fwd.forward_regime(
|
||||
&self.stream,
|
||||
h_s2_ptr,
|
||||
h_s2_aux_ptr,
|
||||
rg_w1, rg_b1, rg_w2, rg_b2,
|
||||
b, sh2, AUX_REGIME_K,
|
||||
self.aux_rg_hidden_buf.raw_ptr(),
|
||||
@@ -17685,28 +17723,50 @@ impl GpuDqnTrainer {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Plan 4 Task 6 Commit B: aux-heads backward orchestrator.
|
||||
/// Plan 4 Task 6 Commit B + SP14 Layer C Phase C.5b (2026-05-08):
|
||||
/// aux-heads + aux-trunk backward orchestrator.
|
||||
///
|
||||
/// Runs both backward kernels (next-bar + regime) over the saved forward
|
||||
/// state, then for each of the 8 aux param tensors:
|
||||
/// Runs both head backward kernels (next-bar + regime) over the saved
|
||||
/// forward state, then for each of the 8 aux head param tensors:
|
||||
/// 1. `aux_param_grad_reduce` collapses `partials [B, P]` into `final [P]`.
|
||||
/// 2. SAXPY `grad_buf[tensor_idx] += aux_weight * final` via the
|
||||
/// existing `dqn_saxpy_f32_kernel`.
|
||||
/// Finally, the per-sample `dh_s2` from both heads is SAXPY-accumulated
|
||||
/// into `bw_d_h_s2 [B, SH2]` with `alpha = aux_weight`.
|
||||
/// Then SP14-C.5b: the per-sample `dh_s2_aux` from both heads is
|
||||
/// SAXPY-accumulated into `dh_s2_aux_accum [B, SH2]` (the AUX TRUNK's
|
||||
/// upstream-grad accumulator, NOT Q's `bw_d_h_s2`). Aux trunk backward
|
||||
/// then reads the accumulator and propagates gradient through w3/w2/w1
|
||||
/// + b3/b2/b1, terminating at the encoder boundary (structural stop-
|
||||
/// grad — kernel set has NO `dx_in` output, so Q's `bw_d_h_s2` is
|
||||
/// unaffected by aux loss). Finally, `launch_aux_trunk_adam_update`
|
||||
/// applies global L2-norm clip + per-tensor Adam updates over the 6
|
||||
/// aux trunk grad buffers.
|
||||
///
|
||||
/// MUST run AFTER `backward_full` + concat-accumulators fill `bw_d_h_s2`
|
||||
/// AND BEFORE `encoder_backward_chain` consumes it.
|
||||
/// AND BEFORE `encoder_backward_chain` consumes `bw_d_h_s2` (so the
|
||||
/// trunk dW/dB chain rule sees the un-augmented Q gradient).
|
||||
///
|
||||
/// `grad_base` is the same value passed to `launch_cublas_backward_to` —
|
||||
/// usually `self.ptrs.grad_buf`, but the vaccine path (g_val computation)
|
||||
/// passes a scratch buffer here. The aux SAXPY targets `grad_base[119..127)`
|
||||
/// at the same `padded_byte_offset` math the rest of the backward uses.
|
||||
/// passes a scratch buffer here. The aux head SAXPY targets
|
||||
/// `grad_base[119..127)` at the same `padded_byte_offset` math the rest
|
||||
/// of the backward uses; the aux trunk Adam launcher writes its 6 grad
|
||||
/// tensors directly (separate slabs, NOT in `params_buf`/`grad_buf`).
|
||||
pub(crate) fn aux_heads_backward(&self, grad_base: u64) -> Result<(), MLError> {
|
||||
let b = self.config.batch_size;
|
||||
let sh2 = self.config.shared_h2;
|
||||
let h_s2_ptr = self.ptrs.save_h_s2;
|
||||
let bw_d_h_s2_ptr = self.bw_d_h_s2.raw_ptr();
|
||||
// SP14 Layer C Phase C.5b (2026-05-08): aux heads now consume the
|
||||
// SEPARATE aux trunk's output `h_s2_aux`, and their per-sample
|
||||
// dh_s2 SAXPYs into `dh_s2_aux_accum` (NOT Q's `bw_d_h_s2`). The
|
||||
// aux trunk's backward kernel then reads `dh_s2_aux_accum` as
|
||||
// `dh_s2_aux_in_ptr` and propagates grad through the aux trunk's
|
||||
// own w1/w2/w3/b1/b2/b3 (terminating at the encoder boundary;
|
||||
// structural stop-grad enforced by missing `dx_in` output).
|
||||
let h_s2_aux_ptr = self.h_s2_aux.raw_ptr();
|
||||
let dh_s2_aux_accum_ptr = self.dh_s2_aux_accum.raw_ptr();
|
||||
let encoder_out_ptr = self.ptrs.save_h_s1;
|
||||
let encoder_out_dim = self.config.shared_h1;
|
||||
let aux_h1 = super::gpu_aux_trunk::AUX_TRUNK_H1;
|
||||
let aux_h2 = super::gpu_aux_trunk::AUX_TRUNK_H2;
|
||||
let aux_w = self.aux_weight;
|
||||
let aux_h = AUX_HIDDEN_DIM;
|
||||
let aux_kr = AUX_REGIME_K;
|
||||
@@ -17725,9 +17785,12 @@ impl GpuDqnTrainer {
|
||||
// SP13 B1.1a (2026-05-05): softmax CE backward. Reads the saved
|
||||
// softmax tile + i32 labels + B_valid scalar (written by loss
|
||||
// reduce) so loss + gradient share the same `1/B_valid` divisor.
|
||||
// SP14-C.5b: input is `h_s2_aux` (aux trunk output). The dh_s2_aux
|
||||
// outputs are written to per-head buffers (overwrite, not
|
||||
// accumulate); the SAXPY-into-`dh_s2_aux_accum` happens below.
|
||||
self.aux_heads_bwd.backward_next_bar(
|
||||
&self.stream,
|
||||
h_s2_ptr,
|
||||
h_s2_aux_ptr,
|
||||
nb_w1, nb_w2,
|
||||
self.aux_nb_hidden_buf.raw_ptr(),
|
||||
self.aux_nb_softmax_buf.raw_ptr(),
|
||||
@@ -17742,7 +17805,7 @@ impl GpuDqnTrainer {
|
||||
)?;
|
||||
self.aux_heads_bwd.backward_regime(
|
||||
&self.stream,
|
||||
h_s2_ptr,
|
||||
h_s2_aux_ptr,
|
||||
rg_w1, rg_w2,
|
||||
self.aux_rg_hidden_buf.raw_ptr(),
|
||||
self.aux_rg_logits_buf.raw_ptr(),
|
||||
@@ -17816,18 +17879,32 @@ impl GpuDqnTrainer {
|
||||
}
|
||||
}
|
||||
|
||||
// SAXPY both per-head dh_s2 buffers into the trunk accumulator
|
||||
// `bw_d_h_s2`. Both source and destination already share the [B, SH2]
|
||||
// shape + unit stride. encoder_backward_chain consumes the augmented
|
||||
// accumulator transparently through the same pointer.
|
||||
// SP14 Layer C Phase C.5b (2026-05-08): SAXPY both per-head
|
||||
// `aux_dh_s2_*_buf` buffers into the AUX TRUNK upstream-grad
|
||||
// accumulator `dh_s2_aux_accum` (NOT Q's `bw_d_h_s2`). The aux
|
||||
// trunk's backward kernel reads `dh_s2_aux_accum` as its
|
||||
// `dh_s2_aux_in_ptr` and propagates gradient through w3/w2/w1 +
|
||||
// b3/b2/b1, terminating at the encoder boundary (no `dx_in`
|
||||
// output param). Pre-zero the accumulator each step so we don't
|
||||
// carry forward last step's gradient — graph-safe via
|
||||
// `cuMemsetD32Async` (matches the d_h_s2 zero pattern in the
|
||||
// backward chain).
|
||||
let n_dh = (b * sh2) as i32;
|
||||
let blocks_dh = ((n_dh as u32 + 255) / 256).max(1);
|
||||
let dh_nb_ptr = self.aux_dh_s2_nb_buf.raw_ptr();
|
||||
let dh_rg_ptr = self.aux_dh_s2_rg_buf.raw_ptr();
|
||||
unsafe {
|
||||
// Pre-zero `dh_s2_aux_accum` so the SAXPY below establishes
|
||||
// the per-step value rather than accumulating across steps.
|
||||
cudarc::driver::sys::cuMemsetD32Async(
|
||||
dh_s2_aux_accum_ptr,
|
||||
0,
|
||||
(b * sh2),
|
||||
self.stream.cu_stream(),
|
||||
);
|
||||
self.stream
|
||||
.launch_builder(&self.saxpy_f32_kernel)
|
||||
.arg(&bw_d_h_s2_ptr)
|
||||
.arg(&dh_s2_aux_accum_ptr)
|
||||
.arg(&dh_nb_ptr)
|
||||
.arg(&aux_w)
|
||||
.arg(&n_dh)
|
||||
@@ -17837,11 +17914,11 @@ impl GpuDqnTrainer {
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"aux_heads_backward saxpy dh_s2 next_bar: {e}"
|
||||
"aux_heads_backward saxpy dh_s2_aux next_bar: {e}"
|
||||
)))?;
|
||||
self.stream
|
||||
.launch_builder(&self.saxpy_f32_kernel)
|
||||
.arg(&bw_d_h_s2_ptr)
|
||||
.arg(&dh_s2_aux_accum_ptr)
|
||||
.arg(&dh_rg_ptr)
|
||||
.arg(&aux_w)
|
||||
.arg(&n_dh)
|
||||
@@ -17851,10 +17928,53 @@ impl GpuDqnTrainer {
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"aux_heads_backward saxpy dh_s2 regime: {e}"
|
||||
"aux_heads_backward saxpy dh_s2_aux regime: {e}"
|
||||
)))?;
|
||||
}
|
||||
|
||||
// SP14 Layer C Phase C.5b: aux trunk backward — propagates the
|
||||
// accumulated upstream gradient `dh_s2_aux_accum` through the
|
||||
// 3-layer aux trunk MLP, writing per-tensor gradients into the
|
||||
// 6 grad buffers (overwrite, not accumulate). Encoder boundary is
|
||||
// structural stop-grad: the kernel set has NO `dx_in` output, so
|
||||
// Q's encoder gradient is unaffected. Pearl
|
||||
// `pearl_no_host_branches_in_captured_graph.md` is preserved —
|
||||
// all kernels pre-loaded at construction; no host dispatch in
|
||||
// the captured region.
|
||||
self.aux_trunk_backward_ops.launch(
|
||||
&self.stream,
|
||||
dh_s2_aux_accum_ptr,
|
||||
encoder_out_ptr,
|
||||
self.h_aux1.raw_ptr(),
|
||||
self.h_aux2.raw_ptr(),
|
||||
self.aux_trunk_w2.raw_ptr(),
|
||||
self.aux_trunk_w3.raw_ptr(),
|
||||
self.dh_aux1_pre_scratch.raw_ptr(),
|
||||
self.dh_aux2_pre_scratch.raw_ptr(),
|
||||
self.aux_trunk_w1_grad.raw_ptr(),
|
||||
self.aux_trunk_b1_grad.raw_ptr(),
|
||||
self.aux_trunk_w2_grad.raw_ptr(),
|
||||
self.aux_trunk_b2_grad.raw_ptr(),
|
||||
self.aux_trunk_w3_grad.raw_ptr(),
|
||||
self.aux_trunk_b3_grad.raw_ptr(),
|
||||
b,
|
||||
encoder_out_dim,
|
||||
aux_h1,
|
||||
aux_h2,
|
||||
sh2,
|
||||
)?;
|
||||
|
||||
// SP14 Layer C Phase C.5b: aux trunk Adam optimiser step —
|
||||
// consumes the 6 grad tensors via global L2-norm clip (single
|
||||
// norm spans all 6 tensors) + per-tensor Adam update. ISV-driven
|
||||
// hyperparams (LR, β1, β2, ε, grad_clip) read inside the launcher;
|
||||
// pinned LR/clip buffers are populated by the caller pre-capture
|
||||
// (host-write-through-mapped-pinned). The aux trunk step counter
|
||||
// `aux_trunk_t_dev_ptr` is incremented by `submit_counter_increments`
|
||||
// before this launch, ensuring Adam bias correction tracks step
|
||||
// count consistently across captured graph replays.
|
||||
self.launch_aux_trunk_adam_update(&self.stream, b as i32)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -29909,13 +30029,49 @@ impl GpuDqnTrainer {
|
||||
self.check_nan_f32(self.ptrs.bw_d_h_s2, b_post * sh2_post, 33)?;
|
||||
}
|
||||
|
||||
// ── SP14 Layer C Phase C.5b (2026-05-08): aux trunk Adam pre-capture ──
|
||||
// Host-write through mapped-pinned to populate Adam hyperparams +
|
||||
// step counter for the captured aux trunk Adam launch (inside
|
||||
// `aux_heads_backward` below). The captured graph reads pointers,
|
||||
// not values — these host writes are picked up at replay time.
|
||||
//
|
||||
// Per `pearl_no_host_branches_in_captured_graph.md` we cannot do
|
||||
// these host writes from inside the captured region, but
|
||||
// `launch_cublas_backward_to` is `&mut self` and runs OUTSIDE the
|
||||
// captured graph (the captured chain is each `cuGraphLaunch`
|
||||
// replay; this Rust function runs each per-step host driver).
|
||||
//
|
||||
// ISV[AUX_TRUNK_LR_INDEX=444] / ISV[AUX_TRUNK_GRAD_CLIP_INDEX=448]
|
||||
// are populated by the StateResetRegistry at fold boundaries
|
||||
// (`AUX_TRUNK_LR_DEFAULT=1e-4` / `AUX_TRUNK_GRAD_CLIP_DEFAULT=1.0`)
|
||||
// and unchanged within a fold; reading once per step lets future
|
||||
// adaptive controllers drive these slots without re-wiring.
|
||||
{
|
||||
use crate::cuda_pipeline::sp14_isv_slots::{
|
||||
AUX_TRUNK_GRAD_CLIP_INDEX, AUX_TRUNK_LR_INDEX,
|
||||
};
|
||||
let aux_lr = self.read_isv_signal_at(AUX_TRUNK_LR_INDEX);
|
||||
let aux_clip = self.read_isv_signal_at(AUX_TRUNK_GRAD_CLIP_INDEX);
|
||||
// Step counter is host-tracked (mirrors `step_ofi_embed_adam`
|
||||
// pattern): increment then write through to the pinned buffer.
|
||||
// The captured Adam kernel reads `*t_dev_ptr` at replay time.
|
||||
self.aux_trunk_adam_step += 1;
|
||||
unsafe {
|
||||
*self.aux_trunk_lr_pinned = aux_lr;
|
||||
*self.aux_trunk_grad_clip_pinned = aux_clip;
|
||||
*self.aux_trunk_t_pinned = self.aux_trunk_adam_step;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Plan 4 Task 6 Commit B: aux-heads backward ──
|
||||
// Runs the per-head backward kernels, reduces 8 per-tensor partials
|
||||
// into final dW/dB, SAXPYs them into `grad_base[119..127)` with
|
||||
// alpha = `aux_weight`, and SAXPYs both per-head dh_s2 buffers into
|
||||
// the trunk accumulator `bw_d_h_s2` (= `d_h_s2_ptr`). MUST run BEFORE
|
||||
// `encoder_backward_chain` so the trunk dW/dB chain rule sees the
|
||||
// augmented dh_s2 (= main + aux contributions).
|
||||
// alpha = `aux_weight`. SP14-C.5b: SAXPYs both per-head dh_s2_aux
|
||||
// buffers into `dh_s2_aux_accum` (NOT `bw_d_h_s2`); aux trunk
|
||||
// backward propagates that gradient to the aux trunk's own
|
||||
// params; aux trunk Adam launches at the end of the chain.
|
||||
// Q's `bw_d_h_s2` is unaffected by aux loss — encoder boundary
|
||||
// structural stop-grad enforced by `aux_trunk_backward`.
|
||||
self.aux_heads_backward(grad_base)?;
|
||||
|
||||
// SP1 Phase B (slot 34): bw_d_h_s2 snapshot AFTER aux_heads_backward
|
||||
|
||||
@@ -4737,19 +4737,54 @@ impl GpuExperienceCollector {
|
||||
}
|
||||
}
|
||||
|
||||
// Step B: aux next-bar forward. Reads `exp_h_s2_f32`
|
||||
// (n_episodes × shared_h2) populated by `forward_online_f32`
|
||||
// above; writes hidden + logits + softmax tiles. The
|
||||
// softmax tile (`exp_aux_nb_softmax_buf [n_episodes, 2]`)
|
||||
// is the production output consumed by step 4's SP14 EGF
|
||||
// chain. Per `pearl_no_host_branches_in_captured_graph`
|
||||
// this launch is OUTSIDE the captured graph (post
|
||||
// Step B-1 (SP14-C.5b, 2026-05-08): aux trunk forward.
|
||||
// Reads encoder layer-1 output `exp_h_s1_f32`
|
||||
// (n_episodes × shared_h1) and writes the aux trunk
|
||||
// hidden caches + final `exp_h_s2_aux` (n_episodes × SH2).
|
||||
// This SEPARATES the aux head input from Q's `exp_h_s2_f32`
|
||||
// — the aux trunk shapes a representation specifically for
|
||||
// the next-bar prediction objective without leaking
|
||||
// gradient back into Q's encoder. Param tensors come from
|
||||
// the trainer's dedicated aux trunk slabs via
|
||||
// `aux_trunk_{w,b}{1,2,3}_ptr` (set by
|
||||
// `set_trainer_aux_trunk_param_ptrs` once the fused
|
||||
// training context is created).
|
||||
let enc_out_dim = self.network_dims.0;
|
||||
let aux_h1 = super::gpu_aux_trunk::AUX_TRUNK_H1;
|
||||
let aux_h2 = super::gpu_aux_trunk::AUX_TRUNK_H2;
|
||||
self.exp_aux_trunk_forward_ops.launch(
|
||||
&self.stream,
|
||||
self.exp_h_s1_f32.raw_ptr(),
|
||||
self.aux_trunk_w1_ptr,
|
||||
self.aux_trunk_b1_ptr,
|
||||
self.aux_trunk_w2_ptr,
|
||||
self.aux_trunk_b2_ptr,
|
||||
self.aux_trunk_w3_ptr,
|
||||
self.aux_trunk_b3_ptr,
|
||||
self.exp_h_aux1.raw_ptr(),
|
||||
self.exp_h_aux2.raw_ptr(),
|
||||
self.exp_h_s2_aux.raw_ptr(),
|
||||
n,
|
||||
enc_out_dim,
|
||||
aux_h1,
|
||||
aux_h2,
|
||||
sh2,
|
||||
)?;
|
||||
|
||||
// Step B-2: aux next-bar forward. SP14-C.5b: input redirected
|
||||
// from `exp_h_s2_f32` (Q trunk output) to `exp_h_s2_aux`
|
||||
// (aux trunk output, written above). Writes hidden + logits
|
||||
// + softmax tiles. The softmax tile
|
||||
// (`exp_aux_nb_softmax_buf [n_episodes, 2]`) is the
|
||||
// production output consumed by step 4's SP14 EGF chain.
|
||||
// Per `pearl_no_host_branches_in_captured_graph` this
|
||||
// launch is OUTSIDE the captured graph (post
|
||||
// `cuGraphLaunch`) — same placement as the trainer's
|
||||
// `aux_heads_forward` call relative to the trainer's
|
||||
// captured graph.
|
||||
self.exp_aux_heads_fwd.forward_next_bar(
|
||||
&self.stream,
|
||||
self.exp_h_s2_f32.raw_ptr(),
|
||||
self.exp_h_s2_aux.raw_ptr(),
|
||||
nb_w1, nb_b1, nb_w2, nb_b2,
|
||||
n, // batch dim = n_episodes
|
||||
sh2,
|
||||
|
||||
@@ -7504,3 +7504,121 @@ With C.5a + C.5a-fixup providing all scaffolding (buffers, scratch, ptrs, access
|
||||
5. Add aux trunk `t_pinned` to `submit_counters_ops` chain (GPU-side step counter increment).
|
||||
6. Pre-capture host write of ISV-driven `aux_trunk_lr_pinned` / `aux_trunk_grad_clip_pinned` (cold-path per-epoch).
|
||||
7. Redirect `aux_heads_fwd` input pointer from Q's `h_s2` to `h_s2_aux`.
|
||||
|
||||
## SP14 Layer C Phase C.5b — atomic contract migration COMPLETE (2026-05-08)
|
||||
|
||||
> Atomic contract migration: aux heads consume the SEPARATE aux trunk's
|
||||
> `h_s2_aux` (NOT Q's GRN trunk `save_h_s2`); aux head backward writes
|
||||
> `dh_s2_aux` into a NEW accumulator `dh_s2_aux_accum`; aux trunk
|
||||
> backward propagates that gradient through w3/w2/w1 + b3/b2/b1, and aux
|
||||
> trunk Adam updates 6 grad tensors via global L2-norm clip. Encoder
|
||||
> boundary is structural stop-grad (kernel set has NO `dx_in` output).
|
||||
|
||||
### Reverted commits (stop-grad band-aids retired)
|
||||
|
||||
| Commit | What it did | C.5b reversal |
|
||||
|---------|-----------------------------------------------------------|---------------------------------------------------------|
|
||||
| `872bd7392` | Zero-fill `dh_s2_out` in `aux_next_bar_backward` (Step 3) | Restored genuine SAXPY: `dh_s2_aux[b,j] = sum_k sh_dh_pre[k] * w1[k,j]` |
|
||||
| `411a30473` | Same zero-fill in `aux_regime_backward` (Step 3) | Symmetric restoration with same gradient formula |
|
||||
|
||||
The stop-grad audit blocks (lines ~599-638 next-bar + ~757-790 regime)
|
||||
deleted; replaced by short C.5b comments documenting the atomic
|
||||
contract flip.
|
||||
|
||||
### Kernel signature renames (4 functions)
|
||||
|
||||
| Kernel | Old param | New param |
|
||||
|-------------------------|-------------------|-------------------|
|
||||
| `aux_next_bar_forward` | `const float* h_s2` | `const float* h_s2_aux` |
|
||||
| `aux_regime_forward` | `const float* h_s2` | `const float* h_s2_aux` |
|
||||
| `aux_next_bar_backward` | `float* dh_s2_out` | `float* dh_s2_aux_out` |
|
||||
| `aux_regime_backward` | `float* dh_s2_out` | `float* dh_s2_aux_out` |
|
||||
|
||||
The Rust `gpu_aux_heads.rs` wrapper renamed call-site argument names to
|
||||
match (`h_s2_ptr → h_s2_aux_ptr`, `dh_s2_out_ptr → dh_s2_aux_out_ptr`)
|
||||
so the param-name → kernel-arg mapping stays self-documenting at
|
||||
launch sites.
|
||||
|
||||
### Wire sites lit (4)
|
||||
|
||||
| Site | Function | Insert point | What it does |
|
||||
|------|----------|--------------|--------------|
|
||||
| Trainer fwd | `aux_heads_forward` | New Step 0, before regime label builder | `aux_trunk_forward_ops.launch(... encoder_out=save_h_s1, h_s2_aux_out=h_s2_aux)` populates the aux trunk's representation |
|
||||
| Trainer bwd | `aux_heads_backward` | After per-head SAXPY into `dh_s2_aux_accum` | `aux_trunk_backward_ops.launch(... dh_s2_aux_in=accum, w/b grads...)` propagates grad through aux trunk; then `launch_aux_trunk_adam_update` applies the optimizer step |
|
||||
| Trainer bwd | `aux_heads_backward` | After both head backwards | Pre-zero `dh_s2_aux_accum` via `cuMemsetD32Async` then SAXPY both `aux_dh_s2_*_buf` into it (REPLACES the old SAXPY into `bw_d_h_s2`) |
|
||||
| Collector fwd | `submit_dqn_step_loop_cublas` step 2b | Post-`forward_online_f32`, before `forward_next_bar` | `exp_aux_trunk_forward_ops.launch(... encoder_out=exp_h_s1_f32, h_s2_aux_out=exp_h_s2_aux)` then `forward_next_bar` consumes `exp_h_s2_aux` (NOT `exp_h_s2_f32`) |
|
||||
|
||||
### Pre-capture host writes (per-step, `&mut self`)
|
||||
|
||||
In `launch_cublas_backward_to` (BEFORE `aux_heads_backward(grad_base)`):
|
||||
|
||||
```rust
|
||||
let aux_lr = self.read_isv_signal_at(AUX_TRUNK_LR_INDEX); // ISV[444]
|
||||
let aux_clip = self.read_isv_signal_at(AUX_TRUNK_GRAD_CLIP_INDEX); // ISV[448]
|
||||
self.aux_trunk_adam_step += 1;
|
||||
unsafe {
|
||||
*self.aux_trunk_lr_pinned = aux_lr;
|
||||
*self.aux_trunk_grad_clip_pinned = aux_clip;
|
||||
*self.aux_trunk_t_pinned = self.aux_trunk_adam_step;
|
||||
}
|
||||
```
|
||||
|
||||
These writes are SAFE in graph-capture context: `launch_cublas_backward_to`
|
||||
is `&mut self` and runs OUTSIDE the captured region (the captured
|
||||
chain is each `cuGraphLaunch`-replay; this Rust function is the
|
||||
per-step host driver). The Adam kernel inside the captured graph
|
||||
reads the pinned device pointer, which resolves to the host's
|
||||
mapped-pinned memory at replay time — picking up whatever value the
|
||||
host last wrote. Same pattern as `step_ofi_embed_adam`.
|
||||
|
||||
The step counter is host-tracked (mirrors OFI embed) rather than
|
||||
extending `increment_step_counters` because the aux trunk lives
|
||||
outside the SP4 8-group taxonomy and a per-host-call increment is
|
||||
sufficient (no coupling to other Adam schedulers' bias correction).
|
||||
|
||||
### SAXPY redirect — straightforward
|
||||
|
||||
The previous `aux_heads_backward` SAXPYed `aux_dh_s2_nb_buf` +
|
||||
`aux_dh_s2_rg_buf` into `bw_d_h_s2_ptr`. C.5b changes the destination
|
||||
ptr from `bw_d_h_s2_ptr` → `dh_s2_aux_accum_ptr`. The 256-byte SAXPY
|
||||
launch geometry is identical — only the destination buffer changes.
|
||||
A `cuMemsetD32Async` pre-zeroes `dh_s2_aux_accum` each step so the
|
||||
SAXPY establishes per-step value (mirrors the existing `d_h_s2_ptr`
|
||||
zero pattern in the backward chain).
|
||||
|
||||
### Encoder out pointer
|
||||
|
||||
Both trainer + collector use `save_h_s1` / `exp_h_s1_f32` (config.shared_h1
|
||||
= 256) as the aux trunk's input — matches the aux trunk's
|
||||
`ENCODER_OUT_DIM = 256` allocation from C.2. The trainer reads via
|
||||
`self.ptrs.save_h_s1`; the collector reads via `self.exp_h_s1_f32.raw_ptr()`.
|
||||
Both buffers are populated by their respective `forward_online` calls
|
||||
BEFORE the aux trunk forward runs.
|
||||
|
||||
### Verification
|
||||
|
||||
- `SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo check -p ml --tests` — clean (1m02s, only pre-existing warnings).
|
||||
- `SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo test -p ml --test aux_trunk_oracle_tests --test sp14_oracle_tests --release -- --ignored --nocapture` — 12/12 pass (8 aux_trunk + 4 sp14). The aux_trunk gradient check still validates within tolerance (max_rel_err=1.33e-2 vs 2e-2 tol), and `aux_trunk_backward_does_not_write_dx` still passes (kernel source clean of `dx_in`/`dx_in_out`).
|
||||
|
||||
### Wire status — POST-C.5b (production-ready)
|
||||
|
||||
| Component | Status |
|
||||
|-----------|--------|
|
||||
| `aux_trunk_forward_ops` (trainer) | LIVE — called from `aux_heads_forward` before label builder |
|
||||
| `aux_trunk_backward_ops` (trainer) | LIVE — called from `aux_heads_backward` post-head-SAXPY |
|
||||
| `launch_aux_trunk_adam_update` (trainer) | LIVE — called from `aux_heads_backward` post-trunk-backward |
|
||||
| `dh_aux1_pre_scratch` / `dh_aux2_pre_scratch` (trainer) | LIVE — consumed by `aux_trunk_backward.launch` |
|
||||
| `dh_s2_aux_accum` (trainer) | LIVE — pre-zeroed + SAXPY-target each step |
|
||||
| `aux_trunk_lr_pinned` / `aux_trunk_grad_clip_pinned` (trainer) | LIVE — host-written from ISV per step |
|
||||
| `aux_trunk_t_pinned` (trainer) | LIVE — host-incremented per step (OFI embed pattern) |
|
||||
| `exp_aux_trunk_forward_ops` (collector) | LIVE — called post-`forward_online_f32` before next-bar forward |
|
||||
| `exp_h_s2_aux` / `exp_h_aux1` / `exp_h_aux2` (collector) | LIVE — populated by `aux_trunk_forward.launch` |
|
||||
|
||||
### Outstanding for C.6+
|
||||
|
||||
C.5b only wires the trunk forward+backward+Adam mechanics. Phases C.6+
|
||||
(separate commits) layer on:
|
||||
- C.6: `h_s2_aux_rms_ema` producer kernel feeding ISV[449].
|
||||
- C.7-C.8: ISV-driven aux trunk Adam β1/β2/ε.
|
||||
- C.9: synthetic-data smoke + audit close-out + memory pearls.
|
||||
- C.10: L40S 30-epoch validation.
|
||||
|
||||
@@ -962,9 +962,115 @@ EOF
|
||||
|
||||
DO NOT push.
|
||||
|
||||
## Phase C.5a-fixup: Add Missing Scaffolding (DEAD CODE)
|
||||
|
||||
> Authorized 2026-05-08 after C.5b implementer flagged 5 specific gaps (most importantly a graph-capture hazard) plus 2 missing buffer allocations from C.5a. C.5a-fixup is **additive-only** — same discipline as C.5a. Goal: by end of fixup, C.5b is a genuine ~300 LOC atomic flip with zero scaffolding work mixed in.
|
||||
|
||||
### Task C.5a-fixup.1: Allocate missing scratch buffers + collector ptr scaffolding + graph-capture audit
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (allocate `dh_aux1_pre_scratch`, `dh_aux2_pre_scratch`; add `aux_trunk_param_ptrs()` accessor)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` (6× `u64` aux_trunk param ptr fields + `set_trainer_aux_trunk_param_ptrs` setter; new `exp_aux_trunk_forward_ops: AuxTrunkForwardOps` field; constructor populates it; **leave wire sites still calling `save_h_s2` — no semantic change**)
|
||||
- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs` (call new `set_trainer_aux_trunk_param_ptrs` setter at the existing setup point near `set_trainer_params_ptr` — ~lines 1703 + 2317 per implementer report)
|
||||
- Modify: `docs/dqn-wire-up-audit.md` (audit entry + graph-capture verification result)
|
||||
|
||||
- [ ] **Step 1: Allocate missing scratch buffers in trainer**
|
||||
|
||||
`aux_trunk_backward.launch` requires `dh_aux1_pre_scratch_ptr [B_max, AUX_TRUNK_H1=256]` and `dh_aux2_pre_scratch_ptr [B_max, AUX_TRUNK_H2=128]` per `gpu_aux_trunk.rs:266-267`. C.5a missed these. Add as `CudaSlice<f32>` zero-init alongside existing C.5a buffers.
|
||||
|
||||
- [ ] **Step 2: Collector aux trunk ptr scaffolding** (mirrors existing `set_trainer_params_ptr` pattern — find via grep)
|
||||
|
||||
```rust
|
||||
// In GpuExperienceCollector struct:
|
||||
aux_trunk_w1_ptr: u64,
|
||||
aux_trunk_b1_ptr: u64,
|
||||
aux_trunk_w2_ptr: u64,
|
||||
aux_trunk_b2_ptr: u64,
|
||||
aux_trunk_w3_ptr: u64,
|
||||
aux_trunk_b3_ptr: u64,
|
||||
exp_aux_trunk_forward_ops: AuxTrunkForwardOps,
|
||||
|
||||
// Setter:
|
||||
pub fn set_trainer_aux_trunk_param_ptrs(&mut self, w1: u64, b1: u64, w2: u64, b2: u64, w3: u64, b3: u64) {
|
||||
self.aux_trunk_w1_ptr = w1;
|
||||
// ... etc
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Trainer-side accessor**
|
||||
|
||||
```rust
|
||||
// In GpuDqnTrainer:
|
||||
pub fn aux_trunk_param_ptrs(&self) -> (u64, u64, u64, u64, u64, u64) {
|
||||
use cudarc::driver::DevicePtr;
|
||||
(
|
||||
*self.aux_trunk_w1.device_ptr(),
|
||||
*self.aux_trunk_b1.device_ptr(),
|
||||
*self.aux_trunk_w2.device_ptr(),
|
||||
*self.aux_trunk_b2.device_ptr(),
|
||||
*self.aux_trunk_w3.device_ptr(),
|
||||
*self.aux_trunk_b3.device_ptr(),
|
||||
)
|
||||
}
|
||||
```
|
||||
|
||||
(Adjust to match the actual cudarc `device_ptr()` API — find an existing accessor like `embed_w_ptr()` and mirror.)
|
||||
|
||||
- [ ] **Step 4: Wire setter call in training_loop**
|
||||
|
||||
Find the existing call to `set_trainer_params_ptr` (~lines 1703 + 2317 per implementer report). Insert call to `set_trainer_aux_trunk_param_ptrs` at the same site, passing the trainer's accessor result.
|
||||
|
||||
- [ ] **Step 5: Graph-capture audit** (the showstopper concern)
|
||||
|
||||
Read `gpu_dqn_trainer.rs` aux_heads_backward orchestrator (around line 29865 per implementer report). Trace: does the call sit INSIDE a `cuStreamBeginCapture`/`end_capture` region?
|
||||
|
||||
If INSIDE: the existing aux backward CANNOT have host-side writes (step counter, pinned buffer). Verify the existing code by inspecting it — the existing aux_heads_backward should have a graph-capture-compatible structure (probably a separate function called pre-capture for any host writes). Document the existing pattern.
|
||||
|
||||
If OUTSIDE: the new aux_trunk Adam launch can sit alongside existing aux_heads_backward; host writes are fine.
|
||||
|
||||
Document the verification result in audit doc — this determines C.5b's wiring strategy.
|
||||
|
||||
- [ ] **Step 6: Verify**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo check -p ml --tests --all-targets 2>&1 | tail -10
|
||||
SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo test -p ml --test aux_trunk_oracle_tests --test sp14_oracle_tests --release -- --ignored --nocapture 2>&1 | tail -15
|
||||
```
|
||||
|
||||
Expected: clean build; 12 tests still pass (no regression — pure scaffolding).
|
||||
|
||||
- [ ] **Step 7: Commit**
|
||||
|
||||
```bash
|
||||
git add -A
|
||||
git commit -m "$(cat <<'EOF'
|
||||
feat(sp14-c.5a-fixup): missing scratch buffers + collector ptr scaffolding (dead code)
|
||||
|
||||
Phase C.5a-fixup — completes C.5a's additive infrastructure:
|
||||
- Allocate dh_aux1_pre_scratch + dh_aux2_pre_scratch (missed in C.5a)
|
||||
- Collector struct gains 6× aux_trunk param ptr fields + setter
|
||||
- Trainer gains aux_trunk_param_ptrs() accessor
|
||||
- training_loop wires the setter at the existing set_trainer_params_ptr site
|
||||
- Collector constructor populates AuxTrunkForwardOps
|
||||
|
||||
NO contract change. Wire sites still call save_h_s2. The setter and ptrs
|
||||
are populated but unused by aux head consumers (still wired through Q's
|
||||
trunk).
|
||||
|
||||
Graph-capture audit: <result documented in audit doc>
|
||||
|
||||
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
||||
EOF
|
||||
)"
|
||||
```
|
||||
|
||||
DO NOT push.
|
||||
|
||||
---
|
||||
|
||||
## Phase C.5b: Atomic Contract Migration — Aux Head Input Switches h_s2 → h_s2_aux
|
||||
|
||||
> ONE atomic commit per `feedback_no_partial_refactor`. ~300 LOC. THIS is the genuine contract migration. C.5a's additive code becomes live here.
|
||||
> ONE atomic commit per `feedback_no_partial_refactor`. With C.5a + C.5a-fixup providing all scaffolding (buffers, ptrs, accessor, setter, AuxTrunkForwardOps in collector), C.5b is now a genuine ~300 LOC atomic flip: 4 kernel param renames + 2 zero-fill reverts + 4 wire sites + Adam launch + ISV pinned-buffer population.
|
||||
|
||||
### Task C.5b.1: Atomic kernel rename + revert zero-fills + wire-up
|
||||
|
||||
|
||||
Reference in New Issue
Block a user