1. Introduction
In this codelab, you define a small decoder transformer in JAX with Flax NNX, train it on Shakespeare text across both GPUs on your node, save and restore it with Orbax, and generate new text from the trained weights.
The model is deliberately tiny with 4 layers, 256-dimensional embeddings, a byte-level vocabulary so that it trains in under a minute on two L4 GPUs. The architecture and the training patterns are the same ones used in much larger models.
What you'll do
- Define a decoder transformer with Flax NNX using
nnx.Embed,nnx.MultiHeadAttention,nnx.Linear, andnnx.LayerNorm - Plug causal
jax.nn.dot_product_attentionin as the attention backend - Train on byte-level TinyShakespeare with
nnx.Optimizerand Optax AdamW, first on one GPU and then on all of them - Measure throughput in tokens/sec and compare the two runs
- Save and restore model parameters with the Orbax
StandardCheckpointer - Generate Shakespeare-like text from the trained model
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)
- Codelabs 1 through 6 completed, or an equivalent JAX GPU environment with at least two GPUs
- Outbound internet access from the Pod, so the first run can download the TinyShakespeare text file
Estimated time to complete: 70 minutes.
The architecture you are building
The model is a stack of transformer blocks. Each block has two sub-layers, each wrapped in a residual connection:
- Self-attention — each position attends to all earlier positions (causal mask).
- Feed-forward network (FFN) — two linear layers with a GELU activation, expanding and then compressing the representation.
Both sub-layers use pre-norm with LayerNorm is applied before the sub-layer. Pre-norm is more stable for training and is the standard in modern transformers.
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://, 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 flax optax orbax-checkpoint matplotlib
Set up and verify the GPU
Import JAX, Flax NNX, Optax, and Orbax, and check how many GPUs the container can see.
import os
os.environ["LD_LIBRARY_PATH"] = "/usr/local/nvidia/lib64:" + os.environ.get(
"LD_LIBRARY_PATH", ""
)
import hashlib
import html
import math
import pathlib
import time
import urllib.request
import warnings
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 flax import nnx
import orbax.checkpoint as ocp
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 lesson needs at least 2 GPUs. Found {len(gpu_devices)}. "
f"Available devices: {devices}"
)
def block_tree(tree):
return jax.block_until_ready(tree)
def show_table(headers, rows, title=None, aligns=None):
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):
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):,.0f} {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, a list of two CUDA devices, and GPU count: 2. The block_tree, show_table, and show_bars helpers render the tables and bar charts used in later steps.
3. Prepare byte-level TinyShakespeare data
TinyShakespeare is a single text file of about 1 MB containing Shakespeare's works concatenated together. This codelab uses byte-level tokenization where every byte of the UTF-8 text becomes one token. That fixes the vocabulary at 256 possible values and removes any tokenizer dependency.
The text is chunked into non-overlapping sequences of length SEQ_LEN. Each sequence is one training example, and the model learns to predict the next byte at every position.
Run the code to download the file, verify its checksum, split it into train and validation sequences, and shuffle the training set:
SHAKESPEARE_URL = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
SHAKESPEARE_MD5 = "d015dc5942f9b2908e24d4827a3e7a5e"
DATA_DIR = pathlib.Path.home() / ".cache" / "jax-course"
DATA_DIR.mkdir(parents=True, exist_ok=True)
DATA_FILE = DATA_DIR / "tinyshakespeare.txt"
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()
if DATA_FILE.exists() and md5sum(DATA_FILE) == SHAKESPEARE_MD5:
print("Using cached tinyshakespeare.txt")
else:
print(f"Downloading tinyshakespeare.txt")
urllib.request.urlretrieve(SHAKESPEARE_URL, DATA_FILE)
raw_text = DATA_FILE.read_text()
data = np.frombuffer(raw_text.encode("utf-8"), dtype=np.uint8).astype(np.int32)
print()
VOCAB_SIZE = 256
SEQ_LEN = 256
PER_DEVICE_BATCH = 32
num_sequences = len(data) // SEQ_LEN
data = data[: num_sequences * SEQ_LEN].reshape(num_sequences, SEQ_LEN)
num_train = int(0.9 * num_sequences)
train_data = data[:num_train]
val_data = data[num_train:]
rng = np.random.default_rng(0)
train_data = train_data[rng.permutation(num_train)]
def make_batches(data, batch_size):
usable = (len(data) // batch_size) * batch_size
return data[:usable].reshape(-1, batch_size, SEQ_LEN)
show_table(
["", "Value"],
[
("Total bytes", f"{len(raw_text):,}"),
("Vocabulary", f"{VOCAB_SIZE} (raw bytes)"),
("Sequence length", SEQ_LEN),
("Training sequences", f"{num_train:,}"),
("Validation sequences", f"{len(val_data):,}"),
],
title="TinyShakespeare — byte-level tokenization",
)
print()
print("Sample text (first 200 bytes):")
print(raw_text[:200])
You should see either Using cached tinyshakespeare.txt or a download message, followed by a table of the dataset shape with total bytes, a vocabulary of 256 raw bytes, the sequence length, and the training and validation sequence counts and then the first 200 bytes of the text so you can see what the model is learning from.
Three constants set here matter for the rest of the codelab. VOCAB_SIZE is 256 because a byte has 256 possible values, and SEQ_LEN is 256 positions per training example.
PER_DEVICE_BATCH is 32 and stays fixed for both the single-GPU and the multi-GPU run. That fixed per-GPU batch is what makes the throughput comparison later a weak-scaling comparison.
4. Define the transformer with Flax NNX
Flax NNX is a simplified API for neural networks in JAX. You define layers as Python objects that own their weight initialization and their forward pass. This step introduces every NNX piece the transformer needs, so no earlier NNX experience is assumed.
The model uses four building blocks:
nnx.Embedis a lookup table that maps a token index to a vector.nnx.Linearis a dense matrix multiply plus an optional bias.nnx.LayerNormnormalizes features before the attention and FFN sub-layers.nnx.MultiHeadAttentionhandles the Q/K/V projections, the attention, and the output projection.
Wire up causal attention
nnx.MultiHeadAttention receives hidden states shaped (B, T, D_MODEL), creates Q, K, and V internally, and splits them into heads. The attention_fn hook controls only the core attention operation that runs after those projections.
NNX passes Flax-style optional arguments such as a dropout rng, a dtype, and a precision into attention_fn. jax.nn.dot_product_attention does not accept those, so the wrapper below takes them with a catch-all and forwards only what the JAX function needs.
D_MODEL = 256
NUM_HEADS = 4
FFN_DIM = 1024
NUM_LAYERS = 4
MAX_SEQ_LEN = 256
LR = 3e-4
WEIGHT_DECAY = 1e-4
def causal_sdpa(query, key, value, **_):
return jax.nn.dot_product_attention(query, key, value, is_causal=True)
Define the block and the model
There are two classes. TransformerBlock is one attention sub-layer plus one FFN sub-layer, and TinyTransformer stacks num_layers of them between the embeddings and the LM head. Each class extends nnx.Module and creates all of its layers in __init__.
class TransformerBlock(nnx.Module):
def __init__(self, d_model: int, num_heads: int, ffn_dim: int, rngs: nnx.Rngs):
self.ln1 = nnx.LayerNorm(d_model, rngs=rngs)
self.attn = nnx.MultiHeadAttention(
num_heads=num_heads,
in_features=d_model,
decode=False,
attention_fn=causal_sdpa,
rngs=rngs,
)
self.ln2 = nnx.LayerNorm(d_model, rngs=rngs)
self.fc_up = nnx.Linear(d_model, ffn_dim, rngs=rngs)
self.fc_down = nnx.Linear(ffn_dim, d_model, rngs=rngs)
def __call__(self, x):
x = x + self.attn(self.ln1(x))
h = jax.nn.gelu(self.fc_up(self.ln2(x)))
x = x + self.fc_down(h)
return x
The pre-norm arrangement is visible in __call__. x = x + self.attn(self.ln1(x)) normalizes before attention and adds the result back to the residual stream, and the FFN branch does the same thing with self.ln2.
class TinyTransformer(nnx.Module):
def __init__(
self,
vocab_size: int,
d_model: int,
num_heads: int,
ffn_dim: int,
num_layers: int,
max_seq_len: int,
rngs: nnx.Rngs,
):
self.token_embed = nnx.Embed(vocab_size, d_model, rngs=rngs)
self.pos_embed = nnx.Embed(max_seq_len, d_model, rngs=rngs)
self.blocks = nnx.List(
[
TransformerBlock(d_model, num_heads, ffn_dim, rngs=rngs)
for _ in range(num_layers)
]
)
self.final_norm = nnx.LayerNorm(d_model, rngs=rngs)
self.lm_head = nnx.Linear(d_model, vocab_size, use_bias=False, rngs=rngs)
def __call__(self, tokens):
B, T = tokens.shape
x = self.token_embed(tokens) + self.pos_embed(jnp.arange(T))
for block in self.blocks:
x = block(x)
x = self.final_norm(x)
return self.lm_head(x)
TinyTransformer.__call__ adds the token embedding and the position embedding, runs the blocks in order, applies a final LayerNorm, and projects to vocab_size logits.
Instantiate and inspect the model
Creating the model is one call. The nnx.Rngs handles all the random state needed for parameter initialization. Once it exists you can count its parameters and run a forward pass through it.
model = TinyTransformer(
VOCAB_SIZE,
D_MODEL,
NUM_HEADS,
FFN_DIM,
NUM_LAYERS,
MAX_SEQ_LEN,
rngs=nnx.Rngs(0),
)
param_count = sum(x.size for x in jax.tree.leaves(nnx.state(model, nnx.Param)))
show_table(
["", "Value"],
[
("Architecture", f"Decoder-only transformer"),
("Layers", NUM_LAYERS),
("Model dimension", D_MODEL),
("Attention heads", f"{NUM_HEADS} (head dim = {D_MODEL // NUM_HEADS})"),
("FFN dimension", FFN_DIM),
("Vocabulary", f"{VOCAB_SIZE} (byte-level)"),
("Max sequence length", MAX_SEQ_LEN),
("Parameters", f"{param_count:,}"),
],
title="TinyTransformer",
)
logits = model(jnp.zeros((1, 16), dtype=jnp.int32))
print(f"Test forward pass: input (1, 16) \u2192 logits {logits.shape}")
You should see a table describing the architecture: 4 layers, model dimension 256, 4 attention heads with head dimension 64, a byte-level vocabulary, and a parameter count. The last line reports the test forward pass, with an input of shape (1, 16) producing logits of shape (1, 16, 256) — one distribution over the 256 byte values for each of the 16 input positions.
5. Write the NNX training step
Here you use @nnx.jit to handle NNX module state automatically. It splits modules into structure plus arrays for JIT compilation, then merges the updated arrays back, so you write the step as if the modules were regular Python objects.
The loss is next-token prediction. At each position the model predicts the following token, so you compare logits[:, :-1] (the predictions at positions 0 to T-2) with tokens[:, 1:] (the actual tokens at positions 1 to T-1).
@nnx.jit
def train_step(model, optimizer, tokens):
def loss_fn(model):
logits = model(tokens)
pred = logits[:, :-1].reshape(-1, VOCAB_SIZE)
target = tokens[:, 1:].reshape(-1)
return optax.softmax_cross_entropy_with_integer_labels(pred, target).mean()
loss, grads = nnx.value_and_grad(loss_fn)(model)
optimizer.update(model, grads)
return {"loss": loss, "perplexity": jnp.exp(loss)}
The loop around that step warms up once so compilation is not timed, then runs a fixed number of steps and converts elapsed time into tokens/sec.
def train_loop(model, optimizer, batches, steps=1000, log_every=100):
num_batches = batches.shape[0]
history = []
# Warmup: compile the training step
warmup_metrics = train_step(model, optimizer, batches[0])
block_tree(warmup_metrics)
start = time.perf_counter()
for step in range(steps):
tokens = batches[step % num_batches]
metrics = train_step(model, optimizer, tokens)
if step % log_every == 0 or step == steps - 1:
metrics = block_tree(metrics)
history.append(
{
"step": step,
"loss": float(metrics["loss"]),
"perplexity": float(metrics["perplexity"]),
}
)
block_tree(metrics)
elapsed = time.perf_counter() - start
batch_size = int(batches.shape[1])
tokens_per_step = batch_size * (SEQ_LEN - 1)
tokens_per_sec = steps * tokens_per_step / elapsed
return history, elapsed, tokens_per_sec
This cell only defines two functions, so it produces no output. The next step is where they run.
6. Train on a single GPU
Start with one GPU to establish a baseline. jax.device_put pins the whole batch array to gpu_devices[0], so nothing spreads across devices before you get to the deliberate multi-GPU comparison.
BENCHMARK_STEPS = 500
STEPS_1GPU = BENCHMARK_STEPS
batches_1gpu = make_batches(train_data, PER_DEVICE_BATCH)
single_device = gpu_devices[0]
batches_1gpu = jax.device_put(batches_1gpu, single_device)
model_1gpu = TinyTransformer(
VOCAB_SIZE,
D_MODEL,
NUM_HEADS,
FFN_DIM,
NUM_LAYERS,
MAX_SEQ_LEN,
rngs=nnx.Rngs(1),
)
optimizer_1gpu = nnx.Optimizer(
model_1gpu, optax.adamw(LR, weight_decay=WEIGHT_DECAY), wrt=nnx.Param
)
history_1gpu, elapsed_1gpu, tps_1gpu = train_loop(
model_1gpu, optimizer_1gpu, batches_1gpu, steps=STEPS_1GPU
)
show_table(
["Step", "Loss", "Perplexity"],
[(h["step"], f"{h['loss']:.3f}", f"{h['perplexity']:.1f}") for h in history_1gpu],
title=f"Single-GPU training \u2014 {tps_1gpu:,.0f} tokens/sec",
aligns=["right", "right", "right"],
)
The very first call compiles the step, which is why train_loop warms up before it starts the clock. When the run finishes you should see a table with one row per logged step, showing loss and perplexity falling as training progresses. The table title reports the measured tokens/sec for this run.
7. Scale the same step to all GPUs
This is a data-parallel pattern where you create a mesh, replicate the model, shard the data along the batch dimension. The training step code does not change at all with @nnx.jit handles the parallelism based on how the arrays are placed.
To keep the throughput comparison fair, the single-GPU and multi-GPU runs use the same number of timed steps. Each GPU still processes PER_DEVICE_BATCH sequences per step, so the multi-GPU run processes a larger global batch.
Replicating an NNX model takes three calls. nnx.state extracts the module's state as a PyTree, jax.device_put places that PyTree on every device with the replicated sharding, and nnx.update writes it back into the module. The optimizer state gets the same treatment.
STEPS_MULTI = BENCHMARK_STEPS
GLOBAL_BATCH = PER_DEVICE_BATCH * NUM_DEVICES
mesh = Mesh(np.array(gpu_devices), ("data",))
replicated = NamedSharding(mesh, P())
data_sharding = NamedSharding(mesh, P(None, "data", None))
batches_multi = make_batches(train_data, GLOBAL_BATCH)
batches_multi = jax.device_put(batches_multi, data_sharding)
model_multi = TinyTransformer(
VOCAB_SIZE,
D_MODEL,
NUM_HEADS,
FFN_DIM,
NUM_LAYERS,
MAX_SEQ_LEN,
rngs=nnx.Rngs(1),
)
optimizer_multi = nnx.Optimizer(
model_multi, optax.adamw(LR, weight_decay=WEIGHT_DECAY), wrt=nnx.Param
)
# Replicate model and optimizer state across all GPUs
model_state = nnx.state(model_multi)
nnx.update(model_multi, jax.device_put(model_state, replicated))
opt_state = nnx.state(optimizer_multi)
nnx.update(optimizer_multi, jax.device_put(opt_state, replicated))
print(
f"Global batch: {GLOBAL_BATCH} ({PER_DEVICE_BATCH} per GPU \u00d7 {NUM_DEVICES} GPUs)"
)
print(f"Training batches: {batches_multi.shape}")
print()
history_multi, elapsed_multi, tps_multi = train_loop(
model_multi, optimizer_multi, batches_multi, steps=STEPS_MULTI
)
show_table(
["Step", "Loss", "Perplexity"],
[(h["step"], f"{h['loss']:.3f}", f"{h['perplexity']:.1f}") for h in history_multi],
title=f"Multi-GPU training \u2014 {tps_multi:,.0f} tokens/sec",
aligns=["right", "right", "right"],
)
You should see the global batch line reporting PER_DEVICE_BATCH sequences per GPU times the number of GPUs, the shape of the sharded batch array, and then a second loss table with the same columns as the single-GPU run. Its title reports the multi-GPU tokens/sec.
8. Compare throughput and plot the curves
The model and the training step were identical in both runs. Only the data placement changed. Both runs used the same number of timed steps, and each GPU still processed PER_DEVICE_BATCH sequences per step, so the multi-GPU run has a larger global batch.
ms_per_step_1gpu = elapsed_1gpu / STEPS_1GPU * 1e3
ms_per_step_multi = elapsed_multi / STEPS_MULTI * 1e3
speedup = tps_multi / tps_1gpu
show_table(
["", "1 GPU", f"{NUM_DEVICES} GPUs", "Ratio"],
[
("Batch size", PER_DEVICE_BATCH, GLOBAL_BATCH, f"{NUM_DEVICES}×"),
("Per-GPU batch", PER_DEVICE_BATCH, PER_DEVICE_BATCH, "same"),
("Timed steps", STEPS_1GPU, STEPS_MULTI, "same"),
("ms/step", f"{ms_per_step_1gpu:.2f}", f"{ms_per_step_multi:.2f}", f"{ms_per_step_1gpu / ms_per_step_multi:.2f}×"),
("Tokens/sec", f"{tps_1gpu:,.0f}", f"{tps_multi:,.0f}", f"{speedup:.2f}×"),
],
title="Throughput comparison",
aligns=["left", "right", "right", "right"],
)
if speedup > NUM_DEVICES * 1.25:
print(
f"Note: the measured speedup is superlinear (> {NUM_DEVICES}x). "
"For this small benchmark, treat that as a measurement artifact rather "
"than a general hardware-scaling claim."
)
show_bars(
[("1 GPU", tps_1gpu), (f"{NUM_DEVICES} GPUs", tps_multi)],
"Training throughput",
"tokens/s",
)
The multi-GPU run is faster in tokens/sec because it processes a larger global batch while keeping the per-GPU batch fixed. If the measured speedup comes out greater than the number of GPUs, treat it as a benchmark artifact from compiler, layout, and kernel differences rather than a general scaling guarantee.
Next, plot training progress for the two runs as the model sees more tokens. The x-axis is tokens processed rather than raw training steps, because the multi-GPU run uses a larger global batch and therefore sees more data per step. Lower loss and perplexity are better, so the curves show how quickly each setup improves for the amount of text processed.
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
for label, history, batch_size in [
("1 GPU", history_1gpu, PER_DEVICE_BATCH),
(f"{NUM_DEVICES} GPUs", history_multi, GLOBAL_BATCH),
]:
tokens_m = [
(h["step"] + 1) * batch_size * (SEQ_LEN - 1) / 1e6 for h in history
]
losses = [h["loss"] for h in history]
perps = [h["perplexity"] for h in history]
ax1.plot(tokens_m, losses, "o-", label=label, markersize=4)
ax2.plot(tokens_m, perps, "o-", label=label, markersize=4)
ax1.set_xlabel("Tokens processed (millions)")
ax1.set_ylabel("Loss")
ax1.set_title("Training loss")
ax1.legend()
ax1.grid(True, alpha=0.25)
ax2.set_xlabel("Tokens processed (millions)")
ax2.set_ylabel("Perplexity")
ax2.set_title("Training perplexity")
ax2.legend()
ax2.grid(True, alpha=0.25)
fig.suptitle("Training curves vs tokens processed")
fig.tight_layout()
plt.show()
You should see two panels side by side with loss on the left and perplexity on the right. Each with one curve per run, both trending downward as tokens processed increases.
9. Save and restore a checkpoint with Orbax
Orbax saves model state as a directory of array files. StandardCheckpointer is its simplest API: one call to save, one call to restore.
For NNX models you extract the parameters with nnx.state(model, nnx.Param), save that PyTree, and later restore it into a fresh model with nnx.update. The graph structure (nnx.GraphDef) is not saved and it comes from the Python class definition, so you need TinyTransformer in scope to rebuild the model before you can load weights into it.
The code below runs the whole lifecycle. It extracts only the trained model parameters, saves them to disk, builds a ShapeDtypeStruct tree telling Orbax what shapes and dtypes to expect, restores into a newly initialized TinyTransformer, and checks that the restored model produces the same logits as the original on a small test input.
ckpt_dir = pathlib.Path("/tmp/jax-course/l7-checkpoints")
# Extract model parameters (not optimizer state)
model_params = nnx.state(model_multi, nnx.Param)
# Save
checkpointer = ocp.StandardCheckpointer()
if (ckpt_dir / "trained").exists():
import shutil
shutil.rmtree(ckpt_dir / "trained")
checkpointer.save(ckpt_dir / "trained", model_params)
print(f"Checkpoint saved to {ckpt_dir / 'trained'}")
# Create abstract target for restore
abstract_params = jax.tree.map(
lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype),
model_params,
)
# Restore into a fresh model
model_restored = TinyTransformer(
VOCAB_SIZE,
D_MODEL,
NUM_HEADS,
FFN_DIM,
NUM_LAYERS,
MAX_SEQ_LEN,
rngs=nnx.Rngs(99),
)
restored_params = checkpointer.restore(ckpt_dir / "trained", abstract_params)
nnx.update(model_restored, restored_params)
# Test
test_input = jnp.zeros((1, 16), dtype=jnp.int32)
logits_original = model_multi(test_input)
logits_restored = model_restored(test_input)
max_diff = float(jnp.max(jnp.abs(logits_original - logits_restored)))
show_table(
["", "Value"],
[
("Checkpoint path", str(ckpt_dir / "trained")),
("Parameters saved", f"{sum(x.size for x in jax.tree.leaves(model_params)):,}"),
("Max |original \u2212 restored|", f"{max_diff:.2e}"),
("Match", "\u2713" if max_diff < 1e-5 else "\u2717"),
],
title="Orbax checkpoint save and restore",
)
You should see the checkpoint path, the number of parameters saved, the maximum absolute difference between the original and restored logits, and a check mark when that difference is below 1e-5.
10. Generate Shakespeare-like text
The trained model predicts the next byte at every position. To generate text, you feed it a prompt, take the logits at the last position, sample a token, append it, and repeat.
The generate function pads its input to MAX_SEQ_LEN so the JIT-compiled forward pass always sees the same input shape with no recompilation as the sequence grows. With causal attention, padding after the real tokens does not affect the output at earlier positions.
nnx.split separates the module into a graphdef and a state PyTree so it can be passed through jax.jit, and nnx.merge rebuilds the module inside the compiled function.
@jax.jit
def get_logits_jit(graphdef, model_state, tokens):
model = nnx.merge(graphdef, model_state)
return model(tokens)
def generate(model, prompt_text, max_new_tokens=300, temperature=0.8):
graphdef, model_state = nnx.split(model)
tokens = list(prompt_text.encode("utf-8"))
key = jax.random.key(42)
for _ in range(max_new_tokens):
context = tokens[-MAX_SEQ_LEN:]
padded = context + [0] * (MAX_SEQ_LEN - len(context))
input_arr = jnp.array([padded], dtype=jnp.int32)
logits = get_logits_jit(graphdef, model_state, input_arr)
next_logit = logits[0, len(context) - 1]
if temperature <= 0:
next_token = int(jnp.argmax(next_logit))
else:
key, subkey = jax.random.split(key)
next_token = int(jax.random.categorical(subkey, next_logit / temperature))
tokens.append(next_token)
return bytes(tokens).decode("utf-8", errors="replace")
# Put model on a single device for generation
gen_model = TinyTransformer(
VOCAB_SIZE,
D_MODEL,
NUM_HEADS,
FFN_DIM,
NUM_LAYERS,
MAX_SEQ_LEN,
rngs=nnx.Rngs(99),
)
nnx.update(gen_model, checkpointer.restore(ckpt_dir / "trained", abstract_params))
print("=== Prompt: 'ROMEO:' | temperature=0.8 ===")
print()
print(generate(gen_model, "ROMEO:", max_new_tokens=300, temperature=0.8))
print()
print("=== Prompt: 'To be, or not' | temperature=0.6 ===")
print()
print(generate(gen_model, "To be, or not", max_new_tokens=300, temperature=0.6))
You should see two blocks of generated text, one per prompt.
Do not expect shakespeare-level text. Give the model, you will see line breaks in roughly the right places, capitalized speaker names, English-looking runs of letters, and very little meaning.
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 built a decoder transformer with Flax NNX, trained it across both GPUs on your node, checkpointed it with Orbax, and generated text from the restored weights.
What you've learned
- How Flax NNX organizes a model into reusable modules:
nnx.Embed,nnx.Linear,nnx.LayerNorm, andnnx.MultiHeadAttentionhandle parameter creation and the forward pass, and the training step usesnnx.value_and_gradwithnnx.Optimizer - How causal attention with
is_causal=Truemasks future positions so the model can only look backward, and why theattention_fnhook needs a**_wrapper - How data-parallel training replicates the model and shards the batch across GPUs exactly as in codelab 6, without changing the training step
- How Orbax saves and restores model parameters, with
nnx.stateandnnx.updatebridging between NNX modules and plain PyTrees - How to measure throughput in tokens/sec, the natural unit for language models, with warmup and
block_until_readystill doing the work of keeping the numbers honest - How text generation feeds the model its own predictions one token at a time, and how temperature controls the randomness of that sampling
Next steps
- Codelab 8: Export and serve a trained JAX model where you take this checkpoint and prepare it for serving with JIT inference, AOT compilation, and export to portable formats
- Change
D_MODEL(try 128 or 512) andNUM_LAYERS(try 2 or 6), then re-run. More capacity means slower steps, and more layers means more memory - Change the generation temperature (try 0.0, 0.5, and 1.2). Zero is greedy, and values above 1 produce more random output
- Switch the attention backend: change the body of
causal_sdpato calljax.nn.dot_product_attention(..., is_causal=True, implementation="cudnn")and use bf16-compatible activations, then watch how cuDNN fused attention behaves