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

fix: all precommit LOWs — pinned sel_t, per-branch Q-gap detection

Browse Source 
LOW 1: sel_t_buf CudaSlice + per-step memcpy_htod → sel_t_pinned
device-mapped. Zero copies, GPU reads directly from host memory.

LOW 2: is_q_gap_frozen uniform [q_gap; 4] → actual per-branch Q-gaps
from 48B DtoH readback in reduce_current_q_stats. Each branch now
independently detected as frozen/unfrozen. Liquid tau and trajectory
backtracking use real per-branch velocities.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-15 07:47:40 +02:00
parent 77fcdb2657
commit 3f4f0c7038
3 changed files with 64 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -642,11 +642,14 @@ pub struct GpuDqnTrainer {
sel_norm_partials: CudaSlice<f32>, // [1]
/// [1] max grad norm for sel Adam clipping (fixed 1.0).
sel_clip_buf: CudaSlice<f32>, // [1]
/// [1] sel Adam step counter on device (for CUDA kernel compatibility).
sel_t_buf: CudaSlice<i32>, // [1]
/// Selectivity Adam step counter — pinned device-mapped (no per-step HtoD copy).
sel_t_pinned: *mut i32,
sel_t_dev_ptr: u64,
// ── Liquid tau ──
per_branch_q_gap_ema: [f32; 4],
/// Last computed per-branch Q-gaps from reduce_current_q_stats.
last_per_branch_q_gaps: [f32; 4],
liquid_mod_pinned: *mut f32, // 4 floats pinned device-mapped
liquid_mod_dev_ptr: u64,
@@ -1069,6 +1072,9 @@ impl Drop for GpuDqnTrainer {
if !self.liquid_mod_pinned.is_null() {
let _ = unsafe { cudarc::driver::result::free_host(self.liquid_mod_pinned.cast()) };
}
if !self.sel_t_pinned.is_null() {
let _ = unsafe { cudarc::driver::result::free_host(self.sel_t_pinned.cast()) };
}
if !self.t_pinned.is_null() {
let _ = unsafe { cudarc::driver::result::free_host(self.t_pinned.cast()) };
}
@@ -2728,8 +2734,23 @@ impl GpuDqnTrainer {
let mut sel_clip_buf = stream.alloc_zeros::<f32>(1)
.map_err(|e| MLError::ModelError(format!("alloc sel_clip_buf: {e}")))?;
super::htod_f32(&stream, &[1.0_f32], &mut sel_clip_buf)?;
// sel_t_buf: Adam step counter for selectivity gate, initialized to 0.
let sel_t_buf = alloc_i32(&stream, 1, "sel_t_buf")?;
// Selectivity Adam step counter — pinned device-mapped (no per-step HtoD copy)
let sel_t_pinned: *mut i32 = unsafe {
let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP;
cudarc::driver::result::malloc_host(std::mem::size_of::<i32>(), flags)
.map_err(|e| MLError::ModelError(format!("pinned sel_t alloc: {e}")))?
as *mut i32
};
unsafe { *sel_t_pinned = 0; }
let sel_t_dev_ptr = unsafe {
let mut dev_ptr: u64 = 0;
cudarc::driver::sys::cuMemHostGetDevicePointer_v2(
&mut dev_ptr as *mut u64,
sel_t_pinned.cast(),
0,
);
dev_ptr
};
// ── Liquid tau: 4 pinned device-mapped floats ─────────────────
let liquid_mod_pinned: *mut f32 = unsafe {
@@ -3312,8 +3333,10 @@ impl GpuDqnTrainer {
sel_norm_buf,
sel_norm_partials,
sel_clip_buf,
sel_t_buf,
sel_t_pinned,
sel_t_dev_ptr,
per_branch_q_gap_ema: [0.0; 4],
last_per_branch_q_gaps: [0.0; 4],
liquid_mod_pinned,
liquid_mod_dev_ptr,
vsn_masked_buf,
@@ -4756,8 +4779,7 @@ impl GpuDqnTrainer {
// Synchronous readback — Q-stats drive eval_v_range which drives
// action selection. Stale values cause 2-state oscillation where
// the system flips between cached results from different steps.
// Cost: ~5µs sync + 28B DtoH. Acceptable since this runs every
// 50 training steps, not per-step.
// Cost: ~5µs sync + 28B + 48B DtoH. Runs every 50 training steps.
unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }
let mut host = [0.0_f32; 7];
unsafe {
@@ -4768,6 +4790,30 @@ impl GpuDqnTrainer {
);
}
// Per-branch Q-gap: read first sample's 12 Q-values from q_out_buf.
// Stream already synced. Branch layout: [dir(3), mag(3), ord(3), urg(3)].
let mut q12 = [0.0_f32; 12];
unsafe {
cudarc::driver::sys::cuMemcpyDtoH_v2(
q12.as_mut_ptr().cast(),
self.q_out_buf.raw_ptr(),
12 * std::mem::size_of::<f32>(),
);
}
let branch_sizes = [
self.config.branch_0_size, self.config.branch_1_size,
self.config.branch_2_size, self.config.branch_3_size,
];
let mut offset = 0;
for d in 0..4 {
let bs = branch_sizes[d];
let slice = &q12[offset..offset + bs];
let max_q = slice.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let mean_q = slice.iter().sum::<f32>() / bs as f32;
self.last_per_branch_q_gaps[d] = (max_q - mean_q).max(0.0);
offset += bs;
}
Ok(QValueStatsResult {
avg_max_q: host[0] as f64,
q_min: host[1],
@@ -4779,6 +4825,9 @@ impl GpuDqnTrainer {
})
}
/// Last computed per-branch Q-gaps from reduce_current_q_stats.
pub fn get_per_branch_q_gaps(&self) -> [f32; 4] { self.last_per_branch_q_gaps }
/// Flush the last in-flight Q-stats readback at epoch end — non-blocking.
///
/// Reads directly from the pinned host buffer without synchronizing.
@@ -6802,16 +6851,15 @@ impl GpuDqnTrainer {
/// Run Adam update on selectivity gate parameters.
///
/// Computes L2 grad norm from sel_grad, then runs dqn_adam_update_kernel at LR=1e-4.
/// Increments sel_adam_step and writes the updated step counter to sel_t_buf.
/// Increments sel_adam_step and writes to pinned device-mapped sel_t_pinned.
pub(crate) fn step_selectivity_adam(&mut self) -> Result<(), MLError> {
self.sel_adam_step += 1;
let step_val = self.sel_adam_step;
let sel_dim = self.config.shared_h2 + 1;
let sel_n = sel_dim as i32;
// Write step counter to device buffer (memcpy_htod — small 4-byte copy, not on hot path).
self.stream.memcpy_htod(&[step_val], &mut self.sel_t_buf)
.map_err(|e| MLError::ModelError(format!("sel_t_buf htod: {e}")))?;
// Write step counter to pinned device-mapped memory (no HtoD copy needed)
unsafe { *self.sel_t_pinned = step_val; }
// Phase 1: grad_norm_kernel on sel_grad → sel_norm_partials [1 block].
let grad_ptr = self.sel_grad.raw_ptr();
@@ -6858,7 +6906,7 @@ impl GpuDqnTrainer {
let m_ptr = self.sel_adam_m.raw_ptr();
let v_ptr = self.sel_adam_v.raw_ptr();
let clip_ptr = self.sel_clip_buf.raw_ptr();
let t_ptr = self.sel_t_buf.raw_ptr();
let t_ptr = self.sel_t_dev_ptr;
let blocks = ((sel_dim as u32 + 255) / 256).max(1);
unsafe {
self.stream

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

@@ -2153,6 +2153,7 @@ impl FusedTrainingCtx {
/// Read per_branch_q_gap_ema.
pub(crate) fn per_branch_q_gap_ema(&self) -> [f32; 4] { self.trainer.per_branch_q_gap_ema() }
pub(crate) fn get_per_branch_q_gaps(&self) -> [f32; 4] { self.trainer.get_per_branch_q_gaps() }
/// Set eval_q_mean_ema (trajectory backtracking restore).
pub(crate) fn set_eval_q_mean_ema(&mut self, v: f32) { self.trainer.set_eval_q_mean_ema(v); }

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

@@ -1404,9 +1404,8 @@ impl DQNTrainer {
// Update eval v_range from observed Q-stats + Q-gap
let q_std = stats.q_variance.max(0.0).sqrt();
let q_gap = (stats.avg_max_q as f32 - stats.q_mean).max(0.0);
// Simple approximation: use global q_gap for all branches initially.
// Per-branch extraction requires reading 12 Q-values from GPU.
let per_branch_q_gaps = [q_gap; 4]; // uniform until per-branch extraction is added
// Per-branch Q-gaps computed in reduce_current_q_stats (48B DtoH readback).
let per_branch_q_gaps = fused.get_per_branch_q_gaps();
fused.update_eval_v_range(stats.q_mean, q_std, q_gap, per_branch_q_gaps);
}
}
@@ -2304,9 +2303,7 @@ impl DQNTrainer {
fn is_q_gap_frozen(&self) -> bool {
if let Some(ref fused) = self.fused_ctx {
let emas = fused.per_branch_q_gap_ema();
// Current q_gap approximated as uniform across branches (same as update_eval_v_range)
let q_gap = self.epoch_q_gap;
let current = [q_gap; 4];
let current = fused.get_per_branch_q_gaps();
let threshold = 0.001_f32;
current.iter().zip(emas.iter()).all(|(gap, ema)| (gap - ema).abs() < threshold)
} else {