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

refactor(sp22): H6 Phase 3 α cleanup + runbook revision for atom-shift design

Browse Source 
Same-session cleanup of Phase 3b scalar-bias residue (commit cb80b74ce's
additions that became architecturally redundant under the 2026-05-13
atom-shift revision). The scalar-bias α design is mathematically
ineffective in C51 distributional Q-learning (softmax-shift-invariance
→ W gradient = 0 → never trains). The revised design threads
`aux_atom_shift[b, a] = W[a] * state_121[b]` through compute_expected_q
+ c51_loss_kernel + c51_grad_kernel + mag_concat_qdir; dW + dstate
gradients integrate into c51_grad_kernel's projection backward.

Files
─────
- crates/ml/build.rs:
    Removed `aux_to_q_dir_bias_kernel.cu` + `aux_to_q_dir_bias_backward_kernel.cu`
    from kernels_with_common. Replaced with comment block documenting
    C51 softmax-invariance reason + spec/audit pointers.

- crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:
    Removed `pub(crate) static SP22_AUX_TO_Q_DIR_BIAS_CUBIN` and
    `SP22_AUX_TO_Q_DIR_BIAS_BWD_CUBIN` static declarations.
    Removed 3 struct fields (aux_to_q_dir_bias_kernel,
    aux_to_q_dir_bias_backward_dw_kernel,
    aux_to_q_dir_bias_backward_dstate_kernel) and their new()
    loading blocks + struct construction entries.
    KEPT: w_aux_to_q_dir + adam_m_w_aux + adam_v_w_aux + dw_aux_buf
    (still needed for atom-shift Adam-trained W).
    Updated W doc-comment to describe atom-shift threading (was
    scalar-bias) and structural-prior initialization `[-0.5, 0.0,
    +0.5, 0.0]`.

- crates/ml/src/cuda_pipeline/aux_to_q_dir_bias_kernel.cu +
  crates/ml/src/cuda_pipeline/aux_to_q_dir_bias_backward_kernel.cu:
    Source files stay on disk as committed dead code (commit
    464bc5f7a history preserved). nvcc no longer compiles them
    (build.rs unregistered). No `include_bytes!` references remain.

- docs/plans/2026-05-13-sp22-h6-phase3-alpha-beta-runbook.md:
    Tasks A3/A4/A5 marked SUPERSEDED; A5 retains cleanup instructions.
    Task B9 Steps 4-10 rewritten as Steps 4 (cleanup), 5 (W structural
    prior init), 6 (compute_expected_q atom-shift threading),
    7 (c51_loss_kernel), 8 (c51_grad_kernel backward dW + dstate),
    9 (mag_concat_qdir), 10 (quantile_q_select/iqn_dual_head
    investigation), 11 (Adam wireup), 12 (verify + capture).
    Task C1 rewritten: collector passes 4 more args through
    existing compute_expected_q launcher (NO new kernel).

- docs/dqn-wire-up-audit.md:
    Cleanup-commit entry documenting the runbook revision and the
    code-cleanup actions.

Verification
────────────
- cargo check -p ml --features cuda: 0 errors, 21 pre-existing
  warnings (Phase 3b baseline parity).
- Build script no longer compiles the deleted-registration kernels.

Phase 3-final scope (~40-60 hr engineering — unchanged)
───────────────────────────────────────────────────────
Implementation per the revised runbook: atom-shift threading
through 4-5 kernels (compute_expected_q, c51_loss_kernel,
c51_grad_kernel, mag_concat_qdir, possibly quantile_q_select/
iqn_dual_head) + Adam wireup + collector launcher arg passing +
A2 eval-side + SP11 controller extension + HEALTH_DIAG telemetry
+ verification gates + atomic Phase F commit + smoke + verdict.

Refs
────
- docs/plans/2026-05-12-sp22-h6-phase3-alpha-beta.md (spec α section
  revised 2026-05-13 — commit 648078ce2)
- docs/plans/2026-05-13-sp22-h6-phase3-alpha-beta-runbook.md (runbook
  α tasks revised in this commit)
- 464bc5f7a (Phase A — α kernels added; now committed dead code)
- cb80b74ce (Phase 3b — α struct fields + cubin statics added; cleaned
  up in this commit)
- pearl_no_partial_refactor (atom-shift threading is the new atomic
  contract for direction-branch atom positions)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-13 01:24:15 +02:00
parent 648078ce20
commit 08fd5803c4
4 changed files with 168 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -684,29 +684,23 @@ pub(crate) static SP14_Q_DISAGREEMENT_CUBIN: &[u8] =
pub(crate) static SP22_AUX_SOFTMAX_TO_PER_ENV_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/aux_softmax_to_per_env_kernel.cubin"));
/// SP22 H6 Phase 3 α forward (2026-05-13): post-encoder Q_dir bias.
/// Forward kernel adds `W_aux_to_Q_dir[a] * state_121[b]` to Q_dir[b, a]
/// for each (b, a). Reads state_121 from the encoder INPUT
/// (`batch_states[b * state_dim + AUX_DIR_PROB_INDEX]`), bypassing the
/// cold encoder weight for slot 121. Loaded from
/// `aux_to_q_dir_bias_kernel.cubin`. Consumed by both the trainer's
/// captured forward graph AND the collector's rollout-time captured
/// forward graph (shared W via setter).
pub(crate) static SP22_AUX_TO_Q_DIR_BIAS_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/aux_to_q_dir_bias_kernel.cubin"));
/// SP22 H6 Phase 3 α backward (2026-05-13): gradients for the α forward.
/// Contains TWO kernels:
/// - `aux_to_q_dir_bias_backward_dw`: per-action block tree-reduce
/// computing `dW[a] = Σ_b state_121[b] * dq_dir[b, a]` (no atomicAdd).
/// - `aux_to_q_dir_bias_backward_dstate`: per-sample sum accumulating
/// `dstate_121[b] += Σ_a W[a] * dq_dir[b, a]` into
/// `dbatch_states[b, AUX_DIR_PROB_INDEX]` (option-(i) gradient routing
/// per spec: both α and encoder paths get signal).
/// Loaded from `aux_to_q_dir_bias_backward_kernel.cubin`. Consumed only
/// by the trainer's captured backward graph (eval doesn't update params).
pub(crate) static SP22_AUX_TO_Q_DIR_BIAS_BWD_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/aux_to_q_dir_bias_backward_kernel.cubin"));
// SP22 H6 Phase 3 α SCALAR-BIAS DESIGN — DELETED 2026-05-13.
// The scalar-bias `Q_dir[b, a] += W[a] * state_121[b]` approach is
// mathematically ineffective in C51 distributional Q-learning (softmax-
// shift-invariance → W gradient = 0 → never trains). Replaced by the
// per-(b, a) atom-position shift design that threads `aux_atom_shift[b, a]
// = W[a] * state_121[b]` through compute_expected_q + c51_loss_kernel +
// c51_grad_kernel + mag_concat_qdir + (investigate quantile_q_select /
// iqn_dual_head). The dW + dstate gradient computations integrate into
// c51_grad_kernel's projection backward — no separate α forward/backward
// kernel needed. See audit doc Phase 3c entry + spec doc α section
// for the revised design.
//
// The .cu files `aux_to_q_dir_bias_kernel.cu` and
// `aux_to_q_dir_bias_backward_kernel.cu` are removed from build.rs
// registration in the same commit. The Adam-trained W buffer
// (`w_aux_to_q_dir`) + Adam moments + dW accumulator remain — they
// still hold the per-action W vector used by the atom-shift threading.
/// SP17 Task PP.3 (2026-05-08): pre-centering V/A mean diagnostic.
/// Single-block tree-reduce kernel reading the existing `on_v_logits_buf
@@ -6892,15 +6886,20 @@ pub struct GpuDqnTrainer {
/// launch); no FoldReset registry entry. Bound is structural (no
/// runtime tanh). Consumed by `launch_aux_pred_to_isv_tanh`.
aux_pred_to_isv_tanh_kernel: CudaFunction,
// ── SP22 H6 Phase 3 α (2026-05-13): post-encoder Q_dir bias ──────────
// ── SP22 H6 Phase 3 α (2026-05-13): atom-position shift Adam-trained W ──
/// SP22 H6 Phase 3 α — learnable weight vector [b0_size=4] f32.
/// Adam-trained. Read by both the trainer (training-time forward +
/// backward) and the collector (rollout-time forward via shared
/// device pointer). Forward kernel adds `W[a] * state_121[b]` to
/// `Q_dir[b, a]`. Cold-start zero (Adam-driven growth tracks
/// aux→Q correlation). Per `pearl_first_observation_bootstrap`
/// the W ∈ R^4 starts at sentinel zero; Adam moves it from there
/// based on the gradient signal computed by the backward kernels.
/// Adam-trained. Read by both the trainer (training-time atom-shift
/// threading through compute_expected_q + c51_loss_kernel +
/// c51_grad_kernel + mag_concat_qdir) and the collector (rollout-time
/// atom-shift in compute_expected_q for action selection, via shared
/// device pointer). Atom-shift formula:
/// `aux_atom_shift[b, a] = W[a] * state_121[b]`
/// applied per atom of direction-branch action a as
/// `z_n_effective[b, a, n] = z_n_base + aux_atom_shift[b, a]`.
/// **Initialization: structural prior** `[-0.5, 0.0, +0.5, 0.0]` for
/// `[Short, Hold, Long, Flat]` — domain belief that aux conviction ×
/// position-sign biases toward aligned trades. Adam refines from
/// the prior via projection backward in c51_grad_kernel.
w_aux_to_q_dir: cudarc::driver::CudaSlice<f32>,
/// SP22 H6 Phase 3 α — Adam first moment (m) for `w_aux_to_q_dir`.
/// Same shape [b0_size=4] f32; zero-init per Adam convention.
@@ -6908,23 +6907,11 @@ pub struct GpuDqnTrainer {
/// SP22 H6 Phase 3 α — Adam second moment (v) for `w_aux_to_q_dir`.
adam_v_w_aux: cudarc::driver::CudaSlice<f32>,
/// SP22 H6 Phase 3 α — gradient accumulator for `w_aux_to_q_dir`.
/// Written by `aux_to_q_dir_bias_backward_dw_kernel`; read by the
/// Adam-step update. [b0_size=4] f32.
/// Written by `c51_grad_kernel`'s projection backward (dW[a] =
/// Σ_b state_121[b] * dz_shift[b, a] where dz_shift comes from
/// dloss/dz_n_effective); read by the Adam-step update.
/// [b0_size=4] f32.
dw_aux_buf: cudarc::driver::CudaSlice<f32>,
/// SP22 H6 Phase 3 α — forward kernel handle. Loaded from
/// `SP22_AUX_TO_Q_DIR_BIAS_CUBIN`. Launched inside the captured
/// training forward graph right after the Q-head produces Q_dir,
/// before Q_dir is consumed by downstream branches.
aux_to_q_dir_bias_kernel: CudaFunction,
/// SP22 H6 Phase 3 α — backward dW kernel handle. Per-action block
/// tree-reduce computing `dW[a] = Σ_b state_121[b] * dq_dir[b, a]`.
/// No atomicAdd per `feedback_no_atomicadd`.
aux_to_q_dir_bias_backward_dw_kernel: CudaFunction,
/// SP22 H6 Phase 3 α — backward dstate kernel handle. Per-sample
/// sum accumulating `dstate_121[b] += Σ_a W[a] * dq_dir[b, a]`
/// into the encoder-input gradient slot 121. Option-(i) gradient
/// routing per spec: both α and encoder paths get signal.
aux_to_q_dir_bias_backward_dstate_kernel: CudaFunction,
// ── SP14 Earned Gradient Flow kernels ─────────────────────────────────
/// SP14 B.3 (2026-05-05): per-step Q-head ↔ aux argmax disagreement EMA
/// producer. Reads online Q logits + aux softmax outputs; computes the
@@ -20848,34 +20835,24 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!(
"sp22-h6-phase3 α: alloc dw_aux_buf: {e}"
)))?;
let aux_to_q_dir_bias_kernel = {
let module = stream.context()
.load_cubin(SP22_AUX_TO_Q_DIR_BIAS_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!(
"sp22-h6-phase3 α: aux_to_q_dir_bias cubin: {e}"
)))?;
module.load_function("aux_to_q_dir_bias")
.map_err(|e| MLError::ModelError(format!(
"sp22-h6-phase3 α: aux_to_q_dir_bias load: {e}"
)))?
};
let (aux_to_q_dir_bias_backward_dw_kernel,
aux_to_q_dir_bias_backward_dstate_kernel) = {
let module = stream.context()
.load_cubin(SP22_AUX_TO_Q_DIR_BIAS_BWD_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!(
"sp22-h6-phase3 α: aux_to_q_dir_bias_bwd cubin: {e}"
)))?;
let dw = module.load_function("aux_to_q_dir_bias_backward_dw")
.map_err(|e| MLError::ModelError(format!(
"sp22-h6-phase3 α: aux_to_q_dir_bias_backward_dw load: {e}"
)))?;
let dstate = module.load_function("aux_to_q_dir_bias_backward_dstate")
.map_err(|e| MLError::ModelError(format!(
"sp22-h6-phase3 α: aux_to_q_dir_bias_backward_dstate load: {e}"
)))?;
(dw, dstate)
};
// SP22 H6 Phase 3 α scalar-bias kernel loads — DELETED 2026-05-13.
// Per the architectural revision (audit doc "Phase 3c — α
// architectural finding"), the scalar-bias α design is replaced
// by atom-position shift threading through compute_expected_q +
// c51_loss_kernel + c51_grad_kernel + mag_concat_qdir. No
// separate α forward / backward kernel is launched — the
// gradient flows via c51_grad_kernel's existing projection
// backward. The kernel handles (aux_to_q_dir_bias_kernel,
// aux_to_q_dir_bias_backward_dw_kernel,
// aux_to_q_dir_bias_backward_dstate_kernel) and the CUBIN
// statics (SP22_AUX_TO_Q_DIR_BIAS_CUBIN,
// SP22_AUX_TO_Q_DIR_BIAS_BWD_CUBIN) are deleted in this same
// commit. The W + Adam moments + dW buffer survive — they
// still hold the Adam-trained W vector used by atom-shift
// threading. Phase 3-final implementation writes the
// structural prior init `[-0.5, 0.0, +0.5, 0.0]` into
// `w_aux_to_q_dir` after `alloc_zeros` above (via a small
// fill-kernel launch or repurposed fill_f32 sequence).
// SP14 EGF kernel loads (B.3–B.6).
// Handles are held in struct fields but not launched until B.9/B.11;
@@ -25226,17 +25203,18 @@ impl GpuDqnTrainer {
// `gpu_experience_collector.rs` per-step action-select site.
apply_fixed_alpha_ema_kernel,
aux_pred_to_isv_tanh_kernel,
// SP22 H6 Phase 3 α (2026-05-13): post-encoder Q_dir bias.
// 4-element W + Adam moments + grad accumulator + fwd/bwd
// kernel handles. Cold-start zero; Adam moves W from there
// based on the backward-emitted dW from the loss path.
// SP22 H6 Phase 3 α (2026-05-13 revised — atom-position shift).
// 4-element Adam-trained W + Adam moments + grad accumulator.
// No separate forward/backward kernel handles — the
// scalar-bias design was replaced by atom-shift threading
// through compute_expected_q + c51_loss_kernel + c51_grad_kernel
// + mag_concat_qdir (Phase 3-final implementation). W is
// initialized to the structural prior `[-0.5, 0.0, +0.5, 0.0]`
// (Phase 3-final step in `new()`); Adam refines from there.
w_aux_to_q_dir,
adam_m_w_aux,
adam_v_w_aux,
dw_aux_buf,
aux_to_q_dir_bias_kernel,
aux_to_q_dir_bias_backward_dw_kernel,
aux_to_q_dir_bias_backward_dstate_kernel,
// SP14 q_disagreement diagnostic + Coupling A forward feature
// wire. Phase C.1 (2026-05-08) deleted the α-machinery struct
// entries (sp14_alpha_grad_compute_kernel,