Multi-process JAX (jax.distributed) on 8 GPUs

super create --checkpoint --sku gpu_8x_l4
super run . --script run.sh --on <session>

#!/usr/bin/env bash
set -euo pipefail
nvidia-smi --query-gpu=index,name,uuid,memory.total --format=csv
cd "$(dirname "$0")"
unset CUDA_VISIBLE_DEVICES
: "${NUM_PROCESSES:?ENV_MISSING: NUM_PROCESSES}"
[[ "$NUM_PROCESSES" =~ ^[1-9][0-9]*$ ]] || { echo "ENV_INVALID: NUM_PROCESSES" >&2; exit 2; }
pids=()
stop_all() {
  for pid in "${pids[@]}"; do kill "$pid" 2>/dev/null || true; done
}
trap stop_all EXIT
for ((rank = 0; rank < NUM_PROCESSES; rank++)); do
  PROCESS_ID=$rank python3 train.py &
  pids+=("$!")
done
# The first rank to exit non-zero ends the job: the other ranks would otherwise
# block in a collective until the session is stopped.
declare -A running
for pid in "${pids[@]}"; do running[$pid]=1; done
# wait PID returns the saved status of a rank that has already ended, so a
# poll cannot miss an exit the way wait -n does for ranks that ended earlier.
while [ "${#running[@]}" -gt 0 ]; do
  for pid in "${!running[@]}"; do
    if kill -0 "$pid" 2>/dev/null; then continue; fi
    rc=0
    wait "$pid" || rc=$?
    unset "running[$pid]"
    if [ "$rc" -ne 0 ]; then
      echo "RANK_EXIT_NONZERO: pid=$pid rc=$rc" >&2
      exit "$rc"
    fi
  done
  if [ "${#running[@]}" -gt 0 ]; then sleep 1; fi
done
trap - EXIT

"""Required configuration and exact state receipts shared by the JAX examples."""

from dataclasses import dataclass, asdict
import hashlib
import json
import math
import os
from pathlib import Path

import jax
import numpy as np


def required(name, convert=str):
    try:
        text = os.environ[name]
    except KeyError as exc:
        raise RuntimeError(f"ENV_MISSING: {name}") from exc
    if not text.strip():
        raise RuntimeError(f"ENV_EMPTY: {name}")
    try:
        return convert(text)
    except ValueError as exc:
        raise RuntimeError(f"ENV_INVALID: {name}={text!r}") from exc


def integer(name, minimum=1):
    value = required(name, int)
    if value < minimum:
        raise RuntimeError(f"ENV_INVALID: {name} must be >= {minimum}")
    return value


def number(name, minimum=0):
    value = required(name, float)
    if not math.isfinite(value) or value <= minimum:
        raise RuntimeError(f"ENV_INVALID: {name} must be finite and > {minimum}")
    return value


@dataclass(frozen=True)
class Plan:
    steps: int
    seed: int
    batch: int
    microbatches: int
    learning_rate: float
    out_dir: Path
    platform: str
    xla_flags: str

    @classmethod
    def read(cls):
        flags = required("XLA_FLAGS")
        if flags.split() != ["--xla_gpu_deterministic_ops=true"]:
            raise RuntimeError("ENV_INVALID: XLA_FLAGS must select deterministic GPU ops only")
        for name in ("XLA_PYTHON_CLIENT_PREALLOCATE", "XLA_PYTHON_CLIENT_MEM_FRACTION",
                     "XLA_CLIENT_MEM_FRACTION", "XLA_PYTHON_CLIENT_ALLOCATOR"):
            if name in os.environ:
                raise RuntimeError(f"MEMORY_DEFAULT_OVERRIDDEN: {name}")
        platform = required("JAX_PLATFORMS")
        if platform not in ("cuda", "cpu"):
            raise RuntimeError("ENV_INVALID: JAX_PLATFORMS must be cuda or cpu")
        return cls(integer("STEPS"), integer("SEED", 0), integer("BATCH"),
                   integer("MICROBATCHES"), number("LEARNING_RATE"),
                   Path(required("OUT_DIR")).absolute(), platform, flags)

    def check_devices(self, count):
        devices = jax.devices()
        expected = "gpu" if self.platform == "cuda" else "cpu"
        if len(devices) != count or any(d.platform != expected for d in devices):
            raise RuntimeError(f"DEVICE_LAYOUT_INVALID: expected={count}x{expected} got={devices}")
        print(f"JAX_DEVICES {devices}", flush=True)

    def document(self):
        doc = asdict(self)
        doc["out_dir"] = str(self.out_dir)
        return doc


def digest(tree):
    # JAX pytree order fixes dict keys and Optax tuple order; empty states add no leaves.
    leaves = jax.tree_util.tree_leaves(jax.device_get(tree))
    if not leaves:
        raise RuntimeError("STATE_EMPTY")
    result = hashlib.sha256()
    for leaf in leaves:
        array = np.asarray(leaf)
        if array.dtype.hasobject:
            raise RuntimeError("STATE_DTYPE_INVALID: object")
        result.update(array.tobytes(order="C"))
    return result.hexdigest()


def metric(step, loss, params, opt_state):
    value = float(jax.device_get(loss))
    if not math.isfinite(value):
        raise RuntimeError(f"LOSS_NONFINITE: step={step} loss={value!r}")
    return {"step": step, "loss": repr(value), "params_sha256": digest((params, opt_state))}


def append(stream, record):
    stream.write(json.dumps(record, sort_keys=True) + "\n")
    stream.flush()


class Output:
    def __init__(self, plan):
        self.plan = plan
        plan.out_dir.mkdir(parents=True, exist_ok=True)
        if (plan.out_dir / "result.json").exists():
            raise RuntimeError("OUTPUT_ALREADY_COMPLETE")
        self.stream = (plan.out_dir / "metrics.jsonl").open("x")
        self.last = None

    def record(self, record):
        previous = 0 if self.last is None else self.last["step"]
        if record["step"] != previous + 1:
            raise RuntimeError("METRIC_STEP_NOT_CONTIGUOUS")
        append(self.stream, record)
        print(f"METRIC step={record['step']} loss={record['loss']} "
              f"params_sha256={record['params_sha256']}", flush=True)
        self.last = record

    def complete(self, params, opt_state, details):
        if self.last is None or self.last["step"] != self.plan.steps:
            raise RuntimeError("TRAINING_INCOMPLETE")
        self.stream.close()
        document = {**details, "complete": True, "steps": self.plan.steps,
                    "final_loss": self.last["loss"],
                    "params_sha256": self.last["params_sha256"],
                    "final_params_sha256": digest(params),
                    "final_opt_state_sha256": digest(opt_state),
                    "plan": self.plan.document(), "jax_version": jax.__version__,
                    "digest_policy": "every step: device_get(params, opt_state), ordered C-order leaf bytes",
                    "digest_cost": "synchronizes devices; transfers and hashes all parameters and optimizer leaves each step"}
        temporary = self.plan.out_dir / "result.json.tmp"
        temporary.write_text(json.dumps(document, sort_keys=True, indent=2) + "\n")
        temporary.replace(self.plan.out_dir / "result.json")

{
  "sku": "gpu_8x_l4",
  "checkpoint": true,
  "min_disk_gb": 64,
  "image": {
    "base": "nvidia/cuda:12.8.1-devel-ubuntu24.04",
    "apt": [
      "python3",
      "python3-pip",
      "python3-venv"
    ],
    "pip": [
      "jax[cuda12]==0.10.2",
      "jaxlib==0.10.2",
      "jax-cuda12-plugin==0.10.2",
      "jax-cuda12-pjrt==0.10.2",
      "optax==0.2.6",
      "numpy==2.2.6",
      "flax==0.12.5",
      "orbax-checkpoint==0.12.6"
    ],
    "env": {
      "XLA_FLAGS": "--xla_gpu_deterministic_ops=true",
      "JAX_PLATFORMS": "cuda",
      "PYTHONUNBUFFERED": "1",
      "PYTHONHASHSEED": "0",
      "STEPS": "1200",
      "SEED": "42",
      "OUT_DIR": "/workspace/out",
      "LEARNING_RATE": "0.001",
      "BATCH": "4096",
      "MICROBATCHES": "8",
      "NUM_PROCESSES": "8",
      "COORDINATOR_PORT": "12355",
      "FEATURES": "512",
      "WIDTH": "2048",
      "LAYERS": "3",
      "CLASSES": "10",
      "NCCL_SHM_DISABLE": "1",
      "NCCL_SOCKET_IFNAME": "lo",
      "HEARTBEAT_TIMEOUT_SECONDS": "86400"
    }
  }
}

"""One controller per device, globally sharded MLP training and per-rank receipts."""

from dataclasses import dataclass, asdict

import jax
import jax.numpy as jnp
import numpy as np
import optax
from flax import linen as nn
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
from jax.experimental import multihost_utils

from runtime import Plan, Output, integer, append, metric


@dataclass(frozen=True)
class DistributedPlan:
    processes: int
    process_id: int
    port: int
    features: int
    width: int
    layers: int
    classes: int
    heartbeat_timeout_s: int

    @classmethod
    def read(cls, plan):
        result = cls(integer("NUM_PROCESSES"), integer("PROCESS_ID", 0),
                     integer("COORDINATOR_PORT"), integer("FEATURES"), integer("WIDTH"),
                     integer("LAYERS"), integer("CLASSES"),
                     integer("HEARTBEAT_TIMEOUT_SECONDS"))
        if result.process_id >= result.processes:
            raise RuntimeError("ENV_INVALID: PROCESS_ID must be less than NUM_PROCESSES")
        if result.port > 65535:
            raise RuntimeError("ENV_INVALID: COORDINATOR_PORT outside TCP port range")
        if plan.batch % result.processes:
            raise RuntimeError("ENV_INVALID: BATCH must be divisible by NUM_PROCESSES")
        return result


class MLP(nn.Module):
    config: DistributedPlan

    @nn.compact
    def __call__(self, inputs):
        for _ in range(self.config.layers):
            inputs = nn.relu(nn.Dense(self.config.width)(inputs))
        return nn.Dense(self.config.classes)(inputs)


def main():
    plan = Plan.read()
    config = DistributedPlan.read(plan)
    # Hopium freezes every process during an eviction, and JAX's coordination service
    # counts the frozen time as missed heartbeats when the job resumes. With JAX's
    # default (100 s) a restored job aborts; the timeout must exceed the whole
    # freeze-to-resume time, which can reach the checkpoint lifetime (1 day by
    # default) plus restore attempts; a shorter timeout makes the restore fail.
    jax.distributed.initialize(coordinator_address=f"127.0.0.1:{config.port}",
                               num_processes=config.processes, process_id=config.process_id,
                               local_device_ids=[config.process_id],
                               heartbeat_timeout_seconds=config.heartbeat_timeout_s)
    try:
        plan.check_devices(config.processes)
        local_devices = jax.local_devices()
        if len(local_devices) != 1:
            raise RuntimeError(f"LOCAL_DEVICE_COUNT_INVALID: rank={config.process_id} got={local_devices}")
        mesh = Mesh(np.asarray(jax.devices()), ("data",))
        replicated, batched = NamedSharding(mesh, P()), NamedSharding(mesh, P("data"))
        model = MLP(config)
        init_key, data_key = jax.random.split(jax.random.key(plan.seed))
        # Initialization uses one local device; each controller gets the same seeded values.
        with jax.default_device(local_devices[0]):
            params = model.init(init_key, jnp.zeros((1, config.features)))["params"]
        params = jax.tree.map(
            lambda array: jax.make_array_from_process_local_data(
                replicated, np.asarray(jax.device_get(array))), params)
        optimizer = optax.adam(plan.learning_rate)
        opt_state = jax.jit(optimizer.init, out_shardings=replicated)(params)

        def loss_fn(params, inputs, labels):
            logits = model.apply({"params": params}, inputs)
            return optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()

        @jax.jit
        def local_batch(step, microbatch):
            key = jax.random.fold_in(jax.random.fold_in(data_key, step), microbatch)
            indices = config.process_id * (plan.batch // config.processes) + jnp.arange(
                plan.batch // config.processes)

            def sample(index):
                image_key, label_key = jax.random.split(jax.random.fold_in(key, index))
                return (jax.random.normal(image_key, (config.features,)),
                        jax.random.randint(label_key, (), 0, config.classes))

            return jax.vmap(sample)(indices)

        def update(params, opt_state, inputs, labels):
            loss, grads = jax.value_and_grad(loss_fn)(params, inputs, labels)
            updates, opt_state = optimizer.update(grads, opt_state, params)
            return optax.apply_updates(params, updates), opt_state, loss

        update = jax.jit(update, in_shardings=(replicated, replicated, batched, batched),
                         out_shardings=(replicated, replicated, replicated))
        plan.out_dir.mkdir(parents=True, exist_ok=True)
        output = Output(plan) if config.process_id == 0 else None
        with (plan.out_dir / f"rank-{config.process_id}.jsonl").open("x") as rank_stream:
            for step in range(1, plan.steps + 1):
                for microbatch in range(plan.microbatches):
                    with jax.default_device(local_devices[0]):
                        local_inputs, local_labels = local_batch(step, microbatch)
                    inputs = jax.make_array_from_process_local_data(batched, np.asarray(local_inputs))
                    labels = jax.make_array_from_process_local_data(batched, np.asarray(local_labels))
                    params, opt_state, loss = update(params, opt_state, inputs, labels)
                # Global replicated arrays are not fully addressable across controllers.
                # Each local replica represents the entire params/optimizer leaf.
                local_params, local_opt = jax.tree.map(
                    lambda array: array.addressable_shards[0].data, (params, opt_state))
                record = metric(step, loss.addressable_shards[0].data, local_params, local_opt)
                append(rank_stream, {**record, "rank": config.process_id})
                if output is not None:
                    output.record(record)
        multihost_utils.sync_global_devices("training-complete")
        if output is not None:
            output.complete(local_params, local_opt, {
                "dataset": "deterministic globally indexed PRNG features and labels",
                "model": asdict(config), "optimizer": "adam",
                "updates_per_metric_step": plan.microbatches,
                "processes": config.processes,
                "heartbeat_timeout_seconds": config.heartbeat_timeout_s,
                "rank_files": [f"rank-{i}.jsonl" for i in range(config.processes)]})
        multihost_utils.sync_global_devices("artifact-complete")
    finally:
        jax.distributed.shutdown()


if __name__ == "__main__":
    main()