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:
@@ -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.
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user