SMARTCAMP Engineer Blog

スマートキャンプ株式会社(SMARTCAMP Co., Ltd.)のエンジニアブログです。業務で取り入れた新しい技術や試行錯誤を知見として共有していきます。

Mojo🔥でllama2を実行してPythonと速度比較するモジョよ

バックエンド 機械学習

挨拶

京都開発拠点でインターンをしてるぱんちです(a.k.a 田中 大貴)
拙い文章ですが初めて記事を書かせてもらいました！ 業務でAIの調査などをしておりその過程でMojoでLLMを動作させたので最後までお付き合いください！

初めに

京都開発拠点のインターンではAIの調査、AIを用いた開発を行っています。
デバッグ時はローカルでAIのモデルを動作させることも多い上にPythonなので、高速に動作させる方法を模索しています。
その過程で速度の問題を根本的に解決できそうなプログラミング言語Mojoを見つけたので、LLMを動作させてPythonと速度を比較したのでまとめてみたいと思います。

対象読者

  • Mojoの速度が気になる人
  • Local LLMに触れたいけど重いから実行できない人

実行環境

Intel MacBook Pro 2.3GHzクアッドコアIntelCorei7　メモリ32GB
Docker Ubuntu(イメージ:mcr.microsoft.com/devcontainers/base:jammy)

Mojoとは

MojoはPythonの代替のAIを拡張するプログラミング言語。Pythonと比較して68000倍の速度が出るらしい……。
Pythonと互換がありPythonのライブラリをMojoで使うことができます。ただし、Pythonのライブラリはインタプリで動作するので高速化はしないようです。
型は強く構文などはPythonに似ています、ありがたい。

現状のMojoの導入方法

現状MojoをインストールするにはModulerにアカウント登録、ログインが必要です。
https://www.modular.com/mojo
OSによって導入方法が違いますがIntel Macでは動作しないようなので、今回はDockerで導入します。
DockerでUbuntuのコンテナを立てて必要なものをインストールします。

apt-get install -y apt-transport-https &&
keyring_location=/usr/share/keyrings/modular-installer-archive-keyring.gpg &&
curl -1sLf 'https://dl.modular.com/{hash}/installer/gpg.0E4925737A3895AD.key' |  gpg --dearmor >> ${keyring_location} &&
curl -1sLf 'https://dl.modular.com/{hash}installer/config.deb.txt?distro=debian&codename=wheezy' > /etc/apt/sources.list.d/modular-installer.list &&
apt-get update &&
apt-get install python3.10-venv &&
apt-get install -y modular # Moduler CIのインストール

Moduler CIを認証してMojoをインストールします。

modular auth {token} &&
modular install mojo

最後にパスを通します。

echo 'export MODULAR_HOME="/home/ubuntu/.modular"' >> ~/.bashrc
echo 'export PATH="/home/ubuntu/.modular/pkg/packages.modular.com_mojo/bin:$PATH"' >> ~/.bashrc
source ~/.bashrc

MojoとPythonの実行時間の比較

MojoとPythonの実行時間を試し割り法を使い、素数を一億個まで求めてるコードの実行時間を比較してみます。
参考:https://transparent-to-radiation.blogspot.com/2023/05/20235.html

Pythonのコード

def calc(upper: int):
    prime_numbers = []

    def is_prime_number(n):
        for pn in prime_numbers:
            if pn * pn > n:
                break
            if n % pn == 0:
                return False
        return True

    for n in range(2, upper + 1):
        if is_prime_number(n):
            prime_numbers.append(n)
calc(100_000_000)

Mojoのコード

from utils.vector import InlinedFixedVector
fn calc(upper: Int):
    var prime_numbers = InlinedFixedVector[Int](100_000_000)
    for n in range(2, upper + 1):
        if is_prime_number(n, prime_numbers):
            prime_numbers.append(n)


fn is_prime_number(n: Int, prime_numbers: InlinedFixedVector[Int]) -> Bool:
    for i in range(0, len(prime_numbers)):
        let pn = prime_numbers[i]
        if pn * pn > n:
            return True
        if n % pn == 0:
            return False
    return True


fn main():
  calc(100000000)

結果

実行時間はtimeコマンドのrealの数値です。 | | Python | Mojo | | ---- | ---- | --- | | 実行時間(一回目) | 7m28.549s | 0m12.865s | 実行時間(二回目) | 7m42.270s | 0m12.370s

7分ほどかかっていたものが12秒に短縮されました。 思っていた以上に速い〜！

Local LLMの実行

Hugging Face(モデル,ライブラリ)などは現状インタプリタで動作するためllama2.mojoを実行してみます。 llama2.mojoはMetaのllama2をmojoで実装したものになります。 まずllma2.mojoをgit cloneしてきます。

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

次にモデルをダウロードします。

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

モデルのダウンロードが終われば以下のコマンドで実行できるようになります。

mojo llama2.mojo stories15M.bin -s 100 -n 256 -t 0.5 -i "Mojo is a language"

llama2.py llama2.c llama.mojoの実行時間の比較

llama2は言語ごとに実装されているのでPythonとC,Mojoでの実行時間を比較します。

llama2.py(Python)

実行コマンド

time python3 llama2.py stories15M.bin 0.8 256 "Dream comes true this day"

生成された文章

<s>
Dream comes true this day. Behinda's eyes, there is a boo-boo from a bird. It is happy. It has a big black stuff on its head. It can move in the night and that is when it is full. It can use necks and sticks and touch its wings. It can make sounds and pull its hair. Be careful, and be gentle. Beak to it. Beak to be gentle. Beak's friend is a firefighter. He helps put out fires with a hose. Beak and his dog are safe. Beak is very brave.
<s>
Once upon a time, there was a little girl named Lily. She loved to play games with her friends. One day, they decided to play a game of hide and seek. Lily was very good at hiding and she won the game. 
After the game, Lily's friends noticed that she didn't have many toys to play with. They asked if she could share her toys with them. Lily was happy to share and said yes. 
But then, Lily's mom came and told her that it was important to share and be kind to others. Lily
achieved tok/s: 0.6090875867949811

llama2.c(C)

実行コマンド

time ./run stories15M.bin -t 0.8 -n 256 -i "Dream comes true this day"

生成された文章

Dream comes true this day. A young girl named Amy and her best friend John are playing together in the park. Amy has her doll and John has a different doll. He has a short doll. Amy is the one who comes to the park and wears her doll.
"Wow, Lucy, look at my doll. She is very pretty and Anna is very nice. She can sing and dance," Amy says.
"Hi, Anna. You are very pretty and smart. I like your doll. She is very kind. She can sing and dance too," Lucy says.
They smile and hug each other. They are happy to meet each other and their doll. They decide to play a game with their doll. They take turns holding their doll and making her sing and dance. They have fun.
achieved tok/s: 30.661249

llama2.mojo(Mojo)

実行コマンド

time mojo llama2.mojo stories15M.bin -n 256 -t 0.8 -i "Dream comes true this day"

生成された文章

num parallel workers: 8  SIMD width: 64
checkpoint size:  60816029 [ 57 MB ] | n layers: 6 | vocab size: 32000
Dream comes true this day! Someone had managed to use the magic language. A prince appeared with a smile on his face and a beautiful cape. The prince smiled and said, "You are a very special prince."
Dream and the prince were in a big fight. But no one wanted to fight and they all said, "You are too lazy!" 
They kept arguing until the prince was tired. He said, "If you don't fight, I will come and get you!" 
The prince said, "Ok, I will fight you and you cannot do it!" So they all fought and laughed. 
But the prince was too lazy to fight and said: "No, I won't fight you!" The prince was so mad and he left and never talked to the prince again. 
The prince was sad but he was glad he was safe. He still had a tiny piece of the magic language that he wanted to remember, and he was never lazy again.
achieved tok/s:  34.066479245907061

結果

Python C Mojo
実行時間(一回目) 7m0.722s 0m6.186s 0m6.902s
実行時間(二回目) 7m13.986s 0m3.542s 0m7.099s

結論

生成された文字数にもよるがPythonと比較するとかなり速いです。約35倍速。 流石にCよりは遅いです。 今回は言語モデルがパラメーターが少なく軽い物なので全体的に短時間ですが、他のモデルがMojoで動作すればローカルLLMが軽快に動作しそうというロマンを感じます!!! Ptyhonの速度に満足できない人はぜひ使ってみてください。