ca7201a1ec
Two remaining issues: (1) ~400ms/step instead of ~1-7ms on RTX 3050, (2) Q-value explosion when C51 kicks in at epoch 2. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
108 lines
6.0 KiB
Markdown
108 lines
6.0 KiB
Markdown
# 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:
|
||
```rust
|
||
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
|