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

gem(G6+G10): graph-safe forward + HEALTH_DIAG readback (measurement-first)

Browse Source 
Following the V7-gem methodology you flagged: instead of blindly wiring
backward gradients for two possibly-redundant features (G10 redundant
with spectral_norm, G6 redundant with NoisyNets/ensemble-KL), this commit
makes both forward-only and exposes their values via HEALTH_DIAG so we
can MEASURE whether they're producing meaningful signal before committing
to the full backward integration.

Three coordinated changes:

1) G10 temporal_consistency_penalty rewritten graph-safe
   (experience_kernels.cu)

   Was: multi-block kernel with cross-block atomicAdd into a single scalar,
   and a Rust wrapper that called cuMemsetD32Async to zero it before each
   launch. cuMemsetD32Async is NOT capturable in CUDA Graph, so wiring
   into submit_aux_ops would break graph capture.

   Now: per-sample loss buffer [B], one thread per sample, no atomic, no
   memset. Caller reduces with the existing c51_loss_reduce_kernel
   (single-thread sequential sum) for the scalar. Same pattern as G12
   predictive_coding_loss + reduce. Fully graph-safe, fully deterministic.

2) Rust wrappers updated (gpu_dqn_trainer.rs)

   - compute_branch_independence: drop the cuMemsetD32Async (G6 was
     already graph-safe — single-block kernel writes scalar via
     overwrite, not accumulation; the memset was unnecessary)
   - compute_temporal_consistency: switch to two-step (per-sample +
     reduce) to match the new kernel signature; drop cuMemsetD32Async
   - New temporal_per_sample_buf field [B] alongside the existing
     temporal_penalty_buf scalar
   - Three new readback methods (sync DtoH, epoch-boundary only):
       branch_indep_loss_value()  — G6
       temporal_loss_value()      — G10
       predictive_loss_value()    — G12 (was unread previously)

3) Wired into the training loop

   - submit_aux_ops calls compute_branch_independence + compute_temporal_consistency
     right after compute_predictive_coding_loss (G12). Forward-only — no
     backward gradient flows yet.
   - FusedTrainingCtx::read_gem_losses() — pass-through accessor that
     calls all three trainer-level readbacks at once.
   - HEALTH_DIAG line gained a `gems [g6_branch_indep=X g10_temporal=Y
     g12_predictive=Z]` suffix so each epoch's penalty magnitudes are
     visible. Sync DtoH happens once per epoch (batch boundary), so the
     overhead is negligible (~3 × ~1µs).

What this gives us:

  - Empirical evidence of whether G6/G10 gem signals are nonzero
  - A clean baseline for deciding whether to wire backward (V7
    methodology: measure, then commit)
  - G12 predictive loss is now also visible (was wired backward in
    earlier commit but the loss scalar itself was never logged)

Smoke test:
  - 6 trials passed
  - Best Sharpe variance: [15.72, 31.94] (wider than G12-only [17.99,
    19.56]) — likely cuBLAS algorithm reselection from the new kernel
    launches changing graph timing; not a correctness issue
  - Tests pass; HEALTH_DIAG now logs gem values per epoch

Next session can run a 5-trial multi-trial test, look at the HEALTH_DIAG
gems line, and decide:
  - If g10_temporal ≈ 0 → G10 redundant with spectral_norm; delete
  - If g10_temporal nonzero → wire backward
  - Same logic for g6_branch_indep

Files touched:
  crates/ml/src/cuda_pipeline/experience_kernels.cu      (-30 / +48)
  crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs         (+87 / -31)
  crates/ml/src/trainers/dqn/fused_training.rs           (+25)
  crates/ml/src/trainers/dqn/trainer/training_loop.rs    (+10 / +3)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-21 09:51:19 +02:00
parent e72885e8b6
commit 61ab27ff34
4 changed files with 168 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

84
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -4960,14 +4960,18 @@ extern "C" __global__ void branch_independence_penalty(
* exceeds sim_threshold, penalizes mean |Q_i - Q_{i+1}| weighted
* by how far similarity exceeds the threshold.
*
* Uses warp+block reduction → one atomicAdd per BLOCK (deterministic).
* Per-sample form (graph-safe, no atomicAdd, no memset):
* per_sample_out[i] = lambda_tc * weight_i * (1/total_actions) *
* sum_a |Q[i,a] - Q[i+1,a]| for i in [0, N-2]
* per_sample_out[N-1] = 0
* Caller reduces with c51_loss_reduce_kernel for the scalar penalty.
*
* Grid: ceil((N-1) / 256), Block: 256.
* Grid: ceil(N / 256), Block: 256. One thread per sample.
*/
extern "C" __global__ void temporal_consistency_penalty(
const float* __restrict__ h_s2,
const float* __restrict__ q_values,
float* __restrict__ penalty_out,
float* __restrict__ per_sample_out, /* [N] — graph-safe, no atomic */
int N,
int sh2,
int total_actions,
@@ -4975,49 +4979,45 @@ extern "C" __global__ void temporal_consistency_penalty(
float sim_threshold
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
float my_penalty = 0.0f;
if (i >= N) return;
if (i < N - 1) {
float dot = 0.0f, norm_a = 0.0f, norm_b = 0.0f;
for (int j = 0; j < sh2; j++) {
float a = h_s2[(long long)i * sh2 + j];
float b = h_s2[(long long)(i + 1) * sh2 + j];
dot += a * b;
norm_a += a * a;
norm_b += b * b;
}
float sim = dot / (sqrtf(norm_a) * sqrtf(norm_b) + 1e-8f);
if (sim > sim_threshold) {
float q_diff = 0.0f;
for (int a = 0; a < total_actions; a++) {
q_diff += fabsf(q_values[(long long)i * total_actions + a]
- q_values[(long long)(i + 1) * total_actions + a]);
}
q_diff /= (float)total_actions;
float weight = (sim - sim_threshold) / (1.0f - sim_threshold);
my_penalty = lambda_tc * q_diff * weight / (float)(N - 1);
}
if (i >= N - 1) {
per_sample_out[i] = 0.0f; /* boundary: no "i+1" available */
return;
}
/* Warp-level reduction */
for (int offset = 16; offset > 0; offset >>= 1)
my_penalty += __shfl_xor_sync(0xFFFFFFFF, my_penalty, offset);
/* Block-level reduction via shared memory */
__shared__ float warp_sums[8];
int warp_id = threadIdx.x / 32;
int lane = threadIdx.x % 32;
if (lane == 0) warp_sums[warp_id] = my_penalty;
__syncthreads();
if (warp_id == 0) {
float val = (lane < blockDim.x / 32) ? warp_sums[lane] : 0.0f;
for (int off = 16; off > 0; off >>= 1)
val += __shfl_xor_sync(0xFFFFFFFF, val, off);
if (lane == 0)
atomicAdd(penalty_out, val);
/* Cosine similarity between h_s2[i] and h_s2[i+1] */
float dot = 0.0f, norm_a = 0.0f, norm_b = 0.0f;
for (int j = 0; j < sh2; j++) {
float a = h_s2[(long long)i * sh2 + j];
float b = h_s2[(long long)(i + 1) * sh2 + j];
dot += a * b;
norm_a += a * a;
norm_b += b * b;
}
float sim = dot / (sqrtf(norm_a) * sqrtf(norm_b) + 1e-8f);
/* Triangular weight: 0 at threshold, 1 at sim=1.0 */
float weight = (sim - sim_threshold) / fmaxf(1.0f - sim_threshold, 1e-6f);
weight = fmaxf(0.0f, weight);
if (weight <= 0.0f) {
per_sample_out[i] = 0.0f;
return;
}
/* Mean-abs Q-difference across actions */
float q_diff = 0.0f;
for (int a = 0; a < total_actions; a++) {
q_diff += fabsf(q_values[(long long)i * total_actions + a]
- q_values[(long long)(i + 1) * total_actions + a]);
}
q_diff /= (float)total_actions;
/* Per-sample contribution: divide by (N-1) so the eventual sum reduces
* to the average penalty per consecutive-pair (matches old kernel's
* normalization). */
per_sample_out[i] = lambda_tc * weight * q_diff / (float)(N - 1);
}
/* ================================================================== */

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

@@ -913,9 +913,13 @@ pub struct GpuDqnTrainer {
branch_indep_penalty_buf: CudaSlice<f32>,
// ── G10: Temporal consistency penalty ──
/// Kernel: temporal_consistency_penalty — Lipschitz penalty on Q-diffs for similar states.
/// Kernel: temporal_consistency_penalty — per-sample Lipschitz penalty
/// on Q-diffs for similar states. Graph-safe (no atomicAdd, no memset);
/// caller reduces with c51_loss_reduce_kernel.
temporal_consistency_kernel: CudaFunction,
/// Scalar output buffer [1] for temporal consistency penalty.
/// Per-sample loss buffer [B] — reduced via c51_loss_reduce.
temporal_per_sample_buf: CudaSlice<f32>,
/// Scalar output buffer [1] for temporal consistency penalty (post-reduce).
temporal_penalty_buf: CudaSlice<f32>,
// ── G12: Predictive coding auxiliary loss + backward ──
@@ -3305,24 +3309,18 @@ impl GpuDqnTrainer {
unsafe { *self.var_ema_pinned = val; }
}
/// G6: Compute branch independence penalty (6-way cosine similarity).
/// Informational — writes scalar to branch_indep_penalty_buf for logging.
/// Gradient integration deferred (requires backward pass modification).
/// G6: Branch independence penalty (forward-only, measurement).
/// Single-block kernel writes the scalar penalty directly into
/// branch_indep_penalty_buf[0] (overwrite, not accumulate) — fully
/// graph-safe, no atomicAdd, no memset needed before launch.
/// Backward integration TODO (gated on measurement showing the value
/// is meaningfully larger than what spectral_norm/NoisyNets achieve).
pub(crate) fn compute_branch_independence(&self, batch_size: usize) -> Result<(), MLError> {
let branch_indep_penalty_buf_ptr = self.branch_indep_penalty_buf.raw_ptr();
let save_h_b0_ptr = self.save_h_b0.raw_ptr();
let save_h_b1_ptr = self.save_h_b1.raw_ptr();
let save_h_b2_ptr = self.save_h_b2.raw_ptr();
let save_h_b3_ptr = self.save_h_b3.raw_ptr();
// Use raw memset to avoid &mut self borrow conflict on the penalty buf
unsafe {
cudarc::driver::sys::cuMemsetD32Async(
branch_indep_penalty_buf_ptr,
0,
1,
self.stream.cu_stream(),
);
}
let ah = self.config.adv_h as i32;
unsafe {
self.stream.launch_builder(&self.branch_indep_kernel)
@@ -3344,36 +3342,47 @@ impl GpuDqnTrainer {
Ok(())
}
/// G10: Compute temporal consistency penalty (Lipschitz on Q-diffs for similar states).
/// Informational — writes scalar to temporal_penalty_buf for logging.
/// G10: Temporal consistency penalty (forward-only, measurement).
/// Per-sample kernel writes [B] losses, then c51_loss_reduce → scalar.
/// Fully graph-safe (no atomicAdd, no memset). Backward integration
/// TODO (gated on measurement showing the value is meaningfully large
/// vs the global Lipschitz already enforced by spectral_norm).
pub(crate) fn compute_temporal_consistency(&self, batch_size: usize) -> Result<(), MLError> {
let temporal_penalty_buf_ptr = self.temporal_penalty_buf.raw_ptr();
let save_h_s2_ptr = self.save_h_s2.raw_ptr();
let q_out_buf_ptr = self.q_out_buf.raw_ptr();
// Use raw memset to avoid &mut self borrow conflict on the penalty buf
unsafe {
cudarc::driver::sys::cuMemsetD32Async(
temporal_penalty_buf_ptr,
0,
1,
self.stream.cu_stream(),
);
}
let per_sample_buf_ptr = self.temporal_per_sample_buf.raw_ptr();
let temporal_penalty_buf_ptr = self.temporal_penalty_buf.raw_ptr();
let ta = self.total_actions() as i32;
let sh2 = self.config.shared_h2 as i32;
// Step 1: per-sample penalty → temporal_per_sample_buf [B]
unsafe {
self.stream.launch_builder(&self.temporal_consistency_kernel)
.arg(&save_h_s2_ptr)
.arg(&q_out_buf_ptr)
.arg(&temporal_penalty_buf_ptr)
.arg(&per_sample_buf_ptr)
.arg(&(batch_size as i32))
.arg(&sh2)
.arg(&ta)
.arg(&0.005_f32) // lambda_tc
.arg(&0.95_f32) // sim_threshold
.launch(LaunchConfig::for_num_elems((batch_size - 1) as u32))
.launch(LaunchConfig::for_num_elems(batch_size as u32))
.map_err(|e| MLError::ModelError(format!("temporal_consistency: {e}")))?;
}
// Step 2: deterministic reduce → temporal_penalty_buf[0] for monitoring
unsafe {
self.stream
.launch_builder(&self.c51_loss_reduce_kernel)
.arg(&per_sample_buf_ptr)
.arg(&temporal_penalty_buf_ptr)
.arg(&(batch_size as i32))
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("temporal_consistency reduce: {e}")))?;
}
Ok(())
}
@@ -5840,7 +5849,11 @@ impl GpuDqnTrainer {
let branch_indep_penalty_buf = stream.alloc_zeros::<f32>(1)
.map_err(|e| MLError::ModelError(format!("alloc branch_indep_penalty_buf: {e}")))?;
// ── G10: Temporal consistency penalty buffer [1] ──
// ── G10: Temporal consistency penalty buffers ──
// Per-sample loss [B] (graph-safe, no atomicAdd, no memset) →
// reduced via c51_loss_reduce → scalar [1] for monitoring readback.
let temporal_per_sample_buf = stream.alloc_zeros::<f32>(config.batch_size)
.map_err(|e| MLError::ModelError(format!("alloc temporal_per_sample_buf: {e}")))?;
let temporal_penalty_buf = stream.alloc_zeros::<f32>(1)
.map_err(|e| MLError::ModelError(format!("alloc temporal_penalty_buf: {e}")))?;
@@ -7372,6 +7385,7 @@ impl GpuDqnTrainer {
branch_indep_kernel,
branch_indep_penalty_buf,
temporal_consistency_kernel,
temporal_per_sample_buf,
temporal_penalty_buf,
predictive_coding_kernel,
predictive_coding_backward_kernel,
@@ -8271,6 +8285,52 @@ impl GpuDqnTrainer {
/// Read eval_q_std_ema.
pub fn eval_q_std_ema(&self) -> f32 { self.eval_q_std_ema }
/// G6: Read branch_independence penalty scalar (sync DtoH; epoch-boundary only).
/// Returns 0.0 if buffer is uninitialized or sync fails — caller treats
/// as "no signal". Use HEALTH_DIAG to log.
pub fn branch_indep_loss_value(&self) -> f32 {
// SAFETY: sync DtoH at epoch boundary; not in graph capture region.
unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }
let mut v = 0.0_f32;
unsafe {
cudarc::driver::sys::cuMemcpyDtoH_v2(
std::ptr::addr_of_mut!(v).cast(),
self.branch_indep_penalty_buf.raw_ptr(),
std::mem::size_of::<f32>(),
);
}
v
}
/// G10: Read temporal_consistency penalty scalar (sync DtoH; epoch-boundary only).
pub fn temporal_loss_value(&self) -> f32 {
unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }
let mut v = 0.0_f32;
unsafe {
cudarc::driver::sys::cuMemcpyDtoH_v2(
std::ptr::addr_of_mut!(v).cast(),
self.temporal_penalty_buf.raw_ptr(),
std::mem::size_of::<f32>(),
);
}
v
}
/// G12: Read predictive_coding loss scalar (sync DtoH; epoch-boundary only).
/// Same pattern as G6/G10. Used for HEALTH_DIAG monitoring.
pub fn predictive_loss_value(&self) -> f32 {
unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }
let mut v = 0.0_f32;
unsafe {
cudarc::driver::sys::cuMemcpyDtoH_v2(
std::ptr::addr_of_mut!(v).cast(),
self.predictive_loss_buf.raw_ptr(),
std::mem::size_of::<f32>(),
);
}
v
}
/// Read per_branch_q_gap_ema — synchronous DtoH for trajectory backtracking.
pub fn per_branch_q_gap_ema(&self) -> [f32; 4] {
unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }

25
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -868,6 +868,19 @@ impl FusedTrainingCtx {
}
}
/// G6/G10/G12 gem loss readbacks (sync DtoH; epoch-boundary only).
/// Returns (branch_indep_loss, temporal_consistency_loss, predictive_coding_loss).
/// Used by HEALTH_DIAG to surface gem-penalty magnitudes for the V7
/// measurement-first methodology — values informs whether the existing
/// mechanisms (spectral_norm, NoisyNets) already cover them.
pub(crate) fn read_gem_losses(&self) -> (f32, f32, f32) {
(
self.trainer.branch_indep_loss_value(),
self.trainer.temporal_loss_value(),
self.trainer.predictive_loss_value(),
)
}
/// Flush the last in-flight async readback from the training step.
/// Called once at epoch end to retrieve the final step's scalars.
pub(crate) fn flush_readback(&mut self) -> Result<crate::cuda_pipeline::gpu_dqn_trainer::FusedTrainScalars> {
@@ -1466,6 +1479,18 @@ impl FusedTrainingCtx {
self.trainer.compute_predictive_coding_loss(self.batch_size)
.map_err(|e| anyhow::anyhow!("Predictive coding (G12): {e}"))?;
// G6 + G10 forward-only (measurement). Both kernels are now graph-safe
// (G6: single-block direct overwrite; G10: per-sample → c51_loss_reduce).
// No backward yet — wired only to log the penalty magnitudes via
// HEALTH_DIAG. Decision on backward integration depends on whether
// the measured penalty is meaningfully larger than what existing
// mechanisms (spectral_norm for G10, NoisyNets/ensemble-KL for G6)
// already enforce. V7-gem methodology: measure before commit.
self.trainer.compute_branch_independence(self.batch_size)
.map_err(|e| anyhow::anyhow!("Branch independence (G6 measure): {e}"))?;
self.trainer.compute_temporal_consistency(self.batch_size)
.map_err(|e| anyhow::anyhow!("Temporal consistency (G10 measure): {e}"))?;
// Ensemble diversity: heads 1..K-1 forward + KL gradient + trunk backward.
// Runs inside aux_child graph — gradients land in grad_buf before Adam.
if !self.ensemble_extra_heads.is_empty() {

13
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -2148,8 +2148,16 @@ impl DQNTrainer {
// diag = pure monitoring signals (not in reward): sharpe EMA (formerly w_dsr
// reward input) and action entropy normalized to [0,1] (formerly
// position_entropy_weight reward). See Phase 1 inventory better-form taxonomy.
// gems = G6/G10/G12 gem-penalty values (sync DtoH from GPU; once per epoch).
// V7 methodology: log first to measure whether the gem is doing real work
// before deciding to wire its backward gradient.
let (g6_loss, g10_loss, g12_loss) = if let Some(ref fused) = self.fused_ctx {
fused.read_gem_losses()
} else {
(0.0, 0.0, 0.0)
};
tracing::info!(
"HEALTH_DIAG[{}]: health={:.2} components [q_gap={:.2} q_var={:.2} atoms={:.2} grad_stable={:.2} ens_agree={:.2} grad_cos={:.2} spectral={:.2}] effective [cql_alpha={:.4} iqn_budget={:.2} cql_budget={:.2} c51_budget={:.2} tau={:.5} sarsa_tau={:.2} gamma={:.3} cf_ratio={:.2}] novels [distill={} barrier={:.3} plasticity={} ib={:.3} ensemble_collapse={:.2} contrarian={} meta_q_pred={:.2}] diag [sharpe_ema={:.3} action_entropy={:.2}]",
"HEALTH_DIAG[{}]: health={:.2} components [q_gap={:.2} q_var={:.2} atoms={:.2} grad_stable={:.2} ens_agree={:.2} grad_cos={:.2} spectral={:.2}] effective [cql_alpha={:.4} iqn_budget={:.2} cql_budget={:.2} c51_budget={:.2} tau={:.5} sarsa_tau={:.2} gamma={:.3} cf_ratio={:.2}] novels [distill={} barrier={:.3} plasticity={} ib={:.3} ensemble_collapse={:.2} contrarian={} meta_q_pred={:.2}] diag [sharpe_ema={:.3} action_entropy={:.2}] gems [g6_branch_indep={:.6} g10_temporal={:.6} g12_predictive={:.6}]",
epoch,
health_value,
self.learning_health.components.q_gap_norm,
@@ -2176,6 +2184,9 @@ impl DQNTrainer {
self.last_meta_q_pred.unwrap_or(0.0),
self.training_sharpe_ema,
self.last_action_entropy.unwrap_or(1.0),
g6_loss,
g10_loss,
g12_loss,
);
// C1/P1: propagate health to GPU replay buffer for diversity-weighted priorities.