foxhunt/docs/superpowers/specs/2026-04-04-training-step-performance-design.md

Training Step Performance + Q-Value Explosion Fix

Goal: Fix two remaining issues blocking integration tests: (1) each training step takes ~400ms instead of ~1-7ms on RTX 3050, (2) Q-values explode to -1e30 when C51 kicks in at epoch 2.

---

Issue 1: Per-Step Performance (~400ms/step)

Symptoms

• RTX 3050, [64,64] network, batch_size=64, buffer_size=256
• Each training step takes ~400ms (measured by: 468 steps × 400ms = 187s per epoch matches observed ~3 min/epoch)
• Expected: ~1-7ms per step (GPU at 100%, small network, small batch)
• GPU and CPU both at 100% — not a deadlock, just slow
• compute-sanitizer: 0 OOB errors

What we know

• ensemble_heads=1 → ensemble sync NOT the cause for this test
• CUDA graphs captured correctly (log confirms)
• No memory errors
• 3-graph architecture (graph_forward + graph_aux + graph_adam) is sound per audit

Investigation needed

Add per-operation timing to run_full_step() in fused_training.rs:

1. Spectral norm (10 ungraphed kernel launches)
2. train_step_gpu / graph_forward replay
3. HER donor gen
4. graph_aux replay (or ungraphed submit_aux_ops on steps 0-1)
5. Ensemble diversity step (skipped when ensemble_heads=1)
6. Causal intervention (168 kernels every 50 steps)
7. Gradient vaccine (30 kernels every 10 steps)
8. compute_grad_norm_outside_graph
9. Pruning mask
10. replay_adam_and_readback
11. PER priority update
12. RwLock acquire/release (agent.read().await / agent.write().await in training_loop.rs)

Log timing at INFO level for the first 5 steps, then every 50 steps. Pattern:

if step < 5 || step % 50 == 0 {
    info!("Step {} timing: spectral={:.1}ms fwd={:.1}ms aux={:.1}ms adam={:.1}ms total={:.1}ms",
          step, t_spectral, t_fwd, t_aux, t_adam, t_total);
}

Possible root causes

• Spectral norm: 10 ungraphed kernel launches per step. Small kernels → launch overhead dominates. On RTX 3050, each kernel launch has ~5-10µs overhead. 10 × 10µs = 100µs. Not 400ms.
• Causal intervention: 168 kernel launches every 50 steps. Average 3.4/step. Not enough.
• graph_aux capture on step 2: sync_all_streams() at line 916 of capture_graph_aux() forces a full GPU sync. This is a one-time cost, not per-step.
• RwLock contention: agent.read().await and agent.write().await in the step loop. If another task holds the lock, the step blocks. But there's only one thread in this test.
• cudarc overhead: Each .launch_builder() call may have overhead in the cudarc wrapper. With ~30 ungraphed launches per step, this adds up.
• GPU stream synchronization hidden in cudarc: Some cudarc operations may internally synchronize the stream. Need to check.

Fix approach

1. Add timing instrumentation → identify the slow operation
2. Fix the root cause (not speculation)

---

Issue 2: Q-Value Explosion (C51 epoch 2)

Symptoms

• Epoch 1 (alpha=0, pure MSE): Q-values normal (~4.5)
• Epoch 2 (alpha=0.2, 20% C51): Q-values explode to -1e30
• q_clip guard correctly catches this and halts training
• Only affects walk-forward fold 2+ (fold 1 sometimes OK)

What we know

• C51 log_softmax is numerically stable (subtracts max, verified in code)
• C51 gradient: d = isw * (exp(lp) - proj). Log-probs (lp) should be ≤ 0 from log_softmax. exp(lp) should be in [0,1].
• Bellman projection clamps targets to [v_min, v_max] and does bilinear interpolation
• Label smoothing is applied (LABEL_SMOOTHING_EPS)
• Diagnostic printf added to C51 loss kernel (logit check) and grad kernel (gradient check) but untested because test too slow

Investigation needed

1. Profile first (Issue 1) — make the test fast enough to reach epoch 2
2. Then check C51 diagnostics — the printf in kernels WILL fire (they're outside the CUDA graph in the loss/grad kernels which run as part of graph_forward's captured sequence — actually, ARE they inside the graph or not?)
3. Verify: are C51 loss/grad kernels inside graph_forward? If yes, printf won't work (graphed). If no, printf will fire.
4. Alternative: add GPU-side NaN/extreme detection — a guard kernel AFTER C51 loss that checks save_current_lp and save_projected for extreme values, writes a flag to a GPU buffer, read back at epoch boundary.

Possible root causes

• Bellman target outside support range: With support [-240, 240] and reward + γ * Q_next potentially > 240, the projection creates a Dirac at the edge atom. Cross-entropy on a spike distribution has extreme gradients.
• bf16 precision loss in save_current_lp: Log-probs stored as bf16 (line 359 of c51_loss_kernel.cu). bf16 has ~3 decimal digits. A log-prob of -0.001 might round to 0.0 in bf16. exp(0.0) = 1.0 instead of exp(-0.001) ≈ 0.999. Small error but accumulated across atoms...
• Gradient budget clip interacting badly with C51 alpha blend: The gradient budget allocates 70% to C51 when alpha > 0. With alpha=0.2, the C51 grad gets 70% budget × 20% blend = 14% effective weight. But the MSE grad gets 100% × 80% = 80%. The budget math might not account for the alpha blend correctly.

Fix approach

1. Make test fast (Issue 1 first)
2. Verify printf fires in C51 kernels (are they graphed?)
3. If printf doesn't fire, add guard kernel after C51 loss
4. Identify exact numerical operation producing the explosion
5. Fix the root cause (clamping, precision, or budget math)

---

Files to modify

File	Change
crates/ml/src/trainers/dqn/fused_training.rs	Add per-operation timing in run_full_step()
crates/ml/src/trainers/dqn/trainer/training_loop.rs	Add per-step timing in run_training_steps_slices()
crates/ml/src/cuda_pipeline/c51_loss_kernel.cu	Already has diagnostic printf (verify it fires)
crates/ml/src/cuda_pipeline/c51_grad_kernel.cu	Already has diagnostic printf (verify it fires)

Success criteria

1. Training step takes <10ms on RTX 3050 with smoketest profile
2. 2 epochs on 60K bars completes in <30 seconds
3. No Q-value explosion through all 3 walk-forward folds
4. test_dqn_trains_on_es_fut passes