feat(sp22-vnext): Phase B1 — trade-outcome head weight tensors + Xavier init
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user