PyTorch 2.0 is here! Quick tutorial + NVIDIA RTX 4080 giveaway

Daniel Bourke walks through PyTorch 2.0's torch.compile API, which delivered a 43% average training speedup across 163 open-source models on NVIDIA A100s. The new feature uses operator fusion and graph capture to cut memory overhead. He also unboxes two RTX 4080 GPUs from NVIDIA and announces a giveaway.
Chapters
Welcome to PyTorch 2.0!
It's an exciting day.
Today is the official release of PyTorch 2.0, with the main improvement being speed. This comes via a single backward-compatible line: torch.compile. In other words, after you create your model, you can pass it to torch.compile to receive a compiled model and in turn expect speedups in training and inference on newer GPUs. But this is just words. How much faster is it actually? Well, the PyTorch team ran tests across 163 open-source models from Hugging Face Transformers, TIMM, or Torch Image Models, and TorchBench, a curated set of popular codebases from across GitHub, using a mixture of AMP, Automatic Mixed Precision, or float16 training, and float32 precision training, the PyTorch team found that torch.compile provides an average speedup of 43% in training on an NVIDIA A100 GPU. That is incredible. And naturally, I tried torch.compile in one of my own benchmarks. I found that it definitely seems that torch.compile is much better when you're using as much of the GPU as possible, e.g., larger batch sizes and data sizes. I also noticed that the first epoch was slower, but subsequent epochs went faster across both my local RTX 4080 GPU and an A100 GPU on Google Colab. My speedups weren't as large as reported in the PyTorch 2.0 release notes. However, they were still around a 10% improvement on training epochs. I'm yet to explore mixed precision training or transformer-like architectures, longer training runs, and this is where I think torch.compile would probably be most effective. All right. So what's happening behind the scenes? Well, torch.compile is designed to just work, but there are a few technologies backing it: TorchDynamo, AOT Autograd, PrimTorch, and TorchInductor. The PyTorch 2.0 release notes explain these in more detail, but from a high level, the two main improvements torch.compile offers are: number one, fusion, And number 2, graph capture. Fusion, also known as operator fusion, is one of the best ways to make deep learning models go brrr. Operator fusion is just like fusion in Dragon Ball Z. It condenses many operations into one, or many to less.
Why?
Modern GPUs have so much compute power that they are often not compute limited, as in the main bottleneck to training models is how fast can you get data from your CPU to your GPU? This is known as bandwidth or memory bandwidth. You want to reduce your bandwidth costs as much as possible and feed your data-hungry GPUs with as much data as possible. So instead of performing an operation on a piece of data and then saving the result to memory, computation, memory, computation, memory, you chain together as many operations as possible via Fusion. A rough analogy would be using a blender to make a smoothie. Most blenders are good at blending things, like GPUs are good at performing matrix multiplications. Using a blender without operator fusion would be like adding each ingredient one by one and blending each time a new ingredient is added. Not only is this insane, it increases your bandwidth costs. The actual blending is fast each time for each single ingredient, like GPU computations generally are, but you lose a bunch of time adding each ingredient one by one. Using a blender with operator fusion is akin to using a blender by adding all the ingredients in at the start and then performing the blend once. Of course, you lose a little time adding at the start, but you gain all of the lost memory bandwidth time back. Graph capture, I'm less confident explaining. But the way I think about it is that graph capture or graph tracing is going through a series of operations that need to happen, such as the operations in a neural network, network and capturing or tracing what needs to happen ahead of time. Computing without graph capture is like going to a new area and following GPS directions turn by turn. As a good human driver, you can follow the turns quite easily, but you still have to think about each turn that you take. This is the equivalent to PyTorch having to look up what each operation does as it does it. As in, in order to perform an addition, it has to look up what an addition does before it can perform it. It does this quickly, but there's still non-zero overhead. Computing with GraphCapture is like driving through your own neighborhood. You barely think about what turns to make. However, it took you some time up front to remember how to drive to your house. This is a caveat of GraphCapture. It takes a little time up front to memorize the operations that need to happen, but subsequent computations should be faster. For more on the concepts of fusion, memory bandwidth, and overhead, I'd recommend reading "Making Deep Learning Go Brrr." From First Principles by Horace Heh. Outside of torch.compile, there are plenty of other notable PyTorch 2.0 releases, such as: Faster Transformers is now in stable condition, which means that models that use the transformer style of architecture or scaled dot product attention should be faster. Universal device setup is now possible thanks to a context manager or global setting. And a bunch of updated MPS, or Metal Performance Shaders, operations improves PyTorch's coverage for Mac. There's also a bunch of domain library upgrades to go along with PyTorch 2.0 in Torchvision, TorchText, TorchRec, TorchAudio, and more. To learn more about all of the updates to PyTorch 2.0, and there's a lot, I'd recommend checking out all of the links I've left in the description and the blog post attached to this video.
Okay, it's a very fun day at the Nerdify HQ. Well, Not really, Josh and I are both stuck on pretty in-depth errors, but something fun just arrived.
Ooh!
I don't actually know what this is. So it's from NVIDIA. I've been talking to them and they said they're going to send me something celebrating upcoming GTC conference. However, This is pretty big packaging, and the best part is my little brother, who's a gaming enthusiast, is coming over and he doesn't know that this is here. What's up, brother? I need your help unpackaging something. Why are you recording? What's going on? I don't know. Have a look at who that box is from.
NVIDIA?
Did you know this was coming? No, what is it? You know how I said, can you help me film something? Yeah, I don't know what it is, but feel how heavy it is. It's pretty heavy. What do you reckon it is? I don't know, GPU or something. Apparently this is 15K. I don't know what's in it.
I think I know.
You do know. Oh, one package and second package. It's come all the way from America. Oh no way, Josh, can you get a close-up of this? You got two? All right, so NVIDIA were very kind to send me not one, but threes, but two. Oh my God, GeForce RTX 4080s! This packaging is ridiculous. Quality. Okay, there's some serious first proofs. Oh! You ready? Cue the epic music— Dumdum dum dumdum dum dum dum dum dum dum... Oh my goodness! Look at that box! What in the world?! Got to handle it with care— oh my gosh! That is some GPU power right there! So the goal will be to get this installed in my deep learning PC and then test it out. Oh, now this is the beautiful Titan RTX Gold Edition with some dust on it because it's been working.
This is a very capable card.
It's got about 24GB of VRAM, and I think it came out in maybe 2019, somewhere around there. Put some stats up there. The 4080 has 16GB of VRAM but is generally faster across every other metric.
Now ideally this would slot right in.
Why do I think that's not going to happen? This is a side-by-side comparison of the Titan RTX, which is a big card, and there's the RTX 4080 side by side. Look at the size of your head. Oh my God. Here we go. You know, I read a good blog post once that the, the brain might actually just be a 3090, and I'm probably getting that completely wrong, but I think it's just gonna— what? No, that can't be it. Oh yeah, let me adjust the camera. Gotta line it up here with the screws. That just clicked in. Okay, I don't think there's a chance that we're getting two of these in there. Maybe one day. 3 PCIe's in. It's like a cooking show for deep learning. Now let's go plug this back in. We got Dinty V2, a little bit Frankenstein, but that's okay. Got the moment of truth. Let's bring it to life. Oh, oh, it lights up! I didn't know that was going to light up. That looks epic. Oh, need to fix this fan. Now I believe we have to install some drivers. Many hours later. Hello and welcome to Daniel's giveaway garage. Spoiler alert right there. Now, after about a morning or so of troubleshooting, I got the RTX working with CUDA. Everything works, can train PyTorch models like a charm. The trick is uninstall previously installed NVIDIA drivers before installing new ones. Took me a while to figure that out. Now you might be wondering why NVIDIA sent me one— well, not one, but two— one second— two RTX 4080s. Well, the spoiler alert is it's a giveaway for you! Now, to celebrate NVIDIA's GTC March 2023 conference, they've been so kind to send me two RTX 4080s, one to keep for myself and use and train deep learning models, and one to give away, as well as five Deep Learning Institute credits. The details of how to enter the giveaway will be in the blog post attached to this video. As for GTC, it's NVIDIA's GPU Technology Conference from March 20 to March 2023. It's free to sign up and there's hundreds of different sessions on everything from AI to deep learning hardware to deep learning software. It's a playground for people who are into AI. I've got 5 sessions I'm going to. Number 1: Compile and train with 43% speed up using PyTorch 2.0. Number 2: Fireside chat with Ilya Sakskeeva. I might be saying that wrong, who's the chief scientist of OpenAI, and Jensen, who's the CEO of NVIDIA. AI today and vision of the future. Number 3, exploring next generation methods for optimizing PyTorch models. PyTorch with TorchTensor RT. Number 4, FP8 mixed precision training, which was introduced on the latest GPUs with Hugging Face Accelerate. And number 5, generative AI demystified. So big shout out to NVIDIA. Thank you for partnering with me for GTC and sending me two RTX 4080s. If you'd like to enter with a chance to win one, as well as one of the five Deep Learning Institute credits, follow the instructions in the blog post. Go and check out NVIDIA GTC. And as for me, it is way too hot in Australia to be wearing this for much longer, so I'm going to take this off. And I've got some models to train.
More from Daniel

Small Language Models (SLMs) Are the Future: Fine-Tuning AI That Runs on Your iPhone

Introducing Sunny - A MedGemma-Powered Skin Health Tracking App (Kaggle Competition Entry)

Benchmarking: NVIDIA DGX Spark & NVIDIA RTX 4090 | Part 4 - Comparing the results

Benchmarking: NVIDIA DGX Spark & NVIDIA RTX 4090 | Part 3 - Running LLMs on the RTX 4090

Benchmarking: NVIDIA DGX Spark & NVIDIA RTX 4090 | Part 2 - Formalizing the tests

Benchmarking: NVIDIA DGX Spark & NVIDIA RTX 4090 | Part 1 - Creating a series of tests

Local Multimodal RAG on the NVIDIA DGX Spark | Part 4 - Publishing our demo to Hugging Face

Local Multimodal RAG on the NVIDIA DGX Spark | Part 3 - Putting it all together


