Lightning Talk: Scaling Recommendation Systems To 2K GPUs and Beyond - Zain Huda, Meta

This talk covers the scaling of recommendation systems using PyTorch, specifically focusing on techniques to extend these systems to over 1,000 GPUs. The speaker explains the importance of efficiently scaling recommendation models, drawing parallels with large language models (LLMs) that benefit from higher parameter counts and GPU usage. They highlight the challenges faced in scaling, such as imbalances in GPU loads, communication overhead from embedding lookups, and increased memory requirements. By introducing a data parallelism layer on top of existing model parallel approaches, they successfully address these challenges and achieve significant performance improvements, including reduced lookup latencies and memory consumption. The speaker also shares details on optimizing embedding weight synchronization in distributed settings, contributing to enhanced training efficiency. Overall, this session presents a comprehensive solution to effectively scaling recommendation systems in large multi-GPU environments.

Hi everyone. My name is Zen. Today I'll be talking about some of the work we've done around scaling recommendation systems to 1000 GPU plus scale. And yeah, I work on PyTorch. And yeah, thank you for coming. This is a lot more people than I expected. So, excited to talk to you guys about our work. So, you might want to wonder like why do we want to scale

Rex models? Like in LLMs, we know of like these scaling laws and we know that like when you scale LLMs to like you know, higher parameter counts or a larger number of GPUs, you get like better performance. And you get the same thing for Rex models as well. And so, we can notice that like for LLMs, like we kind of have this like very good understanding of

like what dense parallel and scaling looks like. For example, on the right hand side of like this is like a very, you know, intense example I would say of like a six-dimensional mesh parallelism. Like you can think of like your context, expert parallelism, tensor, data parallel, like pipeline parallel, all these things kind of stacked on top of each other to enable like, you know, these 10,000 plus

GPU scale runs. On Rex, we're nowhere near close. Like before this work, we were training at around 500 GPU scale. And then with this work, we're able to scale to like kind of 4,000 scale. Um, and so, in for dense modeling, like we have pretty good kind of infra already built up like DDP, FSDP, HSDP, etc. But for Rex models on the sparse side, we don't really

have that and I'll show you kind of why that's the case. And so, you know, like our infra before was only really able to scale to like 500 GPUs by like naive model parallelism. And so, what does like, you know, what what how do these recommendation systems like differ from LLMs? If you're kind of not familiar with how recommendation systems kind of like how that model looks

like. Basically, you can kind of like on the right hand side of this like DLRM, which is stands for like deep learning recommendation model. And this is known as like a two tower kind of model. So, on the left hand side, you have like your bottom MLP. On the right-hand side you have these things called like embeddings, embedding tables. These basically represent things, you know, we different

reels or different movies, user interactions, user IDs, and you just have these are very, very massive in nature. On some of the largest models this can go up to like, you know, 1 TB in size across, you know, hundreds of embedding tables for a given model. And then basically you do like a dot product of these, you know, embeddings plus like whatever the bottom MLP will give

you. And then there's a top MLP which will eventually give you some sort of prediction for whatever your model for whatever the model is being used for. And so, on the right-hand side is what we're really going to be focusing on like are these embedding tables. This is like the sparse kind of component that we'll talk about. And these recommendation models are what comprise of like the

majority of the parameters. And so, when we're thinking about, you like recommendation systems in terms of like the parameter count, embedding tables are, you know, you can almost think of like north of like 80 to 90% of that comprise of that model. And so, these embedding tables are known for like heavy memory consumption. And like these tables we shard across GPUs using model parallelism. And each training

step only uses like a subset of these embeddings, so you're only looking up some portion of these embeddings for a given forward pass. And so, that's what, you know, where the sparsity kind of comes in. And so, the challenge of scaling these we can kind of bring it into like three different parts. So, on the right-hand side before I go into kind of what these challenges are

is kind of what you'll see today or, you know, what model parallel like what or what parallelism strategies for recommendation systems looks like is that for everything besides the embedding tables we'll just do like simple data parallel because the dense parameters aren't that like aren't there isn't that many of them and so, you can do just simple data DDP on them. You you even need FSDP actually,

just DDP enough. DDP is enough. Uh, and then you do model parallel for these embedding tables. Uh, and I'll show you kind of what that looks like a little bit later, but uh, you can think of something like around just taking an embedding table and like column-wise or row-wise sharding it across GPUs. Um, and so the three main challenges we have of scaling just embedding tables through

model parallel to 1,000 plus scale is one, like we have this imbalance uh, and straggler issue. Uh, so because you have um, like s- sparse uh, lookups in nature, um, what happens is like certain GPUs will have more lookups than others, and so certain GPUs are doing more lookup time than others, and so they're waiting for each other. The second one is like these lookup communications. So,

as you're scaling these up, you know, for a fixed batch size per GPU, you now have even more lookups as you're scaling, you know, from 500 to 1,000 to 2,000. Uh, and so your embedding computation time also takes longer. And then also lastly is your embedding activation memory. So, it's kind of in hand in hand with two. Um, what happens is that like because you're now looking

up more lookups, you have to store more on a given GPU, and so you're also bounded by how much memory you have on that given GPU. Uh, and this is kind of example of like the latency and memory bottlenecks we'll see. And you can start And you see like the embedding lookup time increase significantly as you go from like, you know, 256 to 1,000 GPUs. Uh, and

then same thing for, you know, we can like look at our embedding computation, which also goes up. So, our solution is effectively just adding like a data parallelism layer on top of our model parallel. And this helps us go beyond like this 1,000 kind of GPU barrier. Uh, and the reason we had this 1,000 GPU barrier for traditional model parallelism is that communications like these all-to-alls around

on over 2,000 GPUs fail like very often, and so you can't actually get like stable training. a- a- and so that's why we added like this data parallel layer to basically get like this hierarchical sharding, where we know that we can model parallel shard on like 256 or 512, and we can kind of kind of that up with data parallel, too. Effectively, like 8,000 GPUs, which I'll

kind of show why we have this 8,000 GPU kind of um wall. Um this is the setup. I'll kind of go over this cuz I'm running a little bit of time, but but how does like this data parallel plus model parallel like together solve, you know, these challenges that I talked about earlier? Is the first one as well, with the imbalances and irregular issues, we can model

parallel across, let's say, 256 GPUs very well with get like balanced sharding. And then with data parallel, we can scale it up to 1,000. And so, we'll actually get like better balance and sharding than we would if we had just done model parallel across 1,000. For lookup communications, because now we have these smaller smaller model replicas, the lookup communication time is also now lower because for that

given model replica, you're doing a smaller uh batch size for that for that given model. Uh and then same thing with the activation memory, your global batch size is now just of that 256 GPUs instead of 1,000, and that helps you actually reduce the activation memory on a given GPU, as well. And so, we basically through just a very simple kind of concept of just adding data

parallel, we're able to like solve these big three issues. And now we can see that like how does, you know, these uh what are the costs look like with 2D parallelism versus 90 model parallel? And we'll see a significant reduction in the combination of our latencies here. So, the computational lookup uh uh lookup latencies. And we notice and also I guess the other And the other thing

I want you to notice is also like there's no kind of like for for a given model, there's no like right solution of like you want to maximize the number of groups. Like we see on like the left-hand side, the CTR model, four groups is much better than eight. But on the right-hand side for XFM model, we see that like four and eight tend to be pretty

good. And so, that really depends on like what how your model and like the batch the batch size and what the dense computation looks like and as well as what the lookups look like, etc. So, there's a little bit of tuning that's required. The one thing I want to do notice is also the embedding table all reduce. As you know in DDP like you want to like

synchronize the gradients at every step to make sure you have equivalent training. We do the same thing here except we do the weights. The reason we do the weights is because we have an optimization in the backward where our backward kernel and optimizer are fused so we don't have to actually materialize the gradients for these embedding tables in memory so we we reduce our memory consumption significantly.

And so instead we have to do the weights. You might be like, "Well, how do we do uh embedding weights like with the embedding weights like how do you get equivalent training?" We do some learning rate stuff which I won't go into today but there's a paper we have that will that talks about it. Simply put if you just multiply the learning rate by the by the

number of groups you actually get equivalent training. It's a pretty neat trick we found uh and it works and it works pretty well at scale and that's something we use in production today. Uh and that's yeah and so we can look at this traces as well that kind of shows us like this QPS gain. This QPS gain is what basically talks about like or for us is

what training efficiency uh is is what we use for training efficiency. And we see the embedding look up decrease significantly. Same as well like with the embedding all on the right hand side and the parameter sync is like insignificant relatively speaking and so you can get something like 10 to 20% QPS gain on certain models which is significant especially on a thousand plus GPU scale. Um I

unfortunately don't have time to go through the rest of this but uh if you have any questions I'd love to talk to you about all this stuff. This is basically we add like an FSDP kind of like um approach to what we talked about earlier to reduce like the uh the the overhead the memory overhead in data Uh which works really well and the the solution we

have is pretty cool. We're able to basically go from like a peak of 130 gigabytes uh to like 110 gigabytes so we're saving like 20 gigabytes of like peak memory which is which is huge for some of these largest models. Um and yeah and uh thank you so much for listening. I want to shout out my my co-authors with this work. We have a PyTorch blog about

this. We also have a paper on archive if you if you want to learn more and happy to answer any questions. But yeah, thank you so much for listening. I appreciate