diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 625b82563..1dc6a0c43 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -4273,6 +4273,14 @@ impl GpuExperienceCollector { &self.curiosity_weights } + /// Access the GPU curiosity trainer (grad + Adam state). Returns `None` + /// when curiosity is disabled at collector init or after `disable_curiosity`. + /// SP4 Layer A param-group oracle wiring pairs this with `curiosity_weight_set` + /// to populate the curiosity sub-buffer descriptor in `SP4AuxBuffers`. + pub fn curiosity_trainer(&self) -> Option<&GpuCuriosityTrainer> { + self.curiosity_trainer.as_ref() + } + /// Disable curiosity training (e.g. after an async kernel crash). pub fn disable_curiosity(&mut self) { self.curiosity_trainer = None; diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index a6ecfd7bc..d3b29d89a 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -3498,6 +3498,56 @@ impl DQNTrainer { tracing::warn!("launch_moe_lambda_eff_update failed: {e}"); } } + + // SP4 Layer A Tasks A10 + A11: per-step ISV-bound producer + // launches for the 40 new SP4 slots (TARGET_Q_BOUND, ATOM_POS + // ×4 branches, GRAD_CLIP_BOUND, H_S2_BOUND, and 8 param-group + // oracles spanning WEIGHT/ADAM_M/ADAM_V/WD_RATE per group plus + // L1_LAMBDA[trunk]). Layer A is observability-only — no + // consumer reads these slots yet, so behavior is unchanged. + // All five launches go cold-path here next to the rest of + // the per-step ISV producers; Layer B may relocate the + // captured-graph eligible ones (target_q_p99 in particular) + // once the Mech 1 pre-clamp consumer lands. + // + // Producer ordering — like the surrounding ISV producers, + // these must run BEFORE the HEALTH_DIAG line emit (commit + // a5f23b28f post-cascade-fix invariant) and BEFORE any + // future consumer kernel reads the bound slots. + if let Some(ref fused) = self.fused_ctx { + if let Err(e) = fused.trainer().launch_sp4_target_q_p99() { + tracing::warn!("SP4 A10 target_q_p99 launch failed: {e}"); + } + if let Err(e) = fused.trainer().launch_sp4_atom_pos_p99_all_branches() { + tracing::warn!("SP4 A10 atom_pos_p99_all_branches launch failed: {e}"); + } + if let Err(e) = fused.trainer().launch_sp4_grad_norm_p99() { + tracing::warn!("SP4 A10 grad_norm_p99 launch failed: {e}"); + } + if let Err(e) = fused.trainer().launch_sp4_h_s2_p99() { + tracing::warn!("SP4 A10 h_s2_p99 launch failed: {e}"); + } + + // SP4 A11: param_group_oracle Pearl B over 8 groups + // (DQN trunk + 4 branches groups 0-2; IQN, IQL hi/lo, + // Attn, Curiosity groups 3-7). Curiosity (group 7) + // descriptor is built from `gpu_experience_collector`'s + // `curiosity_weight_set` + `curiosity_trainer` pair — + // both `None` together (curiosity disabled at collector + // init) yields an empty group descriptor and the + // launcher silently skips it. + let curiosity_weights = self.gpu_experience_collector.as_ref() + .map(|c| c.curiosity_weight_set()); + let curiosity_trainer = self.gpu_experience_collector.as_ref() + .and_then(|c| c.curiosity_trainer()); + let aux_buffers = fused.build_sp4_aux_buffers( + curiosity_weights, + curiosity_trainer, + ); + if let Err(e) = fused.trainer().launch_sp4_param_group_oracles_all_groups(&aux_buffers) { + tracing::warn!("SP4 A11 param_group_oracles_all_groups launch failed: {e}"); + } + } } // Plan 4 Task 6 Commit B: refresh the aux-loss weight before the diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 5f4a5665d..58d38dad9 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2267,6 +2267,8 @@ F0 risk: low — guards activate at 1e6× ISV magnitude (several orders of magni Permanent diagnostic infrastructure (Phase B slots 24-35) stays as regression sentinel — will detect any future regression of the cuBLAS-overflow NaN class. Combined RELATED commit per `feedback_no_partial_refactor` — both kernels patched in one commit since they share the same unsafe-pattern (cuBLAS GEMM extreme-intermediate-product overflow) + the same fix template. +SP4 Layer A Tasks A10 + A11 — wire all 5 SP4 producer launches (2026-04-30): adds the producer-launch call sites in `crates/ml/src/trainers/dqn/trainer/training_loop.rs` for the 5 SP4 producers built in Tasks A5-A9 (with A7 fix-up + fix-up #2 closing Curiosity). Five new launches added cold-path next to the existing per-step ISV producers (`launch_h_s2_rms_ema`, `launch_fold_warmup_factor`, `launch_iqn_quantile_ema`, `launch_vsn_mask_ema`, `launch_aux_heads_loss_ema`, `launch_moe_expert_util_ema`, `launch_moe_lambda_eff_update`): (1) `launch_sp4_target_q_p99` → ISV[TARGET_Q_BOUND=131], (2) `launch_sp4_atom_pos_p99_all_branches` → ISV[ATOM_POS_BOUND[0..4]=132..135], (3) `launch_sp4_grad_norm_p99` → ISV[GRAD_CLIP_BOUND=168], (4) `launch_sp4_h_s2_p99` → ISV[H_S2_BOUND=169], and (5) `launch_sp4_param_group_oracles_all_groups` → ISV[WEIGHT_BOUND/ADAM_M_BOUND/ADAM_V_BOUND/WD_RATE × 8 + L1_LAMBDA_TRUNK = 33 slots]. All five fire ONCE per training step; cold-path placement matches the surrounding ISV-producer block which the post-cascade-fix invariant (`a5f23b28f`) requires to run BEFORE the HEALTH_DIAG line emit. **For the param-group oracle**, the `SP4AuxBuffers` descriptor is built each step via `FusedTrainingCtx::build_sp4_aux_buffers(curiosity_weights, curiosity_trainer)` (added by A7 fix-up #2), threading curiosity state from `gpu_experience_collector`'s `curiosity_weight_set()` + new `curiosity_trainer()` accessor. **New accessor** in `gpu_experience_collector.rs`: `pub fn curiosity_trainer(&self) -> Option<&GpuCuriosityTrainer>` — mirrors the existing `has_curiosity_trainer`/`curiosity_weight_set` pair. When the collector is absent (init-failure graceful-degrade) or curiosity is disabled (`disable_curiosity()` was called after async kernel crash), the curiosity descriptor empties and the launcher's `count == 0` short-circuit silently skips group 7; groups 0-6 still produce real outputs. **Layer A semantics — observability only**: no consumer reads any of the 40 SP4 ISV slots yet. Mech 1 (`target_q` clamp) still uses `±10 × ISV[Q_ABS_REF=16].max(1.0)`; Mech 2 (`atoms_update_kernel` clamp) likewise; Mech 6 (`adaptive_clip` upper bound) still uses the existing fast/slow grad-norm EMA path; Mech 9 (`weight_clamp_max_abs`) still uses `100 × q_abs_ref_eff` config arg; Mech 10 (`h_s2` clamp) uses the existing slot-96 RMS EMA. Behavior unchanged from before this commit. **Producer ordering**: each launch uses the same `if let Some(ref fused) = self.fused_ctx { ... tracing::warn! on Err }` guard pattern as `launch_h_s2_rms_ema` — errors never propagate (warn-and-continue) since Layer A is observability-only. Launches are sequenced after `launch_moe_lambda_eff_update` and before the closing `}` of the per-step ISV-producer block, ensuring all 40 slots are populated before HEALTH_DIAG reads ISV state. **Layer B path-forward**: the cold-path placement is a Layer A choice — once Mech 1's pre-clamp consumer lands, Layer B may relocate `target_q_p99` (and possibly other captured-graph-eligible producers) into the captured graph for tighter ordering. The 5 producers' Pearls A+D (Task A3) host-side state and Pearl B's per-launch table-packing sync are unaffected by relocation. **No new ISV slots, no new HtoD/DtoD/HtoH copies, no behavioural change.** `cargo check -p ml --lib --tests` clean (11 pre-existing warnings; one warning previously reported as 12 has resolved). `state_reset_registry` test suite passes (3/3). Files touched: `crates/ml/src/trainers/dqn/trainer/training_loop.rs` (+~50 LoC five SP4 launches + curiosity-arg threading next to `launch_moe_lambda_eff_update`), `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` (+~8 LoC `curiosity_trainer()` accessor). + SP4 Layer A Task A7 fix-up #2 (2026-04-30): closes the Curiosity hold-out from fix-up #1 — extends the `param_group_oracle_kernel` to accept an array of sub-buffer pointers + counts and iterate them within each pass. Groups 0-6 launch with `n_sub=1` (behaviour identical to the original single-pointer signature); group 7 (Curiosity) launches with `n_sub=4` describing `[w1, b1, w2, b2]`, and the kernel treats the union as one logical group for p99/WD_RATE. Pass E (L1 trunk lambda) is dispatched only for group 0 which is always single-sub-buffer. **Kernel signature** (`param_group_oracle_kernel.cu`) changed from `(const float* params, const float* grads, const float* adam_m, const float* adam_v, int count, ...)` to `(const u64* params_ptrs, const u64* grads_ptrs, const u64* adam_m_ptrs, const u64* adam_v_ptrs, const i32* sub_counts, int n_sub, int total_count, ...)`. **Histogram device fn** (`sp4_histogram_p99.cuh`) gains a sibling `sp4_histogram_p99_multi(sub_buf_ptrs, sub_counts, n_sub, total_count)` that mirrors the original three-pass structure but iterates sub-buffers within Pass 1 (max-reduce) and Pass 2 (binning); Pass 3 (cumulative-from-top) is unchanged — divides by `total_count`. The single-buffer template stays for tests/single-buffer producers (A4/A5/A6/A8/A9). **Pass D** in the oracle kernel iterates sub-buffers in the per-thread accumulator loop; `inv_count = 1/total_count`. Pass E reads the contiguous trunk grads sub-buffer at `grads_ptrs[0]` (group 0 has `n_sub==1` so the indexed access is canonical). **Sub-buffer table** lives in two persistent mapped-pinned buffers allocated at construction: `oracle_subbuf_table_buf: MappedU64Buffer` of `4 × SP4_ORACLE_TABLE_MAX_SUB = 16` u64s (params/grads/adam_m/adam_v ptr-arrays packed contiguously) and `oracle_subbuf_counts_buf: MappedI32Buffer` of `SP4_ORACLE_TABLE_MAX_SUB = 4` i32s. The launcher overwrites entries `[0..n_sub)` per group launch (`std::ptr::write_volatile`) and zeros the unused tail (defence-in-depth so a kernel bug reading past `n_sub` lands on count=0 no-op). **Inter-launch sync** added (one `stream.synchronize()` per group launch, before the next host overwrite races with in-flight kernel coalesced loads). Cold-path producer; per-launch sync cost is negligible vs the kernel work. The end-of-loop sync is retained for symmetry with the Phase 2 contract. **`Sp4ParamGroupBufs` redefined** in `gpu_dqn_trainer.rs` from a flat `(params_ptr, grads_ptr, adam_m_ptr, adam_v_ptr, count)` quartet to a `Vec` (1 entry for groups 0-6, 4 entries for Curiosity). Convenience constructors `Sp4ParamGroupBufs::single(...)` and `Sp4ParamGroupBufs::empty()` keep the call sites in `param_group_buffers` and `build_sp4_aux_buffers` ergonomic. `total_count()` accessor for the kernel arg. Switched `Copy` derive to `Clone` since the descriptor now owns a `Vec`. **`SP4AuxBuffers` extended** with `curiosity: Sp4ParamGroupBufs` (was 4-tuple, now 5-tuple). **Accessors added** to `CuriosityWeightSet` (gpu_weights.rs): `w1_ptr/w1_len/b1_ptr/b1_len/w2_ptr/w2_len/b2_ptr/b2_len`. To `GpuCuriosityTrainer` (gpu_curiosity_trainer.rs): full set across grad + Adam state — `grad_w1/b1/w2/b2_ptr/_len`, `adam_m_w1/b1/w2/b2_ptr`, `adam_v_w1/b1/w2/b2_ptr`. **`build_sp4_aux_buffers` signature** changed from `&self -> SP4AuxBuffers` to `&self, curiosity_weights: Option<&CuriosityWeightSet>, curiosity_trainer: Option<&GpuCuriosityTrainer> -> SP4AuxBuffers` because Curiosity state lives outside `FusedTrainingCtx` (owned by `GpuExperienceCollector`). Layer B's training-loop caller threads them in from `collector.curiosity_weight_set()` + the collector's `curiosity_trainer` field; both must be `Some` together (caller responsibility — they live on the same collector) and the descriptor enumerates all four `[w1, b1, w2, b2]` sub-buffers. Passing `None` for either yields an empty curiosity descriptor and the launcher silently skips group 7. **`param_group_buffers` return type** changed from `Option<(u64, u64, u64, u64, usize, i32, i32)>` to `Option<(Sp4ParamGroupBufs, i32, i32)>`. All groups now return `Some(...)` (Curiosity included); `None` reserved for forward-compat. **GPU unit test** `sp4_param_group_oracle_per_group_writes_distinct_isv_slots` extended: group 7 now exercises 4 sub-buffers of distinct shapes `[1024, 32, 1024, 32]` (w1/b1/w2/b2-like sizes scaled to keep test runtime small while still exercising the multi-sub-buffer iteration), each with its own seed offset (`(s as u64) << 16` mixed into the per-buffer seed) so distinct sub-buffers have distinct distributions — catches buffer-mixup bugs in the kernel's sub-buffer iteration. Reference computation builds `union_*` vectors and computes p99/WD_RATE over the union; `union_total_count` replaces the old `N` constant for groups other than 0/1..6 (where `N=4096` is still the single-sub-buffer total). Existing tolerances unchanged (5% rel_err for p99, 2% for WD_RATE, 5% for L1). Test launcher refactored: takes `&[TestSubBuffer]` slice (replaces 4 separate host arrays), constructs the same mapped-pinned ptr-table layout used in production, packs `n_sub` entries per call. **All 8 SP4 param-groups now produce real outputs in Layer A.** The launcher's `count == 0` short-circuit is retained — `Sp4ParamGroupBufs::empty()` still exists for the optional-aux-trainer fallback (gpu_iqn / gpu_attention / curiosity init failures). `cargo check -p ml --lib --tests` clean. `MappedU64Buffer` gained a manual `Debug` impl (warn `missing_debug_implementations`) for parity with the existing `MappedU32Buffer`. Files touched: `sp4_histogram_p99.cuh` (+~80 LoC `_multi` template), `param_group_oracle_kernel.cu` (kernel signature + sub-buffer iteration loops + Pass E `grads_ptrs[0]` access), `gpu_dqn_trainer.rs` (`Sp4SubBuffer` type + `Sp4ParamGroupBufs::sub_buffers` reshape + `oracle_subbuf_table_buf` / `oracle_subbuf_counts_buf` fields + constructor allocs + struct init + launcher's table-packing/launch loop + per-launch sync + `param_group_buffers` return-type change), `gpu_weights.rs` (+~25 LoC Curiosity accessor block), `gpu_curiosity_trainer.rs` (+~45 LoC grad + Adam-state accessor block), `fused_training.rs` (`build_sp4_aux_buffers` signature change + Curiosity descriptor block, replacing the empty placeholder), `mapped_pinned.rs` (+10 LoC `MappedU64Buffer` Debug), `tests/sp4_producer_unit_tests.rs` (multi-sub-buffer test launcher + group-7 4-shape exercise + union reference computation). SP4 Layer A Task A7 fix-up (2026-04-30): wires 4 of the 5 deferred aux param-groups to the Pearl B per-param-group statistics oracle. A7 (commit `4f13e2ca3`) launched 3 of 8 groups (DqnTrunk/DqnValue/DqnBranches via main DQN params slicing); the 5 aux groups (Iqn/IqlHigh/IqlLow/Attn/Curiosity) returned `None` from `param_group_buffers` and silently skipped. Without this fix-up, Layer B's atomic flip would route IQN/IQL/Attn Adam clamps through `ISV[WEIGHT_BOUND[3..7)]` which would still be 0 → consumer `.max(EPS_CLAMP_FLOOR=1.0)` → 1.0 clamp → catastrophic over-clamp on aux-trainer params. Architectural decision per A7's DONE_WITH_CONCERNS report: chose Option 3 (thread aux-trainer buffers through the launcher's signature, `FusedTrainingCtx` supplies them) over hoisting the launcher onto `FusedTrainingCtx` — least invasive and matches existing patterns where `FusedTrainingCtx` orchestrates aux trainers. **Accessors added** mirroring SP3 close-out v2's IQN `online_params_ptr` template: IQN gained `online_grad_ptr/len` (the only one missing — `online_params_ptr/len`, `adam_m_ptr/len`, `adam_v_ptr/len` were already public from SP3 slot-48); `GpuIqlTrainer` gained the full `params_ptr/len`, `grads_ptr/len`, `adam_m_ptr/len`, `adam_v_ptr/len` set; `GpuAttention` gained the same set. **New types** in `gpu_dqn_trainer.rs`: `Sp4ParamGroupBufs { params_ptr, grads_ptr, adam_m_ptr, adam_v_ptr, count }` (one trainer's quartet, all four buffers same length per Adam-mirrors-params), `SP4AuxBuffers { iqn, iql_high, iql_low, attn }` (4-tuple of `Sp4ParamGroupBufs`). **Launcher signature** changed from `fn launch_sp4_param_group_oracles_all_groups(&self) -> Result<(), MLError>` to `fn (&self, aux_buffers: &SP4AuxBuffers) -> Result<(), MLError>`; `param_group_buffers(group, aux)` consults `aux.{iqn|iql_high|iql_low|attn}` for `ParamGroup::{Iqn|IqlHigh|IqlLow|Attn}` and the existing main-DQN slicing for groups 0-2. **`FusedTrainingCtx::build_sp4_aux_buffers()`** — anticipatory helper (Layer B will consume it; lint-suppressed until then) that constructs `SP4AuxBuffers` from `self.gpu_iqn`/`self.gpu_iql`/`self.gpu_iql_low`/`self.gpu_attention`. `gpu_iqn` and `gpu_attention` are still `Option<_>` (graceful-degrade init fallback); `None` produces a zero-count placeholder so the launcher's existing `count == 0 { continue; }` short-circuit silently skips the kernel launch — matches the existing skip-then-no-op pattern. **Curiosity is the architectural hold-out** documented in `param_group_buffers` doc-comment: `GpuCuriosityTrainer` stores params/grads/Adam state as four separate `[w1, b1, w2, b2]` sub-buffers (non-contiguous) — a single `(params_ptr, count)` tuple cannot describe the slice the kernel reads. Resolving this requires either a per-layer launch loop (4× the kernel cost) or a contiguous-flat re-layout of the trainer; both are deeper architectural changes scoped beyond Task A7's fix-up. Until then, `ParamGroup::Curiosity` still returns `None` from `param_group_buffers` and the launcher silently skips it; Layer B must guard against `ISV[WEIGHT_BOUND[7]=143]` / `ISV[ADAM_M_BOUND[7]=151]` / `ISV[ADAM_V_BOUND[7]=159]` / `ISV[WD_RATE[7]=167]` being the natural-zero floor for Curiosity-related clamps via `.max(EPS_CLAMP_FLOOR=1.0)`. **Test coverage** is unchanged — the existing kernel-direct unit test `sp4_param_group_oracle_per_group_writes_distinct_isv_slots` already iterates `g_idx ∈ 0..SP4_PARAM_GROUP_COUNT=8` with synthetic Box-Muller buffers, validating all 8 groups at the kernel level (the test does not call the production launcher; it exercises the kernel's per-group passes directly). The fix-up is therefore a wiring change with no test surface change. **No production callsite yet** — the launcher is still wired only as a public method on `GpuDqnTrainer`; Layer B's atomic-flip will add the call to the training-step pipeline. **Unchanged behaviour today** — the launcher is not invoked in production; the fix-up only changes its API to accept `&SP4AuxBuffers` so Layer B can integrate without further refactoring of the launcher. `cargo check -p ml --lib --tests` clean (12 pre-existing warnings, no new warnings). Files touched: `gpu_iqn_head.rs` (+~12 LoC accessor pair), `gpu_iql_trainer.rs` (+~50 LoC accessor block), `gpu_attention.rs` (+~50 LoC accessor block), `gpu_dqn_trainer.rs` (+~70 LoC `Sp4ParamGroupBufs`/`SP4AuxBuffers` definitions + launcher signature change + `param_group_buffers` aux-arm wiring), `fused_training.rs` (+~60 LoC `build_sp4_aux_buffers` helper).