News for hackers. Stuff that matters.

A comment on We fine-tuned Llama 405B on AMD GPUs. Parent comment

Re:We fine-tuned Llama 405B on AMD GPUs by anthonix1 on Monday September 23, 2024 @08:04PM (#41631816)
Does JAX have its own implementations of matmul, flash attention etc? Or does it use the ROCm implementations like PyTorch does? (e.g,. hipblaslt, Composable Kernel FA etc)

Not too familiar with JAX, but the abysmal PyTorch training perf on MI300x is in large part attributable to the slow perf of the ROCm libraries it is using under the hood.

[ Parent | on HN ]
Re:We fine-tuned Llama 405B on AMD GPUs by jdeaton on Tuesday September 24, 2024 @12:47PM (#41638410)
JAX has a sub-system called Pallas[1] with a Triton-like programming model and an example implementation of Flash Attention [2]. It is quite fast. On TPUs I've heard that the XLA compiler already emits a flash-attention-like computation graph for a regular JAX implementation of attention so there's no need to have some specialized kernel in that case.

1. https://jax.readthedocs.io/en/latest/pallas/index.html

2. https://github.com/jax-ml/jax/blob/main/jax/experimental/pal...

[ Parent | on HN ]