How to switch the mesh in jax.jit() ?

[Prayer3th]

Hey, guys, I'm trying to figure something out: I'm using jax.jit for the entire model forward pass, but I want to switch the mesh within the Moe layer—for example, attention weights tp=4, Moe weights ep=2, tp=2. The goal is to decouple the mesh used by the attention mechanism. What's the best approach to achieve this?

[yashk2810]

Great question!

Yes, you can switch the abstract mesh inside jit. For example:

def f():
    mesh = jax.make_mesh((2, 2), ('x', 'y'), axis_types=(jax.sharding.AxisType.Explicit,) * 2)
    jax.set_mesh(mesh)
    np_inp = np.arange(16.).reshape(8, 2)
    abstract_mesh = mesh.abstract_mesh

    @jax.jit
    def f(x):
      x = jnp.sin(x)

     # We are switching the abstract mesh here!!
      with jax.sharding.use_abstract_mesh(
          abstract_mesh.update(axis_sizes=(4, 1), axis_names=('a', 'b'))):

        # Once you switch the abstract mesh, you need to reshard the inputs to the new mesh in that context.
        x = reshard(x, P(('a', 'b'), None))

        out = x * 2
        assert out.aval.sharding.spec == P(('a', 'b'), None)
        return out

    arr = jax.device_put(np_inp, P('x', 'y'))
    out = f(arr)
    assert out.sharding.spec == P(('a', 'b'), None)

f()

Note that this kind of switching is best done under Explicit sharding mode: https://docs.jax.dev/en/latest/notebooks/explicit-sharding.html

[Prayer3th]

Thank you so much for your help! The detailed examples were extremely helpful. I've already implemented the decoupling of Moe parallelism and attention parallelism in the sglang-jax project based on this approach.

[yashk2810]

Awesome! Glad it worked out for you!