Tensor Model Parallelism - xmap vs pjit

[siddk]

Hey folks,

I’m trying to write a model parallel Transformer implementation but have come across the seemingly similar xmap and pjit APIs; xmap is semi-documented in the Jax docs, but it seems pjit has been used more in practice (T5X).

Which API is the current best practice/what are the trade offs of each? Are there minimal examples that might be good to work through (happy to help create one if not, if y’all can just point me in the right direction)?

Thanks!

[hawkinsp]

(I converted your issue to a discussion, since it seems more like a discussion to me!)

I'd recommend pjit for most users writing model-parallel transformers.

With pjit, you program as if you have a single gigantic machine, and tell the compiler to shard the arrays according to a sharding policy.

The xmap programming model is a bit more like pmap in that you program with explict axes of parallelization in your program.

I think all folks working on large-scale transformers at the moment are using pjit. xmap hasn't been very popular by comparison. Hope that helps!

[xiaoda99]

There is a project on training on large-scale transformers which uses xmap: https://github.com/kingoflolz/mesh-transformer-jax
The very good thing about mesh-transformer-jax is that the codebase is very small and clean.
In comparision, t5x's codebase is very big and complicated. To do partitioning, it relies on named axis support implemented in flax (https://github.com/google/flax/blob/main/flax/linen/partitioning.py). This complexity makes it difficult for me to do things like implementing and experimenting with my own model variants, or loading checkpoints from mesh-transformer-jax.
As @siddk , I'm also wondering which one to choose between xmap and pjit for distributed transformer training. If there is a codebase as simple and clean as mesh-transformer-jax, or xmap will ultimately be deprecated in favor of pjit, I think I'd choose pjit.

[xiaoda99]

@siddk FYI, I've just found a "minimal" example of pjit model parallel: https://github.com/huggingface/transformers/tree/main/examples/research_projects/jax-projects/model_parallel
It does not rely on the flax.linen.partitioning named axis but instead specify partioning in a param by param fashion. I'm not sure if this approach is as general as t5x but the code is indeed much simpler.

[Sea-Snell]

I've trying to learn pjit as well and in addition to the links @xiaoda99 sent, I've also found the DALL-E mini codebase to be a good, easy to understand reference for how to use pjit: https://github.com/borisdayma/dalle-mini. Just putting it here for reference.

[siddk]

Revisiting this in 2023 with the new Jax.Array API — are there good examples for doing data and batch parallelism at scale (say for a GPT-esque model)?

[cgarciae]

Hey everyone!

1. Here is a guide the Flax team made on how to use pjit to scale up models, hopefully the concepts apply to other frameworks: https://flax.readthedocs.io/en/latest/guides/flax_on_pjit.html
2. As of jax 0.4.1 jit == pjit (almost), jax.Array lowers the barrier of entry to device parallelism, check out the Distributed arrays and automatic parallelization guide (shameless plug: I also covered this topic in a live stream 😅).
3. jit/pjit rely on the compiler deriving the optimal distributed program, is even if you help it by adding some jax.lax.with_sharding_constraint it might not be optimal. There are some efforts using using xmap that hand code communication between devices that achieve greater performance that pjit at the cost its code being way more complicated. The appeal of jit/pjit is that it requires very few changes to your code, you can get things done pretty quickly, where pmap/xmap require a more effort but give you more control.