Simulated batches for a JAX training loop

Simulated batches for a JAX training loop#

Your network learns from a simulator: eagle runs a fresh batch of physics every step, and JAX trains on it.

planes = [xp.array(xp.from_dlpack(a)) for a in (speed, angle, cd)]     # JAX arrays in, by DLPack
landing = jnp.from_dlpack(sim.simulate(*planes, xp))                   # eagle runs the batch, JAX gets it back
params, state = update(params, state, features, landing)               # your usual jitted training step

Time: ~6 min · Runs on: CPU, GPU if visible (picked up automatically) · You need: eagle, hawk and JAX

Colab or Kaggle: run this first cell (it installs the packages and downloads the helper files this notebook uses). Elsewhere it only checks the imports.

Hide code cell source

 
# Colab setup, generated by python -m raptor_docs.nb colab_setup (do not edit).
import os

try:
    import google.colab  # noqa: F401
    HOSTED = True  # a Colab runtime
except ImportError:
    HOSTED = "KAGGLE_KERNEL_RUN_TYPE" in os.environ  # a Kaggle kernel

if HOSTED:
    import importlib.metadata, pathlib, re, shutil, subprocess, sys, urllib.request

    gpu = bool(shutil.which("nvidia-smi")) and subprocess.run(["nvidia-smi", "-L"], capture_output=True).returncode == 0
    extra = ""
    if gpu:
        # the CUDA major the runtime already has: a CuPy that is installed, else torch's, else the driver's
        major = None
        for m in ("13", "12"):
            try:
                importlib.metadata.version(f"cupy-cuda{m}x")
                major = m
                break
            except importlib.metadata.PackageNotFoundError:
                pass
        if major is None:
            try:
                import torch
                major = "13" if str(torch.version.cuda).startswith("13") else None
            except ImportError:
                pass
        if major is None:
            smi = re.search(r"CUDA Version: (\d+)", subprocess.run(["nvidia-smi"], capture_output=True, text=True).stdout)
            major = "13" if smi and int(smi.group(1)) >= 13 else "12"
        extra = f"[cuda{major}]"
        print(f"GPU found: installing the {extra} extras to match this runtime.")
    pre = ["--pre"] if os.environ.get("RAPTOR_PRE") else []  # optional: test a release candidate
    ref = os.environ.get("RAPTOR_REF", "main")  # optional: fetch the helper files from this git ref
    if pre or "RAPTOR_REF" in os.environ:
        print(f"Pre-release test: installing with --pre, helper files from {ref}.")
    # our packages with --upgrade, so a runtime that kept an older release gets the current one
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--upgrade", *pre, f"raptor-eagle{extra}", f"raptor-hawk{extra}"])
    from packaging.version import Version

    try:
        old = Version(importlib.metadata.version("jax")) < Version("0.11.2")
    except importlib.metadata.PackageNotFoundError:
        old = True
    if old:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", f"jax{extra}>=0.11.2"])
    base = os.environ.get("RAPTOR_RAW_BASE", "https://raw.githubusercontent.com")  # optional: a local mirror
    top = os.environ.get("RAPTOR_WORKDIR") or ("/kaggle/working" if "KAGGLE_KERNEL_RUN_TYPE" in os.environ else "/content")
    root = pathlib.Path(top) / "eagle"
    for rel in (
        "examples/python/datagen_shards.py",
    ):
        (root / rel).parent.mkdir(parents=True, exist_ok=True)
        urllib.request.urlretrieve(f"{base}/amasat01/eagle/{ref}/{rel}", root / rel)
    (root / "docs/content/userguide/tutorials").mkdir(parents=True, exist_ok=True)
    os.chdir(root / "docs/content/userguide/tutorials")
    if not gpu:
        print("Running on the CPU (no GPU in this runtime). For a GPU: " + ("Settings → Accelerator → GPU T4 x2." if "KAGGLE_KERNEL_RUN_TYPE" in os.environ else "Runtime → Change runtime type → T4 GPU."))

The simulator#

A ball thrown at some speed and angle, with air drag that slows it down. The kernel is one time step; eagle repeats it until every ball has landed. The data you want is where each one lands.

import os

os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")  # leave GPU memory for your other libraries
try:  # CUDA 13: CuPy loads NVRTC before JAX does (see Limits on the JAX page)
    import cupy
    cupy.cuda.nvrtc.getVersion()
except Exception:  # no CuPy or no CUDA here
    pass
import inspect
import pathlib
import sys

import jax
import jax.numpy as jnp
import numpy as np

jax.config.update("jax_enable_x64", True)  # the kernel is float64

# the simulator lives in examples/python, so it also runs as a script (see the last section)
sys.path.insert(0, str(pathlib.Path("../../../../examples/python").resolve()))
import datagen_shards as sim

text = inspect.getsource(sim)
print(text[text.index("@hawk.kernel"):text.index("def array_module")].rstrip())

@hawk.kernel(steps=1)
def flight(dt: Param, g: Param, cd: Scalar, terminated: Terminated,
           x: Mutable[Scalar], z: Mutable[Scalar], vx: Mutable[Scalar], vz: Mutable[Scalar]):
    s = sqrt(vx * vx + vz * vz)
    vxn = vx - dt * cd * s * vx
    vzn = vz - dt * (g + cd * s * vz)
    x = x + dt * vxn
    zn = z + dt * vzn
    z = zn
    vx = vxn
    vz = vzn
    terminated = zn < 0.0

With no drag a ball lands at v² sin(2θ) / g. A quick check that eagle’s answer agrees:

xp = sim.array_module() if jax.default_backend() == "gpu" else np   # eagle's planes live where JAX's arrays do
speed, angle = np.array([10.0, 20.0, 30.0]), np.array([0.5, 0.8, 1.0])
got = sim.simulate(xp.asarray(speed), xp.asarray(angle), xp.zeros(3), xp)
got = got.get() if xp is not np else got
exact = speed**2 * np.sin(2 * angle) / sim.G
print("eagle:", got.round(2), " formula:", exact.round(2))
assert np.allclose(got, exact, rtol=0.03)  # the last step overshoots the ground by a fraction of a step

eagle: [ 8.42 40.69 83.31]  formula: [ 8.58 40.76 83.42]

JAX arrays in, JAX arrays out#

simulate takes planes: NumPy arrays, or CuPy arrays on a GPU. A JAX array crosses to one through DLPack, with no host round trip on a GPU. eagle writes into its planes, and a JAX buffer is read-only, so the plane is a copy; the landing distances come back to JAX the same way.

def to_plane(a):
    return xp.array(xp.from_dlpack(a))


def simulated_batch(key, n):
    k1, k2, k3 = jax.random.split(key, 3)
    speed = jax.random.uniform(k1, (n,), minval=5.0, maxval=30.0)
    angle = jax.random.uniform(k2, (n,), minval=0.1, maxval=1.4)
    drag = jax.random.uniform(k3, (n,), minval=0.0, maxval=0.05)
    landing = jnp.from_dlpack(sim.simulate(to_plane(speed), to_plane(angle), to_plane(drag), xp))
    return jnp.stack([speed / 30.0, angle, 20.0 * drag], axis=1), landing / 100.0


features, target = simulated_batch(jax.random.key(0), 1024)
print(features.shape, target.shape, type(target).__name__)

(1024, 3) (1024,) ArrayImpl

Train on a fresh batch each step#

A small network learns the landing distance from speed, angle and drag. Every step draws new throws and has eagle land them, so the network never sees the same ball twice. The loop is ordinary JAX: the simulator sits outside jit, and the update step is jitted.

def init(key, sizes=(3, 64, 64, 1)):
    keys = jax.random.split(key, len(sizes) - 1)
    return [(jax.random.normal(k, (i, o)) / np.sqrt(i), jnp.zeros(o)) for k, i, o in zip(keys, sizes, sizes[1:])]


def predict(params, x):
    for w, b in params[:-1]:
        x = jnp.tanh(x @ w + b)
    w, b = params[-1]
    return (x @ w + b)[:, 0]


def loss(params, x, y):
    return jnp.mean((predict(params, x) - y) ** 2)


@jax.jit
def update(params, m, v, t, x, y, lr=3e-3, b1=0.9, b2=0.999):  # Adam, written out
    l, g = jax.value_and_grad(loss)(params, x, y)
    m = jax.tree.map(lambda m, g: b1 * m + (1 - b1) * g, m, g)
    v = jax.tree.map(lambda v, g: b2 * v + (1 - b2) * g * g, v, g)
    params = jax.tree.map(
        lambda p, m, v: p - lr * (m / (1 - b1**t)) / (jnp.sqrt(v / (1 - b2**t)) + 1e-8), params, m, v)
    return params, m, v, l


params = init(jax.random.key(1))
m = jax.tree.map(jnp.zeros_like, params)
v = jax.tree.map(jnp.zeros_like, params)
key = jax.random.key(2)
for step in range(1, 601):
    key, sub = jax.random.split(key)
    x, y = simulated_batch(sub, 1024)
    params, m, v, l = update(params, m, v, step, x, y)
    if step in (1, 200, 400, 600):
        print(f"step {step:3d}: loss {float(l):.5f}")

step   1: loss 0.12164
step 200: loss 0.00100
step 400: loss 0.00054
step 600: loss 0.00037

x_test, y_test = simulated_batch(jax.random.key(99), 4096)   # throws the network has not seen
err = jnp.abs(predict(params, x_test) - y_test) * 100.0
print(f"mean error {float(err.mean()):.2f} m on landing distances averaging {float(y_test.mean() * 100.0):.1f} m")
assert float(err.mean()) < 3.0

mean error 1.40 m on landing distances averaging 16.4 m

What just happened#

  • eagle ran each batch to the last landing and JAX read the result, with the network trained on distances nobody wrote down by hand.

  • Training ran outside jit around the simulator, and inside jit for the update. For a simulator inside a jitted program see tutorial 18.

One process per GPU#

A bigger dataset can be made by several processes at once, each with its own GPU and its own shard of the data (one process per GPU is one way to use several GPUs; a call also runs on the GPU its arrays live on). datagen_shards.py makes one shard per run (its seed comes from its rank, so no two ranks repeat a throw), and the shards are joined afterwards. The processes share nothing, so a launcher is all you need: here, subprocess.

import subprocess
import tempfile

try:
    ngpu = cupy.cuda.runtime.getDeviceCount()
except Exception:
    ngpu = 0
world = 2
shards = tempfile.mkdtemp()
procs = []
for rank in range(world):
    # one GPU per rank when there are enough of them, else every rank runs on the CPU
    env = {**os.environ, "CUDA_VISIBLE_DEVICES": str(rank) if ngpu >= world else ""}
    cmd = [sys.executable, sim.__file__, "--rank", str(rank), "--world", str(world), "--n", "8192", "--out", shards]
    procs.append(subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE, text=True))
for p in procs:
    print(p.communicate()[0].strip())
    assert p.returncode == 0

parts = [np.load(f"{shards}/shard_{r}.npz") for r in range(world)]
dataset = {k: np.concatenate([p[k] for p in parts]) for k in ("speed", "angle", "cd", "range")}
print({k: v.shape for k, v in dataset.items()})

rank 0/2: 4096 samples on CPU
rank 1/2: 4096 samples on CPU
{'speed': (8192,), 'angle': (8192,), 'cd': (8192,), 'range': (8192,)}

Try this#

Raise world to the number of GPUs you have. Each rank takes --n // --world samples, so the dataset stays the same size and the wall time drops.

Limits#

  • A call uses one GPU: the one its arrays live on (plans, graphs and reductions follow their planes; a JAX function runs on the JAX device it was built under, one function per GPU). Arrays on two GPUs in one call are refused by name.

Next#