Mojo ๐ฅ - Is it Python's faster cousin or just hype?
Background & Motivation
The introduction of Mojo by Chris Lattner captured the attention of the Python community with the allure of dramatic performance enhancements and a syntax that would not alienate current Python developers. As Mojo progresses beyond its infancy, it's critical to assess its evolution and its capacity to disrupt the programming ecosystem, particularly within artificial intelligence and machine learning domains.
Objective & Scope
This presentation will share findings from an AI Safety Camp project which used Mojo to build a Large Language Model Mechanistic Interpretatability and Activation Engineering library. Through our exploration, we aim to provide a candid narrative of Mojo's strengths and limitations, judge its performance claims, and probe its likelihood of adoption for AI development.
Content Overview
Introduction to Mojo: Brief overview of Mojo's conception, ethos, and intended use-cases. Performance Claims: An further look at the purported 68,000x speed increase over Python, including benchmark comparisons and real-world application data. Language Design: An analysis of Mojo's syntax and semantics, drawing parallels and contrasts with Python, and the implications for developers transitioning to or adopting Mojo. Case Study: Detailed account of the process of writing a Large Language Model Interpretation library in Mojo, highlighting the challenges and breakthroughs experienced. Ecosystem Overview: Examination of the current state of Mojo's ecosystem, its community support, and the availability of tooling and libraries. Discussion: Engaging the audience in a discussion about Mojo's potential future, its fit within existing projects, and the propensity for it to become the primary language for AI development.
Conclusion We'll wrap up with predictions for Mojo's trajectory based on our experiences and broader industry trends, potentially setting the stage for Mojo to capture the "Mojo" it needs to triumph or to become a footnote in the annals of programming language history.
This session took place in track Python Language & Ecosystem and was classified suitable for intermediate domain / intermediate python by the speaker.
Transcript (auto)
Auto-generated from the recording utilizing Open-Source AI. Speaker labels (Speaker 1, Speaker 2) reflect diarization, not identity. Timestamps refer to the recording.
Speaker 1 [00:05]
I'd love to get an answer to you now. Good morning, Berlin. Let's jump right in. If you go to this website, the Mojo Playground, and you type in the code PRINTHELLOMOJO, the console will output Hello Mojo. Similarly, if you go to the Playground and you download the file, you'll get a quick Hello Mojo message. It's very easy to get started and jump into Mojo, but should you? This is what this talk is aiming to find out. During our talk today, we're going to cover who and what is Mojo, and then we're going to be looking at three key factors to help determine whether this is Python's faster cousin or just hype. We'll be looking at the performance, the usability, and the governance and community around the Mojo language. Lastly, I'll finish with some conclusions and a little call to action for our community here today. This is the Playground. Getting started with Mojo is faster than ever. You're welcome, if you want to, to pull this up on your phone and just start playing around during the talk so that you can really get an introductory, beginning feel for how this language works. Okay. So you can see that print Hello Mojo is just one line of code. We've got a def function here. It's exactly the same. Who or what is Mojo? Let's start from the beginning. Mojo is a systems level language designed to have very similar syntax to Python. Developed by the highly experienced team of compiler engineers over at Modular. Now you can see that version one was released a little bit over a year ago and there have been eight minor versions later. They switched in 2024 in January to using this year-month-based naming system. AI development suffers heavily from fragmentation at every layer of the stack. It requires domain expertise in C++, Python, and hardware-specific code such as Assembly and CUDA. now modular's vision was to simplify the portability of code by developing a common programming model across ai hardware to fully unlock low-level systems programming and high level application development merger is the first programming language targeted at ai developers when it was released in may 2023 and it's built atop the next generation of compiler technology This is called the Mid-Level Intermediate Representation, or MLIR. And they built an entirely new end-to-end compiler for this. This is why Mojo is trying to become the programming language for all AI developers. Now, within Modular, they have Mojo, the programming language, and they have Max. Max has two parts to it, the Max Engine and Max Serving. Max completes Mojo's AI story. It stands for Modular Accelerated Execution, and it has a commercially licensed Enterprise Edition in addition to the Developer Edition. The idea behind it, they're trying to get the Enterprise Edition to allow testing and deployment of the same Max components from developer edition in a production environment. They're looking to support production services such as EC2, Elastic Container Service, and SageMaker on AWS, and as a result, they've signed a deal with AWS to ensure that they can do that. Now, Modular Max Engine is a high-performance AI compiler and runtime that's designed to optimize and accelerate AI models for inference. It's built atop the MLIR compiler infrastructure, and it has graph-level optimizations to improve the performance and the efficiency of AI models. What does this mean? It means that if you've trained an AI model, then you can use MaxEngine to run these models very quickly. As you can see, it's got input from three different kinds of models, TensorFlow, PyTorch, and Onyx format. and then you can take these uh very fast inference uh pods from max engine and deploy them as microservices uh so this fills out mojo's product line as i understand it and i see it i'm thinking about it in terms of um mojo is the core language and then max parts of this is going to be open sourced and parts of this is going to say commercial it feels a bit like uh chris latner and team are trying to head as fast as they can towards getting working AI that can be deployed. Okay. So then, who is Mojo? We have Chris Lattner, the founder and CEO. The language is being developed by the team at Modular, and Chris Lattner is one of the co-founders, and he has experience across big tech, leading teams at Apple, Google, and Tesla. In the open, he's been a core contributor for Clang, OpenCL, LLVM, and Swift. He's worked at a variety of different companies, including the development of MLIR, Circuit, TPUs, and Mojo. Now, when Chris Lightner built Swift, he built the successor to Objective-C on Apple. It adds a bunch of features, and it did change a variety of syntax, adding new concurrency models. He was quite effective at the time at getting a migration across from Objective-C to Swift, so although there were inevitably going to be some people who stayed with the old language many moved okay and then lastly his team modular what have they done well they're well funded now we saw yesterday in the keynote um the building an open source orchard you have a variety of different organizations uh who they can choose to take vc money or not so crunchbase lists them as a company of 11 to 50 and last year they took 100 million dollars of venture capital money looking at their careers page at modular you can see the kind of priorities they've got so they're looking at ai compiler engineering cpu performance gpu performance improving the runtime and the cloud services there across the company to be very clear i'm not affiliated with modular i'm just doing this to demonstrate the priority and what they're trying to build next okay so that gives you a nice little grounding to start you off what is mojo what is modular what is max who is chris latner and the team building this new faster cousin of python right so first off mojo is very fast i'll demonstrate this with three examples this first example is a matrix multiplication optimized example you can see that it's got this big bumper headline figure here which is talking about being 68 000 times faster than python so there's a jupyter notebook here that they walk through and they talk about the difference between naive code which becomes vectorized code which becomes parallelized code which gets tiled and then becomes unrolled now the naive matrix multiplication performs a matmul using nested loops without any optimizations at all the vectorized approach uses what's called simd which is single instruction multiple data and these instructions can perform multiple operations simultaneously across a cpu you can parallelize it across threads and processes to give you parallel execution. You can see that massively increases the speed from 6 gigaflops to 19. And then this is further accelerated by tiling, which is dividing the matrix into smaller tiles. Now, that has a property that it improves the cache locality and it reduces your memory accesses. Lastly, you can unroll it. This will replicate the loop body multiple times and it will reduce the loop overhead, giving you better instruction level parallelism. So this is their bumper headline figure provided by one of the compiler engineers, and again, this can be accessed from the Mojo Playground. What about real code, though? Recently, there was a series of blog posts, two blog posts, in fact, between driving the conversation forward between Mojo and Rustations and the Rust community. The key example they used was this DNA benchmark, which was looking at passing long strings of DNA. Now, when you have long arrays and the process doesn't actually vectorize very easily, you can find that tools like NumPy don't accelerate very well on these kind of tools because you can't have single instruction run across multiple data. the the intermediate outputs require the entire array up to this point so far this dna benchmark has been prepared by using two tools you can see they're called neo fetch and hyperfine and i genuinely would recommend hyperfine actually it's it's really excellent at giving you some good benchmarks and it has a variety of features there to ensure that you're creating these standardized results the mojo community as well have blog posts talking about best practice when you're doing your benchmarks making sure you've got a fair test as an ex scientist I'm sure I could go on to you at length about exactly how to create this fair test but basically you imagine yourself with the skill of six months in each programming language at least and then you do everything in the idiomatic way to make sure it's fast and you assume you have the same features across your processor the same features across your hardware the same compiler flags and the same kind of code implementation as well. And then you only measure the interior part that you care about rather than relying on screen IO, which can be very slow and highly dependent. This stops you from having a fair test. Okay. So you've put in all this effort to create an excellently fair test. How fast is it now? And so the result that they obtained from the blog post is that mojo fast trim is a hundred times faster than python uh which feels like a relatively um straightforward goal to try and achieve when you're comparing performance beating python is one of the first um steps on the ladder next up is it faster than rust well in this one particular case yes now the reason this got a lot of pushback from the rust community they were saying they were doing things uh wrong there were there were bits and bobs they could do to improve it and this started a really productive conversation between the developers of Mojo and Rust where they were talking through the differences. In particular, because Rust was made in 2010 by Graydon Hall, they are built off LLVM. Because Mojo was built off MLIR, this is this 2019 technology developed at Google, again by Chris Lackner, and this tech has wide tutorials everywhere. It's a compiler that's built from the ground up to do AI tasks quickly. And then lastly, you have this industry-standard bioinformatic tool called CutAdapt, and they were still able to make it faster than this. So you can see that for a variety of benchmarks, Mojo is very fast. Here's one I made earlier. So this is talking about language models, and we're looking not at different programming languages, but different ways to serve it with Python as your base. So along the X, you can see that we're comparing PyTorch, PyTorch compile, a library called nn-site, and our Mojo implementation. And we're just looking at how quickly can you generate 20 new tokens of the tiny stories model, and you can see that Mojo is much faster. So this is relatively uncontroversial. On CPU, Mojo seems faster. Indeed, GPU supports on the roadmap for Mac's engine. Modular have got a partnership with NVIDIA to bring NVIDIA GPUs to the Mac's platform in time. Now, a really key consideration here, if you're running AI models, you're not running them on CPU. So right now, Mojo's performance is fast on CPU. That's great. but until we get this GPU support implemented and done, we can't actually fulfill the core promise of Mojo. This couldn't be more important. So performance, real or hype? We've got our real or hype leaderboard here. We can see that performance, yes, it's definitely faster on CPU. All with me so far? Brilliant. Let's look at usability now. There's this quote from Russ Streamer. He says, if Mojo is legitimate, I think Mojo will win, hands down. And the reason it will win is that you do not change the paradigm of any acclimated or proficient individual. You just have to learn a bit more, and you get amazing performance. If Mojo compiles fast, and it can be ahead of time compiled, AOT, or it can be jitted, just in time compiled, and it looks like a language you're already familiar with, it's really close to being the same speed, I don't see how you're going to make that cell. Now, Mojo, when they say you just have to learn a little bit more, unfortunately it's bigger than the eight minutes I have to try and teach you the entire new language here. So you can see that here's from the language reference. You've got the introduction to Mojo. You've got value ownership. You've got value lifecycle, traits and parameters. So you're working with structs and traits. In that regard, it's quite similar to Rust. you have this fn function syntax in that regard quite similar to rust and then you have deeper python integration as well so in order to show you what i can show you in the time i've got i prepared a small exercise talking about glimpses of mojo now this is using the advent of code solutions advent of code is an advent calendar of small programming puzzles and a variety of skill sets and skill levels that can be solved in any programming language you like. Here we've chosen Mojo so we can explain Mojo. People use them as interview prep, company training, course work, practice problems and a speed contest. So let's get a feel for the syntax and let's get a feel for some of the unique features of Mojo. Here's the day one code. Let's see what we can pull out here. We can see the function signature uses FN rather than death. We can see that the fn declaration enforces type checking and memory safe behaviors like in rust whereas def does not do this neither type declarations nor dynamic behaviors when using fn the function name is allowed by parameters in the return type next up we have variables in mojo you have two types of variables in this code which are mutable variables defined with var and immutable variables designed defined with let note that uh in working with the community chris latina actually removed this feature in december so that um there is no let keyword anymore there is no let mutt coming from rust so there's no way sorry there's no let coming from rust there's no way to make immutable variables and um it's still possible to have immutability as a feature but this is encoded within the function parameters and within the pointers rather than having it as a core keyword syntax so this is demonstrating mojo's commitment to become a python superset next up we have collections and generics they're originally that was called dynamic vector back in i believe the 2024 february release they changed this to becoming list to becoming more pythonic you can see how this interface is changing over time. There are some differences. There are some similarities. We have structs and value semantics. Now, merger introduces the concept of structs, which are value types that have multiple fields. The at value decorator can then define structs such as part, workflow item, workflow range, and state. The value types have value semantics. This means arguments are assigned and parsed as arguments, ensuring data immutability. Lastly, you have generics, and the Mojo team are working towards having a structural typing system rather than a nominal typing system and using algebraic data types, for those in the audience familiar with this. Usability. How much has changed from the core Python? See for yourself whether you like the syntax. Unfortunately, this language is far too big for this margin to contain. Despite the change in date style versioning, the latest version, 2024.2, is still very much 0.x, still very much pre-version 1. The interface moves regularly. Lastly, we're talking about community and governance. So Mojo is a hybrid community, which means that it's an open source language being created by a closed source company. And importantly, this product is a key part of their value proposition. This isn't like Google opening Go because this is not part of their core business product. Mojo's open source journey started off as an initial closed invite-only beta. The current state of their open source journey kind of developed and opened as they went. So we've now got to a point where the open source components such as the standard lib and parts of the max engine are open source, and proprietary components are being developed by modular. This is a really key part where the community can step in here because they're still on their journey. They don't know what to open source and what to close. They're crossing the river by feeling the stones. And so here we can assist with this part. Here we have a graph of the number of GitHub stars over time. This is showing you they had a big launch in May of 2023. And then with additional incremental updates and news coverage, they've been growing over time to a point when they're a little bit over 20,000 stars for the main library. And that tiny blue sliver you see over near March is the release of the Max engine, which is somewhat behind. Now, within the community, there's a bunch of things to point out. The style guide for Mojo, up until perhaps two weeks ago, they had blank case and Pascal case in places where I think a Pythonista might not put them there. Now, these have been changed to snake case recently. Be aware there are some issues where people seem to be falsely claiming that things are pythonic when they are not. So it's really key to make sure that we're upholding the best of the community here. I'll skip this point about governance here. The Mojo vision differs from the Zen of Python. Instead of make it fast, make it work, make it fast, they are doing it a different way so they're focusing on performance by default they're prioritizing performance and then safety and then portability and then debuggability in that order there's a bunch of interesting and strange things that as you might expect when you're coming to a new language so in some sense the development philosophy if you're working on the standard library here is going to be a bit different in order to um uh the philosophy is going to be different here in order to best result with that this is my favorite slide of the whole talk will this partially open software be closed or re-licensed now probably not but look at the words that chris latner says in interviews he's previously talked about being a monk on the hill which is not great for community engagement no doubt but uh does probably serve to kind of show you how he's thinking about this new language he's building look at his actions and his cv i've got copied it down there and then look at his current actions and the way he's currently integrating with the community community and governance so i've got a little scales down here and these scales are showing you the um if you're an optimist you think it's going to be fine if you're a pessimist you think you might try and close it down later on uh i think it's probably going to be good and it's looking good so far from the information we've got. Lastly, quickly conclusions. Mojo is one to watch and probably not just hype. Ensuring the Python community engages with the development of Mojo will ensure that this Python superset continues to grow to become the fastest Python yet. Thank you very much.
Speaker 2 [21:54]
Thank you so much for the interesting talk, Jamie. We have a lot of questions here from the audience. The first one is, does Mojo have compatibility with popular frameworks like PyTorch and TensorFlow?
Speaker 1 [22:05]
Yes, so the core Mojo language doesn't yet, and indeed, in their vision document, they talk about not building and enhancing the tensor data type. They are building a machine learning library called Basalt, Basalt in American, and this is going to be their way of doing machine learning going forward. Note that the Mac's engine can do inference using PyTorch models, using TensorFlow saved models, and with ONNX format as well. So, it's not good for training yet, but it will be good for inference quite soon.
Speaker 2 [22:39]
Thank you. And how does matrix multiplication in Mojo compare to using NumPy in Python?
Speaker 1 [22:46]
so their headline figure of 68,000 times faster I believe that's not comparing it to base Python I believe they're comparing that to NumPy if I'm not mistaken I'm not sure that's right actually so it's faster but I can't tell you how much faster currently sorry
Speaker 2 [23:09]
Is there anything interesting you have written in Mojo and what was your experience like? What is your favorite, least favorite feature?
Speaker 1 [23:16]
Yes, absolutely. So this benchmark comes from, I've cut this out for time from the talk,
Speaker 2 [23:16]
Yes.
Speaker 1 [23:22]
but I was building a mechanistic interpretability library. This is used to understand what's going on inside large language models. So we were using what's called linear probes to look at the intermediate activation state. And we were also using patch scopes. Now, when I started building stuff in Mojo, I found it quite nice to work with overall. I liked some of the, I mean, my favorite feature is the performance, of course. But otherwise, I found working with it quite reasonable, and I found it not too bad so far.
Speaker 2 [23:55]
Thank you. And is the performance because it's a good language or a good implementation or would a version of Rust based on MIR be competitive?
Speaker 1 [24:04]
Great question. This is my favorite question of the whole talk so far. Yes, MLIR seems great. There are loads of tutorials online If you want to avoid the risk of a proprietary company building your open source language Rust with MLIR sounds like a good idea. Julia with MLIR sounds like a good idea. We should get this technology into our code ASAP
Speaker 2 [24:24]
Can we mix mojo and python code in the same file? Yeah
Speaker 1 [24:27]
file yes uh so there's a bit about interop here and so you have this python.from capital python import python you've got python.import module and then you can import your matplotlibs your numpy's and you can have all the data science goodness of the python community and the pydata community of course um inside major
Speaker 2 [24:50]
Are there any disadvantages of Modro as compared to Python or Rust code?
Speaker 1 [24:57]
so it's a different language philosophically it thinks and works differently so this is not the zen of python thinking and coding in this language is different and a bit fiddlier um personally i i still admit i kind of start off by writing python and then converting it in my head rather than idiomatically writing mojo straight away um so if you want developer velocity if you want to build stuff quickly yes keep using python yes this language is one year old it's still immature sure, still want to watch, but Python is still here for a reason, you know? It's sticking around.
Speaker 2 [25:34]
And is it only for data machine learning or do you see, say, the possibility to speed up good old Django?
Speaker 1 [25:41]
Absolutely. There's Lightbug, which is a, Lightbug does web routing. It's, yes, I think Mojo can speed up Django.
Speaker 2 [25:57]
Thank you, and then another question is what can Mojo give me what I cannot achieve with a framework like PyTorch?
Speaker 1 [26:05]
So, the core hypothesis they were de-risking, can they use next generation compiler technology to give you kernels into specific hardware, your CPU, your GPU, your TPU at Google, which is where this whole thing was built, your NPU, your IPU. and it seems to be that despite aiming at a much broader compilation target series they are still outperforming them and that's that's truly quite surprising to me
Speaker 2 [26:36]
Thank you so much, Jamie Coombs.