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.
|
||||
|
||||
@@ -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<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.
|
||||
|
||||
Reference in New Issue
Block a user