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

feat(c51): per-sample support + branch scales in kernel launches, delete v_range

Browse Source 
Replace global v_range infrastructure with per-sample C51 support and
per-branch gradient scales. The IQL trainer sets these pointers after
construction, enabling sample-level distributional RL customization.

- Add per_sample_support_ptr and branch_scales_ptr fields + setters
- Update all kernel launches (C51 loss, C51 grad, CQL, MSE, expected_q)
  to pass per_sample_support_ptr instead of v_range_dev_ptr
- Add branch_scales_ptr arg to c51_grad_kernel launch
- Delete adapt_v_range_full, adapt_v_range, adapt_v_range_with_gap,
  v_range getter, v_range_buf_ptr getter
- Delete v_range_pinned/v_range_dev_ptr/reward_std_ema fields + alloc/free

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-13 15:18:12 +02:00
parent 5fbb39b0b5
commit b2e695fe08
1 changed files with 69 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -25,7 +25,7 @@
//! (IQN trunk, attention, ensemble) into `grad_buf` via SAXPY. The Adam graph
//! then sees the combined C51 + auxiliary gradients.
//!
//! Per-step scalars (adam_step, tau, v_range, adaptive_clip) use pinned
//! Per-step scalars (adam_step, tau, adaptive_clip) use pinned
//! device-mapped memory — GPU reads via pointer, host writes directly,
//! no HtoD copies. Only batch input data needs explicit upload before replay.
//!
@@ -626,9 +626,10 @@ pub struct GpuDqnTrainer {
tau_pinned: *mut f32,
tau_dev_ptr: u64,
/// Adaptive C51 z-support [v_min, v_max] — pinned device-mapped, no HtoD copy.
v_range_pinned: *mut f32,
v_range_dev_ptr: u64,
/// Per-sample C51 support [B, 3] — pointer set by IQL trainer.
per_sample_support_ptr: u64,
/// Per-branch gradient scales [B, 4] — pointer set by IQL trainer.
branch_scales_ptr: u64,
/// Adaptive gradient clip norm — pinned device-mapped host memory.
/// GPU reads via device pointer (graph-safe), host writes directly (no HtoD copy).
@@ -639,8 +640,6 @@ pub struct GpuDqnTrainer {
grad_norm_ema: f32,
/// EMA of Q-divergence for adaptive tau computation.
q_div_ema: f32,
/// EMA of observed reward std for adaptive v_range floor.
reward_std_ema: f32,
// ── Training state ──────────────────────────────────────────────
pub(crate) adam_step: i32,
@@ -963,9 +962,6 @@ impl GpuDqnTrainer {
self.stream.memset_zeros(&mut self.popart_count)
.map_err(|e| MLError::ModelError(format!("reset popart_count: {e}")))?;
// Reset adaptive C51 z-support — fold 2's Q-scale may differ from fold 1.
unsafe { *self.v_range_pinned = -0.005; *self.v_range_pinned.add(1) = 0.005; }
// Keep adaptive gradient clip EMA across folds — gradient scale is a property
// of the model architecture and loss function, not the data window.
// Resetting it would leave the first fold-2 epochs unprotected.
@@ -975,7 +971,7 @@ impl GpuDqnTrainer {
self.stream.memset_zeros(&mut self.q_divergence_buf)
.map_err(|e| MLError::ModelError(format!("reset q_divergence: {e}")))?;
tracing::info!("Adam optimizer + PopArt + v_range + grad_clip state reset for new fold");
tracing::info!("Adam optimizer + PopArt + grad_clip state reset for new fold");
Ok(())
}
}
@@ -1000,9 +996,6 @@ impl Drop for GpuDqnTrainer {
if !self.tau_pinned.is_null() {
let _ = unsafe { cudarc::driver::result::free_host(self.tau_pinned.cast()) };
}
if !self.v_range_pinned.is_null() {
let _ = unsafe { cudarc::driver::result::free_host(self.v_range_pinned.cast()) };
}
}
}
@@ -1012,6 +1005,13 @@ impl GpuDqnTrainer {
self.config.batch_size
}
pub fn set_per_sample_support_ptr(&mut self, ptr: u64) {
self.per_sample_support_ptr = ptr;
}
pub fn set_branch_scales_ptr(&mut self, ptr: u64) {
self.branch_scales_ptr = ptr;
}
/// Raw pointer to the pinned host readback buffer [16 × f32].
/// Used by `FusedTrainingCtx` for async diversity loss readback at offset 9.
pub fn readback_pinned_ptr(&self) -> *mut f32 {
@@ -1727,7 +1727,7 @@ impl GpuDqnTrainer {
.arg(&b1_i32)
.arg(&b2_i32)
.arg(&b3_i32)
.arg(&self.v_range_dev_ptr)
.arg(&self.per_sample_support_ptr)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
@@ -2454,24 +2454,6 @@ impl GpuDqnTrainer {
dev_ptr
};
// Adaptive C51 z-support — pinned device-mapped, no HtoD copy.
let v_range_pinned: *mut f32 = unsafe {
let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP;
cudarc::driver::result::malloc_host(2 * std::mem::size_of::<f32>(), flags)
.map_err(|e| MLError::ModelError(format!("pinned v_range alloc: {e}")))?
as *mut f32
};
unsafe { *v_range_pinned = -0.005_f32; *v_range_pinned.add(1) = 0.005_f32; }
let v_range_dev_ptr = unsafe {
let mut dev_ptr: u64 = 0;
cudarc::driver::sys::cuMemHostGetDevicePointer_v2(
&mut dev_ptr as *mut u64,
v_range_pinned.cast(),
0,
);
dev_ptr
};
// Adaptive gradient clip norm — pinned device-mapped host memory.
// GPU reads via device pointer (zero-copy, no HtoD), host writes directly.
let adaptive_clip_init = config.max_grad_norm;
@@ -2982,7 +2964,7 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("alloc popart_count: {e}")))?;
// v8: Pessimistic Q-value initialization — shift value head bias to -0.1
// Pessimistic Q-init REMOVED — incompatible with adaptive v_range.
// Pessimistic Q-init REMOVED — incompatible with per-sample support.
// Xavier init gives near-zero value head output, correct for adaptive C51.
Ok(Self {
@@ -3066,13 +3048,12 @@ impl GpuDqnTrainer {
t_dev_ptr,
tau_pinned,
tau_dev_ptr,
v_range_pinned,
v_range_dev_ptr,
per_sample_support_ptr: 0,
branch_scales_ptr: 0,
adaptive_clip_pinned,
adaptive_clip_dev_ptr,
grad_norm_ema: 0.0,
q_div_ema: 0.0,
reward_std_ema: 0.0,
adam_step: 0,
total_params,
params_initialized: false,
@@ -4131,13 +4112,55 @@ impl GpuDqnTrainer {
self.config.branch_0_size + self.config.branch_1_size + self.config.branch_2_size + self.config.branch_3_size
}
/// Reference to the Q-value output buffer from the last `forward_only_q()` call.
/// Reference to the Q-value output buffer.
///
/// Shape: `[B, total_actions]`. Only valid after `forward_only_q()`.
/// Shape: `[B, total_actions]`. Valid after `populate_q_out()` or `replay_forward_for_q_values()`.
pub fn q_out_buf(&self) -> &CudaSlice<f32> {
&self.q_out_buf
}
/// Convert current logits → Q-values in q_out_buf.
///
/// Runs compute_expected_q on the logits already in on_v_logits_buf / on_b_logits_buf
/// (populated by graph_forward replay). Does NOT replay any graph — just the
/// lightweight expected_q kernel.
///
/// Used by IQL to get Q(s, a_taken) for expectile regression target.
pub fn populate_q_out(&self, batch_size: usize) -> Result<&CudaSlice<f32>, MLError> {
let n = batch_size as i32;
let na = self.config.num_atoms 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 block_dim = 256_u32;
let grid_dim = ((batch_size as u32 + block_dim - 1) / block_dim).max(1);
let q_out_ptr = self.q_out_buf.raw_ptr();
let null_atom_stats = 0u64;
unsafe {
self.stream
.launch_builder(&self.expected_q_kernel)
.arg(&self.ptrs.on_v_logits_buf)
.arg(&self.ptrs.on_b_logits_buf)
.arg(&q_out_ptr)
.arg(&n)
.arg(&na)
.arg(&b0)
.arg(&b1)
.arg(&b2)
.arg(&b3)
.arg(&self.per_sample_support_ptr)
.arg(&null_atom_stats)
.launch(LaunchConfig {
grid_dim: (grid_dim, 1, 1),
block_dim: (block_dim, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("populate_q_out expected_q: {e}")))?;
}
Ok(&self.q_out_buf)
}
/// Run the cuBLAS forward pass without graph (used during graph capture).
fn replay_forward_ungraphed(&self) -> Result<(), MLError> {
let param_sizes = compute_param_sizes(&self.config);
@@ -4233,7 +4256,7 @@ impl GpuDqnTrainer {
.arg(&b1)
.arg(&b2)
.arg(&b3)
.arg(&self.v_range_dev_ptr)
.arg(&self.per_sample_support_ptr)
.arg(&null_atom_stats)
.launch(LaunchConfig {
grid_dim: (grid_dim, 1, 1),
@@ -4385,7 +4408,7 @@ impl GpuDqnTrainer {
.arg(&b1)
.arg(&b2)
.arg(&b3)
.arg(&self.v_range_dev_ptr)
.arg(&self.per_sample_support_ptr)
.arg(&atom_stats_ptr)
.launch(LaunchConfig {
grid_dim: (eq_blocks, 1, 1),
@@ -5334,11 +5357,11 @@ impl GpuDqnTrainer {
// ── Curiosity Q-penalty (2) ──
.arg(&self.curiosity_error_buf)
.arg(&self.config.curiosity_q_penalty_lambda)
// ── Config (8 — v_range_buf replaces v_min+v_max) ──
// ── Config (8 — per_sample_support replaces v_min+v_max) ──
.arg(&gamma)
.arg(&batch_i32)
.arg(&na_i32)
.arg(&self.v_range_dev_ptr)
.arg(&self.per_sample_support_ptr)
.arg(&b0_i32)
.arg(&b1_i32)
.arg(&b2_i32)
@@ -5478,6 +5501,7 @@ impl GpuDqnTrainer {
.arg(&b3_i32)
.arg(&total_branch_atoms_i32)
.arg(&entropy_coeff)
.arg(&self.branch_scales_ptr)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
@@ -5577,11 +5601,11 @@ impl GpuDqnTrainer {
// ── Curiosity Q-penalty (2) ──
.arg(&self.curiosity_error_buf)
.arg(&self.config.curiosity_q_penalty_lambda)
// ── Config (8 — v_range_buf replaces v_min+v_max) ──
// ── Config (8 — per_sample_support replaces v_min+v_max) ──
.arg(&gamma)
.arg(&batch_i32)
.arg(&na_i32)
.arg(&self.v_range_dev_ptr)
.arg(&self.per_sample_support_ptr)
.arg(&b0_i32)
.arg(&b1_i32)
.arg(&b2_i32)
@@ -5666,7 +5690,7 @@ impl GpuDqnTrainer {
.arg(&b2_i32)
.arg(&b3_i32)
.arg(&total_branch_atoms_i32)
.arg(&self.v_range_dev_ptr)
.arg(&self.per_sample_support_ptr)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
@@ -6359,92 +6383,6 @@ impl GpuDqnTrainer {
Ok(())
}
// ── Adaptive C51 z-support ─────────────────────────────────────
/// Q-stats-driven v_range: center at q_mean, width = 3σ + Bellman headroom.
///
/// Replaces the old heuristic expand/contract with a single data-driven formula.
/// - 3σ coverage captures 99.7% of Q-distribution
/// - Bellman headroom (γ * 3σ) prevents projection clipping on reward shift
/// - MIN_RANGE floor prevents atom collapse when Q-variance is near-zero
/// - Mean-centered: no wasted atoms on impossible Q-values
pub fn adapt_v_range(&mut self, q_mean: f32, q_variance: f32) -> bool {
self.adapt_v_range_full(q_mean, q_variance, 0.0, 0.0)
}
pub fn adapt_v_range_with_gap(&mut self, q_mean: f32, q_variance: f32, q_gap: f32) -> bool {
self.adapt_v_range_full(q_mean, q_variance, q_gap, 0.0)
}
/// Fully adaptive v_range from Q-stats + reward scale.
///
/// The reward_std parameter comes from the experience collector's observed
/// reward statistics. It varies 940× between smoketest (0.007) and production
/// (6.57) — hardcoding it was fundamentally wrong.
///
/// The v_range floor = reward_std (single Bellman step shift magnitude).
/// This ensures the C51 projection can shift atoms by at least one reward
/// unit, breaking the stable fixed point at any scale.
pub fn adapt_v_range_full(&mut self, q_mean: f32, q_variance: f32, q_gap: f32, reward_std: f32) -> bool {
const SIGMA_COVERAGE: f32 = 3.0;
const GAP_ATOM_TARGET: f32 = 10.0;
let q_std = (q_variance.max(0.0) + 1e-10).sqrt();
let half_width = SIGMA_COVERAGE * q_std;
let bellman_headroom = self.config.gamma * half_width;
let total_half = half_width + bellman_headroom;
let target_min = q_mean - total_half;
let target_max = q_mean + total_half;
// Reward-scale floor: derived from ACTUAL observed reward_std.
// Smoketest: reward_std ≈ 0.007 → floor = 0.007
// Production: reward_std ≈ 6.57 → floor = 6.57
// Falls back to EMA of observed reward_std (self.reward_std_ema)
// when per-step reward_std is not available.
let effective_reward_std = if reward_std > 1e-8 {
// Update EMA with fresh observation
const R_BETA: f32 = 0.99;
if self.reward_std_ema <= 1e-8 {
self.reward_std_ema = reward_std;
} else {
self.reward_std_ema = R_BETA * self.reward_std_ema + (1.0 - R_BETA) * reward_std;
}
self.reward_std_ema
} else if self.reward_std_ema > 1e-8 {
self.reward_std_ema // use cached EMA
} else {
0.01 // absolute fallback before first observation
};
// Decay floor as Q-values mature (same exp decay as before, but
// using the ACTUAL reward scale instead of hardcoded 0.01)
let maturity = (q_mean.abs() / effective_reward_std.max(1e-8)).min(5.0);
let num_atoms = self.config.num_atoms.max(2) as f32;
let gap_based_range = q_gap * (num_atoms / GAP_ATOM_TARGET);
let decayed_floor = effective_reward_std * (-maturity).exp();
let adaptive_min = decayed_floor
.max(gap_based_range)
.max(target_max - target_min);
let (final_min, final_max) = if (target_max - target_min) < adaptive_min {
(q_mean - adaptive_min * 0.5, q_mean + adaptive_min * 0.5)
} else {
(target_min, target_max)
};
let cur = self.v_range();
let changed = (final_min - cur[0]).abs() > 1e-6 || (final_max - cur[1]).abs() > 1e-6;
if changed {
unsafe { *self.v_range_pinned = final_min; *self.v_range_pinned.add(1) = final_max; }
}
changed
}
pub fn v_range(&self) -> [f32; 2] {
unsafe { [*self.v_range_pinned, *self.v_range_pinned.add(1)] }
}
pub fn v_range_buf_ptr(&self) -> u64 { self.v_range_dev_ptr }
// ── Adaptive gradient clipping ──────────────────────────────────
/// Update EMA-based adaptive gradient clip norm.