diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index ffe8e823c..2f4e14fbc 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -61,7 +61,8 @@ use super::batched_forward::{CublasForward, CublasGemmSet, f32_weight_ptrs_from_ use super::batched_backward::{CublasBackward, CublasBackwardSet, alloc_backward_scratch, raw_f32_ptr as bw_raw_f32_ptr}; use super::shared_cublas_handle::PerStreamCublasHandles; use super::gpu_aux_heads::{ - AuxHeadsBackwardOps, AuxHeadsForwardOps, AUX_HIDDEN_DIM, AUX_NEXT_BAR_K, AUX_REGIME_K, + AuxHeadsBackwardOps, AuxHeadsForwardOps, AuxTradeOutcomeBackwardOps, + AuxTradeOutcomeForwardOps, AUX_HIDDEN_DIM, AUX_NEXT_BAR_K, AUX_OUTCOME_K, AUX_REGIME_K, }; use super::gpu_health_diag::GpuHealthDiag; use super::gpu_moe_head::GpuMoeHead; @@ -6315,6 +6316,89 @@ 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, + + // ── SP22 H6 vNext Phase B0+B2: trade-outcome aux head orchestrator + buffers ── + // + // Parallel to `aux_heads_fwd` / `aux_heads_bwd` + `aux_nb_*` buffers + // but for the K=3 trade-outcome head. The trainer holds the forward + // ops (forward + loss_reduce + label producer) AND the backward ops + // (backward only — param_grad_reduce is reused from `aux_heads_bwd`). + // Collector-side rollout buffers + ops handle land in Phase B3 wireup. + // + /// Phase B0 (2026-05-14): forward orchestrator for the K=3 + /// trade-outcome head. Holds the `aux_trade_outcome_forward`, + /// `aux_trade_outcome_loss_reduce`, and `trade_outcome_label_kernel` + /// handles. Phase B3 wires the launch chain into the trainer's + /// replay-batch path; for B2 it's instantiated but unused. + aux_to_fwd: AuxTradeOutcomeForwardOps, + /// Phase B0 (2026-05-14): backward orchestrator for the K=3 + /// trade-outcome head. Holds the `aux_trade_outcome_backward` + /// handle. The `aux_param_grad_reduce` K-generic reducer is reused + /// from `aux_heads_bwd`. + aux_to_bwd: AuxTradeOutcomeBackwardOps, + /// `[B, AUX_HIDDEN_DIM]` saved post-ELU activation of the + /// trade-outcome head's Linear_1. Consumed by + /// `aux_trade_outcome_backward` for the dW1 / dh_s2_aux chain. + aux_to_hidden_buf: CudaSlice, + /// `[B, AUX_OUTCOME_K=3]` saved logits — written by + /// `aux_trade_outcome_forward`. Kept for parity with the K=2 sibling + /// (backward reads the softmax tile, not the logits — the softmax + /// is already materialized in the forward). + aux_to_logits_buf: CudaSlice, + /// `[B, AUX_OUTCOME_K=3]` saved softmax tile — written by + /// `aux_trade_outcome_forward`. Three downstream consumers will + /// read this tile: + /// - `aux_trade_outcome_loss_reduce` (CE numerator) + /// - `aux_trade_outcome_backward` (d_logits = softmax - one_hot) + /// - Phase C 3-slot state assembly producer (state[121..124]) + aux_to_softmax_buf: CudaSlice, + /// `[B] i32` per-sample trade-outcome label in `{-1, 0, 1, 2}`: + /// `-1` = mask (no trade close at this bar), `0` = Profit, `1` = + /// Stop, `2` = Timeout. Sparse: most bars (~95-99%) are `-1`. + /// Producer: `trade_outcome_label_kernel` per-step (reads + /// `trade_close_per_sample[i*L+t]` and the save-for-backward + /// `pnl_vs_target_at_close_per_env` / `pnl_vs_stop_at_close_per_env` + /// from Phase A3). Direct-to-trainer PER gather will populate this + /// on every replay sample step in Phase B3 (mirrors the `aux_nb_ + /// label_buf` direct-to-trainer pattern). Consumed by + /// `aux_trade_outcome_loss_reduce` and `aux_trade_outcome_backward`. + /// Cold-start `alloc_zeros` initialization yields label 0 (Profit) + /// for every sample → degraded "model learns to predict Profit + /// everywhere" until the PER gather wires up; same known-degraded + /// behaviour the K=2 head went through between B1.1a and B1.1b. + aux_to_label_buf: CudaSlice, + /// `[1]` scalar mean cross-entropy — written by + /// `aux_trade_outcome_loss_reduce`, will be read by Phase F's + /// HEALTH_DIAG aux-outcome-CE EMA producer (parallel to + /// `aux_heads_loss_ema_update`). + aux_to_loss_scalar_buf: CudaSlice, + /// `[1]` scalar `B_valid` count (number of non-mask labels in the + /// batch) — written by `aux_trade_outcome_loss_reduce`, read by + /// `aux_trade_outcome_backward` so loss + backward share the same + /// `1/B_valid` divisor (derivatives of the same scalar function). + /// Sparse-label arithmetic: typical `B_valid` ≈ 1-5% of nominal B. + aux_to_valid_count_buf: CudaSlice, + /// `[B, SH2]` per-sample dh_s2_aux emitted by + /// `aux_trade_outcome_backward`. SAXPY-accumulated into the aux + /// trunk's `dh_s2_aux_accum` before `aux_trunk_backward` runs. + /// Encoder stop-grad enforced structurally by `aux_trunk_backward` + /// (no `dx_in` output param), so Q's encoder is protected from aux + /// contamination per SP14 Phase C.5b. + aux_dh_s2_to_buf: CudaSlice, + /// `[B, AUX_HIDDEN_DIM, SH2]` per-sample dW1 partial. Reduced along + /// the batch dim by the K-generic `aux_param_grad_reduce` (reused + /// from `aux_heads_bwd`) into the final dW1 written into the flat + /// `params_buf` Adam slot [163]. + aux_partial_to_w1: CudaSlice, + /// `[B, AUX_HIDDEN_DIM]` per-sample db1 partial. Reduced into the + /// flat `params_buf` Adam slot [164]. + aux_partial_to_b1: CudaSlice, + /// `[B, AUX_OUTCOME_K=3, AUX_HIDDEN_DIM]` per-sample dW2 partial. + /// Reduced into the flat `params_buf` Adam slot [165]. + aux_partial_to_w2: CudaSlice, + /// `[B, AUX_OUTCOME_K=3]` per-sample db2 partial. Reduced into the + /// flat `params_buf` Adam slot [166]. + aux_partial_to_b2: 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 @@ -22935,6 +23019,52 @@ 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")?; + + // ── SP22 H6 vNext Phase B0+B2 (2026-05-14): trade-outcome head ── + // orchestrator (load both fwd + bwd) + 11 buffers (mirror of + // aux_nb_* pattern at K=3 instead of K=2). + // + // `aux_to_fwd` holds 3 kernel handles (forward, loss_reduce, + // label producer); `aux_to_bwd` holds the backward kernel — + // param_grad_reduce is reused from `aux_heads_bwd`. + // + // Buffer sizes (B=batch_size, H=128, K=3, SH2=256): + // aux_to_hidden_buf [B, H] = B × 128 + // aux_to_logits_buf [B, K=3] = B × 3 + // aux_to_softmax_buf [B, K=3] = B × 3 + // aux_to_label_buf [B] i32 + // aux_to_loss_scalar_buf [1] + // aux_to_valid_count_buf [1] + // aux_dh_s2_to_buf [B, SH2] = B × 256 + // aux_partial_to_w1 [B, H, SH2] = B × 32,768 (largest) + // aux_partial_to_b1 [B, H] = B × 128 + // aux_partial_to_w2 [B, K=3, H] = B × 384 + // aux_partial_to_b2 [B, K=3] = B × 3 + // + // At B=batch_size=2048 and SH2=256: aux_partial_to_w1 dominates + // at 2048 × 32,768 × 4 bytes = 256 MB. Matches the K=2 head's + // partial size (same SH2, same H) — no new memory pressure. + let aux_to_fwd = AuxTradeOutcomeForwardOps::new(&stream)?; + let aux_to_bwd = AuxTradeOutcomeBackwardOps::new(&stream)?; + let aux_kto = AUX_OUTCOME_K; + let aux_to_hidden_buf = alloc_f32(&stream, aux_b * aux_h, "aux_to_hidden_buf")?; + let aux_to_logits_buf = alloc_f32(&stream, aux_b * aux_kto, "aux_to_logits_buf")?; + let aux_to_softmax_buf = alloc_f32(&stream, aux_b * aux_kto, "aux_to_softmax_buf")?; + let aux_to_label_buf = stream.alloc_zeros::(aux_b) + .map_err(|e| MLError::ModelError(format!( + "alloc aux_to_label_buf (sp22-vnext B2 i32): {e}" + )))?; + let aux_to_loss_scalar_buf = alloc_f32(&stream, 1, "aux_to_loss_scalar_buf")?; + let aux_to_valid_count_buf = alloc_f32(&stream, 1, "aux_to_valid_count_buf")?; + let aux_dh_s2_to_buf = alloc_f32(&stream, aux_b * aux_sh2, "aux_dh_s2_to_buf")?; + let aux_partial_to_w1 = alloc_f32(&stream, aux_b * aux_h * aux_sh2, "aux_partial_to_w1")?; + let aux_partial_to_b1 = alloc_f32(&stream, aux_b * aux_h, "aux_partial_to_b1")?; + let aux_partial_to_w2 = alloc_f32(&stream, aux_b * aux_kto * aux_h, "aux_partial_to_w2")?; + let aux_partial_to_b2 = alloc_f32(&stream, aux_b * aux_kto, "aux_partial_to_b2")?; + // NOTE: aux_param_grad_final_buf is sized to handle the largest + // tensor across BOTH K=2 (next_bar) and K=3 (trade_outcome) + // heads — both heads' largest tensor is W1 [H, SH2] = 32,768 + // floats, identical. No re-sizing needed. // 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 @@ -25352,6 +25482,22 @@ impl GpuDqnTrainer { aux_partial_rg_w2, aux_partial_rg_b2, aux_param_grad_final_buf, + // SP22 H6 vNext Phase B0+B2 (2026-05-14): trade-outcome head + // orchestrator + 11 buffers. Mirrors aux_heads_fwd/bwd + + // aux_nb_* / aux_partial_nb_* placement. + aux_to_fwd, + aux_to_bwd, + aux_to_hidden_buf, + aux_to_logits_buf, + aux_to_softmax_buf, + aux_to_label_buf, + aux_to_loss_scalar_buf, + aux_to_valid_count_buf, + aux_dh_s2_to_buf, + aux_partial_to_w1, + aux_partial_to_b1, + aux_partial_to_w2, + aux_partial_to_b2, aux_conf_at_state_buf, aux_weight, stochastic_depth_scale_buf, diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 4853abc3e..cff65fb5e 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -17809,3 +17809,34 @@ Second Rust-side commit of Phase B. Adds the 4 weight tensors (W1, b1, W2, b2) f **Phase B2 next**: saved-tensor + per-sample partial buffers (hidden_post [B, H], logits [B, K], softmax [B, K], valid_count [1], dW1_partial [B, H, SH2], db1_partial [B, H], dW2_partial [B, K, H], db2_partial [B, K], dh_s2_aux_out [B, SH2]). These get allocated as `CudaSlice` fields on the trainer or collector struct, mirroring the existing `aux_nb_hidden_post / aux_nb_logits / aux_nb_softmax / aux_nb_valid_count / aux_nb_dW1_partial / ...` pattern. Cargo check clean. Adam machinery uniformly extends. Ready for B2. + +#### Phase B2 — Trade-outcome trainer saved-tensor + per-sample partial buffers (2026-05-14) + +Third Rust-side commit of Phase B. Adds 11 buffer fields + 2 orchestrator ops handles to the trainer struct, mirroring the existing `aux_nb_*` / `aux_partial_nb_*` pattern at K=3 instead of K=2. + +**Trainer struct additions** (in `gpu_dqn_trainer.rs`): +- `aux_to_fwd: AuxTradeOutcomeForwardOps` — holds the 3 forward-side kernel handles (Phase B0 scaffold) +- `aux_to_bwd: AuxTradeOutcomeBackwardOps` — holds the backward kernel handle +- `aux_to_hidden_buf [B, H=128]` — saved post-ELU activation +- `aux_to_logits_buf [B, K=3]` — saved logits (parity with K=2 sibling) +- `aux_to_softmax_buf [B, K=3]` — saved softmax (3 future consumers: loss reduce, backward, Phase C state slot producer) +- `aux_to_label_buf [B] i32` — sparse labels in `{-1, 0=Profit, 1=Stop, 2=Timeout}` (~95-99% will be `-1` in production) +- `aux_to_loss_scalar_buf [1]` — mean CE +- `aux_to_valid_count_buf [1]` — `B_valid` for backward's `inv_B = 1/B_valid` divisor +- `aux_dh_s2_to_buf [B, SH2]` — backward emits this; SAXPYs into aux trunk's `dh_s2_aux_accum` +- `aux_partial_to_w1 [B, H, SH2]` — per-sample dW1 partial (32,768 × B floats — largest single buffer) +- `aux_partial_to_b1 [B, H]` — per-sample db1 partial +- `aux_partial_to_w2 [B, K=3, H]` — per-sample dW2 partial +- `aux_partial_to_b2 [B, K=3]` — per-sample db2 partial + +**Memory note**: at B=batch_size=2048, SH2=256, H=128 → `aux_partial_to_w1 = 2048 × 32,768 × 4 bytes = 256 MB`. Matches the K=2 head's partial buffer size exactly (same SH2, same H) — no new memory pressure beyond what's already validated. Total new aux-to buffer footprint ≈ 260 MB at B=2048. + +**param_grad_final scratch reuse**: The existing `aux_param_grad_final_buf` is sized to `max(aux_h × aux_sh2, aux_kr × aux_h, aux_knb × aux_h, ...)`. The trade-outcome head's largest tensor is W1 [H, SH2] = 32,768 floats — identical to the K=2/K=5 heads' W1. The scratch is already big enough; no resize. + +**Cold-start label semantics**: `alloc_zeros` yields label 0 (Profit) for every sample. Until the producer wires in (Phase B3), the trainer's CE loss will treat every sample as "should have predicted Profit", which is a degraded training signal but well-defined (no NaN). Mirrors the K=2 head's known-degraded state between B1.1a (forward landed) and B1.1b (label producer wired). + +**No FoldReset registration**: These buffers are overwritten every batch — no stale-state-leak risk across folds. The existing `aux_nb_*` pattern does the same (no reset registry entries). + +**Phase B3 next**: collector-side rollout buffers (forward + label producer launches on every step alongside the existing K=2 head's launches). Also wires the per-env softmax → per-(i, t) softmax fan-out scatter so the trainer's `aux_to_softmax_buf` gets populated from the rollout's `exp_aux_to_softmax_buf`. + +Cargo check clean (21 warnings, none new on the new fields).