Lightning Talk: Step-Aligned Telemetry for Distributed PyTorch Training (Time... Abhinav Srivastav

This talk discusses step aligned telemetrics for PyTorch distributed training, focusing on how to identify bottlenecks that traditional system-level metrics often miss. The speaker explains that while typical metrics may show high GPU utilization and memory usage, the real issue can lie in unseen training failures due to time-sampled data. By using step aligned telemetrics, they aim to analyze metrics upon the completion of each training step, comparing performance across different ranks. They share two demos using the TraceML library, highlighting how it captures issues such as memory creep and rank stragglers in distributed settings. The speaker emphasizes the importance of understanding phase-attributed metrics to improve efficiency in training workflows.

My name is Abhinav. And I'm going to talk about step aligned telemetrics for PyTorch distributed training. All right, it's one of the work that we have been doing in trying to understand how to catch or detect bottlenecks in Um so we all know this problem where we have a healthy infrastructure dashboard, right? But your training is slowing down. You have a high GPU utilization. You have high

memory, but your throughput is going down, right? Now, the reason is that system level metrics that you have often hide training failures, right? And the primary reason for this is the system metrics are actually sampled by time. Whereas training, it processes step-by-step or in an iterative way, right? In an epoch way. So going more into it, uh what we are looking at is in a step aligned

telemetrics. What you get right now is a time sampled telemetrics where you're sampling telemetrics across different phases of training. So that may be data loading, forward, backward, and also across ranks. So you're averaging across a lot of places which kind of hides the noise. In the case of step aligned telemetrics, what we are looking at is at one completed step at a time. So we do the

phase attribution. So we associate what metrics we collect back to the phases coming from data loading or from forward or backward. We're also comparing in distributed settings across worst and median rank. So you get the skew. You understand how the training is differing across ranks. And this helps us to um preserve the step jitters that you get, right? Okay. So what are the failure modes that standard

um dashboards often miss, right? So one which is foremost known as a rank straggler where you don't get to see very easily or at least find out that one of the rank is slowing down the whole job, Another is data loading stalling. You can use profilers, but that's post hoc, Um you also don't get to see the memory creeping up slowly in your run, which can lead

to at some point um training crash. we are proposing is essentially looking at step aligned, which means that we look at metrics once a particular step or an iteration has completed, right? And then we compare how these metrics look for median rank. Median means how the performance along the median one and compare it to the worst rank. So we get the idea that what is actually how

the distribution of the training phases are. We try to tie the time and the memory back to the training phases. So we understand how much timing that we are spending in memory and memory in different training phases. So I'm going to show you two demos. Um quick one how to use the library. And the library is called TraceML. >> [snorts] >> So the first demo is about

um a simple settings of memory creep. What we have done is we're taking a distill fine-tuning on a single uh GPU. Purposefully we're injecting uh issue there. So we're retaining the loss and logits. And this will cause the graph to save, actually. And the expected behavior is that we will see the rise in the memory as the steps are going forward. Um so this is the prerecorded

video. So the line 16, what you see is the API from the TraceML trace step, and which is basically a context manager. you call this, everything else remains the training as it is. You go down to your training loop. And line 177, yeah. You have to wrap the training loop across this context manager. That's the only change that is required for TraceML to work. And now uh

below we are going to see the Yeah, so this is the injected issue. Now we are purposefully retrain retaining these loss and logits vectors. now we are going to run the training loop using TraceML to see how that looks like. So TraceML comes with its own launcher. It captures the torch run, essentially, and launches bunch of other stuff to actually aggregate the information. so it actually is

going to launch an aggregator, a single process on your rank zero to collect the information. In this case, there's only one. And then it's going to show you on terminal how the training is progressing. So you get to see system metrics, uh process metrics. But um you're also going to see So it's basically downloading the uh model, but you're going to also see how the steps are

progressing uh with respect to memory and with respect to time. So you'll get all the split with just a single context And since we have injected the issue, you are going to see the memory creep. So in this case, yeah, you can see that on the right Yeah, on the right-hand side, there is an issue of memory creep happening, and it's giving you what it is actually

doing there. So it catches the uh memory that is being used inside your loop. It also has um step time, and we're going to talk about in the next uh demo about where what does it actually useful for. Oh. So I'm going to cancel the training. Oh, I think I went ahead. Uh Yeah. This particular one, I think yeah, it's going to go back. It'll take a

time. Okay, I'm going to cancel this one and going to go to the next demo. So this one is about spotting a DDR rank straggler. So it's the same training that we had before. Except one node, we are going to run it or we are still running on one node, but we have a four GPUs. So we are going to launch four different processes. We are going

to inject the issue. So in one of the ranks, we are going to purposefully slow slow down the data loading. So we are just putting a time sleep there. And then we are going to see the worst rank and the median rank and how the latency diverges immediately, So again, the same steps to follow here. We have trace step, which you have to loop your context manager.

Um this is the slow rank collator. So what it is actually going to do is in line 77, you can see um there's a time sleep loop, and on rank zero, this is going to inject sleep This is going to cause the data loading issue on rank one only, but on other ranks, it's going to be normal. and as [snorts] I said before, you have to use

this context manager. So in line 217, we are putting that context manager across the loop. And we're going to run it again. Since we are using uh torch run, we are wrapping over torch run. You have to give the same commands. So here we have to mention the number of processes that we are running. In this case, um yeah, using the end proc per node. It's again

going to launch the same um which is going to collect this information from different um ranks and then show you on the terminal what is happening. It should be faster than the previous So you have the same um thing, and it figures out that there is an issue in the data loading. You can see that the median data loading and the worst one is far apart. And

it correctly identifies the rank where it is this happening. So you can see that it identifies there's an issue in the um there's an issue in the data loading. And this is while running your code in the terminal. Now I am canceling it, and I'm going to remove that part. So basically comment out and rerun the tr- rerun the code with the TraceML just to see that

if um this works in the other way around. So now it's not going to show you the data loading issue anymore because you can see the worst and the median time across the different ranks are actually same or similar in nature. Though it tells you that it's compute bound because your backward is taking a lot of time. But that's the only information it tells you. Okay. So

how does it work? As I told you, for PyTorch, you have to use this TraceML uh trace step context manager. But for hugging face trainer or pilot PyTorch lightning, we have um essentially extended the trainer to have TraceML trainer. And in PyTorch lightning, you have a callback that you can use, and you can get the same view. Uh that is all. Uh any questions if you may.

Any questions? >> [applause]