feat(sp14-c.5a): allocate aux trunk fwd/bwd buffers + Adam launcher (dead code)
Phase C.5a — additive infrastructure for the aux trunk wire-up. Allocates saved-fwd buffers (h_s2_aux, h_aux1, h_aux2), gradient buffers (6× aux_trunk_*_grad), accumulator (dh_s2_aux_accum), and dedicated Adam launcher (launch_aux_trunk_adam_update) reading β1/β2/ε/LR/grad-clip from ISV[444..449). No contract change. No call sites for the new launcher yet. C.5b atomically wires these in. Phase C.5 was split (authorized 2026-05-08) after the original implementer flagged ~600-800 LOC scope across 4 files with correctness windows. C.5a is purely additive; C.5b is the genuine ~300 LOC atomic migration. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -7311,3 +7311,102 @@ Both slots use Pearl-A first-observation bootstrap. Slot 451 has α=0.05 EMA; sl
|
||||
**Open follow-ups:**
|
||||
- Add a target-variance EMA for `AVG_WIN_HOLD_TIME_BARS_INDEX` (Welford-style, single new ISV slot) in a future phase if the fixed α=0.01 blend is too slow/fast on real-data validation. The producer kernel already has the variance-slot path wired; pass `isv_h_target_var_idx >= 0` to enable Wiener-α.
|
||||
- Pearl note for `pearl_event_driven_reward_density_alignment.md` cross-reference: this phase aligns the AUX *label density* to event density (winning trade close), the same principle as the SP12 reward-density alignment.
|
||||
|
||||
## SP14 Layer C Phase C.5a — additive infrastructure (DEAD CODE) (2026-05-08)
|
||||
|
||||
> Phase C.5 was split (authorized 2026-05-08) after the prior implementer flagged ~600-800 LOC scope across 4 files with three correctness windows (gradient pointer aliasing, kernel rename + revert atomicity, Adam launcher dependency on grad buffer layout). C.5a is purely additive — buffers + launcher allocated as dead code. Build is clean before, build is clean after. C.5b atomically wires every consumer per `feedback_no_partial_refactor`.
|
||||
|
||||
**No contract change in this commit.** Buffers exist, launcher compiles, no production caller invokes either.
|
||||
|
||||
### Buffers allocated
|
||||
|
||||
In `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`:
|
||||
|
||||
**Saved-fwd tiles (replay-batch-sized):**
|
||||
- `h_s2_aux: CudaSlice<f32>` — `[B, SH2]` final aux trunk output. Shape matches `aux_dh_s2_nb_buf` so the C.5b SAXPY-back step type-matches.
|
||||
- `h_aux1: CudaSlice<f32>` — `[B, AUX_TRUNK_H1=256]` Layer-1 post-ELU activation.
|
||||
- `h_aux2: CudaSlice<f32>` — `[B, AUX_TRUNK_H2=128]` Layer-2 post-ELU activation.
|
||||
|
||||
**Upstream-grad accumulator:**
|
||||
- `dh_s2_aux_accum: CudaSlice<f32>` — `[B, SH2]`. Mirrors `bw_d_h_s2`'s role for Q's path. C.5b reverts the zero-fills at `aux_next_bar_backward:599-638` + `aux_regime_backward:757-790` (commits `872bd7392` + `411a30473`) so the two aux head backwards SAXPY their per-sample dh into THIS buffer instead of the trunk's `bw_d_h_s2`.
|
||||
|
||||
**Aux trunk gradient tensors (6, mirroring C.2 params):**
|
||||
- `aux_trunk_w1_grad` `[SH1 × 256]`, `aux_trunk_b1_grad` `[256]`,
|
||||
- `aux_trunk_w2_grad` `[256 × 128]`, `aux_trunk_b2_grad` `[128]`,
|
||||
- `aux_trunk_w3_grad` `[128 × SH2]`, `aux_trunk_b3_grad` `[SH2]`.
|
||||
- All `alloc_zeros`-initialised. Overwritten (not accumulated) per training step by `AuxTrunkBackwardOps::launch` in C.5b — matching the documented backward-API contract.
|
||||
|
||||
**Adam launcher scratch + ISV-driven hparams:**
|
||||
- `aux_trunk_grad_norm_partials: CudaSlice<f32>` `[total_blocks]` — per-block sum-of-squares slots laid out per `aux_trunk_grad_block_offsets[7]` (one offset per tensor + total).
|
||||
- `aux_trunk_grad_norm_buf: CudaSlice<f32>` `[1]` — finalised L2 norm read by every Adam launch as `grad_norm_f32`.
|
||||
- `aux_trunk_grad_clip_pinned: *mut f32` (mapped device-pinned) + `aux_trunk_grad_clip_dev_ptr` — host writes `ISV[AUX_TRUNK_GRAD_CLIP_INDEX=448]` per-step.
|
||||
- `aux_trunk_lr_pinned: *mut f32` + `aux_trunk_lr_dev_ptr` — host writes `ISV[AUX_TRUNK_LR_INDEX=444]` per-step.
|
||||
- `aux_trunk_t_pinned: *mut i32` + `aux_trunk_t_dev_ptr` — host increments before each launch.
|
||||
- `aux_trunk_adam_step: i32` — host-side mirror.
|
||||
- `Drop` impl frees the three new pinned host allocations.
|
||||
|
||||
In `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`:
|
||||
- `exp_h_s2_aux: CudaSlice<f32>` — `[alloc_episodes × shared_h2]`.
|
||||
- `exp_h_aux1: CudaSlice<f32>` — `[alloc_episodes × AUX_TRUNK_H1=256]`.
|
||||
- `exp_h_aux2: CudaSlice<f32>` — `[alloc_episodes × AUX_TRUNK_H2=128]`.
|
||||
|
||||
Mirrors the trainer's saved-fwd buffers but rollout-sized — same β-migration pattern as `exp_aux_nb_hidden_buf`.
|
||||
|
||||
### Adam launcher signature
|
||||
|
||||
```rust
|
||||
pub(crate) fn launch_aux_trunk_adam_update(
|
||||
&self,
|
||||
stream: &Arc<CudaStream>,
|
||||
_batch_size: i32,
|
||||
) -> Result<(), MLError>
|
||||
```
|
||||
|
||||
Writes 6 per-tensor `dqn_grad_norm_kernel` partials → 1 `dqn_grad_norm_finalize` (sqrt of sum across all 132K params) → 6 `dqn_adam_update_kernel` launches consuming the same global L2 norm. ISV-driven β1 (slot 445) / β2 (446) / ε (447) with numerical-stability floors per `pearl_first_observation_bootstrap`. SP4 Mech 9 weight clamp + Pearl C engagement counter disabled (aux trunk is outside the SP4 8-group taxonomy — same convention as OFI embed Adam).
|
||||
|
||||
Mirrored launcher: `step_ofi_embed_adam` (single-flat-buffer Adam path). Adapted to 6 separate (params, grad, m, v) tensors with byte-offset-free layout.
|
||||
|
||||
### Choice: 6-buffer (vs flat-buffer)
|
||||
|
||||
Selected **6-buffer with 6 separate Adam launches** over flat-buffer with byte-offset views. Rationale:
|
||||
|
||||
1. C.2 already allocated 6 separate param tensors (`aux_trunk_{w1,b1,w2,b2,w3,b3}`) and 12 separate Adam m/v tensors — a flat-buffer Adam would have required restructuring C.2's allocations to be views into a flat tensor, which is a contract change C.5a explicitly avoids.
|
||||
2. The existing `dqn_adam_update_kernel` works on a single contiguous (params, grads, m, v) tuple per launch. 6 launches consuming 6 tuples is a smaller Δ than refactoring C.2 to a flat layout.
|
||||
3. Global L2-norm clipping is preserved: the kernel reads `*grad_norm_f32` on every launch and computes its own per-launch `clip_scale = (norm > clip) ? clip/norm : 1`. All 6 launches see the same global norm + same clip threshold → consistent global clipping.
|
||||
4. Storage cost is identical (same total f32 footprint).
|
||||
5. Per-launch overhead (≤6 launches × ~1µs ≈ ≤6µs/step) is negligible relative to the trainer step.
|
||||
|
||||
The 7th entry in `aux_trunk_grad_block_offsets` records the total block count, passed as `num_blocks` to the finalize kernel — capturing all 6 tensors' partial blocks in a single reduction.
|
||||
|
||||
### NO contract change
|
||||
|
||||
- `launch_aux_trunk_adam_update` is `#[allow(dead_code)]` because no production caller invokes it in C.5a. C.5b adds the call site atomically with the kernel rename + zero-fill revert.
|
||||
- The 6 grad buffers, 4 saved-fwd tiles, accumulator, partials, norm buf, and pinned LR/clip/t are all `#[allow(dead_code)]` for the same reason.
|
||||
- `aux_trunk_adam_step` initialises to 0 (counter is incremented by C.5b's call site).
|
||||
|
||||
### Compiler warnings accepted (no `#[allow(dead_code)]` workaround)
|
||||
|
||||
`cargo check -p ml --tests --all-targets` is clean. The new fields are `#[allow(dead_code)]`-tagged on the struct definitions because every struct field with no consumer otherwise warns; tagging them inline at the field is the canonical pattern this codebase uses for additive fields awaiting a wire-up commit (matches `aux_trunk_forward_ops`'s convention from C.3).
|
||||
|
||||
### Files modified (this commit)
|
||||
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (+11 fields on `GpuDqnTrainer`; +allocations after C.2's site; +3 pinned host pointers freed in `Drop`; +`launch_aux_trunk_adam_update` method)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` (+3 fields on `GpuExperienceCollector`; +allocations after `exp_aux_dir_acc_buf`'s site)
|
||||
- Modify: `docs/dqn-wire-up-audit.md` (this section)
|
||||
|
||||
### Verification
|
||||
|
||||
- `SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo check -p ml --tests --all-targets` — clean (only pre-existing warnings).
|
||||
- `cargo test -p ml --test aux_trunk_oracle_tests --test sp14_oracle_tests --release -- --ignored --nocapture` — 12/12 pass (8 aux_trunk + 4 sp14, no regression).
|
||||
- `grep -rn "launch_aux_trunk_adam_update" crates/ml/src/` — only the definition appears, no callers.
|
||||
|
||||
### Follow-up: C.5b
|
||||
|
||||
C.5b is the genuine atomic contract migration:
|
||||
1. Insert `aux_trunk_forward.launch(...)` post-`forward_online_f32` in collector + post-`encoder_forward_chain` in trainer.
|
||||
2. Redirect `aux_heads_fwd` input pointer from Q's `h_s2` to `h_s2_aux`.
|
||||
3. Revert zero-fills at `aux_next_bar_backward:599-638` + `aux_regime_backward:757-790` so the SAXPY-into-`dh_s2_aux_accum` path works.
|
||||
4. Param-rename `h_s2 → h_s2_aux` in the 4 aux_heads kernels.
|
||||
5. Insert `aux_trunk_backward.launch(...)` + `launch_aux_trunk_adam_update(...)` calls.
|
||||
6. Initialize `aux_trunk_lr_pinned` + `aux_trunk_grad_clip_pinned` from ISV per-step (host writes before each launch).
|
||||
7. Increment `aux_trunk_adam_step` per Adam call.
|
||||
|
||||
Reference in New Issue
Block a user