8 ms·
We fine-tuned Llama 405B on AMD GPUs
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!
- felarof 2y agoHey 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.
- latchkey 2y agoNice work! I was just playing with the inference side of things with 405B myself this weekend [0]. 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 https://x.com/HotAisle/status/1837580046732874026
- felarof 2y agoNice! I need to calculate MFU. GPU, VRAM details can be found in the repo: https://dub.sh/amd-405b-res 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
- deleted 2y ago[deleted]
- abalaji 2y ago@dang: could we get url to include the username since this isn't about Obsidian itself, but rather a user generated blog?
- m00x 2y agoIt's strange that HN didn't include the full domain "publish.obsidian.cmd".
- sieabahlpark 2y ago[dead]
- meiraleal 2y agoThat's something obsidian should fix if they care about not looking like they are being impersonated on HN.
- viraptor 2y agoObsidian can't do anything about it. It's HN chopping up the url
- 3abiton 2y agoFirstly great work! I dabbled with AMD GPUs and ROCm support a year ago, and it was obvious AMD still a long way from catch ling up with Nvidia. While opting for JAX is in an interesting approach, what were the challenges for you deviating from pytorch (being the standard library for ML)?
- 6y56h56 2y agoI cannot get AMD ROCm running on my debian 12 system which is what I think is causing Ollama to use CPU instead of GPU. So I guess there is still a long way to go.
- ants_everywhere 2y agoI've had more luck with the ROCm docker container. I run it via k8s. It was pretty painless to set up and has been mostly painless since. Prior to that it was nearly impossible to get Jax running reliably on ROCm. Even with the container, you have to be careful installing Python libraries because they can still break things.
- lenova 2y agoI just recently went down the AMD GPU + ROCm rabbit hole as well. ROCm 6.2 was just released in August of this year and introduces a lot better support, though as the above poster mentioned, isn't merged into most recent OSes. This Github repo is good for tracking the latest Ubuntu + ROCm install process: https://github.com/nktice/AMD-AI https://github.com/nktice/AMD-AI
- latchkey 2y agoThat's a nice repo of random installation notes. Very helpful, thanks!
- superkuh 2y agoYou'd probably have a lot better luck using Vulkan acceleration (not ROCm) of llama.cpp as backend to ollama. It is incomparibly easier to set up and maintain compared to ROCm. You can actually do it on your computer's normal OS instead of inside a bunch of container/vms where the system libs are entirely customized to running just that one application. 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.
- yeahwhatever10 2y agoWhere is the performance data?
- felarof 2y ago(author here, sorry for the delay in replying, was stuck in back-to-back meetings) I updated our github repo to include GPU, VRAM utilization data (https://github.com/felafax/felafax?tab=readme-ov-file#amd-405b-fine-tuning-run 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%.
- dahart 2y agoI’d be interested to hear how long it takes, wall clock, to train the same model to the same loss with same number of batches. I don’t trust utilization to say anything useful.
- manojlds 2y agoThought this was a post from Obsidian at first. Why haven't they done the GitHub.com vs GitHub.io thing yet.
- gbraad 2y agoSame thought here. Why would Obsidian bother with AI? Oh wait, this is publish? So this is what $8 per month gets you? I am amazed, as I would have at least expected a subhost: [username].publish.obsidian.md
- codetrotter 2y agoLooking at the URL has me thinking that this confusion would be resolved if HN adds a small piece of logic to treat the domain publish.obsidian.md specially, just like how HN already does for pages served under forbes.com/sites which is not written by the Forbes staff themselves. 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?
- system2 2y agoWhy is obsidian (a note-taking app) doing this?
- Zerot 2y agoThey aren't. This company is using obsidian publish to publish documents.
- system2 2y agoWeird.
- JonChesterfield 2y agoIt's a very easy way to write some markdown that renders as a website. If the URL didn't have obsidian in the title I wouldn't have guessed it was involved.
- yieldcrv 2y agoIs AMD any closer to extracting value from this with large orders of their GPUs causing a shortage? I’m getting the impression of “no”
- latchkey 2y agoI get your sarcasm. But at this point, unless you are willing to give up all of AI to a single source for hardware and software, we've got to start to work towards alternative solutions. People are working against a massive head start and there is clearly a lot to work on on the software side. Give it time.
- Stem0037 2y agoIf possible, it would be interesting to explore ways to overcome the memory constraints and run a JIT-compiled version. This could potentially lead to further performance improvements.
- felarof 2y ago+1, we still have a lot of performance we can extract! JIT-compiled train steps, more optimized data loading and sharding, gradient accumulation, and activation checkpointing. We will continue building and will do another blog soon after implementing all the improvements!
- steeve 2y agoWe (ZML) measured MI300X at 30% faster than H100. These are great chips!
- chillee 2y agoTo be clear, this performance is quite bad (presumably because you didn't manage to get compilation working). 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.
- jamalaramala 2y agoThat's exactly what I wanted to ask: 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.
- brutus1213 2y agoDoes any Cloud provider have a 8xAMD MI300 host that one can rent? I use AWS for a lot of my professional work, and was hoping to try out an AMD GPU.
- latchkey 2y agoDisclosure: my company rents 8xMI300x, contact us.
- JonChesterfield 2y agoOracle do. Others will probably follow, though I expect the smaller players are more reasonable to interact with.
- latchkey 2y agoHas this changed? Previously, my understanding was that they only deal with large quantities.
- JonChesterfield 2y agoThat's a very good point. I know they've got boxes because I've used them, but didn't deal with any of the setup myself and we're probably a large quantities case.
- latchkey 2y agoStill hard to tell, but pushing 16k gpus on people doesn't seem be encouraging in this regard. https://ir.amd.com/news-events/press-releases/detail/1217/amd-instinct-mi300x-accelerators-available-on-oracle-cloud https://ir.amd.com/news-events/press-releases/detail/1217/am... I'm happy to take the stragglers though. ;-)
- lerpgame 2y ago[dead]
- oliver_jack 2y ago[flagged]
- moritonal 2y agoI guess it doesn't matter if this is human or bot, an advert's an advert.
- varispeed 2y agoHow do you buy such a GPU or is it still only reserved to the rich so they can get ahead of the game once the pleb gets their unwashed hands on these cards?
- latchkey 2y agoBuying them is pretty easy. You call up Dell and place an order. Let me know if you want to talk to a sales person, or we can help source them for you. 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.