From 205e46c171643ea927ddbfb74d45b980461bb746 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 14 May 2026 00:33:38 +0200 Subject: [PATCH] =?UTF-8?q?feat(sp22-vnext):=20Phase=20B1=20=E2=80=94=20tr?= =?UTF-8?q?ade-outcome=20head=20weight=20tensors=20+=20Xavier=20init?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the 4 weight tensors (W1, b1, W2, b2) for the SP22 H6 vNext trade-outcome aux head into the trainer's flat params_buf at indices [163..167). Adam machinery (m/v moment buffers, SAXPY iteration over 0..NUM_WEIGHT_TENSORS) picks up the new tensors uniformly — no per-tensor wiring needed. Changes: - NUM_WEIGHT_TENSORS bumped 163 → 167. Most of the 54 references are &[u64; NUM_WEIGHT_TENSORS] array-size generics that resize uniformly with the constant. - compute_param_sizes() adds 4 new size entries: [163] aux_to_w1 [H=128, SH2=256] = 32,768 floats [164] aux_to_b1 [H=128] = 128 floats [165] aux_to_w2 [K=3, H=128] = 384 floats [166] aux_to_b2 [K=3] = 3 floats Total: 33,283 floats = ~133 KB params, ~266 KB Adam state. - compute_param_sizes() debug_assert updated 163 → 167. - Xavier fan_dims added: (H, SH2) for W1, (0, 0) for biases (zero-init), (K=3, H) for W2. Cold-start: logits ≈ 0 → softmax ≈ uniform 1/3 → no Profit/Stop/Timeout preference per pearl_first_observation_bootstrap. SH2 stays at 256 in this commit (mirrors K=2 head exactly). The spec's Phase B input concat (256 → 262 with plan_params 6-dim) will re-shape slot [163] to 128 × 262 = 33,536 floats later — small touch-up vs the full B1 commit. Verification: - cargo check -p ml clean (21 warnings, none new). - 14 pre-existing test failures stashed-verified unrelated (OFI features missing, ISV slot count drift — independent of NUM_WEIGHT_TENSORS). Phase B2 next: saved-tensor + per-sample partial buffers (hidden_post, logits, softmax, valid_count, dW*_partial, dh_s2_aux_out). Audit: docs/dqn-wire-up-audit.md Phase B1 section. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 70 ++++++++++++++++++- docs/dqn-wire-up-audit.md | 29 ++++++++ 2 files changed, 96 insertions(+), 3 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index efc68ec9f..ffe8e823c 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -4484,7 +4484,22 @@ impl Default for CausalInterventionConfig { /// → indices [127..163). Gate: gate_w1[SD,64], gate_b1[64], gate_w2[64,8], /// gate_b2[8]. Per expert: w1[SH2,64], b1[64], w2[64,SH2], b2[SH2]. /// Allocated but not yet wired into forward/backward — Phase 3 wire-up. -pub(crate) const NUM_WEIGHT_TENSORS: usize = 163; +/// SP22 H6 vNext Phase B1 (2026-05-14): bumped 163 → 167 to host the +/// trade-outcome aux head's 4 weight tensors at indices [163..167): +/// [163] aux_to_w1 [H=128, SH2=256] Xavier +/// [164] aux_to_b1 [H=128] zero +/// [165] aux_to_w2 [K=3, H=128] Xavier +/// [166] aux_to_b2 [K=3] zero +/// +/// Adam m/v moment buffers and the SAXPY iteration pick up the new +/// tensors uniformly via `for i in 0..NUM_WEIGHT_TENSORS` loops — no +/// per-tensor Adam wiring needed. +/// +/// Sized with Phase A3 SH2=256 (mirrors the K=2 next-bar head). The +/// spec's Phase B input concat (256 → 262 with `plan_params`) will +/// land later as a small change to `compute_param_sizes()` slot [163] +/// shape + the kernel-launch SH2 arg. +pub(crate) const NUM_WEIGHT_TENSORS: usize = 167; /// Compute the size (element count) of each weight tensor. /// @@ -4735,8 +4750,38 @@ pub(crate) fn compute_param_sizes(cfg: &GpuDqnTrainConfig) -> [usize; NUM_WEIGHT "compute_param_sizes: MoE expert fill ended at {} but expected 163", eidx); } - debug_assert!(NUM_WEIGHT_TENSORS == 163, - "compute_param_sizes: NUM_WEIGHT_TENSORS expected 163 after MoE extension, got {}", + + // ── SP22 H6 vNext Phase B1 (2026-05-14): trade-outcome aux head ── + // Architecture: `Linear(SH2 → AUX_HIDDEN_DIM) → ELU → Linear(H → K=3)`. + // Same shape as `aux_nb_*` (K=2 sibling) but K_OUT = AUX_OUTCOME_K = 3. + // + // Indices: aux_to_w1=163, aux_to_b1=164, aux_to_w2=165, aux_to_b2=166. + // + // Per-tensor sizes (sh2 = cfg.shared_h2 = 256, h_dim = 128, k_to = 3): + // [163] aux_to_w1 [H, SH2] → 128 × 256 = 32,768 floats + // [164] aux_to_b1 [H] → 128 floats + // [165] aux_to_w2 [K=3, H] → 3 × 128 = 384 floats + // [166] aux_to_b2 [K=3] → 3 floats + // + // Total: 33,283 floats = ~133 KB. Adam m/v moments double this to + // ~266 KB additional VRAM — negligible vs the existing aux trunk + // tensor budget. + // + // SH2=256 mirrors the K=2 head exactly at this commit. The spec's + // Phase B (input concat with plan_params 6-dim → SH2=262) will + // re-shape slot [163] later: 128 × 262 = 33,536 floats (+128 from + // the 6 extra input columns × 128 hidden lanes). + { + let h_dim = crate::cuda_pipeline::gpu_aux_heads::AUX_HIDDEN_DIM; + let k_to = crate::cuda_pipeline::gpu_aux_heads::AUX_OUTCOME_K; + sizes[163] = h_dim * sh2; // aux_to_w1 [H, SH2] + sizes[164] = h_dim; // aux_to_b1 [H] + sizes[165] = k_to * h_dim; // aux_to_w2 [K=3, H] + sizes[166] = k_to; // aux_to_b2 [K=3] + } + + debug_assert!(NUM_WEIGHT_TENSORS == 167, + "compute_param_sizes: NUM_WEIGHT_TENSORS expected 167 after vNext extension, got {}", NUM_WEIGHT_TENSORS); sizes @@ -33341,6 +33386,25 @@ impl GpuDqnTrainer { fan_dims[126] = (0, 0); // aux_rg_b2 (zero) } + // ── SP22 H6 vNext Phase B1 (2026-05-14): trade-outcome head fan dims ── + // Same shape pattern as aux_nb_* (K=2 sibling) — Xavier on W + // matrices, zero on biases. Cold-start: logits ≈ 0 → softmax ≈ + // uniform 1/3 → no class preference for {Profit, Stop, Timeout}. + // The label producer's sparse-label semantics (most bars masked) + // mean the head trains only on actual trade-close events, so + // cold-start neutrality is the correct starting point per + // pearl_first_observation_bootstrap. + // + // Indices: aux_to_w1=163, aux_to_b1=164, aux_to_w2=165, aux_to_b2=166. + { + let h_dim = crate::cuda_pipeline::gpu_aux_heads::AUX_HIDDEN_DIM; + let k_to = crate::cuda_pipeline::gpu_aux_heads::AUX_OUTCOME_K; + fan_dims[163] = (h_dim, cfg.shared_h2); // aux_to_w1 [H, SH2] (Xavier) + fan_dims[164] = (0, 0); // aux_to_b1 (zero) + fan_dims[165] = (k_to, h_dim); // aux_to_w2 [K=3, H] (Xavier) + fan_dims[166] = (0, 0); // aux_to_b2 (zero) + } + // ── Phase 1 Task 1.6: MoE fan dims (36 tensors at [127..163)) ───────── // Gate: zero-init (fan_out=0,fan_in=0) so g(s) = uniform 1/K at cold start. // Experts: Xavier init on w1/w2 for initial differentiation; zero on biases. diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 016a7eae1..4853abc3e 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -17780,3 +17780,32 @@ First Rust-side commit of Phase B. Pure scaffolding: adds orchestrator structs i **Discipline**: `feedback_no_partial_refactor` honored — cubin embed + ops struct + load function names are atomic; the ops methods can't be called without the trainer struct fields, but having the methods land first lets the field-allocation commits compile-check the wireup endpoint by endpoint. Cargo check clean (21 warnings, none new — all pre-existing `DevicePtrMut` unused imports + sp14 unused isv_slots). + +#### Phase B1 — Trade-outcome head weight tensors (2026-05-14) + +Second Rust-side commit of Phase B. Adds the 4 weight tensors (W1, b1, W2, b2) for the trade-outcome aux head into the trainer's flat `params_buf` at indices `[163..167)`. The Adam machinery (m/v moment buffers, SAXPY iteration over `0..NUM_WEIGHT_TENSORS`) picks up the new tensors uniformly — no per-tensor wiring needed. + +**Changes**: +- `NUM_WEIGHT_TENSORS`: bumped 163 → 167. Affects 54 references throughout `crates/ml/src/` but most are `&[u64; NUM_WEIGHT_TENSORS]` array-size generics that resize uniformly with the constant. +- `compute_param_sizes()`: added 4 new size entries: + ``` + sizes[163] = AUX_HIDDEN_DIM × SH2 // aux_to_w1 [H, SH2] = 128 × 256 = 32,768 floats + sizes[164] = AUX_HIDDEN_DIM // aux_to_b1 [H] = 128 floats + sizes[165] = AUX_OUTCOME_K × AUX_HIDDEN_DIM // aux_to_w2 [K=3, H] = 384 floats + sizes[166] = AUX_OUTCOME_K // aux_to_b2 [K=3] = 3 floats + ``` + Total: 33,283 floats = ~133 KB params, ~266 KB Adam state. Negligible vs the existing aux trunk budget. +- `compute_param_sizes()` debug_assert updated 163 → 167. +- Xavier `fan_dims` added for slots [163..167): `(H, SH2)` for W1, `(0, 0)` for biases (zero-init), `(K=3, H)` for W2. Cold-start: logits ≈ 0 → softmax ≈ uniform 1/3 → no Profit/Stop/Timeout preference, matching `pearl_first_observation_bootstrap`. + +**SH2 note**: stays at 256 in this commit (mirrors K=2 head exactly). The spec's Phase B input concat (256 → 262 with `plan_params`) will re-shape slot [163] to `128 × 262 = 33,536 floats` later — small touch-up vs the full B1 commit. + +**Adam machinery**: The existing m/v moment buffers are allocated as `[total_params]` arrays (sum over all NUM_WEIGHT_TENSORS sizes), so bumping the constant uniformly extends them. The SAXPY kernel iterates 0..NUM_WEIGHT_TENSORS so new tensors get processed automatically. No explicit Adam wiring needed. + +**Verification**: +- `cargo check -p ml` clean (21 warnings, none new). +- 14 pre-existing test failures stashed-verified to be unrelated (OFI features missing, ISV slot count assertion drift from SP22 H6 — all touch areas independent of NUM_WEIGHT_TENSORS). + +**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.