Sponsored Session: Lightning Talk: Accelerated Software for a Post-Moore World - Jay Dawani

Key Takeaways

Hello everyone. My name is Jay. I am the co-founder and CEO of Frian Labs. Some of you may have heard of us, some of you may have not. Now, I'm going to be talking to you today about accelerated software for Post Morris World. It's going to make sense in a second. It's a whole lot of words. Now, most of you here have seen something along this lines. It is the modern software stack. it I tried to fit everything I could in one slide but it started getting too messy and in fact I don't think I could actually fit the current state of software into one slide because that is how brittle and fragmented and messy everything is right and it's why we have this conference with a few hundred companies trying to solve this problem now the interesting thing is how we got here right that is the actual problem what caused us to get to this world. Now, if you look at all the companies in the bottom box, it is every single chip company that is out there today. Every single one of those chip companies is building a software stack. Every single one of those companies is building a software stack that looks something like CUDA. And CUDA is 350 libraries, right? And that is what their entire ecosystem is. And everyone needs to go and replicate this. And then you have all the AI frameworks on top like PyTorch and Jax and TensorFlow and MXNet. And then you've got all the serving engines and you collectively look at all of this. There's about a thousand different permutations of choices you have to make. That's pretty nuts. But the more important and the most nutty thing is those hundreds of libraries that you need to get coverage across all of the models and all of the hardware that exists. It is roughly around 10 to the 21 That is an intractable problem, right? Even if you throw all of Elm compute at this problem, you still will not get coverage today. Now, what drove this? It was three primary shifts. You had the end of Mo's law, which means transistors stop becoming economical and they stop getting faster. That forces you into a world of heterogeneity, right? So in order to make up for the performance gap, you have special function units. You still many of them together into a single package and then you program them as if they were one chip. On top of that, you've got the memory wall, which means data movement is the most costly aspect of a system to the point where it is two to three orders of magnitude more expensive than doing the computation itself. Right? And as you start scaling AI models, that problem only gets worse. In when I started an AI, I could train a state-of-the-art model on two GPUs, right? Two GPUs. It's been 12 years now, and today to train a model, I need 200,000 GPUs, right? That amount of scale is obscene. But the more important part, the software we have was not built for it. The kernel model we have been relying on was designed for programming a single workload on a single chip that was embarrassingly parallel as in you had to be computebound in order to make best use of a kernel. And when I looked at all of this and how performance was coming I came to this conclusion that kernels have basically become the new assembly. Right? If that is the case, the failure is in compilers. And that's where this idea of a postmorte world came and how accelerated software was born. It was accelerated computing came from making certain parts of your workload faster. But when that stops, you need to relegate the important things of performance to the compiler. Right? Right? So the compiler helps you get the most performance out of your chip at whatever scale. Right? So why do I call kernels in new assembly? Firstly, you have multiple different tiers. Back in the 70s, they used to say real programmers write assembly. Today it's real ML engineers write kernels, right? And we write CUDA, we write PTX. And then you've got actually programming every physical part of the system if you want to get performance. But if you do that, you are locked in to that ecosystem. You are locked into that device. Changing anything means you have to do a complete rewrite. And most companies today are taking about 6 months to deploy a model because they're handwriting kernels. And so where do you lose the performance today? Why are kernels an actual problem? They don't help if you're memory bound. Today workloads don't fit in a single chip. You cannot get data in and out of a system fast enough. In fact, if you do an amazing job of writing a kernel, an absolutely tremendous job, you will actually make your hardware look worse. Let this sink in. You will make your hardware look worse with better kernels because you are exposing the latency of your system. The right way to get performance is to overlap communication with computation and try to essentially create this world where you're doing communication optimal linear algebra. But the problem is that is very hard or near impossible for a human to reason about right because you are shaping your code around the data movement. Takes a minute to wrap your head around. So you are spending almost 60% of your time slushing data around and only 15 to 20% of your time actually doing useful work. Most AI companies when I go and talk to them, they tell me they're somewhere between 3 and 45% utilization, right? So for anywhere between 55 to 97% of the time, your GPU is sitting there doing nothing, waiting for data to arrive so it can do something. That's a huge problem because these GPUs are expensive, right? And it's not just one of them. We're talking about hundreds and thousands of them just sitting idle. So what are we going to do about it? We believe this time for something different. So we asked ourselves a question. Could you build a hardware parallel optimizing compiler that can emit communication optimal code for any system at any scale, right? So can you program a system that contains many chips of different topologies as if it were one chip instead of programming each chip in that system individually? That was our challenge. After almost a year and a half of bunking our head against the wall, asking a whole bunch of questions, breaking things, we came up with an answer. We call it Tachion, as in faster than speed of light. That is our goal. Can you take PyTorch to the limit of device physics? So we have an integration with PyTorch. So you can take that model. We do a whole bunch of graph optimizations. We will emit code for the system you are targeting from any vendor and then feed that into our runtime which manages the execution. But it does fine grain management of the data movement to stage data to optimally use the compute units. So the goal is to keep all the compute units on the system as used as possible as frequently as possible. And we can do this. What happened here? >> Yeah. So we can do this because the compiler itself doesn't really care about the hardware. It queries this hardware parameters that so it understands what are the feeds and speeds of the system. What are the core count? What does the uncor look like? How many of these systems do I have? And then it uses that to do all the partitioning, all the fusion, all the loop enrolling and vectorizations and pick the right set of instructions that it can pass on. On early tests, we've been able to outperform at the kernel level some vendor software stacks by anywhere between 1.4 to 2x. Now, I do have to put a disclaimer. It's why it's in the bottom. If you do use the stack, a few things are likely to happen. You will get more performance. You will start hating kernel writing and you will realize that you've been gaslit all these years. So that is what we've been up to. A single abstraction that allows you to program any system at any scale for training and inference. So what does it do differently? You have a DSL that allows you to have fine grain control of the system and get the full performance of the compiler runtime. You have a graph IR that is explicitly parallel so it can capture more information than compilers normally do. You've got a graph optimizer that uses the hardware information to do all the different kind of optimizations to get you the most performance out of your cluster. And the code generator emits that object code specific to that device for the runtime to then execute for you. So it is coming soon to a cloud near you in 26 in the summer. If you want to get early access, come find us at our booth. We would love to sign you up and let you see what's actually possible. Thank you. [Applause]