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

feat(dqn): Branching DQN + KernelWeightPack + metrics propagation

Browse Source 
Branching DQN (Tavakoli et al., 2018):
- 3 independent advantage heads (exposure=5, order=3, urgency=3) in CUDA kernel
- `use_branching` as hyperopt-tunable parameter (index 11 in 29D space)
- Branch weight pointers packed into KernelWeightPack struct

KernelWeightPack refactor:
- 87→38 CUDA kernel params by packing 48 weight pointers into 384-byte #[repr(C)] struct
- UNPACK_WEIGHT_PTRS macro in CUDA header for clean kernel-side access
- Single .arg(&weight_pack) replaces 48 individual .arg() calls

Hyperopt metrics in JSON output:
- Added `metrics: Option<serde_json::Value>` to TrialResult<P>
- `extract_metrics()` trait method with default None (DQN overrides)
- JSON output now includes per-trial backtest metrics (Sharpe, Sortino,
  Calmar, Omega, drawdown, win_rate, trades) + top-level best_metrics

Prometheus backtest gauges (Rust-side, no CUDA):
- 8 new gauges: foxhunt_hyperopt_best_{sharpe,sortino,calmar,omega,
  max_drawdown_pct,win_rate,total_trades,total_return_pct}

Dynamic episode scaling:
- Removed hardcoded 256 episode cap on small GPUs
- VRAM budget calculation in optimal_n_episodes() handles scaling naturally
- Removed stale .min(256) in PPO trainer

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-11 00:46:53 +01:00
parent 9278561ec5
commit a5ea8713b6
29 changed files with 1510 additions and 307 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
crates/common/src/metrics/training_metrics.rs
View File

@@ -578,6 +578,29 @@ pub fn set_hyperopt_elapsed(model: &str, secs: f64) {
set_gauge_vec("foxhunt_hyperopt_elapsed_seconds", &[model], secs);
}
/// Push per-trial backtest metrics to Prometheus (Rust-side only, no CUDA).
#[allow(clippy::too_many_arguments)]
pub fn set_hyperopt_backtest_metrics(
model: &str,
sharpe: f64,
sortino: f64,
max_drawdown_pct: f64,
win_rate: f64,
total_trades: f64,
total_return_pct: f64,
calmar: f64,
omega: f64,
) {
set_gauge_vec("foxhunt_hyperopt_best_sharpe", &[model], sharpe);
set_gauge_vec("foxhunt_hyperopt_best_sortino", &[model], sortino);
set_gauge_vec("foxhunt_hyperopt_best_max_drawdown_pct", &[model], max_drawdown_pct);
set_gauge_vec("foxhunt_hyperopt_best_win_rate", &[model], win_rate);
set_gauge_vec("foxhunt_hyperopt_best_total_trades", &[model], total_trades);
set_gauge_vec("foxhunt_hyperopt_best_total_return_pct", &[model], total_return_pct);
set_gauge_vec("foxhunt_hyperopt_best_calmar", &[model], calmar);
set_gauge_vec("foxhunt_hyperopt_best_omega", &[model], omega);
}
// ---------------------------------------------------------------------------
// Tier 2: Verbose metrics (no-op when FOXHUNT_VERBOSE_METRICS not set)
// ---------------------------------------------------------------------------

56
crates/ml-core/src/memory_optimization/auto_batch_size.rs
View File

@@ -485,15 +485,32 @@ impl GpuHardwareInfo {
};
let optimal = sm_target.min(vram_cap).clamp(MIN_EPISODES, MAX_EPISODES);
// On small GPUs (≤8 GB), the experience collector output tensors compete
// with model backward graph, candle block allocator overhead, and GPU PER
// buffers. Hard-cap to 256 episodes to keep output under ~25 MB.
let small_gpu_cap = if self.free_memory_mb <= 8192.0 { 256 } else { MAX_EPISODES };
let capped = optimal.min(small_gpu_cap);
// Round up to nearest multiple of 256 for clean block scheduling
let aligned = ((capped + 255) / 256) * 256;
// Round up to nearest multiple of 256 for clean block scheduling.
// The VRAM budget cap above already scales dynamically with free memory,
// so no additional hard-cap is needed for small GPUs.
let aligned = ((optimal + 255) / 256) * 256;
aligned.min(MAX_EPISODES)
}
/// Calculate VRAM-optimal replay buffer capacity.
///
/// Allocates `vram_fraction` of free GPU memory to the replay buffer.
/// Each transition stores: 2 × state_dim × f32 (state + next_state)
/// + 4 scalars (action, reward, done, priority) = (2 * state_dim + 4) * 4 bytes.
///
/// Enforces guardrails:
/// - MIN_REPLAY_CAPACITY = 100_000 (below this, sample efficiency degrades)
/// - MAX_REPLAY_CAPACITY = 10_000_000 (above this, priority staleness dominates)
pub fn optimal_replay_capacity(&self, state_dim: usize, vram_fraction: f64) -> usize {
const MIN_REPLAY_CAPACITY: usize = 100_000;
const MAX_REPLAY_CAPACITY: usize = 10_000_000;
let bytes_per_transition = (2 * state_dim + 4) * 4; // f32 each
let budget_bytes = (self.free_memory_mb * vram_fraction * 1024.0 * 1024.0) as usize;
let raw_capacity = budget_bytes / bytes_per_transition.max(1);
raw_capacity.clamp(MIN_REPLAY_CAPACITY, MAX_REPLAY_CAPACITY)
}
}
/// Detect available GPU memory using nvidia-smi
@@ -1047,4 +1064,29 @@ mod tests {
assert!(n >= 128, "Should hit minimum 128 even with low VRAM, got {}", n);
assert!(n <= 1024, "Should be VRAM-capped with only 500MB free, got {}", n);
}
#[test]
fn test_optimal_replay_capacity_h100() {
let hw = GpuHardwareInfo {
device_name: "NVIDIA H100".to_owned(),
total_memory_mb: 81920.0,
free_memory_mb: 72_000.0,
sm_count: 132,
};
let cap = hw.optimal_replay_capacity(56, 0.70);
assert!(cap >= 500_000, "H100 should get at least 500K capacity, got {}", cap);
assert!(cap <= 10_000_000, "Should be capped at max, got {}", cap);
}
#[test]
fn test_optimal_replay_capacity_small_gpu() {
let hw = GpuHardwareInfo {
device_name: "RTX 3050 Ti".to_owned(),
total_memory_mb: 4096.0,
free_memory_mb: 2_000.0,
sm_count: 20,
};
let cap = hw.optimal_replay_capacity(48, 0.70);
assert!(cap >= 100_000, "Small GPU should get at least min capacity");
}
}

12
crates/ml-dqn/src/dqn.rs
View File

@@ -1150,7 +1150,7 @@ pub struct DQN {
/// UCB count-based exploration bonus for action diversity
count_bonus: super::count_bonus::CountBonus,
/// When true, skip monitoring in forward() to avoid GPU->CPU sync in training hot path
/// When true, skip monitoring in `forward()` to avoid GPU->CPU sync in training hot path
training_forward_active: bool,
}
@@ -3713,6 +3713,7 @@ impl DQN {
/// Same divergence detection and early-stopping logic as [`log_q_values`],
/// but takes pre-computed min/max/mean/variance from the GPU reduction
/// kernel instead of downloading the full Q-matrix to CPU.
#[allow(clippy::cognitive_complexity)]
pub fn log_q_values_from_stats(
&mut self,
q_min: f32,
@@ -3772,6 +3773,8 @@ impl DQN {
self.q_value_divergence_counter
);
self.q_value_divergence_counter = 0;
} else {
// Q-values within normal range, no action needed
}
Ok(())
@@ -3849,8 +3852,9 @@ impl DQN {
Ok(())
}
/// Suppress or enable forward() monitoring (GPU→CPU sync avoidance).
/// Set to `true` before forward() calls in the training hot path, `false` after.
/// Suppress or enable `forward()` monitoring (GPU→CPU sync avoidance).
/// Set to `true` before `forward()` calls in the training hot path, `false` after.
#[allow(clippy::missing_const_for_fn)] // intentionally non-const for future mutation logic
pub fn set_training_forward_active(&mut self, active: bool) {
self.training_forward_active = active;
}
@@ -3904,6 +3908,8 @@ impl DQN {
self.gradient_collapse_counter
);
self.gradient_collapse_counter = 0;
} else {
// Gradient norm within normal range, no action needed
}
Ok(())

6
crates/ml-hyperopt/src/optimizer.rs
View File

@@ -493,6 +493,7 @@ impl ArgminOptimizer {
warn!("Trial {} produced non-finite objective ({:?}), using penalty 1e6", trial_num, raw_objective);
1e6
};
let trial_metrics = M::extract_metrics(&metrics);
let duration_secs = start_time.elapsed().as_secs_f64();
info!("✓ Trial {} completed in {:.1}s", trial_num, duration_secs);
@@ -508,6 +509,7 @@ impl ArgminOptimizer {
params,
objective,
duration_secs,
metrics: trial_metrics,
});
}
@@ -1124,6 +1126,7 @@ where
warn!("Trial {} produced non-finite objective ({:?}), using penalty 1e6", trial_num, raw_objective);
1e6
};
let trial_metrics = M::extract_metrics(&metrics);
let duration_secs = start_time.elapsed().as_secs_f64();
info!("✓ Trial {} completed in {:.1}s", trial_num, duration_secs);
@@ -1139,6 +1142,7 @@ where
params,
objective,
duration_secs,
metrics: trial_metrics,
});
}
@@ -1264,6 +1268,7 @@ where
warn!("Trial {} produced non-finite objective ({:?}), using penalty 1e6", trial_num, raw_objective);
1e6
};
let trial_metrics = M::extract_metrics(&metrics);
let duration_secs = start_time.elapsed().as_secs_f64();
info!("Trial {} completed in {:.1}s", trial_num, duration_secs);
@@ -1278,6 +1283,7 @@ where
params,
objective,
duration_secs,
metrics: trial_metrics,
});
}

14
crates/ml-hyperopt/src/traits.rs
View File

@@ -731,6 +731,14 @@ pub trait HyperparameterOptimizable {
/// # }
/// ```
fn extract_objective(metrics: &Self::Metrics) -> f64;
/// Extract per-trial metrics as opaque JSON for reporting.
///
/// Override this to attach domain-specific metrics (Sharpe, drawdown, etc.)
/// to each `TrialResult`. The default returns `None` (no extra metrics).
fn extract_metrics(_metrics: &Self::Metrics) -> Option<serde_json::Value> {
None
}
}
/// Trait for parameter space definition
@@ -941,6 +949,9 @@ pub struct TrialResult<P> {
pub objective: f64,
/// Wall-clock time for training (seconds)
pub duration_secs: f64,
/// Optional domain-specific metrics (Sharpe, drawdown, etc.)
#[serde(skip_serializing_if = "Option::is_none")]
pub metrics: Option<serde_json::Value>,
}
/// Complete optimization result with best parameters and trial history
@@ -1106,6 +1117,7 @@ mod tests {
},
objective: 1.5,
duration_secs: 10.0,
metrics: None,
},
TrialResult {
trial_num: 2,
@@ -1115,6 +1127,7 @@ mod tests {
},
objective: 0.8,
duration_secs: 12.0,
metrics: None,
},
TrialResult {
trial_num: 3,
@@ -1124,6 +1137,7 @@ mod tests {
},
objective: 0.5,
duration_secs: 15.0,
metrics: None,
},
];

56
crates/ml/examples/hyperopt_baseline_rl.rs
View File

@@ -135,14 +135,19 @@ fn build_model_result(
top_k_params: Vec<Value>,
num_trials: usize,
elapsed_secs: f64,
best_metrics: Option<Value>,
) -> Value {
serde_json::json!({
let mut result = serde_json::json!({
"best_objective": best_objective,
"best_params": best_params_json,
"top_k_params": top_k_params,
"trials": num_trials,
"elapsed_secs": elapsed_secs,
})
});
if let Some(m) = best_metrics {
result["best_metrics"] = m;
}
result
}
#[allow(clippy::cognitive_complexity)]
@@ -210,6 +215,25 @@ fn run_dqn_hyperopt(args: &Args, parallel: usize, gpu_devices: &[candle_core::De
training_metrics::set_hyperopt_best_objective("dqn", result.best_objective);
training_metrics::set_hyperopt_elapsed("dqn", elapsed);
// Push best trial's backtest metrics to Prometheus (Rust-side, no CUDA)
if let Some(best) = result.all_trials.iter().min_by(|a, b| a.objective.partial_cmp(&b.objective).unwrap_or(std::cmp::Ordering::Equal)) {
if let Some(ref m) = best.metrics {
if let Some(bt) = m.get("backtest") {
training_metrics::set_hyperopt_backtest_metrics(
"dqn",
bt.get("sharpe_ratio").and_then(|v| v.as_f64()).unwrap_or(0.0),
bt.get("sortino_ratio").and_then(|v| v.as_f64()).unwrap_or(0.0),
bt.get("max_drawdown_pct").and_then(|v| v.as_f64()).unwrap_or(0.0),
bt.get("win_rate").and_then(|v| v.as_f64()).unwrap_or(0.0),
bt.get("total_trades").and_then(|v| v.as_f64()).unwrap_or(0.0),
bt.get("total_return_pct").and_then(|v| v.as_f64()).unwrap_or(0.0),
bt.get("calmar_ratio").and_then(|v| v.as_f64()).unwrap_or(0.0),
bt.get("omega_ratio").and_then(|v| v.as_f64()).unwrap_or(0.0),
);
}
}
}
info!("DQN hyperopt complete:");
info!(" Best objective: {:.6}", result.best_objective);
info!(" Total trials: {}", result.all_trials.len());
@@ -235,21 +259,28 @@ fn run_dqn_hyperopt(args: &Args, parallel: usize, gpu_devices: &[candle_core::De
.take(5)
.filter_map(|t| {
serde_json::to_value(&t.params).ok().map(|params| {
serde_json::json!({
let mut entry = serde_json::json!({
"params": params,
"objective": t.objective,
"trial_num": t.trial_num,
})
});
if let Some(ref m) = t.metrics {
entry["metrics"] = m.clone();
}
entry
})
})
.collect();
let best_metrics = sorted_trials.first().and_then(|t| t.metrics.clone());
Ok(build_model_result(
result.best_objective,
best_params_json,
top_k,
result.all_trials.len(),
elapsed,
best_metrics,
))
}
@@ -338,21 +369,28 @@ fn run_ppo_hyperopt(args: &Args, parallel: usize, gpu_devices: &[candle_core::De
.take(5)
.filter_map(|t| {
serde_json::to_value(&t.params).ok().map(|params| {
serde_json::json!({
let mut entry = serde_json::json!({
"params": params,
"objective": t.objective,
"trial_num": t.trial_num,
})
});
if let Some(ref m) = t.metrics {
entry["metrics"] = m.clone();
}
entry
})
})
.collect();
let best_metrics = sorted_trials.first().and_then(|t| t.metrics.clone());
Ok(build_model_result(
result.best_objective,
best_params_json,
top_k,
result.all_trials.len(),
elapsed,
best_metrics,
))
}
@@ -372,8 +410,12 @@ fn main() -> Result<()> {
training_metrics::set_active_workers(1.0);
// Pre-allocate CUBLAS workspace for deterministic + faster tensor core ops.
// Enable TF32 for all FP32 matmuls — ~8x throughput on H100 tensor cores.
// SAFETY: called once at startup before any multi-threading or CUDA work begins.
unsafe { std::env::set_var("CUBLAS_WORKSPACE_CONFIG", ":4096:8"); }
unsafe {
std::env::set_var("CUBLAS_WORKSPACE_CONFIG", ":4096:8");
std::env::set_var("NVIDIA_TF32_OVERRIDE", "1");
}
let args = Args::parse();

13
crates/ml/examples/train_baseline_rl.rs
View File

@@ -167,6 +167,10 @@ struct Args {
/// Use this to generate datasets for offline training.
#[arg(long)]
collect_dataset: Option<PathBuf>,
/// Disable Branching DQN (3-head: exposure, order, urgency). Enabled by default.
#[arg(long)]
no_branching: bool,
}
// ---------------------------------------------------------------------------
@@ -469,6 +473,7 @@ fn train_dqn_fold(
trades_data_dir: args.trades_data_dir.as_ref().map(|p| p.to_string_lossy().into_owned()),
offline_mode: args.offline,
dataset_path: args.dataset_path.as_ref().map(|p| p.to_string_lossy().into_owned()),
use_branching: !args.no_branching,
..DQNHyperparameters::default()
};
@@ -1030,6 +1035,14 @@ fn main() -> Result<()> {
eprintln!("Observability init failed (non-fatal): {e}");
}
// Pre-allocate CUBLAS workspace for deterministic + faster tensor core ops.
// Enable TF32 for all FP32 matmuls — ~8x throughput on H100 tensor cores.
// SAFETY: called once at startup before any multi-threading or CUDA work begins.
unsafe {
std::env::set_var("CUBLAS_WORKSPACE_CONFIG", ":4096:8");
std::env::set_var("NVIDIA_TF32_OVERRIDE", "1");
}
metrics::init();
metrics_server::start_metrics_server(9094);
common::metrics::questdb_sink::init(None);

4
crates/ml/hyperparams/dqn_best.toml
View File

@@ -22,6 +22,10 @@ epsilon_decay = 0.995 # Decay rate per episode
buffer_size = 100000 # Experience replay capacity
min_replay_size = 1000 # Minimum experiences before training
# Branching DQN (Tavakoli et al., 2018)
use_branching = true
branch_hidden_dim = 128
# Training Configuration
epochs = 100 # Training epochs (default)
checkpoint_frequency = 10 # Save every 10 epochs

128
crates/ml/src/cuda_pipeline/common_device_functions.cuh
View File

@@ -25,6 +25,14 @@
#define PPO_NUM_ACTIONS 45
#define NUM_ACTIONS DQN_NUM_ACTIONS /* default for DQN kernels */
/* Branching DQN head sizes (Tavakoli et al., 2018).
* 3 independent advantage heads: exposure (5), order (3), urgency (3).
* Factored action index = exposure * 9 + order * 3 + urgency (0-44). */
#define BRANCH_EXPOSURE_ACTIONS 5
#define BRANCH_ORDER_ACTIONS 3
#define BRANCH_URGENCY_ACTIONS 3
#define BRANCH_MAX_ACTIONS BRANCH_EXPOSURE_ACTIONS /* largest branch */
/* Curiosity forward model */
#define CUR_INPUT 35 /* 32 state features + 3-class one-hot */
#define CUR_HIDDEN 64
@@ -550,12 +558,12 @@ __device__ __forceinline__ float action_to_tx_cost(int action_idx) {
}
/**
* PPO: Map factored action index (0-44) to target exposure fraction.
* Map factored action index (0-44) to target exposure fraction.
*
* PPO uses the full 45-action factored space (5 exposures x 3 order x 3 urgency).
* Both PPO and Branching DQN use the 45-action factored space (5 exposures x 3 order x 3 urgency).
* Exposure = action_index / 9.
*/
__device__ __forceinline__ float ppo_action_to_exposure(int action_idx) {
__device__ __forceinline__ float factored_action_to_exposure(int action_idx) {
int exposure_level = action_idx / 9; /* 0-4 */
switch (exposure_level) {
case 0: return -1.0f; /* Short100 */
@@ -995,6 +1003,120 @@ __device__ __forceinline__ float splitmix64_fill(int step, int action_idx) {
return (float)d;
}
/* ------------------------------------------------------------------ */
/* Kernel Weight Pack — reduces ~48 weight pointer params to 1 struct */
/* ------------------------------------------------------------------ */
/**
* Packed weight pointer structs for the DQN experience collection kernel.
*
* Passed by value in CUDA constant memory (384 bytes total, well within
* the 4 KB per-kernel limit). Each field is a device pointer to a
* pre-uploaded weight buffer on GPU.
*
* Layout must match the Rust #[repr(C)] structs in gpu_weights.rs EXACTLY.
*/
struct DuelingWeightPtrs {
const float* w_s1; const float* b_s1; /* shared layer 0 */
const float* w_s2; const float* b_s2; /* shared layer 1 */
const float* w_v1; const float* b_v1; /* value FC */
const float* w_v2; const float* b_v2; /* value out */
const float* w_a1; const float* b_a1; /* advantage FC */
const float* w_a2; const float* b_a2; /* advantage out */
};
struct CuriosityWeightPtrs {
const float* w1; const float* b1; /* FC layer 1 */
const float* w2; const float* b2; /* FC layer 2 */
};
struct RMSNormWeightPtrs {
const float* s0; /* shared_rmsnorm_0 gamma */
const float* s1; /* shared_rmsnorm_1 gamma */
const float* v; /* value_rmsnorm gamma */
const float* a; /* advantage_rmsnorm gamma */
};
struct BranchWeightPtrs {
const float* w_o1; const float* b_o1; /* order FC */
const float* w_o2; const float* b_o2; /* order out */
const float* w_u1; const float* b_u1; /* urgency FC */
const float* w_u2; const float* b_u2; /* urgency out */
};
struct KernelWeightPack {
struct DuelingWeightPtrs online; /* 12 ptrs = 96 B */
struct DuelingWeightPtrs target; /* 12 ptrs = 96 B */
struct CuriosityWeightPtrs curiosity; /* 4 ptrs = 32 B */
struct RMSNormWeightPtrs rmsnorm; /* 4 ptrs = 32 B */
struct BranchWeightPtrs online_br; /* 8 ptrs = 64 B */
struct BranchWeightPtrs target_br; /* 8 ptrs = 64 B */
}; /* total = 384 B */
/**
* Unpack KernelWeightPack into local __restrict__ aliases matching the
* variable names used by q_forward_* helper functions.
*
* Call at the top of each kernel body after receiving `struct KernelWeightPack wp`.
* The __restrict__ qualifier enables the compiler to assume no aliasing.
*/
#define UNPACK_WEIGHT_PTRS(wp) \
/* Online dueling (12) */ \
const float* __restrict__ on_w_s1 = wp.online.w_s1; \
const float* __restrict__ on_b_s1 = wp.online.b_s1; \
const float* __restrict__ on_w_s2 = wp.online.w_s2; \
const float* __restrict__ on_b_s2 = wp.online.b_s2; \
const float* __restrict__ on_w_v1 = wp.online.w_v1; \
const float* __restrict__ on_b_v1 = wp.online.b_v1; \
const float* __restrict__ on_w_v2 = wp.online.w_v2; \
const float* __restrict__ on_b_v2 = wp.online.b_v2; \
const float* __restrict__ on_w_a1 = wp.online.w_a1; \
const float* __restrict__ on_b_a1 = wp.online.b_a1; \
const float* __restrict__ on_w_a2 = wp.online.w_a2; \
const float* __restrict__ on_b_a2 = wp.online.b_a2; \
/* Target dueling (12) */ \
const float* __restrict__ tg_w_s1 = wp.target.w_s1; \
const float* __restrict__ tg_b_s1 = wp.target.b_s1; \
const float* __restrict__ tg_w_s2 = wp.target.w_s2; \
const float* __restrict__ tg_b_s2 = wp.target.b_s2; \
const float* __restrict__ tg_w_v1 = wp.target.w_v1; \
const float* __restrict__ tg_b_v1 = wp.target.b_v1; \
const float* __restrict__ tg_w_v2 = wp.target.w_v2; \
const float* __restrict__ tg_b_v2 = wp.target.b_v2; \
const float* __restrict__ tg_w_a1 = wp.target.w_a1; \
const float* __restrict__ tg_b_a1 = wp.target.b_a1; \
const float* __restrict__ tg_w_a2 = wp.target.w_a2; \
const float* __restrict__ tg_b_a2 = wp.target.b_a2; \
/* Curiosity (4) */ \
const float* __restrict__ cur_w1 = wp.curiosity.w1; \
const float* __restrict__ cur_b1 = wp.curiosity.b1; \
const float* __restrict__ cur_w2 = wp.curiosity.w2; \
const float* __restrict__ cur_b2 = wp.curiosity.b2; \
/* RMSNorm (4) */ \
const float* __restrict__ rms_s0_gamma = wp.rmsnorm.s0; \
const float* __restrict__ rms_s1_gamma = wp.rmsnorm.s1; \
const float* __restrict__ rms_v_gamma = wp.rmsnorm.v; \
const float* __restrict__ rms_a_gamma = wp.rmsnorm.a; \
/* Online branching (8) */ \
const float* __restrict__ on_w_bo1 = wp.online_br.w_o1; \
const float* __restrict__ on_b_bo1 = wp.online_br.b_o1; \
const float* __restrict__ on_w_bo2 = wp.online_br.w_o2; \
const float* __restrict__ on_b_bo2 = wp.online_br.b_o2; \
const float* __restrict__ on_w_bu1 = wp.online_br.w_u1; \
const float* __restrict__ on_b_bu1 = wp.online_br.b_u1; \
const float* __restrict__ on_w_bu2 = wp.online_br.w_u2; \
const float* __restrict__ on_b_bu2 = wp.online_br.b_u2; \
/* Target branching (8) */ \
const float* __restrict__ tg_w_bo1 = wp.target_br.w_o1; \
const float* __restrict__ tg_b_bo1 = wp.target_br.b_o1; \
const float* __restrict__ tg_w_bo2 = wp.target_br.w_o2; \
const float* __restrict__ tg_b_bo2 = wp.target_br.b_o2; \
const float* __restrict__ tg_w_bu1 = wp.target_br.w_u1; \
const float* __restrict__ tg_b_bu1 = wp.target_br.b_u1; \
const float* __restrict__ tg_w_bu2 = wp.target_br.w_u2; \
const float* __restrict__ tg_b_bu2 = wp.target_br.b_u2;
/**
* Compute fill probability for a given order type and market conditions.
*

724
crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu
View File

@@ -104,6 +104,121 @@ __device__ void q_forward_dueling(
}
}
/**
* Branching Dueling forward pass — 3 independent advantage heads.
* (Tavakoli et al., 2018: "Action Branching Architectures for Deep RL")
*
* Same shared encoder + value head as standard dueling, but 3 separate
* advantage streams: exposure (5), order (3), urgency (3).
*
* Per-branch Q-values: Q_d(s,a_d) = V(s) + A_d(s,a_d) - mean(A_d)
* Aggregate Q(s,a) = (1/3) * [Q_e(a_e) + Q_o(a_o) + Q_u(a_u)]
*
* Register-based (no shared memory tiling) for initial implementation.
*/
__device__ void q_forward_branching(
const float* state,
/* shared layer weights (4 pointers) */
const float* __restrict__ w_s1,
const float* __restrict__ b_s1,
const float* __restrict__ w_s2,
const float* __restrict__ b_s2,
/* value head weights (4 pointers) */
const float* __restrict__ w_v1,
const float* __restrict__ b_v1,
const float* __restrict__ w_v2,
const float* __restrict__ b_v2,
/* exposure advantage head (4 pointers — reuses standard advantage slot) */
const float* __restrict__ w_ae1,
const float* __restrict__ b_ae1,
const float* __restrict__ w_ae2,
const float* __restrict__ b_ae2,
/* order advantage head (4 pointers) */
const float* __restrict__ w_ao1,
const float* __restrict__ b_ao1,
const float* __restrict__ w_ao2,
const float* __restrict__ b_ao2,
/* urgency advantage head (4 pointers) */
const float* __restrict__ w_au1,
const float* __restrict__ b_au1,
const float* __restrict__ w_au2,
const float* __restrict__ b_au2,
/* scratch buffers */
float* scratch1, /* [SHARED_H1] */
float* scratch2, /* [SHARED_H2] */
float* scratch_v, /* [VALUE_H] */
float* scratch_a, /* [ADV_H] */
/* outputs */
float* q_exposure, /* [BRANCH_EXPOSURE_ACTIONS] */
float* q_order, /* [BRANCH_ORDER_ACTIONS] */
float* q_urgency /* [BRANCH_URGENCY_ACTIONS] */
) {
/* Shared encoder */
matvec_leaky_relu(w_s1, b_s1, state, scratch1, STATE_DIM, SHARED_H1, 1);
matvec_leaky_relu(w_s2, b_s2, scratch1, scratch2, SHARED_H1, SHARED_H2, 1);
/* Value head */
matvec_leaky_relu(w_v1, b_v1, scratch2, scratch_v, SHARED_H2, VALUE_H, 1);
float value;
{
float acc = b_v2[0];
for (int i = 0; i < VALUE_H; i++) acc += w_v2[i] * scratch_v[i];
value = acc;
}
/* Exposure advantage head (branch 0) */
matvec_leaky_relu(w_ae1, b_ae1, scratch2, scratch_a, SHARED_H2, ADV_H, 1);
float adv_e[BRANCH_EXPOSURE_ACTIONS];
matvec_leaky_relu(w_ae2, b_ae2, scratch_a, adv_e, ADV_H, BRANCH_EXPOSURE_ACTIONS, 0);
float mean_e = 0.0f;
for (int i = 0; i < BRANCH_EXPOSURE_ACTIONS; i++) mean_e += adv_e[i];
mean_e /= (float)BRANCH_EXPOSURE_ACTIONS;
for (int i = 0; i < BRANCH_EXPOSURE_ACTIONS; i++)
q_exposure[i] = value + adv_e[i] - mean_e;
/* Order advantage head (branch 1) */
matvec_leaky_relu(w_ao1, b_ao1, scratch2, scratch_a, SHARED_H2, ADV_H, 1);
float adv_o[BRANCH_ORDER_ACTIONS];
matvec_leaky_relu(w_ao2, b_ao2, scratch_a, adv_o, ADV_H, BRANCH_ORDER_ACTIONS, 0);
float mean_o = 0.0f;
for (int i = 0; i < BRANCH_ORDER_ACTIONS; i++) mean_o += adv_o[i];
mean_o /= (float)BRANCH_ORDER_ACTIONS;
for (int i = 0; i < BRANCH_ORDER_ACTIONS; i++)
q_order[i] = value + adv_o[i] - mean_o;
/* Urgency advantage head (branch 2) */
matvec_leaky_relu(w_au1, b_au1, scratch2, scratch_a, SHARED_H2, ADV_H, 1);
float adv_u[BRANCH_URGENCY_ACTIONS];
matvec_leaky_relu(w_au2, b_au2, scratch_a, adv_u, ADV_H, BRANCH_URGENCY_ACTIONS, 0);
float mean_u = 0.0f;
for (int i = 0; i < BRANCH_URGENCY_ACTIONS; i++) mean_u += adv_u[i];
mean_u /= (float)BRANCH_URGENCY_ACTIONS;
for (int i = 0; i < BRANCH_URGENCY_ACTIONS; i++)
q_urgency[i] = value + adv_u[i] - mean_u;
}
/** Argmax over an array of given length. */
__device__ __forceinline__ int argmax_arr(const float* vals, int len) {
int best = 0;
float best_val = vals[0];
for (int i = 1; i < len; i++) {
if (vals[i] > best_val) {
best_val = vals[i];
best = i;
}
}
return best;
}
/** Max value over an array of given length. */
__device__ __forceinline__ float max_arr(const float* vals, int len) {
float m = vals[0];
for (int i = 1; i < len; i++) {
if (vals[i] > m) m = vals[i];
}
return m;
}
/* ------------------------------------------------------------------ */
/* Shared Memory Dueling Forward Passes */
/* ------------------------------------------------------------------ */
@@ -1004,39 +1119,8 @@ extern "C" __global__ void dqn_full_experience_kernel(
/* Episode start indices [N] — global bar index where each episode begins */
const int* __restrict__ episode_starts,
/* ---- Online Q-network weights (12 pointers) ---- */
const float* __restrict__ on_w_s1,
const float* __restrict__ on_b_s1,
const float* __restrict__ on_w_s2,
const float* __restrict__ on_b_s2,
const float* __restrict__ on_w_v1,
const float* __restrict__ on_b_v1,
const float* __restrict__ on_w_v2,
const float* __restrict__ on_b_v2,
const float* __restrict__ on_w_a1,
const float* __restrict__ on_b_a1,
const float* __restrict__ on_w_a2,
const float* __restrict__ on_b_a2,
/* ---- Target Q-network weights (12 pointers) ---- */
const float* __restrict__ tg_w_s1,
const float* __restrict__ tg_b_s1,
const float* __restrict__ tg_w_s2,
const float* __restrict__ tg_b_s2,
const float* __restrict__ tg_w_v1,
const float* __restrict__ tg_b_v1,
const float* __restrict__ tg_w_v2,
const float* __restrict__ tg_b_v2,
const float* __restrict__ tg_w_a1,
const float* __restrict__ tg_b_a1,
const float* __restrict__ tg_w_a2,
const float* __restrict__ tg_b_a2,
/* ---- Curiosity model weights (4 pointers) ---- */
const float* __restrict__ cur_w1,
const float* __restrict__ cur_b1,
const float* __restrict__ cur_w2,
const float* __restrict__ cur_b2,
/* ---- All network weights packed (384 bytes, constant memory) ---- */
struct KernelWeightPack wp,
/* ---- Per-episode mutable state arrays ---- */
float* portfolio_states, /* [N, PORTFOLIO_STATE_SIZE] */
@@ -1079,12 +1163,6 @@ extern "C" __global__ void dqn_full_experience_kernel(
float reward_norm_alpha,
int use_rmsnorm,
/* ---- D6: RMSNorm gamma weights (4 pointers) ---- */
const float* __restrict__ rms_s0_gamma, /* [SHARED_H1] */
const float* __restrict__ rms_s1_gamma, /* [SHARED_H2] */
const float* __restrict__ rms_v_gamma, /* [VALUE_H] */
const float* __restrict__ rms_a_gamma, /* [ADV_H] */
/* ---- Fill simulation config ---- */
float fill_median_spread,
float fill_median_vol,
@@ -1102,6 +1180,9 @@ extern "C" __global__ void dqn_full_experience_kernel(
/* ---- N-step returns config ---- */
int n_steps,
/* ---- Branching DQN config (Tavakoli et al., 2018) ---- */
int use_branching,
/* ---- RNG states [N] ---- */
unsigned int* rng_states,
@@ -1113,6 +1194,9 @@ extern "C" __global__ void dqn_full_experience_kernel(
float* out_target_q, /* [N, L] */
float* out_td_error /* [N, L] */
) {
/* ---- Unpack weight pointers from struct into local __restrict__ aliases ---- */
UNPACK_WEIGHT_PTRS(wp);
/* ---- Shared memory workspace for weight tiling ---- */
extern __shared__ float shmem[];
float* shmem_weights = shmem;
@@ -1168,6 +1252,13 @@ extern "C" __global__ void dqn_full_experience_kernel(
float* scratch_a = scratch1 + VALUE_H; /* reuse: advantage head */
float q_values[NUM_ACTIONS];
float tgt_q_values[NUM_ACTIONS];
/* Branching DQN per-head Q-value arrays (used only when use_branching) */
float br_q_exposure[BRANCH_EXPOSURE_ACTIONS];
float br_q_order[BRANCH_ORDER_ACTIONS];
float br_q_urgency[BRANCH_URGENCY_ACTIONS];
float br_tgt_exposure[BRANCH_EXPOSURE_ACTIONS];
float br_tgt_order[BRANCH_ORDER_ACTIONS];
float br_tgt_urgency[BRANCH_URGENCY_ACTIONS];
float cur_scratch[CUR_HIDDEN];
float state[STATE_DIM];
@@ -1257,7 +1348,19 @@ extern "C" __global__ void dqn_full_experience_kernel(
/* ---- Step 4: Online Q-network forward (shmem tiling) ----
* All threads participate in cooperative loads + __syncthreads().
* Invalid/skipped threads still compute into scratch (harmlessly). */
if (use_distributional && num_atoms > 1) {
if (use_branching) {
/* Branching DQN: 3 independent advantage heads (register-based). */
q_forward_branching(
state,
on_w_s1, on_b_s1, on_w_s2, on_b_s2,
on_w_v1, on_b_v1, on_w_v2, on_b_v2,
on_w_a1, on_b_a1, on_w_a2, on_b_a2, /* exposure = standard adv slot */
on_w_bo1, on_b_bo1, on_w_bo2, on_b_bo2, /* order head */
on_w_bu1, on_b_bu1, on_w_bu2, on_b_bu2, /* urgency head */
scratch1, scratch2, scratch_v, scratch_a,
br_q_exposure, br_q_order, br_q_urgency
);
} else if (use_distributional && num_atoms > 1) {
q_forward_distributional_shmem(
state,
on_w_s1, on_b_s1, on_w_s2, on_b_s2,
@@ -1306,76 +1409,158 @@ extern "C" __global__ void dqn_full_experience_kernel(
float next_state[STATE_DIM];
if (!skip_data) {
/* ---- Step 4b: Q-value clipping (D2) ---- */
for (int i = 0; i < NUM_ACTIONS; i++) {
if (q_values[i] < q_clip_min) q_values[i] = q_clip_min;
if (q_values[i] > q_clip_max) q_values[i] = q_clip_max;
}
/* ---- Step 4b + Step 5: Q-value clipping + action selection ---- */
int br_exposure_sel = 0, br_order_sel = 0, br_urgency_sel = 0;
/* ---- Step 5: Epsilon-greedy + count-bonus action selection ---- */
float r = gpu_random(&rng);
if (r < epsilon) {
/* Random action [0, 4] (5 DQN exposure actions) */
action_idx = (int)(gpu_random(&rng) * (float)DQN_NUM_ACTIONS);
if (action_idx >= DQN_NUM_ACTIONS) action_idx = DQN_NUM_ACTIONS - 1;
} else {
/* D1: Add UCB count bonus before greedy selection. */
if (count_bonus_coefficient > 0.0f && total_action_count > 0) {
float log_n = logf((float)total_action_count);
float q_bonus[NUM_ACTIONS];
for (int i = 0; i < NUM_ACTIONS; i++) {
float bonus = count_bonus_coefficient
* sqrtf(log_n / (1.0f + (float)action_counts[i]));
q_bonus[i] = q_values[i] + bonus;
}
action_idx = argmax_q(q_bonus);
} else {
action_idx = argmax_q(q_values);
if (use_branching) {
/* Branching: clip per-head Q-values */
for (int i = 0; i < BRANCH_EXPOSURE_ACTIONS; i++) {
if (br_q_exposure[i] < q_clip_min) br_q_exposure[i] = q_clip_min;
if (br_q_exposure[i] > q_clip_max) br_q_exposure[i] = q_clip_max;
}
for (int i = 0; i < BRANCH_ORDER_ACTIONS; i++) {
if (br_q_order[i] < q_clip_min) br_q_order[i] = q_clip_min;
if (br_q_order[i] > q_clip_max) br_q_order[i] = q_clip_max;
}
for (int i = 0; i < BRANCH_URGENCY_ACTIONS; i++) {
if (br_q_urgency[i] < q_clip_min) br_q_urgency[i] = q_clip_min;
if (br_q_urgency[i] > q_clip_max) br_q_urgency[i] = q_clip_max;
}
}
/* D1: Update count state */
action_counts[action_idx]++;
total_action_count++;
/* 3 independent epsilon-greedy selections */
/* Exposure head */
if (gpu_random(&rng) < epsilon) {
br_exposure_sel = (int)(gpu_random(&rng) * (float)BRANCH_EXPOSURE_ACTIONS);
if (br_exposure_sel >= BRANCH_EXPOSURE_ACTIONS) br_exposure_sel = BRANCH_EXPOSURE_ACTIONS - 1;
} else {
br_exposure_sel = argmax_arr(br_q_exposure, BRANCH_EXPOSURE_ACTIONS);
}
/* Order head */
if (gpu_random(&rng) < epsilon) {
br_order_sel = (int)(gpu_random(&rng) * (float)BRANCH_ORDER_ACTIONS);
if (br_order_sel >= BRANCH_ORDER_ACTIONS) br_order_sel = BRANCH_ORDER_ACTIONS - 1;
} else {
br_order_sel = argmax_arr(br_q_order, BRANCH_ORDER_ACTIONS);
}
/* Urgency head */
if (gpu_random(&rng) < epsilon) {
br_urgency_sel = (int)(gpu_random(&rng) * (float)BRANCH_URGENCY_ACTIONS);
if (br_urgency_sel >= BRANCH_URGENCY_ACTIONS) br_urgency_sel = BRANCH_URGENCY_ACTIONS - 1;
} else {
br_urgency_sel = argmax_arr(br_q_urgency, BRANCH_URGENCY_ACTIONS);
}
online_q_selected = q_values[action_idx];
/* Compose factored action: exposure * 9 + order * 3 + urgency */
action_idx = br_exposure_sel * 9 + br_order_sel * 3 + br_urgency_sel;
/* ---- Step 5b: Order routing + fill simulation ---- */
/* Route action through spread/vol-dependent order type selection,
* then check fill probability. Unfilled orders override to Flat. */
if (fill_simulation_enabled) {
int order_type, urgency;
route_order(spread, fill_median_spread, 0.0f /* vol from price */, fill_median_vol,
&order_type, &urgency);
/* Aggregate Q = mean of per-branch Q for selected actions (Tavakoli et al.) */
online_q_selected = (br_q_exposure[br_exposure_sel]
+ br_q_order[br_order_sel]
+ br_q_urgency[br_urgency_sel]) / 3.0f;
float norm_vol = (fill_median_vol > 0.0f) ? (spread / fill_median_spread) : 1.0f;
/* Clamp normalized vol to a reasonable range for fill probability */
norm_vol = fminf(fmaxf(norm_vol, 0.0f), 3.0f);
/* D1: Count bonus tracks exposure dimension only (arrays sized NUM_ACTIONS=5) */
action_counts[br_exposure_sel]++;
total_action_count++;
float cost_adj;
int filled = simulate_fill_check(
order_type, urgency, norm_vol,
spread * 10000.0f, /* convert spread fraction to bps */
global_bar, action_idx,
fill_ioc_fill_prob, fill_limit_fill_min, fill_limit_fill_max,
fill_spread_cost_frac, fill_spread_capture_frac,
&cost_adj
);
/* Fill simulation: branching directly selects order type and urgency.
* br_order_sel maps to: 0=Market, 1=LimitMaker, 2=IoC
* br_urgency_sel maps to: 0=Patient, 1=Normal, 2=Aggressive */
if (fill_simulation_enabled) {
float norm_vol = (fill_median_vol > 0.0f) ? (spread / fill_median_spread) : 1.0f;
norm_vol = fminf(fmaxf(norm_vol, 0.0f), 3.0f);
if (!filled) {
action_idx = 2; /* Override to Flat — order didn't fill */
float cost_adj;
int filled = simulate_fill_check(
br_order_sel, br_urgency_sel, norm_vol,
spread * 10000.0f,
global_bar, action_idx,
fill_ioc_fill_prob, fill_limit_fill_min, fill_limit_fill_max,
fill_spread_cost_frac, fill_spread_capture_frac,
&cost_adj
);
if (!filled) {
/* Override exposure to Flat (2), keep order/urgency learned selections */
br_exposure_sel = 2;
action_idx = br_exposure_sel * 9 + br_order_sel * 3 + br_urgency_sel;
}
}
} else {
/* Standard: single advantage head */
for (int i = 0; i < NUM_ACTIONS; i++) {
if (q_values[i] < q_clip_min) q_values[i] = q_clip_min;
if (q_values[i] > q_clip_max) q_values[i] = q_clip_max;
}
float r = gpu_random(&rng);
if (r < epsilon) {
action_idx = (int)(gpu_random(&rng) * (float)DQN_NUM_ACTIONS);
if (action_idx >= DQN_NUM_ACTIONS) action_idx = DQN_NUM_ACTIONS - 1;
} else {
if (count_bonus_coefficient > 0.0f && total_action_count > 0) {
float log_n = logf((float)total_action_count);
float q_bonus[NUM_ACTIONS];
for (int i = 0; i < NUM_ACTIONS; i++) {
float bonus = count_bonus_coefficient
* sqrtf(log_n / (1.0f + (float)action_counts[i]));
q_bonus[i] = q_values[i] + bonus;
}
action_idx = argmax_q(q_bonus);
} else {
action_idx = argmax_q(q_values);
}
}
action_counts[action_idx]++;
total_action_count++;
online_q_selected = q_values[action_idx];
/* ---- Step 5b: Order routing + fill simulation ---- */
if (fill_simulation_enabled) {
int order_type, urgency;
route_order(spread, fill_median_spread, 0.0f, fill_median_vol,
&order_type, &urgency);
float norm_vol = (fill_median_vol > 0.0f) ? (spread / fill_median_spread) : 1.0f;
norm_vol = fminf(fmaxf(norm_vol, 0.0f), 3.0f);
float cost_adj;
int filled = simulate_fill_check(
order_type, urgency, norm_vol,
spread * 10000.0f,
global_bar, action_idx,
fill_ioc_fill_prob, fill_limit_fill_min, fill_limit_fill_max,
fill_spread_cost_frac, fill_spread_capture_frac,
&cost_adj
);
if (!filled) {
action_idx = 2; /* Override to Flat */
}
}
}
out_actions[out_off] = action_idx;
/* ---- Step 6: Portfolio simulation ---- */
float target_exposure = action_to_exposure(action_idx);
float target_exposure;
if (use_branching) {
/* Branching: factored action 0-44, extract exposure via division */
target_exposure = factored_action_to_exposure(action_idx);
} else {
/* Standard: action 0-4 maps directly to exposure */
target_exposure = action_to_exposure(action_idx);
}
float target_position = target_exposure * max_position;
/* Use order-type-dependent tx cost when fill sim enabled,
* otherwise fall back to fixed market rate. */
float tx_rate;
if (fill_simulation_enabled) {
if (use_branching) {
/* Branching: order type directly from order head selection */
float spread_bps = spread * 10000.0f;
tx_rate = fabsf(order_type_tx_cost(br_order_sel, spread_bps,
fill_spread_cost_frac, fill_spread_capture_frac));
tx_rate *= tx_cost_multiplier;
} else if (fill_simulation_enabled) {
int order_type_for_cost, urgency_unused;
route_order(spread, fill_median_spread, 0.0f, fill_median_vol,
&order_type_for_cost, &urgency_unused);
@@ -1492,7 +1677,37 @@ extern "C" __global__ void dqn_full_experience_kernel(
/* ---- Step 11: Target Q-network forward (shmem tiling) ----
* All threads participate in cooperative loads, same as online forward. */
if (use_distributional && num_atoms > 1) {
float max_target_q = 0.0f;
if (use_branching) {
/* Branching target: 3 independent heads, aggregate max */
q_forward_branching(
next_state,
tg_w_s1, tg_b_s1, tg_w_s2, tg_b_s2,
tg_w_v1, tg_b_v1, tg_w_v2, tg_b_v2,
tg_w_a1, tg_b_a1, tg_w_a2, tg_b_a2, /* exposure */
tg_w_bo1, tg_b_bo1, tg_w_bo2, tg_b_bo2, /* order */
tg_w_bu1, tg_b_bu1, tg_w_bu2, tg_b_bu2, /* urgency */
scratch1, scratch2, scratch_v, scratch_a,
br_tgt_exposure, br_tgt_order, br_tgt_urgency
);
/* Clip per-head target Q-values */
for (int i = 0; i < BRANCH_EXPOSURE_ACTIONS; i++) {
if (br_tgt_exposure[i] < q_clip_min) br_tgt_exposure[i] = q_clip_min;
if (br_tgt_exposure[i] > q_clip_max) br_tgt_exposure[i] = q_clip_max;
}
for (int i = 0; i < BRANCH_ORDER_ACTIONS; i++) {
if (br_tgt_order[i] < q_clip_min) br_tgt_order[i] = q_clip_min;
if (br_tgt_order[i] > q_clip_max) br_tgt_order[i] = q_clip_max;
}
for (int i = 0; i < BRANCH_URGENCY_ACTIONS; i++) {
if (br_tgt_urgency[i] < q_clip_min) br_tgt_urgency[i] = q_clip_min;
if (br_tgt_urgency[i] > q_clip_max) br_tgt_urgency[i] = q_clip_max;
}
/* Target max = mean of per-branch maxes (Tavakoli aggregate) */
max_target_q = (max_arr(br_tgt_exposure, BRANCH_EXPOSURE_ACTIONS)
+ max_arr(br_tgt_order, BRANCH_ORDER_ACTIONS)
+ max_arr(br_tgt_urgency, BRANCH_URGENCY_ACTIONS)) / 3.0f;
} else if (use_distributional && num_atoms > 1) {
q_forward_distributional_shmem(
next_state,
tg_w_s1, tg_b_s1, tg_w_s2, tg_b_s2,
@@ -1506,6 +1721,14 @@ extern "C" __global__ void dqn_full_experience_kernel(
&rng,
shmem_weights, shmem_bias
);
for (int i = 0; i < NUM_ACTIONS; i++) {
if (tgt_q_values[i] < q_clip_min) tgt_q_values[i] = q_clip_min;
if (tgt_q_values[i] > q_clip_max) tgt_q_values[i] = q_clip_max;
}
max_target_q = tgt_q_values[0];
for (int i = 1; i < NUM_ACTIONS; i++) {
if (tgt_q_values[i] > max_target_q) max_target_q = tgt_q_values[i];
}
} else {
q_forward_dueling_shmem(
next_state,
@@ -1516,16 +1739,14 @@ extern "C" __global__ void dqn_full_experience_kernel(
tgt_q_values,
shmem_weights, shmem_bias
);
}
/* D2: Clip target Q-values too */
for (int i = 0; i < NUM_ACTIONS; i++) {
if (tgt_q_values[i] < q_clip_min) tgt_q_values[i] = q_clip_min;
if (tgt_q_values[i] > q_clip_max) tgt_q_values[i] = q_clip_max;
}
float max_target_q = tgt_q_values[0];
for (int i = 1; i < NUM_ACTIONS; i++) {
if (tgt_q_values[i] > max_target_q)
max_target_q = tgt_q_values[i];
for (int i = 0; i < NUM_ACTIONS; i++) {
if (tgt_q_values[i] < q_clip_min) tgt_q_values[i] = q_clip_min;
if (tgt_q_values[i] > q_clip_max) tgt_q_values[i] = q_clip_max;
}
max_target_q = tgt_q_values[0];
for (int i = 1; i < NUM_ACTIONS; i++) {
if (tgt_q_values[i] > max_target_q) max_target_q = tgt_q_values[i];
}
}
/* Steps 12-16 only for valid threads with data */
@@ -1717,39 +1938,8 @@ extern "C" __global__ void dqn_full_experience_kernel_warp(
/* Episode start indices [N] — global bar index where each episode begins */
const int* __restrict__ episode_starts,
/* ---- Online Q-network weights (12 pointers) ---- */
const float* __restrict__ on_w_s1,
const float* __restrict__ on_b_s1,
const float* __restrict__ on_w_s2,
const float* __restrict__ on_b_s2,
const float* __restrict__ on_w_v1,
const float* __restrict__ on_b_v1,
const float* __restrict__ on_w_v2,
const float* __restrict__ on_b_v2,
const float* __restrict__ on_w_a1,
const float* __restrict__ on_b_a1,
const float* __restrict__ on_w_a2,
const float* __restrict__ on_b_a2,
/* ---- Target Q-network weights (12 pointers) ---- */
const float* __restrict__ tg_w_s1,
const float* __restrict__ tg_b_s1,
const float* __restrict__ tg_w_s2,
const float* __restrict__ tg_b_s2,
const float* __restrict__ tg_w_v1,
const float* __restrict__ tg_b_v1,
const float* __restrict__ tg_w_v2,
const float* __restrict__ tg_b_v2,
const float* __restrict__ tg_w_a1,
const float* __restrict__ tg_b_a1,
const float* __restrict__ tg_w_a2,
const float* __restrict__ tg_b_a2,
/* ---- Curiosity model weights (4 pointers) ---- */
const float* __restrict__ cur_w1,
const float* __restrict__ cur_b1,
const float* __restrict__ cur_w2,
const float* __restrict__ cur_b2,
/* ---- All network weights packed (384 bytes, constant memory) ---- */
struct KernelWeightPack wp,
/* ---- Per-episode mutable state arrays ---- */
float* portfolio_states, /* [N, PORTFOLIO_STATE_SIZE] */
@@ -1792,12 +1982,6 @@ extern "C" __global__ void dqn_full_experience_kernel_warp(
float reward_norm_alpha,
int use_rmsnorm,
/* ---- D6: RMSNorm gamma weights (4 pointers) ---- */
const float* __restrict__ rms_s0_gamma, /* [SHARED_H1] */
const float* __restrict__ rms_s1_gamma, /* [SHARED_H2] */
const float* __restrict__ rms_v_gamma, /* [VALUE_H] */
const float* __restrict__ rms_a_gamma, /* [ADV_H] */
/* ---- Fill simulation config ---- */
float fill_median_spread,
float fill_median_vol,
@@ -1815,6 +1999,9 @@ extern "C" __global__ void dqn_full_experience_kernel_warp(
/* ---- N-step returns config ---- */
int n_steps,
/* ---- Branching DQN config (Tavakoli et al., 2018) ---- */
int use_branching,
/* ---- RNG states [N] ---- */
unsigned int* rng_states,
@@ -1826,6 +2013,9 @@ extern "C" __global__ void dqn_full_experience_kernel_warp(
float* out_target_q, /* [N, L] */
float* out_td_error /* [N, L] */
) {
/* ---- Unpack weight pointers from struct into local __restrict__ aliases ---- */
UNPACK_WEIGHT_PTRS(wp);
/* ---- Shared memory workspace for weight tiling ---- */
extern __shared__ float shmem[];
float* shmem_weights = shmem;
@@ -1873,6 +2063,13 @@ extern "C" __global__ void dqn_full_experience_kernel_warp(
float scratch2_dist[DIST_SIZE(SHARED_H2)];
float q_values[NUM_ACTIONS];
float tgt_q_values[NUM_ACTIONS];
/* Branching DQN per-head Q-value arrays (used only when use_branching) */
float br_q_exposure[BRANCH_EXPOSURE_ACTIONS];
float br_q_order[BRANCH_ORDER_ACTIONS];
float br_q_urgency[BRANCH_URGENCY_ACTIONS];
float br_tgt_exposure[BRANCH_EXPOSURE_ACTIONS];
float br_tgt_order[BRANCH_ORDER_ACTIONS];
float br_tgt_urgency[BRANCH_URGENCY_ACTIONS];
/* Zero-init distributed state */
for (int i = 0; i < DIST_SIZE(STATE_DIM); i++) state_dist[i] = 0.0f;
@@ -1984,7 +2181,20 @@ extern "C" __global__ void dqn_full_experience_kernel_warp(
/* ---- Step 4: Online Q-network forward (warp-cooperative) ----
* All lanes participate in cooperative loads + __syncthreads(). */
if (use_distributional && num_atoms > 1) {
if (use_branching) {
/* Branching DQN: register-based forward (all lanes compute identically) */
q_forward_branching(
state_dist,
on_w_s1, on_b_s1, on_w_s2, on_b_s2,
on_w_v1, on_b_v1, on_w_v2, on_b_v2,
on_w_a1, on_b_a1, on_w_a2, on_b_a2,
on_w_bo1, on_b_bo1, on_w_bo2, on_b_bo2,
on_w_bu1, on_b_bu1, on_w_bu2, on_b_bu2,
scratch1_dist, scratch2_dist,
scratch1_dist, scratch1_dist + VALUE_H, /* reuse as scratch_v, scratch_a */
br_q_exposure, br_q_order, br_q_urgency
);
} else if (use_distributional && num_atoms > 1) {
q_forward_distributional_warp_shmem(
state_dist,
on_w_s1, on_b_s1, on_w_s2, on_b_s2,
@@ -2035,76 +2245,145 @@ extern "C" __global__ void dqn_full_experience_kernel_warp(
int barrier_label = 0;
float div_penalty = 0.0f;
int br_exposure_sel = 0, br_order_sel = 0, br_urgency_sel = 0;
if (!skip_data) {
/* ---- Step 4b: Q-value clipping (D2) — all lanes, identical data ---- */
for (int i = 0; i < NUM_ACTIONS; i++) {
if (q_values[i] < q_clip_min) q_values[i] = q_clip_min;
if (q_values[i] > q_clip_max) q_values[i] = q_clip_max;
}
/* ---- Step 5: Epsilon-greedy + count-bonus action selection (lane 0) ---- */
if (lane_id == 0) {
float r = gpu_random(&rng);
if (r < epsilon) {
action_idx = (int)(gpu_random(&rng) * (float)DQN_NUM_ACTIONS);
if (action_idx >= DQN_NUM_ACTIONS) action_idx = DQN_NUM_ACTIONS - 1;
} else {
if (count_bonus_coefficient > 0.0f && total_action_count > 0) {
float log_n = logf((float)total_action_count);
float q_bonus[NUM_ACTIONS];
for (int i = 0; i < NUM_ACTIONS; i++) {
float bonus = count_bonus_coefficient
* sqrtf(log_n / (1.0f + (float)action_counts[i]));
q_bonus[i] = q_values[i] + bonus;
}
action_idx = argmax_q(q_bonus);
} else {
action_idx = argmax_q(q_values);
}
/* ---- Step 4b + Step 5: Q-value clipping + action selection ---- */
if (use_branching) {
/* Branching: clip per-head, all lanes identical */
for (int i = 0; i < BRANCH_EXPOSURE_ACTIONS; i++) {
if (br_q_exposure[i] < q_clip_min) br_q_exposure[i] = q_clip_min;
if (br_q_exposure[i] > q_clip_max) br_q_exposure[i] = q_clip_max;
}
for (int i = 0; i < BRANCH_ORDER_ACTIONS; i++) {
if (br_q_order[i] < q_clip_min) br_q_order[i] = q_clip_min;
if (br_q_order[i] > q_clip_max) br_q_order[i] = q_clip_max;
}
for (int i = 0; i < BRANCH_URGENCY_ACTIONS; i++) {
if (br_q_urgency[i] < q_clip_min) br_q_urgency[i] = q_clip_min;
if (br_q_urgency[i] > q_clip_max) br_q_urgency[i] = q_clip_max;
}
/* D1: Update count state */
action_counts[action_idx]++;
total_action_count++;
if (lane_id == 0) {
/* 3 independent epsilon-greedy */
if (gpu_random(&rng) < epsilon) {
br_exposure_sel = (int)(gpu_random(&rng) * (float)BRANCH_EXPOSURE_ACTIONS);
if (br_exposure_sel >= BRANCH_EXPOSURE_ACTIONS) br_exposure_sel = BRANCH_EXPOSURE_ACTIONS - 1;
} else {
br_exposure_sel = argmax_arr(br_q_exposure, BRANCH_EXPOSURE_ACTIONS);
}
if (gpu_random(&rng) < epsilon) {
br_order_sel = (int)(gpu_random(&rng) * (float)BRANCH_ORDER_ACTIONS);
if (br_order_sel >= BRANCH_ORDER_ACTIONS) br_order_sel = BRANCH_ORDER_ACTIONS - 1;
} else {
br_order_sel = argmax_arr(br_q_order, BRANCH_ORDER_ACTIONS);
}
if (gpu_random(&rng) < epsilon) {
br_urgency_sel = (int)(gpu_random(&rng) * (float)BRANCH_URGENCY_ACTIONS);
if (br_urgency_sel >= BRANCH_URGENCY_ACTIONS) br_urgency_sel = BRANCH_URGENCY_ACTIONS - 1;
} else {
br_urgency_sel = argmax_arr(br_q_urgency, BRANCH_URGENCY_ACTIONS);
}
action_idx = br_exposure_sel * 9 + br_order_sel * 3 + br_urgency_sel;
online_q_selected = (br_q_exposure[br_exposure_sel]
+ br_q_order[br_order_sel]
+ br_q_urgency[br_urgency_sel]) / 3.0f;
action_counts[br_exposure_sel]++;
total_action_count++;
if (fill_simulation_enabled) {
float norm_vol = (fill_median_vol > 0.0f) ? (spread / fill_median_spread) : 1.0f;
norm_vol = fminf(fmaxf(norm_vol, 0.0f), 3.0f);
float cost_adj;
int filled = simulate_fill_check(
br_order_sel, br_urgency_sel, norm_vol,
spread * 10000.0f, global_bar, action_idx,
fill_ioc_fill_prob, fill_limit_fill_min, fill_limit_fill_max,
fill_spread_cost_frac, fill_spread_capture_frac, &cost_adj);
if (!filled) {
br_exposure_sel = 2;
action_idx = br_exposure_sel * 9 + br_order_sel * 3 + br_urgency_sel;
}
}
}
} else {
for (int i = 0; i < NUM_ACTIONS; i++) {
if (q_values[i] < q_clip_min) q_values[i] = q_clip_min;
if (q_values[i] > q_clip_max) q_values[i] = q_clip_max;
}
if (lane_id == 0) {
float r = gpu_random(&rng);
if (r < epsilon) {
action_idx = (int)(gpu_random(&rng) * (float)DQN_NUM_ACTIONS);
if (action_idx >= DQN_NUM_ACTIONS) action_idx = DQN_NUM_ACTIONS - 1;
} else {
if (count_bonus_coefficient > 0.0f && total_action_count > 0) {
float log_n = logf((float)total_action_count);
float q_bonus[NUM_ACTIONS];
for (int i = 0; i < NUM_ACTIONS; i++) {
float bonus = count_bonus_coefficient
* sqrtf(log_n / (1.0f + (float)action_counts[i]));
q_bonus[i] = q_values[i] + bonus;
}
action_idx = argmax_q(q_bonus);
} else {
action_idx = argmax_q(q_values);
}
}
action_counts[action_idx]++;
total_action_count++;
}
}
/* Broadcast action from lane 0 to all lanes */
action_idx = __shfl_sync(0xFFFFFFFF, action_idx, 0);
if (lane_id == 0) {
online_q_selected = q_values[action_idx];
if (!use_branching) {
online_q_selected = q_values[action_idx];
/* ---- Step 5b: Order routing + fill simulation ---- */
if (fill_simulation_enabled) {
int order_type, urgency;
route_order(spread, fill_median_spread, 0.0f, fill_median_vol,
&order_type, &urgency);
/* ---- Step 5b: Order routing + fill simulation ---- */
if (fill_simulation_enabled) {
int order_type, urgency;
route_order(spread, fill_median_spread, 0.0f, fill_median_vol,
&order_type, &urgency);
float norm_vol = (fill_median_vol > 0.0f) ? (spread / fill_median_spread) : 1.0f;
norm_vol = fminf(fmaxf(norm_vol, 0.0f), 3.0f);
float norm_vol = (fill_median_vol > 0.0f) ? (spread / fill_median_spread) : 1.0f;
norm_vol = fminf(fmaxf(norm_vol, 0.0f), 3.0f);
float cost_adj;
int filled = simulate_fill_check(
order_type, urgency, norm_vol,
spread * 10000.0f,
global_bar, action_idx,
fill_ioc_fill_prob, fill_limit_fill_min, fill_limit_fill_max,
fill_spread_cost_frac, fill_spread_capture_frac,
&cost_adj
);
float cost_adj;
int filled = simulate_fill_check(
order_type, urgency, norm_vol,
spread * 10000.0f,
global_bar, action_idx,
fill_ioc_fill_prob, fill_limit_fill_min, fill_limit_fill_max,
fill_spread_cost_frac, fill_spread_capture_frac,
&cost_adj
);
if (!filled) {
action_idx = 2; /* Override to Flat */
if (!filled) {
action_idx = 2;
}
}
}
out_actions[out_off] = action_idx;
/* ---- Step 6: Portfolio simulation ---- */
float target_exposure = action_to_exposure(action_idx);
float target_exposure;
if (use_branching) {
target_exposure = factored_action_to_exposure(action_idx);
} else {
target_exposure = action_to_exposure(action_idx);
}
float target_position = target_exposure * max_position;
float tx_rate;
if (fill_simulation_enabled) {
if (use_branching) {
float spread_bps = spread * 10000.0f;
tx_rate = fabsf(order_type_tx_cost(br_order_sel, spread_bps,
fill_spread_cost_frac, fill_spread_capture_frac));
tx_rate *= tx_cost_multiplier;
} else if (fill_simulation_enabled) {
int order_type_for_cost, urgency_unused;
route_order(spread, fill_median_spread, 0.0f, fill_median_vol,
&order_type_for_cost, &urgency_unused);
@@ -2292,7 +2571,35 @@ extern "C" __global__ void dqn_full_experience_kernel_warp(
}
/* ---- Step 11: Target Q-network forward (warp-cooperative, always clean) ---- */
if (use_distributional && num_atoms > 1) {
float max_target_q = 0.0f;
if (use_branching) {
q_forward_branching(
next_state_dist,
tg_w_s1, tg_b_s1, tg_w_s2, tg_b_s2,
tg_w_v1, tg_b_v1, tg_w_v2, tg_b_v2,
tg_w_a1, tg_b_a1, tg_w_a2, tg_b_a2,
tg_w_bo1, tg_b_bo1, tg_w_bo2, tg_b_bo2,
tg_w_bu1, tg_b_bu1, tg_w_bu2, tg_b_bu2,
scratch1_dist, scratch2_dist,
scratch1_dist, scratch1_dist + VALUE_H,
br_tgt_exposure, br_tgt_order, br_tgt_urgency
);
for (int i = 0; i < BRANCH_EXPOSURE_ACTIONS; i++) {
if (br_tgt_exposure[i] < q_clip_min) br_tgt_exposure[i] = q_clip_min;
if (br_tgt_exposure[i] > q_clip_max) br_tgt_exposure[i] = q_clip_max;
}
for (int i = 0; i < BRANCH_ORDER_ACTIONS; i++) {
if (br_tgt_order[i] < q_clip_min) br_tgt_order[i] = q_clip_min;
if (br_tgt_order[i] > q_clip_max) br_tgt_order[i] = q_clip_max;
}
for (int i = 0; i < BRANCH_URGENCY_ACTIONS; i++) {
if (br_tgt_urgency[i] < q_clip_min) br_tgt_urgency[i] = q_clip_min;
if (br_tgt_urgency[i] > q_clip_max) br_tgt_urgency[i] = q_clip_max;
}
max_target_q = (max_arr(br_tgt_exposure, BRANCH_EXPOSURE_ACTIONS)
+ max_arr(br_tgt_order, BRANCH_ORDER_ACTIONS)
+ max_arr(br_tgt_urgency, BRANCH_URGENCY_ACTIONS)) / 3.0f;
} else if (use_distributional && num_atoms > 1) {
q_forward_distributional_warp_shmem(
next_state_dist,
tg_w_s1, tg_b_s1, tg_w_s2, tg_b_s2,
@@ -2306,6 +2613,14 @@ extern "C" __global__ void dqn_full_experience_kernel_warp(
&rng, lane_id,
shmem_weights, shmem_bias
);
for (int i = 0; i < NUM_ACTIONS; i++) {
if (tgt_q_values[i] < q_clip_min) tgt_q_values[i] = q_clip_min;
if (tgt_q_values[i] > q_clip_max) tgt_q_values[i] = q_clip_max;
}
max_target_q = tgt_q_values[0];
for (int i = 1; i < NUM_ACTIONS; i++) {
if (tgt_q_values[i] > max_target_q) max_target_q = tgt_q_values[i];
}
} else {
q_forward_dueling_warp_shmem(
next_state_dist,
@@ -2317,17 +2632,14 @@ extern "C" __global__ void dqn_full_experience_kernel_warp(
lane_id,
shmem_weights, shmem_bias
);
}
/* D2: Clip target Q-values (all lanes, identical) */
for (int i = 0; i < NUM_ACTIONS; i++) {
if (tgt_q_values[i] < q_clip_min) tgt_q_values[i] = q_clip_min;
if (tgt_q_values[i] > q_clip_max) tgt_q_values[i] = q_clip_max;
}
float max_target_q = tgt_q_values[0];
for (int i = 1; i < NUM_ACTIONS; i++) {
if (tgt_q_values[i] > max_target_q)
max_target_q = tgt_q_values[i];
for (int i = 0; i < NUM_ACTIONS; i++) {
if (tgt_q_values[i] < q_clip_min) tgt_q_values[i] = q_clip_min;
if (tgt_q_values[i] > q_clip_max) tgt_q_values[i] = q_clip_max;
}
max_target_q = tgt_q_values[0];
for (int i = 1; i < NUM_ACTIONS; i++) {
if (tgt_q_values[i] > max_target_q) max_target_q = tgt_q_values[i];
}
}
/* ---- Steps 12-16: reward, done, TD error, output, reset (lane 0 only) ---- */

143
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -22,9 +22,12 @@ use tracing::{debug, info};
use crate::MLError;
use super::gpu_weights::{
CuriosityWeightSet, DuelingWeightSet, RmsNormWeightSet,
extract_curiosity_weights, extract_dueling_weights, extract_rmsnorm_weights,
sync_curiosity_weights, sync_dueling_weights, sync_rmsnorm_weights,
BranchingWeightSet, CuriosityWeightSet, DuelingWeightSet, KernelWeightPack,
RmsNormWeightSet,
extract_branching_weights, extract_curiosity_weights, extract_dueling_weights,
extract_dueling_weights_branching, extract_rmsnorm_weights,
sync_branching_weights, sync_curiosity_weights, sync_dueling_weights,
sync_dueling_weights_branching, sync_rmsnorm_weights,
};
// ---------------------------------------------------------------------------
@@ -277,6 +280,11 @@ pub struct GpuExperienceCollector {
online_rmsnorm: RmsNormWeightSet,
target_rmsnorm: RmsNormWeightSet,
// Branching DQN extra advantage heads (zero-filled when not branching)
online_branching: BranchingWeightSet,
target_branching: BranchingWeightSet,
use_branching: bool,
// Per-episode state buffers [alloc_episodes, ...]
portfolio_states: CudaSlice<f32>, // [alloc_episodes * PORTFOLIO_STATE_SIZE]
barrier_states: CudaSlice<f32>, // [alloc_episodes * BARRIER_STATE_SIZE]
@@ -316,7 +324,7 @@ impl GpuExperienceCollector {
/// `state_dim` must match the network's input_dim (e.g. 56 with OFI+alignment).
/// `market_dim` must match the feature buffer stride (e.g. 42 basic features).
/// `num_atoms_max` must be >= the actual C51 atom count.
/// * `n_episodes` - Number of episodes to allocate buffers for (clamped to 256 max).
/// * `n_episodes` - Number of episodes to allocate buffers for (clamped to 32768 max).
/// * `timesteps_per_episode` - Timesteps per episode to allocate for (clamped to 1000 max).
pub fn new(
stream: Arc<CudaStream>,
@@ -330,6 +338,7 @@ impl GpuExperienceCollector {
kernel_dims: (usize, usize, usize),
n_episodes: usize,
timesteps_per_episode: usize,
use_branching: bool,
) -> Result<Self, MLError> {
let alloc_episodes = n_episodes.min(MAX_EPISODES_LIMIT);
let alloc_timesteps = timesteps_per_episode.min(MAX_TIMESTEPS_LIMIT);
@@ -437,8 +446,18 @@ impl GpuExperienceCollector {
}
// ---- Step 2: Extract weights ----
let online_weights = extract_dueling_weights(online_vars, &stream)?;
let target_weights = extract_dueling_weights(target_vars, &stream)?;
// When branching is active, the VarMap uses `branch_0_fc.*` for the exposure
// head instead of `advantage_fc.*`. Use the branching-aware extractor.
let online_weights = if use_branching {
extract_dueling_weights_branching(online_vars, &stream)?
} else {
extract_dueling_weights(online_vars, &stream)?
};
let target_weights = if use_branching {
extract_dueling_weights_branching(target_vars, &stream)?
} else {
extract_dueling_weights(target_vars, &stream)?
};
let curiosity_weights = match curiosity_vars {
Some(vars) => extract_curiosity_weights(vars, &stream)?,
None => {
@@ -467,6 +486,21 @@ impl GpuExperienceCollector {
}
};
// Branching DQN: extract order + urgency heads (branches 1+2), or zeros when disabled
let (_, shared_h2, _, adv_h) = network_dims;
let (online_branching, target_branching) = if use_branching {
info!("GPU experience collector: extracting branching weights (order + urgency heads)");
(
extract_branching_weights(online_vars, &stream)?,
extract_branching_weights(target_vars, &stream)?,
)
} else {
(
BranchingWeightSet::zeros(&stream, shared_h2, adv_h)?,
BranchingWeightSet::zeros(&stream, shared_h2, adv_h)?,
)
};
// ---- Step 3: Allocate per-episode buffers (sized to actual config) ----
let portfolio_states = stream
.alloc_zeros::<f32>(alloc_episodes * PORTFOLIO_STATE_SIZE)
@@ -600,6 +634,9 @@ impl GpuExperienceCollector {
curiosity_weights,
online_rmsnorm,
target_rmsnorm,
online_branching,
target_branching,
use_branching,
portfolio_states,
barrier_states,
diversity_windows,
@@ -928,49 +965,32 @@ impl GpuExperienceCollector {
// ---- Step 5: Launch kernel ----
let n_i32 = n_episodes as i32;
// Safety: kernel parameter order matches dqn_experience_kernel.cu lines 394-468 exactly.
// Pack all 48 weight pointers into a single struct (384 bytes).
// Passed by value in CUDA constant memory — zero extra allocations.
let weight_pack = KernelWeightPack::build(
&self.online_weights,
&self.target_weights,
&self.curiosity_weights,
&self.online_rmsnorm,
&self.online_branching,
&self.target_branching,
&self.stream,
);
// Safety: kernel parameter order matches dqn_experience_kernel.cu exactly.
// All buffer sizes are pre-validated (MAX_EPISODES, MAX_TIMESTEPS).
// All CudaSlice lifetimes are valid (owned by self).
// Both standard and warp kernels share the same parameter signature.
unsafe {
self.stream
.launch_builder(active_kernel)
// Market data
// Market data (3 ptrs)
.arg(market_features_buf)
.arg(targets_buf)
.arg(&self.episode_starts_buf)
// Online Q-network weights (12)
.arg(&self.online_weights.w_s1)
.arg(&self.online_weights.b_s1)
.arg(&self.online_weights.w_s2)
.arg(&self.online_weights.b_s2)
.arg(&self.online_weights.w_v1)
.arg(&self.online_weights.b_v1)
.arg(&self.online_weights.w_v2)
.arg(&self.online_weights.b_v2)
.arg(&self.online_weights.w_a1)
.arg(&self.online_weights.b_a1)
.arg(&self.online_weights.w_a2)
.arg(&self.online_weights.b_a2)
// Target Q-network weights (12)
.arg(&self.target_weights.w_s1)
.arg(&self.target_weights.b_s1)
.arg(&self.target_weights.w_s2)
.arg(&self.target_weights.b_s2)
.arg(&self.target_weights.w_v1)
.arg(&self.target_weights.b_v1)
.arg(&self.target_weights.w_v2)
.arg(&self.target_weights.b_v2)
.arg(&self.target_weights.w_a1)
.arg(&self.target_weights.b_a1)
.arg(&self.target_weights.w_a2)
.arg(&self.target_weights.b_a2)
// Curiosity model weights (4)
.arg(&self.curiosity_weights.w1)
.arg(&self.curiosity_weights.b1)
.arg(&self.curiosity_weights.w2)
.arg(&self.curiosity_weights.b2)
// Per-episode mutable state (4)
// All network weights packed (48 ptrs → 1 struct, 384 bytes)
.arg(&weight_pack)
// Per-episode mutable state (4 ptrs)
.arg(&mut self.portfolio_states)
.arg(&mut self.barrier_states)
.arg(&mut self.diversity_windows)
@@ -1006,11 +1026,6 @@ impl GpuExperienceCollector {
.arg(&config.v_max)
.arg(&config.reward_norm_alpha)
.arg(&(config.use_distributional as i32)) // use_rmsnorm = same as distributional
// D6: RMSNorm gamma weights (4 pointers)
.arg(&self.online_rmsnorm.gamma_s0)
.arg(&self.online_rmsnorm.gamma_s1)
.arg(&self.online_rmsnorm.gamma_v)
.arg(&self.online_rmsnorm.gamma_a)
// Fill simulation config
.arg(&config.fill_median_spread)
.arg(&config.fill_median_vol)
@@ -1024,6 +1039,8 @@ impl GpuExperienceCollector {
.arg(&(config.use_dsr as i32))
.arg(&config.dsr_eta)
.arg(&config.n_steps)
// Branching DQN flag
.arg(&(self.use_branching as i32))
// RNG
.arg(&mut self.rng_states)
// Outputs
@@ -1043,13 +1060,45 @@ impl GpuExperienceCollector {
}
/// Re-upload online Q-network weights from a `VarMap` (after training step).
///
/// When branching is active, uses `branch_0_fc/branch_0_out` keys for the
/// advantage slot instead of `advantage_fc/advantage_out`.
pub fn sync_online_weights(&mut self, vars: &VarMap) -> Result<(), MLError> {
sync_dueling_weights(vars, &mut self.online_weights, &self.stream)
if self.use_branching {
sync_dueling_weights_branching(vars, &mut self.online_weights, &self.stream)
} else {
sync_dueling_weights(vars, &mut self.online_weights, &self.stream)
}
}
/// Re-upload target Q-network weights from a `VarMap` (after soft update).
///
/// When branching is active, uses `branch_0_fc/branch_0_out` keys for the
/// advantage slot instead of `advantage_fc/advantage_out`.
pub fn sync_target_weights(&mut self, vars: &VarMap) -> Result<(), MLError> {
sync_dueling_weights(vars, &mut self.target_weights, &self.stream)
if self.use_branching {
sync_dueling_weights_branching(vars, &mut self.target_weights, &self.stream)
} else {
sync_dueling_weights(vars, &mut self.target_weights, &self.stream)
}
}
/// Re-upload branching DQN extra head weights for online network (branches 1+2).
pub fn sync_online_branching(&mut self, vars: &VarMap) -> Result<(), MLError> {
if self.use_branching {
sync_branching_weights(vars, &mut self.online_branching, &self.stream)
} else {
Ok(())
}
}
/// Re-upload branching DQN extra head weights for target network (branches 1+2).
pub fn sync_target_branching(&mut self, vars: &VarMap) -> Result<(), MLError> {
if self.use_branching {
sync_branching_weights(vars, &mut self.target_branching, &self.stream)
} else {
Ok(())
}
}
/// Re-upload curiosity forward model weights from a `VarMap`.

405
crates/ml/src/cuda_pipeline/gpu_weights.rs
View File

@@ -15,7 +15,7 @@ use std::sync::Arc;
use candle_core::cuda_backend::cudarc;
use candle_core::Var;
use cudarc::driver::{CudaSlice, CudaStream};
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DeviceRepr};
use candle_nn::VarMap;
use tracing::info;
@@ -49,6 +49,21 @@ const CURIOSITY_WEIGHT_NAMES: [&str; 4] = [
"fc2.bias",
];
/// Branching DQN extra head weight tensor names (8 total).
///
/// In branching mode the standard advantage head becomes the exposure head (branch 0).
/// These names cover the order head (branch 1) and urgency head (branch 2).
const BRANCHING_WEIGHT_NAMES: [&str; 8] = [
"branch_1_fc.weight",
"branch_1_fc.bias",
"branch_1_out.weight",
"branch_1_out.bias",
"branch_2_fc.weight",
"branch_2_fc.bias",
"branch_2_out.weight",
"branch_2_out.bias",
];
/// PPO Actor weight tensor names (6 total).
const PPO_ACTOR_WEIGHT_NAMES: [&str; 6] = [
"policy_layer_0.weight",
@@ -104,6 +119,73 @@ pub struct DuelingWeightSet {
pub b_a2: CudaSlice<f32>, // advantage_out.bias [45]
}
/// GPU buffers holding the 2 extra advantage head weight sets for Branching DQN.
///
/// In branching mode, the standard advantage head (w_a1/b_a1/w_a2/b_a2 in `DuelingWeightSet`)
/// becomes the exposure head (branch 0). This struct holds the order head (branch 1)
/// and urgency head (branch 2).
///
/// Layout matches `BranchingDuelingQNetwork` in `crates/ml-dqn/src/branching.rs`:
/// - Order FC: `[ADV_H, SHARED_H2]`, `[ADV_H]`
/// - Order Out: `[3, ADV_H]`, `[3]`
/// - Urgency FC: `[ADV_H, SHARED_H2]`, `[ADV_H]`
/// - Urgency Out: `[3, ADV_H]`, `[3]`
#[allow(missing_debug_implementations)] // CudaSlice does not implement Debug
pub struct BranchingWeightSet {
// Order head (branch 1)
pub w_bo1: CudaSlice<f32>, // branch_1_fc.weight [adv_h, shared_h2]
pub b_bo1: CudaSlice<f32>, // branch_1_fc.bias [adv_h]
pub w_bo2: CudaSlice<f32>, // branch_1_out.weight [3, adv_h]
pub b_bo2: CudaSlice<f32>, // branch_1_out.bias [3]
// Urgency head (branch 2)
pub w_bu1: CudaSlice<f32>, // branch_2_fc.weight [adv_h, shared_h2]
pub b_bu1: CudaSlice<f32>, // branch_2_fc.bias [adv_h]
pub w_bu2: CudaSlice<f32>, // branch_2_out.weight [3, adv_h]
pub b_bu2: CudaSlice<f32>, // branch_2_out.bias [3]
}
impl BranchingWeightSet {
/// Create zero-filled branching weights on GPU.
///
/// Used when branching mode is disabled. The CUDA kernel may still read
/// these buffers but the branching output is ignored when `num_branches <= 1`.
///
/// `shared_h2`: width of the second shared hidden layer (input to each branch FC).
/// `adv_h`: hidden width of each advantage/branch head FC layer.
pub fn zeros(
stream: &Arc<CudaStream>,
shared_h2: usize,
adv_h: usize,
) -> Result<Self, MLError> {
Ok(Self {
w_bo1: stream.alloc_zeros::<f32>(adv_h * shared_h2).map_err(|e| {
MLError::ModelError(format!("Failed to alloc branching w_bo1 zeros: {e}"))
})?,
b_bo1: stream.alloc_zeros::<f32>(adv_h).map_err(|e| {
MLError::ModelError(format!("Failed to alloc branching b_bo1 zeros: {e}"))
})?,
w_bo2: stream.alloc_zeros::<f32>(3 * adv_h).map_err(|e| {
MLError::ModelError(format!("Failed to alloc branching w_bo2 zeros: {e}"))
})?,
b_bo2: stream.alloc_zeros::<f32>(3).map_err(|e| {
MLError::ModelError(format!("Failed to alloc branching b_bo2 zeros: {e}"))
})?,
w_bu1: stream.alloc_zeros::<f32>(adv_h * shared_h2).map_err(|e| {
MLError::ModelError(format!("Failed to alloc branching w_bu1 zeros: {e}"))
})?,
b_bu1: stream.alloc_zeros::<f32>(adv_h).map_err(|e| {
MLError::ModelError(format!("Failed to alloc branching b_bu1 zeros: {e}"))
})?,
w_bu2: stream.alloc_zeros::<f32>(3 * adv_h).map_err(|e| {
MLError::ModelError(format!("Failed to alloc branching w_bu2 zeros: {e}"))
})?,
b_bu2: stream.alloc_zeros::<f32>(3).map_err(|e| {
MLError::ModelError(format!("Failed to alloc branching b_bu2 zeros: {e}"))
})?,
})
}
}
/// GPU buffers holding all 4 weight tensors of the Curiosity Forward Model.
///
/// Layout matches `ForwardDynamicsModel` in `crates/ml/src/dqn/curiosity.rs`:
@@ -235,6 +317,175 @@ impl RmsNormWeightSet {
}
}
// ---------------------------------------------------------------------------
// KernelWeightPack — packed device pointers for kernel launch
// ---------------------------------------------------------------------------
/// Dueling Q-network weight device pointers (12 × u64).
///
/// Layout matches `struct DuelingWeightPtrs` in `common_device_functions.cuh`.
#[repr(C)]
#[derive(Clone, Copy)]
pub(super) struct GpuDuelingPtrs {
pub w_s1: u64, pub b_s1: u64,
pub w_s2: u64, pub b_s2: u64,
pub w_v1: u64, pub b_v1: u64,
pub w_v2: u64, pub b_v2: u64,
pub w_a1: u64, pub b_a1: u64,
pub w_a2: u64, pub b_a2: u64,
}
/// Curiosity forward model weight device pointers (4 × u64).
///
/// Layout matches `struct CuriosityWeightPtrs` in `common_device_functions.cuh`.
#[repr(C)]
#[derive(Clone, Copy)]
pub(super) struct GpuCuriosityPtrs {
pub w1: u64, pub b1: u64,
pub w2: u64, pub b2: u64,
}
/// RMSNorm gamma weight device pointers (4 × u64).
///
/// Layout matches `struct RMSNormWeightPtrs` in `common_device_functions.cuh`.
#[repr(C)]
#[derive(Clone, Copy)]
pub(super) struct GpuRMSNormPtrs {
pub s0: u64,
pub s1: u64,
pub v: u64,
pub a: u64,
}
/// Branching DQN extra head weight device pointers (8 × u64).
///
/// Layout matches `struct BranchWeightPtrs` in `common_device_functions.cuh`.
#[repr(C)]
#[derive(Clone, Copy)]
pub(super) struct GpuBranchPtrs {
pub w_o1: u64, pub b_o1: u64,
pub w_o2: u64, pub b_o2: u64,
pub w_u1: u64, pub b_u1: u64,
pub w_u2: u64, pub b_u2: u64,
}
/// Packed weight pointers for the DQN experience collection kernel.
///
/// Passed by value as a single kernel argument via `DeviceRepr`.
/// Total size: 48 × 8 = 384 bytes (fits within CUDA's 4 KB parameter limit).
///
/// Layout matches `struct KernelWeightPack` in `common_device_functions.cuh` EXACTLY.
/// Both use `#[repr(C)]` / C struct layout with identical field ordering.
#[repr(C)]
#[derive(Clone, Copy)]
pub(super) struct KernelWeightPack {
pub online: GpuDuelingPtrs,
pub target: GpuDuelingPtrs,
pub curiosity: GpuCuriosityPtrs,
pub rmsnorm: GpuRMSNormPtrs,
pub online_br: GpuBranchPtrs,
pub target_br: GpuBranchPtrs,
}
// Safety: KernelWeightPack is #[repr(C)] with only u64 fields.
// All u64 values are valid CUDA device pointers (CUdeviceptr).
#[allow(unsafe_code)]
unsafe impl DeviceRepr for KernelWeightPack {}
/// Extract raw CUDA device pointer (CUdeviceptr = u64) from a CudaSlice.
///
/// Uses `DevicePtr::device_ptr()` and wraps the `SyncOnDrop` guard in
/// `ManuallyDrop` to skip recording a read event. This is safe because:
/// - The CudaSlice is owned by the caller and outlives the kernel launch
/// - The kernel launches on the same stream, so CUDA ordering is guaranteed
/// - `SyncOnDrop::Record` only records events for cross-stream sync (not needed)
fn raw_device_ptr(slice: &CudaSlice<f32>, stream: &CudaStream) -> u64 {
let (ptr, guard) = slice.device_ptr(stream);
// Keep guard alive without running its Drop (skip read event recording).
// Named binding prevents `let_underscore_must_use`; ManuallyDrop prevents drop.
let _no_drop = std::mem::ManuallyDrop::new(guard);
ptr
}
impl KernelWeightPack {
/// Build a `KernelWeightPack` from the GPU weight sets.
///
/// Extracts raw device pointers from each `CudaSlice` and packs them
/// into the struct for a single kernel argument pass.
pub(super) fn build(
online: &DuelingWeightSet,
target: &DuelingWeightSet,
curiosity: &CuriosityWeightSet,
rmsnorm: &RmsNormWeightSet,
online_br: &BranchingWeightSet,
target_br: &BranchingWeightSet,
stream: &CudaStream,
) -> Self {
Self {
online: GpuDuelingPtrs {
w_s1: raw_device_ptr(&online.w_s1, stream),
b_s1: raw_device_ptr(&online.b_s1, stream),
w_s2: raw_device_ptr(&online.w_s2, stream),
b_s2: raw_device_ptr(&online.b_s2, stream),
w_v1: raw_device_ptr(&online.w_v1, stream),
b_v1: raw_device_ptr(&online.b_v1, stream),
w_v2: raw_device_ptr(&online.w_v2, stream),
b_v2: raw_device_ptr(&online.b_v2, stream),
w_a1: raw_device_ptr(&online.w_a1, stream),
b_a1: raw_device_ptr(&online.b_a1, stream),
w_a2: raw_device_ptr(&online.w_a2, stream),
b_a2: raw_device_ptr(&online.b_a2, stream),
},
target: GpuDuelingPtrs {
w_s1: raw_device_ptr(&target.w_s1, stream),
b_s1: raw_device_ptr(&target.b_s1, stream),
w_s2: raw_device_ptr(&target.w_s2, stream),
b_s2: raw_device_ptr(&target.b_s2, stream),
w_v1: raw_device_ptr(&target.w_v1, stream),
b_v1: raw_device_ptr(&target.b_v1, stream),
w_v2: raw_device_ptr(&target.w_v2, stream),
b_v2: raw_device_ptr(&target.b_v2, stream),
w_a1: raw_device_ptr(&target.w_a1, stream),
b_a1: raw_device_ptr(&target.b_a1, stream),
w_a2: raw_device_ptr(&target.w_a2, stream),
b_a2: raw_device_ptr(&target.b_a2, stream),
},
curiosity: GpuCuriosityPtrs {
w1: raw_device_ptr(&curiosity.w1, stream),
b1: raw_device_ptr(&curiosity.b1, stream),
w2: raw_device_ptr(&curiosity.w2, stream),
b2: raw_device_ptr(&curiosity.b2, stream),
},
rmsnorm: GpuRMSNormPtrs {
s0: raw_device_ptr(&rmsnorm.gamma_s0, stream),
s1: raw_device_ptr(&rmsnorm.gamma_s1, stream),
v: raw_device_ptr(&rmsnorm.gamma_v, stream),
a: raw_device_ptr(&rmsnorm.gamma_a, stream),
},
online_br: GpuBranchPtrs {
w_o1: raw_device_ptr(&online_br.w_bo1, stream),
b_o1: raw_device_ptr(&online_br.b_bo1, stream),
w_o2: raw_device_ptr(&online_br.w_bo2, stream),
b_o2: raw_device_ptr(&online_br.b_bo2, stream),
w_u1: raw_device_ptr(&online_br.w_bu1, stream),
b_u1: raw_device_ptr(&online_br.b_bu1, stream),
w_u2: raw_device_ptr(&online_br.w_bu2, stream),
b_u2: raw_device_ptr(&online_br.b_bu2, stream),
},
target_br: GpuBranchPtrs {
w_o1: raw_device_ptr(&target_br.w_bo1, stream),
b_o1: raw_device_ptr(&target_br.b_bo1, stream),
w_o2: raw_device_ptr(&target_br.w_bo2, stream),
b_o2: raw_device_ptr(&target_br.b_bo2, stream),
w_u1: raw_device_ptr(&target_br.w_bu1, stream),
b_u1: raw_device_ptr(&target_br.b_bu1, stream),
w_u2: raw_device_ptr(&target_br.w_bu2, stream),
b_u2: raw_device_ptr(&target_br.b_bu2, stream),
},
}
}
}
// ---------------------------------------------------------------------------
// Extraction helpers
// ---------------------------------------------------------------------------
@@ -356,6 +607,132 @@ pub fn extract_dueling_weights(
})
}
/// Extract all 8 Branching DQN extra head weight tensors from a `VarMap`
/// and upload them to GPU as flat `CudaSlice<f32>` buffers.
///
/// The `VarMap` must contain tensor keys matching `BranchingDuelingQNetwork`:
/// `branch_1_fc.weight`, `branch_1_fc.bias`, `branch_1_out.weight`, `branch_1_out.bias`,
/// `branch_2_fc.weight`, `branch_2_fc.bias`, `branch_2_out.weight`, `branch_2_out.bias`.
pub fn extract_branching_weights(
vars: &VarMap,
stream: &Arc<CudaStream>,
) -> Result<BranchingWeightSet, MLError> {
let vars_data = vars.data().lock().map_err(|e| {
MLError::ModelError(format!("Failed to lock VarMap: {e}"))
})?;
let extract = |name: &str| -> Result<CudaSlice<f32>, MLError> {
extract_one(&vars_data, name, stream)
};
let w_bo1 = extract("branch_1_fc.weight")?;
let b_bo1 = extract("branch_1_fc.bias")?;
let w_bo2 = extract("branch_1_out.weight")?;
let b_bo2 = extract("branch_1_out.bias")?;
let w_bu1 = extract("branch_2_fc.weight")?;
let b_bu1 = extract("branch_2_fc.bias")?;
let w_bu2 = extract("branch_2_out.weight")?;
let b_bu2 = extract("branch_2_out.bias")?;
let total_params: usize = BRANCHING_WEIGHT_NAMES
.iter()
.filter_map(|name| {
vars_data.get(*name).map(|v| v.as_tensor().elem_count())
})
.sum();
info!(
total_params,
total_bytes = total_params * std::mem::size_of::<f32>(),
"Branching DQN weights extracted and uploaded to GPU"
);
Ok(BranchingWeightSet {
w_bo1,
b_bo1,
w_bo2,
b_bo2,
w_bu1,
b_bu1,
w_bu2,
b_bu2,
})
}
/// Extract Dueling Q-Network weights using branching VarMap key names.
///
/// In `BranchingDuelingQNetwork`, the exposure head (branch 0) uses VarMap keys
/// `branch_0_fc.*` / `branch_0_out.*` instead of `advantage_fc.*` / `advantage_out.*`.
/// Shared layers and value head keys are identical to standard dueling.
///
/// The returned `DuelingWeightSet` maps:
/// - `w_a1`/`b_a1` ← `branch_0_fc.weight`/`branch_0_fc.bias`
/// - `w_a2`/`b_a2` ← `branch_0_out.weight`/`branch_0_out.bias`
pub fn extract_dueling_weights_branching(
vars: &VarMap,
stream: &Arc<CudaStream>,
) -> Result<DuelingWeightSet, MLError> {
let vars_data = vars.data().lock().map_err(|e| {
MLError::ModelError(format!("Failed to lock VarMap: {e}"))
})?;
let extract = |name: &str| -> Result<CudaSlice<f32>, MLError> {
extract_one(&vars_data, name, stream)
};
// Shared + value heads: same key names
let w_s1 = extract("shared_0.weight")?;
let b_s1 = extract("shared_0.bias")?;
let w_s2 = extract("shared_1.weight")?;
let b_s2 = extract("shared_1.bias")?;
let w_v1 = extract("value_fc.weight")?;
let b_v1 = extract("value_fc.bias")?;
let w_v2 = extract("value_out.weight")?;
let b_v2 = extract("value_out.bias")?;
// Advantage/exposure head: branch_0_* in branching VarMap
let w_a1 = extract("branch_0_fc.weight")?;
let b_a1 = extract("branch_0_fc.bias")?;
let w_a2 = extract("branch_0_out.weight")?;
let b_a2 = extract("branch_0_out.bias")?;
info!("Branching DQN dueling weights (branch_0=exposure) extracted and uploaded to GPU");
Ok(DuelingWeightSet {
w_s1, b_s1, w_s2, b_s2,
w_v1, b_v1, w_v2, b_v2,
w_a1, b_a1, w_a2, b_a2,
})
}
/// Re-upload Dueling Q-Network weights using branching VarMap key names.
///
/// Counterpart to [`extract_dueling_weights_branching`] for weight sync.
pub fn sync_dueling_weights_branching(
vars: &VarMap,
weights: &mut DuelingWeightSet,
stream: &Arc<CudaStream>,
) -> Result<(), MLError> {
let vars_data = vars.data().lock().map_err(|e| {
MLError::ModelError(format!("Failed to lock VarMap: {e}"))
})?;
sync_one(&vars_data, "shared_0.weight", &mut weights.w_s1, stream)?;
sync_one(&vars_data, "shared_0.bias", &mut weights.b_s1, stream)?;
sync_one(&vars_data, "shared_1.weight", &mut weights.w_s2, stream)?;
sync_one(&vars_data, "shared_1.bias", &mut weights.b_s2, stream)?;
sync_one(&vars_data, "value_fc.weight", &mut weights.w_v1, stream)?;
sync_one(&vars_data, "value_fc.bias", &mut weights.b_v1, stream)?;
sync_one(&vars_data, "value_out.weight", &mut weights.w_v2, stream)?;
sync_one(&vars_data, "value_out.bias", &mut weights.b_v2, stream)?;
// Advantage/exposure: branch_0_* keys in branching VarMap
sync_one(&vars_data, "branch_0_fc.weight", &mut weights.w_a1, stream)?;
sync_one(&vars_data, "branch_0_fc.bias", &mut weights.b_a1, stream)?;
sync_one(&vars_data, "branch_0_out.weight", &mut weights.w_a2, stream)?;
sync_one(&vars_data, "branch_0_out.bias", &mut weights.b_a2, stream)?;
Ok(())
}
/// Extract all 4 Curiosity Forward Model weight tensors from a `VarMap`
/// and upload them to GPU as flat `CudaSlice<f32>` buffers.
///
@@ -424,6 +801,32 @@ pub fn sync_dueling_weights(
Ok(())
}
/// Re-upload Branching DQN extra head weights from a `VarMap` into existing
/// GPU buffers without re-allocating.
///
/// Call this after each training step (or target-network soft update)
/// to keep the GPU weight copies in sync with the Candle model.
pub fn sync_branching_weights(
vars: &VarMap,
weights: &mut BranchingWeightSet,
stream: &Arc<CudaStream>,
) -> Result<(), MLError> {
let vars_data = vars.data().lock().map_err(|e| {
MLError::ModelError(format!("Failed to lock VarMap: {e}"))
})?;
sync_one(&vars_data, "branch_1_fc.weight", &mut weights.w_bo1, stream)?;
sync_one(&vars_data, "branch_1_fc.bias", &mut weights.b_bo1, stream)?;
sync_one(&vars_data, "branch_1_out.weight", &mut weights.w_bo2, stream)?;
sync_one(&vars_data, "branch_1_out.bias", &mut weights.b_bo2, stream)?;
sync_one(&vars_data, "branch_2_fc.weight", &mut weights.w_bu1, stream)?;
sync_one(&vars_data, "branch_2_fc.bias", &mut weights.b_bu1, stream)?;
sync_one(&vars_data, "branch_2_out.weight", &mut weights.w_bu2, stream)?;
sync_one(&vars_data, "branch_2_out.bias", &mut weights.b_bu2, stream)?;
Ok(())
}
/// Extract RMSNorm gamma weights from a distributional dueling `VarMap`.
///
/// Looks for keys: `shared_rmsnorm_0.weight`, `shared_rmsnorm_1.weight`,

2
crates/ml/src/cuda_pipeline/ppo_experience_kernel.cu
View File

@@ -429,7 +429,7 @@ extern "C" __global__ void ppo_full_experience_kernel(
gae_values[t] = value;
/* ---- Step 5: Portfolio simulation ---- */
float target_exposure = ppo_action_to_exposure(action_idx);
float target_exposure = factored_action_to_exposure(action_idx);
float target_position = target_exposure * max_position;
float tx_rate = ppo_action_to_tx_cost(action_idx);

1
crates/ml/src/hyperopt/adapters/continuous_ppo.rs
View File

@@ -525,6 +525,7 @@ impl HyperparameterOptimizable for ContinuousPPOTrainer {
params,
objective: Self::extract_objective(&metrics),
duration_secs,
metrics: None,
};
std::fs::create_dir_all(self.training_paths.hyperopt_dir()).ok();

1
crates/ml/src/hyperopt/adapters/diffusion.rs
View File

@@ -398,6 +398,7 @@ impl HyperparameterOptimizable for DiffusionTrainer {
params,
objective: Self::extract_objective(&metrics),
duration_secs,
metrics: None,
};
std::fs::create_dir_all(self.training_paths.hyperopt_dir()).ok();
crate::hyperopt::shared_data::write_trial_result_json(

66
crates/ml/src/hyperopt/adapters/dqn.rs
View File

@@ -364,6 +364,10 @@ pub struct DQNParams {
/// Controls capacity of the per-dimension advantage streams in Branching DQN.
/// Only used when `use_branching` is true.
pub branch_hidden_dim: usize,
/// Whether to use Branching DQN (3 independent heads) vs standard Dueling DQN.
/// Hyperopt-tunable via index 11 (repurposed from dead v_range_mirror slot).
pub use_branching: bool,
}
impl Default for DQNParams {
@@ -441,6 +445,7 @@ impl Default for DQNParams {
eval_softmax_temp: 0.1, // Default: nearly greedy (hyperopt will tune)
dsr_eta: 0.01, // Default: ~100-step DSR lookback window
branch_hidden_dim: 128, // Default: 128 hidden units per branching head
use_branching: true, // Default: branching enabled (3 heads)
}
}
}
@@ -468,7 +473,7 @@ impl ParameterSpace for DQNParams {
// Q_max = max_reward / (1-gamma): gamma=0.88→3.75, gamma=0.99→45.
// Range [10, 50] covers the full gamma search space [0.88, 0.99].
(10.0, 50.0), // 10: v_range (symmetric: v_min=-v_range, v_max=+v_range)
(10.0, 50.0), // 11: v_range_mirror (UNUSED — kept for compat, mirrors index 10)
(0.0, 1.0), // 11: use_branching (discrete: 0=standard dueling, 1=branching 3-head)
(0.1_f64.ln(), 1.0_f64.ln()), // 12: noisy_sigma_init (log scale)
(128.0, 512.0), // 13: dueling_hidden_dim (linear, step=128)
(3.0, 5.0), // 14: n_steps (raised min: n=3 is Rainbow standard)
@@ -532,6 +537,7 @@ impl ParameterSpace for DQNParams {
let v_range = x[10].clamp(10.0, 50.0);
let v_min = -v_range;
let v_max = v_range;
let use_branching = x[11].round() >= 1.0;
let noisy_sigma_init = x[12].exp().clamp(0.1, 1.0);
let dueling_hidden_dim = (x[13].round() / 128.0).round() * 128.0;
let dueling_hidden_dim = dueling_hidden_dim.clamp(128.0, 512.0) as usize;
@@ -655,6 +661,7 @@ impl ParameterSpace for DQNParams {
eval_softmax_temp,
dsr_eta,
branch_hidden_dim,
use_branching,
};
Ok(params)
@@ -674,7 +681,7 @@ impl ParameterSpace for DQNParams {
self.per_alpha, // 8
self.per_beta_start, // 9
self.v_max, // 10: v_range (symmetric: v_min=-v_range)
self.v_max, // 11: v_range_mirror (same as 10)
if self.use_branching { 1.0 } else { 0.0 }, // 11: use_branching
self.noisy_sigma_init.ln(), // 12
self.dueling_hidden_dim as f64, // 13
self.n_steps as f64, // 14
@@ -709,7 +716,7 @@ impl ParameterSpace for DQNParams {
"per_alpha", // 8
"per_beta_start", // 9
"v_range", // 10: symmetric support ±v_range
"v_range_mirror", // 11: UNUSED (mirrors index 10)
"use_branching", // 11: branching vs standard dueling
"noisy_sigma_init", // 12
"dueling_hidden_dim", // 13
"n_steps", // 14
@@ -2608,9 +2615,14 @@ impl HyperparameterOptimizable for DQNTrainer {
iql_expectile_tau: 0.7,
iql_advantage_temperature: 3.0,
use_iql: false,
// Phase C+: branching disabled in hyperopt (use_branching not in search space)
use_branching: false,
use_branching: params.use_branching, // Phase C: hyperopt-tunable (index 11)
branch_hidden_dim: params.branch_hidden_dim,
// AutoReplaySizer VRAM fraction: scale linearly from 0% at ≤8GB to 70% at ≥48GB.
// Below 8GB the static buffer_size is used (regime heads + C51 need the headroom).
replay_buffer_vram_fraction: ((budget.gpu_memory_mb as f64 - 8192.0)
/ (49152.0 - 8192.0))
.clamp(0.0, 1.0)
* 0.70,
// Mixed precision: auto-detect from GPU hardware
mixed_precision: {
@@ -2666,7 +2678,9 @@ impl HyperparameterOptimizable for DQNTrainer {
// Training data is borrowed via train_with_shared_data (large: ~80%).
let training_metrics = if let Some(handle) = &self.runtime_handle {
if use_preloaded {
let train_ref = cached_train.as_ref().unwrap();
let train_ref = cached_train.as_ref().ok_or_else(|| {
MLError::TrainingError("cached_train is None but use_preloaded is true".to_owned())
})?;
let val_data = cached_val.as_ref().map(|a| a.as_ref().clone()).unwrap_or_default();
info!("Training DQN with shared data ({} train, {} val samples)",
train_ref.len(), val_data.len());
@@ -2682,7 +2696,9 @@ impl HyperparameterOptimizable for DQNTrainer {
MLError::TrainingError(format!("Failed to create runtime: {}", e))
})?;
if use_preloaded {
let train_ref = cached_train.as_ref().unwrap();
let train_ref = cached_train.as_ref().ok_or_else(|| {
MLError::TrainingError("cached_train is None but use_preloaded is true".to_owned())
})?;
let val_data = cached_val.as_ref().map(|a| a.as_ref().clone()).unwrap_or_default();
info!("Training DQN with shared data ({} train, {} val samples)",
train_ref.len(), val_data.len());
@@ -3340,6 +3356,7 @@ impl HyperparameterOptimizable for DQNTrainer {
params: params.clone(), // Clone params since we need it later for best trial tracking
objective: Self::extract_objective(&metrics),
duration_secs,
metrics: Self::extract_metrics(&metrics),
};
std::fs::create_dir_all(self.training_paths.hyperopt_dir()).ok();
@@ -3573,6 +3590,33 @@ impl HyperparameterOptimizable for DQNTrainer {
objective_total
}
fn extract_metrics(metrics: &Self::Metrics) -> Option<serde_json::Value> {
let mut m = serde_json::json!({
"avg_episode_reward": metrics.avg_episode_reward,
"train_loss": metrics.train_loss,
"val_loss": metrics.val_loss,
"epochs_completed": metrics.epochs_completed,
"buy_action_pct": metrics.buy_action_pct,
"sell_action_pct": metrics.sell_action_pct,
"hold_action_pct": metrics.hold_action_pct,
});
if let Some(ref bt) = metrics.backtest_metrics {
m["backtest"] = serde_json::json!({
"sharpe_ratio": bt.sharpe_ratio,
"sortino_ratio": bt.sortino_ratio,
"calmar_ratio": bt.calmar_ratio,
"omega_ratio": bt.omega_ratio,
"win_rate": bt.win_rate,
"max_drawdown_pct": bt.max_drawdown_pct,
"total_return_pct": bt.total_return_pct,
"total_trades": bt.total_trades,
"var_95": bt.var_95,
"cvar_95": bt.cvar_95,
});
}
Some(m)
}
}
#[cfg(test)]
@@ -3639,6 +3683,7 @@ mod tests {
cql_alpha: 0.05,
eval_softmax_temp: 0.8,
branch_hidden_dim: 128,
use_branching: true,
};
let continuous = params.to_continuous();
@@ -3680,6 +3725,9 @@ mod tests {
assert!((recovered.norm_type - 1.0).abs() < 1e-6);
assert!((recovered.activation_type - 1.0).abs() < 1e-6);
assert!((recovered.qr_kappa - 1.0).abs() < 1e-6);
// C6: use_branching roundtrips via index 11
assert_eq!(recovered.use_branching, params.use_branching);
}
#[test]
@@ -3711,7 +3759,7 @@ mod tests {
assert_eq!(bounds[8], (0.4, 0.8)); // per_alpha
assert_eq!(bounds[9], (0.2, 0.6)); // per_beta_start
assert_eq!(bounds[10], (10.0, 50.0)); // v_range (symmetric support, covers gamma 0.88-0.99)
assert_eq!(bounds[11], (10.0, 50.0)); // v_range_mirror (unused, same as 10)
assert_eq!(bounds[11], (0.0, 1.0)); // use_branching (C6: repurposed from v_range_mirror)
assert_eq!(bounds[13], (128.0, 512.0)); // dueling_hidden_dim
assert_eq!(bounds[14], (3.0, 5.0)); // n_steps (raised min: Rainbow standard)
assert_eq!(bounds[15], (51.0, 201.0)); // num_atoms
@@ -3744,7 +3792,7 @@ mod tests {
assert_eq!(names[8], "per_alpha");
assert_eq!(names[9], "per_beta_start");
assert_eq!(names[10], "v_range");
assert_eq!(names[11], "v_range_mirror");
assert_eq!(names[11], "use_branching");
assert_eq!(names[12], "noisy_sigma_init");
assert_eq!(names[13], "dueling_hidden_dim");
assert_eq!(names[14], "n_steps");

1
crates/ml/src/hyperopt/adapters/kan.rs
View File

@@ -386,6 +386,7 @@ impl HyperparameterOptimizable for KANTrainer {
params,
objective: Self::extract_objective(&metrics),
duration_secs,
metrics: None,
};
std::fs::create_dir_all(self.training_paths.hyperopt_dir()).ok();
crate::hyperopt::shared_data::write_trial_result_json(

1
crates/ml/src/hyperopt/adapters/liquid.rs
View File

@@ -462,6 +462,7 @@ impl HyperparameterOptimizable for LiquidTrainer {
params,
objective: Self::extract_objective(&metrics),
duration_secs,
metrics: None,
};
std::fs::create_dir_all(self.training_paths.hyperopt_dir()).ok();
crate::hyperopt::shared_data::write_trial_result_json(

1
crates/ml/src/hyperopt/adapters/mamba2.rs
View File

@@ -1008,6 +1008,7 @@ impl HyperparameterOptimizable for Mamba2Trainer {
params,
objective: Self::extract_objective(&metrics),
duration_secs,
metrics: None,
};
std::fs::create_dir_all(self.training_paths.hyperopt_dir()).ok();

1
crates/ml/src/hyperopt/adapters/ppo.rs
View File

@@ -1198,6 +1198,7 @@ impl HyperparameterOptimizable for PPOTrainer {
params,
objective: Self::extract_objective(&metrics),
duration_secs,
metrics: None,
};
std::fs::create_dir_all(self.training_paths.hyperopt_dir()).ok();

1
crates/ml/src/hyperopt/adapters/tft.rs
View File

@@ -513,6 +513,7 @@ impl HyperparameterOptimizable for TFTTrainer {
params,
objective: Self::extract_objective(&metrics),
duration_secs,
metrics: None,
};
std::fs::create_dir_all(self.training_paths.hyperopt_dir()).ok();

1
crates/ml/src/hyperopt/adapters/tggn.rs
View File

@@ -414,6 +414,7 @@ impl HyperparameterOptimizable for TGGNTrainer {
params,
objective: Self::extract_objective(&metrics),
duration_secs,
metrics: None,
};
std::fs::create_dir_all(self.training_paths.hyperopt_dir()).ok();
crate::hyperopt::shared_data::write_trial_result_json(

1
crates/ml/src/hyperopt/adapters/tlob.rs
View File

@@ -394,6 +394,7 @@ impl HyperparameterOptimizable for TLOBTrainer {
params,
objective: Self::extract_objective(&metrics),
duration_secs,
metrics: None,
};
std::fs::create_dir_all(self.training_paths.hyperopt_dir()).ok();
crate::hyperopt::shared_data::write_trial_result_json(

1
crates/ml/src/hyperopt/adapters/xlstm.rs
View File

@@ -377,6 +377,7 @@ impl HyperparameterOptimizable for XLSTMTrainer {
params,
objective: Self::extract_objective(&metrics),
duration_secs,
metrics: None,
};
std::fs::create_dir_all(self.training_paths.hyperopt_dir()).ok();
crate::hyperopt::shared_data::write_trial_result_json(

11
crates/ml/src/trainers/dqn/config.rs
View File

@@ -640,6 +640,10 @@ pub struct DQNHyperparameters {
pub epsilon_decay: f64,
/// Replay buffer capacity
pub buffer_size: usize,
/// Fraction of free VRAM to allocate for replay buffer (0.0 = use static buffer_size).
/// When > 0.0, overrides buffer_size with VRAM-optimal capacity.
/// Typical: 0.70 (70% of free VRAM after model load).
pub replay_buffer_vram_fraction: f64,
/// Minimum replay buffer size before training starts
pub min_replay_size: usize,
/// Number of training epochs
@@ -929,8 +933,8 @@ pub struct DQNHyperparameters {
/// When false, experience collection uses the batched CPU path while
/// the training step (forward/backward/optimizer) still runs on CUDA.
pub enable_gpu_experience_collector: bool,
/// Number of parallel episodes per GPU kernel launch (default: 128, max: 256)
/// Higher values improve GPU utilization on larger GPUs (e.g. H100: 256)
/// Number of parallel episodes per GPU kernel launch (default: 128, scaled dynamically)
/// Higher values improve GPU utilization on larger GPUs (e.g. H100: 8192)
pub gpu_n_episodes: usize,
/// Timesteps per episode in GPU kernel (default: 500, max: 1000)
/// Higher values collect more experience per launch at the cost of VRAM
@@ -1021,6 +1025,7 @@ impl DQNHyperparameters {
epsilon_end: 0.01,
epsilon_decay: 0.995,
buffer_size: 500000, // WAVE 24: Increased from 100K for better sample diversity (anti-overfitting)
replay_buffer_vram_fraction: 0.70,
min_replay_size: 1000,
epochs: 100,
checkpoint_frequency: 10,
@@ -1211,7 +1216,7 @@ impl DQNHyperparameters {
iql_expectile_tau: 0.7,
iql_advantage_temperature: 3.0,
use_iql: false,
use_branching: false,
use_branching: true,
branch_hidden_dim: 128,
}
}

128
crates/ml/src/trainers/dqn/trainer.rs
View File

@@ -392,6 +392,38 @@ impl DQNTrainer {
.map_err(|e| anyhow::anyhow!("Failed to initialize device: {}", e))?
};
// Dynamic replay buffer sizing: scale replay capacity to available VRAM.
// Only activates when replay_buffer_vram_fraction > 0 and GPU is detected.
if hyperparams.replay_buffer_vram_fraction > 0.0 && device.is_cuda() {
use ml_core::memory_optimization::detect_gpu_hardware;
match detect_gpu_hardware() {
Ok(hw) => {
let raw_sd = if hyperparams.mbp10_data_dir.is_some() { 53 } else { 45 };
let aligned_sd = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_sd, &device);
let optimal = hw.optimal_replay_capacity(
aligned_sd,
hyperparams.replay_buffer_vram_fraction,
);
if optimal != hyperparams.buffer_size {
info!(
"AutoReplaySizer: replay buffer {} -> {} (VRAM={:.0}MB, fraction={:.0}%)",
hyperparams.buffer_size,
optimal,
hw.free_memory_mb,
hyperparams.replay_buffer_vram_fraction * 100.0,
);
hyperparams.buffer_size = optimal;
}
}
Err(e) => {
info!(
"AutoReplaySizer unavailable ({}), using static buffer_size: {}",
e, hyperparams.buffer_size
);
}
}
}
info!(
"Initializing DQN trainer on device: {:?}, using 5 exposure actions + OrderRouter",
if device.is_cuda() { "CUDA GPU" } else { "CPU" },
@@ -1832,12 +1864,10 @@ impl DQNTrainer {
// Phase 1c: Initialize zero-roundtrip GPU experience collector (once)
// RegimeConditional agents now supported — GPU tensor batches are inserted
// into all head buffers, regime routing happens at train time via GPU masks.
// Branching DQN (3 independent heads, 11 total outputs) is incompatible with the
// fused CUDA kernel which hardcodes DQN_NUM_ACTIONS=5. When use_branching=true,
// skip initialization and fall through to the Candle tensor ops path.
// Branching DQN: the fused CUDA kernel now supports branching mode via
// use_branching flag + 16 extra weight pointers for order/urgency heads.
#[cfg(feature = "cuda")]
if self.hyperparams.enable_gpu_experience_collector
&& !self.hyperparams.use_branching
&& self.gpu_experience_collector.is_none()
{
if let candle_core::Device::Cuda(ref cuda_dev) = self.device {
@@ -1868,8 +1898,54 @@ impl DQNTrainer {
let num_atoms_max = (self.hyperparams.num_atoms as usize).max(51);
let kernel_dims = (state_dim, market_dim, num_atoms_max);
// Try plain dueling first, then hybrid (distributional+dueling)
if let (Some(online), Some(target)) = (
let use_branching = self.hyperparams.use_branching;
// Priority: branching > plain dueling > hybrid (distributional+dueling)
if let (true, Some(online_br), Some(target_br)) = (
use_branching,
dqn.branching_q_network.as_ref(),
dqn.branching_target_network.as_ref(),
) {
// Branching DQN: use branching network VarMaps.
// The collector's extract_dueling_weights_branching maps
// branch_0_fc → advantage slot; branches 1+2 go to BranchingWeightSet.
let cfg = online_br.config();
let dims = (
*cfg.shared_hidden_dims.first().unwrap_or(&256),
*cfg.shared_hidden_dims.get(1).unwrap_or(&256),
cfg.value_hidden_dim,
cfg.branch_hidden_dim,
);
let alloc_episodes = {
use ml_core::memory_optimization::detect_gpu_hardware;
let configured = self.hyperparams.gpu_n_episodes;
if configured >= 128 {
match detect_gpu_hardware() {
Ok(hw) => configured.max(hw.optimal_n_episodes(
state_dim,
self.hyperparams.gpu_timesteps_per_episode,
)).min(0x8000),
Err(_) => configured,
}
} else {
configured
}
};
Some(GpuExperienceCollector::new(
stream,
online_br.vars(),
target_br.vars(),
curiosity_vars,
self.hyperparams.initial_capital as f32,
self.hyperparams.avg_spread as f32,
self.hyperparams.cash_reserve_percent as f32,
dims,
kernel_dims,
alloc_episodes,
self.hyperparams.gpu_timesteps_per_episode,
true, // use_branching
))
} else if let (Some(online), Some(target)) = (
dqn.dueling_q_network.as_ref(),
dqn.dueling_target_network.as_ref(),
) {
@@ -1911,6 +1987,7 @@ impl DQNTrainer {
kernel_dims,
alloc_episodes,
self.hyperparams.gpu_timesteps_per_episode,
false, // use_branching
))
} else if let (Some(online), Some(target)) = (
dqn.dist_dueling_q_network.as_ref(),
@@ -1954,9 +2031,10 @@ impl DQNTrainer {
kernel_dims,
alloc_episodes,
self.hyperparams.gpu_timesteps_per_episode,
false, // use_branching
))
} else {
warn!("GPU experience collector skipped: no dueling or hybrid networks");
warn!("GPU experience collector skipped: no dueling, hybrid, or branching networks");
None
}
} else {
@@ -2007,14 +2085,10 @@ impl DQNTrainer {
// Phase 3: GPU experience collection via zero-roundtrip CUDA kernel
// Collects N episodes × L timesteps entirely on GPU, then bulk-inserts into replay buffer.
// Falls back to CPU path if GPU collection fails or is unavailable.
// Branching DQN: fused kernel hardcodes NUM_ACTIONS=5, incompatible with
// 3-head architecture (exposure=5, order=3, urgency=3 = 11 outputs).
// Candle tensor ops path is already GPU-resident, just not fused.
// Branching DQN: the fused kernel now supports branching via use_branching flag
// + 16 extra weight pointers for order/urgency heads.
#[cfg(feature = "cuda")]
let gpu_experiences_collected = if self.hyperparams.use_branching {
debug!("Branching DQN: skipping fused CUDA experience kernel, using Candle tensor ops");
false
} else if let (
let gpu_experiences_collected = if let (
Some(ref mut collector),
Some(ref features_buf),
Some(ref targets_buf),
@@ -3189,8 +3263,12 @@ impl DQNTrainer {
DQNAgentType::RegimeConditional(ref regime_dqn) => Some(regime_dqn.primary_head()),
};
if let Some(dqn) = dqn_ref {
// Sync online weights: plain dueling or hybrid (distributional+dueling)
let online_synced = if let Some(ref online) = dqn.dueling_q_network {
// Sync online weights: branching > plain dueling > hybrid (distributional+dueling)
// When branching, sync_online_weights internally uses branch_0 key names
// for the advantage slot, so we must pass the branching VarMap.
let online_synced = if let Some(ref bn) = dqn.branching_q_network {
collector.sync_online_weights(bn.vars()).is_ok()
} else if let Some(ref online) = dqn.dueling_q_network {
collector.sync_online_weights(online.vars()).is_ok()
} else if let Some(ref online) = dqn.dist_dueling_q_network {
collector.sync_online_weights(online.vars()).is_ok()
@@ -3200,8 +3278,10 @@ impl DQNTrainer {
if !online_synced {
warn!("GPU online weight sync failed or no network available");
}
// Sync target weights: plain dueling or hybrid
let target_synced = if let Some(ref target) = dqn.dueling_target_network {
// Sync target weights: branching > plain dueling > hybrid
let target_synced = if let Some(ref bn) = dqn.branching_target_network {
collector.sync_target_weights(bn.vars()).is_ok()
} else if let Some(ref target) = dqn.dueling_target_network {
collector.sync_target_weights(target.vars()).is_ok()
} else if let Some(ref target) = dqn.dist_dueling_target_network {
collector.sync_target_weights(target.vars()).is_ok()
@@ -3212,6 +3292,18 @@ impl DQNTrainer {
warn!("GPU target weight sync failed or no network available");
}
// Sync branching DQN extra heads (order + urgency, branches 1+2)
if let Some(ref bn) = dqn.branching_q_network {
if let Err(e) = collector.sync_online_branching(bn.vars()) {
warn!("GPU online branching weight sync failed: {}", e);
}
}
if let Some(ref bn) = dqn.branching_target_network {
if let Err(e) = collector.sync_target_branching(bn.vars()) {
warn!("GPU target branching weight sync failed: {}", e);
}
}
// Sync RMSNorm weights for distributional dueling networks (D6)
if let Some(ref online) = dqn.dist_dueling_q_network {
if let Err(e) = collector.sync_online_rmsnorm(online.vars()) {

4
crates/ml/src/trainers/ppo.rs
View File

@@ -68,7 +68,7 @@ pub struct PpoHyperparameters {
pub max_grad_norm_lstm: f64,
// Phase 3: GPU experience collection kernel configuration
/// Number of parallel episodes per GPU kernel launch (default: 128, max: 256)
/// Number of parallel episodes per GPU kernel launch (default: 128, scaled dynamically)
pub gpu_n_episodes: usize,
/// Timesteps per episode in GPU kernel (default: 500, max: 1000)
pub gpu_timesteps_per_episode: usize,
@@ -555,7 +555,7 @@ impl PpoTrainer {
) {
use crate::cuda_pipeline::gpu_ppo_collector::PpoCollectorConfig;
let n_episodes = self.hyperparams.gpu_n_episodes.min(256) as i32;
let n_episodes = self.hyperparams.gpu_n_episodes as i32;
let timesteps = self.hyperparams.gpu_timesteps_per_episode.min(1000) as i32;
let total_bars = self.raw_data_num_bars as i32;
let usable_bars = (total_bars - timesteps).max(1);

9
crates/ml/tests/gpu_kernel_parity_test.rs
View File

@@ -109,6 +109,7 @@ mod gpu_parity {
None,
100_000.0,
0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
false,
).unwrap();
let total_bars = 2000;
@@ -178,6 +179,7 @@ mod gpu_parity {
None,
100_000.0,
0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
false,
).unwrap();
let total_bars = 2000;
@@ -257,6 +259,7 @@ mod gpu_parity {
None,
100_000.0,
0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
false,
).unwrap();
// Market data with test state as every bar
@@ -376,6 +379,7 @@ mod gpu_parity {
let mut collector_clean = GpuExperienceCollector::new(
stream.clone(), network.vars(), target.vars(), None,
100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
false,
).unwrap();
let config_clean = ExperienceCollectorConfig {
@@ -399,6 +403,7 @@ mod gpu_parity {
let mut collector_noisy = GpuExperienceCollector::new(
stream.clone(), network.vars(), target.vars(), None,
100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
false,
).unwrap();
let config_noisy = ExperienceCollectorConfig {
@@ -463,6 +468,7 @@ mod gpu_parity {
let mut collector = GpuExperienceCollector::new(
stream.clone(), network.vars(), target.vars(), None,
100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
false,
).unwrap();
collector.sync_online_weights(network.vars()).expect("Online sync failed");
@@ -553,6 +559,7 @@ mod gpu_parity {
let mut collector = GpuExperienceCollector::new(
stream.clone(), network.vars(), target.vars(), None,
100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
false,
).unwrap();
for run in 0..5 {
@@ -592,6 +599,7 @@ mod gpu_parity {
let mut std_collector = GpuExperienceCollector::new(
stream.clone(), std_net.vars(), std_target.vars(), None,
100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
false,
).unwrap();
let std_batch = std_collector
@@ -611,6 +619,7 @@ mod gpu_parity {
let mut dist_collector = GpuExperienceCollector::new(
stream.clone(), dist_net.vars(), dist_target.vars(), None,
100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
false,
).unwrap();
let dist_batch = dist_collector

3
crates/ml/tests/smoke_test_real_data.rs
View File

@@ -464,6 +464,7 @@ mod gpu_smoke {
(256, 256, 128, 128),
(54, 51, 51),
4, 50,
false,
).unwrap();
// Run 4 episodes of 50 timesteps each
@@ -549,6 +550,7 @@ mod gpu_smoke {
(256, 256, 128, 128),
(54, 51, 51),
4, 50,
false,
).unwrap();
let n_episodes = 4_usize;
@@ -619,6 +621,7 @@ mod gpu_smoke {
(256, 256, 128, 128),
(54, 51, 51),
4, 50,
false,
).unwrap();
// Sync RMSNorm weights (C51 distributional)