# Run course scripts and multi-host TPU Pod slices

Phase 20: TPU Workflows: Launching Jobs & Checkpointed Experiments · about 20 minutes · CPU

## What you will be able to do

- Classify course scripts using a CPU-only / single-host-compatible / multi-host-qualified execution matrix before staging.
- Distinguish single-host TPU slices (`1` host, `1–8` chips) from multi-host TPU Pod slices (`2+` hosts, `--worker=all`).
- Generate exact `gcloud compute tpus tpu-vm scp` and `ssh` command pipelines for single-host and multi-host runs.
- Retrieve `events.jsonl`, stdout logs, and checkpoints via `scp --recurse` before deleting the TPU VM.

## The problem

A course script may contain CPU-specific device selection or assumptions that do not transfer to TPU. Before staging, classify the exact script as CPU-only, single-host accelerator-compatible, multi-host-qualified, or not yet reviewed. Multi-host discovery is a topology receipt; it is not evidence of distributed training.

## The idea

Use a per-script execution matrix. `distributed-02` is CPU-only because it explicitly selects CPU devices; route it to CPU. Treat other scripts as unreviewed until their source is checked. A TPU-compatible single-host script may run on one TPU VM after its backend contract is verified. Multi-host requires simultaneous worker launch and distributed initialization, followed by a program that assigns each process data and uses the intended collective. Keep discovery receipts separate from training receipts.

## Why multi-host TPU Pod slices require `--worker=all` and `jax.distributed.initialize()`

`--worker=all` starts a process on each host, and `jax.distributed.initialize()` lets those processes discover the global topology. These steps establish process and device discovery only. They do not create a training mesh, distribute data, perform a collective update, or prove checkpoint recovery. A multi-host training receipt must separately record each process identity, its owned input range, collective result, synchronized update, and per-process logs.

**Execution matrix:** `distributed-02` → CPU-only (its source requests `jax.devices("cpu")`); do not set `JAX_PLATFORMS=tpu`. Other lesson scripts → inspect-required until a source review identifies device selection and assumptions. Single-host TPU-compatible → allowed only after that review and a backend assertion. Multi-host-qualified → allowed only after the program explicitly initializes all hosts, assigns distinct process-local inputs, and checks a global result. No generic “run any script” command is valid.

$$
g_{\text{global}} = \frac{1}{P}\sum_{p=0}^{P-1} g^{(p)} = \operatorname{psum}(g^{(p)}, \texttt{'data'}) / P
$$

### Pause and reason

If a TPU slice has `hosts = 4` and `chips_per_host = 4` (`v5litepod-16`), what are `jax.local_device_count()` and `jax.device_count()` on each host after `jax.distributed.initialize()`?

<details><summary>Compare your reasoning</summary>

`jax.local_device_count()` is `4` (the chips attached to that host VM), while `jax.device_count()` is `16` (all chips across all 4 hosts in the Pod slice).

</details>

## 1. Route a script using the execution matrix

For `distributed-02`, use its CPU mode and retain its CPU receipt. For a script reviewed as single-host TPU-compatible, set `LESSON_ID` to that reviewed ID and run on one TPU VM. Use `set -o pipefail` so `tee` cannot hide a failed Python process. Do not infer TPU compatibility from a successful import or from `JAX_PLATFORMS=tpu`.

**Run the known CPU-only distributed lesson in CPU mode**

```bash
LESSON_ID="distributed-02"
set -o pipefail
python "public/exercises/${LESSON_ID}.py" | tee "${LESSON_ID}-cpu.txt"
```

**Expected:** Executes the source-selected CPU path and saves its output. `pipefail` preserves a failed lesson exit status through `tee`; this receipt makes no TPU execution claim.

## 2. Separate multi-host discovery from a qualified training run

Launching with `--worker=all` and calling `jax.distributed.initialize()` verifies that processes can discover the global device topology. The command below is a discovery check only. A training claim additionally requires distinct per-process data, an explicit global reduction/update, and synchronized output/checkpoint evidence; no TPU training has been qualified by this lesson.

The following one-step example performs an actual global reduction and parameter update: every process contributes a distinct scalar input, `lax.psum` computes a global mean gradient, each replica applies the same update, and process 0 publishes a checkpoint after a global barrier. It is executable code, not a claimed result. Qualify it on a named TPU slice and retain all per-worker receipts before describing it as verified.

**Record multi-host process and device discovery only**

```bash
gcloud compute tpus tpu-vm ssh "$TPU_NAME" --zone="$ZONE" --worker=all --command="set -euo pipefail; JAX_PLATFORMS=tpu ~/jax-tpu-lab/.venv/bin/python -c 'import jax; jax.distributed.initialize(); print(dict(process=jax.process_index(), process_count=jax.process_count(), global_devices=jax.device_count(), local_devices=jax.local_device_count()))'"
```

**Expected:** Each worker reports process and device counts. This verifies topology discovery only; it does not establish data ownership, a collective training update, or checkpoint recovery.

**Run a one-step multi-host collective update after discovery**

```bash
cat > ./multihost_collective.py <<'PY'
import json
import os
from pathlib import Path
from functools import partial
import jax
import jax.numpy as jnp
import numpy as np
from jax.experimental import multihost_utils

jax.distributed.initialize()
process_index = jax.process_index()
process_count = jax.process_count()
local_devices = jax.local_device_count()
# Each process owns a distinct local batch value. pmap distributes each local value to one local device.
x_local = jnp.full((local_devices,), float(process_index + 1), dtype=jnp.float32)
@partial(jax.pmap, axis_name="devices")
def update(theta, x):
    local_grad = 2.0 * (theta - x)
    global_mean_grad = jax.lax.psum(local_grad, "devices") / jax.device_count()
    return theta - 0.1 * global_mean_grad, global_mean_grad

theta0 = jnp.zeros((local_devices,), dtype=jnp.float32)
theta1, grad = update(theta0, x_local)
jax.block_until_ready((theta1, grad))
expected_mean_gradient = -float(process_count + 1)
expected_theta = 0.1 * float(process_count + 1)
if not np.allclose(np.asarray(grad), expected_mean_gradient):
    raise RuntimeError("global reduction does not match the independently derived mean gradient")
if not np.allclose(np.asarray(theta1), expected_theta):
    raise RuntimeError("updated parameter does not match the expected global update")
if not np.allclose(np.asarray(theta1), np.asarray(theta1)[0]):
    raise RuntimeError("replicas produced different parameter updates")

out_dir = Path(os.environ.get("TPU_RUN_DIR", "~/jax-tpu-lab/multihost-run")).expanduser()
out_dir.mkdir(parents=True, exist_ok=True)
receipt = {"process_index": process_index, "process_count": process_count,
           "local_devices": local_devices, "global_devices": jax.device_count(),
           "owned_input": float(process_index + 1),
           "global_mean_gradient": float(np.asarray(grad)[0]),
           "updated_theta": float(np.asarray(theta1)[0]),
           "backend": jax.default_backend()}
(out_dir / f"worker-{process_index}.json").write_text(json.dumps(receipt, sort_keys=True) + "\n")
# Publish a single checkpoint only after all workers finish the same collective update.
multihost_utils.sync_global_devices("after-global-update")
if process_index == 0:
    tmp = out_dir / "checkpoint.json.tmp"
    tmp.write_text(json.dumps({"theta": receipt["updated_theta"], "global_devices": jax.device_count()}, sort_keys=True) + "\n")
    os.replace(tmp, out_dir / "checkpoint.json")
print(json.dumps(receipt, sort_keys=True))
PY
gcloud compute tpus tpu-vm scp ./multihost_collective.py "$TPU_NAME":~/jax-tpu-lab/ --zone="$ZONE" --worker=all
gcloud compute tpus tpu-vm ssh "$TPU_NAME" --zone="$ZONE" --worker=all --command="set -euo pipefail; mkdir -p ~/jax-tpu-lab/multihost-run; export TPU_RUN_DIR=~/jax-tpu-lab/multihost-run; JAX_PLATFORMS=tpu ~/jax-tpu-lab/.venv/bin/python ~/jax-tpu-lab/multihost_collective.py | tee ~/jax-tpu-lab/multihost-run/stdout-\$(hostname).jsonl"
```

**Expected:** Each process writes a worker-indexed receipt with its distinct input, global gradient, updated parameter, and observed backend. The independent expected values are mean gradient = -(process_count + 1) and updated theta = 0.1 * (process_count + 1); these checks reject a local-only reduction. Process 0 publishes checkpoint.json after the global barrier. Treat these as TPU results only after running on and naming a multi-host slice.

**Copy each worker receipt and the process-0 checkpoint back before teardown**

```bash
mkdir -p ./tpu-runs
gcloud compute tpus tpu-vm scp --recurse "$TPU_NAME":~/jax-tpu-lab/multihost-run ./tpu-runs/ --zone="$ZONE" --worker=all
```

**Expected:** Copies each host-local multihost-run directory before TPU deletion. Verify one worker-N.json receipt per process, process-0 checkpoint.json, and one stdout log per host.

## Tabulate single-host and multi-host TPU slice topologies

Create main.py with SLICE_CATALOG and compute global_chips, total_hbm_gib, and multi_host.

```python
# Step 1 — Tabulate single-host and multi-host TPU slice topologies: Separating hosts from chips_per_host shows immediately when a...
# Import json for this computation.
import json
import jax
import numpy as np

# Initialize list `SLICE_CATALOG` for the stage values.
SLICE_CATALOG = [
    {"slice": "v5litepod-4", "hosts": 1, "chips_per_host": 4, "hbm_per_chip_gib": 16},
    {"slice": "v5litepod-8", "hosts": 1, "chips_per_host": 8, "hbm_per_chip_gib": 16},
    {"slice": "v5litepod-16", "hosts": 4, "chips_per_host": 4, "hbm_per_chip_gib": 16},
    {"slice": "v6e-16", "hosts": 4, "chips_per_host": 4, "hbm_per_chip_gib": 32},
]
# Iterate over `entry` to step through the computation:
for entry in SLICE_CATALOG:
    # Compute `entry["global_chips"]` from `entry["hosts"] * entry["chips_per_host"]`
    entry["global_chips"] = entry["hosts"] * entry["chips_per_host"]
    # Compute `entry["total_hbm_gib"]` from `entry["global_chips"] * entry["hbm_per_chip_gib"]`
    entry["total_hbm_gib"] = entry["global_chips"] * entry["hbm_per_chip_gib"]
    # Compute `entry["multi_host"]` from `entry["hosts"] > 1`
    entry["multi_host"] = entry["hosts"] > 1
```

Separating hosts from chips_per_host shows immediately when a slice crosses the single-host boundary into multi-host execution.

## Generate single-host and multi-host gcloud command pipelines

Use an explicit reviewed execution class; this fixture authorizes only the CPU route for distributed-02.

```python
def build_tpu_experiment_commands(tpu_name, zone, lesson_id, hosts=1, execution_class="review-required"):
    if execution_class == "review-required":
        raise ValueError("classify and review the script before generating a TPU launch command")
    if execution_class == "cpu-only":
        if hosts != 1:
            raise ValueError("CPU-only scripts must use the CPU route")
        return {"run": f"python public/exercises/{lesson_id}.py", "multi_host": False}
    if execution_class not in {"single-host-compatible", "multi-host-qualified"}:
        raise ValueError("unknown execution class")
    if execution_class == "single-host-compatible" and hosts != 1:
        raise ValueError("single-host scripts cannot run across a TPU Pod")
    if hosts > 1 and execution_class != "multi-host-qualified":
        raise ValueError("multi-host launch requires a qualified training program")
    if hosts < 1:
        raise ValueError("hosts must be positive")
    worker_flag = " --worker=all" if hosts > 1 else ""
    scp_stage = f"gcloud compute tpus tpu-vm scp public/exercises/{lesson_id}.py {tpu_name}:~/jax-tpu-lab/ --zone={zone}{worker_flag}"
    ssh_run = f"gcloud compute tpus tpu-vm ssh {tpu_name} --zone={zone}{worker_flag} --command='set -o pipefail; JAX_PLATFORMS=tpu ~/jax-tpu-lab/.venv/bin/python ~/jax-tpu-lab/{lesson_id}.py | tee ~/jax-tpu-lab/{lesson_id}-tpu.txt'"
    scp_fetch = f"gcloud compute tpus tpu-vm scp {tpu_name}:~/jax-tpu-lab/{lesson_id}-tpu.txt ./ --zone={zone}"
    return {"scp_stage": scp_stage, "ssh_run": ssh_run, "scp_fetch": scp_fetch, "multi_host": hosts > 1}

single_cmds = build_tpu_experiment_commands("jax-tpu-course", "us-west1-c", "distributed-02", hosts=1, execution_class="cpu-only")
assert single_cmds["run"] == "python public/exercises/distributed-02.py"
for kwargs in ({}, {"hosts": 4, "execution_class": "cpu-only"}):
    try:
        build_tpu_experiment_commands("jax-tpu-pod", "us-west1-c", "distributed-02", **kwargs)
    except ValueError:
        pass
    else:
        raise AssertionError("unreviewed or CPU-only work must not be routed to a TPU Pod")
print("CPU route:", single_cmds["run"])
```

The helper requires an explicit reviewed execution class. The only named course script in this example is CPU-only, so it has no TPU or Pod command.

## Step 3: Verify invariants on the completed state

Verify that the CPU route is preserved and unsafe TPU or Pod routes remain rejected.

```python
assert single_cmds["run"] == "python public/exercises/distributed-02.py"
assert single_cmds["multi_host"] is False
try:
    build_tpu_experiment_commands("jax-tpu-pod", "us-west1-c", "distributed-02", hosts=4, execution_class="cpu-only")
except ValueError:
    print("CPU-only route remains excluded from TPU Pod execution")
else:
    raise AssertionError("CPU-only work must not be routed to a TPU Pod")
```

The completed checks retain the CPU route and reject a Pod launch for this CPU-only source.

## Run the example

```python
# Run course scripts and multi-host TPU Pod slices: Every lesson in this course generates a self-contained script in...
# Import json for this computation.
import json
import jax
import numpy as np

# Initialize list `SLICE_CATALOG` for the stage values.
SLICE_CATALOG = [
    {"slice": "v5litepod-4", "hosts": 1, "chips_per_host": 4, "hbm_per_chip_gib": 16},
    {"slice": "v5litepod-8", "hosts": 1, "chips_per_host": 8, "hbm_per_chip_gib": 16},
    {"slice": "v5litepod-16", "hosts": 4, "chips_per_host": 4, "hbm_per_chip_gib": 16},
    {"slice": "v6e-16", "hosts": 4, "chips_per_host": 4, "hbm_per_chip_gib": 32},
]
# Iterate over `entry` to step through the computation:
for entry in SLICE_CATALOG:
    # Compute `entry["global_chips"]` from `entry["hosts"] * entry["chips_per_host"]`
    entry["global_chips"] = entry["hosts"] * entry["chips_per_host"]
    # Compute `entry["total_hbm_gib"]` from `entry["global_chips"] * entry["hbm_per_chip_gib"]`
    entry["total_hbm_gib"] = entry["global_chips"] * entry["hbm_per_chip_gib"]
    # Compute `entry["multi_host"]` from `entry["hosts"] > 1`
    entry["multi_host"] = entry["hosts"] > 1


# Function `build_tpu_experiment_commands(tpu_name, zone, lesson_id, hosts)` implementing this stage's computation:
def build_tpu_experiment_commands(tpu_name, zone, lesson_id, hosts=1, execution_class="review-required"):
    if execution_class == "review-required":
        raise ValueError("classify and review the script before generating a TPU launch command")
    if execution_class == "cpu-only":
        if hosts != 1:
            raise ValueError("CPU-only scripts must use the CPU route")
        return {"run": f"python public/exercises/{lesson_id}.py", "multi_host": False}
    if execution_class == "single-host-compatible" and hosts != 1:
        raise ValueError("single-host scripts cannot run across a TPU Pod")
    if hosts > 1 and execution_class != "multi-host-qualified":
        raise ValueError("multi-host launch requires a qualified training program")
    worker_flag = " --worker=all" if hosts > 1 else ""
    scp_stage = f"gcloud compute tpus tpu-vm scp public/exercises/{lesson_id}.py {tpu_name}:~/jax-tpu-lab/ --zone={zone}{worker_flag}"
    ssh_run = f"gcloud compute tpus tpu-vm ssh {tpu_name} --zone={zone}{worker_flag} --command='set -o pipefail; JAX_PLATFORMS=tpu ~/jax-tpu-lab/.venv/bin/python ~/jax-tpu-lab/{lesson_id}.py | tee ~/jax-tpu-lab/{lesson_id}-tpu.txt'"
    scp_fetch = f"gcloud compute tpus tpu-vm scp {tpu_name}:~/jax-tpu-lab/{lesson_id}-tpu.txt ./ --zone={zone}"
    return {"scp_stage": scp_stage, "ssh_run": ssh_run, "scp_fetch": scp_fetch, "multi_host": hosts > 1}

# Run `build_tpu_experiment_commands` to compute `single_cmds`.
single_cmds = build_tpu_experiment_commands("jax-tpu-course", "us-west1-c", "distributed-02", hosts=1, execution_class="cpu-only")
# Run `build_tpu_experiment_commands` to compute `pod_cmds`.
pod_cmds = None  # This lesson has no multi-host-qualified training program.
# Assert the CPU route does not request TPU execution.
assert single_cmds["run"] == "python public/exercises/distributed-02.py"
# Assert that `"--worker=all" in pod_cmds["scp_stage"] and "--worker=all" in pod_cmds["ssh_run"]`.
try:
    build_tpu_experiment_commands("jax-tpu-pod", "us-west1-c", "distributed-02", hosts=4, execution_class="cpu-only")
except ValueError:
    print("Correctly rejected CPU-only lesson for TPU Pod execution")
else:
    raise AssertionError("CPU-only lesson must not be routed to TPU Pod")
# Print the observed values to compare against the expected result.
print("Slice catalog global chips:", {s["slice"]: s["global_chips"] for s in SLICE_CATALOG})
# Print diagnostic summary of the computed outputs.
print("CPU route:", single_cmds["run"])
```

Expected: Prints the global chip counts across v5litepod-4 (4), v5litepod-8 (8), v5litepod-16 (16), and v6e-16 (16) and the verified multi-host Pod SSH command.

## Host count, per-host local chips, and global TPU chips across single-host and Pod slices

**Predict:** How do `hosts`, `chips_per_host`, and `global_chips` change when moving from `v5litepod-4` and `v5litepod-8` to `v5litepod-16` and `v6e-16`?

![Host count, per-host local chips, and global TPU chips across single-host and Pod slices](../figures/tpu-05.svg)

**Recorded CPU computation**

### Read the figure

The horizontal axis compares four Cloud TPU slices (`v5litepod-4`, `v5litepod-8`, `v5litepod-16`, `v6e-16`). For each slice, the three bars show the number of Linux VM hosts, the number of local TPU chips per host, and the total global TPU chips.

### Connect it to the computation

`v5litepod-4` and `v5litepod-8` keep `hosts = 1`, whereas `v5litepod-16` and `v6e-16` use `4` hosts with `4` local chips each (`16` global chips), visually showing why `--worker=all` becomes mandatory at 16 chips.

```python
# Compute figure data for: Host count, per-host local chips, and global TPU chips across single-host and Pod slices
# Construct dictionary `visual_data` with the structured fields for this stage.
visual_data = {
    'kind': 'bar',
    'labels': [s['slice'] for s in SLICE_CATALOG],
    'xlabel': 'Cloud TPU slice topology',
    'ylabel': 'host / chip count',
    'series': [
        {'label': 'Linux VM hosts', 'y': [float(s['hosts']) for s in SLICE_CATALOG]},
        {'label': 'local chips per host', 'y': [float(s['chips_per_host']) for s in SLICE_CATALOG]},
        {'label': 'global TPU chips', 'y': [float(s['global_chips']) for s in SLICE_CATALOG]},
    ],
}
```

## Recorded reference execution

CPU run: 2026-10-20T10:58:54.082998+00:00. JAX 0.9.2.

```text
CPU route: python public/exercises/distributed-02.py
CPU-only route remains excluded from TPU Pod execution
Correctly rejected CPU-only lesson for TPU Pod execution
Slice catalog global chips: {'v5litepod-4': 4, 'v5litepod-8': 8, 'v5litepod-16': 16, 'v6e-16': 16}
CPU route: python public/exercises/distributed-02.py
v5litepod-16 global chips: 16 Idle chips without --worker=all: 12
Total HBM GiB by slice: {'v5litepod-4': 64, 'v5litepod-8': 128, 'v5litepod-16': 256, 'v6e-16': 512}
CPU route: python public/exercises/distributed-02.py
Rejected unsafe route: {}
Rejected unsafe route: {'hosts': 4, 'execution_class': 'cpu-only'}
Multi-host Pod slices: ['v5litepod-16', 'v6e-16']
Per-host batch: 32 Global batch: 128
PASS: tpu-05

```

## Compare local vs global device counts across slices

**Predict before running:** How many global chips are left idle if you launch a script without `--worker=all` on `v5litepod-16`?

```python
# Experiment — Compare local vs global device counts across slices: Worker 0 only controls its own 4 local chips; the other 3 hosts...
v5e16 = next(s for s in SLICE_CATALOG if s["slice"] == "v5litepod-16")
# Compute `idle_without_worker_all` from `v5e16["global_chips"] - v5e16["chips_per_host"]`
idle_without_worker_all = v5e16["global_chips"] - v5e16["chips_per_host"]
# Print the observed values to compare against the expected result.
print("v5litepod-16 global chips:", v5e16["global_chips"], "Idle chips without --worker=all:", idle_without_worker_all)
# Assert invariant `idle_without_worker_all == 12` holds
assert idle_without_worker_all == 12
```

**Expected:** Without `--worker=all`, 12 of the 16 TPU chips on `v5litepod-16` never receive a process.

Worker 0 only controls its own 4 local chips; the other 3 hosts (`3 * 4 = 12` chips) require `--worker=all`.

## Verify total HBM scaling from v5litepod-16 to v6e-16

**Predict before running:** How much total HBM is available across a 16-chip `v6e-16` Pod slice compared to `v5litepod-16`?

```python
# Experiment — Verify total HBM scaling from v5litepod-16 to v6e-16: Doubling per-chip HBM on Trillium (v6e) doubles the total...
hbm_by_slice = {s["slice"]: s["total_hbm_gib"] for s in SLICE_CATALOG}
# Print the observed values to compare against the expected result.
print("Total HBM GiB by slice:", hbm_by_slice)
# Assert that `hbm_by_slice["v6e-16"] == 512 and hbm_by_slice["v5litepod-16"] == 256`.
assert hbm_by_slice["v6e-16"] == 512 and hbm_by_slice["v5litepod-16"] == 256
```

**Expected:** v6e-16 provides 512 GiB of total HBM across 4 hosts (32 GiB/chip), double v5litepod-16's 256 GiB (16 GiB/chip).

Doubling per-chip HBM on Trillium (`v6e`) doubles the total sharded model state capacity for the same 4-host Pod topology.

## Make it yours

Use `build_tpu_experiment_commands` for the reviewed CPU-only `distributed-02` route. Then verify that its default unreviewed class and a four-host CPU-only request both raise `ValueError`.

### How to write this exercise — Step-by-step recipe & starter scaffold

**Key functions & syntax to use:**
- `build_tpu_experiment_commands(...)` — Call `build_tpu_experiment_commands` with your updated parameters or inputs from this lesson's workspace.
- `in(...)` — Call `in` with your updated parameters or inputs from this lesson's workspace.
- `assert condition` — Verify that the observed output shape, status, or numerical value satisfies the contract.

**Step-by-step implementation plan:**
1. Set up the changed input or configuration parameters requested in the prompt.
2. Run the target function or transformation and assert that the output matches your prediction.

**Starter code scaffold (fill in the TODOs):**

```python
cpu_cmds = build_tpu_experiment_commands(...)  # TODO: compute cpu_cmds
print("CPU route:", cpu_cmds["run"])
assert cpu_cmds["run"]  # TODO: complete assertion check
for kwargs in ({}, {"hosts": 4, "execution_class": "cpu-only"}):
    try:
        build_tpu_experiment_commands("jax-tpu-pod", "us-west1-c", "distributed-02", **kwargs)
    except ValueError:
        print("Rejected unsafe route:", kwargs)
    else:
        raise AssertionError("unreviewed or CPU-only work must not be routed to a TPU Pod")
```

<details><summary>Reference solution</summary>

```python
cpu_cmds = build_tpu_experiment_commands("jax-tpu-course", "us-west1-c", "distributed-02", hosts=1, execution_class="cpu-only")
print("CPU route:", cpu_cmds["run"])
assert cpu_cmds["run"] == "python public/exercises/distributed-02.py"
for kwargs in ({}, {"hosts": 4, "execution_class": "cpu-only"}):
    try:
        build_tpu_experiment_commands("jax-tpu-pod", "us-west1-c", "distributed-02", **kwargs)
    except ValueError:
        print("Rejected unsafe route:", kwargs)
    else:
        raise AssertionError("unreviewed or CPU-only work must not be routed to a TPU Pod")
```

</details>

## Filter multi-host slices from the catalog

**Foundations**

List the slice names in `SLICE_CATALOG` that require `--worker=all` (`multi_host == True`).

<details><summary>Hint</summary>

Filter by `s['multi_host']`.

</details>

### How to write: Filter multi-host slices from the catalog — Step-by-step recipe & starter scaffold

**Key functions & syntax to use:**
- `catalog(...)` — Call `catalog` with your updated parameters or inputs from this lesson's workspace.
- `hosts(...)` — Call `hosts` with your updated parameters or inputs from this lesson's workspace.
- `assert condition` — Verify that the observed output shape, status, or numerical value satisfies the contract.

**Step-by-step implementation plan:**
1. Filter multi-host slices from the catalog (Foundations): Both 16-chip slices span 4 hosts (4 chips per host) and...
2. Print the observed values to compare against the expected result.
3. Assert invariant `pod_slices == ["v5litepod-16"` holds

**Starter code scaffold (fill in the TODOs):**

```python
# Filter multi-host slices from the catalog (Foundations): Both 16-chip slices span 4 hosts (4 chips per host) and...
pod_slices = ...  # TODO: compute pod_slices
# Print the observed values to compare against the expected result.
print("Multi-host Pod slices:", pod_slices)
# Assert invariant `pod_slices == ["v5litepod-16"` holds
assert pod_slices  # TODO: complete assertion check
```

<details><summary>Reference solution and reasoning</summary>

```python
# Filter multi-host slices from the catalog (Foundations): Both 16-chip slices span 4 hosts (4 chips per host) and...
pod_slices = [s["slice"] for s in SLICE_CATALOG if s["multi_host"]]
# Print the observed values to compare against the expected result.
print("Multi-host Pod slices:", pod_slices)
# Assert invariant `pod_slices == ["v5litepod-16"` holds
assert pod_slices == ["v5litepod-16", "v6e-16"]
```

Both 16-chip slices span 4 hosts (`4` chips per host) and therefore require multi-host coordination.

</details>

## Compute per-host vs global batch size on a 4-host Pod slice

**Transfer / diagnosis**

If each TPU chip processes a microbatch of $8$ sequences on `v5litepod-16` (`4` hosts, `4` chips/host), compute the per-host batch size and the global batch size per step.

<details><summary>Hint</summary>

Multiply `8 * chips_per_host` and `8 * global_chips`.

</details>

### How to write: Compute per-host vs global batch size on a 4-host Pod slice — Step-by-step recipe & starter scaffold

**Key functions & syntax to use:**
- `slice(...)` — Call `slice` with your updated parameters or inputs from this lesson's workspace.
- `chips(...)` — Call `chips` with your updated parameters or inputs from this lesson's workspace.
- `assert condition` — Verify that the observed output shape, status, or numerical value satisfies the contract.

**Step-by-step implementation plan:**
1. Compute `per_host_batch` from `microbatch_per_chip * 4`
2. Compute `global_batch` from `microbatch_per_chip * 16`
3. Print the observed values to compare against the expected result.
4. Assert invariant `per_host_batch == 32 and global_batch == 128` holds

**Starter code scaffold (fill in the TODOs):**

```python
# Compute per-host vs global batch size on a 4-host Pod slice (Transfer / diagnosis): Each host feeds its 4 local chips (32 sequences per host),...
microbatch_per_chip = ...  # TODO: compute microbatch_per_chip
# Compute `per_host_batch` from `microbatch_per_chip * 4`
per_host_batch = ...  # TODO: compute per_host_batch
# Compute `global_batch` from `microbatch_per_chip * 16`
global_batch = ...  # TODO: compute global_batch
# Print the observed values to compare against the expected result.
print("Per-host batch:", per_host_batch, "Global batch:", global_batch)
# Assert invariant `per_host_batch == 32 and global_batch == 128` holds
assert per_host_batch  # TODO: complete assertion check
```

<details><summary>Reference solution and reasoning</summary>

```python
# Compute per-host vs global batch size on a 4-host Pod slice (Transfer / diagnosis): Each host feeds its 4 local chips (32 sequences per host),...
microbatch_per_chip = 8
# Compute `per_host_batch` from `microbatch_per_chip * 4`
per_host_batch = microbatch_per_chip * 4
# Compute `global_batch` from `microbatch_per_chip * 16`
global_batch = microbatch_per_chip * 16
# Print the observed values to compare against the expected result.
print("Per-host batch:", per_host_batch, "Global batch:", global_batch)
# Assert invariant `per_host_batch == 32 and global_batch == 128` holds
assert per_host_batch == 32 and global_batch == 128
```

Each host feeds its 4 local chips (`32` sequences per host), yielding a synchronized global batch of `128` sequences across the 16-chip mesh.

</details>

## Check your understanding

What happens if you launch a distributed JAX script on a 4-host `v5litepod-16` slice without `--worker=all`?

1. Google Cloud automatically merges the 4 Linux hosts into 1 host
2. Only worker 0 starts Python, so `jax.distributed.initialize()` hangs waiting for the other 3 hosts (or 12 of the 16 TPU chips sit idle)
3. The script runs four times faster

<details><summary>Answer and explanation</summary>

Only worker 0 starts Python, so `jax.distributed.initialize()` hangs waiting for the other 3 hosts (or 12 of the 16 TPU chips sit idle)

Every host in a multi-host TPU Pod slice runs its own OS and Python process; `--worker=all` is required so all hosts join the collective mesh.

</details>

## Diagnose the result

If `gcloud compute tpus tpu-vm ssh --worker=all` fails on one worker because `.venv` is missing, run your `.venv` bootstrap command with `--worker=all` as well so every host has `~/jax-tpu-lab/.venv` installed.

## Carry forward

- Stage any `public/exercises/<lesson-id>.py` script with `gcloud compute tpus tpu-vm scp` to run course experiments on TPU in parallel.
- Use single-host commands for `1–8` chip slices (`v5litepod-4`, `v5litepod-8`, `v6e-4`) and add `--worker=all` + `jax.distributed.initialize()` for multi-host slices (`v5litepod-16`, `v6e-16`).

## Keep your evidence

Keep the generated scp_in_cmd, ssh_run_cmd, and scp_out_cmd workflow for single-host and multi-host slices plus the process-0 checkpoint writer check.

Save your predictions, modified code, terminal outputs, and reasoning in your engineering portfolio.

## Primary references

- [Multi-host and distributed JAX on TPU Pods](https://docs.jax.dev/en/latest/multi_process.html)
- [Run JAX code on Cloud TPU Pod slices](https://cloud.google.com/tpu/docs/run-calculation-jax)