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

fix: training guard host-side accumulator for epoch metrics

Browse Source 
Guard's check_host uses host-side loss/grad_norm accumulation instead of
GPU kernel. Reads from pinned device-mapped readback (one-step lag).
Fixes 8 smoke test failures from stale zero values.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-19 09:57:23 +02:00
parent e0846a5335
commit 8a58f8410d
2 changed files with 24 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
crates/ml/src/cuda_pipeline/gpu_training_guard.rs
View File

@@ -175,6 +175,9 @@ pub struct GpuTrainingGuard {
guard_mapped: [MappedBuffer; 2],
/// Current write buffer index (alternates 0/1).
guard_buf_idx: usize,
host_loss_sum: f64,
host_grad_sum: f64,
host_steps: u64,
/// Whether at least one step has completed (first step has no previous data).
guard_has_prev: bool,
@@ -239,6 +242,9 @@ impl GpuTrainingGuard {
qvalue_div_func,
guard_mapped,
guard_buf_idx: 0,
host_loss_sum: 0.0,
host_grad_sum: 0.0,
host_steps: 0,
guard_has_prev: false,
qstats_mapped,
qstats_buf_idx: 0,
@@ -354,6 +360,13 @@ impl GpuTrainingGuard {
let clipped_loss = loss.min(clip_threshold);
let halt_loss_clip = loss > clip_threshold && !warmup;
let halt_grad_collapse = grad_norm < collapse_threshold && !warmup;
// Host-side accumulation (replaces GPU kernel accumulator)
if !halt_nan && loss.is_finite() && grad_norm.is_finite() {
self.host_loss_sum += loss as f64;
self.host_grad_sum += grad_norm as f64;
self.host_steps += 1;
}
GuardResult {
halt_nan,
halt_loss_clip,
@@ -369,28 +382,22 @@ impl GpuTrainingGuard {
/// Direct `read_volatile` from mapped pinned memory — no memcpy needed.
/// The kernel's `__threadfence_system()` ensures CPU-visible writes.
pub fn read_accumulators(&mut self) -> Result<(f64, f64), MLError> {
let loss_sum = self.acc_buf.read(0) as f64;
let grad_sum = self.acc_buf.read(1) as f64;
let steps = self.acc_buf.read(2) as f64;
if steps <= 0.0 {
// Host-side accumulators (no GPU kernel needed)
if self.host_steps == 0 {
return Ok((0.0, 0.0));
}
Ok((loss_sum / steps, grad_sum / steps))
let avg_loss = self.host_loss_sum / self.host_steps as f64;
let avg_grad = self.host_grad_sum / self.host_steps as f64;
Ok((avg_loss, avg_grad))
}
/// Zero the accumulator — write zeros directly to mapped host memory.
/// No GPU command needed; the kernel uses `volatile` reads so it will
/// see the zeroed values on the next launch.
pub fn reset_accumulators(&mut self) -> Result<(), MLError> {
// Safety: host_ptr is valid for 3 floats, and volatile kernel reads
// will pick up the new zeros without any explicit cache flush.
unsafe {
std::ptr::write_volatile(self.acc_buf.host_ptr.add(0), 0.0_f32);
std::ptr::write_volatile(self.acc_buf.host_ptr.add(1), 0.0_f32);
std::ptr::write_volatile(self.acc_buf.host_ptr.add(2), 0.0_f32);
}
self.host_loss_sum = 0.0;
self.host_grad_sum = 0.0;
self.host_steps = 0;
Ok(())
}

6
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -1543,9 +1543,9 @@ impl DQNTrainer {
let guard_start = std::time::Instant::now();
// Run guard every 5th step — host-side pinned reads only, zero GPU kernel.
// Loss + grad_norm are in pinned device-mapped memory (written by GPU,
// visible to host after graph replay completes on next step's launch).
if train_step_count % 5 == 0 {
// Loss + grad_norm are from pinned readback (one-step delay due to
// double-buffering). Skip first 5 steps to avoid reading zeros.
if train_step_count >= 5 && train_step_count % 5 == 0 {
if let Some(ref mut guard) = self.training_guard {
let fused = self.fused_ctx.as_ref().unwrap();
let scalars = fused.last_readback_scalars();