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
@@ -827,6 +827,17 @@ fn main() {
|
||||
// projection backward (Phase 3-final B9 Step 8). Pure 4-thread
|
||||
// single-block; trivial cost.
|
||||
"aux_w_prior_init_kernel.cu",
|
||||
// SP22 H6 Phase 3 α Step 8 (2026-05-13): backward dW kernel for
|
||||
// w_aux_to_q_dir [4]. Reads scratch 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))
|
||||
// — no atomicAdd per pearl_no_atomicadd. Per-block writes dw_aux[a].
|
||||
"c51_aux_dw_kernel.cu",
|
||||
// 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
|
||||
// the captured forward/backward graphs.
|
||||
"adam_w_aux_kernel.cu",
|
||||
// SP14 Layer C Phase C.4b (2026-05-08): adaptive aux prediction
|
||||
// horizon producer. Single-thread kernel writing
|
||||
// `ISV[AUX_PRED_HORIZON_BARS_INDEX=450]` from
|
||||
|
||||
72
crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu
Normal file
72
crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu
Normal file
@@ -0,0 +1,72 @@
|
||||
// crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu
|
||||
//
|
||||
// SP22 H6 Phase 3 α Step 11 — Adam optimizer update for w_aux_to_q_dir [4].
|
||||
//
|
||||
// Standard Adam (Kingma & Ba 2015) with bias correction:
|
||||
// m_t = β1 × m_{t-1} + (1 - β1) × g_t
|
||||
// v_t = β2 × v_{t-1} + (1 - β2) × g_t²
|
||||
// m_hat = m_t / (1 - β1^t)
|
||||
// v_hat = v_t / (1 - β2^t)
|
||||
// W_t = W_{t-1} - lr × m_hat / (√v_hat + ε)
|
||||
//
|
||||
// Discipline
|
||||
// ──────────
|
||||
// - 4 threads, single block — trivial cost (<1µs). Launched once per
|
||||
// training step OUTSIDE the captured forward/backward graph (mirrors
|
||||
// the existing main-params Adam launch).
|
||||
// - β1, β2, ε, lr passed by value from host — host-stable across captured
|
||||
// launches per pearl_no_host_branches_in_captured_graph (this kernel is
|
||||
// not captured anyway).
|
||||
// - step is i32 (1-indexed) per the trainer's existing Adam step counter
|
||||
// semantics. `__powf` for the bias-correction denominator (sufficient
|
||||
// precision for the few-thousand-step horizon).
|
||||
// - Per pearl_first_observation_bootstrap.md: m and v default to 0 at
|
||||
// alloc_zeros init; the first step's bias-correction divides by
|
||||
// (1 - β^1) = (1 - β), not (1 - β^0) = 0 → no NaN.
|
||||
// - No HtoD inside kernel per feedback_no_htod_htoh_only_mapped_pinned.md.
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#define W_AUX_DIM 4
|
||||
|
||||
/* Graph-capture-safe signature: `step` and `lr` are POINTERS to
|
||||
* device-mapped pinned scratch the host updates each training step
|
||||
* (mirrors the main Adam launcher's `t_buf` + `lr_dev_ptr` pattern).
|
||||
* The captured launch records the pointers; the kernel reads the
|
||||
* current values at replay time. Beta/eps are constants — host-stable
|
||||
* across all replays per pearl_no_host_branches_in_captured_graph. */
|
||||
extern "C" __global__ void adam_w_aux_update(
|
||||
float* __restrict__ w, /* [W_AUX_DIM=4] in-place update */
|
||||
const float* __restrict__ dw, /* [W_AUX_DIM=4] gradient from c51_aux_dw_kernel */
|
||||
float* __restrict__ m, /* [W_AUX_DIM=4] first moment in-place */
|
||||
float* __restrict__ v, /* [W_AUX_DIM=4] second moment in-place */
|
||||
const float* __restrict__ lr_ptr, /* [1] device-mapped pinned f32 */
|
||||
float beta1,
|
||||
float beta2,
|
||||
float eps,
|
||||
const int* __restrict__ step_ptr /* [1] device-mapped pinned i32, 1-indexed */
|
||||
) {
|
||||
int a = threadIdx.x;
|
||||
if (a >= W_AUX_DIM) return;
|
||||
|
||||
float lr = lr_ptr[0];
|
||||
int step = step_ptr[0];
|
||||
|
||||
float g = dw[a];
|
||||
float m_new = beta1 * m[a] + (1.0f - beta1) * g;
|
||||
float v_new = beta2 * v[a] + (1.0f - beta2) * g * g;
|
||||
m[a] = m_new;
|
||||
v[a] = v_new;
|
||||
|
||||
/* Bias-corrected moments. step is at least 1 on the first invocation
|
||||
* (trainer increments before submit_adam_ops), so (1 - β^step) > 0. */
|
||||
float bc1 = 1.0f - __powf(beta1, (float)step);
|
||||
float bc2 = 1.0f - __powf(beta2, (float)step);
|
||||
/* Floors at fp32 epsilon to avoid /0 in degenerate first-step
|
||||
* configurations (β1 ≈ 0 or β2 ≈ 0). */
|
||||
bc1 = fmaxf(bc1, 1e-30f);
|
||||
bc2 = fmaxf(bc2, 1e-30f);
|
||||
float m_hat = m_new / bc1;
|
||||
float v_hat = v_new / bc2;
|
||||
w[a] = w[a] - lr * m_hat / (sqrtf(v_hat) + eps);
|
||||
}
|
||||
132
crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu
Normal file
132
crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu
Normal file
@@ -0,0 +1,132 @@
|
||||
// crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu
|
||||
//
|
||||
// SP22 H6 Phase 3 α Step 8 — backward dW kernel for w_aux_to_q_dir [4].
|
||||
//
|
||||
// Computes gradient of C51 loss w.r.t. the per-action atom-shift weights:
|
||||
// dW[a] = Σ_b indicator(a_d_b == a) × isw_b × inv_batch × (SP_b / Δz_b)
|
||||
// × state_121_b
|
||||
// + Σ_b indicator(a*_b == a) × isw_b × inv_batch × (-γ_eff_b × (1-done_b))
|
||||
// × (SP_b / Δz_b) × next_state_121_b
|
||||
//
|
||||
// where:
|
||||
// - SP_b = projection log-diff sum (saved by c51_loss_kernel forward):
|
||||
// SP_b = Σ_n p_target_n × (current_lp[upper_n] - current_lp[lower_n])
|
||||
// - a_d_b = taken direction action (= a0, computed from actions[b])
|
||||
// - a*_b = sampled target action (saved by c51_loss_kernel as
|
||||
// aux_target_a_dir[b])
|
||||
// - Δz_b = per_sample_support[b * 12 + 0 * 3 + 2] (dir branch dz)
|
||||
// - γ_eff_b = effective gamma for dir branch from gamma_buf[b] (γ^n_steps,
|
||||
// ISV gamma_dir read inside)
|
||||
// - state_121_b = batch_states[b * state_dim + AUX_DIR_PROB_INDEX]
|
||||
// - next_state_121_b = next_batch_states[b * state_dim + AUX_DIR_PROB_INDEX]
|
||||
//
|
||||
// Output: dw_aux_buf[4] f32 (zeroed before launch — overwrite semantics).
|
||||
//
|
||||
// Discipline
|
||||
// ──────────
|
||||
// - Per pearl_no_atomicadd: ONE BLOCK PER ACTION (grid=(4,1,1), block=(256,1,1));
|
||||
// each block tree-reduces across batch samples via shmem. No atomicAdd.
|
||||
// - Per feedback_no_htod_htoh_only_mapped_pinned.md: all reads from device /
|
||||
// mapped pinned buffers; no HtoD inside kernel.
|
||||
// - Per pearl_no_host_branches_in_captured_graph: launch happens OUTSIDE the
|
||||
// captured training graph (mirrors the existing post-c51_grad launches).
|
||||
// - NULL-safety: when forward didn't populate scratch (aux_shift_active=false),
|
||||
// both aux_target_a_dir and aux_proj_logdiff_dir stay at alloc_zeros
|
||||
// sentinels; gradient stays at 0 → Adam step on W is a no-op.
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#define BLOCK_THREADS 256
|
||||
#define NUM_WARPS (BLOCK_THREADS / 32)
|
||||
#define MAX_DIR_ACTIONS 4
|
||||
|
||||
extern "C" __global__ void c51_aux_dw_kernel(
|
||||
/* Inputs from c51_loss_kernel forward scratch */
|
||||
const int* __restrict__ aux_target_a_dir, /* [B] sampled a* per sample */
|
||||
const float* __restrict__ aux_proj_logdiff_dir,/* [B] SP_b per sample */
|
||||
|
||||
/* Inputs that the gradient formula needs */
|
||||
const int* __restrict__ actions, /* [B] factored action codes */
|
||||
const float* __restrict__ is_weights, /* [B] PER importance weights (clamped to 10) */
|
||||
const float* __restrict__ dones, /* [B] 0 or 1 */
|
||||
const float* __restrict__ gamma_buf, /* [B] effective γ^n_steps */
|
||||
const float* __restrict__ per_sample_support, /* [B, 4, 3] for Δz */
|
||||
const float* __restrict__ batch_states, /* [B, state_dim] for state_121 */
|
||||
const float* __restrict__ next_batch_states, /* [B, state_dim] for next_state_121 */
|
||||
|
||||
/* Output */
|
||||
float* __restrict__ dw_aux, /* [b0_size=4] gradient — written by tid==0 of each block */
|
||||
|
||||
/* Config */
|
||||
int batch_size,
|
||||
int b0_size,
|
||||
int b1_size,
|
||||
int b2_size,
|
||||
int b3_size,
|
||||
int aux_dir_prob_index, /* = 121 */
|
||||
int state_dim /* = 128 */
|
||||
) {
|
||||
int a = blockIdx.x; /* which W index this block accumulates */
|
||||
int tid = threadIdx.x;
|
||||
if (a >= b0_size) return;
|
||||
if (a >= MAX_DIR_ACTIONS) return;
|
||||
|
||||
float inv_batch = 1.0f / (float)batch_size;
|
||||
|
||||
/* Decode a_d (taken dir action = a0) inverse to c51_loss_kernel's decode:
|
||||
* a0 = factored / (b1 * b2 * b3)
|
||||
* Note: b0_size=4 in production but read from kernel arg for ABI safety. */
|
||||
int divisor_a0 = b1_size * b2_size * b3_size;
|
||||
if (divisor_a0 <= 0) divisor_a0 = 1;
|
||||
|
||||
/* Per-thread partial sum across samples assigned to this thread. */
|
||||
float local_sum = 0.0f;
|
||||
for (int b = tid; b < batch_size; b += BLOCK_THREADS) {
|
||||
int factored = actions[b];
|
||||
if (factored < 0) factored = 0;
|
||||
int a_d = factored / divisor_a0;
|
||||
if (a_d < 0 || a_d >= b0_size) a_d = 0;
|
||||
|
||||
int a_star = aux_target_a_dir[b];
|
||||
if (a_star < 0 || a_star >= b0_size) a_star = 0;
|
||||
|
||||
float sp = aux_proj_logdiff_dir[b];
|
||||
float isw = fminf(is_weights[b], 10.0f);
|
||||
float dz = per_sample_support[b * 12 + 0 * 3 + 2];
|
||||
if (dz < 1e-7f) continue; /* degenerate dir support: no gradient (matches forward skip) */
|
||||
|
||||
float dL_dDelta_online = sp / dz;
|
||||
float gamma_eff = gamma_buf[b];
|
||||
float done = dones[b];
|
||||
float dL_dDelta_target = -gamma_eff * (1.0f - done) * dL_dDelta_online;
|
||||
|
||||
if (a == a_d) {
|
||||
float s121 = batch_states[(long long)b * state_dim + aux_dir_prob_index];
|
||||
local_sum += inv_batch * isw * dL_dDelta_online * s121;
|
||||
}
|
||||
if (a == a_star) {
|
||||
float ns121 = next_batch_states[(long long)b * state_dim + aux_dir_prob_index];
|
||||
local_sum += inv_batch * isw * dL_dDelta_target * ns121;
|
||||
}
|
||||
}
|
||||
|
||||
/* Block tree-reduce (per pearl_no_atomicadd): warp shuffle + shmem. */
|
||||
__shared__ float shmem_warp[NUM_WARPS];
|
||||
unsigned int mask = 0xFFFFFFFF;
|
||||
int lane = tid & 31;
|
||||
int warp = tid >> 5;
|
||||
|
||||
/* Warp-level reduce. */
|
||||
for (int offset = 16; offset > 0; offset >>= 1) {
|
||||
local_sum += __shfl_xor_sync(mask, local_sum, offset);
|
||||
}
|
||||
if (lane == 0) shmem_warp[warp] = local_sum;
|
||||
__syncthreads();
|
||||
|
||||
/* Final reduce across warps (sequential in tid==0). */
|
||||
if (tid == 0) {
|
||||
float total = 0.0f;
|
||||
for (int w = 0; w < NUM_WARPS; w++) total += shmem_warp[w];
|
||||
dw_aux[a] = total;
|
||||
}
|
||||
}
|
||||
@@ -648,7 +648,28 @@ extern "C" __global__ void c51_loss_batched(
|
||||
const float* __restrict__ batch_states, /* [B, state_dim] f32, or NULL */
|
||||
const float* __restrict__ next_batch_states, /* [B, state_dim] f32, or NULL */
|
||||
int aux_dir_prob_index, /* = SL_PADDING_START = 121 */
|
||||
int state_dim /* = STATE_DIM = 128 */
|
||||
int state_dim, /* = STATE_DIM = 128 */
|
||||
|
||||
/* ── SP22 H6 Phase 3 α Step 8 backward scratch outputs ──────────────
|
||||
* Per-sample scratch produced by the FORWARD pass for the dir-branch
|
||||
* (d == 0) only; consumed by `c51_aux_dW_kernel` to compute W's
|
||||
* gradient without re-running the projection arithmetic.
|
||||
*
|
||||
* `aux_target_a_dir_out[b]` (i32, [B]): best_next_a sampled in Step c
|
||||
* for d == 0. The dW[a*] contribution accumulates per this index.
|
||||
*
|
||||
* `aux_proj_logdiff_dir_out[b]` (f32, [B]): SP_b = Σ_n p_target_n ×
|
||||
* (shmem_current_lp[upper_n] - shmem_current_lp[lower_n]) computed
|
||||
* during the d==0 Bellman projection. dW gradient = (SP_b / Δz) ×
|
||||
* state_121 for the online side, × (-γ*(1-done)*next_state_121) for
|
||||
* the target side.
|
||||
*
|
||||
* NULL-tolerant: when either is NULL OR aux_shift_active is false,
|
||||
* no write occurs. Both default to alloc_zeros at trainer init so
|
||||
* test scaffolds without Step 8 wiring see 0 → bit-identical pre-
|
||||
* Phase-3-α gradient flow downstream. */
|
||||
int* __restrict__ aux_target_a_dir_out, /* [B] i32, or NULL */
|
||||
float* __restrict__ aux_proj_logdiff_dir_out /* [B] f32, or NULL */
|
||||
) {
|
||||
extern __shared__ float shmem_f[];
|
||||
|
||||
@@ -1166,6 +1187,16 @@ extern "C" __global__ void c51_loss_batched(
|
||||
__syncthreads();
|
||||
int best_next_a = sampled_action;
|
||||
|
||||
/* SP22 H6 Phase 3 α Step 8 backward scratch save:
|
||||
* record sampled target action for d == 0 only. Consumed by
|
||||
* c51_aux_dW_kernel to look up W[best_next_a] for the target
|
||||
* side dW gradient. */
|
||||
if (d == 0 && tid == 0
|
||||
&& aux_shift_active
|
||||
&& aux_target_a_dir_out != NULL) {
|
||||
aux_target_a_dir_out[sample_id] = best_next_a;
|
||||
}
|
||||
|
||||
/* Target distribution for sampled action */
|
||||
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
|
||||
shmem_lp[j] = shmem_val[j] + shmem_adv[best_next_a * num_atoms + j] - shmem_proj[j];
|
||||
@@ -1207,6 +1238,57 @@ extern "C" __global__ void c51_loss_batched(
|
||||
d, a0, isv_signals);
|
||||
__syncthreads();
|
||||
|
||||
/* SP22 H6 Phase 3 α Step 8 backward scratch save:
|
||||
* SP_b = Σ_n p_target_n × (current_lp[upper_n] - current_lp[lower_n])
|
||||
* for d == 0 only. Inputs:
|
||||
* - shmem_proj[n] = target probabilities (input to projection,
|
||||
* unchanged by block_bellman_project_f)
|
||||
* - shmem_current_lp[k] = online log-probs for taken action a_d,
|
||||
* intact from Step a (lines 925-929; Steps b/c/d don't touch
|
||||
* shmem_current_lp by name — they use shmem_lp / shmem_proj as
|
||||
* scratch).
|
||||
* - effective_reward already includes Δ_target + Δ_online so
|
||||
* b_pos_n maps target atom n to the correct bin on the
|
||||
* UNSHIFTED support — matching the projection's actual bin
|
||||
* selection.
|
||||
*
|
||||
* Per-thread local accumulation; final block_reduce_sum_f
|
||||
* produces sample-scalar SP_b written by tid==0. Re-derives
|
||||
* lower_n / upper_n by re-running the projection's per-atom
|
||||
* arithmetic (cheap: ~10 fmul/atom per thread). */
|
||||
if (d == 0 && aux_shift_active
|
||||
&& aux_proj_logdiff_dir_out != NULL) {
|
||||
float local_sp = 0.0f;
|
||||
float v_range_for_clip = v_max - v_min;
|
||||
(void)v_range_for_clip;
|
||||
float gamma_d = gamma_eff;
|
||||
for (int n = tid; n < num_atoms; n += BLOCK_THREADS) {
|
||||
float z_n = shmem_support[n];
|
||||
float t_z = effective_reward + gamma_d * z_n * (1.0f - done);
|
||||
/* Match block_bellman_project_f's Huber compression + clamp. */
|
||||
if (t_z < 0.0f) {
|
||||
t_z = -10.0f * (1.0f - expf(t_z / 10.0f));
|
||||
}
|
||||
t_z = fminf(fmaxf(t_z, v_min), v_max);
|
||||
float b_pos = (t_z - v_min) / delta_z;
|
||||
b_pos = fminf(b_pos, (float)(num_atoms - 1) - 0.001f);
|
||||
b_pos = fmaxf(b_pos, 0.001f);
|
||||
int lower = (int)floorf(b_pos);
|
||||
int upper = lower + 1;
|
||||
lower = max(min(lower, num_atoms - 1), 0);
|
||||
upper = max(min(upper, num_atoms - 1), 0);
|
||||
float p_target_n = shmem_proj[n];
|
||||
float lp_upper = shmem_current_lp[upper];
|
||||
float lp_lower = shmem_current_lp[lower];
|
||||
local_sp += p_target_n * (lp_upper - lp_lower);
|
||||
}
|
||||
float SP_b = block_reduce_sum_f(local_sp, shmem_reduce, tid);
|
||||
if (tid == 0) {
|
||||
aux_proj_logdiff_dir_out[sample_id] = SP_b;
|
||||
}
|
||||
__syncthreads(); /* barrier before downstream label-smoothing reads */
|
||||
}
|
||||
|
||||
/* Label smoothing — health-coupled.
|
||||
*
|
||||
* eps_eff = LABEL_SMOOTHING_BASE × (1 − health)
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -17096,3 +17096,84 @@ W stays at the structural prior `[-0.5, 0.0, +0.5, 0.0]` (Step 5 init). All grad
|
||||
**Defensible intermediate checkpoint rationale**: Steps 7+9+10 land all FORWARD-path consumers of atom positions. Steps 8+11 add the dW + Adam path. Without Step 8/11, W stays at the structural prior — functionally equivalent to the "fixed-W version" the user previously characterized as "would activate α immediately for action selection but never learn state-dependent corrections." The user explicitly chose the adaptive version, but this checkpoint is structurally honest: all forward consumers see consistent shifts; W is just frozen. A smoke at this checkpoint can answer "does the structural prior alone move WR?" before investing the Step 8+11 effort. The math derivation in `7eae832f2` makes Steps 8+11 a transcription job for the next session.
|
||||
|
||||
Verification: cargo check -p ml --lib: 0 errors, 21 pre-existing warnings (baseline parity).
|
||||
|
||||
#### Atom-shift backward + Adam complete — Step 8 (c51_aux_dw_kernel) + Step 11 (adam_w_aux_kernel) — 2026-05-13
|
||||
|
||||
**Step 8 — dW backward kernel** (`c51_aux_dw_kernel.cu`, new):
|
||||
|
||||
Grid (b0_size=4, 1, 1), block (256, 1, 1). One block per W index a; each block tree-reduces dW[a] across all batch samples via warp shuffle + shmem. No atomicAdd per `pearl_no_atomicadd.md`.
|
||||
|
||||
Per-sample contributions accumulated by each block (only the matching action contributes):
|
||||
- Online side: when `a == a_d_b` (taken direction action a0 from `actions[b]` decode), `dW[a] += inv_batch × isw_b × (SP_b / Δz_b) × state_121_b`.
|
||||
- Target side: when `a == a*_b` (sampled target action from `aux_target_a_dir[b]`), `dW[a] += inv_batch × isw_b × (-γ_eff_b × (1-done_b)) × (SP_b / Δz_b) × next_state_121_b`.
|
||||
|
||||
Where:
|
||||
- `SP_b` = projection log-diff sum saved by c51_loss_kernel forward in `aux_proj_logdiff_dir[b]`. Computed as `Σ_n p_target_n × (current_lp[upper_n] - current_lp[lower_n])` for d == 0 only.
|
||||
- `Δz_b` = `per_sample_support[b * 12 + 0*3 + 2]` (dir branch dz).
|
||||
- `isw_b` = `min(is_weights[b], 10.0)` matching c51_grad_kernel's PER IS-weight clamp.
|
||||
- `inv_batch = 1/B` matching c51_grad_kernel's per-sample scaling.
|
||||
|
||||
Degenerate-support skip: when `Δz_b < 1e-7`, contributions for that sample are skipped (matches c51_loss_kernel forward's `branch_degenerate` skip).
|
||||
|
||||
**c51_loss_kernel forward scratch writes** (added in this commit):
|
||||
|
||||
- After Step c's Expected SARSA sampling sets `best_next_a` for d == 0: `aux_target_a_dir_out[sample_id] = best_next_a` (single tid==0 write).
|
||||
- After Step d's `block_bellman_project_f` returns: re-derive `b_pos_n → lower_n / upper_n` for each target atom n (re-running the projection's per-atom arithmetic; cheap ~10 fmul/atom) and accumulate `local_sp += p_target_n × (current_lp[upper_n] - current_lp[lower_n])`. Block-reduce → `aux_proj_logdiff_dir_out[sample_id]` (tid==0 write).
|
||||
|
||||
The re-derivation matches `block_bellman_project_f`'s exact arithmetic (Huber compression on `t_z < 0`, clamp to [v_min, v_max], floor/ceil index extraction) — guarantees the lower_n/upper_n used in SP computation are the SAME bins the projection wrote mass into.
|
||||
|
||||
**Step 11 — Adam-update kernel** (`adam_w_aux_kernel.cu`, new):
|
||||
|
||||
Grid (1, 1, 1), block (4, 1, 1). One thread per W slot. Standard Adam (Kingma & Ba 2015) with bias correction:
|
||||
```
|
||||
m_t = β1 m_{t-1} + (1-β1) g_t
|
||||
v_t = β2 v_{t-1} + (1-β2) g_t²
|
||||
m_hat = m_t / (1 - β1^t)
|
||||
v_hat = v_t / (1 - β2^t)
|
||||
W_t = W_{t-1} - lr × m_hat / (√v_hat + ε)
|
||||
```
|
||||
|
||||
Graph-capture-safe via device-mapped pointer args:
|
||||
- `lr_ptr` = `self.lr_dev_ptr` (matches main Adam's lr pointer pattern).
|
||||
- `step_ptr` = `self.ptrs.t_buf` (device-mapped i32 step counter; host writes via `t_pinned` each step).
|
||||
- `beta1` / `beta2` / `eps` passed by value from `sp5_isv_slots::{ADAM_BETA1_BASE, ADAM_BETA2_BASE, ADAM_EPS_BASE}` (host-stable constants per `pearl_no_host_branches_in_captured_graph`).
|
||||
|
||||
Bias-correction denominator floored at fp32 epsilon (`max(bc, 1e-30)`) to avoid /0 if β1 ≈ 0.
|
||||
|
||||
**Trainer wiring** (`gpu_dqn_trainer.rs::submit_adam_ops`):
|
||||
|
||||
Both launches added inside `submit_adam_ops` right after `launch_adam_update`:
|
||||
```rust
|
||||
self.launch_adam_update()?;
|
||||
self.launch_c51_aux_dw()?;
|
||||
self.launch_adam_w_aux()?;
|
||||
```
|
||||
|
||||
This places W gradient computation + W Adam update INSIDE the captured `adam_child` graph alongside the main Adam step. All inputs to both kernels are stable pointers (the scratch buffers, params, states are updated each step but the pointers don't change).
|
||||
|
||||
**c51_loss_kernel launcher updates** (`launch_c51_loss`): 2 new args at the end (`aux_target_a_dir_buf.raw_ptr()`, `aux_proj_logdiff_dir_buf.raw_ptr()`).
|
||||
|
||||
**New trainer fields**:
|
||||
- `aux_target_a_dir_buf: CudaSlice<i32>` shape `[batch_size]`. alloc_zeros default = 0.
|
||||
- `aux_proj_logdiff_dir_buf: CudaSlice<f32>` shape `[batch_size]`. alloc_zeros default = 0.
|
||||
- `c51_aux_dw_kernel: CudaFunction`.
|
||||
- `adam_w_aux_kernel: CudaFunction`.
|
||||
|
||||
Both scratch buffers default to 0 → when `aux_shift_active` is false in c51_loss_kernel forward (NULL W or NULL states), nothing is written → next step's dW kernel sees 0 SP → dW stays at 0 → Adam W is a no-op. So Steps 8+11 are NULL-safe at every level.
|
||||
|
||||
**Scope NOT covered in this commit (deliberate deferral)**:
|
||||
|
||||
- `dL/dstate_121`: the gradient path c51_loss → state_121 → aux head weights. This would refine the aux head to maximize C51 task performance. Currently aux head trains via its OWN supervised CE loss (next-bar direction); the missing path is a REFINEMENT, not a correctness issue. Adding it requires writing to a `dbatch_states` buffer + integrating with the existing state gradient infrastructure — non-trivial. Deferred per scope discipline.
|
||||
- Phase C1 (collector W ptr setter): rollout-time `compute_expected_q` + `quantile_q_select` still pass NULL W. Until C1, W's effect on training is internal to the trainer (no rollout-time action distribution shift). Phase D (eval-side) needs C1's setter pattern.
|
||||
- Phase D (A2 eval-side aux infrastructure): 7 tasks. Eval-time action selection doesn't see atom-shift yet.
|
||||
|
||||
**Verification**:
|
||||
- `cargo build -p ml --lib`: 0 errors, 21 pre-existing warnings (baseline parity).
|
||||
- nvcc compiles all new .cu files cleanly (no warnings/errors on c51_aux_dw_kernel.cu, adam_w_aux_kernel.cu, c51_loss_kernel.cu).
|
||||
- Full kernel recompile took 1m 05s — within expected wall-time for sm_89 target.
|
||||
|
||||
**End-state for trainer-side adaptive α**:
|
||||
|
||||
W now trains via the c51 loss gradient. Starting from structural prior `[-0.5, 0.0, +0.5, 0.0]`, Adam refines W per-action based on observed `state_121 × (loss reduction from atom-shift)`. The aux head trains independently via its supervised next-bar CE; together they form a learned cross-coupling from aux direction predictions to dir-branch Q distribution shifts.
|
||||
|
||||
Next session can dispatch a smoke at this checkpoint to measure whether the adaptive W (vs. fixed structural prior) lifts WR off 50%. If WR moves on the trainer side without rollout-side activation, Phase C1 is the next priority. If not, the hypothesis needs reconsideration.
|
||||
|
||||
Reference in New Issue
Block a user