Lightning Talk: Debugging the Undebuggable: Introducing Torch.distributed.debug - Tristan Rice

This talk presents improvements to debugging in PyTorch distributed, emphasizing the new 'torch distributed debug' module introduced by the speaker, Tristan, who works on PyTorch distributed at Meta. The current debugging process is explained as error-prone and lacking standardization, leading to challenges in identifying issues such as timeouts and hangs in distributed jobs. The new debug server is designed to be easily integrated into any trainer and runs automatically, offering tools for aggregating results and analyzing data across multiple ranks. Key features include stack traces for both C++ and Python, an overview of ongoing collective operations, and integration with the PyTorch profiler to pinpoint performance issues. The speaker also discusses the lightweight nature of the debugger, highlighting its design to minimize overhead during training, and invites contributions to enhance its functionality.

Hi folks. I'm Tristan. I'm I work on PyTorch distributed at Meta. I'm going to be talking about kind of a small feature, but something I'm pretty excited about improving debuggability in PyTorch distributed. So, just kind of to start off, I'm going to motivate the problem a little bit, though I'm not sure it needs to be motivated given how hard debugging some of these distributed jobs is. But,

generally right now debugging a lot of this is pretty kind of error-prone. There's not a whole lot of standard tooling. Everyone kind of has their own custom paths for figuring out like, you know, your job timed out or you have a hang or you have some like network error. And, you know, there's a lot of these individual tools like things like logs, things like PySpy, debuggers, things

like that. But, they're all kind of fragmented. And then, the kind of other biggest part of this is even when you get the error messages or have something like that, it's really hard to figure out like what actually went wrong. Like, you know, you might have a bug in your code that leads to a hang somewhere due to like say a data loader issue, and you know,

you'll get a timeout, and that's pretty unhelpful in most cases. Um, you know, cuz you can't actually figure out what we're going into. we basically decided to try and tackle this problem in a pretty lightweight and sort of simple solution. The key idea behind this is like how about we bake a lot of this sort of pretty bare-bones standard tooling into PyTorch distributed itself. So, we have

this new module here called torch distributed debug with this debug server. And the kind of key idea is you can just put this in any trainer, basically in, you know, wherever, and it will just start up this debug server that at least gives you something. So, this here, you run this one line, this will start the debug server running on uh rank zero worker, um and then

it also spin up all of the like worker server and background infrastructure um just sort of automatic. this right here is sort of the high-level overview of this. Um it's all done using pretty standard HTTP and Python. we will run this interactive server that's designed to basically host all of the analyzers and all the code to actually basically aggregate all the results. And then we have this

simple like debug worker server that runs within every single one of your um PyTorch workers across all the different ranks. And then we can automatically do discovery for this, and the front-end server can basically just go in and fetch all that data so you can access it across your entire job. Um we currently have a bunch of different analyzers. Um I'll show the you to these to

you in just a little bit. Uh it's a lot of the common stuff you might want. So things like being able to get stack traces both for C++ and Python, being able to see like which collectives for PyTorch distributed are in flight, being able to access the PyTorch profiler, um and some other lower-level stuff that people probably didn't even know existed, such as like weight counters. So,

um I'm just going to go over and show you guys this using a little bit of a demo here. it's probably easier just for you guys to see how this actually works. Um So yeah, so this right here is the PyTorch distributed debug server. It doesn't look like too much, but it has a lot of the the UI um and all the actual data there, which is

what matters. Uh so right here, um you know, first thing here, like this is a small training job that I'm just running on uh one of my dev machines. Um but basically, this is a very simple stack trace handler. So this calls into Python on every single one of the ranks, um and will basically just give you all the Python stack traces. So if you're seeing something

stuck, you can go in and just very quickly scan through all your ranks and see like, is there something weird going on? Is there something that's like mismatched or things like that? Um I will say that this visualization is very, very basic and new. And this is actually something that we're interested in getting contributions for if people are interested in adding fancier like stack trace diffing and

things like Um but yeah, so we've got that. Um we have PySpy as well here. Um you can be able to get actual like native stack traces when you have issues. and then uh something that some people might be familiar with, um but maybe isn't super accessible until now, um has been flight recorder. Um so, this is a feature within PyTorch distributed um that basically tracks all

of the in-flight operations using a ring buffer across all of the different uh workers you have. So, what this does is it basically goes and fetches it across all of them. So, in this case, we have a world size four, um and we can see which collectives are actually in flight. So, right now, this is running a nickel all reduce, um and we basically see under missing

ranks that like rank zero hasn't actually joined the job here. Um so, this tells you at least which machine you're likely having some issue on for something like a hang. Um and then we can basically go in and actually look at some of the other stuff here. So, in this case, this isn't a hang, but we're seeing rank zero being consistently slow. So, we can basically fire

go to the uh PyTorch profiler. Um and this just pulled a 1-second profile across all of the different uh ranks here. So, we can go in and, you know, take a look at say rank one here. Um and this just opens it up in Perfetto. and we can basically take a look and see where we're spending time. So, in this case, on the CPU, we've got the

CUDA device synchronize. Um if we look at the actual GPU traces here, we can basically see this worker is waiting in nickel all reduce. So, this is something you know, we're waiting on some other worker. Um, we're not quite sure which. Um, but using that flight recorder, we know the issue is on rank zero. So you can basically fire this up. Um, go in and take a

look and we can basically see that I put in a dummy sleep on this case. So you can actually just kind of pinpoint it here. Um, so yeah, that's kind of the basics there. Um, I'm going to go over and talk a little bit about how this works under the hood um, and some of the details here. Um, so yeah, I'm going to just going to plug

flight recorder again cuz this is something that has been quite helpful for debugging a lot of these issues. when you turn it on, it's basically this always on ring buffer that you can either call via the um, like debug server, but you can also just go in and actually call this from your own code. So if you have your own custom tooling, but it basically will dump

this ring buffer and you can use the flight recorder um, FRTrace analyzer um, to go in and actually pull out those insights on like which ranks are having issues and things like that. Um, anyway, if you're curious on how to enable it, there's a little bit different depending on whether you're using torch comms or process group, but it's pretty straightforward. Um, so yeah, so like the server

is pretty bare bones, um, but we've basically tested this up to about 10,000 um, ranks. Um, this all very lightweight HTTP both on the actual handler side as well as on the aggregation side. So like it's designed to be very minimal low overhead. Um, we do two main things. One is the actual analyzer that's fetching all of it is running as a sub process. So it won't

have any impact to your training loop. And then the actual handlers that are running inside each of the ranks basically run as like a C++ background thread. So you don't have any Python GIL contention. Um, so there shouldn't be any major concern there. Now, the analyzers themselves, if you're writing something custom, might have some more overhead if you're turning on the um but the kind of key

idea is when you're running this in interactive mode at least, like there's zero overhead until you actually use it and, you know, you're actually having errors or issues Um we also have sort of a new feature that we're testing out. We're still kind of experimenting some with this, um but we have a feature for doing basically uh periodic dumps. So, for like flight recorder, well, like you

can set it up, so, you know, if you have a 10-minute timeout for say nickel, you can have flight recorder fire every 5 minutes and basically, you know, if you have an error, you can go and be like, "Oh, well, this was the rank that was stuck the last time." Um so, yeah. We're also looking at doing dumping on errors, so if you can actually go in

and configure like, "Oh, if you hit some error, we can go in and actually have the server fetch everything right before it exits." Um and then um yeah, so in terms of like extensibility, um all of this is designed using a pretty generic handler system. Right now, all of the handlers are built into PyTorch, so if you're interested in contributing, you can we're happy to take, you

know, contributions for any new data or things like that, um but you can also like implement these and register these handlers yourself. It's using a little bit of a lower-level APIs, but that's sort of the intention there. Um and basically what this handler does, um is it just registers a very simple HTTP endpoint. So, you can basically dump any data you want. So, in this case, this

is using um torch profile and we're just writing it out as JSON. So, if you want to actually just use the underlying debug servers directly from your own code, you don't actually have to use our analyzer, you can just turn this on and hit these handlers directly. Um and then this right here is an example of how you can actually write a um basically an analyzer. So,

it's also Basically, there's one call you can and it will basically fan out the requests to every single one of the ranks. And then you can write whatever analysis you want to be able to, you know, put it into the actual sort of pretty HTML or things like Perfetto or so on. So. That's the debug server. Um hope you guys try it out. Happy to chat about

it. It's still quite new, so uh yeah. Eager to get some feedback. Thank you. Woo!