Reshapable meshes

[jfc4050]

Hello!

is there a recommended way to have meshes that can change logical shape dynamically? or put another way, can have multiple axis name mappings?

to make it a little more concrete, user might want to have mesh of shape (PPa, DPa, TPa) in certain model regions, but (PPa, DPb, TPb) in others, where PPa * DPa * TPa == PPb * DPb * TPb.

then in sharding annotations etc, code would be able to reference either XXa or XXb

one possible application of this would be DeepSpeed-TED style expert parallelism, where in expert regions DP is "traded" for EP. Megatron expands on this even further and makes it so users can do arbitrary reshapes of the (DP, EP, CP, TP) submesh for the expert region (i.e. also reduce CP or TP in favor of EP).

another possible use case would be "hybrid sharded data parallel" like in MiCS, where we want to shard parameters only along a subset of DP.

[yashk2810]

Yes, I think there is a way to do this. Let me write an example to show you (in a little bit).

[yashk2810]

Actually we need to do some more work to make this possible. But this is roughly how you can do it (note this code raises an error right now which I'll fix):

In [19]: mesh = jax.make_mesh((4, 2), ('x', 'y'), axis_types=(AxisType.Explicit, AxisType.Explicit))

In [20]: a = jax.device_put(np.arange(8).reshape(4, 2), NamedSharding(mesh, P('x', 'y')))

In [21]: b = jax.device_put(np.arange(8).reshape(4, 2), NamedSharding(mesh, P('x', 'y')))

In [22]: @jax.jit
    ...: def f(a, b):
    ...:     c = a * b
    ...:     print(f'{jax.typeof(c)}=')
    ...:     with use_abstract_mesh(AbstractMesh((4, 2), ('x1', 'y1'), axis_types=(AxisType.Explicit, AxisType.Explicit))):
    ...:         c = mesh_cast(c, P('x1', 'y1'))
    ...:         d = c * 2
    ...:         print(f'{jax.typeof(d)}=')
    ...:     return d
    ...:

In [23]: with use_mesh(mesh):
    ...:     f(a, b)

See https://docs.jax.dev/en/latest/notebooks/explicit-sharding.html for more information on Explicit mode.

[jfc4050]

thanks! sorry for the late response, took some time to mess around with it.

whats the relationship between this and the explicit sharding? is it possible to have multiple active meshes but continue to leave it to GSPMD? my mental model was that it was orthogonal.

second, is there any interest for having this kind of programming model in JAX?

mesh0 = jax.make_mesh((4, 2), ("x0", "y0"))
mesh1 = jax.make_mesh((2, 4), ("x1", "y1"))

x = jax.device_put(np.arange(16).reshape(4, 4), NamedSharding(mesh0, P('x0', 'y0')))
with mesh1, mesh2:
   x = jax.lax.with_sharding_constraint(x, P("x0", "y0"))  # reference axes in mesh0
   x = jax.lax.with_sharding_constraint(x, P("x1", "y1"))  # reference axes in mesh1
   x = jax.lax.with_sharding_constraint(x, P("x0", "y1"))  # error

the idea being that the modeling code could continue to not care about meshes, it just specifies some parallelism axes that it wants to use.

[yashk2810]

is it possible to have multiple active meshes but continue to leave it to GSPMD

Yes, sorry my terminal history only has Explicit now so that's what I ended up creating. You can get the same behavior without Explicit i.e. Auto only.

with mesh1, mesh2:

I don't think you want to enter into 2 meshes like this. It should be scoped to each wsc or just pass NamedSharding(mesh1.abstract_mesh, pspec) to wsc.

[jfc4050]

yeah maybe that double mesh idea isn't ideal. the main problem i'm trying to solve is the need to have fine grained control of the mesh swapping in model code, where traditionally we can just have an outer mesh context, then model code for the most part can just specify pspecs without worrying about the mesh)

other than the code cleanliness type of arguments, in some cases, like this common pattern with flax where we're applying transformations to a big pytree of sharded objects, it would require a big rearrangement of the code if we had some of those objects sharded with different interpretations of the mesh

    # partition the model.
    # https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html#initialize-a-sharded-model
    state = flax.nnx.state(model)
    sharded_state = jax.lax.with_sharding_constraint(
        state, shardings=flax.nnx.get_partition_spec(state)  # these pspecs might have a mix of mesh interpretations
    )
    flax.nnx.update(model, sharded_state)

maybe another programming model could be something like this

with ReinterpretableMesh(
  (4, 2), ("x0", "y0"),
  (2, 4), ("x1", "y1"),
):
   x = jax.lax.with_sharding_constraint(x, P("x0", "y0"))  # reference axes in mesh0
   x = jax.lax.with_sharding_constraint(x, P("x1", "y1"))  # reference axes in mesh1
   x = jax.lax.with_sharding_constraint(x, P("x0", "y1"))  # error

[yashk2810]

where traditionally we can just have an outer mesh context

Yeah this is possible today. Let me write the first reply code in Auto mode and show you.

a = jax.device_put(np.arange(8).reshape(4, 2), NamedSharding(mesh, P('x', 'y')))    
b = jax.device_put(np.arange(8).reshape(4, 2), NamedSharding(mesh, P('x', 'y')))
    
@jax.jit
def f(a, b):
    c = a * b
    print(f'{jax.typeof(c).sharding}')
    with use_abstract_mesh(AbstractMesh((4, 2), ('x1', 'y1'))):
        d = c * 2
        print(f'{jax.typeof(d).sharding}')
    return d
         
    
with use_mesh(mesh):
  f(a, b)
 
# Output:
NamedSharding(mesh=AbstractMesh('x': 4, 'y': 2, axis_types=(Auto, Auto)), spec=PartitionSpec(None, None))
NamedSharding(mesh=AbstractMesh('x1': 4, 'y1': 2, axis_types=(Auto, Auto)), spec=PartitionSpec(None, None))

Once you enter into a context with use_abstract_mesh then you can specify partitionSpecs everywhere.