perf(dqn): replace O(n²) cumsum with custom CUDA prefix-sum kernel
Candle's Tensor::cumsum(0) internally allocates an [n,n] upper-triangular matrix (9.3 GB for n=50K) causing OOM on GPUs ≤48 GB. Replace with a block-parallel Hillis-Steele scan kernel that is O(n) in time and memory. Additional fixes in this commit: - Break autograd chain leak in loss/grad accumulation via .detach() (was leaking ~32 MB/step across entire training run) - Release features_raw_cuda/targets_raw_cuda after GPU experience collection (~164 MB VRAM reclaimed) - Use softmax eval (temp=0.3) in walk-forward backtest to prevent action collapse causing trades=0 on early-stage models Validated: 1286 tests pass (408 ml-dqn + 878 ml), 0 failures. VRAM stable at 754 MB across 45K+ training steps on RTX 3050 Ti. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -151,6 +151,132 @@ impl candle_core::CustomOp2 for SearchSorted {
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// CumsumOp — O(n) prefix sum replacing candle's O(n²) triu-based cumsum
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Candle `CustomOp1` that computes inclusive prefix sum (cumulative sum) in O(n).
|
||||
///
|
||||
/// Candle's built-in `Tensor::cumsum` creates an [n, n] upper-triangular matrix
|
||||
/// and does matmul — O(n²) in both time and memory. For n=50,000 that's 9.3 GB
|
||||
/// in F32. This kernel uses a sequential scan: O(n) time, O(n) output, zero
|
||||
/// scratch memory.
|
||||
struct CumsumOp {
|
||||
n: usize,
|
||||
}
|
||||
|
||||
impl candle_core::CustomOp1 for CumsumOp {
|
||||
fn name(&self) -> &'static str {
|
||||
"cumsum_sequential"
|
||||
}
|
||||
|
||||
fn cpu_fwd(
|
||||
&self,
|
||||
s: &CpuStorage,
|
||||
l: &Layout,
|
||||
) -> candle_core::Result<(CpuStorage, Shape)> {
|
||||
let data = s.as_slice::<f32>()?;
|
||||
let offset = l.start_offset();
|
||||
let n = self.n;
|
||||
let mut out = Vec::with_capacity(n);
|
||||
let mut sum = 0.0_f32;
|
||||
for i in 0..n {
|
||||
sum += data.get(offset + i).copied().unwrap_or(0.0);
|
||||
out.push(sum);
|
||||
}
|
||||
Ok((CpuStorage::F32(out), Shape::from_dims(&[n])))
|
||||
}
|
||||
|
||||
#[cfg(feature = "cuda")]
|
||||
fn cuda_fwd(
|
||||
&self,
|
||||
s: &candle_core::CudaStorage,
|
||||
l: &Layout,
|
||||
) -> candle_core::Result<(candle_core::CudaStorage, Shape)> {
|
||||
use candle_core::cuda_backend::cudarc;
|
||||
use cudarc::driver::{LaunchConfig, PushKernelArg};
|
||||
use std::sync::OnceLock;
|
||||
|
||||
static PREFIX_SUM_PTX: OnceLock<Result<cudarc::nvrtc::Ptx, String>> = OnceLock::new();
|
||||
let ptx_result = PREFIX_SUM_PTX.get_or_init(|| {
|
||||
let src = include_str!("prefix_sum_kernel.cu");
|
||||
cudarc::nvrtc::compile_ptx(src)
|
||||
.map_err(|e| format!("prefix_sum CUDA kernel compilation failed: {e}"))
|
||||
});
|
||||
let ptx = ptx_result.as_ref().map_err(|e| {
|
||||
candle_core::Error::Cuda(e.clone().into())
|
||||
})?;
|
||||
|
||||
let n = self.n;
|
||||
let dev = &s.device;
|
||||
let stream = dev.cuda_stream();
|
||||
|
||||
let input_slice = s.as_cuda_slice::<f32>()?;
|
||||
let input_view = input_slice.slice(l.start_offset()..);
|
||||
|
||||
let module = stream.context().load_module(ptx.clone()).map_err(|e| {
|
||||
candle_core::Error::Cuda(format!("prefix_sum module load: {e}").into())
|
||||
})?;
|
||||
let func = module.load_function("prefix_sum_kernel").map_err(|e| {
|
||||
candle_core::Error::Cuda(format!("prefix_sum function load: {e}").into())
|
||||
})?;
|
||||
|
||||
let mut output = stream.alloc_zeros::<f32>(n).map_err(|e| {
|
||||
candle_core::Error::Cuda(format!("alloc prefix_sum output: {e}").into())
|
||||
})?;
|
||||
|
||||
// Query actual device limits for block size. Use the maximum threads
|
||||
// per block the GPU supports (typically 1024), clamped to n so we don't
|
||||
// launch idle threads. Shared memory = 4 bytes per thread (float partial sum).
|
||||
let max_threads: u32 = {
|
||||
use candle_core::cuda_backend::cudarc::driver::{result, sys};
|
||||
let ordinal = stream.context().ordinal();
|
||||
// Safety: ordinal comes from an already-initialized CudaContext.
|
||||
unsafe {
|
||||
let cu_dev = result::device::get(ordinal as i32).unwrap_or(0);
|
||||
result::device::get_attribute(
|
||||
cu_dev,
|
||||
sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK,
|
||||
)
|
||||
.map(|v| (v as u32).min(1024))
|
||||
.unwrap_or(256)
|
||||
}
|
||||
};
|
||||
let block_dim = max_threads.min(n as u32).max(1);
|
||||
let shared_bytes = block_dim * 4; // float per thread
|
||||
let config = LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (block_dim, 1, 1),
|
||||
shared_mem_bytes: shared_bytes,
|
||||
};
|
||||
let n_i32 = n as i32;
|
||||
|
||||
// Safety: input_view has >= n elements, output has n elements.
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(&func)
|
||||
.arg(&input_view)
|
||||
.arg(&mut output)
|
||||
.arg(&n_i32)
|
||||
.launch(config)
|
||||
.map_err(|e| {
|
||||
candle_core::Error::Cuda(
|
||||
format!("prefix_sum kernel launch: {e}").into(),
|
||||
)
|
||||
})?;
|
||||
}
|
||||
|
||||
let out_storage = candle_core::CudaStorage::wrap_cuda_slice(output, dev.clone());
|
||||
Ok((out_storage, Shape::from_dims(&[n])))
|
||||
}
|
||||
}
|
||||
|
||||
/// O(n) cumulative sum that stays on GPU. Replaces `tensor.cumsum(0)` which
|
||||
/// internally creates an [n, n] upper-triangular matrix (O(n²) memory).
|
||||
fn gpu_cumsum(tensor: &Tensor, n: usize) -> Result<Tensor, MLError> {
|
||||
Ok(tensor.apply_op1_no_bwd(&CumsumOp { n })?)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// GpuReplayBuffer
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -441,8 +567,9 @@ impl GpuReplayBuffer {
|
||||
let alpha_tensor = Tensor::full(self.config.alpha, &[n], &self.device)?;
|
||||
let prios_alpha = active_prios.pow(&alpha_tensor)?;
|
||||
|
||||
// Cumulative sum for prefix-sum sampling
|
||||
let cumsum = prios_alpha.cumsum(0)?;
|
||||
// O(n) prefix sum via custom CUDA kernel — replaces candle's cumsum(0)
|
||||
// which internally builds an [n,n] upper-triangular matrix (9.3 GB for n=50K).
|
||||
let cumsum = gpu_cumsum(&prios_alpha, n)?;
|
||||
|
||||
// total_sum stays on GPU as a [1] tensor — NO to_vec0 readback.
|
||||
// Priorities are always >= epsilon > 0, so sum is guaranteed positive & finite.
|
||||
@@ -522,8 +649,8 @@ impl GpuReplayBuffer {
|
||||
let alpha_tensor = Tensor::full(self.config.alpha, &[n], &self.device)?;
|
||||
let rank_probs = ranks_tensor.pow(&alpha_tensor)?.recip()?;
|
||||
|
||||
// Cumulative sum — total_sum stays on GPU as [1] tensor
|
||||
let cumsum = rank_probs.cumsum(0)?;
|
||||
// O(n) prefix sum via custom CUDA kernel (same fix as sample_proportional)
|
||||
let cumsum = gpu_cumsum(&rank_probs, n)?;
|
||||
let total_sum_tensor = cumsum.narrow(0, n - 1, 1)?;
|
||||
|
||||
// Random targets in [0, total_sum) — scaled on GPU, no CPU readback
|
||||
|
||||
68
crates/ml-dqn/src/prefix_sum_kernel.cu
Normal file
68
crates/ml-dqn/src/prefix_sum_kernel.cu
Normal file
@@ -0,0 +1,68 @@
|
||||
// GPU inclusive prefix sum (cumulative sum) kernel for PER sampling.
|
||||
//
|
||||
// Replaces candle's O(n²) triu-based cumsum which allocates an [n,n] matrix
|
||||
// (e.g. [50000, 50000] = 9.3 GB in F32). This kernel is O(n) in time and
|
||||
// O(n) in output space with zero scratch allocations.
|
||||
//
|
||||
// Algorithm: block-parallel chunked scan (1 block, up to 1024 threads).
|
||||
// Phase 1: Each thread sequentially scans its chunk → local prefix sums.
|
||||
// Phase 2: Inclusive scan of per-thread partial sums via shared memory
|
||||
// (Hillis-Steele on up to 1024 elements — fits in 4 KB smem).
|
||||
// Phase 3: Each thread adds its predecessor's total to its chunk.
|
||||
//
|
||||
// Scales from RTX 3050 Ti (n=50K, 256 threads) to H100 (n=1M+, 1024 threads)
|
||||
// without requiring multi-block coordination or atomic ops.
|
||||
//
|
||||
// Grid: 1 block × BLOCK_DIM threads. BLOCK_DIM chosen by host (256 or 1024).
|
||||
|
||||
extern "C" __global__ void prefix_sum_kernel(
|
||||
const float* __restrict__ input, // [n] input values
|
||||
float* __restrict__ output, // [n] cumulative sums (inclusive)
|
||||
int n // number of elements
|
||||
) {
|
||||
extern __shared__ float sdata[]; // [blockDim.x] partial sums
|
||||
|
||||
const int tid = threadIdx.x;
|
||||
const int nthreads = blockDim.x;
|
||||
|
||||
// --- Phase 1: each thread sequentially scans its chunk ---
|
||||
// Divide n into nthreads chunks. Last thread handles remainder.
|
||||
const int chunk = (n + nthreads - 1) / nthreads;
|
||||
const int start = tid * chunk;
|
||||
const int end = min(start + chunk, n);
|
||||
|
||||
float local_sum = 0.0f;
|
||||
for (int i = start; i < end; i++) {
|
||||
local_sum += input[i];
|
||||
output[i] = local_sum; // local prefix sum (not yet globally correct)
|
||||
}
|
||||
|
||||
// Store each thread's total into shared memory
|
||||
sdata[tid] = local_sum;
|
||||
__syncthreads();
|
||||
|
||||
// --- Phase 2: inclusive scan of partial sums (Hillis-Steele, O(T log T)) ---
|
||||
// T = nthreads ≤ 1024, so log2(T) ≤ 10 iterations.
|
||||
for (int stride = 1; stride < nthreads; stride <<= 1) {
|
||||
float val = 0.0f;
|
||||
if (tid >= stride) {
|
||||
val = sdata[tid - stride];
|
||||
}
|
||||
__syncthreads();
|
||||
if (tid >= stride) {
|
||||
sdata[tid] += val;
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// sdata[tid] now holds the inclusive prefix sum of all thread totals up to tid.
|
||||
// The prefix for thread tid's chunk = sdata[tid - 1] (or 0 for tid == 0).
|
||||
|
||||
// --- Phase 3: propagate prefix to each thread's local results ---
|
||||
if (tid > 0) {
|
||||
float prefix = sdata[tid - 1];
|
||||
for (int i = start; i < end; i++) {
|
||||
output[i] += prefix;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -3002,8 +3002,10 @@ impl HyperparameterOptimizable for DQNTrainer {
|
||||
))
|
||||
})?;
|
||||
|
||||
// Softmax with low temperature: near-greedy but avoids
|
||||
// action collapse that causes trades=0 on early-stage models.
|
||||
let action_indices =
|
||||
agent_guard.batch_greedy_actions(&batch_tensor)?;
|
||||
agent_guard.batch_softmax_actions(&batch_tensor, 0.3)?;
|
||||
|
||||
// 3. Sequential trade simulation on CPU
|
||||
for (i, &action_idx) in action_indices.iter().enumerate() {
|
||||
|
||||
@@ -2210,6 +2210,28 @@ impl DQNTrainer {
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
let gpu_experiences_collected = false;
|
||||
|
||||
// Free DqnGpuData when GPU experience collector is active — it's only used
|
||||
// in the CPU fallback path. On a 4 GB card this reclaims ~90 MB of VRAM
|
||||
// (BF16 features + F32 targets) plus the candle-retained F32→BF16 conversion
|
||||
// block (~150 MB). The raw cudarc buffers (features_raw_cuda / targets_raw_cuda)
|
||||
// remain for the GPU experience collector.
|
||||
if gpu_experiences_collected && self.gpu_data.is_some() {
|
||||
info!("GPU experience collector active — releasing DqnGpuData to reclaim VRAM");
|
||||
self.gpu_data = None;
|
||||
// Also release the buffer pool's CPU staging vecs (up to 40 MB RSS)
|
||||
if let Some(pool) = self.buffer_pool.take() {
|
||||
drop(pool);
|
||||
}
|
||||
// Release raw cudarc buffers — experiences are in the replay buffer now,
|
||||
// these are only used by the GPU experience collector kernel.
|
||||
// Frees ~164 MB (features: 895K×48×4 + targets: 895K×4×4).
|
||||
if self.features_raw_cuda.is_some() {
|
||||
info!("Releasing features_raw_cuda + targets_raw_cuda (~164 MB VRAM)");
|
||||
self.features_raw_cuda = None;
|
||||
self.targets_raw_cuda = None;
|
||||
}
|
||||
}
|
||||
|
||||
// Skip CPU experience collection if GPU already filled the replay buffer
|
||||
if !gpu_experiences_collected {
|
||||
|
||||
@@ -2910,11 +2932,13 @@ impl DQNTrainer {
|
||||
}
|
||||
};
|
||||
|
||||
// GPU-side accumulation — no readback
|
||||
loss_acc = (&loss_acc + &gpu_result.loss_gpu)
|
||||
.map_err(|e| anyhow::anyhow!("loss acc: {e}"))?;
|
||||
grad_acc = (&grad_acc + &gpu_result.grad_norm_gpu)
|
||||
.map_err(|e| anyhow::anyhow!("grad acc: {e}"))?;
|
||||
// GPU-side accumulation — detach() breaks the autograd chain so
|
||||
// candle doesn't keep every forward/backward graph alive across steps.
|
||||
// Without detach: ~32 MB/step leak (entire computation graph retained).
|
||||
loss_acc = (&loss_acc + &gpu_result.loss_gpu.detach())
|
||||
.map_err(|e| anyhow::anyhow!("loss acc: {e}"))?.detach();
|
||||
grad_acc = (&grad_acc + &gpu_result.grad_norm_gpu.detach())
|
||||
.map_err(|e| anyhow::anyhow!("grad acc: {e}"))?.detach();
|
||||
train_step_count += 1;
|
||||
}
|
||||
|
||||
@@ -3032,11 +3056,11 @@ impl DQNTrainer {
|
||||
}
|
||||
};
|
||||
|
||||
// GPU-side accumulation — no readback
|
||||
loss_acc = (&loss_acc + &gpu_result.loss_gpu)
|
||||
.map_err(|e| anyhow::anyhow!("loss acc: {e}"))?;
|
||||
grad_acc = (&grad_acc + &gpu_result.grad_norm_gpu)
|
||||
.map_err(|e| anyhow::anyhow!("grad acc: {e}"))?;
|
||||
// GPU-side accumulation — detach to prevent autograd chain growth
|
||||
loss_acc = (&loss_acc + &gpu_result.loss_gpu.detach())
|
||||
.map_err(|e| anyhow::anyhow!("loss acc: {e}"))?.detach();
|
||||
grad_acc = (&grad_acc + &gpu_result.grad_norm_gpu.detach())
|
||||
.map_err(|e| anyhow::anyhow!("grad acc: {e}"))?.detach();
|
||||
train_step_count += 1;
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user