feat(sp22-vnext): Phase B2 — trade-outcome trainer saved-tensor + partial buffers
Adds 11 buffer fields + 2 orchestrator ops handles (fwd + bwd) to the
trainer struct, mirroring the existing aux_nb_* / aux_partial_nb_*
pattern at K=3 instead of K=2.
Trainer struct additions:
- aux_to_fwd: AuxTradeOutcomeForwardOps (Phase B0 scaffold)
- aux_to_bwd: AuxTradeOutcomeBackwardOps
- aux_to_hidden_buf [B, H=128] saved post-ELU
- aux_to_logits_buf [B, K=3] saved logits
- aux_to_softmax_buf [B, K=3] saved softmax (3 future consumers)
- aux_to_label_buf [B] i32 sparse {-1, 0, 1, 2}
- aux_to_loss_scalar_buf [1] mean CE
- aux_to_valid_count_buf [1] B_valid for backward
- aux_dh_s2_to_buf [B, SH2] SAXPYs into dh_s2_aux_accum
- aux_partial_to_w1 [B, H, SH2] per-sample dW1
- aux_partial_to_b1 [B, H] per-sample db1
- aux_partial_to_w2 [B, K=3, H] per-sample dW2
- aux_partial_to_b2 [B, K=3] per-sample db2
Memory: aux_partial_to_w1 = 256 MB at B=2048 — identical to K=2 head's
partial size (same SH2, same H). Total new aux-to footprint ≈ 260 MB.
The existing aux_param_grad_final_buf scratch is sized to the largest
tensor across all aux heads; trade-outcome head's largest is W1 [H, SH2]
= 32,768 floats — identical to K=2/K=5 W1s. No resize needed.
Cold-start label semantics: alloc_zeros yields label 0 (Profit) for
every sample. Until the producer wires in (B3), the trainer's CE loss
treats every sample as "should have predicted Profit" — degraded but
well-defined (no NaN). Mirrors the K=2 head's known-degraded state
between B1.1a and B1.1b.
No FoldReset registration: these buffers are overwritten every batch
— no stale-state-leak risk across folds (matches the existing aux_nb_*
pattern).
Phase B3 next: collector-side rollout buffers + forward chain wireup
into collect_experiences_gpu (per-env softmax → per-(i, t) fan-out
scatter for trainer's aux_to_softmax_buf population).
Audit: docs/dqn-wire-up-audit.md Phase B2 section.
Cargo check clean (21 warnings, none new on aux_to_* fields).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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<f32>,
|
||||
|
||||
// ── 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<f32>,
|
||||
/// `[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<f32>,
|
||||
/// `[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<f32>,
|
||||
/// `[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<i32>,
|
||||
/// `[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<f32>,
|
||||
/// `[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<f32>,
|
||||
/// `[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<f32>,
|
||||
/// `[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<f32>,
|
||||
/// `[B, AUX_HIDDEN_DIM]` per-sample db1 partial. Reduced into the
|
||||
/// flat `params_buf` Adam slot [164].
|
||||
aux_partial_to_b1: CudaSlice<f32>,
|
||||
/// `[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<f32>,
|
||||
/// `[B, AUX_OUTCOME_K=3]` per-sample db2 partial. Reduced into the
|
||||
/// flat `params_buf` Adam slot [166].
|
||||
aux_partial_to_b2: CudaSlice<f32>,
|
||||
/// 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::<i32>(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,
|
||||
|
||||
@@ -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<f32>` 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).
|
||||
|
||||
Reference in New Issue
Block a user