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

fix: wire max_training_steps_per_epoch to DQN + per-step profiling

Browse Source 
- DQN CLI --max-steps-per-epoch was only wired to PPO (bug: 1394 steps ran instead of 10)
- Added per-substep timing: sample/fused/guard breakdown per epoch
- Phase 1a results: PER sampling now 0.1ms/step (was ~100ms with CPU roundtrips)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-21 11:08:38 +01:00
parent 2c39d100a0
commit f97f4b23e6
2 changed files with 23 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/ml/examples/train_baseline_rl.rs
View File

@@ -560,6 +560,7 @@ fn train_dqn_fold(
// to leave room for training kernels + replay buffer.
gpu_n_episodes: if detect_vram_mb() < 8192 { 16 } else { 256 },
gpu_timesteps_per_episode: if detect_vram_mb() < 8192 { 100 } else { 500 },
max_training_steps_per_epoch: args.max_steps_per_epoch,
..DQNHyperparameters::default()
};

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

@@ -947,9 +947,13 @@ impl DQNTrainer {
// Batch pre-sampling: sample K batches under one READ lock
const PREFETCH_K: usize = 32;
let mut sample_total_us = 0_u64;
let mut fused_total_us = 0_u64;
let mut guard_total_us = 0_u64;
for chunk_start in (0..num_training_steps).step_by(PREFETCH_K) {
let chunk_end = (chunk_start + PREFETCH_K).min(num_training_steps);
let sample_start = std::time::Instant::now();
let batches = {
let agent = self.agent.read().await;
let buffer = agent.memory();
@@ -972,6 +976,7 @@ impl DQNTrainer {
}
b
};
sample_total_us += sample_start.elapsed().as_micros() as u64;
// GPU train steps (single WRITE lock)
{
@@ -980,8 +985,7 @@ impl DQNTrainer {
if accum_steps <= 1 {
for explicit_batch in batches {
// Fused CUDA path: 3 kernels + CUDA Graph, no Candle dispatch.
// Fused training failures are hard errors — no silent fallback.
let fused_start = std::time::Instant::now();
let _gpu_result = if let Some(ref mut fused) = self.fused_ctx {
let batch_data = explicit_batch.as_ref().ok_or_else(|| {
anyhow::anyhow!("No batch data for fused training step")
@@ -993,6 +997,9 @@ impl DQNTrainer {
.map_err(|e| anyhow::anyhow!("Train step FAILED: {e}"))?
};
fused_total_us += fused_start.elapsed().as_micros() as u64;
let guard_start = std::time::Instant::now();
if let Some(ref mut guard) = self.training_guard {
let loss_slice = _gpu_result.loss_cuda_slice()
.map_err(|e| anyhow::anyhow!("guard loss CudaSlice: {e}"))?;
@@ -1018,6 +1025,7 @@ impl DQNTrainer {
})?;
}
}
guard_total_us += guard_start.elapsed().as_micros() as u64;
train_step_count += 1;
self.gradient_logging_step += 1;
}
@@ -1153,6 +1161,18 @@ impl DQNTrainer {
}
}
if train_step_count > 0 {
info!(
"Training step breakdown ({} steps): sample={:.0}ms fused={:.0}ms guard={:.0}ms (per-step: sample={:.1}ms fused={:.1}ms guard={:.1}ms)",
train_step_count,
sample_total_us as f64 / 1000.0,
fused_total_us as f64 / 1000.0,
guard_total_us as f64 / 1000.0,
sample_total_us as f64 / 1000.0 / train_step_count as f64,
fused_total_us as f64 / 1000.0 / train_step_count as f64,
guard_total_us as f64 / 1000.0 / train_step_count as f64,
);
}
Ok(train_step_count)
}