orbax save/restore with 8 devices

[jecampagne]

Hello,
I have setup a snippet on Colab Here

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.

[orbax-dev]

I see you're using pmap - the Orbax support for this is not great. It's preferable to use jit if possible.

Before saving, you can also try converting the model using ocp.utils.fully_replicated_host_local_array_to_global_array (and skip flax.jax_utils). This should convert the array into one represented using NamedSharding. At that point, flax.jax_utils are not needed to save, and the restored model should have the same shardings as when it was saved (NamedSharding).

[jecampagne]

Thanks for your interest.

1. have you the possibility to update the snippet?
2. what would be changed if I use FLAX NXX upto date versions as the plateform where I develop would liks to update JAX to the new 0.5.0 version

[orbax-dev]

I don't think you need to fully update JAX (although your Orbax version is 0.6 while head is at 0.11, so it would be ideal to bring things more up to date). The more important consideration is using jit vs. pmap, which is the source of the issue. It should definitely be possible to switch to jit without changing any versions.

That said, we are looking into a better handling of the case you've outlined. I'm hoping to provide seamless support since this issue can arise even without using pmap.

[jecampagne]

Ok thank, nice that I've tigger a new dev.
Now, looking at Flax parallel training they are showing exemple using pmap. I've found an other exemple using Keras and the replica mechanism. I will try to change my code.