Optimization/gpu utilization

[directedLink]

This PR significantly optimizes the training throughput of the HOPE model (from ~20 to ~370 tokens/s on an RTX 4090) by resolving a critical CPU-GPU synchronization bottleneck in the surprise computation path. All core mechanisms of the Nested Learning (NL) architecture are preserved.

The Problem:

The original implementation triggered over 150 synchronous .item() calls per training step (one per block per forward pass).

Root Cause: Each .item() call forces the CPU to wait for the GPU to finish, resulting in extremely low GPU utilization (5-10%) and "stuttering" training.

Impact: Training was 18x slower than optimal, making experimentation on consumer GPUs like the 4090 nearly impossible.

The Solution:

 Surprise Pre-computation: Refactored training.py to calculate the L2​ norm (surprise value) once per step in the main training loop and pass it as an override to the model.

Redundancy Elimination: This allows the model blocks to skip internal .item() calls, maintaining mathematical equivalence while removing all per-block synchronization points.

Optimized Config: Added configs/pilot_paper_faithful_optimized.yaml which increases batch_size and tunes chunk_size for high-memory GPUs.

Performance Gains (Results on RTX 4090)

Throughput: 20 tokens/s → 370 tokens/s (18.5x improvement).

Training Time: 200 steps in 2.5h (previously 6-7h).

GPU Utilization: 5-10% → 20-30%.

Convergence: Verified that loss drops from 95.5 to 24.0 over 200 steps (logic is intact).

What's Included:

Core Fix: Refactored _compute_surprise_override in src/nested_learning/training.py.

New Config: pilot_paper_faithful_optimized.yaml.

Comprehensive Docs: PERFORMANCE_OPTIMIZATION_SUMMARY.md and EVAL_RESULTS_ANALYSIS.md for detailed technical evidence.

Eval Tools: eval_correct.sh for easy verification.

[kmccleary3301]

Thanks for putting work into performance. Unfortunately, I don't think I can merge this PR in its current form.

I want to keep the optimization effort, but I need it rebased and narrowed:

• only the minimal training-path code changes needed for the speedup
• one small config override for reproduction
• one benchmark script with fixed command and machine details
• before/after metrics captured by the repo’s current logging/eval flow

Right now this PR mixes core changes with large narrative docs and broad benchmark claims, which makes regression review and reproducibility checks harder.

If you open a minimal rebased PR on current main, I will prioritize review.