How to scan a sharded array efficiently?

[HeavyCrab]

Directly applying a scan method on a sharded array seems to consume more memory than shard_map. Is there a way to do something like shard_scan in an efficient way?

Basically what I want to do is like:

import jax
import jax.numpy as jnp
from jax.sharding import NamedSharding, PartitionSpec as P


mesh = jax.make_mesh((2,), ('d',))
a = jax.random.normal(jax.random.key(0), (2, 8192, 8192))
a = jax.device_put(a, NamedSharding(mesh, P('d')))

def scan_fn(carry, x):
    return carry + x.sum(), carry + x

carry = jnp.array(0)
carry = jax.device_put(carry, NamedSharding(mesh, P()))
carry, res = jax.lax.scan(scan_fn, carry, a)

res = jax.device_put(res, NamedSharding(mesh, P('d')))

But more efficiently.

[yashk2810]

You shouldn't shard the 0th dim which is being scanned over i.e. change the pspec of a to P(None, 'd')` and see if that helps. Sharding the scanned over dim doesn't really make sense to me.

[HeavyCrab]

The reason I want to scan over the sharding dimension is that I have a big batch of training data distributed across different devices (pspec = P('d')). I also have a model with pspec = P(). Normally, we would use a shard_map to let the model compute gradients on this distributed batch (a classical paradigm of DDP). However, I want to split this batch into K mini-batches and let the model update sequentially on these mini-batches, so I use scan.

This big batch is dynamically generated by a distributed prepare() function. The reason I don’t simply place it on a single device and make it generate data of size mini_batch_size each time is:

1. prepare is very time-consuming. This way, prepare and model weight updates can be executed asynchronously, saving time, so that the model can perform K updates in the time it takes to wait for one prepare.

2. On a single device, preparing data of size mini_batch_size already reaches the parallel capacity limit; enlarging the batch leads to linear growth in runtime or OOM error.

I’m not sure if this makes sense to you.