diff --git a/crates/common/src/metrics/training_metrics.rs b/crates/common/src/metrics/training_metrics.rs index ae32fc502..bd1a4d95e 100644 --- a/crates/common/src/metrics/training_metrics.rs +++ b/crates/common/src/metrics/training_metrics.rs @@ -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) // --------------------------------------------------------------------------- diff --git a/crates/ml-core/src/memory_optimization/auto_batch_size.rs b/crates/ml-core/src/memory_optimization/auto_batch_size.rs index bc23d587f..415581a48 100644 --- a/crates/ml-core/src/memory_optimization/auto_batch_size.rs +++ b/crates/ml-core/src/memory_optimization/auto_batch_size.rs @@ -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"); + } } diff --git a/crates/ml-dqn/src/dqn.rs b/crates/ml-dqn/src/dqn.rs index f4075b8a3..1822ff8d0 100644 --- a/crates/ml-dqn/src/dqn.rs +++ b/crates/ml-dqn/src/dqn.rs @@ -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(()) diff --git a/crates/ml-hyperopt/src/optimizer.rs b/crates/ml-hyperopt/src/optimizer.rs index b2551b81c..840dd7ca5 100644 --- a/crates/ml-hyperopt/src/optimizer.rs +++ b/crates/ml-hyperopt/src/optimizer.rs @@ -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, }); } diff --git a/crates/ml-hyperopt/src/traits.rs b/crates/ml-hyperopt/src/traits.rs index 1dbe42900..51faac9c9 100644 --- a/crates/ml-hyperopt/src/traits.rs +++ b/crates/ml-hyperopt/src/traits.rs @@ -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 { + None + } } /// Trait for parameter space definition @@ -941,6 +949,9 @@ pub struct TrialResult

{ 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, } /// 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, }, ]; diff --git a/crates/ml/examples/hyperopt_baseline_rl.rs b/crates/ml/examples/hyperopt_baseline_rl.rs index d93779c58..47d4676c2 100644 --- a/crates/ml/examples/hyperopt_baseline_rl.rs +++ b/crates/ml/examples/hyperopt_baseline_rl.rs @@ -135,14 +135,19 @@ fn build_model_result( top_k_params: Vec, num_trials: usize, elapsed_secs: f64, + best_metrics: Option, ) -> 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(); diff --git a/crates/ml/examples/train_baseline_rl.rs b/crates/ml/examples/train_baseline_rl.rs index 6e0ddb7b3..dfe8ecf80 100644 --- a/crates/ml/examples/train_baseline_rl.rs +++ b/crates/ml/examples/train_baseline_rl.rs @@ -167,6 +167,10 @@ struct Args { /// Use this to generate datasets for offline training. #[arg(long)] collect_dataset: Option, + + /// 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); diff --git a/crates/ml/hyperparams/dqn_best.toml b/crates/ml/hyperparams/dqn_best.toml index 08ae8e4d1..9008c9bab 100644 --- a/crates/ml/hyperparams/dqn_best.toml +++ b/crates/ml/hyperparams/dqn_best.toml @@ -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 diff --git a/crates/ml/src/cuda_pipeline/common_device_functions.cuh b/crates/ml/src/cuda_pipeline/common_device_functions.cuh index 3c35b733b..e558413f2 100644 --- a/crates/ml/src/cuda_pipeline/common_device_functions.cuh +++ b/crates/ml/src/cuda_pipeline/common_device_functions.cuh @@ -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. * diff --git a/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu b/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu index ed4a89c24..241cec71a 100644 --- a/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu +++ b/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu @@ -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) ---- */ diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 9be07075f..002c7352d 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -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, // [alloc_episodes * PORTFOLIO_STATE_SIZE] barrier_states: CudaSlice, // [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, @@ -330,6 +338,7 @@ impl GpuExperienceCollector { kernel_dims: (usize, usize, usize), n_episodes: usize, timesteps_per_episode: usize, + use_branching: bool, ) -> Result { 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::(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`. diff --git a/crates/ml/src/cuda_pipeline/gpu_weights.rs b/crates/ml/src/cuda_pipeline/gpu_weights.rs index bbe02d44a..e067fb350 100644 --- a/crates/ml/src/cuda_pipeline/gpu_weights.rs +++ b/crates/ml/src/cuda_pipeline/gpu_weights.rs @@ -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, // 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, // branch_1_fc.weight [adv_h, shared_h2] + pub b_bo1: CudaSlice, // branch_1_fc.bias [adv_h] + pub w_bo2: CudaSlice, // branch_1_out.weight [3, adv_h] + pub b_bo2: CudaSlice, // branch_1_out.bias [3] + // Urgency head (branch 2) + pub w_bu1: CudaSlice, // branch_2_fc.weight [adv_h, shared_h2] + pub b_bu1: CudaSlice, // branch_2_fc.bias [adv_h] + pub w_bu2: CudaSlice, // branch_2_out.weight [3, adv_h] + pub b_bu2: CudaSlice, // 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, + shared_h2: usize, + adv_h: usize, + ) -> Result { + Ok(Self { + w_bo1: stream.alloc_zeros::(adv_h * shared_h2).map_err(|e| { + MLError::ModelError(format!("Failed to alloc branching w_bo1 zeros: {e}")) + })?, + b_bo1: stream.alloc_zeros::(adv_h).map_err(|e| { + MLError::ModelError(format!("Failed to alloc branching b_bo1 zeros: {e}")) + })?, + w_bo2: stream.alloc_zeros::(3 * adv_h).map_err(|e| { + MLError::ModelError(format!("Failed to alloc branching w_bo2 zeros: {e}")) + })?, + b_bo2: stream.alloc_zeros::(3).map_err(|e| { + MLError::ModelError(format!("Failed to alloc branching b_bo2 zeros: {e}")) + })?, + w_bu1: stream.alloc_zeros::(adv_h * shared_h2).map_err(|e| { + MLError::ModelError(format!("Failed to alloc branching w_bu1 zeros: {e}")) + })?, + b_bu1: stream.alloc_zeros::(adv_h).map_err(|e| { + MLError::ModelError(format!("Failed to alloc branching b_bu1 zeros: {e}")) + })?, + w_bu2: stream.alloc_zeros::(3 * adv_h).map_err(|e| { + MLError::ModelError(format!("Failed to alloc branching w_bu2 zeros: {e}")) + })?, + b_bu2: stream.alloc_zeros::(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, 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` 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, +) -> Result { + let vars_data = vars.data().lock().map_err(|e| { + MLError::ModelError(format!("Failed to lock VarMap: {e}")) + })?; + + let extract = |name: &str| -> Result, 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::(), + "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, +) -> Result { + let vars_data = vars.data().lock().map_err(|e| { + MLError::ModelError(format!("Failed to lock VarMap: {e}")) + })?; + + let extract = |name: &str| -> Result, 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, +) -> 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` 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, +) -> 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`, diff --git a/crates/ml/src/cuda_pipeline/ppo_experience_kernel.cu b/crates/ml/src/cuda_pipeline/ppo_experience_kernel.cu index 2e6ae24d7..d0076ad45 100644 --- a/crates/ml/src/cuda_pipeline/ppo_experience_kernel.cu +++ b/crates/ml/src/cuda_pipeline/ppo_experience_kernel.cu @@ -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); diff --git a/crates/ml/src/hyperopt/adapters/continuous_ppo.rs b/crates/ml/src/hyperopt/adapters/continuous_ppo.rs index 67d11b8a0..1401bf32e 100644 --- a/crates/ml/src/hyperopt/adapters/continuous_ppo.rs +++ b/crates/ml/src/hyperopt/adapters/continuous_ppo.rs @@ -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(); diff --git a/crates/ml/src/hyperopt/adapters/diffusion.rs b/crates/ml/src/hyperopt/adapters/diffusion.rs index 6f7ecf77c..07df9986b 100644 --- a/crates/ml/src/hyperopt/adapters/diffusion.rs +++ b/crates/ml/src/hyperopt/adapters/diffusion.rs @@ -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( diff --git a/crates/ml/src/hyperopt/adapters/dqn.rs b/crates/ml/src/hyperopt/adapters/dqn.rs index 48653e665..eac9697cb 100644 --- a/crates/ml/src/hyperopt/adapters/dqn.rs +++ b/crates/ml/src/hyperopt/adapters/dqn.rs @@ -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 { + 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"); diff --git a/crates/ml/src/hyperopt/adapters/kan.rs b/crates/ml/src/hyperopt/adapters/kan.rs index 323703480..0b8a3d9ee 100644 --- a/crates/ml/src/hyperopt/adapters/kan.rs +++ b/crates/ml/src/hyperopt/adapters/kan.rs @@ -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( diff --git a/crates/ml/src/hyperopt/adapters/liquid.rs b/crates/ml/src/hyperopt/adapters/liquid.rs index dafaf6bbf..4194130af 100644 --- a/crates/ml/src/hyperopt/adapters/liquid.rs +++ b/crates/ml/src/hyperopt/adapters/liquid.rs @@ -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( diff --git a/crates/ml/src/hyperopt/adapters/mamba2.rs b/crates/ml/src/hyperopt/adapters/mamba2.rs index 408886ed3..4b592820a 100644 --- a/crates/ml/src/hyperopt/adapters/mamba2.rs +++ b/crates/ml/src/hyperopt/adapters/mamba2.rs @@ -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(); diff --git a/crates/ml/src/hyperopt/adapters/ppo.rs b/crates/ml/src/hyperopt/adapters/ppo.rs index 2dd0a0f46..35cc3200e 100644 --- a/crates/ml/src/hyperopt/adapters/ppo.rs +++ b/crates/ml/src/hyperopt/adapters/ppo.rs @@ -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(); diff --git a/crates/ml/src/hyperopt/adapters/tft.rs b/crates/ml/src/hyperopt/adapters/tft.rs index 9a29e48c2..569dcfa18 100644 --- a/crates/ml/src/hyperopt/adapters/tft.rs +++ b/crates/ml/src/hyperopt/adapters/tft.rs @@ -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(); diff --git a/crates/ml/src/hyperopt/adapters/tggn.rs b/crates/ml/src/hyperopt/adapters/tggn.rs index 7e8e3a422..39ce59757 100644 --- a/crates/ml/src/hyperopt/adapters/tggn.rs +++ b/crates/ml/src/hyperopt/adapters/tggn.rs @@ -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( diff --git a/crates/ml/src/hyperopt/adapters/tlob.rs b/crates/ml/src/hyperopt/adapters/tlob.rs index b5addb888..c6e6f5de3 100644 --- a/crates/ml/src/hyperopt/adapters/tlob.rs +++ b/crates/ml/src/hyperopt/adapters/tlob.rs @@ -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( diff --git a/crates/ml/src/hyperopt/adapters/xlstm.rs b/crates/ml/src/hyperopt/adapters/xlstm.rs index ca5987087..e288e12bf 100644 --- a/crates/ml/src/hyperopt/adapters/xlstm.rs +++ b/crates/ml/src/hyperopt/adapters/xlstm.rs @@ -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( diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index cd2a54a1b..d3581b8da 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -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, } } diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index 2879921bd..e6f21144c 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/crates/ml/src/trainers/dqn/trainer.rs @@ -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()) { diff --git a/crates/ml/src/trainers/ppo.rs b/crates/ml/src/trainers/ppo.rs index 180bab3bc..e64263f01 100644 --- a/crates/ml/src/trainers/ppo.rs +++ b/crates/ml/src/trainers/ppo.rs @@ -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); diff --git a/crates/ml/tests/gpu_kernel_parity_test.rs b/crates/ml/tests/gpu_kernel_parity_test.rs index 33cc1291e..ba6f0147e 100644 --- a/crates/ml/tests/gpu_kernel_parity_test.rs +++ b/crates/ml/tests/gpu_kernel_parity_test.rs @@ -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 diff --git a/crates/ml/tests/smoke_test_real_data.rs b/crates/ml/tests/smoke_test_real_data.rs index 0653bacad..16967184b 100644 --- a/crates/ml/tests/smoke_test_real_data.rs +++ b/crates/ml/tests/smoke_test_real_data.rs @@ -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)