How do I do this rather trivial parallelization?

[sokol11]

Hi. I'm trying to do what I believe is rather trivial parallelization. I have an fn that accepts four arguments:

1. configs of shape (devices, batches, hyperparams)
2. rng;
3. 2d dataset;
4. optional network params pytree;

I'm trying to transform the fn, so it knows to shard configs to individual devices, replicate rng and the 2d dataset, and shard the network params, I also need vmap configs and network params across the batch dimension. For the life of me, I can't figure out how. I think part of the problem is the sheer number of available parallelization APIs, with overlapping but slightly different syntax, and my difficulty of understanding the docs (not sure if it's me, or the docs can be improved, particularly the experimental sections).

Anyone providing a basic example I can follow would be a hero!

I should note that the actual fn is somewhat complicated with mutiple nested lax.scans and lax.conds. I think I may have specified the parallelization correctly at some point with shard_map but then it told me something called check_rep has not been implemented for lax_cond and needs to be set to False, then setting it to False failed in some other way I think. shard_map feels like an overkill for what I need anyway, but any advice, clarity would be greatly appreciated! Thanks!

[yashk2810]

Have you checked this doc out: https://jax.readthedocs.io/en/latest/sharded-computation.html?

[sokol11]

Thanks, I don't understand so many things lol. I guess let's start with the basics, why do we need to specify both the Mesh and PartitionSpec, what is the difference between the two?

[sokol11]

Also, I don't really understand what a Sharding is...

[sokol11]

So, if I can summarize, I don't understand the relationship between Mesh, PartitionSpec, and NamedSharding (is there also an unnamed kind?). I also don't understand why we need to specify all three, if in pmap only in/out_axes were needed

[yashk2810]

Mesh defines an nd grid of devices that you can shard your array on. With pmap, you could only shard 1 dimension at a time and to do multi-dimension sharding, you would need to nest pmaps (which became very complicated).

With LLMs, you often shard multiple dimensions different ways which is why you need an abstraction that can define such a sharding in a very straightforward way. This is where Mesh comes into play.

PartitionSpec defines how an array should be sharded given the axis of the meshes. For example:

# Here you have a 1d grid of devices with an axis name `x`. Imagine there are 8 devices.
mesh = Mesh(jax.devices(), ('x',))

# Here you are sharding the numpy array on mesh axis `x`.
# Which means that the numpy array on dimension 1 will be sharded 8 ways.
sharding = jax.NamedSharding(mesh, PartitionSpec('x'))
arr = jax.device_put(np.arange(8), sharding)

So a jax.sharding.Sharding just clubs these 2 concepts (mesh and PartitionSpec) into 1 data structure to keep things together and saner.

You can also create a 2D mesh and shard you array on both dimensions. For example:

# Here you have a 2d grid of devices with an axis name `x` and `y`. Imagine there are 8 devices.
mesh = jax.make_mesh((4, 2), ('x', 'y'))

# Here you are sharding the numpy array on mesh axis `x` and `y`.
# Which means that the numpy array on dimension 1 will be sharded 4 ways and dimension 2 will be sharded 2 ways
sharding = jax.NamedSharding(mesh, PartitionSpec('x', 'y'))
arr = jax.device_put(np.arange(16).reshape(4, 2), sharding)

# You can also shard dimension 1, 8 ways here by specifying both mesh axes together like this:
# NOTE: the tuple inside PartitionSpec combining both mesh axes. Here you just multiply 4 * 2 == 8 to determine how many ways
# each dimension is sharded.
sharding = jax.NamedSharding(mesh, PartitionSpec(('x', 'y'), None))
arr = jax.device_put(np.arange(16).reshape(4, 2), sharding)

There are multiple utilities you can use to figure out how your array is sharded

• arr.addressable_shards (try using this and you'll see the output)
• sharding.shard_shape(global_shape) (this will tell you what's the shape of each shard on each device given the global shape of your array)

I would appreciate if you ask all the questions at once so I can explain everything at once.

[yashk2810]

(is there also an unnamed kind?

Yes, it's called PositionalSharding but we are moving away from that. We encourage users to use NamedSharding because almost 99% of JAX users prefer that because it's simple and provides a good mental model to think about sharded arrays. This is not JAX but users showing us that they prefer NamedSharding because no one really uses PositionalSharding anymore :)

[sokol11]

Hey @yashk2810 thanks for all the guidance. I got it to work after internalizing some of your answers and spending more time with the docs. I think the biggest hurdle for me was not fully understanding how vmap works (I was manually adding an extra batch dimension before passing to shmap and vmap, which I later realized was not necessary). I'm still a little fuzzy on the difference between pjit and shard_map APIs, but I got shard_map to work and it sounds like that may be the stable API going forward, so I should be all set.

My last question is why do we need to manually shard arrays with jax.device_put if we have these wrappers like jit and shard_map that will do that automatically given the shardings? Thank you!

[yashk2810]

My last question is why do we need to manually shard arrays with jax.device_put if we have these wrappers like jit and shard_map that will do that automatically given the shardings?

Fantastic question!

The reason has to do with dispatch speed. jax.jit has a C++ implementation where after the first call, we go via C++ to do super duper fast dispatch. If your arrays are pre-sharded correctly, we can take this fast dispatch path.

If they are not, we need to drop out to python to shard them correct as specified by in_shardings (on jit) or in_specs (on shard_map). This ends up being slow so it's recommended to pre-shard your arrays before passing them to jit.

I'm still a little fuzzy on the difference between pjit and shard_map APIs

pjit or jax.jit (they are the same) does automatic parallelization for you. It shards the intermediates and the outputs for you using the GSPMD partitioner (https://arxiv.org/abs/2105.04663).

But sometimes, GSPMD may not do what you want. So you can drop down into device local mode (via shard_map) and write your own collectives.

Basically, with_sharding_constraint and shard_map are all escape hatches designed to give you control back from the compiler. You can use them as you see fit :)