Load and save distributed array best practice

[sh0416]

Is there any reference or tutorial for load and save distributed array?

I want to do this in a distributed manner.

[sh0416]

import jax
import jax.numpy as jnp
from jax.sharding import Mesh, PartitionSpec, NamedSharding
from jax.experimental import mesh_utils

mesh = Mesh(mesh_utils.create_device_mesh((2, 1)), axis_names=('dp', 'mp'))

@jax.jit
def create_tensor():
    x = jnp.zeros((131072, 16384))
    x = jax.lax.with_sharding_constraint(x, NamedSharding(mesh, PartitionSpec('dp', None)))
    return x

x = create_tensor()
jnp.save("test.npy", x)
y = jnp.load("test.npy")

Save doesn't incure the GPU memory.. but the loaded tensor is located in a single GPU.

[sh0416]

I found that when we use jnp.savez, the array is moved to cpu. Is there any other clever way to save array?

[sh0416]

@jax.jit
def load_tensor():
    x = jnp.load("test.npy")
    x = jax.lax.with_sharding_constraint(x, NamedSharding(mesh, PartitionSpec('dp', None)))
    return x

This kind of approach works.