We fine-tuned Llama 405B on AMD GPUs (publish.obsidian.md)
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!
106 comments
[ 3.4 ms ] story [ 188 ms ] threadWe'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!
Also any technical issues encountered?
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.
Were you using activation checkpointing? https://jax.readthedocs.io/en/latest/_autosummary/jax.checkp... is very important for keeping memory usage reasonable when training large models.
There are good 1p [a] and 3p [b] benchmarks comparing TPUs vs NVIDIA GPUs.
[a] - https://github.com/GoogleCloudPlatform/vertex-ai-samples/blo...
[b] - https://arxiv.org/pdf/2309.07181
But my take performance per dollar of TPU > AMD > NVIDIA.
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?
And by work I don't mean: spend two weeks trying to get the drivers set up and never update the server again.
[1] https://wandb.ai/augmxnt/train-bench/reports/Trainer-perform...
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.)
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.
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?
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.
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.
1. https://jax.readthedocs.io/en/latest/pallas/index.html
2. https://github.com/jax-ml/jax/blob/main/jax/experimental/pal...
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?
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.
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.
I don't quite understand this argument. Lack of transparency from running PyTorch so instead we're gonna leave it all to XLA? How does this solve the "transparency" issue?
Moreover, this will get worse as more CUDA specific features are added to PyTorch with ad-hoc fallback functions.
I guess OP is saying that XLA is more transparent in this regard, because it wouldn’t use functions like these and the generated comparable code would be on-pare performance wise?
Perhaps if XLA generated all functions from scratch, this would be more compelling. But XLA relies very heavily on pattern-matching to common library functions (e.g. CuDNN), and these patterns will certainly work better on Nvidia GPUs than AMD GPUs.
In this way, I actually think explicitly calling the common library functions is actually much more transparent.
All you have to do is pip install the ROCm version of PyTorch (or run the docker image) and it's seamless (the ROCm version just treats torch.cuda as calling ROCm).
I've used axolotl (trl/accelerate based), torchtune, and LLaMA-Factory, which are all PyTorch-based without any issues for training.
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).
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)
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)
/edit I think it did alright…
==========================================================
Let's break down the results described in the post.
Context: The post is discussing the performance of a large language model on a MI300X GPU, which is a high-performance computing (HPC) system. The model has approximately 405 billion parameters and is trained using a batch size of 16 and sequence length of 64.
Key metrics:
MFU (Million Floating-Point Operations per Second): This is a measure of the model's performance, specifically the number of floating-point operations (FLOPs) it can perform per second.
FLOPs: The number of floating-point operations required to perform a matrix multiplication, which is a fundamental operation in deep learning.
GPU performance: The MI300X GPU is capable of 1.3 petaflops (1.3 x 10^15 FLOPs) per second in bfloat16 (a 16-bit floating-point format).
Calculations:
The author provides two calculations to estimate the MFU of the model:
Initial calculation: Assuming full weight training (not LoRA), the author estimates the MFU as:
405 billion parameters
2 FLOPs per matrix multiply per parameter
3 matrix multiplies (forward, backward parameter, and backward activation)
Batch size 16
Sequence length 64
30 seconds to complete the calculation
1.3 petaflops per second per GPU
8 GPUs
The calculation yields an MFU of approximately 0.8%.
Revised calculation: After correcting the assumption to use LoRA (a technique that reduces the number of FLOPs), the author revises the calculation by removing the backward parameter pass, which is only applied to adaptor weights (much smaller than the main weights). This yields an MFU of approximately 0.53%.
Interpretation:
The results indicate that the MI300X GPU is not yet optimized for this large language model, with an MFU of only 0.53% (or 0.8% in the initial calculation). This is a relatively low performance compared to the theoretical maximum of 1.3 petaflops per second per GPU. The author notes that the biggest issues are related to scaling to multiple nodes, suggesting that the performance may improve when running on a larger cluster.
The revised calculation using LoRA reduces the MFU by about 33%, indicating that using this technique can lead to a more efficient use of the GPU resources.
==========================================================
MFU means model flops utilitization. it is a measure of efficiency from 0% to 100% . 100% means that the model is running at maximum possible efficiency, i.e. 1.3 petaflops per gpu.
==========================================================
In that case, the results indicate that the MI300X GPU is running the large language model at a relatively low efficiency, with an MFU of 0.53% (or 0.8% in the initial calculation).
This means that the model is only utilizing a tiny fraction of the GPU's maximum theoretical performance of 1.3 petaflops per second. In other words, the model is not fully utilizing the GPU's capabilities, and there is a significant amount of headroom for optimization.
To put this into perspective, an MFU of 100% would mean that the model is running at the maximum possible efficiency, using 1.3 petaflops per second per GPU. An MFU of 0.53% or 0.8% is extremely low, indicating that the model is running at a tiny fraction of its potential performance.
The author's comment that the MI300X results are "promising functionally" suggests that the model is able to run, but the low MFU indicates that there are significant opportunities for optimization and performance improvement.
Backward pass still runs on the non adapter weights. But yeah 10 TFlops/GPU specially on tiny sequence size is very bad compared to what you can get on Nvidia. And I believe the difference would be even higher with large sequence length.
Why compute gradients with regards to weights that aren't going to be updated?
Maybe challenging for Darrick, but in reality, you just need to follow the directions and configure things correctly...
https://techdocs.broadcom.com/us/en/storage-and-ethernet-con...
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.
I'm not convinced that 'torch.cuda' is really that bad since the AMD version of PyTorch just translates that for you. More like a naming problem, than anything. Fact is that it is just as easy to grab the rocm:pytorch container, as it is the rocm:jax container.
I don't see very many numbers posted. What MFU did you get?
[0] https://x.com/HotAisle/status/1837580046732874026
I need to calculate MFU. GPU, VRAM details can be found in the repo: https://dub.sh/amd-405b-res.
I plan to reattempt the training run next weekend and JIT the entire training step to calculate MFU then
Even with the container, you have to be careful installing Python libraries because they can still break things.
This Github repo is good for tracking the latest Ubuntu + ROCm install process: https://github.com/nktice/AMD-AI
AMD's support of consumer cards is very, very short. By the time it's stable enough for a new card to run the card is no longer supported. In 2021 I bought an AMD GPU that came out 3 years before and 1 year after I bought it (4 years since release) they dropped ROCm support.
llama.cpp and stable-diffusion.cpp offer Vulkan backends but generally you can run most models on Vulkan if you use IREE[1].
[1] <https://iree.dev/guides/ml-frameworks/>
Note: if you're building llama.cpp, all you have to do is swap GGML_HIPBLAS=1 and GGML_VULKAN=1 so the extra effort is just installing ROCm? (vs the Vulkan devtools)
ROCm:
Vulkan: EDIT: HN should really support markdown...Vulkan only exposes the raw compute capabilities of the hardware and any well optimized Vulkan application can reach the full performance, but you need to write such optimized code.
On the other hand, ROCm, like CUDA, includes optimized libraries for certain applications, like rocBLAS.
It is likely that here the ROCm backend uses optimized library functions, perhaps from rocBLAS, while the Vulkan backend might use some generic functions for linear algebra, which are not optimized for the AMD GPUs.
The problem with ROCm is that for non-bleeding edge AMD cards you have to install an out of date unsupported version of it because the $current version does not support your card. And that means containerization woes. If you're going to spend $800 on a top of the line current generation video card anyway then you'll have fewer problems (for a few years).
Also, the vulkan vs. rocm performance difference for non-bleeding edge non-top of the line cards is smaller.
I have not noticed any remarkable differences between Vulkan and ROCm when using IREE but it's not a turnkey solution yet[1].
[1] <https://github.com/nod-ai/sharktank/blob/main/docs/model_coo...>
If there were improvements on the AMDGPU DRM driver side, you would not see them in Debian any time soon, as the 6.1 LTS kernel will be stuck with roughly whatever shipped January of last year. This is just a shortcoming in the Linux kernel, due to its lack of any kind of stable ABI for drivers.
Of course it is possible this would help nothing or even hurt. My experience running stable (or even newer) kernels has been quite good, though. I run stable or newer across a few devices and run into hiccups not more than once every few years, which is definitely worth it to be able to get new driver improvements years in advance.
(FWIW Debian is not even supported by ROCm[1]... although distros with even older kernels are. But, even if ROCm works, I can't imagine you will get ideal hardware support when running older kernels. I am not sure if ROCm has some workaround for enterprise Linux distributions specifically, but it feels like they must, given how many of their customers in the datacenter are likely to want to use them.)
[1]: https://rocm.docs.amd.com/en/latest/compatibility/compatibil...
The firmware-amd-graphics package in stable is too old to properly support RDNA 3. It kind of works, but it is quite buggy. All RDNA 3 users on Debian 12 should be sure to install the kernel and firmware from bookworm-backports.
There is full support for RDNA 3 hardware enabled on Debian Testing (both in the drivers and runtime libraries). The Debian ROCm Team intended to backport all the ROCm packages from Testing into Bookworm, but have been held up as LLVM 17 is not available in bookworm-backports (yet?).
> FWIW Debian is not even supported by ROCm
ROCm does not support Debian, but Debian supports ROCm. Most of the libraries that comprise ROCm have been directly packaged by the distribution.
Old kernel + old kernel driver + new rocm => the driver doesn't really know what the userspace is doing and you get the bugs which have been fixed since
Old kernel + new kernel driver => very ymmv, the internel kernel api is not stable
New kernel + matching driver, old rocm is probably OK, unless you're using upstream clang in which case it's all bad once more
ROCm was designed and implemented in the HPC environment, where you know the exact kernel in use and the whole stack is deployed as one self consistent lump. Driver, compiler, libraries and so forth. It's not having such a good time in the Linux world of mix and match because the aggressive internal testing structure assumes you're using a consistent system. Backwards/forwards ABI and API compatibility is difficult, expensive and slow so it's not where the money is being spent. Rightly so, probably.
We initially started with the goal of fine-tuning LLaMA 3 on TPUs, but PyTorch XLA was clunky, so we decided to rewrite the model in JAX. That said, as mentioned earlier in the thread, we also believe JAX is a better platform for non-NVIDIA GPUs and want to build on JAX+openXLA for building infra for non-NVIDIA GPUs.
I updated our github repo to include GPU, VRAM utilization data (https://github.com/felafax/felafax?tab=readme-ov-file#amd-40...)
Note: we couldn't run the JIT-compiled version of the 405B model due to our code/VRAM constraints (we need to investigate this further). The entire training run was executed in JAX eager mode, so there is significant potential for performance improvements.
GPU utilization across the board was still ~30-40% even with eager mode, which is quite good! With JIT, I think the GPU util can easily shoot up to ~50-60%.
But struggling to get custom domain to work with it (have emailed support).
So instead of showing the domain as obsidian.md, HN would show the domain for this link as publish.obsidian.md
Maybe something for dang to consider if he sees this comment?
I’m getting the impression of “no”
You're getting 35 tokens/s for a 405B model, which comes out to about 85 Teraflops. 8 MI300x GPUs comes out to 10.4 Petaflops, so you're getting about 0.8% MFU (which is about 40-50x worse than decent training performance of 30-40% MFU).
For AMD's sake, I hope that it's your software stack that's limiting perf.
Their github page claims that it is possible to "tune LLaMa3.1 on Google Cloud TPUs for 30% lower cost", but they don't mention performance.
https://ir.amd.com/news-events/press-releases/detail/1217/am...
I'm happy to take the stragglers though. ;-)
The hard parts are that you have to buy 8 of them at a time, you need cooling, each server takes up around 10-11kWh and the fans sound like a 747 taking off. Oh and they cost a small fortune.
Disclosure: I handle the hard parts and you can rent them from me.