gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(sp14-c.5b): atomic contract migration — h_s2 → h_s2_aux + revert zero-fills

Browse Source 
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
2026-05-08 03:01:13 +02:00
parent 1edd71a2c1
commit b26b189925
6 changed files with 529 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files

210
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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, &param_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