17611538698
webmaster@21cto.com

Mojo 刚刚融了1个亿,号称在 AI 负载方面比 Python 快 3.6万倍

编程语言 0 1201 2023-08-30 01:23:22

图片

Modular Mojo 是一种专为 AI 开发者设计的新型编程语言,据该语言的研发团队说它将 Python 的可用性与 C 的性能完美地结合起来。

此外,他们在一段矩阵乘法代码运行负载上的性能是 Python 的 3.6万倍以上。

起初,模块化 Mojo 编程语言并不在Modular AI公司的最初计划中,而是当时公司的创始人Chris Lattner 和 Tim Davis 专注于构建一个平台来统一世界的 ML/AI 基础设施时的附加品。

图片

在开发AI平台过程时,他们意识到跨整个技术堆栈的编程过于复杂,最后也的确出现了这种情况,需要手工编写大量 MLIR(多级中间层表示)。

“加速超过Python 36,000 倍”的说法是这样的:有一个matmul.py 脚本,它在 Python 运行时中执行 128×128 矩阵乘法,吞吐量为 0.00215 GFLOP/s,另一个脚本在 Mojo 中以 79.636 GFLOP/s 执行 512×512 矢量化 + 并行矩阵乘法。

这个说法看起来很有趣,奇怪的是他们用了不同的矩阵大小。

有些投资机构已经准备好在该公司上投入大量资金,有一家称为 Monitor 的机构声称已经筹集了 1 亿美元,并已经打入Modular的帐户。

Twitter/X 上的一条讨论线索显示。

图片

Mojo创始人回复此贴子说:

图片



Mojo 没有对 Python 造成威胁,它不仅提升了 Python 的地位,还为 Python 程序员提供了 “超能力”。

如果真的要说感到恐惧,那应该是 C++ 和其他难以提升速度的语言。

其实正如Sachin Shekhar所说:所有机器学习库都在底层使用 C。

而Python 也只是C语言一个接口。鉴于 C 语言仍在使用,替代这个编程语言并不容易。你需要给予一些特别的东西。很少有人在生产中使用原生 Python 运行超大循环体。

目前开发者还不能重现上面的测试,因为 Mojo SDK在 9 月份才会发布,但是它的技术文档已经可用,并且里面包含一个关于矩阵乘法的文档。

有网友根据文档做了一部分复现。可以看到 Mojo 是 Python 的超集,因此两个模拟脚本是相同的:

%%python
def matmul_python(C, A, B):
for m in range(C.rows):
for k in range(A.cols):
for n in range(C.cols):
C[m, n] += A[m, k] * B[k, n]

他们只是在 Mojo 脚本中将函数名称更改为“matmul_untyped”。以下是 Python 矩阵乘法脚本代码:

import numpy as np
from timeit import timeit
class Matrix:
def __init__(self, value, rows, cols):
self.value = value
self.rows = rows
self.cols = cols


def __getitem__(self, idxs):
return self.value[idxs[0]][idxs[1]]


def __setitem__(self, idxs, value):
self.value[idxs[0]][idxs[1]] = value
def benchmark_matmul_python(M, N, K):
A = Matrix(list(np.random.rand(M, K)), M, K)
B = Matrix(list(np.random.rand(K, N)), K, N)
C = Matrix(list(np.zeros((M, N))), M, N)
secs = timeit(lambda: matmul_python(C, A, B), number=2)/2
gflops = ((2*M*N*K)/secs) / 1e9
print(gflops, "GFLOP/s")
return gflops
python_gflops = benchmark_matmul_python(128, 128, 128).to_float64()

其文档中显示吞吐量为 0.0016717199881536883 GFLOP/s。

下面我们来看 Mojo 脚本,它还导入特定于 Mojo 语言的模块:

from benchmark import Benchmark
from sys.intrinsics import strided_load
from utils.list import VariadicList
from math import div_ceil, min
from memory import memset_zero
from memory.unsafe import DTypePointer
from random import rand, random_float64
from sys.info import simdwidthof
fn matrix_getitem(self: object, i: object) raises -> object:
return self.value[i]
fn matrix_setitem(self: object, i: object, value: object) raises -> object:
self.value[i] = value
return None
fn matrix_append(self: object, value: object) raises -> object:
self.value.append(value)
return None
fn matrix_init(rows: Int, cols: Int) raises -> object:
let value = object([])
return object(
Attr("value", value), Attr("__getitem__", matrix_getitem), Attr("__setitem__", matrix_setitem),
Attr("rows", rows), Attr("cols", cols), Attr("append", matrix_append),
)
def benchmark_matmul_untyped(M: Int, N: Int, K: Int, python_gflops: Float64):
C = matrix_init(M, N)
A = matrix_init(M, K)
B = matrix_init(K, N)
for i in range(M):
c_row = object([])
b_row = object([])
a_row = object([])
for j in range(N):
c_row.append(0.0)
b_row.append(random_float64(-5, 5))
a_row.append(random_float64(-5, 5))
C.append(c_row)
B.append(b_row)
A.append(a_row)
@parameter
fn test_fn():
try:
_ = matmul_untyped(C, A, B)
except:
pass
let secs = Float64(Benchmark().run[test_fn]()) / 1_000_000_000
_ = (A, B, C)
let gflops = ((2*M*N*K)/secs) / 1e9
let speedup : Float64 = gflops / python_gflops
print(gflops, "GFLOP/s, a", speedup.value, "x speedup over Python")
benchmark_matmul_untyped(128, 128, 128, python_gflops)

该脚本输出结果为:

0.029258 GFLOP/s, a 17.501798 x speedup over Python


可以看到,当使用相同的 128 x 128 矩阵时,Mojo 比 Python 快了 17.5 倍,它说的 3.6万 倍似乎有点不诚实……不过,17.5 倍的加速仍然相当不错,这可能是由于Mojo是多线程的原因。

由于 Mojo 是 Python 的超集,因此它在任何 Linux 平台上应该没问题。

目前 Mojo 编程语言仍处于 Beta 尝试阶段,Chris Lattner 表示它在未来几周内将会正式 GA。
Mojo 官网:https://www.modular.com/mojo

作者:场长

评论