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

fix(F4/F5/D6): kernel buffer-stride bug + D6 warmup gate

Browse Source 
Root cause of L40S segfault (train-tgdlq workflow, commit 0d639dfe9):

1. F4 (IB) and F5 (barrier) gradient kernels indexed cql_d_adv_logits and
   on_b_logits_buf with stride b0_size (4) instead of total_actions (13).
   The buffer is sized [B, total_actions, num_atoms]. Writes for batch i>0
   landed in other actions'/batches' gradient slots — corrupting CQL gradient
   for every batch beyond the first. After cuBLAS backward, weights diverged
   in undefined ways; validation forward then segfaulted reading NaN-laced
   parameters in GpuBacktestEvaluator init.

   Fix: add `total_actions` kernel parameter (int), use it as the full stride
   for adv_row and d_adv_a pointer arithmetic; keep b0_size as the loop
   bound (direction-branch-only update). All three launch sites updated:
   inlined F5 launch in apply_cql_gradient, inlined F4 launch alongside,
   and the standalone inject_barrier_into_cql_d_logits method.

2. D6 ensemble oracle fired at epoch 0 because per-branch Q-gaps are all 0
   at random init (range = 0 → score = 1.0 → plasticity trigger).
   Shrink-and-perturb ran immediately, then again next epoch, etc. Gate
   behind `learning_health.epoch > 5` (3 warmup + 2 buffer epochs for Q to
   move) so the oracle only fires on post-warmup real collapse, not
   untrained networks.

Both bugs are regressions from today's work — F4/F5 introduced yesterday,
D6 became live after the ens_disagreement real-signal fix in d9d35b6fa.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-20 22:03:45 +02:00
parent e45e58a41a
commit cc96dd7fd3
3 changed files with 34 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
View File

@@ -759,15 +759,16 @@ extern "C" __global__ void c51_loss_reduce(
* Launch: grid=(ceil(B/256)), block=(256). One thread per sample.
* ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void barrier_gradient_direction(
const float* __restrict__ adv_logits, /* [B, b0_size, num_atoms] f32 */
const float* __restrict__ adv_logits, /* [B, total_actions, num_atoms] f32 (direction is first b0_size slots) */
const float* __restrict__ v_logits, /* [B, num_atoms] f32 */
const float* __restrict__ z_vals, /* [num_atoms] f32 */
const float* __restrict__ isv_signals, /* [ISV_DIM=13] pinned; [12]=health */
float* __restrict__ d_adv_logits, /* [B, b0_size, num_atoms] — atomicAdd here */
float* __restrict__ d_adv_logits, /* [B, total_actions, num_atoms] — atomicAdd direction slice only */
float* __restrict__ d_v_logits, /* [B, num_atoms] — atomicAdd here */
int batch_size,
int num_atoms,
int b0_size,
int total_actions, /* true stride of adv / d_adv buffers */
float barrier_weight
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
@@ -779,13 +780,15 @@ extern "C" __global__ void barrier_gradient_direction(
/* Safety guard: b0_size must be in a reasonable range */
if (b0_size <= 1 || b0_size > 16) return;
if (total_actions < b0_size) return; /* must be at least direction branch */
/* Compute Q(a) for each direction action.
* Dueling: combined_logit(a, z) = v_logits[i, z] + adv_logits[i, a, z]
* (skip advantage-mean centering — small epsilon vs correctness simplicity). */
float q_vals[16];
const float* v_row = v_logits + (long long)i * num_atoms;
const float* adv_row = adv_logits + (long long)i * b0_size * num_atoms;
/* adv_logits is [B, total_actions, num_atoms]; direction actions are 0..b0_size. */
const float* adv_row = adv_logits + (long long)i * total_actions * num_atoms;
for (int a = 0; a < b0_size; a++) {
const float* adv_a = adv_row + a * num_atoms;
@@ -848,7 +851,7 @@ extern "C" __global__ void barrier_gradient_direction(
for (int z = 0; z < num_atoms; z++) {
sum += expf(v_row[z] + adv_a[z] - max_l);
}
float* d_adv_a = d_adv_logits + (long long)i * b0_size * num_atoms + a * num_atoms;
float* d_adv_a = d_adv_logits + (long long)i * total_actions * num_atoms + a * num_atoms;
float* d_v_row = d_v_logits + (long long)i * num_atoms;
for (int z = 0; z < num_atoms; z++) {
float p = expf(v_row[z] + adv_a[z] - max_l) / (sum + 1e-8f);
@@ -891,15 +894,16 @@ extern "C" __global__ void barrier_gradient_direction(
* Launch: grid=(ceil(B/256)), block=(256). One thread per sample.
* ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void ib_gradient_direction(
const float* __restrict__ adv_logits, /* [B, b0_size, num_atoms] f32 */
const float* __restrict__ adv_logits, /* [B, total_actions, num_atoms] f32 */
const float* __restrict__ v_logits, /* [B, num_atoms] f32 */
const float* __restrict__ z_vals, /* [num_atoms] f32 */
const float* __restrict__ isv_signals, /* [ISV_DIM=13] pinned; [12]=health */
float* __restrict__ d_adv_logits, /* [B, b0_size, num_atoms] — atomicAdd */
float* __restrict__ d_adv_logits, /* [B, total_actions, num_atoms] — atomicAdd direction slice */
float* __restrict__ d_v_logits, /* [B, num_atoms] — atomicAdd */
int batch_size,
int num_atoms,
int b0_size,
int total_actions, /* true stride of adv / d_adv buffers */
float min_var
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
@@ -907,13 +911,15 @@ extern "C" __global__ void ib_gradient_direction(
/* Safety guard */
if (b0_size <= 1 || b0_size > 16) return;
if (total_actions < b0_size) return;
float health = (isv_signals != NULL) ? isv_signals[12] : 0.5f;
float ib_weight = 0.05f * (1.0f - health);
if (ib_weight < 1e-6f) return;
const float* v_row = v_logits + (long long)i * num_atoms;
const float* adv_row = adv_logits + (long long)i * b0_size * num_atoms;
/* adv_logits is [B, total_actions, num_atoms]; direction actions are 0..b0_size. */
const float* adv_row = adv_logits + (long long)i * total_actions * num_atoms;
/* Compute Q(a) for each direction action via numerically-stable softmax. */
float q_vals[16];
@@ -970,7 +976,7 @@ extern "C" __global__ void ib_gradient_direction(
for (int z = 0; z < num_atoms; z++)
sum += expf(v_row[z] + adv_a[z] - max_l);
float* d_adv_a = d_adv_logits + (long long)i * b0_size * num_atoms + a * num_atoms;
float* d_adv_a = d_adv_logits + (long long)i * total_actions * num_atoms + a * num_atoms;
float q_a = q_vals[a];
for (int z = 0; z < num_atoms; z++) {
float p = expf(v_row[z] + adv_a[z] - max_l) / (sum + 1e-8f);

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

@@ -4690,6 +4690,7 @@ impl GpuDqnTrainer {
let b_i32 = b as i32;
let na_i32_b = na as i32;
let b0_i32_b = b0 as i32;
let ta_i32_b = total_actions as i32;
let blocks_b = blocks;
unsafe {
let _ = self.stream
@@ -4703,6 +4704,7 @@ impl GpuDqnTrainer {
.arg(&b_i32)
.arg(&na_i32_b)
.arg(&b0_i32_b)
.arg(&ta_i32_b)
.arg(&barrier_weight)
.launch(LaunchConfig {
grid_dim: (blocks_b, 1, 1),
@@ -4727,6 +4729,7 @@ impl GpuDqnTrainer {
let b_i32_ib = b as i32;
let na_i32_ib = na as i32;
let b0_i32_ib = b0 as i32;
let ta_i32_ib = total_actions as i32;
let blocks_ib = blocks;
let min_var = 0.01_f32;
unsafe {
@@ -4741,6 +4744,7 @@ impl GpuDqnTrainer {
.arg(&b_i32_ib)
.arg(&na_i32_ib)
.arg(&b0_i32_ib)
.arg(&ta_i32_ib)
.arg(&min_var)
.launch(LaunchConfig {
grid_dim: (blocks_ib, 1, 1),
@@ -4968,13 +4972,18 @@ impl GpuDqnTrainer {
let b = self.config.batch_size;
let na = self.config.num_atoms;
let b0 = self.config.branch_0_size;
let b1 = self.config.branch_1_size;
let b2 = self.config.branch_2_size;
let b3 = self.config.branch_3_size;
let total_actions = b0 + b1 + b2 + b3;
let n_i32 = b as i32;
let na_i32 = na as i32;
let b0_i32 = b0 as i32;
let ta_i32 = total_actions as i32;
let blocks = ((b + 255) / 256) as u32;
// adv_logits for b0 branch is at the START of on_b_logits_buf (branch-major layout).
// cql_d_adv_logits is also branch-major: [B, b0_size, NA] at offset 0.
// adv_logits and cql_d_adv_logits are both [B, total_actions, num_atoms];
// we only update the direction slice (first b0 action slots).
let adv_ptr = self.on_b_logits_buf.raw_ptr();
let v_ptr = self.on_v_logits_buf.raw_ptr();
// z_vals: use atom_positions_buf[0..num_atoms] (branch 0 adaptive positions).
@@ -4994,6 +5003,7 @@ impl GpuDqnTrainer {
.arg(&n_i32)
.arg(&na_i32)
.arg(&b0_i32)
.arg(&ta_i32)
.arg(&barrier_weight)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),

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

@@ -1987,7 +1987,14 @@ impl DQNTrainer {
let ensemble_collapse_score = 1.0 - smoothstep(0.01, 0.1, mean_pairwise_gap);
self.last_ensemble_collapse_score = Some(ensemble_collapse_score);
if ensemble_collapse_score > 0.8 && self.unhealthy_epoch_count < 3 {
// D6 warmup guard: don't trigger during the first 5 epochs. At random
// init, per-branch Q-gaps are all ~0 → ensemble_collapse_score = 1.0, which
// would fire plasticity injection every epoch before training has any
// chance to produce real Q differentiation. Wait until learning_health
// has exited its own warmup (3 epochs) plus 2 buffer epochs for Q to move.
if ensemble_collapse_score > 0.8
&& self.unhealthy_epoch_count < 3
&& self.learning_health.epoch > 5 {
tracing::info!(
"Ensemble collapse detected (score={:.2}): forcing plasticity injection",
ensemble_collapse_score