Checkpoints and randomness

Save and restore model, optimizer, and random-stream state together. Verify the next update is reproduced, rather than checking only that weights load.

These outputs were saved from local execution. The documentation renderer does not execute notebook cells.

Show code cell source

Hide code cell source

from pathlib import Path
import sys
ROOT = Path.cwd() if (Path.cwd() / 'examples').exists() else Path.cwd().parent
sys.path.insert(0, str(ROOT))
import jax
import jax.numpy as jnp
import numpy as np
import jaxsaw as js
from jaxsaw import nn, optim, data, checkpoint, inspection
from jaxsaw.nn import functional as F
print(f"JAXsaw {js.__version__} · JAX {jax.__version__}")
print(jax.devices())

JAXsaw 0.1.0 · JAX 0.5.3
[CudaDevice(id=0)]

Own random streams explicitly

Initialization and runtime randomness use separate streams. The random stream is an explicit step argument and part of the checkpoint.

model = nn.Linear(2,1,rngs=js.Rngs(params=15))
optimizer = optim.Adam(model.parameters(),lr=.01)
runtime = js.Rngs(sample=16)
@js.jit
def step(model,optimizer,rngs):
    x = js.random.normal(rngs.next('sample'),(8,2))
    target = x.sum(-1,keepdims=True)
    loss,grads = js.value_and_grad(lambda m:F.mse_loss(m(x),target))(model)
    optimizer.step(grads)
    return loss
print('Initial update:',float(step(model,optimizer,runtime)))

Initial update: 6.502008438110352

Save a complete snapshot

NPZ checkpoints contain arrays and a JSON manifest; loading does not execute embedded Python. Reconstruct the same architecture and optimizer configuration before restoration.

directory = ROOT / 'runs' / 'api-checkpoint'
directory.mkdir(parents=True,exist_ok=True)
path = directory / 'state.npz'
checkpoint.save(path,model=model,optimizer=optimizer,rngs=runtime,
                progress={'next_update':int(optimizer.step_count)},metadata={'seed':15})
expected_loss = step(model,optimizer,runtime)
expected_weights = model.state_dict()
print('Next update before restore:',float(expected_loss))

Next update before restore: 2.25146222114563

Verify exact continuation on this backend

Restoring RNG counters and optimizer moments is necessary to reproduce the next update. Array comparisons below verify actual continuation on this GPU/runtime.

restored = checkpoint.restore(path,model=model,optimizer=optimizer,rngs=runtime)
actual_loss = step(model,optimizer,runtime)
np.testing.assert_allclose(actual_loss,expected_loss,atol=1e-7)
for name,value in model.state_dict().items():
    np.testing.assert_allclose(value,expected_weights[name],atol=1e-7)
print('Next update after restore:',float(actual_loss))
print('Saved progress:',restored.progress)
print('Continuation verified.')

Next update after restore: 2.25146222114563
Saved progress: {'next_update': 1}
Continuation verified.

Weights-only snapshots

Use state_dict when transferring weights into an already constructed matching model. It does not contain optimizer or RNG state, so it is insufficient for exact training continuation.

weights = model.state_dict()
copy = nn.Linear(2,1,rngs=js.Rngs(params=99))
copy.load_state_dict(weights)
probe = jnp.ones((2,2))
np.testing.assert_allclose(model(probe),copy(probe))
print('Matching inference:',np.asarray(copy(probe)))

Matching inference: [[-0.7157132]
 [-0.7157132]]