jax.lax.with_sharding_constraint within a function that will be vmaped

[vvvm23]

Suppose I have the following line within a function:

...
x = jax.lax.with_sharding_constraint(x, NamedSharding(mesh, PartitionSpec(None, 'tensor')))
...

where x is a tensor of shape (sequence_length, dim). This will apply the constraint to the GSPMD partioner to shard the tensor along the hidden dimension and not along the sequence dimension.

However, suppose this function is then vmaped to add a batch dimension, which we want to also shard across 'data'. Basically, the sharding pattern ('mesh', None, 'tensor').

How can we apply this sharding constraint to the vmapped function?

[mattjj]

Thanks for the question! Just to make sure I understand:

However, suppose this function is then vmaped to add a batch dimension, which we want to also shard across 'data'. Basically, the sharding pattern ('mesh', None, 'tensor').

Did you mean "Basically the sharding pattern ('data', None, 'tensor')? That is, the first element of the tuple is different.

[vvvm23]

Yes sorry :,)

Basically the non vmap version only has a notion of tensor parallelism, and non batched inputs.

I want the vmap version to have a constraint that also included the data parallelism dim over the new batch dimension.

[mattjj]

This is basically what the experimental and incomplete spmd_axis_name argument to vmap is for: just pass spmd_axis_name='data' to your vmap call.

It's incomplete in one particular way: if the function you're vmapping also has a shard_map in it, and the body of that shard_map has any collectives over the axis_name of the vmap, then all hell will break loose.

WDYT?

[vvvm23]

That works for me :) Thanks for the pointer!

One more question, if I passed in tensors with sharding pattern ('data', None, 'tensor') to the vmaped function without using spmd_axis_name, I am guessing it the compiler would likely automatically use this data axis even without the constraint applied?

Sorry if the question doesn't make sense :D