Scale JAX training across multiple GPUs

1. Introduction

In this codelab you keep that same JAX training step and split the work across both GPUs on your node.

There are several ways to split training across devices: data parallelism, tensor or model parallelism, and pipeline parallelism. This codelab covers data parallelism, the simplest starting point. Each GPU receives a different slice of the batch, runs the same model and the same training step, and contributes gradients to a shared update.

The important part is that the training step code barely changes. You change how arrays are placed on devices, and JAX handles the distributed execution.

What you'll do

• Create a Mesh, a logical grid of GPUs with named axes
• Shard training batches and replicate parameters with NamedSharding and PartitionSpec
• Inspect the resulting placement with jax.debug.visualize_array_sharding
• Run the same jax.jit training step on sharded arrays and let JAX parallelize it
• Rewrite the gradient computation with shard_map for explicit per-shard control
• Measure single-GPU against multi-GPU throughput and sweep the global batch size

What you'll need

• A Google Cloud project with billing enabled, and workshop credits or a reservation covering GPU usage
• Quota for at least 2 NVIDIA L4 GPUs in your chosen region (how to check GPU quota)
• An environment where JAX sees two or more GPUs. This codelab stops on the first cell if only one is visible.
• Completion of codelabs 1 through 5, or an equivalent JAX GPU environment. Codelab 4 in particular sets up the Fashion-MNIST cache and the optax training loop reused here.

Estimated time to complete: 60 minutes.

How data-parallel training works

Data parallelism is the simplest multi-GPU strategy. It splits the batch, replicate the model. Here you have the process:

1. Replicate the model parameters so every GPU holds a full copy of the weights.
2. Shard the data batch along the batch dimension and each GPU gets a different slice.
3. Forward + backward on each GPU independently with each computes gradients on its local slice.
4. All-reduce the gradients to average across GPUs so every copy sees the same update.
5. Update parameters identically on every GPU with same gradients means same new weights.

When each GPU has enough local compute, data parallelism can process per_device_batch * num_gpus examples with only a modest increase in step time compared with one GPU processing per_device_batch.

That is the upside. The tradeoff is that each GPU must store a full copy of the model, so data parallelism does not help when the model itself is too large for one device. It also requires gradient synchronization across devices on every step, which can become a bottleneck for very large models, small batches, or slower interconnects.

2. Before you begin

Select your project

In the Google Cloud console, select or create a project with billing enabled.

Open Cloud Shell

Click Activate Cloud Shell (the terminal icon at the top right of the console) to start a Cloud Shell session, then point it at your project:

gcloud config set project <YOUR_PROJECT_ID>

Provision the GPU environment

Run the following in Cloud Shell. Codelab 1 explains each command in detail, including GPU quota and zone requirements.

gcloud services enable \
  container.googleapis.com \
  compute.googleapis.com \
  iam.googleapis.com \
  cloudresourcemanager.googleapis.com \
  logging.googleapis.com \
  monitoring.googleapis.com

git clone https://github.com/Google-Cloud-AI/partner-ai-nvidia.git
cd partner-ai-nvidia/05-workshops/jax-on-gpu/terraform

cp terraform.tfvars.example terraform.tfvars

Edit terraform.tfvars and set project_id = "", then provision the cluster and deploy JupyterLab:

terraform init
terraform apply
$(terraform output -raw get_credentials_command)

cd ..
kubectl apply -f deploy/jupyter.yaml

terraform apply takes about 12 minutes. When it finishes, wait for the Pod and the LoadBalancer, then read the one-time JupyterLab token from the Pod log:

kubectl get pod jax-jupyter -w         # wait for Running, then Ctrl+C
kubectl get svc jax-jupyter-svc -w     # wait for EXTERNAL-IP, then Ctrl+C
kubectl logs jax-jupyter | grep -o 'token=[a-z0-9]*' | head -1

Open http://:8884, paste the token, and create a new Python 3 notebook in /workspace. Every code block in this codelab goes into a cell of that notebook.

Install what this codelab needs

!pip install --quiet optax matplotlib

Set up and verify the GPU

This code imports JAX along with the three sharding primitives you need, Mesh, PartitionSpec, and NamedSharding, all from jax.sharding and prints what JAX can see. The assert len(gpu_devices) >= 2 line is the gate for this codelab: everything after it assumes more than one device, so if only one GPU is visible the code stops here rather than letting later steps fail in confusing ways.

import os
os.environ["LD_LIBRARY_PATH"] = "/usr/local/nvidia/lib64:" + os.environ.get("LD_LIBRARY_PATH", "")

import gzip
import gc
import hashlib
import shutil
import subprocess
import html
import math
import pathlib
import struct
import time
import urllib.request
import warnings
from functools import partial

from IPython.display import HTML, display
import matplotlib.pyplot as plt
import numpy as np

warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", message=".*ml_dtypes.*")
warnings.filterwarnings("ignore", message=".*JAX_PLATFORMS.*")

import jax
import jax.numpy as jnp
import optax
from jax.sharding import Mesh, PartitionSpec as P, NamedSharding

devices = jax.devices()
gpu_devices = [d for d in devices if d.platform == "gpu"]
NUM_DEVICES = len(gpu_devices)

print(f"JAX version:     {jax.__version__}")
print(f"Default backend: {jax.default_backend()}")
print(f"GPU devices:     {gpu_devices}")
print(f"GPU count:       {NUM_DEVICES}")

assert len(gpu_devices) >= 2, (
    f"This lab needs at least 2 GPUs. Found {len(gpu_devices)}. "
    f"Available devices: {devices}"
)


def block_tree(tree):
    """Wait until a PyTree of JAX arrays is ready on device."""
    return jax.block_until_ready(tree)


def drop_device_refs(*names, clear_compilation_cache=False):
    """Drop global references that may hold device buffers, then run cleanup."""
    for name in names:
        globals().pop(name, None)
    gc.collect()
    if clear_compilation_cache and hasattr(jax, "clear_caches"):
        jax.clear_caches()


def show_table(headers, rows, title=None, aligns=None):
    """Render rows as an HTML table."""
    aligns = aligns or ["left"] * len(headers)
    parts = ["<div style='font-family: system-ui; max-width: 980px;'>"]
    if title:
        parts.append(f"<h4 style='margin: 0 0 8px 0;'>{html.escape(title)}</h4>")
    parts.append("<table style='border-collapse: collapse; width: 100%; font-size: 13px;'>")
    parts.append("<thead><tr>")
    for h, a in zip(headers, aligns):
        parts.append(
            f"<th style='text-align:{a}; border-bottom:1px solid #d0d7de; padding:6px;'>"
            f"{html.escape(str(h))}</th>"
        )
    parts.append("</tr></thead><tbody>")
    for row in rows:
        parts.append("<tr>")
        for cell, a in zip(row, aligns):
            parts.append(
                f"<td style='text-align:{a}; border-bottom:1px solid #eef1f4; padding:6px;'>"
                f"{html.escape(str(cell))}</td>"
            )
        parts.append("</tr>")
    parts.append("</tbody></table></div>")
    display(HTML("".join(parts)))


def show_bars(rows, title, unit="", lower_is_better=False):
    """Render (label, value) pairs as a horizontal bar chart in HTML."""
    max_value = max(float(value) for _, value in rows) or 1.0
    color = "#1a7f37" if not lower_is_better else "#0969da"
    parts = ["<div style='font-family: Arial, sans-serif; max-width: 760px;'>"]
    parts.append(f"<h4 style='margin: 0 0 8px 0;'>{html.escape(title)}</h4>")
    for label, value in rows:
        width = max(3, 100 * float(value) / max_value)
        parts.append(
            "<div style='display:grid; grid-template-columns: 190px 1fr 130px; gap: 8px; "
            "align-items:center; margin: 6px 0;'>"
            f"<div style='font-size:13px;'>{html.escape(str(label))}</div>"
            "<div style='background:#f6f8fa; border-radius:6px; overflow:hidden; height:22px;'>"
            f"<div style='height:22px; width:{width:.1f}%; background:{color};'></div></div>"
            f"<div style='font-size:13px; font-variant-numeric: tabular-nums;'>{float(value):,.1f} {html.escape(unit)}</div>"
            "</div>"
        )
    parts.append(
        f"<div style='font-size:12px; color:#57606a;'>"
        f"{'Lower' if lower_is_better else 'Higher'} is better.</div></div>"
    )
    display(HTML("".join(parts)))

You should see a JAX version, gpu as the default backend, and a list of two CUDA devices with a GPU count of 2.

3. Load Fashion-MNIST and define a compute-heavy MLP

This step uses the same Fashion-MNIST dataset as previous lab, but with a compute-heavy model so the multi-GPU effect is easier to see. Dataset loading and host-side batch preparation happen once, before any timing. The benchmark later in this codelab measures only the compiled GPU training step.

Download and prepare the data

DATA_DIR = pathlib.Path.home() / ".cache" / "jax-course" / "fashion-mnist"
DATA_DIR.mkdir(parents=True, exist_ok=True)

FILES = {
    "train-images-idx3-ubyte.gz": "8d4fb7e6c68d591d4c3dfef9ec88bf0d",
    "train-labels-idx1-ubyte.gz": "25c81989df183df01b3e8a0aad5dffbe",
}
PRIMARY_BASE_URL = "https://github.com/zalandoresearch/fashion-mnist/raw/master/data/fashion"

def md5sum(path):
    digest = hashlib.md5()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def download_if_needed(filename, expected_md5):
    path = DATA_DIR / filename
    if path.exists() and md5sum(path) == expected_md5:
        return path
    for base in [PRIMARY_BASE_URL]:
        try:
            print(f"Downloading {filename}")
            urllib.request.urlretrieve(f"{base}/{filename}", path)
            if md5sum(path) != expected_md5:
                raise ValueError("MD5 mismatch")
            return path
        except Exception:
            if path.exists():
                path.unlink()
    raise RuntimeError(f"Could not download {filename}")


def read_idx_images(path):
    with gzip.open(path, "rb") as f:
        _, n, rows, cols = struct.unpack(">IIII", f.read(16))
        return np.frombuffer(f.read(), dtype=np.uint8).reshape(n, rows, cols)


def read_idx_labels(path):
    with gzip.open(path, "rb") as f:
        _, n = struct.unpack(">II", f.read(8))
        return np.frombuffer(f.read(), dtype=np.uint8).reshape(n)


paths = {name: download_if_needed(name, cs) for name, cs in FILES.items()}
train_images = read_idx_images(paths["train-images-idx3-ubyte.gz"])
train_labels = read_idx_labels(paths["train-labels-idx1-ubyte.gz"])

# Shuffle once on the host
perm = np.random.default_rng(0).permutation(len(train_images))
x_train_all = (train_images[perm].astype(np.float32) / 255.0).reshape(len(train_images), -1)
y_train_all = train_labels[perm].astype(np.int32)

drop_device_refs("train_images", "train_labels", "perm")

Define the shared-block model

The model applies one shared hidden block many times. Reusing the same block increases the local compute each GPU does per step without increasing the number of gradient values that must be synchronized across GPUs.

The design turns on that split. Gradient synchronization cost scales with parameter count, while compute scales with how much arithmetic you do per example. BLOCK_REPEATS raises the second without touching the first.

INPUT_DIM = 28 * 28
WIDTH = 1024
NUM_CLASSES = 10
BLOCK_REPEATS = 128
BLOCK_MIX = 0.10
LEARNING_RATE = 3e-4

PER_DEVICE_BATCH = 1024
GLOBAL_BATCH = PER_DEVICE_BATCH * NUM_DEVICES
NUM_TRAIN_BATCHES = 8
BENCHMARK_WARMUP = 4
BENCHMARK_STEPS = 15
BENCHMARK_REPEATS = 3


def init_params(seed=0):
    rng = np.random.default_rng(seed)

    def normal(shape, scale):
        return rng.standard_normal(shape).astype(np.float32) * scale

    return {
        "w_in": normal((INPUT_DIM, WIDTH), math.sqrt(2.0 / INPUT_DIM)),
        "b_in": np.zeros((WIDTH,), dtype=np.float32),
        "w_block": normal((WIDTH, WIDTH), math.sqrt(2.0 / WIDTH)),
        "b_block": np.zeros((WIDTH,), dtype=np.float32),
        "w_out": normal((WIDTH, NUM_CLASSES), math.sqrt(2.0 / WIDTH)),
        "b_out": np.zeros((NUM_CLASSES,), dtype=np.float32),
    }


def make_fashion_batches(batch_size, num_batches=NUM_TRAIN_BATCHES):
    needed = batch_size * num_batches
    if needed > len(x_train_all):
        raise ValueError(
            f"Need {needed:,} examples, but Fashion-MNIST has {len(x_train_all):,}."
        )
    x = x_train_all[:needed].reshape(num_batches, batch_size, INPUT_DIM)
    y = y_train_all[:needed].reshape(num_batches, batch_size)
    return x, y


def model(params, x):
    h = jax.nn.gelu(x @ params["w_in"] + params["b_in"])

    def block(h, _):
        z = jax.nn.gelu(h @ params["w_block"] + params["b_block"])
        h = (1.0 - BLOCK_MIX) * h + BLOCK_MIX * z
        return h, None

    h, _ = jax.lax.scan(block, h, xs=None, length=BLOCK_REPEATS)
    return h @ params["w_out"] + params["b_out"]


def loss_with_metrics(params, batch):
    x, y = batch
    logits = model(params, x)
    loss = optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()
    accuracy = jnp.mean(jnp.argmax(logits, axis=-1) == y)
    return loss, {"accuracy": accuracy}


optimizer = optax.adamw(learning_rate=LEARNING_RATE, weight_decay=1e-4)
param_template = init_params(seed=1)
PARAM_COUNT = sum(x.size for x in param_template.values())
GRADIENT_MB = PARAM_COUNT * np.dtype(np.float32).itemsize / 1e6

drop_device_refs("param_template")

show_table(
    ["", "Value"],
    [
        ("Dataset", f"Fashion-MNIST train ({len(x_train_all):,} examples)"),
        ("Input shape", "28 x 28 grayscale, flattened to 784"),
        ("Model", f"shared-block MLP, width={WIDTH}, repeats={BLOCK_REPEATS}"),
        ("Parameters", f"{PARAM_COUNT:,}"),
        ("Gradient size", f"{GRADIENT_MB:.1f} MB per step"),
        ("Per-GPU batch", PER_DEVICE_BATCH),
        ("Global batch on all GPUs", GLOBAL_BATCH),
        ("Benchmark", f"median of {BENCHMARK_REPEATS} x {BENCHMARK_STEPS} steps"),
    ],
    title="Fashion-MNIST compute-heavy workload",
)

You should see a summary table describing the workload, including the parameter count, the size of the gradients that will be synchronized on every step, and the per-GPU and global batch sizes. On a two-GPU node the global batch is twice the per-GPU batch.

4. Measure the single-GPU baseline

Before adding a second GPU you need a number to compare against. The baseline runs on one GPU with batch size PER_DEVICE_BATCH — exactly the amount of work each GPU will get in the multi-GPU run.

Start by pinning the data and the parameters to a single device:

single_device = gpu_devices[0]

x_batches_1gpu, y_batches_1gpu = make_fashion_batches(PER_DEVICE_BATCH)
x_batches_1gpu = jax.device_put(x_batches_1gpu, single_device)
y_batches_1gpu = jax.device_put(y_batches_1gpu, single_device)

params_1gpu = jax.device_put(init_params(seed=1), single_device)
opt_state_1gpu = optimizer.init(params_1gpu)

Define the training step and the benchmark

train_step is the ordinary single-GPU step with value and gradient, an optax update, and new parameters. And there is not devices or sharding, which is part of this codelab.

benchmark_training warms up first so compilation is not counted, then times three repeats of fifteen steps and reports the median. block_tree is what makes the timing fair with JAX dispatches asynchronously.

@jax.jit
def train_step(params, opt_state, batch):
    (loss, metrics), grads = jax.value_and_grad(loss_with_metrics, has_aux=True)(
        params,
        batch,
    )
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    return params, opt_state, {"loss": loss, "accuracy": metrics["accuracy"]}


def benchmark_training(
    step_fn,
    params,
    opt_state,
    x_batches,
    y_batches,
    warmup=BENCHMARK_WARMUP,
    steps=BENCHMARK_STEPS,
    repeats=BENCHMARK_REPEATS,
):
    """Warm up, then report median steady-state throughput."""
    num_batches = x_batches.shape[0]
    for i in range(warmup):
        batch = (x_batches[i % num_batches], y_batches[i % num_batches])
        params, opt_state, _ = step_fn(params, opt_state, batch)
    block_tree((params, opt_state))

    batch_size = x_batches.shape[1]
    timings = []
    metrics = None
    for repeat in range(repeats):
        start = time.perf_counter()
        for i in range(steps):
            batch_index = (repeat * steps + i) % num_batches
            batch = (x_batches[batch_index], y_batches[batch_index])
            params, opt_state, metrics = step_fn(params, opt_state, batch)
        params, opt_state, metrics = block_tree((params, opt_state, metrics))
        timings.append(time.perf_counter() - start)

    elapsed = float(np.median(timings))
    return {
        "examples_per_sec": steps * batch_size / elapsed,
        "ms_per_step": 1000 * elapsed / steps,
        "final_loss": float(metrics["loss"]),
        "final_accuracy": float(metrics["accuracy"]),
    }

Run the baseline

The last two lines of code below call drop_device_refs, which recurs throughout this codelab: each run allocates parameters, optimizer state, and batches on the GPU, and those buffers stay alive for as long as a Python global still refers to them. Dropping the names and running gc.collect() releases the device memory before the next run allocates its own, so a later step does not fail with an out-of-memory error caused by a run you already finished.

result_1gpu = benchmark_training(
    train_step,
    params_1gpu,
    opt_state_1gpu,
    x_batches_1gpu,
    y_batches_1gpu,
)

show_table(
    ["Metric", "Value"],
    [
        ("GPUs used", "1"),
        ("Batch per step", PER_DEVICE_BATCH),
        ("Throughput", f"{result_1gpu['examples_per_sec']:,.0f} examples/sec"),
        ("Step time", f"{result_1gpu['ms_per_step']:.2f} ms"),
        ("Final loss", f"{result_1gpu['final_loss']:.4f}"),
        ("Final accuracy", f"{100 * result_1gpu['final_accuracy']:.1f}%"),
    ],
    title="Single-GPU baseline",
)

# Keep scalar timing results, but free device buffers from the single-GPU run.
drop_device_refs(
    "params_1gpu",
    "opt_state_1gpu",
    "x_batches_1gpu",
    "y_batches_1gpu",
)

The code takes a little while because it compiles the step, warms up, and then runs 45 timed steps. You should end up with a Single-GPU baseline table reporting throughput in examples per second, step time in milliseconds, and a loss and accuracy from the last step. result_1gpu survives the cleanup because it holds plain Python floats, not device arrays.

5. Create a device mesh

A Mesh maps physical GPUs to a logical grid with named axes. For data parallelism you create a one-dimensional mesh with all GPUs along a single 'data' axis.

Axis names to choose

This codelab calls the mesh axis 'data' because it shards the batch for data parallelism. In larger models you might use names like 'model', 'tensor', 'pipeline', or 'fsdp' to describe other kinds of parallelism. A two-dimensional mesh might use ('data', 'model'), where one axis shards the batch and the other shards model weights or activations.

The names have no special meaning to JAX. They become meaningful only through the PartitionSpecs and the collectives that refer to them.

In practice the three primitives are used together: PartitionSpec describes the layout, NamedSharding attaches that layout to a mesh of devices, and jax.device_put moves an array into that layout.

mesh = Mesh(np.array(gpu_devices), ("data",))

show_table(
    ["", "Value"],
    [
        ("Mesh shape", str(mesh.shape)),
        ("Axis names", str(mesh.axis_names)),
        ("Devices", ", ".join(str(d) for d in mesh.devices.flat)),
    ],
    title="Device mesh",
)

The table should report a single axis named data whose size equals your GPU count, and list both CUDA devices.

Check the GPU topology

Data-parallel training all-reduces gradients on every single step, so the path between the two GPUs sits directly on the critical path. NVLink paths, which nvidia-smi reports as NV*, are much better for this than PHB paths that travel through the host bridge and PCIe.

if shutil.which("nvidia-smi"):
    topo = subprocess.run(
        ["nvidia-smi", "topo", "-m"],
        check=False,
        text=True,
        capture_output=True,
    )
    print(topo.stdout or topo.stderr)
else:
    print("nvidia-smi is not available in this environment.")

You should see a matrix with a row and a column per GPU. On a g2-standard-24 the two L4s are connected over PCIe, so expect the GPU0-to-GPU1 code to report a PHB-class path. That is what this machine type gives you, not a misconfiguration but it is a constraint on how far data-parallel training scales here, and it explains any result you get in the comparison step.

6. Shard the data and replicate the parameters

In data-parallel training there are exactly two placements:

• Data is sharded along the batch dimension and each GPU gets a different slice of the batch
• Parameters are replicated with every GPU holds a full copy so the forward pass runs identically

PartitionSpec('data', None) splits the first dimension across the 'data' mesh axis and replicates the second dimension. PartitionSpec() with no arguments just replicates everything.

The batch arrays have a leading batch dimension, because make_fashion_batches returns all training batches stacked together. That is why they use P(None, "data", None) so we leave dimension 0 whole, shard the examples in dimension 1 across GPUs, and replicate the features.

batch_data_sharding = NamedSharding(mesh, P("data", None))
batch_label_sharding = NamedSharding(mesh, P("data"))
all_data_sharding = NamedSharding(mesh, P(None, "data", None))
all_label_sharding = NamedSharding(mesh, P(None, "data"))
replicated = NamedSharding(mesh, P())

x_batches_multi, y_batches_multi = make_fashion_batches(GLOBAL_BATCH)
x_batches_multi = jax.device_put(x_batches_multi, all_data_sharding)
y_batches_multi = jax.device_put(y_batches_multi, all_label_sharding)

params_multi = jax.device_put(init_params(seed=1), replicated)
opt_state_multi = optimizer.init(params_multi)

print(
    f"Global batch: {GLOBAL_BATCH} examples "
    f"({PER_DEVICE_BATCH} per GPU x {NUM_DEVICES} GPUs)"
)
print(f"Training batches shape: {x_batches_multi.shape}")
print()

print("One training batch: sharded along the batch dimension")
jax.debug.visualize_array_sharding(x_batches_multi[0])

print()
print("Weight w_block: replicated on all GPUs")
jax.debug.visualize_array_sharding(params_multi["w_block"])

jax.debug.visualize_array_sharding prints a text grid showing which GPU holds which portion of the array. On a two-GPU node the first lines of output should read:

Global batch: 2048 examples (1024 per GPU x 2 GPUs)
Training batches shape: (8, 2048, 784)

Below that you should see the batch drawn as two stacked blocks, one labelled for each GPU, and w_block drawn as a single block annotated with both GPUs, sharded data, replicated weights.

7. Run the same jitted step on sharded arrays

The training step code does not change at all. It uses the same compiled train_step you used for the single-GPU baseline works on sharded inputs.

When JAX sees that the batch is sharded across GPUs and the parameters are replicated, it automatically:

1. Runs the forward pass on each GPU's data slice
2. Computes per-shard gradients
3. Inserts an all-reduce to average gradients across GPUs
4. Updates parameters identically on every GPU

You do not write any communication code. The parallelism comes entirely from how the arrays are placed.

result_multi = benchmark_training(
    train_step,
    params_multi,
    opt_state_multi,
    x_batches_multi,
    y_batches_multi,
)

show_table(
    ["Metric", "Value"],
    [
        ("GPUs used", NUM_DEVICES),
        ("Global batch", GLOBAL_BATCH),
        ("Per-GPU batch", PER_DEVICE_BATCH),
        ("Throughput", f"{result_multi['examples_per_sec']:,.0f} examples/sec"),
        ("Step time", f"{result_multi['ms_per_step']:.2f} ms"),
        ("Final loss", f"{result_multi['final_loss']:.4f}"),
        ("Final accuracy", f"{100 * result_multi['final_accuracy']:.1f}%"),
    ],
    title=f"Data-parallel training on {NUM_DEVICES} GPUs",
)

You should get a table shaped like the baseline one, now reporting both a global batch and a per-GPU batch. Resist comparing the throughput numbers by eye — the next step does that properly, and the ratio is the only figure that means anything.

8. Compare single-GPU and multi-GPU throughput

At this point, the per-GPU batch is the same in both runs and the multi-GPU run processes more examples per step because each GPU receives its own shard.

You add a GPU and add workload at the same time, then ask whether throughput keeps up. It is not the same question as "does a fixed batch finish twice as fast". Here faster is about higher training throughput in examples per second.

speed_ratio = result_multi["examples_per_sec"] / result_1gpu["examples_per_sec"]

show_table(
    ["", "1 GPU", f"{NUM_DEVICES} GPUs", "Throughput ratio"],
    [
        ("Per-GPU batch", PER_DEVICE_BATCH, PER_DEVICE_BATCH, "same"),
        ("Global batch", PER_DEVICE_BATCH, GLOBAL_BATCH, f"{NUM_DEVICES}x"),
        (
            "Examples/sec",
            f"{result_1gpu['examples_per_sec']:,.0f}",
            f"{result_multi['examples_per_sec']:,.0f}",
            f"{speed_ratio:.2f}x",
        ),
        (
            "ms/step",
            f"{result_1gpu['ms_per_step']:.2f}",
            f"{result_multi['ms_per_step']:.2f}",
            "",
        ),
    ],
    title="Throughput: same per-GPU batch",
    aligns=["left", "right", "right", "right"],
)

show_bars(
    [
        ("1 GPU", result_1gpu["examples_per_sec"]),
        (f"{NUM_DEVICES} GPUs", result_multi["examples_per_sec"]),
    ],
    "Training throughput (examples/sec)",
    "examples/s",
)

Read the result honestly

The next code branches on what you actually measured. Run it and look at what you get.

step_ratio = result_multi["ms_per_step"] / result_1gpu["ms_per_step"]
if speed_ratio >= 1.0:
    message = (
        f"The multi-GPU run is faster for this Fashion-MNIST workload: "
        f"throughput improves by {speed_ratio:.2f}x. Each GPU still processes "
        f"{PER_DEVICE_BATCH} examples, while the global batch increases from "
        f"{PER_DEVICE_BATCH} to {GLOBAL_BATCH}. Step time changes by {step_ratio:.2f}x, "
        f"so the larger batch translates into higher examples/sec."
    )
else:
    message = (
        f"This run is still communication-bound: throughput changes by {speed_ratio:.2f}x. "
        f"Increase BLOCK_REPEATS or PER_DEVICE_BATCH to give each GPU more local work."
    )

border_color = "#1a7f37" if speed_ratio >= 1.0 else "#d1242f"
display(HTML(
    "<div style='font-family: system-ui; max-width: 900px; "
    f"border-left: 4px solid {border_color}; padding: 10px 12px; "
    "background: #f6f8fa; margin: 12px 0;'>"
    f"{html.escape(message)}"
    "</div>"
))

If the ratio is at or above 1.0, each GPU runs the same local workload, and the step time grew by less than the batch did. If the ratio is below 1.0, the run is communication-bound where the gradient all-reduce over the PHB path you saw in the topology check costs more than the extra GPU.

9. Take explicit control with shard_map

The default approach covers most data-parallel workloads. Sometimes, though, you want to control exactly what each GPU computes. shard_map lets you write a function that operates on per-shard arrays and use explicit collectives to communicate across devices.

Inside a shard_map function:

• Each GPU receives its local shard such as (1024, 784)
• in_specs declares how inputs are sliced
• out_specs declares how outputs are reassembled
• jax.lax.pmean(x, 'data') averages x across all GPUs along the 'data' axis

Note that now the jax.lax.pmean calls are the all-reduce that jax.jit inserted for you in the previous step.

@partial(
    jax.shard_map,
    mesh=mesh,
    in_specs=(P(), P("data", None), P("data",)),
    out_specs=(P(), P(), P()),
)
def compute_grads_shardmap(params, x_shard, y_shard):
    (loss, metrics), grads = jax.value_and_grad(loss_with_metrics, has_aux=True)(
        params,
        (x_shard, y_shard),
    )
    grads = jax.lax.pmean(grads, "data")
    loss = jax.lax.pmean(loss, "data")
    accuracy = jax.lax.pmean(metrics["accuracy"], "data")
    return grads, loss, accuracy


@jax.jit
def train_step_explicit(params, opt_state, batch):
    x, y = batch
    grads, loss, accuracy = compute_grads_shardmap(params, x, y)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    return params, opt_state, {"loss": loss, "accuracy": accuracy}

The optimizer update stays outside shard_map. Gradients are already averaged when they come out, and the parameters are replicated, so every GPU applies the identical update.

Now benchmark it against the automatic version:

params_explicit = jax.device_put(init_params(seed=1), replicated)
opt_state_explicit = optimizer.init(params_explicit)

result_explicit = benchmark_training(
    train_step_explicit,
    params_explicit,
    opt_state_explicit,
    x_batches_multi,
    y_batches_multi,
)

show_table(
    ["Approach", "Examples/sec", "ms/step"],
    [
        (
            "jit on sharded arrays",
            f"{result_multi['examples_per_sec']:,.0f}",
            f"{result_multi['ms_per_step']:.2f}",
        ),
        (
            "shard_map explicit",
            f"{result_explicit['examples_per_sec']:,.0f}",
            f"{result_explicit['ms_per_step']:.2f}",
        ),
    ],
    title="Automatic vs explicit data parallelism",
    aligns=["left", "right", "right"],
)

drop_device_refs(
    "params_multi",
    "opt_state_multi",
    "params_explicit",
    "opt_state_explicit",
    "x_batches_multi",
    "y_batches_multi",
)

You should see two rows describing the same computation expressed two ways. Treat them as confirmation that the explicit version reproduces the automatic one, not as a race — they do the same work and end at the same place.

10. Sweep the global batch size

Data parallelism lets you scale the global batch size with the number of GPUs. Larger batches amortize kernel launch overhead and improve GPU utilization up to the point where per-device memory or communication becomes the bottleneck.

The sweep below tests several global batch sizes across all GPUs. Sizes that do not divide evenly by your GPU count are skipped, and a batch that fails, for example by running out of memory, is reported without stopping the loop.

BATCH_SIZES = [256, 512, 1024, 2048, 4096]

scaling_results = []

for bs in BATCH_SIZES:
    if bs % NUM_DEVICES != 0:
        print(f"Skipping global batch {bs}: not divisible by {NUM_DEVICES} GPUs.")
        continue

    try:
        x_bs, y_bs = make_fashion_batches(bs)

        x_bs = jax.device_put(x_bs, all_data_sharding)
        y_bs = jax.device_put(y_bs, all_label_sharding)

        params_bs = jax.device_put(init_params(seed=1), replicated)
        opt_bs = optimizer.init(params_bs)

        result = benchmark_training(
            train_step,
            params_bs,
            opt_bs,
            x_bs,
            y_bs,
        )

        scaling_results.append(
            {
                "batch_size": bs,
                "per_device": bs // NUM_DEVICES,
                "examples_per_sec": result["examples_per_sec"],
                "ms_per_step": result["ms_per_step"],
            }
        )

    except Exception as e:
        print(f"Batch size {bs}: {e}")

    finally:
        drop_device_refs("x_bs", "y_bs", "params_bs", "opt_bs", "result")

This is the longest-running code in the codelab where every batch size triggers its own compilation, warm-up, and timed repeats. Now plot what you measured:

show_table(
    ["Global batch", "Per GPU", "Examples/sec", "ms/step"],
    [
        (
            r["batch_size"],
            r["per_device"],
            f"{r['examples_per_sec']:,.0f}",
            f"{r['ms_per_step']:.2f}",
        )
        for r in scaling_results
    ],
    title=f"Batch-size scaling on {NUM_DEVICES} GPUs",
    aligns=["right", "right", "right", "right"],
)


fig, ax = plt.subplots(figsize=(8, 5))

batches = [r["batch_size"] for r in scaling_results]
throughputs = [r["examples_per_sec"] for r in scaling_results]

ax.plot(
    batches,
    throughputs,
    "o-",
    color="#0969da",
    linewidth=2,
    markersize=8,
)

ax.set_xlabel("Global batch size")
ax.set_ylabel("Examples per second")
ax.set_title(f"Throughput vs batch size — {NUM_DEVICES} GPUs data-parallel")
ax.set_xscale("log", base=2)
ax.set_xticks(batches)
ax.set_xticklabels([str(b) for b in batches])
ax.grid(True, alpha=0.25)

fig.tight_layout()
plt.show()

You should get one table row and one point on the curve for every batch size that completed. You should see that throughput rises as the batch grows and the fixed per-step overheads get amortized, then flattens once the GPUs are saturated or the all-reduce starts to dominate. Where that flattening happens is a property of this model on this interconnect, and it is the number worth knowing before you scale to more GPUs.

11. Clean up

Delete the Jupyter workload, including the LoadBalancer and the persistent volume:

kubectl delete -f deploy/jupyter.yaml

Destroy the cluster, node pool, VPC, and service account:

cd terraform
terraform destroy

Type yes when prompted, then confirm nothing is left behind:

gcloud container clusters list
gcloud compute instances list

Both should be empty for this project. If you created a project just for this series, you can instead delete the whole project from the Cloud console.

12. Congratulations

You moved a JAX training loop from one GPU to two by changing where the arrays live, not by rewriting the training step.

What you've learned

• How Mesh(devices, axis_names) maps physical GPUs to a logical grid with named axes, and that those names are yours to choose
• How PartitionSpec declares which array dimensions map to which mesh axes — P('data', None) shards the batch dimension and replicates the features
• How NamedSharding(mesh, spec) combines a mesh and a spec into a placement plan for jax.device_put
• How jax.debug.visualize_array_sharding shows which GPU holds which slice, and why to run it after every placement change
• How automatic parallelism works: jax.jit on sharded inputs inserts the all-reduce and the per-shard computation for you, with no code changes
• How shard_map gives explicit per-shard control, using jax.lax.pmean for gradient averaging when you need to customize the communication pattern
• How batch-size scaling behaves: larger global batches can improve throughput until GPU utilization, memory, or communication becomes the bottleneck

Next steps

• Codelab 7: Train a transformer end to end with Flax NNX and Orbax combines the attention mechanism from codelab 5 with the multi-GPU training from this codelab
• Raise BLOCK_REPEATS or PER_DEVICE_BATCH to give each GPU more local work, then re-run the comparison step and watch the throughput ratio move
• Scale the node pool to g2-standard-48 with 4 L4s — set gpu_count = 4 in terraform.tfvars and nvidia.com/gpu: "4" in deploy/jupyter.yaml — and re-run the batch-size sweep on four devices
• Try a two-dimensional mesh with a model axis alongside data, and shard w_block along it instead of replicating it

Reference docs

• Distributed arrays and automatic parallelization
• Introduction to shard_map
• jax.debug.visualize_array_sharding
• NVIDIA JAX container on NGC