diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 93f5e65de..552478d8e 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -6160,6 +6160,27 @@ pub struct GpuDqnTrainer { /// `[max_aux_param_len]` reusable per-tensor reduce target for /// `aux_param_grad_reduce`. Sized to the largest of the 8 aux tensors. aux_param_grad_final_buf: CudaSlice, + /// SP20 Phase 5 (2026-05-10): per-sample aux confidence at the SAMPLED + /// state — `[batch_size]` f32, range `[0, 0.5]` (K=2 peak-softmax-above- + /// uniform; 0.0 = uniform / no information sentinel; 0.5 = peak softmax + /// at 1.0). Populated by direct-to-trainer PER gather from the replay + /// buffer's `aux_conf` ring (filled by `experience_env_step` per (i, t) + /// in Phase 3 Task 3.4). Wired in at trainer init via + /// `GpuReplayBuffer::set_trainer_aux_conf_ptr(self.aux_conf_at_state_buf + /// .raw_ptr())` so PER's `gather_f32_scalar` writes the SAMPLED bar's + /// per-batch aux_conf directly into here on every step (mirrors the + /// SP13 B1.1b `aux_nb_label_buf` direct-to-trainer pattern). Consumed + /// by `c51_loss_batched`'s reward gate at the Bellman projection: + /// `gate = sigmoid((aux_conf - ISV[AUX_CONF_THRESHOLD]) / ISV[AUX_GATE_TEMP])`, + /// `r_used = gate × reward`. NULL-tolerant (gate=1.0 = no-op) at the + /// kernel level so test scaffolds without a wired aux head still work. + /// `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 — the "graceful degradation" semantic + /// from the Phase 3 Task 3.4 audit doc spec §4.4. Once the PER + /// gather populates it on the first sample step, the per-bar values + /// from the producer drive the gate as designed. + aux_conf_at_state_buf: CudaSlice, /// Per-step aux-loss weight, baked into the SAXPY alpha at graph-capture /// time. ISV-driven via `set_aux_weight`: /// `aux_weight = clamp(0.1 × LEARNING_HEALTH × (1 - tanh(SHARPE_SCALE × sharpe_ema)), 0.05, 0.3)`. @@ -22564,6 +22585,12 @@ impl GpuDqnTrainer { .max(aux_kr) .max(aux_knb); let aux_param_grad_final_buf = alloc_f32(&stream, max_aux_tensor_len, "aux_param_grad_final_buf")?; + // SP20 Phase 5: per-sample aux_conf at sampled state ([B] f32). PER's + // direct-to-trainer `gather_f32_scalar` writes here on every sample + // step once `set_trainer_aux_conf_ptr` is wired; until then, the + // alloc_zeros 0.0 sentinel drives the c51_loss reward-gate to ~0 + // (graceful degradation per Phase 3 Task 3.4 audit doc spec §4.4). + let aux_conf_at_state_buf = alloc_f32(&stream, aux_b, "aux_conf_at_state_buf")?; // Initial aux_weight at the lower numerical-stability clamp (0.05). // training_loop's per-step `set_aux_weight` recomputes from ISV every // step. The first graph capture bakes 0.05 — an honest cold-start @@ -24975,6 +25002,7 @@ impl GpuDqnTrainer { aux_partial_rg_w2, aux_partial_rg_b2, aux_param_grad_final_buf, + aux_conf_at_state_buf, aux_weight, stochastic_depth_scale_buf, stochastic_depth_kernel, @@ -33800,6 +33828,17 @@ impl GpuDqnTrainer { /// in `gpu_experience_collector.rs`) into here on every step. pub(crate) fn aux_nb_label_buf_ptr(&self) -> u64 { self.aux_nb_label_buf.raw_ptr() } + /// SP20 Phase 5 (2026-05-10): raw pointer to `aux_conf_at_state_buf` + /// for direct-to-trainer PER gather. Mirrors the SP13 B1.1b + /// `aux_nb_label_buf_ptr` direct-to-trainer pattern. The buffer is + /// `CudaSlice` `[batch_size]`; the c51 reward-gate consumer + /// (`c51_loss_batched`) reads it as f32 to compute + /// `gate = sigmoid((aux_conf - ISV[AUX_CONF_THRESHOLD]) / ISV[AUX_GATE_TEMP])`. + /// `GpuReplayBuffer::set_trainer_aux_conf_ptr` is invoked once at + /// trainer init in `training_loop.rs` so PER's `gather_f32_scalar` + /// writes the SAMPLED bar's aux_conf into here on every step. + pub(crate) fn aux_conf_at_state_buf_ptr(&self) -> u64 { self.aux_conf_at_state_buf.raw_ptr() } + /// Padded state dimension (128-byte aligned) for direct-to-trainer gather. pub(crate) fn state_dim_padded(&self) -> usize { ml_core::state_layout::STATE_DIM_PADDED } diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 828d128bc..cc2605f34 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -4160,6 +4160,19 @@ impl FusedTrainingCtx { self.trainer.aux_nb_label_buf_ptr() } + /// SP20 Phase 5 (2026-05-10): raw pointer to trainer's + /// `aux_conf_at_state_buf` for direct-to-trainer PER gather. Wraps + /// the inner `GpuDqnTrainer::aux_conf_at_state_buf_ptr()` accessor — + /// same delegation pattern as `trainer_aux_sign_labels_buf_ptr`. + /// Caller (`training_loop.rs`) pipes this into + /// `GpuReplayBuffer::set_trainer_aux_conf_ptr` so PER's + /// `gather_f32_scalar` writes the per-batch sampled aux_conf + /// straight into the trainer's f32 buffer the c51_loss reward-gate + /// consumer reads. + pub(crate) fn trainer_aux_conf_at_state_buf_ptr(&self) -> u64 { + self.trainer.aux_conf_at_state_buf_ptr() + } + /// Padded state dimension for direct-to-trainer gather. pub(crate) fn trainer_state_dim_padded(&self) -> usize { self.trainer.state_dim_padded() diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 6d3ed6f72..b64e1f5eb 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -822,6 +822,19 @@ impl DQNTrainer { fused.trainer_aux_sign_labels_buf_ptr(), fused.trainer_state_dim_padded(), ); + // SP20 Phase 5 (2026-05-10): wire trainer's + // `aux_conf_at_state_buf` raw_ptr so PER's + // `gather_f32_scalar` writes the SAMPLED bar's + // per-batch aux_conf directly into the trainer's + // f32 buffer the c51_loss reward-gate consumer + // reads. Same direct-to-trainer pattern as the + // SP13 B1.1b `aux_nb_label_buf` (i32) wire-up + // immediately above; the only difference is the + // element type (f32 vs i32) and the consumer + // (c51_loss_batched gate vs aux_next_bar CE). + agent_w.memory_mut().gpu.set_trainer_aux_conf_ptr( + fused.trainer_aux_conf_at_state_buf_ptr(), + ); fused.set_per_rng_step_dev_ptr(agent_w.memory().gpu.rng_step_dev_ptr()); } } @@ -2629,6 +2642,14 @@ impl DQNTrainer { fused.trainer_aux_sign_labels_buf_ptr(), fused.trainer_state_dim_padded(), ); + // SP20 Phase 5 (2026-05-10): re-wire trainer's + // `aux_conf_at_state_buf` raw_ptr after fused_ctx + // re-init. Mirrors the init site above; both call + // sites must stay in sync per + // `feedback_no_partial_refactor`. + agent_w.memory_mut().gpu.set_trainer_aux_conf_ptr( + fused.trainer_aux_conf_at_state_buf_ptr(), + ); fused.set_per_rng_step_dev_ptr(agent_w.memory().gpu.rng_step_dev_ptr()); } } diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 07121479c..e7e1c3fae 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -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` — `[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.