[RFC] Shardy:MPMD – an ergonomic and performant MLIR-based MPMD partitioning system

[JoelWee]

We've just published our RFC on Shardy:MPMD, and would love comments.

RFC: https://github.com/openxla/shardy/blob/main/rfcs/2025-06-18-mpmd-rfc.md
Read-only demo colab: https://github.com/openxla/shardy/blob/main/rfcs/2025-06-18-mpmd-rfc-read-only-colab.ipynb

Highlights from the RFC and colab:

• Shardy:MPMD is an MPMD partitioning system, aimed at enabling ergonomic and performant MPMD-partitioning while retaining single-device programming semantics.
• It comes with an experimental JAX API which compiles to a bunch of XLA programs, orchestrated by Pathways on Cloud.
  • See the colab for an example with different pipeline schedules applied.

Let us know what you think!

Feel free to comment here, or on PR #583.

[patrick-toulme]

Very impressive. Will this be supported on XLA GPU as well?

[JoelWee]

Yes, it is supported with XLA GPU! What Shardy:MPMD does is that it creates a bunch of StableHLO programs which can then be orchestrated. I.e. so it'll work with XLA GPU / anything that would work with StableHLO :)

[MoFHeka]

Any relationship between this and torch_xla?

[MoFHeka]

Compared with 1F1B, reinforcement learning is the best and important application scenario. It's possible to deploy multiple roles in RL system easily by Shardy:MPMD?

[JoelWee]

Any relationship between this and torch_xla?

Not that I'm aware of – we've discussed briefly with some of the folks there about using Shardy:MPMD, but from our discussions, it seems that IFRT IR would be a better fit (unless they need the MPMD partitioning/propagation that Shardy:MPMD gives).

Compared with 1F1B, reinforcement learning is the best and important application scenario. It's possible to deploy multiple roles in RL system easily by Shardy:MPMD?

What do you mean by RL system in this case? Do you mean e.g. the training + sampling for RL?

[MoFHeka]

@JoelWee In RL system, we usually create several Ray actors for policy model, reference model, reward model and so on. Each actor run in a training pipeline synchronously.
Maybe I misunderstood the Shardy:MPMD, is it just for building a Local SPMD (PJRT) for pipeline parallelism?
If I want to create a multiple model training pipeline, I still need to use IFRT IR? What's the relationship between IFRT IR and Shardy:MPMD? Will IFRT IR regard Shardy:MPMD as a SPMD component? But even in single node, JAX use IFRT not PJRT.

[JoelWee]

I see, thanks for explaining.

What's the relationship between IFRT IR and Shardy:MPMD? Will IFRT IR regard Shardy:MPMD as a SPMD component?

Shardy:MPMD via JAX will lower to IFRT IR, and execute via IFRT IR. I.e. once we go to IFRT IR, there's no more Shardy:MPMD, everything is IFRT IR + StableHLO + Shardy shardings since the MPMD program is encoded in the IFRT IR as Ifrt.Calls (Local SPMD) and Ifrt.CopyArrays (Cross SPMD data transfer).

But even in single node, JAX use IFRT not PJRT.

Yes, so it's worth clarifying that IFRT IR != IFRT. It's built as a layer above IFRT, e.g. an IFRT IR call will use IFRT to execute the SPMD program, and we can have multiple IFRT IR Calls in an MLIR program. E.g.

func.func ... {
%x = ifrt.Call "@stage0" devices=[0,1] (%arg0)  // IfrtIR Op executed with IFRT 
%y = ifrt.CopyArrays %x : devices [0,1] -> [2,3]  // IfrtIR Op executed with IFRT 
%z = ifrt.Call "@stage1" devices=[2,3] (%y)  // IfrtIR Op executed with IFRT 
}

In RL system, we usually create several Ray actors for policy model, reference model, reward model and so on

We can express multiple models as an MPMD program, so we could define the multiple model training pipeline with Shardy:MPMD.

E.g. we could have a simple "Pipeline of models" where we have

Policy model --> Reference model --> Reward model

encoded in the IFRT IR graph (and use Shardy:MPMD to obtain that graph).

However, IFRT IR is purely static, and doesn't currently allow running dynamic amounts of compute at the moment. And RL is typically quite dynamic. E.g. RL for LLMs involves sampling which can involve an unknown amount of steps.

If we want this sort of dynamism, it would probably make sense to have multiple programs (one for each model, MPMD or SPMD), and orchestrate it manually. It also sounds like your different models are already different MLIR programs (since they're in different Ray actors).

Where Shardy:MPMD gives a benefit, is transforming an existing MLIR module and splitting it. The most common case is pipelining, but there are opportunities for purely parallel execution too.

E.g.

[Angelogeb]

Thanks for publishing this, really exciting direction!
I understand that the frontend/user-driven scheduling APIs are still work in progress, but I think it might still be beneficial to discuss this.

I was wondering if you could briefly describe your considerations for using the following pattern:

  for i, x in enumerate(xs):
    # Note the mpmd.call here, with call counter, wrapping the accumulation
    # function.
    carry = mpmd.call(microbatch_step, call_counter=i)(carry, state.params, x)

Over a more conventional approach that avoids call_counter, such as:

carry = mpmd.scan(microbatch_step, carry, xs, ..., schedule=OneFOneB)

Using a scan-like primitive makes explicit that the schedule only affects the loop, and it could be unrolled to a given schedule even in the absence of pipeline parallelism.
This approach is also closer to how users currently write gradient accumulation in JAX and enables having multiple loops, each with a different schedule, within the same computation.

In JaxPP we have a treduce (accumulate_grads in the paper) variant that translates to a scan internally.
We hide the carried state to prevent users from producing back-edges that would defeat pipeline parallelism (i.e., iteration 1 at rank 0 cannot start until iteration 0 at rank k has finished).
I really like that all the remaining concepts such as Fragment scheduling and merging are close to those in JaxPP so it would be nice to consider sharing the frontend in the future.

[JoelWee]

Thanks! And great question!

We do have a version of the scan-like primitive (i.e. an mpmd.fori_loop, which is basically a scan) which we unroll to apply a schedule; although we apply the schedule globally. We've been thinking about applying the schedule locally, but haven't got to yet.

We opted for something more explicit in exposing the call_counters directly, since it exposed more functionality and flexibility. Obviously this is at the risk of giving people a footgun (like the example you give of hiding the carried state). We don't have any opposition to having a scan, but we'd want users to be able to manually do the equivalent without having to use the scan. E.g. exposing mpmd.call allows the user to write a scan-like wrapper around that.

it would be nice to consider sharing the frontend in the future.

+1 🙂

[MoFHeka]

One more question, it seems that Shardy:MPMD need to be run in close-source Pathway? So it may be a new runtime in the future? And what's different between Pathway and McJax. There are several runtimes confusing me.

[JoelWee]

it seems that Shardy:MPMD need to be run in close-source Pathway

For now yes, but we're working on an open-source runtime :)

And what's different between Pathway and McJax

The main difference is that Pathways is a single-controller runtime orchestrator, whereas McJAX is a multi-controller open-source runtime orchestrator. Aside from that, they have different performance and resilience characteristics, but that's true of any two systems.

[MoFHeka]

but we're working on an open-source runtime

Brilliant! Is this on-coming runtime a single controller runtime? What main differences and key solutions between this runtime and the existing Ray runtime(assuming GPU objects has been supported)? Is it able to run SPMD with dynamic shape, I mean uncertain numbers of output tensors such as RL.