From b7c4f84ea04a38a7346c1c68ea06fb9f07a73a32 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 10 May 2026 22:44:43 +0200 Subject: [PATCH] =?UTF-8?q?feat(sp21):=20T2.2=20Phase=204=20=E2=80=94=20wi?= =?UTF-8?q?re=20E4=20branch=5Flr=5Fscale=20=E2=86=92=20per-branch=20Adam?= =?UTF-8?q?=20(atomic)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires EnrichmentResult::branch_lr_scale (E4's [f32; 4] clamped [0.5, 2.0] per-branch LR multiplier) into the DQN Adam optimizer. Splits the existing single DqnBranches Adam sub-launch into 4 per-branch sub-launches, each consuming its own LR scale from ISV[521..525). Plan amendment from on-paper design: - Plan said "via existing per-group Adam infrastructure" but per-group operates at PARAM_GROUP granularity (8 groups, all 4 action branches lumped into DqnBranches). E4 needs per-action-BRANCH granularity. - Resolution: split DqnBranches sub-launch into 4 per-branch sub-launches using the already-canonical branch byte ranges (4 param tensors per branch: Dir [17..21), Mag [21..25), Order [25..29), Urgency [29..33)). Coverage invariant updated to 7-way (was 4-way). - Pearl C engagement tracking on branches DEFERRED to Phase 4.5: the shared DqnBranches engagement counter offset would collide on writes if 4 sub-launches use the same offset. Pearl C is a diagnostic system (not load-bearing) so its temporary unavailability for branches is acceptable; Trunk + Value + Trunk-extras Pearl C still active. Phase 4.5 follow-up will re-instate via ParamGroup expansion or sub-block offsetting scheme. ISV slot allocation: - BRANCH_LR_SCALE_{DIR,MAG,ORDER,URGENCY}_INDEX = 521..525 - ISV_TOTAL_DIM 521 → 525 (bus extension) - Layout fingerprint adds 4 SLOT entries - New branch_lr_scale_index(branch_idx) accessor for clean mapping - Cold-start floor: launcher reads ISV, floors at 1.0 if at sentinel 0.0 (per pearl_first_observation_bootstrap — first emit replaces directly with no intermediate state) ABI surgery: - dqn_adam_update_kernel: ONE new arg float lr_scale at end of signature; lr = *lr_ptr * lr_scale inside kernel - 5 callers migrated atomically per feedback_no_partial_refactor: - launch_adam_update (main DQN): 7 sub-launches with per-branch lr_scale (Trunk + Value + 4 branch + Trunk-extras) - 4 post-aux launchers (ofi_embed/aux_trunk/denoise/sel) pass 1.0 - decision_transformer Adam launch passes 1.0 Producer wireup (training_loop.rs post-enrichment block): ```rust for (branch_idx, &scale) in result.branch_lr_scale.iter().enumerate() { fused.trainer().write_isv_signal_at( branch_lr_scale_index(branch_idx), scale, ); } ``` Files changed: - crates/ml/src/cuda_pipeline/sp21_isv_slots.rs: +4 slots + accessor + tests - crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs: ISV_TOTAL_DIM bump + fingerprint + 7-way Adam split with per-branch lr_scale - crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu: lr_scale arg + apply - crates/ml/src/cuda_pipeline/decision_transformer.rs: lr_scale=1.0 - crates/ml/src/trainers/dqn/trainer/training_loop.rs: write 4 ISV slots - docs/dqn-wire-up-audit.md: 2026-05-11 audit entry Verification (passing): - cargo check -p ml --tests --features cuda: 0 errors - cargo test -p ml --lib sp21_isv_slots --features cuda: 3/3 (bus bounds + slot uniqueness + branch index mapping) - sp20_aggregate_inputs_test: 12/12 - sp20_phase1_4_wireup_test: 2/2 - sp20_emas_compute_test: 4/4 - sp20_controllers_compute_test: 7/7 - sp21_per_trade_predicted_q_test: 3/3 Total: 31 tests, 0 failures. Behavioral gate (per-branch LR divergence) is the upcoming smoke training run. After this commit (T2.2 Phases 5-7 + 8 + 4.5): - Phase 5: E6 winner indices → PER priority bumps - Phase 6: E7 hindsight → replay buffer injection - Phase 7: E8 curriculum weights → segment sampling - Phase 8: signal-drive remaining controller GAINS - Phase 4.5: re-instate Pearl C engagement tracking for branches Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/cuda_pipeline/decision_transformer.rs | 1 + .../src/cuda_pipeline/dqn_utility_kernels.cu | 17 +- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 104 ++++++++--- crates/ml/src/cuda_pipeline/sp21_isv_slots.rs | 69 ++++++- .../src/trainers/dqn/trainer/training_loop.rs | 21 ++- docs/dqn-wire-up-audit.md | 174 ++++++++++++++++++ 6 files changed, 358 insertions(+), 28 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/decision_transformer.rs b/crates/ml/src/cuda_pipeline/decision_transformer.rs index 9229db795..41577f9d6 100644 --- a/crates/ml/src/cuda_pipeline/decision_transformer.rs +++ b/crates/ml/src/cuda_pipeline/decision_transformer.rs @@ -1011,6 +1011,7 @@ impl DecisionTransformer { .arg(&dt_diag_slot_disabled) .arg(&dt_engage_buf_null) .arg(&dt_engage_off_disabled) + .arg(&1.0_f32) // SP21 Phase 4: lr_scale (DT is non-branch, no-op = 1.0) .launch(lc) .map_err(|e| MLError::ModelError(format!("DT adam: {e}")))?; } diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index 7b49eb64b..81c28e41f 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -147,7 +147,19 @@ extern "C" __global__ void dqn_adam_update_kernel( int* __restrict__ nan_flags_buf, /* sticky overflow flag store; nan_flags_buf[diag_slot]=1 on first thread that overshoots */ int diag_slot, /* index into nan_flags_buf (one per Adam kernel) */ int* __restrict__ clamp_engage_per_block_buf, /* [SP4_PARAM_GROUP_COUNT × MAX_BLOCKS_PER_ADAM] per-block engagement counts */ - int engage_buf_offset /* per-group offset (group*MAX_BLOCKS); -1 skips writeback (aux trainers) */ + int engage_buf_offset, /* per-group offset (group*MAX_BLOCKS); -1 skips writeback (aux trainers) */ + /* SP21 T2.2 Phase 4 (2026-05-11) — per-launch LR multiplier from + * E4 (`enrichment::compute_branch_lr_scale`). Each branch sub-launch + * passes its own `BRANCH_LR_SCALE_*_INDEX` ISV value here, scaling + * the effective Adam LR for that branch's params. Non-branch + * sub-launches (Trunk, Value, Trunk-extras, IQN, Attn, Curiosity) + * pass 1.0 (no-op). The kernel multiplies `*lr_ptr * lr_scale` + * before applying the Adam update; combined with the existing + * `lr_ptr` (host-controlled global LR schedule), this gives + * `effective_lr = global_lr × branch_scale`. Cold-start sentinel + * (ISV reads 0.0 before E4 emits) is handled HOST-side by the + * launcher (floor at 1.0) — kernel sees only validated scales. */ + float lr_scale ) { int idx = blockIdx.x * blockDim.x + threadIdx.x; /* Out-of-range threads still participate in the block reduce (with @@ -155,7 +167,8 @@ extern "C" __global__ void dqn_adam_update_kernel( bool active = (idx < total_params); int local_engage = 0; - float lr = active ? *lr_ptr : 0.0f; + /* SP21 Phase 4: scale effective LR by per-branch multiplier. */ + float lr = active ? (*lr_ptr * lr_scale) : 0.0f; int t = active ? *t_ptr : 1; float g_f = active ? grads[idx] : 0.0f; diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 17a4ac836..2006c33ea 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -2622,7 +2622,7 @@ const ISV_NETWORK_DIM: usize = 23; /// (and the C-side mirrors in `state_layout.cuh`). Bump this constant in the /// SAME commit that adds the new range, and update `layout_fingerprint_seed` /// to register the new slot names. -pub(crate) const ISV_TOTAL_DIM: usize = 521; // was 520, SP21 Phase 3 added 1 slot [520..521) for Q_CORRECTION_INDEX +pub(crate) const ISV_TOTAL_DIM: usize = 525; // SP21 Phase 4 adds 4 slots [521..525) for BRANCH_LR_SCALE_{DIR,MAG,ORDER,URGENCY}; Phase 3 added 1 slot [520..521) for Q_CORRECTION /// Legacy alias preserved for call sites that haven't been audited for the /// network-vs-total split. New code should pick `ISV_NETWORK_DIM` (for weight /// tensor sizing) or `ISV_TOTAL_DIM` (for the broadcast bus buffer). @@ -3749,7 +3749,11 @@ const fn layout_fingerprint_seed() -> &'static [u8] { SLOT_518_AUX_CONF_THRESHOLD=518;\ SLOT_519_AUX_GATE_TEMP=519;\ SLOT_520_Q_CORRECTION=520;\ - ISV_TOTAL_DIM=521;\ + SLOT_521_BRANCH_LR_SCALE_DIR=521;\ + SLOT_522_BRANCH_LR_SCALE_MAG=522;\ + SLOT_523_BRANCH_LR_SCALE_ORDER=523;\ + SLOT_524_BRANCH_LR_SCALE_URGENCY=524;\ + ISV_TOTAL_DIM=525;\ SP19_PRODUCER_HARDCODED_HORIZON_BLEND=sp19_path_b;\ SP14_C_AUX_TRUNK_CONTROL_PLANE=sp14_c_phase_1;\ ALPHA_MACHINERY_DELETED=sp14_c_phase_1;\ @@ -10877,6 +10881,7 @@ impl GpuDqnTrainer { .arg(&_aux_diag_slot_disabled) .arg(&_aux_engage_buf_null) .arg(&_aux_engage_off_disabled) + .arg(&1.0_f32) // SP21 Phase 4: lr_scale (aux trainer, non-branch = 1.0 no-op) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), @@ -11118,6 +11123,7 @@ impl GpuDqnTrainer { .arg(&diag_slot_disabled) .arg(&engage_buf_null) .arg(&engage_off_disabled) + .arg(&1.0_f32) // SP21 Phase 4: lr_scale (aux trunk Adam, non-branch = 1.0 no-op) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), @@ -29399,6 +29405,7 @@ impl GpuDqnTrainer { .arg(&_aux_diag_slot_disabled) .arg(&_aux_engage_buf_null) .arg(&_aux_engage_off_disabled) + .arg(&1.0_f32) // SP21 Phase 4: lr_scale (denoise aux trainer, non-branch = 1.0 no-op) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), @@ -32139,49 +32146,102 @@ impl GpuDqnTrainer { let value_byte_off = padded_byte_offset(¶m_sizes, 13); let value_byte_end = padded_byte_offset(¶m_sizes, 17); let value_count = ((value_byte_end - value_byte_off) / f32_sz) as usize; - let branches_byte_off = padded_byte_offset(¶m_sizes, 17); - let branches_byte_end = padded_byte_offset(¶m_sizes, 33); - let branches_count = ((branches_byte_end - branches_byte_off) / f32_sz) as usize; + + // SP21 T2.2 Phase 4 (2026-05-11): per-branch byte ranges within the + // canonical DqnBranches range [17..33). Each branch occupies 4 param + // tensors (W_FC, B_FC, W_OUT, B_OUT). Splitting the single + // DqnBranches sub-launch into 4 per-branch sub-launches lets each + // branch consume its own E4 LR scale from + // ISV[BRANCH_LR_SCALE_{DIR,MAG,ORDER,URGENCY}_INDEX]. Coverage + // invariant unchanged — total span is still [17..33) without gaps. + let b_dir_off = padded_byte_offset(¶m_sizes, 17); + let b_dir_end = padded_byte_offset(¶m_sizes, 21); + let b_mag_off = b_dir_end; + let b_mag_end = padded_byte_offset(¶m_sizes, 25); + let b_order_off = b_mag_end; + let b_order_end = padded_byte_offset(¶m_sizes, 29); + let b_urgency_off = b_order_end; + let b_urgency_end = padded_byte_offset(¶m_sizes, 33); + let b_dir_count = ((b_dir_end - b_dir_off) / f32_sz) as usize; + let b_mag_count = ((b_mag_end - b_mag_off) / f32_sz) as usize; + let b_order_count = ((b_order_end - b_order_off) / f32_sz) as usize; + let b_urgency_count = ((b_urgency_end - b_urgency_off) / f32_sz) as usize; + let branches_count = b_dir_count + b_mag_count + b_order_count + b_urgency_count; + // Trunk-extras tail: tensors [33..163). Spans from end of branches // to end of params buffer. let trunk_extras_byte_off = padded_byte_offset(¶m_sizes, 33); let trunk_extras_count = self.total_params .saturating_sub((trunk_extras_byte_off / f32_sz) as usize); - // (group_tag, byte_off, count, engage_offset). The 4th entry is - // tagged as DqnTrunk (pragma: shares trunk's WEIGHT_BOUND/WD_RATE - // contract) but uses SP4_ENGAGE_OFFSET_DISABLED so its block - // engagement does not overwrite trunk's slot in the Pearl C buf. + // SP21 T2.2 Phase 4 (2026-05-11): read per-branch LR scales from + // ISV[521..525). Cold-start sentinel (slot at 0.0 before E4 + // first emits) is floored at 1.0 so the cold-start Adam update + // doesn't degenerate to lr=0. Once E4 writes a real value + // (range [0.5, 2.0] per its `compute_branch_lr_scale` clamp), + // the floor is a no-op. Order matches branch indexing convention + // (0=Dir, 1=Mag, 2=Order, 3=Urgency). + let read_branch_lr_scale = |branch_idx: usize| -> f32 { + use crate::cuda_pipeline::sp21_isv_slots::branch_lr_scale_index; + let v = self.read_isv_signal_at(branch_lr_scale_index(branch_idx)); + // Bootstrap floor: ISV at sentinel 0.0 → 1.0 (no-op). + // Per `pearl_first_observation_bootstrap` discipline. + if v.abs() <= 1e-6 { 1.0 } else { v } + }; + let lr_scale_dir = read_branch_lr_scale(0); + let lr_scale_mag = read_branch_lr_scale(1); + let lr_scale_order = read_branch_lr_scale(2); + let lr_scale_urgency = read_branch_lr_scale(3); + + // (group_tag, byte_off, count, engage_offset, lr_scale). All + // non-branch entries pass lr_scale=1.0 (no-op for trunk / value + // / trunk-extras / aux trainers). The 4 branch sub-launches use + // SP4_ENGAGE_OFFSET_DISABLED for Pearl C — splitting the single + // DqnBranches engagement-tracking buffer across 4 sub-launches + // would require adding 3 ParamGroup variants OR a sub-block + // offsetting scheme; both deferred to a Phase 4.5 follow-up + // commit. Per `feedback_no_partial_refactor` Pearl C is a + // diagnostic system (not a load-bearing controller), so its + // temporary unavailability for branches is acceptable as long + // as Trunk/Value/Trunk-extras Pearl C still fires. let trunk_engage_off: i32 = (SP4Group::DqnTrunk.idx() * MAX_BLOCKS_PER_ADAM) as i32; let value_engage_off: i32 = (SP4Group::DqnValue.idx() * MAX_BLOCKS_PER_ADAM) as i32; - let branches_engage_off: i32 = (SP4Group::DqnBranches.idx() * MAX_BLOCKS_PER_ADAM) as i32; - let groups: [(SP4Group, u64, usize, i32); 4] = [ - (SP4Group::DqnTrunk, trunk_byte_off, trunk_count, trunk_engage_off), - (SP4Group::DqnValue, value_byte_off, value_count, value_engage_off), - (SP4Group::DqnBranches, branches_byte_off, branches_count, branches_engage_off), + let groups: [(SP4Group, u64, usize, i32, f32); 7] = [ + (SP4Group::DqnTrunk, trunk_byte_off, trunk_count, trunk_engage_off, 1.0), + (SP4Group::DqnValue, value_byte_off, value_count, value_engage_off, 1.0), + // SP21 Phase 4: 4 per-branch sub-launches. + (SP4Group::DqnBranches, b_dir_off, b_dir_count, SP4_ENGAGE_OFFSET_DISABLED, lr_scale_dir), + (SP4Group::DqnBranches, b_mag_off, b_mag_count, SP4_ENGAGE_OFFSET_DISABLED, lr_scale_mag), + (SP4Group::DqnBranches, b_order_off, b_order_count, SP4_ENGAGE_OFFSET_DISABLED, lr_scale_order), + (SP4Group::DqnBranches, b_urgency_off, b_urgency_count, SP4_ENGAGE_OFFSET_DISABLED, lr_scale_urgency), // Trunk-extras tail — shares DqnTrunk bounds; Pearl C disabled // for this sub-launch (engagement folds into trunk's // accounting since the same WEIGHT_BOUND drives both clamps). - (SP4Group::DqnTrunk, trunk_extras_byte_off, trunk_extras_count, SP4_ENGAGE_OFFSET_DISABLED), + (SP4Group::DqnTrunk, trunk_extras_byte_off, trunk_extras_count, SP4_ENGAGE_OFFSET_DISABLED, 1.0), ]; // Coverage invariant: every of the 163 weight tensors must be - // covered by exactly one of the 4 sub-launches above. Layer B + // covered by exactly one of the 7 sub-launches above. Layer B // fix-up: prevents regression of the silent ~70% Adam freeze - // that the original 3-way split introduced. + // that the original 3-way split introduced. SP21 Phase 4 + // expanded from 4 to 7 sub-launches (DqnBranches split into 4 + // per-branch); the coverage sum remains + // trunk + value + branches_total + trunk_extras. debug_assert_eq!( trunk_count + value_count + branches_count + trunk_extras_count, self.total_params, - "SP4 Layer B fix-up: 4-way Adam coverage must span all 163 weight tensors \ - (got trunk={} + value={} + branches={} + trunk_extras={} = {}, expected total_params={})", - trunk_count, value_count, branches_count, trunk_extras_count, + "SP21 Phase 4 7-way Adam coverage must span all 163 weight tensors \ + (got trunk={} + value={} + branches_total={} (dir={}+mag={}+order={}+urgency={}) + trunk_extras={} = {}, expected total_params={})", + trunk_count, value_count, branches_count, + b_dir_count, b_mag_count, b_order_count, b_urgency_count, + trunk_extras_count, trunk_count + value_count + branches_count + trunk_extras_count, self.total_params ); // Safety: argument order matches the extern "C" kernel signature exactly. // grad_norm_buf contains L2 norm from launch_grad_norm_finalize(). - for (group, byte_off, count, engage_buf_offset) in groups { + for (group, byte_off, count, engage_buf_offset, lr_scale) in groups { if count == 0 { continue; } @@ -32263,6 +32323,7 @@ impl GpuDqnTrainer { .arg(&diag_slot) // SP4 A14: diag slot (50 for dqn_adam) .arg(&engage_buf_ptr) // SP4 A14: per-block engagement buf .arg(&engage_buf_offset) // SP4 Layer B: per-group offset; SP4_ENGAGE_OFFSET_DISABLED for trunk-extras tail + .arg(&lr_scale) // SP21 Phase 4: per-branch LR scale (1.0 for non-branch groups) .launch(cfg) .map_err(|e| { MLError::ModelError(format!( @@ -34109,6 +34170,7 @@ impl GpuDqnTrainer { .arg(&_aux_diag_slot_disabled) .arg(&_aux_engage_buf_null) .arg(&_aux_engage_off_disabled) + .arg(&1.0_f32) // SP21 Phase 4: lr_scale (sel/recursive_conf aux trainer, non-branch = 1.0 no-op) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), diff --git a/crates/ml/src/cuda_pipeline/sp21_isv_slots.rs b/crates/ml/src/cuda_pipeline/sp21_isv_slots.rs index 9d11e845d..19e5c83e5 100644 --- a/crates/ml/src/cuda_pipeline/sp21_isv_slots.rs +++ b/crates/ml/src/cuda_pipeline/sp21_isv_slots.rs @@ -11,6 +11,10 @@ //! | Index | Name | Producer | Consumer | //! |-------|------|----------|----------| //! | 520 | `Q_CORRECTION_INDEX` | `enrichment::compute_q_correction` (host, post-val) | `mse_loss_kernel` + `c51_loss_kernel::block_bellman_project_f` (Bellman target offset) | +//! | 521 | `BRANCH_LR_SCALE_DIR_INDEX` | `enrichment::compute_branch_lr_scale[0]` | `dqn_adam_update_kernel` per-branch sub-launch (dir tensors 17-20) | +//! | 522 | `BRANCH_LR_SCALE_MAG_INDEX` | `enrichment::compute_branch_lr_scale[1]` | `dqn_adam_update_kernel` per-branch sub-launch (mag tensors 21-24) | +//! | 523 | `BRANCH_LR_SCALE_ORDER_INDEX` | `enrichment::compute_branch_lr_scale[2]` | `dqn_adam_update_kernel` per-branch sub-launch (order tensors 25-28) | +//! | 524 | `BRANCH_LR_SCALE_URGENCY_INDEX` | `enrichment::compute_branch_lr_scale[3]` | `dqn_adam_update_kernel` per-branch sub-launch (urgency tensors 29-32) | //! //! ## Bus extension //! @@ -47,12 +51,57 @@ /// → no-op until the first val pass writes a real correction. pub const Q_CORRECTION_INDEX: usize = 520; +/// SP21 T2.2 Phase 4 (2026-05-11) — per-branch Adam LR multiplicative +/// scale from E4 (`enrichment::compute_branch_lr_scale`). +/// +/// Producer: host-side after each validation backtest pass. E4 emits +/// `[f32; 4]` indexed by branch (0=Dir, 1=Mag, 2=Order, 3=Urgency) +/// with each element clamped to `[0.5, 2.0]` based on the per-branch +/// loss / mismatch rate inferred from the per-trade tape's +/// (direction, magnitude, holding_bars, pnl) tuples. Branches with +/// higher loss rates get larger LR scales (faster correction); +/// branches with lower rates get smaller scales (less perturbation +/// of already-good behaviour). +/// +/// Consumer: `dqn_adam_update_kernel` via `launch_adam_update`'s +/// per-branch sub-launches (Phase 4 split DqnBranches' single +/// sub-launch into 4 per-branch sub-launches, each reading its own +/// `BRANCH_LR_SCALE_*_INDEX` slot and passing it as the new +/// `float lr_scale` kernel arg). The kernel multiplies its +/// effective LR by `lr_scale` before the Adam update step. +/// +/// Cold-start: E4 returns `[1.0; 4]` when no losers are observed +/// yet (empty `Vec` or all-winning trades), so the +/// initial bus values stay at sentinel 0.0 → kernel sees lr_scale +/// 0.0 which would zero out the update; we therefore floor the +/// kernel-side scale at 1.0 when the slot is at sentinel +/// (`if lr_scale_isv <= 1e-6 then 1.0 else lr_scale_isv`). +/// Matches `pearl_first_observation_bootstrap` discipline. +pub const BRANCH_LR_SCALE_DIR_INDEX: usize = 521; +pub const BRANCH_LR_SCALE_MAG_INDEX: usize = 522; +pub const BRANCH_LR_SCALE_ORDER_INDEX: usize = 523; +pub const BRANCH_LR_SCALE_URGENCY_INDEX: usize = 524; + +/// Convenience accessor: branch index (0..4) → ISV slot index. Used +/// by the launcher's per-branch sub-launch loop to keep the +/// branch-index ↔ ISV-slot mapping in one place. +#[inline] +pub fn branch_lr_scale_index(branch: usize) -> usize { + match branch { + 0 => BRANCH_LR_SCALE_DIR_INDEX, + 1 => BRANCH_LR_SCALE_MAG_INDEX, + 2 => BRANCH_LR_SCALE_ORDER_INDEX, + 3 => BRANCH_LR_SCALE_URGENCY_INDEX, + _ => panic!("branch_lr_scale_index: branch must be 0..4 (got {branch})"), + } +} + /// First SP21 slot. Used by the bus-bounds compile-time check below /// to assert SP21's slot range is reachable. pub const SP21_SLOT_START: usize = 520; /// One past the last SP21 slot. -pub const SP21_SLOT_END: usize = 521; +pub const SP21_SLOT_END: usize = 525; #[cfg(test)] mod tests { @@ -76,12 +125,24 @@ mod tests { #[test] fn sp21_slot_indices_are_unique() { - // One slot today; placeholder for when SP21 grows. The check - // becomes load-bearing once a second slot is added. - let slots = [Q_CORRECTION_INDEX]; + let slots = [ + Q_CORRECTION_INDEX, + BRANCH_LR_SCALE_DIR_INDEX, + BRANCH_LR_SCALE_MAG_INDEX, + BRANCH_LR_SCALE_ORDER_INDEX, + BRANCH_LR_SCALE_URGENCY_INDEX, + ]; let mut sorted = slots.to_vec(); sorted.sort(); sorted.dedup(); assert_eq!(slots.len(), sorted.len(), "SP21 ISV slot indices must be unique"); } + + #[test] + fn branch_lr_scale_index_matches_constants() { + assert_eq!(branch_lr_scale_index(0), BRANCH_LR_SCALE_DIR_INDEX); + assert_eq!(branch_lr_scale_index(1), BRANCH_LR_SCALE_MAG_INDEX); + assert_eq!(branch_lr_scale_index(2), BRANCH_LR_SCALE_ORDER_INDEX); + assert_eq!(branch_lr_scale_index(3), BRANCH_LR_SCALE_URGENCY_INDEX); + } } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 9d621bd9d..e6adf3dd5 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1588,12 +1588,31 @@ impl DQNTrainer { // Both kernels apply it as additive shift on the // Bellman target. Cold-start: `result.q_correction = // 0.0` (no trades yet) is the no-op sentinel. + // + // SP21 T2.2 Phase 4 (2026-05-11): apply E4 branch_lr_scale + // by writing to ISV[521..525) (DIR/MAG/ORDER/URGENCY). + // `launch_adam_update` reads these slots host-side and + // passes per-branch scalar lr_scale args to the 4 split + // DqnBranches sub-launches. Cold-start: enrichment + // returns `[1.0; 4]` when no losers observed; if the + // ISV bus is at sentinel 0.0 (pre-first-emit), the + // launcher floors at 1.0 → kernel sees no-op scale. if let Some(ref fused) = self.fused_ctx { - use crate::cuda_pipeline::sp21_isv_slots::Q_CORRECTION_INDEX; + use crate::cuda_pipeline::sp21_isv_slots::{ + Q_CORRECTION_INDEX, branch_lr_scale_index, + }; fused.trainer().write_isv_signal_at( Q_CORRECTION_INDEX, result.q_correction, ); + for (branch_idx, &scale) in + result.branch_lr_scale.iter().enumerate() + { + fused.trainer().write_isv_signal_at( + branch_lr_scale_index(branch_idx), + scale, + ); + } } } } diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index e5814b66c..a7242977f 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -13870,3 +13870,177 @@ T2.2 multi-phase scope continues with Phases 4-7 + Phase 8: - Phase 7: wire E8 (curriculum weights) → segment sampling. - Phase 8: signal-drive remaining controller GAINS (E5 0.9/1.1 steps; E2 2.0/0.5/-0.5 epsilon thresholds; etc.). + +## 2026-05-11 — SP21 T2.2 Phase 4: E4 branch_lr_scale → per-branch Adam sub-launches + +### Scope (atomic single commit) + +Wires `EnrichmentResult::branch_lr_scale` (E4's `[f32; 4]` clamped +`[0.5, 2.0]` per-branch LR multiplier from +`enrichment::compute_branch_lr_scale`) into the DQN Adam optimizer. +Splits the existing single `DqnBranches` Adam sub-launch into 4 +per-branch sub-launches, each consuming its own LR scale from +ISV[521..525). The kernel ABI gains one new `float lr_scale` arg +on `dqn_adam_update_kernel`; all 5 callers (main DQN + +ofi_embed/aux_trunk/denoise/sel post-aux trainers + decision +transformer) updated atomically per `feedback_no_partial_refactor`. + +### Plan amendment from on-paper design + +The 2026-05-11 plan said *"Wire E4 (per-branch LR scaling) via +existing per-group Adam infrastructure (`pearl_per_group_adam_ +hyperparams`)."* The pearl is about per-PARAM-GROUP β1/β2/ε +(8 groups: DqnTrunk, DqnValue, DqnBranches, Iqn, IqlHigh, IqlLow, +Attn, Curiosity), but E4 is per-action-BRANCH (4 branches: Dir, +Mag, Order, Urgency). All 4 action branches are lumped into the +single `DqnBranches` param group. + +**Resolution**: split the `DqnBranches` Adam sub-launch into 4 +per-branch sub-launches. Each branch occupies 4 param tensors +(W_FC, B_FC, W_OUT, B_OUT) at consecutive `param_sizes` indices: + +| Branch | Index | Param tensors | +|--------|-------|---------------| +| Dir (b0) | 0 | 17, 18, 19, 20 | +| Mag (b1) | 1 | 21, 22, 23, 24 | +| Order (b2) | 2 | 25, 26, 27, 28 | +| Urgency (b3) | 3 | 29, 30, 31, 32 | + +Coverage invariant unchanged: the 4 branch sub-launches still +span `[17..33)` byte range without gaps, just at finer granularity. +The total `groups: [(group, byte_off, count, engage_off, lr_scale); 7]` +array grows from 4 to 7 entries (3 unchanged + 4 branch); the +runtime 4-way coverage assertion is updated to enforce 7-way +coverage with branches summed. + +### Pearl C engagement tracking deferral + +The single `DqnBranches` engagement counter (offset +`group_idx * MAX_BLOCKS_PER_ADAM` in `clamp_engage_per_block_buf`) +is shared across all 4 branches in the prior layout. Splitting +into 4 sub-launches with the SAME engagement offset would cause +write collisions on overlapping block indices. + +For Phase 4, all 4 branch sub-launches use +`SP4_ENGAGE_OFFSET_DISABLED` (matches the existing trunk-extras +tail pattern). This means Pearl C engagement tracking is +TEMPORARILY UNAVAILABLE for branch params. Pearl C remains active +for `Trunk`, `Value`, `Trunk-extras` (folds into Trunk), and the +8 group-level `WEIGHT_BOUND[group]` consumers — only the per-block +engagement-rate-deficit EMA is paused for branches. + +Pearl C is a DIAGNOSTIC system per its docstring (rate-deficit +EMA logged when sustained at >0.5%; not a load-bearing +controller). Its branch coverage will be re-instated in a Phase +4.5 follow-up commit either by (a) growing `ParamGroup` enum to +add 4 sub-variants `DqnBranchDir/Mag/Order/Urgency` (allocates +4 distinct engagement offsets), or (b) introducing a sub-block +offsetting scheme within `DqnBranches`'s `MAX_BLOCKS_PER_ADAM=256` +range. Per `feedback_no_partial_refactor`, Pearl C re-instatement +will be its own atomic commit when the design is chosen. + +### ISV slot allocation + +| Index | Name | Producer | Consumer | +|-------|------|----------|----------| +| 521 | `BRANCH_LR_SCALE_DIR_INDEX` | `enrichment::compute_branch_lr_scale[0]` | `launch_adam_update` per-branch sub-launch (Dir tensors 17-20) | +| 522 | `BRANCH_LR_SCALE_MAG_INDEX` | `enrichment::compute_branch_lr_scale[1]` | `launch_adam_update` per-branch sub-launch (Mag tensors 21-24) | +| 523 | `BRANCH_LR_SCALE_ORDER_INDEX` | `enrichment::compute_branch_lr_scale[2]` | `launch_adam_update` per-branch sub-launch (Order tensors 25-28) | +| 524 | `BRANCH_LR_SCALE_URGENCY_INDEX` | `enrichment::compute_branch_lr_scale[3]` | `launch_adam_update` per-branch sub-launch (Urgency tensors 29-32) | + +`ISV_TOTAL_DIM` bumped 521 → 525. Layout fingerprint adds +`SLOT_521_BRANCH_LR_SCALE_DIR`, `SLOT_522_BRANCH_LR_SCALE_MAG`, +`SLOT_523_BRANCH_LR_SCALE_ORDER`, `SLOT_524_BRANCH_LR_SCALE_URGENCY`. +New `branch_lr_scale_index(branch_idx)` accessor for clean mapping. + +### ABI surgery (kernel-side) + +**dqn_adam_update_kernel.cu**: ONE new arg `float lr_scale` at the +end. Inside the kernel: + +```c +float lr = active ? (*lr_ptr * lr_scale) : 0.0f; +``` + +The host-controlled global LR (`*lr_ptr`, mapped-pinned for +graph-safety) is multiplied by the per-launch `lr_scale`. All +non-branch sub-launches pass `lr_scale = 1.0` (no-op). + +### Cold-start bootstrap + +E4 returns `[1.0; 4]` when `losers.is_empty()` (empty +`Vec` or all winning trades). The first call to +`run_enrichments` writes `[1.0; 4]` to ISV[521..525). However, +the FIRST training step before any val pass reads ISV at sentinel +0.0 — directly multiplying lr_ptr × 0.0 would zero the update. + +**Mitigation**: the launcher reads `read_isv_signal_at(branch_lr_ +scale_index(b))` and floors at 1.0 if the ISV value is at sentinel: + +```rust +let v = self.read_isv_signal_at(branch_lr_scale_index(branch_idx)); +if v.abs() <= 1e-6 { 1.0 } else { v } +``` + +Once E4 emits a real value (clamped `[0.5, 2.0]`), the floor is a +no-op. Matches `pearl_first_observation_bootstrap` discipline — +sentinel 0 → first observation replaces directly with no +intermediate state. + +### Files changed + +| File | Status | Purpose | +|------|--------|---------| +| `crates/ml/src/cuda_pipeline/sp21_isv_slots.rs` | +4 slots | `BRANCH_LR_SCALE_{DIR,MAG,ORDER,URGENCY}_INDEX` constants; `branch_lr_scale_index` accessor; tests for uniqueness + accessor mapping | +| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | ISV bump + 7-way Adam split | `ISV_TOTAL_DIM` 521→525; fingerprint adds 4 SLOT entries; `launch_adam_update` splits DqnBranches into 4 per-branch sub-launches with per-branch lr_scale; coverage invariant updated to 7-way | +| `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu` | ABI bump | `dqn_adam_update_kernel` gains `float lr_scale` arg; `lr = *lr_ptr * lr_scale` | +| `crates/ml/src/cuda_pipeline/decision_transformer.rs` | ABI alignment | DT Adam launcher passes `lr_scale=1.0` (non-branch, no-op) | +| `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | Producer wireup | After enrichment block, writes `result.branch_lr_scale[i]` to ISV[521..525) | +| `docs/dqn-wire-up-audit.md` | This entry | 2026-05-11 audit log | + +### Pearls + invariants honoured + +- `feedback_no_partial_refactor` — kernel ABI change ships with + ALL 5 callers migrated atomically (main DQN + 4 post-aux + DT). +- `pearl_per_group_adam_hyperparams` — extends the per-group + pattern with per-branch granularity; Adam β1/β2/ε remain + per-group (no change to that pearl's contract); only LR scale + goes per-branch. +- `pearl_controller_anchors_isv_driven` — branch LR multipliers + live in ISV bus (slots 521-524), not as config constants. +- `feedback_isv_for_adaptive_bounds` — adaptive per-branch LR + IS adaptive bound on Adam updates; lives in ISV. +- `pearl_first_observation_bootstrap` — sentinel 0.0 → launcher + floor at 1.0 (no-op until first E4 emit). +- `feedback_no_atomicadd` — kernel still uses block tree-reduce + for grad clip + engagement (per existing convention). +- `feedback_no_stubs` — `branch_lr_scale[4]` consumed for real + by 4 distinct sub-launches; the prior commit's E1 wireup left + E4's output unconsumed, this commit closes that gap. + +### Verification + +``` +SQLX_OFFLINE=true cargo check -p ml --tests --features cuda # 0 errors +cargo test -p ml --lib sp21_isv_slots --features cuda # 3/3 (bus bounds + uniqueness + branch index) +cargo test -p ml --test sp20_aggregate_inputs_test ... # 12/12 +cargo test -p ml --test sp20_phase1_4_wireup_test ... # 2/2 +cargo test -p ml --test sp20_emas_compute_test ... # 4/4 +cargo test -p ml --test sp20_controllers_compute_test ... # 7/7 +cargo test -p ml --test sp21_per_trade_predicted_q_test ... # 3/3 +``` + +Total: 31 tests, 0 failures. Behavioral gate (per-branch LR +divergence — Dir branch trains 2× faster than Mag when E4 +identifies Dir as the high-loss branch) is the upcoming smoke +training run. + +### After this commit lands + +T2.2 multi-phase scope continues with Phases 5-7 + Phase 8: +- Phase 5: wire E6 (winner indices) → PER priority bumps. +- Phase 6: wire E7 (`HindsightExperience`) → replay buffer. +- Phase 7: wire E8 (curriculum weights) → segment sampling. +- Phase 8: signal-drive remaining controller GAINS. +- Phase 4.5 (deferred): re-instate Pearl C engagement tracking + for branches via ParamGroup expansion or sub-block offsetting.