🚀 JAX + XLA Training Speed Benchmark

A comprehensive benchmark comparing neural network training speed with and without XLA JIT compilation in JAX, using the MNIST dataset.

---

📋 Table of Contents

• Overview
• Project Structure
• Setup
• Usage
• Benchmark Results
• Epoch-by-Epoch Breakdown
• How It Works
• Key Findings
• Further Reading

---

Overview

This project demonstrates the performance impact of XLA (Accelerated Linear Algebra) JIT compilation in JAX by training the same MLP model on MNIST in two modes:

Mode	Description
Without JIT	Eager execution — each op dispatches individually
With JIT	@jax.jit compiled — entire training step fused by XLA

We measure and compare:

• ⏱️ Time per epoch and total training time
• 📉 Training loss convergence
• 🎯 Test accuracy over epochs
• 💾 Memory usage (peak and delta)

---

Project Structure

jax-xla-training-speed/
│
├── README.md                 ← You are here
├── requirements.txt          ← Python dependencies
├── model.py                  ← MLP model definition (pure JAX, no Flax)
├── data_loader.py            ← MNIST loading (pure NumPy, no TensorFlow needed)
├── train_no_jit.py           ← Training loop WITHOUT XLA JIT
├── train_with_jit.py         ← Training loop WITH XLA JIT
├── benchmark.py              ← Full benchmark runner + visualizations
├── utils.py                  ← Timing, memory tracking, logging utilities
│
├── results/
│   ├── training_time.png     ← Per-epoch & total time bar chart
│   ├── loss_curve.png        ← Training loss comparison
│   ├── accuracy_curve.png    ← Test accuracy comparison
│   ├── no_jit_results.json   ← Raw metrics (no JIT)
│   ├── jit_results.json      ← Raw metrics (JIT)
│   └── benchmark_results.json← Combined results
│
├── notebooks/
│   └── experiment.ipynb      ← Interactive Jupyter notebook
│
└── report/
    └── analysis.md           ← Detailed analysis report

---

Setup

Prerequisites

• Python 3.9+
• pip
• Internet connection (first run only — MNIST downloads ~12 MB, cached after)

Installation

# Navigate to the project folder
cd jax-xla-training-speed

# (Recommended) Create a virtual environment
python -m venv venv
venv\Scripts\activate        # Windows
# source venv/bin/activate   # macOS/Linux

# Install dependencies
pip install -r requirements.txt

Dependencies

Package	Version Used	Purpose
jax[cpu]	0.4.38	Core JAX framework
jaxlib	0.4.38	JAX backend (XLA compiler)
tensorflow-datasets	4.9.9	Optional MNIST loader
matplotlib	≥ 3.7.0	Plots & visualizations
numpy	1.26.4	Numerical operations
scipy	1.15.3	JAX dependency
psutil	≥ 5.9.0	Memory usage tracking
tabulate	≥ 0.9.0	Formatted console tables
tqdm	≥ 4.65.0	Progress bars

Note: TensorFlow is NOT required. MNIST is downloaded directly via urllib and cached in ~/.mnist_cache/.

---

Usage

▶️ Run the Full Benchmark (Recommended)

python benchmark.py

This trains both modes back-to-back, prints a side-by-side comparison table, and saves all plots + JSON results to results/.

Run Each Script Individually

python train_no_jit.py    # Without JIT only → saves results/no_jit_results.json
python train_with_jit.py  # With JIT only    → saves results/jit_results.json

Interactive Notebook

pip install jupyter
jupyter notebook notebooks/experiment.ipynb

---

Benchmark Results

Measured on CPU · Python 3.11 · JAX 0.4.38 · 10 epochs · batch size 128 · lr 0.01

⏱️ Summary Table

╒═════════════════════╤═══════════════╤════════════╤═══════════╕
│       Metric        │  Without JIT  │  With JIT  │  Speedup  │
╞═════════════════════╪═══════════════╪════════════╪═══════════╡
│ Total Training Time │    50.21 s    │  29.34 s   │   1.71×   │
├─────────────────────┼───────────────┼────────────┼───────────┤
│ Avg. Time per Epoch │    5.02 s     │   2.93 s   │   1.71×   │
├─────────────────────┼───────────────┼────────────┼───────────┤
│     Final Loss      │    0.2022     │   0.2027   │     —     │
├─────────────────────┼───────────────┼────────────┼───────────┤
│ Final Test Accuracy │    94.48%     │   94.47%   │     —     │
├─────────────────────┼───────────────┼────────────┼───────────┤
│  Peak Memory (MB)   │   1131.2 MB   │  1161.6 MB │     —     │
├─────────────────────┼───────────────┼────────────┼───────────┤
│  Memory Delta (MB)  │   354.0 MB    │  252.0 MB  │     —     │
├─────────────────────┼───────────────┼────────────┼───────────┤
│   JIT Warm-up Time  │      —        │   89.1 ms  │  one-time │
╘═════════════════════╧═══════════════╧════════════╧═══════════╛

📊 Visualizations

Training Time	Loss Curve	Accuracy Curve

---

Epoch-by-Epoch Breakdown

Without JIT

Epoch	Loss	Accuracy	Time
1	1.0003	87.92%	4.137 s
2	0.4245	90.80%	3.369 s
3	0.3403	91.76%	3.697 s
4	0.3003	92.50%	4.809 s
5	0.2740	93.07%	5.602 s
6	0.2545	93.42%	5.182 s
7	0.2381	93.71%	5.655 s
8	0.2246	93.88%	5.700 s
9	0.2131	94.21%	5.840 s
10	0.2022	94.48%	5.815 s

With XLA JIT

Epoch	Loss	Accuracy	Time
1	0.9996	87.97%	2.810 s
2	0.4244	90.76%	2.830 s
3	0.3405	91.88%	3.056 s
4	0.3003	92.52%	2.830 s
5	0.2743	93.04%	2.909 s
6	0.2544	93.31%	2.832 s
7	0.2382	93.61%	2.901 s
8	0.2248	94.09%	2.848 s
9	0.2131	94.29%	3.054 s
10	0.2027	94.47%	2.950 s

JIT warm-up: 89.1 ms (one-time, not included in epoch times)

---

How It Works

Model Architecture

Input Image (28×28×1)
       ↓  flatten
    784 inputs
       ↓  Dense + ReLU
    512 neurons
       ↓  Dense + ReLU
    256 neurons
       ↓  Dense (no activation)
     10 logits  →  Softmax  →  class (0–9)

Total trainable parameters: 535,818

Without JIT — Eager Execution

# Every op dispatches to hardware separately, Python overhead between each
grads = jax.grad(cross_entropy_loss)(params, images, labels)
loss  = cross_entropy_loss(params, images, labels)
params = sgd_update(params, grads, learning_rate)

With JIT — XLA Compiled

@jax.jit
def train_step(params, images, labels, learning_rate):
    # Entire block compiled into one fused XLA kernel
    loss, grads = jax.value_and_grad(cross_entropy_loss)(params, images, labels)
    new_params  = sgd_update(params, grads, learning_rate)
    return new_params, loss

# First call: Python → JAX trace → HLO IR → XLA compile → execute  (~116 ms)
# Every subsequent call: execute cached compiled code               (~3.7 s/epoch)

XLA Optimizations Applied

Without JIT (Eager):
┌──────────┐  ┌──────────┐  ┌──────────┐  ┌──────────┐
│  MatMul  │→ │ BiasAdd  │→ │   ReLU   │→ │  MatMul  │→ ...
│ +dispatch│  │ +dispatch│  │ +dispatch│  │ +dispatch│
└──────────┘  └──────────┘  └──────────┘  └──────────┘
   ↕Python       ↕Python       ↕Python       ↕Python

With JIT (XLA fused):
┌──────────────────────────────────────────────────────────┐
│   MatMul → BiasAdd → ReLU → MatMul → BiasAdd → Softmax  │
│            Single compiled kernel — zero Python overhead  │
└──────────────────────────────────────────────────────────┘

---

Key Findings

1. 🚀 1.71× faster on CPU — XLA JIT reduced total training time from 50.21 s → 29.34 s, saving ~21 seconds over 10 epochs. Speedup is even larger on GPU/TPU.

2. 🎯 Identical model quality — Both modes reached ~94.47% accuracy with nearly identical loss curves, confirming JIT is a pure performance optimization with no effect on numerical results.

3. ⚡ Near-instant JIT warm-up — Compilation took only 89.1 ms — negligible compared to the ~21 s saved. After the first call, every subsequent epoch uses the cached compiled code.

4. 💾 JIT uses slightly more peak memory (1161.6 MB vs 1131.2 MB) to cache compiled graphs, but has a lower memory delta (252.0 MB vs 354.0 MB) during training due to better memory reuse by XLA.

5. ✅ Best practice — Always use @jax.jit in production. Only disable it during debugging via with jax.disable_jit(): ....

---

📖 Further Reading

• JAX Documentation — Just-in-Time Compilation
• XLA: Optimizing Compiler for Machine Learning
• JAX — The Sharp Bits
• How JAX's Autodiff Works

---

License

This project is for educational and benchmarking purposes.