From d9a4d98a3d9dd3d559eca83e251b60c77e957255 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 30 Apr 2026 11:27:54 +0200 Subject: [PATCH] =?UTF-8?q?feat(dqn):=20SP3=20Mech=207=20=E2=80=94=20per-e?= =?UTF-8?q?lement=20gradient=20clip=20in=20Adam=20kernel?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Smoke smoke-test-fxvkk (commit 48c25d999, Mech 6 anchored clip) F1-NaN'd at step 3720 with slots 36-38, 40-42 still firing — Adam EMAs saturated despite Mech 6 capping the global clip threshold. Diagnosis: Mech 6 bounds the AGGREGATE L2 gradient norm, but Adam m/v EMAs are PER-ELEMENT. A gradient with one large element (e.g., element_X = 100, rest small) has L2 norm ≈ 100, passing a clip threshold of 1000 untouched. Adam m_X EMA accumulates the large element; β1=0.9 steady-state gives m_X ≈ 1000, exceeding slot 36 threshold (100 × isv). Fix: in dqn_adam_update_kernel, after the global L2 clip is applied, clip EACH gradient element to ±per_element_cap where: per_element_cap = 10 × sqrt(adaptive_clip / total_params) Rationale: - sqrt(adaptive_clip / N) is the average per-element contribution to the L2 norm budget - 10× allows legitimate per-element deviations up to 10× average - Scales with adaptive_clip — when global clip is doing its job (Mech 6), per_element_cap is tight enough to prevent saturation - When global clip is loose (post-fold warmup), per_element_cap scales with it — never tighter than the global clip's intent Together with Mech 6, prevents Adam saturation at both the aggregate (L2 norm) and per-element levels. Closes the residual pathology after Mech 6's partial 3060→3720 improvement. --- .../src/cuda_pipeline/dqn_utility_kernels.cu | 21 +++++++++++++++++++ docs/dqn-wire-up-audit.md | 2 ++ 2 files changed, 23 insertions(+) diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index 546feee48..c245257c2 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -136,6 +136,27 @@ extern "C" __global__ void dqn_adam_update_kernel( float clip_scale_f = (norm_f > max_grad_norm) ? (max_grad_norm / norm_f) : 1.0f; float clipped_g_f = g_f * clip_scale_f; + /* SP3 Mech 7: per-element gradient clip — prevents Adam EMA saturation + * at single-element extreme gradients that pass the global L2 norm clip. + * Mech 6 (anchored upper-bound on adaptive_clip) bounds the AGGREGATE + * gradient magnitude, but a gradient with one large element + rest small + * can satisfy the L2 budget while poisoning the Adam m EMA at that + * element's index. Smoke smoke-test-fxvkk (commit 48c25d999) F1-NaN'd at + * step 3720 with slots 36-38, 40-42 firing despite Mech 6 — confirming + * per-element saturation is the residual pathology. + * + * per_element_cap = 10 × sqrt(adaptive_clip / total_params) + * - sqrt(adaptive_clip / N) is the AVERAGE per-element L2 budget + * (if all elements equal, each contributes sqrt(clip²/N) = clip/sqrt(N)) + * - 10× allows legitimate per-element deviations up to 10× average + * (some elements legitimately have larger gradients than others) + * - When the global clip is doing its job (most steps), per_element_cap + * is tight enough to prevent saturation + * - When global clip is loose (e.g., post-fold-warmup), per_element_cap + * scales with it — never tighter than the global clip's intent */ + const float per_element_cap = 10.0f * sqrtf(max_grad_norm / (float)total_params); + clipped_g_f = copysignf(fminf(fabsf(clipped_g_f), per_element_cap), clipped_g_f); + /* Adam bias correction: native float — no bf16 rounding issues. */ float beta1_t = fmaxf(1.0f - powf(beta1, (float)t), 1e-8f); float beta2_t = fmaxf(1.0f - powf(beta2, (float)t), 1e-8f); diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 34d6ea30d..38fa34be5 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2285,4 +2285,6 @@ SP3 Task B6 — fused kernel extended for slots 36-47 threshold checks (2026-04- 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 7 — per-element gradient clip in Adam kernel (2026-04-29): added a per-element gradient cap inside `dqn_adam_update_kernel` (`crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu`) immediately after the existing global L2 norm clip (`clipped_g_f = g_f * clip_scale_f`) and before the Adam m/v moment updates. **Why Mech 6 wasn't enough**: smoke `smoke-test-fxvkk` (commit `48c25d999`, Mech 6 anchored upper-bound clip) F1-NaN'd at step 3720 (vs step 3060 in prior smoke; Mech 6 helped marginally, ~22% later) with slots 36-38, 40-42 still firing — DQN main Adam m + v saturated despite Mech 6 capping the AGGREGATE clip threshold. The L2 norm is an aggregate quantity but Adam EMAs are PER-ELEMENT: a gradient with one large element (e.g. `element_X = 100`, rest small) has `||g||₂ ≈ 100` and passes a clip threshold of 1000 untouched, but the Adam m_X EMA accumulates the large element. With β1=0.9 steady-state, `m_X ≈ 100 / (1 - 0.9) = 1000`, exceeding slot 36's threshold (`100 × ISV[Q_ABS_REF].max(1.0)`). **Cap formula**: `per_element_cap = 10 × sqrt(max_grad_norm / total_params)` where `max_grad_norm` is the existing kernel arg sourced from `max_grad_norm_buf` (= the Mech-6-bounded `adaptive_clip` mapped-pinned scalar), and `total_params` is the existing kernel arg counting all DQN main params. **Rationale**: `sqrt(adaptive_clip / N)` is the AVERAGE per-element contribution to the L2 norm budget — if all elements equal in magnitude, each contributes `sqrt(clip² / N) = clip / sqrt(N)`, and the L2 norm equals `clip` exactly. The 10× multiplier allows legitimate per-element deviations up to 10× average (some elements legitimately have larger gradients than others — e.g. branch heads vs trunk parameters); together this means the per-element cap kicks in only for the genuinely outlying elements that would saturate Adam m EMAs. **Coupling with Mech 6**: the per-element cap scales with `adaptive_clip`, so when Mech 6's anchored bound keeps `adaptive_clip` reasonable (the typical case), `per_element_cap` is correspondingly tight (e.g. `adaptive_clip = 10`, `total_params ≈ 200k` → `per_element_cap = 10 × sqrt(10 / 200000) ≈ 0.0707` — a tight bound on per-element gradient contribution). When Mech 6's bound is loose (e.g. immediately post-fold-warmup before the slow-EMA stabilises), `per_element_cap` scales up with it — never tighter than the global clip's intent. **Implementation**: pure kernel modification using existing args; no new kernel arg, no new buffer, no new ISV slot, no new launch site, no graph recapture. The cap is applied via `clipped_g_f = copysignf(fminf(fabsf(clipped_g_f), per_element_cap), clipped_g_f)` — preserves sign, bounds magnitude. **What stays unchanged**: the global L2 norm clip itself (computed identically), Adam bias correction, m/v EMA β1/β2 coefficients, AdamW decoupled weight decay, L1 proximal step on `w_s1`, the NaN/Inf gradient skip guard, every kernel launch site and graph layout. 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), Mech 6 (anchored upper bound on aggregate adaptive grad clip). Mechs 6+7 together prevent Adam saturation at BOTH the aggregate (L2 norm) and per-element levels — closes the residual pathology after Mech 6's partial 3060→3720 improvement. F0 risk is low: F0-typical `adaptive_clip` is on the order of 5-20 with `total_params ≈ 200k`, giving `per_element_cap ≈ 10 × sqrt(10/200000) = 0.0707`, well above F0-typical per-element gradients (~1e-3 to 1e-2 for a converged trunk on stable input), so Mech 7 is invisible in normal F0 operation. F1+F2 benefit by preventing single-element Adam EMA poisoning. Zero new HtoD/DtoD/HtoH copies; zero new ISV slots; zero new buffers; zero new kernels. `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.