feat(sp14-c.5a): allocate aux trunk fwd/bwd buffers + Adam launcher (dead code)
Phase C.5a — additive infrastructure for the aux trunk wire-up. Allocates saved-fwd buffers (h_s2_aux, h_aux1, h_aux2), gradient buffers (6× aux_trunk_*_grad), accumulator (dh_s2_aux_accum), and dedicated Adam launcher (launch_aux_trunk_adam_update) reading β1/β2/ε/LR/grad-clip from ISV[444..449). No contract change. No call sites for the new launcher yet. C.5b atomically wires these in. Phase C.5 was split (authorized 2026-05-08) after the original implementer flagged ~600-800 LOC scope across 4 files with correctness windows. C.5a is purely additive; C.5b is the genuine ~300 LOC atomic migration. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -7327,6 +7327,114 @@ pub struct GpuDqnTrainer {
|
||||
/// Adam second moment for `aux_trunk_b3` [SH2].
|
||||
aux_trunk_b3_v: CudaSlice<f32>,
|
||||
|
||||
// ── SP14 Layer C Phase C.5a saved-fwd + grad buffers (2026-05-08) ──
|
||||
// Allocated as DEAD CODE: these buffers are real device allocations
|
||||
// but no production caller reads / writes them yet. Phase C.5b
|
||||
// atomically wires them via the aux trunk forward (writes
|
||||
// h_s2_aux/h_aux1/h_aux2), aux heads backward (SAXPY-accumulates into
|
||||
// dh_s2_aux_accum), aux trunk backward (overwrites the 6 grad
|
||||
// tensors), and `launch_aux_trunk_adam_update` (consumes grads via
|
||||
// global L2-norm clip + per-tensor Adam step).
|
||||
//
|
||||
// 6-buffer choice (vs flat-buffer): mirrors C.2's separate-tensor
|
||||
// params/m/v allocation. Lets each Adam launch consume one
|
||||
// (param, grad, m, v) tuple of matching extent — no byte-offset
|
||||
// arithmetic. Per-tensor grad_norm partial sums into a unified
|
||||
// `aux_trunk_norm_partials` buffer at fixed offsets so the global
|
||||
// L2 norm spans all 6 tensors.
|
||||
/// `[B_max, SH2]` final aux trunk output (saved post-fwd; consumed by
|
||||
/// aux heads forward + backward in C.5b). Same shape contract as
|
||||
/// `aux_dh_s2_nb_buf` so the SAXPY accumulator type-matches.
|
||||
#[allow(dead_code)]
|
||||
h_s2_aux: CudaSlice<f32>,
|
||||
/// `[B_max, AUX_TRUNK_H1=256]` Layer-1 post-ELU activation. Consumed
|
||||
/// by `aux_trunk_bwd_dh_pre` (`ELU'(h_aux1)`) + `aux_trunk_bwd_dW_reduce`
|
||||
/// (Layer-2 `dW2 = h_aux1.T @ dh_aux2_pre`).
|
||||
#[allow(dead_code)]
|
||||
h_aux1: CudaSlice<f32>,
|
||||
/// `[B_max, AUX_TRUNK_H2=128]` Layer-2 post-ELU activation. Consumed
|
||||
/// by `aux_trunk_bwd_dh_pre` (`ELU'(h_aux2)`) + `aux_trunk_bwd_dW_reduce`
|
||||
/// (Layer-3 `dW3 = h_aux2.T @ d_logits`).
|
||||
#[allow(dead_code)]
|
||||
h_aux2: CudaSlice<f32>,
|
||||
/// `[B_max, SH2]` SAXPY accumulator for upstream gradient feeding the
|
||||
/// aux trunk's Layer-3 input (`dh_s2_aux`). C.5b reverts the zero-fill
|
||||
/// hack at `aux_next_bar_backward:599-638` + `aux_regime_backward:757-790`
|
||||
/// so the two aux head backwards SAXPY their per-sample dh into this
|
||||
/// buffer instead of the trunk's `bw_d_h_s2`. Consumed by aux trunk
|
||||
/// backward (`launch(... dh_s2_aux_in_ptr=accum_ptr ...)`).
|
||||
#[allow(dead_code)]
|
||||
dh_s2_aux_accum: CudaSlice<f32>,
|
||||
|
||||
/// `[SH1 × 256]` aux trunk weight-1 gradient accumulator (overwritten
|
||||
/// per training step by `aux_trunk_bwd_dW_reduce`).
|
||||
#[allow(dead_code)]
|
||||
aux_trunk_w1_grad: CudaSlice<f32>,
|
||||
/// `[256]` aux trunk bias-1 gradient (overwritten by `aux_trunk_bwd_db_reduce`).
|
||||
#[allow(dead_code)]
|
||||
aux_trunk_b1_grad: CudaSlice<f32>,
|
||||
/// `[256 × 128]` aux trunk weight-2 gradient.
|
||||
#[allow(dead_code)]
|
||||
aux_trunk_w2_grad: CudaSlice<f32>,
|
||||
/// `[128]` aux trunk bias-2 gradient.
|
||||
#[allow(dead_code)]
|
||||
aux_trunk_b2_grad: CudaSlice<f32>,
|
||||
/// `[128 × SH2]` aux trunk weight-3 gradient.
|
||||
#[allow(dead_code)]
|
||||
aux_trunk_w3_grad: CudaSlice<f32>,
|
||||
/// `[SH2]` aux trunk bias-3 gradient.
|
||||
#[allow(dead_code)]
|
||||
aux_trunk_b3_grad: CudaSlice<f32>,
|
||||
|
||||
/// `[num_blocks_total]` per-block partial sum-of-squares written by
|
||||
/// `dqn_grad_norm_kernel` across all 6 grad tensors (offsets matching
|
||||
/// `aux_trunk_grad_block_offsets`). Sized to the maximum total block
|
||||
/// count any of the 6 grads would produce.
|
||||
#[allow(dead_code)]
|
||||
aux_trunk_grad_norm_partials: CudaSlice<f32>,
|
||||
/// Per-tensor block offsets into `aux_trunk_grad_norm_partials`. The
|
||||
/// 7th entry is the total block count (passed as `num_blocks` to
|
||||
/// `dqn_grad_norm_finalize`). Captured at construction so the launcher
|
||||
/// has no host-side branching during its launches.
|
||||
#[allow(dead_code)]
|
||||
aux_trunk_grad_block_offsets: [usize; 7],
|
||||
/// `[1]` finalised L2 grad norm written by `dqn_grad_norm_finalize`.
|
||||
/// Read by 6 Adam launches as the `grad_norm_f32` arg.
|
||||
#[allow(dead_code)]
|
||||
aux_trunk_grad_norm_buf: CudaSlice<f32>,
|
||||
|
||||
/// Pinned device-mapped grad-clip threshold. Host writes
|
||||
/// `ISV[AUX_TRUNK_GRAD_CLIP_INDEX=448]` per-step before launch; GPU
|
||||
/// reads via `aux_trunk_grad_clip_dev_ptr`. The Adam kernel's
|
||||
/// `max_grad_norm_buf` arg.
|
||||
#[allow(dead_code)]
|
||||
aux_trunk_grad_clip_pinned: *mut f32,
|
||||
/// Device pointer for `aux_trunk_grad_clip_pinned` (one
|
||||
/// `cuMemHostGetDevicePointer_v2` per construction; reusable across
|
||||
/// launches inside captured graphs).
|
||||
#[allow(dead_code)]
|
||||
aux_trunk_grad_clip_dev_ptr: u64,
|
||||
|
||||
/// Pinned device-mapped LR. Host writes `ISV[AUX_TRUNK_LR_INDEX=444]`
|
||||
/// per-step. The Adam kernel's `lr_ptr` arg.
|
||||
#[allow(dead_code)]
|
||||
aux_trunk_lr_pinned: *mut f32,
|
||||
/// Device pointer for `aux_trunk_lr_pinned`.
|
||||
#[allow(dead_code)]
|
||||
aux_trunk_lr_dev_ptr: u64,
|
||||
|
||||
/// Pinned device-mapped Adam step counter (i32). Incremented by host
|
||||
/// before each launch; the kernel's `t_ptr` arg.
|
||||
#[allow(dead_code)]
|
||||
aux_trunk_t_pinned: *mut i32,
|
||||
/// Device pointer for `aux_trunk_t_pinned`.
|
||||
#[allow(dead_code)]
|
||||
aux_trunk_t_dev_ptr: u64,
|
||||
/// Host-side mirror of the Adam step counter. Production caller in
|
||||
/// C.5b increments this and writes through to the pinned buffer.
|
||||
#[allow(dead_code)]
|
||||
aux_trunk_adam_step: i32,
|
||||
|
||||
/// SP14 Layer C Phase C.3 (2026-05-08): aux trunk forward orchestrator.
|
||||
/// Owns the pre-loaded `CudaFunction` handle for `aux_trunk_forward`
|
||||
/// per `pearl_no_host_branches_in_captured_graph.md` — the launch
|
||||
@@ -7911,6 +8019,16 @@ impl Drop for GpuDqnTrainer {
|
||||
if !self.ofi_embed_t_pinned.is_null() {
|
||||
let _ = unsafe { cudarc::driver::result::free_host(self.ofi_embed_t_pinned.cast()) };
|
||||
}
|
||||
// SP14 Layer C Phase C.5a (2026-05-08): aux trunk Adam pinned ptrs.
|
||||
if !self.aux_trunk_grad_clip_pinned.is_null() {
|
||||
let _ = unsafe { cudarc::driver::result::free_host(self.aux_trunk_grad_clip_pinned.cast()) };
|
||||
}
|
||||
if !self.aux_trunk_lr_pinned.is_null() {
|
||||
let _ = unsafe { cudarc::driver::result::free_host(self.aux_trunk_lr_pinned.cast()) };
|
||||
}
|
||||
if !self.aux_trunk_t_pinned.is_null() {
|
||||
let _ = unsafe { cudarc::driver::result::free_host(self.aux_trunk_t_pinned.cast()) };
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -9691,6 +9809,244 @@ impl GpuDqnTrainer {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// SP14 Layer C Phase C.5a (2026-05-08): aux trunk Adam optimizer step.
|
||||
///
|
||||
/// **DEAD CODE.** No production caller invokes this in C.5a — the
|
||||
/// launcher exists so the build stays green while C.5b is the
|
||||
/// genuine atomic contract migration that wires the aux head
|
||||
/// backward → `dh_s2_aux_accum` → aux trunk backward → THIS launcher
|
||||
/// chain. Per `feedback_no_partial_refactor.md` the wire-up is one
|
||||
/// commit. Per `feedback_no_stubs.md` the body is fully real (no
|
||||
/// placeholder panic); a synthetic call from a unit test would
|
||||
/// execute end-to-end correctly given populated grad buffers.
|
||||
///
|
||||
/// Reads ISV-driven hyper-parameters (no hardcoded constants):
|
||||
///
|
||||
/// | ISV slot | Use |
|
||||
/// |----------------------------------|---------------------------|
|
||||
/// | `AUX_TRUNK_LR_INDEX = 444` | Adam learning rate |
|
||||
/// | `AUX_TRUNK_BETA1_INDEX = 445` | Adam β1 (first moment) |
|
||||
/// | `AUX_TRUNK_BETA2_INDEX = 446` | Adam β2 (second moment) |
|
||||
/// | `AUX_TRUNK_EPS_INDEX = 447` | Adam ε (denom stability) |
|
||||
/// | `AUX_TRUNK_GRAD_CLIP_INDEX = 448`| Global L2 clip threshold |
|
||||
///
|
||||
/// Mirrors the OFI embed Adam pattern (`step_ofi_embed_adam`) but
|
||||
/// over 6 separate (params, grad, m, v) tensor tuples instead of one
|
||||
/// flat buffer. The aux trunk lives outside the SP4 8-group
|
||||
/// taxonomy (Q/Target/CQL/Aux per-branch + auxiliaries) so a
|
||||
/// dedicated launcher is appropriate; SP4 Mech 9 weight clamp
|
||||
/// (`weight_clamp_max_abs`) and Pearl C engagement counter
|
||||
/// (`engage_buf_offset`) are disabled — same convention as OFI
|
||||
/// embed.
|
||||
///
|
||||
/// Phase 1: per-tensor partial sum-of-squares into the unified
|
||||
/// `aux_trunk_grad_norm_partials` buffer at fixed offsets
|
||||
/// (`aux_trunk_grad_block_offsets[i..i+1]` slice covers tensor `i`'s
|
||||
/// blocks).
|
||||
///
|
||||
/// Phase 2: a single `dqn_grad_norm_finalize` reduces all blocks to
|
||||
/// `aux_trunk_grad_norm_buf[0] = sqrt(sum_b partial_b)`.
|
||||
///
|
||||
/// Phase 3: 6 `dqn_adam_update_kernel` launches, each consuming one
|
||||
/// (params, grad, m, v) tuple. The kernel's internal
|
||||
/// `clip_scale = (norm > clip) ? clip/norm : 1` divides by the
|
||||
/// SAME global norm (read by every launch), so clipping is
|
||||
/// consistent across tensors.
|
||||
///
|
||||
/// `batch_size` is unused by this launcher (Adam is per-parameter,
|
||||
/// not per-batch). Kept in the signature so call-site code reads
|
||||
/// uniformly with other Adam launchers in this trainer.
|
||||
///
|
||||
/// Per `pearl_no_host_branches_in_captured_graph.md` the launcher
|
||||
/// only touches pre-loaded `CudaFunction` handles
|
||||
/// (`grad_norm_standalone_post_aux`, `grad_norm_finalize_post_aux`,
|
||||
/// `adam_update_post_aux`) and pre-allocated buffers — no
|
||||
/// `load_cubin` / `load_function` on the hot path.
|
||||
#[allow(dead_code)]
|
||||
pub(crate) fn launch_aux_trunk_adam_update(
|
||||
&self,
|
||||
stream: &Arc<CudaStream>,
|
||||
_batch_size: i32,
|
||||
) -> Result<(), MLError> {
|
||||
use crate::cuda_pipeline::sp14_isv_slots::{
|
||||
AUX_TRUNK_BETA1_INDEX, AUX_TRUNK_BETA2_INDEX, AUX_TRUNK_EPS_INDEX,
|
||||
};
|
||||
|
||||
// Adam hparams — ISV-driven with cold-start floors per
|
||||
// `pearl_first_observation_bootstrap.md` numerical-stability
|
||||
// carve-outs (β1≥0.5 avoids momentum inversion; β2≥0.9 prevents
|
||||
// exploding second-moment; ε≥1e-12 avoids div-by-zero).
|
||||
let beta1 = self
|
||||
.read_isv_signal_at(AUX_TRUNK_BETA1_INDEX)
|
||||
.max(0.5_f32)
|
||||
.min(0.9999_f32);
|
||||
let beta2 = self
|
||||
.read_isv_signal_at(AUX_TRUNK_BETA2_INDEX)
|
||||
.max(0.9_f32)
|
||||
.min(0.99999_f32);
|
||||
let epsilon = self.read_isv_signal_at(AUX_TRUNK_EPS_INDEX).max(1e-12_f32);
|
||||
let weight_decay: f32 = 0.0; // aux trunk is outside AdamW WD policy
|
||||
|
||||
// Per-tensor (params, grad, m, v, extent) tuples in the same
|
||||
// order as `aux_trunk_grad_block_offsets`.
|
||||
let tensors: [(u64, u64, u64, u64, usize); 6] = [
|
||||
(
|
||||
self.aux_trunk_w1.raw_ptr(),
|
||||
self.aux_trunk_w1_grad.raw_ptr(),
|
||||
self.aux_trunk_w1_m.raw_ptr(),
|
||||
self.aux_trunk_w1_v.raw_ptr(),
|
||||
self.aux_trunk_w1.len(),
|
||||
),
|
||||
(
|
||||
self.aux_trunk_b1.raw_ptr(),
|
||||
self.aux_trunk_b1_grad.raw_ptr(),
|
||||
self.aux_trunk_b1_m.raw_ptr(),
|
||||
self.aux_trunk_b1_v.raw_ptr(),
|
||||
self.aux_trunk_b1.len(),
|
||||
),
|
||||
(
|
||||
self.aux_trunk_w2.raw_ptr(),
|
||||
self.aux_trunk_w2_grad.raw_ptr(),
|
||||
self.aux_trunk_w2_m.raw_ptr(),
|
||||
self.aux_trunk_w2_v.raw_ptr(),
|
||||
self.aux_trunk_w2.len(),
|
||||
),
|
||||
(
|
||||
self.aux_trunk_b2.raw_ptr(),
|
||||
self.aux_trunk_b2_grad.raw_ptr(),
|
||||
self.aux_trunk_b2_m.raw_ptr(),
|
||||
self.aux_trunk_b2_v.raw_ptr(),
|
||||
self.aux_trunk_b2.len(),
|
||||
),
|
||||
(
|
||||
self.aux_trunk_w3.raw_ptr(),
|
||||
self.aux_trunk_w3_grad.raw_ptr(),
|
||||
self.aux_trunk_w3_m.raw_ptr(),
|
||||
self.aux_trunk_w3_v.raw_ptr(),
|
||||
self.aux_trunk_w3.len(),
|
||||
),
|
||||
(
|
||||
self.aux_trunk_b3.raw_ptr(),
|
||||
self.aux_trunk_b3_grad.raw_ptr(),
|
||||
self.aux_trunk_b3_m.raw_ptr(),
|
||||
self.aux_trunk_b3_v.raw_ptr(),
|
||||
self.aux_trunk_b3.len(),
|
||||
),
|
||||
];
|
||||
|
||||
let partials_base = self.aux_trunk_grad_norm_partials.raw_ptr();
|
||||
let f32_sz = std::mem::size_of::<f32>() as u64;
|
||||
|
||||
// ── Phase 1: per-tensor partial sum-of-squares ────────────────
|
||||
// Each launch writes `extent.div_ceil(256)` blocks into the
|
||||
// partials buffer at the per-tensor offset.
|
||||
for (i, &(_p, grad_ptr, _m, _v, extent)) in tensors.iter().enumerate() {
|
||||
let blocks = extent.div_ceil(256).max(1) as u32;
|
||||
let off = self.aux_trunk_grad_block_offsets[i] as u64;
|
||||
let slot_ptr = partials_base + off * f32_sz;
|
||||
let n_i32 = extent as i32;
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(&self.grad_norm_standalone_post_aux)
|
||||
.arg(&grad_ptr)
|
||||
.arg(&slot_ptr)
|
||||
.arg(&n_i32)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| {
|
||||
MLError::ModelError(format!(
|
||||
"aux_trunk grad_norm[tensor={}, blocks={}]: {e}",
|
||||
i, blocks
|
||||
))
|
||||
})?;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Phase 2: finalize across all 6 tensors' blocks ────────────
|
||||
let total_blocks = self.aux_trunk_grad_block_offsets[6] as i32;
|
||||
let norm_out_ptr = self.aux_trunk_grad_norm_buf.raw_ptr();
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(&self.grad_norm_finalize_post_aux)
|
||||
.arg(&partials_base)
|
||||
.arg(&norm_out_ptr)
|
||||
.arg(&total_blocks)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| {
|
||||
MLError::ModelError(format!("aux_trunk grad_norm_finalize: {e}"))
|
||||
})?;
|
||||
}
|
||||
|
||||
// ── Phase 3: per-tensor Adam updates ──────────────────────────
|
||||
// All 6 launches consume the SAME `norm_out_ptr` so the kernel's
|
||||
// internal `clip_scale = (norm > max_grad_norm) ? max_grad_norm/norm : 1`
|
||||
// applies a globally consistent scale. `aux_trunk_grad_clip_dev_ptr`
|
||||
// is the ISV-driven threshold (host writes per-step).
|
||||
let lr_ptr = self.aux_trunk_lr_dev_ptr;
|
||||
let clip_ptr = self.aux_trunk_grad_clip_dev_ptr;
|
||||
let t_ptr = self.aux_trunk_t_dev_ptr;
|
||||
let wd_mask_ptr = self.weight_decay_mask.raw_ptr();
|
||||
let l1_end: i32 = 0;
|
||||
let l1_lambda: f32 = 0.0;
|
||||
let weight_clamp_max_abs: f32 = 0.0; // SP4 outside-taxonomy: disabled
|
||||
// SP4 Task A14: aux trunk lives outside the 8-group taxonomy —
|
||||
// pass null nan_flags + disabled engage offset (mirrors OFI embed).
|
||||
let nan_flags_null: u64 = 0;
|
||||
let diag_slot_disabled: i32 = -1;
|
||||
let engage_buf_null: u64 = 0;
|
||||
let engage_off_disabled: i32 = SP4_ENGAGE_OFFSET_DISABLED;
|
||||
|
||||
for (i, &(params_ptr, grad_ptr, m_ptr, v_ptr, extent)) in tensors.iter().enumerate() {
|
||||
let n_i32 = extent as i32;
|
||||
let blocks = (extent.div_ceil(256)).max(1) as u32;
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(&self.adam_update_post_aux)
|
||||
.arg(¶ms_ptr)
|
||||
.arg(&grad_ptr)
|
||||
.arg(&m_ptr)
|
||||
.arg(&v_ptr)
|
||||
.arg(&norm_out_ptr) // global L2 norm (shared across all 6)
|
||||
.arg(&lr_ptr)
|
||||
.arg(&beta1)
|
||||
.arg(&beta2)
|
||||
.arg(&epsilon)
|
||||
.arg(&weight_decay)
|
||||
.arg(&clip_ptr) // ISV-driven global clip threshold
|
||||
.arg(&t_ptr)
|
||||
.arg(&n_i32)
|
||||
.arg(&wd_mask_ptr) // unused (wd=0); kernel still reads — pass real ptr
|
||||
.arg(&l1_end)
|
||||
.arg(&l1_lambda)
|
||||
.arg(&weight_clamp_max_abs)
|
||||
.arg(&nan_flags_null)
|
||||
.arg(&diag_slot_disabled)
|
||||
.arg(&engage_buf_null)
|
||||
.arg(&engage_off_disabled)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| {
|
||||
MLError::ModelError(format!(
|
||||
"aux_trunk adam_update[tensor={}, n={}]: {e}",
|
||||
i, extent
|
||||
))
|
||||
})?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Recursive confidence backward: MSE loss gradient into trunk + conf weight gradients.
|
||||
/// Accumulates into grad_buf (same buffer Adam reads) and bw_d_h_s2 trunk gradient.
|
||||
pub(crate) fn launch_recursive_confidence_backward(&self, batch_size: usize) -> Result<(), MLError> {
|
||||
@@ -22282,6 +22638,121 @@ impl GpuDqnTrainer {
|
||||
aux_trunk_in_dim, aux_trunk_h1, aux_trunk_h2, aux_trunk_out_dim, aux_trunk_total_params,
|
||||
);
|
||||
|
||||
// ── SP14 Layer C Phase C.5a additive infrastructure (2026-05-08) ──
|
||||
// Saved-fwd buffers (h_s2_aux, h_aux1, h_aux2), upstream-grad
|
||||
// accumulator (dh_s2_aux_accum), 6 grad tensors mirroring C.2
|
||||
// params, plus grad-norm scratch + ISV-driven LR/clip pinned
|
||||
// pointers + Adam step counter. ALL DEAD until C.5b atomically
|
||||
// wires the aux head backward → accumulator → aux trunk backward
|
||||
// → Adam launcher chain. No production caller reads these yet.
|
||||
let h_s2_aux = alloc_f32(&stream, b * aux_trunk_out_dim, "h_s2_aux")?;
|
||||
let h_aux1 = alloc_f32(&stream, b * aux_trunk_h1, "h_aux1")?;
|
||||
let h_aux2 = alloc_f32(&stream, b * aux_trunk_h2, "h_aux2")?;
|
||||
let dh_s2_aux_accum = alloc_f32(&stream, b * aux_trunk_out_dim, "dh_s2_aux_accum")?;
|
||||
|
||||
let aux_trunk_w1_grad = alloc_f32(&stream, aux_trunk_w1_count, "aux_trunk_w1_grad")?;
|
||||
let aux_trunk_b1_grad = alloc_f32(&stream, aux_trunk_h1, "aux_trunk_b1_grad")?;
|
||||
let aux_trunk_w2_grad = alloc_f32(&stream, aux_trunk_w2_count, "aux_trunk_w2_grad")?;
|
||||
let aux_trunk_b2_grad = alloc_f32(&stream, aux_trunk_h2, "aux_trunk_b2_grad")?;
|
||||
let aux_trunk_w3_grad = alloc_f32(&stream, aux_trunk_w3_count, "aux_trunk_w3_grad")?;
|
||||
let aux_trunk_b3_grad = alloc_f32(&stream, aux_trunk_out_dim, "aux_trunk_b3_grad")?;
|
||||
|
||||
// Per-tensor block counts for the dqn_grad_norm_kernel launches
|
||||
// (block_dim=256). One partial slot per (block × tensor) pair —
|
||||
// we lay them out in a contiguous buffer with fixed offsets so
|
||||
// the finalize kernel sees them as a single `[total_blocks]`
|
||||
// array. Order matches the 6 (param, grad, m, v) tuples consumed
|
||||
// by the Adam launcher: w1, b1, w2, b2, w3, b3.
|
||||
let aux_trunk_grad_extents: [usize; 6] = [
|
||||
aux_trunk_w1_count,
|
||||
aux_trunk_h1,
|
||||
aux_trunk_w2_count,
|
||||
aux_trunk_h2,
|
||||
aux_trunk_w3_count,
|
||||
aux_trunk_out_dim,
|
||||
];
|
||||
let mut aux_trunk_grad_block_offsets: [usize; 7] = [0; 7];
|
||||
let mut total_grad_blocks: usize = 0;
|
||||
for (i, &ext) in aux_trunk_grad_extents.iter().enumerate() {
|
||||
aux_trunk_grad_block_offsets[i] = total_grad_blocks;
|
||||
let blocks = ext.div_ceil(256);
|
||||
total_grad_blocks += blocks;
|
||||
}
|
||||
aux_trunk_grad_block_offsets[6] = total_grad_blocks;
|
||||
let aux_trunk_grad_norm_partials =
|
||||
alloc_f32(&stream, total_grad_blocks.max(1), "aux_trunk_grad_norm_partials")?;
|
||||
let aux_trunk_grad_norm_buf =
|
||||
alloc_f32(&stream, 1, "aux_trunk_grad_norm_buf")?;
|
||||
|
||||
// Pinned device-mapped grad-clip threshold. Host writes
|
||||
// ISV[AUX_TRUNK_GRAD_CLIP_INDEX] before each launch; the Adam
|
||||
// kernel reads via the device pointer. Cold-start value zero
|
||||
// means clip disabled (kernel's `(norm > clip) ? clip/norm : 1.0`
|
||||
// path; with clip=0 the second branch fires, dividing by zero —
|
||||
// C.5b is responsible for writing a real ISV value before the
|
||||
// first launch goes live).
|
||||
let aux_trunk_grad_clip_pinned: *mut f32 = unsafe {
|
||||
let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP;
|
||||
cudarc::driver::result::malloc_host(std::mem::size_of::<f32>(), flags)
|
||||
.map_err(|e| MLError::ModelError(format!("pinned aux_trunk_grad_clip alloc: {e}")))?
|
||||
as *mut f32
|
||||
};
|
||||
unsafe { *aux_trunk_grad_clip_pinned = 0.0_f32; }
|
||||
let aux_trunk_grad_clip_dev_ptr = unsafe {
|
||||
let mut dp: u64 = 0;
|
||||
cudarc::driver::sys::cuMemHostGetDevicePointer_v2(
|
||||
&mut dp as *mut u64,
|
||||
aux_trunk_grad_clip_pinned.cast(),
|
||||
0,
|
||||
);
|
||||
dp
|
||||
};
|
||||
|
||||
// Pinned device-mapped LR (host writes ISV[AUX_TRUNK_LR_INDEX]
|
||||
// each step). Cold-start zero — C.5b initialises before first
|
||||
// launch.
|
||||
let aux_trunk_lr_pinned: *mut f32 = unsafe {
|
||||
let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP;
|
||||
cudarc::driver::result::malloc_host(std::mem::size_of::<f32>(), flags)
|
||||
.map_err(|e| MLError::ModelError(format!("pinned aux_trunk_lr alloc: {e}")))?
|
||||
as *mut f32
|
||||
};
|
||||
unsafe { *aux_trunk_lr_pinned = 0.0_f32; }
|
||||
let aux_trunk_lr_dev_ptr = unsafe {
|
||||
let mut dp: u64 = 0;
|
||||
cudarc::driver::sys::cuMemHostGetDevicePointer_v2(
|
||||
&mut dp as *mut u64,
|
||||
aux_trunk_lr_pinned.cast(),
|
||||
0,
|
||||
);
|
||||
dp
|
||||
};
|
||||
|
||||
// Pinned device-mapped Adam step counter. Host increments before
|
||||
// each Adam launch; the kernel's bias-correction reads `*t_ptr`.
|
||||
let aux_trunk_t_pinned: *mut i32 = unsafe {
|
||||
let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP;
|
||||
cudarc::driver::result::malloc_host(std::mem::size_of::<i32>(), flags)
|
||||
.map_err(|e| MLError::ModelError(format!("pinned aux_trunk_t alloc: {e}")))?
|
||||
as *mut i32
|
||||
};
|
||||
unsafe { *aux_trunk_t_pinned = 0; }
|
||||
let aux_trunk_t_dev_ptr = unsafe {
|
||||
let mut dp: u64 = 0;
|
||||
cudarc::driver::sys::cuMemHostGetDevicePointer_v2(
|
||||
&mut dp as *mut u64,
|
||||
aux_trunk_t_pinned.cast(),
|
||||
0,
|
||||
);
|
||||
dp
|
||||
};
|
||||
info!(
|
||||
"SP14-C.5a: aux trunk fwd/bwd buffers + Adam infra allocated \
|
||||
(h_s2_aux/h_aux1/h_aux2 + dh_s2_aux_accum, 6 grad tensors, \
|
||||
{} grad-norm partial slots, ISV-driven LR/clip + step counter; DEAD until C.5b)",
|
||||
total_grad_blocks
|
||||
);
|
||||
|
||||
// SP14 Layer C Phase C.3 (2026-05-08): aux trunk forward orchestrator.
|
||||
// Pre-loads the `aux_trunk_forward` CudaFunction handle once at
|
||||
// construction (per `pearl_no_host_branches_in_captured_graph.md`).
|
||||
@@ -23111,6 +23582,27 @@ impl GpuDqnTrainer {
|
||||
aux_trunk_w3_v,
|
||||
aux_trunk_b3_m,
|
||||
aux_trunk_b3_v,
|
||||
// SP14 Layer C Phase C.5a (2026-05-08): saved-fwd + accumulator + grad buffers + Adam infra (DEAD).
|
||||
h_s2_aux,
|
||||
h_aux1,
|
||||
h_aux2,
|
||||
dh_s2_aux_accum,
|
||||
aux_trunk_w1_grad,
|
||||
aux_trunk_b1_grad,
|
||||
aux_trunk_w2_grad,
|
||||
aux_trunk_b2_grad,
|
||||
aux_trunk_w3_grad,
|
||||
aux_trunk_b3_grad,
|
||||
aux_trunk_grad_norm_partials,
|
||||
aux_trunk_grad_block_offsets,
|
||||
aux_trunk_grad_norm_buf,
|
||||
aux_trunk_grad_clip_pinned,
|
||||
aux_trunk_grad_clip_dev_ptr,
|
||||
aux_trunk_lr_pinned,
|
||||
aux_trunk_lr_dev_ptr,
|
||||
aux_trunk_t_pinned,
|
||||
aux_trunk_t_dev_ptr,
|
||||
aux_trunk_adam_step: 0,
|
||||
// SP14 Layer C Phase C.3 (2026-05-08): aux trunk forward orchestrator.
|
||||
aux_trunk_forward_ops,
|
||||
// SP14 Layer C Phase C.4 (2026-05-08): aux trunk backward orchestrator.
|
||||
|
||||
@@ -992,6 +992,29 @@ pub struct GpuExperienceCollector {
|
||||
/// observability of the rollout-time label distribution).
|
||||
exp_aux_dir_acc_buf: super::mapped_pinned::MappedF32Buffer,
|
||||
|
||||
/// SP14 Layer C Phase C.5a (2026-05-08): rollout-sized aux trunk
|
||||
/// final output (`h_s2_aux`). Sized `[alloc_episodes × shared_h2]`
|
||||
/// to mirror the trainer's `h_s2_aux [B_max, SH2]`. C.5b atomically
|
||||
/// inserts `aux_trunk_forward.launch(... h_s2_aux_out=this ...)`
|
||||
/// post-`forward_online_f32` so the collector-native EGF chain
|
||||
/// reads from the aux representation pipeline instead of the
|
||||
/// Q-trunk's `h_s2`. DEAD until C.5b.
|
||||
#[allow(dead_code)]
|
||||
exp_h_s2_aux: cudarc::driver::CudaSlice<f32>,
|
||||
/// SP14 Layer C Phase C.5a: rollout-sized aux trunk Layer-1 post-ELU
|
||||
/// activation. Sized `[alloc_episodes × AUX_TRUNK_H1=256]`. Saved
|
||||
/// for parity with the trainer's per-step backward path; the
|
||||
/// collector itself only runs forward (no backward), but the buffer
|
||||
/// is allocated so the same `aux_trunk_forward.launch` signature
|
||||
/// works in both contexts (collector vs trainer). DEAD until C.5b.
|
||||
#[allow(dead_code)]
|
||||
exp_h_aux1: cudarc::driver::CudaSlice<f32>,
|
||||
/// SP14 Layer C Phase C.5a: rollout-sized aux trunk Layer-2 post-ELU
|
||||
/// activation. Sized `[alloc_episodes × AUX_TRUNK_H2=128]`. Saved
|
||||
/// for forward-API parity; see `exp_h_aux1` rationale.
|
||||
#[allow(dead_code)]
|
||||
exp_h_aux2: cudarc::driver::CudaSlice<f32>,
|
||||
|
||||
/// SP14 β-migration step 3 (2026-05-07): per-rollout-step label
|
||||
/// producer kernel handle. Sibling of the trajectory-wide
|
||||
/// `aux_sign_label_kernel` (line 3776 above) — same per-thread O(1)
|
||||
@@ -2022,6 +2045,34 @@ impl GpuExperienceCollector {
|
||||
"sp14-β: alloc exp_aux_dir_acc_buf (6 f32, collector): {e}"
|
||||
)))?;
|
||||
|
||||
// SP14 Layer C Phase C.5a (2026-05-08): rollout-sized aux trunk
|
||||
// saved-fwd buffers (h_s2_aux, h_aux1, h_aux2). Sized to
|
||||
// `alloc_episodes` × per-layer dim (mirrors the β-migration
|
||||
// pattern: replay-batch-sized in trainer, rollout-batch-sized
|
||||
// here). DEAD until C.5b atomically inserts the
|
||||
// `aux_trunk_forward.launch` post-`forward_online_f32` and
|
||||
// redirects the EGF chain from Q's `h_s2` to `h_s2_aux`. No
|
||||
// production caller reads these in C.5a.
|
||||
use super::gpu_aux_trunk::{AUX_TRUNK_H1, AUX_TRUNK_H2};
|
||||
let exp_h_s2_aux = stream
|
||||
.alloc_zeros::<f32>(alloc_episodes * shared_h2)
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"sp14-c.5a: alloc exp_h_s2_aux ({} f32, collector): {e}",
|
||||
alloc_episodes * shared_h2
|
||||
)))?;
|
||||
let exp_h_aux1 = stream
|
||||
.alloc_zeros::<f32>(alloc_episodes * AUX_TRUNK_H1)
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"sp14-c.5a: alloc exp_h_aux1 ({} f32, collector): {e}",
|
||||
alloc_episodes * AUX_TRUNK_H1
|
||||
)))?;
|
||||
let exp_h_aux2 = stream
|
||||
.alloc_zeros::<f32>(alloc_episodes * AUX_TRUNK_H2)
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"sp14-c.5a: alloc exp_h_aux2 ({} f32, collector): {e}",
|
||||
alloc_episodes * AUX_TRUNK_H2
|
||||
)))?;
|
||||
|
||||
// SP14 β-migration step 3 (2026-05-07): per-rollout-step label
|
||||
// producer kernel. Loaded on the collector's stream so its launch
|
||||
// chains directly after `forward_online_f32` (which populates
|
||||
@@ -2472,6 +2523,11 @@ impl GpuExperienceCollector {
|
||||
exp_aux_nb_softmax_buf,
|
||||
exp_aux_nb_label_buf,
|
||||
exp_aux_dir_acc_buf,
|
||||
// SP14 Layer C Phase C.5a (2026-05-08): rollout-sized aux trunk
|
||||
// saved-fwd buffers (DEAD until C.5b atomic wire-up).
|
||||
exp_h_s2_aux,
|
||||
exp_h_aux1,
|
||||
exp_h_aux2,
|
||||
exp_aux_sign_label_per_step_kernel,
|
||||
// SP15 Wave 5 follow-up (2026-05-07): pre-loaded SP15 kernel handles.
|
||||
exp_sp15_dd_state_kernel,
|
||||
|
||||
@@ -7311,3 +7311,102 @@ Both slots use Pearl-A first-observation bootstrap. Slot 451 has α=0.05 EMA; sl
|
||||
**Open follow-ups:**
|
||||
- Add a target-variance EMA for `AVG_WIN_HOLD_TIME_BARS_INDEX` (Welford-style, single new ISV slot) in a future phase if the fixed α=0.01 blend is too slow/fast on real-data validation. The producer kernel already has the variance-slot path wired; pass `isv_h_target_var_idx >= 0` to enable Wiener-α.
|
||||
- Pearl note for `pearl_event_driven_reward_density_alignment.md` cross-reference: this phase aligns the AUX *label density* to event density (winning trade close), the same principle as the SP12 reward-density alignment.
|
||||
|
||||
## SP14 Layer C Phase C.5a — additive infrastructure (DEAD CODE) (2026-05-08)
|
||||
|
||||
> Phase C.5 was split (authorized 2026-05-08) after the prior implementer flagged ~600-800 LOC scope across 4 files with three correctness windows (gradient pointer aliasing, kernel rename + revert atomicity, Adam launcher dependency on grad buffer layout). C.5a is purely additive — buffers + launcher allocated as dead code. Build is clean before, build is clean after. C.5b atomically wires every consumer per `feedback_no_partial_refactor`.
|
||||
|
||||
**No contract change in this commit.** Buffers exist, launcher compiles, no production caller invokes either.
|
||||
|
||||
### Buffers allocated
|
||||
|
||||
In `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`:
|
||||
|
||||
**Saved-fwd tiles (replay-batch-sized):**
|
||||
- `h_s2_aux: CudaSlice<f32>` — `[B, SH2]` final aux trunk output. Shape matches `aux_dh_s2_nb_buf` so the C.5b SAXPY-back step type-matches.
|
||||
- `h_aux1: CudaSlice<f32>` — `[B, AUX_TRUNK_H1=256]` Layer-1 post-ELU activation.
|
||||
- `h_aux2: CudaSlice<f32>` — `[B, AUX_TRUNK_H2=128]` Layer-2 post-ELU activation.
|
||||
|
||||
**Upstream-grad accumulator:**
|
||||
- `dh_s2_aux_accum: CudaSlice<f32>` — `[B, SH2]`. Mirrors `bw_d_h_s2`'s role for Q's path. C.5b reverts the zero-fills at `aux_next_bar_backward:599-638` + `aux_regime_backward:757-790` (commits `872bd7392` + `411a30473`) so the two aux head backwards SAXPY their per-sample dh into THIS buffer instead of the trunk's `bw_d_h_s2`.
|
||||
|
||||
**Aux trunk gradient tensors (6, mirroring C.2 params):**
|
||||
- `aux_trunk_w1_grad` `[SH1 × 256]`, `aux_trunk_b1_grad` `[256]`,
|
||||
- `aux_trunk_w2_grad` `[256 × 128]`, `aux_trunk_b2_grad` `[128]`,
|
||||
- `aux_trunk_w3_grad` `[128 × SH2]`, `aux_trunk_b3_grad` `[SH2]`.
|
||||
- All `alloc_zeros`-initialised. Overwritten (not accumulated) per training step by `AuxTrunkBackwardOps::launch` in C.5b — matching the documented backward-API contract.
|
||||
|
||||
**Adam launcher scratch + ISV-driven hparams:**
|
||||
- `aux_trunk_grad_norm_partials: CudaSlice<f32>` `[total_blocks]` — per-block sum-of-squares slots laid out per `aux_trunk_grad_block_offsets[7]` (one offset per tensor + total).
|
||||
- `aux_trunk_grad_norm_buf: CudaSlice<f32>` `[1]` — finalised L2 norm read by every Adam launch as `grad_norm_f32`.
|
||||
- `aux_trunk_grad_clip_pinned: *mut f32` (mapped device-pinned) + `aux_trunk_grad_clip_dev_ptr` — host writes `ISV[AUX_TRUNK_GRAD_CLIP_INDEX=448]` per-step.
|
||||
- `aux_trunk_lr_pinned: *mut f32` + `aux_trunk_lr_dev_ptr` — host writes `ISV[AUX_TRUNK_LR_INDEX=444]` per-step.
|
||||
- `aux_trunk_t_pinned: *mut i32` + `aux_trunk_t_dev_ptr` — host increments before each launch.
|
||||
- `aux_trunk_adam_step: i32` — host-side mirror.
|
||||
- `Drop` impl frees the three new pinned host allocations.
|
||||
|
||||
In `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`:
|
||||
- `exp_h_s2_aux: CudaSlice<f32>` — `[alloc_episodes × shared_h2]`.
|
||||
- `exp_h_aux1: CudaSlice<f32>` — `[alloc_episodes × AUX_TRUNK_H1=256]`.
|
||||
- `exp_h_aux2: CudaSlice<f32>` — `[alloc_episodes × AUX_TRUNK_H2=128]`.
|
||||
|
||||
Mirrors the trainer's saved-fwd buffers but rollout-sized — same β-migration pattern as `exp_aux_nb_hidden_buf`.
|
||||
|
||||
### Adam launcher signature
|
||||
|
||||
```rust
|
||||
pub(crate) fn launch_aux_trunk_adam_update(
|
||||
&self,
|
||||
stream: &Arc<CudaStream>,
|
||||
_batch_size: i32,
|
||||
) -> Result<(), MLError>
|
||||
```
|
||||
|
||||
Writes 6 per-tensor `dqn_grad_norm_kernel` partials → 1 `dqn_grad_norm_finalize` (sqrt of sum across all 132K params) → 6 `dqn_adam_update_kernel` launches consuming the same global L2 norm. ISV-driven β1 (slot 445) / β2 (446) / ε (447) with numerical-stability floors per `pearl_first_observation_bootstrap`. SP4 Mech 9 weight clamp + Pearl C engagement counter disabled (aux trunk is outside the SP4 8-group taxonomy — same convention as OFI embed Adam).
|
||||
|
||||
Mirrored launcher: `step_ofi_embed_adam` (single-flat-buffer Adam path). Adapted to 6 separate (params, grad, m, v) tensors with byte-offset-free layout.
|
||||
|
||||
### Choice: 6-buffer (vs flat-buffer)
|
||||
|
||||
Selected **6-buffer with 6 separate Adam launches** over flat-buffer with byte-offset views. Rationale:
|
||||
|
||||
1. C.2 already allocated 6 separate param tensors (`aux_trunk_{w1,b1,w2,b2,w3,b3}`) and 12 separate Adam m/v tensors — a flat-buffer Adam would have required restructuring C.2's allocations to be views into a flat tensor, which is a contract change C.5a explicitly avoids.
|
||||
2. The existing `dqn_adam_update_kernel` works on a single contiguous (params, grads, m, v) tuple per launch. 6 launches consuming 6 tuples is a smaller Δ than refactoring C.2 to a flat layout.
|
||||
3. Global L2-norm clipping is preserved: the kernel reads `*grad_norm_f32` on every launch and computes its own per-launch `clip_scale = (norm > clip) ? clip/norm : 1`. All 6 launches see the same global norm + same clip threshold → consistent global clipping.
|
||||
4. Storage cost is identical (same total f32 footprint).
|
||||
5. Per-launch overhead (≤6 launches × ~1µs ≈ ≤6µs/step) is negligible relative to the trainer step.
|
||||
|
||||
The 7th entry in `aux_trunk_grad_block_offsets` records the total block count, passed as `num_blocks` to the finalize kernel — capturing all 6 tensors' partial blocks in a single reduction.
|
||||
|
||||
### NO contract change
|
||||
|
||||
- `launch_aux_trunk_adam_update` is `#[allow(dead_code)]` because no production caller invokes it in C.5a. C.5b adds the call site atomically with the kernel rename + zero-fill revert.
|
||||
- The 6 grad buffers, 4 saved-fwd tiles, accumulator, partials, norm buf, and pinned LR/clip/t are all `#[allow(dead_code)]` for the same reason.
|
||||
- `aux_trunk_adam_step` initialises to 0 (counter is incremented by C.5b's call site).
|
||||
|
||||
### Compiler warnings accepted (no `#[allow(dead_code)]` workaround)
|
||||
|
||||
`cargo check -p ml --tests --all-targets` is clean. The new fields are `#[allow(dead_code)]`-tagged on the struct definitions because every struct field with no consumer otherwise warns; tagging them inline at the field is the canonical pattern this codebase uses for additive fields awaiting a wire-up commit (matches `aux_trunk_forward_ops`'s convention from C.3).
|
||||
|
||||
### Files modified (this commit)
|
||||
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (+11 fields on `GpuDqnTrainer`; +allocations after C.2's site; +3 pinned host pointers freed in `Drop`; +`launch_aux_trunk_adam_update` method)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` (+3 fields on `GpuExperienceCollector`; +allocations after `exp_aux_dir_acc_buf`'s site)
|
||||
- Modify: `docs/dqn-wire-up-audit.md` (this section)
|
||||
|
||||
### Verification
|
||||
|
||||
- `SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo check -p ml --tests --all-targets` — clean (only pre-existing warnings).
|
||||
- `cargo test -p ml --test aux_trunk_oracle_tests --test sp14_oracle_tests --release -- --ignored --nocapture` — 12/12 pass (8 aux_trunk + 4 sp14, no regression).
|
||||
- `grep -rn "launch_aux_trunk_adam_update" crates/ml/src/` — only the definition appears, no callers.
|
||||
|
||||
### Follow-up: C.5b
|
||||
|
||||
C.5b is the genuine atomic contract migration:
|
||||
1. Insert `aux_trunk_forward.launch(...)` post-`forward_online_f32` in collector + post-`encoder_forward_chain` in trainer.
|
||||
2. Redirect `aux_heads_fwd` input pointer from Q's `h_s2` to `h_s2_aux`.
|
||||
3. Revert zero-fills at `aux_next_bar_backward:599-638` + `aux_regime_backward:757-790` so the SAXPY-into-`dh_s2_aux_accum` path works.
|
||||
4. Param-rename `h_s2 → h_s2_aux` in the 4 aux_heads kernels.
|
||||
5. Insert `aux_trunk_backward.launch(...)` + `launch_aux_trunk_adam_update(...)` calls.
|
||||
6. Initialize `aux_trunk_lr_pinned` + `aux_trunk_grad_clip_pinned` from ISV per-step (host writes before each launch).
|
||||
7. Increment `aux_trunk_adam_step` per Adam call.
|
||||
|
||||
@@ -869,17 +869,113 @@ If C.10 L40S validation shows aux loss STILL flatlines at ln(2) with H=60:
|
||||
|
||||
---
|
||||
|
||||
## Phase C.5: Atomic Forward+Backward Wire-Up
|
||||
## Phase C.5: Wire-Up — Split into C.5a (additive infra) + C.5b (atomic contract migration)
|
||||
|
||||
> **CRITICAL: this is ONE atomic commit per `feedback_no_partial_refactor`.** All consumers of aux head's input migrate from `h_s2` to `h_s2_aux` simultaneously with the new kernel launches in collector. If the build is broken between forward-only and backward-wired, the commit is incomplete.
|
||||
> Split authorized 2026-05-08 after C.5 implementer reported ~600-800 LOC scope across 4 files with subtle gradient-correctness windows. Split is architecturally legitimate (mirrors SP14 Layer A→B pattern):
|
||||
> - **C.5a** is *additive* — allocates buffers + adds Adam launcher as dead code. No contract change; nothing in production calls the new code yet. Cannot break the build's behavior.
|
||||
> - **C.5b** is *atomic* — kernel rename + revert zero-fills + wire-up redirects in ONE commit. ~300 LOC. This is the genuine contract migration.
|
||||
|
||||
### Task C.5.1: Wire aux trunk forward + backward into collector chain
|
||||
## Phase C.5a: Additive Infrastructure — Buffers + Adam Launcher (DEAD CODE)
|
||||
|
||||
> No contract change. Build clean before, build clean after. Not yet wired.
|
||||
|
||||
### Task C.5a.1: Allocate fwd/bwd buffers + Adam launcher
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (if backward chain lives here)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/aux_heads_kernel.cu` (input parameter rename `h_s2` → `h_s2_aux`)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_aux_heads.rs` (Rust wrapper input plumbing)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (allocate `h_s2_aux`, `h_aux1`, `h_aux2` saved-fwd; allocate aux trunk grad buffers if separate from params; add `launch_aux_trunk_adam_update` fn — dead, no caller yet)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` (allocate rollout-sized `h_s2_aux`, `h_aux1`, `h_aux2`)
|
||||
- Modify: `docs/dqn-wire-up-audit.md` (audit entry — additive infra, no behavior change)
|
||||
|
||||
- [ ] **Step 1: Allocate trainer-side saved-fwd buffers**
|
||||
|
||||
In `gpu_dqn_trainer.rs`, alongside existing aux buffers (find `aux_next_bar_logits` or similar), allocate:
|
||||
- `h_s2_aux: CudaSlice<f32>` — `[B_max, AUX_HIDDEN_DIM]`
|
||||
- `h_aux1: CudaSlice<f32>` — `[B_max, H1=256]` saved-for-backward intermediate
|
||||
- `h_aux2: CudaSlice<f32>` — `[B_max, H2=128]` saved-for-backward intermediate
|
||||
- `dh_s2_aux_accum: CudaSlice<f32>` — `[B_max, AUX_HIDDEN_DIM]` accumulator for upstream gradient (where aux_next_bar_backward + aux_regime_backward will write SAXPY-style after C.5b's revert)
|
||||
|
||||
These are dead until C.5b wires them up.
|
||||
|
||||
- [ ] **Step 2: Allocate collector-side rollout-sized buffers**
|
||||
|
||||
In `gpu_experience_collector.rs`, mirror trainer buffers but rollout-sized (`B * T_rollout` * dim).
|
||||
|
||||
- [ ] **Step 3: Allocate aux trunk gradient buffers**
|
||||
|
||||
C.2 allocated 6 separate weight/bias param tensors + 12 separate Adam m/v tensors. C.5a needs corresponding gradient buffers:
|
||||
- `aux_trunk_w1_grad: CudaSlice<f32>` matching `aux_trunk_w1` shape
|
||||
- ... same for b1, w2, b2, w3, b3 (6 grad tensors total)
|
||||
|
||||
OR, if you prefer the flat-buffer Adam pattern that OFI embed uses: introduce flat `aux_trunk_params_flat`, `aux_trunk_grad_flat`, `aux_trunk_m_flat`, `aux_trunk_v_flat` (each ~132K f32) with byte-offset views into w1/b1/w2/b2/w3/b3 sub-regions. This requires re-declaring C.2's separate buffers as views into the flat buffer.
|
||||
|
||||
**Implementer choice:** whichever fits cleanly into existing patterns. If 6-buffer approach is straightforward, prefer it (smaller refactor of C.2's work). If existing Adam launcher pattern strictly requires flat buffer, restructure to flat. Document the choice in audit entry.
|
||||
|
||||
- [ ] **Step 4: Write `launch_aux_trunk_adam_update`** (dead, no caller yet)
|
||||
|
||||
Mirror existing Adam launcher pattern. Search for `OFI` embed Adam path or similar single-param-group launcher (e.g., `launch_ofi_embed_adam`). The aux trunk lives outside the SP4-canonical 4-way Adam taxonomy (Q/Target/CQL/Aux per-branch), so a dedicated launcher is appropriate.
|
||||
|
||||
Reads from ISV:
|
||||
- `lr` from `ISV[AUX_TRUNK_LR_INDEX]` (444)
|
||||
- `β1` from `ISV[AUX_TRUNK_BETA1_INDEX]` (445)
|
||||
- `β2` from `ISV[AUX_TRUNK_BETA2_INDEX]` (446)
|
||||
- `ε` from `ISV[AUX_TRUNK_EPS_INDEX]` (447)
|
||||
- `grad_clip_norm` from `ISV[AUX_TRUNK_GRAD_CLIP_INDEX]` (448)
|
||||
|
||||
Pre-step: global L2 norm of all 6 grad tensors → if exceeds clip threshold, scale all by `clip_norm / norm`. Then Adam step.
|
||||
|
||||
- [ ] **Step 5: Verify**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo check -p ml --tests --all-targets 2>&1 | tail -10
|
||||
```
|
||||
|
||||
Expected: clean build. Dead code may warn about unused fields/functions; that's OK because C.5b consumes them. Add `#[allow(dead_code)]` only if clippy promotes the warning to an error.
|
||||
|
||||
Run all oracle tests to confirm nothing regressed:
|
||||
```bash
|
||||
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 -25
|
||||
```
|
||||
|
||||
Expected: 12 tests pass (no regression — the new buffers + launcher are inert).
|
||||
|
||||
- [ ] **Step 6: Audit doc + commit**
|
||||
|
||||
```bash
|
||||
git add -A
|
||||
git commit -m "$(cat <<'EOF'
|
||||
feat(sp14-c.5a): allocate aux trunk fwd/bwd buffers + Adam launcher (dead code)
|
||||
|
||||
Phase C.5a — additive infrastructure for the aux trunk wire-up.
|
||||
Allocates saved-fwd buffers (h_s2_aux, h_aux1, h_aux2),
|
||||
gradient buffers (6× aux_trunk_*_grad), accumulator
|
||||
(dh_s2_aux_accum), and dedicated Adam launcher
|
||||
(launch_aux_trunk_adam_update) reading β1/β2/ε/LR/grad-clip from
|
||||
ISV[444..449).
|
||||
|
||||
No contract change. No call sites for the new launcher yet.
|
||||
C.5b atomically wires these in.
|
||||
|
||||
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.
|
||||
|
||||
### Task C.5b.1: Atomic kernel rename + revert zero-fills + wire-up
|
||||
|
||||
> **CRITICAL: this is ONE atomic commit per `feedback_no_partial_refactor`.** All consumers of aux head's input migrate from `h_s2` to `h_s2_aux` simultaneously with the new kernel launches in collector. C.5a already allocated the buffers + Adam launcher (dead). C.5b makes them live atomically.
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` (insert aux_trunk_forward post-graph; redirect aux_heads_fwd input ptr to h_s2_aux)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (insert aux_trunk_forward in trainer fwd; insert aux_trunk_backward + Adam launch in bwd; redirect aux_heads_fwd/bwd input/grad ptrs)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/aux_heads_kernel.cu` (param rename `h_s2` → `h_s2_aux` in 4 kernels; **REVERT** zero-fills in `aux_next_bar_backward:599-638` and `aux_regime_backward:757-790` from commits `872bd7392` + `411a30473` — restore proper SAXPY-back-to-input-grad)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_aux_heads.rs` (Rust wrapper input/grad plumbing)
|
||||
- Modify: `docs/dqn-wire-up-audit.md` (atomic migration entry; note C.0 stop-grad audit comments are now stale and removed)
|
||||
|
||||
- [ ] **Step 1: Locate aux head's current forward call site in collector**
|
||||
|
||||
|
||||
Reference in New Issue
Block a user