Kill it and prove it

A saver for a pytree is short enough to read in one sitting. Flatten with paths, turn each path into a name, and write the leaves into a .npz. Restore by flattening a freshly initialized state to recover the structure, then unflattening the loaded arrays back into it.

Using the key path as the archive name is what makes the file survive a change to the model. Position would not: lesson one showed that dict keys flatten sorted, so renaming a parameter reorders the leaves while the count stays the same, and a positional restore would load silently and wrongly.

The archive holds thirteen arrays and the step comes back as int32 with weak_type False, which matters for the reason lesson one measured: restore that counter as a Python int and the resumed run compiles a second and third executable for no reason.

import numpy as np


def save(path, state):
    flat, _ = jax.tree_util.tree_flatten_with_path(state)
    np.savez(path, **{jax.tree_util.keystr(p): np.asarray(leaf) for p, leaf in flat})


def load(path, like):
    z = np.load(path)
    flat, treedef = jax.tree_util.tree_flatten_with_path(like)
    return jax.tree.unflatten(treedef, [jnp.asarray(z[jax.tree_util.keystr(p)]) for p, _ in flat])


state = init(jax.random.key(0))
for step in range(10):
    state, _ = train_step(state, data, jax.random.fold_in(base, step))

save("run.npz", state)
z = np.load("run.npz")
print(len(z.files), z[".step"], z[".step"].dtype)

restored = load("run.npz", init(jax.random.key(0)))
print(jax.tree.structure(restored) == jax.tree.structure(state))
print(all(jnp.array_equal(a, b) for a, b in zip(jax.tree.leaves(state), jax.tree.leaves(restored))))
print(restored.step.aval.weak_type)

# 13 10 int32
# True
# True
# False

Notice what is not in that archive. There is no key in the state at all, and if there were, np.asarray would refuse it: a typed key has dtype key<fry>, and the error says so and names the function that gets you out.

So a state that carries a key is still serializable. key_data on the way out, wrap_key_data on the way in, two extra lines and one decision about which form the archive stores. Chapter 8's arc owns both calls and the comparison this arc leans on: which key discipline survives a restore, and what a resumed run looks like when the wrong one was used.

The reason there is no key here is that the batch key is fold_in(base, step), so the base key is a constant of the program rather than state, and the counter carries the stream. That is what makes the archive thirteen arrays with nothing left to decide, and the next two sections measure what it is worth.

Restore the step number and the randomness comes back with it.

import jax
import numpy as np

key = jax.random.key(7)
try:
    np.asarray(key)
except TypeError as err:
    print(err)

np.savez("key.npz", key=jax.random.key_data(key))
words = np.load("key.npz")["key"]
print(words.dtype, words.shape)

# JAX array with PRNGKey dtype cannot be converted to a NumPy array. Use
# jax.random.key_data(arr) if you wish to extract the underlying integer array.
# uint32 (2,)

Twenty steps straight through. The same run killed at ten, saved, restored into objects built from scratch, and continued. The claim is that the last ten losses are the same list, and the check is == on the list, not a plot.

They match to six decimals as printed, and the whole final state matches leaf for leaf as array_equal. Both halves are worth running: the losses agreeing tells you the forward pass lined up, and the state agreeing tells you the moments and the counter did too, which is what the eleventh step depends on.

The PyTorch path's own resume lesson runs the same experiment against a global generator and a scheduler, and LAB·P4 draws the overlaid curves on a TPU bridge. What is different here is how little the checkpoint has to hold: no generator block, no scheduler object, thirteen arrays and a step number, because the randomness is a function of the step and the learning rate is a constant in the program.

def run_from(state, start, n):
    losses = []
    for step in range(start, start + n):
        state, loss = train_step(state, data, jax.random.fold_in(base, step))
        losses.append(round(float(loss), 6))
    return state, losses


straight, uninterrupted = run_from(init(jax.random.key(0)), 0, 20)
after, resumed = run_from(restored, 10, 10)

print(uninterrupted[10:])
print(resumed)
print(resumed == uninterrupted[10:])
print(all(jnp.array_equal(a, b) for a, b in zip(jax.tree.leaves(straight), jax.tree.leaves(after))))

# [5.617665, 7.358051, 6.792553, 5.592591, 2.720668, 3.390179, 4.089592,
#  3.357401, 3.01731, 1.720861]
# [5.617665, 7.358051, 6.792553, 5.592591, 2.720668, 3.390179, 4.089592,
#  3.357401, 3.01731, 1.720861]
# True
# True

Break the resume two ways and watch how long each one takes to show. Restore the state correctly but restart the step index at zero, and the first resumed loss is 8.487146 against the correct 5.617665, because a different 32 rows were drawn. That failure is loud and lands immediately.

The second one shows up a step later. Restore the parameters and the counter, zero the two Adam moments, and the first resumed loss is 5.617665, exactly right. The loss is measured before the update, so missing moments cannot change the number until the step after. Step twelve reads 7.644585 where the uninterrupted run had 7.358051.

One correct step is enough to fool a check that only looks at the first number after a restore, and this is a small model on synthetic data where the gap opens fast. On a real run the divergence is a slightly worse curve that nobody attributes to the restore. Compare the list, not the first element.

zeros = jax.tree.map(jnp.zeros_like, restored.params)

_, from_zero = run_from(restored, 0, 10)
_, no_moments = run_from(restored._replace(m=zeros, v=zeros), 10, 10)

print(uninterrupted[10:13])
print(from_zero[:3])
print(no_moments[:3])

# [5.617665, 7.358051, 6.792553]
# [8.487146, 5.50882, 5.776051]
# [5.617665, 7.644585, 7.207603]

The eight-line saver above is enough to prove a resume and not enough to run one. It writes the whole file synchronously, so a training step waits on the disk. It writes in place, so a process killed mid-write leaves an archive that loads and is wrong. And it has no notion of a mesh, so a sharded run has nothing to restore onto.

Orbax is the library the chapter names for those three problems: asynchronous saves so the step does not block, atomic directory commits so a partial write is never mistaken for a checkpoint, and sharding-aware restore so a tree written from one device layout comes back onto another. The tree is the same tree either way, which is why the saver above is a good thing to have written once before adopting a library that hides it.

The bar does not change when the library does. Kill the run, restore, continue, and compare the list. A checkpoint that exists is not evidence; a resumed loss list equal to the uninterrupted one is.