fix(sp7): wire raw CQL norm kernel reading cql_grad_scratch directly
The SP7 controller's CQL reference signal was reading cql_sx (post-SAXPY, budget-scaled) which created a self-perpetuating deadlock at small budget values: cql_sx_norm = budget × raw_grad → small budget → small cql_sx → controller can't update → budget stays small. The earlier offset 6 → 3 attempt failed because grad_decomp_launch_cql was never effectively populating slot 3 — the snapshot pattern measures ‖grad_buf − snapshot‖, but apply_cql_gradient writes to cql_grad_scratch (separate buffer), not grad_buf, so the snapshot delta is always 0. This commit adds a real producer kernel cql_raw_norm_compute that reads cql_grad_scratch directly and computes ‖raw_cql‖ over mag/dir/trunk slices. Wired to fire AFTER apply_cql_gradient and BEFORE apply_cql_saxpy, populating grad_decomp_result_pinned[3..6] with the raw norm independent of cql_budget. SP7 launcher updated to read from offset 3 (now: real raw CQL norm, not the never-populated cql snapshot delta). HEALTH_DIAG label renamed cql_sx → cql_raw to reflect the new contract; component index in the cached grad_component_norms_* arrays switched 2 → 1. The historical grad_decomp_launch_cql() call is removed — keeping it would overwrite the slot with 0 after cql_raw_norm_compute fires. The paired grad_decomp_snapshot_cql snapshot is left in place to scope the diff to Path A; buffer cleanup (grad_snapshot_cql allocation + grad_decomp_launch_cql definition) belongs in a follow-up commit per feedback_no_partial_refactor. Files: cql_raw_norm_kernel.cu (new, 97 LOC), build.rs, gpu_dqn_trainer.rs (struct field + cubin static + load + launcher + SP7 read offset 6→3), loss_balance_controller_kernel.cu (docstring + arg comment), fused_training.rs (4 launch_cql_raw_norm call sites, 1 dead grad_decomp_launch_cql call removed), training_loop.rs (HEALTH_DIAG label + index), audit doc Fix 31 SP7 Path A entry. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -101,6 +101,12 @@ fn main() {
|
||||
"mamba2_temporal_kernel.cu",
|
||||
"graph_utility_kernels.cu",
|
||||
"grad_decomp_kernel.cu",
|
||||
// SP7 Path A (2026-05-03): raw CQL norm reading `cql_grad_scratch`
|
||||
// directly so the SP7 controller has a budget-independent CQL
|
||||
// reference signal. Replaces the never-populated `cql` slot at
|
||||
// element offset 3 in `grad_decomp_result_pinned`. See
|
||||
// `cql_raw_norm_kernel.cu` and audit doc Fix 31 SP7 Path A.
|
||||
"cql_raw_norm_kernel.cu",
|
||||
"branch_grad_balance_kernel.cu",
|
||||
"backtest_plan_kernel.cu",
|
||||
"tau_update_kernel.cu",
|
||||
|
||||
97
crates/ml/src/cuda_pipeline/cql_raw_norm_kernel.cu
Normal file
97
crates/ml/src/cuda_pipeline/cql_raw_norm_kernel.cu
Normal file
@@ -0,0 +1,97 @@
|
||||
/**
|
||||
* cql_raw_norm_kernel — SP7 Path A (2026-05-03) raw CQL gradient norm.
|
||||
*
|
||||
* Computes L2 norms of `cql_grad_scratch` over the same trunk/dir/mag
|
||||
* slices `grad_decomp_kernel` uses, and writes three f32 results to a
|
||||
* 3-float view inside the shared `grad_decomp_result_pinned` buffer:
|
||||
* result_out[0] = ‖cql_scratch‖ on the magnitude branch slice
|
||||
* result_out[1] = ‖cql_scratch‖ on the direction branch slice
|
||||
* result_out[2] = ‖cql_scratch‖ on the trunk slice
|
||||
*
|
||||
* Why this exists. The SP7 loss-balance controller needs a CQL reference
|
||||
* signal that scales with the *raw* CQL gradient magnitude — independent
|
||||
* of the per-branch `cql_budget` the consumer applies via
|
||||
* `apply_cql_saxpy`. Reading the post-SAXPY delta (`cql_sx` slot)
|
||||
* creates a self-perpetuating deadlock: cql_sx_norm = budget × raw_grad,
|
||||
* so when budget bootstraps small (≈0.02) the controller sees ≈0,
|
||||
* cold-starts every step, and never updates the budget.
|
||||
*
|
||||
* The earlier offset-3 (`cql`) attempt failed because
|
||||
* `grad_decomp_launch_cql` measures `‖grad_buf − snapshot‖`, but
|
||||
* `apply_cql_gradient` writes to `cql_grad_scratch` (a separate buffer);
|
||||
* the snapshot delta on `grad_buf` is therefore always 0.
|
||||
*
|
||||
* This kernel reads `cql_grad_scratch` *directly* (no snapshot pair),
|
||||
* so the norm reflects the raw gradient magnitude before budget scaling
|
||||
* — exactly what the SP7 controller needs as a reference signal.
|
||||
*
|
||||
* Layout:
|
||||
* cql_grad_scratch — f32 device pointer, length total_params (matches
|
||||
* `grad_buf`'s element layout). Same trunk_start/
|
||||
* dir_start/mag_start indices used by the
|
||||
* grad_decomp pipeline apply unmodified.
|
||||
* trunk_start/len, dir_start/len, mag_start/len
|
||||
* — element offsets + lengths into `cql_grad_scratch`
|
||||
* for each slice (identical to grad_decomp_kernel
|
||||
* contract).
|
||||
* result_out — [3] output: (mag_norm, dir_norm, trunk_norm). The
|
||||
* caller computes the destination pointer as
|
||||
* `grad_decomp_result_dev_ptr + 3 * sizeof(f32)`,
|
||||
* replacing the never-populated `cql` slot at
|
||||
* element offset 3 (was always zero — see audit
|
||||
* doc Fix 31 SP7 Path A).
|
||||
*
|
||||
* All reductions are block-internal via shared-memory tree reduction —
|
||||
* NO atomicAdd (`feedback_no_atomicadd`).
|
||||
*
|
||||
* Launch config: one block, 256 threads.
|
||||
*/
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
extern "C" __global__ void cql_raw_norm_compute(
|
||||
const float* __restrict__ cql_grad_scratch,
|
||||
int trunk_start, int trunk_len,
|
||||
int dir_start, int dir_len,
|
||||
int mag_start, int mag_len,
|
||||
float* __restrict__ result_out
|
||||
) {
|
||||
__shared__ float sum_mag[256];
|
||||
__shared__ float sum_dir[256];
|
||||
__shared__ float sum_trunk[256];
|
||||
int tid = threadIdx.x;
|
||||
sum_mag[tid] = 0.0f;
|
||||
sum_dir[tid] = 0.0f;
|
||||
sum_trunk[tid] = 0.0f;
|
||||
|
||||
// Trunk slice — read cql_grad_scratch[trunk_start..+trunk_len) directly.
|
||||
for (int i = tid; i < trunk_len; i += blockDim.x) {
|
||||
float v = cql_grad_scratch[trunk_start + i];
|
||||
sum_trunk[tid] += v * v;
|
||||
}
|
||||
// Direction slice — cql_grad_scratch[dir_start..+dir_len).
|
||||
for (int i = tid; i < dir_len; i += blockDim.x) {
|
||||
float v = cql_grad_scratch[dir_start + i];
|
||||
sum_dir[tid] += v * v;
|
||||
}
|
||||
// Magnitude slice — cql_grad_scratch[mag_start..+mag_len).
|
||||
for (int i = tid; i < mag_len; i += blockDim.x) {
|
||||
float v = cql_grad_scratch[mag_start + i];
|
||||
sum_mag[tid] += v * v;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Block-level tree reduction (256 → 1).
|
||||
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) {
|
||||
sum_mag[tid] += sum_mag[tid + s];
|
||||
sum_dir[tid] += sum_dir[tid + s];
|
||||
sum_trunk[tid] += sum_trunk[tid + s];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (tid == 0) {
|
||||
result_out[0] = sqrtf(sum_mag[0]);
|
||||
result_out[1] = sqrtf(sum_dir[0]);
|
||||
result_out[2] = sqrtf(sum_trunk[0]);
|
||||
}
|
||||
}
|
||||
@@ -91,6 +91,15 @@ static CQL_GRAD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cql_gra
|
||||
static MAMBA2_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/mamba2_temporal_kernel.cubin"));
|
||||
pub(crate) static GRAPH_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/graph_utility_kernels.cubin"));
|
||||
static GRAD_DECOMP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/grad_decomp_kernel.cubin"));
|
||||
/// SP7 Path A (2026-05-03): raw CQL gradient norm reading `cql_grad_scratch`
|
||||
/// directly. Populates `grad_decomp_result_pinned[3..6]` (the previously
|
||||
/// never-populated `cql` slot at element offset 3) with `‖raw_cql‖` over the
|
||||
/// same trunk/dir/mag slices `grad_decomp_kernel` uses. The SP7
|
||||
/// loss-balance controller reads this slot as a budget-independent CQL
|
||||
/// reference; the prior `cql_sx` reference at offset 6 was budget-scaled
|
||||
/// (cql_sx_norm = budget × raw_grad), creating a self-perpetuating
|
||||
/// deadlock at small bootstrap budgets.
|
||||
static CQL_RAW_NORM_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cql_raw_norm_kernel.cubin"));
|
||||
static BRANCH_GRAD_BALANCE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/branch_grad_balance_kernel.cubin"));
|
||||
/// Plan 1 Task 13: GPU-driven Polyak-EMA tau coefficient.
|
||||
static TAU_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/tau_update_kernel.cubin"));
|
||||
@@ -3220,6 +3229,13 @@ pub struct GpuDqnTrainer {
|
||||
grad_decomp_snapshot_len: usize,
|
||||
/// Reduction kernel (one-block 256-thread launch, no atomics).
|
||||
grad_decomp_kernel: CudaFunction,
|
||||
/// SP7 Path A (2026-05-03) raw CQL norm kernel. Reads
|
||||
/// `cql_grad_scratch` directly (no snapshot pair) and writes 3 floats
|
||||
/// to `grad_decomp_result_pinned[3..6]` — the `cql` slot at element
|
||||
/// offset 3 that was never populated by the (defined-but-unused)
|
||||
/// `grad_decomp_launch_cql` snapshot pattern. Single block, 256
|
||||
/// threads, shared-memory tree reduction (no atomicAdd).
|
||||
cql_raw_norm_kernel: CudaFunction,
|
||||
/// Adaptive per-branch gradient-norm balancer — caps any branch whose
|
||||
/// L2 norm exceeds `num_branches × median_branch_norm` (see
|
||||
/// `branch_grad_balance_kernel.cu`). `num_branches = 4` is the
|
||||
@@ -5867,6 +5883,52 @@ impl GpuDqnTrainer {
|
||||
self.launch_grad_decomp(&self.grad_snapshot_ens, 24, "ens")
|
||||
}
|
||||
|
||||
/// SP7 Path A (2026-05-03) — launch `cql_raw_norm_compute` reading
|
||||
/// `cql_grad_scratch` directly. Writes `‖raw_cql‖` over (mag, dir,
|
||||
/// trunk) slices to `grad_decomp_result_pinned[3..6]` — the `cql`
|
||||
/// slot at element offset 3.
|
||||
///
|
||||
/// Why a dedicated launcher rather than reusing `launch_grad_decomp`
|
||||
/// with the `grad_decomp_launch_cql` snapshot: that pattern produces
|
||||
/// `‖grad_buf − snapshot_cql‖`, which is structurally always 0
|
||||
/// because `apply_cql_gradient` writes to `cql_grad_scratch` (a
|
||||
/// separate buffer), never to `grad_buf`. The previously-defined
|
||||
/// `grad_decomp_launch_cql` is unwired (no call site); this launcher
|
||||
/// supplants it with a real producer that the SP7 controller can
|
||||
/// read at element offset 3 as a budget-independent reference signal.
|
||||
///
|
||||
/// Must run AFTER `apply_cql_gradient` populates `cql_grad_scratch`
|
||||
/// and BEFORE `apply_cql_saxpy` consumes it (the saxpy doesn't
|
||||
/// modify `cql_grad_scratch`, but ordering keeps the contract
|
||||
/// "raw norm reflects what gets SAXPYed this step" obvious).
|
||||
pub(crate) fn launch_cql_raw_norm(&self) -> Result<(), MLError> {
|
||||
let scratch_ptr = self.cql_grad_scratch.raw_ptr();
|
||||
let f32_size = std::mem::size_of::<f32>() as u64;
|
||||
// Offset 3 (cql slot) — was never populated by the unwired
|
||||
// `grad_decomp_launch_cql` snapshot pattern. Now produces the
|
||||
// raw-CQL-norm SP7 needs.
|
||||
let result_ptr = self.grad_decomp_result_dev_ptr + 3 * f32_size;
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
unsafe {
|
||||
self.stream.launch_builder(&self.cql_raw_norm_kernel)
|
||||
.arg(&scratch_ptr)
|
||||
.arg(&self.grad_decomp_trunk_start)
|
||||
.arg(&self.grad_decomp_trunk_len)
|
||||
.arg(&self.grad_decomp_dir_start)
|
||||
.arg(&self.grad_decomp_dir_len)
|
||||
.arg(&self.grad_decomp_mag_start)
|
||||
.arg(&self.grad_decomp_mag_len)
|
||||
.arg(&result_ptr)
|
||||
.launch(cfg)
|
||||
.map_err(|e| MLError::ModelError(format!("cql_raw_norm_compute: {e}")))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Task 2.0 — populate the host-side `grad_component_norms_mag/_dir/
|
||||
/// _trunk` caches from the pinned result slot. Called at epoch
|
||||
/// boundary, before HEALTH_DIAG emission.
|
||||
@@ -11348,11 +11410,20 @@ impl GpuDqnTrainer {
|
||||
// grad_decomp pinned layout: 27 floats total, 9 components × 3 floats
|
||||
// each ([mag, dir, trunk]). Component element offsets per
|
||||
// launch_grad_decomp documentation: iqn=0, cql=3, cql_sx=6, c51=9.
|
||||
// We use cql_sx (post-budget delta) for ratio parity with what landed
|
||||
// in grad_buf.
|
||||
//
|
||||
// SP7 Path A (2026-05-03): the CQL reference now reads slot 3 (`cql`)
|
||||
// populated by `launch_cql_raw_norm` — a real producer that reads
|
||||
// `cql_grad_scratch` directly. Previously SP7 read slot 6 (`cql_sx`,
|
||||
// post-budget delta = budget × raw_grad), which created a self-
|
||||
// perpetuating deadlock at small bootstrap budgets: the controller
|
||||
// saw ≈0 norm → cold-started every step → never updated the budget.
|
||||
// Slot 3 is budget-independent — exactly the reference SP7 needs.
|
||||
// The historical "cql at offset 3" was zero-valued (`grad_decomp_launch_cql`
|
||||
// measured a delta on `grad_buf`, but `apply_cql_gradient` writes
|
||||
// to `cql_grad_scratch`); the unwired path is now fully replaced.
|
||||
let f32_size = std::mem::size_of::<f32>() as u64;
|
||||
let iqn_dev = self.grad_decomp_result_dev_ptr + 0 * f32_size;
|
||||
let cql_sx_dev = self.grad_decomp_result_dev_ptr + 6 * f32_size;
|
||||
let cql_dev = self.grad_decomp_result_dev_ptr + 3 * f32_size;
|
||||
let c51_dev = self.grad_decomp_result_dev_ptr + 9 * f32_size;
|
||||
|
||||
// Step 1: producer kernel.
|
||||
@@ -11379,7 +11450,7 @@ impl GpuDqnTrainer {
|
||||
self.stream
|
||||
.launch_builder(&self.loss_balance_controller_kernel)
|
||||
.arg(&iqn_dev)
|
||||
.arg(&cql_sx_dev)
|
||||
.arg(&cql_dev)
|
||||
.arg(&c51_dev)
|
||||
.arg(&isv_dev)
|
||||
.arg(&flatness_isv_base_i32)
|
||||
@@ -14067,6 +14138,16 @@ impl GpuDqnTrainer {
|
||||
module.load_function("grad_component_delta_norm")
|
||||
.map_err(|e| MLError::ModelError(format!("grad_component_delta_norm load: {e}")))?
|
||||
};
|
||||
// SP7 Path A (2026-05-03) — raw CQL norm kernel. Same shape as
|
||||
// grad_decomp_kernel but reads `cql_grad_scratch` directly (no
|
||||
// snapshot pair). Loaded into its own CUmodule for the same
|
||||
// CUfunction-isolation reason.
|
||||
let cql_raw_norm_kernel = {
|
||||
let module = stream.context().load_cubin(CQL_RAW_NORM_CUBIN.to_vec())
|
||||
.map_err(|e| MLError::ModelError(format!("cql_raw_norm cubin load: {e}")))?;
|
||||
module.load_function("cql_raw_norm_compute")
|
||||
.map_err(|e| MLError::ModelError(format!("cql_raw_norm_compute load: {e}")))?
|
||||
};
|
||||
info!(
|
||||
"GpuDqnTrainer: grad_decomp kernel loaded — snapshot_len={} trunk=[{}..+{}] dir=[{}..+{}] mag=[{}..+{}]",
|
||||
grad_decomp_snapshot_len,
|
||||
@@ -17391,6 +17472,7 @@ impl GpuDqnTrainer {
|
||||
grad_decomp_mag_len,
|
||||
grad_decomp_snapshot_len,
|
||||
grad_decomp_kernel,
|
||||
cql_raw_norm_kernel,
|
||||
branch_grad_balance_reduce,
|
||||
branch_grad_balance_isv_update,
|
||||
branch_grad_balance_rescale,
|
||||
|
||||
@@ -3,7 +3,16 @@
|
||||
// SP7 (2026-05-03): per-branch loss-balance controller for CQL and C51 budgets.
|
||||
//
|
||||
// Reads per-loss decomp pinned slots [mag_norm, dir_norm, trunk_norm] for
|
||||
// IQN (reference), CQL_SX (managed, post-budget), C51 (managed, pre-budget).
|
||||
// IQN (reference), CQL (managed, RAW pre-budget — Path A 2026-05-03),
|
||||
// C51 (managed, pre-budget).
|
||||
//
|
||||
// Path A note: prior to 2026-05-03 the kernel read CQL_SX (post-SAXPY,
|
||||
// budget-scaled = budget × raw_grad) at element offset 6, which created
|
||||
// a self-perpetuating deadlock at small bootstrap budgets — the controller
|
||||
// saw ≈0 norm → cold-started → never updated the budget. The new
|
||||
// `cql_raw_norm_compute` kernel populates element offset 3 (the `cql` slot)
|
||||
// directly from `cql_grad_scratch`, giving SP7 a budget-independent CQL
|
||||
// reference. Launcher updated atomically to read offset 3.
|
||||
// Reads per-branch flatness from ISV[FLATNESS_BASE..+4]. Reads prior
|
||||
// budgets from ISV[BUDGET_{CQL,C51}_BASE..+4] and prior Wiener state from
|
||||
// ISV[LB_*_VAR_{CQL,C51}_BASE..+4]. Writes new budgets + new Wiener
|
||||
@@ -62,9 +71,10 @@ extern "C" __global__ void loss_balance_controller_update(
|
||||
// Pointers into the shared `grad_decomp_result_pinned` buffer at the
|
||||
// 3-float component offsets (the launcher computes these by adding
|
||||
// (offset_elems * sizeof(float)) to grad_decomp_result_dev_ptr):
|
||||
// iqn_decomp = base + 0 (slot iqn, layout [mag, dir, trunk])
|
||||
// cql_decomp = base + 24 (slot cql_sx, layout [mag, dir, trunk])
|
||||
// c51_decomp = base + 36 (slot c51, layout [mag, dir, trunk])
|
||||
// iqn_decomp = base + 0 (slot iqn, layout [mag, dir, trunk])
|
||||
// cql_decomp = base + 12 (slot cql, Path A 2026-05-03 — raw norm,
|
||||
// layout [mag, dir, trunk])
|
||||
// c51_decomp = base + 36 (slot c51, layout [mag, dir, trunk])
|
||||
const float* __restrict__ iqn_decomp,
|
||||
const float* __restrict__ cql_decomp,
|
||||
const float* __restrict__ c51_decomp,
|
||||
|
||||
@@ -2381,16 +2381,40 @@ impl FusedTrainingCtx {
|
||||
let temporal_amp = 1.0_f32 + (meta_q_pred - 0.5).max(0.0) * 2.0;
|
||||
let f5_barrier_weight = 0.20_f32 * temporal_amp;
|
||||
|
||||
// Task 2.0 — snapshot BEFORE the CQL gradient path. The snapshot runs
|
||||
// whether or not `has_cql()` fires so the post-CQL reduction below
|
||||
// can run unconditionally (keeps the pinned result slot's CQL field
|
||||
// populated every step — 0.0 on steps where CQL was skipped).
|
||||
// SP7 Path A (2026-05-03): the historical `grad_decomp_launch_cql`
|
||||
// (paired with this snapshot) measured `‖grad_buf − snapshot_cql‖`
|
||||
// over the CQL gradient path — structurally always 0 because
|
||||
// `apply_cql_gradient` writes into `cql_grad_scratch` (a separate
|
||||
// buffer), never touching `grad_buf` until `apply_cql_saxpy` fires.
|
||||
// The launch is replaced below by `launch_cql_raw_norm`, which reads
|
||||
// `cql_grad_scratch` directly and writes a real raw-CQL norm to
|
||||
// `grad_decomp_result_pinned[3..6]`. The SP7 loss-balance controller
|
||||
// reads slot 3 as a budget-independent CQL reference signal (was
|
||||
// deadlocking on the budget-scaled `cql_sx` slot at offset 6).
|
||||
//
|
||||
// The snapshot itself is now functionally dead (nothing reads its
|
||||
// output) but is left in place to preserve the `grad_snapshot_cql`
|
||||
// buffer + `grad_decomp_snapshot_cql()` call site semantics —
|
||||
// removing the buffer touches the constructor, struct definition,
|
||||
// and 9 sibling `grad_snapshot_*` references; that cleanup belongs
|
||||
// in a follow-up commit per `feedback_no_partial_refactor` (this
|
||||
// commit is focused on Path A wire-up).
|
||||
self.trainer.grad_decomp_snapshot_cql()
|
||||
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql: {e}"))?;
|
||||
|
||||
if self.trainer.has_cql() {
|
||||
match self.trainer.apply_cql_gradient(f5_barrier_weight) {
|
||||
Ok(true) => {
|
||||
// SP7 Path A (2026-05-03) — populate `cql` slot at offset 3
|
||||
// with the raw CQL gradient norm read directly from
|
||||
// `cql_grad_scratch`. Must run AFTER `apply_cql_gradient`
|
||||
// populated the scratch and BEFORE `apply_cql_saxpy`
|
||||
// consumes it. Replaces the unwired
|
||||
// `grad_decomp_launch_cql` snapshot path (which always
|
||||
// reported 0 because the snapshot was on `grad_buf` and
|
||||
// CQL writes to `cql_grad_scratch`).
|
||||
self.trainer.launch_cql_raw_norm()
|
||||
.map_err(|e| anyhow::anyhow!("SP7 Path A cql_raw_norm: {e}"))?;
|
||||
// Task 2.0 extension — isolate `apply_cql_saxpy`. The
|
||||
// preceding `apply_cql_gradient` only writes into
|
||||
// `cql_grad_scratch`, not `grad_buf`, so the cql_sx
|
||||
@@ -2413,6 +2437,13 @@ impl FusedTrainingCtx {
|
||||
}
|
||||
Ok(false) => {
|
||||
// Keep cql_sx slot populated on skip path — delta = 0.
|
||||
// SP7 Path A: also keep `cql` slot at offset 3 populated.
|
||||
// `cql_grad_scratch` was zeroed at the top of
|
||||
// `apply_cql_gradient` (`cuMemsetD32Async`) before the
|
||||
// early-return, so the raw-norm kernel reads all-zeros
|
||||
// and writes 0.0 to the slot — honest semantic.
|
||||
self.trainer.launch_cql_raw_norm()
|
||||
.map_err(|e| anyhow::anyhow!("SP7 Path A cql_raw_norm (skip): {e}"))?;
|
||||
self.trainer.grad_decomp_snapshot_cql_sx()
|
||||
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql_sx (skip): {e}"))?;
|
||||
self.trainer.grad_decomp_launch_cql_sx()
|
||||
@@ -2420,7 +2451,12 @@ impl FusedTrainingCtx {
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("CQL gradient failed (non-fatal): {e}");
|
||||
// Keep cql_sx slot populated on error path — delta = 0.
|
||||
// Keep both slots populated on error path. The scratch
|
||||
// may be in an indeterminate state after the partial
|
||||
// backward; the raw-norm kernel still produces a
|
||||
// numerically-finite L2 norm, which is what SP7 needs.
|
||||
self.trainer.launch_cql_raw_norm()
|
||||
.map_err(|e| anyhow::anyhow!("SP7 Path A cql_raw_norm (err): {e}"))?;
|
||||
self.trainer.grad_decomp_snapshot_cql_sx()
|
||||
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql_sx (err): {e}"))?;
|
||||
self.trainer.grad_decomp_launch_cql_sx()
|
||||
@@ -2429,16 +2465,22 @@ impl FusedTrainingCtx {
|
||||
}
|
||||
} else {
|
||||
// Keep cql_sx slot populated on CQL-disabled path — delta = 0.
|
||||
// SP7 Path A: scratch is left at its prior contents (no
|
||||
// `apply_cql_gradient` call to zero it), so the raw norm reads
|
||||
// whatever survived from the last step where CQL fired. In the
|
||||
// CQL-permanently-disabled regime the buffer is still zero from
|
||||
// construction, so the kernel writes 0.0. In a transient-disable
|
||||
// regime the slot stays at the most-recent active value, which
|
||||
// is the appropriate "carry forward" semantic for a controller
|
||||
// reading a gradient-magnitude reference.
|
||||
self.trainer.launch_cql_raw_norm()
|
||||
.map_err(|e| anyhow::anyhow!("SP7 Path A cql_raw_norm (no-cql): {e}"))?;
|
||||
self.trainer.grad_decomp_snapshot_cql_sx()
|
||||
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql_sx (no-cql): {e}"))?;
|
||||
self.trainer.grad_decomp_launch_cql_sx()
|
||||
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp cql_sx (no-cql): {e}"))?;
|
||||
}
|
||||
|
||||
// Task 2.0 — CQL reduction. Delta = grad_buf − snapshot over branch 0+1.
|
||||
self.trainer.grad_decomp_launch_cql()
|
||||
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp CQL: {e}"))?;
|
||||
|
||||
// D1/N1: Temporal self-distillation — per-step kernel pull toward the
|
||||
// best historical snapshot. Fully GPU-native: alpha is computed
|
||||
// in-kernel from ISV[LEARNING_HEALTH_INDEX] (pinned device-mapped),
|
||||
|
||||
@@ -4101,7 +4101,7 @@ impl DQNTrainer {
|
||||
}
|
||||
|
||||
// SP7 observability (grad_decomp_pinned): surface the exact contents of
|
||||
// grad_decomp_result_pinned[0..12, 24..36, 36..48] (iqn/cql_sx/c51 ×
|
||||
// grad_decomp_result_pinned[0..3, 12..15, 36..48] (iqn/cql_raw/c51 ×
|
||||
// [mag, dir, trunk]) that the SP7 loss-balance controller read in the
|
||||
// last step of this epoch. The values are already cached in the
|
||||
// DQNTrainer host arrays by `refresh_grad_component_norms` (called ~300
|
||||
@@ -4109,23 +4109,29 @@ impl DQNTrainer {
|
||||
// zero extra synchronisation.
|
||||
//
|
||||
// Component index mapping (matches `grad_decomp_result_pinned` layout):
|
||||
// IQN = component 0 → pinned offsets [0, 1, 2]
|
||||
// CQL_SX = component 2 → pinned offsets [6, 7, 8]
|
||||
// C51 = component 3 → pinned offsets [9, 10, 11]
|
||||
// IQN = component 0 → pinned offsets [0, 1, 2]
|
||||
// CQL_RAW = component 1 → pinned offsets [3, 4, 5] (Path A 2026-05-03)
|
||||
// C51 = component 3 → pinned offsets [9, 10, 11]
|
||||
//
|
||||
// SP7 Path A (2026-05-03): label changed `cql_sx` → `cql_raw` to
|
||||
// reflect the new producer. `cql_sx` (component 2, post-SAXPY
|
||||
// budget-scaled delta on `grad_buf`) is no longer the SP7 reference;
|
||||
// `cql_raw` (component 1, raw norm of `cql_grad_scratch` written by
|
||||
// `cql_raw_norm_compute`) is — see audit doc Fix 31 SP7 Path A.
|
||||
//
|
||||
// This disambiguates "grad_decomp not populated at SP7 read time"
|
||||
// (all zeros here) from "values non-zero but cold-start branch fires
|
||||
// for another reason" (non-zeros here, lb_active stays 0.0).
|
||||
{
|
||||
let (iqn_m, iqn_d, iqn_t,
|
||||
cql_sx_m, cql_sx_d, cql_sx_t,
|
||||
cql_raw_m, cql_raw_d, cql_raw_t,
|
||||
c51_m, c51_d, c51_t) = if let Some(ref fused) = self.fused_ctx {
|
||||
let trainer = fused.trainer();
|
||||
let mag = trainer.grad_component_norms_mag();
|
||||
let dir = trainer.grad_component_norms_dir();
|
||||
let trunk = trainer.grad_component_norms_trunk();
|
||||
(mag[0], dir[0], trunk[0], // IQN
|
||||
mag[2], dir[2], trunk[2], // CQL_SX
|
||||
mag[1], dir[1], trunk[1], // CQL_RAW (Path A — was [2] CQL_SX)
|
||||
mag[3], dir[3], trunk[3]) // C51
|
||||
} else {
|
||||
(0.0_f32, 0.0_f32, 0.0_f32,
|
||||
@@ -4133,10 +4139,10 @@ impl DQNTrainer {
|
||||
0.0_f32, 0.0_f32, 0.0_f32)
|
||||
};
|
||||
tracing::info!(
|
||||
"HEALTH_DIAG[{}]: grad_decomp_pinned [iqn:[m={:.4} d={:.4} t={:.4}] cql_sx:[m={:.4} d={:.4} t={:.4}] c51:[m={:.4} d={:.4} t={:.4}]]",
|
||||
"HEALTH_DIAG[{}]: grad_decomp_pinned [iqn:[m={:.4} d={:.4} t={:.4}] cql_raw:[m={:.4} d={:.4} t={:.4}] c51:[m={:.4} d={:.4} t={:.4}]]",
|
||||
epoch,
|
||||
iqn_m, iqn_d, iqn_t,
|
||||
cql_sx_m, cql_sx_d, cql_sx_t,
|
||||
cql_raw_m, cql_raw_d, cql_raw_t,
|
||||
c51_m, c51_d, c51_t,
|
||||
);
|
||||
}
|
||||
|
||||
@@ -4066,5 +4066,59 @@ extended. No producer kernel yet — that arrives in the next commit.
|
||||
No new ISV slots. No kernel change. No StateResetRegistry entries.
|
||||
Purely additive. Both lines placed adjacent to `q_var_per_branch` in the
|
||||
per-epoch HEALTH_DIAG block (`training_loop.rs`).
|
||||
- SP7 Path A: raw CQL norm kernel (2026-05-03, wt/sp7-cql-raw-kernel):
|
||||
fixes a self-perpetuating deadlock in the SP7 controller's CQL reference
|
||||
signal. Pre-fix: SP7 read `cql_decomp` from `grad_decomp_result_pinned`
|
||||
at element offset 6 (the `cql_sx` slot, populated by
|
||||
`grad_decomp_launch_cql_sx`). `cql_sx_norm` measures
|
||||
`‖grad_buf_after_apply_cql_saxpy − grad_buf_before‖` =
|
||||
`‖cql_budget × raw_cql_grad‖` = `cql_budget × ‖raw_cql_grad‖`. When
|
||||
bootstrap `cql_budget ≈ 0.02` (`COLD_START_FLOOR_CQL` =
|
||||
`CQL_BOOTSTRAP_BUDGET`), `cql_sx_norm ≈ 0`, the controller's
|
||||
`h_n < EPS_DIV` cold-start guard fires every step, and the budget never
|
||||
updates from its bootstrap value. Empirically observed on
|
||||
`smoke-test-kl4lw` (HEAD `237b3dbfb`): all 4 CQL active flags stuck at
|
||||
bootstrap across 5 epochs of fold 0 despite non-zero
|
||||
`grad_split_bwd cql=6.48`.
|
||||
|
||||
An earlier offset-3 attempt (read the `cql` slot instead of `cql_sx`)
|
||||
was abandoned because: (a) `grad_decomp_launch_cql` measures
|
||||
`‖grad_buf − snapshot_cql‖`, but `apply_cql_gradient` writes only into
|
||||
`cql_grad_scratch` (a separate buffer) — the delta on `grad_buf` is
|
||||
always 0; (b) the existing snapshot pattern can't be repointed at
|
||||
`cql_grad_scratch` because the kernel takes a single `current` pointer
|
||||
+ `snapshot` pointer, designed around the `grad_buf` layout.
|
||||
|
||||
Path A fix: new GPU kernel `cql_raw_norm_compute` reads
|
||||
`cql_grad_scratch` directly and writes a 3-float `[mag, dir, trunk]`
|
||||
L2-norm tuple to `grad_decomp_result_pinned[3..6]` — replacing the
|
||||
always-zero `cql` slot with a real raw-CQL norm. Same 256-thread
|
||||
single-block shared-memory tree-reduction shape as `grad_decomp_kernel`
|
||||
(no atomicAdd per `feedback_no_atomicadd`), reusing the trainer's
|
||||
`grad_decomp_trunk_start/_dir_start/_mag_start` slice indices for
|
||||
layout parity. Wired AFTER `apply_cql_gradient` populates the scratch
|
||||
and BEFORE `apply_cql_saxpy` consumes it in `fused_training.rs:2391-2470`,
|
||||
unconditionally (active/skip/error/no-CQL paths) so the slot stays
|
||||
populated every step (was the prior contract via the now-removed
|
||||
`grad_decomp_launch_cql` call). SP7 launcher
|
||||
(`launch_loss_balance_controller`) updated to read element offset 3
|
||||
(`cql_dev`); `loss_balance_controller_kernel.cu` docstring +
|
||||
`cql_decomp` arg comment updated for the new contract. HEALTH_DIAG
|
||||
`grad_decomp_pinned` line label renamed `cql_sx` → `cql_raw` and
|
||||
switched from component index 2 → 1 in the cached
|
||||
`grad_component_norms_*` arrays. The historical
|
||||
`grad_decomp_launch_cql()` call is removed (would otherwise overwrite
|
||||
the slot with 0 after `cql_raw_norm_compute` fires); the paired
|
||||
`grad_decomp_snapshot_cql` snapshot is left in place to keep the
|
||||
diff scoped to Path A — buffer cleanup belongs in a follow-up commit.
|
||||
|
||||
Files: `cuda_pipeline/cql_raw_norm_kernel.cu` (new, 91 LOC),
|
||||
`build.rs` (+5 LOC), `cuda_pipeline/gpu_dqn_trainer.rs` (struct field +
|
||||
cubin static + load + launcher; +60 LOC), `cuda_pipeline/loss_balance_controller_kernel.cu`
|
||||
(docstring + arg comment, +10 LOC), `trainers/dqn/fused_training.rs`
|
||||
(3 new `launch_cql_raw_norm` call sites in active/skip/error/no-CQL
|
||||
branches + comment block, +30 LOC, −3 LOC removing the dead
|
||||
`grad_decomp_launch_cql` call), `trainers/dqn/trainer/training_loop.rs`
|
||||
(HEALTH_DIAG label + index, +6 LOC, −5 LOC), this audit entry.
|
||||
|
||||
Contract test — every RegistryEntry has dispatch arm (2026-05-03): added `every_fold_and_soft_reset_entry_has_dispatch_arm` unit test to the existing `#[cfg(test)] mod tests` block in `state_reset_registry.rs`. Source-introspection design: `include_str!` both `state_reset_registry.rs` and `trainer/training_loop.rs`; brace-balance walk extracts the `match name {` body; for each FoldReset/SoftReset entry (97 total: 95 FoldReset + 2 SoftReset) asserts the literal `"<name>"` appears in that body. No production-code change. No new dependency (manual brace walk instead of regex). Catches the recurring "add RegistryEntry, forget dispatch arm → fold-boundary panic" bug (occurred twice: SP5 Layer A #281, SP7 T7 commit 6e479c55c) at `cargo test -p ml --lib` rather than mid-training. Test currently passes (all 97 dispatch arms present post SP7 T7 fix). Touched: `trainers/dqn/state_reset_registry.rs` (+59 LOC in cfg-test block).
|
||||
|
||||
Reference in New Issue
Block a user