{
"cells": [
{
"cell_type": "markdown",
"id": "cell-00",
"metadata": {},
"source": [
"# Scale one training run across many devices\n",
"\n",
"When a model or its batch outgrows one accelerator, we spread the work over several. In Dew that is one argument to the `Trainer`: a `MeshSpec` that arranges the devices into named axes. The training code does not change. In this notebook we train the language model from [notebook 05](05-train-a-language-model.ipynb) on eight devices in two layouts:\n",
"\n",
"- data parallel, `MeshSpec(fsdp=1)`: every device holds the whole model and a different slice of the batch;\n",
"- fully sharded, `MeshSpec(fsdp=8)`: every device holds one eighth of each large weight, and the devices fetch the other pieces when they need them.\n",
"\n",
"We look at where one weight lives in each layout, and we restore a checkpoint written in one layout into the other.\n",
"\n",
"Most readers do not have eight accelerators at hand, so this notebook runs on eight simulated devices on the CPU. XLA can split one CPU into several devices with a flag, and everything Dew does with them is the same as on eight GPUs or TPU chips. The flag has to be set before JAX starts, which is why it comes first. The model is tiny, and the whole notebook runs in under a minute on a desktop CPU."
]
},
{
"cell_type": "markdown",
"id": "cell-01",
"metadata": {},
"source": [
"## Install"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "cell-02",
"metadata": {
"execution": {
"iopub.execute_input": "2026-09-22T19:36:46.458259Z",
"iopub.status.busy": "2026-09-22T19:36:46.458027Z",
"iopub.status.idle": "2026-09-22T19:36:53.976005Z",
"shell.execute_reply": "2026-09-22T19:36:53.975341Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Note: you may need to restart the kernel to use updated packages.\n"
]
}
],
"source": [
"%pip install -q \"dew-ml @ git+https://github.com/AshishKumar4/dew\""
]
},
{
"cell_type": "markdown",
"id": "cell-03",
"metadata": {},
"source": [
"## Eight simulated devices\n",
"\n",
"`--xla_force_host_platform_device_count=8` gives the CPU backend eight devices, and `JAX_PLATFORMS=cpu` makes JAX use that backend even on a machine with a GPU."
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "cell-04",
"metadata": {
"execution": {
"iopub.execute_input": "2026-09-22T19:36:53.980356Z",
"iopub.status.busy": "2026-09-22T19:36:53.980130Z",
"iopub.status.idle": "2026-09-22T19:36:54.998804Z",
"shell.execute_reply": "2026-09-22T19:36:54.998324Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[CpuDevice(id=0), CpuDevice(id=1), CpuDevice(id=2), CpuDevice(id=3), CpuDevice(id=4), CpuDevice(id=5), CpuDevice(id=6), CpuDevice(id=7)]\n"
]
}
],
"source": [
"import os\n",
"\n",
"os.environ[\"XLA_FLAGS\"] = \"--xla_force_host_platform_device_count=8\"\n",
"os.environ[\"JAX_PLATFORMS\"] = \"cpu\"\n",
"\n",
"import jax\n",
"\n",
"print(jax.devices())"
]
},
{
"cell_type": "markdown",
"id": "cell-05",
"metadata": {},
"source": [
"## Settings\n",
"\n",
"The batch of 16 has to split evenly over the eight devices."
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "cell-06",
"metadata": {
"execution": {
"iopub.execute_input": "2026-09-22T19:36:55.009663Z",
"iopub.status.busy": "2026-09-22T19:36:55.008693Z",
"iopub.status.idle": "2026-09-22T19:36:55.027386Z",
"shell.execute_reply": "2026-09-22T19:36:55.017369Z"
}
},
"outputs": [],
"source": [
"STEPS = 20\n",
"BATCH_SIZE = 16\n",
"SEQUENCE_LENGTH = 32\n",
"EMB_FEATURES = 64\n",
"NUM_LAYERS = 2\n",
"NUM_HEADS = 4\n",
"DATA_DIR = \"data/07-tokens\"\n",
"RUN_DIR = \"runs/07-scaling\""
]
},
{
"cell_type": "markdown",
"id": "cell-07",
"metadata": {},
"source": [
"## The mesh\n",
"\n",
"`build_mesh(MeshSpec(...))` arranges the devices into a mesh with six named axes: `data`, `expert`, `fsdp`, `tensor`, `sequence` and `stage`. Axes we do not set have size 1, and the `data` axis takes whatever devices are left. A batch splits over the `data` and `fsdp` axes together, so both meshes below split the batch eight ways. They differ only in what happens to the weights."
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "cell-08",
"metadata": {
"execution": {
"iopub.execute_input": "2026-09-22T19:36:55.039620Z",
"iopub.status.busy": "2026-09-22T19:36:55.035642Z",
"iopub.status.idle": "2026-09-22T19:36:55.879210Z",
"shell.execute_reply": "2026-09-22T19:36:55.878403Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"data parallel: {'data': 8, 'expert': 1, 'fsdp': 1, 'tensor': 1, 'sequence': 1, 'stage': 1}\n",
"fully sharded: {'data': 1, 'expert': 1, 'fsdp': 8, 'tensor': 1, 'sequence': 1, 'stage': 1}\n"
]
}
],
"source": [
"from dew import MeshSpec\n",
"from dew.training.distributed import build_mesh\n",
"\n",
"mesh_dp = build_mesh(MeshSpec(fsdp=1))\n",
"mesh_fsdp = build_mesh(MeshSpec(fsdp=8))\n",
"print(\"data parallel:\", dict(mesh_dp.shape))\n",
"print(\"fully sharded:\", dict(mesh_fsdp.shape))"
]
},
{
"cell_type": "markdown",
"id": "cell-09",
"metadata": {},
"source": [
"## Where a weight lives\n",
"\n",
"`Layout` decides how each weight maps onto the mesh, from the logical axis names the model's layers declare. Weights smaller than `min_shard` elements stay whole on every device, because splitting them costs more in communication than it saves in memory. The default of 65,536 would keep every weight of this small model whole, so we lower it to 1 to see the effect.\n",
"\n",
"Here is the token embedding table, 256 rows by 64 features, in each layout. In the data-parallel mesh every device holds all of it. In the sharded mesh each device holds 32 rows."
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "cell-10",
"metadata": {
"execution": {
"iopub.execute_input": "2026-09-22T19:36:55.880768Z",
"iopub.status.busy": "2026-09-22T19:36:55.880489Z",
"iopub.status.idle": "2026-09-22T19:36:56.251378Z",
"shell.execute_reply": "2026-09-22T19:36:56.250951Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"data parallel P()\n"
]
},
{
"data": {
"text/html": [
" \n",
" \n",
" \n",
" \n",
" \n",
"CPU 0,1,2,3,4,5,6,7\n",
" \n",
" \n",
" \n",
" \n",
" \n",
"\n"
],
"text/plain": [
" \n",
" \n",
" \n",
" \n",
" \n",
"CPU 0,1,2,3,4,5,6,7\n",
" \n",
" \n",
" \n",
" \n",
" \n"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"fully sharded P('fsdp',)\n"
]
},
{
"data": {
"text/html": [
" CPU 0 \n",
" \n",
" CPU 1 \n",
" \n",
" CPU 2 \n",
" \n",
" CPU 3 \n",
" \n",
" CPU 4 \n",
" \n",
" CPU 5 \n",
" \n",
" CPU 6 \n",
" \n",
" CPU 7 \n",
" \n",
"\n"
],
"text/plain": [
" CPU 0 \n",
" \n",
" CPU 1 \n",
" \n",
" CPU 2 \n",
" \n",
" CPU 3 \n",
" \n",
" CPU 4 \n",
" \n",
" CPU 5 \n",
" \n",
" CPU 6 \n",
" \n",
" CPU 7 \n",
" \n"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"import jax.numpy as jnp\n",
"from dew import Layout\n",
"\n",
"layout = Layout(min_shard=1)\n",
"table = jnp.ones((256, EMB_FEATURES))\n",
"for name, mesh in ((\"data parallel\", mesh_dp), (\"fully sharded\", mesh_fsdp)):\n",
" sharding = layout.shardings(mesh, {\"params\": {\"embed_tokens\": {\"embedding\": table}}})\n",
" placed = jax.device_put(table, sharding[\"params\"][\"embed_tokens\"][\"embedding\"])\n",
" print(name, placed.sharding.spec)\n",
" jax.debug.visualize_array_sharding(placed)"
]
},
{
"cell_type": "markdown",
"id": "cell-11",
"metadata": {},
"source": [
"## The data and the model\n",
"\n",
"The data is a small generated corpus of short sentences, written in the token layout from notebook 05. The model is a two-layer version of the notebook 05 decoder."
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "cell-12",
"metadata": {
"execution": {
"iopub.execute_input": "2026-09-22T19:36:56.254254Z",
"iopub.status.busy": "2026-09-22T19:36:56.253853Z",
"iopub.status.idle": "2026-09-22T19:36:56.429905Z",
"shell.execute_reply": "2026-09-22T19:36:56.428924Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"training windows: 2987\n"
]
}
],
"source": [
"import json\n",
"from pathlib import Path\n",
"\n",
"import numpy as np\n",
"from dew.data import ByteTokenizer, Loading, TokenWindows\n",
"\n",
"rng = np.random.default_rng(0)\n",
"subjects = [\"the cat\", \"a dog\", \"the bird\", \"my friend\", \"the child\"]\n",
"verbs = [\"sees\", \"likes\", \"finds\", \"wants\", \"hears\"]\n",
"objects = [\"the ball\", \"a tree\", \"the river\", \"some food\", \"the moon\"]\n",
"text = \"\".join(f\"{rng.choice(subjects)} {rng.choice(verbs)} {rng.choice(objects)}.\\n\" for _ in range(4000))\n",
"\n",
"data_dir = Path(DATA_DIR)\n",
"data_dir.mkdir(parents=True, exist_ok=True)\n",
"ids = np.asarray(ByteTokenizer().encode(text), np.uint8)\n",
"val_len = len(ids) // 50\n",
"ids[:val_len].tofile(data_dir / \"val.bin\")\n",
"ids[val_len:].tofile(data_dir / \"train.bin\")\n",
"(data_dir / \"meta.json\").write_text(json.dumps(\n",
" {\"tokenizer\": \"byte\", \"vocab_size\": 256, \"dtype\": \"uint8\",\n",
" \"train_tokens\": len(ids) - val_len, \"val_tokens\": val_len, \"eos_id\": None}))\n",
"\n",
"data = TokenWindows(path=DATA_DIR, seq_len=SEQUENCE_LENGTH, val_batches=2,\n",
" loading=Loading(workers=0, threads=1, read_buffer=2)).load(batch=BATCH_SIZE)\n",
"print(\"training windows:\", data.records)"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "cell-13",
"metadata": {
"execution": {
"iopub.execute_input": "2026-09-22T19:36:56.431606Z",
"iopub.status.busy": "2026-09-22T19:36:56.431369Z",
"iopub.status.idle": "2026-09-22T19:37:00.238139Z",
"shell.execute_reply": "2026-09-22T19:37:00.237404Z"
}
},
"outputs": [],
"source": [
"import optax\n",
"from dew import Checkpoints, Trainer, models\n",
"from dew.objectives.lm import LMObjective\n",
"\n",
"model = models.build(\"causal_transformer\", vocab_size=256, emb_features=EMB_FEATURES,\n",
" num_layers=NUM_LAYERS, num_heads=NUM_HEADS, max_seq_len=SEQUENCE_LENGTH,\n",
" dtype=\"float32\", attention_impl=\"xla\")\n",
"objective = LMObjective(model, SEQUENCE_LENGTH, ema_decay=None)"
]
},
{
"cell_type": "markdown",
"id": "cell-14",
"metadata": {},
"source": [
"## Training in both layouts\n",
"\n",
"The two trainers below differ only in `mesh` and in the folder they write checkpoints to. The trainer builds the mesh, works out a sharding for every array in the training state, and compiles the step with those shardings; XLA inserts the communication between devices. The state is created directly in its final layout, so a model too large for one device never has to fit on one.\n",
"\n",
"After each run we read the embedding table's placement off the returned state. `sharding.spec` names the mesh axis each dimension is split over, and `addressable_shards` lists the piece each device holds."
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "cell-15",
"metadata": {
"execution": {
"iopub.execute_input": "2026-09-22T19:37:00.241161Z",
"iopub.status.busy": "2026-09-22T19:37:00.240847Z",
"iopub.status.idle": "2026-09-22T19:37:07.926118Z",
"shell.execute_reply": "2026-09-22T19:37:07.925397Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Training from step 0 to 20 on {'data': 8, 'expert': 1, 'fsdp': 1, 'tensor': 1, 'sequence': 1, 'stage': 1} (1 process(es))\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"step 10: loss 3.5792\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"step 20: loss 2.6488\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Goodput: first step after 5.21 s, 11.5% of the wall time in steps\n",
"spec: P()\n",
" cpu:0: rows slice(None, None, None), local shape (256, 64)\n",
" cpu:1: rows slice(None, None, None), local shape (256, 64)\n",
" cpu:2: rows slice(None, None, None), local shape (256, 64)\n"
]
}
],
"source": [
"def describe(state):\n",
" table = state.params[\"params\"][\"embed_tokens\"][\"embedding\"]\n",
" print(\"spec:\", table.sharding.spec)\n",
" for shard in table.addressable_shards[:3]:\n",
" print(f\" {shard.device}: rows {shard.index[0]}, local shape {shard.data.shape}\")\n",
"\n",
"replicated = Trainer(objective, optax.adamw(1e-3), key=jax.random.key(0),\n",
" mesh=MeshSpec(fsdp=1), layout=Layout(min_shard=1),\n",
" checkpoints=Checkpoints(f\"{RUN_DIR}/replicated\"))\n",
"replicated_state = replicated.fit(data, steps=STEPS, log_every=10)\n",
"describe(replicated_state)"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "cell-16",
"metadata": {
"execution": {
"iopub.execute_input": "2026-09-22T19:37:07.927959Z",
"iopub.status.busy": "2026-09-22T19:37:07.927630Z",
"iopub.status.idle": "2026-09-22T19:37:14.299302Z",
"shell.execute_reply": "2026-09-22T19:37:14.298751Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Training from step 0 to 20 on {'data': 1, 'expert': 1, 'fsdp': 8, 'tensor': 1, 'sequence': 1, 'stage': 1} (1 process(es))\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"step 10: loss 3.5792\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"step 20: loss 2.6488\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Goodput: first step after 4.97 s, 15.0% of the wall time in steps\n",
"spec: P('fsdp',)\n",
" cpu:0: rows slice(0, 32, None), local shape (32, 64)\n",
" cpu:1: rows slice(32, 64, None), local shape (32, 64)\n",
" cpu:2: rows slice(64, 96, None), local shape (32, 64)\n"
]
}
],
"source": [
"sharded = Trainer(objective, optax.adamw(1e-3), key=jax.random.key(0),\n",
" mesh=MeshSpec(fsdp=8), layout=Layout(min_shard=1),\n",
" checkpoints=Checkpoints(f\"{RUN_DIR}/sharded\"))\n",
"sharded_state = sharded.fit(data, steps=STEPS, log_every=10)\n",
"describe(sharded_state)"
]
},
{
"cell_type": "markdown",
"id": "cell-17",
"metadata": {},
"source": [
"Both runs start from the same key and read the same batches, so they compute the same thing and the losses match. Only the placement of the weights differs.\n",
"\n",
"## A checkpoint moves between layouts\n",
"\n",
"Dew can restore a checkpoint into a different layout from the one that wrote it. The trainer below uses the sharded layout but points at the data-parallel run's checkpoints. `place()` reads each array and lays it out the way this trainer's mesh asks. The restored weights equal the data-parallel run's, and training continues from step 20."
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "cell-18",
"metadata": {
"execution": {
"iopub.execute_input": "2026-09-22T19:37:14.309444Z",
"iopub.status.busy": "2026-09-22T19:37:14.309186Z",
"iopub.status.idle": "2026-09-22T19:37:19.144844Z",
"shell.execute_reply": "2026-09-22T19:37:19.142355Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Resumed from step 20 in /tmp/nbwork/run/runs/07-scaling/replicated\n",
"spec: P('fsdp',)\n",
" cpu:0: rows slice(0, 32, None), local shape (32, 64)\n",
" cpu:1: rows slice(32, 64, None), local shape (32, 64)\n",
" cpu:2: rows slice(64, 96, None), local shape (32, 64)\n",
"restored weights equal the data-parallel run's: True\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Resumed from step 20 in /tmp/nbwork/run/runs/07-scaling/replicated\n",
"Training from step 20 to 30 on {'data': 1, 'expert': 1, 'fsdp': 8, 'tensor': 1, 'sequence': 1, 'stage': 1} (1 process(es))\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"step 30: loss 1.9283\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Goodput: first step after 2.96 s, 12.2% of the wall time in steps\n",
"continued to step 30\n"
]
}
],
"source": [
"crossed = Trainer(objective, optax.adamw(1e-3), key=jax.random.key(0),\n",
" mesh=MeshSpec(fsdp=8), layout=Layout(min_shard=1),\n",
" checkpoints=Checkpoints(f\"{RUN_DIR}/replicated\"))\n",
"restored, _, _ = crossed.place()\n",
"describe(restored)\n",
"same = all(np.array_equal(np.asarray(a), np.asarray(b))\n",
" for a, b in zip(jax.tree_util.tree_leaves(restored.params),\n",
" jax.tree_util.tree_leaves(replicated_state.params)))\n",
"print(\"restored weights equal the data-parallel run's:\", same)\n",
"\n",
"continued = crossed.fit(data, steps=STEPS + 10, log_every=10)\n",
"print(\"continued to step\", int(continued.step))"
]
},
{
"cell_type": "markdown",
"id": "cell-19",
"metadata": {},
"source": [
"## Several hosts\n",
"\n",
"Everything above ran in one process. On a TPU pod slice every host runs the same script, and each one first joins the group:\n",
"\n",
"```python\n",
"from dew.training.runtime import prepare_process\n",
"\n",
"prepare_process(multi_host=True)\n",
"```\n",
"\n",
"`prepare_process` calls `jax.distributed.initialize()`, which finds the coordinator from the environment the TPU pod provides. The token loaders give each host its own share of the records, and the checkpoint folder has to be one every host can write to, such as a `gs://` bucket. The `dew-tpu` command creates a slice, installs Dew on every worker and starts a recipe on all of them; the [TPU guide](../docs/tpu.md) describes it. None of that ran in this notebook.\n",
"\n",
"## Where to go next\n",
"\n",
"`Layout.min_shard` decides which weights are worth splitting; in a decoder with a large vocabulary the embedding table is usually the first. `Trainer(accumulation=k)` adds up gradients over `k` smaller batches when a full batch does not fit in memory. The [distributed training guide](../docs/concepts/distributed.md) covers the `expert`, `tensor`, `sequence` and `stage` axes."
]
}
],
"metadata": {
"accelerator": "GPU",
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.13"
}
},
"nbformat": 4,
"nbformat_minor": 5
}