Scale one training run across many devices

Open in ColabView on GitHub DownloadCPU

When I trained text-to-image models from scratch with FlaxDiff, Dew’s predecessor, it was on 128 TPU v4 chips. Once a model or its batch outgrows one accelerator, there is no way around spreading 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 itself does not change. In this notebook, we train the language model from notebook 05 on eight devices, in two layouts:

data parallel, MeshSpec(fsdp=1): every device holds the whole model and a different slice of the batch;
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. The idea comes from the paper “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models” by Rajbhandari et al. (2019).

We look at where one weight lives in each layout, and then restore a checkpoint written in one layout into the other.

Most readers don’t 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.

Install Dependencies

Just Dew from GitHub this time, no extras, since everything here runs on the CPU.

%pip install -q "dew-ml @ git+https://github.com/AshishKumar4/dew"

install log

Note: you may need to restart the kernel to use updated packages.

Eight simulated devices

--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.

import os

os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
os.environ["JAX_PLATFORMS"] = "cpu"

import jax

print(jax.devices())

[CpuDevice(id=0), CpuDevice(id=1), CpuDevice(id=2), CpuDevice(id=3), CpuDevice(id=4), CpuDevice(id=5), CpuDevice(id=6), CpuDevice(id=7)]

Settings

The batch of 16 has to split evenly over the eight devices.

STEPS = 20
BATCH_SIZE = 16
SEQUENCE_LENGTH = 32
EMB_FEATURES = 64
NUM_LAYERS = 2
NUM_HEADS = 4
DATA_DIR = "data/07-tokens"
RUN_DIR = "runs/07-scaling"

The mesh

build_mesh(MeshSpec(...)) arranges the devices into a mesh with six named axes: data, expert, fsdp, tensor, sequence and stage. Axes we don’t 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.

from dew import MeshSpec
from dew.training.distributed import build_mesh

mesh_dp = build_mesh(MeshSpec(fsdp=1))
mesh_fsdp = build_mesh(MeshSpec(fsdp=8))
print("data parallel:", dict(mesh_dp.shape))
print("fully sharded:", dict(mesh_fsdp.shape))

data parallel: {'data': 8, 'expert': 1, 'fsdp': 1, 'tensor': 1, 'sequence': 1, 'stage': 1}
fully sharded: {'data': 1, 'expert': 1, 'fsdp': 8, 'tensor': 1, 'sequence': 1, 'stage': 1}

Where a weight lives

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.

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.

import jax.numpy as jnp
from dew import Layout

layout = Layout(min_shard=1)
table = jnp.ones((256, EMB_FEATURES))
for name, mesh in (("data parallel", mesh_dp), ("fully sharded", mesh_fsdp)):
    sharding = layout.shardings(mesh, {"params": {"embed_tokens": {"embedding": table}}})
    placed = jax.device_put(table, sharding["params"]["embed_tokens"]["embedding"])
    print(name, placed.sharding.spec)
    jax.debug.visualize_array_sharding(placed)

data parallel P()

                   
                   
                   
                   
                   
CPU 0,1,2,3,4,5,6,7

fully sharded P('fsdp',)

The data and the model

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.

import json
from pathlib import Path

import numpy as np
from dew.data import ByteTokenizer, Loading, TokenWindows

rng = np.random.default_rng(0)
subjects = ["the cat", "a dog", "the bird", "my friend", "the child"]
verbs = ["sees", "likes", "finds", "wants", "hears"]
objects = ["the ball", "a tree", "the river", "some food", "the moon"]
text = "".join(f"{rng.choice(subjects)} {rng.choice(verbs)} {rng.choice(objects)}.\n" for _ in range(4000))

data_dir = Path(DATA_DIR)
data_dir.mkdir(parents=True, exist_ok=True)
ids = np.asarray(ByteTokenizer().encode(text), np.uint8)
val_len = len(ids) // 50
ids[:val_len].tofile(data_dir / "val.bin")
ids[val_len:].tofile(data_dir / "train.bin")
(data_dir / "meta.json").write_text(json.dumps(
    {"tokenizer": "byte", "vocab_size": 256, "dtype": "uint8",
     "train_tokens": len(ids) - val_len, "val_tokens": val_len, "eos_id": None}))

data = TokenWindows(path=DATA_DIR, seq_len=SEQUENCE_LENGTH, val_batches=2,
                    loading=Loading(workers=0, threads=1, read_buffer=2)).load(batch=BATCH_SIZE)
print("training windows:", data.records)

training windows: 2987

import optax
from dew import Checkpoints, Trainer, models
from dew.objectives.lm import LMObjective

model = models.build("causal_transformer", vocab_size=256, emb_features=EMB_FEATURES,
                     num_layers=NUM_LAYERS, num_heads=NUM_HEADS, max_seq_len=SEQUENCE_LENGTH,
                     dtype="float32", attention_impl="xla")
objective = LMObjective(model, SEQUENCE_LENGTH, ema_decay=None)

Training in both layouts

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 the devices for us. The state is created directly in its final layout, so a model too large for one device never has to fit on one.

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.

def describe(state):
    table = state.params["params"]["embed_tokens"]["embedding"]
    print("spec:", table.sharding.spec)
    for shard in table.addressable_shards[:3]:
        print(f"  {shard.device}: rows {shard.index[0]}, local shape {shard.data.shape}")

replicated = Trainer(objective, optax.adamw(1e-3), key=jax.random.key(0),
                     mesh=MeshSpec(fsdp=1), layout=Layout(min_shard=1),
                     checkpoints=Checkpoints(f"{RUN_DIR}/replicated"))
replicated_state = replicated.fit(data, steps=STEPS, log_every=10)
describe(replicated_state)

Training from step 0 to 20 on {'data': 8, 'expert': 1, 'fsdp': 1, 'tensor': 1, 'sequence': 1, 'stage': 1} (1 process(es))
step 10: loss 3.5792
step 20: loss 2.6488
Goodput: first step after 5.21 s, 11.5% of the wall time in steps
spec: P()
  cpu:0: rows slice(None, None, None), local shape (256, 64)
  cpu:1: rows slice(None, None, None), local shape (256, 64)
  cpu:2: rows slice(None, None, None), local shape (256, 64)

sharded = Trainer(objective, optax.adamw(1e-3), key=jax.random.key(0),
                  mesh=MeshSpec(fsdp=8), layout=Layout(min_shard=1),
                  checkpoints=Checkpoints(f"{RUN_DIR}/sharded"))
sharded_state = sharded.fit(data, steps=STEPS, log_every=10)
describe(sharded_state)

Training from step 0 to 20 on {'data': 1, 'expert': 1, 'fsdp': 8, 'tensor': 1, 'sequence': 1, 'stage': 1} (1 process(es))
step 10: loss 3.5792
step 20: loss 2.6488
Goodput: first step after 4.97 s, 15.0% of the wall time in steps
spec: P('fsdp',)
  cpu:0: rows slice(0, 32, None), local shape (32, 64)
  cpu:1: rows slice(32, 64, None), local shape (32, 64)
  cpu:2: rows slice(64, 96, None), local shape (32, 64)

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.

A checkpoint moves between layouts

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.

crossed = Trainer(objective, optax.adamw(1e-3), key=jax.random.key(0),
                  mesh=MeshSpec(fsdp=8), layout=Layout(min_shard=1),
                  checkpoints=Checkpoints(f"{RUN_DIR}/replicated"))
restored, _, _ = crossed.place()
describe(restored)
same = all(np.array_equal(np.asarray(a), np.asarray(b))
           for a, b in zip(jax.tree_util.tree_leaves(restored.params),
                           jax.tree_util.tree_leaves(replicated_state.params)))
print("restored weights equal the data-parallel run's:", same)

continued = crossed.fit(data, steps=STEPS + 10, log_every=10)
print("continued to step", int(continued.step))

Resumed from step 20 in /tmp/nbwork/run/runs/07-scaling/replicated
spec: P('fsdp',)
  cpu:0: rows slice(0, 32, None), local shape (32, 64)
  cpu:1: rows slice(32, 64, None), local shape (32, 64)
  cpu:2: rows slice(64, 96, None), local shape (32, 64)
restored weights equal the data-parallel run's: True
Resumed from step 20 in /tmp/nbwork/run/runs/07-scaling/replicated
Training from step 20 to 30 on {'data': 1, 'expert': 1, 'fsdp': 8, 'tensor': 1, 'sequence': 1, 'stage': 1} (1 process(es))
step 30: loss 1.9283
Goodput: first step after 2.96 s, 12.2% of the wall time in steps
continued to step 30

Several hosts

Everything above ran in one process. On a TPU pod slice, every host runs the same script, and each one first joins the group:

from dew.training.runtime import prepare_process

prepare_process(multi_host=True)

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 describes it. None of that ran in this notebook.

Where to go next

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 doesn’t fit in memory. The distributed training guide covers the expert, tensor, sequence and stage axes.