jax.Array must be fully replicated to be saved in aggregate file about orbax HOT 6 CLOSED

I'm told that this is a known bug. In the meantime, you should use a different method to replicate the state - perhaps just use pjit.

from orbax.

If replicating the arrays, they can be safely saved into the msgpack file. If calling unreplicate, I believe flax's behavior is to instruct Orbax to save using Tensorstore, which supports sharded arrays.

from orbax.

To use the aggregate option, you should either have numpy arrays, basic scalar types, or you can reshard your jax.Arrays to be replicated across all devices. This is pretty easy to do - just supply a sharding of None instead. JAX documentation has lots of details

from orbax.

Hi @cpgaffney1, thanks for the response. I am replicating the state prior to training with state = jax_utils.replicate(state). I also shard the batches during training with flax.training.common_utils.shard. This error goes away if I call flax.jax_utils.unreplicate() on the state before saving like so : ckpt = {'state': unreplicate(state), 'config': model.config}. Is this supposed to happen?

from orbax.

jax_utils uses jax.device_put_replicated. When I run the following

import jax

replicated = jax.device_put_replicated(np.arange(32), jax.devices())
replicated.is_fully_replicated

is_fully_replicated is actually False, which is surprising. Checking to see if this is expected behavior.

from orbax.

Thanks. What would be the difference between saving the replicated state after pjit vs calling unreplicate() and saving as I'm doing now?

from orbax.