JAX MNIST CNN (1×L4)

cd examples/jax-mnist
super create --sku gpu_1x_l4      # → s-xxxx
super run . --script run.sh --on s-xxxx
super download s-xxxx /workspace/out/jax_mnist.json ./jax_mnist.json
super stop s-xxxx

#!/usr/bin/env bash
# Train a small CNN on MNIST with JAX/Flax INSIDE a supercomputer session — a smoke test
# that the JAX CUDA stack (jit/XLA/GPU) works on supercomputer. See README.md.
set -euo pipefail

# Fail loudly if JAX can't see the GPU (the whole point of this example).
python -c "import jax; d=jax.devices(); print('jax devices:', d); assert d[0].platform=='gpu', 'JAX is on CPU'"

exec python "$(dirname "$0")/train.py"

{
    "sku": "gpu_1x_l4",
    "image": {
        "base": "pytorch/pytorch:2.8.0-cuda12.8-cudnn9-devel",
        "apt": ["build-essential"],
        "pip": ["jax[cuda12]", "flax", "optax"],
        "env": {
            "PYTHONUNBUFFERED": "1",
            "SSL_CERT_FILE": "/etc/ssl/certs/ca-certificates.crt",
            "PIP_CERT": "/etc/ssl/certs/ca-certificates.crt",
            "REQUESTS_CA_BUNDLE": "/etc/ssl/certs/ca-certificates.crt"
        }
    }
}

"""Train a small CNN on MNIST with JAX/Flax to verify the JAX GPU stack.

Downloads MNIST as a single npz, trains a LeNet-style CNN for a few epochs with
Flax + Optax, then writes final test accuracy and per-epoch metrics to a JSON
artifact. A couple of minutes on one L4 — the point is to prove the JAX CUDA
path (jit, XLA, GPU) works end to end, not to chase accuracy.
"""

from __future__ import annotations

import json
import os
import time
import urllib.request
from pathlib import Path

import jax
import jax.numpy as jnp
import numpy as np
import optax
from flax import linen as nn

MNIST_URL = "https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz"
OUT = Path("/workspace/out/jax_mnist.json")
EPOCHS = int(os.environ.get("EPOCHS", "3"))
BATCH = int(os.environ.get("BATCH", "128"))
# Soak knob (stress tier): keep training whole epochs until this many seconds
# have elapsed; the artifact is written only after the loop. 0 = exactly EPOCHS.
STRESS_DURATION_S = int(os.environ.get("SUPERCOMPUTER_STRESS_DURATION_S", "0") or "0")


class CNN(nn.Module):
    """LeNet-style convnet: two conv/pool blocks then two dense layers."""

    @nn.compact
    def __call__(self, x):
        x = nn.relu(nn.Conv(32, (3, 3))(x))
        x = nn.avg_pool(x, (2, 2), (2, 2))
        x = nn.relu(nn.Conv(64, (3, 3))(x))
        x = nn.avg_pool(x, (2, 2), (2, 2))
        x = x.reshape((x.shape[0], -1))
        x = nn.relu(nn.Dense(128)(x))
        return nn.Dense(10)(x)


def load_mnist() -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Fetch the keras MNIST npz and return (x_train, y_train, x_test, y_test)."""
    path = Path("/tmp/mnist.npz")
    if not path.exists():
        urllib.request.urlretrieve(MNIST_URL, path)
    with np.load(path) as d:
        x_train = d["x_train"].astype("float32")[..., None] / 255.0
        x_test = d["x_test"].astype("float32")[..., None] / 255.0
        return x_train, d["y_train"], x_test, d["y_test"]


def make_steps(model: CNN, optimizer: optax.GradientTransformation):
    """Build the jitted train step and accuracy fn closing over model/optimizer."""

    def loss_fn(params, x, y):
        logits = model.apply(params, x)
        return optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()

    @jax.jit
    def train_step(params, opt_state, x, y):
        loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
        updates, opt_state = optimizer.update(grads, opt_state, params)
        return optax.apply_updates(params, updates), opt_state, loss

    @jax.jit
    def accuracy(params, x, y):
        preds = jnp.argmax(model.apply(params, x), axis=-1)
        return (preds == y).mean()

    return train_step, accuracy


def main() -> None:
    x_train, y_train, x_test, y_test = load_mnist()
    model = CNN()
    params = model.init(jax.random.key(0), jnp.ones((1, 28, 28, 1)))
    optimizer = optax.adam(1e-3)
    opt_state = optimizer.init(params)
    train_step, accuracy = make_steps(model, optimizer)

    n = x_train.shape[0]
    history = []
    deadline = time.monotonic() + STRESS_DURATION_S
    epoch = 0
    while epoch < EPOCHS or time.monotonic() < deadline:
        perm = np.random.permutation(n)
        losses = []
        for i in range(0, n - BATCH + 1, BATCH):
            idx = perm[i:i + BATCH]
            params, opt_state, loss = train_step(
                params, opt_state, jnp.asarray(x_train[idx]), jnp.asarray(y_train[idx]))
            losses.append(float(loss))
        acc = float(accuracy(params, jnp.asarray(x_test), jnp.asarray(y_test)))
        history.append({"epoch": epoch, "train_loss": sum(losses) / len(losses), "test_acc": acc})
        print(f"epoch {epoch}: loss={history[-1]['train_loss']:.4f} test_acc={acc:.4f}", flush=True)
        epoch += 1

    final_acc = history[-1]["test_acc"]
    OUT.parent.mkdir(parents=True, exist_ok=True)
    OUT.write_text(json.dumps({
        "framework": "jax", "device": str(jax.devices()[0]),
        "epochs": len(history), "final_test_acc": final_acc, "history": history,
    }, indent=2))
    print(f"wrote {OUT} (final test acc {final_acc:.4f})", flush=True)
    assert final_acc > 0.95, f"accuracy too low for a healthy JAX run: {final_acc}"


if __name__ == "__main__":
    main()