feat(dqn-v2): Plan 4 Task 6 Commit B — aux heads behavioral wiring
Activates the dormant Commit A scaffolding: forward + backward + ISV
producer + aux loss combination + HEALTH_DIAG observability. Two
auxiliary heads (next-bar return MSE + 5-class regime CE) train end-
to-end off the trunk's h_s2 activation, with their gradient flowing
into the trunk's bw_d_h_s2 accumulator (ISV-scaled by aux_weight)
before encoder_backward_chain consumes it.
Wire-points:
- Forward (launch_cublas_forward, after forward_online_raw saves
save_h_s2, BEFORE stochastic depth): regime label builder →
strided_gather (col 0 of next_states → dedicated aux_nb_label_buf,
NOT aliased) → next-bar forward → regime forward → next-bar MSE
reduce → regime CE reduce.
- Backward (launch_cublas_backward_to, AFTER backward_full fills
bw_d_h_s2, BEFORE encoder_backward_chain): per-head backward
emits per-sample partials + per-sample dh_s2; aux_param_grad_reduce
collapses partials → final, then SAXPY into grad_buf[119..127);
both dh_s2 buffers SAXPY into bw_d_h_s2 with alpha=aux_weight.
- ISV producer (training_loop.rs per-step): launch_aux_heads_loss_ema
alongside launch_h_s2_rms_ema / launch_vsn_mask_ema; updates
ISV[113][114] EMAs.
- Aux-weight refresh (training_loop.rs per-step): clamp(0.1 ×
LEARNING_HEALTH × (1 − tanh(0.1 × sharpe_ema)), 0.05, 0.3); pushed
via FusedTrainingCtx::set_aux_weight delegate.
- HEALTH_DIAG aux clause: emits per-epoch aux[next_bar_mse=…
regime_ce=… w=…] from the ISV slots populated by the producer.
Lessons applied from previous WIP attempt (stashed):
- Dedicated aux_nb_label_buf [B] f32 field (NOT aliased over
aux_partial_nb_b2) — partial-refactor invariant honoured cleanly.
- Init log "AuxHeadsForwardOps initialized: K_nb=1 K_rg=5 aux_h=32 …"
for observability sanity check.
- HEALTH_DIAG aux clause verified firing in smoke (5 lines, one per
epoch, with non-zero loss EMAs and ISV-driven aux_weight evolution
0.050 → 0.105 over 5 epochs).
Smoke (multi_fold_convergence --release, 3 folds × 5 epochs on RTX
3050 Ti, 684.37s, all 3 dqn_fold{N}_best.safetensors written):
F0=-9.7831 F1=71.5327 F2=65.9598
Within baseline noise band [F1: 61.10–74.56, F2: 61.57–88.20] from
two clean Commit A baseline runs. F0 is bit-reproducible across all
runs (deterministic cold-start collapse under --release profile —
pre-existing, not from this task).
HEALTH_DIAG observability (5 epochs):
aux [next_bar_mse=9.962e-5 regime_ce=2.099e-5 w=0.050]
aux [next_bar_mse=2.295e-4 regime_ce=3.924e-5 w=0.088]
aux [next_bar_mse=3.261e-4 regime_ce=6.893e-5 w=0.099]
aux [next_bar_mse=4.307e-4 regime_ce=6.819e-5 w=0.102]
aux [next_bar_mse=4.824e-4 regime_ce=6.634e-5 w=0.105]
Aux losses grow as expected (cold-start Xavier predictions diverge
from labels until trunk learns); aux_weight rises with LEARNING_HEALTH
and sharpe_ema; nothing pinned at clamp boundaries — formula working.
Constraints honoured: GPU-only (every reduce + SAXPY + EMA on-device,
zero DtoH in hot path; only HEALTH_DIAG cold-path read on CPU); no
atomicAdd; no stubs; no // ok: band-aids; no buffer aliasing tricks;
no tuned constants beyond the documented 0.1 aux base + [0.05, 0.3]
numerical-stability clamp; partial-refactor invariant honoured
(bw_d_h_s2 accumulator semantics extended, not replaced;
encoder_backward_chain consumes through unchanged pointer).
target_ema_update NOT extended for aux heads (Commit A design
decision preserved): aux heads are online-only supervised heads, NOT
used in Bellman bootstrapping. target_params_buf[119..127) remains at
Xavier init forever.
cargo check clean at 11 warnings (workspace baseline preserved); no
fingerprint change (Commit A already shifted to 0x26f7b1deb94cb226).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -59,6 +59,9 @@ use super::gpu_weights::{DuelingWeightSet, BranchingWeightSet};
|
||||
use super::batched_forward::{CublasForward, CublasGemmSet, f32_weight_ptrs_from_base};
|
||||
use super::batched_backward::{CublasBackward, CublasBackwardSet, alloc_backward_scratch, raw_f32_ptr as bw_raw_f32_ptr};
|
||||
use super::shared_cublas_handle::PerStreamCublasHandles;
|
||||
use super::gpu_aux_heads::{
|
||||
AuxHeadsBackwardOps, AuxHeadsForwardOps, AUX_HIDDEN_DIM, AUX_NEXT_BAR_K, AUX_REGIME_K,
|
||||
};
|
||||
|
||||
// ── Precompiled cubins (build.rs → include_bytes! → ZERO runtime nvcc) ──────
|
||||
pub(crate) static DQN_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dqn_utility_kernels.cubin"));
|
||||
@@ -2473,6 +2476,91 @@ pub struct GpuDqnTrainer {
|
||||
/// launch wired in `training_loop.rs` next to `launch_h_s2_rms_ema`.
|
||||
vsn_mask_ema_kernel: CudaFunction,
|
||||
|
||||
// ── Plan 4 Task 6 Commit B: aux-heads forward + backward orchestrators ──
|
||||
/// Forward orchestrator (next-bar regression + 5-class regime CE).
|
||||
/// Owns the 5 forward / loss-reduce / label-builder kernel handles + the
|
||||
/// `aux_heads_loss_ema_update` ISV producer handle. Construction loads
|
||||
/// from `AUX_HEADS_CUBIN` + `AUX_HEADS_LOSS_EMA_CUBIN`.
|
||||
aux_heads_fwd: AuxHeadsForwardOps,
|
||||
/// Backward orchestrator — 3 kernel handles (next_bar_backward,
|
||||
/// regime_backward, param_grad_reduce). Consumed by `aux_heads_backward`
|
||||
/// (called from `launch_cublas_backward_to` BEFORE `encoder_backward_chain`
|
||||
/// so the SAXPY accumulating `aux_dh_s2_*` into `bw_d_h_s2` lands before
|
||||
/// the trunk backward consumes that accumulator).
|
||||
aux_heads_bwd: AuxHeadsBackwardOps,
|
||||
/// `[B, AUX_HIDDEN_DIM]` saved post-ELU activation of the next-bar head's
|
||||
/// Linear_1 (consumed by `aux_next_bar_backward` for the dW1 / dh_s2 chain).
|
||||
aux_nb_hidden_buf: CudaSlice<f32>,
|
||||
/// `[B, AUX_NEXT_BAR_K=1]` scalar prediction per sample. Read by
|
||||
/// `aux_next_bar_loss_reduce` (vs `next_states[:, 0]` label) AND by
|
||||
/// `aux_next_bar_backward` (for the `(2/B)*(pred-label)` derivative).
|
||||
aux_nb_pred_buf: CudaSlice<f32>,
|
||||
/// `[B]` per-sample next-bar return label, gathered from
|
||||
/// `next_states_buf[:, 0]` via `strided_gather` at the start of every
|
||||
/// forward call. DEDICATED buffer (no aliasing — the previous WIP
|
||||
/// over-aliased `aux_partial_nb_b2` here and triggered a regression).
|
||||
/// Consumed by both `aux_next_bar_loss_reduce` and `aux_next_bar_backward`.
|
||||
aux_nb_label_buf: CudaSlice<f32>,
|
||||
/// `[1]` scalar mean MSE — written by `aux_next_bar_loss_reduce`, read by
|
||||
/// `aux_heads_loss_ema_update` (the ISV producer).
|
||||
aux_nb_loss_scalar_buf: CudaSlice<f32>,
|
||||
/// `[B, AUX_HIDDEN_DIM]` saved post-ELU activation of the regime head's
|
||||
/// Linear_1 (consumed by `aux_regime_backward`).
|
||||
aux_rg_hidden_buf: CudaSlice<f32>,
|
||||
/// `[B, AUX_REGIME_K=5]` pre-softmax logits — written by
|
||||
/// `aux_regime_forward`, read by `aux_regime_loss_reduce` AND by
|
||||
/// `aux_regime_backward` (which re-softmaxes on the fly).
|
||||
aux_rg_logits_buf: CudaSlice<f32>,
|
||||
/// `[B] uint8` per-sample regime labels in `[0, AUX_REGIME_K)`. Produced
|
||||
/// per-step by `aux_regime_label_from_states` (reads
|
||||
/// `states_buf[:, OFI_START + 29]` = `regime_score` ∈ [0,1], bucketed
|
||||
/// into 5 uniform bins). Consumed by both `aux_regime_loss_reduce` and
|
||||
/// `aux_regime_backward`.
|
||||
aux_rg_label_buf: CudaSlice<u8>,
|
||||
/// `[1]` scalar mean cross-entropy — written by `aux_regime_loss_reduce`.
|
||||
aux_rg_loss_scalar_buf: CudaSlice<f32>,
|
||||
/// `[1]` scalar mean classification accuracy — written by
|
||||
/// `aux_regime_loss_reduce` alongside the CE loss. Currently unread by
|
||||
/// downstream consumers; kept as a real device buffer since the kernel's
|
||||
/// output ABI requires it.
|
||||
#[allow(dead_code)]
|
||||
aux_rg_correct_scalar_buf: CudaSlice<f32>,
|
||||
/// `[B, SH2]` per-sample dh_s2 emitted by `aux_next_bar_backward`. SAXPY-
|
||||
/// accumulated into `bw_d_h_s2` before `encoder_backward_chain` runs.
|
||||
aux_dh_s2_nb_buf: CudaSlice<f32>,
|
||||
/// `[B, SH2]` per-sample dh_s2 emitted by `aux_regime_backward`. Same
|
||||
/// SAXPY-accumulation pattern as `aux_dh_s2_nb_buf`.
|
||||
aux_dh_s2_rg_buf: CudaSlice<f32>,
|
||||
/// `[B, AUX_HIDDEN_DIM, SH2]` per-sample dW1 partial — next-bar head.
|
||||
aux_partial_nb_w1: CudaSlice<f32>,
|
||||
/// `[B, AUX_HIDDEN_DIM]` per-sample db1 partial — next-bar head.
|
||||
aux_partial_nb_b1: CudaSlice<f32>,
|
||||
/// `[B, AUX_HIDDEN_DIM]` per-sample dW2 partial — next-bar head (W2 is
|
||||
/// `[1, H]` flat-stored).
|
||||
aux_partial_nb_w2: CudaSlice<f32>,
|
||||
/// `[B]` per-sample db2 partial — next-bar head (b2 is scalar).
|
||||
/// DEDICATED — does NOT alias `aux_nb_label_buf`.
|
||||
aux_partial_nb_b2: CudaSlice<f32>,
|
||||
/// `[B, AUX_HIDDEN_DIM, SH2]` per-sample dW1 partial — regime head.
|
||||
aux_partial_rg_w1: CudaSlice<f32>,
|
||||
/// `[B, AUX_HIDDEN_DIM]` per-sample db1 partial — regime head.
|
||||
aux_partial_rg_b1: CudaSlice<f32>,
|
||||
/// `[B, AUX_REGIME_K, AUX_HIDDEN_DIM]` per-sample dW2 partial — regime head.
|
||||
aux_partial_rg_w2: CudaSlice<f32>,
|
||||
/// `[B, AUX_REGIME_K]` per-sample db2 partial — regime head.
|
||||
aux_partial_rg_b2: CudaSlice<f32>,
|
||||
/// `[max_aux_param_len]` reusable per-tensor reduce target for
|
||||
/// `aux_param_grad_reduce`. Sized to the largest of the 8 aux tensors.
|
||||
aux_param_grad_final_buf: CudaSlice<f32>,
|
||||
/// Per-step aux-loss weight, baked into the SAXPY alpha at graph-capture
|
||||
/// time. ISV-driven via `set_aux_weight`:
|
||||
/// `aux_weight = clamp(0.1 × LEARNING_HEALTH × (1 - tanh(SHARPE_SCALE × sharpe_ema)), 0.05, 0.3)`.
|
||||
/// The clamp `[0.05, 0.3]` is a numerical-stability bound (Invariant 1
|
||||
/// carve-out per `feedback_isv_for_adaptive_bounds.md`) — `0.05` keeps
|
||||
/// the gradient signal alive when `learning_health == 0`, `0.3` caps the
|
||||
/// share so aux can never dominate the policy gradient.
|
||||
aux_weight: f32,
|
||||
|
||||
/// #21 Stochastic depth: per-layer scale buffer [3] (h_s1, h_s2, h_v).
|
||||
/// Written by GPU RNG kernel before each graph replay.
|
||||
stochastic_depth_scale_buf: CudaSlice<f32>,
|
||||
@@ -7519,6 +7607,335 @@ impl GpuDqnTrainer {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Plan 4 Task 6 Commit B: set the per-step ISV-driven aux-loss weight.
|
||||
///
|
||||
/// `aux_weight = clamp(0.1 × ISV[LEARNING_HEALTH] × (1 - tanh(SHARPE_SCALE × sharpe_ema)), 0.05, 0.3)`
|
||||
/// is computed CPU-side in the training loop. Same staleness contract as
|
||||
/// `set_c51_alpha`: the value is baked into kernel SAXPY alphas at
|
||||
/// graph-capture time. The clamp `[0.05, 0.3]` is a numerical-stability
|
||||
/// bound (Invariant 1 carve-out per `feedback_isv_for_adaptive_bounds.md`).
|
||||
pub fn set_aux_weight(&mut self, aux_weight: f32) {
|
||||
// Defensive re-clamp at the kernel boundary in case the caller forgot.
|
||||
self.aux_weight = aux_weight.clamp(0.05, 0.3);
|
||||
}
|
||||
|
||||
/// Plan 4 Task 6 Commit B: current aux-loss weight (the value baked into
|
||||
/// the most-recently-captured graph). HEALTH_DIAG reads this to surface
|
||||
/// the per-epoch `aux[w=…]` field.
|
||||
pub fn aux_weight(&self) -> f32 {
|
||||
self.aux_weight
|
||||
}
|
||||
|
||||
/// Plan 4 Task 6 Commit B: launch `aux_heads_loss_ema_update` (single-
|
||||
/// thread, single-block ISV producer kernel). Reads `aux_nb_loss_scalar_buf`
|
||||
/// + `aux_rg_loss_scalar_buf` (both populated by the corresponding
|
||||
/// `aux_*_loss_reduce` kernels during the training-step forward) and EMA-
|
||||
/// updates ISV[AUX_NEXT_BAR_MSE_EMA_INDEX=113] + ISV[AUX_REGIME_CE_EMA_INDEX=114].
|
||||
pub fn launch_aux_heads_loss_ema(&self, ema_alpha: f32) -> Result<(), MLError> {
|
||||
debug_assert!(self.isv_signals_dev_ptr != 0,
|
||||
"launch_aux_heads_loss_ema: isv_signals_dev_ptr must be allocated by constructor");
|
||||
self.aux_heads_fwd.launch_loss_ema(
|
||||
&self.stream,
|
||||
self.aux_nb_loss_scalar_buf.raw_ptr(),
|
||||
self.aux_rg_loss_scalar_buf.raw_ptr(),
|
||||
self.isv_signals_dev_ptr,
|
||||
AUX_NEXT_BAR_MSE_EMA_INDEX as i32,
|
||||
AUX_REGIME_CE_EMA_INDEX as i32,
|
||||
ema_alpha,
|
||||
)
|
||||
}
|
||||
|
||||
/// Plan 4 Task 6 Commit B: aux-heads forward orchestrator.
|
||||
///
|
||||
/// Runs the next-bar regression head + 5-class regime classification head
|
||||
/// off the trunk's just-produced `save_h_s2 [B, SH2]`. Sequence:
|
||||
/// 1. regime label builder: `states[:, OFI_START + 29]` → uint8 [B]
|
||||
/// 2. next-bar label gather: `next_states[:, 0]` → f32 [B] dedicated buf
|
||||
/// 3. next-bar forward: Linear → ELU → Linear → pred[B, 1]
|
||||
/// 4. regime forward: Linear → ELU → Linear → logits[B, 5]
|
||||
/// 5. next-bar MSE reduce → loss_scalar[1]
|
||||
/// 6. regime CE reduce → loss_scalar[1] + correct_scalar[1]
|
||||
///
|
||||
/// Online-only — aux heads NEVER enter Bellman bootstrapping.
|
||||
pub(crate) fn aux_heads_forward(&self) -> Result<(), MLError> {
|
||||
let b = self.config.batch_size;
|
||||
let sh2 = self.config.shared_h2;
|
||||
let h_s2_ptr = self.ptrs.save_h_s2;
|
||||
let states_ptr = self.ptrs.states_buf;
|
||||
let next_states_ptr = self.ptrs.next_states_buf;
|
||||
|
||||
let param_sizes = compute_param_sizes(&self.config);
|
||||
let on_w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes);
|
||||
// Aux param tensors at indices [119..127): nb_{w1, b1, w2, b2} + rg_{w1, b1, w2, b2}.
|
||||
let nb_w1 = on_w_ptrs[119];
|
||||
let nb_b1 = on_w_ptrs[120];
|
||||
let nb_w2 = on_w_ptrs[121];
|
||||
let nb_b2 = on_w_ptrs[122];
|
||||
let rg_w1 = on_w_ptrs[123];
|
||||
let rg_b1 = on_w_ptrs[124];
|
||||
let rg_w2 = on_w_ptrs[125];
|
||||
let rg_b2 = on_w_ptrs[126];
|
||||
|
||||
let state_dim_padded = ml_core::state_layout::STATE_DIM_PADDED;
|
||||
let regime_score_idx = ml_core::state_layout::OFI_START + 29;
|
||||
|
||||
// Step 1: regime-label builder. `regime_score` is at OFI offset 29
|
||||
// (= absolute index 71 = OFI_START + 29). It's a sigmoid-bounded
|
||||
// [0, 1] scalar per docs. The kernel discretises into K=5 uniform bins
|
||||
// with internal clamp.
|
||||
self.aux_heads_fwd.regime_label_from_states(
|
||||
&self.stream,
|
||||
states_ptr,
|
||||
b,
|
||||
state_dim_padded,
|
||||
regime_score_idx,
|
||||
AUX_REGIME_K,
|
||||
self.aux_rg_label_buf.raw_ptr(),
|
||||
)?;
|
||||
|
||||
// Step 2: next-bar regression label = column 0 of next_states_buf
|
||||
// (= `log_return` per `pipeline.rs`; first feature in `MARKET` group).
|
||||
// Use `strided_gather` to extract into the DEDICATED `aux_nb_label_buf`
|
||||
// (NOT aliased over a backward partial — the previous WIP aliased
|
||||
// `aux_partial_nb_b2` here and triggered the F1/F2 regression).
|
||||
// strided_gather signature: (src, dst, src_lda, col, dst_stride, total).
|
||||
{
|
||||
let label_dst_ptr = self.aux_nb_label_buf.raw_ptr();
|
||||
let src_lda: i32 = state_dim_padded as i32;
|
||||
let col: i32 = 0;
|
||||
let dst_stride: i32 = 1;
|
||||
let total: i32 = b as i32;
|
||||
let blocks = ((b as u32 + 255) / 256).max(1);
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.vsn_logit_gather_kernel)
|
||||
.arg(&next_states_ptr)
|
||||
.arg(&label_dst_ptr)
|
||||
.arg(&src_lda)
|
||||
.arg(&col)
|
||||
.arg(&dst_stride)
|
||||
.arg(&total)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"aux_heads next-bar label gather (col 0 of next_states): {e}"
|
||||
)))?;
|
||||
}
|
||||
}
|
||||
|
||||
// Step 3: next-bar regression forward.
|
||||
self.aux_heads_fwd.forward_next_bar(
|
||||
&self.stream,
|
||||
h_s2_ptr,
|
||||
nb_w1, nb_b1, nb_w2, nb_b2,
|
||||
b, sh2,
|
||||
self.aux_nb_hidden_buf.raw_ptr(),
|
||||
self.aux_nb_pred_buf.raw_ptr(),
|
||||
)?;
|
||||
|
||||
// Step 4: regime classification forward.
|
||||
self.aux_heads_fwd.forward_regime(
|
||||
&self.stream,
|
||||
h_s2_ptr,
|
||||
rg_w1, rg_b1, rg_w2, rg_b2,
|
||||
b, sh2, AUX_REGIME_K,
|
||||
self.aux_rg_hidden_buf.raw_ptr(),
|
||||
self.aux_rg_logits_buf.raw_ptr(),
|
||||
)?;
|
||||
|
||||
// Step 5: next-bar MSE reduce against the dedicated label buffer.
|
||||
self.aux_heads_fwd.next_bar_loss_reduce(
|
||||
&self.stream,
|
||||
self.aux_nb_pred_buf.raw_ptr(),
|
||||
self.aux_nb_label_buf.raw_ptr(),
|
||||
b,
|
||||
self.aux_nb_loss_scalar_buf.raw_ptr(),
|
||||
)?;
|
||||
|
||||
// Step 6: regime CE reduce.
|
||||
self.aux_heads_fwd.regime_loss_reduce(
|
||||
&self.stream,
|
||||
self.aux_rg_logits_buf.raw_ptr(),
|
||||
self.aux_rg_label_buf.raw_ptr(),
|
||||
b,
|
||||
AUX_REGIME_K,
|
||||
self.aux_rg_loss_scalar_buf.raw_ptr(),
|
||||
self.aux_rg_correct_scalar_buf.raw_ptr(),
|
||||
)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Plan 4 Task 6 Commit B: aux-heads backward orchestrator.
|
||||
///
|
||||
/// Runs both backward kernels (next-bar + regime) over the saved forward
|
||||
/// state, then for each of the 8 aux param tensors:
|
||||
/// 1. `aux_param_grad_reduce` collapses `partials [B, P]` into `final [P]`.
|
||||
/// 2. SAXPY `grad_buf[tensor_idx] += aux_weight * final` via the
|
||||
/// existing `dqn_saxpy_f32_kernel`.
|
||||
/// Finally, the per-sample `dh_s2` from both heads is SAXPY-accumulated
|
||||
/// into `bw_d_h_s2 [B, SH2]` with `alpha = aux_weight`.
|
||||
///
|
||||
/// MUST run AFTER `backward_full` + concat-accumulators fill `bw_d_h_s2`
|
||||
/// AND BEFORE `encoder_backward_chain` consumes it.
|
||||
///
|
||||
/// `grad_base` is the same value passed to `launch_cublas_backward_to` —
|
||||
/// usually `self.ptrs.grad_buf`, but the vaccine path (g_val computation)
|
||||
/// passes a scratch buffer here. The aux SAXPY targets `grad_base[119..127)`
|
||||
/// at the same `padded_byte_offset` math the rest of the backward uses.
|
||||
pub(crate) fn aux_heads_backward(&self, grad_base: u64) -> Result<(), MLError> {
|
||||
let b = self.config.batch_size;
|
||||
let sh2 = self.config.shared_h2;
|
||||
let h_s2_ptr = self.ptrs.save_h_s2;
|
||||
let bw_d_h_s2_ptr = self.bw_d_h_s2.raw_ptr();
|
||||
let aux_w = self.aux_weight;
|
||||
let aux_h = AUX_HIDDEN_DIM;
|
||||
let aux_kr = AUX_REGIME_K;
|
||||
|
||||
let param_sizes = compute_param_sizes(&self.config);
|
||||
let on_w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes);
|
||||
let nb_w1 = on_w_ptrs[119];
|
||||
let nb_w2 = on_w_ptrs[121];
|
||||
let rg_w1 = on_w_ptrs[123];
|
||||
let rg_w2 = on_w_ptrs[125];
|
||||
|
||||
// Run BOTH backward kernels. They WRITE to disjoint output buffers
|
||||
// (per-head partials + per-head dh_s2), so no inter-dependency beyond
|
||||
// the just-completed forward.
|
||||
self.aux_heads_bwd.backward_next_bar(
|
||||
&self.stream,
|
||||
h_s2_ptr,
|
||||
nb_w1, nb_w2,
|
||||
self.aux_nb_hidden_buf.raw_ptr(),
|
||||
self.aux_nb_pred_buf.raw_ptr(),
|
||||
self.aux_nb_label_buf.raw_ptr(),
|
||||
b, sh2,
|
||||
self.aux_partial_nb_w1.raw_ptr(),
|
||||
self.aux_partial_nb_b1.raw_ptr(),
|
||||
self.aux_partial_nb_w2.raw_ptr(),
|
||||
self.aux_partial_nb_b2.raw_ptr(),
|
||||
self.aux_dh_s2_nb_buf.raw_ptr(),
|
||||
)?;
|
||||
self.aux_heads_bwd.backward_regime(
|
||||
&self.stream,
|
||||
h_s2_ptr,
|
||||
rg_w1, rg_w2,
|
||||
self.aux_rg_hidden_buf.raw_ptr(),
|
||||
self.aux_rg_logits_buf.raw_ptr(),
|
||||
self.aux_rg_label_buf.raw_ptr(),
|
||||
b, sh2, aux_kr,
|
||||
self.aux_partial_rg_w1.raw_ptr(),
|
||||
self.aux_partial_rg_b1.raw_ptr(),
|
||||
self.aux_partial_rg_w2.raw_ptr(),
|
||||
self.aux_partial_rg_b2.raw_ptr(),
|
||||
self.aux_dh_s2_rg_buf.raw_ptr(),
|
||||
)?;
|
||||
|
||||
// Reduce + SAXPY each of the 8 aux param tensors. Tensor layout (Commit A):
|
||||
// [119] aux_nb_w1 [H, SH2] ← partial [B, H*SH2]
|
||||
// [120] aux_nb_b1 [H] ← partial [B, H]
|
||||
// [121] aux_nb_w2 [1, H] ← partial [B, H]
|
||||
// [122] aux_nb_b2 [1] ← partial [B]
|
||||
// [123] aux_rg_w1 [H, SH2]
|
||||
// [124] aux_rg_b1 [H]
|
||||
// [125] aux_rg_w2 [K, H]
|
||||
// [126] aux_rg_b2 [K]
|
||||
// Note: kernels emit (1/B)*sum_b grad — `aux_param_grad_reduce`
|
||||
// sums per-sample partials, no extra division. SAXPY alpha is
|
||||
// `aux_weight` (small positive, [0.05, 0.3]).
|
||||
let final_ptr = self.aux_param_grad_final_buf.raw_ptr();
|
||||
let aux_param_specs: [(usize, u64, usize); 8] = [
|
||||
(119, self.aux_partial_nb_w1.raw_ptr(), aux_h * sh2),
|
||||
(120, self.aux_partial_nb_b1.raw_ptr(), aux_h),
|
||||
(121, self.aux_partial_nb_w2.raw_ptr(), aux_h),
|
||||
(122, self.aux_partial_nb_b2.raw_ptr(), 1),
|
||||
(123, self.aux_partial_rg_w1.raw_ptr(), aux_h * sh2),
|
||||
(124, self.aux_partial_rg_b1.raw_ptr(), aux_h),
|
||||
(125, self.aux_partial_rg_w2.raw_ptr(), aux_kr * aux_h),
|
||||
(126, self.aux_partial_rg_b2.raw_ptr(), aux_kr),
|
||||
];
|
||||
for &(tensor_idx, partial_ptr, p_len) in &aux_param_specs {
|
||||
// Step (a): reduce per-sample partials [B, P] → final [P].
|
||||
self.aux_heads_bwd.param_grad_reduce(
|
||||
&self.stream,
|
||||
partial_ptr,
|
||||
b,
|
||||
p_len,
|
||||
final_ptr,
|
||||
)?;
|
||||
// Step (b): SAXPY into `grad_base[tensor_idx]`, scaled by aux_weight.
|
||||
// `grad_base[119..127)` is otherwise untouched by any earlier kernel
|
||||
// in this pipeline (Commit A's design: aux params are dormant
|
||||
// without this commit's wire-up). For the main backward call site
|
||||
// `grad_base == self.ptrs.grad_buf`; the vaccine path uses a
|
||||
// scratch buffer (`g_val`) but that sub-region also stays unwritten
|
||||
// for [119..127) — both code paths exercise the same aux backward.
|
||||
let dst_ptr = grad_base + padded_byte_offset(¶m_sizes, tensor_idx);
|
||||
let n = p_len as i32;
|
||||
let blocks = ((p_len as u32 + 255) / 256).max(1);
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.saxpy_f32_kernel)
|
||||
.arg(&dst_ptr)
|
||||
.arg(&final_ptr)
|
||||
.arg(&aux_w)
|
||||
.arg(&n)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"aux_heads_backward saxpy tensor {tensor_idx}: {e}"
|
||||
)))?;
|
||||
}
|
||||
}
|
||||
|
||||
// SAXPY both per-head dh_s2 buffers into the trunk accumulator
|
||||
// `bw_d_h_s2`. Both source and destination already share the [B, SH2]
|
||||
// shape + unit stride. encoder_backward_chain consumes the augmented
|
||||
// accumulator transparently through the same pointer.
|
||||
let n_dh = (b * sh2) as i32;
|
||||
let blocks_dh = ((n_dh as u32 + 255) / 256).max(1);
|
||||
let dh_nb_ptr = self.aux_dh_s2_nb_buf.raw_ptr();
|
||||
let dh_rg_ptr = self.aux_dh_s2_rg_buf.raw_ptr();
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.saxpy_f32_kernel)
|
||||
.arg(&bw_d_h_s2_ptr)
|
||||
.arg(&dh_nb_ptr)
|
||||
.arg(&aux_w)
|
||||
.arg(&n_dh)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks_dh, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"aux_heads_backward saxpy dh_s2 next_bar: {e}"
|
||||
)))?;
|
||||
self.stream
|
||||
.launch_builder(&self.saxpy_f32_kernel)
|
||||
.arg(&bw_d_h_s2_ptr)
|
||||
.arg(&dh_rg_ptr)
|
||||
.arg(&aux_w)
|
||||
.arg(&n_dh)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks_dh, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"aux_heads_backward saxpy dh_s2 regime: {e}"
|
||||
)))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Plan 4 Task 3 (E.3): launch `iqn_quantile_ema_update` (4-block,
|
||||
/// 256-thread shmem-reduction kernel). Reads the IQN online forward's
|
||||
/// per-quantile Q surface (`save_q_online [TBA, B*Q]`, populated by
|
||||
@@ -9495,6 +9912,56 @@ impl GpuDqnTrainer {
|
||||
module.load_function("vsn_mask_ema_update")
|
||||
.map_err(|e| MLError::ModelError(format!("vsn_mask_ema_update load: {e}")))?
|
||||
};
|
||||
|
||||
// ── Plan 4 Task 6 Commit B: aux-heads orchestrator + buffers ─────────
|
||||
// Construct the forward + backward orchestrators (loads 11 kernel
|
||||
// handles total from `AUX_HEADS_CUBIN` + `AUX_HEADS_LOSS_EMA_CUBIN`).
|
||||
// Buffers are sized off the runtime config (B, SH2, AUX_HIDDEN_DIM=32,
|
||||
// AUX_REGIME_K=5, AUX_NEXT_BAR_K=1) and zero-init'd so cold-start
|
||||
// launches are well-defined.
|
||||
let aux_heads_fwd = AuxHeadsForwardOps::new(&stream)?;
|
||||
let aux_heads_bwd = AuxHeadsBackwardOps::new(&stream)?;
|
||||
let aux_b = b;
|
||||
let aux_sh2 = config.shared_h2;
|
||||
let aux_h = AUX_HIDDEN_DIM;
|
||||
let aux_kr = AUX_REGIME_K;
|
||||
let aux_knb = AUX_NEXT_BAR_K;
|
||||
let aux_nb_hidden_buf = alloc_f32(&stream, aux_b * aux_h, "aux_nb_hidden_buf")?;
|
||||
let aux_nb_pred_buf = alloc_f32(&stream, aux_b * aux_knb, "aux_nb_pred_buf")?;
|
||||
let aux_nb_label_buf = alloc_f32(&stream, aux_b, "aux_nb_label_buf")?;
|
||||
let aux_nb_loss_scalar_buf = alloc_f32(&stream, 1, "aux_nb_loss_scalar_buf")?;
|
||||
let aux_rg_hidden_buf = alloc_f32(&stream, aux_b * aux_h, "aux_rg_hidden_buf")?;
|
||||
let aux_rg_logits_buf = alloc_f32(&stream, aux_b * aux_kr, "aux_rg_logits_buf")?;
|
||||
let aux_rg_label_buf = stream.alloc_zeros::<u8>(aux_b)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc aux_rg_label_buf: {e}")))?;
|
||||
let aux_rg_loss_scalar_buf = alloc_f32(&stream, 1, "aux_rg_loss_scalar_buf")?;
|
||||
let aux_rg_correct_scalar_buf = alloc_f32(&stream, 1, "aux_rg_correct_scalar_buf")?;
|
||||
let aux_dh_s2_nb_buf = alloc_f32(&stream, aux_b * aux_sh2, "aux_dh_s2_nb_buf")?;
|
||||
let aux_dh_s2_rg_buf = alloc_f32(&stream, aux_b * aux_sh2, "aux_dh_s2_rg_buf")?;
|
||||
let aux_partial_nb_w1 = alloc_f32(&stream, aux_b * aux_h * aux_sh2, "aux_partial_nb_w1")?;
|
||||
let aux_partial_nb_b1 = alloc_f32(&stream, aux_b * aux_h, "aux_partial_nb_b1")?;
|
||||
let aux_partial_nb_w2 = alloc_f32(&stream, aux_b * aux_h, "aux_partial_nb_w2")?;
|
||||
let aux_partial_nb_b2 = alloc_f32(&stream, aux_b, "aux_partial_nb_b2")?;
|
||||
let aux_partial_rg_w1 = alloc_f32(&stream, aux_b * aux_h * aux_sh2, "aux_partial_rg_w1")?;
|
||||
let aux_partial_rg_b1 = alloc_f32(&stream, aux_b * aux_h, "aux_partial_rg_b1")?;
|
||||
let aux_partial_rg_w2 = alloc_f32(&stream, aux_b * aux_kr * aux_h, "aux_partial_rg_w2")?;
|
||||
let aux_partial_rg_b2 = alloc_f32(&stream, aux_b * aux_kr, "aux_partial_rg_b2")?;
|
||||
let max_aux_tensor_len = (aux_h * aux_sh2)
|
||||
.max(aux_kr * aux_h)
|
||||
.max(aux_h)
|
||||
.max(aux_kr)
|
||||
.max(aux_knb);
|
||||
let aux_param_grad_final_buf = alloc_f32(&stream, max_aux_tensor_len, "aux_param_grad_final_buf")?;
|
||||
// Initial aux_weight at the lower numerical-stability clamp (0.05).
|
||||
// training_loop's per-step `set_aux_weight` recomputes from ISV every
|
||||
// step. The first graph capture bakes 0.05 — an honest cold-start
|
||||
// value, not a stub.
|
||||
let aux_weight: f32 = 0.05;
|
||||
tracing::info!(
|
||||
"AuxHeadsForwardOps initialized: K_nb={} K_rg={} aux_h={} max_aux_tensor_len={} (params [119..127), ISV[113..115))",
|
||||
aux_knb, aux_kr, aux_h, max_aux_tensor_len,
|
||||
);
|
||||
|
||||
let nan_flags_buf = stream.alloc_zeros::<i32>(8)
|
||||
.map_err(|e| MLError::ModelError(format!("nan_flags alloc: {e}")))?;
|
||||
|
||||
@@ -10896,6 +11363,30 @@ impl GpuDqnTrainer {
|
||||
vsn_dw_iqn_aux_scratch,
|
||||
vsn_dw_ensemble_aux_scratch,
|
||||
vsn_mask_ema_kernel,
|
||||
// Plan 4 Task 6 Commit B — aux-heads orchestrator + buffers (24 fields)
|
||||
aux_heads_fwd,
|
||||
aux_heads_bwd,
|
||||
aux_nb_hidden_buf,
|
||||
aux_nb_pred_buf,
|
||||
aux_nb_label_buf,
|
||||
aux_nb_loss_scalar_buf,
|
||||
aux_rg_hidden_buf,
|
||||
aux_rg_logits_buf,
|
||||
aux_rg_label_buf,
|
||||
aux_rg_loss_scalar_buf,
|
||||
aux_rg_correct_scalar_buf,
|
||||
aux_dh_s2_nb_buf,
|
||||
aux_dh_s2_rg_buf,
|
||||
aux_partial_nb_w1,
|
||||
aux_partial_nb_b1,
|
||||
aux_partial_nb_w2,
|
||||
aux_partial_nb_b2,
|
||||
aux_partial_rg_w1,
|
||||
aux_partial_rg_b1,
|
||||
aux_partial_rg_w2,
|
||||
aux_partial_rg_b2,
|
||||
aux_param_grad_final_buf,
|
||||
aux_weight,
|
||||
stochastic_depth_scale_buf,
|
||||
stochastic_depth_kernel,
|
||||
stochastic_depth_rng_kernel,
|
||||
@@ -15245,6 +15736,17 @@ impl GpuDqnTrainer {
|
||||
self.ptrs.mag_concat_buf,
|
||||
)?;
|
||||
|
||||
// ── Plan 4 Task 6 Commit B: aux-heads forward (online-only) ──
|
||||
// Runs the next-bar regression + 5-class regime CE heads off the
|
||||
// freshly-saved `save_h_s2` activation. MUST run BEFORE stochastic
|
||||
// depth (which scales `save_h_s2` in-place); the aux backward
|
||||
// re-reads `save_h_s2` from the SAME buffer in
|
||||
// `launch_cublas_backward_to`, so the aux forward + backward see the
|
||||
// pre-dropout activation symmetrically. Stochastic depth still
|
||||
// applies to the trunk's downstream branch consumers; aux heads are
|
||||
// an off-branch supervised loss family per spec §4.E.6.
|
||||
self.aux_heads_forward()?;
|
||||
|
||||
// ── #21 Stochastic depth: scale hidden activations by per-layer mask ──
|
||||
// Only applied to Pass 1 (online on current states). Target and Double-DQN
|
||||
// passes use full network (no stochastic depth at inference).
|
||||
@@ -16441,6 +16943,15 @@ impl GpuDqnTrainer {
|
||||
)?;
|
||||
}
|
||||
|
||||
// ── Plan 4 Task 6 Commit B: aux-heads backward ──
|
||||
// Runs the per-head backward kernels, reduces 8 per-tensor partials
|
||||
// into final dW/dB, SAXPYs them into `grad_base[119..127)` with
|
||||
// alpha = `aux_weight`, and SAXPYs both per-head dh_s2 buffers into
|
||||
// the trunk accumulator `bw_d_h_s2` (= `d_h_s2_ptr`). MUST run BEFORE
|
||||
// `encoder_backward_chain` so the trunk dW/dB chain rule sees the
|
||||
// augmented dh_s2 (= main + aux contributions).
|
||||
self.aux_heads_backward(grad_base)?;
|
||||
|
||||
// ── GRN trunk backward (Plan 4 Task 2c.3c.4) ──
|
||||
// d_h_s2_ptr now holds the fully-accumulated gradient flowing into
|
||||
// the h_s2 GRN output. encoder_backward_chain unrolls both GRN
|
||||
|
||||
@@ -2922,6 +2922,14 @@ impl FusedTrainingCtx {
|
||||
self.trainer.set_c51_alpha(alpha);
|
||||
}
|
||||
|
||||
/// Plan 4 Task 6 Commit B: thin delegate for the aux-loss weight setter.
|
||||
/// `aux_weight` is the per-step ISV-driven scalar that scales aux head
|
||||
/// gradient SAXPYs into `grad_buf[119..127)` and `bw_d_h_s2`. See
|
||||
/// `GpuDqnTrainer::set_aux_weight` for the [0.05, 0.3] clamp rationale.
|
||||
pub(crate) fn set_aux_weight(&mut self, aux_weight: f32) {
|
||||
self.trainer.set_aux_weight(aux_weight);
|
||||
}
|
||||
|
||||
/// Return a reference to the GPU-resident loss scalar (f32) for the training guard.
|
||||
///
|
||||
/// During warmup (c51_alpha near 0), C51 loss may be NaN from random logits.
|
||||
|
||||
@@ -2845,6 +2845,44 @@ impl DQNTrainer {
|
||||
tracing::warn!("Plan 4 1B-iii vsn_mask_ema launch failed: {e}");
|
||||
}
|
||||
}
|
||||
|
||||
// Plan 4 Task 6 Commit B: per-step aux-heads loss EMA into
|
||||
// ISV[AUX_NEXT_BAR_MSE_EMA_INDEX=113] +
|
||||
// ISV[AUX_REGIME_CE_EMA_INDEX=114]. Single-thread single-block
|
||||
// ISV producer reads the just-populated `aux_nb_loss_scalar_buf`
|
||||
// + `aux_rg_loss_scalar_buf` (filled by `aux_*_loss_reduce`
|
||||
// inside the captured forward graph) and EMAs them into the
|
||||
// two ISV slots. Same launch cadence as the other producers
|
||||
// above. Honors the post-cascade-fix invariant (commit
|
||||
// a5f23b28f): ISV producer launches MUST run BEFORE the
|
||||
// HEALTH_DIAG line emit further down in this loop body.
|
||||
if let Some(ref fused) = self.fused_ctx {
|
||||
if let Err(e) = fused.trainer().launch_aux_heads_loss_ema(ema_alpha) {
|
||||
tracing::warn!("Plan 4 Task 6 Commit B aux_heads_loss_ema launch failed: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Plan 4 Task 6 Commit B: refresh the aux-loss weight before the
|
||||
// next training-step graph capture picks it up.
|
||||
//
|
||||
// aux_weight = clamp(0.1 × ISV[LEARNING_HEALTH] × (1 - tanh(SHARPE_SCALE × sharpe_ema)),
|
||||
// 0.05, 0.3)
|
||||
//
|
||||
// SHARPE_SCALE = 0.1 — matches Plan 3's controller convention for
|
||||
// sharpe-driven α adaptation. Numerical-stability bounds [0.05, 0.3]
|
||||
// per Invariant 1 carve-out. Coefficient `0.1` is the standard
|
||||
// aux-loss base weight per spec §4.E.6 (only tuned multiplicand).
|
||||
// The setter is idempotent and does NOT invalidate the captured
|
||||
// graph — the new value takes effect on the next graph re-capture
|
||||
// (fold boundary, lr change, etc.).
|
||||
if let Some(ref mut fused) = self.fused_ctx {
|
||||
use crate::cuda_pipeline::gpu_dqn_trainer::LEARNING_HEALTH_INDEX;
|
||||
let learning_health = fused.read_isv_signal_at(LEARNING_HEALTH_INDEX);
|
||||
const SHARPE_SCALE: f32 = 0.1;
|
||||
let sharpe_tanh = (SHARPE_SCALE * self.training_sharpe_ema).tanh();
|
||||
let aux_w = (0.1_f32 * learning_health * (1.0 - sharpe_tanh)).clamp(0.05, 0.3);
|
||||
fused.set_aux_weight(aux_w);
|
||||
}
|
||||
|
||||
// B.2 Plan 3 Task 3: freeze TRADE_TARGET_RATE at the configured
|
||||
@@ -3114,6 +3152,33 @@ impl DQNTrainer {
|
||||
);
|
||||
}
|
||||
|
||||
// Plan 4 Task 6 Commit B: aux-heads diagnostic line. Reads
|
||||
// ISV[AUX_NEXT_BAR_MSE_EMA_INDEX=113] +
|
||||
// ISV[AUX_REGIME_CE_EMA_INDEX=114] (populated by
|
||||
// `aux_heads_loss_ema_update` launched per-step above), and the
|
||||
// host-side `aux_weight` baked into the most-recently-captured
|
||||
// graph. Cold-start (epoch 0 of each fold): both EMAs at 0.0
|
||||
// (FoldReset), aux_weight at 0.05 (constructor cold-start).
|
||||
{
|
||||
use crate::cuda_pipeline::gpu_dqn_trainer::{
|
||||
AUX_NEXT_BAR_MSE_EMA_INDEX, AUX_REGIME_CE_EMA_INDEX,
|
||||
};
|
||||
let (aux_nb_mse, aux_rg_ce, aux_w) = if let Some(ref fused) = self.fused_ctx {
|
||||
let trainer = fused.trainer();
|
||||
(
|
||||
trainer.read_isv_signal_at(AUX_NEXT_BAR_MSE_EMA_INDEX),
|
||||
trainer.read_isv_signal_at(AUX_REGIME_CE_EMA_INDEX),
|
||||
trainer.aux_weight(),
|
||||
)
|
||||
} else {
|
||||
(0.0, 0.0, 0.0)
|
||||
};
|
||||
tracing::info!(
|
||||
"HEALTH_DIAG[{}]: aux [next_bar_mse={:.3e} regime_ce={:.3e} w={:.3}]",
|
||||
epoch, aux_nb_mse, aux_rg_ce, aux_w,
|
||||
);
|
||||
}
|
||||
|
||||
// C1/P1: propagate health to GPU replay buffer for diversity-weighted priorities.
|
||||
{
|
||||
let mut agent = self.agent.write().await;
|
||||
|
||||
File diff suppressed because one or more lines are too long
Reference in New Issue
Block a user