Keynote: PyTorch Updates - Edward Yang, Research Engineer, Meta

In this keynote presentation at the first PyTorch conference in Europe, Edward, a core maintainer from the PyTorch team, discusses the recent advancements in PyTorch, particularly focusing on distributed training features. He highlights impressive community contributions, with over 15,000 contributions and 1,250 new contributors in the past year. Edward introduces two significant features for distributed training: pre-compilation support for 'torch.compile' and SPMD (Single Program Multiple Data) types. Pre-compilation enables ahead-of-time compilation to streamline distributed workloads and improve performance, while SPMD types provide a type system that helps manage tensor distributions across different nodes. This approach aims to ensure correctness in gradient calculations, simplifying the process of implementing parallelism in machine learning models. Edward emphasizes the commitment of the PyTorch team to support researchers and enhance the overall user experience in machine learning.

Thank you so much for that introduction. Wow, it's great to see all of you here in Paris for the very first PyTorch conference Europe. My name is Edward. I'm a core maintainer on the PyTorch team, and I'm thrilled to give you an update about what's been going on in PyTorch, and I also want I'm excited to spotlight a few new features that I really want to tell

you about. Um it's something of a tradition to do some stats about the main repository. So, here are some stats as of the beginning of March 2026. Um since this is our first PyTorch conference Europe, all of the all of the bars are shifted. So, you know, it's every March instead of every you know, November, which is how we used to do it um during regular PDC.

Um but you can see that PyTorch has continued to steadily grow, not just in stars, but also in the number of contributions. I want to remind you that this chart is the number of contributions this year, not cumulative contributions. So, we've had 15,000 contributions to the PyTorch core library. Thank you so much. You know, without the help of the community, we would not be able to cover

so much ground. It's actually a little bit of, you know, like how do we keep on top of everyone who wants to contribute to the library? That's something that I'd be very interested to talk to people about at this conference. And it's also great to see that, you know, we've continued to grow the set of people contributing to the library. Once again, this is 1,250 new contributors

that have contributed within the last 12 months. So, thank you once again. I also want to shout out to our, you know, core audience, right? Researchers who have trusted their workflows to PyTorch. Thank you for continuing to trust us. Um we're still being used nine out of 10 research projects in the top ML conferences, and 90% of open source machine learning projects are using PyTorch. So, you

know, once again, thank you for putting your trust in us. Um we want to continue to be the operating system for machine learning in the community. Um one particular thing I want to call out is that um you know, D tensor and device mesh continue to rise in usage. We had a We saw over 1,000 open source projects use these new capabilities. To me, it really shows

that you know, uh we have been working on making distributed training better and there's a lot of interest in how can I, you know, do distributed training more easily, more effectively. And so, this is a good segue into two of the things that I want to spotlight for you. Now, there are a lot of new things that have happened in PyTorch and you'll get to learn all

about them in various talks during the day. But, for this keynote, I want to talk about two things that are specifically about distributed training. One is pre-compilation support and the other is SPMD types. And I'd love to chat with all of you at the meet maintainers at 10:30 today. So, you know, uh give me any questions there cuz I can't answer them after this talk. All right.

So, let's talk about pre-compilation support for torch.compile. So, what is pre-compilation? Well, to answer that question, I have to tell you a little bit about how torch.compile works today, right? So, torch.compile, you put it on a function, it's a just-in-time mechanism. The first time you run that function, we'll go ahead, look at the inputs, figure out how to compile your model, and give you that sort of

thing. And so, when you think about what happens when you're doing distributed training with torch.compile, well, every node is running the same Python program. They are running this just-in-time compilation. So, it turns out that everybody is doing the compilation over all of these nodes. And well, this works. Like, you can do this in production. We do this in production. But, there are some things that sort of

make this not so great. So, for one, it's kind of slow, right? Not only do you have to worry about, "Oh, is my cache warmed up? You know, am I going to have to actually go ahead and do, you know, minutes of compilation every time I run my job?" But, also, you know, to actually do a cache look up, we have to do the entire tracing process

in Dynamo to figure out if the cache product is still valid. So, you know, people complain about how long it takes to restart jobs in the regime. It's also kind of complicated, right? Like, you can have situations where your model is actually doing different things on different ranks. Um in recommendation systems, it's very common to have dynamic shapes, meaning that, you know, actually every rank is seeing

a different number of inputs on every compile. And this can lead to problems. For example, if you have some uh you know, a setup where one node is recompiling because it just happened to get unlucky and see a different distribution of shapes than the other ranks, this can just lead to a nickel time out because everyone else is waiting for that one rank to finish compiling. Similarly,

if your graphs aren't the same, even if your compiler is deterministic, it may make different choices on different ranks. And if it reorders your collectives in the wrong way, well, that's another way to have a deadlock. We don't have these optimizations by default for this reason. So, how do you solve this problem? Well, instead of just-in-time compiling, let's just go ahead and compile it ahead of time.

Let's run the compilation process on one node and then distribute the compiled artifact to all the ranks. And now you're guaranteed not to do any compilation at run time, and you're guaranteed to be running the same code on every rank. So, I've been like on people for this for a very long time. I was like, "Oh, we really need pre-compilations to support it." And it hasn't been

easy to implement it. Why hasn't it been easy? Well, there are two problems you need to solve to actually make this happen. One is you need to be able to actually compile ahead of time in the first place. So, this is not so easy because one of the things that makes torch compile so great and so easy to use is you can just slap it on whatever

Python code you want and it will always work, right? Like if there's some code that doesn't work, we'll just put in a graph break, you know, we'll compile part of it, we won't compile the part that we don't understand, and then we'll keep going. And so, if you want to precompile something, well, we kind of need to know everything that needs to be compiled and that's a

lot easier if you give us an entire graph. So, full model capture is one of the things that, you know, like you need to do this. The other problem is that you need to actually be in a situation where you can actually compile once and use it everywhere. And so, if you've got some code that's like looking at, you know, what the rank of a particular process

is, well, that actually needs to be written in a generic way so that it works no matter which process you're doing. So, what have we got? Um so, this is still very early days for this, but we have a new graph trainer in torch Titan, which is all about this workflow, this full model capture workflow. And we've shown that this works with bitwise equivalence and exactly matching

performance on llama 3. Here's what it looks like. So, you take um we have a uh interface called the precompile uh main and what you can do is you pass it in the training arguments you want to run for real and it goes ahead and dumps the artifacts in some directory. And then when you do actual training, you do your normal training code, but you just say,

"Hey, here are where the precompile artifacts are and we now guarantee you that we're just going to use those as is without having to retrace your model or do anything like that." So, what do we still need to do? Well, we want to make sure this works for all the models in torch Titan. And we also want to make it easier for people to do things where

they do actually do want different behavior under different ranks. For example, it's very common when using FSDP in PyTorch to just have a situation where your parameter doesn't divide it exactly the same way you want. And you know, we want this to work. We want it to be okay to have your last rank have a little bit less data than everyone else. Similarly, sometimes you want to

do things like only log on a single rank. And like that's something we also want to be able to support. It's very early. Um, you know, we did land the PR to Torch Titan last week. Go check it out. That's all I want to say about pre-compilation. The second thing I want to talk about is a sort of different topic um, that's uh, also targeted at this

dis- distributed space, which is SPMD types. So, what do I mean by SPMD types? Well, let's just break this name down. It comes in two parts. SPMD, short for single program multiple data, saying you're running the same code on every node, but they have different data. They have different tensors on every node. And types is the good old-fashioned type system, right? The lightweight formal method for, you

know, annotating your program with some properties about it, and then checking it mechanically with a type checker. So, SPMD types is all about describing how the tensors, how the data are distributed across the ranks on your program. Do you have the same value on every rank or is it different? Do you have a bunch of values which need to do a reduction? And also, do you have

gradients that need to do a reduction? We say this is local SPMD because we don't actually say how to assemble the full tensor in the situation. So, what problem are we trying to solve with SPMD types? Well, if we think about what PyTorch traditionally gives to users, we have this foundational guarantee that if you write a program in PyTorch, we can differentiate it. You can run backward

on it, get gradients for these things. Every op in PyTorch has a derivative for it, except for collectives. We don't have differentiable collectives in PyTorch. You might ask, "Why not?" And the answer is, well, you know, you kind of can do it, right? If you look at a framework like Megatron, they in fact do have custom autograd functions for doing differentiation through collectives. And the problem is

is that when you write these autograd functions, you end up with these kind of weird autograd functions that, you know, do things that aren't very intuitive. So, this uh diagram is from the Megatron LM paper. It's very old at this point. It shows how to do what we say Megatron tensor parallelism. And the thing I want to point out is there is this green box called F.

What exactly is this F function? Well, it's this weird op where you do nothing in forwards, and in backwards, you do an all reduce. And so, this is like pretty standard, right? Like, if you've ever implemented tensor parallelism, you know, you know you have to do this sort of thing, but it's like you have to do it right. If you forget to put this call to F

in your program, you're just going to get silently incorrect gradients. So, people are like of two minds of this. Some people are like, well, you know, this is just how tensor parallelism works. Once you know this is the thing you need to do, like whatever, like you do it. There's not that many ways to do this sort of thing. On the other hand, you have lots of

people who are like, "Ugh, I hate tensor parallelism. It's so terrible because of things like this." And you know, one of the things we always are thinking about in PyTorch is how do we make things like this easier for people? So, another question you might have is, "Hey, isn't D tensor good?" Like, D tensor also gives a guarantee for giving differential gradients. And this is true, and

this is one of the reasons why I think D tensor is very popular. And as VMD types are like actually trying to solve a slightly different problem than D tensor. D tensor is all about imagine you have a full tensor and you know, it's distributed across all the nodes but like semantically you still have the full tensor. You're doing operations on this full tensor. We gave this

to a bunch of power users and they're like, "Oh man, why are you doing this global semantics? The way I think about what I want my programs to do is I want to think about what I'm running locally on every rank, right? Like this is like doing things with plain tensors, plain collectives. This is the way I think about my programs. I want to be able to

write my programs in this way." And D tensor doesn't let you do that because D tensor is all about, you know, giving you this global semantics, letting you sort of forget the parallelism exists at all. So, SPMD types is about, well, let's say you want to do local semantics. Let's say you want to explicitly put your collectives in your program because there aren't that many collectives cuz

they they're expensive. So, you know, you want to keep track of where they all are. And maybe, you know, our fine friends at ByteDance wrote a paper about this. Maybe you also don't like D tensor's eager overhead because, you know, it is kind of expensive. There are ways around it like CUDA graphs or compiling, but it is, you know, it is a problem when you're using it

in eager workflows. So, here is the value proposition of SPMD types, right? Give me your Megatron style plain tensor code with collectives. All I need you to do is to add some type annotations to your inputs and your weights. And I also need you to use some of our special collectives that know how to propagate these types. And in return, I will give you a guarantee that

your type check program is guaranteed to give correct gradients. You cannot get it wrong. The type checker will yell at you if you forget one of those weird calls. And these type errors will tell you where you are missing collectives so you can go ahead and add them in and enforce contracts across module boundaries. And we're going to do this without any runtime overhead because the type

checking is strictly optional. You can run it all the time or you can just run it in a unit test. And all of your communication patterns are going to work out of the box as they do today. I would be remiss at this point not to say where this idea comes from. Um so, we are heavily inspired. In fact, this type system is just a direct implementation

of JAX's sharding and types if you include their undocumented reduced unreduced types. What we've done is we've adapted it for the PyTorch ecosystem, changed some defaults, you know, made it work in this case. And I owe a huge intellectual, um you know, gratitude for the JAX team convincing us that, "Hey, it's a good idea to talk about explicitly what the placements of your gradients are even in

forwards cuz that's, I think, one of the key things here." So, let's just look at a little bit of code, right? So, what exactly is a local sharded type? Well, there's a bunch of different types describing different situations for how your data is laid out over ranks, right? So, like, let's say varying. When something is varying, we say it's different across the ranks. And, you know, this

is different from, you know, a sharded D tensor because all I'm saying is they're different. How do you, you know, stack them back together? I don't know. Um they're just different. That's the only thing we know about the tensor in this case. Similarly, you can have tensors which are partial, which just means that, "Hey, they not only are they differing over the ranks, but you are in

obligated to do a reduction to find what the real value is." And you're not allowed to do operations that can see into the individual pieces because semantically they are already summed together. So, you can do linear operations because linear operations distribute over sums, but you're not allowed to do non-linear operations on these things. The next two examples are very interesting because they are the same in forwards

and different in backwards. So, I'm first going to introduce this concept of invariant because it is simpler. Invariant tensors, in forwards they have exactly the same values across And in backwards, they also all have exactly the same in a value. This is kind of sort of like intuitively what you would think as yeah, this is what I would expect to see out of a tensor that's all

the same in forwards. But it turns out when you're doing most of your code, you want a different type. You want this type which we call replicate where in forwards all the ranks have the same values, but in backwards you have a partial tensor. So, there is a pending reduction you need to remember to do in this situation. So, putting it all together, right? What is a

local SPMD type? Well, for every axis in your device mesh, so let's say you're doing data parallel and context parallel and tensor parallel, we just annotate it with you know what exactly is the type in this case. So, in this example, we say that well, the data is varying over the data parallel axis and varying over the context parallel axis, which is like what you'd expect for

you know the inputs to your model. And it is invariant on the tensor parallel axis, which is something that you might expect outside of a tensor parallel region if you weren't doing sequence parallel. So, let's just look at a little example. So, here is the very classic DDP FSDP case, right? The magic of DDP and FSDP is you don't need to change your model code at all,

right? We go ahead and we say the inputs are sharded on the DP axis and the CP axis. So, you know, you're running in a data parallel fashion, but in the case of FSDP, you've gone ahead and already all gathered all of your weights together, so that you have the entire weight. The weight is replicated. Every rank has a replica of the weight in this case. And

I just want to point out that the type here for the weight is replicate, not invariant, because when you run backwards on a data parallel program, you end up with gradients that are for the local batch on your rank, and you have to do the, you know, good old fashion all reduce at the end to actually get the final gradient. And DDP and FSDP traditionally do these

via hooks, sort of behind your back, so your module code doesn't have to worry about this reduction. The framework takes care of it for you. And so, when you write a type like replicate, you're saying, "Hey, I'm not doing the reduction. There is a pending reduction. I need someone else to take care of the reduction." In this case, the hooks from your parallelism modules. Here's another example.

So, this is the norm in an LLM, and there's no sequence parallel involved. So, the norms are outside of the tensor parallel region, and so typically you will just run the same compute on every single node. And so, in this case, in fact, the gradients are also exactly the same on every node, and we've written we've written everything with invariant. If you say then, "Oh, let's do

sequence parallelism. Let's go ahead and shard the input on the TP dimension." You'll go ahead and change the annotation on the input to be varying. But actually, this is an error. This will error on the type checker, saying that, "Hey, you can't mix invariant and varying outputs." In fact, the full error message will even suggest, "Hey, maybe you needed to do a collective translating the invariant tensor

into a replicated tensor. And the reason for this is the rule is that you can actually only mix, replicate, and varying tensor inputs. Invariant inputs can only mix with themselves. So, the type checker is telling you, "Hey, you're missing a collective. Either you need to change the annotation on the weight to replicate, saying that, 'Hey, I have an unreduced gradient. Someone else needs to take care of

it,' which is very common. Megatron, in fact, has its own hooks for dealing with the sequence parallel reduction. Or, you can keep the weight the same way it is before and force the module to do the reduction. And that's what the That's what the invariant to replicate here is. It's the F function. It's the no op in forwards and the all reduce in So, what do you

get from SPMD types? So, you get a higher API. It says, "Here is the type and how to type check it." You get a bunch of collectives and you get some weird local operations, like invariant to replicate. So, we've open-sourced SPMD types. You can go ahead and check them out on Meta PyTorch. We do plan to put it into PyTorch at some point, but we're still like

working out some of the API details. Um, that's everything I wanted to say about SPMD types, and thank you. >> All right.