Multi-node training using sharding API

[neel04]

I'm using TPUs, and wished to do training across multi-node/hosts. The official docs suggest using xmap/pmap.

However, I'm using the sharding API to shard across multiple local devices.

So how can we extend the sharding to accomodate a multi-node setup?

AIUI, we should be able to to provide a sort of 3D sharding like (2, 8, 1) for 2x TPUs with 8 local devices each, DDP styled.
This would allow us to switch between n-way data parallelism and m-way model parallelism as outlined here.

But this doesn't seem to be the case?

Related: SO bounty thread I started here.

[wen020]

Do you currently know the best practices for implementing multi-host data parallelism with JAX?

[neel04]

I ended up using scalax

[wen020]

Can Scalax be used for distributed training on multiple machines?

[neel04]

yes

[ayaka14732]

However, I'm using the sharding API to shard across multiple local devices. So how can we extend the sharding to accomodate a multi-node setup?

Just change the code from:

arr = jax.device_put(arr, NamedSharding(mesh, P(*partition_spec)))

To:

arr = jax.make_array_from_callback(arr.shape, NamedSharding(mesh, P(*partition_spec)), lambda idx: arr[idx]) 

See #20041

[neel04]

is idx the host_id?

[yashk2810]

If your input pipeline is fully data parallel, take a look at the docstring of: https://jax.readthedocs.io/en/latest/_autosummary/jax.make_array_from_single_device_arrays.html

If am also planning to expose a helper function to do just this! Are you also asking on how to shard the weights?

[neel04]

Yes please. I'd prefer a complete jax example for TP, DDP and FSDP 🙏