jax.Array must be fully replicated to be saved in aggregate file

[tatami-galaxy]

I'm trying to save a checkpoint and getting this error message. Saving code :

ckpt = {'state': state, 'config': model.config} save_args = orbax_utils.save_args_from_target(ckpt) checkpoint_manager.save(global_step + 1, ckpt, save_kwargs={'save_args': save_args})

state is an instance of flax.training.train_state. What could be causing this? I tried disabling jax.Array with jax.config.update('jax_array', False) but that does not work with jax and jaxlib 0.4.7.

[cgarciae]

Hey @tatami-galaxy, can you give a minimal example?

[tatami-galaxy]

The error goes away if I do flax.jax_utils.unreplicate() on the state before saving like so :

ckpt = {'state': unreplicate(state), 'config': model.config}

I'm calling replicate on the state before training starts : state = jax_utils.replicate(state). Is this expected?

https://github.com/tatami-galaxy/speech-processing/blob/main/experiments/python/jax/whisper_flax.py#L697

This is the full code (which now has a probably unrelated CUDA error if I do data processing with multiple processes with huggingface map). I will try to replicate the jax.Array error with a minimal example. Thanks for taking a look.