Mojo 可能是几十年来最重要的编程语言进步

Mojo 是一种基于 Python 的新型编程语言，它解决了 Python 的性能和部署问题。
技术
作者

Jeremy Howard

发布

2023 年 5 月 4 日

我记得第一次使用 Visual Basic v1.0 的情景。当时，它是一个用于 DOS 的程序。在此之前，编写程序极其复杂，我从未能在最基本的玩具应用程序之外取得多少进展。但有了 VB，我在屏幕上画了一个按钮，输入了一行代码，让它在按钮被点击时运行，然后我就有了一个现在可以运行的完整应用程序。那是一次如此令人惊叹的体验，我永远不会忘记那种感觉。

感觉编码从此将不再一样。

使用 Mojo 编写代码——这是 Modular1 推出的一种新编程语言——是我人生中第二次有这种感觉。它看起来是这样的

为什么不直接使用 Python？

在我解释为何我对 Mojo 如此兴奋之前，我首先需要谈谈 Python。

Python 是我过去几年几乎所有工作所使用的语言。它是一门优美的语言。它有一个优雅的核心，一切都建立在其之上。这种方法意味着 Python 可以（而且确实）做任何事情。但这带来了一个缺点：性能。

这里或那里的几个百分点无关紧要。但 Python 比 C++ 等语言慢了许多几千倍。这使得将 Python 用于对性能敏感的代码部分（即性能至关重要的内循环）变得不切实际。

然而，Python 有一个锦囊妙计：它可以调用用快速语言编写的代码。因此 Python 程序员学会了避免使用 Python 来实现性能关键的部分，而是使用 Python 对 C、FORTRAN、Rust 等代码进行封装。Numpy 和 PyTorch 等库为高性能代码提供了“Pythonic”接口，让 Python 程序员即使在使用高度优化的数值库时也能得心应手。

得益于其灵活优雅的编程语言、出色的工具和生态系统以及高性能的编译库，如今几乎所有 AI 模型都在 Python 中开发。

但这种“双语言”方法有严重的缺点。例如，AI 模型通常必须从 Python 转换为更快的实现，例如 ONNX 或 torchscript。但这些部署方法无法支持 Python 的所有功能，因此 Python 程序员必须学习使用与其部署目标相匹配的语言子集。对代码的部署版本进行性能分析或调试非常困难，而且不能保证它会与 Python 版本完全相同地运行。

双语言问题阻碍了学习。运行时，你无法单步进入算法实现，或跳转到感兴趣的方法定义，而是发现自己深陷 C 库和二进制代码的细节中。所有程序员都是学习者（或者至少应该是），因为该领域不断发展，没有人能完全理解所有内容。因此，学习困难和经验丰富的开发者遇到的问题与初学者一样多。

在尝试调试代码或查找和解决性能问题时，也会出现同样的问题。双语言问题意味着 Python 程序员熟悉的工具在我们跳转到后端实现语言时就不再适用了。

即使使用更快的编译实现语言作为库，也存在一些不可避免的性能问题。一个主要问题是缺乏“融合”（fusion）——即连续调用一系列编译函数会产生大量开销，因为数据需要在 Python 格式之间来回转换，并且必须支付反复从 Python 切换到 C 再切换回来的成本。因此，我们不得不为常见的函数组合（例如神经网络中的线性层后跟一个 ReLU 层）编写特殊的“融合”版本，并从 Python 中调用这些融合版本。这意味着需要实现和记住更多的库函数，而且如果你做的任何事情哪怕只是一点点非标准，你就很不走运了，因为不会有为你准备的融合版本。

我们还必须应对 Python 中缺乏有效的并行处理的问题。如今，我们都拥有多核计算机，但 Python 通常一次只使用一个核心。有一些笨拙的方法可以编写使用多个核心的并行代码，但它们要么必须在完全独立的内存上工作（并且启动开销很大），要么必须轮流访问内存（可怕的“全局解释器锁”，它常常使得并行代码实际上比单线程代码还要慢！)

PyTorch 等库一直在开发越来越巧妙的方法来解决这些性能问题，新发布的 PyTorch 2 甚至包含一个 compile() 函数，它使用复杂的编译后端来创建 Python 代码的高性能实现。然而，这样的功能并不能变魔术：基于 Python 本身的语言设计方式，它能实现的功能存在根本性的限制。

你可能会认为在实践中，AI 模型只有少量的构建块，因此我们是否必须用 C 实现它们并不重要。而且，它们总体上都是相当基础的算法，对吧？例如，Transformer 模型几乎完全由两部分的多层实现：多层感知器 (MLP) 和注意力机制，使用 PyTorch 只需几行 Python 代码即可实现。这是一个 MLP 的实现

nn.Sequential(nn.Linear(ni,nh), nn.GELU(), nn.LayerNorm(nh), nn.Linear(nh,ni))

…这是一个自注意力层

def forward(self, x):
    x = self.qkv(self.norm(x))
    x = rearrange(x, 'n s (h d) -> (n h) s d', h=self.nheads)
    q,k,v = torch.chunk(x, 3, dim=-1)
    s = (q@k.transpose(1,2))/self.scale
    x = s.softmax(dim=-1)@v
    x = rearrange(x, '(n h) s d -> n s (h d)', h=self.nheads)
    return self.proj(x)

但这隐藏了一个事实，即这些操作的实际实现要复杂得多。例如，请查看 CUDA C 中经过内存优化的“Flash Attention”实现。它也隐藏了通过这些通用方法构建模型会损失大量性能的事实。例如，“块稀疏”（block sparse）方法可以显著提高速度和内存使用率。研究人员正在对常见架构的几乎每个部分进行调整，并提出新的架构（以及 SGD 优化器、数据增强方法等）——我们离拥有一个每个人都会永远使用的整洁打包系统还差得很远。

实际上，当今用于语言模型的许多最快代码都是用 C 和 C++ 编写的。例如，Fabrice Bellard 的 TextSynth 和 Georgi Gerganov 的 ggml 都使用 C，因此能够充分利用完全编译语言的性能优势。

Mojo 登场

Chris Lattner 负责创建了我们今天都依赖的许多项目——尽管我们可能甚至没听说过他构建的所有东西！作为他博士论文的一部分，他开始了 LLVM 的开发，这从根本上改变了编译器的创建方式，如今构成了世界上许多最广泛使用的语言生态系统的基础。然后他继续推出了 Clang，一个基于 LLVM 的 C 和 C++ 编译器，被世界上大多数重要的软件开发人员使用（包括为 Google 的性能关键代码提供支持）。LLVM 包含一个“中间表示”（IR），这是一种专门为机器读写（而不是为人）设计的语言，它使庞大的软件社区能够协同工作，在更广泛的硬件范围内提供更好的编程语言功能。

然而，Chris 发现 C 和 C++ 并未能真正充分利用 LLVM 的力量，因此他在苹果工作期间设计了一种新语言，名为“Swift”，他将其描述为“LLVM 的语法糖”。Swift 已成为世界上最广泛使用的编程语言之一，尤其因为它现在是为 iPhone、iPad、MacOS 和 Apple TV 创建 iOS 应用的主要方式。

不幸的是，苹果对 Swift 的控制意味着它未能在封闭的苹果世界之外真正大放异彩。Chris 在 Google 曾领导一个团队，试图将 Swift 移出其苹果舒适区，成为 AI 模型开发中 Python 的替代品。我对这个项目非常兴奋，但遗憾的是，它没有得到苹果或 Google 的必要支持，最终未能成功。

话虽如此，Chris 在 Google 期间确实开发了另一个非常成功的项目：MLIR。MLIR 是 LLVM 的 IR 在多核计算和 AI 工作负载现代时代的替代品。它对于充分利用 GPU、TPU 以及越来越多添加到服务器级 CPU 中的矢量单元等硬件的能力至关重要。

那么，如果 Swift 是“LLVM 的语法糖”，那么“MLIR 的语法糖”是什么？答案是：Mojo！Mojo 是一种全新的语言，旨在充分利用 MLIR。而且 Mojo 就是 Python。

等等，什么？

好吧，让我解释一下。或许更准确地说，Mojo 是 Python++。它（完成后）将是 Python 语言的严格超集。但它也具有附加功能，因此我们可以编写利用现代加速器的高性能代码。

在我看来，Mojo 比 Swift 更务实。Swift 是一种全新的语言，包含了基于最新编程语言设计研究的各种酷炫功能，而 Mojo 的核心只是 Python。这似乎很明智，不仅因为数百万程序员已经非常了解 Python，而且因为经过数十年的使用，它的能力和局限性现在已得到很好的理解。依赖最新的编程语言研究非常酷，但这是一种潜在危险的猜测，因为你永远不知道结果会怎样。（我个人承认，例如，我经常被 Swift 强大但古怪的类型系统搞糊涂，有时甚至能把 Swift 编译器搞崩溃！)

Mojo 的一个关键技巧是，作为开发者，你可以随时选择进入一种更快的“模式”，只需使用“fn”而不是“def”来创建函数。在这种模式下，你必须准确声明每个变量的类型，因此 Mojo 可以生成优化的机器代码来实现你的函数。此外，如果你使用“struct”而不是“class”，你的属性将紧密地打包到内存中，这样它们甚至可以在数据结构中使用而无需四处追踪指针。这些是使 C 等语言如此快速的功能，现在 Python 程序员也可以访问这些功能了——只需学习一点点新语法即可。

这怎么可能？

截至目前，几十年来有数百次尝试创建简洁、灵活、快速、实用且易于使用的编程语言——但收效甚微。但不知何故，Modular 似乎做到了。这怎么可能呢？我们可以提出几个假设

  1. Mojo 实际上并未实现这些目标，而那些花哨的演示隐藏了令人失望的实际性能，或者
  2. Modular 是一家拥有数百名开发人员、工作多年的大型公司，投入了更多时间来取得前所未有的成就。

这些都不是真的。事实上，这个演示是在我录制视频前几天才创建的。我们给出的两个例子（matmul 和 mandelbrot）并非经过仔细选择，作为尝试了几十种方法后唯一恰好能用的东西；相反，它们是我们为演示尝试的仅有的东西，而且它们第一次就奏效了！尽管在早期阶段还有许多缺失的功能（Mojo 尚未公开发布，除了在线“playground”），但你看到的演示确实像你看到的那样工作。而且，你现在确实可以在 playground 中自己运行它。

Modular 是一家相当小的初创公司，成立才一年，而且公司只有一个部门在开发 Mojo 语言。Mojo 的开发也是最近才开始的。这是一个小团队，工作时间不长，那他们是怎么取得这么多成就的呢？

关键在于 Mojo 建立在一些非常强大的基础之上。我见过的软件项目中，很少有足够的时间来构建正确的基础，结果往往会积累大量的技术债。随着时间的推移，添加新功能和修复错误变得越来越困难。然而，在一个精心设计的系统中，每一个新功能都比上一个更容易添加，速度更快，且错误更少，因为每个功能所依赖的基础都在不断改进。Mojo 就是一个精心设计的系统。

其核心是 MLIR，它已经开发了很多年，最初由 Chris Lattner 在 Google 启动。他认识到“AI 时代编程语言”所需的核心基础是什么，并专注于构建它们。MLIR 是关键的一环。正如 LLVM 在过去十年中极大地简化了强大新编程语言（如 Rust、Julia 和 Swift，它们都基于 LLVM）的开发一样，MLIR 为基于其构建的语言提供了更强大的核心。

Mojo 快速开发的另一个关键推动因素是决定使用 Python 作为语法。开发和迭代语法是语言开发中最容易出错、最复杂、最具争议的部分之一。通过简单地将这一部分外包给一种现有语言（它恰好也是当今使用最广泛的语言），整个问题就消失了！然后在 Python 之上所需的相对少量的新语法很大程度上都能自然地融入，因为基础已经到位。

下一步是创建一个最小的 Pythonic 方式来直接调用 MLIR。这根本不是什么大工程，但却是构建 Mojo 的所有基础——以及直接在Mojo 中进行所有其他工作所需的全部。这意味着 Mojo 的开发者们几乎从一开始就可以在编写 Mojo 时“狗食”（dog-food）Mojo。每当他们在开发 Mojo 时发现有什么不太好用的地方，他们就可以在 Mojo 本身中添加所需的功能，从而使他们更容易开发 Mojo 的下一个部分！

这与 Julia 非常相似，Julia 是在一个最小的类 LISP 核心上开发的，该核心提供 Julia 语言元素，然后这些元素与基本的 LLVM 操作绑定。Julia 中的几乎所有内容都是在此基础上构建的，使用 Julia 本身。

我无法开始描述 Mojo 设计和实现中的所有小（和大！）想法——这是 Chris 和他的团队数十年编译器和语言设计工作的成果，包含了那段时间所有的技巧和来之不易的经验——但我可以描述的是我亲眼见证的惊人结果。

Modular 团队内部宣布，他们决定通过一个视频（包括一个演示）来发布 Mojo——并设定了一个仅在几周后的日期。但当时 Mojo 还只是一个最基本的语言。没有可用的 Notebook 内核，几乎没有实现任何 Python 语法，也没有进行任何优化。我无法理解他们如何希望在几周内实现所有这些——更不用说把它做好！在这段时间里我看到的情况令人震惊。每隔一两天，全新的语言功能就会被实现，一旦有足够的准备可以尝试运行算法，通常它们立刻就能达到或接近最先进的性能水平！我意识到，这一切之所以发生，是因为所有基础都已到位，而且它们被明确设计用于构建当前正在开发的东西。所以，一切都能正常工作并且运行良好也就不足为奇了——毕竟，这从一开始就是他们的计划！

这是对 Mojo 未来保持乐观的理由。尽管这个项目尚处于早期阶段，但根据我过去几周的观察，我的猜测是它将比我们大多数人预期的发展得更快、更远……

部署

我把最让我兴奋的一点留到了最后：部署。目前，如果你想把你的很酷的 Python 程序送给朋友，那么你将不得不告诉他们首先安装 Python！或者，你可以给他们一个巨大的文件，里面包含了整个 Python 以及你使用的所有库，打包在一起，当他们运行你的程序时会被解压和加载。

由于 Python 是一种解释型语言，你的程序行为将取决于安装的 Python 的确切版本、存在哪些库的哪些版本以及这一切是如何配置的。为了避免这种维护噩梦，Python 社区选择了几种安装 Python 应用程序的方案：环境，它为每个程序提供独立的 Python 安装；或容器，它为每个应用程序设置了几乎整个操作系统。这两种方法都会在开发和部署 Python 应用程序时带来大量困惑和开销。

将此与部署静态编译的 C 应用程序进行比较：你可以字面上将编译好的程序直接提供下载。它的大小可能只有 100k 左右，并且会快速启动和运行。

还有 Go 采用的方法，它无法像 C 一样生成小型应用程序，而是将一个“运行时”（runtime）整合到每个打包好的应用程序中。这种方法是 Python 和 C 之间的折衷，仍然需要几十兆字节的二进制文件，但提供了比 Python 更容易的部署方式。

作为一种编译型语言，Mojo 的部署情况基本上与 C 相同。例如，一个从头编写的包含 matmul 版本的程序大小约为 100k。

这意味着 Mojo 不仅仅是用于 AI/ML 应用程序的语言。它实际上是一种 Python 版本，允许我们编写快速、小型、易于部署的应用程序，充分利用所有可用的核心和加速器！

Mojo 的替代方案

Mojo 并不是唯一试图解决 Python 性能和部署问题的尝试。在语言方面，Julia 可能是目前最强的替代方案。它拥有 Mojo 的许多优点，并且已经有许多很棒的项目用它构建。Julia 的开发者们非常友好地邀请我在他们最近的会议上发表主题演讲，我借此机会描述了我认为 Julia 当前的缺点（和机遇）

正如这段视频中讨论的，Julia 最大的挑战源于其庞大的运行时，这反过来又源于在该语言中使用垃圾回收的决定。此外，Julia 中使用的多重分派（multi-dispatch）方法是一个相当不寻常的选择，它既为在该语言中做酷炫的事情打开了很多大门，但也可能使开发人员的工作变得相当复杂。（我对此方法非常热衷，甚至构建了一个它的Python 版本——但因此我也特别清楚它的局限性！）

在 Python 中，当前最突出的解决方案可能是 Jax，它有效地使用 Python 创建了一个领域特定语言（DSL）。这种语言的输出是 XLA，这是一个早于 MLIR 的机器学习编译器（我相信它正在逐步移植到 MLIR）。Jax 继承了 Python（例如该语言无法表示结构体、直接分配内存或创建快速循环）和 XLA（主要限于机器学习特定概念且主要针对 TPU）的局限性，但其巨大优点在于它不需要新的语言或新的编译器。

如前所述，还有新的 PyTorch 编译器，以及 TensorFlow 也能生成 XLA 代码。就我个人而言，我发现以这种方式使用 Python 最终并不令人满意。我实际上无法使用 Python 的全部能力，而必须使用与我目标后端兼容的子集。我无法轻易调试和分析编译后的代码，而且有太多的“魔法”在其中，以至于很难知道实际执行的是什么。我甚至最终得不到一个独立的二进制文件，而是必须使用特殊的运行时并处理复杂的 API。（我不是唯一这样想的人——我认识的每一个使用 PyTorch 或 TensorFlow 瞄准边缘设备或优化服务基础设施的人都说这是他们尝试过的最复杂、最令人沮丧的任务之一！而且我不确定我是否认识任何真正使用 Jax 完成了这些事情的人。）

Python 的另一个有趣方向是 Numba 和 Cython。我是这些项目的忠实拥趸，并在我的教学和软件开发中都使用过它们。Numba 使用一个特殊的装饰器，使 Python 函数通过 LLVM 编译成优化的机器代码。Cython 与之类似，但也提供了一种类似 Python 的语言，它具有 Mojo 的一些特性，并将这种 Python 方言转换为 C，然后进行编译。这两种语言都未能解决部署挑战，但它们可以很大程度上帮助解决性能问题。

两者都无法用通用的跨平台代码针对一系列加速器，尽管 Numba 确实提供了一种非常有用的编写 CUDA 代码的方式（因此可以针对 NVIDIA GPU）。

我非常感谢 Numba 和 Cython 的存在，并且从中获益匪浅。然而，它们与使用一种完整的语言和编译器生成独立的二进制文件完全不同。它们是 Python 性能问题的创可贴解决方案，对于只需要这些情况来说很好。

但我更喜欢使用一种既像 Python 一样优雅又像专家编写的 C 一样快的语言，它允许我使用一种语言编写从应用程序服务器到模型架构再到安装程序的一切，并且让我可以直接在我编写代码的语言中调试和分析我的代码。

你想要一种那样的语言吗？

脚注

  1. 我是 Modular 的顾问。↩︎