plan: direction-conditioned magnitude head (Layer 3) — 7 tasks
Tasks: CUDA kernel, param widening, forward conditioning, backward split, zero-init, experience collector, H100 integration test. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -0,0 +1,740 @@
|
||||
# Direction-Conditioned Magnitude Head Implementation Plan
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||
|
||||
**Goal:** Feed direction E[Q] (3 scalars) as extra input to the magnitude branch FC, so the model can distinguish Long+Full from Flat+Full.
|
||||
|
||||
**Architecture:** A small CUDA kernel computes E[Q] from direction logits and concatenates with h_s2 into a scratch buffer `[B, SH2+3]`. The magnitude branch FC GEMM uses this wider input (K=SH2+3). Backward splits dX back into d_h_s2 and d_Q_dir (discarded — detached, no backprop to direction). All other branches unchanged.
|
||||
|
||||
**Tech Stack:** CUDA 12.4, Rust (cudarc), cuBLASLt SGEMM, nvcc via build.rs
|
||||
|
||||
**Baseline:** Commits through b2eb1b774 — Layers 1+2+4 + asymmetric spread + per-branch scaling.
|
||||
|
||||
---
|
||||
|
||||
### Task 1: CUDA Kernel — mag_concat_qdir
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu`
|
||||
|
||||
- [ ] **Step 1: Add the kernel to experience_kernels.cu**
|
||||
|
||||
Append at the end of the file (before any trailing `#endif` or after the last kernel):
|
||||
|
||||
```cuda
|
||||
/* ================================================================== */
|
||||
/* Kernel: mag_concat_qdir */
|
||||
/* ================================================================== */
|
||||
|
||||
/**
|
||||
* Concatenate h_s2 with direction E[Q] for magnitude branch conditioning.
|
||||
*
|
||||
* For each sample b:
|
||||
* 1. Copy h_s2[b, 0..SH2] → concat_out[b, 0..SH2]
|
||||
* 2. Compute E[Q_dir_a] for a=0,1,2 from dir_logits via softmax × support
|
||||
* 3. Write E[Q_dir_0..2] → concat_out[b, SH2..SH2+3]
|
||||
*
|
||||
* Grid: ceil(B/256), Block: 256. One thread per sample.
|
||||
*/
|
||||
extern "C" __global__ void mag_concat_qdir(
|
||||
const float* __restrict__ h_s2, /* [B, SH2] trunk activation */
|
||||
const float* __restrict__ dir_logits_v, /* [B, NA] value logits (shared head) */
|
||||
const float* __restrict__ dir_logits_b, /* [B*b0*NA] branch 0 logits (direction) */
|
||||
float* __restrict__ concat_out, /* [B, SH2+3] output */
|
||||
int B,
|
||||
int SH2,
|
||||
int NA, /* num_atoms */
|
||||
int b0_size, /* direction branch size (3) */
|
||||
const float* __restrict__ per_sample_support /* [B, 3]: v_min, v_max, delta_z */
|
||||
) {
|
||||
int b = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (b >= B) return;
|
||||
|
||||
/* Step 1: copy h_s2 row */
|
||||
int out_stride = SH2 + 3;
|
||||
for (int i = 0; i < SH2; i++) {
|
||||
concat_out[b * out_stride + i] = h_s2[b * SH2 + i];
|
||||
}
|
||||
|
||||
/* Step 2: compute E[Q_dir] for each direction action */
|
||||
float v_min = per_sample_support[b * 3 + 0];
|
||||
float dz = per_sample_support[b * 3 + 2];
|
||||
|
||||
for (int a = 0; a < b0_size && a < 3; a++) {
|
||||
/* Combined value + advantage logits for direction action a */
|
||||
/* Branch logits layout: BRANCH-MAJOR [N × b0 × NA], so action a for sample b is:
|
||||
* dir_logits_b[b * b0_size * NA + a * NA + z] */
|
||||
float max_logit = -1e30f;
|
||||
for (int z = 0; z < NA; z++) {
|
||||
float logit = dir_logits_v[b * NA + z] + dir_logits_b[b * b0_size * NA + a * NA + z];
|
||||
if (logit > max_logit) max_logit = logit;
|
||||
}
|
||||
float sum_exp = 0.0f;
|
||||
for (int z = 0; z < NA; z++) {
|
||||
float logit = dir_logits_v[b * NA + z] + dir_logits_b[b * b0_size * NA + a * NA + z];
|
||||
sum_exp += expf(logit - max_logit);
|
||||
}
|
||||
float eq = 0.0f;
|
||||
for (int z = 0; z < NA; z++) {
|
||||
float logit = dir_logits_v[b * NA + z] + dir_logits_b[b * b0_size * NA + a * NA + z];
|
||||
float prob = expf(logit - max_logit) / sum_exp;
|
||||
float z_val = v_min + (float)z * dz;
|
||||
eq += prob * z_val;
|
||||
}
|
||||
concat_out[b * out_stride + SH2 + a] = eq;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Verify nvcc compilation**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
|
||||
```
|
||||
|
||||
Expected: `Finished` with no nvcc errors (experience_kernels.cu compiles via build.rs with common header).
|
||||
|
||||
- [ ] **Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/experience_kernels.cu
|
||||
git commit -m "feat: mag_concat_qdir kernel — direction E[Q] conditioning for magnitude branch
|
||||
|
||||
Computes softmax E[Q] for 3 direction actions from combined value+advantage
|
||||
logits, concatenates with h_s2 trunk activation to produce [B, SH2+3] input
|
||||
for the magnitude branch FC layer.
|
||||
|
||||
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: Widen w_b1fc in param_sizes and allocate concat buffers
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
|
||||
|
||||
- [ ] **Step 1: Change param_sizes for w_b1fc**
|
||||
|
||||
In `compute_param_sizes` (around line 392), change:
|
||||
|
||||
```rust
|
||||
cfg.adv_h * cfg.shared_h2, // [12] w_b1fc
|
||||
```
|
||||
|
||||
To:
|
||||
|
||||
```rust
|
||||
cfg.adv_h * (cfg.shared_h2 + 3), // [12] w_b1fc (direction-conditioned: SH2+3 input)
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Add mag_concat_buf and d_mag_concat_buf fields to GpuDqnTrainer struct**
|
||||
|
||||
Find the struct fields near `save_h_b1` (around line 587). Add after the branch hidden buffers:
|
||||
|
||||
```rust
|
||||
/// Magnitude branch concat input [B, SH2+3]: [h_s2; Q_dir] for direction conditioning.
|
||||
mag_concat_buf: CudaSlice<f32>,
|
||||
/// Backward scratch for magnitude concat dX [B, SH2+3].
|
||||
d_mag_concat_buf: CudaSlice<f32>,
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Add mag_concat pointers to CachedPtrs struct**
|
||||
|
||||
Find the `CachedPtrs` struct (search for `pub(crate) struct CachedPtrs`). Add:
|
||||
|
||||
```rust
|
||||
pub(crate) mag_concat_buf: u64,
|
||||
pub(crate) d_mag_concat_buf: u64,
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Allocate buffers in constructor**
|
||||
|
||||
Find where `save_h_b1` is allocated (around line 2370). Add after:
|
||||
|
||||
```rust
|
||||
let mag_concat_dim = config.shared_h2 + 3;
|
||||
let mag_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "mag_concat_buf")?;
|
||||
let d_mag_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "d_mag_concat_buf")?;
|
||||
```
|
||||
|
||||
(`kt` is the padding tail already used by surrounding allocations — find its value from the `save_h_b0` line.)
|
||||
|
||||
- [ ] **Step 5: Initialize in CachedPtrs**
|
||||
|
||||
In the CachedPtrs initialization block (around line 2907), add:
|
||||
|
||||
```rust
|
||||
mag_concat_buf: mag_concat_buf.raw_ptr(),
|
||||
d_mag_concat_buf: d_mag_concat_buf.raw_ptr(),
|
||||
```
|
||||
|
||||
- [ ] **Step 6: Add to struct literal**
|
||||
|
||||
In the struct literal (around line 3059), add after `save_h_b1`:
|
||||
|
||||
```rust
|
||||
mag_concat_buf,
|
||||
d_mag_concat_buf,
|
||||
```
|
||||
|
||||
- [ ] **Step 7: Load mag_concat_qdir kernel**
|
||||
|
||||
Add a field to the struct:
|
||||
|
||||
```rust
|
||||
mag_concat_kernel: CudaFunction,
|
||||
```
|
||||
|
||||
In the constructor, after loading `expected_q_kernel` (search for `compile_expected_q_kernel`), load from the experience kernels cubin:
|
||||
|
||||
```rust
|
||||
let mag_concat_kernel = exp_cubin_module.load_function("mag_concat_qdir")
|
||||
.map_err(|e| MLError::ModelError(format!("mag_concat_qdir load: {e}")))?;
|
||||
```
|
||||
|
||||
NOTE: You'll need to get a reference to the experience kernels module. The pattern varies — the experience kernels cubin is loaded in `gpu_experience_collector.rs`, not `gpu_dqn_trainer.rs`. The trainer loads kernels from `DQN_UTILITY_CUBIN` and `experience_kernels.cubin`. Search for how `expected_q_kernel` is loaded — it uses `compile_expected_q_kernel` which loads `experience_kernels.cubin`:
|
||||
|
||||
```rust
|
||||
fn compile_expected_q_kernel(stream: &Arc<CudaStream>) -> Result<CudaFunction, MLError> {
|
||||
```
|
||||
|
||||
Add a similar function or extend `compile_expected_q_kernel` to also return the mag_concat function.
|
||||
|
||||
Add `mag_concat_kernel` to the struct literal.
|
||||
|
||||
- [ ] **Step 8: Add a method to launch the concat kernel**
|
||||
|
||||
```rust
|
||||
/// Build [h_s2; Q_dir] concat buffer for magnitude branch conditioning.
|
||||
/// Must be called AFTER direction branch forward (branch 0 logits are ready).
|
||||
fn launch_mag_concat(&self, v_logits_ptr: u64, b_logits_ptr: u64) -> Result<(), MLError> {
|
||||
let b = self.config.batch_size;
|
||||
let sh2 = self.config.shared_h2 as i32;
|
||||
let na = self.config.num_atoms as i32;
|
||||
let b0 = self.config.branch_0_size as i32;
|
||||
let n = b as i32;
|
||||
let blocks = ((b as u32 + 255) / 256).max(1);
|
||||
let h_s2_ptr = self.ptrs.save_h_s2;
|
||||
let concat_ptr = self.ptrs.mag_concat_buf;
|
||||
let support_ptr = self.per_sample_support_ptr;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.mag_concat_kernel)
|
||||
.arg(&h_s2_ptr)
|
||||
.arg(&v_logits_ptr)
|
||||
.arg(&b_logits_ptr)
|
||||
.arg(&concat_ptr)
|
||||
.arg(&n)
|
||||
.arg(&sh2)
|
||||
.arg(&na)
|
||||
.arg(&b0)
|
||||
.arg(&support_ptr)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("mag_concat_qdir: {e}")))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 9: Verify build**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
|
||||
```
|
||||
|
||||
Expected: clean build.
|
||||
|
||||
- [ ] **Step 10: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
|
||||
git commit -m "feat: widen w_b1fc param_sizes + allocate mag_concat buffers
|
||||
|
||||
w_b1fc grows from [AH, SH2] to [AH, SH2+3] for direction conditioning.
|
||||
mag_concat_buf [B, SH2+3] and d_mag_concat_buf allocated for forward/backward.
|
||||
mag_concat_qdir kernel loaded and launch method added.
|
||||
|
||||
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Forward Pass — Magnitude Branch Uses Concat Input
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs`
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (callers)
|
||||
|
||||
- [ ] **Step 1: Add mag_concat_dim field to CublasForward**
|
||||
|
||||
In `crates/ml/src/cuda_pipeline/batched_forward.rs`, add a field to `CublasForward`:
|
||||
|
||||
```rust
|
||||
/// Magnitude branch FC input dimension: SH2+3 (direction conditioned).
|
||||
mag_concat_dim: usize,
|
||||
```
|
||||
|
||||
Initialize it in the constructor:
|
||||
|
||||
```rust
|
||||
mag_concat_dim: shared_h2 + 3,
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Add wider GEMM shape to cache**
|
||||
|
||||
In the constructor of `CublasForward`, after the existing `unique_shapes.push((adv_h, ...))` for branch FC (around line 263), add:
|
||||
|
||||
```rust
|
||||
// Magnitude branch FC uses wider input: K=SH2+3 (direction-conditioned)
|
||||
unique_shapes.push((adv_h, batch_size, shared_h2 + 3, shared_h2 + 3));
|
||||
```
|
||||
|
||||
Also add to `relu_bias_shapes`:
|
||||
|
||||
```rust
|
||||
(adv_h, batch_size, shared_h2 + 3, shared_h2 + 3), // h_b1 (mag, conditioned)
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Add mag_concat_ptr parameter to forward_online_raw**
|
||||
|
||||
Change the signature of `forward_online_raw` to accept the concat pointer:
|
||||
|
||||
```rust
|
||||
pub fn forward_online_raw(
|
||||
&self,
|
||||
stream: &Arc<CudaStream>,
|
||||
states_ptr: u64,
|
||||
w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
|
||||
h_s1_ptr: u64, h_s2_ptr: u64, h_v_ptr: u64,
|
||||
h_b0_ptr: u64, h_b1_ptr: u64, h_b2_ptr: u64, h_b3_ptr: u64,
|
||||
v_logits_ptr: u64, b_logits_ptr: u64,
|
||||
mag_concat_ptr: u64, // [B, SH2+3] or 0 for no conditioning
|
||||
) -> Result<(), MLError> {
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Branch 1 uses concat input in multi-stream path**
|
||||
|
||||
In `forward_online_raw`, inside the multi-stream branch loop (`for d in 0..4`), replace the magnitude branch FC GEMM. Change:
|
||||
|
||||
```rust
|
||||
if self.sgemm_f32_fused_relu_bias(bs, w_ptrs[w_fc_idx], h_s2_ptr, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1],
|
||||
self.adv_h, b, self.shared_h2, self.shared_h2, bws, bwss, "h_bd").is_err()
|
||||
{
|
||||
self.sgemm_f32_branch(bs, w_ptrs[w_fc_idx], h_s2_ptr, branch_h_ptrs[d], self.adv_h, b, self.shared_h2, d, "h_bd")?;
|
||||
self.launch_add_bias_relu_f32_raw(bs, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], self.adv_h, b)?;
|
||||
}
|
||||
```
|
||||
|
||||
To use conditional input and K dimension:
|
||||
|
||||
```rust
|
||||
let (fc_input, fc_k) = if d == 1 && mag_concat_ptr != 0 {
|
||||
(mag_concat_ptr, self.mag_concat_dim)
|
||||
} else {
|
||||
(h_s2_ptr, self.shared_h2)
|
||||
};
|
||||
if self.sgemm_f32_fused_relu_bias(bs, w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1],
|
||||
self.adv_h, b, fc_k, fc_k, bws, bwss, "h_bd").is_err()
|
||||
{
|
||||
self.sgemm_f32_branch(bs, w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], self.adv_h, b, fc_k, d, "h_bd")?;
|
||||
self.launch_add_bias_relu_f32_raw(bs, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], self.adv_h, b)?;
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Branch 1 uses concat input in sequential fallback**
|
||||
|
||||
Apply the same change to the sequential fallback path (`else` branch):
|
||||
|
||||
```rust
|
||||
let (fc_input, fc_k) = if d == 1 && mag_concat_ptr != 0 {
|
||||
(mag_concat_ptr, self.mag_concat_dim)
|
||||
} else {
|
||||
(h_s2_ptr, self.shared_h2)
|
||||
};
|
||||
self.sgemm_f32(stream, w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], self.adv_h, b, fc_k, "h_bd")?;
|
||||
```
|
||||
|
||||
- [ ] **Step 6: Apply same changes to forward_online_f32**
|
||||
|
||||
Add `mag_concat_ptr: u64` parameter. Same conditional logic for `d == 1`.
|
||||
|
||||
- [ ] **Step 7: Apply same changes to forward_target_raw**
|
||||
|
||||
Add `mag_concat_ptr: u64` parameter. Same conditional logic for `d == 1`.
|
||||
|
||||
IMPORTANT: The target network also needs conditioning. The concat must be computed from the TARGET network's direction logits. The caller (gpu_dqn_trainer) must call `launch_mag_concat` with target v_logits and b_logits for this path too. Use a separate target concat buffer or reuse the same one (target forward runs after online forward).
|
||||
|
||||
- [ ] **Step 8: Update forward_online and forward_target wrappers**
|
||||
|
||||
The CudaSlice-based wrappers (`forward_online`, `forward_target`) call the raw versions. Add the `mag_concat_ptr` parameter (pass 0 to disable, or pass the actual ptr).
|
||||
|
||||
- [ ] **Step 9: Update ALL callers in gpu_dqn_trainer.rs**
|
||||
|
||||
Search for all calls to `forward_online_raw`, `forward_online_f32`, `forward_target_raw`. Each needs the new `mag_concat_ptr` argument. Pass `self.ptrs.mag_concat_buf` for online forward, and a target concat ptr for target forward.
|
||||
|
||||
The training graph capture and the causal sensitivity forward also call these — they need the parameter too (can pass `0u64` for causal sensitivity since it's a perturbation analysis, or `self.ptrs.mag_concat_buf` for correctness).
|
||||
|
||||
- [ ] **Step 10: Add mag_concat call before branch forward**
|
||||
|
||||
In the training forward (the CUDA graph path), add `self.launch_mag_concat(v_logits_ptr, b_logits_ptr)?;` AFTER the value head forward completes and AFTER direction branch (d=0) forward completes, but BEFORE the magnitude branch (d=1) forward.
|
||||
|
||||
The challenge: with multi-stream parallel branches, direction (d=0) and magnitude (d=1) run in parallel. For conditioning, magnitude must wait for direction. The simplest fix: run direction first (d=0), sync, launch concat, then run magnitude (d=1) + order (d=2) + urgency (d=3) in parallel.
|
||||
|
||||
In the forward, change the multi-stream loop from `for d in 0..4` to:
|
||||
1. Submit direction (d=0) on branch_streams[0]
|
||||
2. Record branch_done_events[0] on branch_streams[0]
|
||||
3. Main stream waits for branch_done_events[0]
|
||||
4. Launch mag_concat on main stream
|
||||
5. Submit magnitude (d=1), order (d=2), urgency (d=3) on branch_streams 1-3
|
||||
|
||||
- [ ] **Step 11: Verify build**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
|
||||
```
|
||||
|
||||
Expected: clean build.
|
||||
|
||||
- [ ] **Step 12: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/batched_forward.rs \
|
||||
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
|
||||
git commit -m "feat: magnitude branch forward uses [h_s2; Q_dir] concat input
|
||||
|
||||
Direction branch runs first, then mag_concat_qdir produces [B, SH2+3].
|
||||
Magnitude FC GEMM uses K=SH2+3. Order and urgency branches unchanged.
|
||||
All 3 forward paths updated: online_raw, online_f32, target_raw.
|
||||
|
||||
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Backward Pass — Magnitude Branch Uses Wider dX
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/batched_backward.rs`
|
||||
|
||||
- [ ] **Step 1: Add mag_concat_dim to backward context**
|
||||
|
||||
The backward struct needs `mag_concat_dim` (or it can reference the forward's). Add the field if needed (check if backward has its own `shared_h2` field).
|
||||
|
||||
- [ ] **Step 2: Modify branch FC backward for d==1**
|
||||
|
||||
In the backward pass's branch loop (around line 809), the branch FC backward for d==1 must use the wider input. Currently:
|
||||
|
||||
```rust
|
||||
// dW for branch FC: dW[AH, SH2] += dY^T @ h_s2
|
||||
self.launch_dw_only(
|
||||
stream,
|
||||
scratch_d_h_b[d], // dY [B, AH]
|
||||
save_h_s2, // X [B, SH2]
|
||||
grad_buf_base + goff_w_bfc[d], // dW [AH, SH2]
|
||||
grad_buf_base + goff_b_bfc[d], // db [AH]
|
||||
self.adv_h, // out_dim
|
||||
self.shared_h2, // in_dim
|
||||
b,
|
||||
)?;
|
||||
```
|
||||
|
||||
For d==1, change to:
|
||||
|
||||
```rust
|
||||
let (fc_input, fc_in_dim) = if d == 1 {
|
||||
(mag_concat_ptr, self.shared_h2 + 3) // [B, SH2+3]
|
||||
} else {
|
||||
(save_h_s2, self.shared_h2)
|
||||
};
|
||||
|
||||
self.launch_dw_only(
|
||||
stream,
|
||||
scratch_d_h_b[d],
|
||||
fc_input,
|
||||
grad_buf_base + goff_w_bfc[d],
|
||||
grad_buf_base + goff_b_bfc[d],
|
||||
self.adv_h,
|
||||
fc_in_dim,
|
||||
b,
|
||||
)?;
|
||||
```
|
||||
|
||||
The `mag_concat_ptr` is the saved forward concat buffer (not the backward one).
|
||||
|
||||
- [ ] **Step 3: dX for d==1 uses wider W and outputs to d_mag_concat**
|
||||
|
||||
Currently:
|
||||
|
||||
```rust
|
||||
let beta_s2 = if d == 0 { 0.0_f32 } else { 1.0_f32 };
|
||||
self.launch_dx_only(
|
||||
stream,
|
||||
scratch_d_h_b[d], // dY [B, AH]
|
||||
w_fc, // W [AH, SH2]
|
||||
scratch_d_h_s2, // dX [B, SH2]
|
||||
self.adv_h, // out_dim
|
||||
self.shared_h2, // in_dim
|
||||
b,
|
||||
beta_s2,
|
||||
)?;
|
||||
```
|
||||
|
||||
For d==1, output to d_mag_concat_buf then extract the first SH2 columns into scratch_d_h_s2:
|
||||
|
||||
```rust
|
||||
if d == 1 {
|
||||
// Magnitude: dX = dY @ W_b1fc → [B, SH2+3]
|
||||
self.launch_dx_only(
|
||||
stream,
|
||||
scratch_d_h_b[d],
|
||||
w_fc, // W [AH, SH2+3]
|
||||
d_mag_concat_ptr, // dX [B, SH2+3]
|
||||
self.adv_h,
|
||||
self.shared_h2 + 3,
|
||||
b,
|
||||
0.0, // beta=0 (fresh write)
|
||||
)?;
|
||||
// Accumulate first SH2 columns of d_mag_concat into scratch_d_h_s2
|
||||
// Use a SAXPY or memcpy+add. Simplest: launch a small kernel or
|
||||
// use cuBLAS SAXPY. The last 3 columns (d_Q_dir) are discarded (detached).
|
||||
self.accumulate_d_h_s2_from_concat(
|
||||
stream, d_mag_concat_ptr, scratch_d_h_s2, b, beta_s2,
|
||||
)?;
|
||||
} else {
|
||||
self.launch_dx_only(
|
||||
stream,
|
||||
scratch_d_h_b[d],
|
||||
w_fc,
|
||||
scratch_d_h_s2,
|
||||
self.adv_h,
|
||||
self.shared_h2,
|
||||
b,
|
||||
beta_s2,
|
||||
)?;
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Implement accumulate_d_h_s2_from_concat**
|
||||
|
||||
Add a method that copies `d_mag_concat[:, 0:SH2]` into `scratch_d_h_s2` with beta accumulation. This can use `cuMemcpy2DAsync` (strided copy) or a simple CUDA kernel. The simplest approach: a SAXPY-like kernel that reads stride `SH2+3` and writes stride `SH2`:
|
||||
|
||||
```rust
|
||||
/// Accumulate the first SH2 columns of d_mag_concat [B, SH2+3] into d_h_s2 [B, SH2].
|
||||
/// beta=0: overwrite. beta=1: add.
|
||||
fn accumulate_d_h_s2_from_concat(
|
||||
&self,
|
||||
stream: &Arc<CudaStream>,
|
||||
d_concat_ptr: u64, // [B, SH2+3]
|
||||
d_h_s2_ptr: u64, // [B, SH2]
|
||||
batch: usize,
|
||||
beta: f32,
|
||||
) -> Result<(), MLError> {
|
||||
// Use a simple kernel: for each (b, i) where i < SH2:
|
||||
// d_h_s2[b*SH2+i] = beta * d_h_s2[b*SH2+i] + d_concat[b*(SH2+3)+i]
|
||||
// Total elements = B * SH2, launched as 1D grid.
|
||||
let total = batch * self.shared_h2;
|
||||
let blocks = ((total + 255) / 256) as u32;
|
||||
let sh2 = self.shared_h2 as i32;
|
||||
let sh2_plus_3 = (self.shared_h2 + 3) as i32;
|
||||
let total_i32 = total as i32;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.strided_accumulate_kernel)
|
||||
.arg(&d_concat_ptr)
|
||||
.arg(&d_h_s2_ptr)
|
||||
.arg(&sh2)
|
||||
.arg(&sh2_plus_3)
|
||||
.arg(&total_i32)
|
||||
.arg(&beta)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("strided_accumulate: {e}")))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
The `strided_accumulate_kernel` needs to be added to experience_kernels.cu:
|
||||
|
||||
```cuda
|
||||
extern "C" __global__ void strided_accumulate(
|
||||
const float* __restrict__ src, /* [B, src_stride] */
|
||||
float* __restrict__ dst, /* [B, dst_stride] */
|
||||
int dst_stride, /* SH2 */
|
||||
int src_stride, /* SH2+3 */
|
||||
int total, /* B * dst_stride */
|
||||
float beta /* 0=overwrite, 1=add */
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= total) return;
|
||||
int b = idx / dst_stride;
|
||||
int i = idx % dst_stride;
|
||||
float val = src[b * src_stride + i];
|
||||
if (beta > 0.0f) {
|
||||
dst[idx] += val;
|
||||
} else {
|
||||
dst[idx] = val;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Add mag_concat_ptr and d_mag_concat_ptr parameters to backward**
|
||||
|
||||
The backward function signature needs these two extra pointers.
|
||||
|
||||
- [ ] **Step 6: Update all backward callers**
|
||||
|
||||
Search for all calls to the backward function and add the new parameters.
|
||||
|
||||
- [ ] **Step 7: Verify build**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
|
||||
```
|
||||
|
||||
Expected: clean build.
|
||||
|
||||
- [ ] **Step 8: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/batched_backward.rs \
|
||||
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs \
|
||||
crates/ml/src/cuda_pipeline/experience_kernels.cu
|
||||
git commit -m "feat: magnitude backward uses wider [SH2+3] dX with strided accumulate
|
||||
|
||||
Branch 1 backward: dW uses [h_s2; Q_dir] as X, dX writes to d_mag_concat
|
||||
[B, SH2+3], first SH2 columns accumulated into d_h_s2. d_Q_dir (last 3
|
||||
columns) discarded — detached, no backprop to direction head.
|
||||
|
||||
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: Zero-Init New Weight Columns and Build Verification
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
|
||||
|
||||
- [ ] **Step 1: Zero the extra 3 columns of w_b1fc after xavier init**
|
||||
|
||||
The xavier init fills the entire `w_b1fc[AH, SH2+3]` buffer with random values. The extra 3 columns (Q_dir conditioning) should start at zero so the model doesn't rely on random direction signals at epoch 0.
|
||||
|
||||
Find where weights are initialized (search for `xavier_init` or `param_init`). After the full buffer is initialized, zero the last 3 columns of w_b1fc:
|
||||
|
||||
```rust
|
||||
// Zero the Q_dir conditioning columns of w_b1fc so training starts
|
||||
// as if magnitude is unconditioned. The model learns to use Q_dir gradually.
|
||||
let w_b1fc_offset = padded_byte_offset(¶m_sizes, 12);
|
||||
let sh2_plus_3 = config.shared_h2 + 3;
|
||||
for row in 0..config.adv_h {
|
||||
let col_offset = w_b1fc_offset + (row * sh2_plus_3 + config.shared_h2) as u64 * 4;
|
||||
stream.memset_zeros_at(params_ptr + col_offset, 3 * std::mem::size_of::<f32>())
|
||||
.map_err(|e| MLError::ModelError(format!("zero w_b1fc qdir cols row {row}: {e}")))?;
|
||||
}
|
||||
```
|
||||
|
||||
NOTE: The exact API for partial memset may differ. Alternative: use `cuMemsetD32Async` for 3 floats at a time, or write a small kernel.
|
||||
|
||||
- [ ] **Step 2: Run smoke tests**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep "avg_grad\|test result" | tail -5
|
||||
```
|
||||
|
||||
Expected: 19 tests pass, avg_grad non-zero.
|
||||
|
||||
- [ ] **Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
|
||||
git commit -m "feat: zero-init Q_dir conditioning columns of w_b1fc
|
||||
|
||||
New 3 columns of magnitude FC weight start at zero so training begins
|
||||
as if unconditioned. Model learns to use direction E[Q] gradually.
|
||||
|
||||
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 6: Experience Collector Forward — Add Conditioning
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`
|
||||
|
||||
The experience collector has its own cuBLAS forward pass (`forward_online_f32`). It also needs the mag_concat conditioning.
|
||||
|
||||
- [ ] **Step 1: Allocate mag_concat buffer in collector**
|
||||
|
||||
Add `mag_concat_exp_buf: CudaSlice<f32>` field, allocate `[N, SH2+3]`, and load the `mag_concat_qdir` kernel.
|
||||
|
||||
- [ ] **Step 2: Call mag_concat before action selection**
|
||||
|
||||
In the experience collection timestep loop, after the forward pass produces v_logits and b_logits, call `mag_concat_qdir` to build the concat buffer, then the forward pass uses it for branch 1.
|
||||
|
||||
Since the collector calls `forward_online_f32` which now takes `mag_concat_ptr`, pass the buffer pointer.
|
||||
|
||||
- [ ] **Step 3: Verify build and smoke tests**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep "avg_grad\|test result" | tail -5
|
||||
```
|
||||
|
||||
Expected: clean build, 19 tests pass.
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
|
||||
git commit -m "feat: experience collector forward uses direction-conditioned magnitude
|
||||
|
||||
mag_concat_qdir runs after forward pass produces direction logits.
|
||||
Magnitude branch sees [h_s2; Q_dir] during experience collection.
|
||||
|
||||
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 7: Integration Test on H100
|
||||
|
||||
**Files:**
|
||||
- No code changes — deployment test
|
||||
|
||||
- [ ] **Step 1: Push and launch**
|
||||
|
||||
```bash
|
||||
git push origin main
|
||||
./scripts/argo-train.sh dqn --epochs 50 --baseline --gpu-pool ci-training-h100
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Monitor first 10 epochs**
|
||||
|
||||
Check:
|
||||
- `avg_grad` non-zero (~1.0-2.5)
|
||||
- `val_Sharpe` positive (target: >15)
|
||||
- Q-gap growing (target: >0.05 by epoch 5)
|
||||
- No CUDA errors, no graph capture failures
|
||||
- Magnitude Q-values should diverge for different direction actions
|
||||
|
||||
```bash
|
||||
argo logs <workflow-name> -n foxhunt -c main | grep -E "val_Sharpe|Q-gap|avg_grad|EPOCH_DIAG"
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Monitor epoch 20-30**
|
||||
|
||||
The key test: does the Q-gap freeze break? With reward rank normalization + Q-gap momentum + direction conditioning, the plateau at epoch 22 should not recur.
|
||||
|
||||
- [ ] **Step 4: Document results**
|
||||
|
||||
Record: peak val_Sharpe, sustained val_Sharpe range, Q-gap trajectory, any anomalies. Compare to train-w6qfd baseline (frozen at val_Sharpe=22.18, Q-gap=0.133 from epoch 22).
|
||||
Reference in New Issue
Block a user