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

fix(ml-alpha): C25 per-horizon path graph-capture safety

Browse Source 
Four code paths in PerHorizonTrainState violated CUDA Graph capture
invariants and tripped CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED on smoke
alpha-perception-qk2p9:

1. stream.synchronize() inside forward_with_blend (illegal in capture)
2. stream.synchronize() inside backward_through_blend (idem)
3. reduce_per_batch_scratches_to_shared used host-side vec allocations
   + memcpy_dtoh + CPU summation + memcpy_htod (forbidden during
   capture per pearl_no_host_branches_in_captured_graph)
4. zero_grads allocated host zero vectors + memcpy_htod each step
   (idem)

Fix:
- Remove both synchronizes; same-stream issue order is sufficient.
- Replace host-side reduction with three reduce_axis0 GPU kernel
  launches (q_h, w_res, bias_res). PerHorizonTrainState now owns its
  own reduce_axis0 cubin handle.
- Replace host-zero memcpy with stream.memset_zeros for all nine
  gradient buffers plus d_alpha_reduced.

Verified locally: per_horizon_full_pipeline_smoke (2/2) and both
numgrad parity tests (attention pool + residual head) pass on RTX 3050.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-18 11:23:10 +02:00
parent f50b466f77
commit adaf275af3
2 changed files with 151 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

67
Cargo.lock generated
View File

@@ -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"

125
crates/ml-alpha/src/trainer/per_horizon_state.rs
View File

@@ -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<f32>,
// 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<CudaModule>,
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(())
}
}