How to get started with Pallas?

[srush]

I would like to start trying out Pallas for GPU (ideally just on colab). However I can't seem to figure out how to get started. It does not seem to be in 0.4.14. I installed from main and it is there but I get the error

JaxStackTraceBeforeTransformation: NotImplementedError: MLIR translation rule for primitive 'pallas_call' not found for platform cuda

What is the recommend installation procedure?

[sharadmv]

Thanks for the interest!

TLDR; Here is Pallas on GPU working on colab: https://colab.research.google.com/drive/1P0d9S3lQHiDCyWvtXf8pdBWAyY1uBptZ?usp=sharing

(No Pallas on TPU colab ready yet unfortunately.)

The installation story is a bit complicated at the moment which is why we don't have an official "Installation guide" yet. Pallas is now part of JAX but depends on

1. JAX-Triton: this is a plugin that allows Pallas to lower to Triton IR and compile that Triton to PTX. JAX and JAX-Triton are synced so JAX at HEAD is compatible with JAX-Triton at HEAD.
2. Triton: Pallas depends on Triton transitively through JAX-Triton. However, JAX-Triton depends on the Triton that is currently being mirrored internally at Google. This version of Triton usually lags Triton HEAD by a couple weeks. We thus have to externally pin to a specific Triton nightly that matches our internal Google version.
3. Jaxlib GPU: The PTX we generate in JAX-Triton is passed into a custom call that lives in Jaxlib GPU and when generating this custom call, we use bindings in the jaxlib. Sometimes we change the signatures of these bindings so if there is a JAX/JAXlib version mismatch, we will see errors. For a given JAX version, we need to find the corresponding nightly jaxlib GPU.

The solution I am using in the colab is to find a set of pinned versions for everything that I know works. When Triton cuts a release on PyPI, we can do a JAX and JAX-Triton release simultaneously with pinned dependencies. Until then, the solution I recommend is to:

1. Install JAX at HEAD. Install a JAXlib nightly matching your desired CUDA version.
2. Install JAX-Triton at HEAD. Ideally we have a pinned Triton version but we haven't been good about keeping it up to date. Otherwise, try to binary search to find a Triton nightly that works. Eventually you should find one.

This is a not so great user experience for now, but we have some long term solutions to this. The most compelling one is to have Pallas not depend on JAX-Triton but depend on Triton bindings that ship with jaxlib itself (JAXlib contains Triton as well because XLA GPU uses Triton).

[srush]

Thanks!

Think I'll interpret this as "wait 6 months". Excited to use it at some point.

[srush]

Oh actually this demo looks pretty easy to get started with.

Is there a command to export PTX?

[sharadmv]

Yes. To dump all the IRs, you can pass debug=True into pallas_call. (You can also pass interpret=True into pallas_call to emulate execution using XLA + HLO).

I've updated the colab to dump the PTX for matmul.

[srush]

(Nvm. Figured it out!)

Running into this issue:

XlaRuntimeError: INTERNAL: Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.custom_call' failed: ptxas exited with non-zero error code 65280, output: ptxas /tmp/tempfile-8e729b259b50-49a0586e-7706-60464c0f887be, line 5; fatal   : Unsupported .version 8.0; current version is '7.8'
ptxas fatal   : Ptx assembly aborted due to errors

Is this a jax / cuda versioning issue?

[sharadmv]

Right now, Triton ships with it's own CUDA and emits PTX at 8.0. With JAX-Triton, jaxlib compiles the PTX with a bundled ptxas that might be older and can't handle ptx 8.0.

So, you will need to install a jaxlib that has cuda12 support.

Alternatively, if jaxlib detects an older CUDA locally installed it will use its ptxas. So you'll have to remove existing CUDA installs (this is annoying and we should fix this). In the colab, we rm -rf the local CUDA 11 which should fix the issue.

Are you running into it locally?

[srush]

Thanks, I was able to fix this error by reading the issues.

However then I tried to start coding and immediately ran into this jax-ml/jax-triton#228

[andportnoy]

"wait 6 months"

#18603 (comment)

[jakevdp]

Stable Pallas installation is still tricky; work on this is being tracked in #18603