Files
foxhunt/crates/ml/tests
jgrusewski 9a96691001 perf(dqn): eliminate all GPU→CPU roundtrips from training hot path
Replace per-step to_scalar()/to_vec1() readbacks with GPU-resident
tensor accumulation and single epoch-boundary sync. On H100 this
removes ~12-15 pipeline flushes per training step (~5μs each),
enabling full GPU saturation with zero CPU sync in the inner loop.

Key changes:
- GpuTrainResult: train_step() returns GPU scalar tensors (loss_gpu,
  grad_norm_gpu) instead of f32 — zero readback per step
- Regime conditional: train all 3 heads unconditionally with
  zero-masked weights (mathematical no-op) instead of 3-6
  to_scalar() mask count checks per step
- GPU PER: max_priority as GPU tensor with flush_max_priority(),
  delta-based priority update via index_add (no CPU dedup)
- Deferred diagnostics: NaN detection, CQL logging, Q-value
  estimation all moved to epoch boundary where pipeline is
  already synced
- Legacy CPU-readback wrappers removed entirely

9 files changed, +393/-295 lines. 520 DQN tests passing.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-10 10:05:18 +01:00
..