Investigation: JAX GPU Performance in Jupyter Notebook Execution

[mmcky]

Summary

During the implementation of RunsOn GPU support for lecture-python-programming.myst (PR QuantEcon/lecture-python-programming.myst#437), we discovered unexpected performance differences between GPU and CPU runners.

Key Finding: A lecture that took 15.56 seconds on CPU took 176.67 seconds on GPU - despite the GPU being correctly detected and working.

Hardware Benchmark Results

We ran identical benchmarks on both GitHub Actions (CPU) and RunsOn (GPU) to diagnose the performance differences.

System Information

Metric	GitHub Actions (CPU)	RunsOn (GPU)
Platform	Linux Azure	Linux AWS
CPU	AMD EPYC 7763 64-Core @ 3281 MHz	Intel Xeon Platinum 8259CL @ 2500 MHz
CPU Cores	4	8
GPU	None	Tesla T4 (15360 MiB)

CPU Performance (Pure Python)

Benchmark	GitHub Actions	RunsOn	Winner
Integer sum (10M)	0.599 sec	0.881 sec	GitHub Actions 1.5x faster
Float sqrt (1M)	0.077 sec	0.107 sec	GitHub Actions 1.4x faster

CPU Performance (NumPy)

Benchmark	GitHub Actions	RunsOn	Winner
Matrix multiply (3000x3000)	0.642 sec	0.224 sec	RunsOn 2.9x faster
Element-wise (50M)	1.686 sec	1.768 sec	~Same

CPU Performance (Numba)

Benchmark	GitHub Actions	RunsOn	Winner
Integer sum warm-up	0.320 sec	0.300 sec	~Same
Integer sum compiled	0.000 sec	0.000 sec	Same
Parallel sum warm-up	0.344 sec	0.382 sec	~Same
Parallel sum compiled	0.012 sec	0.010 sec	~Same

JAX Performance (Direct Script Execution)

Benchmark	GitHub Actions (CPU)	RunsOn (GPU)	Winner
1000x1000 warm-up	0.030 sec	0.079 sec	GitHub Actions
1000x1000 compiled	0.011 sec	0.001 sec	GPU 11x faster
3000x3000 warm-up	0.426 sec	0.645 sec	GitHub Actions
3000x3000 compiled	0.276 sec	0.009 sec	GPU 30x faster
50M element-wise warm-up	0.816 sec	0.118 sec	GPU 7x faster
50M element-wise compiled	0.381 sec	0.002 sec	GPU 190x faster

Key Findings

1. Pure Python is slower on RunsOn - The Intel Xeon @ 2.5 GHz is slower than the AMD EPYC @ 3.3 GHz for single-threaded Python (1.4-1.5x slower).

2. NumPy matrix ops are faster on RunsOn - Likely due to 8 cores vs 4 cores for BLAS parallelization (2.9x faster).

3. GPU (JAX compiled) is massively faster in direct execution - 30-190x faster for compiled operations!

4. JIT compilation overhead is higher on GPU - Warm-up times are longer on GPU due to CUDA kernel compilation.

The Mystery: Why is Jupyter Execution So Much Slower?

The benchmark script (direct Python execution) shows GPU is working correctly with massive speedups. However, when the same JAX code runs through Jupyter Book / Jupyter kernel execution, the total time is ~11x slower than CPU.

Possible Causes to Investigate

1. Jupyter kernel overhead - Each cell execution may trigger additional overhead
2. JIT cache not persisting - JAX compiled kernels may not persist between cells
3. Multiple recompilations - Different array sizes in lectures trigger recompilation
4. Memory management - Jupyter may handle GPU memory differently
5. Cell isolation - Each cell may create new JAX contexts

Relevant Lectures

• numpy_vs_numba_vs_jax.md - 176.67 sec (GPU) vs 15.56 sec (CPU)
• jax_intro.md - Similar pattern expected

Benchmark Script

The benchmark script is available at:

• scripts/benchmark-hardware.py

Related PRs

• QuantEcon/lecture-python-programming.myst#437 - RunsOn GPU support
• QuantEcon/lecture-python-programming.myst#439 - GitHub Actions benchmark comparison

Next Steps

1. Investigate Jupyter kernel execution overhead with JAX
2. Consider if JAX JIT cache can be preserved across cells
3. Test if running JAX lectures as scripts (vs notebooks) shows different performance
4. Evaluate if this affects other lecture series using JAX

/cc @mmcky

[mmcky]

Benchmark Results: CPU vs GPU Comparison

We've completed comprehensive benchmarking across three execution pathways (bare metal, Jupyter kernel, jupyter-book) to understand the GPU performance characteristics.

Test Environment

Environment	Runner	CPU	Cores	GPU
CPU	GitHub Actions ubuntu-latest	AMD EPYC 7763	4	None
GPU	RunsOn g4dn.2xlarge	Intel Xeon Platinum 8259CL	8	Tesla T4 (15GB)

Bare Metal Results (Python Script)

Benchmark	CPU	GPU	Speedup
JAX matmul 3000×3000 (compiled)	0.275s	0.010s	28x faster
JAX elementwise 50M (compiled)	0.381s	0.002s	207x faster
NumPy matmul 3000×3000	0.624s	0.199s	3x faster

Jupyter Kernel Results (nbconvert)

Benchmark	CPU	GPU	Speedup
JAX matmul 3000×3000 (compiled)	0.277s	0.010s	29x faster
JAX elementwise 50M (compiled)	0.382s	0.002s	196x faster

Jupyter-Book Results

Benchmark	CPU	GPU	Speedup
JAX matmul 3000×3000 (compiled)	0.277s	0.009s	30x faster
JAX elementwise 50M (compiled)	0.381s	0.002s	195x faster

Multi-Operation Benchmark (Simulating Lecture Cells)

This test runs 5 operations with different array sizes (100², 500², 1000², 2000², 3000²), each requiring fresh JIT compilation:

Size	CPU	GPU
100×100	0.061s	0.116s
500×500	0.064s	0.170s
1000×1000	0.078s	0.168s
2000×2000	0.185s	0.368s
3000×3000	0.476s	0.176s ✅
Total	1.51s	1.82s

Key Findings

1. ✅ GPU JAX compiled operations are 28-207x faster for large arrays
2. ⚠️ GPU JIT compilation overhead is higher (~2x slower warmup times)
3. ⚠️ Small operations are slower on GPU due to:
   • CUDA kernel launch overhead (~10-50μs per operation)
   • Slower XLA→CUDA compilation vs XLA→CPU
   • Data transfer overhead not amortized for small arrays
4. ✅ All three execution pathways show consistent behavior

Why Lecture Builds Are Slower on GPU

The lectures contain many small code cells with varying array shapes. Each cell triggers:

• Fresh JIT compilation (GPU compilation is slower)
• CUDA kernel launch overhead
• Insufficient parallelism to benefit from GPU architecture

The GPU excels when: Large arrays, repeated operations with warm JIT cache, sustained computation.

The GPU struggles when: Many small operations, constantly changing array shapes, first-time JIT compilations.

Recommendation

For this lecture series (Python Programming), which has many small educational examples, CPU-only builds may be more efficient. GPU builds would be more beneficial for lecture series with larger computational workloads (e.g., quantitative economics lectures with substantive numerical simulations).

Related PRs

• GPU branch: ENH: Enable RunsOn GPU support for lecture builds lecture-python-programming#437
• CPU benchmark branch: DEBUG: Hardware benchmark on GitHub Actions (CPU) lecture-python-programming#439

[mmcky]

Sourcefiles for these benchmarks can be found in https://github.com/QuantEcon/benchmarks