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

fix(dqn-v2): bottleneck Linear runtime indices + missing target/DDQN bn paths (2c.3a follow-up)

Browse Source 
Three related bugs from the 2c.3a GRN trunk reshuffle that the
bottleneck-Linear runtime sites silently inherited:

1. **on_w_ptrs[24/25] -> [33/34]** (4 sites in gpu_dqn_trainer.rs).
   2c.3a's +9 shift migrated branch/value/VSN/GLU/KAN consumers but
   missed the bottleneck Linear's forward GEMM (line ~14098), forward
   bias-add (~14107), and backward dW/db offsets (~15109-15110).
   Post-2c.3a, indices 24/25 point at b_b1out (153 floats) and start
   of w_b2fc (4288 floats) — not w_bn (672 floats) / b_bn (16 floats).
   Forward + backward were self-consistent on wrong tensors; the actual
   w_bn/b_bn at documented indices 33/34 sat untouched as zeros for
   ~10 commits. Smoke kept passing because GRN's Linear_residual
   projection ran in parallel and provided a clean residual-only
   pathway around the corrupted bottleneck slot.

2. **Target net missing bottleneck path.** forward_target_raw passed
   raw next_states_buf (128-padded) to an encoder expecting
   [B, s1_input_dim=102]; it was reading
   [market|ofi|tlob|mtf] of next_states instead of
   [bn_market_proj|portfolio]. Fix: build tg_bn_concat_buf inline
   before forward_target_raw using tg_w_ptrs[33]/[34] (target's w_bn)
   on next_states_buf[market]; new buffer pair tg_bn_hidden_buf +
   tg_bn_concat_buf.

3. **DDQN argmax missing bottleneck path.** submit_forward_ops_ddqn
   had the same shape mismatch on the online-on-next_states pass.
   Fix: build on_next_bn_concat_buf using on_w_ptrs[33]/[34] (online
   weights, since DDQN argmax uses online net); new buffer pair
   on_next_bn_hidden_buf + on_next_bn_concat_buf.

Three call sites of the existing bn_tanh_concat_kernel now: online-on-
states (states + on_w_bn), target-on-next_states (next + tg_w_bn), and
ddqn-online-on-next_states (next + on_w_bn). Each combination of
weights × input states produces distinct features; sharing the kernel
across distinct buffer pairs preserves the GPU-only cold-path contract.

No fingerprint change (no ISV slot or param tensor added).

Smoke validation (multi_fold_convergence, 700.43s, 3 folds x 5 epochs):
  fold 0 best Sharpe  6.53 at epoch 2
  fold 1 best Sharpe 80.11 at epoch 4
  fold 2 best Sharpe 66.72 at epoch 4
  geom-mean: 39.27 (vs Task 3 20.03, 2c.3c.6 18.80)

The Sharpe lift is consistent with the bottleneck Linear now seeing
its actual weights and the target/DDQN nets seeing input-shape-
consistent features for TD-target / argmax-action computations.

+166 / -9 LOC all in gpu_dqn_trainer.rs.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-25 18:24:06 +02:00
parent 475101054d
commit f3e3ac3477
2 changed files with 168 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -2053,6 +2053,22 @@ pub struct GpuDqnTrainer {
/// #31 Bottleneck + portfolio concatenated [B, bottleneck_dim + portfolio_dim] f32.
/// This replaces states_buf as input to h_s1 when bottleneck is active.
bn_concat_buf: CudaSlice<f32>,
/// Target-net mirror of `bn_hidden_buf` — fed by `tg_w_ptrs[33]` (target's
/// w_bn) on `next_states_buf[market]`. Required so the target trunk sees
/// the same compressed-features representation the online trunk does;
/// without this the target encoder reads raw `next_states[0..s1_input_dim)`
/// = `[market | ofi | tlob | mtf]` instead of `[bn_market_proj | portfolio]`,
/// which silently mismatches the online encoder's input semantics.
tg_bn_hidden_buf: CudaSlice<f32>,
/// Target-net mirror of `bn_concat_buf`.
tg_bn_concat_buf: CudaSlice<f32>,
/// DDQN argmax-pass mirror — fed by **online** weights `on_w_ptrs[33]`
/// on `next_states_buf[market]`. The DDQN argmax forward runs the online
/// network on next_states, so its bottleneck input is online-w_bn @
/// next_states (distinct from `tg_bn_hidden_buf` which uses target-w_bn).
on_next_bn_hidden_buf: CudaSlice<f32>,
/// DDQN argmax-pass mirror of `bn_concat_buf` (online-on-next_states).
on_next_bn_concat_buf: CudaSlice<f32>,
/// #31 Bottleneck tanh + concat kernel.
bn_tanh_concat_kernel: CudaFunction,
/// #31 Backward scratch: d_loss/d_bn_concat [B, concat_dim] f32.
@@ -8594,6 +8610,20 @@ impl GpuDqnTrainer {
let concat_dim = bn_alloc_dim + portfolio_dim;
let bn_concat_buf = stream.alloc_zeros::<f32>(b * concat_dim + kt)
.map_err(|e| MLError::ModelError(format!("alloc bn_concat: {e}")))?;
// Target-net mirrors of bn_hidden / bn_concat — same shapes; populated
// each forward by `submit_forward_ops_target` from `next_states_buf`
// and the target weights `tg_w_ptrs[33]` / `tg_w_ptrs[34]`. Same
// K-tile padding rationale as the online buffers.
let tg_bn_hidden_buf = stream.alloc_zeros::<f32>(b * bn_alloc_dim)
.map_err(|e| MLError::ModelError(format!("alloc tg_bn_hidden: {e}")))?;
let tg_bn_concat_buf = stream.alloc_zeros::<f32>(b * concat_dim + kt)
.map_err(|e| MLError::ModelError(format!("alloc tg_bn_concat: {e}")))?;
// DDQN argmax-pass mirrors — populated by `submit_forward_ops_ddqn`
// using online weights on `next_states_buf`.
let on_next_bn_hidden_buf = stream.alloc_zeros::<f32>(b * bn_alloc_dim)
.map_err(|e| MLError::ModelError(format!("alloc on_next_bn_hidden: {e}")))?;
let on_next_bn_concat_buf = stream.alloc_zeros::<f32>(b * concat_dim + kt)
.map_err(|e| MLError::ModelError(format!("alloc on_next_bn_concat: {e}")))?;
let bn_d_concat_buf = stream.alloc_zeros::<f32>(b * concat_dim + kt)
.map_err(|e| MLError::ModelError(format!("alloc bn_d_concat: {e}")))?;
let bn_d_hidden_buf = stream.alloc_zeros::<f32>(b * bn_alloc_dim)
@@ -9942,6 +9972,10 @@ impl GpuDqnTrainer {
vaccine_project_kernel,
bn_hidden_buf,
bn_concat_buf,
tg_bn_hidden_buf,
tg_bn_concat_buf,
on_next_bn_hidden_buf,
on_next_bn_concat_buf,
bn_d_concat_buf,
bn_d_hidden_buf,
bn_tanh_concat_kernel: bn_tanh_concat_kernel_fn,
@@ -13860,8 +13894,66 @@ impl GpuDqnTrainer {
pub(crate) fn submit_forward_ops_ddqn(&mut self) -> Result<(), MLError> {
let param_sizes = compute_param_sizes(&self.config);
let on_w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, &param_sizes);
// Bottleneck for DDQN argmax pass: online weights on next_states.
// Mirrors the online-on-states block in `launch_cublas_forward` and
// the target-on-next_states block; without this the DDQN argmax
// encoder reads raw `next_states[0..s1_input_dim)` and produces
// argmax actions inconsistent with the online net's training-time
// input semantics.
let on_next_s1_input_ptr = if self.config.bottleneck_dim > 0 {
let bn_dim = self.config.bottleneck_dim;
let b = self.config.batch_size;
let market_dim = self.config.market_dim;
let portfolio_dim = ml_core::state_layout::STATE_DIM.saturating_sub(market_dim);
let concat_dim = bn_dim + portfolio_dim;
let bn_hidden_ptr = self.on_next_bn_hidden_buf.raw_ptr();
let state_dim_padded_usize = self.cublas_forward.state_dim_padded;
self.cublas_forward.sgemm_f32_ldb(
&self.stream,
on_w_ptrs[33], // online w_bn
self.ptrs.next_states_buf,
bn_hidden_ptr,
bn_dim, b, market_dim,
state_dim_padded_usize,
"ddqn_h_bn",
)?;
self.cublas_forward.launch_add_bias_f32_raw(
&self.stream, bn_hidden_ptr, on_w_ptrs[34], bn_dim, b,
)?;
let concat_ptr = self.on_next_bn_concat_buf.raw_ptr();
let state_dim_padded = state_dim_padded_usize as i32;
let total_elems = (b * concat_dim) as i32;
let blocks = ((total_elems as u32 + 255) / 256) as u32;
let bn_dim_i32 = bn_dim as i32;
let market_dim_i32 = market_dim as i32;
let concat_dim_i32 = concat_dim as i32;
let b_i32 = b as i32;
unsafe {
self.stream
.launch_builder(&self.bn_tanh_concat_kernel)
.arg(&bn_hidden_ptr)
.arg(&self.ptrs.next_states_buf)
.arg(&concat_ptr)
.arg(&b_i32)
.arg(&bn_dim_i32)
.arg(&market_dim_i32)
.arg(&state_dim_padded)
.arg(&concat_dim_i32)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("ddqn_bn_tanh_concat: {e}")))?;
}
concat_ptr
} else {
self.ptrs.next_states_buf
};
self.cublas_forward_ddqn.forward_online_raw(
&self.stream, self.ptrs.next_states_buf, &on_w_ptrs,
&self.stream, on_next_s1_input_ptr, &on_w_ptrs,
self.ptrs.on_next_h_s1_scratch, self.ptrs.on_next_h_s2_scratch,
self.ptrs.on_next_h_v_scratch,
self.ptrs.on_next_h_b_scratch, self.ptrs.on_next_h_b_scratch,
@@ -14090,21 +14182,25 @@ impl GpuDqnTrainer {
let concat_dim = bn_dim + portfolio_dim;
// Bottleneck GEMM: states[B, market_dim] @ w_bn^T → h_bn[B, bn_dim]
// w_bn is at on_w_ptrs[24], b_bn is at on_w_ptrs[25]
// w_bn is at on_w_ptrs[33], b_bn is at on_w_ptrs[34] (post Plan 4
// Task 2c.3a, the GRN trunk reshuffle moved every legacy
// tensor-index ≥ 4 by +9 — these two runtime sites were missed in
// that migration; 24/25 in the post-2c.3a layout point at b_b1out
// / w_b2fc which would silently corrupt the bottleneck weights).
let bn_hidden_ptr = self.bn_hidden_buf.raw_ptr();
let state_dim_padded_usize = self.cublas_forward.state_dim_padded;
self.cublas_forward.sgemm_f32_ldb(
&self.stream,
on_w_ptrs[24], // w_bn [bn_dim, market_dim]
on_w_ptrs[33], // w_bn [bn_dim, market_dim]
self.ptrs.states_buf, // states [B, state_dim_padded]
bn_hidden_ptr, // h_bn [B, bn_dim]
bn_dim, b, market_dim,
state_dim_padded_usize, // ldb = padded state_dim
"h_bn",
)?;
// Add bias (on_w_ptrs[25]) — no ReLU, tanh applied by concat kernel
// Add bias (on_w_ptrs[34]) — no ReLU, tanh applied by concat kernel
self.cublas_forward.launch_add_bias_f32_raw(
&self.stream, bn_hidden_ptr, on_w_ptrs[25], bn_dim, b,
&self.stream, bn_hidden_ptr, on_w_ptrs[34], bn_dim, b,
)?;
// Tanh + concat: [tanh(h_bn), portfolio_from_states] → bn_concat
@@ -14210,8 +14306,67 @@ impl GpuDqnTrainer {
// with next_states features (target forward sees next_states).
self.launch_concat_ofi(2, self.ptrs.next_states_buf)?;
self.launch_concat_ofi(3, self.ptrs.next_states_buf)?;
// Target-net bottleneck: mirror of the online block above. Without this,
// the target encoder reads raw `next_states_buf[0..s1_input_dim)` =
// `[market | ofi | tlob | mtf]` while the online encoder reads
// `[bn_market_proj | portfolio]` from `bn_concat`. The two paths must
// share encoder-input semantics for the TD target Q-values to be
// comparable to the online Q-values. Uses target weights
// `tg_w_ptrs[33]` / `[34]` (Polyak EMA of online's `w_bn` / `b_bn`).
let tg_s1_input_ptr = if self.config.bottleneck_dim > 0 {
let bn_dim = self.config.bottleneck_dim;
let b = self.config.batch_size;
let market_dim = self.config.market_dim;
let portfolio_dim = ml_core::state_layout::STATE_DIM.saturating_sub(market_dim);
let concat_dim = bn_dim + portfolio_dim;
let tg_bn_hidden_ptr = self.tg_bn_hidden_buf.raw_ptr();
let state_dim_padded_usize = self.cublas_forward.state_dim_padded;
self.cublas_forward.sgemm_f32_ldb(
&self.stream,
tg_w_ptrs[33], // tg_w_bn [bn_dim, market_dim]
self.ptrs.next_states_buf, // next_states [B, state_dim_padded]
tg_bn_hidden_ptr, // tg_h_bn [B, bn_dim]
bn_dim, b, market_dim,
state_dim_padded_usize,
"tg_h_bn",
)?;
self.cublas_forward.launch_add_bias_f32_raw(
&self.stream, tg_bn_hidden_ptr, tg_w_ptrs[34], bn_dim, b,
)?;
let tg_concat_ptr = self.tg_bn_concat_buf.raw_ptr();
let state_dim_padded = state_dim_padded_usize as i32;
let total_elems = (b * concat_dim) as i32;
let blocks = ((total_elems as u32 + 255) / 256) as u32;
let bn_dim_i32 = bn_dim as i32;
let market_dim_i32 = market_dim as i32;
let concat_dim_i32 = concat_dim as i32;
let b_i32 = b as i32;
unsafe {
self.stream
.launch_builder(&self.bn_tanh_concat_kernel)
.arg(&tg_bn_hidden_ptr)
.arg(&self.ptrs.next_states_buf)
.arg(&tg_concat_ptr)
.arg(&b_i32)
.arg(&bn_dim_i32)
.arg(&market_dim_i32)
.arg(&state_dim_padded)
.arg(&concat_dim_i32)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("tg_bn_tanh_concat: {e}")))?;
}
tg_concat_ptr
} else {
self.ptrs.next_states_buf
};
self.cublas_forward.forward_target_raw(
&self.stream, self.ptrs.next_states_buf, &tg_w_ptrs,
&self.stream, tg_s1_input_ptr, &tg_w_ptrs,
self.ptrs.tg_h_s1_scratch, self.ptrs.tg_h_s2_buf,
self.ptrs.tg_h_v_scratch,
self.ptrs.tg_h_b0_scratch, self.ptrs.tg_h_b1_scratch,
@@ -15105,9 +15260,11 @@ impl GpuDqnTrainer {
}
// Step 2: dW_bn[bn_dim, market_dim] += d_bn^T @ states[:, :market_dim]
// Grad offset for w_bn (tensor 24)
let goff_w_bn = padded_byte_offset(&param_sizes, 24);
let goff_b_bn = padded_byte_offset(&param_sizes, 25);
// Grad offsets for w_bn (tensor 33) / b_bn (tensor 34) — see the
// matching forward-site comment for the post-2c.3a +9 shift that
// these two sites were missing.
let goff_w_bn = padded_byte_offset(&param_sizes, 33);
let goff_b_bn = padded_byte_offset(&param_sizes, 34);
// dW_bn = d_bn^T[bn_dim, B] @ states[B, market_dim_padded]
// Use launch_dw_only (d_bn=dY f32, states=X, grad=dW f32)

2
docs/dqn-wire-up-audit.md
View File

File diff suppressed because one or more lines are too long