Lightning Talk: Why Logging Isn’t Enough: Making PyTorch Training Regressions Vi... Sahana Venkatesh

In this session, Sahana from Wave discusses the challenges of training models at scale, particularly focusing on the issues of monitoring and catching regressions. She defines regressions as changes that detrimental affect a system's performance, including increased training time and deteriorating model metrics, which can lead to wasted resources. Sahana elaborates on the importance of understanding both the model's performance and runtime signals to ensure a healthy training process. She introduces various strategies to detect regressions, such as using trajectory checks for monitoring the loss curves and implementing runtime budgets to catch training stalls. Throughout the talk, she emphasizes the necessity of comparing current runs against historical data and baselines to identify performance issues. Sahana concludes by highlighting the ongoing development of automated tools for regression checks, which can support future training processes.

Hi everyone. Thank you so much for coming. I'm Sahana. I work at training tech team at Wave. At Wave we are training driving models at scale. I'm going to be speaking about a topic that's close to me. I'm going to be exploring the the different things that trips the teams that are training models at scale. Basically gap between having logs and catching regressions and this is my

topic of the talk. Why logging is not enough and making PyTorch regressions visible in practice. So the agenda of the talk is going to be as follows. We're going to be exploring what healthy training looks like. Why regressions are expensive, how they manifest and the patterns we follow at Wave to catch them. Before moving on forwards, let's explore what regressions are. Regressions are basically any change that

makes a system worse than a previously known good state. And this could be things like training slowing down, model metric getting worse and both of these could make the GPU hours like more expensive and developer time going to waste. Some of the regressions are not visible and might be missed if you are not actively monitoring them and this make them harmful. So let's basically explore like what

a happy training state looks like. It's a two-story thing. You have the model story, you also have the run time signal story. The model story is something we are familiar with. Basically the training loss goes down, the validation metric looks aligned and repeated runs don't have much difference between them or the differences are like acceptable. The run time signals on the other hand could be something like

the step time, the time it time it takes for a step update. Throughput, which could be either iterations per second or output samples tokens per second are all within range. Memory stays within bounds, there are no training stalls, and the failures are caught early before they become a problem. So, a good healthy workflow should have both the stories both both the stories in sync and both the

stories should have should be within bounds. So, now we explore how a healthy Let's actually look into how these regressions come up in practice. They there could be like train large train val divergence. This is when the training loss is going down, but the validation loss slightly diverges. There could be delayed breakdown. Here is when the training loss looks fine for like in the early epoxies, but

then in the later one it starts diverging. There could be training stalls when the basically nothing's going on for a while, but the GPU hours are still wasted. And finally, there could be quiet performance regressions. These are ones when um model metrics looks the fine, but the data loader is suddenly consuming more memory or the training system itself is somehow consuming more memory, data loading time takes

a lot more than expected. These are all like quiet performance regressions, which we would miss if we are not monitoring them. So, now just let's think about couple of ways we can catch these regressions we discussed in the earlier slide. For the train loss train val divergence, as well as the delayed breakdown where the training loss suddenly diverges, Uh a fix could be traject- trajectory checks. We

don't just look at the final metric, but we kind of study how the curve is changing and have alerts on the shape of the curve. Uh basically, these alerts could be is it diverging? Uh is it overfitting? And uh these changes could basically be detectable only if we are looking for Uh for the training stalls, we could have something like runtime budgets. Uh basically, measuring the step

time or the or the glo- or basically, somehow recording the global step. The key insight here is if the training stalls or training slows down, the model quality might be the same, but the training has got little more expensive. And finally, these are these quiet regression performance regressions are are quite interesting. These are like questions like, is my train training 10% slower? We won't be able to

catch them if you're just comparing your run against itself. You need to compare the run against a reference run. And in some cases, these are known as like baseline comparison. Every run needs to be compared against a known reference. There are some exceptions, but uh this is a good rule to have. So, uh we basically spoke about how regressions manifest. We also spoke about like some uh

some ways we can catch them, but let's bring the concepts together more in practice. Uh you can think about regre- uh regression checks through two lenses. You have uh you can see them through timing lens. You can also see them through fidelity or cost. I'm going to start with the right side first. Uh a software engineer could find these terms uh familiar like unit test testing the

functionality, testing the logic, integration test, slightly more expensive than unit test, but both of these are basically run before a change is actually merged. For the ML systems, which are more complicated than the software systems, you have something called federated These are basically try to capture the real world behavior. You You run the entire end-to-end training run in a system that's close to production, and here you're

monitoring for training as well as like runtime regressions. In the timing lens, there are two kind of metrics you can think about measuring. There's online metrics and there's offline metrics. Online metrics are those metrics that that are measured during the training. This could be step time, heartbeat, which is the current global step, etc., that can be logged during training. Offline metrics could be could be those metrics

that can be gathered at the end of the training. For example, end of the run summaries, comparing against baseline comparisons, or like comparing against historical trends, etc. I'm going to be speaking more about the rest of the slides. First, I'm going to try to present an architecture before diving in further with concrete examples. This is basically a user workflow. Let me know if this is clear. The

user user basically pushes their get commit. A training image is created, and then it's pushed to the cluster, and the training starts. You have You basically log the usual training metrics such as loss, etc. You also monitor the training heartbeat. Here, this could be the global step. You monitor You monitor flops. this is basically the first line of defense and these metrics are gathered during training. The

second line of defense here is like as soon as the training is finished, you basically check for throughput, peak memory, final loss, etc. And you log these metrics to a database. Here you make a comparison to the to historical trends and then flag for a regression if you see something off bounds. Let me show you some concrete examples that we have at Wave. This is one of

the Heartbeat base stall detection. Here training starts and during on fit start using on fit start call back, you basically start a background thread and you start logging the current global step to a database. There's a third-party actor that basically constantly looks at the value coming in from the training job. And if the actor sees this this training job has not done any any new additions in

a prefixed amount of time like an hour for example, then the actor cancels the job. And you can make this entire process transparent to the user by basically sending having a Slack workflow and sending this information to the user like the image attached basically saying that this training job has been canceled because it did not make any progress in in an hour or something. And this is

helpful because this basically surfaces when the training job has started stalling and they can look into the logs better with this view now to understand that why it's stalling. And the second concrete example I want highlight here is like making use of historical trends for regression detection. Here we run that a trusted baseline at a schedule and compare it against historical data, 30 days average. When a

metric is an outlier or not computed anymore or is invalid, then you basically raise as a regression. In this example, you can see that the baseline job B, the training time is is out of out of bounds in sense that it's greater than the standard deviation we accounted for, then we we basically flag it as an outlier. In the example A, you can see that the job

used to compute MFU, but it's not computing anymore, something has happened, then you flag it as an outlier again. The basic takeaway here is a regularly regularly scheduled a baseline can provide like quick quick subtle hints whether something has gone wrong with your code base. And this can help you to like highlight high cost regressions. And when these signals can be then tied back to get history

and we can like narrow down to that commit that started causing the regression and reverted. This this has been like quite useful in addressing some incidents we had in the past. Finally, answering why does it all make sense right now? Well, we have been seeing like some blog posts about agents like writing the training code, running training, and even monitoring the training for us. These tools that

we are developing for regression checks could be the tools that agents can use to monitor training. So, these automated regression checks right now can become can help help to have automated reliable training in future. And finally, for the takeaway, if there is three things that I would love for you guys to take away is like always have some sort of a baseline to compare your your most

important runs. Check for invariants and also surface not just the model, but also system as runtime regressions early. And finally, this is such a cool problem and we are still solving them. Not just within the company, but within the industry. And we are a cool team and we are also hiring. These are some of the opportunities I have within our own team. I would love for you

guys to apply. Thank you so much.