I have setup a snippet on Colab here with jax.__version__ # 0.4.33 9Feb2025 orbax.checkpoint.__version__ # 0.6.4 9Feb2025 It quite difficult to follow the flax/orbax changes in the save/restore (simple) model even with following the "latest" documentation of these two packages. I have managed to cooked something but I was wandering if I'm doing the right thing using 8TPUs on Colab; for instance it semms that one can save a single instance of the Model among the 8 existing ones (ie. the use of flax.jax_utils.unreplicate seems necessary ckpt = {'model': flax.jax_utils.unreplicate(model_state)} ) At restoration in the same environment after target={'model': abstract_state} # a Training State quite dummy chpt_restored = checkpoint_manager.restore(checkpoint_manager.latest_step(), items=target) one restaure 8 vesions using new_model_state = flax.jax_utils.replicate(chpt_restored['model']) but this is 8 replicated version of the model with the same instance. It may be foreseen to behaves like that, but I wandering how one can resume a first training session, to continue the training as one may use s unique instance at the second training session? Hope that I have been clear. Any comment on the Colab snippet is welcome. Continue reading...
