fix: memcpy_dtoh buffer size mismatch in Q-value readback

q_out_buf is [config.batch_size, total_actions] but host readback buffers
were sized for [sample_size, total_actions]. When sample_size < batch_size,
cudarc's memcpy_dtoh assertion (dst.len >= src.len) panics with SIGABRT.

Fix: allocate host buffer to match q_out.len(), then truncate to sample_size.
Affects: compute_epoch_q_diagnostics, compute_validation_loss, PER refresh.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-21 23:58:01 +01:00
parent a75a98bd0d
commit 49444c19d8
2 changed files with 11 additions and 11 deletions

View File

@@ -224,13 +224,13 @@ impl DQNTrainer {
Err(_) => return None,
};
// Download Q-values to host (cold path, epoch boundary only)
// host_q has exactly sample_size * total_actions elements;
// cudarc memcpy_dtoh copies dst.len() elements from device buffer.
// Download Q-values to host (cold path, epoch boundary only).
// q_out is [config.batch_size, total_actions] but we only need [sample_size, total_actions].
// Allocate host buffer to match GPU buffer size, then truncate.
let stream = self.cuda_stream.as_ref()?;
let total_elems = sample_size * total_actions;
let mut host_q = vec![0.0_f32; total_elems];
let mut host_q = vec![0.0_f32; q_out.len()];
stream.memcpy_dtoh(q_out, &mut host_q).ok()?;
host_q.truncate(sample_size * total_actions);
// Compute gap (best - second_best) and per-action averages on host
let cols = total_actions;
@@ -474,11 +474,11 @@ impl DQNTrainer {
.map_err(|e| anyhow::anyhow!("cuBLAS validation forward: {e}"))?;
// Download Q-values to CPU for action selection + Sharpe computation.
// All remaining computation stays on CPU — epoch boundary cold path.
let q_total_elems = sample_size * total_actions;
let mut host_q = vec![0.0_f32; q_total_elems];
// q_out is [config.batch_size, total_actions]; allocate to match, then truncate.
let mut host_q = vec![0.0_f32; q_out.len()];
stream.memcpy_dtoh(q_out, &mut host_q)
.map_err(|e| anyhow::anyhow!("Q-val DtoH: {e}"))?;
host_q.truncate(sample_size * total_actions);
// ── CPU greedy action selection ──
// Direction LUT maps exposure head index → position size direction.

View File

@@ -1392,10 +1392,10 @@ impl DQNTrainer {
let q_result = fused.compute_q_values(&states_gpu, batch_size_refresh);
if let Ok(q_out) = q_result {
// Download Q-values for max-per-row computation (small batch, epoch boundary)
let q_total = batch_size_refresh * total_actions;
let mut host_q = vec![0.0_f32; q_total];
// Download Q-values (q_out may be larger than needed — match GPU buffer size)
let mut host_q = vec![0.0_f32; q_out.len()];
if refresh_stream.memcpy_dtoh(q_out, &mut host_q).is_ok() {
host_q.truncate(batch_size_refresh * total_actions);
let mut max_vals = Vec::with_capacity(batch_size_refresh);
for row in 0..batch_size_refresh {
let offset = row * total_actions;