Lightning Talk: Monarch: An API To Your Supercomputer - Marius Eriksen, Meta

In this talk, Marius Eriksen presents Monarch, a project from Meta that provides a Pythonic API for managing distributed computing resources, particularly clusters of GPUs. He explains how Monarch allows developers to program large-scale machine learning applications as if they are operating on local resources. The speaker details how the architecture of Monarch includes meshes of processes for distributed tasks and emphasizes its imperative nature, which gives developers control over resource allocation and task management. He provides examples of implementing process jobs and actor communication patterns, highlighting features like RDMA buffers for efficient memory transfers. Throughout the session, Eriksen demonstrates the capabilities of Monarch through a live coding session, illustrating how it can be used to orchestrate data loading and training across multiple nodes in an efficient manner. He concludes with a practical demonstration of how agents can leverage these APIs to automate and debug distributed trainings effectively.

All right, hi. I'm Marius Eriksen, and I work at Meta on a project called Monarch, which is what we'll just talk about today. So, what is Monarch? It is a project that we introduced in the last PyTorch, and today I'm going to show you a sort of slightly different angle, I guess, of Monarch than we did last time. And in short, what Monarch is is fundamentally an

API to your supercomputer. Meaning, we want to make your cluster of GPUs, whether that's you know, attached to a single machine or attached to 100,000 nodes, programmable through ordinary Pythonic APIs. Very much like PyTorch makes, you know, your local GPUs programmable directly. And Monarch accomplishes this by representing all these resources as meshes of, you know, different parts of that distributed computing setup. So, we have meshes of

hosts and processes, and then actors that represent computation that runs within these processes. Now, importantly, Monarch is imperative, meaning that when you're programming Monarch, you have direct control over precisely where all the code runs, how resources are allocated, the life cycle of all of the hosts and processes and actors and so on. And fundamentally, what that gives you the ability to do is to write complete and

self-contained distributed machine learning applications as if they are a local Python program, and that's what I'm going to be demonstrating today. So, I'm first I'm just going to go over very, very briefly the sort of high-level concepts of the Monarch APIs, and then I want to spend most of the time actually showing you a demo of how what the kind of magic that happens when you combine

this with using agents actually to help you sort orchestrate and manage your jobs. So, what we have here is fundamentally fairly simple, and I think fairly expected in terms of how you think about and represent distributed computing tasks. So, first of all, we have a job. In this case, there's a process job, which means it's just a job that's represented by locally running processes. But you could

substitute this with like a Kubernetes job for example to run your your job on a Kubernetes cluster or within meta we have a scheduler called mast and so we have a mast job implementation, slurm and so on and so forth. And um once I have a job, I can create a bunch of processes that uh are running on these on this job. And what you're seeing here

is that I'm spotting two meshes of processes. One's one that represents a set of trainers and another that represents a set of data loaders. And then once I have these process meshes, I'm spawning actors onto them. And so I have a data loader actor that in this case, you know, uh loads some data from uh some data source that I give it. And then I pass those

loaders to the trainer actors which may now interact directly with with those data And then once I have these actor meshes, I can do the sort of normally expected actor communication patterns like broadcast and accumulate or direct calls and so on and so forth. And uh the broadcast here uh at the at the very last line is actually doing a lot of lifting. So uh uh monarch

implements scalable um multicasting effectively through um distribution trees and so on. Uh and also uh supports scalable accumulation of results from from calls as well. And so these are the fundamental building blocks that let you represent very very large clusters of um of computers uh you know, using simple imperative APIs but then be able to actually uh invoke, you know, things like calls in a very efficient

and cost-effective way. Um another really really important part of monarch is that we take great care to sort of separate uh control plane from from the data plane. And the the primary mechanism or building block we have to manage data movement within monarch is something we call RDMA buffers. And what an RDMA buffer is is the ability to effectively just represent some region of memory on a

machine as the kind of handle that I can pass around and do one-sided IOs with. So, I generate an RDMA buffer that represents some GPU memory on one machine, pass it to another machine, and that gives that other machine the capability to read and write directly into that GPU memory. And so, here again, I'm sort of demonstrating again, you know, fairly fairly short in here, but how

you might use RDMA buffers to represent batches of data in a data loader and then access those batches of data in a trainer. And by just using these APIs, underneath Monarch figures out how to, you know, use the underlying RDMA fabric to um uh bypass, you know, everything else in the data path and do direct memory-to-memory transfers, basically. And when you combine all these things together, uh

this gives you the ability to do single control orchestration. And so, what I've kind of built up here is sort of very obviously very simple sketch of a training application where I have a data loader and a and a set of trainer actors, and I want to, you know, effectively plumb the data from the data loaders to the training actors, and then orchestrate that just through a

normal ordinary Python script. So, another thing that we kind of, you know, honestly kind of discovered along the way is that one of the things that these building blocks enable us to do is to provide fairly sophisticated services as simple Python libraries. So, for example, one of the things we built was again just using the primitives that I just showed you, the building blocks, um is uh

fully RDMA capable remote mounts. And so, I can have, you know, a set of files or a path on one machine, like maybe the client machine that I'm running my script from, and then mount that file system on, you know, a bunch of machines in my cluster, and then the the fuse implementation is just implemented using a Monarch actor that passes RDMA buffers around. And so, we

can have extremely fast, you know, just ordinary POSIX uh file transfers in this The same thing uh applies to things like observability tools, like, for example, running PySpy's over your whole cluster. And so, uh what this kind of gives us at the end of the day is what you might call agentic legibility. So, uh by having these APIs sort of presented in this way, and by kind

of making your computation sort of tractable as just a local set of Python APIs that I can run, you know, literally using UV run on a uh on a local machine, this gives your agents sort of the capability to manage, you know, large computations across a cluster. And that brings us to our demo. So, what I want to show you, you know, almost running out of time

here, so what I want to what I want to show you really, really briefly is how uh we effectively built a very, very simple sort of what you might call a serverless SPMD tool using Monarch. And uh the way this works is that I run a tool from a single, you know, from from a client machine, and I can instantiate a bunch of different worker machines, again,

using whatever job implementation you prefer. And um uh what it uh what that does is that it allows me to just run arbitrary code on those machines, but then crucially, uh I mount my local directory on all of those remote machines, and then conversely, I uh gather mount, so I mount all of the remote directories uh directly on my machine as well. And so, this gives your

um agents a kind of like I/O to the whole cluster. I can make changes locally and have them instantly be propagated, you know, do printf debugging, run stuff, etc., etc. And then, in order to debug what's going on, the agent can also access all the files on the whole cluster directly. I now want to show you a brief demo. And so, uh what we did is to

use this setup to basically give Claude a buggy training script. And we told Claude, "Hey, make this go on a small cluster." In this case, I think it's four nodes of eight GPUs each, and just make it work. Um what you're seeing here is basically uh you're seeing live execution, but Claude is kind of out of the loop. We sort of captured in a transcript everything that

Claude did uh just to eliminate one variable for this presentation. Uh but first, what Claude figured out is that actually there's a missing dependency. There's TQDM. And again, this is actually live running. So, this deployed onto a cluster. Um and uh when when I first tried to run my my my training script, we discovered that TQDM was was broken uh or was not there. What Claude now

did was just to my local UV virtual environment, add TQDM. Um and now, I guess we're waiting for UV to kind of complete. But the idea now is that um uh Claude simply added this as a dependency, and then now it's going to try to rerun the script. And the kind of magic of this is that it doesn't have to do anything else because I'm using the

same virtual environment uh on on all the other nodes. And by using this RDMA mounted fuse file system or fuse mounted RDMA file system, uh those changes were just immediately propagated to all the remote nodes. Now, the second thing that we discovered is that I'm getting nan losses. And so, something is going on. Uh what Claude did in practice here was actually insert a few printf statements

and discovered a bug in the in the CUDA kernel that we were using. It fixed that kernel. Uh is now recompiling it uh and doing the same thing again. uh you know, keep in mind that this is a transcript that was extracted from Claude. Uh this is completely hands-free from the operator perspective. Claude did all the work here. Uh and was able to again, act as if

the compute was local, uh but then uh uh you know, distributed over uh the the the the cluster. In this case, a fairly small one, but in principle it could be arbitrarily large. Um and then um again redeploy the script. And so now finally we're getting losses, looks good. So now that it's got a single node training to work, it's going to try to do distributed training.

And so now we're using torch distributed run instead of a single training script. And we're running on all of the all of the nodes in the and and the job. And I think what's going to happen now is that it's discovering a uh deadlock. And so uh it discovered that there's a deadlock. It used a monarch service uh that is again just a library. It's an actor

that can invoke uh PySpy and discovered that there's a uh deadlock and discovered exactly where it is. It fixed it. Uh it was just an off-by-one error and finally got distributed training to run. And and that's it. So here you can see uh I'm I'm basically out of time. I've got 20 20 seconds left, but here you can see how powerful these kinds of APIs are to

sort of agentic computing if you will. Um and we've had a lot of great success actually just giving Claude a pro- a problem and and kind of running with it. Um and and finally I'm going to kind of leave one more kind of presentation as an as a as an outro here, which is uh this is sort of uh taking this to uh you know the kind

of maximal degree. We gave Claude just a monarch tool that I that I showed you this sort of serverless SPMD setup and the prompt to train a GPT-2 model using openly available data sets. And it was able to do uh to basically create a uh working um uh application in about 20 minutes time. So this is obviously highly sped up, but this is a a transcript that

we captured again from Claude completely hands-off from uh from the operator's perspective. And within about 20 minutes time it was able to uh get the basic kind of GPT training setup working in a distributed fashion. I think this is over eight Uh again in principle it scale arbitrarily. But again, I think it's another neat demonstration of the power you get when you combine all these kinds of

building blocks together and give them to agents and let them go. And so with that, I'm I'm out of time. It's a pleasure to talk to you all and please come and ask questions if you have any after.