feat(sp20): allocate trainer aux_conf_at_state_buf + wire PER direct-gather

Phase 5 plumbing — consumer-side wire-up. Allocates `aux_conf_at_state_buf:
CudaSlice<f32>` ([batch_size]) on `GpuDqnTrainer`, exposes the raw_ptr via
`aux_conf_at_state_buf_ptr()` and the `FusedTrainerCtx` delegating accessor
`trainer_aux_conf_at_state_buf_ptr()`, and invokes
`GpuReplayBuffer::set_trainer_aux_conf_ptr` at both fused_ctx init sites in
training_loop.rs (init + re-init, atomic per `feedback_no_partial_refactor`).

The PER `gather_f32_scalar` now writes the SAMPLED bar's per-batch aux_conf
directly into the trainer's f32 buffer on every step — same direct-to-trainer
pattern as the SP13 B1.1b `aux_nb_label_buf` (i32) wire-up immediately above
the new call. The c51_loss_batched reward gate (lands in the next commit)
reads this buffer to compute `gate = sigmoid((aux_conf - ISV[AUX_CONF_THRESHOLD])
/ ISV[AUX_GATE_TEMP])` and applies `r_used = gate * reward` at the Bellman
projection.

`alloc_zeros` cold-start: 0.0 sentinel → at threshold ≈ 0.10 and temp ≈ 0.05
the gate is `sigmoid(-2) ≈ 0.12` → reward is mostly suppressed pre-population.
This is the "graceful degradation" semantic from the Phase 3 Task 3.4 audit
doc spec §4.4. Once PER's direct-gather populates from the producer ring on
the first sample step, the per-bar aux_conf values drive the gate as designed.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-10 14:43:31 +02:00
parent d1a8ec206f
commit d3a057af4f
4 changed files with 156 additions and 0 deletions

View File

@@ -2,6 +2,89 @@
**Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7.
## 2026-05-10 — SP20 Phase 5 (commit 1/4): trainer aux_conf_at_state_buf + PER direct-gather wire-up
Phase 5 consumer-side plumbing — atomic across all struct boundaries per
`feedback_no_partial_refactor`. Lands the trainer-side buffer the c51_loss
reward-gate kernel will consume in commit 2/4.
### Trainer-side allocation
`GpuDqnTrainer.aux_conf_at_state_buf: CudaSlice<f32>``[batch_size]` f32.
Allocated via `alloc_f32(&stream, aux_b, "aux_conf_at_state_buf")` in the
constructor, alongside the existing aux-heads buffers. Sentinel: `0.0`
(alloc_zeros) — matches the K=2 uniform-prior default the Phase 3 Task 3.4
producer writes when `aux_logits_per_env == NULL`.
### Trainer-side accessor
`pub(crate) fn aux_conf_at_state_buf_ptr(&self) -> u64` — mirrors the
`aux_nb_label_buf_ptr` accessor used by SP13 B1.1b's i32 direct-gather. The
wrapping `FusedTrainerCtx::trainer_aux_conf_at_state_buf_ptr` delegates to
the inner trainer accessor (same delegation pattern as the 6 sibling getters).
### Replay-buffer wire-up (already plumbed in Phase 3 Task 3.4)
`GpuReplayBuffer::set_trainer_aux_conf_ptr(u64)` was added in Phase 3 Task 3.4
(comment at `gpu_replay_buffer.rs:1163-1165` explicitly says: "Phase 5 will
call this once after the trainer's `aux_conf_at_state_buf` is allocated").
This commit calls it.
### Training loop wire-up (atomic — both fused_ctx init sites)
`crates/ml/src/trainers/dqn/trainer/training_loop.rs` — two call sites mirror
the existing 6-arg `set_trainer_buffers` block:
- Line ~826 (init site): `agent_w.memory_mut().gpu.set_trainer_aux_conf_ptr(
fused.trainer_aux_conf_at_state_buf_ptr())` immediately after
`set_trainer_buffers(...)`.
- Line ~2645 (re-init site after fused_ctx tear-down): same call. Both
sites stay in sync per `feedback_no_partial_refactor`.
### Default-state semantics ("graceful degradation")
`alloc_zeros` cold-start: 0.0 sentinel. The Phase 5 gate kernel (commit 2/4)
will compute:
```
gate = sigmoid((aux_conf - threshold) / temp)
= sigmoid((0 - 0.10) / 0.05) // typical ISV values per Phase 1.3
= sigmoid(-2)
≈ 0.12
```
so `r_used = 0.12 × reward` ≈ "reward mostly suppressed". This is the
"uncertain-state neutralizer" semantic from Phase 3 Task 3.4 audit doc
spec §4.4 — pre-population (no aux signal yet), the model gets minimal
reward feedback at the Bellman target → Q collapses to `γ × Q(s', a')`,
i.e., "don't update Q on uncertain transitions." Once PER's
`gather_f32_scalar` populates from the producer ring on the first sample
step, the per-bar aux_conf values drive the gate as designed.
### Files modified
| File | Status | Purpose |
|------|--------|---------|
| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | +field, +alloc, +accessor | Trainer-side buffer ownership |
| `crates/ml/src/trainers/dqn/fused_training.rs` | +delegating accessor | FusedTrainerCtx wrapper |
| `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | +2 call sites | PER direct-gather wire-up |
### Verification
```
SQLX_OFFLINE=true cargo check --workspace --examples # green
```
Forward-references (commits 2-4 of this Phase 5 lands in this branch):
- Commit 2/4: `c51_loss_kernel.cu` adds `aux_conf_at_state` arg + gate
computation at the Bellman projection. NULL-tolerant (gate=1.0
no-op).
- Commit 3/4: launcher in `gpu_dqn_trainer.rs::launch_c51_loss` threads
`self.aux_conf_at_state_buf.raw_ptr()` as the new last arg; unit
tests for the gate formula.
- Commit 4/4: audit-doc consolidation + close-out.
## 2026-05-10 — MAX_UPLOAD_BYTES 2GB → 8GB
`crates/ml/src/cuda_pipeline/mod.rs:136`. L40S (48GB) and H100 (80GB) have plenty of headroom; the 2GB cap was a conservative leftover that tripped on workflow `zgjgc` with 17.8M imbalance bars (3.2GB needed). 8GB accommodates richer bar densities while leaving ~28GB free on L40S after model + activations + workspace. Updates both `DqnGpuData::upload_slices` (training data) and `PpoGpuData::upload` (market data) — same constant, same error message format.