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

feat(dqn): SP3 Mech 6 — anchored upper bound on adaptive grad clip

Browse Source 
Adds an upper bound to update_adaptive_clip's new_clip formula to
prevent clip-drift pathology: consecutive elevated samples were
ratcheting the EMA-driven clip threshold upward without bound,
eventually rendering clipping a no-op against in-distribution drift.

Smoke smoke-test-5rqzs (commit b9edccfc1) F1-NaN'd at step 3060 with
Mech 5 diagnostic slots 36-38, 40-42 firing — DQN main Adam m + v
EMAs saturated. Diagnostic confirmed Adam saturation as the root cause.

Diagnostic chain identified: grad_norm_ema ratchets up via repeat-
sample compounding (winsorizer caps single-sample but not consecutive
elevated samples) -> clip threshold > actual grad_norm -> clip no-op
-> Adam m/v poisoned -> saturate at slot 36-42 thresholds -> cuBLAS
overflow at slots 26+32 -> loss=34.68 grad=inf.

Fix: cap new_clip by grad_norm_slow_ema * 100 * isv[Q_ABS_REF=16].max(1).
- grad_norm_slow_ema (existing alpha=0.001 slow EMA): anchors against
  legitimate steady-state grad norm
- 100x: headroom for per-step deviations
- isv[Q_ABS_REF].max(1): adaptive scale per
  feedback_isv_for_adaptive_bounds; epsilon on multiplier per SP1 pearl
- .max(MIN_CLIP): cold-start epsilon-floor

Standard DL practice — bounds the runaway EMA clip-drift while
preserving adaptive behavior. F0 risk: low — slow_ema * 100 is
broad headroom; F0-typical clip thresholds are well below this
cap. F1+F2 benefit by preventing the ratchet pathology.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-30 11:06:36 +02:00
parent b9edccfc19
commit 48c25d9997
2 changed files with 35 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -20316,6 +20316,39 @@ impl GpuDqnTrainer {
self.grad_norm_ema = EMA_BETA * self.grad_norm_ema + (1.0 - EMA_BETA) * clamped_grad_norm;
}
let new_clip = (self.grad_norm_ema * CLIP_MULTIPLIER).max(MIN_CLIP);
// SP3 Mech 6 (2026-04-29): anchored upper bound on `new_clip`.
// The winsorizer above caps a SINGLE sample at K × prev_clip but
// does NOT prevent CONSECUTIVE elevated samples from compounding
// the EMA upward without bound. Over hundreds of steps the clip
// threshold ratchets to thousands while actual grad_norm tracks
// it from below → clipping becomes a no-op against in-distribution
// drift → Adam m/v EMAs poisoned → saturate. This is the F1 NaN
// root cause from smoke-test-5rqzs (commit b9edccfc1) at step 3060
// (Mech 5 diag slots 36–38, 40–42 firing).
//
// Bound: 100 × grad_norm_slow_ema × ISV[Q_ABS_REF=16].max(1.0)
// - `grad_norm_slow_ema` (existing α=0.001 slow EMA, updated below
// in this same function) anchors against the legitimate
// steady-state grad norm; here we read the PREVIOUS step's slow
// EMA from pinned memory (it gets overwritten further down).
// - 100×: headroom for legitimate per-step deviations; single
// steps can have norm 10–100× the slow average without being
// pathological.
// - ISV[Q_ABS_REF=16] multiplier: scales the cap with the
// Q-magnitude regime per `feedback_isv_for_adaptive_bounds.md`.
// Same ISV slot used by SP3 Mechs 1+2+5 — no new slot.
// - `.max(1.0_f32)` ε on the multiplier per the SP1 ε-floor pearl:
// cold-start ISV[16] ≈ 0 must not collapse the bound.
// - `.max(MIN_CLIP)` ε-floor on the bound itself: cold-start
// `grad_norm_slow_ema` ≈ 0 must not pin upper_bound at 0 (which
// would force `new_clip` down to MIN_CLIP every step until the
// slow EMA warms up).
let abs_mult = self.read_isv_signal_at(Q_ABS_REF_INDEX).max(1.0_f32);
let prev_slow_ema = unsafe { *self.grad_norm_slow_ema_pinned };
let upper_bound = (prev_slow_ema * 100.0_f32 * abs_mult).max(MIN_CLIP);
let new_clip = new_clip.min(upper_bound);
// Write directly to pinned memory — GPU sees it on next kernel read
unsafe { *self.adaptive_clip_pinned = new_clip; }

2
docs/dqn-wire-up-audit.md
View File

@@ -2284,3 +2284,5 @@ SP3 Task B5 — comprehensive Adam EMA reset (2026-04-29): audited every Adam op
SP3 Task B6 — fused kernel extended for slots 36-47 threshold checks (2026-04-30): extended `dqn_nan_check_fused_f32_kernel` (SP2 Task A4) from 12-slot NaN-only coverage (slots 24-35) to 24-slot NaN + threshold coverage (slots 24-47). Kernel signature gained one f32 by-value arg `q_abs_ref_eff` (passed as `f32` to match the kernel's `float` param per `feedback_cudarc_f64_f32_abi`); per-slot threshold computed inline via the relative `slot = blockIdx.x` index. Slots 0-11 (= absolute 24-35) keep `thr = INFINITY` so the magnitude check is trivially false (NaN-only behaviour preserved bit-exactly). Slots 12-15 (= absolute 36-39) check Adam m max-abs ≥ `100 × q_abs_ref_eff`; 16-19 (= 40-43) check Adam v max-abs ≥ `1e6 × q_abs_ref_eff²`; 20-21 (= 44-45) check weight max-abs ≥ `1e3 × q_abs_ref_eff`; 22 (= 46) checks `target_q` post-clip ≥ `95 × q_abs_ref_eff` (= `9.5 × max_abs_target_q` where `max_abs_target_q = 10 × q_abs_ref_eff` per B2); 23 (= 47) checks atom-positions span ≥ `190 × q_abs_ref_eff` (= `9.5 × max_atom_abs × 2` where `max_atom_abs = 10 × q_abs_ref_eff` per B3). Threshold logic is inline `if (slot >= 12 && slot < 16) thr = 100*q...` — no per-slot threshold buffer, no per-step HtoD. `q_abs_ref_eff = max(ISV[Q_ABS_REF=16], 1.0)` computed host-side at launch via `read_isv_signal_at` and floored at 1.0 per the SP1 ε-floor pearl (cold-start ISV ~0 yields q_abs_ref_eff = 1, so all thresholds floor at their ε-floor multiplier × 1). **Buffer + populate + launch wiring**: `nan_check_buf_ptrs` / `nan_check_buf_lens` resized 12 → 24 entries (mapped-pinned host write at construction; no HtoD copy per `feedback_no_htod_htoh_only_mapped_pinned`); `populate_nan_check_meta` extended with 4 new args `iqn_adam_m_ptr / iqn_adam_m_len / iqn_adam_v_ptr / iqn_adam_v_len: Option<u64>/Option<usize>` (None when IQN inactive — entry becomes (0, 0), kernel's null-pointer guard skips the block; symmetric with slots 27/28); 12 new entries appended for slots 36-47 calling the SP3 Task B1 accessors (`trunk_adam_{m,v}_{ptr,len}`, `value_adam_{m,v}_{ptr,len}`, `branch_adam_{m,v}_{ptr,len}`, `trunk_params_{ptr,len}`, `heads_params_{ptr,len}`, `target_q_{ptr,len}`, `atom_positions_{ptr,len}`); caller `fused_training.rs::FusedTrainingCtx::new` updated to pass the IQN Adam ptrs via `gpu_iqn.as_ref().map(|h| h.adam_m_ptr())` etc. **Launch wrapper** `launch_nan_check_fused_f32` reads `q_abs_ref` via `read_isv_signal_at(Q_ABS_REF_INDEX)`, computes `q_abs_ref_eff: f32 = q_abs_ref.max(1.0)`, and grows the grid from 12 → 24 blocks; the 5-arg launch builder now passes `q_abs_ref_eff` between `buf_lens_dev` and `BASE_FLAG_IDX`. Sticky-flag semantics preserved (kernel writes 1 only). Closes the diagnostic instrumentation loop for SP3 Mech 5 — slots 36-47 fire when their threshold is exceeded, providing observability for the other 4 SP3 mechanisms' effectiveness; if the SP3 fix doesn't fully resolve F1 NaN, the slot 36-47 firing pattern guides the next iteration. Zero new HtoD/DtoD/HtoH copies (one f32 by-value kernel arg = register pressure only, identical to existing `BASE_FLAG_IDX`). Zero new ISV slots (reuses `Q_ABS_REF_INDEX = 16`). Zero new buffers (per-slot threshold computed inline). Single launch covers all 24 slots; null-pointer guard handles slots 27/28/31/33-35/39/43 (deferred or IQN-inactive) without per-step Rust branching. `cargo check -p ml --lib` clean.
SP3 Task B7 — name-table entries for slots 36-47 readback log (2026-04-30): replaced the SP1 Phase B placeholder strings (`"rsv36"`-`"rsv47"`) in both `training_loop.rs` `let names = [...]` tables (the `halt_nan` block ~L1992 and the `halt_grad_collapse` block ~L2089) with the SP3 Mech 5 diagnostic slot names per the B1/B6 accessor allocation: 36 `trunk_adam_m_max`, 37 `value_adam_m_max`, 38 `branch_adam_m_max`, 39 `iqn_adam_m_max`, 40 `trunk_adam_v_max`, 41 `value_adam_v_max`, 42 `branch_adam_v_max`, 43 `iqn_adam_v_max`, 44 `trunk_weight_max`, 45 `heads_weight_max`, 46 `target_q_post_clip`, 47 `atom_span_max`. Both name tables receive byte-identical replacement content (modulo the indentation difference between the two enclosing blocks) per `feedback_no_partial_refactor` — the post-SP2 stale-doc cleanup commit `387335e2b` already established that the two tables must remain identical, and the same shared-contract migration principle applies here. When the fused kernel (B6) sets a slot bit, the readback log line names the buffer that exceeded its ISV-derived threshold (e.g. `flagged=[42=branch_adam_v_max, 46=target_q_post_clip]`) instead of the opaque `rsv*` placeholder, providing direct observability for SP3 mechanism effectiveness. Pure name-string replacement — no logic changes, no kernel changes, no buffer changes, no ISV changes. `cargo check -p ml --lib` clean.
SP3 Mech 6 — anchored upper bound on adaptive grad clip (2026-04-29): added an upper bound to `GpuDqnTrainer::update_adaptive_clip`'s `new_clip` formula in `gpu_dqn_trainer.rs`. The existing winsorizer (Plan C T11 follow-up N) caps a SINGLE input sample at `K=100 × prev_clip` before the EMA absorbs it, but does NOT prevent CONSECUTIVE elevated samples from compounding the EMA upward without bound. Over hundreds of steps the clip threshold ratchets to thousands while actual `grad_norm` tracks it from below — clipping becomes a no-op against in-distribution drift, Adam m/v EMAs are poisoned, and they saturate at the SP3 Mech 5 slot 36-43 thresholds. This is the F1 NaN root cause from `smoke-test-5rqzs` (commit `b9edccfc1`) at step 3060: Mech 5 diagnostic flags `[36=trunk_adam_m_max, 37=value_adam_m_max, 38=branch_adam_m_max, 40=trunk_adam_v_max, 41=value_adam_v_max, 42=branch_adam_v_max]` fired with target_q + atoms bounded (Mechs 1+2 working) and weights still finite — narrowing the divergence to the Adam state itself. **Bound formula**: `upper_bound = (grad_norm_slow_ema × 100 × ISV[Q_ABS_REF=16].max(1.0)).max(MIN_CLIP=1.0)`; final `new_clip = (grad_norm_ema × CLIP_MULTIPLIER).max(MIN_CLIP).min(upper_bound)`. **Anchor**: `grad_norm_slow_ema` is the existing α=0.001 slow-EMA scalar (mapped-pinned, updated later in this same function via `*self.grad_norm_slow_ema_pinned`) — read from pinned memory BEFORE the slow-EMA update on this step, so it reflects the previous step's slow EMA (the legitimate steady-state grad norm at the time of clip computation). **Headroom 100×**: legitimate per-step deviations can be 10-100× the slow average without being pathological, so `100 ×` keeps Mech 6 invisible in normal training and only kicks in when the EMA-driven clip ratchets past plausible-deviation bounds. **ISV-adaptive multiplier**: `ISV[Q_ABS_REF_INDEX = 16].max(1.0_f32)` scales the cap with the Q-magnitude regime per `feedback_isv_for_adaptive_bounds` — same ISV slot used by SP3 Mechs 1, 2, 3, and 5 (no new slot). ε on the multiplier (`.max(1.0)`) per the SP1 ε-floor pearl — cold-start `ISV[16] ≈ 0` would otherwise collapse `upper_bound` toward zero. **ε-floor on the bound itself** (`.max(MIN_CLIP=1.0)`): cold-start `grad_norm_slow_ema ≈ 0` (first ~200 steps before the slow EMA warms up) would otherwise pin `upper_bound` at 0, which combined with `new_clip.min(upper_bound)` would force `new_clip = MIN_CLIP=1.0` every step until the slow EMA established a meaningful baseline — exactly the cold-start ratcheting the upper bound is meant to PREVENT. The `MIN_CLIP` floor on the bound aligns with the existing `MIN_CLIP` floor on `new_clip` itself, so the bound is at minimum a no-op (matching the floor below) until the slow EMA warms up. **What stays unchanged**: the existing winsorizer (single-sample input cap before EMA update), the `EMA_BETA = 0.95` adaptive_clip EMA, the `CLIP_MULTIPLIER = 2.0` and `MIN_CLIP = 1.0` constants, the `grad_norm_fast_ema` / `grad_norm_slow_ema` updates further down in the same function (still receive the RAW `observed_grad_norm` per follow-up K's "fast/slow EMAs are stability signal that should respond to outliers" rationale), the `fold_warmup_factor_update` kernel that consumes the fast/slow EMAs, and every ISV slot. Pure formula change in one function — no new buffer, no new kernel, no new ISV slot, no new launch site, no graph recapture. F0 risk is low: F0-typical `grad_norm_slow_ema` is on the order of 1-10 → `upper_bound = 100-1000 × ISV[16].max(1)`; F0-typical `new_clip` (= `grad_norm_ema × 2`) is single-digits to low tens, well below the cap, so Mech 6 is invisible in normal F0 operation. F1+F2 benefit by preventing the EMA-ratchet pathology. Composes with: Mech 1 (target_q clipping at single-point source, B2), Mech 2 (atom-position growth bounds, B3), Mech 3 (CQL aux loss clamping, deferred), Mech 4 (Adam EMA reset comprehensive, B5), Mech 5 (fused threshold-check kernel for diagnostic slots 36-47, B6+B7). Zero new HtoD/DtoD/HtoH copies; zero new ISV slots; zero new buffers; zero new kernels. `cargo check -p ml --lib` clean.