Bringing PyTorch Monarch to AMD GPUs: Single-Controller Distributed Tra... Liz Li & Zachary Streeter

This talk covers the collaboration between AMD and Meta to enable single-controller distributed training using Monarch on AMD GPUs. The speaker presents an overview of the Monarch framework, highlighting its fault tolerance features, including the use of a supervision tree and actor isolation. They discuss the integration of Monarch with Torch FT and Torch Titan, showcasing how these technologies work together to manage failures without requiring a complete job restart. The talk includes a real-world case study demonstrating the effectiveness of this approach, wherein failures during training are isolated to individual replicas, allowing the rest of the training processes to continue seamlessly. The session concludes with discussions on future plans to enhance the Monarch framework and expand its capabilities in various distributed training environments.

Good morning everyone. So, welcome to our talk. I'm lately from AMD. So, today I'm going to talk about a very exciting work with my colleague Zach on bringing Monarch AMD GPUs. So, how we enable single controller distributed training on Rock M. So, this is a joint work from AMD and Meta. So, huge thanks for all the contributors and sponsors. So, let me quickly walk through what we

are going to talk today. We will do a quick recap of our PyTorch conference last year to scale to scale Torch Titan to 1K AMD GPUs. Then, that will set the stage of introducing Monarch and how we support Monarch on AMD GPUs like pipification, RDMA pass that will cover. Then, I will also walk through the real real world case study of 4 terabyte using Monarch, Torch FT,

and Torch Titan. So, then we will conclude for next steps. Okay. Yeah, this is a this is a very quick recap from PyTorch 2025. We presented the scaling Torch Titan to 1K GPUs on AMD. Achieved linear scaling of FP8 training at a scale. So, after that we upstreamed all our key optimizations on open source project like Torch Titan and Torch AO. So, but the key challenge of

training is that we have a lot of training failure that can cause like restart of the whole training job. So, that led us to think about how can we how can we think about improving the resilience by design. So, this is where Monarch comes into the picture. So, we want to leverage the supervision tree and fault isolation feature of Monarch. Um so, now hand over to Zach

to talk about the support. >> And let me get into position. Uh bonjour. Is my audible to everyone? In the back, can you hear us? Okay. Kind of noticed that was an issue last time. And kind of a nod to Christopher Hitchens. Okay. Um so, I believe there was a Monarch talk previously. Um so, if you want to go into the gory details of Monarch, I suggest

you check out the recording whenever that's released. But for our purposes, I will go over at a high level what Monarch is. Um So, here on the left, I'm showing the programming model. Uh the gist of it is you define an actor uh class. Within the actor class, you define some endpoints. and then you can spawn You can spawn it across a mesh of GPUs where you

can call those endpoints. Pretty simple. Um the aesthetic is it looks like it's running locally, but it is running distributed. So, that's That's pretty nice. Um so, how does Monarch actually allow for this aesthetic? Uh you got to look at what's going on under the hood. So, you have, like I mentioned, uh the top layer, uh really nice Python API interface. Um beneath that is the Monarch

runtime. So, you have actors and meshes. Uh you have something called a supervisor tree, which I'll go into in the next slide. Um then you have tensor sharding. All of this is actually defined in Rust. So, that's pretty awesome. You have performance uh and not degraded performance. And you also have safety. So, that's Um so, you have a Python uh Rust binding in that instance. In Monarch,

you also have an infrastructure layer. Um and kind of a layer. Uh you have RDMA, nickel for collective communications on video, rickel for collective communications respectively on on uh AMD, um and then you have some schedulers defined in there as well. You have Slurm, K8s, and a new one that I wasn't aware of, uh Sky Pilot. But for our case for the uh fault tolerance, all we

need to uh look into a little bit deeper is actors and supervisor trees. So, let's check out those two data structures. And to the best way to look into those is to ask the question, I mean, we are presenting on how Monarch allows for um fault tolerance. So, natural question is, okay, well, how does Monarch actually handle failures? Um it's pretty captured well with these three bullet

points. So, failures are isolated, failures are hierarchical, and because of those two, you have recovery, it's quite fast. yeah. How are failures isolated? Um so, the actor data structure I mentioned before has a private state. So, um so, no shared memory. When one actor crashes, um it doesn't disrupt the uh the state of the other actors. So, this is fundamental. I'm sure many of you know in

distributed computing, um if one crash if one process crashes, that leaves the state in kind of a um you know, inconsistent state. So, you have to have a full restart. And with actors, you don't have to deal with that. So, that's super nice. Um second, the failures are hierarchical. So, if you look at this diagram on the right, this is kind of a uh an example of

the supervisor tree. Um this is a computer science tree, so the leaves are on the bottom. Um but so, you have in this example trainer actors, the parent of that's replica actor, and then orchestration manager. If a trainer actor fails, the error doesn't propagate sideways. It actually propagates up to And once the error propagates up to the replica actor kind of has a midlife crisis decision of

what do I do with this with this error? And that I'm actually showing you here in this pseudo code. Um and so it can either try to handle it locally or it can send it up to its uh respected parent. like I mentioned, when a child fails, Monarch calls the parent supervisor method. If the failure count um is under some limit shown here with uh max retries,

um then the replica actor restarts the trainers in place. So, and returns true. So, it's handled. But if the failure keeps happening, it will eventually return false and then it will escalate to its parent. And in this case, the escalation is uh uh to the orchestration manager and it will allocate a new Slurm allocation, right? So, I hope this hierarchy kind of demonstrates how recoveries are fast.

So, uh the replica actor handles most failures in place on restart. That's that's much quicker. Only when you've exhausted your local recovery here with um kind of a buffer, if you will, of of max retries, number of iterations, um do you actually escalate up to the orchestration manager. And that takes a much longer cuz of course you're, you know, allocating new resources. So, in summary, um most

failures recoveries are are much faster at a lower level um and training uh continues on healthy replicas and only catastrophic failures require heavier intervention. That's what the hierarchy kind of gives you. what I've gone over is briefly what Monarch is, the two pieces of Monarch that we need to understand for um uh fault tolerance and why, you know, the fault tolerance is actually really nice um with

Monarch pieces. I'm going to go on a little side quest on enabling um Monarch on AMD cuz that's represent. initially Monarch is built on top of Nvidia and CUDA. That's the the stack here on the left. On the right is where we needed to get to ROCm and AMD. And in the middle is the porting strategy. Um so initially the for example the nickel collectives we uh

oops Excuse me. Um we integrated hip by torch to handle the CUDA to hip source conversions automatically. Um that's kind of sounds like some mystical tool that just does everything for you magically and it's a bit hand-wavy and I'm I'm waving my hands literally. But uh that is actually a battle-tested tool developed by AMD. Um it was originally integrated into PyTorch PyTorch proper to enable um seamless

AMD hardware adoption. It's a bit of a strategy on our end. Um and it really helped and paid off to allow um power users to just use our hardware and they didn't have to rewrite any of their code. So with that precedence we applied the same approach to Monarch. Automatic conversion, same user code. You don't have to do anything and it works automatically on AMD hardware. Um

the GPU operations we needed to translate to the CUDA driver API uh for hip. This uh build system now auto detects which platform you're on um and configures accordingly. Uh the RDMA path um strictly talking about the data control path um we updated lib by verse integrations to work with hip instead of CUDA. It's all integrated with this um environment variable. Also gave uh the Monarch team

uh type aliases for Rust so they have this exact same Rust code but it compiles both for um uh CUDA. So that's really nice. So at the time of making this slide all unit tests passed. I'm going to leave it at that. Um and we also support uh from now on ROCm 7.0 and up. And we, I mean jointly with Meta. This is upstreamed into the main

fork of ROCm. These two uh PR references here, if you would like to go into the gory details instead of this nice abstract high-level uh presentation slide, you can go to those PRs and check out everything. given that uh went over what Monarch is, the two pieces we need to understand to really understand how Monarch helps with the fault tolerance that Liz is about to present, um

and then the side quest of how we enabled everything on ROCm. Let's zoom out and see where we are. Um so, kind of like what I mentioned earlier, um at the top you see what Monarch runs. You either have Slurm for traditional HPC clusters, Kubernetes for cloud native, and Skypilot for multi-cloud services or scenarios. Monarch abstracts away the scheduler, which is really nice if any of you

had to like, "Okay, now I'm Slurm. I got to do K8 now." Um you just write one piece of code and it runs on any of those schedulers. That's that's a nice little um convenience to have. In the middle layer is is what we contributed. So, this Monarch uh on ROCm. This include um everything in the runtime listed there. And again, those are the PRs that you

can go in and check out all of the porting um that we did. and at the bottom you see what this enables. So, this bleeds into the actual demo Liz is about to show. So, Torch Titan if you're uh as your training engine, Torch FT for uh fault tolerance on the trainer level, um and these these are the tools that the PyTorch community is already using. Um

and now they work on and or Monarch on AMD. Just And the tagline at the bottom captures it. So, it's full ecosystem support. This is ready for production workloads. You're not adopting any experimental software. Uh um Super, super nice. Um and with that, we will go into the the demo for fault tolerance, and you can see a real application of this. >> Yeah, thank you, Zach. Okay,

yeah, let me walk you through a real-world case study for fault-tolerant training at a scale. So, as we probably everybody will know, like training at a scale will have always have like frequent failures from across the stack. It can be like GPU memory errors, it can be network issues, and hardware faults, or even software crashes. So, our solution is to build Monarch Torch Titan and Torch FT

stack. So, we want to leverage Monarch's supervision tree for fault-tolerant detection and isolation. We also want to leverage Torch FT to enable the replicas to continue So, if we compare the traditional training and the new solution, so traditional training we can see just with one failure, we will need to restart the entire job, and it has a heavy dependency of checkpointing, and with very heavy computer waste,

also pretty hard for debugging. With this new solution, we failure is isolated, and the recovery is local. Also, it's very minimal disruption, and the iteration is pretty fast. Okay, let's look into the architecture. So, basically high-level wise, it's a three-layer architecture. On top is Monarch, which is which is handles processes and the cluster orchestration. Torch FT is in the middle layer, which is to handle the fault-tolerant

at the step level, and Torch Titan is our training engine at the bottom level. So, how they interact with each other? Monarch will communicate communicate with Torch FD through Lighthouse, which can provide the membership information for fault-tolerant coordination, and Torch FD will integrate with Torch Titan. When a failure is injected into the training, Torch FD will detect it detect the failure. So, from the training loops perspective,

we can see uh if there is a failure, it will be detected by Lighthouse and handled by Torch FD, and then the training will continue without global interruption. Okay. So, here we built a demo to showcase how Monarch and Torch FD can isolate the um replica level failure from global restart to just a local recovery, and how to rejoin the workflow. So, if we look at this

uh diagram, the orchestration manager will launch and control the job, and the Lighthouse on the left is uh controlling the uh the Torch FD quorum coordination. The replica actors uh is a Monarch supervisor, and the replica uh uh on the bottom is like Torch Titan trainers on eight GPU processes. So, the gradient will sync every 20 steps. So, this is a basic layout. Then, Okay. Once the

training starts, orchestration uh manager will launch replica actor, um and then the replica actor will also launch the replica, which will run the Torch Titan training on eight GPUs. At this stage, Lighthouse has the quorum ID one with all four replicas ready. Now, we are going to inject a failure on replica zero here. So, the the failure will be detected upwards by replica actor zero from Monarch

supervisor. And but the rest of the replicas uh uh uh like continue training. So, this failure is isolated only on replica zero. Okay. So, the Monarch supervisor will mark this replica zero as a restarting um replica. So, it will stop the old process match and spawn a new match, also start the training again. So, it will take some time to restart. Here, Lighthouse will mark only three

replicas continue syncing. Uh here is replica uh quorum ID two. So, once the replica is um uh is complete the the failure uh recovery is complete. So, the Lighthouse will identify a donor, like in this case, a healthy replica one to do peer check checkpoint uh transfer. So, it will do like model optimizer scheduler and trainer state transfer. So, it can rejoin the training uh group. So,

now reforming with uh four replicas in this stage. So, lastly, um once the replica zero is synced up, we will see like all four replicas uh join the quorum again. So, now it's going to be going back to the global training again. So, above is the demo flow. Here, we are showing the experimental result. So, we are uh training Lama 38B and we uh inject the recall

failure every 180 180 seconds and the quorum sync is every 20 steps. If you look at the left a figure, this is the number of participants. It looks like um every time the failure injected, the number of participants will drop out, then it will rejoin after some time. So, um but if you look at the right figure, this is the loss loss loss curve. The green one

is the normal loss curve and the orange one is like um the loss curve with uh failure injection. So, when we can see with even with failure injection, the loss is steadily uh convergent. Um so, this shows like with with like uh frequent failure injection, we can still uh achieve global training without like restart every training job globally. We did the We did the similar uh without

uh Kubernetes 32 nodes with MI 355 cluster. So, similarly, if you look at the left a figure, it's the number of participants, and it fluctuates, but the like uh the magnitude is relatively smaller compared to the MI 300 cluster. we also have the very steady convergent loss curve on the right side. So, the key takeaway is from this experiment is like um the scaling is not enough

by itself. We have to recover from failures um automatically, and Monarch gives uh this us the cleaner way to do that. So, this approach works on both uh Slurm and the Kubernetes environment. Okay, that concludes to our next steps. So, for next step, we want to extend more NIC support of Monarch and improve the runtime performance. And we also want to expand potentially Monarch to pre-trained and

RL frameworks on Rock M. we also want to further optimize the fault tolerant performance like what to optimize the region reload latency and overlap the recovery and the compute time. And lastly and importantly, we want to continue contribute to open source community. Okay, thanks everyone. If you have any questions, yeah, please ask. >> Okay. >> And of course this is a conference. You can just talk to

us. I'm not jet-lagged anymore, so I welcome it. Not yesterday. That wouldn't have been I was not ready. >> Thank you so much. >> Thank you.