Pjit: What is the effect of defining in-compatible in_axis_resources/out_axis_resources?

[mattiasmar]

How does JAX (pjit) act if the in_axis_resources of one layer don't match the out_axis_resources of the previous layer?
If I leave in_axis_resources as None, would it be adjusted to match the out_axis_resources of the previous layer?

[mattiasmar]

...or if I leave out_axis_resources as None, would it be adjusted to match the in_axis_resources of the following layer?
Ping @skye ?

[skye]

Thanks for the ping. In general, the input to a pjit will be resharded to match in_axis_resources, unless the input is already correctly sharded (e.g. if it comes from a matching out_axis_resources). None just means fully replicated, so that won't automatically be adjusted.

We're working on a new feature called GlobalDeviceArray (https://github.com/google/jax/blob/main/jax/experimental/global_device_array.py#L134), which "knows" it's own sharding instead of using in_axis_resources when passed to a pjit. So if you set up pjit to output GDAs (see the docstring for how to do this), the output GDA shardings will be determined by out_axis_resources, and then you can use FROM_GDA (https://github.com/google/jax/blob/main/jax/experimental/pjit.py#L51) as a special input to in_axis_resources, and pass in the GDAs that were the output of the previous pjit. We're still working on documentation and ironing out any bugs, but this functionality is roughly ready to use if you're willing to deal with bleeding-edge features. You can look at the pjit unit tests (https://github.com/google/jax/blob/main/tests/pjit_test.py) for basic usage examples. Feel free to file issues if you run into anything with GDA.

[mattiasmar]

@skye I would be grateful for and ready to test and debug any example applying pjit GDA on a Flax train_state object.
With regular pjit I had very limited success. My best attempt is expressed in this collab. Though the code runs from beginning to end without failure, I cannot see any particular impact on runtimes when changing the PartionSpecs.

[mattiasmar]

@skye P.S. A pure JAX example (without FLAX) could also do me well, if it demonstrates the application of Pjit/GDA externally from the training script (that is: not through adding with_sharding_constraints or axis definitions inside the training functions). My vision is that a user should be able to pjit any model, without change the core code of the workload (already developed). If it is not already supported by JAX (I think it is though) I believe that kind of decoupling would be very useful for the community to have. That decoupling expresses two different roles: That of the DNN developer (the researcher) and that of the code optimizer (the SW/ML engineer). Thanks!

[skye]

Check out T5X for a real-world example of using pjit with Flax. We unfortunately don't have a more minimal example for educational purposes, but T5X definitely drives peak performance using pjit. You can also look at their partitioning guide, although this uses constructs they've built on top of pjit (still might be interesting!).

Lingvo-JAX is another example of real-world usage of pjit, which is starting to switch to Flax.

I believe both these frameworks still use sharding constraints in some places. I agree the dream is for the compiler to figure all this out by itself, but we're not quite there yet.

Hope this helps!

[mattiasmar]

I have a solution to solve the partitioning problem. All my questions is how about how to express those partition strategies (at one point of contact) using JAX/Pjit. If the later could be sorted out, we would have the compiler that you mentioned.

Thanks for the links. I'll have a look.