Compare commits
27 Commits
ml-alpha-a
...
ml-alpha-p
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d8247034f8 | ||
|
|
346e6670f5 | ||
|
|
54de55d4bc | ||
|
|
d76919d6a2 | ||
|
|
9c1b70edec | ||
|
|
7c96504155 | ||
|
|
f7427b27de | ||
|
|
5717bc07fe | ||
|
|
ff8cacb0e9 | ||
|
|
de378c5f62 | ||
|
|
6fca6a1d9b | ||
|
|
dfbc916227 | ||
|
|
90f178ae9e | ||
|
|
2bdb55cc5b | ||
|
|
fa36d55384 | ||
|
|
fab9c0e324 | ||
|
|
35d21cb42f | ||
|
|
49dd4146ee | ||
|
|
1c0c246ddf | ||
|
|
25ec8c7bcf | ||
|
|
f385558fdb | ||
|
|
7f4cd86421 | ||
|
|
66ec7f75f4 | ||
|
|
a1277af6c7 | ||
|
|
4fa9da9fc1 | ||
|
|
e730b5cbd1 | ||
|
|
6e641b934c |
1
Cargo.lock
generated
1
Cargo.lock
generated
@@ -6090,6 +6090,7 @@ dependencies = [
|
||||
"approx",
|
||||
"arrow 56.2.0",
|
||||
"bincode",
|
||||
"bytemuck",
|
||||
"clap",
|
||||
"cudarc",
|
||||
"data",
|
||||
|
||||
@@ -33,6 +33,7 @@ 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,6 +73,7 @@ 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
|
||||
@@ -111,6 +112,7 @@ const KERNELS: &[&str] = &[
|
||||
"rl_hindsight_track", // HER Phase 1: per-step mid-price ring + peak tracking for backward hindsight
|
||||
"rl_hindsight_inject", // HER Phase 2: backward inject — synthetic replay push on done if peak >> actual
|
||||
"rl_hindsight_forward", // HER Phase 3: forward continuation — evaluates closed trades after lookahead
|
||||
"rl_lr_from_mapped_pinned", // Mega-graph: LR controller reads loss from mapped-pinned dev_ptrs (eliminates host loss readback between reward + training graphs)
|
||||
"rl_increment_step", // device-resident step counter bump (ISV[548] += 1.0); graph-safe prereq — removes scalar current_step from all downstream kernel args
|
||||
"rl_fused_controllers", // fused kernel: 10 RL ISV controllers in one launch (gamma, tau, ppo_clip, entropy_coef, rollout_steps, per_alpha, reward_scale, ppo_ratio_clamp, gate_threshold, q_distill_lambda) — saves 9 kernel launch overheads (~40-80μs/step)
|
||||
"rl_popart_normalize", // PopArt: Welford-EMA reward normalization (replaces apply_reward_scale)
|
||||
|
||||
@@ -44,3 +44,41 @@ extern "C" __global__ void adamw_step(
|
||||
|
||||
theta[i] -= lr * (m_hat / (sqrtf(v_hat) + eps) + wd * theta[i]);
|
||||
}
|
||||
|
||||
// Mega-graph variant: reads LR from an ISV device pointer at a given
|
||||
// slot index instead of taking a scalar argument. This allows the LR
|
||||
// to vary across graph replays (the ISV slot is modified in-place by
|
||||
// the rl_lr_from_mapped_pinned controller kernel which is captured
|
||||
// earlier in the same graph). All other args are identical to adamw_step.
|
||||
extern "C" __global__ void adamw_step_isv_lr(
|
||||
float* __restrict__ theta,
|
||||
const float* __restrict__ grad,
|
||||
float* __restrict__ m,
|
||||
float* __restrict__ v,
|
||||
int n_params,
|
||||
const float* __restrict__ isv,
|
||||
int lr_slot,
|
||||
float beta1,
|
||||
float beta2,
|
||||
float eps,
|
||||
float wd,
|
||||
int step
|
||||
) {
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= n_params) return;
|
||||
|
||||
const float lr = isv[lr_slot];
|
||||
|
||||
const float g = grad[i];
|
||||
const float m_n = beta1 * m[i] + (1.0f - beta1) * g;
|
||||
const float v_n = beta2 * v[i] + (1.0f - beta2) * g * g;
|
||||
m[i] = m_n;
|
||||
v[i] = v_n;
|
||||
|
||||
const float bc1 = 1.0f - powf(beta1, (float) step);
|
||||
const float bc2 = 1.0f - powf(beta2, (float) step);
|
||||
const float m_hat = m_n / fmaxf(bc1, 1e-12f);
|
||||
const float v_hat = v_n / fmaxf(bc2, 1e-12f);
|
||||
|
||||
theta[i] -= lr * (m_hat / (sqrtf(v_hat) + eps) + wd * theta[i]);
|
||||
}
|
||||
|
||||
@@ -24,6 +24,8 @@
|
||||
|
||||
#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 {
|
||||
@@ -80,11 +82,26 @@ 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, gpu_gather_frd_labels,
|
||||
// and gpu_gather_bce_labels to reuse the same sampling decision.
|
||||
// per batch element. Read by gpu_gather_next and gpu_gather_frd_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;
|
||||
@@ -148,9 +165,26 @@ 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 ──────────────────────────────────────────────
|
||||
// ── Label gather kernel (STANDALONE — kept for backward compat) ─────
|
||||
//
|
||||
// Runs AFTER gpu_sample_and_gather. Gathers per-horizon FRD labels at
|
||||
// the anchor position (rightmost K position = newest snapshot in the
|
||||
|
||||
@@ -36,3 +36,20 @@ extern "C" __global__ void grad_h_accumulate_scaled(
|
||||
if (i >= n) return;
|
||||
grad_h_encoder[i] += lambda * grad_h_head[i];
|
||||
}
|
||||
|
||||
// Mega-graph variant: reads lambda from ISV device pointer at given slot.
|
||||
// This allows the lambda to vary across graph replays (ISV is modified
|
||||
// in-place by controller kernels captured earlier in the same graph).
|
||||
extern "C" __global__ void grad_h_accumulate_scaled_isv(
|
||||
const float* __restrict__ grad_h_head,
|
||||
const float* __restrict__ isv,
|
||||
int lambda_slot,
|
||||
float batch_inv, // 1.0 / b_size
|
||||
int n,
|
||||
float* __restrict__ grad_h_encoder
|
||||
) {
|
||||
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= n) return;
|
||||
const float lambda = isv[lambda_slot] * batch_inv;
|
||||
grad_h_encoder[i] += lambda * grad_h_head[i];
|
||||
}
|
||||
|
||||
@@ -34,6 +34,8 @@
|
||||
#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
|
||||
@@ -51,6 +53,35 @@ 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;
|
||||
|
||||
49
crates/ml-alpha/cuda/rl_drawdown_stop.cu
Normal file
49
crates/ml-alpha/cuda/rl_drawdown_stop.cu
Normal file
@@ -0,0 +1,49 @@
|
||||
// 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;
|
||||
}
|
||||
134
crates/ml-alpha/cuda/rl_lr_from_mapped_pinned.cu
Normal file
134
crates/ml-alpha/cuda/rl_lr_from_mapped_pinned.cu
Normal file
@@ -0,0 +1,134 @@
|
||||
// rl_lr_from_mapped_pinned.cu — GPU-side LR controller that reads loss
|
||||
// observations from mapped-pinned device pointers instead of host-
|
||||
// supplied scalar arguments. This eliminates the host-side loss readback
|
||||
// between the reward graph and training graph, enabling mega-graph
|
||||
// capture of the entire per-step pipeline.
|
||||
//
|
||||
// The mapped-pinned buffers (ss_q_loss_mapped, ss_pi_loss_mapped,
|
||||
// ss_v_loss_sum_mapped) have stable device pointers. The GPU writes
|
||||
// loss values during backward kernels; this kernel reads them from the
|
||||
// same physical memory via the device-side pointer. The reads are
|
||||
// coherent because this kernel launches AFTER the backward kernels on
|
||||
// the same stream (stream ordering guarantees visibility).
|
||||
//
|
||||
// Internally delegates to the same plateau_decay_head logic as
|
||||
// rl_lr_controller.cu. The only difference is the loss observation
|
||||
// source: pointer dereference instead of scalar argument.
|
||||
|
||||
#define RL_LR_BCE_INDEX 412
|
||||
#define RL_LR_Q_INDEX 413
|
||||
#define RL_LR_PI_INDEX 414
|
||||
#define RL_LR_V_INDEX 415
|
||||
#define RL_LR_AUX_INDEX 416
|
||||
|
||||
#define RL_LR_BOOTSTRAP_INDEX 462
|
||||
#define RL_LR_MIN_INDEX 463
|
||||
#define RL_LR_MAX_INDEX 464
|
||||
#define RL_LR_LOSS_EMA_ALPHA_INDEX 465
|
||||
#define RL_LR_DECAY_FACTOR_INDEX 466
|
||||
|
||||
#define RL_IMPROVEMENT_THRESHOLD_INDEX 455
|
||||
#define RL_PLATEAU_PATIENCE_INDEX 456
|
||||
#define RL_LR_WARMUP_STEPS_INDEX 461
|
||||
|
||||
__device__ __forceinline__ void plateau_decay_head(
|
||||
float* isv,
|
||||
int lr_idx,
|
||||
float observed_loss,
|
||||
int loss_ema_slot,
|
||||
int best_slot,
|
||||
int counter_slot,
|
||||
int warmup_slot
|
||||
) {
|
||||
const float lr_prev = isv[lr_idx];
|
||||
const float lr_bootstrap = isv[RL_LR_BOOTSTRAP_INDEX];
|
||||
|
||||
if (lr_prev == 0.0f) {
|
||||
isv[lr_idx] = lr_bootstrap;
|
||||
return;
|
||||
}
|
||||
|
||||
if (observed_loss == 0.0f) return;
|
||||
if (loss_ema_slot < 0) return;
|
||||
|
||||
const float loss_ema_prev = isv[loss_ema_slot];
|
||||
const float loss_ema_alpha = isv[RL_LR_LOSS_EMA_ALPHA_INDEX];
|
||||
float loss_ema_new;
|
||||
if (loss_ema_prev == 0.0f) {
|
||||
loss_ema_new = observed_loss;
|
||||
} else {
|
||||
loss_ema_new = (1.0f - loss_ema_alpha) * loss_ema_prev
|
||||
+ loss_ema_alpha * observed_loss;
|
||||
}
|
||||
isv[loss_ema_slot] = loss_ema_new;
|
||||
|
||||
const float warmup_prev = isv[warmup_slot];
|
||||
const float warmup_target = isv[RL_LR_WARMUP_STEPS_INDEX];
|
||||
if (warmup_prev < warmup_target) {
|
||||
isv[best_slot] = loss_ema_new;
|
||||
isv[counter_slot] = 0.0f;
|
||||
isv[warmup_slot] = warmup_prev + 1.0f;
|
||||
return;
|
||||
}
|
||||
|
||||
const float improvement_threshold = isv[RL_IMPROVEMENT_THRESHOLD_INDEX];
|
||||
const float plateau_patience = isv[RL_PLATEAU_PATIENCE_INDEX];
|
||||
const float best_prev = isv[best_slot];
|
||||
|
||||
if (loss_ema_new < best_prev * improvement_threshold) {
|
||||
isv[best_slot] = loss_ema_new;
|
||||
isv[counter_slot] = 0.0f;
|
||||
return;
|
||||
}
|
||||
|
||||
const float counter_next = isv[counter_slot] + 1.0f;
|
||||
if (counter_next >= plateau_patience) {
|
||||
const float lr_min = isv[RL_LR_MIN_INDEX];
|
||||
const float decay_factor = isv[RL_LR_DECAY_FACTOR_INDEX];
|
||||
const float lr_new = fmaxf(lr_min, lr_prev * decay_factor);
|
||||
isv[lr_idx] = lr_new;
|
||||
isv[counter_slot] = 0.0f;
|
||||
} else {
|
||||
isv[counter_slot] = counter_next;
|
||||
}
|
||||
}
|
||||
|
||||
extern "C" __global__ void rl_lr_from_mapped_pinned(
|
||||
float* __restrict__ isv,
|
||||
const float* __restrict__ q_loss_ptr, // mapped-pinned dev_ptr (1 float)
|
||||
const float* __restrict__ pi_loss_ptr, // mapped-pinned dev_ptr (1 float)
|
||||
const float* __restrict__ v_loss_sum_ptr, // mapped-pinned dev_ptr (1 float)
|
||||
int b_size, // batch size for V loss normalization
|
||||
int q_loss_ema_slot,
|
||||
int q_best_slot,
|
||||
int q_counter_slot,
|
||||
int q_warmup_slot,
|
||||
int pi_loss_ema_slot,
|
||||
int pi_best_slot,
|
||||
int pi_counter_slot,
|
||||
int pi_warmup_slot,
|
||||
int v_loss_ema_slot,
|
||||
int v_best_slot,
|
||||
int v_counter_slot,
|
||||
int v_warmup_slot
|
||||
) {
|
||||
if (threadIdx.x != 0 || blockIdx.x != 0) return;
|
||||
|
||||
// Read losses from mapped-pinned device pointers. These are the
|
||||
// same physical pages the backward kernels wrote to — coherent
|
||||
// because this kernel is stream-ordered after them.
|
||||
const float observed_loss_q = *q_loss_ptr;
|
||||
const float observed_loss_pi = *pi_loss_ptr;
|
||||
// V loss is stored as a sum across the batch; normalize to per-sample.
|
||||
const float observed_loss_v = (b_size > 0) ? (*v_loss_sum_ptr / (float)b_size) : 0.0f;
|
||||
|
||||
// BCE and AUX: perception-owned heads, pass slot=-1 to skip.
|
||||
plateau_decay_head(isv, RL_LR_BCE_INDEX, 0.0f, -1, -1, -1, -1);
|
||||
plateau_decay_head(isv, RL_LR_Q_INDEX, observed_loss_q,
|
||||
q_loss_ema_slot, q_best_slot, q_counter_slot, q_warmup_slot);
|
||||
plateau_decay_head(isv, RL_LR_PI_INDEX, observed_loss_pi,
|
||||
pi_loss_ema_slot, pi_best_slot, pi_counter_slot, pi_warmup_slot);
|
||||
plateau_decay_head(isv, RL_LR_V_INDEX, observed_loss_v,
|
||||
v_loss_ema_slot, v_best_slot, v_counter_slot, v_warmup_slot);
|
||||
plateau_decay_head(isv, RL_LR_AUX_INDEX, 0.0f, -1, -1, -1, -1);
|
||||
}
|
||||
@@ -1,25 +1,23 @@
|
||||
/**
|
||||
* rl_per_push_flush.cu — Coordinated coalesced replay write.
|
||||
* rl_per_push_flush.cu — Two-kernel coordinated coalesced replay write.
|
||||
*
|
||||
* Second kernel in the split rl_per_push pipeline. After rl_per_push_ring has
|
||||
* written data to the n-step ring and set flush_flags[b] for completed n-step
|
||||
* transitions, this kernel coordinates slot allocation and performs the actual
|
||||
* coalesced write to the replay buffer.
|
||||
* Split from the original single-kernel volatile-spin design into two
|
||||
* kernels with stream-ordered dependency (graph-safe):
|
||||
*
|
||||
* Architecture:
|
||||
* - Grid=(b_size, 1, 1), Block=(128, 1, 1). One block per batch element.
|
||||
* - Block 0 runs the prefix-sum to assign write slots (single-threaded, O(B)).
|
||||
* - All other blocks spin on a volatile global-memory ready flag (no atomics).
|
||||
* - Once ready, each flushing block does a coalesced 128-thread copy of h_t
|
||||
* and h_tp1, then thread 0 computes the n-step return and writes scalars +
|
||||
* priority + resets the ring count.
|
||||
* Kernel 1: rl_per_flush_prefix_sum
|
||||
* Grid=(1), Block=(1). Single thread computes exclusive prefix-sum of
|
||||
* flush_flags, advances write_head and replay_len, writes per-batch
|
||||
* offsets + base into write_offsets[0..b_size+1].
|
||||
*
|
||||
* Coordination: block-0 prefix-sum + volatile ready flag.
|
||||
* No atomicAdd — slot assignment is deterministic serial scan (feedback_no_atomicadd).
|
||||
* Kernel 2: rl_per_flush_write
|
||||
* Grid=(b_size), Block=(128). Each block reads its pre-computed offset
|
||||
* from write_offsets, then does the coalesced h_t/h_tp1 copy and
|
||||
* scalar write. No spin — the prefix-sum kernel completed before this
|
||||
* launch (stream ordering).
|
||||
*
|
||||
* No atomicAdd — slot assignment is deterministic serial scan.
|
||||
* No volatile spin — correctness from stream ordering / graph node deps.
|
||||
* Pre-compiled cubin only (feedback_no_nvrtc).
|
||||
*
|
||||
* Launch: CUDA_HOME=/usr/local/cuda nvcc -cubin -arch sm_86 -O3 \
|
||||
* --generate-line-info rl_per_push_flush.cu -o rl_per_push_flush.cubin
|
||||
*/
|
||||
|
||||
#define HIDDEN_DIM 128
|
||||
@@ -30,105 +28,89 @@
|
||||
#define RL_GAMMA_INDEX 400
|
||||
#define RL_N_STEP_INDEX 403
|
||||
|
||||
extern "C" __global__ void rl_per_push_flush(
|
||||
/* ── Flush coordination ── */
|
||||
const int* __restrict__ flush_flags, /* [B] from ring kernel: 1=flush, 0=skip */
|
||||
int* __restrict__ write_offsets, /* [B+2] output: per-batch offsets, base, ready flag */
|
||||
/* =====================================================================
|
||||
* Kernel 1: prefix-sum + write_head advance.
|
||||
*
|
||||
* Grid=(1,1,1), Block=(1,1,1). Single thread.
|
||||
*
|
||||
* Reads flush_flags[0..b_size], computes exclusive prefix-sum into
|
||||
* write_offsets[0..b_size], stores base write_head at write_offsets[b_size],
|
||||
* and advances the replay buffer's write_head + replay_len.
|
||||
* ===================================================================== */
|
||||
extern "C" __global__ void rl_per_flush_prefix_sum(
|
||||
const int* __restrict__ flush_flags, /* [B] */
|
||||
int* __restrict__ write_offsets, /* [B+1] output: offsets + base */
|
||||
unsigned int* __restrict__ write_head, /* [1] */
|
||||
unsigned int* __restrict__ replay_len, /* [1] */
|
||||
int b_size,
|
||||
int capacity
|
||||
)
|
||||
{
|
||||
unsigned int base = write_head[0];
|
||||
|
||||
/* ── N-step ring (read source) ── */
|
||||
const float* __restrict__ nstep_scalars, /* [B * N_STEP_MAX * 3]: (r_scaled, r_raw, done) per ring slot */
|
||||
/* Exclusive prefix-sum */
|
||||
int total = 0;
|
||||
for (int i = 0; i < b_size; ++i) {
|
||||
write_offsets[i] = total;
|
||||
total += flush_flags[i];
|
||||
}
|
||||
|
||||
/* Advance write head and replay length */
|
||||
if (total > 0) {
|
||||
write_head[0] = (base + (unsigned int)total) % (unsigned int)capacity;
|
||||
unsigned int len = replay_len[0] + (unsigned int)total;
|
||||
if (len > (unsigned int)capacity) len = (unsigned int)capacity;
|
||||
replay_len[0] = len;
|
||||
}
|
||||
|
||||
/* Store base write head for the write kernel */
|
||||
write_offsets[b_size] = (int)base;
|
||||
}
|
||||
|
||||
/* =====================================================================
|
||||
* Kernel 2: coalesced replay write.
|
||||
*
|
||||
* Grid=(b_size,1,1), Block=(128,1,1). One block per batch element.
|
||||
*
|
||||
* Reads pre-computed offset from write_offsets (populated by kernel 1).
|
||||
* No spin — stream ordering guarantees kernel 1 completed before this.
|
||||
* ===================================================================== */
|
||||
extern "C" __global__ void rl_per_flush_write(
|
||||
const int* __restrict__ flush_flags, /* [B] */
|
||||
const int* __restrict__ write_offsets, /* [B+1] */
|
||||
|
||||
/* N-step ring (read source) */
|
||||
const float* __restrict__ nstep_scalars, /* [B * N_STEP_MAX * 3] */
|
||||
const float* __restrict__ nstep_h_t, /* [B * N_STEP_MAX * HIDDEN_DIM] */
|
||||
const unsigned int* __restrict__ nstep_write_idx, /* [B] ring write cursor */
|
||||
const unsigned int* __restrict__ nstep_count, /* [B] items in ring */
|
||||
|
||||
/* ── Current step data (for h_tp1) ── */
|
||||
/* Current step data (for h_tp1) */
|
||||
const float* __restrict__ h_tp1_current, /* [B * HIDDEN_DIM] */
|
||||
const float* __restrict__ log_pi_old, /* [B] */
|
||||
const int* __restrict__ actions, /* [B] */
|
||||
const float* __restrict__ isv, /* ISV array */
|
||||
|
||||
/* ── Replay storage (write destination) ── */
|
||||
/* Replay storage (write destination) */
|
||||
float* __restrict__ replay_h_t, /* [capacity * HIDDEN_DIM] */
|
||||
float* __restrict__ replay_h_tp1, /* [capacity * HIDDEN_DIM] */
|
||||
float* __restrict__ replay_scalars, /* [capacity * SCALARS_PER_TRANSITION] */
|
||||
float* __restrict__ priority_tree, /* [2 * capacity] (leaf layer at offset capacity) */
|
||||
unsigned int* __restrict__ write_head, /* [1] circular write cursor */
|
||||
unsigned int* __restrict__ replay_len, /* [1] current buffer occupancy */
|
||||
float* __restrict__ max_priority, /* [1] running max priority for new inserts */
|
||||
float* __restrict__ max_priority, /* [1] running max priority */
|
||||
|
||||
/* ── N-step count reset (written on flush) ── */
|
||||
unsigned int* __restrict__ nstep_count_out, /* [B] — set to 0 for flushed batches */
|
||||
/* N-step count reset (written on flush) */
|
||||
unsigned int* __restrict__ nstep_count_out, /* [B] */
|
||||
|
||||
int b_size,
|
||||
int capacity
|
||||
)
|
||||
{
|
||||
/* ====================================================================
|
||||
* PHASE 1: Block 0 coordination (thread 0 only).
|
||||
*
|
||||
* Computes prefix-sum of flush_flags to assign contiguous write slots,
|
||||
* advances write_head and replay_len, stores base for other blocks,
|
||||
* and signals ready via volatile write.
|
||||
* ==================================================================== */
|
||||
if (blockIdx.x == 0 && threadIdx.x == 0) {
|
||||
unsigned int base = write_head[0];
|
||||
|
||||
/* Exclusive prefix-sum: write_offsets[i] = number of flushes before batch i */
|
||||
int total = 0;
|
||||
for (int i = 0; i < b_size; ++i) {
|
||||
write_offsets[i] = total;
|
||||
total += flush_flags[i];
|
||||
}
|
||||
|
||||
/* Advance write head and replay length */
|
||||
if (total > 0) {
|
||||
write_head[0] = (base + (unsigned int)total) % (unsigned int)capacity;
|
||||
unsigned int len = replay_len[0] + (unsigned int)total;
|
||||
if (len > (unsigned int)capacity) len = (unsigned int)capacity;
|
||||
replay_len[0] = len;
|
||||
}
|
||||
|
||||
/* Store base write head for other blocks at slot [b_size] */
|
||||
write_offsets[b_size] = (int)base;
|
||||
|
||||
/* Fence: ensure all writes above are globally visible before signaling */
|
||||
__threadfence();
|
||||
|
||||
/* Signal ready at slot [b_size + 1] */
|
||||
write_offsets[b_size + 1] = 1;
|
||||
__threadfence();
|
||||
}
|
||||
|
||||
/* ====================================================================
|
||||
* PHASE 2: All blocks wait for ready signal.
|
||||
*
|
||||
* Non-block-0 blocks spin on a volatile read of the ready flag.
|
||||
* Block 0 already wrote it, so it proceeds immediately.
|
||||
* ==================================================================== */
|
||||
if (threadIdx.x == 0 && blockIdx.x != 0) {
|
||||
volatile int* ready = (volatile int*)&write_offsets[b_size + 1];
|
||||
while (*ready == 0) {
|
||||
/* spin — volatile load, no atomic */
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
/* ====================================================================
|
||||
* PHASE 3: Coalesced replay write (128 threads per block).
|
||||
*
|
||||
* Each flushing block writes:
|
||||
* - h_t[HIDDEN_DIM]: from oldest ring entry (coalesced, all 128 threads)
|
||||
* - h_tp1[HIDDEN_DIM]: from current step (coalesced, all 128 threads)
|
||||
* - scalars[7]: n-step return + metadata (thread 0 only)
|
||||
* - priority: max_priority into leaf layer (thread 0 only)
|
||||
* - reset: nstep_count_out[b] = 0 (thread 0 only)
|
||||
* ==================================================================== */
|
||||
const int b = blockIdx.x;
|
||||
|
||||
/* Early exit: nothing to flush for this batch */
|
||||
if (flush_flags[b] == 0) return;
|
||||
|
||||
/* Compute target replay slot */
|
||||
/* Compute target replay slot from pre-computed offsets */
|
||||
const unsigned int base = (unsigned int)write_offsets[b_size];
|
||||
const unsigned int my_offset = (unsigned int)write_offsets[b];
|
||||
const unsigned int my_slot = (base + my_offset) % (unsigned int)capacity;
|
||||
|
||||
@@ -1,11 +1,11 @@
|
||||
/* =====================================================================
|
||||
* rl_per_sample.cu — GPU-resident PER: proportional sampling via sum-tree
|
||||
*
|
||||
* Grid=(b_size), Block=(1). One thread per sample.
|
||||
* Grid=(b_size), Block=(128). One block per sample.
|
||||
*
|
||||
* Each thread samples a leaf from the priority sum-tree using stratified
|
||||
* sampling (segment per thread) with xorshift32 PRNG, then gathers the
|
||||
* transition data from replay storage.
|
||||
* Thread 0 does the tree walk + scalar unpack. All 128 threads cooperate
|
||||
* on the coalesced h_t and h_tp1 gathers (128 floats = HIDDEN_DIM per
|
||||
* sample, one float per thread, single 512-byte coalesced transaction).
|
||||
*
|
||||
* ISV reads: none (alpha used only at priority-update time)
|
||||
* ===================================================================== */
|
||||
@@ -44,75 +44,87 @@ extern "C" __global__ void rl_per_sample(
|
||||
int capacity
|
||||
)
|
||||
{
|
||||
const int b = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int b = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
if (b >= b_size) return;
|
||||
|
||||
const unsigned int len = replay_len[0];
|
||||
const float total_priority = priority_tree[1]; /* root */
|
||||
/* ── Shared memory for broadcasting the sampled leaf index ────────── */
|
||||
__shared__ int s_safe_leaf;
|
||||
|
||||
/* Guard: empty or zero-priority replay */
|
||||
if (total_priority < 1e-9f || len == 0) {
|
||||
for (int i = 0; i < HIDDEN_DIM; ++i) {
|
||||
sampled_h_t[b * HIDDEN_DIM + i] = 0.0f;
|
||||
sampled_h_tp1[b * HIDDEN_DIM + i] = 0.0f;
|
||||
/* ── Thread 0: tree walk + scalar reads ──────────────────────────── */
|
||||
if (tid == 0) {
|
||||
const unsigned int len = replay_len[0];
|
||||
const float total_priority = priority_tree[1]; /* root */
|
||||
|
||||
/* Guard: empty or zero-priority replay */
|
||||
if (total_priority < 1e-9f || len == 0) {
|
||||
s_safe_leaf = -1; /* sentinel: zero-fill outputs */
|
||||
} else {
|
||||
/* Seed PRNG if cold */
|
||||
if (prng_state[b] == 0) {
|
||||
prng_state[b] = (unsigned int)(b + 1) * 2654435761u;
|
||||
}
|
||||
|
||||
/* Stratified sampling: draw u in [b*segment, (b+1)*segment) */
|
||||
const float segment = total_priority / (float)b_size;
|
||||
unsigned int rng = xorshift32(&prng_state[b]);
|
||||
const float u_frac = (float)(rng & 0x00FFFFFFu) / (float)0x01000000u;
|
||||
float u = ((float)b + u_frac) * segment;
|
||||
|
||||
/* Clamp to avoid floating-point overshoot */
|
||||
if (u >= total_priority) u = total_priority - 1e-6f;
|
||||
if (u < 0.0f) u = 0.0f;
|
||||
|
||||
/* Walk tree top-down */
|
||||
int idx = 1;
|
||||
while (idx < capacity) {
|
||||
const int left = 2 * idx;
|
||||
const float left_val = priority_tree[left];
|
||||
if (u <= left_val) {
|
||||
idx = left;
|
||||
} else {
|
||||
u -= left_val;
|
||||
idx = left + 1;
|
||||
}
|
||||
}
|
||||
const int leaf = idx - capacity;
|
||||
|
||||
/* Clamp leaf to valid range */
|
||||
const int safe_leaf = (leaf >= 0 && leaf < (int)len) ? leaf : 0;
|
||||
s_safe_leaf = safe_leaf;
|
||||
sample_indices[b] = (unsigned int)safe_leaf;
|
||||
|
||||
/* Unpack scalars (thread 0 only — small data, not worth parallelising) */
|
||||
const int sc_base = safe_leaf * SCALARS_PER_TRANSITION;
|
||||
sampled_actions[b] = (int)replay_scalars[sc_base + 0];
|
||||
sampled_rewards[b] = replay_scalars[sc_base + 1];
|
||||
sampled_dones[b] = replay_scalars[sc_base + 3];
|
||||
sampled_log_pi_old[b] = replay_scalars[sc_base + 4];
|
||||
sampled_n_step_gammas[b] = replay_scalars[sc_base + 5];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const int safe_leaf = s_safe_leaf;
|
||||
|
||||
if (safe_leaf < 0) {
|
||||
/* Empty replay — zero-fill h_t and h_tp1 (coalesced, all 128 threads) */
|
||||
sampled_h_t[b * HIDDEN_DIM + tid] = 0.0f;
|
||||
sampled_h_tp1[b * HIDDEN_DIM + tid] = 0.0f;
|
||||
if (tid == 0) {
|
||||
sampled_rewards[b] = 0.0f;
|
||||
sampled_dones[b] = 0.0f;
|
||||
sampled_log_pi_old[b] = 0.0f;
|
||||
sampled_n_step_gammas[b] = 0.0f;
|
||||
sampled_actions[b] = 0;
|
||||
sample_indices[b] = 0;
|
||||
}
|
||||
sampled_rewards[b] = 0.0f;
|
||||
sampled_dones[b] = 0.0f;
|
||||
sampled_log_pi_old[b] = 0.0f;
|
||||
sampled_n_step_gammas[b] = 0.0f;
|
||||
sampled_actions[b] = 0;
|
||||
sample_indices[b] = 0;
|
||||
return;
|
||||
}
|
||||
|
||||
/* ── Seed PRNG if cold ────────────────────────────────────────────── */
|
||||
if (prng_state[b] == 0) {
|
||||
prng_state[b] = (unsigned int)(b + 1) * 2654435761u;
|
||||
}
|
||||
/* ── Coalesced h_t gather (128 threads, 1 float each) ────────────── */
|
||||
sampled_h_t[b * HIDDEN_DIM + tid] = replay_h_t[safe_leaf * HIDDEN_DIM + tid];
|
||||
|
||||
/* ── Stratified sampling: draw u in [b*segment, (b+1)*segment) ────── */
|
||||
const float segment = total_priority / (float)b_size;
|
||||
unsigned int rng = xorshift32(&prng_state[b]);
|
||||
const float u_frac = (float)(rng & 0x00FFFFFFu) / (float)0x01000000u; /* [0, 1) */
|
||||
float u = ((float)b + u_frac) * segment;
|
||||
|
||||
/* Clamp to avoid floating-point overshoot */
|
||||
if (u >= total_priority) u = total_priority - 1e-6f;
|
||||
if (u < 0.0f) u = 0.0f;
|
||||
|
||||
/* ── Walk tree top-down ───────────────────────────────────────────── */
|
||||
int idx = 1;
|
||||
while (idx < capacity) {
|
||||
const int left = 2 * idx;
|
||||
const float left_val = priority_tree[left];
|
||||
if (u <= left_val) {
|
||||
idx = left;
|
||||
} else {
|
||||
u -= left_val;
|
||||
idx = left + 1;
|
||||
}
|
||||
}
|
||||
const int leaf = idx - capacity;
|
||||
|
||||
/* Clamp leaf to valid range */
|
||||
const int safe_leaf = (leaf >= 0 && leaf < (int)len) ? leaf : 0;
|
||||
sample_indices[b] = (unsigned int)safe_leaf;
|
||||
|
||||
/* ── Gather h_t ──────────────────────────────────────────────────── */
|
||||
for (int i = 0; i < HIDDEN_DIM; ++i) {
|
||||
sampled_h_t[b * HIDDEN_DIM + i] = replay_h_t[safe_leaf * HIDDEN_DIM + i];
|
||||
}
|
||||
|
||||
/* ── Gather h_tp1 ────────────────────────────────────────────────── */
|
||||
for (int i = 0; i < HIDDEN_DIM; ++i) {
|
||||
sampled_h_tp1[b * HIDDEN_DIM + i] = replay_h_tp1[safe_leaf * HIDDEN_DIM + i];
|
||||
}
|
||||
|
||||
/* ── Unpack scalars ──────────────────────────────────────────────── */
|
||||
const int sc_base = safe_leaf * SCALARS_PER_TRANSITION;
|
||||
sampled_actions[b] = (int)replay_scalars[sc_base + 0];
|
||||
sampled_rewards[b] = replay_scalars[sc_base + 1]; /* n-step scaled return */
|
||||
sampled_dones[b] = replay_scalars[sc_base + 3];
|
||||
sampled_log_pi_old[b] = replay_scalars[sc_base + 4];
|
||||
sampled_n_step_gammas[b] = replay_scalars[sc_base + 5];
|
||||
/* ── Coalesced h_tp1 gather (128 threads, 1 float each) ──────────── */
|
||||
sampled_h_tp1[b * HIDDEN_DIM + tid] = replay_h_tp1[safe_leaf * HIDDEN_DIM + tid];
|
||||
}
|
||||
|
||||
@@ -1,44 +1,30 @@
|
||||
/* =====================================================================
|
||||
* rl_per_tree_rebuild.cu — GPU-resident PER: bottom-up parallel sum-tree rebuild
|
||||
* rl_per_tree_rebuild.cu — GPU-resident PER: per-level sum-tree rebuild
|
||||
*
|
||||
* Grid=(128), Block=(256). Total 32768 threads = capacity.
|
||||
* One kernel invocation per tree level. The host launches log2(capacity)
|
||||
* times, from bottom (level 0 = parents of leaves) to top (root).
|
||||
* Stream ordering between launches provides the device-wide barrier that
|
||||
* the previous __threadfence() approach lacked — each level's writes are
|
||||
* fully complete before the next level reads them.
|
||||
*
|
||||
* Rebuilds the entire sum-tree from leaf priorities (at indices
|
||||
* [capacity..2*capacity)) up to the root (index 1). Each level is
|
||||
* processed in parallel with __threadfence() device-wide barriers
|
||||
* between levels.
|
||||
* Graph-safe: each level is a separate kernel node in the CUDA graph.
|
||||
* The graph engine respects stream-order dependencies between nodes.
|
||||
*
|
||||
* No atomicAdd — each internal node is written by exactly one thread.
|
||||
* ===================================================================== */
|
||||
|
||||
extern "C" __global__ void rl_per_tree_rebuild(
|
||||
extern "C" __global__ void rl_per_tree_rebuild_level(
|
||||
float* __restrict__ priority_tree, /* [2 * capacity] */
|
||||
int capacity
|
||||
int start, /* first node index at this level */
|
||||
int nodes_at_level /* number of nodes to process at this level */
|
||||
)
|
||||
{
|
||||
const int tid_global = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int total_threads = gridDim.x * blockDim.x;
|
||||
|
||||
/* Bottom-up: level 0 = parents of leaves, level (log2(cap)-1) = root */
|
||||
/* At level L, there are capacity >> (L+1) internal nodes. */
|
||||
/* Node indices at level L: [capacity >> (L+1) .. capacity >> L) */
|
||||
|
||||
int nodes_at_level = capacity >> 1; /* level 0: cap/2 nodes */
|
||||
int start = nodes_at_level; /* first node index at this level */
|
||||
|
||||
while (nodes_at_level >= 1) {
|
||||
/* Each thread processes multiple nodes via grid-stride loop */
|
||||
for (int i = tid_global; i < nodes_at_level; i += total_threads) {
|
||||
const int node = start + i;
|
||||
priority_tree[node] = priority_tree[2 * node] + priority_tree[2 * node + 1];
|
||||
}
|
||||
|
||||
/* Device-wide fence: all writes at this level visible before
|
||||
* any thread reads them at the next level */
|
||||
__threadfence();
|
||||
|
||||
/* Move up one level */
|
||||
nodes_at_level >>= 1;
|
||||
start >>= 1;
|
||||
/* Grid-stride loop: each thread processes one or more nodes */
|
||||
for (int i = tid_global; i < nodes_at_level; i += total_threads) {
|
||||
const int node = start + i;
|
||||
priority_tree[node] = priority_tree[2 * node] + priority_tree[2 * node + 1];
|
||||
}
|
||||
}
|
||||
|
||||
@@ -230,6 +230,15 @@ 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,6 +188,7 @@ 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());
|
||||
@@ -210,6 +211,18 @@ 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,6 +547,12 @@ 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,6 +311,12 @@ 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.
|
||||
|
||||
@@ -47,7 +47,7 @@ impl GpuReplayBuffer {
|
||||
nstep_write_idx_d: stream.alloc_zeros(b_size).context("gpu_replay nstep_write_idx")?,
|
||||
nstep_count_d: stream.alloc_zeros(b_size).context("gpu_replay nstep_count")?,
|
||||
flush_flags_d: stream.alloc_zeros(b_size).context("gpu_replay flush_flags")?,
|
||||
write_offsets_d: stream.alloc_zeros(b_size + 2).context("gpu_replay write_offsets")?,
|
||||
write_offsets_d: stream.alloc_zeros(b_size + 1).context("gpu_replay write_offsets")?,
|
||||
sample_indices_d: stream.alloc_zeros(b_size).context("gpu_replay sample_indices")?,
|
||||
sample_prng_d: stream.alloc_zeros(b_size).context("gpu_replay sample_prng")?,
|
||||
})
|
||||
|
||||
@@ -1111,5 +1111,16 @@ 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 = 585;
|
||||
pub const RL_SLOTS_END: usize = 588;
|
||||
|
||||
@@ -46,9 +46,57 @@
|
||||
|
||||
use anyhow::Result;
|
||||
use cudarc::driver::CudaSlice;
|
||||
use cudarc::driver::sys::CUfunction;
|
||||
|
||||
use crate::cfc::snap_features::Mbp10RawInput;
|
||||
|
||||
/// Cached raw device pointers and kernel function handles for the lobsim.
|
||||
/// Extracted once per step to bypass trait dispatch + CudaSlice borrow
|
||||
/// overhead in the mega-graph hot path. All pointers are stable (pre-
|
||||
/// allocated at lobsim construction time, never reallocated).
|
||||
pub struct LobSimRawPtrs {
|
||||
// ── Book update (apply_snapshot_from_device) ────────────────────
|
||||
pub bid_px: u64,
|
||||
pub bid_sz: u64,
|
||||
pub ask_px: u64,
|
||||
pub ask_sz: u64,
|
||||
pub books: u64,
|
||||
pub prev_mid: u64,
|
||||
pub atr_mid_ema: u64,
|
||||
pub snapshots_skipped: u64,
|
||||
pub min_reasonable_px: u64,
|
||||
pub max_reasonable_px: u64,
|
||||
pub book_update_fn: CUfunction,
|
||||
|
||||
// ── Fill (step_fill_from_market_targets) ────────────────────────
|
||||
pub market_targets: u64,
|
||||
pub pos: u64,
|
||||
pub cost_per_lot_per_side: u64,
|
||||
pub total_fees_per_b: u64,
|
||||
pub submit_market_fn: CUfunction,
|
||||
|
||||
// ── PnL track (step_pnl_track) ─────────────────────────────────
|
||||
pub open_trade_state: u64,
|
||||
pub trade_log: u64,
|
||||
pub trade_log_head: u64,
|
||||
pub trail_hwm: u64,
|
||||
pub zero_vwap_at_open: u64,
|
||||
pub saturated_vwap_at_open: u64,
|
||||
pub defensive_exit_clamp: u64,
|
||||
pub conv_signed_ema: u64,
|
||||
pub diag_hold_hist: u64,
|
||||
pub diag_outcome_n: u64,
|
||||
pub diag_outcome_sum_pnl: u64,
|
||||
pub diag_outcome_n_wins: u64,
|
||||
pub pnl_track_fn: CUfunction,
|
||||
|
||||
// ── Dimensions ─────────────────────────────────────────────────
|
||||
pub n_backtests: i32,
|
||||
pub pos_bytes: i32,
|
||||
/// `TRADE_LOG_CAP` from ml-backtesting — passed to `pnl_track` kernel.
|
||||
pub trade_log_cap: i32,
|
||||
}
|
||||
|
||||
/// Narrow device-oriented backend the integrated RL trainer's
|
||||
/// `step_with_lobsim` invokes for its LOB-simulator interaction.
|
||||
/// Implemented for `LobSimCuda` in `ml-backtesting`. See module-level
|
||||
@@ -126,4 +174,10 @@ pub trait RlLobBackend {
|
||||
ask_px_src: u64,
|
||||
ask_sz_src: u64,
|
||||
) -> Result<()>;
|
||||
|
||||
/// Extract all raw device pointers and kernel function handles needed
|
||||
/// for mega-graph capture. Returns a [`LobSimRawPtrs`] struct that the
|
||||
/// trainer caches to bypass trait dispatch inside the captured section.
|
||||
/// All pointers are stable for the lobsim's lifetime.
|
||||
fn raw_ptrs(&self) -> LobSimRawPtrs;
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -10,6 +10,7 @@
|
||||
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;
|
||||
@@ -121,9 +122,155 @@ impl AdamW {
|
||||
&mut self.v
|
||||
}
|
||||
|
||||
/// Mega-graph variant of `step()`. Reads LR from an ISV device
|
||||
/// pointer at `lr_slot` instead of the host-side `self.lr` field.
|
||||
/// This allows the LR to vary across graph replays because the ISV
|
||||
/// is modified in-place by the `rl_lr_from_mapped_pinned` controller
|
||||
/// kernel captured earlier in the same graph.
|
||||
///
|
||||
/// `isv_lr_fn` is the `adamw_step_isv_lr` kernel function handle.
|
||||
/// `isv_ptr` is the stable ISV device pointer.
|
||||
pub fn step_isv_lr(
|
||||
&mut self,
|
||||
theta: &mut CudaSlice<f32>,
|
||||
grad: &CudaSlice<f32>,
|
||||
isv_lr_fn: cudarc::driver::sys::CUfunction,
|
||||
isv_ptr: u64,
|
||||
lr_slot: i32,
|
||||
) -> Result<()> {
|
||||
let n = theta.len();
|
||||
assert_eq!(grad.len(), n, "grad/theta size mismatch");
|
||||
assert_eq!(self.m.len(), n, "m size mismatch");
|
||||
assert_eq!(self.v.len(), n, "v size mismatch");
|
||||
|
||||
self.step_count_host += 1;
|
||||
|
||||
let n_params_i = n as i32;
|
||||
let grid_x = (n as u32).div_ceil(256);
|
||||
{
|
||||
let mut args = RawArgs::new();
|
||||
args.push_ptr(theta.raw_ptr());
|
||||
args.push_ptr(grad.raw_ptr());
|
||||
args.push_ptr(self.m.raw_ptr());
|
||||
args.push_ptr(self.v.raw_ptr());
|
||||
args.push_i32(n_params_i);
|
||||
args.push_ptr(isv_ptr);
|
||||
args.push_i32(lr_slot);
|
||||
args.push_f32(self.beta1);
|
||||
args.push_f32(self.beta2);
|
||||
args.push_f32(self.eps);
|
||||
args.push_f32(self.wd);
|
||||
args.push_i32(self.step_count_host);
|
||||
let mut ptrs = args.build_arg_ptrs();
|
||||
unsafe {
|
||||
raw_launch(
|
||||
isv_lr_fn,
|
||||
(grid_x, 1, 1), (256, 1, 1), 0,
|
||||
self.raw_stream,
|
||||
&mut ptrs[..args.len()],
|
||||
).map_err(|e| anyhow::anyhow!("adamw_step_isv_lr launch: {:?}", e))?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Return the current step counter value. No GPU sync needed —
|
||||
/// the counter is host-resident.
|
||||
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(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -718,6 +718,13 @@ pub struct PerceptionTrainer {
|
||||
/// include cuBLAS calls.
|
||||
cublas_warmed: bool,
|
||||
|
||||
/// Mega-graph mode: when true, all perception sub-graph state
|
||||
/// machines (train_graph, forward_graph, forward_graph_no_scatter)
|
||||
/// run eagerly (no sub-capture, no sub-replay). The mega-graph in
|
||||
/// the integrated trainer captures the entire pipeline including
|
||||
/// perception kernels.
|
||||
pub mega_graph_enabled: bool,
|
||||
|
||||
/// Event recorded after every training graph launch in
|
||||
/// `step_batched_from_device`. The integrated trainer syncs on
|
||||
/// this event at the start of the NEXT step
|
||||
@@ -2316,6 +2323,7 @@ impl PerceptionTrainer {
|
||||
forward_graph_no_scatter: None,
|
||||
forward_no_scatter_warmed: false,
|
||||
cublas_warmed: false,
|
||||
mega_graph_enabled: false,
|
||||
training_done_event,
|
||||
|
||||
// AoS staging — single mapped-pinned buffer for B*K Mbp10RawInput.
|
||||
@@ -3235,7 +3243,7 @@ impl PerceptionTrainer {
|
||||
// replay (third+). The captured graph records all
|
||||
// in-graph kernel decisions at capture time per
|
||||
// `pearl_no_host_branches_in_captured_graph`.
|
||||
if self.train_graph.is_some() {
|
||||
if self.train_graph.is_some() && !self.mega_graph_enabled {
|
||||
self.train_graph
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
@@ -3245,6 +3253,9 @@ impl PerceptionTrainer {
|
||||
self.dispatch_train_step(b_sz, k_seq, total_snaps)
|
||||
.context("train warmup dispatch")?;
|
||||
self.cublas_warmed = true;
|
||||
} else if self.mega_graph_enabled {
|
||||
self.dispatch_train_step(b_sz, k_seq, total_snaps)
|
||||
.context("train eager dispatch (mega-graph mode)")?;
|
||||
} else {
|
||||
// Event tracking was disabled at trainer construction; the
|
||||
// trainer's CudaSlices have no read/write events, so neither
|
||||
@@ -3886,7 +3897,7 @@ impl PerceptionTrainer {
|
||||
// Three-state machine for the training graph — identical to
|
||||
// step_batched but dispatches `dispatch_train_step_no_scatter`
|
||||
// (skips the AoS→SoA scatter since SoA buffers are pre-filled).
|
||||
if self.train_graph.is_some() {
|
||||
if self.train_graph.is_some() && !self.mega_graph_enabled {
|
||||
self.train_graph
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
@@ -3896,6 +3907,11 @@ impl PerceptionTrainer {
|
||||
self.dispatch_train_step_no_scatter(b_sz, k_seq, total_snaps)
|
||||
.context("train warmup dispatch (from_device)")?;
|
||||
self.cublas_warmed = true;
|
||||
} else if self.mega_graph_enabled {
|
||||
// Mega-graph mode: always dispatch eagerly — the mega-graph
|
||||
// in the integrated trainer captures these kernel launches.
|
||||
self.dispatch_train_step_no_scatter(b_sz, k_seq, total_snaps)
|
||||
.context("train eager dispatch (mega-graph mode)")?;
|
||||
} else {
|
||||
let begin = self.stream.begin_capture(
|
||||
CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED,
|
||||
@@ -4090,7 +4106,7 @@ impl PerceptionTrainer {
|
||||
// warmup → capture → replay with a new graph that excludes the
|
||||
// scatter. Let's use this approach since it's zero overhead
|
||||
// after the second call.
|
||||
if self.forward_graph_no_scatter.is_some() {
|
||||
if self.forward_graph_no_scatter.is_some() && !self.mega_graph_enabled {
|
||||
self.forward_graph_no_scatter
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
@@ -4100,6 +4116,10 @@ impl PerceptionTrainer {
|
||||
self.dispatch_forward_kernels_no_scatter(b_sz, k_seq, total_snaps)
|
||||
.context("forward_no_scatter warmup dispatch")?;
|
||||
self.forward_no_scatter_warmed = true;
|
||||
} else if self.mega_graph_enabled {
|
||||
// Mega-graph mode: always dispatch eagerly.
|
||||
self.dispatch_forward_kernels_no_scatter(b_sz, k_seq, total_snaps)
|
||||
.context("forward_no_scatter eager dispatch (mega-graph mode)")?;
|
||||
} else {
|
||||
use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
|
||||
let begin = self.stream.begin_capture(
|
||||
|
||||
@@ -1143,7 +1143,6 @@ impl LobSimCuda {
|
||||
.arg(&n)
|
||||
.launch(cfg)?;
|
||||
}
|
||||
self.stream.synchronize()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -1434,6 +1433,43 @@ impl ml_alpha::rl::reward::RlLobBackend for LobSimCuda {
|
||||
ask_sz_src,
|
||||
)
|
||||
}
|
||||
|
||||
fn raw_ptrs(&self) -> ml_alpha::rl::reward::LobSimRawPtrs {
|
||||
ml_alpha::rl::reward::LobSimRawPtrs {
|
||||
bid_px: self.bid_px_d.raw_ptr(),
|
||||
bid_sz: self.bid_sz_d.raw_ptr(),
|
||||
ask_px: self.ask_px_d.raw_ptr(),
|
||||
ask_sz: self.ask_sz_d.raw_ptr(),
|
||||
books: self.books_d.raw_ptr(),
|
||||
prev_mid: self.prev_mid_d.raw_ptr(),
|
||||
atr_mid_ema: self.atr_mid_ema_d.raw_ptr(),
|
||||
snapshots_skipped: self.snapshots_skipped_d.raw_ptr(),
|
||||
min_reasonable_px: self.min_reasonable_px_d.raw_ptr(),
|
||||
max_reasonable_px: self.max_reasonable_px_d.raw_ptr(),
|
||||
book_update_fn: self.book_update_fn.cu_function(),
|
||||
market_targets: self.market_targets_d.raw_ptr(),
|
||||
pos: self.pos_d.raw_ptr(),
|
||||
cost_per_lot_per_side: self.cost_per_lot_per_side_d.raw_ptr(),
|
||||
total_fees_per_b: self.total_fees_per_b_d.raw_ptr(),
|
||||
submit_market_fn: self.submit_market_fn.cu_function(),
|
||||
open_trade_state: self.open_trade_state_d.raw_ptr(),
|
||||
trade_log: self.trade_log_d.raw_ptr(),
|
||||
trade_log_head: self.trade_log_head_d.raw_ptr(),
|
||||
trail_hwm: self.trail_hwm_d.raw_ptr(),
|
||||
zero_vwap_at_open: self.zero_vwap_at_open_d.raw_ptr(),
|
||||
saturated_vwap_at_open: self.saturated_vwap_at_open_d.raw_ptr(),
|
||||
defensive_exit_clamp: self.defensive_exit_clamp_d.raw_ptr(),
|
||||
conv_signed_ema: self.conv_signed_ema_d.raw_ptr(),
|
||||
diag_hold_hist: self.diag_hold_hist_d.raw_ptr(),
|
||||
diag_outcome_n: self.diag_outcome_n_d.raw_ptr(),
|
||||
diag_outcome_sum_pnl: self.diag_outcome_sum_pnl_d.raw_ptr(),
|
||||
diag_outcome_n_wins: self.diag_outcome_n_wins_d.raw_ptr(),
|
||||
pnl_track_fn: self.pnl_track_fn.cu_function(),
|
||||
n_backtests: self.n_backtests as i32,
|
||||
pos_bytes: std::mem::size_of::<crate::lob::PosFlat>() as i32,
|
||||
trade_log_cap: crate::lob::TRADE_LOG_CAP as i32,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Re-open the impl block for `LobSimCuda` so subsequent methods (if any
|
||||
|
||||
@@ -0,0 +1,333 @@
|
||||
# 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%
|
||||
@@ -0,0 +1,152 @@
|
||||
# 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)
|
||||
129
docs/superpowers/specs/2026-05-27-mega-graph-pipeline.md
Normal file
129
docs/superpowers/specs/2026-05-27-mega-graph-pipeline.md
Normal file
@@ -0,0 +1,129 @@
|
||||
# Mega-Graph CUDA Pipeline
|
||||
|
||||
**Date**: 2026-05-27
|
||||
**Status**: Draft
|
||||
**Goal**: Capture the entire per-step training pipeline into one CUDA graph launch, eliminating ~150 individual kernel launches and ~140ms of host-side overhead. Target: 50+ sps (from 6.4).
|
||||
|
||||
## Problem
|
||||
|
||||
nsys profiling shows GPU utilization at 8.5% (13ms kernel time per 156ms step). The host spends 143ms between kernel launches on: Rust arg assembly, Result unwrapping, conditional logic, 5 graph launches + ~20 individual raw_launch calls. The CUDA APIs themselves are fast (<1ms total); the overhead is Rust code between them.
|
||||
|
||||
## Current Per-Step Pipeline
|
||||
|
||||
```
|
||||
HOST GPU (13ms total)
|
||||
───────────────────────────── ────────────────────────────
|
||||
build 3 gather args (50μs) → sample_and_gather (13ms)
|
||||
gather_next (0.04ms)
|
||||
gather_current (0.04ms)
|
||||
graph_launch prefill (50μs) → GRAPH A: encoder + Q/V/π + actions (~2ms)
|
||||
build 4 gate args (100μs) → conf_gate + frd_gate + action_ent + log_pi (0.1ms)
|
||||
lobsim.apply_snapshot (20μs) → book_update (0.01ms)
|
||||
graph_launch postfill (50μs) → GRAPH A2: risk + trail + actions_to_market (~1ms)
|
||||
lobsim.step_fill (20μs) → fill + pnl_track (0.05ms)
|
||||
graph_launch reward (50μs) → GRAPH C: reward + scale + clamp + context (~2ms)
|
||||
build 2 PER args (20μs) → per_push_ring + per_push_flush (0.05ms)
|
||||
perception.step_batched (10μs)→ PERCEPTION GRAPH: encoder train (~2ms)
|
||||
graph_launch training (50μs) → GRAPH D: Q/π/V backward + Adam (~3ms)
|
||||
build 4 target args (50μs) → soft_update + q_divergence (0.1ms)
|
||||
DiagFrame memcpy (100μs) → (async, background thread)
|
||||
───────────────────────────
|
||||
Total host: ~570μs arg work
|
||||
But Rust overhead: ~155ms (!)
|
||||
```
|
||||
|
||||
The 570μs of pure arg work can't explain 155ms. The remaining ~154ms is:
|
||||
- `step_synthetic()` → reads ISV from mapped-pinned → launches LR controller → builds training graph args → launches GRAPH D → reads ISV for LR plateau
|
||||
- `step_with_lobsim_reward_and_train()` → builds reward graph args → conditional K-loop → event-based cross-stream sync → PER sample on train_stream
|
||||
|
||||
Each of these involves deep Rust function call chains (integrated.rs is 8400 lines) with Result unwrapping, debug_asserts, conditional branches, and method dispatch through trait objects.
|
||||
|
||||
## Approaches
|
||||
|
||||
### Approach A: Mega-Graph (Recommended)
|
||||
|
||||
Capture ALL kernels from a single step into ONE graph during warmup. On subsequent steps, launch the mega-graph with a single `cuGraphLaunch` call. Zero host work between kernels.
|
||||
|
||||
**Requirements**:
|
||||
1. ALL kernel arguments use stable device pointers (pre-allocated at init)
|
||||
2. No host-side conditional logic during the captured section
|
||||
3. No memory allocation/deallocation during capture
|
||||
4. K-loop fixed at K=1 (current config)
|
||||
5. cuBLAS workspace pre-allocated (already done)
|
||||
|
||||
**What changes per step** (must be handled):
|
||||
- `sample_and_gather` uses per-step PRNG state → device-side, graph-safe
|
||||
- Perception's `ts_ns` value → write to mapped-pinned staging, kernel reads from stable pointer
|
||||
- ISV bus → modified by controller kernels in-place, graph-safe
|
||||
- LR controller → reads host-side loss values → **BLOCKER**: needs mapped-pinned reads
|
||||
|
||||
**LR controller blocker**: The host reads `last_q_loss`, `last_pi_loss`, `last_v_loss` from mapped-pinned memory between the reward graph and the training graph. These feed the per-head LR scaling. Fix: move LR controller to a GPU kernel that reads from mapped-pinned device pointers directly. The host doesn't need to see the LR values — the Adam kernel reads LR from ISV.
|
||||
|
||||
**Lobsim blocker**: `apply_snapshot_from_device` and `step_fill_from_market_targets` are lobsim methods with internal state. These need to be either: (a) converted to raw_launch with cached pointers, or (b) kept outside the mega-graph as a "lobsim mini-graph."
|
||||
|
||||
**Estimated result**: 1 graph launch per step. Host does: graph launch (50μs) + DiagFrame memcpy (100μs) = 150μs/step → **~6600 sps** (GPU-bound at 13ms/step → ~77 sps with GPU saturation).
|
||||
|
||||
### Approach B: Coalesce Existing Graphs
|
||||
|
||||
Merge the 4 existing graphs + loose kernels into 2 larger graphs: "env_graph" (A + gates + A2 + fill + C) and "train_graph" (PER + perception + D + target).
|
||||
|
||||
**Simpler**: doesn't require solving the LR controller blocker or lobsim abstraction. Just expand the existing capture boundaries.
|
||||
|
||||
**Host work**: 2 graph launches + cross-stream event + DiagFrame = ~300μs. But the Rust overhead between the two launches is still significant (K-loop, LR reads, etc.).
|
||||
|
||||
**Estimated result**: ~20-30 sps (2× improvement, not 10×).
|
||||
|
||||
### Approach C: Async Host Pipeline
|
||||
|
||||
Keep the current graph structure but pipeline: host launches step N+1's gathers while step N's training graph runs. Double-buffer all state.
|
||||
|
||||
**Complex**: requires double-buffering 50+ device buffers. Weight updates from step N must be visible to step N+1's forward. Correctness is hard to verify.
|
||||
|
||||
**Estimated result**: ~12 sps (2× from pipelining).
|
||||
|
||||
## Recommendation: Approach A (Mega-Graph)
|
||||
|
||||
The only approach that reaches 50+ sps. The blockers are solvable:
|
||||
|
||||
1. **LR controller → GPU kernel**: already ISV-driven; just move the host-side `AdamW.lr` mutation to a kernel that writes the LR value into the Adam kernel's arg buffer. ~30 lines.
|
||||
|
||||
2. **Lobsim → raw_launch**: `apply_snapshot_from_device` is 2 DtoD copies + 1 kernel. `step_fill_from_market_targets` is 2 kernels. Convert both to raw_launch with cached pointers. ~100 lines.
|
||||
|
||||
3. **K-loop fixed at K=1**: already the case in production config. The mega-graph captures one iteration.
|
||||
|
||||
4. **Cross-stream PER**: PER sample runs on train_stream, but at K=1 it's just 1 sample per step. Can be moved to the main stream (the separate stream was for K>1 pipelining).
|
||||
|
||||
## Implementation Phases
|
||||
|
||||
### Phase 1: Convert lobsim methods to raw_launch (prerequisite)
|
||||
- `apply_snapshot_from_device` → 2 DtoD + 1 kernel via raw_launch
|
||||
- `step_fill_from_market_targets` → 2 kernels via raw_launch
|
||||
- Cache all lobsim internal pointers in LobPtrs
|
||||
|
||||
### Phase 2: Move LR controller to GPU kernel
|
||||
- New `rl_lr_from_mapped_pinned.cu` kernel reads loss from mapped-pinned device pointers
|
||||
- Writes per-head LR to ISV slots
|
||||
- Adam kernel reads LR from ISV (already supported)
|
||||
|
||||
### Phase 3: Move PER sample to main stream
|
||||
- At K=1, PER sample + priority_update + tree_rebuild run once per step
|
||||
- Move from train_stream to self.stream
|
||||
- Eliminate cross-stream events
|
||||
|
||||
### Phase 4: Mega-graph capture
|
||||
- Warmup step 0: eager execution (all kernels launch individually)
|
||||
- Warmup step 1: `begin_capture` → entire pipeline → `end_capture`
|
||||
- Step 2+: single `cuGraphLaunch` per step
|
||||
- Host per-step: ts_ns write to mapped-pinned + graph launch + DiagFrame
|
||||
|
||||
### Phase 5: Validation
|
||||
- nsys: confirm 1 cuGraphLaunch per step, GPU utilization >80%
|
||||
- Functional: wr/pnl match baseline within 1%
|
||||
- Smoke: 1k steps local, 5k steps L40S
|
||||
|
||||
## Anti-patterns to avoid
|
||||
|
||||
- No `cuGraphExecUpdate` (overkill — our topology is fixed)
|
||||
- No conditional graph nodes (CUDA 12.4 feature, complex, fragile)
|
||||
- No multi-stream graphs (correctness nightmare)
|
||||
- No host reads inside the captured section (use mapped-pinned reads OUTSIDE or defer)
|
||||
@@ -158,6 +158,7 @@ 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(false),
|
||||
has_async_alloc: AtomicBool::new(true),
|
||||
num_streams: AtomicUsize::new(0),
|
||||
event_tracking: AtomicBool::new(true),
|
||||
error_state: AtomicU32::new(0),
|
||||
|
||||
Reference in New Issue
Block a user