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

feat(cuda): wire GpuTrainingGuard into train_step_with_accumulation (zero-sync accumulators)

Browse Source 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-10 14:28:52 +01:00
parent 38f0bdec00
commit 6bbd2eb5ce
1 changed files with 66 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

91
crates/ml/src/trainers/dqn/trainer.rs
View File

@@ -4758,6 +4758,37 @@ impl DQNTrainer {
all_td_errors.extend(result.td_errors);
all_indices.extend(result.indices);
// GPU guard: check + accumulate loss/grad for this sub-step (borrows
// tensors before the move into gpu_*_tensors below).
#[cfg(feature = "cuda")]
{
if let (Some(ref loss_gpu), Some(ref gn_gpu)) =
(&result.loss_tensor_gpu, &result.grad_norm_gpu)
{
if let Some(ref mut guard) = self.training_guard {
let grad_collapse_threshold =
self.hyperparams.learning_rate as f32
* self.hyperparams.gradient_collapse_multiplier as f32;
let warmup_steps = (self.hyperparams.buffer_size as f64 * 0.2) as u64;
let past_warmup = self.gradient_logging_step as u64 > warmup_steps;
let guard_result = guard.check_and_accumulate(
loss_gpu,
gn_gpu,
1e6_f32,
grad_collapse_threshold,
!past_warmup,
).map_err(|e| anyhow::anyhow!("GPU guard sub-step {}: {e}", step))?;
if guard_result.halt_nan {
return Err(anyhow::anyhow!(
"NaN/Inf at accumulation sub-step {}: loss={}, grad={}",
step, guard_result.raw_loss, guard_result.raw_grad_norm
));
}
}
}
}
#[cfg(feature = "cuda")]
{
if let Some(td_gpu) = result.td_errors_gpu {
@@ -4830,34 +4861,44 @@ impl DQNTrainer {
}
agent.step_replay_buffer();
// Single readback at accumulation boundary
// Single readback at accumulation boundary — prefer GPU guard accumulators
// (zero extra sync), fall back to cat+mean+to_scalar if guard absent.
#[cfg(feature = "cuda")]
let (avg_loss, final_grad_norm_f64) = {
let avg_l = if !gpu_loss_tensors.is_empty() {
let stacked = candle_core::Tensor::cat(&gpu_loss_tensors, 0)
.map_err(|e| anyhow::anyhow!("GPU loss concat: {e}"))?;
let mean_loss: f32 = stacked
.mean_all()
.map_err(|e| anyhow::anyhow!("GPU loss mean: {e}"))?
.to_scalar()
.map_err(|e| anyhow::anyhow!("GPU loss readback: {e}"))?;
mean_loss as f64
if let Some(ref mut guard) = self.training_guard {
let (avg_l, avg_gn) = guard.read_accumulators()
.map_err(|e| anyhow::anyhow!("GPU guard read_accumulators: {e}"))?;
guard.reset_accumulators()
.map_err(|e| anyhow::anyhow!("GPU guard reset: {e}"))?;
(avg_l, avg_gn)
} else {
total_loss / accumulation_steps as f64
};
let avg_gn = if !gpu_grad_tensors.is_empty() {
let stacked = candle_core::Tensor::cat(&gpu_grad_tensors, 0)
.map_err(|e| anyhow::anyhow!("GPU grad concat: {e}"))?;
let mean_gn: f32 = stacked
.mean_all()
.map_err(|e| anyhow::anyhow!("GPU grad mean: {e}"))?
.to_scalar()
.map_err(|e| anyhow::anyhow!("GPU grad readback: {e}"))?;
mean_gn as f64
} else {
final_grad_norm as f64
};
(avg_l, avg_gn)
// CPU fallback: original cat+mean+to_scalar path
let avg_l = if !gpu_loss_tensors.is_empty() {
let stacked = candle_core::Tensor::cat(&gpu_loss_tensors, 0)
.map_err(|e| anyhow::anyhow!("GPU loss concat: {e}"))?;
let mean_loss: f32 = stacked
.mean_all()
.map_err(|e| anyhow::anyhow!("GPU loss mean: {e}"))?
.to_scalar()
.map_err(|e| anyhow::anyhow!("GPU loss readback: {e}"))?;
mean_loss as f64
} else {
total_loss / accumulation_steps as f64
};
let avg_gn = if !gpu_grad_tensors.is_empty() {
let stacked = candle_core::Tensor::cat(&gpu_grad_tensors, 0)
.map_err(|e| anyhow::anyhow!("GPU grad concat: {e}"))?;
let mean_gn: f32 = stacked
.mean_all()
.map_err(|e| anyhow::anyhow!("GPU grad mean: {e}"))?
.to_scalar()
.map_err(|e| anyhow::anyhow!("GPU grad readback: {e}"))?;
mean_gn as f64
} else {
final_grad_norm as f64
};
(avg_l, avg_gn)
}
};
#[cfg(not(feature = "cuda"))]
let (avg_loss, final_grad_norm_f64) = (