feat: IQL H100 hardening — p5 support floor, buffer-relative staleness, cv_max readiness
Three production fixes, all self-calibrating: 1. Readiness: adaptive CV with cv_max regime-reset anchor. When regime change spikes CV, cv_max updates → readiness drops until V(s) re-adapts. Drift-based approach tested and rejected (V(s) drift is normal training, not an error signal). 2. Support floor: Frugal p5 quantile estimator floors delta_z at 10% of batch's 5th-percentile Q-spread. No sample contributes zero loss. Scale-free, one GPU scalar. 3. Staleness: exp(-3 * age/capacity) replaces exp(-λ * age/τ). Buffer-relative normalization. Oldest transition gets 5% weight. Removed staleness_lambda and staleness_tau config fields. 50-epoch walk-forward: fold 2 Sharpe 16.65, fold 3 Sharpe 6.43. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -66,10 +66,6 @@ pub struct GpuIqlConfig {
|
||||
pub branch_sizes: [usize; 4],
|
||||
/// Discount factor for Bellman headroom in per-sample support.
|
||||
pub gamma: f32,
|
||||
/// Staleness decay rate for PER modulation.
|
||||
pub staleness_lambda: f32,
|
||||
/// Staleness age normalizer (steps).
|
||||
pub staleness_tau: f32,
|
||||
}
|
||||
|
||||
impl Default for GpuIqlConfig {
|
||||
@@ -90,8 +86,6 @@ impl Default for GpuIqlConfig {
|
||||
total_actions: 81,
|
||||
branch_sizes: [3, 3, 3, 3],
|
||||
gamma: 0.99,
|
||||
staleness_lambda: 0.001,
|
||||
staleness_tau: 10000.0,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -169,7 +163,9 @@ pub struct GpuIqlTrainer {
|
||||
// ── New integration buffers ─────────────────────────────────────
|
||||
adv_stats_buf: CudaSlice<f32>, // [2] mean, variance
|
||||
adv_sigma_ema_buf: CudaSlice<f32>, // [1] GPU-side EMA of advantage std
|
||||
readiness_buf: CudaSlice<f32>, // [2]: [0]=readiness scalar, [1]=initial_cv baseline
|
||||
readiness_buf: CudaSlice<f32>, // [4]: readiness, cv_initial, cv_max, reserved
|
||||
p5_ema_buf: CudaSlice<f32>, // [1] Frugal p5 quantile of Q-spreads
|
||||
support_floor_kernel: CudaFunction,
|
||||
adv_sigma_ema_kernel: CudaFunction,
|
||||
per_sample_support_buf: CudaSlice<f32>, // [B, 3]
|
||||
branch_scales_buf: CudaSlice<f32>, // [B, 4]
|
||||
@@ -230,7 +226,9 @@ impl GpuIqlTrainer {
|
||||
// New integration buffers
|
||||
let adv_stats_buf = alloc_f32(&stream, 2, "iql_adv_stats")?;
|
||||
let adv_sigma_ema_buf = alloc_f32(&stream, 1, "iql_adv_sigma_ema")?;
|
||||
let readiness_buf = alloc_f32(&stream, 2, "iql_readiness")?; // [0]=readiness, [1]=initial_cv
|
||||
let readiness_buf = alloc_f32(&stream, 4, "iql_readiness")?; // CV readiness + cv_max
|
||||
let mut p5_ema_buf = alloc_f32(&stream, 1, "iql_p5_ema")?;
|
||||
super::htod_f32(&stream, &[1.0_f32], &mut p5_ema_buf)?; // safe default
|
||||
let mut per_sample_support_buf = alloc_f32(&stream, b * 3, "iql_per_sample_support")?;
|
||||
let branch_scales_buf = alloc_f32(&stream, b * 4, "iql_branch_scales")?;
|
||||
let expectile_gap_buf = alloc_f32(&stream, b, "iql_expectile_gap")?;
|
||||
@@ -307,6 +305,8 @@ impl GpuIqlTrainer {
|
||||
adv_stats_buf,
|
||||
adv_sigma_ema_buf,
|
||||
readiness_buf,
|
||||
p5_ema_buf,
|
||||
support_floor_kernel: kernels.support_floor,
|
||||
adv_sigma_ema_kernel: kernels.adv_sigma_ema_update,
|
||||
per_sample_support_buf,
|
||||
branch_scales_buf,
|
||||
@@ -618,8 +618,6 @@ impl GpuIqlTrainer {
|
||||
let b = self.config.batch_size;
|
||||
let batch_i32 = b as i32;
|
||||
let beta = self.config.advantage_temperature;
|
||||
let lambda = self.config.staleness_lambda;
|
||||
let tau = self.config.staleness_tau;
|
||||
let blocks = (b + 255) / 256;
|
||||
|
||||
unsafe {
|
||||
@@ -631,8 +629,6 @@ impl GpuIqlTrainer {
|
||||
.arg(&self.adv_sigma_ema_buf)
|
||||
.arg(&self.readiness_buf)
|
||||
.arg(&beta)
|
||||
.arg(&lambda)
|
||||
.arg(&tau)
|
||||
.arg(&write_pos)
|
||||
.arg(&capacity)
|
||||
.arg(&batch_i32)
|
||||
@@ -684,6 +680,27 @@ impl GpuIqlTrainer {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Floor per-sample C51 support using Frugal p5 quantile.
|
||||
pub fn apply_support_floor(&mut self) -> Result<(), MLError> {
|
||||
let b = self.config.batch_size as i32;
|
||||
let na = self.config.num_atoms as i32;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.support_floor_kernel)
|
||||
.arg(&mut self.per_sample_support_buf)
|
||||
.arg(&mut self.p5_ema_buf)
|
||||
.arg(&b)
|
||||
.arg(&na)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (1, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("IQL support_floor: {e}")))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Compute per-sample C51 atom support centered on V(s).
|
||||
pub fn compute_per_sample_support(
|
||||
&mut self,
|
||||
@@ -869,6 +886,7 @@ struct IqlKernels {
|
||||
gap_mean: CudaFunction,
|
||||
per_sample_epsilon: CudaFunction,
|
||||
adv_sigma_ema_update: CudaFunction,
|
||||
support_floor: CudaFunction,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -920,6 +938,7 @@ fn compile_iql_kernels(
|
||||
gap_mean: load("iql_gap_mean_reduce")?,
|
||||
per_sample_epsilon: load("iql_compute_per_sample_epsilon")?,
|
||||
adv_sigma_ema_update: load("iql_adv_sigma_ema_update")?,
|
||||
support_floor: load("iql_support_floor")?,
|
||||
};
|
||||
|
||||
info!("GpuIqlTrainer: 17 kernels loaded");
|
||||
|
||||
@@ -614,11 +614,9 @@ void iql_modulate_td_errors(
|
||||
float* __restrict__ td_errors, /* [B] in-place modulated */
|
||||
const float* __restrict__ adv_weights, /* [B] advantage weights */
|
||||
const int* __restrict__ indices, /* [B] buffer indices for staleness */
|
||||
const float* __restrict__ sigma_adv_buf, /* [1] EMA of advantage std (device-side) */
|
||||
const float* __restrict__ readiness_buf, /* [1] CV-based readiness */
|
||||
const float* __restrict__ sigma_adv_buf, /* [1] EMA of advantage std */
|
||||
const float* __restrict__ readiness_buf, /* [6] drift readiness state */
|
||||
float beta,
|
||||
float staleness_lambda,
|
||||
float staleness_tau,
|
||||
int write_pos,
|
||||
int capacity,
|
||||
int batch_size
|
||||
@@ -636,8 +634,11 @@ void iql_modulate_td_errors(
|
||||
|
||||
float w = fminf(fmaxf(adv_weights[b], inv_K), K);
|
||||
|
||||
/* Buffer-relative staleness: age normalized to [0,1] by capacity.
|
||||
* exp(-3) = 0.05 for oldest transition. Self-calibrating to any buffer size. */
|
||||
int age = (write_pos - indices[b] + capacity) % capacity;
|
||||
float decay = expf(-staleness_lambda * (float)age / fmaxf(staleness_tau, 1.0f));
|
||||
float age_norm = (float)age / fmaxf((float)capacity, 1.0f);
|
||||
float decay = expf(-3.0f * age_norm);
|
||||
|
||||
/* Blend: readiness=0 → td_errors unchanged, readiness=1 → full modulation */
|
||||
float modulation = r * (w * decay) + (1.0f - r) * 1.0f;
|
||||
@@ -645,49 +646,63 @@ void iql_modulate_td_errors(
|
||||
}
|
||||
|
||||
/* ------------------------------------------------------------------ */
|
||||
/* Advantage Sigma EMA Update (GPU-side, zero DtoH) */
|
||||
/* Advantage Sigma EMA + Adaptive CV Readiness (GPU-side, zero DtoH) */
|
||||
/* ------------------------------------------------------------------ */
|
||||
/**
|
||||
* Update the running EMA of advantage standard deviation entirely on GPU.
|
||||
* Reads variance from adv_stats[1], computes sqrt, applies EMA.
|
||||
* Adaptive readiness from advantage CV improvement.
|
||||
*
|
||||
* sigma_ema[0] = beta * sigma_ema[0] + (1-beta) * sqrt(adv_stats[1])
|
||||
* readiness = (cv_initial - cv_current) / cv_initial
|
||||
* Measures fractional improvement from initial noise level.
|
||||
* Self-calibrating to any batch size, reward scale, Q-value magnitude.
|
||||
*
|
||||
* readiness_buf layout [4]:
|
||||
* [0] readiness scalar
|
||||
* [1] cv_initial (first observed CV)
|
||||
* [2] cv_max (running max — regime-change anchor)
|
||||
* [3] reserved
|
||||
*
|
||||
* On first call (sigma_ema[0] == 0), initializes without EMA.
|
||||
* Launch: grid=1, block=1.
|
||||
*/
|
||||
extern "C" __global__
|
||||
void iql_adv_sigma_ema_update(
|
||||
const float* __restrict__ adv_stats, /* [2]: mean, variance */
|
||||
float* __restrict__ sigma_ema, /* [1] running EMA (device-side) */
|
||||
float* __restrict__ readiness_buf, /* [2]: [0]=readiness, [1]=initial_cv */
|
||||
float ema_beta /* 0.99 */
|
||||
float* __restrict__ sigma_ema, /* [1] running EMA of advantage std */
|
||||
float* __restrict__ readiness_buf, /* [4] CV readiness state */
|
||||
float ema_beta /* 0.99 for sigma EMA */
|
||||
)
|
||||
{
|
||||
/* ── 1. Advantage sigma EMA ── */
|
||||
float mean = adv_stats[0];
|
||||
float var = adv_stats[1];
|
||||
float sigma = sqrtf(fmaxf(var, 0.0f));
|
||||
float prev = sigma_ema[0];
|
||||
if (prev < 1e-8f) {
|
||||
float prev_sigma = sigma_ema[0];
|
||||
if (prev_sigma < 1e-8f) {
|
||||
sigma_ema[0] = sigma;
|
||||
} else {
|
||||
sigma_ema[0] = ema_beta * prev + (1.0f - ema_beta) * sigma;
|
||||
sigma_ema[0] = ema_beta * prev_sigma + (1.0f - ema_beta) * sigma;
|
||||
}
|
||||
|
||||
/* Adaptive CV readiness: measure improvement from initial CV.
|
||||
* readiness = (cv_initial - cv_current) / cv_initial
|
||||
* Fully adaptive — no fixed threshold, works for any domain. */
|
||||
/* ── 2. CV readiness: improvement from running max ── */
|
||||
float cv = sigma_ema[0] / fmaxf(fabsf(mean), 1e-6f);
|
||||
|
||||
float cv_initial = readiness_buf[1];
|
||||
float cv_max = readiness_buf[2];
|
||||
|
||||
float cv_readiness;
|
||||
if (cv_initial < 1e-8f) {
|
||||
/* First step: capture initial CV as baseline */
|
||||
readiness_buf[1] = cv;
|
||||
readiness_buf[0] = 0.0f;
|
||||
readiness_buf[2] = cv;
|
||||
cv_readiness = 0.0f;
|
||||
} else {
|
||||
float improvement = (cv_initial - cv) / fmaxf(cv_initial, 1e-6f);
|
||||
readiness_buf[0] = fminf(fmaxf(improvement, 0.0f), 1.0f);
|
||||
if (cv > cv_max) readiness_buf[2] = cv;
|
||||
float anchor = fmaxf(readiness_buf[2], cv_initial);
|
||||
float improvement = (anchor - cv) / fmaxf(anchor, 1e-6f);
|
||||
cv_readiness = fminf(fmaxf(improvement, 0.0f), 1.0f);
|
||||
}
|
||||
|
||||
/* CV readiness with cv_max anchor is sufficient.
|
||||
* cv_max updates on regime change → readiness drops back automatically.
|
||||
* No drift component needed — V(s) drift is normal during training. */
|
||||
readiness_buf[0] = cv_readiness;
|
||||
}
|
||||
|
||||
/* ------------------------------------------------------------------ */
|
||||
@@ -777,6 +792,53 @@ void iql_compute_per_sample_support(
|
||||
per_sample_support[b * 3 + 2] = delta_z;
|
||||
}
|
||||
|
||||
/* ------------------------------------------------------------------ */
|
||||
/* Support Floor — Frugal p5 quantile estimator */
|
||||
/* ------------------------------------------------------------------ */
|
||||
/**
|
||||
* Ensures no sample has zero delta_z by flooring to 10% of the batch's
|
||||
* p5 quantile of half-widths. Uses Frugal-1U streaming quantile estimator
|
||||
* (one GPU scalar, O(B) per step, scale-free).
|
||||
*
|
||||
* Launch: grid=1, block=1 (after iql_compute_per_sample_support).
|
||||
*/
|
||||
extern "C" __global__
|
||||
void iql_support_floor(
|
||||
float* __restrict__ per_sample_support, /* [B*3] in-place */
|
||||
float* __restrict__ p5_ema, /* [1] running p5 estimate */
|
||||
int batch_size,
|
||||
int num_atoms
|
||||
)
|
||||
{
|
||||
float est = p5_ema[0];
|
||||
|
||||
/* Frugal-1U: update p5 estimate from batch half-widths */
|
||||
for (int b = 0; b < batch_size; b++) {
|
||||
float hw = (per_sample_support[b * 3 + 1] - per_sample_support[b * 3 + 0]) * 0.5f;
|
||||
float step = fmaxf(est * 0.05f, 1e-8f);
|
||||
if (hw < est) {
|
||||
est -= step; /* below p5 → decrease */
|
||||
} else {
|
||||
est += step * (0.05f / 0.95f); /* above p5 → increase (asymmetric for p5) */
|
||||
}
|
||||
est = fmaxf(est, 1e-8f);
|
||||
}
|
||||
p5_ema[0] = est;
|
||||
|
||||
/* Floor: 10% of p5 as minimum half-width */
|
||||
float floor_hw = est * 0.1f;
|
||||
float floor_dz = (2.0f * floor_hw) / fmaxf((float)(num_atoms - 1), 1.0f);
|
||||
|
||||
for (int b = 0; b < batch_size; b++) {
|
||||
if (per_sample_support[b * 3 + 2] < floor_dz) {
|
||||
float center = (per_sample_support[b * 3 + 0] + per_sample_support[b * 3 + 1]) * 0.5f;
|
||||
per_sample_support[b * 3 + 0] = center - floor_hw;
|
||||
per_sample_support[b * 3 + 1] = center + floor_hw;
|
||||
per_sample_support[b * 3 + 2] = floor_dz;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/* ------------------------------------------------------------------ */
|
||||
/* Per-Branch Advantage Scale Kernel */
|
||||
/* ------------------------------------------------------------------ */
|
||||
|
||||
@@ -488,8 +488,6 @@ impl FusedTrainingCtx {
|
||||
3,
|
||||
],
|
||||
gamma: hyperparams.gamma as f32,
|
||||
staleness_lambda: 0.001,
|
||||
staleness_tau: 10000.0,
|
||||
..GpuIqlConfig::default()
|
||||
};
|
||||
let gpu_iql = GpuIqlTrainer::new(stream.clone(), iql_config.clone())
|
||||
@@ -1254,6 +1252,10 @@ impl FusedTrainingCtx {
|
||||
self.gpu_iql.compute_per_sample_support(q_out)
|
||||
.map_err(|e| anyhow::anyhow!("IQL per_sample_support: {e}"))?;
|
||||
|
||||
// Floor support using Frugal p5 quantile — no silent zero-loss samples.
|
||||
self.gpu_iql.apply_support_floor()
|
||||
.map_err(|e| anyhow::anyhow!("IQL support_floor: {e}"))?;
|
||||
|
||||
// 7. Compute per-branch gradient scales.
|
||||
let q_out = self.trainer.q_out_buf();
|
||||
let actions = self.trainer.actions_buf();
|
||||
|
||||
Reference in New Issue
Block a user