5 ms·
Hey HN, we recently fine-tuned the llama3.1 405B model on 8xAMD MI300x GPUs using JAX instead of PyTorch. JAX's advanced sharding APIs allowed us to achieve gre
by felarof 2y ago
Hey HN, we recently fine-tuned the llama3.1 405B model on 8xAMD MI300x GPUs using JAX instead of PyTorch. JAX's advanced sharding APIs allowed us to achieve great performance. Check out our blog post to learn about the cool sharding tricks we used. We've also open-sourced the code: https://github.com/felafax/felafax https://github.com/felafax/felafax
We're a small startup building AI infra for fine-tuning and serving LLMs on non-NVIDIA hardware (TPUs, AMD, Trainium).
Problem: Many companies are trying to get PyTorch working on AMD GPUs, but we believe this is a treacherous path. PyTorch is deeply intertwined with the NVIDIA ecosystem in a lot of ways (e.g., `torch.cuda` or scaled_dot_product_attention is an NVIDIA CUDA kernel exposed as a PyTorch function). So, to get PyTorch code running on non-NVIDIA hardware, there's a lot of "de-NVIDIAfying" that needs to be done.
Solution: We believe JAX is a better fit for non-NVIDIA hardware. In JAX, ML model code compiles to hardware-independent HLO graphs, which are then optimized by the XLA compiler before hardware-specific optimization. This clean separation allowed us to run the same LLaMA3 JAX code both on Google TPUs and AMD GPUs with no changes.
Our strategy as a company is to invest upfront in porting models to JAX, then leverage its framework and XLA kernels to extract maximum performance from non-NVIDIA backends. This is why we first ported Llama 3.1 from PyTorch to JAX, and now the same JAX model works great on TPUs and runs perfectly on AMD GPUs.
We'd love to hear your thoughts on our vision and repo!
- ngcc_hk 2y agoGiven it is a migration, is there actual comparison of the same model on PyTorch vs your version. The comparison table there seems to be on technical side. Also any technical issues encountered?
- felarof 2y agoWe have a few technical issues that we still need to address: 1) This entire fine-tuning run was done in JAX eager mode. I kept running out of memory (OOM) when trying to `jax.jit` the entire training step. Even gradual `jax.jit` didn't work. 2) The current version doesn't have gradient accumulation, and with a batch size of just 16, that’s not ideal. I'm working on implementing gradient accumulation next. 3) We still haven't found a good way to load large sequence-length data (like 32k sequence length). Currently, before sharding the training batch across GPUs, it ends up loading the entire batch onto a single GPU’s VRAM and causes OOM issues.
- logicchains 2y ago> I kept running out of memory (OOM) when trying to `jax.jit` the entire training step. Even gradual `jax.jit` didn't work. Were you using activation checkpointing? https://jax.readthedocs.io/en/latest/_autosummary/jax.checkpoint.html https://jax.readthedocs.io/en/latest/_autosummary/jax.checkp... is very important for keeping memory usage reasonable when training large models.
- jgalt212 2y agoIs there some cost rule of thumb to compare Nvidia, AMD, and Google TPU?
- CuriouslyC 2y agoTPUs are slow but efficient and AMD has bugs but for some things works quite well. Nvidia is obviously the gold standard.
- felarof 2y agoHaha, TPUs are not slow :) All of Google's training (including Gemini models) is done on TPUs. There are good 1p [a] and 3p [b] benchmarks comparing TPUs vs NVIDIA GPUs. [a] - https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/community-content/vertex_model_garden/benchmarking_reports/jax_vit_benchmarking_report.md https://github.com/GoogleCloudPlatform/vertex-ai-samples/blo... [b] - https://arxiv.org/pdf/2309.07181 https://arxiv.org/pdf/2309.07181
- felarof 2y agoGood question. No good metric give performance depends on software stack (JAX vs PyTorch) + optimizations. But my take performance per dollar of TPU > AMD > NVIDIA.
- cameron_b 2y agoI'm glad to see a full implementation on AMD hardware. I'm not familiar with JAX, but the idea of providing an abstraction layer to more easily get to work on what hardware is available seems really valuable. Bringing back some competitiveness to the ecosystem will be a big win for workload mobility. I suspect that price/performance across implementations will be highly dependent on contract details, but do you intend to publish some comparisons in the future?
- llm_trw 2y agoDoes this work on the consumer grade cards like the 7090 XTX? And by work I don't mean: spend two weeks trying to get the drivers set up and never update the server again.
- lhl 2y agoA couple months ago I did some testing on some consumer cards. [1] I think you should be able to use torchtune or axolotl without anything besides installing the ROCm version of PyTorch. [1] https://wandb.ai/augmxnt/train-bench/reports/Trainer-performance-comparison-torchtune-vs-axolotl-vs-Unsloth---Vmlldzo4MzU3NTAx https://wandb.ai/augmxnt/train-bench/reports/Trainer-perform...
- tommiegannert 2y agoAm I reading that right that the 7900 XTX is on a par with 3090, and 4090 is twice as fast?
- lhl 2y agoYeah, those numbers are correct as of their testing (in June) although people who are really interested should check out the linked repo and do their own runs as software/optimizations have continued to change a lot and the RDNA3 side has a lot of untapped potential. Eg, the 7900 XTX has a huge theoretical FLOPS advantage over the 3090 but the results totally don't reflect that. One example of this hobbling is that RDNA3 only recently got backpass FA via a still under-optimized aotriton implementation: https://github.com/ROCm/aotriton/pull/39 https://github.com/ROCm/aotriton/pull/39 There are also still ongoing optimizations on the Nvidia side as well. In the beginning of the year the 7900 XTX and 3090 were pretty close on llama.cpp inference performance, but a few months ago llama.cpp got CUDA graph and FA support implemented that boosted perf significantly for both my 3090 and 4090. (For AI/ML, a used 3090 remains I think the best bang/buck for both inference and small training runs. You can pay twice as much for the twice as fast 4090, but at the end of the day you'll still wish you had more VRAM, so it's hard to really recommend unless you're going to use mixed precision. The RDNA3 cards are not as bad to work with as the Internet would have you believe, but they'd have to be a lot cheaper if your main use case was AI/ML for both the PITA factor and just from pure real-world performance.)
- germanjoey 2y agoHow are you verifying accuracy for your JAX port of Llama 3.1? IMHO, the main reason to use pytorch is actually that the original model used pytorch. What can seem to be identical logic between different model versions may actually cause model drift when infinitesimal floating point errors accumulate due to the huge scale of the data. My experience is that debugging an accuracy mismatches like this in a big model is a torturous ordeal beyond the 10th circle of hell.
- srcreigh 2y agoVery fascinating, can you explain more about a time when this happened? Like what area was affected by fp errors, why were they introduced (was it like refactoring of pytorch code?), how was this determined to be the cause?
- credit_guy 2y agoWhen you say "model versions", do you mean different quantizations of the model? Then it's not floating point errors that accumulate. Different quantizations of the model are different models. People will call such a model something like Meta-Llama-3.1-8B-Instruct--q4_0, claiming that it's just a "version" of the Meta-Llama-3.1-8B-Instruct. But it's just a lie. It's not the same model, and you should not expect the same results. There is no reason to debug the differences, what exactly would you expect to find, and what action would you envision to take once you find what you are looking for? However, is the quantized version still a useful LLM? Absolutely. Most people don't have an A100 to run the original model, so a quantized version is better than nothing.
- felarof 2y agoGood question. We used a new AI+math-based testing tool (benchify.com) to run comparison tests, but we are working on building more robust infrastructure for this. Translating models from PyTorch to JAX is core to our strategy. That said, this path is not uncommon (translating from one framework to another). HuggingFace translates Google's Gemma family models from JAX to PyTorch, and a ton of people use it.
- anthonix1 2y agoDoes 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.
- jdeaton 2y agoJAX 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 https://jax.readthedocs.io/en/latest/pallas/index.html 2. https://github.com/jax-ml/jax/blob/main/jax/experimental/pallas/ops/gpu/attention.py https://github.com/jax-ml/jax/blob/main/jax/experimental/pal...
- ipsum2 2y agoI, and several others, had no problem running on PyTorch on AMD GPUs, with no code changes from CUDA. Check out MosaicML's blog posts: https://www.databricks.com/blog/training-llms-scale-amd-mi250-gpus https://www.databricks.com/blog/training-llms-scale-amd-mi25...
- felarof 2y agoAhh, interesting, will take a look! Curious what are the steps to run PyTorch on AMD (does it work out-of-box with PyTorch+rocm docker image)? Does torch.compile work smoothly?
- anthonix1 2y agoYeah I would suggest taking a look at PyTorch on AMD before saying stuff like "scaled_dot_product_attention is an NVIDIA CUDA kernel exposed as a PyTorch function", because that is demonstrably false. Also, FWIW, I would suggest getting a small Llama 3.1 model training fast before trying to do a big 405B model -- faster to iterate and almost everything you'll learn on the small models will scale to the 405B.
- felarof 2y agoThanks for the feedback! I appreciate you pointing that out. My understanding was based on the PyTorch documentation for scaled_dot_product_attention (https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html https://pytorch.org/docs/stable/generated/torch.nn.functiona...). - "The function may call optimized kernels for improved performance when using the CUDA backend. For all other backends, the PyTorch implementation will be used." And was trying to make a broader point about the lack of transparency (in performance, lower-level impl) in PyTorch when running on NVIDIA vs. non-NVIDIA hardware.
- anthonix1 2y ago[flagged]
- anthonix1 2y agoAny direct comparisons to 8xH100? 2 toks/sec seems very slow! I haven't done any LoRA training on MI300x myself, but I have done LLama 3.1 full training on 8xMI300x and got pretty close to 8xH100 performance with my own kernels (ROCm is just too slow).
- felarof 2y agoOops, my calculation was wrong. Let me add an edit to the blog, thanks for pointing it out! My train step was taking 30s. And I was using a batch size of 16 and seq length of 64, making the training speed as (16*64/30) tokens per sec == 35 tokens per second (for fine-tuning in JAX eager mode). (I haven't done comparison with 8XH100)
- gdiamos 2y agoThat’s approximately 0.8% MFU - h100 would get more like 30% or 40% MFU if well tuned 405e9 parameters 2 flops per matrix multiply per parameter 3 matrix multiplies for (forward, backward param, and backward activation) passes batch size 16 seq length 64 1.3 petaflops per second per GPU in bfloat16 8 GPUs 30 seconds So that’s 0.8% = (405e9 * 2 * 3 * 16 * 64 / 30) / (1.3e15 * 8) Note that I’m ignoring the attention flops in this simplified calculation, but they would be a second order effect at this sequence length Also note that I’m assuming full weight training, not LoRA . The result would be lower MFU if using LoRA These MI300X results are promising functionally (it's tough to get any model this big running) but they have a long way to go on perf. It's also single node. The biggest issues I've seen on MI300X are related to scaling to multiple nodes. EDIT: The blog seems to indicate it is using LoRA. So we should remove the backward param pass from the equation above. Backward param only applies to adaptor weights, which are much more than 10x smaller, so we set it to 0 in the approximation. So we get 0.53% = (405e9 * 2 * 2 * 16 * 64 / 30) / (1.3e15 * 8)
- darrick_horton 2y agoWe've significantly optimized multinode on AMD MI300X for both stability and performance at TensorWave. There were certainly a lot of challenges, but we've become experts at multinode on AMD. We'd be happy to show you the latest results! They are quite compelling
- logicchains 2y agoDid you consider using https://github.com/AI-Hypercomputer/maxtext https://github.com/AI-Hypercomputer/maxtext ? It has a Jax llama implementation, and gets decent MFU on TPU and GPU (I've only tried it on NVidia GPU, not AMD).
- lostmsu 2y agoCould you share performance so we could compare?
- upbeat_general 2y agoscaled_dot_product_attention isn’t CUDA specific, it even works on TPUs.
- ewalk153 2y agoDo you see tinygrad as a useful lower level abstraction or is JAX sufficient to get pref out of AMD GPUs?
- felarof 2y agoTinygrad is great, but still in early stages I believe. JAX has matured a lot over last 6 years and XLA has been around for lot longer. We believe we can extract good perf from AMD with JAX + XLA kernels.