Compare commits
3 Commits
worktree-a
...
ml-alpha-d
| Author | SHA1 | Date | |
|---|---|---|---|
|
251b871294 | ||
|
|
b7bbb700f8 | ||
|
|
84caa99b65 |
1
Cargo.lock
generated
1
Cargo.lock
generated
@@ -6090,7 +6090,6 @@ dependencies = [
|
||||
"approx",
|
||||
"arrow 56.2.0",
|
||||
"bincode",
|
||||
"bytemuck",
|
||||
"clap",
|
||||
"cudarc",
|
||||
"data",
|
||||
|
||||
@@ -33,7 +33,6 @@ tracing-subscriber.workspace = true
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
serde_json.workspace = true
|
||||
bincode.workspace = true
|
||||
bytemuck = { workspace = true }
|
||||
thiserror.workspace = true
|
||||
clap = { workspace = true, features = ["derive"] }
|
||||
|
||||
|
||||
@@ -73,7 +73,6 @@ const KERNELS: &[&str] = &[
|
||||
"rl_kl_reference_grad",
|
||||
"rl_q_pi_distill_grad", // audit 2026-05-24 vj5f6 followup: KL(softmax(E_Q/τ) || π_new) gradient ADDED to pi_grad_logits — couples Q's improved C51 calibration to action selection (was decoupled per Option B)
|
||||
"action_entropy_per_step", // POST-gate action entropy EMA for SAC α/τ co-tuning
|
||||
"rl_drawdown_stop", // per-step drawdown penalty + hard stop-loss
|
||||
"rl_q_distill_lambda_controller",// audit 2026-05-24 rljzl followup: adaptive λ_distill via Schulman bounded step on KL_EMA vs target
|
||||
"rl_unit_state_update", // SP20 P1+P5 audit fix: per-unit trade state machine — detects open/close/reverse transitions, sets up unit slot 0 entry+trail
|
||||
"rl_trail_mutate", // SP20 P1+P5 audit fix (a7/a8 dead): TrailTighten/TrailLoosen mutate unit_trail_distance bounded MIN/MAX, symmetric reciprocal adjust
|
||||
|
||||
@@ -24,8 +24,6 @@
|
||||
|
||||
#define BOOK_LEVELS 10
|
||||
#define REGIME_DIM 6
|
||||
#define N_HORIZONS 3
|
||||
#define N_LABEL_SOURCES 5
|
||||
|
||||
// Must match #[repr(C)] Mbp10RawInput in snap_features.rs (216 bytes).
|
||||
struct __align__(8) Mbp10Raw {
|
||||
@@ -82,26 +80,11 @@ extern "C" __global__ void gpu_sample_and_gather(
|
||||
long long* __restrict__ ts_ns_soa, // [B*K]
|
||||
long long* __restrict__ prev_ts_ns_soa, // [B*K]
|
||||
// Global-memory outputs for the sampled (file_offset, anchor, file_idx)
|
||||
// per batch element. Read by gpu_gather_next and gpu_gather_frd_labels
|
||||
// to reuse the same sampling decision.
|
||||
// per batch element. Read by gpu_gather_next, gpu_gather_frd_labels,
|
||||
// and gpu_gather_bce_labels to reuse the same sampling decision.
|
||||
int* __restrict__ out_file_offset, // [B]
|
||||
int* __restrict__ out_anchor, // [B]
|
||||
int* __restrict__ out_file_idx, // [B]
|
||||
// Fused label gather — 5 label sources (BCE, prof_long, prof_short,
|
||||
// size_long, size_short), each [N_HORIZONS × total_snaps]. Reads in
|
||||
// the same pass as snapshot gather to avoid 5 separate random-access
|
||||
// kernel launches that consumed 95.6% of GPU time.
|
||||
int total_snaps,
|
||||
const float* __restrict__ labels_src_0, // [N_HORIZONS × total_snaps]
|
||||
const float* __restrict__ labels_src_1,
|
||||
const float* __restrict__ labels_src_2,
|
||||
const float* __restrict__ labels_src_3,
|
||||
const float* __restrict__ labels_src_4,
|
||||
float* __restrict__ labels_out_0, // [K, B, N_HORIZONS]
|
||||
float* __restrict__ labels_out_1,
|
||||
float* __restrict__ labels_out_2,
|
||||
float* __restrict__ labels_out_3,
|
||||
float* __restrict__ labels_out_4,
|
||||
int B
|
||||
) {
|
||||
int b = blockIdx.x;
|
||||
@@ -165,26 +148,9 @@ extern "C" __global__ void gpu_sample_and_gather(
|
||||
tc_soa[n] = (int)snap.trade_count;
|
||||
ts_ns_soa[n] = (long long)snap.ts_ns;
|
||||
prev_ts_ns_soa[n] = (long long)snap.prev_ts_ns;
|
||||
|
||||
// Fused label gather — same global_idx, zero extra random access.
|
||||
// Output layout: [K, B, N_HORIZONS] row-major.
|
||||
int lbl_base = (k * B + b) * N_HORIZONS;
|
||||
const float* srcs[N_LABEL_SOURCES] = {
|
||||
labels_src_0, labels_src_1, labels_src_2, labels_src_3, labels_src_4
|
||||
};
|
||||
float* outs[N_LABEL_SOURCES] = {
|
||||
labels_out_0, labels_out_1, labels_out_2, labels_out_3, labels_out_4
|
||||
};
|
||||
#pragma unroll
|
||||
for (int s = 0; s < N_LABEL_SOURCES; s++) {
|
||||
#pragma unroll
|
||||
for (int h = 0; h < N_HORIZONS; h++) {
|
||||
outs[s][lbl_base + h] = srcs[s][h * total_snaps + global_idx];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Label gather kernel (STANDALONE — kept for backward compat) ─────
|
||||
// ── Label gather kernel ──────────────────────────────────────────────
|
||||
//
|
||||
// Runs AFTER gpu_sample_and_gather. Gathers per-horizon FRD labels at
|
||||
// the anchor position (rightmost K position = newest snapshot in the
|
||||
|
||||
@@ -34,8 +34,6 @@
|
||||
#define RL_HOLD_TARGET_FRAC_INDEX 575
|
||||
#define RL_HOLD_FRAC_EMA_INDEX 576
|
||||
#define RL_CONF_GATE_MAX_HOLD_FRAC_INDEX 584
|
||||
#define RL_STOP_LOSS_THRESHOLD_INDEX 587
|
||||
#define RL_MEAN_ABS_PNL_EMA_INDEX 423
|
||||
#define THRESHOLD_MIN 0.05f
|
||||
#define THRESHOLD_MAX 0.95f
|
||||
#define HOLD_EMA_ALPHA 0.1f
|
||||
@@ -53,35 +51,6 @@ extern "C" __global__ void rl_confidence_gate(
|
||||
const int b = blockIdx.x;
|
||||
if (b >= b_size) return;
|
||||
|
||||
// ── Hard stop-loss: override action to Flat when unrealized loss
|
||||
// exceeds ISV-driven threshold. Fires BEFORE gate logic so the
|
||||
// lobsim actually closes the position (no state mismatch).
|
||||
// unrealized_loss = |vwap - mid| × |lots|; normalized by mean_abs_pnl.
|
||||
{
|
||||
const unsigned char* p = pos_state + b * pos_bytes;
|
||||
const int lots = *(const int*)(p + 0);
|
||||
if (lots != 0) {
|
||||
const float stop_thresh = isv[RL_STOP_LOSS_THRESHOLD_INDEX];
|
||||
const float mean_abs = fmaxf(isv[RL_MEAN_ABS_PNL_EMA_INDEX], 1e-6f);
|
||||
// Approximate mid from the first bid/ask level would need
|
||||
// book data; instead use the realized_pnl delta as a proxy
|
||||
// for position health. Simpler: use vwap vs realized_pnl trend.
|
||||
// For now, use the unit_entry_price-based unrealized_r from
|
||||
// trade_context — but that runs AFTER this kernel.
|
||||
//
|
||||
// Practical approach: read peak_equity (offset 12) and
|
||||
// realized_pnl (offset 8). Drawdown from peak = peak - realized.
|
||||
const float realized = *(const float*)(p + 8);
|
||||
const float peak = *(const float*)(p + 12);
|
||||
const float drawdown = peak - realized;
|
||||
const float normalized_dd = drawdown / mean_abs;
|
||||
if (normalized_dd > stop_thresh) {
|
||||
actions[b] = (lots > 0) ? 3 : 4; // FlatFromLong / FlatFromShort
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const int current_step = (int)isv[RL_STEP_COUNTER_ISV_INDEX];
|
||||
const int warmup = (int)isv[RL_GATE_WARMUP_STEPS_INDEX];
|
||||
if (current_step < warmup) return;
|
||||
@@ -96,7 +65,6 @@ extern "C" __global__ void rl_confidence_gate(
|
||||
|
||||
const int action = actions[b];
|
||||
if (action == ACTION_HOLD) return;
|
||||
if (action == 3 || action == 4) return; // FlatFromLong/Short: never gate exits
|
||||
|
||||
const float threshold = isv[RL_CONF_GATE_THRESHOLD_INDEX];
|
||||
const float lambda = isv[RL_CONF_GATE_LAMBDA_INDEX];
|
||||
|
||||
@@ -1,49 +0,0 @@
|
||||
// rl_drawdown_stop.cu — per-step drawdown penalty + hard stop-loss.
|
||||
//
|
||||
// Runs AFTER rl_trade_context_update (which computes unrealized_r).
|
||||
// Reads unrealized_r from trade_context_d and applies:
|
||||
//
|
||||
// 1. Per-step drawdown penalty: adds min(0, unrealized_r) × PENALTY_RATE
|
||||
// to rewards[b]. Creates continuous exit gradient on losing positions.
|
||||
// Penalty-only (never positive) = no exposure incentive.
|
||||
//
|
||||
// 2. Hard stop-loss: when unrealized_r < -STOP_THRESHOLD, force-close
|
||||
// by setting dones[b] = 1. The realized PnL at this point becomes
|
||||
// the final trade reward. Caps max loss per trade.
|
||||
//
|
||||
// Both thresholds are ISV-driven per feedback_isv_for_adaptive_bounds.
|
||||
// Grid: ceil(B/32), Block: 32.
|
||||
|
||||
#define RL_DRAWDOWN_PENALTY_RATE_INDEX 586
|
||||
#define RL_STOP_LOSS_THRESHOLD_INDEX 587
|
||||
#define TRADE_CONTEXT_STRIDE 5
|
||||
#define UNREALIZED_R_OFFSET 1
|
||||
|
||||
extern "C" __global__ void rl_drawdown_stop(
|
||||
const float* __restrict__ trade_context, // [B, TRADE_CONTEXT_STRIDE]
|
||||
float* __restrict__ rewards, // [B] IN/OUT
|
||||
float* __restrict__ dones, // [B] IN/OUT
|
||||
float* __restrict__ isv,
|
||||
int b_size
|
||||
) {
|
||||
const int b = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (b >= b_size) return;
|
||||
|
||||
const float unrealized_r = trade_context[b * TRADE_CONTEXT_STRIDE + UNREALIZED_R_OFFSET];
|
||||
const float penalty_rate = isv[RL_DRAWDOWN_PENALTY_RATE_INDEX];
|
||||
const float stop_thresh = isv[RL_STOP_LOSS_THRESHOLD_INDEX];
|
||||
|
||||
// Per-step drawdown penalty: only fires when losing (unrealized_r < 0).
|
||||
// min(0, x) * rate is always <= 0.
|
||||
if (unrealized_r < 0.0f) {
|
||||
rewards[b] += unrealized_r * penalty_rate;
|
||||
}
|
||||
|
||||
// Hard stop-loss DISABLED: setting dones=1 here creates a state
|
||||
// mismatch — the lobsim position stays open while Q sees a false
|
||||
// close. The done signal must come from actual position changes
|
||||
// in the lobsim, not synthetic overrides. To implement hard
|
||||
// stop-loss properly, the action must be overridden to FlatL/FlatS
|
||||
// BEFORE the lobsim step, not after.
|
||||
(void) stop_thresh;
|
||||
}
|
||||
@@ -230,15 +230,6 @@ extern "C" __global__ void rl_reward_clamp_controller(
|
||||
// Suppress unused warnings on diagnostic-only state.
|
||||
(void) margin;
|
||||
|
||||
// Adaptive LOSS clamp from observed neg/pos ratio. WIN stays
|
||||
// structural at 1.0; LOSS adapts to actual loss magnitude to
|
||||
// give Q accurate resolution without over-allocating to rare tails.
|
||||
if (ema_new > 0.0f && neg_new > 0.0f) {
|
||||
const float observed_ratio = neg_new / ema_new;
|
||||
const float adaptive_loss = fmaxf(1.0f, fminf(observed_ratio * 1.1f, 3.0f));
|
||||
isv[RL_REWARD_CLAMP_LOSS_INDEX] = adaptive_loss;
|
||||
}
|
||||
|
||||
// ── Step 5: C51 atom span adaptation from observed reward EMAs. ──
|
||||
//
|
||||
// G.2 disabled this because atom-span growth + clamp-bound growth
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -188,7 +188,6 @@ impl GpuDataLoader {
|
||||
soa: &SoaBufferPtrs,
|
||||
seq_len: usize,
|
||||
batch_size: usize,
|
||||
label_outs: [u64; 5],
|
||||
) -> Result<()> {
|
||||
let mut args = RawArgs::new();
|
||||
args.push_ptr(dataset.snapshots_d.raw_ptr());
|
||||
@@ -211,18 +210,6 @@ impl GpuDataLoader {
|
||||
args.push_ptr(self.sample_file_offset_d.raw_ptr());
|
||||
args.push_ptr(self.sample_anchor_d.raw_ptr());
|
||||
args.push_ptr(self.sample_file_idx_d.raw_ptr());
|
||||
// Fused label gather arguments
|
||||
args.push_i32(dataset.total_snapshots as i32);
|
||||
args.push_ptr(dataset.labels_d.raw_ptr());
|
||||
args.push_ptr(dataset.outcome_prof_long_d.raw_ptr());
|
||||
args.push_ptr(dataset.outcome_prof_short_d.raw_ptr());
|
||||
args.push_ptr(dataset.outcome_size_long_d.raw_ptr());
|
||||
args.push_ptr(dataset.outcome_size_short_d.raw_ptr());
|
||||
args.push_ptr(label_outs[0]);
|
||||
args.push_ptr(label_outs[1]);
|
||||
args.push_ptr(label_outs[2]);
|
||||
args.push_ptr(label_outs[3]);
|
||||
args.push_ptr(label_outs[4]);
|
||||
args.push_i32(batch_size as i32);
|
||||
let mut ptrs = args.build_arg_ptrs();
|
||||
unsafe {
|
||||
|
||||
@@ -547,12 +547,6 @@ impl Mamba2Block {
|
||||
)
|
||||
.result()
|
||||
.map_err(|e| anyhow!("Mamba2Block: cublasSetWorkspace_v2: {e:?}"))?;
|
||||
cudarc::cublas::sys::cublasSetMathMode(
|
||||
*cublas.handle(),
|
||||
cudarc::cublas::sys::cublasMath_t::CUBLAS_TF32_TENSOR_OP_MATH,
|
||||
)
|
||||
.result()
|
||||
.map_err(|e| anyhow!("Mamba2Block: cublasSetMathMode TF32: {e:?}"))?;
|
||||
}
|
||||
|
||||
// ── Parameter allocation + initialisation. ──────────────────────
|
||||
|
||||
@@ -311,12 +311,6 @@ impl DqnHead {
|
||||
)
|
||||
.result()
|
||||
.map_err(|e| anyhow::anyhow!("DqnHead: cublasSetWorkspace_v2: {e:?}"))?;
|
||||
cudarc::cublas::sys::cublasSetMathMode(
|
||||
*cublas.handle(),
|
||||
cudarc::cublas::sys::cublasMath_t::CUBLAS_TF32_TENSOR_OP_MATH,
|
||||
)
|
||||
.result()
|
||||
.map_err(|e| anyhow::anyhow!("DqnHead: cublasSetMathMode TF32: {e:?}"))?;
|
||||
}
|
||||
|
||||
// Bias-add and reduce-sum-axis0 kernel handles from ml-core.
|
||||
|
||||
@@ -1111,16 +1111,5 @@ pub const RL_ACTION_ENTROPY_EMA_INDEX: usize = 583;
|
||||
/// Bootstrap: 0.85 (15% exploration floor).
|
||||
pub const RL_CONF_GATE_MAX_HOLD_FRAC_INDEX: usize = 584;
|
||||
|
||||
/// Per-step drawdown penalty rate. Applied as min(0, unrealized_r) × rate
|
||||
/// to the reward each step a position is losing. Penalty-only (never
|
||||
/// positive) creates continuous exit gradient without exposure incentive.
|
||||
/// Bootstrap: 0.01 (1% of unrealized_r magnitude per step).
|
||||
pub const RL_DRAWDOWN_PENALTY_RATE_INDEX: usize = 586;
|
||||
|
||||
/// Hard stop-loss threshold in initial-risk multiples. When unrealized_r
|
||||
/// drops below -threshold, the position force-closes (done=1). Caps max
|
||||
/// loss per trade. Bootstrap: 2.0 (2× initial risk).
|
||||
pub const RL_STOP_LOSS_THRESHOLD_INDEX: usize = 587;
|
||||
|
||||
/// Last RL-allocated slot index (exclusive).
|
||||
pub const RL_SLOTS_END: usize = 588;
|
||||
pub const RL_SLOTS_END: usize = 585;
|
||||
|
||||
@@ -150,8 +150,6 @@ const COMPUTE_ADVANTAGE_RETURN_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/compute_advantage_return.cubin"));
|
||||
const ACTION_ENTROPY_PER_STEP_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/action_entropy_per_step.cubin"));
|
||||
const RL_DRAWDOWN_STOP_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_drawdown_stop.cubin"));
|
||||
|
||||
// Phase R4: GPU-resident action sampling kernels. Replace the host
|
||||
// Thompson + host argmax + host log-softmax loops the flawed Phase F
|
||||
@@ -568,8 +566,6 @@ pub struct IntegratedTrainer {
|
||||
compute_advantage_return_fn: CudaFunction,
|
||||
_action_entropy_per_step_module: Arc<CudaModule>,
|
||||
action_entropy_per_step_fn: CudaFunction,
|
||||
_rl_drawdown_stop_module: Arc<CudaModule>,
|
||||
rl_drawdown_stop_fn: CudaFunction,
|
||||
|
||||
// ── Phase R4: GPU action sampling ─────────────────────────────────
|
||||
_rl_action_kernel_module: Arc<CudaModule>,
|
||||
@@ -1331,12 +1327,6 @@ impl IntegratedTrainer {
|
||||
let action_entropy_per_step_module = ctx
|
||||
.load_cubin(ACTION_ENTROPY_PER_STEP_CUBIN.to_vec())
|
||||
.context("load action_entropy_per_step cubin")?;
|
||||
let rl_drawdown_stop_module = ctx
|
||||
.load_cubin(RL_DRAWDOWN_STOP_CUBIN.to_vec())
|
||||
.context("load rl_drawdown_stop cubin")?;
|
||||
let rl_drawdown_stop_fn = rl_drawdown_stop_module
|
||||
.load_function("rl_drawdown_stop")
|
||||
.context("load rl_drawdown_stop")?;
|
||||
let action_entropy_per_step_fn = action_entropy_per_step_module
|
||||
.load_function("action_entropy_per_step")
|
||||
.context("load action_entropy_per_step")?;
|
||||
@@ -2333,8 +2323,6 @@ impl IntegratedTrainer {
|
||||
ema_update_per_step_fn,
|
||||
_action_entropy_per_step_module: action_entropy_per_step_module,
|
||||
action_entropy_per_step_fn,
|
||||
_rl_drawdown_stop_module: rl_drawdown_stop_module,
|
||||
rl_drawdown_stop_fn,
|
||||
_compute_advantage_return_module: compute_advantage_return_module,
|
||||
compute_advantage_return_fn,
|
||||
_rl_action_kernel_module: rl_action_kernel_module,
|
||||
@@ -2685,7 +2673,7 @@ impl IntegratedTrainer {
|
||||
// (slot, value) pair — pure device write, no HtoD per
|
||||
// `feedback_no_htod_htoh_only_mapped_pinned`.
|
||||
{
|
||||
let isv_constants: [(usize, f32); 113] = [
|
||||
let isv_constants: [(usize, f32); 111] = [
|
||||
// Static seeds for the adaptive reward-clamp controller —
|
||||
// these are the initial values that
|
||||
// `rl_reward_clamp_controller` will replace once it
|
||||
@@ -2852,8 +2840,6 @@ impl IntegratedTrainer {
|
||||
(crate::rl::isv_slots::RL_SAC_ALPHA_INDEX, 0.01),
|
||||
(crate::rl::isv_slots::RL_SAC_ENTROPY_TARGET_INDEX, 1.68),
|
||||
(crate::rl::isv_slots::RL_CONF_GATE_MAX_HOLD_FRAC_INDEX, 0.85),
|
||||
(crate::rl::isv_slots::RL_DRAWDOWN_PENALTY_RATE_INDEX, 0.01),
|
||||
(crate::rl::isv_slots::RL_STOP_LOSS_THRESHOLD_INDEX, 2.0),
|
||||
];
|
||||
for (slot, value) in isv_constants.iter() {
|
||||
let slot_i32 = *slot as i32;
|
||||
@@ -6025,29 +6011,6 @@ impl IntegratedTrainer {
|
||||
}
|
||||
}
|
||||
|
||||
// Per-step drawdown penalty + hard stop-loss. Reads unrealized_r
|
||||
// from trade_context_d (just written above). Modifies rewards_d
|
||||
// and dones_d in place before apply_reward_scale.
|
||||
{
|
||||
let grid_x = ((b_size as u32) + 31) / 32;
|
||||
let b_size_i = b_size as i32;
|
||||
let mut args = RawArgs::new();
|
||||
args.push_ptr(self.trade_context_d.raw_ptr());
|
||||
args.push_ptr(self.rewards_d.raw_ptr());
|
||||
args.push_ptr(self.dones_d.raw_ptr());
|
||||
args.push_ptr(self.isv_dev_ptr);
|
||||
args.push_i32(b_size_i);
|
||||
let mut ptrs = args.build_arg_ptrs();
|
||||
unsafe {
|
||||
raw_launch(
|
||||
self.rl_drawdown_stop_fn.cu_function(),
|
||||
(grid_x.max(1), 1, 1), (32, 1, 1), 0,
|
||||
self.raw_stream,
|
||||
&mut ptrs[..args.len()],
|
||||
).map_err(|e| anyhow::anyhow!("rl_drawdown_stop: {:?}", e))?;
|
||||
}
|
||||
}
|
||||
|
||||
// Multi-resolution streaming features — time-weighted EMA at 3 horizons.
|
||||
{
|
||||
let bid_px_d = lobsim.bid_px_d();
|
||||
@@ -6886,15 +6849,8 @@ impl IntegratedTrainer {
|
||||
// 3. forward_encoder_from_device: runs encoder on SoA (= s_{t+1}).
|
||||
// 4. DtoD h_t -> h_tp1: saves h_{t+1}.
|
||||
let soa = self.perception.soa_buffer_ptrs();
|
||||
let label_stg = self.perception.label_staging_ptrs();
|
||||
gpu_loader
|
||||
.sample_and_gather(gpu_dataset, &soa, seq_len, b_size, [
|
||||
label_stg.labels,
|
||||
label_stg.aux_y_prof_long,
|
||||
label_stg.aux_y_prof_short,
|
||||
label_stg.aux_y_size_long,
|
||||
label_stg.aux_y_size_short,
|
||||
])
|
||||
.sample_and_gather(gpu_dataset, &soa, seq_len, b_size)
|
||||
.context("step_with_lobsim_gpu: sample_and_gather")?;
|
||||
gpu_loader
|
||||
.gather_next_window(gpu_dataset, &soa, seq_len, b_size)
|
||||
@@ -6939,10 +6895,47 @@ impl IntegratedTrainer {
|
||||
// After this call the encoder weights are updated with BCE+aux
|
||||
// gradient signal, so the subsequent RL forward on h_t sees the
|
||||
// jointly-supervised encoder.
|
||||
// Label gathers fused into sample_and_gather above — 5 separate
|
||||
// random-access launches eliminated (was 95.6% of GPU time per nsys).
|
||||
{
|
||||
use crate::heads::N_HORIZONS as PERC_N_HORIZONS;
|
||||
let stg = self.perception.label_staging_ptrs();
|
||||
let total_snaps = gpu_dataset.total_snapshots;
|
||||
|
||||
// BCE labels → stg_labels
|
||||
gpu_loader
|
||||
.gather_bce_labels(
|
||||
gpu_dataset.labels_d.raw_ptr(), total_snaps,
|
||||
PERC_N_HORIZONS, seq_len, b_size, stg.labels,
|
||||
)
|
||||
.context("step_with_lobsim_gpu: gather_bce_labels")?;
|
||||
// Aux prof-long → stg_aux_y_prof_long
|
||||
gpu_loader
|
||||
.gather_bce_labels(
|
||||
gpu_dataset.outcome_prof_long_d.raw_ptr(), total_snaps,
|
||||
PERC_N_HORIZONS, seq_len, b_size, stg.aux_y_prof_long,
|
||||
)
|
||||
.context("step_with_lobsim_gpu: gather outcome_prof_long")?;
|
||||
// Aux prof-short → stg_aux_y_prof_short
|
||||
gpu_loader
|
||||
.gather_bce_labels(
|
||||
gpu_dataset.outcome_prof_short_d.raw_ptr(), total_snaps,
|
||||
PERC_N_HORIZONS, seq_len, b_size, stg.aux_y_prof_short,
|
||||
)
|
||||
.context("step_with_lobsim_gpu: gather outcome_prof_short")?;
|
||||
// Aux size-long → stg_aux_y_size_long
|
||||
gpu_loader
|
||||
.gather_bce_labels(
|
||||
gpu_dataset.outcome_size_long_d.raw_ptr(), total_snaps,
|
||||
PERC_N_HORIZONS, seq_len, b_size, stg.aux_y_size_long,
|
||||
)
|
||||
.context("step_with_lobsim_gpu: gather outcome_size_long")?;
|
||||
// Aux size-short → stg_aux_y_size_short
|
||||
gpu_loader
|
||||
.gather_bce_labels(
|
||||
gpu_dataset.outcome_size_short_d.raw_ptr(), total_snaps,
|
||||
PERC_N_HORIZONS, seq_len, b_size, stg.aux_y_size_short,
|
||||
)
|
||||
.context("step_with_lobsim_gpu: gather outcome_size_short")?;
|
||||
|
||||
// pos_fraction: gathered into a small device scratch, then
|
||||
// read back on the host to compute clamped pos_weight. The
|
||||
// gather writes to a mapped-pinned buffer so we can read it
|
||||
@@ -7726,371 +7719,6 @@ impl IntegratedTrainer {
|
||||
}
|
||||
Ok(probs)
|
||||
}
|
||||
|
||||
// ── Checkpoint save / load ──────────────────────────────────────────
|
||||
|
||||
/// Persist all trainable weights, ISV state, and AdamW optimizer state
|
||||
/// to a binary checkpoint file.
|
||||
///
|
||||
/// Binary format:
|
||||
/// ```text
|
||||
/// [magic: u32 = 0x464F5843] // "FOXC"
|
||||
/// [version: u32 = 1]
|
||||
/// [step: u64]
|
||||
/// [n_sections: u32]
|
||||
/// For each section:
|
||||
/// [name_len: u16][name: UTF-8 bytes][data_len: u64][data: raw bytes]
|
||||
/// ```
|
||||
///
|
||||
/// The encoder (perception trunk) is saved to a sibling file with
|
||||
/// `.encoder` extension via `PerceptionTrainer::save_checkpoint`.
|
||||
///
|
||||
/// All device→host transfers use mapped-pinned staging + DtoD per
|
||||
/// `feedback_no_htod_htoh_only_mapped_pinned`.
|
||||
pub fn save_checkpoint(&self, path: &std::path::Path, step: u64) -> Result<()> {
|
||||
use std::io::Write;
|
||||
|
||||
const MAGIC: u32 = 0x464F5843; // "FOXC"
|
||||
const VERSION: u32 = 1;
|
||||
|
||||
// ── Encoder checkpoint (own format / sibling file) ───────────
|
||||
let encoder_path = path.with_extension("encoder");
|
||||
self.perception
|
||||
.save_checkpoint(&encoder_path)
|
||||
.context("save encoder checkpoint")?;
|
||||
|
||||
// ── Collect device-buffer sections ───────────────────────────
|
||||
// (name, CUdeviceptr, len_in_f32s)
|
||||
let device_sections: Vec<(&str, u64, usize)> = vec![
|
||||
("dqn_w", self.dqn_head.w_d.raw_ptr(), self.dqn_head.w_d.len()),
|
||||
("dqn_b", self.dqn_head.b_d.raw_ptr(), self.dqn_head.b_d.len()),
|
||||
("dqn_w_target", self.dqn_head.w_target_d.raw_ptr(), self.dqn_head.w_target_d.len()),
|
||||
("dqn_b_target", self.dqn_head.b_target_d.raw_ptr(), self.dqn_head.b_target_d.len()),
|
||||
("policy_w", self.policy_head.w_d.raw_ptr(), self.policy_head.w_d.len()),
|
||||
("policy_b", self.policy_head.b_d.raw_ptr(), self.policy_head.b_d.len()),
|
||||
("value_w", self.value_head.w_d.raw_ptr(), self.value_head.w_d.len()),
|
||||
("value_b", self.value_head.b_d.raw_ptr(), self.value_head.b_d.len()),
|
||||
("iqn_w_embed", self.iqn_head.w_embed_d.raw_ptr(), self.iqn_head.w_embed_d.len()),
|
||||
("iqn_b_embed", self.iqn_head.b_embed_d.raw_ptr(), self.iqn_head.b_embed_d.len()),
|
||||
("iqn_w_out", self.iqn_head.w_out_d.raw_ptr(), self.iqn_head.w_out_d.len()),
|
||||
("iqn_b_out", self.iqn_head.b_out_d.raw_ptr(), self.iqn_head.b_out_d.len()),
|
||||
("iqn_w_embed_target", self.iqn_head.w_embed_target_d.raw_ptr(), self.iqn_head.w_embed_target_d.len()),
|
||||
("iqn_b_embed_target", self.iqn_head.b_embed_target_d.raw_ptr(), self.iqn_head.b_embed_target_d.len()),
|
||||
("iqn_w_out_target", self.iqn_head.w_out_target_d.raw_ptr(), self.iqn_head.w_out_target_d.len()),
|
||||
("iqn_b_out_target", self.iqn_head.b_out_target_d.raw_ptr(), self.iqn_head.b_out_target_d.len()),
|
||||
("frd_w1", self.frd_head.w1_d.raw_ptr(), self.frd_head.w1_d.len()),
|
||||
("frd_b1", self.frd_head.b1_d.raw_ptr(), self.frd_head.b1_d.len()),
|
||||
("frd_w2", self.frd_head.w2_d.raw_ptr(), self.frd_head.w2_d.len()),
|
||||
("frd_b2", self.frd_head.b2_d.raw_ptr(), self.frd_head.b2_d.len()),
|
||||
("noisy_mu_w", self.noisy_exploration.mu_w.raw_ptr(), self.noisy_exploration.mu_w.len()),
|
||||
("noisy_sigma_w", self.noisy_exploration.sigma_w.raw_ptr(), self.noisy_exploration.sigma_w.len()),
|
||||
("noisy_mu_b", self.noisy_exploration.mu_b.raw_ptr(), self.noisy_exploration.mu_b.len()),
|
||||
("noisy_sigma_b", self.noisy_exploration.sigma_b.raw_ptr(), self.noisy_exploration.sigma_b.len()),
|
||||
("outcome_w", self.outcome_head.w_d.raw_ptr(), self.outcome_head.w_d.len()),
|
||||
("outcome_b", self.outcome_head.b_d.raw_ptr(), self.outcome_head.b_d.len()),
|
||||
];
|
||||
|
||||
// Find the largest device buffer to size the staging area.
|
||||
let max_len = device_sections
|
||||
.iter()
|
||||
.map(|(_, _, n)| *n)
|
||||
.max()
|
||||
.unwrap_or(0)
|
||||
.max(RL_SLOTS_END);
|
||||
let staging = unsafe { MappedF32Buffer::new(max_len) }
|
||||
.map_err(|e| anyhow::anyhow!("checkpoint staging alloc: {e}"))?;
|
||||
|
||||
// ISV (1) + device buffers (26) + adam (20) = 47 sections.
|
||||
let n_sections: u32 = 1 + device_sections.len() as u32 + 20;
|
||||
|
||||
let mut file = std::io::BufWriter::new(
|
||||
std::fs::File::create(path)
|
||||
.with_context(|| format!("create checkpoint {}", path.display()))?,
|
||||
);
|
||||
|
||||
// ── Header ──────────────────────────────────────────────────
|
||||
file.write_all(&MAGIC.to_le_bytes())?;
|
||||
file.write_all(&VERSION.to_le_bytes())?;
|
||||
file.write_all(&step.to_le_bytes())?;
|
||||
file.write_all(&n_sections.to_le_bytes())?;
|
||||
|
||||
// Helper: write a section header (name_len + name + data_len).
|
||||
let write_section_header =
|
||||
|w: &mut std::io::BufWriter<std::fs::File>, name: &str, data_len: u64| -> Result<()> {
|
||||
let name_bytes = name.as_bytes();
|
||||
w.write_all(&(name_bytes.len() as u16).to_le_bytes())?;
|
||||
w.write_all(name_bytes)?;
|
||||
w.write_all(&data_len.to_le_bytes())?;
|
||||
Ok(())
|
||||
};
|
||||
|
||||
// ── Section 1: ISV (already host-resident via mapped-pinned) ─
|
||||
{
|
||||
let isv_len = RL_SLOTS_END;
|
||||
let isv_bytes = (isv_len * std::mem::size_of::<f32>()) as u64;
|
||||
write_section_header(&mut file, "isv", isv_bytes)?;
|
||||
// Volatile reads for each element.
|
||||
for i in 0..isv_len {
|
||||
let val = unsafe { std::ptr::read_volatile(self.isv_mapped.host_ptr.add(i)) };
|
||||
file.write_all(&val.to_le_bytes())?;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Device-buffer sections (DtoD → staging → file) ──────────
|
||||
let raw_s = self.raw_stream;
|
||||
for (name, src_ptr, n_f32) in &device_sections {
|
||||
let byte_len = (*n_f32 * std::mem::size_of::<f32>()) as u64;
|
||||
write_section_header(&mut file, name, byte_len)?;
|
||||
if *n_f32 > 0 {
|
||||
unsafe {
|
||||
raw_memcpy_dtod_async(
|
||||
staging.dev_ptr,
|
||||
*src_ptr,
|
||||
*n_f32 * std::mem::size_of::<f32>(),
|
||||
raw_s,
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!("checkpoint save DtoD {name}: {e:?}"))?;
|
||||
raw_stream_sync(raw_s)
|
||||
.map_err(|e| anyhow::anyhow!("checkpoint save sync {name}: {e:?}"))?;
|
||||
}
|
||||
let host_data = unsafe {
|
||||
std::slice::from_raw_parts(staging.host_ptr, *n_f32)
|
||||
};
|
||||
file.write_all(bytemuck::cast_slice(host_data))?;
|
||||
}
|
||||
}
|
||||
|
||||
// ── AdamW sections ──────────────────────────────────────────
|
||||
let adam_pairs: Vec<(&str, &AdamW)> = vec![
|
||||
("adam_dqn_w", &self.dqn_w_adam),
|
||||
("adam_dqn_b", &self.dqn_b_adam),
|
||||
("adam_iqn_w_embed", &self.iqn_w_embed_adam),
|
||||
("adam_iqn_b_embed", &self.iqn_b_embed_adam),
|
||||
("adam_iqn_w_out", &self.iqn_w_out_adam),
|
||||
("adam_iqn_b_out", &self.iqn_b_out_adam),
|
||||
("adam_policy_w", &self.policy_w_adam),
|
||||
("adam_policy_b", &self.policy_b_adam),
|
||||
("adam_value_w", &self.value_w_adam),
|
||||
("adam_value_b", &self.value_b_adam),
|
||||
("adam_frd_w1", &self.frd_w1_adam),
|
||||
("adam_frd_b1", &self.frd_b1_adam),
|
||||
("adam_frd_w2", &self.frd_w2_adam),
|
||||
("adam_frd_b2", &self.frd_b2_adam),
|
||||
("adam_noisy_mu_w", &self.noisy_mu_w_adam),
|
||||
("adam_noisy_sigma_w", &self.noisy_sigma_w_adam),
|
||||
("adam_noisy_mu_b", &self.noisy_mu_b_adam),
|
||||
("adam_noisy_sigma_b", &self.noisy_sigma_b_adam),
|
||||
("adam_outcome_w", &self.outcome_w_adam),
|
||||
("adam_outcome_b", &self.outcome_b_adam),
|
||||
];
|
||||
for (name, adam) in &adam_pairs {
|
||||
// Serialize to a temp buffer to measure length, then write header + data.
|
||||
let mut buf = Vec::new();
|
||||
adam.save(&mut buf)
|
||||
.with_context(|| format!("AdamW save {name}"))?;
|
||||
write_section_header(&mut file, name, buf.len() as u64)?;
|
||||
file.write_all(&buf)?;
|
||||
}
|
||||
|
||||
file.flush()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Restore all trainable weights, ISV state, and AdamW optimizer state
|
||||
/// from a binary checkpoint file. Returns the step number embedded in
|
||||
/// the checkpoint header.
|
||||
///
|
||||
/// Unknown section names are silently skipped for forward compatibility.
|
||||
///
|
||||
/// All host→device transfers use mapped-pinned staging + DtoD per
|
||||
/// `feedback_no_htod_htoh_only_mapped_pinned`.
|
||||
pub fn load_checkpoint(&mut self, dev: &MlDevice, path: &std::path::Path) -> Result<u64> {
|
||||
use std::io::Read;
|
||||
|
||||
const MAGIC: u32 = 0x464F5843;
|
||||
const VERSION: u32 = 1;
|
||||
|
||||
// ── Encoder checkpoint (sibling file) ────────────────────────
|
||||
let encoder_path = path.with_extension("encoder");
|
||||
if encoder_path.exists() {
|
||||
let trunk_cfg = crate::cfc::trunk::CfcConfig {
|
||||
n_in: crate::cfc::snap_features::ENCODER_INPUT_DIM,
|
||||
n_hid: HIDDEN_DIM,
|
||||
cfc_n_in: HIDDEN_DIM,
|
||||
mamba2_state_dim: self.perception.config().mamba2_state_dim,
|
||||
n_batch: self.perception.config().n_batch,
|
||||
seq_len: self.perception.config().seq_len,
|
||||
};
|
||||
let new_trunk =
|
||||
crate::cfc::trunk::CfcTrunk::load_checkpoint(dev, &trunk_cfg, &encoder_path)
|
||||
.context("load encoder checkpoint")?;
|
||||
self.perception.trunk = new_trunk;
|
||||
}
|
||||
|
||||
let mut file = std::io::BufReader::new(
|
||||
std::fs::File::open(path)
|
||||
.with_context(|| format!("open checkpoint {}", path.display()))?,
|
||||
);
|
||||
|
||||
// ── Header ──────────────────────────────────────────────────
|
||||
let mut buf4 = [0u8; 4];
|
||||
let mut buf8 = [0u8; 8];
|
||||
let mut buf2 = [0u8; 2];
|
||||
|
||||
file.read_exact(&mut buf4)?;
|
||||
let magic = u32::from_le_bytes(buf4);
|
||||
anyhow::ensure!(magic == MAGIC, "bad checkpoint magic: {magic:#010x}");
|
||||
|
||||
file.read_exact(&mut buf4)?;
|
||||
let version = u32::from_le_bytes(buf4);
|
||||
anyhow::ensure!(version == VERSION, "unsupported checkpoint version: {version}");
|
||||
|
||||
file.read_exact(&mut buf8)?;
|
||||
let step = u64::from_le_bytes(buf8);
|
||||
|
||||
file.read_exact(&mut buf4)?;
|
||||
let n_sections = u32::from_le_bytes(buf4);
|
||||
|
||||
// Build a lookup from section name to (device ptr, f32 count) for
|
||||
// device-buffer restoration.
|
||||
let device_lookup: std::collections::HashMap<&str, (u64, usize)> = [
|
||||
("dqn_w", (self.dqn_head.w_d.raw_ptr(), self.dqn_head.w_d.len())),
|
||||
("dqn_b", (self.dqn_head.b_d.raw_ptr(), self.dqn_head.b_d.len())),
|
||||
("dqn_w_target", (self.dqn_head.w_target_d.raw_ptr(), self.dqn_head.w_target_d.len())),
|
||||
("dqn_b_target", (self.dqn_head.b_target_d.raw_ptr(), self.dqn_head.b_target_d.len())),
|
||||
("policy_w", (self.policy_head.w_d.raw_ptr(), self.policy_head.w_d.len())),
|
||||
("policy_b", (self.policy_head.b_d.raw_ptr(), self.policy_head.b_d.len())),
|
||||
("value_w", (self.value_head.w_d.raw_ptr(), self.value_head.w_d.len())),
|
||||
("value_b", (self.value_head.b_d.raw_ptr(), self.value_head.b_d.len())),
|
||||
("iqn_w_embed", (self.iqn_head.w_embed_d.raw_ptr(), self.iqn_head.w_embed_d.len())),
|
||||
("iqn_b_embed", (self.iqn_head.b_embed_d.raw_ptr(), self.iqn_head.b_embed_d.len())),
|
||||
("iqn_w_out", (self.iqn_head.w_out_d.raw_ptr(), self.iqn_head.w_out_d.len())),
|
||||
("iqn_b_out", (self.iqn_head.b_out_d.raw_ptr(), self.iqn_head.b_out_d.len())),
|
||||
("iqn_w_embed_target", (self.iqn_head.w_embed_target_d.raw_ptr(), self.iqn_head.w_embed_target_d.len())),
|
||||
("iqn_b_embed_target", (self.iqn_head.b_embed_target_d.raw_ptr(), self.iqn_head.b_embed_target_d.len())),
|
||||
("iqn_w_out_target", (self.iqn_head.w_out_target_d.raw_ptr(), self.iqn_head.w_out_target_d.len())),
|
||||
("iqn_b_out_target", (self.iqn_head.b_out_target_d.raw_ptr(), self.iqn_head.b_out_target_d.len())),
|
||||
("frd_w1", (self.frd_head.w1_d.raw_ptr(), self.frd_head.w1_d.len())),
|
||||
("frd_b1", (self.frd_head.b1_d.raw_ptr(), self.frd_head.b1_d.len())),
|
||||
("frd_w2", (self.frd_head.w2_d.raw_ptr(), self.frd_head.w2_d.len())),
|
||||
("frd_b2", (self.frd_head.b2_d.raw_ptr(), self.frd_head.b2_d.len())),
|
||||
("noisy_mu_w", (self.noisy_exploration.mu_w.raw_ptr(), self.noisy_exploration.mu_w.len())),
|
||||
("noisy_sigma_w", (self.noisy_exploration.sigma_w.raw_ptr(), self.noisy_exploration.sigma_w.len())),
|
||||
("noisy_mu_b", (self.noisy_exploration.mu_b.raw_ptr(), self.noisy_exploration.mu_b.len())),
|
||||
("noisy_sigma_b", (self.noisy_exploration.sigma_b.raw_ptr(), self.noisy_exploration.sigma_b.len())),
|
||||
("outcome_w", (self.outcome_head.w_d.raw_ptr(), self.outcome_head.w_d.len())),
|
||||
("outcome_b", (self.outcome_head.b_d.raw_ptr(), self.outcome_head.b_d.len())),
|
||||
]
|
||||
.into_iter()
|
||||
.collect();
|
||||
|
||||
// Staging buffer sized to the largest device section.
|
||||
let max_staging = device_lookup.values().map(|(_, n)| *n).max().unwrap_or(0).max(RL_SLOTS_END);
|
||||
let staging = unsafe { MappedF32Buffer::new(max_staging) }
|
||||
.map_err(|e| anyhow::anyhow!("checkpoint load staging alloc: {e}"))?;
|
||||
|
||||
let raw_s = self.raw_stream;
|
||||
|
||||
for _ in 0..n_sections {
|
||||
// Read section header.
|
||||
file.read_exact(&mut buf2)?;
|
||||
let name_len = u16::from_le_bytes(buf2) as usize;
|
||||
let mut name_buf = vec![0u8; name_len];
|
||||
file.read_exact(&mut name_buf)?;
|
||||
let name = String::from_utf8(name_buf)
|
||||
.context("checkpoint section name not valid UTF-8")?;
|
||||
file.read_exact(&mut buf8)?;
|
||||
let data_len = u64::from_le_bytes(buf8) as usize;
|
||||
|
||||
if name == "isv" {
|
||||
// ISV: write directly to isv_mapped host_ptr via
|
||||
// write_volatile, then DtoD host→device is unnecessary
|
||||
// because isv_mapped IS the device memory (mapped-pinned).
|
||||
let isv_n = RL_SLOTS_END;
|
||||
anyhow::ensure!(
|
||||
data_len == isv_n * std::mem::size_of::<f32>(),
|
||||
"ISV section size mismatch: file {data_len} vs expected {}",
|
||||
isv_n * std::mem::size_of::<f32>()
|
||||
);
|
||||
for i in 0..isv_n {
|
||||
let mut vbuf = [0u8; 4];
|
||||
file.read_exact(&mut vbuf)?;
|
||||
let val = f32::from_le_bytes(vbuf);
|
||||
unsafe {
|
||||
std::ptr::write_volatile(self.isv_mapped.host_ptr.add(i), val);
|
||||
}
|
||||
}
|
||||
} else if let Some(&(dst_ptr, expected_n)) = device_lookup.get(name.as_str()) {
|
||||
// Device buffer: read file → staging host_ptr → DtoD → device.
|
||||
let expected_bytes = expected_n * std::mem::size_of::<f32>();
|
||||
anyhow::ensure!(
|
||||
data_len == expected_bytes,
|
||||
"section '{name}' size mismatch: file {data_len} vs expected {expected_bytes}"
|
||||
);
|
||||
if expected_n > 0 {
|
||||
let host_slice = unsafe {
|
||||
std::slice::from_raw_parts_mut(staging.host_ptr, expected_n)
|
||||
};
|
||||
file.read_exact(bytemuck::cast_slice_mut(host_slice))?;
|
||||
unsafe {
|
||||
raw_memcpy_dtod_async(
|
||||
dst_ptr,
|
||||
staging.dev_ptr,
|
||||
expected_bytes,
|
||||
raw_s,
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!("checkpoint load DtoD {name}: {e:?}"))?;
|
||||
raw_stream_sync(raw_s)
|
||||
.map_err(|e| anyhow::anyhow!("checkpoint load sync {name}: {e:?}"))?;
|
||||
}
|
||||
}
|
||||
} else if name.starts_with("adam_") {
|
||||
// AdamW section: dispatch to the matching instance's load().
|
||||
let mut section_data = vec![0u8; data_len];
|
||||
file.read_exact(&mut section_data)?;
|
||||
let mut cursor = std::io::Cursor::new(§ion_data[..]);
|
||||
match name.as_str() {
|
||||
"adam_dqn_w" => self.dqn_w_adam.load(&mut cursor)?,
|
||||
"adam_dqn_b" => self.dqn_b_adam.load(&mut cursor)?,
|
||||
"adam_iqn_w_embed" => self.iqn_w_embed_adam.load(&mut cursor)?,
|
||||
"adam_iqn_b_embed" => self.iqn_b_embed_adam.load(&mut cursor)?,
|
||||
"adam_iqn_w_out" => self.iqn_w_out_adam.load(&mut cursor)?,
|
||||
"adam_iqn_b_out" => self.iqn_b_out_adam.load(&mut cursor)?,
|
||||
"adam_policy_w" => self.policy_w_adam.load(&mut cursor)?,
|
||||
"adam_policy_b" => self.policy_b_adam.load(&mut cursor)?,
|
||||
"adam_value_w" => self.value_w_adam.load(&mut cursor)?,
|
||||
"adam_value_b" => self.value_b_adam.load(&mut cursor)?,
|
||||
"adam_frd_w1" => self.frd_w1_adam.load(&mut cursor)?,
|
||||
"adam_frd_b1" => self.frd_b1_adam.load(&mut cursor)?,
|
||||
"adam_frd_w2" => self.frd_w2_adam.load(&mut cursor)?,
|
||||
"adam_frd_b2" => self.frd_b2_adam.load(&mut cursor)?,
|
||||
"adam_noisy_mu_w" => self.noisy_mu_w_adam.load(&mut cursor)?,
|
||||
"adam_noisy_sigma_w" => self.noisy_sigma_w_adam.load(&mut cursor)?,
|
||||
"adam_noisy_mu_b" => self.noisy_mu_b_adam.load(&mut cursor)?,
|
||||
"adam_noisy_sigma_b" => self.noisy_sigma_b_adam.load(&mut cursor)?,
|
||||
"adam_outcome_w" => self.outcome_w_adam.load(&mut cursor)?,
|
||||
"adam_outcome_b" => self.outcome_b_adam.load(&mut cursor)?,
|
||||
_ => { /* unknown adam section — skip for forward compat */ }
|
||||
}
|
||||
} else {
|
||||
// Unknown section — skip for forward compatibility.
|
||||
let mut discard = vec![0u8; data_len];
|
||||
file.read_exact(&mut discard)?;
|
||||
}
|
||||
}
|
||||
|
||||
// Invalidate captured CUDA graphs — weight addresses haven't moved
|
||||
// but values changed, so any graph that baked in the old data needs
|
||||
// re-capture on the next step.
|
||||
self.prefill_graph = None;
|
||||
self.postfill_graph = None;
|
||||
self.reward_graph = None;
|
||||
self.training_graph = None;
|
||||
self.graph_warmup_done = false;
|
||||
self.training_warmup_done = false;
|
||||
|
||||
Ok(step)
|
||||
}
|
||||
}
|
||||
|
||||
/// Free-function entry point for the `grad_h_accumulate_scaled` kernel.
|
||||
|
||||
@@ -10,7 +10,6 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use bytemuck;
|
||||
use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
|
||||
use cudarc::driver::sys::CUstream;
|
||||
use ml_core::device::MlDevice;
|
||||
@@ -127,98 +126,4 @@ impl AdamW {
|
||||
pub fn step_count(&self) -> i32 {
|
||||
self.step_count_host
|
||||
}
|
||||
|
||||
/// Serialize Adam state (m, v, step_count, hyperparams) to writer.
|
||||
///
|
||||
/// Wire format (little-endian):
|
||||
/// u64 n — number of parameters
|
||||
/// i32 step_count
|
||||
/// f32 lr, beta1, beta2, eps, wd
|
||||
/// [f32; n] m — first moment
|
||||
/// [f32; n] v — second moment
|
||||
pub fn save(&self, w: &mut impl std::io::Write) -> Result<()> {
|
||||
use crate::pinned_mem::MappedF32Buffer;
|
||||
use crate::trainer::raw_launch::raw_memcpy_dtod_async;
|
||||
let n = self.m.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("AdamW save staging: {e}"))?;
|
||||
let raw_s = self.raw_stream;
|
||||
unsafe {
|
||||
raw_memcpy_dtod_async(staging.dev_ptr, self.m.raw_ptr(), n * 4, raw_s)
|
||||
.map_err(|e| anyhow::anyhow!("AdamW save m dtod: {e:?}"))?;
|
||||
}
|
||||
self._stream.synchronize()
|
||||
.map_err(|e| anyhow::anyhow!("AdamW save sync m: {e}"))?;
|
||||
let m_host = staging.read_all();
|
||||
unsafe {
|
||||
raw_memcpy_dtod_async(staging.dev_ptr, self.v.raw_ptr(), n * 4, raw_s)
|
||||
.map_err(|e| anyhow::anyhow!("AdamW save v dtod: {e:?}"))?;
|
||||
}
|
||||
self._stream.synchronize()
|
||||
.map_err(|e| anyhow::anyhow!("AdamW save sync v: {e}"))?;
|
||||
let v_host = staging.read_all();
|
||||
w.write_all(&(n as u64).to_le_bytes())?;
|
||||
w.write_all(&self.step_count_host.to_le_bytes())?;
|
||||
w.write_all(&self.lr.to_le_bytes())?;
|
||||
w.write_all(&self.beta1.to_le_bytes())?;
|
||||
w.write_all(&self.beta2.to_le_bytes())?;
|
||||
w.write_all(&self.eps.to_le_bytes())?;
|
||||
w.write_all(&self.wd.to_le_bytes())?;
|
||||
w.write_all(bytemuck::cast_slice(&m_host))?;
|
||||
w.write_all(bytemuck::cast_slice(&v_host))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Restore Adam state from reader. Parameter count must match allocation.
|
||||
pub fn load(&mut self, r: &mut impl std::io::Read) -> Result<()> {
|
||||
let mut buf8 = [0u8; 8];
|
||||
let mut buf4 = [0u8; 4];
|
||||
r.read_exact(&mut buf8)?;
|
||||
let n = u64::from_le_bytes(buf8) as usize;
|
||||
anyhow::ensure!(
|
||||
n == self.m.len(),
|
||||
"AdamW load: m size mismatch {n} vs {}",
|
||||
self.m.len()
|
||||
);
|
||||
r.read_exact(&mut buf4)?;
|
||||
self.step_count_host = i32::from_le_bytes(buf4);
|
||||
r.read_exact(&mut buf4)?;
|
||||
self.lr = f32::from_le_bytes(buf4);
|
||||
r.read_exact(&mut buf4)?;
|
||||
self.beta1 = f32::from_le_bytes(buf4);
|
||||
r.read_exact(&mut buf4)?;
|
||||
self.beta2 = f32::from_le_bytes(buf4);
|
||||
r.read_exact(&mut buf4)?;
|
||||
self.eps = f32::from_le_bytes(buf4);
|
||||
r.read_exact(&mut buf4)?;
|
||||
self.wd = f32::from_le_bytes(buf4);
|
||||
use crate::pinned_mem::MappedF32Buffer;
|
||||
use crate::trainer::raw_launch::raw_memcpy_dtod_async;
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("AdamW load staging: {e}"))?;
|
||||
let raw_s = self.raw_stream;
|
||||
// m: read file → mapped-pinned host_ptr → DtoD → device
|
||||
{
|
||||
let host_slice = unsafe { std::slice::from_raw_parts_mut(staging.host_ptr, n) };
|
||||
r.read_exact(bytemuck::cast_slice_mut(host_slice))?;
|
||||
unsafe {
|
||||
raw_memcpy_dtod_async(self.m.raw_ptr(), staging.dev_ptr, n * 4, raw_s)
|
||||
.map_err(|e| anyhow::anyhow!("AdamW load m dtod: {e:?}"))?;
|
||||
}
|
||||
self._stream.synchronize()
|
||||
.map_err(|e| anyhow::anyhow!("AdamW load sync m: {e}"))?;
|
||||
}
|
||||
// v: same pattern
|
||||
{
|
||||
let host_slice = unsafe { std::slice::from_raw_parts_mut(staging.host_ptr, n) };
|
||||
r.read_exact(bytemuck::cast_slice_mut(host_slice))?;
|
||||
unsafe {
|
||||
raw_memcpy_dtod_async(self.v.raw_ptr(), staging.dev_ptr, n * 4, raw_s)
|
||||
.map_err(|e| anyhow::anyhow!("AdamW load v dtod: {e:?}"))?;
|
||||
}
|
||||
self._stream.synchronize()
|
||||
.map_err(|e| anyhow::anyhow!("AdamW load sync v: {e}"))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -11,9 +11,21 @@ fn main() -> Result<(), String> {
|
||||
|
||||
// Per pearl_build_rs_rerun_if_env_changed.md: pair every env::var
|
||||
// with rerun-if-env-changed.
|
||||
// Default sm_86 covers RTX 3050 (local dev). Production Argo workflows
|
||||
// override via FOXHUNT_CUDA_ARCH=sm_89 (L40S Ada) or sm_90 (H100).
|
||||
let arch = std::env::var("FOXHUNT_CUDA_ARCH").unwrap_or_else(|_| "sm_86".into());
|
||||
//
|
||||
// Arch detection (mirrors crates/ml-alpha/build.rs:detect_arch — the
|
||||
// original "mirror" docstring above was aspirational, B-9 cluster
|
||||
// alpha-rl-mbg2n on H100 caught the gap: ml-alpha picked sm_90 via
|
||||
// CUDA_COMPUTE_CAP but ml-backtesting fell through to default sm_86
|
||||
// → no kernel image error at LobSimCuda::new). Priority:
|
||||
// 1. CUDA_COMPUTE_CAP (numeric, e.g. "90") — set by alpha-rl-template
|
||||
// from nvidia-smi inside the compile pod
|
||||
// 2. FOXHUNT_CUDA_ARCH (sm_-prefixed, e.g. "sm_90") — set by
|
||||
// lob-backtest-sweep-template
|
||||
// 3. nvidia-smi --query-gpu=compute_cap at build time
|
||||
// 4. Default sm_86 (RTX 3050 Ti local dev)
|
||||
let arch = detect_arch();
|
||||
eprintln!(" ml-backtesting: compiling kernels for {arch}");
|
||||
println!("cargo:rerun-if-env-changed=CUDA_COMPUTE_CAP");
|
||||
println!("cargo:rerun-if-env-changed=FOXHUNT_CUDA_ARCH");
|
||||
println!("cargo:rerun-if-env-changed=CUDA_PATH");
|
||||
println!("cargo:rerun-if-env-changed=NVCC");
|
||||
@@ -60,3 +72,38 @@ fn main() -> Result<(), String> {
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Returns the `sm_<NN>` arch string, honoring (in order):
|
||||
/// `CUDA_COMPUTE_CAP` numeric env → `FOXHUNT_CUDA_ARCH` sm_-prefixed env
|
||||
/// → `nvidia-smi --query-gpu=compute_cap` → default `sm_86`.
|
||||
fn detect_arch() -> String {
|
||||
// 1. Numeric CUDA_COMPUTE_CAP (e.g. "90") from alpha-rl-template.
|
||||
if let Ok(cap) = std::env::var("CUDA_COMPUTE_CAP") {
|
||||
let trimmed = cap.trim();
|
||||
if !trimmed.is_empty() {
|
||||
return format!("sm_{trimmed}");
|
||||
}
|
||||
}
|
||||
// 2. sm_-prefixed FOXHUNT_CUDA_ARCH (e.g. "sm_90") from lob-backtest-sweep.
|
||||
if let Ok(arch) = std::env::var("FOXHUNT_CUDA_ARCH") {
|
||||
let trimmed = arch.trim();
|
||||
if !trimmed.is_empty() {
|
||||
return trimmed.to_string();
|
||||
}
|
||||
}
|
||||
// 3. Query the GPU on this machine (e.g. "9.0" → "90" → "sm_90").
|
||||
if let Ok(output) = Command::new("nvidia-smi")
|
||||
.args(["--query-gpu=compute_cap", "--format=csv,noheader"])
|
||||
.output()
|
||||
{
|
||||
if output.status.success() {
|
||||
let s = String::from_utf8_lossy(&output.stdout);
|
||||
let cap = s.trim().replace('.', "");
|
||||
if !cap.is_empty() {
|
||||
return format!("sm_{cap}");
|
||||
}
|
||||
}
|
||||
}
|
||||
// 4. Default — RTX 3050 Ti local dev.
|
||||
"sm_86".to_string()
|
||||
}
|
||||
|
||||
@@ -16,8 +16,13 @@ pub const STOP_SLOT_BYTES: usize = 32;
|
||||
/// Bytes per Orders struct (limits[32] + stops[16]).
|
||||
pub const ORDERS_BYTES: usize = MAX_LIMITS * LIMIT_SLOT_BYTES + MAX_STOPS * STOP_SLOT_BYTES;
|
||||
/// Max closed-trade records buffered per backtest before the host
|
||||
/// must drain. Sized for a single fixture day at ~30s decision cadence.
|
||||
pub const TRADE_LOG_CAP: usize = 1024;
|
||||
/// must drain. Bumped from 1024 → 4096 on 2026-05-31 to prevent
|
||||
/// eval-phase ring wrap at cluster scale (b=1024 × 5000 eval steps
|
||||
/// produces ~342 dones/account mean, ~1000 peak; 4096 gives 4×
|
||||
/// headroom). Memory cost at b=1024: 1024 × 4096 × 40 B = 167 MB
|
||||
/// device (was 41 MB) — comfortable on L40S 48 GB / H100 80 GB.
|
||||
/// See spec `docs/superpowers/specs/2026-05-31-eval-summary-trade-aggregation-design.md`.
|
||||
pub const TRADE_LOG_CAP: usize = 4096;
|
||||
/// Bytes per TradeRecord (must match crates/ml-backtesting/src/order.rs).
|
||||
pub const TRADE_RECORD_BYTES: usize = 40;
|
||||
/// Bytes per OpenTradeState (kernel-side tracking of currently-open entry).
|
||||
|
||||
@@ -1143,6 +1143,7 @@ impl LobSimCuda {
|
||||
.arg(&n)
|
||||
.launch(cfg)?;
|
||||
}
|
||||
self.stream.synchronize()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -1514,6 +1515,132 @@ impl LobSimCuda {
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Read per-backtest cumulative trade counters. Returns a `Vec<u32>`
|
||||
/// of length `n_backtests` where entry `i` is the total number of
|
||||
/// closed-trade records ever pushed to backtest `i`'s ring buffer
|
||||
/// (may exceed `TRADE_LOG_CAP` if the ring wrapped).
|
||||
///
|
||||
/// Used by `alpha_rl_train` to snapshot the eval-phase boundary per
|
||||
/// account: `head_before_per_b[i] = read_per_backtest_trade_counts()[i]`
|
||||
/// at eval start; post-eval `head_after_per_b[i] - head_before_per_b[i]`
|
||||
/// gives the true eval-phase trade count for backtest `i`.
|
||||
///
|
||||
/// Uses mapped-pinned staging per `feedback_no_htod_htoh_only_mapped_pinned`:
|
||||
/// allocate a `MappedRecordBuffer<u32>`, DtoD copy from
|
||||
/// `trade_log_head_d` into its dev_ptr, sync, read via host_ptr.
|
||||
///
|
||||
/// Replaces the (`head_before_eval = read_total_trade_count()`,
|
||||
/// `all_records = read_trade_records(0)`) pattern that conflated
|
||||
/// aggregate counts with per-account ring contents and produced
|
||||
/// meaningless eval_summary metrics — see spec
|
||||
/// `docs/superpowers/specs/2026-05-31-eval-summary-trade-aggregation-design.md`.
|
||||
pub fn read_per_backtest_trade_counts(&self) -> Result<Vec<u32>> {
|
||||
use ml_alpha::pinned_mem::MappedRecordBuffer;
|
||||
use ml_alpha::trainer::raw_launch::raw_memcpy_dtod_async;
|
||||
let raw_stream = self.stream.cu_stream();
|
||||
let staging: MappedRecordBuffer<u32> = unsafe { MappedRecordBuffer::new(self.n_backtests) }
|
||||
.map_err(|e| anyhow::anyhow!("alloc trade_head staging: {e}"))?;
|
||||
unsafe {
|
||||
raw_memcpy_dtod_async(
|
||||
staging.dev_ptr,
|
||||
self.trade_log_head_d.raw_ptr(),
|
||||
self.n_backtests * std::mem::size_of::<u32>(),
|
||||
raw_stream,
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!("DtoD trade_log_head: {e:?}"))?;
|
||||
}
|
||||
self.stream
|
||||
.synchronize()
|
||||
.context("sync after trade_log_head DtoD")?;
|
||||
let mut out = Vec::with_capacity(self.n_backtests);
|
||||
for i in 0..self.n_backtests {
|
||||
// SAFETY: mapped-pinned host_ptr is valid for [0, len); sync
|
||||
// above ensures the DtoD writes are visible to host.
|
||||
unsafe {
|
||||
out.push(std::ptr::read_volatile(staging.host_ptr.add(i)));
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Read all backtests' trade record rings. Returns
|
||||
/// `Vec<Vec<TradeRecord>>` of length `n_backtests`; entry `i` is the
|
||||
/// most recent up-to-`TRADE_LOG_CAP` records for backtest `i` (oldest
|
||||
/// dropped on wrap).
|
||||
///
|
||||
/// Uses mapped-pinned staging for both the per-backtest head counters
|
||||
/// and the trade-log payload (per `feedback_no_htod_htoh_only_mapped_pinned`).
|
||||
/// At cluster scale (b=1024, cap=4096) the payload staging buffer
|
||||
/// is 167 MB — allocated once per call (end-of-run). Acceptable
|
||||
/// because this is invoked once per fold-eval, not in the hot path.
|
||||
///
|
||||
/// Callers compose this with
|
||||
/// [`read_per_backtest_trade_counts`](Self::read_per_backtest_trade_counts)
|
||||
/// to identify each account's eval-phase-only slice.
|
||||
pub fn read_trade_records_all(
|
||||
&self,
|
||||
) -> Result<Vec<Vec<crate::order::TradeRecord>>> {
|
||||
use ml_alpha::pinned_mem::MappedRecordBuffer;
|
||||
use ml_alpha::trainer::raw_launch::raw_memcpy_dtod_async;
|
||||
let raw_stream = self.stream.cu_stream();
|
||||
let rec_bytes = crate::lob::TRADE_RECORD_BYTES;
|
||||
let cap = crate::lob::TRADE_LOG_CAP;
|
||||
let payload_len_u8 = self.n_backtests * cap * rec_bytes;
|
||||
|
||||
let head_staging: MappedRecordBuffer<u32> =
|
||||
unsafe { MappedRecordBuffer::new(self.n_backtests) }
|
||||
.map_err(|e| anyhow::anyhow!("alloc head staging: {e}"))?;
|
||||
let payload_staging: MappedRecordBuffer<u8> =
|
||||
unsafe { MappedRecordBuffer::new(payload_len_u8) }
|
||||
.map_err(|e| anyhow::anyhow!("alloc payload staging ({payload_len_u8} B): {e}"))?;
|
||||
|
||||
unsafe {
|
||||
raw_memcpy_dtod_async(
|
||||
head_staging.dev_ptr,
|
||||
self.trade_log_head_d.raw_ptr(),
|
||||
self.n_backtests * std::mem::size_of::<u32>(),
|
||||
raw_stream,
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!("DtoD trade_log_head: {e:?}"))?;
|
||||
raw_memcpy_dtod_async(
|
||||
payload_staging.dev_ptr,
|
||||
self.trade_log_d.raw_ptr(),
|
||||
payload_len_u8,
|
||||
raw_stream,
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!("DtoD trade_log payload: {e:?}"))?;
|
||||
}
|
||||
self.stream
|
||||
.synchronize()
|
||||
.context("sync after trade_log DtoD")?;
|
||||
|
||||
let mut out = Vec::with_capacity(self.n_backtests);
|
||||
for b in 0..self.n_backtests {
|
||||
// SAFETY: mapped-pinned host_ptr is valid + post-sync the
|
||||
// device-side DtoD writes are visible to the host.
|
||||
let head = unsafe { std::ptr::read_volatile(head_staging.host_ptr.add(b)) } as usize;
|
||||
let n_to_read = head.min(cap);
|
||||
let off = b * cap * rec_bytes;
|
||||
let mut records = Vec::with_capacity(n_to_read);
|
||||
for i in 0..n_to_read {
|
||||
let rec_off = off + i * rec_bytes;
|
||||
// SAFETY: bounds-checked above (i < n_to_read ≤ cap; rec_off + rec_bytes ≤ payload_len_u8).
|
||||
let bytes: [u8; crate::lob::TRADE_RECORD_BYTES] = unsafe {
|
||||
let mut buf = [0u8; crate::lob::TRADE_RECORD_BYTES];
|
||||
for k in 0..rec_bytes {
|
||||
buf[k] = std::ptr::read_volatile(payload_staging.host_ptr.add(rec_off + k));
|
||||
}
|
||||
buf
|
||||
};
|
||||
let r: crate::order::TradeRecord =
|
||||
bytemuck::pod_read_unaligned(&bytes);
|
||||
records.push(r);
|
||||
}
|
||||
out.push(records);
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Read back the Pos state for a specific backtest.
|
||||
pub fn read_pos(&self, backtest_idx: usize) -> Result<PosFlat> {
|
||||
anyhow::ensure!(
|
||||
|
||||
@@ -1,333 +0,0 @@
|
||||
# Alpha-RL Performance + Checkpoint + Walk-Forward 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:** 2-3× training throughput via TF32, crash recovery via checkpointing, OOS validation via 3-fold walk-forward on H100.
|
||||
|
||||
**Architecture:** Enable TF32 Tensor Core math on all cuBLAS handles (2 sites). Add checkpoint/resume serialization for model weights + Adam state + ISV to the training binary. Validate with 3-fold walk-forward at 25k steps/fold on H100.
|
||||
|
||||
**Tech Stack:** Rust, CUDA cuBLAS, cudarc, mapped-pinned memory, Argo Workflows
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Enable TF32 on DQN cuBLAS handle
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml-alpha/src/rl/dqn.rs:314` (after `cublasSetWorkspace_v2`)
|
||||
|
||||
- [ ] **Step 1: Add TF32 math mode after workspace setup**
|
||||
|
||||
In `crates/ml-alpha/src/rl/dqn.rs`, after line 313 (the closing of the `cublasSetWorkspace_v2` block), add inside the same `unsafe` block:
|
||||
|
||||
```rust
|
||||
cudarc::cublas::sys::cublasSetMathMode(
|
||||
*cublas.handle(),
|
||||
cudarc::cublas::sys::cublasMath_t::CUBLAS_TF32_TENSOR_OP_MATH,
|
||||
)
|
||||
.result()
|
||||
.map_err(|e| anyhow::anyhow!("DqnHead: cublasSetMathMode TF32: {e:?}"))?;
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Build check**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml-alpha`
|
||||
Expected: compiles cleanly
|
||||
|
||||
- [ ] **Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml-alpha/src/rl/dqn.rs
|
||||
git commit -m "perf(cuda): enable TF32 Tensor Core math on DQN cuBLAS handle"
|
||||
```
|
||||
|
||||
### Task 2: Enable TF32 on Mamba2 cuBLAS handle
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml-alpha/src/mamba2_block.rs:549` (after `cublasSetWorkspace_v2`)
|
||||
|
||||
- [ ] **Step 1: Add TF32 math mode after workspace setup**
|
||||
|
||||
In `crates/ml-alpha/src/mamba2_block.rs`, after line 549 (the closing of the `cublasSetWorkspace_v2` block), add inside the same `unsafe` block:
|
||||
|
||||
```rust
|
||||
cudarc::cublas::sys::cublasSetMathMode(
|
||||
*cublas.handle(),
|
||||
cudarc::cublas::sys::cublasMath_t::CUBLAS_TF32_TENSOR_OP_MATH,
|
||||
)
|
||||
.result()
|
||||
.map_err(|e| anyhow!("Mamba2Block: cublasSetMathMode TF32: {e:?}"))?;
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Build check**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml-alpha`
|
||||
Expected: compiles cleanly
|
||||
|
||||
- [ ] **Step 3: Local smoke test (RTX 3050 Ti, sm_86 supports TF32)**
|
||||
|
||||
Run a 100-step local smoke and verify l_q is within 1% of baseline:
|
||||
```bash
|
||||
FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml-alpha --lib -- smoke_tests --ignored --nocapture 2>&1 | tail -5
|
||||
```
|
||||
Expected: smoke passes, no NaN
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml-alpha/src/mamba2_block.rs
|
||||
git commit -m "perf(cuda): enable TF32 Tensor Core math on Mamba2 cuBLAS handle"
|
||||
```
|
||||
|
||||
### Task 3: Add AdamW save/load methods
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml-alpha/src/trainer/optim.rs`
|
||||
|
||||
- [ ] **Step 1: Add save method to AdamW**
|
||||
|
||||
Add after the existing `impl AdamW` block (or at the end of the impl):
|
||||
|
||||
```rust
|
||||
/// Serialize Adam state (m, v, step_count, hyperparams) to writer.
|
||||
/// Reads device buffers via mapped-pinned DtoH (one-shot, not hot path).
|
||||
pub fn save(&self, w: &mut impl std::io::Write) -> Result<()> {
|
||||
let m_host = self._stream.clone().read_sync(&self.m)
|
||||
.map_err(|e| anyhow::anyhow!("AdamW save m: {e}"))?;
|
||||
let v_host = self._stream.clone().read_sync(&self.v)
|
||||
.map_err(|e| anyhow::anyhow!("AdamW save v: {e}"))?;
|
||||
let n = m_host.len() as u64;
|
||||
w.write_all(&n.to_le_bytes())?;
|
||||
w.write_all(&self.step_count_host.to_le_bytes())?;
|
||||
w.write_all(&self.lr.to_le_bytes())?;
|
||||
w.write_all(&self.beta1.to_le_bytes())?;
|
||||
w.write_all(&self.beta2.to_le_bytes())?;
|
||||
w.write_all(&self.eps.to_le_bytes())?;
|
||||
w.write_all(&self.wd.to_le_bytes())?;
|
||||
let m_bytes: &[u8] = bytemuck::cast_slice(&m_host);
|
||||
w.write_all(m_bytes)?;
|
||||
let v_bytes: &[u8] = bytemuck::cast_slice(&v_host);
|
||||
w.write_all(v_bytes)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Restore Adam state from reader. Sizes must match.
|
||||
pub fn load(&mut self, r: &mut impl std::io::Read) -> Result<()> {
|
||||
let mut buf8 = [0u8; 8];
|
||||
let mut buf4 = [0u8; 4];
|
||||
r.read_exact(&mut buf8)?;
|
||||
let n = u64::from_le_bytes(buf8) as usize;
|
||||
anyhow::ensure!(n == self.m.len(), "AdamW load: m size mismatch {n} vs {}", self.m.len());
|
||||
r.read_exact(&mut buf4)?;
|
||||
self.step_count_host = i32::from_le_bytes(buf4);
|
||||
r.read_exact(&mut buf4)?; self.lr = f32::from_le_bytes(buf4);
|
||||
r.read_exact(&mut buf4)?; self.beta1 = f32::from_le_bytes(buf4);
|
||||
r.read_exact(&mut buf4)?; self.beta2 = f32::from_le_bytes(buf4);
|
||||
r.read_exact(&mut buf4)?; self.eps = f32::from_le_bytes(buf4);
|
||||
r.read_exact(&mut buf4)?; self.wd = f32::from_le_bytes(buf4);
|
||||
let mut m_host = vec![0f32; n];
|
||||
r.read_exact(bytemuck::cast_slice_mut(&mut m_host))?;
|
||||
self._stream.clone().write_sync(&m_host, &mut self.m)
|
||||
.map_err(|e| anyhow::anyhow!("AdamW load m: {e}"))?;
|
||||
let mut v_host = vec![0f32; n];
|
||||
r.read_exact(bytemuck::cast_slice_mut(&mut v_host))?;
|
||||
self._stream.clone().write_sync(&v_host, &mut self.v)
|
||||
.map_err(|e| anyhow::anyhow!("AdamW load v: {e}"))?;
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Verify bytemuck is in dependencies**
|
||||
|
||||
Run: `grep bytemuck crates/ml-alpha/Cargo.toml`
|
||||
If missing, add `bytemuck = { version = "1", features = ["derive"] }`.
|
||||
|
||||
- [ ] **Step 3: Build check**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml-alpha`
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml-alpha/src/trainer/optim.rs crates/ml-alpha/Cargo.toml
|
||||
git commit -m "feat(rl): AdamW save/load for checkpoint persistence"
|
||||
```
|
||||
|
||||
### Task 4: Add IntegratedTrainer checkpoint save/load
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml-alpha/src/trainer/integrated.rs`
|
||||
|
||||
- [ ] **Step 1: Add save_checkpoint method**
|
||||
|
||||
Add a `save_checkpoint` method to `IntegratedTrainer` that serializes:
|
||||
- DQN weights (online + target): `dqn_head.w_d`, `dqn_head.b_d`, `dqn_head.w_target_d`, `dqn_head.b_target_d`
|
||||
- Policy/V head weights
|
||||
- ISV bus (585 f32 from mapped-pinned `isv_mapped.host_ptr`)
|
||||
- All 20 AdamW states via `adam.save()`
|
||||
- Step counter
|
||||
- Encoder via `perception.save_checkpoint()`
|
||||
|
||||
Format: sequential sections with `[name_len: u16][name: bytes][data_len: u64][data: bytes]`.
|
||||
File header: `[magic: u32 = 0x464F5843][version: u32 = 1][step: u64]`.
|
||||
|
||||
Write to `<out_dir>/checkpoint-<step>.bin`. Keep last 2, delete older.
|
||||
|
||||
- [ ] **Step 2: Add load_checkpoint method**
|
||||
|
||||
Inverse of save: read header, validate magic/version, restore each section by name lookup. Skip unknown sections for forward compatibility.
|
||||
|
||||
- [ ] **Step 3: Build check**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml-alpha`
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml-alpha/src/trainer/integrated.rs
|
||||
git commit -m "feat(rl): IntegratedTrainer checkpoint save/load"
|
||||
```
|
||||
|
||||
### Task 5: Wire checkpoint into training loop
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml-alpha/examples/alpha_rl_train.rs`
|
||||
|
||||
- [ ] **Step 1: Add CLI flags**
|
||||
|
||||
Add to the CLI struct (clap):
|
||||
```rust
|
||||
/// Resume from checkpoint file path. Restores model weights, Adam
|
||||
/// state, ISV, and step counter. Training continues from saved step.
|
||||
#[arg(long)]
|
||||
resume_from: Option<PathBuf>,
|
||||
|
||||
/// Checkpoint interval in steps (default: 5000). Set to 0 to disable.
|
||||
#[arg(long, default_value = "5000")]
|
||||
checkpoint_every: usize,
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Add resume logic before training loop**
|
||||
|
||||
After trainer init, before `for step in 0..cli.n_steps`:
|
||||
```rust
|
||||
let start_step = if let Some(ref ckpt_path) = cli.resume_from {
|
||||
trainer.load_checkpoint(ckpt_path)
|
||||
.with_context(|| format!("resume from {}", ckpt_path.display()))?
|
||||
} else {
|
||||
0
|
||||
};
|
||||
```
|
||||
|
||||
Change loop to `for step in start_step..cli.n_steps`.
|
||||
|
||||
- [ ] **Step 3: Add checkpoint save inside training loop**
|
||||
|
||||
After the per-step JSONL write, before the next iteration:
|
||||
```rust
|
||||
if cli.checkpoint_every > 0 && step > 0 && step % cli.checkpoint_every == 0 {
|
||||
let ckpt_path = cli.out.join(format!("checkpoint-{step}.bin"));
|
||||
trainer.save_checkpoint(&ckpt_path, step as u64)
|
||||
.with_context(|| format!("checkpoint at step {step}"))?;
|
||||
eprintln!("checkpoint saved: {}", ckpt_path.display());
|
||||
// Rolling: keep last 2, delete older
|
||||
let old_step = step.saturating_sub(cli.checkpoint_every * 2);
|
||||
if old_step > 0 {
|
||||
let old_path = cli.out.join(format!("checkpoint-{old_step}.bin"));
|
||||
let _ = std::fs::remove_file(&old_path);
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Build check**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml-alpha`
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml-alpha/examples/alpha_rl_train.rs
|
||||
git commit -m "feat(rl): wire checkpoint save/resume into training loop"
|
||||
```
|
||||
|
||||
### Task 6: Local validation smoke
|
||||
|
||||
**Files:** None (test only)
|
||||
|
||||
- [ ] **Step 1: Run 200-step local smoke with TF32**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo build --release -p ml-alpha --example alpha_rl_train
|
||||
# Run 200 steps with checkpoint at 100
|
||||
./target/release/examples/alpha_rl_train \
|
||||
--n-backtests 16 --n-steps 200 --checkpoint-every 100 \
|
||||
--out /tmp/smoke-tf32 \
|
||||
--mbp10-data-dir test_data/futures-baseline \
|
||||
--trades-data-dir test_data/futures-baseline 2>&1 | tail -5
|
||||
```
|
||||
|
||||
Expected: completes without NaN, checkpoint files at `/tmp/smoke-tf32/checkpoint-100.bin`
|
||||
|
||||
- [ ] **Step 2: Test resume**
|
||||
|
||||
```bash
|
||||
./target/release/examples/alpha_rl_train \
|
||||
--n-backtests 16 --n-steps 200 --checkpoint-every 100 \
|
||||
--resume-from /tmp/smoke-tf32/checkpoint-100.bin \
|
||||
--out /tmp/smoke-tf32-resume \
|
||||
--mbp10-data-dir test_data/futures-baseline \
|
||||
--trades-data-dir test_data/futures-baseline 2>&1 | tail -5
|
||||
```
|
||||
|
||||
Expected: starts from step 100, runs to 200, no crash
|
||||
|
||||
- [ ] **Step 3: Commit all together**
|
||||
|
||||
```bash
|
||||
git add -A
|
||||
git commit -m "test: validate TF32 + checkpoint save/resume local smoke"
|
||||
```
|
||||
|
||||
### Task 7: Deploy walk-forward on H100
|
||||
|
||||
**Files:** None (deployment only)
|
||||
|
||||
- [ ] **Step 1: Push branch**
|
||||
|
||||
```bash
|
||||
git push origin ml-alpha-phase-a
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Submit 3-fold walk-forward on H100**
|
||||
|
||||
```bash
|
||||
./scripts/argo-train.sh \
|
||||
--model alpha-rl \
|
||||
--branch ml-alpha-phase-a \
|
||||
--gpu-pool ci-training-h100 \
|
||||
--folds 3
|
||||
```
|
||||
|
||||
If `argo-train.sh` doesn't support `--folds` with step override, submit directly:
|
||||
```bash
|
||||
argo submit --from=wftmpl/alpha-rl -n foxhunt \
|
||||
-p git-branch=ml-alpha-phase-a \
|
||||
-p n-steps=25000 \
|
||||
-p n-backtests=1024 \
|
||||
-p per-capacity=65536 \
|
||||
-p gpu-pool=ci-training-h100
|
||||
```
|
||||
Run 3 times with fold-idx 0, 1, 2.
|
||||
|
||||
- [ ] **Step 3: Monitor with log-only (no extra pods on GPU node)**
|
||||
|
||||
```bash
|
||||
kubectl logs <pod-name> -n foxhunt --tail=1 -f | grep "step.*wr="
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Validate success criteria**
|
||||
|
||||
All 3 folds must show:
|
||||
- wr > 0.55
|
||||
- No fold with wr < 0.50
|
||||
- Entropy stable (no collapse)
|
||||
- hold% between 30-70%
|
||||
@@ -1,152 +0,0 @@
|
||||
# Alpha-RL Performance + Checkpoint + Walk-Forward
|
||||
|
||||
**Date**: 2026-05-27
|
||||
**Status**: Approved
|
||||
**Scope**: TF32 matmuls, checkpoint/resume, controller fusion, walk-forward validation
|
||||
|
||||
## Context
|
||||
|
||||
The alpha-rl pipeline achieves wr=0.567 at b=1024 but takes 4.5h for 100k steps on L40S at 6.2 sps — and the model converges by 5-10k steps. The last 90k steps are wasted. The L40S node OOM'd at 58k from monitoring pod overhead, losing all state (zero checkpoint infrastructure). All cuBLAS SGEMMs force `CUBLAS_COMPUTE_32F`, bypassing Tensor Cores entirely on both L40S and H100.
|
||||
|
||||
## Goals
|
||||
|
||||
1. **2-3× throughput** via TF32 Tensor Core acceleration on cuBLAS SGEMMs
|
||||
2. **Crash recovery** via checkpoint/resume every 5k steps
|
||||
3. **~30% launch overhead reduction** via controller kernel fusion
|
||||
4. **OOS validation** via 3-fold walk-forward at 25k steps/fold on H100
|
||||
|
||||
Target wall clock: 3 folds × 25k steps at ~15 sps on H100 = ~83 min total (vs 4.5h+ single-fold L40S that never finished).
|
||||
|
||||
## P0: TF32 Matmuls
|
||||
|
||||
### What
|
||||
|
||||
Enable TF32 Tensor Core math on all cuBLAS handles. TF32 uses 19-bit mantissa (vs FP32's 23-bit) with identical range. Precision loss is well within RL gradient noise. PyTorch default since 1.7.
|
||||
|
||||
### Where
|
||||
|
||||
Two cuBLAS handle creation sites:
|
||||
|
||||
1. `crates/ml-alpha/src/rl/dqn.rs` — after `cublasCreate_v2()`, add `cublasSetMathMode(handle, CUBLAS_TF32_TENSOR_OP_MATH)`. Affects DQN forward (online + target) and backward (grad_h + grad_w) = 4 SGEMMs/step.
|
||||
|
||||
2. `crates/ml-alpha/src/mamba2_block.rs` — same pattern after handle creation. Affects W_in, W_a, W_b, W_out projection GEMMs = 4+ SGEMMs/step.
|
||||
|
||||
### Impact
|
||||
|
||||
8+ SGEMMs per step go from FP32 scalar to TF32 Tensor Core. On L40S (Ada Lovelace): ~2× throughput. On H100 (Hopper): ~3× throughput. Combined with H100's 2× memory bandwidth: expect 12-20 sps at b=1024 (vs 6.2 current).
|
||||
|
||||
### Verification
|
||||
|
||||
Run 1k-step local smoke on RTX 3050 Ti (sm_86, supports TF32). Compare l_q at step 1000 with and without TF32 — should be within 1%.
|
||||
|
||||
## P1: Checkpoint/Resume
|
||||
|
||||
### What
|
||||
|
||||
Serialize training state every 5k steps to PVC. Resume from last checkpoint on restart.
|
||||
|
||||
### State to persist
|
||||
|
||||
| Component | Size | Source |
|
||||
|-----------|------|--------|
|
||||
| DQN weights (online) | W[128×231] + b[231] = ~120KB | `dqn_head.w_d`, `dqn_head.b_d` |
|
||||
| DQN weights (target) | Same = ~120KB | `dqn_head.w_target_d`, `dqn_head.b_target_d` |
|
||||
| Policy head weights | ~120KB | `policy_head.w_d`, `policy_head.b_d` |
|
||||
| V head weights | ~2KB | `v_head.w_d`, `v_head.b_d` |
|
||||
| Encoder weights | ~2-10MB (Mamba2 + CfC) | `encoder.params_d()` |
|
||||
| Adam state (m, v per param) | 2× model size = ~20MB | Per-group m_d, v_d |
|
||||
| ISV bus | 585 × f32 = 2.3KB | `isv_dev_ptr` |
|
||||
| PER buffer | priorities + tree = ~1MB | `gpu_replay.priorities_d` |
|
||||
| Step counter + RNG | ~100B | `step`, xorshift state |
|
||||
|
||||
Total: ~50MB per checkpoint.
|
||||
|
||||
### Format
|
||||
|
||||
Flat binary: `[magic: u32][version: u32][step: u64][n_sections: u32][sections...]` where each section is `[name_len: u16][name: bytes][data_len: u64][data: bytes]`. No serde, no JSON — raw device→mapped-pinned→file.
|
||||
|
||||
### Path
|
||||
|
||||
`/feature-cache/alpha-rl-runs/<sha>/fold<N>/checkpoint-<step>.bin`
|
||||
|
||||
Keep last 2 checkpoints (rolling). Delete older ones to save PVC space.
|
||||
|
||||
### Resume
|
||||
|
||||
CLI flag `--resume-from <path>`. Loads checkpoint, restores all device buffers, continues from saved step. The `alpha_rl_train.rs` binary checks for the flag before the training loop.
|
||||
|
||||
### Verification
|
||||
|
||||
Save at step 1000, kill, resume, compare l_q/wr at step 2000 with an uninterrupted run. Should be identical (deterministic with scoped_init_seed per pearl).
|
||||
|
||||
## P2: Controller Kernel Fusion
|
||||
|
||||
### What
|
||||
|
||||
Fuse 10+ single-thread ISV controller kernels into one `rl_fused_all_controllers` mega-kernel. Currently each controller is a separate `raw_launch(grid=1, block=1)` — the launch overhead (~5μs each) dominates the actual compute (~1μs each).
|
||||
|
||||
### Controllers to fuse
|
||||
|
||||
- `rl_gamma_controller`
|
||||
- `rl_target_tau_controller`
|
||||
- `rl_ppo_clip_controller` (legacy but still launched)
|
||||
- `rl_entropy_coef_controller`
|
||||
- `rl_per_alpha_controller`
|
||||
- `rl_reward_scale_controller`
|
||||
- `rl_rollout_steps_controller`
|
||||
- `rl_lr_controller` (5 heads)
|
||||
|
||||
All share the same pattern: read ISV input slot, Wiener-α blend, Schulman bounded step, write ISV output slot. The existing `rl_fused_controllers.cu` already handles a subset — extend to cover all.
|
||||
|
||||
### Impact
|
||||
|
||||
~10 kernel launches → 1. Saves ~50μs/step. At 6 sps that's ~0.3ms saved per 160ms step = ~0.2% improvement. Small but free.
|
||||
|
||||
### Verification
|
||||
|
||||
Before/after nsys trace: total GPU kernel launches should drop by ~10 per step.
|
||||
|
||||
## P3: Walk-Forward Validation
|
||||
|
||||
### What
|
||||
|
||||
3-fold temporal walk-forward on H100 with 25k steps per fold. Each fold trains on window [i] and evaluates on window [i+1]. The argo template already supports `--folds 3` which fans out via DAG.
|
||||
|
||||
### Config
|
||||
|
||||
```bash
|
||||
argo submit --from=wftmpl/alpha-rl -n foxhunt \
|
||||
-p git-branch=ml-alpha-phase-a \
|
||||
-p n-steps=25000 \
|
||||
-p n-backtests=1024 \
|
||||
-p per-capacity=65536 \
|
||||
-p gpu-pool=ci-training-h100
|
||||
# --folds 3 via argo-train.sh
|
||||
```
|
||||
|
||||
### Success criteria
|
||||
|
||||
- wr > 0.55 on ALL 3 folds (not just the average)
|
||||
- No fold with wr < 0.50
|
||||
- Entropy stable (no collapse on any fold)
|
||||
- hold% between 30-70% on all folds
|
||||
|
||||
### Wall clock estimate
|
||||
|
||||
3 folds × 25k steps. With TF32 on H100 at ~15 sps: 25k/15 = 28 min/fold. Sequential: 84 min. Parallel (3 H100s): 28 min.
|
||||
|
||||
## Implementation Order
|
||||
|
||||
1. **P0: TF32** — 2 one-line changes, local smoke, deploy
|
||||
2. **P2: Controller fusion** — extend existing fused kernel, local smoke
|
||||
3. **P1: Checkpoint** — new infrastructure, needs careful testing
|
||||
4. **P3: Walk-forward** — deploy after P0+P2 are validated
|
||||
|
||||
P0 and P2 can be implemented in parallel. P1 is independent. P3 runs after all code changes are merged.
|
||||
|
||||
## Anti-patterns to avoid
|
||||
|
||||
- No FP16 matmuls (loss scaling complexity, RL numerical sensitivity)
|
||||
- No multi-GPU (NCCL overhead not justified at 25k steps)
|
||||
- No changes to model architecture (validated at wr=0.567)
|
||||
- No monitoring pods on GPU node (caused OOM — use log-only monitoring)
|
||||
@@ -158,7 +158,6 @@ spec:
|
||||
|
||||
# Auto-detect GPU compute capability for all crates' build.rs
|
||||
export CUDA_COMPUTE_CAP=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | head -1 | tr -d '.')
|
||||
export FOXHUNT_CUDA_ARCH="sm_${CUDA_COMPUTE_CAP}"
|
||||
echo "=== GPU arch: sm_${CUDA_COMPUTE_CAP} ==="
|
||||
|
||||
# Clear stale build artifacts when arch changes
|
||||
|
||||
2
vendor/cudarc/src/driver/safe/core.rs
vendored
2
vendor/cudarc/src/driver/safe/core.rs
vendored
@@ -85,7 +85,7 @@ impl CudaContext {
|
||||
cu_device,
|
||||
cu_ctx,
|
||||
ordinal,
|
||||
has_async_alloc: AtomicBool::new(true),
|
||||
has_async_alloc: AtomicBool::new(false),
|
||||
num_streams: AtomicUsize::new(0),
|
||||
event_tracking: AtomicBool::new(true),
|
||||
error_state: AtomicU32::new(0),
|
||||
|
||||
Reference in New Issue
Block a user