request: 08_save_load_checkpoints update

[jecampagne]

Hello,
I'm running with

• Flax 0.10.3 with the nnx API
• orbax.checkpoint 0.11.5
• and jax 0.5.0

The current version of 08_save_load_checkpoints.py is fine but it would be nice to make a midified version for a real fine use case where one saves during training :
1 - the best model to be used in a second session for evaluation
2 - a check point of the model and optimizer for a second training session

And the orbax checkpointmanager is not so clear how to use with the nnx modules.
I have a use case with a model using some Variables that are not updated (on purpose) by the opimizer, while all the others nnx.Params are.

Thanks to help using NNX.

[cgarciae]

Hi @jecampagne, please check out the Save and load checkpoints guide which is more complete.

[jecampagne]

Thanks for the link. BTW I had first read this pages and find them quite heavy.
And I am still stuck how to save & restore the model AND optimizer states.

Here is a rough sketch of what I succeed to do with a single device to save & restore the model
(looking at #4224)
but I do not know how I can manage the optimizer state.

At training stage

import orbax.checkpoint as ocp

ckpt_dir = '<parth>'
ckpt_dir = ocp.test_utils.erase_and_create_empty(ckpt_dir)

options = ocp.CheckpointManagerOptions(
    save_interval_steps=2,
    max_to_keep=2,
)
mngr = ocp.CheckpointManager(
    ckpt_dir,
    options=options,
)

model =  MyModel(rngs=nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.adam(learning_rate=1e-3)) # update nnx.Params
graphdef, state = nnx.split((model, optimizer)) 

#in epoch loop
state, loss = train_step(graphdef, state, batch_data) 
nnx.update((model, optimizer), state)
 _, model_state = nnx.split(model)
 mngr.save(epoch, args=ocp.args.StandardSave(model_state))

....
mngr.wait_until_finished()

Now at restore time

options = ocp.CheckpointManagerOptions(
    read_only=True, 
    create=False,
)

ckpt_dir = "<the diectory at training>"

mngr = ocp.CheckpointManager(
    ckpt_dir,
    options=options,
)

abstract_model = nnx.eval_shape(lambda: MyModel(rngs=nnx.Rngs(0)) )
graphdef, abstract_state = nnx.split(abstract_model)

state_restored = mngr.restore(mngr.latest_step(), items=abstract_state)
model = nnx.merge(graphdef, state_restored)

[jecampagne]

BTW with orbax-checkpoint in Colab ( v 0.6.4 ) then I succeed with a modified mngr.restore version

state_restored = mngr.restore(mngr.latest_step(), args=ocp.args.StandardSave(abstract_state))

I do not know when the behavior have changed between 0.6.4 and 0.11.5...
I do not know neither when Colba will switch to the 0.5.x Jax version