diff --git a/Cargo.lock b/Cargo.lock index 8ffa4507f..e383b21f6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1862,6 +1862,20 @@ name = "bytemuck" version = "1.25.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] [[package]] name = "byteorder" @@ -2700,8 +2714,11 @@ dependencies = [ name = "cudarc" version = "0.19.3" dependencies = [ + "float4", + "float8", "half", "libloading 0.9.0", + "no-std-compat", ] [[package]] @@ -3634,6 +3651,21 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ce81f49ae8a0482e4c55ea62ebbd7e5a686af544c00b9d090bba3ff9be97b3d" +[[package]] +name = "float4" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5939bac0ef2ad7c83a53e4fb889c1d81f007b07061d648cd271071984d86f257" + +[[package]] +name = "float8" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2d1f04709a8ac06e8e8042875a3c466cc4832d3c1a18dbcb9dba3c6e83046bc" +dependencies = [ + "half", +] + [[package]] name = "flume" version = "0.11.1" @@ -4035,6 +4067,21 @@ dependencies = [ "zeroize", ] +[[package]] +name = "fxt-backtest" +version = "1.0.0" +dependencies = [ + "anyhow", + "clap", + "ml-alpha", + "ml-backtesting", + "ml-core", + "serde", + "serde_yaml", + "tracing", + "tracing-subscriber", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -6032,6 +6079,7 @@ dependencies = [ "anyhow", "approx", "arrow 56.2.0", + "bincode", "clap", "cudarc", "data", @@ -6066,11 +6114,24 @@ name = "ml-backtesting" version = "1.0.0" dependencies = [ "anyhow", + "arrow 56.2.0", + "arrow-array 56.2.0", + "arrow-schema 56.2.0", + "bytemuck", "chrono", "csv", + "cudarc", + "ml-alpha", "ml-core", + "parquet", + "rand 0.8.5", + "rand_chacha 0.3.1", "serde", + "serde_json", + "serde_yaml", "tempfile", + "thiserror 1.0.69", + "tracing", ] [[package]] @@ -6696,6 +6757,12 @@ dependencies = [ "libc", ] +[[package]] +name = "no-std-compat" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b93853da6d84c2e3c7d730d6473e8817692dd89be387eb01b94d7f108ecb5b8c" + [[package]] name = "nom" version = "7.1.3" diff --git a/crates/ml-alpha/src/trainer/per_horizon_state.rs b/crates/ml-alpha/src/trainer/per_horizon_state.rs index e8e114811..78f471882 100644 --- a/crates/ml-alpha/src/trainer/per_horizon_state.rs +++ b/crates/ml-alpha/src/trainer/per_horizon_state.rs @@ -28,6 +28,8 @@ use crate::trainer::optim::AdamW; const PROB_BLEND_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/per_horizon_prob_blend.cubin")); +const REDUCE_AXIS0_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/reduce_axis0.cubin")); pub struct PerHorizonTrainState { // Kernel bindings. @@ -77,6 +79,12 @@ pub struct PerHorizonTrainState { /// d_α [N_HORIZONS] — what AdamW updates for α. Written by the /// reduce kernel after BCE backward has populated grad_probs_per_k. pub d_alpha_reduced_d: CudaSlice, + // reduce_axis0 GPU kernel — collapses [B, n_tail] per-batch + // gradient scratch into shared [n_tail] grads. Must be GPU + // (capture-safe); host-side reduction is forbidden during graph + // capture per pearl_no_host_branches_in_captured_graph. + _reduce_axis0_module: Arc, + reduce_axis0_fn: CudaFunction, } impl PerHorizonTrainState { @@ -102,6 +110,13 @@ impl PerHorizonTrainState { .load_function("per_horizon_prob_blend_reduce_alpha_residual") .context("per_horizon_prob_blend_reduce_alpha_residual")?; + let reduce_axis0_module = ctx + .load_cubin(REDUCE_AXIS0_CUBIN.to_vec()) + .context("load reduce_axis0 cubin")?; + let reduce_axis0_fn = reduce_axis0_module + .load_function("reduce_axis0") + .context("reduce_axis0 symbol")?; + let mut rng = ChaCha8Rng::seed_from_u64(seed); // Q_h: Xavier-style init at 1/sqrt(HIDDEN_DIM). Per the spec @@ -172,6 +187,8 @@ impl PerHorizonTrainState { prob_blend_fwd_fn, prob_blend_reduce_fn, d_alpha_reduced_d, + _reduce_axis0_module: reduce_axis0_module, + reduce_axis0_fn, }) } @@ -222,7 +239,8 @@ impl PerHorizonTrainState { .launch(cfg) .context("per_horizon_prob_blend_fwd")?; } - self.stream.synchronize()?; + // No synchronize: same-stream issue order is sufficient and + // synchronize is illegal during CUDA Graph capture. Ok(()) } @@ -271,7 +289,8 @@ impl PerHorizonTrainState { .launch(cfg_per_h) .context("per_horizon_prob_blend_reduce_alpha_residual")?; } - self.stream.synchronize()?; + // No synchronize: same-stream issue order suffices and + // synchronize is illegal during CUDA Graph capture. // 2. Residual-head bwd: consumes d_residual → produces // d_w_res_scratch, d_bias_res_scratch, d_context. @@ -296,30 +315,60 @@ impl PerHorizonTrainState { Ok(()) } + /// Reduce per-batch scratches into shared grad buffers using the + /// `reduce_axis0` GPU kernel. Capture-safe: no host allocations, + /// no D↔H copies, no host-side arithmetic. One kernel launch per + /// scratch (3 total): grad_q_h, grad_w_res, grad_bias_res. fn reduce_per_batch_scratches_to_shared(&mut self) -> Result<()> { - // Read per-batch scratches; sum over batch axis; push reduced - // shared grads back. For typical n_batch ≤ 64 this is - // negligible compared to the kernel launches. - let n_qh = N_HORIZONS * HIDDEN_DIM; - let n_bres = N_HORIZONS; - let mut qh_scratch = vec![0.0_f32; self.n_batch * n_qh]; - let mut wres_scratch = vec![0.0_f32; self.n_batch * n_qh]; - let mut bres_scratch = vec![0.0_f32; self.n_batch * n_bres]; - self.stream.memcpy_dtoh(&self.grad_q_h_scratch_d, qh_scratch.as_mut_slice())?; - self.stream.memcpy_dtoh(&self.grad_w_res_scratch_d, wres_scratch.as_mut_slice())?; - self.stream.memcpy_dtoh(&self.grad_bias_res_scratch_d, bres_scratch.as_mut_slice())?; + let n_qh = (N_HORIZONS * HIDDEN_DIM) as i32; + let n_bres = N_HORIZONS as i32; + let n_batch_i = self.n_batch as i32; - let mut qh_reduced = vec![0.0_f32; n_qh]; - let mut wres_reduced = vec![0.0_f32; n_qh]; - let mut bres_reduced = vec![0.0_f32; n_bres]; - for b in 0..self.n_batch { - for i in 0..n_qh { qh_reduced[i] += qh_scratch [b * n_qh + i]; } - for i in 0..n_qh { wres_reduced[i] += wres_scratch[b * n_qh + i]; } - for i in 0..n_bres { bres_reduced[i] += bres_scratch[b * n_bres + i]; } + // grad_q_h_scratch [B, n_qh] → grad_q_h_d [n_qh] + { + let cfg = LaunchConfig { + grid_dim: (n_qh as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn); + launch + .arg(&self.grad_q_h_scratch_d) + .arg(&n_batch_i) + .arg(&n_qh) + .arg(&mut self.grad_q_h_d); + unsafe { launch.launch(cfg).context("reduce_axis0 grad_q_h")?; } + } + // grad_w_res_scratch [B, n_qh] → grad_w_res_d [n_qh] + { + let cfg = LaunchConfig { + grid_dim: (n_qh as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn); + launch + .arg(&self.grad_w_res_scratch_d) + .arg(&n_batch_i) + .arg(&n_qh) + .arg(&mut self.grad_w_res_d); + unsafe { launch.launch(cfg).context("reduce_axis0 grad_w_res")?; } + } + // grad_bias_res_scratch [B, n_bres] → grad_bias_res_d [n_bres] + { + let cfg = LaunchConfig { + grid_dim: (n_bres as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn); + launch + .arg(&self.grad_bias_res_scratch_d) + .arg(&n_batch_i) + .arg(&n_bres) + .arg(&mut self.grad_bias_res_d); + unsafe { launch.launch(cfg).context("reduce_axis0 grad_bias_res")?; } } - self.stream.memcpy_htod(&qh_reduced, &mut self.grad_q_h_d)?; - self.stream.memcpy_htod(&wres_reduced, &mut self.grad_w_res_d)?; - self.stream.memcpy_htod(&bres_reduced, &mut self.grad_bias_res_d)?; Ok(()) } @@ -334,26 +383,20 @@ impl PerHorizonTrainState { } /// Zero all gradient scratch buffers between training steps. + /// Uses device-side `memset_zeros` (capture-safe). Host-zero + + /// memcpy_htod is forbidden during CUDA Graph capture. pub fn zero_grads(&mut self) -> Result<()> { let stream = &self.stream; - let zeros_qh = vec![0.0_f32; self.grad_q_h_scratch_d.len()]; - let zeros_wres = vec![0.0_f32; self.grad_w_res_scratch_d.len()]; - let zeros_bres = vec![0.0_f32; self.grad_bias_res_scratch_d.len()]; - let zeros_alp = vec![0.0_f32; self.grad_alpha_d.len()]; - let zeros_ctx = vec![0.0_f32; self.grad_context_d.len()]; - let zeros_res = vec![0.0_f32; self.grad_residual_d.len()]; - let zeros_qhr = vec![0.0_f32; self.grad_q_h_d.len()]; - let zeros_wrr = vec![0.0_f32; self.grad_w_res_d.len()]; - let zeros_brr = vec![0.0_f32; self.grad_bias_res_d.len()]; - stream.memcpy_htod(&zeros_qh, &mut self.grad_q_h_scratch_d)?; - stream.memcpy_htod(&zeros_wres, &mut self.grad_w_res_scratch_d)?; - stream.memcpy_htod(&zeros_bres, &mut self.grad_bias_res_scratch_d)?; - stream.memcpy_htod(&zeros_alp, &mut self.grad_alpha_d)?; - stream.memcpy_htod(&zeros_ctx, &mut self.grad_context_d)?; - stream.memcpy_htod(&zeros_res, &mut self.grad_residual_d)?; - stream.memcpy_htod(&zeros_qhr, &mut self.grad_q_h_d)?; - stream.memcpy_htod(&zeros_wrr, &mut self.grad_w_res_d)?; - stream.memcpy_htod(&zeros_brr, &mut self.grad_bias_res_d)?; + stream.memset_zeros(&mut self.grad_q_h_scratch_d)?; + stream.memset_zeros(&mut self.grad_w_res_scratch_d)?; + stream.memset_zeros(&mut self.grad_bias_res_scratch_d)?; + stream.memset_zeros(&mut self.grad_alpha_d)?; + stream.memset_zeros(&mut self.grad_context_d)?; + stream.memset_zeros(&mut self.grad_residual_d)?; + stream.memset_zeros(&mut self.grad_q_h_d)?; + stream.memset_zeros(&mut self.grad_w_res_d)?; + stream.memset_zeros(&mut self.grad_bias_res_d)?; + stream.memset_zeros(&mut self.d_alpha_reduced_d)?; Ok(()) } }