perf: eliminate 86ms/step PER double-sample + guard overhead

Three performance optimizations targeting the 15.6s/epoch training:

1. Single write lock for sample+step (was: read lock → sample → release
   → write lock → step → release = 43ms lock contention per step)
2. Vaccine batch sampled only when needed (every 10th step, not every
   step — saves 43ms on 90% of steps)
3. Training guard runs every 5th step instead of every step (12ms/step
   → 2.4ms/step amortized). NaN/loss checks persist across steps.
4. Causal intervention interval 10→50 (14 forward passes less often)
5. Removed all remaining H100_DIAG tracing::info in training_loop

Expected impact: ~86ms → ~43ms sample, ~12ms → ~2.4ms guard = ~52ms/step
savings = ~3.6s/epoch → target ~12s/epoch.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-06 22:37:47 +02:00
parent e32afd62fb
commit 9f5b7b58d4
2 changed files with 19 additions and 13 deletions

View File

@@ -265,7 +265,7 @@ impl Default for GpuDqnTrainConfig {
pruning_fraction: 0.7,
enable_causal_intervention: true, // always on
causal_weight: 0.1,
causal_intervention_interval: 10,
causal_intervention_interval: 50,
enable_gradient_vaccine: true, // always on
bottleneck_dim: 2,
market_dim: 42, // Default: 42 base features. Overridden to 50 when OFI (MBP-10) enabled.

View File

@@ -1004,13 +1004,11 @@ impl DQNTrainer {
..Default::default()
};
tracing::info!("H100_DIAG: starting collect_experiences_gpu");
let gpu_batch = collector.collect_experiences_gpu(
features_buf, targets_buf, &episode_starts, &config,
).map_err(|e| anyhow::anyhow!(
"GPU zero-roundtrip collection FAILED (no CPU fallback): {e}"
))?;
tracing::info!("H100_DIAG: collect_experiences_gpu done, inserting to PER");
let count = gpu_batch.n_episodes * gpu_batch.timesteps * 2; // counterfactual doubles experiences
@@ -1067,7 +1065,6 @@ impl DQNTrainer {
let agent = self.agent.read().await;
let mut gpu_buf = agent.memory().as_gpu_buffer()
.ok_or_else(|| anyhow::anyhow!("GPU PER buffer required"))?;
tracing::info!("H100_DIAG: PER insert_batch_bf16 starting (total={})", total);
gpu_buf.gpu.insert_batch_bf16(
&gpu_batch.states,
&gpu_batch.next_states,
@@ -1076,7 +1073,6 @@ impl DQNTrainer {
&gpu_batch.dones,
total,
).map_err(|e| anyhow::anyhow!("GPU PER insert_batch: {e}"))?;
tracing::info!("H100_DIAG: PER insert_batch_bf16 done");
}
match collector.read_epoch_dsr_state() {
@@ -1175,23 +1171,28 @@ impl DQNTrainer {
}
for _step in 0..num_training_steps {
// Single write lock for both sample + fused step — eliminates
// read→write lock cycling overhead (~43ms per step from contention).
let mut agent = self.agent.write().await;
let sample_start = std::time::Instant::now();
let (batch, vaccine_batch) = {
let agent = self.agent.read().await;
{
let buffer = agent.memory();
if !buffer.can_sample(self.current_batch_size) {
break;
}
let b = buffer.sample(self.current_batch_size)
.map_err(|e| anyhow::anyhow!("PER sample: {e}"))?;
let vb = buffer.sample(self.current_batch_size).ok();
(b, vb)
}
let batch = agent.memory().sample(self.current_batch_size)
.map_err(|e| anyhow::anyhow!("PER sample: {e}"))?;
// Vaccine batch: only sample when needed (every 10th step)
let vaccine_batch = if train_step_count % 10 == 0 {
agent.memory().sample(self.current_batch_size).ok()
} else {
None
};
sample_total_us += sample_start.elapsed().as_micros() as u64;
{
let mut agent = self.agent.write().await;
let fused_start = std::time::Instant::now();
let fused = self.fused_ctx.as_mut()
.expect("Fused CUDA training is the only production path");
@@ -1204,6 +1205,10 @@ impl DQNTrainer {
fused_total_us += fused_start.elapsed().as_micros() as u64;
let guard_start = std::time::Instant::now();
// Run guard every 5th step — NaN/loss checks add ~12ms/step overhead.
// Critical failures (NaN, loss explosion) persist across steps,
// so checking every 5th step catches them within <1s.
if train_step_count % 5 == 0 {
if let Some(ref mut guard) = self.training_guard {
let fused = self.fused_ctx.as_ref().unwrap();
let loss_raw = fused.loss_gpu_buf().raw_ptr();
@@ -1241,6 +1246,7 @@ impl DQNTrainer {
})?;
}
}
} // end guard frequency check
guard_total_us += guard_start.elapsed().as_micros() as u64;
if train_step_count % 50 == 0 {