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

fix(sp14-egf): gate q_disagreement EMA update on total_cnt > 0 — fixes training-time decay-to-zero of rollout signal

Browse Source 
Root cause from train-6fcml 5-epoch trajectory (commit 5608b866b after
producer cadence migration): HEALTH_DIAG[0] (post-experience-collection)
showed q_dis_s=0.0595 q_dis_l=0.1329 var_q=0.00091 — meaningful rollout
signal. HEALTH_DIAG[1+] (post-training, per-step launches) all showed
q_dis_s=0.0000 q_dis_l=0.0000 var_q=0.00000 — signal decayed to zero
inside ONE epoch.

The kernel's ISV write block ran unconditionally even when total_cnt
(non-masked-row count after Hold/Flat masking) was 0. Empty-batch
launches blended `batch_mean = 0/1 = 0` into the EMA, decaying the
rollout signal to 0 over ~178 training steps × 0.7^n. Per-step training
launches read replay batches whose Q-direction picks are dominated by
Hold/Flat (the natural distribution); so total_cnt = 0 was the common
case, not a corner case.

Fix (atomic, single kernel):
- Wrap the ISV write block in `if (total_cnt > 0.0f) { ... }`. When the
  training batch has no non-masked rows, the kernel is a no-op for that
  step — EMAs stay at the prior step's values. Stream-ordered launches
  still run; only the ISV write is skipped.
- Remove redundant `&& (total_cnt > 0.0f)` clause from the `is_first`
  bootstrap check (now guaranteed by the outer gate).

Per pearl_first_observation_bootstrap semantics: "no observation"
preserves prior; only "first observation" replaces sentinel. Decay-on-
empty was inconsistent with both rules.

Other EGF-chain kernels audited:
- alpha_grad_compute_kernel.cu — operates on persistent ISV state,
  no batch concept; var_aux/var_alpha Welford updates use `diff` of
  persistent EMAs, not batch means. No empty-batch path. SAFE.
- aux_dir_acc_reduce_kernel.cu — emits out_6[0..3] with sentinel
  fallback (0.5) when denom==0; downstream apply_fixed_alpha_ema then
  blends 0.5 toward EMA. The sentinel is the random-baseline (target
  threshold lies above it), so empty-batch pulls EMA toward harmless
  baseline rather than zero. Different semantics from q_disagreement
  (which has 0 — far below baseline 0.5). SAFE.
- gradient_hack_detect_kernel.cu — single-thread state machine on
  persistent ISV, no batch. SAFE.

Verification:
- 6 existing sp14_oracle_tests pass.
- New q_disagreement_empty_batch_preserves_ema test asserts bit-exact
  preservation of pre-seeded EMAs (0.0595, 0.1329, 0.0009 — the
  train-6fcml HEALTH_DIAG[0] values) across an all-Hold batch. Catches
  the regression the existing all-hold test missed (its bound
  `[0.0, 0.5]` accepted both decay-to-blend and preserve-prior; the new
  test is strict bit-equality).
- L40S smoke validation pending — train-6fcml symptoms (alpha_smoothed
  stuck at 0.0002, gate1 closed forever) expected to resolve.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-07 20:54:07 +02:00
parent 5608b866b6
commit 9d0c124cee
3 changed files with 166 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

88
crates/ml/src/cuda_pipeline/q_disagreement_update_kernel.cu
View File

@@ -184,45 +184,63 @@ void q_disagreement_update_kernel(
if (tid == 0) {
const float total_num = num_smem[0];
const float total_cnt = cnt_smem[0];
// `fmaxf(total_cnt, 1.0f)` keeps the divide finite when every row
// is masked; when total_cnt > 0 it is the natural denominator.
const float batch_mean = total_num / fmaxf(total_cnt, 1.0f);
const float prev_short = isv[Q_DISAGREEMENT_SHORT_EMA_IDX];
const float prev_long = isv[Q_DISAGREEMENT_LONG_EMA_IDX];
const float prev_var = isv[Q_DISAGREEMENT_VARIANCE_EMA_IDX];
// ── Empty-batch gate (post-train-6fcml fix, 2026-05-07) ────────
// When the batch has zero non-masked contributions (every row
// Q-picked Hold or Flat — common for the training replay buffer
// because Hold/Flat dominate naturally), preserve all three EMAs
// at their previous values. Without this gate, the kernel would
// blend `batch_mean = 0/1 = 0` into the EMAs every step, which
// exponentially decayed the post-rollout signal to zero across
// the per-step training launches (commit 5608b866b cadence
// migration). Per `pearl_first_observation_bootstrap.md`:
// "no observation" preserves prior; only "first observation"
// replaces sentinel. Decay-on-empty was inconsistent with both
// rules. Verified via train-6fcml HEALTH_DIAG: post-rollout
// q_dis_s=0.0595 → post-training q_dis_s=0.0000 in one epoch.
if (total_cnt > 0.0f) {
const float batch_mean = total_num / total_cnt;
// Pearl-A first-observation: sentinel = 0.5 exact (per
// `SENTINEL_Q_DISAGREEMENT`). 0.5 is exactly representable in
// IEEE 754 (mantissa 1.0, exponent -1) so the `==` compare is
// safe. The first observation REPLACES the sentinel directly
// — see `pearl_first_observation_bootstrap.md`. Both EMAs are
// gated on `(total_cnt > 0)` to avoid the cold-start case
// where the batch has zero contributions (would replace the
// sentinel with batch_mean = 0, defeating bootstrap).
const bool is_first =
(prev_short == 0.5f) && (prev_long == 0.5f) && (total_cnt > 0.0f);
const float prev_short = isv[Q_DISAGREEMENT_SHORT_EMA_IDX];
const float prev_long = isv[Q_DISAGREEMENT_LONG_EMA_IDX];
const float prev_var = isv[Q_DISAGREEMENT_VARIANCE_EMA_IDX];
const float new_short = is_first
? batch_mean
: (alpha_short * batch_mean + (1.0f - alpha_short) * prev_short);
const float new_long = is_first
? batch_mean
: (alpha_long * batch_mean + (1.0f - alpha_long) * prev_long);
// Pearl-A first-observation: sentinel = 0.5 exact (per
// `SENTINEL_Q_DISAGREEMENT`). 0.5 is exactly representable
// in IEEE 754 (mantissa 1.0, exponent -1) so the `==`
// compare is safe. The first observation REPLACES the
// sentinel directly — see
// `pearl_first_observation_bootstrap.md`. The previous
// `&& (total_cnt > 0.0f)` guard on `is_first` is now
// redundant — the outer empty-batch gate above guarantees
// total_cnt > 0 here.
const bool is_first =
(prev_short == 0.5f) && (prev_long == 0.5f);
// Welford-style variance EMA: var_new = α_var × (x - mean)² +
// (1 - α_var) × var_old. Uses `new_short` as the running mean
// (the fast EMA tracks the mean's current best estimate); the
// alternative would be a separate mean-tracking slot, but the
// adaptive k_q controller in B.4 only consumes the variance
// ratio against `VARIANCE_REF_Q`, so the choice of "mean" here
// is mostly a numerical-stability detail (small bias toward
// the fast-EMA estimate, which is what B.4 wants anyway).
const float diff = batch_mean - new_short;
const float new_var = alpha_var * (diff * diff) + (1.0f - alpha_var) * prev_var;
const float new_short = is_first
? batch_mean
: (alpha_short * batch_mean + (1.0f - alpha_short) * prev_short);
const float new_long = is_first
? batch_mean
: (alpha_long * batch_mean + (1.0f - alpha_long) * prev_long);
isv[Q_DISAGREEMENT_SHORT_EMA_IDX] = new_short;
isv[Q_DISAGREEMENT_LONG_EMA_IDX] = new_long;
isv[Q_DISAGREEMENT_VARIANCE_EMA_IDX] = new_var;
// Welford-style variance EMA: var_new = α_var × (x - mean)²
// + (1 - α_var) × var_old. Uses `new_short` as the running
// mean (the fast EMA tracks the mean's current best
// estimate); the alternative would be a separate
// mean-tracking slot, but the adaptive k_q controller in
// B.4 only consumes the variance ratio against
// `VARIANCE_REF_Q`, so the choice of "mean" here is mostly
// a numerical-stability detail (small bias toward the
// fast-EMA estimate, which is what B.4 wants anyway).
const float diff = batch_mean - new_short;
const float new_var = alpha_var * (diff * diff) + (1.0f - alpha_var) * prev_var;
isv[Q_DISAGREEMENT_SHORT_EMA_IDX] = new_short;
isv[Q_DISAGREEMENT_LONG_EMA_IDX] = new_long;
isv[Q_DISAGREEMENT_VARIANCE_EMA_IDX] = new_var;
}
// else: total_cnt == 0 → preserve EMA state untouched. Kernel
// is a no-op for this step; stream-ordered launches still run.
}
}

111
crates/ml/tests/sp14_oracle_tests.rs
View File

@@ -332,6 +332,117 @@ mod gpu {
);
}
// ── Test B.3c: empty-batch preserves pre-seeded EMAs (post-train-6fcml) ──
/// Regression test for the post-train-6fcml decay-to-zero bug. Pre-seeds
/// the q-disagreement EMAs with non-zero values matching the post-rollout
/// HEALTH_DIAG[0] state observed in train-6fcml (q_dis_s = 0.0595,
/// q_dis_l = 0.1329, var_q = 0.0009 — clearly past the Pearl-A sentinel
/// 0.5 so the bootstrap branch is NOT taken), then launches the kernel
/// against an all-Hold batch (every row masked, total_cnt = 0).
///
/// Pre-fix (commit 5608b866b cadence migration symptom): the kernel's
/// ISV write block ran unconditionally with `batch_mean = 0/1 = 0`,
/// blending zero into the EMAs every step. Across the per-step training
/// launches the rollout signal exponentially decayed to zero in one
/// epoch, leaving the EGF gate1 stuck closed forever.
///
/// Post-fix: when total_cnt == 0 the kernel preserves all three EMAs at
/// their previous values. The `q_disagreement_all_hold_no_contribution`
/// test above only covers the cold-start case (prev = sentinel 0.5)
/// where blend-toward-0 happened to settle in [0.0, 0.5]; this test
/// catches the warm-start regression where blend-toward-0 silently
/// destroys the rollout-collected signal.
#[test]
#[ignore = "requires GPU"]
fn q_disagreement_empty_batch_preserves_ema() {
let stream = make_test_stream();
let kernel = load_q_disagreement(&stream);
const B: usize = 4;
const K_AUX: usize = 2;
const K_DIR: usize = 4;
const ISV_DIM: usize = 1024;
// Uniform aux (won't matter — all rows masked anyway).
let aux_softmax = vec![0.5_f32; B * K_AUX];
// Hold winner everywhere (DIR_HOLD = 1) → every row masked → total_cnt = 0.
let mut q_logits = vec![0.0_f32; B * K_DIR];
for b in 0..B {
q_logits[b * K_DIR + 1] = 1.0;
}
// Pre-seed EMAs with the train-6fcml HEALTH_DIAG[0] values (post-
// rollout signal). These are deliberately NOT the 0.5 sentinel so
// the Pearl-A bootstrap branch does not fire, isolating the
// empty-batch decay path.
const SEED_SHORT: f32 = 0.0595;
const SEED_LONG: f32 = 0.1329;
const SEED_VAR: f32 = 0.0009;
let mut isv = vec![0.0_f32; ISV_DIM];
isv[Q_DISAGREEMENT_SHORT_EMA_INDEX] = SEED_SHORT;
isv[Q_DISAGREEMENT_LONG_EMA_INDEX] = SEED_LONG;
isv[Q_DISAGREEMENT_VARIANCE_EMA_INDEX] = SEED_VAR;
let aux_buf = unsafe { MappedF32Buffer::new(B * K_AUX) }
.expect("alloc aux_softmax buffer");
let q_buf = unsafe { MappedF32Buffer::new(B * K_DIR) }
.expect("alloc q_logits buffer");
let isv_buf = unsafe { MappedF32Buffer::new(ISV_DIM) }
.expect("alloc isv buffer");
aux_buf.write_from_slice(&aux_softmax);
q_buf.write_from_slice(&q_logits);
isv_buf.write_from_slice(&isv);
let bsi: i32 = B as i32;
let q_stride: i32 = K_DIR as i32;
let alpha_short: f32 = 0.3;
let alpha_long: f32 = 0.05;
let alpha_var: f32 = 0.05;
unsafe {
stream
.launch_builder(&kernel)
.arg(&aux_buf.dev_ptr)
.arg(&q_buf.dev_ptr)
.arg(&isv_buf.dev_ptr)
.arg(&bsi)
.arg(&q_stride)
.arg(&alpha_short)
.arg(&alpha_long)
.arg(&alpha_var)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (Q_DIS_BLOCK, 1, 1),
shared_mem_bytes: Q_DIS_SMEM_BYTES,
})
.expect("launch q_disagreement_update_kernel");
}
stream.synchronize().expect("sync after q_disagreement_update_kernel");
let result = isv_buf.read_all();
let short_ema = result[Q_DISAGREEMENT_SHORT_EMA_INDEX];
let long_ema = result[Q_DISAGREEMENT_LONG_EMA_INDEX];
let var_ema = result[Q_DISAGREEMENT_VARIANCE_EMA_INDEX];
// Strict bit-equality on the seed values — the kernel must skip the
// ISV write block entirely when total_cnt == 0. Any deviation
// indicates the empty-batch gate is missing or buggy.
assert_eq!(
short_ema, SEED_SHORT,
"Empty-batch must preserve short EMA bit-exactly; expected {SEED_SHORT}, got {short_ema}",
);
assert_eq!(
long_ema, SEED_LONG,
"Empty-batch must preserve long EMA bit-exactly; expected {SEED_LONG}, got {long_ema}",
);
assert_eq!(
var_ema, SEED_VAR,
"Empty-batch must preserve variance EMA bit-exactly; expected {SEED_VAR}, got {var_ema}",
);
}
// ── B.4 cubin handle ────────────────────────────────────────────────────
const SP14_ALPHA_GRAD_CUBIN: &[u8] =