torchrun reimplemented with Monarch

[ezyang]

@LucasLLC and I discussed this during PTC. A good way for people to dip their toes into Monarch is for us to provide a reimplementation of torchrun on top of Monarch instead of the way it is currently implemented.

[pbontrager]

I think that the downside with torchrun is that it's external to your code, along with the flags that go with it and where monarch can shine is by bringing that into your code but in a way that behaves exactly like torchrun. Something like

from monarch.utils import spmd_launcher

def script(...):
   ...

def __name__ == "__main__":
	launcher = spmd_launcher(...) # same parameter flags as torchrun
	launcher.run(job)

This provides some value add over torchrun and also enables quickly converting existing jupyter notebooks to distributed jobs. It also brings monarch into the code making it easier to start picking up other monarch features over time.

This util would be very simple, but it might be valuable to not hide as much of monarch (job, hosts, meshes) if you want to make it easy for people to slowly step down the abstraction ladder.

[dongreenberg]

I agree this is a good way for users to dip their toes in, but as a "hello world" for Monarch more than a separate utility. It should be basically trivial to take the entrypoint function train of a training script, and instead of calling it once with argparse (wrapped on the outside by an SPMD launcher), it creates an actor mesh and calls trainers.train(cfg). That way you also get the fault tolerance benefit of being able to catch exceptions raised in train. This is not Monarch code but see what I mean here: https://github.com/run-house/kubetorch-examples/blob/main/pytorch_ddp/pytorch_ddp.py (Our Compute.distribute(workers=n is the same as creating a Job+HostMesh+ProcMesh). The behavior should be identical to their existing training but gets the wheels turning about the programming model.

[colin2328]

#1750

cc @ezyang @dongreenberg

[ezyang]

Great! All we need now is a non-torch-ABI dependent monarch pip package so it's easy to install everywhere 🤣