gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(sp21): T2.2 Phase 4 — wire E4 branch_lr_scale → per-branch Adam (atomic)

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-10 22:44:43 +02:00
parent 21f911151a
commit b7c4f84ea0
6 changed files with 358 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

104
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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(&param_sizes, 13);
let value_byte_end = padded_byte_offset(&param_sizes, 17);
let value_count = ((value_byte_end - value_byte_off) / f32_sz) as usize;
let branches_byte_off = padded_byte_offset(&param_sizes, 17);
let branches_byte_end = padded_byte_offset(&param_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(&param_sizes, 17);
let b_dir_end = padded_byte_offset(&param_sizes, 21);
let b_mag_off = b_dir_end;
let b_mag_end = padded_byte_offset(&param_sizes, 25);
let b_order_off = b_mag_end;
let b_order_end = padded_byte_offset(&param_sizes, 29);
let b_urgency_off = b_order_end;
let b_urgency_end = padded_byte_offset(&param_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(&param_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),