Parallelism 2

From inside one GPU to across many

00:05 This is the second week of the systems course, where we are trying to get the most out of the available equipment to speed up model training. Last week we talked about parallelism within a single GPU, and this week we're talking about parallelism between multiple GPUs. So, here is the picture you should imagine. We have a whole series of nodes. These are essentially computers, each with multiple GPUs, usually eight, and each GPU has a set of stream multiprocessors, or SMs, that actually do the work, and the ones marked in green are memory and connectivity. Each SM has a very small L1 cache, the GPU has high-speed HBM memory, which is larger, and there are also these communication channels connecting the different GPUs. So the way to think about it is this: the calculations should happen inside the SM in these ALUs, right? Computations require input data and, a record of output data. Typically, the input and output data may be relatively far apart. If you're lucky, they're in the L1 cache; if you are, let's say, less fortunate, they are in HBM. And now, this week, we're talking about training on many GPUs and many nodes, where the data you might need might be on another GPU, right? So the main challenge is how to structure all the calculations to avoid bottlenecks in

01:39 data transfer, because we want, remember, to maintain high arithmetic intensity. We want to load our GPUs, make them work at full capacity, and data transfer will usually be much slower. So this will be our bottleneck. Last week we looked at a number of different methods for doing this within a single GPU, including fusion and tiling. The idea is that instead of reading and writing from HBM, you can load data into the L1 cache or, I guess, shared memory that's the same type and, you know, the same speed, and just work there in your local workspace, and then write to HBM only when needed. This week we started looking at the connection between GPUs and nodes, where we have to replicate and shard our model parameters and optimizer states. And it is how we do it that will determine the costs. So, I'll take some

The hierarchy: L1 to HBM to NVLink to NVSwitch

02:42 liberties to organize everything into a hierarchy, from small and fast to large and slow. So, the smallest and fastest is a configuration with one node and one GPU. You have the L1 cache, which is extremely fast but very small, then there is HBM on a single GPU, then between GPUs on the same node we have NVLink, and finally we have NVSwitch; and, of course, all of this in the Nvidia ecosystem. So the idea is that many of the basic concepts of minimizing data transfer remain the same, but the mechanics are now a bit different because L1 behaves differently than these NVSwitches. So this lecture will be mostly dedicated to fleshing out the concepts from the previous lecture in code. There will be a few new points, but Tatsu did a great job of reviewing all the different types of parallelism. I'll try to pin this down in code so we can get a deeper understanding of what's going on. And then I'll be accessing this standard output file, which is the result of running this lecture. There were some minor issues that I won't talk about, where this framework didn't work quite correctly when multiprocessing was present. Okay, so this lecture is in two parts. In the first part, we will look at the building blocks, the collective operations that we discussed last time, how

Collective operations, and the words for them

04:21 they are implemented in NCCL and PyTorch, and then we will do some performance testing. And in the second part, we will look at distributed training itself, data parallelism, tensor and pipeline parallelism. Okay, let's start with collective operations. So, collective operations are primitives that are commonly used for distributed programming, and "collective" means that you have many nodes. They are actually quite old, at least since the 80s in the parallel programming literature. And in general, they provide a better abstraction than trying to manage point-to-point communication yourself. So these are truly time-tested primitives. A little terminology. World size essentially refers to the number of devices, for example four, and rank, which can be confusing if you're used to linear algebra, actually just refers to the device. So we have rank 0, rank 1, rank 2, and rank 3 if you have four devices. OK. So, collective operations look like this. So, let's start with broadcast: the idea is that you have t0 on one of the ranks, and you just want to broadcast it to all the other ranks or to all ranks in general. Okay, it's very simple. Scatter is similar, but you have four values and you want to place each of the values at different

05:55 ranks. So each rank gets different values, not the same value. Gather is kind of the opposite of scatter, where each rank has a different value, and then you gather them all together into one rank. Reduce is the same as gather, except instead of combining, you add them. All-gather—This is the same as gather, but you do it for all destinations. Gather was only for rank zero, or say rank one or two, or any individual rank. All-gather is when you do it for all of them. And finally, reduce-scatter. I couldn't find a good image for this, so I'm reusing what I had last time. It's like reduce, where you take a bunch of different values and, you know, add them or do some other commutative operation on them, and place the result at one rank. But, as with scatter, you place different parts of the vector or tensor at different ranks. OK. And remember that all-reduce is equivalent to reduce plus all-gather. So, the way to remember this terminology is this. Since it can be confusing—what is all-gather and what is reduce-scatter—reduce simply means that you are performing some sort of associative and commutative operation, such as sum,

07:35 minimum, maximum, or average. Broadcast and scatter are the opposite of gather, and all means that all devices are destinations. Okay, I hope this was a review of what we learned last time. So are there any questions before I go any further, since we're going to use these primitives as a base? So it's helpful for everyone to understand that. OK. Now let's see how

The wire underneath: PCIe, Ethernet, NVLink, NVSwitch

08:16 this is actually implemented, starting with the hardware. OK. This is what GPU hardware looks like in a classic way. It's like a home computer: you have processors, and typically on the same node are GPUs that communicate over a PCIe bus. And if you need to establish communication between different nodes, then all this is connected via Ethernet. So, this is roughly how machines were usually built. If you're buying a GPU for gaming or something similar, your configuration probably looks like this. As we will see, this is quite suboptimal, as there is a lot of overhead when data needs to be transferred from one GPU to another. They have to go through the kernel, be copied into buffers, and then transmitted over Ethernet, which creates high latency. So in today's environment, with scientific computing and deep learning, if you're going to put a bunch of GPUs together to work together, we just connect them directly. So, in the Nvidia ecosystem, we have NVLink, which connects the GPU directly, bypassing the

09:51 CPU. You don't need to go through the kernel of the host machine. And between nodes we can connect GPUs directly via NVSwitch. This way we bypass Ethernet, which was developed a long time ago and was clearly not intended for such applications. So NVSwitch and NVLink simply skip all of that and optimize directly for the types of workloads we're interested in. If you look at the H100, each GPU has 18 fourth- generation NVLink links. This gives a total bandwidth of 900 gigabytes. In comparison, it is significantly faster than PCIe and much faster than Ethernet. If you consider the cost of moving from SM to high- bandwidth memory reading, it's still about four times faster. Of course, these numbers are constantly changing with new Blackwells. It seems that this number has increased by two or three times. OK. Yes. Regarding PCI, does it go to the CPU and then to another GPU, or directly to the GPU? So, the question regarding PCIe is,

11:32 how exactly is data transferred? I think they still have to go through the CPU. Was there another question? And PCIe, I mean, it's designed for things, and other devices connect to it, like a sound card or an SSD. So it's not really, it's kind of a general purpose bus for communication between devices. Yes. NVLink also has a connection to the CPU. Yes. So the question is whether NVLink also connects to the CPU. We'll see a little later, I think maybe on a slide, how everything is connected. Yes. So you still need to communicate with your CPU, of course. Yes. OK. So, there is a command that you can run, and it produces some output that allows you to see how the GPUs are actually connected. So I ran this on our cluster. There are eight GPUs. I don't think you can get eight GPUs, but if you could, it would look like this. And you see that there is an NVLink connection between each pair of GPUs. There are also network cards and other things. OK. Oh, yes. So network cards are essentially what

NCCL and torch.distributed

13:10 provides the PCIe connection and communication with the CPU. So, okay. This is about hardware. So how to use this equipment? Nvidia has spent a lot of time developing really good software for its apparently very good hardware. There is a collective communication library from Nvidia called NCCL. It essentially converts the collective operations we looked at earlier, such as all-reduce, into low-level packets that need to be sent between GPUs. So this library does a lot of the work because it allows the programmer to just operate at the level of "I need this tensor to appear on all machines," and it just happens. OK. So, you know, a little bit about what happens when you set up NCCL: you connect a bunch of devices, and there's some data exchange to determine the topology of the hardware; this optimizes the path between GPUs, and then when you actually call these collective exchange operations, CUDA kernels are launched to send and receive data. So, this is NCCL. It is provided as a library. But NCCL is still a bit too low-level for us, since most of what we do is, well, Python. Therefore, PyTorch has the torch.distributed library, which essentially provides a convenient

14:42 interface for these collective operations. Now, working in your PyTorch program, you can simply write "all-gather" for the tensor and it will appear at all different ranks. It also has this nice and useful feature that it supports multiple backends for different hardware. In particular, remember that NCCL was for GPUs, but you can also do collective operations. Remember, this is not GPU specific. This is simple for any set of devices. So you can do this for CPU too, using a backend called Gloo. So if you're debugging code on your laptop for an assignment, for example, you can use Gloo and still run your programs even without a GPU. So that's another advantage of having these high-level primitives—they're much more portable than having something purely GPU-specific. Of course, performance will actually depend on the hardware, but at least logically you can verify that your code works. PyTorch distributed also supports other high-level things like FSDP, which Tatsu talked about in the last lecture, but we won't use that in this course, because in the spirit of developing everything from scratch, that's what we'll do. Okay, let's look at some examples of how distributed collective operations work. Here is a helper

Four processes, one function: spawn, init, barrier

16:17 function I wrote, and you can take a look at it in the code if you want. It takes a function and simply runs it; essentially, it's a wrapper around Python multiprocessing that simply launches four processes performing this function. So when you 're in this function, you should think about it like this: there are actually world_size processes running this identical function, where the rank indices go from zero, one, and all the way down to world_size minus one. Okay, now I'm going through only one of the ranks step by step because the lectures aren't parallel. So, usually the first thing a process needs to do is initialize itself. And basically, they need to find each other somehow, right? Since you are running on a multiprocessor system, running many processes, they need to connect to the same host to learn about each other's existence. Note that this is not where all data is sent. Data is transmitted through NCCL, and this is only for some coordination. And since we have a GPU, we can use NCCL. Otherwise, you would use Gloo. Okay, after the setup, we'll move on to performing some actions. There is a useful function called barrier that essentially waits for all processes in your

17:52 group to reach this point. Of course? Remember that everything works asynchronously, and in some cases you just need to have a synchronization point, so barrier does just that. The reason I added it here is pretty trivial: I want all these output statements to be grouped together, but there are other reasons to use barrier, which we'll talk about later. So for each of these groups, I'll create a tensor. A tensor is 0, 1, 2, 3 plus,

The collectives on tiny tensors

18:24 you know, rank. I will output the value for each rank before all-reduce. What does it look like? Okay, here's what it looks like. Can the people in the back rows read this? Yes. Okay, great. OK. So, at rank zero it's 0, 1, 2, 3. At rank one it's 1, 2, 3, 4, and so on. Note that because this is asynchronous, the order of the output is random, depending on how it is printed. OK. So, each rank has its own tensor, and then all-reduce is performed. So, for all-reduce you pass this tensor. You say you want to summarize it. In this case, I won't do it asynchronously, but you can do it. Which is useful for overlapping communication and computing. And then, as expected after all-reduce, for the first component you add them up and get six, then 10, 14, and 18. Okay. So, after all-reduce, this tensor is actually overwritten by the corresponding sum. So it's very, very, you know, convenient and easy to use. OK. So, let's run reduce-scatter. For reduce-scatter, I will create

20:00 input data that has a dimension equal to the number of nodes, in this case four, and I will allocate memory for the output data, since reduce-scatter does not work "in place". It will just be a scalar. So, before the reduce-scatter operation, it looks like this. I have the input data as before, and the output is zero, but it could have been any value since I didn't initialize it. Then, after reduce-scatter, I pass in the input and output data, and I will perform the summation. Then I get, essentially, this: for the first component, I do the summation, and it goes to rank zero, for the second, it goes to rank one, and so on. Well, as you noticed, this performs the same operation as all-reduce, only the original data is spread across all the different ranks. Okay, now let's do an all-gather. So I'm just going to use the reduce-scatter output, which is this, as the input. And then I'll allocate an empty array for the output data. So, for the all-gather operation, the

21:35 input data is this, and the output data is, I think, just arbitrary values. And after I run all-gather, what happens is that all these tensors appear on all devices. Okay, this is also a kind of example. I hope you are now convinced that reduce-scatter plus all-gather is all-reduce, because I calculated the same value as for all-reduce. Okay, the question is, is that clear? Yes. Do we track in reduce-scatter which GPU it is? So the question is, do you track in reduce-scatter which index goes to which GPU? So, by convention, the dimension should be, essentially, the size of the world. So it could be a general tensor, but one of the dimensions is the size of the world, and it just infers what you want to do. The input data must match the size of the world, and then it knows that the appropriate calculations go to each of the outputs. Yes, you have to be a little

23:16 careful to make sure the dimensions match. So going through smaller examples like this can be helpful. Are there any more questions? Okay, so we're finally in this process that works, and when you're done, you just clean up after yourself. OK. So, so far we've talked about these collective

Benchmarking all-reduce

23:47 operations, a little bit about how they're implemented in PyTorch, and that's NCCL, and then PyTorch. Let's do some performance testing in the spirit of what we did in the assignment or in the first lecture, or rather in the second. We will focus on one node for now. So let's do all-reduce. So I'll have this tensor with 100 million elements and a world size of four. OK. I'll just allocate memory for the tensor. And in general, I hope you can already appreciate that you need to be very careful when testing to clear the field, so to speak. In this case, I will do a "warm-up": run the operation once, synchronize, and execute the barrier. Some of this, I think, is probably a bit of a safety net, but just to make sure that all the kernels are loaded and everything that needs to be computed is computed, and then I'll start the timer, do an all-reduce, resync, and stop the timer. Okay. Now I can see how long it took. OK. So if I scroll down, I guess it's not that informative. I probably should have

25:27 outputted the time in microseconds. It was probably very fast, some number of seconds. Now let's measure the throughput, that is, the number of gigabytes that were actually transferred in total per second. So the way we do this is we need to think about what exactly is being conveyed here. Here's a tensor with this element size, and the size of each element is probably float 32. That would be, you know, two or, excuse me, four bytes, and that's the size in bytes. OK. And now this is a bit of a, you know, delicate point. So, how many bytes were actually sent or transmitted, sent, received? So, each tensor in rank has a size in bytes, okay? And he needs to send it to world size minus one, well, other machines, or rather, ranks. So, there is a factor of two there. Why is there a factor of two? Because you are doing all-reduce, remember. So, you need to send all the individual items to essentially one place. They need to be summed up, and then it should come back to everyone. OK. So rank needs to

27:08 send input, so to speak, and then receive output. That's why there's a factor of two. So, the total duration is the world size multiplied by the actual elapsed time. I think we're just assuming that if we have four processors, that's like four times the real time, and throughput is just bytes divided by duration. OK. So what do we get here? That's about 277 gigabytes per second. OK. You know, I think for the H100 above I stated it was something like 900 gigabytes per second. Of course, as we know, your results may vary depending on the size of the tensors, the exact number of devices, and the weather, well, not the weather, but various factors. So

Benchmarking reduce-scatter, and the accounting trap

28:13 your results may vary. So it's always a good idea to test to see how many gigabytes per second you're actually getting. OK. So, reduce-scatter will be very, very similar. So let's go through this really quickly. So, we create input that is equal to the world size multiplied by the number of elements. So each rank will have this matrix, and so we'll do a warm-up, then, start a timer, do a reduce-scatter, stop the timer, and see how long that took. Well, okay, that doesn't help. Now let's look at, bandwidth. So the number of bytes sent here is not a factor of two, because in reduce-scatter, remember, all you're doing is sending, well, your input to, well, one place. If you just think about reduce, all the elements just go into one place, and that's it. And scatter just means that different components of your tensor go to different places, but actually it's, it's like, you know, reduce. Good. So if you do the same calculation, you'll see that, apparently, in this case, I get 70. So I'm not exactly sure why 70 and not some other number. I guess you could say that

29:50 all-reduce usually has more traffic, and, all-reduce is probably more optimized. I think the Nvidia hardware has some kind of accelerated work that actually does some of this, computation in, you know, the network itself, which cuts the load in half, but I don't know if that fully explains the difference here. There's a lot going on at NCCL, so it's a little hard to speculate about performance accurately. Hence, performance testing. Yes. Another question regarding the bytes or bytes of data sent and how exactly that was calculated. It seems like it only calculates, the data that is sent to the output. But what about the input data? So, the reduction stage. I was wondering how it gets the input to perform the reduction. So the question is, it seems like these are just, bytes for output, but what about the input? So for clarity: I'm assuming the input data is already on the device. So I don't take that time into account and just, I just calculate what needs to be done to reduce-scatter. Is this just a scatter operation? This is the reduce-scatter operation. So, you need a reduction stage. This function does reduce-scatter. So, this is one operation. Good. I mean,

31:30 we counted it twice in the previous one, because we were doing a half-by-half reduction. So you're saying that for all-reduce, there was a 2x factor because you had to do the reduction and then, you know, distribute again. For reduce-scatter, I mean, it's just a name. This is called reduce-scatter, but it's really just a summary. Good. You can also see that if you do reduce-scatter and all-gather, each of them does not have a factor of two. So when you add them, you get a factor of two, which is another way to see that all-reduce is being executed twice. Good. And there are several sources you can read about how to benchmark these collective operations. Good. So now let's talk about distributed training. Our general approach is that

Data parallelism: cut the batch

32:39 I will walk through the basic implementation of each strategy on deep multilayer perceptrons (MLPs). Remember that you are usually in a mode where the MLP is the computational bottleneck, not the attention, in Transformers. So in a way, even though it's a very simple architecture, it's pretty indicative of the type of workloads you'll see. Okay, let's start with data parallelism. By the way, I'll note that data parallelism, tensor and pipeline parallelism can be thought of as different ways of partitioning your model or your data, which is what I hope to visually depict here. Good. So, in data parallelism, here's your model, let's say it has four layers. Each layer of an MLP is simply a matrix multiplication, where this is the hidden dimension. So the data is also a matrix that has a batch dimension and a hidden dimension, and data parallelism simply slices it along the batch dimension into smaller pieces. Good. Now each rank will receive a separate piece of data. Let's look at an example. So, I'll generate some sample data. Let's say I have a batch size of 128, a hidden dimension of 1024. And I'll just generate random data. Okay. So I have a

34:19 batch size of the number of dimensions, and I'm going to run this data parallelism algorithm, or DDP. So, here is where I got this data. There is batch size and dimension, as mentioned earlier. Now I divide the batch size by the total number of processes ( world size). This is how I get the local batch size. This shows what the batch size is at a particular rank. And then, based on the rank, I simply determine what starting and ending local-packet-sized indices I need to get to access the corresponding data. So, essentially, I just go in and take a certain subset of rows based on rank. OK. Now I'm setting up MLP here. And it's done very, let's say, minimalistically. So, here I am creating the MLP parameters. Each layer essentially has a matrix of size "number of dimensions" by "number of dimensions", and remember, the number of dimensions is 1024. And I'm going to create an optimizer. So remember that this,...function works asynchronously on all the different ones, on every rank. So, each of the four ranks will do this with rank 0, 1, 2, 3. And now I'll start the

35:57 training. So for a certain number of steps, I'll be doing a forward pass through the layers: matrix multiplication, nonlinearity, matrix multiplication, nonlinearity. There are four layers here. I will calculate some losses. I don't really care what the losses are. It's just a made-up, artificial thing. And I'm going to do a backward pass. So, so far it looks like I'm implementing a regular SGD, right? And that's actually the point. The only difference for the DDP implementation is that you just insert this line here, which synchronizes the gradients between workers. So what you do is for each layer you call all-reduce where you calculate the average, and what you average is param.grad. OK. It's like you took someone's SGD code and said, wait, I'm actually just going to blend all the gradients after the backward pass. And after you do that, you just update the settings as usual. So from an SGD perspective, it seems like nothing much is happening. I'm just running SGD, but, you know, someone just messed up my gradients. OK. So, I think I'll just put something on the screen. So, for data parallelism, I derive the losses. It is worth noting that the losses

37:33 differ between all the different ranks because they have different data. But after all-reduce, all parameters become the same. OK. So, this is kind of a textbook application of all-reduce in a machine learning setting. Yes. When each rank runs this all-reduce, how do they ensure that they are all on the same step? Or maybe it doesn't matter. So the question is how to guarantee, if all these processes are just running asynchronously, how to make sure that each of them is actually, for example, at the same step. That's because all-reduce is a synchronization point. This will stop everyone and perform an all-reduce operation. So you have to be careful, because if one of your ranks misses all-reduce, everything will just hang. Others are waiting for yes oh, why does getting initial parameters depend on rank? The question is why obtaining the initial parameters depends on the rank, they should be the same. They are the same, the only reason is that I guess...the code for this just places them on the corresponding GPU. Okay, are there any more questions? So, DDP is what you

39:20 implement in the second task, which some of you may have already looked at and some of you may not. This will be done in the context of a transformer, but this is kind of the simplest version so you can see very clearly what's going on. Okay, this is DDP. The losses are different for different ranks, but the gradients reduce to the same values; so the parameters of all ranks are the same, right? So effectively you are doing the number of SGD runs equal to world_size, but because they are synchronized, they do the same thing. You can think of it as a kind of analogue of activation checkpoints, where sometimes you do extra calculations because you don't want to store anything. In this case, we could, for example, pass the optimizer state, but that would be a

Tensor parallelism: cut the width

40:23 bad idea, since it is much faster to update the optimizer state than to move its parameters. Good. I tried to do FSDP last year, but it was too difficult, so I'll skip it. And I'll move on to tensor parallelism. The idea here is that we leave the data unchanged, and now slice the model along the latent dimension. Good. So each rank will get each layer, but only a portion of each layer. And as a result, we will be forwarding all data and activations. Okay. So, we generate the same sample data. Let's look at tensor parallelism. Okay. So, my, batch size and number of dimensions are the same as before. And now I'm going to, earlier I was dividing the batch size, and now I'm dividing the number of dimensions. So my local number of dimensions is 1024 divided by the world size, which is 256. So each model essentially, excuse me, each rank gets a portion of the model that is a fraction of the parameters equal to one over the

41:59 world size. Good. And remember why we use parallelism in the first place—because the model won't fit on a single GPU. So we're going to split it between multiple GPUs. So the parameter matrices now have the size: the number of dimensions times the local number of dimensions. And now every rank will be, I'll only implement a forward pass here. Not the entire training loop. So I'll start going through all the layers. So, first I'll calculate the activations. This looks pretty familiar, except remember: activations actually have a shape of batch size by local number of dimensions, not the total number of dimensions, because each rank only has a fraction of the activations. But now, when I get activations, I need to, you know, communicate, and here I have to, allocate memory for all the activations. So at this point, everyone has, an X, but this X, represents a different, portion of the activations. Good. So now I'm just going to take batch size by local number of dimensions times the world size. So basically, each rank will have

43:37 enough, I'll just get, have a number of matrices of size "batch size times local number of dimensions" equal to the size of the world, and then do an "all-gather." Good. So, I'll send all the activations. And this, I mean, is pretty, simple. So, remember, X is batch size by local number of dimensions, but X is different for each rank. So when I do an "all-gather", I'm going to put that in activations that are essentially, well, world size times, you know, the same shape as X. Okay. So now every, rank has the same activations, now has activations of all models, the entire model. Good. And then I'll just combine them together to get, you know, X. Okay. So now X is again batch size by number of dimensions. Good. And I, you know, repeat. As you can see, there's quite a lot of communication going on here, so remember, Tatsu said that tensor parallelism requires fairly fast interconnects, otherwise you'll be constantly forwarding these activations. Okay, and then you do the same for the next

45:13 layer, and the next, and you get the idea. And just to get some results. So tensor parallelism, let's see, the forward pass creates full- size activations, and at the end everyone has the same activations. Good. So, I'll skip the backward pass because it's pretty boring to do. OK. Are there any questions about this? Yes. I was just wondering why it's so hard to do. So why is it difficult to perform a backward pass? I don't think it's necessarily difficult, but in the context of limited time and space it's, well, it's not difficult. It just requires a little more

Pipeline parallelism: cut the depth

46:18 work. Good. So now let's move on to pipeline parallelism. In this case, we break the model down into layers. So all ranks get all the data. And all the ranks, each rank gets the whole layer, but they get different layers. Good. So, we select the data and run this program or this function for all ranks. Good. So here I'm going to figure out how many layers there are for each rank, which is two in this case. So I have a four-layer network. I have two ranks. So each rank gets two layers, just like in this picture. And here I'm going to highlight the parameters for only the layers that I need. Good. So, I'm going to do a forward pass. Remember that there is further optimization you can do, because if you do it naively you get those pipeline "bubbles

47:54 "that Tatsu talked about earlier; One way to mitigate this is to break the batch into micro-batches. So here I'm going to split this batch into packets of size 32. So four packets of 32 elements. And now the idea is that each rank will essentially wait for the previous rank to pass on activations to it. It applies these layers and then forwards them to the next rank. Starting from the base case, we have rank zero. It's just data. So, I'm just breaking the data into a bunch of micro-batches. And I go through each of the micro-batches. I get the tensor first. So now I use these point-to-point primitives instead of collective primitives. And I actually basically get the tensor X and then calculate the layers assigned to that rank. In this case, there are only two, and then I'm going to send this to the next rank. And again, shipping is a point-to- point operation. And then the next micro-batch, I'll

49:31 do the same thing. So okay, I'll skip it. OK. That's all, actually. So pipeline parallelism, at least its very naive version, is relatively simple conceptually. I meant to mention last time that there is a lot missing from this basic implementation. Overlapping communication and computation is something we don't do here at all. For example, receiving and sending are synchronous, but they should actually be made asynchronous. And also the order in which you perform the forward pass. Actually, it's just a forward pass, not even a backward pass, but when the reverse comes along, you'll need to figure out how to alternate forward and backward steps. So, what's interesting, I guess, is that you just mentioned asynchrony, which isn't shown here— in fact, the GPU will kind of "listen" to see if something else is passing it, and it's kind of event-driven, meaning it only starts processing after the previous layer has passed it data. So the question is, is this event-driven programming, where you just wait for something to happen? I think in event-driven programming, you write handlers, and when something happens —like a mouse click or a

51:03 file ready event—a certain piece of code is executed. This is quite different, in my opinion, from a coding style where everything has to work synchronously. It's true that you're kind of waiting for the previous rank to send you the information, but at least in this implementation there's no flexibility as to where it comes from. It's not like it's waiting for arbitrary data from anywhere. I think there

What the toy code leaves out, and the Jax detour

51:39 are ways to do asynchronous learning, which, you know, was quite popular 10-15 years ago, where it's increasingly event-based, where there's a server that sends the data, and as soon as the gradients are ready, it loads them; gradients accumulate, and if workers fail, it's, you know, handled more reliably. But in modern learning, despite significant scaling, everything seems to operate in a synchronous paradigm. Yes. So it's true that when I say that workers and ranks run asynchronously, it's only because they're different processes, but you're still imposing a pretty tight synchronization on how everything runs synchronously. Yes. How do I modify this program to manage configuration overlap? So the question is, how do we change this to block communication and computation? For example, when you send this, there's no point in just waiting for the data to be sent. You just start sending. Remember that the sending actually happens to the GPU by running a specific core. So it's kind of independent. And it can immediately move on to processing another micro-batch. I think this

53:17 can be done using another function called isend, which is asynchronous and returns a handle. So, you just do all the sendings, and at the end you just wait for them all to complete. And for the overlap, when you have a backpropagation step, you actually need to schedule it here. Yes. Sending and receiving are identical, if you have multiple sends and multiple receives, how does it know which is which? So the question is: if you have multiple sends and multiple receives, how do you know where what is? Here the name of the tensor doesn't matter, it's just any variable, and what you specify is the source. So if I'm on a node and I'm receiving data, then the next message that comes from that rank, I'll just put it in this x and continue execution. What if I want to make two shipments from the same rank? If you want to make two shipments from the same rank, to the same destination. So, I'm not entirely sure about this,

55:00 but I think if you have two send operations, they kind of get queued up. So, the order of sending is still preserved. It's just that other things can be happening at the same time. For example, you can send, I think if you have a pair and you make two sends, then that order is preserved, but, the order in which, well, you send and another rank sends to another rank, that can happen at any time. Yes. What would happen if you just made a shipment like that, but no one receives it, would it just get stuck there or something? So what happens if you sent it and no one received it? I think it will just stop; it will just wait because there is no...yeah, I mean because the process can just keep running and you don't know if it will get it. It's just code execution, so you don't know if it will never get there or if it's just a matter of time. Yes. So the question is, what happens to the last rank? So, in the end, the last rank has all the activations. That is, it actually has the results of a complete, forward pass. And then, you know, if you implement a backward pass, then you're now actually calculating the gradient with respect to the loss function, and then you go back and send, from rank to rank minus one and so on. Good. I guess, I

56:44 was afraid I was running out of time, but it looks like I actually do have time. Maybe next year we should consider a backward pass. Okay, actually I'm going to finish pretty early today, so, if you have any other questions, you should ask them. So we've gone through three simple examples of parallelism so far: data, tensor, and pipeline. Of course, this is for simple MLPs. You'd actually like to do this with your own, you know, more complex model, like a transformer. I argued that at least the basic ideas can be understood to some extent through MLP. I think so, but of course, when you want to train, you want to train a transformer, not a deep MLP. So you'll still have to implement the full complexity. What's also missing is the communication and computing overlap, which isn't very, well- crafted here. And overall it's more complex code with record keeping. I, you know, encourage you to check out Megatron-LM or, PyTorch FSDP. It's getting, you know, pretty...well, hot. And one of the things that I think makes some of the accounting, at least for say FSDP, and you'll run into this a little bit in A2, is that if you want something that handles arbitrary architectures, you have to deal with parameters and

58:20 do a bunch of calculations to figure out what layers are there and so on. Whereas in the case of MLP, I just decided that I would break down the model in this, you know, pretty simple way. One more thing I'll just mention in passing: everything we're doing in this course is PyTorch, but it's useful to know about the other ecosystem around Jax and TPU. Which, by the way, is quite good in a way. The idea here is that Jax allows you to simply define a model. You define the sharding strategy and the Jax compiler does the rest. So, there is a toolkit that we developed called Levanter based on Jax. And I'll just show you a snippet of what it looks like. So this is FSDP in 10 lines of code, and essentially you have a model and you just say "shard" with this particular... I mean, I don't expect you to read it in detail, but you just specify what dimension you're going to shard, and that's basically it. Similarly, for tensor parallelism, you just say, "I'm going to shard the model along...", you can shard along the dimensionality of the heads for the attention mechanism, as well as along the dimensionality of the model. So, to some extent, it

60:06 gives you a kind of conceptual simplicity of what you're trying to do. You have this essentially a graph of computations, but it has these dimensions: model dimensions, embedding dimension, attention sequence dimension, and Jax allows you to just specify which dimensions you want to partition along, and also define a mapping of that to the actual TPUs, and then the JAX compiler magically figures out how to compile that into the primitives that move the data. So this is a much higher level than working with collective communications. But you know, we're sticking with PyTorch because it allows you to look under the hood, so to speak, and see what's really going on there. But if you're actually doing this in the real world, then

Summary, then questions

61:08 of course you don't need to and probably shouldn't implement all of this from scratch. Okay, that concludes Jax detour. So, to summarize, we have already seen many ways to parallelize. And each of these parallelization methods can be thought of as distributing the model or data along some dimension: the data dimension, the batch dimension, the width dimension, the depth dimension, or the context length dimension. We also see a recurring theme of overcomputing. You can recalculate something from scratch, or you can store it in memory and incur the cost of data transfer, or now, in multiprocessor or multi-node systems, you can store it in memory on another GPU and then communicate between them, which, you know, is even slower, so there are some tradeoffs here. And often recalculation may be the better option, but of course you can't recalculate everything, and often you're limited by either connectivity or memory. In conclusion, I will say that hardware is getting better. So you might think that maybe none of this is needed because in five years everything will fit in L1 HBM. But that's not the case, because while they can

62:54 grow significantly, there are physical limitations and we will always build larger models that are at the limit of the hardware's capabilities. So this architectural structure has always been with us since the advent of computer systems and always will be. Okay, that's all I wanted to tell you today. So I am ready to answer any questions. Yes, augmenting with the same set of parameters may cost more because your normalization may depend on the entire dataset. For example, in So the question is, are you saying that with data parallelism, even if the parameters are synchronized, there can be other things that depend on the data, like, for example, batch norm. Honestly, I don't know about you, but batch normalization is always a little annoying. So I don't know exactly how to do it, from memory. And I think, at least in the world of language models, this doesn't really arise. Because layer normalization is used. And as long as you initialize all parameters using the same random seed number, you'll be fine. I mean, there might be issues with non-determinism

64:34 on the GPU, but hopefully they're, you know, minor. Well, yes. So, the question is whether PyTorch has any of the same conveniences as Jax. So, I mean, PyTorch has an FSDP library that you should definitely use if you're not on this course. Which is essentially a wrapper. You define any model and it simply applies FSDP to it. I guess if you're asking how well it allows you to do your own sharding, I think some things are coming up, but it's not that developed. I mean, I think there's a spectrum between the Jax world where you define things declaratively, and I think Google's infrastructure, if you stay in the Jax TPU system, is pretty well developed. But then if you look at DeepSeek, which is the opposite end, where you have these GPUs with a really bad interconnect, that means they have to go in and hack, you know, they actually get down to the NCCL level and do a bunch of things that I don't really understand to squeeze out

66:07 performance. Whereas if you write in Jax, you just declare your model, and then, you know, everything happens by itself. So the ways you use hardware, I think, really depends on what ecosystem you're working in. Yes. The amount of recalculation of activations, they can recalculate part of the activations. Is there an API that can. So, the question about checkpointing activations, is there an API that allows you, I mean, in PyTorch and Jax, to specify which parts you want to recalculate, because you definitely don't want to recalculate everything or nothing. Probably every few layers, right after big matrix multiplications, when, for example, you have matrix multiplication and then pointwise linearity, I don't think you need to keep two copies. That is, if you have two things that are fairly trivial to retrieve, you can just keep one version. Yes, continue. Good. Do you think GPUs will be replaced by specific hardware, or something more specialized? So the question is, will GPUs ever be replaced by hardware specifically designed for Transformers? You already see this quite often in the field of inference, for example Groq and Cerebras have specialized equipment that can

67:54 perform inference, and also, apparently, training. Cerebras is engaged in training. So basically this hardware just gives you a lot more memory on the chip. I mean, that's actually the main goal. It seems to me that Cerebras has a huge, essentially efficient L1 cache, so you don't have to constantly move data around. And I think there's a lot of simplification that can be made, because GPUs carry a lot of redundancy, if you think about it, because they were designed in an era where you had to do a lot of branching and different types of special calculations that aren't really needed in deep learning mode. So I think there is quite a lot of room for improvement in the equipment. I think there was a hand there, I am now. I'm not sure if this is the right question I'm thinking about, but in the context of the lecture, it's essentially a model that learns in one go, that optimizes. But I'm wondering if any of the techniques we're talking about can be used to incrementally train a model, like when you get new training data, not just to retrain, but to actually update everything without having to recalculate everything. Yes. So the question is, can these methods be used for continuous learning? Yes, definitely. If you

69:29 think about the unit of what we're working with, it's just doing gradient steps, right? So, if you take a checkpoint, you can just keep doing what you were doing. There is nothing specific about starting from scratch. Here. I think there was a question there. Yes. So, in terms of model-specific equipment, you know, the previous question, there are probably physical or technical reasons why the nodes can't be made much larger than they are now. So what exactly is the change you are talking about? If it were possible to make GPU nodes infinitely large, as much as you like, people would do it. So, there must be a technical or hardware reason why this is not possible. So what is the actual progress for the specialized equipment you mentioned? Yes. The issue is that there are physical limitations, certainly for GPUs. Let me just continue. So you can't make GPUs infinitely large or infinitely dense. I mean, there are also problems with energy consumption. You know, you have to get rid of all the heat, and, you know, there's a limit to the bandwidth that can be accommodated. I don't know the exact details, but

71:08 at least in the case of Cerebras, they have a certain way of manufacturing chips where the memory is on the chip itself. So I guess it's just a way of putting it. And I think there are obviously trade-offs here because it comes at a cost in terms of less flexibility. But overall, I think it's worth thinking about it more broadly: GPUs were developed in the CPU era, where the focus was on control. Like, I have code that I'm executing—it's a priority, and data needs to be moved to execute that code. But the big difference with deep learning workloads is that they are a data stream: the graph of computation, if you look at it, is static. You know exactly from the beginning all the calculations that will be performed by the end of the training, right? So, using this knowledge, you can plan your calculations much more intelligently than having to adapt to the uncertainty of the flexibility of arbitrary calculations. Okay, maybe a few more questions, and that's it. Yes. Is the computational graph usually stored on the CPU or on the GPU? So the question is,

72:42 where is the computational graph stored? Well, the code, actually, all of this code runs on the CPU. But when you call something like a PyTorch function that needs to run on the GPU, it launches kernels in the background, and those kernels are the code that runs on the GPU. Yeah, I'm not sure if that's... I guess another answer is that a computational graph is more of a conceptual thing... You know, it's not like there 's some graph that literally... I guess it exists, but it doesn't mean that the graph is hosted on the GPU, you know. Of course. So these communication primitives that we have, for example...are they actually CPU instructions, or are these programs using the GPU? So, the question is, are the communication primitives CPU- or GPU- dependent? These collective operations are in a sense an abstract specification of what types of operations should be performed, which can happen if you remember PyTorch distributed has different backends. So this can happen on the GPU or on the CPU. But when they are

74:27 executed on the CPU, does the CPU schedule them somehow, or is it like cores running independently? Yes. Well, the CPU actually runs the process, it remains the master, and when you do a collective operation, it calls the NCCL library, which runs—it's still the CPU —and then it runs the kernels that move the data. Yes. Good. Maybe this is a good time to wrap up. OK. See you next Monday, or rather Tuesday.