Zenn
🎧

Llama3をローカルMac(MLX)で使ってみる

2024/04/20に公開2件
Python
LLM
Llama
mlx
Llama 3
tech

Llama3

Meta社のLlama3が出たようで話題です。
https://llama.meta.com/llama3/

8Bと70Bがあるようです。
MLXでも使えるようなのでさっそく試してみました。

8B Instructを動かす

mlx-communityですでに登録されている8bit量子化モデルを使ってみます
(4bit量子化モデルもあるのですが、なんだか壊れ気味でした)
https://huggingface.co/mlx-community/Meta-Llama-3-8B-Instruct-8bit

準備

手元ではmlxが0.10.0以上でないと動作しなかったためmlxをバージョン指定でインストールしました。

pip install mlx-lm==0.10.0

スクリプトと出力

apply_chat_templateで適切なテンプレートを使用してくれます。
eos_tokenがいるような気もしますが、気にしない。

ライフゲームを出力させてみます。

スクリプト

from mlx_lm import load, generate

def main():
    model, tokenizer = load("mlx-community/Meta-Llama-3-8B-Instruct-8bit")

    messages = [
        {"role": "system", "content": "you are helpful and skillfull programmer."},
        {"role": "user", "content": "write game of life in python."}
    ]

    inputs = tokenizer.apply_chat_template(messages,
                                           tokenize=False,
                                           add_generation_prompt=True)
    
    response = generate(
        model, 
        tokenizer,
        prompt=inputs,
        verbose=True,
        temp=0.6,
        top_p=0.9,
        max_tokens=4096,
    )

if __name__ == "__main__":
    main()

出力されるプロンプト

Prompt: <|begin_of_text|><|start_header_id|>system<|end_header_id|>
you are helpful and skillfull programmer.<|eot_id|><|start_header_id|>user<|end_header_id|>
write game of life in python.<|eot_id|><|start_header_id|>assistant<|end_header_id|>

レスポンス

Here's a simple implementation of Conway's Game of Life in Python:

import random
import time
import os

# Game of Life constants
WIDTH = 100
HEIGHT = 100

# Create a 2D list filled with random values
def create_board():
    board = [[random.choice([0,1]) for _ in range(WIDTH)] for _ in range(HEIGHT)]
    return board

# Print the current state of the board
def print_board(board):
    os.system('clear') # Clear the screen
    for row in board:
        print(' '.join(['*' if cell else ' for cell in row]))

# Count the number of live neighbors for a cell
def count_neighbors(board, x, y):
    count = 0
    for i in range(-1, 2):
        for j in range(-1, 2):
            if i == 0 and j == 0:
                continue
            count += board[(x+i)%HEIGHT][(y+j)%WIDTH]
    return count

# Apply the rules of the Game of Life to the board
def apply_rules(board):
    new_board = [[0 for _ in range(WIDTH)] for _ in range(HEIGHT)]
    for x in range(HEIGHT):
        for y in range(WIDTH):
            live_neighbors = count_neighbors(board, x, y)
            if board[x][y] == 1:
                if live_neighbors < 2 or live_neighbors > 3:
                    new_board[x][y] = 0
                else:
                    new_board[x][y] = 1
            else:
                if live_neighbors == 3:
                    new_board[x][y] = 1
    return new_board

# Run the Game of Life
def run_game(board):
    while True:
        print_board(board)
        time.sleep(0.5)
        board = apply_rules(board)

# Start the Game of Life
board = create_board()
run_game(board)

This code creates a 100x100 grid of cells, with each cell either alive (represented by `1`) or dead (represented by `0`). The `create_board` function initializes the board with random values. The `print_board` function prints the current state of the board to the screen. The `count_neighbors` function counts the number of live neighbors for a cell. The `apply_rules` function applies the rules of the Game of Life to the board, and the `run_game` function runs the game in an infinite loop.
==========
Prompt: 116.176 tokens-per-sec
Generation: 34.153 tokens-per-sec

M1Max(64GB)で30tpsぐらいは出るのでけっこう速いです。

わりと日本語も書けます。

Prompt: <|begin_of_text|><|start_header_id|>system<|end_header_id|>
you are helpful AI assistant. you MUST write Japanese language.日本語で書いてください<|eot_id|><|start_header_id|>user<|end_header_id|>
こんにちは！あなたの名前は？<|eot_id|><|start_header_id|>assistant<|end_header_id|>

こんにちは！私は、有用なAIアシスタントです。名前は、AI-chanと呼んでください！
==========
Prompt: 140.172 tokens-per-sec
Generation: 34.790 tokens-per-sec

量子化しない版

以下でコンバートします。

python -m mlx_lm.convert --hf-path meta-llama/Meta-Llama-3-8B-Instruct --mlx-path "your model path"

利用はパスを指定します。

from mlx_lm import load, generate
...
model, tokenizer = load("./your model path")

70B Instructを動かす

現在(2024/04/19)MLX版の4bit量子化版はどうもゴミを出力する問題があるようです。

https://github.com/ml-explore/mlx-examples/issues/692

2024/04/20更新: 以下で4bit量子化版を使えるようです。
https://huggingface.co/mlx-community/Meta-Llama-3-70B-Instruct-4bit-mlx

まだ8bitがなかったため、元のLlama3 70B-InstructをMLX版にコンバートして利用してみます。

準備

Llama3をHuggingFaceから取得可能にする

デフォルトだとアクセス制限がかかっています。
以下などでフォームに入力し、届いたメールから制限を解除します。
https://huggingface.co/meta-llama/Meta-Llama-3-70B-Instruct

HuggingFaceトークンを設定する

HuggingFaceからあらかじめTokenを設定・取得しておきます。
シェルでパスを通しておきます(以下はzshの場合)

echo 'export HF_TOKEN="your_huggingface_token"' >> ~/.zshrc
source ~/.zshrc

8bit量子化コンバートする

以下で簡単に行えます。

mlx_lm.convert --hf-path meta-llama/Meta-Llama-3-70B-Instruct --mlx-path "your model path" --q-bits 8 --quantize

出力してみる

数分かかって5文字。遅すぎて使い物になりませんでした。

Prompt: <|begin_of_text|><|start_header_id|>system<|end_header_id|>

you are helpful and skillfull programmer.<|eot_id|><|start_header_id|>user<|end_header_id|>

write game of life in python.<|eot_id|><|start_header_id|>assistant<|end_header_id|>


Here is a

128GBなどでは8tpsで動く報告もあるようです。

70B（8bit）のMLXモデルが出ていた

わざわざコンバートしましたが、こちら使った方がよさそうです。
https://huggingface.co/mlx-community/Meta-Llama-3-70B-Instruct-8bit

読んでくださりありがとうございました、MLX楽しいです。

Discussion

ぼぶぼぶ

70B（8bit）のMLXモデルのアップロード者です！記事で触れていただき嬉しいです！
追記されている70B（4bit）の更新について、更新版を今触っているのですが、少なくとも日本語については8B版のほうがマシな出力をしており、安定したLLMって難しいなと思う次第です。
https://x.com/Prince_Canuma/status/1781600716668227813

RomotRomot

情報ありがとうございます＆70B(8bit)についてもありがとうございます！
4bit版が出力おかしいのは感じております(日本語ではなくローマ字で出力されるなど…)、むずかしそうです…。
もろもろのちほど追記させていただきます！