plan: unified state layout implementation — 8 tasks
Task 1: Rust constants (ml-core/state_layout.rs) Task 2: CUDA header (state_layout.cuh + assemble_state()) Task 3: Refactor experience_state_gather to use shared assembly Task 4: Delete backtest_gather_kernel, replace with shared function Task 5: Remove configurable state_dim everywhere (131 call sites) Task 6: Update OFI_DIAG indices + checkpoint validation Task 7: Layout match integration test Task 8: Deploy and verify on H100 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
595
docs/superpowers/plans/2026-04-20-unified-state-layout.md
Normal file
595
docs/superpowers/plans/2026-04-20-unified-state-layout.md
Normal file
@@ -0,0 +1,595 @@
|
||||
# Unified State Layout 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:** Eliminate train/validation state mismatch by enforcing a single state layout, single assembly kernel, and compile-time STATE_DIM constant.
|
||||
|
||||
**Architecture:** A CUDA header (`state_layout.cuh`) defines all layout constants and an inline `__device__` `assemble_state()` function. Both the experience collector and backtest evaluator call it. A Rust constant (`STATE_DIM = 96`) replaces all configurable `state_dim` fields across 5 systems. The backtest gather kernel is deleted.
|
||||
|
||||
**Tech Stack:** Rust 1.85, CUDA 12.4, cudarc, cuBLAS, nvcc (cubin compilation)
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Rust State Layout Constants
|
||||
|
||||
**Files:**
|
||||
- Create: `crates/ml-core/src/state_layout.rs`
|
||||
- Modify: `crates/ml-core/src/lib.rs`
|
||||
|
||||
- [ ] **Step 1: Create the constants module**
|
||||
|
||||
```rust
|
||||
// crates/ml-core/src/state_layout.rs
|
||||
|
||||
/// Canonical state layout — single source of truth.
|
||||
///
|
||||
/// Training and validation MUST produce identical state vectors.
|
||||
/// All systems use these constants instead of configurable state_dim fields.
|
||||
///
|
||||
/// Layout (grouped by update frequency, fast-changing first):
|
||||
/// [0..42) Market features (OHLCV + technicals)
|
||||
/// [42..62) OFI (20-dim order flow from MBP-10)
|
||||
/// [62..78) MTF (4 lookbacks x 4 features)
|
||||
/// [78..86) Portfolio base (8 features)
|
||||
/// [86..92) Plan/ISV (6 features)
|
||||
/// [92..96) Zero padding (tensor core alignment)
|
||||
|
||||
pub const STATE_DIM: usize = 96;
|
||||
pub const STATE_DIM_PADDED: usize = 128; // pad128(STATE_DIM) for cuBLAS K-tile alignment
|
||||
|
||||
pub const MARKET_DIM: usize = 42;
|
||||
pub const OFI_DIM: usize = 20;
|
||||
pub const MTF_DIM: usize = 16;
|
||||
pub const PORTFOLIO_BASE_DIM: usize = 8;
|
||||
pub const PORTFOLIO_PLAN_DIM: usize = 6;
|
||||
pub const PADDING_DIM: usize = 4;
|
||||
|
||||
pub const MARKET_START: usize = 0;
|
||||
pub const OFI_START: usize = MARKET_START + MARKET_DIM; // 42
|
||||
pub const MTF_START: usize = OFI_START + OFI_DIM; // 62
|
||||
pub const PORTFOLIO_START: usize = MTF_START + MTF_DIM; // 78
|
||||
pub const PLAN_ISV_START: usize = PORTFOLIO_START + PORTFOLIO_BASE_DIM; // 86
|
||||
pub const PADDING_START: usize = PLAN_ISV_START + PORTFOLIO_PLAN_DIM; // 92
|
||||
|
||||
const _: () = assert!(PADDING_START + PADDING_DIM == STATE_DIM, "layout must sum to STATE_DIM");
|
||||
const _: () = assert!(STATE_DIM % 8 == 0, "STATE_DIM must be 8-aligned for tensor cores");
|
||||
const _: () = assert!(STATE_DIM_PADDED % 128 == 0, "STATE_DIM_PADDED must be 128-aligned for cuBLAS");
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Export from ml-core lib.rs**
|
||||
|
||||
Add to `crates/ml-core/src/lib.rs`:
|
||||
|
||||
```rust
|
||||
pub mod state_layout;
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Verify compilation**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml-core`
|
||||
Expected: PASS with no errors
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml-core/src/state_layout.rs crates/ml-core/src/lib.rs
|
||||
git commit -m "feat: add state_layout constants module — single source of truth for STATE_DIM"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: CUDA State Layout Header
|
||||
|
||||
**Files:**
|
||||
- Create: `crates/ml/src/cuda_pipeline/state_layout.cuh`
|
||||
- Modify: `crates/ml/build.rs`
|
||||
- Modify: `crates/ml-dqn/build.rs`
|
||||
|
||||
- [ ] **Step 1: Create the CUDA header**
|
||||
|
||||
```c
|
||||
// crates/ml/src/cuda_pipeline/state_layout.cuh
|
||||
// Single source of truth for GPU state vector layout.
|
||||
// Included by experience_kernels.cu and any kernel that assembles states.
|
||||
// Must match crates/ml-core/src/state_layout.rs exactly.
|
||||
|
||||
#pragma once
|
||||
|
||||
// ── Dimensions ──
|
||||
#define SL_STATE_DIM 96
|
||||
#define SL_MARKET_DIM 42
|
||||
#define SL_OFI_DIM 20
|
||||
#define SL_MTF_DIM 16
|
||||
#define SL_PORTFOLIO_BASE_DIM 8
|
||||
#define SL_PORTFOLIO_PLAN_DIM 6
|
||||
#define SL_PADDING_DIM 4
|
||||
|
||||
// ── Offsets ──
|
||||
#define SL_MARKET_START 0
|
||||
#define SL_OFI_START (SL_MARKET_START + SL_MARKET_DIM)
|
||||
#define SL_MTF_START (SL_OFI_START + SL_OFI_DIM)
|
||||
#define SL_PORTFOLIO_START (SL_MTF_START + SL_MTF_DIM)
|
||||
#define SL_PLAN_ISV_START (SL_PORTFOLIO_START + SL_PORTFOLIO_BASE_DIM)
|
||||
#define SL_PADDING_START (SL_PLAN_ISV_START + SL_PORTFOLIO_PLAN_DIM)
|
||||
|
||||
// ── cuBLAS alignment ──
|
||||
#define SL_STATE_DIM_PADDED 128
|
||||
|
||||
// ── Compile-time checks ──
|
||||
_Static_assert(SL_PADDING_START + SL_PADDING_DIM == SL_STATE_DIM,
|
||||
"State layout dimensions must sum to SL_STATE_DIM");
|
||||
_Static_assert(SL_STATE_DIM % 8 == 0,
|
||||
"SL_STATE_DIM must be 8-aligned for tensor core cuBLAS");
|
||||
|
||||
// ── Shared state assembly function ─<><E29480><EFBFBD>
|
||||
// Called by both experience_state_gather (training) and backtest_state_gather (validation).
|
||||
// Each caller prepares its own data and calls this to produce identical layouts.
|
||||
//
|
||||
// All pointers are per-thread (one state row). Caller is responsible for
|
||||
// computing the correct source pointers for this thread's episode/window.
|
||||
|
||||
__device__ __forceinline__ void assemble_state(
|
||||
float* __restrict__ out,
|
||||
const float* __restrict__ market, // [SL_MARKET_DIM]
|
||||
const float* __restrict__ ofi, // [SL_OFI_DIM]
|
||||
const float mtf[SL_MTF_DIM], // caller-computed MTF features
|
||||
const float portfolio[SL_PORTFOLIO_BASE_DIM],
|
||||
const float plan_isv[SL_PORTFOLIO_PLAN_DIM]
|
||||
) {
|
||||
// Market features [0..42)
|
||||
for (int k = 0; k < SL_MARKET_DIM; k++)
|
||||
out[SL_MARKET_START + k] = market[k];
|
||||
|
||||
// OFI [42..62)
|
||||
for (int k = 0; k < SL_OFI_DIM; k++)
|
||||
out[SL_OFI_START + k] = ofi[k];
|
||||
|
||||
// MTF [62..78)
|
||||
for (int k = 0; k < SL_MTF_DIM; k++)
|
||||
out[SL_MTF_START + k] = mtf[k];
|
||||
|
||||
// Portfolio base [78..86)
|
||||
for (int k = 0; k < SL_PORTFOLIO_BASE_DIM; k++)
|
||||
out[SL_PORTFOLIO_START + k] = portfolio[k];
|
||||
|
||||
// Plan/ISV [86..92)
|
||||
for (int k = 0; k < SL_PORTFOLIO_PLAN_DIM; k++)
|
||||
out[SL_PLAN_ISV_START + k] = plan_isv[k];
|
||||
|
||||
// Padding [92..96) — zero
|
||||
for (int k = 0; k < SL_PADDING_DIM; k++)
|
||||
out[SL_PADDING_START + k] = 0.0f;
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Update ml build.rs include path**
|
||||
|
||||
In `crates/ml/build.rs`, find the nvcc compilation args (around line 135-146). The kernel dir is already used as an include path (`-I{kernel_dir}`). Since `state_layout.cuh` lives in the same `cuda_pipeline/` directory as the `.cu` files, it will be found by `#include "state_layout.cuh"` without additional flags. Verify the `-I` flag covers it:
|
||||
|
||||
```rust
|
||||
// Existing line in build.rs — confirms kernel_dir is in the include path:
|
||||
format!("-I{}", kernel_dir.display()),
|
||||
```
|
||||
|
||||
No change needed in `ml/build.rs` — `state_layout.cuh` is in `kernel_dir`.
|
||||
|
||||
- [ ] **Step 3: Update ml-dqn build.rs include path**
|
||||
|
||||
In `crates/ml-dqn/build.rs`, the kernels in `ml-dqn/src/` may need access to the header in `ml/src/cuda_pipeline/`. Add an include path:
|
||||
|
||||
Find the nvcc args section (around line 112-125) and add:
|
||||
|
||||
```rust
|
||||
// After existing include args, add ml crate's kernel dir for state_layout.cuh
|
||||
let ml_kernel_dir = manifest_dir.join("../ml/src/cuda_pipeline");
|
||||
// Add to nvcc args:
|
||||
format!("-I{}", ml_kernel_dir.canonicalize().unwrap_or(ml_kernel_dir).display()),
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Verify cubin compilation**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml -p ml-dqn`
|
||||
Expected: PASS (cubins recompile with header available)
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/state_layout.cuh crates/ml/build.rs crates/ml-dqn/build.rs
|
||||
git commit -m "feat: add state_layout.cuh — CUDA header with layout constants and assemble_state()"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Refactor experience_state_gather to Use Shared Assembly
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu`
|
||||
|
||||
This is the highest-risk task. The training kernel must use the new layout via `assemble_state()`.
|
||||
|
||||
- [ ] **Step 1: Add include at top of experience_kernels.cu**
|
||||
|
||||
At the top of the file (after any existing includes), add:
|
||||
|
||||
```c
|
||||
#include "state_layout.cuh"
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Refactor the state assembly section**
|
||||
|
||||
In `experience_state_gather` (lines ~580-703), the current code writes market features, portfolio features, MTF features, and OFI features at hardcoded positions. Replace the inner state assembly with a call to `assemble_state()`.
|
||||
|
||||
The key changes:
|
||||
- Compute `market[SL_MARKET_DIM]` from `market_features_buf` (existing code at lines ~580-590)
|
||||
- Compute `ofi[SL_OFI_DIM]` from `ofi_features` buffer (existing code at lines ~695-702, but fix index from 66 to use the shared function)
|
||||
- Compute `mtf[SL_MTF_DIM]` (existing code at lines ~625-680, move into local array)
|
||||
- Compute `portfolio[SL_PORTFOLIO_BASE_DIM]` (existing code at lines ~581-620, first 8 values)
|
||||
- Compute `plan_isv[SL_PORTFOLIO_PLAN_DIM]` (existing code at lines ~605-620, last 6 values)
|
||||
- Call `assemble_state(out, market, ofi, mtf, portfolio, plan_isv)`
|
||||
- Remove all hardcoded index constants (`ofi_start = 66`, `mtf_base = market_dim + 14`, etc.)
|
||||
|
||||
The outer loop (episode iteration, time-reversal, feature noise) stays unchanged. Only the inner assembly changes.
|
||||
|
||||
- [ ] **Step 3: Run local smoke test**
|
||||
|
||||
Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- test_fxcache_zero_copy_training --ignored --nocapture`
|
||||
Expected: PASS with `OFI_DIAG` showing non-zero values at NEW indices `[42..62)` (will fail until OFI_DIAG is updated in Task 6 — verify training completes without CUDA errors)
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/experience_kernels.cu
|
||||
git commit -m "refactor: experience_state_gather uses assemble_state() from state_layout.cuh"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Backtest Kernel — Delete and Replace
|
||||
|
||||
**Files:**
|
||||
- Delete: `crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu`
|
||||
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu` (add backtest entry point)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs`
|
||||
- Modify: `crates/ml/build.rs` (remove backtest_gather_kernel from compilation list)
|
||||
|
||||
- [ ] **Step 1: Add backtest entry point to experience_kernels.cu**
|
||||
|
||||
Add a new `__global__` kernel at the end of `experience_kernels.cu` that wraps `assemble_state()` with the backtest's outer loop (per-window, not per-episode):
|
||||
|
||||
```c
|
||||
// Backtest state gather — uses same assemble_state() as training.
|
||||
// Outer loop: one thread per window (simpler than episode-based experience collector).
|
||||
extern "C" __global__ void backtest_state_gather(
|
||||
const float* __restrict__ features, // [n_windows * max_len * feat_dim]
|
||||
const float* __restrict__ ofi_features, // [total_bars * SL_OFI_DIM]
|
||||
const float* __restrict__ portfolio, // [n_windows * 8] live portfolio state
|
||||
const float* __restrict__ plan_isv, // [n_windows * 6] plan/ISV features
|
||||
float* __restrict__ states_out, // [n_windows * padded_sd]
|
||||
int n_windows,
|
||||
int max_len,
|
||||
int feat_dim,
|
||||
int current_step,
|
||||
int total_bars,
|
||||
int padded_sd
|
||||
) {
|
||||
int w = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (w >= n_windows) return;
|
||||
|
||||
int feat_base = (w * max_len + current_step) * feat_dim;
|
||||
int out_base = w * padded_sd;
|
||||
float* out = states_out + out_base;
|
||||
|
||||
// Market features from pre-uploaded feature buffer
|
||||
float market[SL_MARKET_DIM];
|
||||
for (int k = 0; k < SL_MARKET_DIM; k++)
|
||||
market[k] = features[feat_base + k];
|
||||
|
||||
// OFI from feature buffer (appended after market features by data pipeline)
|
||||
// feat_dim = market_dim + ofi_dim = 62 when OFI is present
|
||||
int market_dim = feat_dim - SL_OFI_DIM;
|
||||
float ofi[SL_OFI_DIM];
|
||||
for (int k = 0; k < SL_OFI_DIM; k++)
|
||||
ofi[k] = features[feat_base + market_dim + k];
|
||||
|
||||
// MTF: compute from price lookbacks (same logic as training kernel)
|
||||
float mtf[SL_MTF_DIM];
|
||||
const int lookbacks[4] = {5, 15, 60, 240};
|
||||
int window_bar_base = w * max_len;
|
||||
for (int lb = 0; lb < 4; lb++) {
|
||||
int N_lb = lookbacks[lb];
|
||||
int past_step = current_step - N_lb;
|
||||
int slot = lb * 4;
|
||||
|
||||
if (past_step >= 0) {
|
||||
int now_off = (window_bar_base + current_step) * feat_dim;
|
||||
int past_off = (window_bar_base + past_step) * feat_dim;
|
||||
float close_cur = features[now_off];
|
||||
float close_past = features[past_off];
|
||||
|
||||
float ret = (close_past > 0.0f) ? (close_cur - close_past) / close_past : 0.0f;
|
||||
mtf[slot + 0] = fmaxf(-10.0f, fminf(10.0f, ret * 100.0f));
|
||||
|
||||
float max_val = close_cur, min_val = close_cur;
|
||||
float vol_sum = 0.0f;
|
||||
int vol_count = 0;
|
||||
for (int j = past_step; j <= current_step; j++) {
|
||||
int j_off = (window_bar_base + j) * feat_dim;
|
||||
float v = features[j_off];
|
||||
if (v > max_val) max_val = v;
|
||||
if (v < min_val) min_val = v;
|
||||
if (feat_dim > 4) { vol_sum += features[j_off + 4]; vol_count++; }
|
||||
}
|
||||
float range = (close_cur > 0.0f) ? (max_val - min_val) / close_cur : 0.0f;
|
||||
mtf[slot + 1] = fmaxf(0.0f, fminf(10.0f, range * 100.0f));
|
||||
float avg_vol = (vol_count > 0) ? vol_sum / (float)vol_count : 1.0f;
|
||||
float cur_vol = (feat_dim > 4) ? features[now_off + 4] : 1.0f;
|
||||
mtf[slot + 2] = fmaxf(0.0f, fminf(5.0f, (avg_vol > 0.0f) ? cur_vol / avg_vol : 1.0f));
|
||||
float range_size = max_val - min_val;
|
||||
mtf[slot + 3] = (range_size > 0.0f) ? (close_cur - min_val) / range_size : 0.5f;
|
||||
} else {
|
||||
mtf[slot + 0] = 0.0f; mtf[slot + 1] = 0.0f;
|
||||
mtf[slot + 2] = 0.0f; mtf[slot + 3] = 0.0f;
|
||||
}
|
||||
}
|
||||
|
||||
// Portfolio base: read from pre-computed array
|
||||
float pf[SL_PORTFOLIO_BASE_DIM];
|
||||
for (int k = 0; k < SL_PORTFOLIO_BASE_DIM; k++)
|
||||
pf[k] = portfolio[w * SL_PORTFOLIO_BASE_DIM + k];
|
||||
|
||||
// Plan/ISV: read from pre-computed array
|
||||
float pisv[SL_PORTFOLIO_PLAN_DIM];
|
||||
for (int k = 0; k < SL_PORTFOLIO_PLAN_DIM; k++)
|
||||
pisv[k] = plan_isv[w * SL_PORTFOLIO_PLAN_DIM + k];
|
||||
|
||||
assemble_state(out, market, ofi, mtf, pf, pisv);
|
||||
|
||||
// Zero-pad [STATE_DIM .. padded_sd)
|
||||
for (int k = SL_STATE_DIM; k < padded_sd; k++)
|
||||
out[k] = 0.0f;
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Delete backtest_gather_kernel.cu**
|
||||
|
||||
```bash
|
||||
rm crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Update build.rs to remove backtest_gather_kernel from compilation**
|
||||
|
||||
In `crates/ml/build.rs`, find the list of `.cu` files compiled to cubin. Remove `backtest_gather_kernel.cu` from the list.
|
||||
|
||||
- [ ] **Step 4: Update gpu_backtest_evaluator.rs**
|
||||
|
||||
The evaluator currently loads `gather_states` from `backtest_gather_kernel.cubin`. Change it to load `backtest_state_gather` from the experience kernels cubin. Also:
|
||||
- Add `plan_isv_buf: CudaSlice<f32>` field (6 floats per window) to the evaluator
|
||||
- Compute the 6 plan/ISV features during the backtest step loop (conviction from ISV, plan progress from portfolio sim)
|
||||
- Pass `plan_isv_buf` as a kernel argument
|
||||
|
||||
The portfolio computation logic (8 features) already exists in the backtest evaluator's step kernel — keep it. The 6 new plan/ISV features need to be computed alongside it.
|
||||
|
||||
- [ ] **Step 5: Verify local smoke test**
|
||||
|
||||
Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- test_fxcache_zero_copy_training --ignored --nocapture`
|
||||
Expected: PASS — training and validation both use the same layout now
|
||||
|
||||
- [ ] **Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git add -A
|
||||
git commit -m "feat: delete backtest_gather_kernel, backtest uses shared assemble_state()"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: Remove Configurable state_dim — Use Constants Everywhere
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml-dqn/src/dqn.rs` (DQNConfig)
|
||||
- Modify: `crates/ml-dqn/src/gpu_replay_buffer.rs` (GpuReplayBufferConfig)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
|
||||
- Modify: `crates/ml/src/trainers/dqn/trainer/constructor.rs`
|
||||
- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs`
|
||||
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`
|
||||
- Modify: ~28 additional files with 1-3 reads each
|
||||
|
||||
This is a mechanical replacement. The pattern is:
|
||||
- `config.state_dim` → `ml_core::state_layout::STATE_DIM`
|
||||
- `self.state_dim` → `ml_core::state_layout::STATE_DIM`
|
||||
- `pad128(config.state_dim)` → `ml_core::state_layout::STATE_DIM_PADDED`
|
||||
- Remove `state_dim` field from structs where it was stored
|
||||
- Remove `state_dim` parameter from constructors where it was passed
|
||||
|
||||
- [ ] **Step 1: DQNConfig — remove state_dim field**
|
||||
|
||||
In `crates/ml-dqn/src/dqn.rs`:
|
||||
- Remove `pub state_dim: usize` from `DQNConfig` struct (line 30)
|
||||
- Remove `state_dim: 48` from the `Default` impl (line 247)
|
||||
- Replace all `config.state_dim` reads with `ml_core::state_layout::STATE_DIM` (11 sites in this file)
|
||||
- Update `get_state_dim()` to return the constant
|
||||
- Update checkpoint metadata: save `STATE_DIM` instead of config field, validate on load
|
||||
|
||||
- [ ] **Step 2: GpuReplayBufferConfig — remove state_dim field**
|
||||
|
||||
In `crates/ml-dqn/src/gpu_replay_buffer.rs`:
|
||||
- Remove `pub state_dim: usize` from config struct (line 126)
|
||||
- Replace all `self.config.state_dim` reads with `ml_core::state_layout::STATE_DIM` (5 sites)
|
||||
- Remove `state_dim` parameter from `StagedGpuBuffer::try_with_halving` signature
|
||||
- Update all callers of `try_with_halving` (remove the state_dim argument)
|
||||
|
||||
- [ ] **Step 3: GpuExperienceCollector — remove state_dim field**
|
||||
|
||||
In `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`:
|
||||
- Remove `state_dim: usize` field (line 479)
|
||||
- Remove `state_dim` from constructor parameters
|
||||
- Replace all `self.state_dim` reads with `ml_core::state_layout::STATE_DIM` (~10 sites)
|
||||
- Remove `state_dim` from `kernel_dims` tuple in `training_loop.rs`
|
||||
|
||||
- [ ] **Step 4: GpuDqnTrainer — remove config.state_dim reads**
|
||||
|
||||
In `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`:
|
||||
- Replace all `self.config.state_dim` reads with `ml_core::state_layout::STATE_DIM` (~37 sites)
|
||||
- Replace `pad128(self.config.state_dim)` with `ml_core::state_layout::STATE_DIM_PADDED`
|
||||
- Remove the local `pad128` function (line 12625) — use the constant directly
|
||||
|
||||
- [ ] **Step 5: Constructor — remove state_dim computation**
|
||||
|
||||
In `crates/ml/src/trainers/dqn/trainer/constructor.rs`:
|
||||
- Remove `let raw_state_dim: usize = 92;` and `let full_state_dim = ...` (lines 225-230)
|
||||
- Use `ml_core::state_layout::STATE_DIM` directly in the DQNConfig construction
|
||||
- Remove `state_dim` from the `kernel_dims` tuple
|
||||
|
||||
- [ ] **Step 6: Remaining files — bulk mechanical replacement**
|
||||
|
||||
Search all `.rs` files in `crates/ml/` and `crates/ml-dqn/` for remaining `state_dim` reads and replace with the constant. Use:
|
||||
|
||||
```bash
|
||||
grep -rn "config\.state_dim\|self\.state_dim\|\.state_dim()" crates/ml/src crates/ml-dqn/src --include="*.rs" | grep -v test | grep -v "//\|///\|#\["
|
||||
```
|
||||
|
||||
Replace each with `ml_core::state_layout::STATE_DIM` (or `STATE_DIM_PADDED` for padded variants).
|
||||
|
||||
- [ ] **Step 7: Verify compilation**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check --workspace`
|
||||
Expected: PASS — all state_dim references use the constant
|
||||
|
||||
- [ ] **Step 8: Commit**
|
||||
|
||||
```bash
|
||||
git add -A
|
||||
git commit -m "refactor: remove configurable state_dim — use STATE_DIM constant everywhere"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 6: Update OFI_DIAG and Checkpoint Loading
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs` (OFI_DIAG indices)
|
||||
- Modify: `crates/ml-dqn/src/dqn.rs` (checkpoint validation)
|
||||
|
||||
- [ ] **Step 1: Update OFI_DIAG indices**
|
||||
|
||||
In `training_loop.rs` (lines 1949-1952), update the hardcoded indices to match the new layout:
|
||||
|
||||
```rust
|
||||
// Old: sample[66..74], sample[74..82], sample[82], sample[83]
|
||||
// New: OFI at [42..62) per state_layout
|
||||
let ofi_raw: &[f32] = &sample[42..50]; // was [66..74]
|
||||
let ofi_delta: &[f32] = &sample[50..58]; // was [74..82]
|
||||
let book_agg = sample[58]; // was [82]
|
||||
let log_dur = sample[59]; // was [83]
|
||||
```
|
||||
|
||||
Also update the guard: `sample.len() > 83` → `sample.len() >= STATE_DIM` (use the constant).
|
||||
|
||||
- [ ] **Step 2: Update checkpoint loading**
|
||||
|
||||
In `crates/ml-dqn/src/dqn.rs`, the checkpoint metadata saves/loads `dqn.state_dim`. Since state_dim is now a constant:
|
||||
- On save: write `STATE_DIM` to metadata (no change needed if constructor already uses the constant)
|
||||
- On load: validate that saved `state_dim` equals `STATE_DIM`. If not, return error:
|
||||
|
||||
```rust
|
||||
let saved_sd: usize = get("dqn.state_dim")?.parse()
|
||||
.map_err(|e| MLError::CheckpointError(format!("Bad dqn.state_dim: {e}")))?;
|
||||
if saved_sd != ml_core::state_layout::STATE_DIM {
|
||||
return Err(MLError::CheckpointError(format!(
|
||||
"Checkpoint state_dim={} but STATE_DIM={}. Old checkpoints are invalid (PER bug: model trained on garbage).",
|
||||
saved_sd, ml_core::state_layout::STATE_DIM,
|
||||
)));
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Verify smoke test with updated indices**
|
||||
|
||||
Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- test_fxcache_zero_copy_training --ignored --nocapture`
|
||||
Expected: `OFI_DIAG: raw_mean=...` shows non-zero values at the new `[42..62)` positions
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/trainers/dqn/trainer/training_loop.rs crates/ml-dqn/src/dqn.rs
|
||||
git commit -m "fix: update OFI_DIAG indices to new layout, reject old checkpoints"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 7: Layout Match Integration Test
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs`
|
||||
|
||||
- [ ] **Step 1: Add layout match test**
|
||||
|
||||
Add a test that verifies training and validation produce identical state layouts for the same bar. This is the definitive test that the mismatch is fixed:
|
||||
|
||||
```rust
|
||||
#[test]
|
||||
#[ignore] // requires GPU
|
||||
fn test_train_val_state_layout_match() -> anyhow::Result<()> {
|
||||
// 1. Load fxcache
|
||||
// 2. Create trainer, upload data
|
||||
// 3. Run 1 experience collection step → readback batch_states[0]
|
||||
// 4. Run 1 validation backtest step for the same bar → readback states_out[0]
|
||||
// 5. Assert state vectors match at all STATE_DIM positions within f32 tolerance
|
||||
// Specifically: OFI at [42..62), MTF at [62..78), portfolio at [78..86), plan at [86..92)
|
||||
//
|
||||
// This test FAILS if any code path assembles states differently.
|
||||
}
|
||||
```
|
||||
|
||||
The exact implementation depends on the readback APIs available after Tasks 3-4. The test should:
|
||||
- Read back one training state from `batch_states` (via pinned DtoH)
|
||||
- Read back one validation state from the backtest evaluator's `states_out`
|
||||
- Compare all 96 positions element-by-element
|
||||
- Market features [0..42) should match exactly (same source data)
|
||||
- OFI [42..62) should match exactly (same source data)
|
||||
- MTF [62..78) should match within tolerance (same computation, possible float ordering differences)
|
||||
- Portfolio [78..86) may differ (different simulation paths) — assert non-zero in both
|
||||
- Plan/ISV [86..92) may differ — assert non-zero in both
|
||||
|
||||
- [ ] **Step 2: Run the test**
|
||||
|
||||
Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- test_train_val_state_layout_match --ignored --nocapture`
|
||||
Expected: PASS
|
||||
|
||||
- [ ] **Step 3: Run full smoke test suite**
|
||||
|
||||
Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture`
|
||||
Expected: All existing tests pass with the new layout
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs
|
||||
git commit -m "test: add layout match test — training and validation states must be identical"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 8: Deploy and Verify on H100
|
||||
|
||||
**Files:** None (deployment only)
|
||||
|
||||
- [ ] **Step 1: Push and deploy**
|
||||
|
||||
```bash
|
||||
git push
|
||||
./scripts/argo-train.sh dqn --sha HEAD --baseline
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Verify OFI_DIAG at new indices**
|
||||
|
||||
Check logs for: `OFI_DIAG: raw_mean=X.XXXX` where X is non-zero.
|
||||
The indices are now [42..62) instead of [66..84).
|
||||
|
||||
- [ ] **Step 3: Verify val_Sharpe trajectory**
|
||||
|
||||
Watch epochs 1-10. val_Sharpe should no longer catastrophically diverge from training Sharpe. The exact values depend on the model learning, but the distribution mismatch is eliminated.
|
||||
|
||||
- [ ] **Step 4: Monitor for CUDA errors**
|
||||
|
||||
Any CUDA_ERROR_ILLEGAL_ADDRESS or kernel launch failures indicate a layout mismatch between the CUDA header and Rust constants. If seen, verify `state_layout.cuh` constants match `state_layout.rs`.
|
||||
Reference in New Issue
Block a user