專為AI而來(lái)的新語(yǔ)言Mojo，推理LLaMA2性能比Python提升250倍！

比C語(yǔ)言也要快上20%。

上周五才開(kāi)放下載，Mojo這么快就自證實(shí)力了。

要知道，之前官方號(hào)稱Mojo可以比Python快6.8萬(wàn)倍。

而且作者表示，其實(shí)還有進(jìn)一步提升的空間。

這也展示了通過(guò)Mojo進(jìn)行硬件優(yōu)化的潛力。

OpenAI創(chuàng)始成員Karpathy已經(jīng)趕來(lái)圍觀了。

目前，LLaMA.mojo已開(kāi)放下載~

幾行代碼即可下載

帶來(lái)這個(gè)版本的老哥是一位前Meta工程師Aydyn Tairov。

他利用Mojo的SIMD（Single Instruction Multiple Data，單指令多數(shù)據(jù)）和向量化原語(yǔ)，將llama2.py轉(zhuǎn)化為Mojo，性能較Python版本提升了近250倍。

即便在快速運(yùn)行模式下，Mojo版本也比C語(yǔ)言版本性能提升15-20%。

不過(guò)作者嘗試了在Mojo中使用并行模式，速度就慢了很多。

作者進(jìn)行性能比較的系統(tǒng)和硬件情況如下：

如果你也想下載運(yùn)行這個(gè)模型，需要先在環(huán)境中安裝配置Mojo（文檔鏈接見(jiàn)文末）。

首先將存儲(chǔ)庫(kù)保存到保存項(xiàng)目時(shí)的文件夾：

git clone https://github.com/tairov/llama2.mojo.git

然后打開(kāi)存儲(chǔ)文件夾：

cd llama2.mojo

下面就可以下載模型：

wget https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin

然后即可運(yùn)行：

mojo llama2.mojo
num hardware threads:  6  SIMD vector width:  8
checkpoint size:  60816028
<s>
Once upon a time, there was a little girl named Lily. She loved to play outside in the sunshine. One day, she saw a big, red ball in the sky. It was the sun! She thought it was so pretty.
Lily wanted to play with the ball, but it was too high up in the sky. She tried to jump and reach it, but she couldn't. Then, she had an idea. She would use a stick to knock the ball down.Lily found a stick and tried to hit the ball. But the stick was too short. She tried again and again, but she couldn't reach it. She felt sad.
Suddenly, a kind man came by and saw Lily. He asked her what was wrong. Lily told him about the ball. The man smiled and said, "I have a useful idea!" He took out a long stick and used it to knock the ball down. Lily was so happy! She thanked the man and they played together in the sunshine.
<s>
Once upon a time, there was a little girl named Lily. She loved to play outside in the sunshine. One day, she saw a big, red
achieved tok/s:  264.24870466321244

為啥Mojo這么快？

不過(guò)話說(shuō)回來(lái)，為啥Mojo的速度可以這么快？

這還得從Mojo的來(lái)歷說(shuō)起。

它誕生于今年5月，專為AI領(lǐng)域開(kāi)發(fā)，由LLVM之父和Swift之父Chris Lattner帶來(lái)。

它兼顧了Python和C++的優(yōu)點(diǎn)，語(yǔ)法簡(jiǎn)單、運(yùn)行快，而且可以和任何Python庫(kù)無(wú)縫交互。

自從上線以來(lái)，Mojo已經(jīng)吸引了12萬(wàn)開(kāi)發(fā)者，GitHub星標(biāo)達(dá)9K。

今年8月，Mojo背后公司Modular新獲1億美元融資，總?cè)谫Y金額達(dá)1.3億美元。

Mojo語(yǔ)言這么快的原因，可以歸結(jié)為4點(diǎn)。

第1步，通過(guò)類型注釋消除Python動(dòng)態(tài)類型的損失，并做代數(shù)簡(jiǎn)化（algebraic simplifications），避免開(kāi)方運(yùn)算以及簡(jiǎn)化復(fù)數(shù)平方運(yùn)算，達(dá)到89倍加速。

第2步，通過(guò)向量化實(shí)現(xiàn)SIMD（單指令多數(shù)據(jù)）的并行計(jì)算，并讓向量寬度以匹配CPU的FMA（浮點(diǎn)乘法累加單元）數(shù)量，達(dá)到874倍。

第3步，把前兩步開(kāi)發(fā)好的單線程實(shí)現(xiàn)改成多核并行化，對(duì)于88核的系統(tǒng)再獲得30倍加速，與原始Python相比已經(jīng)到了26000倍。

第4步，解決并行化中的加載不均衡問(wèn)題，讓線程從池中動(dòng)態(tài)獲取任務(wù)，得到最終結(jié)果68000倍。

幾天前，Mojo正式開(kāi)放下載。目前支持Linux系統(tǒng)，后續(xù)將陸續(xù)添加Mac和Windows。

同時(shí)支持VSCode插件，可以實(shí)現(xiàn)語(yǔ)法高亮和代碼補(bǔ)全等功能。

以及也能像Python一樣在Jupyter里交互式操作。

感興趣的童鞋，可以去上手體驗(yàn)一下~

GitHub地址：https://github.com/tairov/llama2.mojo

Mojo文檔：https://docs.modular.com/mojo/manual/get-started/index.html