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

fix: eliminate cuGraphClone + recursive_confidence atomicAdd — deterministic graph replay

Browse Source 
- capture_child_graph: use std::mem::forget to take ownership of cudarc's
  CudaGraph handles directly. cuGraphClone corrupts cublasLtMatmul HMMA tile
  scheduling state on Hopper, causing graph replay to hang with temporal ops.
- create_eval_forward_exec: instantiate directly from original graph (no clone).
  Multiple execs from the same CUgraph is safe per NVIDIA docs.
- recursive_confidence_backward: replace atomicAdd into grad_buf with per-block
  partial output + deterministic reduce kernel. Zero atomicAdd on gradient path.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-18 10:50:46 +02:00
parent 80769abb9c
commit 6960a02af4
3 changed files with 113 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

76
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -5344,9 +5344,11 @@ extern "C" __global__ void recursive_confidence_forward(
/* Kernel: recursive_confidence_backward — MSE grad into trunk */
/* ================================================================== */
/**
* Per-sample trunk gradient + warp-reduced weight gradient.
* Per-sample trunk gradient + per-block weight gradient partials.
* d_h_s2[i, k] += d_sigmoid * w_conf[k] (plain write — single writer per (i,k)).
* d_w_conf and d_b_conf use warp+block reduction → one atomicAdd per BLOCK.
* Weight gradients are reduced within each block (warp shuffle + shared mem)
* and written to per-block partial buffers. A separate reduce kernel sums
* the partials deterministically — zero atomicAdd.
*
* Grid: ceil(B/256), Block: 256.
*/
@@ -5355,13 +5357,13 @@ extern "C" __global__ void recursive_confidence_backward(
const float* __restrict__ predicted_error, /* [B] */
const float* __restrict__ lagged_td_error, /* [1] pinned — target */
const float* __restrict__ w_conf, /* [SH2] */
float* __restrict__ d_w_conf, /* [SH2] gradient accumulator */
float* __restrict__ d_b_conf, /* [1] gradient accumulator */
float* __restrict__ d_conf_partials, /* [num_blocks, SH2 + 1] block partials */
float* __restrict__ d_h_s2, /* [B, SH2] trunk gradient (accumulate) */
int B, int SH2,
float loss_weight /* 0.01 */
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
int stride = SH2 + 1; /* per-block output: SH2 weight grads + 1 bias grad */
float d_sigmoid = 0.0f;
if (i < B) {
@@ -5376,25 +5378,29 @@ extern "C" __global__ void recursive_confidence_backward(
}
}
/* Weight gradient reduction: d_b_conf = sum_b(d_sigmoid_b) */
float my_db = d_sigmoid;
for (int offset = 16; offset > 0; offset >>= 1)
my_db += __shfl_xor_sync(0xFFFFFFFF, my_db, offset);
__shared__ float ws_db[8];
int warp_id = threadIdx.x / 32;
int lane = threadIdx.x % 32;
if (lane == 0) ws_db[warp_id] = my_db;
__syncthreads();
if (warp_id == 0) {
float val = (lane < blockDim.x / 32) ? ws_db[lane] : 0.0f;
for (int off = 16; off > 0; off >>= 1)
val += __shfl_xor_sync(0xFFFFFFFF, val, off);
if (lane == 0) atomicAdd(d_b_conf, val);
/* Bias gradient: block-reduced partial → d_conf_partials[block, SH2] */
{
float my_db = d_sigmoid;
for (int offset = 16; offset > 0; offset >>= 1)
my_db += __shfl_xor_sync(0xFFFFFFFF, my_db, offset);
__shared__ float ws_db[8];
if (lane == 0) ws_db[warp_id] = my_db;
__syncthreads();
if (warp_id == 0) {
float val = (lane < blockDim.x / 32) ? ws_db[lane] : 0.0f;
for (int off = 16; off > 0; off >>= 1)
val += __shfl_xor_sync(0xFFFFFFFF, val, off);
if (lane == 0)
d_conf_partials[blockIdx.x * stride + SH2] = val;
}
}
/* Weight gradient reduction: d_w_conf[k] = sum_b(d_sigmoid_b * h_s2[b, k]) */
/* Weight gradient: block-reduced partial → d_conf_partials[block, k] */
for (int k = 0; k < SH2; k++) {
float my_dw = (i < B) ? d_sigmoid * h_s2[(long long)i * SH2 + k] : 0.0f;
@@ -5409,12 +5415,44 @@ extern "C" __global__ void recursive_confidence_backward(
float val = (lane < blockDim.x / 32) ? ws_dw[lane] : 0.0f;
for (int off = 16; off > 0; off >>= 1)
val += __shfl_xor_sync(0xFFFFFFFF, val, off);
if (lane == 0) atomicAdd(&d_w_conf[k], val);
if (lane == 0)
d_conf_partials[blockIdx.x * stride + k] = val;
}
__syncthreads();
}
}
/* ================================================================== */
/* Kernel: recursive_confidence_reduce — deterministic block partial */
/* reduction for weight gradients. Zero atomicAdd. */
/* ================================================================== */
/**
* One thread per output parameter (SH2 + 1 total).
* Loops over block partials in fixed order → fully deterministic.
* Grid: ceil((SH2+1)/256), Block: 256.
*/
extern "C" __global__ void recursive_confidence_reduce(
const float* __restrict__ d_conf_partials, /* [num_blocks, SH2 + 1] */
float* __restrict__ d_w_conf, /* [SH2] output */
float* __restrict__ d_b_conf, /* [1] output */
int num_blocks, int SH2)
{
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = SH2 + 1;
if (idx >= total) return;
int stride = SH2 + 1;
float sum = 0.0f;
for (int blk = 0; blk < num_blocks; blk++) {
sum += d_conf_partials[blk * stride + idx];
}
if (idx < SH2)
d_w_conf[idx] += sum;
else
*d_b_conf += sum;
}
/* ================================================================== */
/* Kernel: isv_feature_gate — ISV embedding modulates h_s2 features */
/* ================================================================== */

34
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -1375,6 +1375,11 @@ pub struct GpuDqnTrainer {
// ── Recursive confidence kernels ──
recursive_conf_fwd_kernel: CudaFunction,
recursive_conf_bwd_kernel: CudaFunction,
recursive_conf_reduce_kernel: CudaFunction,
/// Per-block partial buffer for recursive_confidence_backward deterministic reduce.
/// Shape: [max_conf_blocks, SH2 + 1].
recursive_conf_partials: CudaSlice<f32>,
max_conf_blocks: usize,
// ── Trade plan head ──
plan_params_buf: CudaSlice<f32>, // [B, 6]
@@ -2278,13 +2283,13 @@ impl GpuDqnTrainer {
let loss_weight = 0.01_f32;
unsafe {
// Stage 1: per-block partial reduction (no atomicAdd)
self.stream.launch_builder(&self.recursive_conf_bwd_kernel)
.arg(&self.save_h_s2)
.arg(&self.predicted_error_buf)
.arg(&self.lagged_td_error_dev_ptr)
.arg(&w_conf)
.arg(&d_w_conf)
.arg(&d_b_conf)
.arg(&self.recursive_conf_partials)
.arg(&self.bw_d_h_s2)
.arg(&b_i32)
.arg(&sh2)
@@ -2295,6 +2300,23 @@ impl GpuDqnTrainer {
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("recursive_confidence_backward: {e}")))?;
// Stage 2: deterministic reduce across block partials
let num_blocks_i32 = blocks as i32;
let total_params = sh2 + 1;
let reduce_blocks = ((total_params as u32 + 255) / 256).max(1);
self.stream.launch_builder(&self.recursive_conf_reduce_kernel)
.arg(&self.recursive_conf_partials)
.arg(&d_w_conf)
.arg(&d_b_conf)
.arg(&num_blocks_i32)
.arg(&sh2)
.launch(LaunchConfig {
grid_dim: (reduce_blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("recursive_confidence_reduce: {e}")))?;
}
Ok(())
}
@@ -4425,6 +4447,11 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("recursive_confidence_forward load: {e}")))?;
let recursive_conf_bwd_kernel = exp_module_for_mag.load_function("recursive_confidence_backward")
.map_err(|e| MLError::ModelError(format!("recursive_confidence_backward load: {e}")))?;
let recursive_conf_reduce_kernel = exp_module_for_mag.load_function("recursive_confidence_reduce")
.map_err(|e| MLError::ModelError(format!("recursive_confidence_reduce load: {e}")))?;
let max_conf_blocks = ((b + 255) / 256).max(1);
let recursive_conf_partials = stream.alloc_zeros::<f32>(max_conf_blocks * (config.shared_h2 + 1))
.map_err(|e| MLError::ModelError(format!("recursive_conf_partials alloc: {e}")))?;
let trade_plan_fwd_kernel = exp_module_for_mag.load_function("trade_plan_forward")
.map_err(|e| MLError::ModelError(format!("trade_plan_forward load: {e}")))?;
let plan_noise_kernel = exp_module_for_mag.load_function("plan_noise_inject")
@@ -5970,6 +5997,9 @@ impl GpuDqnTrainer {
predicted_error_buf,
recursive_conf_fwd_kernel,
recursive_conf_bwd_kernel,
recursive_conf_reduce_kernel,
recursive_conf_partials,
max_conf_blocks,
plan_params_buf,
trade_plan_fwd_kernel,
plan_noise_kernel,

75
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -53,25 +53,26 @@ pub(crate) struct FusedStepResult {
pub grad_norm: f32,
}
/// Raw `CUgraph` handle (uninstantiated) for child graph composition.
/// Child sub-graph: owns both the CUgraph topology and an instantiated CUgraphExec.
///
/// Children are NOT individually instantiated -- they are composed into
/// a parent graph via `cuGraphAddChildGraphNode` and instantiated once.
/// The eval forward child also gets a separate instantiated exec for
/// deterministic Q-value replay during evaluation.
/// The exec comes directly from cudarc's `end_capture` instantiation — we do NOT
/// clone + re-instantiate because `cuGraphClone` corrupts cublasLtMatmul HMMA tile
/// scheduling state on Hopper (sm_90), causing graph replay to hang.
struct ChildGraph {
graph: cuda_sys::CUgraph,
exec: cuda_sys::CUgraphExec,
}
impl ChildGraph {
fn new(graph: cuda_sys::CUgraph) -> Result<Self> {
/// Create an additional exec from an existing graph (for eval replays).
/// Does NOT clone — instantiates directly from the original graph.
fn instantiate_extra_exec(graph: cuda_sys::CUgraph) -> Result<cuda_sys::CUgraphExec> {
let mut exec: cuda_sys::CUgraphExec = std::ptr::null_mut();
let result = unsafe { cuda_sys::cuGraphInstantiateWithFlags(&mut exec, graph, 0) };
if result != cuda_sys::cudaError_enum::CUDA_SUCCESS {
anyhow::bail!("cuGraphInstantiateWithFlags failed: {result:?}");
}
Ok(Self { graph, exec })
Ok(exec)
}
fn launch(&self, stream: cuda_sys::CUstream) -> Result<()> {
@@ -1389,9 +1390,9 @@ impl FusedTrainingCtx {
/// Capture a single child sub-graph using cudarc stream capture.
///
/// Returns a `ChildGraph` with the raw `CUgraph` (uninstantiated).
/// Children are NOT individually instantiated -- they are composed
/// into the parent graph via `cuGraphAddChildGraphNode`.
/// Takes ownership of cudarc's CudaGraph handles directly via `mem::forget`.
/// We do NOT use `cuGraphClone` — cloning corrupts cublasLtMatmul HMMA tile
/// state on Hopper, causing graph replay to hang when temporal ops are present.
///
/// Uses `disable_event_tracking` around capture to prevent cudarc corruption.
fn capture_child_graph(
@@ -1415,8 +1416,6 @@ impl FusedTrainingCtx {
let result = ops(self);
// end_capture returns a CudaGraph which instantiates the graph.
// We need the raw CUgraph for child composition, so extract via clone.
let graph_result = self.stream.end_capture(
cuda_sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
);
@@ -1429,17 +1428,14 @@ impl FusedTrainingCtx {
.map_err(|e| anyhow::anyhow!("{label} end_capture: {e}"))?
.ok_or_else(|| anyhow::anyhow!("{label} capture returned None"))?;
// Clone the raw CUgraph from cudarc's CudaGraph. cuGraphClone gives
// an independent copy that survives CudaGraph drop.
let mut cloned: cuda_sys::CUgraph = std::ptr::null_mut();
let r = unsafe { cuda_sys::cuGraphClone(&mut cloned, cudarc_graph.cu_graph) };
if r != cuda_sys::cudaError_enum::CUDA_SUCCESS {
anyhow::bail!("{label} cuGraphClone: {r:?}");
}
// CudaGraph drops here, destroying the exec + original graph.
drop(cudarc_graph);
// Take ownership of the raw handles from cudarc's CudaGraph.
// cudarc already instantiated the exec with AUTO_FREE_ON_LAUNCH — use
// it directly. mem::forget prevents CudaGraph::drop from destroying them.
let graph = cudarc_graph.cu_graph;
let exec = cudarc_graph.cu_graph_exec;
std::mem::forget(cudarc_graph);
ChildGraph::new(cloned)
Ok(ChildGraph { graph, exec })
}
/// Capture all 5 children and compose into parent graph.
@@ -1564,46 +1560,23 @@ impl FusedTrainingCtx {
/// Create eval-only forward graph exec for deterministic Q-value replay.
///
/// Clones the forward child graph and instantiates it separately.
/// Injects the exec handle into GpuDqnTrainer for replay_forward_for_q_values().
/// Create eval-only forward graph exec for deterministic Q-value replay.
///
/// Instantiates directly from the forward child's original CUgraph — no
/// cuGraphClone. Multiple execs from the same graph is safe (NVIDIA docs).
fn create_eval_forward_exec(&mut self) -> Result<()> {
let forward_child = self.forward_child.as_ref()
.ok_or_else(|| anyhow::anyhow!("forward_child not captured"))?;
// Clone the forward child graph for independent instantiation.
let mut eval_graph: cuda_sys::CUgraph = std::ptr::null_mut();
let r = unsafe { cuda_sys::cuGraphClone(&mut eval_graph, forward_child.graph) };
if r != cuda_sys::cudaError_enum::CUDA_SUCCESS {
anyhow::bail!("eval forward cuGraphClone: {r:?}");
}
let eval_exec = ChildGraph::instantiate_extra_exec(forward_child.graph)?;
let mut eval_exec: cuda_sys::CUgraphExec = std::ptr::null_mut();
let r = unsafe { cuda_sys::cuGraphInstantiateWithFlags(&mut eval_exec, eval_graph, 0) };
// Destroy the cloned graph -- only the exec is needed.
unsafe { cuda_sys::cuGraphDestroy(eval_graph); }
if r != cuda_sys::cudaError_enum::CUDA_SUCCESS {
anyhow::bail!("eval forward instantiate: {r:?}");
}
// Destroy previous eval exec.
if let Some(old) = self.eval_forward_exec.take() {
unsafe { cuda_sys::cuGraphExecDestroy(old); }
}
self.eval_forward_exec = Some(eval_exec);
// Inject into GpuDqnTrainer for replay_forward_for_q_values().
// Create a separate instantiation from the forward child graph.
let mut trainer_eval_graph: cuda_sys::CUgraph = std::ptr::null_mut();
let r = unsafe { cuda_sys::cuGraphClone(&mut trainer_eval_graph, forward_child.graph) };
if r != cuda_sys::cudaError_enum::CUDA_SUCCESS {
anyhow::bail!("trainer eval forward cuGraphClone: {r:?}");
}
let mut trainer_eval_exec: cuda_sys::CUgraphExec = std::ptr::null_mut();
let r = unsafe { cuda_sys::cuGraphInstantiateWithFlags(&mut trainer_eval_exec, trainer_eval_graph, 0) };
unsafe { cuda_sys::cuGraphDestroy(trainer_eval_graph); }
if r != cuda_sys::cudaError_enum::CUDA_SUCCESS {
anyhow::bail!("trainer eval forward instantiate: {r:?}");
}
let trainer_eval_exec = ChildGraph::instantiate_extra_exec(forward_child.graph)?;
self.trainer.set_eval_forward_exec(trainer_eval_exec);
Ok(())