feat(sp22): H6 Phase 3 α — adaptive W via dW backward + Adam (B9 Steps 8+11)
Step 8 — c51_aux_dw_kernel (new):
- Per-action block tree-reduce: grid=(4,1,1), block=(256,1,1). One block
per W index a, tree-reduces dW[a] across batch via warp shuffle + shmem.
Zero atomicAdd per pearl_no_atomicadd.
- Per-sample contributions:
a == a_d: dW[a] += inv_batch × isw × (SP_b/dz) × state_121
a == a*: dW[a] += inv_batch × isw × (-γ(1-done)) × (SP_b/dz) × next_state_121
c51_loss_kernel forward — new scratch outputs:
- aux_target_a_dir_buf[B] (i32): saves best_next_a for d==0 after Step c
sampling.
- aux_proj_logdiff_dir_buf[B] (f32): saves SP_b = Σ_n p_target_n ×
(current_lp[upper_n] - current_lp[lower_n]) after Step d's projection
via re-derivation of lower_n/upper_n (matching Huber compression +
clamp arithmetic of block_bellman_project_f).
Step 11 — adam_w_aux_kernel (new):
- Standard Adam with bias correction, grid=(1,1,1), block=(4,1,1).
- Graph-capture-safe: lr via self.lr_dev_ptr pointer arg; step via
self.ptrs.t_buf pointer arg (matches main Adam pattern). beta/eps
as value args from sp5_isv_slots constants.
- Bias-correction denominator floored at 1e-30 to avoid /0.
Trainer wiring (submit_adam_ops):
- launch_c51_aux_dw + launch_adam_w_aux added right after
launch_adam_update. Both inside the captured adam_child graph.
- New trainer fields: aux_target_a_dir_buf, aux_proj_logdiff_dir_buf,
c51_aux_dw_kernel, adam_w_aux_kernel. Cubin statics SP22_C51_AUX_DW_CUBIN
+ SP22_ADAM_W_AUX_CUBIN added.
NULL-safety:
- aux_shift_active=false in c51_loss_kernel forward → both scratch
buffers stay at alloc_zeros 0 → dW reads 0 → Adam W is a no-op.
- aux_target_a_dir_out / aux_proj_logdiff_dir_out are NULL-tolerant.
Deferred (deliberate scope):
- dL/dstate_121 backward (c51 → aux head): refinement, not correctness;
aux head trains via own supervised CE loss.
- Phase C1 collector W ptr setter.
- Phase D (eval-side aux infrastructure).
Verification:
- cargo build -p ml --lib: 0 errors, 21 pre-existing warnings.
- nvcc full recompile clean (1m05s for sm_89 target).
- All forward atom-shift consumers + adaptive W backward + Adam now wired.
End-state: adaptive W trains from structural prior [-0.5, 0, +0.5, 0]
via projection log-diff gradient. Aux head trains independently via
supervised CE. Together they form learned cross-coupling from aux
direction predictions to dir-branch Q distribution shifts. Smoke can
now measure adaptive W's effect on WR.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -693,6 +693,21 @@ pub(crate) static SP22_AUX_SOFTMAX_TO_PER_ENV_CUBIN: &[u8] =
|
||||
pub(crate) static SP22_AUX_W_PRIOR_INIT_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/aux_w_prior_init_kernel.cubin"));
|
||||
|
||||
/// SP22 H6 Phase 3 α Step 8 (2026-05-13): backward dW kernel for
|
||||
/// `w_aux_to_q_dir [4]`. Reads scratch buffers from c51_loss_kernel
|
||||
/// forward (`aux_target_a_dir` + `aux_proj_logdiff_dir`) and computes
|
||||
/// per-action gradient via block tree-reduce (grid=(4,1,1),
|
||||
/// block=(256,1,1)). Launched OUTSIDE captured graphs.
|
||||
pub(crate) static SP22_C51_AUX_DW_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/c51_aux_dw_kernel.cubin"));
|
||||
|
||||
/// SP22 H6 Phase 3 α Step 11 (2026-05-13): Adam update kernel for
|
||||
/// `w_aux_to_q_dir [4]`. Standard Adam with bias correction. 4 threads
|
||||
/// single block. Launched once per training step OUTSIDE captured
|
||||
/// graphs.
|
||||
pub(crate) static SP22_ADAM_W_AUX_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/adam_w_aux_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-
|
||||
@@ -6921,6 +6936,26 @@ pub struct GpuDqnTrainer {
|
||||
/// 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 α Step 8 — per-sample sampled target action for d == 0.
|
||||
/// Written by c51_loss_kernel forward at the Expected SARSA sampling
|
||||
/// site (`best_next_a` for dir branch). Read by `c51_aux_dw_kernel`
|
||||
/// to look up W[a*] for the target-side dW contribution.
|
||||
/// Shape `[batch_size]` i32. alloc_zeros default = 0 (Short action;
|
||||
/// harmless when aux_shift inactive since SP_b stays at 0 too).
|
||||
aux_target_a_dir_buf: cudarc::driver::CudaSlice<i32>,
|
||||
/// SP22 H6 Phase 3 α Step 8 — per-sample projection log-diff sum
|
||||
/// for d == 0. Written by c51_loss_kernel forward right after the
|
||||
/// Bellman projection completes:
|
||||
/// `SP_b = Σ_n p_target_n × (current_lp[upper_n] - current_lp[lower_n])`
|
||||
/// Read by `c51_aux_dw_kernel` to compute
|
||||
/// `dL/dΔ_online = SP_b / Δz`, `dL/dΔ_target = -γ*(1-done)*dL/dΔ_online`
|
||||
/// Shape `[batch_size]` f32. alloc_zeros default = 0 → no gradient
|
||||
/// when aux_shift inactive.
|
||||
aux_proj_logdiff_dir_buf: cudarc::driver::CudaSlice<f32>,
|
||||
/// SP22 H6 Phase 3 α Step 8 — kernel handle for `c51_aux_dw_kernel`.
|
||||
c51_aux_dw_kernel: cudarc::driver::CudaFunction,
|
||||
/// SP22 H6 Phase 3 α Step 11 — kernel handle for `adam_w_aux_update`.
|
||||
adam_w_aux_kernel: cudarc::driver::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
|
||||
@@ -20874,6 +20909,43 @@ impl GpuDqnTrainer {
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"sp22-h6-phase3 α: alloc dw_aux_buf: {e}"
|
||||
)))?;
|
||||
// SP22 H6 Phase 3 α Step 8 (2026-05-13): per-sample scratch buffers
|
||||
// for the W gradient backward. aux_target_a_dir saves best_next_a
|
||||
// for d==0 from the forward; aux_proj_logdiff_dir saves the
|
||||
// per-sample SP scalar.
|
||||
let aux_target_a_dir_buf = stream
|
||||
.alloc_zeros::<i32>(config.batch_size)
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"sp22-h6-phase3 α Step 8: alloc aux_target_a_dir_buf: {e}"
|
||||
)))?;
|
||||
let aux_proj_logdiff_dir_buf = stream
|
||||
.alloc_zeros::<f32>(config.batch_size)
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"sp22-h6-phase3 α Step 8: alloc aux_proj_logdiff_dir_buf: {e}"
|
||||
)))?;
|
||||
// Load Step 8 + Step 11 kernels.
|
||||
let c51_aux_dw_kernel = {
|
||||
let module = stream.context()
|
||||
.load_cubin(SP22_C51_AUX_DW_CUBIN.to_vec())
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"sp22-h6-phase3 α Step 8: c51_aux_dw cubin: {e}"
|
||||
)))?;
|
||||
module.load_function("c51_aux_dw_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"sp22-h6-phase3 α Step 8: c51_aux_dw_kernel load: {e}"
|
||||
)))?
|
||||
};
|
||||
let adam_w_aux_kernel = {
|
||||
let module = stream.context()
|
||||
.load_cubin(SP22_ADAM_W_AUX_CUBIN.to_vec())
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"sp22-h6-phase3 α Step 11: adam_w_aux cubin: {e}"
|
||||
)))?;
|
||||
module.load_function("adam_w_aux_update")
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"sp22-h6-phase3 α Step 11: adam_w_aux_update load: {e}"
|
||||
)))?
|
||||
};
|
||||
|
||||
// SP22 H6 Phase 3 α (atom-shift design 2026-05-13): structural-prior
|
||||
// init for w_aux_to_q_dir. Writes per-action prior
|
||||
@@ -25288,6 +25360,10 @@ impl GpuDqnTrainer {
|
||||
adam_m_w_aux,
|
||||
adam_v_w_aux,
|
||||
dw_aux_buf,
|
||||
aux_target_a_dir_buf,
|
||||
aux_proj_logdiff_dir_buf,
|
||||
c51_aux_dw_kernel,
|
||||
adam_w_aux_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,
|
||||
@@ -30124,6 +30200,114 @@ impl GpuDqnTrainer {
|
||||
}
|
||||
|
||||
|
||||
/// SP22 H6 Phase 3 α Step 8 (2026-05-13): dW kernel launcher.
|
||||
/// Reads scratch buffers populated by c51_loss_kernel forward
|
||||
/// (aux_target_a_dir + aux_proj_logdiff_dir for d == 0 only) and
|
||||
/// writes dw_aux_buf[4] via per-action block tree-reduce. No
|
||||
/// atomicAdd per pearl_no_atomicadd.
|
||||
///
|
||||
/// Graph-capture-safe: all kernel args are device pointers (the
|
||||
/// scratch buffers / params buffer / state buffers are stable across
|
||||
/// captured replays; their CONTENTS update each step via the captured
|
||||
/// forward and the externally-written batch upload).
|
||||
fn launch_c51_aux_dw(&self) -> Result<(), MLError> {
|
||||
let b = self.config.batch_size as i32;
|
||||
let b0 = self.config.branch_0_size as i32;
|
||||
let b1 = self.config.branch_1_size as i32;
|
||||
let b2 = self.config.branch_2_size as i32;
|
||||
let b3 = self.config.branch_3_size as i32;
|
||||
let aux_dir_prob_index = ml_core::state_layout::AUX_DIR_PROB_INDEX as i32;
|
||||
let state_dim_i32 = ml_core::state_layout::STATE_DIM as i32;
|
||||
|
||||
let aux_target_a_ptr = self.aux_target_a_dir_buf.raw_ptr();
|
||||
let aux_proj_logdiff_ptr = self.aux_proj_logdiff_dir_buf.raw_ptr();
|
||||
let actions_ptr = self.ptrs.actions_buf;
|
||||
let is_weights_ptr = self.ptrs.is_weights_buf;
|
||||
let dones_ptr = self.ptrs.dones_buf;
|
||||
let gamma_buf_ptr = self.gamma_buf.raw_ptr();
|
||||
let per_sample_support_ptr = self.per_sample_support_ptr;
|
||||
let states_ptr = self.ptrs.states_buf;
|
||||
let next_states_ptr = self.ptrs.next_states_buf;
|
||||
let dw_aux_ptr = self.dw_aux_buf.raw_ptr();
|
||||
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.c51_aux_dw_kernel)
|
||||
.arg(&aux_target_a_ptr)
|
||||
.arg(&aux_proj_logdiff_ptr)
|
||||
.arg(&actions_ptr)
|
||||
.arg(&is_weights_ptr)
|
||||
.arg(&dones_ptr)
|
||||
.arg(&gamma_buf_ptr)
|
||||
.arg(&per_sample_support_ptr)
|
||||
.arg(&states_ptr)
|
||||
.arg(&next_states_ptr)
|
||||
.arg(&dw_aux_ptr)
|
||||
.arg(&b)
|
||||
.arg(&b0)
|
||||
.arg(&b1)
|
||||
.arg(&b2)
|
||||
.arg(&b3)
|
||||
.arg(&aux_dir_prob_index)
|
||||
.arg(&state_dim_i32)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (4, 1, 1), // one block per W index (b0_size=4)
|
||||
block_dim: (256, 1, 1), // tree-reduce across batch
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"sp22-h6-phase3 α Step 8: c51_aux_dw_kernel launch: {e}"
|
||||
)))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// SP22 H6 Phase 3 α Step 11 (2026-05-13): Adam-update launcher for
|
||||
/// `w_aux_to_q_dir [4]`. Reads `dw_aux_buf` (written by
|
||||
/// `launch_c51_aux_dw`) and updates W in-place using Adam moments.
|
||||
///
|
||||
/// Graph-capture-safe: lr and step counter are read via
|
||||
/// device-mapped pinned pointers (matching the main Adam pattern).
|
||||
/// beta1/beta2/eps are constants from sp5_isv_slots — host-stable.
|
||||
fn launch_adam_w_aux(&self) -> Result<(), MLError> {
|
||||
use crate::cuda_pipeline::sp5_isv_slots::{
|
||||
ADAM_BETA1_BASE, ADAM_BETA2_BASE, ADAM_EPS_BASE,
|
||||
};
|
||||
|
||||
let w_ptr = self.w_aux_to_q_dir.raw_ptr();
|
||||
let dw_ptr = self.dw_aux_buf.raw_ptr();
|
||||
let m_ptr = self.adam_m_w_aux.raw_ptr();
|
||||
let v_ptr = self.adam_v_w_aux.raw_ptr();
|
||||
let lr_ptr = self.lr_dev_ptr;
|
||||
let t_ptr = self.ptrs.t_buf;
|
||||
let beta1 = ADAM_BETA1_BASE as f32;
|
||||
let beta2 = ADAM_BETA2_BASE as f32;
|
||||
let eps = ADAM_EPS_BASE as f32;
|
||||
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.adam_w_aux_kernel)
|
||||
.arg(&w_ptr)
|
||||
.arg(&dw_ptr)
|
||||
.arg(&m_ptr)
|
||||
.arg(&v_ptr)
|
||||
.arg(&lr_ptr)
|
||||
.arg(&beta1)
|
||||
.arg(&beta2)
|
||||
.arg(&eps)
|
||||
.arg(&t_ptr)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (4, 1, 1), // one thread per W slot
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"sp22-h6-phase3 α Step 11: adam_w_aux_kernel launch: {e}"
|
||||
)))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Submit the optimizer phase ops to the stream (used by adam_child capture).
|
||||
///
|
||||
/// Steps: Adam → unflatten.
|
||||
@@ -30138,6 +30322,15 @@ impl GpuDqnTrainer {
|
||||
// ── 6. Adam update (f32 master weights) ─────────────────
|
||||
self.launch_adam_update()?;
|
||||
|
||||
// ── 6a. SP22 H6 Phase 3 α Step 8 + Step 11 (2026-05-13) ──
|
||||
// Compute dW and Adam-update w_aux_to_q_dir [4] alongside the
|
||||
// main Adam step. Both kernels are graph-capture-safe (all
|
||||
// varying inputs via device-mapped pointers; constants via
|
||||
// value args). The dW kernel reads scratch buffers populated
|
||||
// by the captured c51_loss_kernel forward earlier in this step.
|
||||
self.launch_c51_aux_dw()?;
|
||||
self.launch_adam_w_aux()?;
|
||||
|
||||
// ── 6.5. Snapshot grad_buf → prev_grad_buf for next step's vaccine comparison ──
|
||||
// Graph-safe: submit_adam_ops is captured in adam_update child graph.
|
||||
self.graph_safe_copy_f32(
|
||||
@@ -30955,6 +31148,14 @@ impl GpuDqnTrainer {
|
||||
.arg(&self.ptrs.next_states_buf)
|
||||
.arg(&(ml_core::state_layout::AUX_DIR_PROB_INDEX as i32))
|
||||
.arg(&(ml_core::state_layout::STATE_DIM as i32))
|
||||
// ── SP22 H6 Phase 3 α Step 8 backward scratch (2026-05-13) ──
|
||||
// Per-sample scratch outputs for the W gradient kernel:
|
||||
// aux_target_a_dir[B] — best_next_a for d==0 per sample
|
||||
// aux_proj_logdiff_dir[B] — projection log-diff sum per sample
|
||||
// Both populated by c51_loss_kernel forward d==0 path;
|
||||
// consumed by c51_aux_dw_kernel post-loss.
|
||||
.arg(&self.aux_target_a_dir_buf.raw_ptr())
|
||||
.arg(&self.aux_proj_logdiff_dir_buf.raw_ptr())
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (b as u32, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
|
||||
Reference in New Issue
Block a user